gitea_admin/sientia-dataops-laborious_temporal
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

SIENTIAPDE-1231

Browse Source 
Update .gitignore and refactor metrics.py for improved logging and consistency

- Added coverage.xml to .gitignore to prevent tracking of coverage reports.
- Refactored metric labels in metrics.py for consistency in string formatting and improved readability.
- Enhanced logging messages in various activities to ensure uniformity in message formatting.
This commit is contained in:
vitor-aignosi
2025-10-15 16:00:18 -03:00
parent a5d2b0d3fd
commit ac795c7c53
39 changed files with 4122 additions and 2602 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
.gitignore vendored
View File

@@ -37,6 +37,7 @@ __pycache__/
# Ignorar coverage
htmlcov/
.coverage
coverage.xml
# git keys
git_key*

40
laborious/activities/activities.py
View File

@@ -1,13 +1,15 @@
from temporalio import activity, workflow
from temporalio import workflow
with workflow.unsafe.imports_passed_through():
from laborious.activities.storage import Storage
from typing import Any
from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler
from sientia_do.observability.logger import Logger
from laborious.activities.mlflow import MLFlow
from laborious.activities.gates import Gates
from laborious.activities.mlflow import MLFlow
from laborious.activities.opc import OPC
from typing import Any
from laborious.activities.storage import Storage
class Activities(Storage, MLFlow, Gates, OPC):
@@ -32,13 +34,15 @@ class Activities(Storage, MLFlow, Gates, OPC):
notification_handler (NotificationHandler): Notification management instance
"""
def __init__(self,
def __init__(
self,
postgres_config: dict[str, Any],
mlflow_config: dict[str, Any],
minio_config: dict[str, Any],
opc_config: dict[str, Any],
logger: Logger,
notification_handler: NotificationHandler):
notification_handler: NotificationHandler,
):
"""
Initialize the Activities orchestrator with all required configurations.
@@ -59,7 +63,9 @@ class Activities(Storage, MLFlow, Gates, OPC):
Exception: If any parent class initialization fails
"""
# Initialize parent classes
Storage.__init__(self, host=postgres_config['host'],
Storage.__init__(
self,
host=postgres_config['host'],
port=postgres_config['port'],
user=postgres_config['user'],
password=postgres_config['password'],
@@ -68,23 +74,25 @@ class Activities(Storage, MLFlow, Gates, OPC):
max_connections=postgres_config['max_connections'],
minio_config=minio_config,
logger=logger,
notification_handler=notification_handler)
notification_handler=notification_handler,
)
MLFlow.__init__(self, mlflow_host=mlflow_config['host'],
MLFlow.__init__(
self,
mlflow_host=mlflow_config['host'],
mlflow_port=mlflow_config['port'],
mlflow_username=mlflow_config['username'],
mlflow_password=mlflow_config['password'],
minio_config=minio_config,
logger=logger,
notification_handler=notification_handler)
notification_handler=notification_handler,
)
Gates.__init__(self, logger=logger,
notification_handler=notification_handler)
Gates.__init__(self, logger=logger, notification_handler=notification_handler)
OPC.__init__(self,
opc_servers=opc_config,
logger=logger,
notification_handler=notification_handler)
OPC.__init__(
self, opc_servers=opc_config, logger=logger, notification_handler=notification_handler
)
async def shutdown(self):
"""

290
laborious/activities/gates.py
View File

@@ -1,55 +1,66 @@
from temporalio import activity, workflow
with workflow.unsafe.imports_passed_through():
import traceback
from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler
from sientia_do.notifications.models import NotificationLevel
from sientia_do.temporal.activities.base import BaseActivity
from sientia_do.observability.logger import Logger
from sientia_do.temporal.constants import DATETIME_FORMAT_WITH_TZ, now
from sientia_do.formatters import create_sample_dict
from laborious.utils.filters.mlflow_filters import nan_values_filter, api_error_filter
from typing import Any
from laborious.utils.filters.conditional_filters import (
filter_empty_data,
filter_specific_variables_null_values
)
from pandas import DataFrame
from laborious import metrics
from collections.abc import Callable, Mapping
from os import path
from shutil import rmtree
from typing import Any
from pandas import DataFrame
from sientia_do.formatters import create_sample_dict
from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler
from sientia_do.notifications.models import NotificationLevel
from sientia_do.observability.logger import Logger
from sientia_do.temporal.activities.base import BaseActivity
from sientia_do.temporal.constants import DATETIME_FORMAT_WITH_TZ, now
from laborious import metrics
from laborious.utils.filters.conditional_filters import (
filter_empty_data,
filter_specific_variables_null_values,
)
from laborious.utils.filters.mlflow_filters import api_error_filter, nan_values_filter
# Strongly-typed filter function signatures
InputFilterFunc = Callable[[DataFrame, dict[str, Any]], bool]
ResponseFilterFunc = Callable[[dict[str, Any], dict[str, Any]], bool]
ContentFilterFunc = Callable[[DataFrame, dict[str, Any]], bool]
# Input filter function mappings
input_filter_functions = {
input_filter_functions: dict[str, InputFilterFunc] = {
'SPECIFIC_VARIABLES_NULL_VALUES': filter_specific_variables_null_values,
'EMPTY_DATA': filter_empty_data,
'path_confidence': {
}
# Confidence mappings kept separate from function maps to avoid Union types
input_path_confidence: Mapping[str, int] = {
'STOP': -1,
'CONTINUE': 2,
'REPEAT': -1
}
'REPEAT': -1,
}
# MLFlow response filter function mappings
mlflow_response_filter_functions = {
mlflow_response_filter_functions: dict[str, ResponseFilterFunc] = {
'API_ERROR': api_error_filter,
'path_confidence': {
}
mlflow_response_path_confidence: Mapping[str, int] = {
'STOP': -1,
'CONTINUE': 10,
'REPEAT': -1
},
'REPEAT': -1,
}
# MLFlow content filter function mappings
mlflow_content_filter_functions = {
mlflow_content_filter_functions: dict[str, ContentFilterFunc] = {
'NAN_VALUES': nan_values_filter,
'EMPTY_DATA': filter_empty_data,
'path_confidence': {
}
mlflow_content_path_confidence: Mapping[str, int] = {
'STOP': -1,
'CONTINUE': 18,
'REPEAT': -1
}
'REPEAT': -1,
}
@@ -84,10 +95,9 @@ class Gates(BaseActivity):
Raises:
Exception: If BaseActivity initialization fails
"""
BaseActivity.__init__(
self, logger, notification_handler, set_error_counter=True)
BaseActivity.__init__(self, logger, notification_handler, set_error_counter=True)
@activity.defn(name="input_gate")
@activity.defn(name='input_gate')
async def input_gate(self, input_data: dict[str, Any]) -> tuple[str | None, int, str]:
"""
Apply input data quality filters and validation.
@@ -123,7 +133,7 @@ class Gates(BaseActivity):
"""
metadata = input_data['metadata']
self.info("Performing input gate...", metadata)
self.info('Performing input gate...', metadata)
filters = input_data['filters']
data = DataFrame(input_data['data'])
@@ -131,40 +141,38 @@ class Gates(BaseActivity):
filter_output = []
self.debug(f"Input data: {data.head(5).to_string()}", metadata)
self.debug(f"Filters: {filters}", metadata)
self.debug(f'Input data: {data.head(5).to_string()}', metadata)
self.debug(f'Filters: {filters}', metadata)
# Apply each configured filter
for fil, config in filters.items():
if fil not in input_filter_functions:
self.error(f"Filter {fil} not found", metadata)
self.error(f'Filter {fil} not found', metadata)
continue
try:
if input_filter_functions[fil](data, config['config']):
self.debug(
f"Data not passed the input filter {fil}:{config}", metadata)
self.debug(f'Data not passed the input filter {fil}:{config}', metadata)
filter_output.append(config['policy'])
except Exception as e:
trace = traceback.format_exc()
self.send_notification(
metadata=metadata,
notification_id=f"INTPUT_GATE_ERROR__{fil}",
message=f"Error in filter {fil}:{config}: \n {e}",
block="input_gate",
notification_id=f'INTPUT_GATE_ERROR__{fil}',
message=f'Error in filter {fil}:{config}: \n {e}',
block='input_gate',
level=NotificationLevel.ERROR,
attachment_content=trace
attachment_content=trace,
)
for path_flag in path_priority:
if path_flag in filter_output:
self.info(f"Input gate result: {path_flag}", metadata)
return path_flag, input_filter_functions['path_confidence'][path_flag], \
"Input data with bad quality"
self.info(f'Input gate result: {path_flag}', metadata)
return path_flag, input_path_confidence[path_flag], 'Input data with bad quality'
self.info("Nothing was filtered by the input gate", metadata)
return None, 0, ""
self.info('Nothing was filtered by the input gate', metadata)
return None, 0, ''
@activity.defn(name="mlflow_response_gate")
@activity.defn(name='mlflow_response_gate')
async def mlflow_response_gate(self, input_data: dict[str, Any]) -> tuple[str | None, int, str]:
"""
Validate MLFlow API response quality and integrity.
@@ -199,7 +207,7 @@ class Gates(BaseActivity):
Exception: If response validation fails or configuration is invalid
"""
metadata = input_data['metadata']
self.info("Performing mlflow response gate...", metadata)
self.info('Performing mlflow response gate...', metadata)
filters = input_data['filters']
data = input_data['data']
@@ -208,9 +216,8 @@ class Gates(BaseActivity):
filter_output = []
self.debug(
f"Input data: \n {create_sample_dict(data, max_items=5, max_depth=5)}", metadata)
self.debug(f"Filters: {filters}", metadata)
self.debug(f'Input data: \n {create_sample_dict(data, max_items=5, max_depth=5)}', metadata)
self.debug(f'Filters: {filters}', metadata)
comments = []
for fil, config in filters.items():
@@ -222,34 +229,32 @@ class Gates(BaseActivity):
comments.append(data['content']['message'])
self.send_notification(
metadata=metadata,
notification_id=f"{gate_type.upper()}_GATE_RESPONSE_FILTER__{fil}",
notification_id=f'{gate_type.upper()}_GATE_RESPONSE_FILTER__{fil}',
message=data['content']['message'],
block="mlflow_gate",
block='mlflow_gate',
level=NotificationLevel.ERROR,
attachment_content=data['content']['traceback']
attachment_content=data['content']['traceback'],
)
except Exception as e:
trace = traceback.format_exc()
self.send_notification(
metadata=metadata,
notification_id=f"MLFLOW_GATE_RESPONSE_FILTER__{fil}",
message=f"Error in filter {fil}:{config}: \n {e}",
block="mlflow_gate",
notification_id=f'MLFLOW_GATE_RESPONSE_FILTER__{fil}',
message=f'Error in filter {fil}:{config}: \n {e}',
block='mlflow_gate',
level=NotificationLevel.ERROR,
attachment_content=trace
attachment_content=trace,
)
for path_flag in path_priority:
if path_flag in filter_output:
self.info(
f"Mlflow response gate result: {path_flag}", metadata)
return path_flag, mlflow_response_filter_functions['path_confidence'][path_flag], \
", ".join(comments)
self.info(f'Mlflow response gate result: {path_flag}', metadata)
return path_flag, mlflow_response_path_confidence[path_flag], ', '.join(comments)
self.info("Nothing was filtered by the mlflow response gate", metadata)
return None, 0, ""
self.info('Nothing was filtered by the mlflow response gate', metadata)
return None, 0, ''
@activity.defn(name="mlflow_content_gate")
@activity.defn(name='mlflow_content_gate')
async def mlflow_content_gate(self, input_data: dict[str, Any]) -> tuple[str | None, int, str]:
"""
Validate MLFlow prediction content quality and integrity.
@@ -284,7 +289,7 @@ class Gates(BaseActivity):
Exception: If content validation fails or configuration is invalid
"""
metadata = input_data['metadata']
self.info("Performing mlflow content gate...", metadata)
self.info('Performing mlflow content gate...', metadata)
filters = input_data['filters']
data = DataFrame(input_data['data'])
@@ -293,8 +298,8 @@ class Gates(BaseActivity):
filter_output = []
self.debug(f"Input data:\n {data.head(5).to_string()}", metadata)
self.debug(f"Filters: \n {filters}", metadata)
self.debug(f'Input data:\n {data.head(5).to_string()}', metadata)
self.debug(f'Filters: \n {filters}', metadata)
for fil, config in filters.items():
if fil not in mlflow_content_filter_functions:
@@ -304,36 +309,38 @@ class Gates(BaseActivity):
filter_output.append(config['policy'])
self.send_notification(
metadata=metadata,
notification_id=f"{gate_type.upper()}_GATE_CONTENT_FILTER__{fil}",
message=f"Data not passed the content filter {fil}:{config}",
block="mlflow_gate",
notification_id=f'{gate_type.upper()}_GATE_CONTENT_FILTER__{fil}',
message=f'Data not passed the content filter {fil}:{config}',
block='mlflow_gate',
level=NotificationLevel.WARNING,
attachment_content=data.to_string()
attachment_content=data.to_string(),
)
except Exception as e:
trace = traceback.format_exc()
self.send_notification(
metadata=metadata,
notification_id=f"MLFLOW_GATE_CONTENT_FILTER__{fil}",
message=f"Error in filter {fil}:{config}: \n {e}",
block="mlflow_gate",
notification_id=f'MLFLOW_GATE_CONTENT_FILTER__{fil}',
message=f'Error in filter {fil}:{config}: \n {e}',
block='mlflow_gate',
level=NotificationLevel.ERROR,
attachment_content=trace
attachment_content=trace,
)
for path_flag in path_priority:
if path_flag in filter_output:
self.info(
f"Mlflow content gate result: {path_flag}", metadata)
return path_flag, mlflow_content_filter_functions['path_confidence'][path_flag], \
"Transformed data not passed the content filter"
self.info(f'Mlflow content gate result: {path_flag}', metadata)
return (
path_flag,
mlflow_content_path_confidence[path_flag],
'Transformed data not passed the content filter',
)
self.info("Nothing was filtered by the mlflow content gate", metadata)
return None, 0, ""
self.info('Nothing was filtered by the mlflow content gate', metadata)
return None, 0, ''
def get_prediction_store_policy(self,
prediction_store_policy: str,
metadata: dict[str, Any]) -> tuple[str, int]:
def get_prediction_store_policy(
self, prediction_store_policy: str, metadata: dict[str, Any]
) -> tuple[str, int]:
"""
Parse and validate prediction store policy configuration.
@@ -359,7 +366,9 @@ class Gates(BaseActivity):
if len(policy_elements) < 2:
self.error(
f"Invalid prediction store policy: {prediction_store_policy}, using default policy", metadata)
f'Invalid prediction store policy: {prediction_store_policy}, using default policy',
metadata,
)
return 'lts', 1
policy_type = policy_elements[0]
@@ -367,14 +376,20 @@ class Gates(BaseActivity):
# If the policy_type is not lts or erl, we use the default policy
# If the policty_value is not a number or 0, we use the default policy
if policy_type not in ['lts', 'erl'] or not policy_value.isdigit() or int(policy_value) == 0:
if (
policy_type not in ['lts', 'erl']
or not policy_value.isdigit()
or int(policy_value) == 0
):
self.error(
f"Invalid prediction store policy: {prediction_store_policy}, using default policy", metadata)
f'Invalid prediction store policy: {prediction_store_policy}, using default policy',
metadata,
)
return 'lts', 1
return policy_type, int(policy_value)
@activity.defn(name="format_prediction")
@activity.defn(name='format_prediction')
async def format_prediction(self, input_data: dict[str, Any]) -> dict[Any, Any]:
"""
Format prediction data according to configured storage policies.
@@ -401,7 +416,7 @@ class Gates(BaseActivity):
"""
metadata = input_data['metadata']
prediction_store_policy = input_data['prediction_store_policy']
self.info("Formatting prediction...", metadata)
self.info('Formatting prediction...', metadata)
data = DataFrame(input_data['data'])
@@ -409,48 +424,45 @@ class Gates(BaseActivity):
data['timestamp'] = data.index
data = data.reset_index(drop=True)
self.debug(
f"Prediction store policy: {prediction_store_policy}", metadata)
self.debug(f"Prediction data: {data.head(5).to_string()}", metadata)
self.debug(f'Prediction store policy: {prediction_store_policy}', metadata)
self.debug(f'Prediction data: {data.head(5).to_string()}', metadata)
policy_type, policy_value = self.get_prediction_store_policy(
prediction_store_policy, metadata)
prediction_store_policy, metadata
)
# If data has no timestamp, we use the default timestamp and not sort the data
self.info(
f"Sorting data by timestamp and applying policy: {policy_type}:{policy_value}", metadata)
f'Sorting data by timestamp and applying policy: {policy_type}:{policy_value}', metadata
)
# If policy_type is lts, we need to sort the data by timestamp descending and take the first policy_value rows
if policy_type == 'lts':
self.debug(
"Sorting data by timestamp descending", metadata)
self.debug('Sorting data by timestamp descending', metadata)
data = data.sort_values(by='timestamp', ascending=False)
# If policy_type is erl, we need to sort the data by timestamp ascending and take the first policy_value rows
elif policy_type == 'erl':
self.debug(
"Sorting data by timestamp ascending", metadata)
self.debug('Sorting data by timestamp ascending', metadata)
data = data.sort_values(by='timestamp', ascending=True)
else:
self.error(
f"Invalid policy type: {policy_type}, using default policy", metadata)
raise ValueError(
f"Invalid policy type: {policy_type}")
self.error(f'Invalid policy type: {policy_type}, using default policy', metadata)
raise ValueError(f'Invalid policy type: {policy_type}')
data = data.head(int(policy_value))
data['model_id'] = input_data['model_id']
data['prediction_confidence'] = input_data['prediction_confidence']
data['prediction_status'] = 'Good'
data['comments'] = ""
data['comments'] = ''
data = data.sort_values(by='timestamp', ascending=False)
data = data.reset_index(drop=True)
self.info(f"Prediction formatted: {len(data)} rows", metadata)
self.debug(f"Prediction data: {data.head(5).to_string()}", metadata)
self.info(f'Prediction formatted: {len(data)} rows', metadata)
self.debug(f'Prediction data: {data.head(5).to_string()}', metadata)
return data.to_dict()
@activity.defn(name="format_default_prediction")
@activity.defn(name='format_default_prediction')
async def format_default_prediction(self, input_data: dict[str, Any]) -> dict[Any, Any]:
"""
Create and format default prediction data for error conditions.
@@ -478,40 +490,44 @@ class Gates(BaseActivity):
"""
metadata = input_data['metadata']
self.debug("Formatting default prediction...", metadata)
self.debug('Formatting default prediction...', metadata)
data = DataFrame({
data = DataFrame(
{
'prediction': [0],
'response_time': [0],
'timestamp': [input_data['timestamp']],
'model_id': [input_data['model_id']],
'prediction_confidence': [input_data['prediction_confidence']],
'prediction_status': ['Bad'],
'comments': [input_data['comment']]
})
'comments': [input_data['comment']],
}
)
self.info(f"Default prediction formatted: {data.size} rows", metadata)
self.info(f'Default prediction formatted: {data.size} rows', metadata)
return data.to_dict()
@activity.defn(name="format_retrain_report")
@activity.defn(name='format_retrain_report')
async def format_retrain_report(self, input_data: dict[str, Any]) -> dict[Any, Any]:
"""
Format retrain report data according to configured storage policies.
"""
metadata = input_data['metadata']
self.info("Formatting retrain report...", metadata)
self.info('Formatting retrain report...', metadata)
experiment_response = input_data['experiment_response']
update_report = input_data['update_report']
model_id = input_data['model_id']
model_name = input_data['model_name']
report = DataFrame({
report = DataFrame(
{
'model_id': [model_id],
'model_name': [model_name],
'timestamp': [experiment_response['timestamp']],
'status': [experiment_response['message']]
})
'status': [experiment_response['message']],
}
)
if experiment_response['success']:
# Retrain was successfull
@@ -519,11 +535,11 @@ class Gates(BaseActivity):
report['mlflow_run_id'] = update_report['mlflow_run_id']
report['mlflow_experiment_id'] = update_report['mlflow_experiment_id']
self.debug(f"Retrain report: {report.to_csv()}", metadata)
self.debug(f'Retrain report: {report.to_csv()}', metadata)
return report.to_dict()
@activity.defn(name="get_last_timestamp")
@activity.defn(name='get_last_timestamp')
async def get_last_timestamp(self, input_data: dict[str, Any]) -> str:
"""
Extract the most recent timestamp from prediction data.
@@ -548,24 +564,22 @@ class Gates(BaseActivity):
"""
metadata = input_data['metadata']
self.info("Getting last timestamp...", metadata)
self.info('Getting last timestamp...', metadata)
data = DataFrame(input_data['data'])
self.debug(f"Input data: {data.head(5).to_string()}", metadata)
self.debug(f'Input data: {data.head(5).to_string()}', metadata)
if data.empty:
return now().strftime(DATETIME_FORMAT_WITH_TZ)
max_timestamp = max(
data['timestamp'].values.tolist())
max_timestamp = max(data['timestamp'].values.tolist())
self.info(
f"Last timestamp: {max_timestamp}", metadata)
self.info(f'Last timestamp: {max_timestamp}', metadata)
return max_timestamp
@activity.defn(name="write_metrics")
@activity.defn(name='write_metrics')
async def write_metrics(self, input_data: dict[str, Any]):
"""
Write prediction performance metrics to Prometheus monitoring system.
@@ -593,42 +607,40 @@ class Gates(BaseActivity):
prediction_confidence = prediction['prediction_confidence'].values[0]
response_time = prediction['response_time'].values[0]
self.info(
f"Writing metrics for model {metadata['model_name']}", metadata)
self.info(f'Writing metrics for model {metadata["model_name"]}', metadata)
metrics.PREDICTIONS_WRITTEN_COUNT.labels(
pod_id=self.pod_id,
model_name=metadata['model_name'],
pipeline_name=metadata['workflow_name']
pipeline_name=metadata['workflow_name'],
).inc()
metrics.PREDICTION_CONFIDENCE_MONITOR.labels(
pod_id=self.pod_id,
model_name=metadata['model_name'],
pipeline_name=metadata['workflow_name']
pipeline_name=metadata['workflow_name'],
).set(prediction_confidence)
metrics.PREDICTION_RESPONSE_TIME_MONITOR.labels(
pod_id=self.pod_id,
model_name=metadata['model_name'],
pipeline_name=metadata['workflow_name']
pipeline_name=metadata['workflow_name'],
).observe(response_time)
self.info(
f"Metrics written for model {metadata['model_name']}", metadata)
self.info(f'Metrics written for model {metadata["model_name"]}', metadata)
@activity.defn(name="clean_tmp_files")
@activity.defn(name='clean_tmp_files')
async def clean_tmp_files(self, input_data: dict[str, Any]):
"""
Clean temporary files in the tmp directory.
"""
model_name = input_data['model_name']
metadata = input_data['metadata']
self.info(f"Cleaning tmp files for model {model_name}...", metadata)
self.info(f'Cleaning tmp files for model {model_name}...', metadata)
if path.exists(f"tmp/retrain_data/{model_name}"):
rmtree(f"tmp/retrain_data/{model_name}")
if path.exists(f"tmp/artifacts/{model_name}"):
rmtree(f"tmp/artifacts/{model_name}")
if path.exists(f'tmp/retrain_data/{model_name}'):
rmtree(f'tmp/retrain_data/{model_name}')
if path.exists(f'tmp/artifacts/{model_name}'):
rmtree(f'tmp/artifacts/{model_name}')
self.info("Tmp files cleaned", metadata)
self.info('Tmp files cleaned', metadata)

144
laborious/activities/mlflow.py
View File

@@ -1,21 +1,25 @@
from sientia_do.temporal.constants import DATETIME_FORMAT_MS_WITH_TZ, now
from temporalio import activity, workflow
with workflow.unsafe.imports_passed_through():
from pandas import to_datetime
from sientia_do.temporal.constants import DATETIME_FORMAT, DATETIME_FORMAT_WITH_TZ
from sientia_do.temporal.activities.base import BaseActivity
import traceback
from typing import Any
import numpy as np
from pandas import DataFrame, to_datetime
from sientia_do.formatters import create_sample_dict
from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler
from sientia_do.notifications.models import NotificationLevel
from sientia_do.observability.logger import Logger
from sientia_do.formatters import create_sample_dict
from laborious.utils.repository.model_repository import MLFlowRepository
from typing import Any
import numpy as np
from pandas import DataFrame
import traceback
from sientia_do.temporal.activities.base import BaseActivity
from sientia_do.temporal.constants import (
DATETIME_FORMAT,
DATETIME_FORMAT_MS_WITH_TZ,
DATETIME_FORMAT_WITH_TZ,
now,
)
from laborious.utils.repository.minio_repository import MinioRepository
from laborious.utils.repository.model_repository import MLFlowRepository
class MLFlow(BaseActivity):
@@ -37,9 +41,16 @@ class MLFlow(BaseActivity):
model_monitoring_repository (MLFlowRepository): Repository for MLFlow operations
"""
def __init__(self, mlflow_host: str, mlflow_port: int, mlflow_username: str,
minio_config: dict[str, Any], mlflow_password: str,
logger: Logger, notification_handler: NotificationHandler):
def __init__(
self,
mlflow_host: str,
mlflow_port: int,
mlflow_username: str,
minio_config: dict[str, Any],
mlflow_password: str,
logger: Logger,
notification_handler: NotificationHandler,
):
"""
Initialize MLFlow activities with server configuration.
@@ -54,26 +65,18 @@ class MLFlow(BaseActivity):
Raises:
Exception: If MLFlowRepository initialization fails
"""
BaseActivity.__init__(
self, logger, notification_handler, set_error_counter=True)
BaseActivity.__init__(self, logger, notification_handler, set_error_counter=True)
self.mlflow_host = mlflow_host
self.mlflow_port = mlflow_port
self.mlflow_username = mlflow_username
self.mlflow_password = mlflow_password
self.model_monitoring_repository = MLFlowRepository(
f"{mlflow_host}:{mlflow_port}", mlflow_username, mlflow_password, logger
f'{mlflow_host}:{mlflow_port}', mlflow_username, mlflow_password, logger
)
if not hasattr(self, 'minio_repository'):
self.minio_repository = MinioRepository(
logger=logger,
notification_handler=notification_handler,
minio_endpoint_url=minio_config['endpoint_url'],
minio_access_key=minio_config['access_key'],
minio_secret_key=minio_config['secret_key'],
minio_region_name=minio_config['region_name'],
minio_default_bucket=minio_config['default_bucket'])
self.minio_repository: MinioRepository | None = None
if self.minio_repository is None:
self.minio_repository = MinioRepository(
@@ -83,9 +86,10 @@ class MLFlow(BaseActivity):
minio_access_key=minio_config['access_key'],
minio_secret_key=minio_config['secret_key'],
minio_region_name=minio_config['region_name'],
minio_default_bucket=minio_config['default_bucket'])
minio_default_bucket=minio_config['default_bucket'],
)
@activity.defn(name="request_transform")
@activity.defn(name='request_transform')
async def request_transform(self, input_data: dict[str, Any]) -> dict[str, Any]:
"""
Transform input data using MLFlow models.
@@ -121,7 +125,7 @@ class MLFlow(BaseActivity):
model_name = input_data['model_name']
model_config = input_data.get('model_config', {})
self.debug("Raw input data:", metadata)
self.debug('Raw input data:', metadata)
self.debug(data.head(5).to_string(), metadata)
# Sort by created_at in descending order and keep first occurrence of each variable/timestamp pair
@@ -130,14 +134,12 @@ class MLFlow(BaseActivity):
)
# Pivot data for model input format
data = data.pivot(
index='timestamp', columns='variable',
values='value')
data = data.pivot(index='timestamp', columns='variable', values='value')
data.fillna(np.nan, inplace=True)
# data.reset_index(inplace=True)
data.columns.name = None
self.debug("Processed input data:", metadata)
self.debug('Processed input data:', metadata)
self.debug(data.head(5).to_string(), metadata)
# Request transformation from MLFlow model
@@ -146,16 +148,20 @@ class MLFlow(BaseActivity):
)
self.debug(
f"Transform raw response data: \n {create_sample_dict(response_data, max_items=5, max_depth=5)}", metadata)
f'Transform raw response data: \n {create_sample_dict(response_data, max_items=5, max_depth=5)}',
metadata,
)
self.debug(
f"Transform response data: \n {create_sample_dict(response_data, max_items=5, max_depth=5)}", metadata)
f'Transform response data: \n {create_sample_dict(response_data, max_items=5, max_depth=5)}',
metadata,
)
self.info("Data transformed successfully", metadata)
self.info('Data transformed successfully', metadata)
return response_data
@activity.defn(name="request_predict")
@activity.defn(name='request_predict')
async def request_predict(self, input_data: dict[str, Any]) -> dict[str, Any]:
"""
Execute predictions using MLFlow models.
@@ -191,14 +197,15 @@ class MLFlow(BaseActivity):
model_name = input_data['model_name']
model_config = input_data.get('model_config', {})
self.debug(f"Input data for: \n {data.head(5).to_string()}", metadata)
self.debug(f'Input data for: \n {data.head(5).to_string()}', metadata)
# Convert numpy.nan to None for model compatibility
data.replace(np.nan, None, inplace=True)
data['timestamp'] = data.index
data['timestamp'] = to_datetime(
data['timestamp'], format=DATETIME_FORMAT_WITH_TZ).dt.strftime(DATETIME_FORMAT)
data['timestamp'], format=DATETIME_FORMAT_WITH_TZ
).dt.strftime(DATETIME_FORMAT)
# Request prediction from MLFlow model
response_data = self.model_monitoring_repository.predict(
@@ -206,13 +213,15 @@ class MLFlow(BaseActivity):
)
self.debug(
f"Prediction response data: \n {create_sample_dict(response_data, max_items=5, max_depth=5)}", metadata)
f'Prediction response data: \n {create_sample_dict(response_data, max_items=5, max_depth=5)}',
metadata,
)
self.info("Data predicted successfully", metadata)
self.info('Data predicted successfully', metadata)
return response_data
@activity.defn(name="retrain_model")
@activity.defn(name='retrain_model')
async def retrain_model(self, input_data: dict[str, Any]) -> dict[str, Any]:
"""
Retrain MLFlow models with updated training data.
@@ -244,15 +253,19 @@ class MLFlow(BaseActivity):
Raises:
Exception: If retraining fails or encounters critical errors
"""
if self.minio_repository is None:
raise ValueError('Minio repository not initialized')
metadata = input_data['metadata']
object_key = input_data['object_key']
self.info(f'Loading retrain data from Key: {object_key}', metadata)
try:
data = self.minio_repository.get_parquet_as_dataframe(
object_key=object_key, metadata=metadata)
object_key=object_key, metadata=metadata
)
except Exception as e:
trace = traceback.format_exc()
self.send_notification(
@@ -261,18 +274,17 @@ class MLFlow(BaseActivity):
message=f'Error loading retrain data: {e}',
block='retrain_model',
level=NotificationLevel.ERROR,
attachment_content=trace
attachment_content=trace,
)
self.error(trace, metadata)
return {
'success': False,
'message': f'Error loading retrain data: {e}',
'traceback': trace,
'timestamp': now().strftime(DATETIME_FORMAT_MS_WITH_TZ)
'timestamp': now().strftime(DATETIME_FORMAT_MS_WITH_TZ),
}
self.debug(
f'Retrain data loaded successfully: shape {data.shape}', metadata)
self.debug(f'Retrain data loaded successfully: shape {data.shape}', metadata)
model_name = input_data['model_name']
model_config = input_data.get('model_config', {})
@@ -291,47 +303,38 @@ class MLFlow(BaseActivity):
data.drop(columns=['created_at'], inplace=True, errors='ignore')
# Pivot data for model input format
data = data.pivot(
index='timestamp', columns='variable',
values='value')
data = data.pivot(index='timestamp', columns='variable', values='value')
data.fillna(np.nan, inplace=True)
# data.reset_index(inplace=True)
data.columns.name = None
data['timestamp'] = data.index
data['timestamp'] = to_datetime(
data['timestamp'], format=DATETIME_FORMAT_WITH_TZ).dt.strftime(DATETIME_FORMAT)
data['timestamp'] = to_datetime(
data['timestamp'], format=DATETIME_FORMAT)
data['timestamp'], format=DATETIME_FORMAT_WITH_TZ
).dt.strftime(DATETIME_FORMAT)
data['timestamp'] = to_datetime(data['timestamp'], format=DATETIME_FORMAT)
data.columns.name = None
retrain_output = self.model_monitoring_repository.retrain_model(
data=data,
model_name=model_name,
model_config=model_config,
metadata=metadata
data=data, model_name=model_name, model_config=model_config, metadata=metadata
)
if not retrain_output['success']:
trace = retrain_output['traceback']
self.send_notification(
metadata=metadata,
notification_id='RETRAIN_MODEL_ERROR',
message=f"Error retraining model {model_name}: {retrain_output['message']}",
message=f'Error retraining model {model_name}: {retrain_output["message"]}',
block='retrain_model',
level=NotificationLevel.ERROR,
attachment_content=trace
attachment_content=trace,
)
self.error(trace, metadata=metadata)
return {
**retrain_output,
'timestamp': timestamp
}
return {**retrain_output, 'timestamp': timestamp}
@activity.defn(name="update_production_model")
@activity.defn(name='update_production_model')
async def update_production_model(self, input_data: dict[str, Any]) -> dict[Any, Any]:
"""
Update production model with newly trained model version.
@@ -372,16 +375,15 @@ class MLFlow(BaseActivity):
model_name = input_data['model_name']
experiment = input_data['experiment']
self.info(
f'Updating production model {model_name} from experiment {experiment}...', metadata)
f'Updating production model {model_name} from experiment {experiment}...', metadata
)
try:
response = self.model_monitoring_repository.update_production_model(
experiment=experiment,
model_name=model_name
experiment=experiment, model_name=model_name, metadata=metadata
)
self.info(
f'Production model {model_name} updated successfully', metadata)
self.info(f'Production model {model_name} updated successfully', metadata)
return response
except Exception as e:
@@ -392,7 +394,7 @@ class MLFlow(BaseActivity):
message=f'Error updating production model {model_name}: {e}',
block='update_production_model',
level=NotificationLevel.ERROR,
attachment_content=trace
attachment_content=trace,
)
self.error(trace, metadata=metadata)
raise e

119
laborious/activities/opc.py
View File

@@ -1,15 +1,16 @@
from temporalio import activity, workflow
with workflow.unsafe.imports_passed_through():
import traceback
from typing import Any
from pandas import DataFrame
from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler
from sientia_do.notifications.models import NotificationLevel
from sientia_do.temporal.activities.base import BaseActivity
from sientia_do.observability.logger import Logger
from sientia_do.temporal.activities.base import BaseActivity
from laborious.utils.repository.opc_repository import OpcRepository
from typing import Any
import traceback
from pandas import DataFrame
OPC_WRITTING_ERROR_CONFIDENCE = 12
@@ -33,15 +34,17 @@ class OPC(BaseActivity):
notification_handler (NotificationHandler): Notification management instance
"""
def __init__(self, opc_servers: dict[str, dict[str, Any]],
logger: Logger, notification_handler: NotificationHandler):
def __init__(
self,
opc_servers: dict[str, dict[str, Any]],
logger: Logger,
notification_handler: NotificationHandler,
):
self.logger = logger
self.notification_handler = notification_handler
self.opc_servers = opc_servers
BaseActivity.__init__(
self, logger, notification_handler, set_error_counter=True)
BaseActivity.__init__(self, logger, notification_handler, set_error_counter=True)
self.opc_repository: dict[str, OpcRepository] = {}
self.opc_servers = opc_servers
@@ -70,10 +73,10 @@ class OPC(BaseActivity):
the initialization of other OPC servers. Each server is handled
independently to ensure maximum availability.
"""
self.logger.info("Initializing OPC servers...")
for id, server in self.opc_servers.items():
self.opc_repository[id] = OpcRepository(
id=server['id'],
self.logger.info('Initializing OPC servers...')
for opc_id, server in self.opc_servers.items():
self.opc_repository[opc_id] = OpcRepository(
opc_id=server['id'],
url=server['url'],
logger=self.logger,
server_uri=server['server_uri'],
@@ -82,30 +85,35 @@ class OPC(BaseActivity):
server_cert_path=server['server_cert_path'],
notification_handler=self.notification_handler,
reconnection_interval=server['reconnection_interval'],
pod_id=self.pod_id
pod_id=self.pod_id,
)
is_connected, error_data = await self.opc_repository[id].connect()
is_connected, error_data = await self.opc_repository[opc_id].connect()
if not is_connected:
self.send_notification(
metadata={
'model_id': '-',
'model_name': '-',
'workflow_name': '-',
'schedule_name': 'INITIALIZATION'
'schedule_name': 'INITIALIZATION',
},
notification_id=error_data['notification_id'],
message=error_data['message'],
block=error_data['block'],
level=error_data.get('level', NotificationLevel.ERROR),
attachment_content=error_data.get(
'attachment_content', None)
attachment_content=error_data.get('attachment_content', None),
)
else:
self.logger.info(
f"OPC server {id} connected successfully.")
self.logger.info(f'OPC server {opc_id} connected successfully.')
async def write_data(self, server_id: str, tag: str, data: Any,
data_type: str, tag_type: str, metadata: dict[str, Any]) -> bool:
async def write_data(
self,
server_id: str,
tag: str,
data: Any,
data_type: str,
tag_type: str,
metadata: dict[str, Any],
) -> bool:
"""
Write data to a specific OPC server tag with comprehensive error handling.
@@ -127,7 +135,8 @@ class OPC(BaseActivity):
try:
is_success, error_data = await self.opc_repository[server_id].write_data(
tag, data, data_type, self.logger, metadata)
tag, data, data_type, self.logger, metadata
)
if not is_success:
self.send_notification(
metadata=metadata,
@@ -135,8 +144,7 @@ class OPC(BaseActivity):
message=error_data['message'],
block=error_data['block'],
level=error_data.get('level', NotificationLevel.ERROR),
attachment_content=error_data.get(
'attachment_content', None)
attachment_content=error_data.get('attachment_content', None),
)
return False
return True
@@ -144,11 +152,11 @@ class OPC(BaseActivity):
trace = traceback.format_exc()
self.send_notification(
metadata=metadata,
notification_id=f"WRITE_OPC_{tag_type.upper()}_ERROR",
message=f"Error writing data to OPC server: {e}",
block="write_opc_data",
notification_id=f'WRITE_OPC_{tag_type.upper()}_ERROR',
message=f'Error writing data to OPC server: {e}',
block='write_opc_data',
level=NotificationLevel.ERROR,
attachment_content=trace
attachment_content=trace,
)
raise e
@@ -174,21 +182,26 @@ class OPC(BaseActivity):
This helps operators quickly identify configuration issues.
"""
if self.opc_repository.get(server_id) is None:
message = f"OPC server {server_id} not found to perform write operation."
message = f'OPC server {server_id} not found to perform write operation.'
self.send_notification(
metadata=metadata,
notification_id="OPC_SERVER_NOT_FOUND",
notification_id='OPC_SERVER_NOT_FOUND',
message=message,
block="write_opc_data",
block='write_opc_data',
level=NotificationLevel.ERROR,
attachment_content=f"OPC servers: {list(self.opc_repository.keys())}"
attachment_content=f'OPC servers: {list(self.opc_repository.keys())}',
)
return False
return True
async def manage_output_tags(
self, server_id: str, config: dict[str, Any], data: DataFrame,
metadata: dict[str, Any], success: bool) -> tuple[bool, int]:
self,
server_id: str,
config: dict[str, Any],
data: DataFrame,
metadata: dict[str, Any],
success: bool,
) -> tuple[bool, int]:
"""
Manage the writing of prediction and confidence data to OPC server tags.
@@ -225,11 +238,13 @@ class OPC(BaseActivity):
data=data.head(1)['prediction'].values[0],
data_type=tag_config['data_type'],
tag_type='prediction',
metadata=metadata
metadata=metadata,
)
if local_success:
self.info(
f"Prediction data written to OPC server {server_id} for tag {tag}.", metadata)
f'Prediction data written to OPC server {server_id} for tag {tag}.',
metadata,
)
count += 1
success = success and local_success
@@ -241,11 +256,13 @@ class OPC(BaseActivity):
data=data.head(1)['prediction_confidence'].values[0],
data_type=tag_config['data_type'],
tag_type='confidence',
metadata=metadata
metadata=metadata,
)
if local_success:
self.info(
f"Confidence data written to OPC server {server_id} for tag {tag}.", metadata)
f'Confidence data written to OPC server {server_id} for tag {tag}.',
metadata,
)
count += 1
success = success and local_success
@@ -271,29 +288,33 @@ class OPC(BaseActivity):
"""
metadata = input_data['metadata']
self.info("Writing data to OPC servers...", metadata)
self.info('Writing data to OPC servers...', metadata)
data = DataFrame(input_data['data'])
opc_output_config = input_data['opc_output_config']
self.info(f"Data to write: {data.size} rows", metadata)
self.info(f'Data to write: {data.size} rows', metadata)
success = True
for server_id, config in opc_output_config.items():
if not self.validate_server(server_id, metadata):
success = False
continue
local_success, local_count = await self.manage_output_tags(
server_id, config, data, metadata, success)
server_id, config, data, metadata, success
)
success = success and local_success
self.info(
f"Process completed for OPC server {server_id}: {local_count} of {len(config.get('prediction_tags', []))} prediction tags and {len(config.get('confidence_tags', []))} confidence tags", metadata)
f'Process completed for OPC server {server_id}: {local_count} of {len(config.get("prediction_tags", []))} prediction tags and {len(config.get("confidence_tags", []))} confidence tags',
metadata,
)
return self.process_confidence(data, success, metadata)
def process_confidence(self, data: DataFrame, success: bool, metadata: dict[str, Any]) -> dict[Any, Any]:
def process_confidence(
self, data: DataFrame, success: bool, metadata: dict[str, Any]
) -> dict[Any, Any]:
"""
Process prediction confidence based on OPC write operation success.
@@ -323,12 +344,12 @@ class OPC(BaseActivity):
if not success:
data['prediction_confidence'] = OPC_WRITTING_ERROR_CONFIDENCE
self.debug(
f"Some data could not be written to OPC servers, setting confidence to {OPC_WRITTING_ERROR_CONFIDENCE}.",
metadata
f'Some data could not be written to OPC servers, setting confidence to {OPC_WRITTING_ERROR_CONFIDENCE}.',
metadata,
)
else:
self.debug("Data written to OPC servers successfully.", metadata)
self.debug('Data written to OPC servers successfully.', metadata)
return data.to_dict()

89
laborious/activities/storage.py
View File

@@ -1,20 +1,21 @@
from temporalio import activity, workflow
with workflow.unsafe.imports_passed_through():
# Extend the Temporal Postgres activities for convenient query -> MinIO export
import traceback
from typing import Any
import pandas as pd
from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler
from sientia_do.notifications.models import NotificationLevel
from sientia_do.observability.logger import Logger
from sientia_do.temporal.activities.postgres import Postgres
from sientia_do.temporal.constants import now
from laborious.utils.repository.minio_repository import MinioRepository
with workflow.unsafe.imports_passed_through():
# Extend the Temporal Postgres activities for convenient query -> MinIO export
from sientia_do.temporal.activities.postgres import Postgres
from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler
from sientia_do.observability.logger import Logger
from sientia_do.temporal.constants import now
from sientia_do.notifications.models import NotificationLevel
from typing import Any
import traceback
import pandas as pd
DATETIME_FILENAME_FORMAT = "%Y-%m-%d_%H-%M-%S"
DATETIME_FILENAME_FORMAT = '%Y-%m-%d_%H-%M-%S'
class Storage(Postgres):
@@ -23,7 +24,8 @@ class Storage(Postgres):
directly to MinIO as Parquet and return the object name.
"""
def __init__(self,
def __init__(
self,
host: str,
port: int,
user: str,
@@ -33,8 +35,10 @@ class Storage(Postgres):
max_connections: int,
minio_config: dict[str, Any],
logger: Logger,
notification_handler: NotificationHandler):
super().__init__(host=host,
notification_handler: NotificationHandler,
):
super().__init__(
host=host,
port=port,
user=user,
password=password,
@@ -42,17 +46,11 @@ class Storage(Postgres):
min_connections=min_connections,
max_connections=max_connections,
logger=logger,
notification_handler=notification_handler)
notification_handler=notification_handler,
)
if not hasattr(self, 'minio_repository'):
self.minio_repository = MinioRepository(
logger=logger,
notification_handler=notification_handler,
minio_endpoint_url=minio_config['endpoint_url'],
minio_access_key=minio_config['access_key'],
minio_secret_key=minio_config['secret_key'],
minio_region_name=minio_config['region_name'],
minio_default_bucket=minio_config['default_bucket'])
self.minio_repository: MinioRepository | None = None
if self.minio_repository is None:
self.minio_repository = MinioRepository(
@@ -62,7 +60,8 @@ class Storage(Postgres):
minio_access_key=minio_config['access_key'],
minio_secret_key=minio_config['secret_key'],
minio_region_name=minio_config['region_name'],
minio_default_bucket=minio_config['default_bucket'])
minio_default_bucket=minio_config['default_bucket'],
)
@activity.defn(name='query_to_minio')
async def query_to_minio(self, input_data: dict[str, Any]) -> dict[str, Any]:
@@ -79,64 +78,60 @@ class Storage(Postgres):
dict: { success: bool, object_name: str, uri: str }
"""
if self.minio_repository is None:
raise ValueError('Minio repository not initialized')
metadata = input_data.get('metadata', {})
object_prefix = input_data.get('object_prefix', 'datasets/retrain')
timestamp = now().strftime(DATETIME_FILENAME_FORMAT)
object_name = f"{object_prefix}_{timestamp}.parquet"
uri = f"s3://{self.minio_repository.minio_bucket}/{object_name}"
object_name = f'{object_prefix}_{timestamp}.parquet'
uri = f's3://{self.minio_repository.minio_bucket}/{object_name}'
try:
data = await self.load_custom_query(input_data)
if not data:
self.error(
f"query_to_minio failed: No data returned from query", metadata)
return {"success": False, "message": "No data returned from query"}
self.error('query_to_minio failed: No data returned from query', metadata)
return {'success': False, 'message': 'No data returned from query'}
# Ensure we have a DataFrame
data = pd.DataFrame(data)
# Write parquet to memory and upload via persistent client
self.minio_repository.store_dataframe_as_parquet(
dataframe=data,
uri=uri,
object_name=object_name,
metadata=metadata
dataframe=data, uri=uri, object_name=object_name, metadata=metadata
)
return {"success": True, "object_key": object_name, "uri": uri}
return {'success': True, 'object_key': object_name, 'uri': uri}
except Exception as e:
trace = traceback.format_exc()
self.send_notification(
metadata=metadata,
notification_id="ERROR_LOADING_CUSTOM_QUERY",
message=f"Error fetching data from query: {e}",
block="load_custom_query",
notification_id='ERROR_STORING_QUERY_TO_MINIO',
message=f'Error storing query to MinIO: {e}',
block='query_to_minio',
level=NotificationLevel.ERROR,
attachment_content=trace
attachment_content=trace,
)
self.error(trace, metadata)
return {"success": False, "message": str(e)}
return {'success': False, 'message': str(e)}
def close(self) -> None:
"""Close Storage resources (MinIO client and Postgres engine)."""
try:
if hasattr(self, 's3_client') and self.s3_client is not None:
if hasattr(self, 'minio_repository') and self.minio_repository is not None:
try:
self.s3_client.close()
self.minio_repository.close()
finally:
self.s3_client = None
self.minio_repository = None
finally:
# Ensure Postgres resources are disposed as well
try:
super().close()
except Exception:
pass
self.logger.error('Error closing Postgres resources')
def __del__(self):
try:
self.close()
except Exception:
pass

38
laborious/metrics.py
View File

@@ -23,50 +23,50 @@ Metric Labels:
- opc_server_id: Identifier for OPC server operations
"""
from prometheus_client import Gauge, Counter, Histogram
from prometheus_client import Counter, Gauge, Histogram
# Application health metric
APP_UP = Gauge(
"app_up",
"Indicates if the application is running (1) or shutting down (0)",
["pod_id"],
'app_up',
'Indicates if the application is running (1) or shutting down (0)',
['pod_id'],
)
# Core labels used across multiple metrics
CORE_LABELS = ["pod_id", "model_name", "pipeline_name"]
CORE_LABELS = ['pod_id', 'model_name', 'pipeline_name']
# Prediction operation metrics
PREDICTIONS_WRITTEN_COUNT = Counter(
"laborious_predictions_written_count",
"Number of predictions written to the database table predictions",
'laborious_predictions_written_count',
'Number of predictions written to the database table predictions',
CORE_LABELS,
)
# Prediction quality metrics
PREDICTION_CONFIDENCE_MONITOR = Gauge(
"laborious_prediction_confidence_monitor",
"Current confidence of each prediction",
'laborious_prediction_confidence_monitor',
'Current confidence of each prediction',
CORE_LABELS,
)
# Performance monitoring metrics
PREDICTION_RESPONSE_TIME_MONITOR = Histogram(
"laborious_prediction_response_time_monitor",
"Current response time of each prediction",
'laborious_prediction_response_time_monitor',
'Current response time of each prediction',
CORE_LABELS,
buckets=[0.01, 0.05, 0.1, 0.2, 0.5, 1.0, 2.0, 5.0, 10.0]
buckets=[0.01, 0.05, 0.1, 0.2, 0.5, 1.0, 2.0, 5.0, 10.0],
)
# OPC export metrics
PREDICTION_OPC_WRITING_COUNT = Counter(
"laborious_prediction_opc_writing_count",
"Number of predictions written to the OPC server",
[*CORE_LABELS, "opc_server_id"],
'laborious_prediction_opc_writing_count',
'Number of predictions written to the OPC server',
[*CORE_LABELS, 'opc_server_id'],
)
PREDICTION_OPC_WRITING_RESPONSE_TIME_MONITOR = Histogram(
"laborious_prediction_opc_writing_response_time_monitor",
"Current response time of each prediction written to the OPC server",
[*CORE_LABELS, "opc_server_id"],
buckets=[0.01, 0.05, 0.1, 0.2, 0.5, 1.0, 2.0, 5.0, 10.0]
'laborious_prediction_opc_writing_response_time_monitor',
'Current response time of each prediction written to the OPC server',
[*CORE_LABELS, 'opc_server_id'],
buckets=[0.01, 0.05, 0.1, 0.2, 0.5, 1.0, 2.0, 5.0, 10.0],
)

24
laborious/utils/connectors_config.py
View File

@@ -1,9 +1,9 @@
from os import getenv
import json
from typing import Dict, Any
from os import getenv
from typing import Any
def build_postgres_config() -> Dict[str, Any]:
def build_postgres_config() -> dict[str, Any]:
"""
Build PostgreSQL database configuration from environment variables.
@@ -30,11 +30,11 @@ def build_postgres_config() -> Dict[str, Any]:
'password': getenv('POSTGRES_PASSWORD', 'sientia'),
'dbname': getenv('POSTGRES_DBNAME', 'sientia'),
'min_connections': int(getenv('POSTGRES_MIN_CONNECTIONS', '5')),
'max_connections': int(getenv('POSTGRES_MAX_CONNECTIONS', '20'))
'max_connections': int(getenv('POSTGRES_MAX_CONNECTIONS', '20')),
}
def build_mlflow_config() -> Dict[str, Any]:
def build_mlflow_config() -> dict[str, Any]:
"""
Build MLFlow server configuration from environment variables.
@@ -55,11 +55,11 @@ def build_mlflow_config() -> Dict[str, Any]:
'host': getenv('MLFLOW_HOST', 'http://localhost'),
'port': int(getenv('MLFLOW_PORT', '5080')),
'username': getenv('MLFLOW_USERNAME', 'aignosi'),
'password': getenv('MLFLOW_PASSWORD', 'aignosi')
'password': getenv('MLFLOW_PASSWORD', 'aignosi'),
}
def build_opc_config() -> Dict[str, Any]:
def build_opc_config() -> dict[str, Any]:
"""
Build OPC server configuration from environment variables.
@@ -93,12 +93,12 @@ def build_opc_config() -> Dict[str, Any]:
'cert_path': getenv('OPC_CERT_PATH', None),
'private_key_path': getenv('OPC_PRIVATE_KEY_PATH', None),
'server_cert_path': getenv('OPC_SERVER_CERT_PATH', None),
'reconnection_interval': int(getenv('OPC_RECONNECTION_INTERVAL', '120'))
'reconnection_interval': int(getenv('OPC_RECONNECTION_INTERVAL', '120')),
}
}
def build_mongodb_config() -> Dict[str, Any]:
def build_mongodb_config() -> dict[str, Any]:
"""
Build MongoDB configuration from environment variables.
@@ -125,11 +125,11 @@ def build_mongodb_config() -> Dict[str, Any]:
return {
'connection_string': connection_string,
'database_name': getenv('MONGODB_DATABASE_NAME', 'sientia'),
'ttl_index_seconds': int(getenv('MONGODB_TTL_INDEX_HOURS', '1')) * 3600
'ttl_index_seconds': int(getenv('MONGODB_TTL_INDEX_HOURS', '1')) * 3600,
}
def build_minio_config() -> Dict[str, Any]:
def build_minio_config() -> dict[str, Any]:
"""
Build MinIO (S3-compatible) configuration from environment variables.
@@ -148,5 +148,5 @@ def build_minio_config() -> Dict[str, Any]:
'access_key': getenv('MINIO_ACCESS_KEY', 'minioadmin'),
'secret_key': getenv('MINIO_SECRET_KEY', 'minioadmin'),
'region_name': getenv('MINIO_REGION_NAME', 'us-east-1'),
'default_bucket': getenv('MINIO_DEFAULT_BUCKET', 'laborious')
'default_bucket': getenv('MINIO_DEFAULT_BUCKET', 'laborious'),
}

3
laborious/utils/filters/conditional_filters.py
View File

@@ -24,8 +24,7 @@ def filter_specific_variables_null_values(data: DataFrame, config: dict) -> bool
if data.empty:
return False
return not data[
data['variable'].isin(config['variables']) & data['value'].isna()].empty
return not data[data['variable'].isin(config['variables']) & data['value'].isna()].empty
def filter_empty_data(data: DataFrame, _config: dict) -> bool:

7
laborious/utils/filters/mlflow_filters.py
View File

@@ -52,8 +52,11 @@ def nan_values_filter(predictions: DataFrame, _config: dict) -> bool:
bool: True if data should be filtered (too many NaN values), False otherwise
"""
data = predictions.replace({None: np.nan}).drop(
columns=['timestamp'], errors='ignore').infer_objects()
data = (
predictions.replace({None: np.nan})
.drop(columns=['timestamp'], errors='ignore')
.infer_objects()
)
if data.isna().all().all():
return True

76
laborious/utils/repository/minio_repository.py
View File

@@ -1,40 +1,41 @@
from io import BytesIO
import traceback
from typing import Any
import boto3
from botocore.config import Config
from pandas import DataFrame, read_parquet
from sientia_do.observability.logger import Logger
from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler
from sientia_do.notifications.models import NotificationLevel
from typing import Any
from botocore.exceptions import ClientError
from pandas import DataFrame, read_parquet
from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler
from sientia_do.observability.logger import Logger
class MinioRepository():
def __init__(self,
class MinioRepository:
def __init__(
self,
minio_endpoint_url: str,
minio_access_key: str,
minio_secret_key: str,
minio_region_name: str,
minio_default_bucket: str,
logger: Logger,
notification_handler: NotificationHandler):
notification_handler: NotificationHandler,
):
# MinIO settings shared with pandas s3fs
self.storage_options = {
'key': minio_access_key,
'secret': minio_secret_key,
'client_kwargs': {'endpoint_url': minio_endpoint_url}
'client_kwargs': {'endpoint_url': minio_endpoint_url},
}
self.minio_bucket = minio_default_bucket
self.minio_endpoint_url = minio_endpoint_url
self.minio_region_name = minio_region_name
logger.info(
f"Connecting to MinIO at {self.minio_endpoint_url}, default bucket: {self.minio_bucket}")
f'Connecting to MinIO at {self.minio_endpoint_url}, default bucket: {self.minio_bucket}'
)
# Reusable MinIO client
self.s3_client = boto3.client(
self.s3_client: Any = boto3.client(
's3',
endpoint_url=self.minio_endpoint_url,
aws_access_key_id=self.storage_options['key'],
@@ -48,67 +49,46 @@ class MinioRepository():
read_timeout=120,
),
)
self._bucket_checked = False
self.logger = logger
self.notification_handler = notification_handler
def close(self):
self.s3_client.close()
def ensure_bucket_exists(self, metadata: dict[str, Any]) -> bool:
"""
Ensure the MinIO bucket exists; create it if necessary.
"""
if self._bucket_checked:
return True
try:
self.logger.custom_info(
f"Checking if bucket '{self.minio_bucket}' exists", metadata)
self.logger.custom_info(f"Checking if bucket '{self.minio_bucket}' exists", metadata)
self.s3_client.head_bucket(Bucket=self.minio_bucket)
self._bucket_checked = True
return True
except ClientError:
try:
self.logger.custom_info(
f"Creating bucket '{self.minio_bucket}'", metadata)
self.logger.custom_info(f"Creating bucket '{self.minio_bucket}'", metadata)
self.s3_client.create_bucket(Bucket=self.minio_bucket)
self._bucket_checked = True
return True
except ClientError as ce:
trace = traceback.format_exc()
self.notification_handler.send_notification(
metadata=metadata,
notification_id="ERROR_CREATING_MINIO_BUCKET",
message=f"Failed to ensure bucket '{self.minio_bucket}': {ce}",
block="ensure_bucket_exists",
level=NotificationLevel.ERROR,
attachment_content=str(ce)
)
self.logger.custom_error(trace, metadata)
return False
def store_dataframe_as_parquet(self, dataframe: DataFrame, uri: str,
object_name: str, metadata: dict[str, Any]):
def store_dataframe_as_parquet(
self, dataframe: DataFrame, uri: str, object_name: str, metadata: dict[str, Any]
):
self.ensure_bucket_exists(metadata)
self.logger.custom_info(
f"Storing dataframe as parquet in {uri}", metadata)
self.logger.custom_info(f'Storing dataframe as parquet in {uri}', metadata)
buffer = BytesIO()
dataframe.to_parquet(buffer, engine='pyarrow', index=True)
buffer.seek(0)
self.s3_client.put_object(
Bucket=self.minio_bucket, Key=object_name, Body=buffer.getvalue())
self.s3_client.put_object(Bucket=self.minio_bucket, Key=object_name, Body=buffer.getvalue())
self.logger.custom_info(
f"Dataframe stored as parquet in {uri}", metadata)
self.logger.custom_info(f'Dataframe stored as parquet in {uri}', metadata)
def get_parquet_as_dataframe(self, object_key: str, metadata: dict[str, Any]) -> DataFrame:
self.logger.custom_info(
f"Getting parquet as dataframe from {object_key}", metadata)
self.logger.custom_info(f'Getting parquet as dataframe from {object_key}', metadata)
response = self.s3_client.get_object(
Bucket=self.minio_bucket, Key=object_key)
response = self.s3_client.get_object(Bucket=self.minio_bucket, Key=object_key)
# Read the content into a BytesIO buffer to support seek operations
buffer = BytesIO(response['Body'].read())

753
laborious/utils/repository/model_repository.py
View File

File diff suppressed because it is too large Load Diff

180
laborious/utils/repository/opc_repository.py
View File

@@ -1,16 +1,16 @@
import asyncio
import traceback
import time
import traceback
from datetime import datetime
from pathlib import Path
from typing import Any
from asyncua import Client
from asyncua.crypto.security_policies import SecurityPolicyBasic256
from asyncua.ua import DataValue, Variant, VariantType, DateTime
from regex import F
from asyncua.ua import DataValue, DateTime, Variant, VariantType
from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler
from sientia_do.notifications.models import NotificationLevel
from sientia_do.observability.logger import Logger
from laborious import metrics
data_type_map = {
@@ -33,17 +33,26 @@ data_type_map = {
'str': {
'converter': str,
'opc_type': VariantType.String,
}
},
}
class OpcRepository():
def __init__(self, id: str, url: str, logger: Logger,
class OpcRepository:
def __init__(
self,
opc_id: str,
url: str,
logger: Logger,
notification_handler: NotificationHandler,
reconnection_interval: int = 60, server_uri: str = None, cert_path: str = None,
private_key_path: str = None, server_cert_path: str = None, pod_id: str = None):
reconnection_interval: int = 60,
server_uri: str | None = None,
cert_path: str | None = None,
private_key_path: str | None = None,
server_cert_path: str | None = None,
pod_id: str | None = None,
):
self.url = url
self.id = id
self.id = opc_id
self.server_uri = server_uri
self.cert_path = cert_path
self.private_key_path = private_key_path
@@ -51,16 +60,16 @@ class OpcRepository():
self.logger = logger
self.error_count = 0
self.reconnection_interval = reconnection_interval
self.last_reconnection_time = None
self.last_reconnection_time: None | datetime = None
self.notification_handler = notification_handler
self.client = None
self.client: None | Client = None
self.pod_id = pod_id
self.metadata = {
'model_name': '-',
'model_id': '-',
'workflow_name': 'opc_repository',
'schedule_name': '-'
'schedule_name': '-',
}
async def set_security(self):
@@ -85,11 +94,18 @@ class OpcRepository():
if not all([self.cert_path, self.private_key_path]):
raise ValueError(
"Certificate and private key paths must be provided for secure connection.")
'Certificate and private key paths must be provided for secure connection.'
)
if self.cert_path is None or self.private_key_path is None:
raise ValueError('Certificate and private key paths cannot be None')
cert = Path(self.cert_path)
private_key = Path(self.private_key_path)
server_cert = Path(
self.server_cert_path) if self.server_cert_path else None
server_cert = Path(self.server_cert_path) if self.server_cert_path else None
if self.client is None:
raise ValueError('Client must be initialized before setting security')
self.client.application_uri = self.server_uri
self.logger.custom_info('Setting security...', self.metadata)
@@ -97,7 +113,7 @@ class OpcRepository():
SecurityPolicyBasic256,
certificate=str(cert),
private_key=str(private_key),
server_certificate=str(server_cert)
server_certificate=str(server_cert) if server_cert else None,
)
self.client.secure_channel_timeout = 10000000
self.client.session_timeout = 10000000
@@ -115,8 +131,7 @@ class OpcRepository():
self.client = Client(self.url)
if self.cert_path:
await self.set_security()
self.logger.custom_info(
f'Starting connection to OPC server {self.id}...', self.metadata)
self.logger.custom_info(f'Starting connection to OPC server {self.id}...', self.metadata)
return await self.try_connect()
async def try_connect(self) -> tuple[bool, dict[str, Any]]:
@@ -136,6 +151,13 @@ class OpcRepository():
try:
self.last_reconnection_time = datetime.now()
if self.client is None:
return False, {
'notification_id': f'OPC_CONNECTION_ERROR_{self.id}',
'message': 'Client is not initialized',
'block': 'opc_repository',
'level': NotificationLevel.ERROR,
}
await self.client.connect()
return True, {}
except Exception as e:
@@ -143,11 +165,11 @@ class OpcRepository():
self.logger.custom_error(trace, self.metadata)
return False, {
"notification_id": f"OPC_CONNECTION_ERROR_{self.id}",
"message": f"Failed to connect to OPC server: {e}",
"block": "opc_repository",
"level": NotificationLevel.ERROR,
"attachment_content": trace
'notification_id': f'OPC_CONNECTION_ERROR_{self.id}',
'message': f'Failed to connect to OPC server: {e}',
'block': 'opc_repository',
'level': NotificationLevel.ERROR,
'attachment_content': trace,
}
async def disconnect(self):
@@ -162,11 +184,9 @@ class OpcRepository():
return
try:
await self.client.disconnect()
self.logger.custom_info(
'Disconnected from OPC server', self.metadata)
self.logger.custom_info('Disconnected from OPC server', self.metadata)
except Exception as e:
self.logger.custom_error(
f"Failed to disconnect from OPC server: {e}", self.metadata)
self.logger.custom_error(f'Failed to disconnect from OPC server: {e}', self.metadata)
self.client = None
async def validate_connection(self) -> tuple[bool, dict[str, Any]]:
@@ -203,52 +223,62 @@ class OpcRepository():
if self.error_count > 5:
self.logger.custom_warning(
f"OPC server {self.id} will be disconnected due to multiple errors", self.metadata)
f'OPC server {self.id} will be disconnected due to multiple errors', self.metadata
)
try:
await self.disconnect()
except Exception as e:
trace = traceback.format_exc()
self.logger.custom_error(
f"Failed to disconnect from OPC server: {e}", self.metadata)
f'Failed to disconnect from OPC server: {e}', self.metadata
)
self.logger.custom_error(trace, self.metadata)
self.logger.custom_info(
f"Attempting to reconnect to OPC server {self.id}...", self.metadata)
f'Attempting to reconnect to OPC server {self.id}...', self.metadata
)
return await self.connect()
# Check if client is connected using asyncua's connection state
try:
if self.client.uaclient.protocol is None or self.client.uaclient.protocol.state == "closed":
if (
self.client.uaclient.protocol is None
or self.client.uaclient.protocol.state == 'closed'
):
# OPC server is not connected
self.logger.custom_error(
f"OPC server {self.id} is not connected", self.metadata)
if self.last_reconnection_time is None or (datetime.now() - self.last_reconnection_time).total_seconds(
) > self.reconnection_interval:
self.logger.custom_error(f'OPC server {self.id} is not connected', self.metadata)
if (
self.last_reconnection_time is None
or (datetime.now() - self.last_reconnection_time).total_seconds()
> self.reconnection_interval
):
await self.disconnect()
self.logger.custom_info(
f"Trying to reconnect to OPC server {self.id}...", self.metadata)
f'Trying to reconnect to OPC server {self.id}...', self.metadata
)
return await self.connect()
return False, {
"notification_id": f"OPC_CONNECTION_AWAITING_RECONNECTION_WINDOW_{self.id}",
"message": f"OPC server {self.id} is not connected, waiting for next reconnection window...",
"block": "opc_repository",
"level": NotificationLevel.WARNING
'notification_id': f'OPC_CONNECTION_AWAITING_RECONNECTION_WINDOW_{self.id}',
'message': f'OPC server {self.id} is not connected, waiting for next reconnection window...',
'block': 'opc_repository',
'level': NotificationLevel.WARNING,
}
return True, {}
except Exception as e:
trace = traceback.format_exc()
message = f"Failed to validate connection to OPC server: {e}"
message = f'Failed to validate connection to OPC server: {e}'
self.logger.custom_error(message, self.metadata)
return False, {
"notification_id": f"OPC_CONNECTION_CHECK_ERROR_{self.id}",
"message": message,
"block": "opc_repository",
"level": NotificationLevel.ERROR,
"attachment_content": trace
'notification_id': f'OPC_CONNECTION_CHECK_ERROR_{self.id}',
'message': message,
'block': 'opc_repository',
'level': NotificationLevel.ERROR,
'attachment_content': trace,
}
async def write_data(self, node: str, value: Any, data_type: str,
logger: Logger, metadata: dict[str, Any]) -> tuple[bool, dict[str, Any]]:
async def write_data(
self, node: str, value: Any, data_type: str, logger: Logger, metadata: dict[str, Any]
) -> tuple[bool, dict[str, Any]]:
"""
Write data to OPC server with comprehensive validation and monitoring.
@@ -286,42 +316,42 @@ class OpcRepository():
start_time = time.time()
try:
if self.client is None:
return False, {
'notification_id': f'OPC_WRITE_GET_NODE_ERROR_{self.id}',
'message': 'Client is not initialized',
'block': 'opc_repository',
'level': NotificationLevel.ERROR,
}
node_obj = self.client.get_node(node)
except Exception as e:
trace = traceback.format_exc()
logger.custom_error(trace, metadata.get('schedule_name', 'N/A'))
self.error_count += 1
return False, {
"notification_id": f"OPC_WRITE_GET_NODE_ERROR_{self.id}",
"message": f"Failed to get node from OPC server: {e} | metadata: {metadata}",
"block": "opc_repository",
"level": NotificationLevel.ERROR,
"attachment_content": trace
'notification_id': f'OPC_WRITE_GET_NODE_ERROR_{self.id}',
'message': f'Failed to get node from OPC server: {e} | metadata: {metadata}',
'block': 'opc_repository',
'level': NotificationLevel.ERROR,
'attachment_content': trace,
}
if data_type not in data_type_map:
return False, {
"notification_id": f"OPC_WRITE_DATA_TYPE_ERROR_{self.id}",
"message": f"Unsupported data type: {data_type} | metadata: {metadata}",
"block": "opc_repository",
"level": NotificationLevel.ERROR
'notification_id': f'OPC_WRITE_DATA_TYPE_ERROR_{self.id}',
'message': f'Unsupported data type: {data_type} | metadata: {metadata}',
'block': 'opc_repository',
'level': NotificationLevel.ERROR,
}
data = data_type_map[data_type]['converter'](value)
logger.custom_info(
f'Writing {data} - {type(data)} to {node}', metadata)
logger.custom_info(f'Writing {data} - {type(data)} to {node}', metadata)
now = datetime.now()
ua_data = DataValue(
Variant(data, data_type_map[data_type]['opc_type']),
SourceTimestamp=DateTime(
now.year,
now.month,
now.day,
now.hour,
now.minute,
now.second,
now.microsecond
)
now.year, now.month, now.day, now.hour, now.minute, now.second, now.microsecond
),
)
try:
@@ -331,7 +361,7 @@ class OpcRepository():
pod_id=self.pod_id,
model_name=metadata['model_name'],
pipeline_name=metadata['workflow_name'],
opc_server_id=self.id
opc_server_id=self.id,
).inc()
end_time = time.time()
@@ -340,7 +370,7 @@ class OpcRepository():
pod_id=self.pod_id,
model_name=metadata['model_name'],
pipeline_name=metadata['workflow_name'],
opc_server_id=self.id
opc_server_id=self.id,
).observe(response_time)
except Exception as e:
@@ -348,11 +378,11 @@ class OpcRepository():
logger.custom_error(trace, metadata)
self.error_count += 1
return False, {
"notification_id": f"OPC_WRITE_DATA_ERROR_{self.id}",
"message": f"Failed to write data to OPC server: {e} | metadata: {metadata}",
"block": "opc_repository",
"level": NotificationLevel.ERROR,
"attachment_content": trace
'notification_id': f'OPC_WRITE_DATA_ERROR_{self.id}',
'message': f'Failed to write data to OPC server: {e} | metadata: {metadata}',
'block': 'opc_repository',
'level': NotificationLevel.ERROR,
'attachment_content': trace,
}
self.error_count = 0

77
laborious/worker/worker.py
View File

@@ -25,37 +25,37 @@ Environment Variables:
- PROJECT_NAME: Project name for notifications (default: laborious)
"""
from temporalio import workflow, client
from temporalio.worker import Worker, PollerBehaviorAutoscaling
from temporalio.runtime import Runtime, TelemetryConfig, PrometheusConfig
from temporalio import client, workflow
from temporalio.runtime import PrometheusConfig, Runtime, TelemetryConfig
from temporalio.worker import PollerBehaviorAutoscaling, Worker
with workflow.unsafe.imports_passed_through():
import asyncio
import os
import sys
import asyncio
from laborious.workflows.minimal_retrain import MinimalRetrain
from laborious.workflows.predictions_batch import PredictionsBatch
from laborious.workflows.sub_workflows.prediction_process import PredictionProcess
from laborious.workflows.sub_workflows.format_and_export_prediction import \
FormatAndExportPrediction
from laborious.activities.activities import Activities
from laborious.utils.connectors_config import (
build_postgres_config,
build_mlflow_config,
build_minio_config,
build_opc_config,
build_mongodb_config
)
from prometheus_client import start_http_server
from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler
from sientia_do.observability.logger import get_logger
from laborious import metrics
from prometheus_client import start_http_server
import lzma
import dataclasses
from laborious import metrics
from laborious.activities.activities import Activities
from laborious.utils.connectors_config import (
build_minio_config,
build_mlflow_config,
build_mongodb_config,
build_opc_config,
build_postgres_config,
)
from laborious.workflows.minimal_retrain import MinimalRetrain
from laborious.workflows.predictions_batch import PredictionsBatch
from laborious.workflows.sub_workflows.format_and_export_prediction import (
FormatAndExportPrediction,
)
from laborious.workflows.sub_workflows.prediction_process import PredictionProcess
POD_ID = os.getenv('POD_ID')
SDK_METRICS_PORT = int(os.getenv('HTTP_SDK_METRICS_PORT', "9091"))
SDK_METRICS_PORT = int(os.getenv('HTTP_SDK_METRICS_PORT', '9091'))
async def main():
@@ -91,7 +91,7 @@ async def main():
logger.custom_info(f'Starting Worker with POD_ID: {POD_ID}', metadata)
logger.custom_info("Starting prometheus client...", metadata)
logger.custom_info('Starting prometheus client...', metadata)
start_prometheus_server()
logger.custom_info('Starting Notification Handler...', metadata)
@@ -101,7 +101,7 @@ async def main():
connection_string=mongo_config['connection_string'],
database=mongo_config['database_name'],
logger=logger,
project_name=os.getenv('PROJECT_NAME', 'laborious')
project_name=os.getenv('PROJECT_NAME', 'laborious'),
)
logger.custom_info('Starting Activities...', metadata)
@@ -112,19 +112,17 @@ async def main():
minio_config=build_minio_config(),
opc_config=build_opc_config(),
logger=logger,
notification_handler=notification_handler
notification_handler=notification_handler,
)
logger.custom_info('Initializing OPC...', metadata)
await activities.init_opc()
logger.custom_info(
f'Starting SDK Metrics Server on port {SDK_METRICS_PORT}...', metadata)
logger.custom_info(f'Starting SDK Metrics Server on port {SDK_METRICS_PORT}...', metadata)
new_runtime = Runtime(
telemetry=TelemetryConfig(
metrics=PrometheusConfig(
bind_address=f"0.0.0.0:{SDK_METRICS_PORT}")
metrics=PrometheusConfig(bind_address=f'0.0.0.0:{SDK_METRICS_PORT}')
)
)
@@ -133,7 +131,7 @@ async def main():
temporal_client = await client.Client.connect(
target_host=host,
namespace=os.getenv('TEMPORAL_NAMESPACE', 'laborious'),
runtime=new_runtime
runtime=new_runtime,
)
logger.custom_info('Starting Workers...', metadata)
@@ -156,13 +154,12 @@ async def main():
max_concurrent_local_activities=50,
max_cached_workflows=200,
workflow_task_poller_behavior=PollerBehaviorAutoscaling(),
activity_task_poller_behavior=PollerBehaviorAutoscaling()
activity_task_poller_behavior=PollerBehaviorAutoscaling(),
),
Worker(
temporal_client,
task_queue='predictions_batch-queue',
workflows=[PredictionsBatch, PredictionProcess,
FormatAndExportPrediction],
workflows=[PredictionsBatch, PredictionProcess, FormatAndExportPrediction],
activities=[
# MLFlow
activities.request_predict,
@@ -181,15 +178,15 @@ async def main():
activities.load_custom_query,
activities.repeat_last_prediction,
activities.export_data_to_postgres,
activities.write_metrics
activities.write_metrics,
],
max_concurrent_workflow_tasks=50,
max_concurrent_activities=50,
max_concurrent_local_activities=50,
max_cached_workflows=200,
workflow_task_poller_behavior=PollerBehaviorAutoscaling(),
activity_task_poller_behavior=PollerBehaviorAutoscaling()
)
activity_task_poller_behavior=PollerBehaviorAutoscaling(),
),
]
handlers = []
@@ -203,7 +200,7 @@ async def main():
# If an exception occurs in any of the worker handlers, it will be propagated here.
await asyncio.gather(*handlers)
except BaseException as e: # NOSONAR
logger.custom_error(f"An unhandled exception occurred: {e}", metadata)
logger.custom_error(f'An unhandled exception occurred: {e}', metadata)
finally:
if notification_handler:
notification_handler.shutdown()
@@ -232,12 +229,12 @@ def start_prometheus_server():
SystemExit: If the metrics server fails to start
"""
try:
port = int(os.getenv("HTTP_METRICS_PORT", 9090))
port = int(os.getenv('HTTP_METRICS_PORT', 9090))
start_http_server(port)
print(f"Prometheus server started on port {port}.")
print(f'Prometheus server started on port {port}.')
metrics.APP_UP.labels(pod_id=POD_ID).set(1) # Mark app as UP
except Exception as e:
print(f"Failed to start Prometheus server: {e}")
print(f'Failed to start Prometheus server: {e}')
os._exit(1)

44
laborious/workflows/minimal_retrain.py
View File

@@ -1,14 +1,16 @@
from temporalio import workflow
with workflow.unsafe.imports_passed_through():
from laborious.activities.activities import Activities
from typing import Any
from sientia_do.temporal.policies import retry_policy
from datetime import timedelta
from typing import Any
from sientia_do.temporal.policies import retry_policy
from laborious.activities.activities import Activities
@workflow.defn(name="minimal_retrain")
class MinimalRetrain():
@workflow.defn(name='minimal_retrain')
class MinimalRetrain:
"""
Automated model retraining workflow for the Laborious system.
@@ -63,24 +65,24 @@ class MinimalRetrain():
'schedule_name': input_data['schedule_name'],
'model_name': input_data['model_name'],
'model_id': input_data['model_id'],
'workflow_name': 'minimal_retrain'
'workflow_name': 'minimal_retrain',
}
}
model_name = input_data['model_name']
model_config = input_data.get('model_config', {})
storage_result = await workflow.execute_local_activity_method(
storage_result = await workflow.execute_activity_method(
Activities.query_to_minio,
{
**metadata,
'query': input_data['query'],
'datetime_columns': input_data.get('datetime_columns', []),
'model_name': model_name,
'object_prefix': f'retrain_datasets/{model_name}/data'
'object_prefix': f'retrain_datasets/{model_name}/data',
},
retry_policy=retry_policy,
start_to_close_timeout=timedelta(seconds=600)
start_to_close_timeout=timedelta(seconds=600),
)
if not storage_result['success']:
@@ -92,38 +94,32 @@ class MinimalRetrain():
**metadata,
'object_key': storage_result['object_key'],
'model_name': model_name,
'model_config': model_config
'model_config': model_config,
},
retry_policy=retry_policy,
start_to_close_timeout=timedelta(hours=1)
start_to_close_timeout=timedelta(hours=1),
)
if experiment_response['success']:
update_report = await workflow.execute_activity_method(
Activities.update_production_model,
{
**metadata,
'model_name': model_name,
**experiment_response
},
{**metadata, 'model_name': model_name, **experiment_response},
retry_policy=retry_policy,
start_to_close_timeout=timedelta(seconds=60)
start_to_close_timeout=timedelta(seconds=60),
)
else:
update_report = {}
report = await workflow.execute_activity_method(
report = await workflow.execute_local_activity_method(
Activities.format_retrain_report,
{
**metadata,
'experiment_response': experiment_response,
'model_id': input_data['model_id'],
'model_name': model_name,
'update_report': update_report
'update_report': update_report,
},
retry_policy=retry_policy,
start_to_close_timeout=timedelta(seconds=60)
start_to_close_timeout=timedelta(seconds=60),
)
await workflow.execute_activity_method(
@@ -132,8 +128,8 @@ class MinimalRetrain():
**metadata,
'data': report,
'schema': input_data['schema'],
'table_name': input_data['table_name']
'table_name': input_data['table_name'],
},
retry_policy=retry_policy,
start_to_close_timeout=timedelta(seconds=600)
start_to_close_timeout=timedelta(seconds=600),
)

46
laborious/workflows/predictions_batch.py
View File

@@ -1,14 +1,16 @@
from temporalio import workflow
with workflow.unsafe.imports_passed_through():
from laborious.activities.activities import Activities
from typing import Any
from sientia_do.temporal.policies import retry_policy
from datetime import timedelta
from typing import Any
from sientia_do.temporal.policies import retry_policy
from laborious.activities.activities import Activities
@workflow.defn(name="predictions_batch")
class PredictionsBatch():
@workflow.defn(name='predictions_batch')
class PredictionsBatch:
"""
Main batch prediction workflow for the Laborious system.
@@ -74,7 +76,7 @@ class PredictionsBatch():
'schedule_name': input_data['schedule_name'],
'model_name': input_data['model_name'],
'model_id': input_data['model_id'],
'workflow_name': 'predictions_batch'
'workflow_name': 'predictions_batch',
}
}
@@ -84,10 +86,10 @@ class PredictionsBatch():
{
**metadata,
'query': input_data['query'],
'datetime_columns': input_data.get('datetime_columns', [])
'datetime_columns': input_data.get('datetime_columns', []),
},
retry_policy=retry_policy,
start_to_close_timeout=timedelta(seconds=300)
start_to_close_timeout=timedelta(seconds=300),
)
# Prepare input for prediction_process workflow
@@ -98,28 +100,18 @@ class PredictionsBatch():
'table_name': input_data['table_name'],
'model_id': input_data['model_id'],
'model_name': input_data['model_name'],
'input_filters': input_data.get('input_filters', {
'EMPTY_DATA': {
'POLICY': 'STOP'
}
}),
'mlflow_transform_filters': input_data.get('mlflow_transform_filters', {
'API_ERROR': {
'POLICY': 'STOP'
}
}),
'mlflow_predict_filters': input_data.get('mlflow_predict_filters', {
'API_ERROR': {
'POLICY': 'STOP'
}
}),
'input_filters': input_data.get('input_filters', {'EMPTY_DATA': {'POLICY': 'STOP'}}),
'mlflow_transform_filters': input_data.get(
'mlflow_transform_filters', {'API_ERROR': {'POLICY': 'STOP'}}
),
'mlflow_predict_filters': input_data.get(
'mlflow_predict_filters', {'API_ERROR': {'POLICY': 'STOP'}}
),
'model_config': input_data.get('model_config', {}),
'path_priority': input_data.get('path_priority', ['STOP', 'CONTINUE', 'REPEAT']),
'opc_output_config': input_data.get('opc_output_config', {}),
'prediction_store_policy': input_data.get(
'prediction_store_policy', 'lts:1')
'prediction_store_policy': input_data.get('prediction_store_policy', 'lts:1'),
}
# Execute prediction process workflow
await workflow.execute_child_workflow(
'prediction_process', prediction_input)
await workflow.execute_child_workflow('prediction_process', prediction_input)

42
laborious/workflows/sub_workflows/format_and_export_prediction.py
View File

@@ -1,15 +1,17 @@
from temporalio import workflow
with workflow.unsafe.imports_passed_through():
from laborious.activities.activities import Activities
from typing import Any
from datetime import timedelta
from sientia_do.temporal.policies import retry_policy
from typing import Any
from sientia_do.temporal.constants import DATETIME_FORMAT_WITH_TZ
from sientia_do.temporal.policies import retry_policy
from laborious.activities.activities import Activities
@workflow.defn(name="format_and_export_prediction")
class FormatAndExportPrediction():
@workflow.defn(name='format_and_export_prediction')
class FormatAndExportPrediction:
"""
Data formatting and export workflow for prediction results.
@@ -79,10 +81,10 @@ class FormatAndExportPrediction():
'timestamp': input_data['timestamp'],
'model_id': input_data['model_id'],
'prediction_confidence': prediction_confidence,
'prediction_store_policy': input_data['prediction_store_policy']
'prediction_store_policy': input_data['prediction_store_policy'],
},
retry_policy=retry_policy,
start_to_close_timeout=timedelta(seconds=60)
start_to_close_timeout=timedelta(seconds=60),
)
else:
@@ -94,22 +96,18 @@ class FormatAndExportPrediction():
'timestamp': input_data['timestamp'],
'model_id': input_data['model_id'],
'prediction_confidence': prediction_confidence,
'comment': input_data['comment']
'comment': input_data['comment'],
},
retry_policy=retry_policy,
start_to_close_timeout=timedelta(seconds=60)
start_to_close_timeout=timedelta(seconds=60),
)
# write to opc
prediction = await workflow.execute_activity_method(
Activities.write_opc_data,
{
**metadata,
'opc_output_config': input_data['opc_output_config'],
'data': prediction
},
{**metadata, 'opc_output_config': input_data['opc_output_config'], 'data': prediction},
retry_policy=retry_policy,
start_to_close_timeout=timedelta(seconds=60)
start_to_close_timeout=timedelta(seconds=60),
)
# write to postgres
@@ -120,21 +118,15 @@ class FormatAndExportPrediction():
'schema': input_data['schema'],
'table_name': input_data['table_name'],
'data': prediction,
'timestamp_conversion': {
'column': 'timestamp',
'format': DATETIME_FORMAT_WITH_TZ
}
'timestamp_conversion': {'column': 'timestamp', 'format': DATETIME_FORMAT_WITH_TZ},
},
retry_policy=retry_policy,
start_to_close_timeout=timedelta(seconds=60)
start_to_close_timeout=timedelta(seconds=60),
)
await workflow.execute_activity_method(
Activities.write_metrics,
{
**metadata,
'prediction': prediction
},
{**metadata, 'prediction': prediction},
retry_policy=retry_policy,
start_to_close_timeout=timedelta(seconds=60)
start_to_close_timeout=timedelta(seconds=60),
)

55
laborious/workflows/sub_workflows/prediction_process.py
View File

@@ -1,14 +1,16 @@
from temporalio import workflow
with workflow.unsafe.imports_passed_through():
from laborious.activities.activities import Activities
from typing import Any
from sientia_do.temporal.policies import retry_policy
from datetime import timedelta
from typing import Any
from sientia_do.temporal.policies import retry_policy
from laborious.activities.activities import Activities
@workflow.defn(name="prediction_process")
class PredictionProcess():
@workflow.defn(name='prediction_process')
class PredictionProcess:
"""
Core prediction processing workflow for the Laborious system.
@@ -84,10 +86,7 @@ class PredictionProcess():
# Get last timestamp for incremental processing
last_timestamp = await workflow.execute_local_activity_method(
Activities.get_last_timestamp,
{
**metadata,
'data': data
},
{**metadata, 'data': data},
retry_policy=retry_policy,
start_to_close_timeout=timedelta(minutes=1),
)
@@ -97,7 +96,7 @@ class PredictionProcess():
**metadata,
'filters': input_data['input_filters'],
'data': data,
'path_priority': input_data['path_priority']
'path_priority': input_data['path_priority'],
}
path_flag, confidence, comment = await workflow.execute_local_activity_method(
@@ -116,12 +115,7 @@ class PredictionProcess():
# Request MLFlow model transformation
response_data = await workflow.execute_local_activity_method(
Activities.request_transform,
{
**metadata,
'data': data,
'model_name': model_name,
'model_config': model_config
},
{**metadata, 'data': data, 'model_name': model_name, 'model_config': model_config},
retry_policy=retry_policy,
start_to_close_timeout=timedelta(minutes=5),
)
@@ -134,7 +128,7 @@ class PredictionProcess():
'filters': input_data['mlflow_transform_filters'],
'data': response_data,
'type': 'transform',
'path_priority': input_data['path_priority']
'path_priority': input_data['path_priority'],
},
retry_policy=retry_policy,
start_to_close_timeout=timedelta(minutes=1),
@@ -155,7 +149,7 @@ class PredictionProcess():
'filters': input_data['mlflow_transform_filters'],
'data': transformed_data,
'type': 'transform',
'path_priority': input_data['path_priority']
'path_priority': input_data['path_priority'],
},
retry_policy=retry_policy,
start_to_close_timeout=timedelta(minutes=1),
@@ -172,7 +166,7 @@ class PredictionProcess():
**metadata,
'data': transformed_data,
'model_name': model_name,
'model_config': model_config
'model_config': model_config,
},
retry_policy=retry_policy,
start_to_close_timeout=timedelta(minutes=5),
@@ -186,7 +180,7 @@ class PredictionProcess():
'filters': input_data['mlflow_predict_filters'],
'data': response_data,
'type': 'predict',
'path_priority': input_data['path_priority']
'path_priority': input_data['path_priority'],
},
retry_policy=retry_policy,
start_to_close_timeout=timedelta(minutes=1),
@@ -214,12 +208,19 @@ class PredictionProcess():
'schema': input_data['schema'],
'table_name': input_data['table_name'],
'comment': comment,
'prediction_store_policy': input_data['prediction_store_policy']
}
'prediction_store_policy': input_data['prediction_store_policy'],
},
)
async def path_flag_handler(self, data: dict, path_flag: str, input_data: dict,
confidence: int, last_timestamp: str, comment: str) -> bool:
async def path_flag_handler(
self,
data: dict,
path_flag: str,
input_data: dict,
confidence: int,
last_timestamp: str,
comment: str,
) -> bool:
"""
Handle path decisions based on filter results and confidence levels.
@@ -265,7 +266,7 @@ class PredictionProcess():
'schema': schema,
'table_name': table_name,
'model': model_id,
'last_timestamp': last_timestamp
'last_timestamp': last_timestamp,
},
retry_policy=retry_policy,
start_to_close_timeout=timedelta(minutes=1),
@@ -288,8 +289,8 @@ class PredictionProcess():
'table_name': table_name,
'comment': comment,
'opc_output_config': input_data['opc_output_config'],
'prediction_store_policy': input_data['prediction_store_policy']
}
'prediction_store_policy': input_data['prediction_store_policy'],
},
)
return True

6
mlruns/0/meta.yaml Normal file
View File

@@ -0,0 +1,6 @@
artifact_location: file:///home/grezewave/Documents/projects/sientia/sientia-dataops-laborious_temporal/mlruns/0
creation_time: 1760447041053
experiment_id: '0'
last_update_time: 1760447041053
lifecycle_stage: active
name: Default

6
mlruns/586524947870967910/meta.yaml Normal file
View File

@@ -0,0 +1,6 @@
artifact_location: file:///home/grezewave/Documents/projects/sientia/sientia-dataops-laborious_temporal/mlruns/586524947870967910
creation_time: 1760447067255
experiment_id: '586524947870967910'
last_update_time: 1760447067255
lifecycle_stage: active
name: test

5
mlruns/models/test/meta.yaml Normal file
View File

@@ -0,0 +1,5 @@
aliases: {}
creation_timestamp: 1760447068191
description: null
last_updated_timestamp: 1760447068191
name: test

159
pyproject.toml Normal file
View File

@@ -0,0 +1,159 @@
[build-system]
requires = ["setuptools>=61.0"]
build-backend = "setuptools.build_meta"
[project]
name = "laborious"
version = "0.0.0"
description = "Sientia DataOps Laborious - ML Model Orchestration System"
readme = "README.md"
requires-python = ">=3.11"
authors = [
{name = "Aignosi", email = "dev@aignosi.com"}
]
[tool.ruff]
line-length = 100
target-version = "py311"
exclude = [
".git",
".venv",
"venv",
"__pycache__",
"*.pyc",
".pytest_cache",
"htmlcov",
]
[tool.ruff.lint]
select = [
"E", # pycodestyle errors
"W", # pycodestyle warnings
"F", # pyflakes
"I", # isort
"B", # flake8-bugbear
"C4", # flake8-comprehensions
"UP", # pyupgrade
"N", # pep8-naming
"YTT", # flake8-2020
"S", # flake8-bandit
"BLE", # flake8-blind-except
"A", # flake8-builtins
"C90", # mccabe complexity
]
ignore = [
"BLE001", # ignore blind except, we need to send notifications with any error
"E501", # line too long (handled by formatter)
"S101", # use of assert (needed for tests)
"S105", # possible hardcoded password (false positives)
"S106", # possible hardcoded password (false positives)
"N802", # function name should be lowercase (temporal decorators)
"N806", # variable in function should be lowercase
]
[tool.ruff.lint.per-file-ignores]
"tests/**/*.py" = [
"S101", # assert allowed in tests
"S105", # hardcoded passwords ok in tests
"S106", # hardcoded passwords ok in tests
]
[tool.ruff.lint.mccabe]
max-complexity = 15
[tool.ruff.format]
quote-style = "single"
indent-style = "space"
line-ending = "auto"
[tool.mypy]
python_version = "3.11"
warn_return_any = false
warn_unused_configs = true
disallow_untyped_defs = false
disallow_incomplete_defs = false
check_untyped_defs = true
no_implicit_optional = true
warn_redundant_casts = true
warn_unused_ignores = false
warn_no_return = true
strict_equality = true
ignore_missing_imports = true
# Ignore missing imports for external packages
[[tool.mypy.overrides]]
module = "temporalio.*"
ignore_missing_imports = true
[[tool.mypy.overrides]]
module = "sientia_do.*"
ignore_missing_imports = true
[[tool.mypy.overrides]]
module = "mlflow.*"
ignore_missing_imports = true
[[tool.mypy.overrides]]
module = "prometheus_client.*"
ignore_missing_imports = true
[[tool.mypy.overrides]]
module = "sientia.*"
ignore_missing_imports = true
[[tool.mypy.overrides]]
module = "pandas.*"
ignore_missing_imports = true
[tool.pytest.ini_options]
testpaths = ["tests"]
python_files = ["test_*.py"]
python_classes = ["Test*"]
python_functions = ["test_*"]
addopts = [
"-v",
"--strict-markers",
"--cov=model_manager",
"--cov-report=term-missing",
"--cov-report=html",
"--cov-report=xml",
]
markers = [
"asyncio: marks tests as async",
"integration: marks tests as integration tests",
"unit: marks tests as unit tests",
]
[tool.coverage.run]
source = ["model_manager"]
omit = [
"*/tests/*",
"*/venv/*",
"*/__pycache__/*",
"*/site-packages/*",
]
branch = true
[tool.coverage.report]
precision = 2
show_missing = true
skip_covered = false
exclude_lines = [
"pragma: no cover",
"def __repr__",
"def __str__",
"raise AssertionError",
"raise NotImplementedError",
"if __name__ == .__main__.:",
"if TYPE_CHECKING:",
"class .*\\bProtocol\\):",
"@(abc\\.)?abstractmethod",
]
[tool.coverage.html]
directory = "htmlcov"
[tool.bandit]
exclude_dirs = ["tests", "venv", ".venv"]
skips = ["B101", "B601"] # Skip assert and shell injection in controlled environments

19
requirements-dev.txt Normal file
View File

@@ -0,0 +1,19 @@
# Development and Testing Dependencies
# These packages are only needed for development, testing, and code quality checks
# Install with: pip install -r requirements-dev.txt
# Code Quality & Linting
ruff>=0.1.0 # Fast Python linter and formatter (replaces flake8, black, isort)
mypy>=1.7.0 # Static type checker
bandit>=1.7.5 # Security vulnerability scanner
pandas-stubs>=2.0.0 # Type stubs for pandas
types-requests>=2.31.0 # Type stubs for requests
# Testing
pytest>=7.4.0 # Testing framework
pytest-cov>=4.1.0 # Coverage plugin for pytest
pytest-asyncio>=0.21.0 # Async test support (already in main requirements)
# Development Tools
ipython>=8.12.0 # Enhanced Python shell
ipdb>=0.13.13 # IPython debugger

79
tests/laborious/activities/test_activities.py
View File

@@ -1,18 +1,19 @@
from unittest.mock import ANY, MagicMock, patch
from pytest import mark
from unittest.mock import patch, MagicMock, ANY
from sientia_do.temporal.activities.postgres import Postgres
from laborious.activities.activities import Activities
from laborious.activities.mlflow import MLFlow
from laborious.activities.gates import Gates
from laborious.activities.mlflow import MLFlow
from laborious.activities.opc import OPC
from laborious.activities.storage import Storage
@patch('laborious.activities.activities.Postgres.__init__')
@patch('laborious.activities.activities.Storage.__init__')
@patch('laborious.activities.activities.MLFlow.__init__')
@patch('laborious.activities.activities.OPC.__init__')
@patch('laborious.activities.activities.Gates.__init__')
def test___init__(mock_gates_init, mock_opc_init, mock_mlflow_init, mock_postgres_init):
def test___init__(mock_gates_init, mock_opc_init, mock_mlflow_init, mock_storage_init):
postgres_config = {
'host': 'localhost',
'port': 5432,
@@ -20,20 +21,23 @@ def test___init__(mock_gates_init, mock_opc_init, mock_mlflow_init, mock_postgre
'password': 'postgres',
'dbname': 'postgres',
'min_connections': 1,
'max_connections': 10
'max_connections': 10,
}
mlflow_config = {
'host': 'localhost',
'port': 5000,
'username': 'mlflow',
'password': 'mlflow'
minio_config = {
'endpoint_url': 'localhost:9000',
'access_key': 'minio',
'secret_key': 'minio123',
'region_name': 'us-east-1',
'default_bucket': 'test',
}
mlflow_config = {'host': 'localhost', 'port': 5000, 'username': 'mlflow', 'password': 'mlflow'}
opc_config = {
'bootstrap_servers': 'localhost:9092',
'polling_time': 1000,
'group_id': 'test-group'
'group_id': 'test-group',
}
logger = MagicMock()
@@ -42,18 +46,19 @@ def test___init__(mock_gates_init, mock_opc_init, mock_mlflow_init, mock_postgre
activities = Activities(
postgres_config=postgres_config,
mlflow_config=mlflow_config,
minio_config=minio_config,
opc_config=opc_config,
logger=logger,
notification_handler=notification_handler
notification_handler=notification_handler,
)
assert isinstance(activities, Activities)
assert isinstance(activities, Postgres)
assert isinstance(activities, Storage)
assert isinstance(activities, MLFlow)
assert isinstance(activities, OPC)
assert isinstance(activities, Gates)
mock_postgres_init.assert_called_once_with(
mock_storage_init.assert_called_once_with(
ANY,
host=postgres_config['host'],
port=postgres_config['port'],
@@ -62,8 +67,9 @@ def test___init__(mock_gates_init, mock_opc_init, mock_mlflow_init, mock_postgre
dbname=postgres_config['dbname'],
min_connections=postgres_config['min_connections'],
max_connections=postgres_config['max_connections'],
minio_config=minio_config,
logger=logger,
notification_handler=notification_handler
notification_handler=notification_handler,
)
mock_mlflow_init.assert_called_once_with(
@@ -72,30 +78,25 @@ def test___init__(mock_gates_init, mock_opc_init, mock_mlflow_init, mock_postgre
mlflow_port=mlflow_config['port'],
mlflow_username=mlflow_config['username'],
mlflow_password=mlflow_config['password'],
minio_config=minio_config,
logger=logger,
notification_handler=notification_handler
notification_handler=notification_handler,
)
mock_opc_init.assert_called_once_with(
ANY,
opc_servers=opc_config,
logger=logger,
notification_handler=notification_handler
ANY, opc_servers=opc_config, logger=logger, notification_handler=notification_handler
)
mock_gates_init.assert_called_once_with(
ANY,
logger=logger,
notification_handler=notification_handler
ANY, logger=logger, notification_handler=notification_handler
)
@mark.asyncio
@patch('laborious.activities.activities.Postgres', return_value=MagicMock())
@patch('laborious.activities.activities.Storage', return_value=MagicMock())
@patch('laborious.activities.activities.MLFlow', return_value=MagicMock())
@patch('laborious.activities.activities.OPC', return_value=MagicMock())
async def test_shutdown(mock_opc_init,
_mock_mlflow_init, mock_postgres_init):
async def test_shutdown(mock_opc_init, _mock_mlflow_init, mock_storage_init):
postgres_config = {
'host': 'localhost',
'port': 5432,
@@ -103,20 +104,23 @@ async def test_shutdown(mock_opc_init,
'password': 'postgres',
'dbname': 'postgres',
'min_connections': 1,
'max_connections': 10
'max_connections': 10,
}
mlflow_config = {
'host': 'localhost',
'port': 5000,
'username': 'mlflow',
'password': 'mlflow'
minio_config = {
'endpoint_url': 'localhost:9000',
'access_key': 'minio',
'secret_key': 'minio123',
'region_name': 'us-east-1',
'default_bucket': 'test',
}
mlflow_config = {'host': 'localhost', 'port': 5000, 'username': 'mlflow', 'password': 'mlflow'}
opc_config = {
'bootstrap_servers': 'localhost:9092',
'polling_time': 1000,
'group_id': 'test-group'
'group_id': 'test-group',
}
logger = MagicMock()
@@ -125,11 +129,12 @@ async def test_shutdown(mock_opc_init,
activities = Activities(
postgres_config=postgres_config,
mlflow_config=mlflow_config,
minio_config=minio_config,
opc_config=opc_config,
logger=logger,
notification_handler=notification_handler
notification_handler=notification_handler,
)
await activities.shutdown()
mock_opc_init.shutdown.assert_called_once()
mock_postgres_init.close.assert_called_once()
mock_storage_init.close.assert_called_once()

235
tests/laborious/activities/test_gates.py
View File

@@ -1,6 +1,8 @@
from unittest.mock import MagicMock, ANY, patch
from unittest.mock import ANY, MagicMock, patch
from pytest import fixture, mark
from sientia_do.notifications.models import NotificationLevel
from laborious.activities.gates import Gates
@@ -20,11 +22,11 @@ def gates_activity():
metadata = {
"metadata": {
"model_id": "test_model",
"model_name": "test_model",
"workflow_name": "test_workflow",
"schema_name": "test_schedule",
'metadata': {
'model_id': 'test_model',
'model_name': 'test_model',
'workflow_name': 'test_workflow',
'schema_name': 'test_schedule',
},
}
@@ -34,20 +36,18 @@ async def test_input_gate_invalid_filter(gates_activity):
# Arrange
input_data = {
**metadata,
'filters': {
'INVALID_FILTER': {'POLICY': 'STOP'}
},
'filters': {'INVALID_FILTER': {'POLICY': 'STOP'}},
'data': {'value': [1, 2, 3]},
'path_priority': ['STOP', 'CONTINUE', 'REPEAT']
'path_priority': ['STOP', 'CONTINUE', 'REPEAT'],
}
# Act
result = await gates_activity.input_gate(input_data)
# Assert
assert result == (None, 0, "")
assert result == (None, 0, '')
gates_activity.error.assert_called_once_with(
"Filter INVALID_FILTER not found", metadata['metadata']
'Filter INVALID_FILTER not found', metadata['metadata']
)
@@ -57,28 +57,27 @@ async def test_input_gate_filter_exception(mock_input_filter_functions, gates_ac
# Arrange
mock_input_filter_functions.__contains__.return_value = True
mock_input_filter_functions.__getitem__.return_value = MagicMock(
side_effect=Exception("Test error"))
side_effect=Exception('Test error')
)
input_data = {
**metadata,
'filters': {
'EMPTY_DATA': {'policy': 'STOP', 'config': {}}
},
'filters': {'EMPTY_DATA': {'policy': 'STOP', 'config': {}}},
'data': {'value': []},
'path_priority': ['STOP', 'CONTINUE', 'REPEAT']
'path_priority': ['STOP', 'CONTINUE', 'REPEAT'],
}
# Act
result = await gates_activity.input_gate(input_data)
# Assert
assert result == (None, 0, "")
assert result == (None, 0, '')
gates_activity.send_notification.assert_called_once_with(
metadata=metadata['metadata'],
notification_id="INTPUT_GATE_ERROR__EMPTY_DATA",
notification_id='INTPUT_GATE_ERROR__EMPTY_DATA',
message="Error in filter EMPTY_DATA:{'policy': 'STOP', 'config': {}}: \n Test error",
block="input_gate",
block='input_gate',
level=NotificationLevel.ERROR,
attachment_content=ANY
attachment_content=ANY,
)
@@ -89,14 +88,14 @@ async def test_input_gate_no_filters(gates_activity):
**metadata,
'filters': {},
'data': {'value': [1, 2, 3]},
'path_priority': ['CONTINUE', 'STOP', 'REPEAT']
'path_priority': ['CONTINUE', 'STOP', 'REPEAT'],
}
# Act
result = await gates_activity.input_gate(input_data)
# Assert
assert result == (None, 0, "")
assert result == (None, 0, '')
gates_activity.debug.assert_called()
@@ -105,18 +104,16 @@ async def test_input_gate_with_filter(gates_activity):
# Arrange
input_data = {
**metadata,
'filters': {
'EMPTY_DATA': {'policy': 'STOP', 'config': {}}
},
'filters': {'EMPTY_DATA': {'policy': 'STOP', 'config': {}}},
'data': {'value': []},
'path_priority': ['STOP', 'CONTINUE', 'REPEAT']
'path_priority': ['STOP', 'CONTINUE', 'REPEAT'],
}
# Act
result = await gates_activity.input_gate(input_data)
# Assert
assert result == ('STOP', -1, "Input data with bad quality")
assert result == ('STOP', -1, 'Input data with bad quality')
gates_activity.debug.assert_called()
@@ -125,51 +122,49 @@ async def test_mlflow_response_gate_invalid_filter(gates_activity):
# Arrange
input_data = {
**metadata,
'filters': {
'INVALID_FILTER': {'POLICY': 'STOP'}
},
'filters': {'INVALID_FILTER': {'POLICY': 'STOP'}},
'data': {'content': {'message': 'success'}},
'type': 'test',
'path_priority': ['STOP', 'CONTINUE', 'REPEAT']
'path_priority': ['STOP', 'CONTINUE', 'REPEAT'],
}
# Act
result = await gates_activity.mlflow_response_gate(input_data)
# Assert
assert result == (None, 0, "")
assert result == (None, 0, '')
@mark.asyncio
@patch('laborious.activities.gates.mlflow_response_filter_functions')
async def test_mlflow_response_gate_filter_exception(mock_mlflow_response_filter_functions,
gates_activity):
async def test_mlflow_response_gate_filter_exception(
mock_mlflow_response_filter_functions, gates_activity
):
# Arrange
mock_mlflow_response_filter_functions.__contains__.return_value = True
mock_mlflow_response_filter_functions.__getitem__.return_value = MagicMock(
side_effect=Exception("Test error"))
side_effect=Exception('Test error')
)
input_data = {
**metadata,
'filters': {
'INVALID_FILTER': {'POLICY': 'STOP'}
},
'filters': {'INVALID_FILTER': {'POLICY': 'STOP'}},
'data': {'content': {'message': 'success'}},
'type': 'test',
'path_priority': ['STOP', 'CONTINUE', 'REPEAT']
'path_priority': ['STOP', 'CONTINUE', 'REPEAT'],
}
# Act
result = await gates_activity.mlflow_response_gate(input_data)
# Assert
assert result == (None, 0, "")
assert result == (None, 0, '')
gates_activity.send_notification.assert_called_once_with(
metadata=metadata['metadata'],
notification_id="MLFLOW_GATE_RESPONSE_FILTER__INVALID_FILTER",
notification_id='MLFLOW_GATE_RESPONSE_FILTER__INVALID_FILTER',
message="Error in filter INVALID_FILTER:{'POLICY': 'STOP'}: \n Test error",
block="mlflow_gate",
block='mlflow_gate',
level=NotificationLevel.ERROR,
attachment_content=ANY
attachment_content=ANY,
)
@@ -181,14 +176,14 @@ async def test_mlflow_response_gate_no_filters(gates_activity):
'filters': {},
'data': {'content': {'message': 'success'}},
'type': 'test',
'path_priority': ['CONTINUE', 'STOP', 'REPEAT']
'path_priority': ['CONTINUE', 'STOP', 'REPEAT'],
}
# Act
result = await gates_activity.mlflow_response_gate(input_data)
# Assert
assert result == (None, 0, "")
assert result == (None, 0, '')
gates_activity.debug.assert_called()
@@ -197,25 +192,20 @@ async def test_mlflow_response_gate_with_filter(gates_activity):
# Arrange
input_data = {
**metadata,
'filters': {
'API_ERROR': {'policy': 'STOP'}
},
'filters': {'API_ERROR': {'policy': 'STOP'}},
'data': {
'success': False,
'content': {
'message': 'API error occurred',
'traceback': 'error trace'
}
'content': {'message': 'API error occurred', 'traceback': 'error trace'},
},
'type': 'test',
'path_priority': ['STOP', 'CONTINUE', 'REPEAT']
'path_priority': ['STOP', 'CONTINUE', 'REPEAT'],
}
# Act
result = await gates_activity.mlflow_response_gate(input_data)
# Assert
assert result == ('STOP', -1, "API error occurred")
assert result == ('STOP', -1, 'API error occurred')
gates_activity.debug.assert_called()
gates_activity.send_notification.assert_called()
@@ -225,58 +215,53 @@ async def test_mlflow_content_gate_invalid_filter(gates_activity):
# Arrange
input_data = {
**metadata,
'filters': {
'INVALID_FILTER': {'POLICY': 'STOP'}
},
'filters': {'INVALID_FILTER': {'POLICY': 'STOP'}},
'data': {'value': [1, 2, 3]},
'type': 'test',
'path_priority': ['STOP', 'CONTINUE', 'REPEAT']
'path_priority': ['STOP', 'CONTINUE', 'REPEAT'],
}
# Act
result = await gates_activity.mlflow_content_gate(input_data)
# Assert
assert result == (None, 0, "")
assert result == (None, 0, '')
@mark.asyncio
@patch('laborious.activities.gates.mlflow_content_filter_functions')
async def test_mlflow_content_gate_filter_exception(mock_mlflow_content_filter_functions,
gates_activity):
async def test_mlflow_content_gate_filter_exception(
mock_mlflow_content_filter_functions, gates_activity
):
# Arrange
mock_mlflow_content_filter_functions.__contains__.return_value = True
mock_mlflow_content_filter_functions.__getitem__.return_value = MagicMock(
side_effect=Exception("Test error"))
side_effect=Exception('Test error')
)
input_data = {
**metadata,
'filters': {
'API_ERROR': {'POLICY': 'STOP'}
},
'filters': {'API_ERROR': {'POLICY': 'STOP'}},
'data': {
'success': False,
'content': {
'message': 'API error occurred',
'traceback': 'error trace'
}
'content': {'message': 'API error occurred', 'traceback': 'error trace'},
},
'type': 'test',
'path_priority': ['STOP', 'CONTINUE', 'REPEAT']
'path_priority': ['STOP', 'CONTINUE', 'REPEAT'],
}
# Act
result = await gates_activity.mlflow_content_gate(input_data)
# Assert
assert result == (None, 0, "")
assert result == (None, 0, '')
gates_activity.debug.assert_called()
gates_activity.send_notification.assert_called_once_with(
metadata=metadata['metadata'],
notification_id="MLFLOW_GATE_CONTENT_FILTER__API_ERROR",
notification_id='MLFLOW_GATE_CONTENT_FILTER__API_ERROR',
message="Error in filter API_ERROR:{'POLICY': 'STOP'}: \n Test error",
block="mlflow_gate",
block='mlflow_gate',
level=NotificationLevel.ERROR,
attachment_content=ANY
attachment_content=ANY,
)
@@ -288,14 +273,14 @@ async def test_mlflow_content_gate_no_filters(gates_activity):
'filters': {},
'data': {'value': [1, 2, 3]},
'type': 'test',
'path_priority': ['CONTINUE', 'STOP', 'REPEAT']
'path_priority': ['CONTINUE', 'STOP', 'REPEAT'],
}
# Act
result = await gates_activity.mlflow_content_gate(input_data)
# Assert
assert result == (None, 0, "")
assert result == (None, 0, '')
gates_activity.debug.assert_called()
@@ -304,20 +289,17 @@ async def test_mlflow_content_gate_with_filter(gates_activity):
# Arrange
input_data = {
**metadata,
'filters': {
'NAN_VALUES': {'policy': 'STOP', 'config': {}}
},
'filters': {'NAN_VALUES': {'policy': 'STOP', 'config': {}}},
'data': {'value': [None, None, None]},
'type': 'test',
'path_priority': ['STOP', 'CONTINUE', 'REPEAT']
'path_priority': ['STOP', 'CONTINUE', 'REPEAT'],
}
# Act
result = await gates_activity.mlflow_content_gate(input_data)
# Assert
assert result == (
'STOP', -1, "Transformed data not passed the content filter")
assert result == ('STOP', -1, 'Transformed data not passed the content filter')
gates_activity.debug.assert_called()
gates_activity.send_notification.assert_called()
@@ -328,7 +310,8 @@ def test_get_prediction_store_policy_invalid_policy(gates_activity):
# Act
policy_type, policy_value = gates_activity.get_prediction_store_policy(
prediction_store_policy, metadata)
prediction_store_policy, metadata
)
# Assert
assert policy_type == 'lts'
@@ -341,7 +324,8 @@ def test_get_prediction_store_policy_invalid_policy_value(gates_activity):
# Act
policy_type, policy_value = gates_activity.get_prediction_store_policy(
prediction_store_policy, metadata)
prediction_store_policy, metadata
)
# Assert
assert policy_type == 'lts'
@@ -354,7 +338,8 @@ def test_get_prediction_store_policy_valid_policy_type(gates_activity):
# Act
policy_type, policy_value = gates_activity.get_prediction_store_policy(
prediction_store_policy, metadata)
prediction_store_policy, metadata
)
# Assert
assert policy_type == 'lts'
@@ -367,7 +352,8 @@ def test_get_prediction_store_policy_valid_policy(gates_activity):
# Act
policy_type, policy_value = gates_activity.get_prediction_store_policy(
prediction_store_policy, metadata)
prediction_store_policy, metadata
)
# Assert
assert policy_type == 'erl'
@@ -380,16 +366,12 @@ async def test_format_prediction_no_timestamp(gates_activity):
input_data = {
**metadata,
'data': {
'prediction': {
'2023-05-26 11:12:27': 1
},
'response_time': {
'2023-05-26 11:12:27': 0.1
}
'prediction': {'2023-05-26 11:12:27': 1},
'response_time': {'2023-05-26 11:12:27': 0.1},
},
'model_id': 'test_model',
'prediction_confidence': 0.9,
'prediction_store_policy': 'lts:1'
'prediction_store_policy': 'lts:1',
}
# Act
@@ -402,7 +384,7 @@ async def test_format_prediction_no_timestamp(gates_activity):
assert result['model_id'] == {0: 'test_model'}
assert result['prediction_confidence'] == {0: 0.9}
assert result['prediction_status'] == {0: 'Good'}
assert result['comments'] == {0: ""}
assert result['comments'] == {0: ''}
@mark.asyncio
@@ -420,11 +402,11 @@ async def test_format_prediction_with_timestamp_erl(gates_activity):
'2023-05-26 11:12:27': 0.1,
'2023-05-26 11:12:28': 0.2,
'2023-05-26 11:12:29': 0.3,
}
},
},
'model_id': 'test_model',
'prediction_confidence': 0.9,
'prediction_store_policy': 'erl:2'
'prediction_store_policy': 'erl:2',
}
# Act
@@ -433,12 +415,11 @@ async def test_format_prediction_with_timestamp_erl(gates_activity):
# Assert
assert result['prediction'] == {0: 2, 1: 1}
assert result['response_time'] == {0: 0.2, 1: 0.1}
assert result['timestamp'] == {
0: '2023-05-26 11:12:28', 1: '2023-05-26 11:12:27'}
assert result['timestamp'] == {0: '2023-05-26 11:12:28', 1: '2023-05-26 11:12:27'}
assert result['model_id'] == {0: 'test_model', 1: 'test_model'}
assert result['prediction_confidence'] == {0: 0.9, 1: 0.9}
assert result['prediction_status'] == {0: 'Good', 1: 'Good'}
assert result['comments'] == {0: "", 1: ""}
assert result['comments'] == {0: '', 1: ''}
@mark.asyncio
@@ -456,11 +437,11 @@ async def test_format_prediction_with_timestamp_lts(gates_activity):
'2023-05-26 11:12:27': 0.1,
'2023-05-26 11:12:28': 0.2,
'2023-05-26 11:12:29': 0.3,
}
},
},
'model_id': 'test_model',
'prediction_confidence': 0.9,
'prediction_store_policy': 'lts:2'
'prediction_store_policy': 'lts:2',
}
# Act
@@ -469,12 +450,11 @@ async def test_format_prediction_with_timestamp_lts(gates_activity):
# Assert
assert result['prediction'] == {0: 3, 1: 2}
assert result['response_time'] == {0: 0.3, 1: 0.2}
assert result['timestamp'] == {
0: '2023-05-26 11:12:29', 1: '2023-05-26 11:12:28'}
assert result['timestamp'] == {0: '2023-05-26 11:12:29', 1: '2023-05-26 11:12:28'}
assert result['model_id'] == {0: 'test_model', 1: 'test_model'}
assert result['prediction_confidence'] == {0: 0.9, 1: 0.9}
assert result['prediction_status'] == {0: 'Good', 1: 'Good'}
assert result['comments'] == {0: "", 1: ""}
assert result['comments'] == {0: '', 1: ''}
@mark.asyncio
@@ -482,22 +462,23 @@ async def test_format_prediction_with_timestamp_invalid_policy(gates_activity):
# Arrange
input_data = {
**metadata,
'data': {'prediction': [1, 2, 3],
'data': {
'prediction': [1, 2, 3],
'response_time': [0.1, 0.2, 0.3],
'timestamp': ['2023-05-26 11:12:27', '2023-05-26 11:12:28', '2023-05-26 11:12:29']},
'timestamp': ['2023-05-26 11:12:27', '2023-05-26 11:12:28', '2023-05-26 11:12:29'],
},
'model_id': 'test_model',
'prediction_confidence': 0.9,
'prediction_store_policy': 'lts:2'
'prediction_store_policy': 'lts:2',
}
gates_activity.get_prediction_store_policy = MagicMock(
return_value=('invalid', 1))
gates_activity.get_prediction_store_policy = MagicMock(return_value=('invalid', 1))
try:
result = await gates_activity.format_prediction(input_data)
await gates_activity.format_prediction(input_data)
except ValueError as e:
assert str(e) == "Invalid policy type: invalid"
assert str(e) == 'Invalid policy type: invalid'
else:
assert False, "Expected ValueError"
raise AssertionError('Expected ValueError')
@mark.asyncio
@@ -508,7 +489,7 @@ async def test_format_default_prediction(gates_activity):
'timestamp': '2023-05-26 11:12:27',
'model_id': 'test_model',
'prediction_confidence': 0.1,
'comment': 'Test comment'
'comment': 'Test comment',
}
# Act
@@ -528,12 +509,7 @@ async def test_format_default_prediction(gates_activity):
@mark.asyncio
async def test_get_last_timestamp_with_data(gates_activity):
# Arrange
input_data = {
**metadata,
'data': {
'timestamp': ['2023-05-26 11:12:27', '2023-05-26 11:12:28']
}
}
input_data = {**metadata, 'data': {'timestamp': ['2023-05-26 11:12:27', '2023-05-26 11:12:28']}}
# Act
result = await gates_activity.get_last_timestamp(input_data)
@@ -545,10 +521,7 @@ async def test_get_last_timestamp_with_data(gates_activity):
@mark.asyncio
async def test_get_last_timestamp_no_data(gates_activity):
# Arrange
input_data = {
'data': {},
**metadata
}
input_data = {'data': {}, **metadata}
# Act
result = await gates_activity.get_last_timestamp(input_data)
@@ -567,30 +540,28 @@ async def test_write_metrics(mock_metrics, gates_activity):
'prediction': {
'prediction': [1, 2, 3],
'prediction_confidence': [0.9, 0.8, 0.7],
'response_time': [0.1, 0.2, 0.3]
}
'response_time': [0.1, 0.2, 0.3],
},
}
await gates_activity.write_metrics(input_data)
mock_metrics.PREDICTIONS_WRITTEN_COUNT.labels.assert_called_once_with(
pod_id=gates_activity.pod_id,
model_name=metadata['metadata']['model_name'],
pipeline_name=metadata['metadata']['workflow_name']
pipeline_name=metadata['metadata']['workflow_name'],
)
mock_metrics.PREDICTIONS_WRITTEN_COUNT.labels.return_value.inc.assert_called_once_with()
mock_metrics.PREDICTION_CONFIDENCE_MONITOR.labels.assert_called_once_with(
pod_id=gates_activity.pod_id,
model_name=metadata['metadata']['model_name'],
pipeline_name=metadata['metadata']['workflow_name']
)
mock_metrics.PREDICTION_CONFIDENCE_MONITOR.labels.return_value.set.assert_called_once_with(
0.9
pipeline_name=metadata['metadata']['workflow_name'],
)
mock_metrics.PREDICTION_CONFIDENCE_MONITOR.labels.return_value.set.assert_called_once_with(0.9)
mock_metrics.PREDICTION_RESPONSE_TIME_MONITOR.labels.assert_called_once_with(
pod_id=gates_activity.pod_id,
model_name=metadata['metadata']['model_name'],
pipeline_name=metadata['metadata']['workflow_name']
pipeline_name=metadata['metadata']['workflow_name'],
)
mock_metrics.PREDICTION_RESPONSE_TIME_MONITOR.labels.return_value.observe.assert_called_once_with(
0.1

402
tests/laborious/activities/test_mlflow.py
View File

@@ -1,45 +1,68 @@
from datetime import datetime
from unittest.mock import ANY, MagicMock, patch
from unittest.mock import ANY, MagicMock, call, patch
import numpy as np
from pandas import DataFrame, Timestamp
from pytest import fixture, mark, raises
from sientia_do.temporal.constants import DATETIME_FORMAT, DATETIME_FORMAT_WITH_TZ
from laborious.activities.mlflow import MLFlow
from pytest import fixture, mark
from sientia_do.notifications.models import NotificationLevel
from sientia_do.temporal.constants import DATETIME_FORMAT, DATETIME_FORMAT_WITH_TZ
from laborious.activities.mlflow import MLFlow
@patch("laborious.activities.mlflow.MLFlowRepository")
def test___init__(mock_mlflow_repository):
@patch('laborious.activities.mlflow.MLFlowRepository')
@patch('laborious.activities.mlflow.MinioRepository')
def test___init__(mock_minio_repository, mock_mlflow_repository):
mlflow = MLFlow(
mlflow_host="http://localhost",
mlflow_host='http://localhost',
mlflow_port=5000,
mlflow_username="admin",
mlflow_password="admin",
mlflow_username='admin',
mlflow_password='admin',
minio_config={
'endpoint_url': 'http://localhost:9000',
'access_key': 'minio',
'secret_key': 'minio123',
'region_name': 'us-east-1',
'default_bucket': 'test',
},
logger=MagicMock(),
notification_handler=MagicMock()
notification_handler=MagicMock(),
)
assert mlflow.mlflow_host == "http://localhost"
assert mlflow.mlflow_host == 'http://localhost'
assert mlflow.mlflow_port == 5000
assert mlflow.mlflow_username == "admin"
assert mlflow.mlflow_password == "admin"
assert mlflow.mlflow_username == 'admin'
assert mlflow.mlflow_password == 'admin'
mock_mlflow_repository.assert_called_once_with(
"http://localhost:5000", "admin", "admin", ANY
mock_mlflow_repository.assert_called_once_with('http://localhost:5000', 'admin', 'admin', ANY)
mock_minio_repository.assert_called_once_with(
logger=ANY,
notification_handler=ANY,
minio_endpoint_url='http://localhost:9000',
minio_access_key='minio',
minio_secret_key='minio123',
minio_region_name='us-east-1',
minio_default_bucket='test',
)
@fixture
@patch("laborious.activities.mlflow.MLFlowRepository")
def mlflow(mock_mlflow_repository):
@patch('laborious.activities.mlflow.MLFlowRepository')
@patch('laborious.activities.mlflow.MinioRepository')
def mlflow(mock_minio_repository, mock_mlflow_repository):
mlflow = MLFlow(
mlflow_host="http://localhost:5000",
mlflow_host='http://localhost:5000',
mlflow_port=5000,
mlflow_username="admin",
mlflow_password="admin",
mlflow_username='admin',
mlflow_password='admin',
minio_config={
'endpoint_url': 'http://localhost:9000',
'access_key': 'minio',
'secret_key': 'minio123',
'region_name': 'us-east-1',
'default_bucket': 'test',
},
logger=MagicMock(),
notification_handler=MagicMock()
notification_handler=MagicMock(),
)
mlflow.send_notification = MagicMock()
@@ -48,44 +71,67 @@ def mlflow(mock_mlflow_repository):
metadata = {
"metadata": {
"model_id": "test_model",
"model_name": "test_model",
"workflow_name": "test_workflow",
"schema_name": "test_schedule",
'metadata': {
'model_id': 'test_model',
'model_name': 'test_model',
'workflow_name': 'test_workflow',
'schema_name': 'test_schedule',
},
}
@mark.asyncio
@patch("laborious.activities.mlflow.DataFrame")
@patch("laborious.activities.mlflow.max")
@patch('laborious.activities.mlflow.DataFrame')
@patch('laborious.activities.mlflow.max')
async def test_request_transform_success(mock_max, mock_dataframe, mlflow):
mock_max.return_value = '2024-01-02'
# Mock input data
input_data = {
**metadata,
'data': [
{'timestamp': '2024-01-01', 'variable': 'var1',
'value': 1.0, 'created_at': '2024-01-01 12:00:00'},
{'timestamp': '2024-01-01', 'variable': 'var2',
'value': 2.0, 'created_at': '2024-01-01 12:00:00'},
{'timestamp': '2024-01-02', 'variable': 'var1',
'value': 3.0, 'created_at': '2024-01-02 12:00:00'},
{'timestamp': '2024-01-02', 'variable': 'var2',
'value': 4.0, 'created_at': '2024-01-02 12:00:00'},
{'timestamp': '2024-01-02', 'variable': 'var1',
'value': 1.0, 'created_at': '2024-01-01 12:00:00'},
{'timestamp': '2024-01-02', 'variable': 'var2',
'value': 1.0, 'created_at': '2024-01-01 12:00:00'}
{
'timestamp': '2024-01-01',
'variable': 'var1',
'value': 1.0,
'created_at': '2024-01-01 12:00:00',
},
{
'timestamp': '2024-01-01',
'variable': 'var2',
'value': 2.0,
'created_at': '2024-01-01 12:00:00',
},
{
'timestamp': '2024-01-02',
'variable': 'var1',
'value': 3.0,
'created_at': '2024-01-02 12:00:00',
},
{
'timestamp': '2024-01-02',
'variable': 'var2',
'value': 4.0,
'created_at': '2024-01-02 12:00:00',
},
{
'timestamp': '2024-01-02',
'variable': 'var1',
'value': 1.0,
'created_at': '2024-01-01 12:00:00',
},
{
'timestamp': '2024-01-02',
'variable': 'var2',
'value': 1.0,
'created_at': '2024-01-01 12:00:00',
},
],
'model_name': 'test_model',
'model_config': {}
'model_config': {},
}
# Mock the transform response
expected_response = {'prediction': [0.5, 0.6], 'timestamp': [
'2024-01-01', '2024-01-02']}
expected_response = {'prediction': [0.5, 0.6], 'timestamp': ['2024-01-01', '2024-01-02']}
mlflow.model_monitoring_repository.transform.return_value = expected_response
mock_dataframe.return_value.sort_values.return_value = mock_dataframe.return_value
@@ -114,30 +160,25 @@ async def test_request_transform_success(mock_max, mock_dataframe, mlflow):
@mark.asyncio
@patch("laborious.activities.mlflow.DataFrame")
@patch("laborious.activities.mlflow.to_datetime")
@patch("laborious.activities.mlflow.max")
@patch('laborious.activities.mlflow.DataFrame')
@patch('laborious.activities.mlflow.to_datetime')
@patch('laborious.activities.mlflow.max')
async def test_request_predict(mock_max, mock_to_datetime, mock_dataframe, mlflow):
mock_max.return_value = '2024-01-02'
# Mock input data
input_data = {
**metadata,
'data': {
"variable": {
"2024-01-01": "var1",
"2024-01-02": "var2",
"2024-01-03": "var1",
"2024-01-04": "var2"
'variable': {
'2024-01-01': 'var1',
'2024-01-02': 'var2',
'2024-01-03': 'var1',
'2024-01-04': 'var2',
},
"value": {
"2024-01-01": 1.0,
"2024-01-02": 2.0,
"2024-01-03": 3.0,
"2024-01-04": 4.0
}
'value': {'2024-01-01': 1.0, '2024-01-02': 2.0, '2024-01-03': 3.0, '2024-01-04': 4.0},
},
'model_name': 'test_model',
'model_config': {}
'model_config': {},
}
# Mock the predict response
@@ -148,9 +189,7 @@ async def test_request_predict(mock_max, mock_to_datetime, mock_dataframe, mlflo
response_data = await mlflow.request_predict(input_data)
mock_dataframe.assert_called_once_with(input_data['data'])
mock_dataframe.return_value.replace.assert_called_once_with(
np.nan, None, inplace=True
)
mock_dataframe.return_value.replace.assert_called_once_with(np.nan, None, inplace=True)
mock_dataframe.return_value.__setitem__.assert_any_call(
'timestamp', mock_to_datetime.return_value.dt.strftime.return_value
)
@@ -158,9 +197,7 @@ async def test_request_predict(mock_max, mock_to_datetime, mock_dataframe, mlflo
mock_to_datetime.assert_called_once_with(
mock_dataframe.return_value.__getitem__.return_value, format=DATETIME_FORMAT_WITH_TZ
)
mock_to_datetime.return_value.dt.strftime.assert_called_once_with(
DATETIME_FORMAT
)
mock_to_datetime.return_value.dt.strftime.assert_called_once_with(DATETIME_FORMAT)
# Verify the response
assert response_data == expected_response
@@ -172,98 +209,211 @@ async def test_request_predict(mock_max, mock_to_datetime, mock_dataframe, mlflo
@mark.asyncio
async def test_retrain_model(mlflow):
data = {
"model_id": [4, 5, 6, 7],
"created_at": [1, 2, 3, 4],
"timestamp": [1, 1, 2, 2],
"variable": ["var1", "var2", "var1", "var2"],
"value": [1, 2, 3, 4]
@patch('laborious.activities.mlflow.to_datetime')
async def test_retrain_model_success_data_success_retrain(mock_to_datetime, mlflow):
mlflow.model_monitoring_repository.retrain_model.return_value = {
'success': True,
'experiment': 'test_experiment',
'message': 'Model retrained successfully.',
}
mlflow.model_monitoring_repository.retrain_model.return_value = (
'Model retrained successfully', 'test')
response = await mlflow.retrain_model({
response = await mlflow.retrain_model(
{
**metadata,
'data': data,
'model_name': 'test_model'
})
'object_key': 'test_object_key',
'model_name': 'test_model',
'model_config': {
'target': 'target',
'transform_flavor': 'sklearn',
'predict_flavor': 'pyfunc',
},
}
)
mlflow.model_monitoring_repository.retrain_model.assert_called_once()
raw_data = mlflow.minio_repository.get_parquet_as_dataframe.return_value
timestamp = raw_data.__getitem__.return_value.max.return_value
raw_data.sort_values.assert_called_once_with('created_at', ascending=False)
raw_data.sort_values.return_value.drop_duplicates.assert_called_once_with(
subset=['variable', 'timestamp'], keep='first'
)
raw_data = raw_data.sort_values.return_value.drop_duplicates.return_value
raw_data.drop.assert_has_calls(
[
call(columns=['model_id'], inplace=True, errors='ignore'),
call(columns=['created_at'], inplace=True, errors='ignore'),
]
)
raw_data.pivot.assert_called_once_with(index='timestamp', columns='variable', values='value')
raw_data.pivot.return_value.fillna.assert_called_once_with(np.nan, inplace=True)
raw_data = raw_data.pivot.return_value
raw_data.__setitem__.assert_has_calls(
[
call('timestamp', raw_data.index),
call('timestamp', mock_to_datetime.return_value.dt.strftime.return_value),
call('timestamp', mock_to_datetime.return_value),
]
)
mock_to_datetime.assert_has_calls(
[call(raw_data.__getitem__.return_value, format=DATETIME_FORMAT_WITH_TZ)]
)
mock_to_datetime.assert_has_calls(
[call(raw_data.__getitem__.return_value, format=DATETIME_FORMAT)]
)
mlflow.model_monitoring_repository.retrain_model.assert_called_once_with(
data=raw_data,
model_name='test_model',
model_config={
'target': 'target',
'transform_flavor': 'sklearn',
'predict_flavor': 'pyfunc',
},
metadata=metadata['metadata'],
)
assert response == {
"status": 'Model retrained successfully',
"timestamp": 2,
"experiment": 'test'
'success': True,
'experiment': 'test_experiment',
'message': 'Model retrained successfully.',
'timestamp': timestamp,
}
@mark.asyncio
async def test_retrain_model_error(mlflow):
mlflow.model_monitoring_repository.retrain_model.side_effect = Exception(
'Error retraining model'
)
data = {
"model_id": [4, 5, 6, 7],
"created_at": [1, 2, 3, 4],
"timestamp": [1, 1, 2, 2],
"variable": ["var1", "var2", "var1", "var2"],
"value": [1, 2, 3, 4]
@patch('laborious.activities.mlflow.to_datetime')
async def test_retrain_model_success_data_fail_retrain(mock_to_datetime, mlflow):
mlflow.model_monitoring_repository.retrain_model.return_value = {
'success': False,
'traceback': 'test_traceback',
'message': 'Model retrained failed.',
}
try:
await mlflow.retrain_model({
response = await mlflow.retrain_model(
{
**metadata,
'data': data,
'model_name': 'test_model'
})
except Exception as e:
assert str(e) == 'Error retraining model'
'object_key': 'test_object_key',
'model_name': 'test_model',
'model_config': {
'target': 'target',
'transform_flavor': 'sklearn',
'predict_flavor': 'pyfunc',
},
}
)
raw_data = mlflow.minio_repository.get_parquet_as_dataframe.return_value
timestamp = raw_data.__getitem__.return_value.max.return_value
raw_data.sort_values.assert_called_once_with('created_at', ascending=False)
raw_data.sort_values.return_value.drop_duplicates.assert_called_once_with(
subset=['variable', 'timestamp'], keep='first'
)
raw_data = raw_data.sort_values.return_value.drop_duplicates.return_value
raw_data.drop.assert_has_calls(
[
call(columns=['model_id'], inplace=True, errors='ignore'),
call(columns=['created_at'], inplace=True, errors='ignore'),
]
)
raw_data.pivot.assert_called_once_with(index='timestamp', columns='variable', values='value')
raw_data.pivot.return_value.fillna.assert_called_once_with(np.nan, inplace=True)
raw_data = raw_data.pivot.return_value
raw_data.__setitem__.assert_has_calls(
[
call('timestamp', raw_data.index),
call('timestamp', mock_to_datetime.return_value.dt.strftime.return_value),
call('timestamp', mock_to_datetime.return_value),
]
)
mock_to_datetime.assert_has_calls(
[call(raw_data.__getitem__.return_value, format=DATETIME_FORMAT_WITH_TZ)]
)
mock_to_datetime.assert_has_calls(
[call(raw_data.__getitem__.return_value, format=DATETIME_FORMAT)]
)
mlflow.model_monitoring_repository.retrain_model.assert_called_once_with(
data=raw_data,
model_name='test_model',
model_config={
'target': 'target',
'transform_flavor': 'sklearn',
'predict_flavor': 'pyfunc',
},
metadata=metadata['metadata'],
)
mlflow.send_notification.assert_called_once_with(
metadata=metadata['metadata'],
notification_id='RETRAIN_MODEL_ERROR',
message='Error retraining model test_model: Error retraining model',
message='Error retraining model test_model: Model retrained failed.',
block='retrain_model',
level=NotificationLevel.ERROR,
attachment_content=ANY
attachment_content=ANY,
)
else:
assert False, "No exception raised"
assert response == {
'success': False,
'traceback': 'test_traceback',
'message': 'Model retrained failed.',
'timestamp': timestamp,
}
@mark.asyncio
async def test_retrain_model_data_error(mlflow):
mlflow.minio_repository.get_parquet_as_dataframe.side_effect = Exception(
'Error loading retrain data'
)
response = await mlflow.retrain_model(
{
**metadata,
'object_key': 'test_object_key',
'model_name': 'test_model',
'model_config': {
'target': 'target',
'transform_flavor': 'sklearn',
'predict_flavor': 'pyfunc',
},
}
)
assert response == {
'success': False,
'message': 'Error loading retrain data: Error loading retrain data',
'traceback': ANY,
'timestamp': ANY,
}
@mark.asyncio
async def test_update_production_model(mlflow):
mlflow.model_monitoring_repository.update_production_model.return_value = (
{
"data1": 1,
"data2": 2
}
)
input_data = {
**metadata,
'model_name': 'test_model',
'model_id': 1,
'experiment': 'test',
'timestamp': 2,
'status': 'success'
'status': 'success',
}
response = await mlflow.update_production_model(input_data)
mlflow.model_monitoring_repository.update_production_model.assert_called_once_with(
experiment='test', model_name='test_model')
experiment='test', model_name='test_model', metadata=metadata['metadata']
)
assert response == {
'data1': {0: 1},
'data2': {0: 2},
'model_id': {0: 1},
'model_name': {0: 'test_model'},
'timestamp': {0: 2},
'status': {0: 'success'}
}
assert response == mlflow.model_monitoring_repository.update_production_model.return_value
@mark.asyncio
@@ -278,7 +428,7 @@ async def test_update_production_model_error(mlflow):
'model_id': 1,
'experiment': 'test',
'timestamp': 2,
'status': 'success'
'status': 'success',
}
try:
@@ -291,7 +441,7 @@ async def test_update_production_model_error(mlflow):
message='Error updating production model test_model: Error updating production model',
block='update_production_model',
level=NotificationLevel.ERROR,
attachment_content=ANY
attachment_content=ANY,
)
else:
assert False, "No exception raised"
raise AssertionError('No exception raised')

257
tests/laborious/activities/test_opc.py
View File

@@ -1,57 +1,55 @@
from unittest.mock import patch, MagicMock, ANY, call, AsyncMock
from pandas import DataFrame
from pytest import fixture, mark
from unittest.mock import ANY, AsyncMock, MagicMock, call, patch
import pytest_asyncio
from pandas import DataFrame
from pytest import mark
from sientia_do.notifications.models import NotificationLevel
from laborious.activities.opc import OPC
metadata = {
"metadata": {
"model_id": "test_model",
"model_name": "test_model",
"workflow_name": "test_workflow",
"schema_name": "test_schedule",
'metadata': {
'model_id': 'test_model',
'model_name': 'test_model',
'workflow_name': 'test_workflow',
'schema_name': 'test_schedule',
},
}
def test__init__():
servers = {
'server1': 'config'
}
opc = OPC(
opc_servers=servers,
logger=MagicMock(),
notification_handler=MagicMock()
)
servers = {'server1': 'config'}
opc = OPC(opc_servers=servers, logger=MagicMock(), notification_handler=MagicMock())
assert opc.opc_servers == servers
assert opc.opc_repository == {}
@mark.asyncio
@patch("laborious.activities.opc.OpcRepository")
@patch("laborious.activities.opc.OPC.send_notification")
@patch('laborious.activities.opc.OpcRepository')
@patch('laborious.activities.opc.OPC.send_notification')
async def test_init_opc(mock_send_notification, mock_opc_repository):
mock_logger = MagicMock()
server1 = MagicMock(
connect=AsyncMock(return_value=(True, {})),
write_data=AsyncMock(return_value=(True, {}))
connect=AsyncMock(return_value=(True, {})), write_data=AsyncMock(return_value=(True, {}))
)
server2 = MagicMock(
connect=AsyncMock(return_value=(True, {})),
write_data=AsyncMock(return_value=(True, {}))
connect=AsyncMock(return_value=(True, {})), write_data=AsyncMock(return_value=(True, {}))
)
server3 = MagicMock(
connect=AsyncMock(return_value=(False, {
connect=AsyncMock(
return_value=(
False,
{
'notification_id': 'OPC_CONNECTION_ERROR_server3',
'message': 'Failed to connect to OPC server: Test error',
'block': 'opc_repository',
'level': NotificationLevel.ERROR,
'attachment_content': 'Test error'
})),
write_data=AsyncMock(return_value=(True, {}))
'attachment_content': 'Test error',
},
)
),
write_data=AsyncMock(return_value=(True, {})),
)
mock_opc_repository.side_effect = [server1, server2, server3]
mock_notification_handler = MagicMock()
@@ -82,12 +80,10 @@ async def test_init_opc(mock_send_notification, mock_opc_repository):
'private_key_path': '',
'server_cert_path': '',
'reconnection_interval': 60,
}
},
}
opc = OPC(
opc_servers=servers,
logger=mock_logger,
notification_handler=mock_notification_handler
opc_servers=servers, logger=mock_logger, notification_handler=mock_notification_handler
)
await opc.init_opc()
@@ -97,57 +93,63 @@ async def test_init_opc(mock_send_notification, mock_opc_repository):
assert opc.opc_repository['server1'] == server1
assert opc.opc_repository['server2'] == server2
mock_opc_repository.assert_has_calls([
mock_opc_repository.assert_has_calls(
[
call(
id="server1",
url="http://localhost:8080",
opc_id='server1',
url='http://localhost:8080',
logger=mock_logger,
server_uri="opc.tcp://localhost:4840",
cert_path="",
private_key_path="",
server_cert_path="",
server_uri='opc.tcp://localhost:4840',
cert_path='',
private_key_path='',
server_cert_path='',
notification_handler=mock_notification_handler,
reconnection_interval=60,
pod_id='localhost'
pod_id='localhost',
),
])
mock_opc_repository.assert_has_calls([
]
)
mock_opc_repository.assert_has_calls(
[
call(
id="server2",
url="http://localhost:8080",
opc_id='server2',
url='http://localhost:8080',
logger=mock_logger,
server_uri="opc.tcp://localhost:4840",
cert_path="",
private_key_path="",
server_cert_path="",
server_uri='opc.tcp://localhost:4840',
cert_path='',
private_key_path='',
server_cert_path='',
notification_handler=mock_notification_handler,
reconnection_interval=60,
pod_id='localhost'
pod_id='localhost',
)
]
)
])
server1.connect.assert_called_once()
server2.connect.assert_called_once()
mock_send_notification.assert_has_calls([
mock_send_notification.assert_has_calls(
[
call(
metadata={
'model_id': '-',
'model_name': '-',
'workflow_name': '-',
'schedule_name': 'INITIALIZATION'
'schedule_name': 'INITIALIZATION',
},
notification_id="OPC_CONNECTION_ERROR_server3",
message="Failed to connect to OPC server: Test error",
block="opc_repository",
notification_id='OPC_CONNECTION_ERROR_server3',
message='Failed to connect to OPC server: Test error',
block='opc_repository',
level=NotificationLevel.ERROR,
attachment_content=ANY
attachment_content=ANY,
)
]
)
])
@pytest_asyncio.fixture
@patch("laborious.activities.opc.OpcRepository")
@patch('laborious.activities.opc.OpcRepository')
async def opc(mock_opc_repository):
servers = {
'server1': {
@@ -161,17 +163,9 @@ async def opc(mock_opc_repository):
}
}
mock_opc_repository.return_value.write_data = AsyncMock(
return_value=(True, {})
)
mock_opc_repository.return_value.connect = AsyncMock(
return_value=(True, {})
)
opc = OPC(
opc_servers=servers,
logger=MagicMock(),
notification_handler=MagicMock()
)
mock_opc_repository.return_value.write_data = AsyncMock(return_value=(True, {}))
mock_opc_repository.return_value.connect = AsyncMock(return_value=(True, {}))
opc = OPC(opc_servers=servers, logger=MagicMock(), notification_handler=MagicMock())
await opc.init_opc()
opc.send_notification = MagicMock()
return opc
@@ -188,58 +182,79 @@ WRITE_DATA_CASES = [
@mark.parametrize('tag,data_type,data', WRITE_DATA_CASES)
@mark.asyncio
async def test_write_data_success(opc, tag, data_type, data):
result = await opc.write_data(server_id='server1', tag=tag, data=data,
data_type=data_type, tag_type='prediction', metadata=metadata)
result = await opc.write_data(
server_id='server1',
tag=tag,
data=data,
data_type=data_type,
tag_type='prediction',
metadata=metadata,
)
assert result is True
opc.opc_repository['server1'].write_data.assert_called_once_with(
tag, data, data_type, opc.logger, metadata)
tag, data, data_type, opc.logger, metadata
)
@mark.asyncio
async def test_write_data_failed(opc):
opc.opc_repository['server1'].write_data.return_value = (False, {
opc.opc_repository['server1'].write_data.return_value = (
False,
{
'notification_id': 'OPC_WRITE_DATA_ERROR_server1',
'message': 'Failed to write data to OPC server: Test error',
'block': 'opc_repository',
'level': NotificationLevel.ERROR,
'attachment_content': 'Test error'
})
'attachment_content': 'Test error',
},
)
result = await opc.write_data(server_id='server1', tag='tag1', data=50,
data_type='int', tag_type='prediction', metadata=metadata)
result = await opc.write_data(
server_id='server1',
tag='tag1',
data=50,
data_type='int',
tag_type='prediction',
metadata=metadata,
)
assert result is False
opc.send_notification.assert_called_once_with(
metadata=metadata,
notification_id="OPC_WRITE_DATA_ERROR_server1",
message="Failed to write data to OPC server: Test error",
block="opc_repository",
notification_id='OPC_WRITE_DATA_ERROR_server1',
message='Failed to write data to OPC server: Test error',
block='opc_repository',
level=NotificationLevel.ERROR,
attachment_content=ANY
attachment_content=ANY,
)
@mark.asyncio
async def test_write_data_exception(opc):
opc.opc_repository['server1'].write_data.side_effect = Exception(
"Test error")
opc.opc_repository['server1'].write_data.side_effect = Exception('Test error')
try:
await opc.write_data(server_id='server1', tag='tag1', data=50,
data_type='int', tag_type='prediction', metadata=metadata)
await opc.write_data(
server_id='server1',
tag='tag1',
data=50,
data_type='int',
tag_type='prediction',
metadata=metadata,
)
except Exception:
opc.send_notification.assert_called_once_with(
metadata=metadata,
notification_id="WRITE_OPC_PREDICTION_ERROR",
message="Error writing data to OPC server: Test error",
block="write_opc_data",
notification_id='WRITE_OPC_PREDICTION_ERROR',
message='Error writing data to OPC server: Test error',
block='write_opc_data',
level=NotificationLevel.ERROR,
attachment_content=ANY
attachment_content=ANY,
)
else:
assert False, "Expected an exception to be raised"
raise AssertionError('Expected an exception to be raised')
@mark.asyncio
@@ -247,20 +262,13 @@ async def test_write_opc_data_success(opc):
# Arrange
input_data = {
**metadata,
'data': {
'prediction': [0.75],
'prediction_confidence': [0.95]
},
'data': {'prediction': [0.75], 'prediction_confidence': [0.95]},
'opc_output_config': {
'server1': {
'prediction_tags': {
'tag1': {'data_type': 'float'}
'prediction_tags': {'tag1': {'data_type': 'float'}},
'confidence_tags': {'tag2': {'data_type': 'float'}},
}
},
'confidence_tags': {
'tag2': {'data_type': 'float'}
}
}
}
}
# Act
@@ -270,25 +278,30 @@ async def test_write_opc_data_success(opc):
# Assert
assert output == {'data': 'data'}
opc.write_data.assert_has_calls([
opc.write_data.assert_has_calls(
[
call(
server_id='server1',
tag='tag1',
data=0.75,
data_type='float',
tag_type='prediction',
metadata=metadata['metadata']
)])
opc.write_data.assert_has_calls([
metadata=metadata['metadata'],
)
]
)
opc.write_data.assert_has_calls(
[
call(
server_id='server1',
tag='tag2',
data=0.95,
data_type='float',
tag_type='confidence',
metadata=metadata['metadata']
metadata=metadata['metadata'],
)
]
)
])
assert opc.write_data.call_count == 2
@@ -297,17 +310,9 @@ async def test_write_opc_data_empty_config(opc):
# Arrange
input_data = {
**metadata,
'data': {
'prediction': [0.75],
'prediction_confidence': [0.95]
},
'data': {'prediction': [0.75], 'prediction_confidence': [0.95]},
'opc_servers': ['server1'],
'opc_output_config': {
'server1': {
'prediction_tags': {},
'confidence_tags': {}
}
}
'opc_output_config': {'server1': {'prediction_tags': {}, 'confidence_tags': {}}},
}
# Act
@@ -322,20 +327,13 @@ async def test_write_opc_data_no_validate_server(opc):
opc.validate_server = MagicMock(return_value=False)
input_data = {
**metadata,
'data': {
'prediction': [0.75],
'prediction_confidence': [0.95]
},
'data': {'prediction': [0.75], 'prediction_confidence': [0.95]},
'opc_output_config': {
'server1': {
'prediction_tags': {
'tag1': {'data_type': 'float'}
'prediction_tags': {'tag1': {'data_type': 'float'}},
'confidence_tags': {'tag2': {'data_type': 'float'}},
}
},
'confidence_tags': {
'tag2': {'data_type': 'float'}
}
}
}
}
# Act
@@ -345,10 +343,13 @@ async def test_write_opc_data_no_validate_server(opc):
opc.opc_repository['server1'].write_data.assert_not_called()
@mark.parametrize('data,success,expected', [
@mark.parametrize(
'data,success,expected',
[
(DataFrame({'prediction_confidence': [0]}), True, 0),
(DataFrame({'prediction_confidence': [0]}), False, 12),
])
],
)
def test_process_confidence(opc, data, success, expected):
# Act
result = opc.process_confidence(data, success, metadata)

203
tests/laborious/activities/test_storage.py Normal file
View File

@@ -0,0 +1,203 @@
import datetime
from unittest.mock import ANY, AsyncMock, MagicMock, patch
from pytest import fixture, mark
from sientia_do.notifications.models import NotificationLevel
from sientia_do.temporal.activities.postgres import Postgres
from laborious.activities.storage import Storage
metadata = {
'metadata': {
'model_id': 'test_model_id',
'model_name': 'test_model',
'workflow_name': 'test_workflow',
'schedule_name': 'test_schedule',
}
}
@fixture
@patch('laborious.activities.storage.MinioRepository')
def storage(mock_minio_repository):
return Storage(
host='localhost',
port=5432,
user='postgres',
password='postgres',
dbname='postgres',
min_connections=1,
max_connections=10,
minio_config={
'endpoint_url': 'localhost:9000',
'access_key': 'minio',
'secret_key': 'minio123',
'region_name': 'us-east-1',
'default_bucket': 'test',
},
logger=MagicMock(),
notification_handler=MagicMock(),
)
@patch('laborious.activities.storage.MinioRepository')
def test___init___not_hasattr(mock_minio_repository):
logger = MagicMock()
notification_handler = MagicMock()
storage = Storage(
host='localhost',
port=5432,
user='postgres',
password='postgres',
dbname='postgres',
min_connections=1,
max_connections=10,
minio_config={
'endpoint_url': 'localhost:9000',
'access_key': 'minio',
'secret_key': 'minio123',
'region_name': 'us-east-1',
'default_bucket': 'test',
},
logger=logger,
notification_handler=notification_handler,
)
assert isinstance(storage, Postgres)
mock_minio_repository.assert_called_once_with(
logger=logger,
notification_handler=notification_handler,
minio_endpoint_url='localhost:9000',
minio_access_key='minio',
minio_secret_key='minio123',
minio_region_name='us-east-1',
minio_default_bucket='test',
)
@patch('laborious.activities.storage.MinioRepository')
def test___init___none_minio_repository(mock_minio_repository, storage):
storage.minio_repository = None
logger = MagicMock()
notification_handler = MagicMock()
storage.__init__(
host='localhost',
port=5432,
user='postgres',
password='postgres',
dbname='postgres',
min_connections=1,
max_connections=10,
minio_config={
'endpoint_url': 'localhost:9000',
'access_key': 'minio',
'secret_key': 'minio123',
'region_name': 'us-east-1',
'default_bucket': 'test',
},
logger=logger,
notification_handler=notification_handler,
)
mock_minio_repository.assert_called_once_with(
logger=logger,
notification_handler=notification_handler,
minio_endpoint_url='localhost:9000',
minio_access_key='minio',
minio_secret_key='minio123',
minio_region_name='us-east-1',
minio_default_bucket='test',
)
@patch('laborious.activities.storage.MinioRepository')
def test___init___done_repository(mock_minio_repository, storage):
storage.__init__(
host='localhost',
port=5432,
user='postgres',
password='postgres',
dbname='postgres',
min_connections=1,
max_connections=10,
minio_config={
'endpoint_url': 'localhost:9000',
'access_key': 'minio',
'secret_key': 'minio123',
'region_name': 'us-east-1',
'default_bucket': 'test',
},
logger=MagicMock(),
notification_handler=MagicMock(),
)
mock_minio_repository.assert_not_called()
assert storage.minio_repository is not None
@mark.asyncio
async def test_query_to_minio_not_data(storage):
storage.load_custom_query = AsyncMock(return_value=None)
result = await storage.query_to_minio({})
storage.load_custom_query.assert_called_once_with({})
assert result['success'] is False
assert result['message'] == 'No data returned from query'
@mark.asyncio
@patch('laborious.activities.storage.pd.DataFrame')
@patch('laborious.activities.storage.now')
async def test_query_to_minio_success(now, dataframe, storage):
data = [{'a': 1}, {'a': 2}, {'a': 3}]
storage.load_custom_query = AsyncMock(return_value=data)
now.return_value = datetime.datetime(2024, 1, 1, 0, 0, 0)
storage.minio_repository.minio_bucket = 'test'
result = await storage.query_to_minio({'object_prefix': 'test', **metadata})
dataframe.assert_called_once_with(data)
storage.minio_repository.store_dataframe_as_parquet.assert_called_once_with(
dataframe=dataframe.return_value,
uri='s3://test/test_2024-01-01_00-00-00.parquet',
object_name='test_2024-01-01_00-00-00.parquet',
metadata=metadata['metadata'],
)
assert result['success'] is True
assert result['object_key'] == 'test_2024-01-01_00-00-00.parquet'
assert result['uri'] == 's3://test/test_2024-01-01_00-00-00.parquet'
@mark.asyncio
async def test_query_to_minio_error(storage):
storage.send_notification = MagicMock()
storage.load_custom_query = AsyncMock(side_effect=Exception('test'))
result = await storage.query_to_minio({**metadata, 'object_prefix': 'test'})
assert result['success'] is False
assert result['message'] == 'test'
storage.send_notification.assert_called_once_with(
metadata=metadata['metadata'],
notification_id='ERROR_STORING_QUERY_TO_MINIO',
message='Error storing query to MinIO: test',
block='query_to_minio',
level=NotificationLevel.ERROR,
attachment_content=ANY,
)
def test_close(storage):
storage.minio_repository = MagicMock()
storage.close()
assert storage.minio_repository is None
def test___del__(storage):
storage.close = MagicMock()
storage.__del__()
storage.close.assert_called_once()

31
tests/laborious/utils/filters/test_conditional_filters.py
View File

@@ -1,23 +1,29 @@
from pandas import DataFrame
from laborious.utils.filters.conditional_filters import (
filter_empty_data,
filter_specific_variables_null_values,
filter_empty_data
)
def test_filter_specific_variables_null_values():
assert filter_specific_variables_null_values(
DataFrame(
{'variable': ['variable1', 'variable2'], 'value': [1, 2]}),
config={'variables': ['variable2']}) is False
assert (
filter_specific_variables_null_values(
DataFrame({'variable': ['variable1', 'variable2'], 'value': [1, 2]}),
config={'variables': ['variable2']},
)
is False
)
def test_filter_specific_variables_null_values_with_null_values():
assert filter_specific_variables_null_values(
DataFrame(
{'variable': ['variable1', 'variable2'], 'value': [1, None]}),
config={'variables': ['variable2']}) is True
assert (
filter_specific_variables_null_values(
DataFrame({'variable': ['variable1', 'variable2'], 'value': [1, None]}),
config={'variables': ['variable2']},
)
is True
)
def test_filter_empty_data():
@@ -25,6 +31,7 @@ def test_filter_empty_data():
def test_filter_empty_data_with_data():
assert filter_empty_data(
DataFrame({'variable': ['variable1', 'variable2'], 'value': [1, 2]}),
{}) is False
assert (
filter_empty_data(DataFrame({'variable': ['variable1', 'variable2'], 'value': [1, 2]}), {})
is False
)

11
tests/laborious/utils/filters/test_mlflow_filters.py
View File

@@ -1,22 +1,23 @@
from pandas import DataFrame
from laborious.utils.filters.mlflow_filters import api_error_filter, nan_values_filter
def test_api_error_filter_invalid_response():
assert api_error_filter(None, {}) == True # NOSONAR
assert api_error_filter(None, {}) is True # NOSONAR
def test_api_error_filter_valid_response_fail():
assert api_error_filter({'success': False}, {}) == True
assert api_error_filter({'success': False}, {}) is True
def test_api_error_filter_valid_response_success():
assert api_error_filter({'success': True}, {}) == False
assert api_error_filter({'success': True}, {}) is False
def test_nan_values_filter_all_nan_values():
assert nan_values_filter(DataFrame({'variable': [None, None]}), {}) == True
assert nan_values_filter(DataFrame({'variable': [None, None]}), {}) is True
def test_nan_values_filter_no_nan_values():
assert nan_values_filter(DataFrame({'variable': [1, 2]}), {}) == False
assert nan_values_filter(DataFrame({'variable': [1, 2]}), {}) is False

134
tests/laborious/utils/repository/test_minio_repository.py Normal file
View File

@@ -0,0 +1,134 @@
from unittest.mock import MagicMock, patch
from botocore.utils import ClientError
from pytest import fixture, raises
from laborious.utils.repository.minio_repository import MinioRepository
@patch('laborious.utils.repository.minio_repository.boto3')
@patch('laborious.utils.repository.minio_repository.Config')
def test___init___(mock_config, mock_boto3):
minio_repository = MinioRepository(
minio_endpoint_url='localhost:9000',
minio_access_key='minio',
minio_secret_key='minio123',
minio_region_name='us-east-1',
minio_default_bucket='test',
logger=MagicMock(),
notification_handler=MagicMock(),
)
assert minio_repository.storage_options == {
'key': 'minio',
'secret': 'minio123',
'client_kwargs': {'endpoint_url': 'localhost:9000'},
}
assert minio_repository.minio_bucket == 'test'
assert minio_repository.minio_endpoint_url == 'localhost:9000'
assert minio_repository.minio_region_name == 'us-east-1'
mock_config.assert_called_once_with(
signature_version='s3v4',
s3={'addressing_style': 'path'},
retries={'max_attempts': 5, 'mode': 'standard'},
connect_timeout=5,
read_timeout=120,
)
mock_boto3.client.assert_called_once_with(
's3',
endpoint_url='localhost:9000',
aws_access_key_id='minio',
aws_secret_access_key='minio123',
region_name='us-east-1',
config=mock_config.return_value,
)
@fixture
@patch('laborious.utils.repository.minio_repository.Config')
@patch('laborious.utils.repository.minio_repository.boto3')
def minio_repository(mock_boto3, mock_config):
return MinioRepository(
minio_endpoint_url='localhost:9000',
minio_access_key='minio',
minio_secret_key='minio123',
minio_region_name='us-east-1',
minio_default_bucket='test',
logger=MagicMock(),
notification_handler=MagicMock(),
)
def test_ensure_bucket_exists_bucket_exists(minio_repository):
assert minio_repository.ensure_bucket_exists({}) is True
minio_repository.s3_client.head_bucket.assert_called_once_with(Bucket='test')
def test_ensure_bucket_exists_bucket_not_exists_create_success(minio_repository):
minio_repository.s3_client.head_bucket.side_effect = ClientError(
error_response={'Error': {'Code': '404'}}, operation_name='head_bucket'
)
assert minio_repository.ensure_bucket_exists({}) is True
minio_repository.s3_client.head_bucket.assert_called_once_with(Bucket='test')
minio_repository.s3_client.create_bucket.assert_called_once_with(Bucket='test')
def test_ensure_bucket_exists_bucket_not_exists_create_error(minio_repository):
minio_repository.send_notification = MagicMock()
minio_repository.s3_client.head_bucket.side_effect = ClientError(
error_response={'Error': {'Code': '404'}}, operation_name='head_bucket'
)
minio_repository.s3_client.create_bucket.side_effect = ClientError(
error_response={'Error': {'Code': '404'}}, operation_name='create_bucket'
)
with raises(ClientError):
minio_repository.ensure_bucket_exists({})
minio_repository.s3_client.head_bucket.assert_called_once_with(Bucket='test')
minio_repository.s3_client.create_bucket.assert_called_once_with(Bucket='test')
@patch('laborious.utils.repository.minio_repository.BytesIO')
def test_store_dataframe_as_parquet(mock_bytesio, minio_repository):
input_data = MagicMock()
minio_repository.ensure_bucket_exists = MagicMock(return_value=True)
minio_repository.store_dataframe_as_parquet(
dataframe=input_data, uri='s3://test/test.parquet', object_name='test.parquet', metadata={}
)
minio_repository.ensure_bucket_exists.assert_called_once_with({})
mock_bytesio.assert_called_once()
input_data.to_parquet.assert_called_once_with(
mock_bytesio.return_value, engine='pyarrow', index=True
)
mock_bytesio.return_value.seek.assert_called_once_with(0)
minio_repository.s3_client.put_object.assert_called_once_with(
Bucket='test', Key='test.parquet', Body=mock_bytesio.return_value.getvalue.return_value
)
@patch('laborious.utils.repository.minio_repository.BytesIO')
@patch('laborious.utils.repository.minio_repository.read_parquet')
def test_get_parquet_as_dataframe(mock_read_parquet, mock_bytesio, minio_repository):
input_data = {'Body': MagicMock(read=MagicMock(return_value=b'test'))}
minio_repository.s3_client.get_object.return_value = input_data
output = minio_repository.get_parquet_as_dataframe(object_key='test.parquet', metadata={})
minio_repository.s3_client.get_object.assert_called_once_with(Bucket='test', Key='test.parquet')
mock_bytesio.assert_called_once_with(input_data['Body'].read.return_value)
mock_read_parquet.assert_called_once_with(mock_bytesio.return_value)
assert output == mock_read_parquet.return_value

1233
tests/laborious/utils/repository/test_model_repository.py
View File

File diff suppressed because it is too large Load Diff

203
tests/laborious/utils/repository/test_opc_repository.py
View File

@@ -1,9 +1,11 @@
import pytest
from unittest.mock import AsyncMock, Mock, patch, MagicMock, ANY, call
from asyncua.crypto.security_policies import SecurityPolicyBasic256
from laborious.utils.repository.opc_repository import OpcRepository
from sientia_do.notifications.models import NotificationLevel
from datetime import datetime
from unittest.mock import ANY, AsyncMock, MagicMock, Mock, call, patch
import pytest
from asyncua.crypto.security_policies import SecurityPolicyBasic256
from sientia_do.notifications.models import NotificationLevel
from laborious.utils.repository.opc_repository import OpcRepository
@pytest.fixture
@@ -14,15 +16,15 @@ def mock_logger():
@pytest.fixture
def opc_repository(mock_logger):
return OpcRepository(
id="test_repo",
url="opc.tcp://localhost:4840",
opc_id='test_repo',
url='opc.tcp://localhost:4840',
logger=mock_logger,
notification_handler=Mock(),
reconnection_interval=60,
server_uri="urn:test:server",
cert_path="/path/to/cert.pem",
private_key_path="/path/to/key.pem",
server_cert_path="/path/to/server_cert.pem"
server_uri='urn:test:server',
cert_path='/path/to/cert.pem',
private_key_path='/path/to/key.pem',
server_cert_path='/path/to/server_cert.pem',
)
@@ -35,22 +37,22 @@ def mock_client():
metadata = {
"metadata": {
"model_id": "test_model",
"model_name": "test_model",
"workflow_name": "test_workflow",
"schema_name": "test_schedule",
'metadata': {
'model_id': 'test_model',
'model_name': 'test_model',
'workflow_name': 'test_workflow',
'schema_name': 'test_schedule',
},
}
def test_init(opc_repository):
assert opc_repository.id == "test_repo"
assert opc_repository.url == "opc.tcp://localhost:4840"
assert opc_repository.server_uri == "urn:test:server"
assert opc_repository.cert_path == "/path/to/cert.pem"
assert opc_repository.private_key_path == "/path/to/key.pem"
assert opc_repository.server_cert_path == "/path/to/server_cert.pem"
assert opc_repository.id == 'test_repo'
assert opc_repository.url == 'opc.tcp://localhost:4840'
assert opc_repository.server_uri == 'urn:test:server'
assert opc_repository.cert_path == '/path/to/cert.pem'
assert opc_repository.private_key_path == '/path/to/key.pem'
assert opc_repository.server_cert_path == '/path/to/server_cert.pem'
assert opc_repository.reconnection_interval == 60
assert opc_repository.client is None
assert opc_repository.last_reconnection_time is None
@@ -62,12 +64,12 @@ async def test_set_security(opc_repository, mock_client):
opc_repository.client = mock_client
await opc_repository.set_security()
mock_client.application_uri = "urn:test:server"
mock_client.application_uri = 'urn:test:server'
mock_client.set_security.assert_called_once_with(
SecurityPolicyBasic256,
certificate="/path/to/cert.pem",
private_key="/path/to/key.pem",
server_certificate="/path/to/server_cert.pem"
certificate='/path/to/cert.pem',
private_key='/path/to/key.pem',
server_certificate='/path/to/server_cert.pem',
)
assert mock_client.secure_channel_timeout == 10000000
assert mock_client.session_timeout == 10000000
@@ -81,8 +83,7 @@ async def test_set_security_missing_certificates(opc_repository):
try:
await opc_repository.set_security()
except ValueError as e:
assert str(
e) == "Certificate and private key paths must be provided for secure connection."
assert str(e) == 'Certificate and private key paths must be provided for secure connection.'
@pytest.mark.asyncio
@@ -123,15 +124,15 @@ async def test_try_connect_success(opc_repository):
async def test_try_connect_fail(opc_repository):
opc_repository.last_reconnection_time = None
opc_repository.client = MagicMock()
opc_repository.client.connect.side_effect = Exception("Test error")
opc_repository.client.connect.side_effect = Exception('Test error')
is_connected, error_data = await opc_repository.try_connect()
opc_repository.client.connect.assert_called_once()
assert is_connected is False
assert error_data['notification_id'] == f"OPC_CONNECTION_ERROR_{opc_repository.id}"
assert error_data['message'] == "Failed to connect to OPC server: Test error"
assert error_data['block'] == "opc_repository"
assert error_data['notification_id'] == f'OPC_CONNECTION_ERROR_{opc_repository.id}'
assert error_data['message'] == 'Failed to connect to OPC server: Test error'
assert error_data['block'] == 'opc_repository'
assert error_data['level'] == NotificationLevel.ERROR
assert error_data['attachment_content'] is not None
@@ -154,12 +155,11 @@ async def test_disconnect_no_client(opc_repository):
@pytest.mark.asyncio
async def test_disconnect_error(opc_repository, mock_client):
opc_repository.client = mock_client
mock_client.disconnect.side_effect = Exception("Test error")
mock_client.disconnect.side_effect = Exception('Test error')
await opc_repository.disconnect()
opc_repository.logger.custom_error.assert_called_once_with(
"Failed to disconnect from OPC server: Test error",
ANY
'Failed to disconnect from OPC server: Test error', ANY
)
assert opc_repository.client is None
@@ -177,9 +177,7 @@ async def test_validate_connection_none_client(opc_repository):
async def test_validate_connection_error_count_disconnect_error(opc_repository):
opc_repository.error_count = 6
opc_repository.client = AsyncMock()
opc_repository.disconnect = AsyncMock(
side_effect=Exception("Test error")
)
opc_repository.disconnect = AsyncMock(side_effect=Exception('Test error'))
opc_repository.connect = AsyncMock(return_value=(True, {}))
response = await opc_repository.validate_connection()
@@ -188,34 +186,34 @@ async def test_validate_connection_error_count_disconnect_error(opc_repository):
opc_repository.connect.assert_called_once()
opc_repository.logger.custom_error.assert_has_calls(
[
call("Failed to disconnect from OPC server: Test error", ANY),
call('Failed to disconnect from OPC server: Test error', ANY),
]
)
@pytest.mark.asyncio
async def test_validate_connection_error_validate_connection_error(opc_repository):
opc_repository.client = MagicMock(
uaclient=Exception("Test error")
)
opc_repository.client = MagicMock(uaclient=Exception('Test error'))
opc_repository.error_count = 0
response = await opc_repository.validate_connection()
assert response == (False, {
"notification_id": f"OPC_CONNECTION_CHECK_ERROR_{opc_repository.id}",
"message": "Failed to validate connection to OPC server: 'Exception' object has no attribute 'protocol'",
"block": "opc_repository",
"level": NotificationLevel.ERROR,
"attachment_content": ANY
})
assert response == (
False,
{
'notification_id': f'OPC_CONNECTION_CHECK_ERROR_{opc_repository.id}',
'message': "Failed to validate connection to OPC server: 'Exception' object has no attribute 'protocol'",
'block': 'opc_repository',
'level': NotificationLevel.ERROR,
'attachment_content': ANY,
},
)
@pytest.mark.asyncio
@patch('laborious.utils.repository.opc_repository.datetime')
async def test_validate_connection_lost_not_time_to_reconnect(_mock_datetime, opc_repository):
_mock_datetime.now = MagicMock(
return_value=datetime(2025, 1, 1, 0, 0, 0))
_mock_datetime.now = MagicMock(return_value=datetime(2025, 1, 1, 0, 0, 0))
opc_repository.error_count = 0
opc_repository.client = MagicMock()
opc_repository.client.uaclient.protocol = None
@@ -224,19 +222,21 @@ async def test_validate_connection_lost_not_time_to_reconnect(_mock_datetime, op
response = await opc_repository.validate_connection()
opc_repository.connect.assert_not_called()
assert response == (False, {
"notification_id": f"OPC_CONNECTION_AWAITING_RECONNECTION_WINDOW_{opc_repository.id}",
"message": f"OPC server {opc_repository.id} is not connected, waiting for next reconnection window...",
"block": "opc_repository",
"level": NotificationLevel.WARNING
})
assert response == (
False,
{
'notification_id': f'OPC_CONNECTION_AWAITING_RECONNECTION_WINDOW_{opc_repository.id}',
'message': f'OPC server {opc_repository.id} is not connected, waiting for next reconnection window...',
'block': 'opc_repository',
'level': NotificationLevel.WARNING,
},
)
@pytest.mark.asyncio
@patch('laborious.utils.repository.opc_repository.datetime')
async def test_validate_connection_lost_time_to_reconnect(mock_datetime, opc_repository):
mock_datetime.now = MagicMock(
return_value=datetime(2025, 1, 1, 1, 0, 0))
mock_datetime.now = MagicMock(return_value=datetime(2025, 1, 1, 1, 0, 0))
opc_repository.error_count = 0
opc_repository.client = AsyncMock()
opc_repository.client.uaclient.protocol = None
@@ -253,7 +253,7 @@ async def test_validate_connection_success(opc_repository):
opc_repository.client = MagicMock()
opc_repository.error_count = 0
opc_repository.client.uaclient.protocol = MagicMock()
opc_repository.client.uaclient.protocol.state = "open"
opc_repository.client.uaclient.protocol.state = 'open'
output = await opc_repository.validate_connection()
assert output == (True, {})
@@ -262,17 +262,16 @@ async def test_validate_connection_success(opc_repository):
@pytest.mark.asyncio
async def test_write_data_validate_connection_do_nothing(opc_repository):
opc_repository.validate_connection = AsyncMock(return_value=(True, {}))
opc_repository.client = AsyncMock(
get_node=MagicMock()
)
opc_repository.client = AsyncMock(get_node=MagicMock())
mock_node = AsyncMock()
opc_repository.client.get_node.return_value = mock_node
result = await opc_repository.write_data("ns=2;s=TestNode", 42.0,
"float", opc_repository.logger, metadata['metadata'])
result = await opc_repository.write_data(
'ns=2;s=TestNode', 42.0, 'float', opc_repository.logger, metadata['metadata']
)
opc_repository.validate_connection.assert_called_once()
opc_repository.client.get_node.assert_called_once_with("ns=2;s=TestNode")
opc_repository.client.get_node.assert_called_once_with('ns=2;s=TestNode')
assert result == (True, {})
@@ -282,8 +281,9 @@ async def test_write_data_validate_connection_failed(opc_repository):
opc_repository.client = AsyncMock()
opc_repository.error_count = 0
result = await opc_repository.write_data("ns=2;s=TestNode", 42.0,
"float", opc_repository.logger, metadata['metadata'])
result = await opc_repository.write_data(
'ns=2;s=TestNode', 42.0, 'float', opc_repository.logger, metadata['metadata']
)
opc_repository.validate_connection.assert_called_once()
opc_repository.client.get_node.assert_not_called()
@@ -295,18 +295,21 @@ async def test_write_data_get_node_failed(opc_repository):
opc_repository.validate_connection = AsyncMock(return_value=(True, {}))
opc_repository.client = AsyncMock()
opc_repository.error_count = 0
opc_repository.client.get_node = MagicMock(
side_effect=Exception("Test error"))
opc_repository.client.get_node = MagicMock(side_effect=Exception('Test error'))
is_success, error_data = await opc_repository.write_data("ns=2;s=TestNode", 42.0,
"float", opc_repository.logger, metadata['metadata'])
is_success, error_data = await opc_repository.write_data(
'ns=2;s=TestNode', 42.0, 'float', opc_repository.logger, metadata['metadata']
)
opc_repository.validate_connection.assert_called_once()
opc_repository.client.get_node.assert_called_once_with("ns=2;s=TestNode")
opc_repository.client.get_node.assert_called_once_with('ns=2;s=TestNode')
assert is_success is False
assert error_data['notification_id'] == f"OPC_WRITE_GET_NODE_ERROR_{opc_repository.id}"
assert error_data['message'] == "Failed to get node from OPC server: Test error | metadata: {'model_id': 'test_model', 'model_name': 'test_model', 'workflow_name': 'test_workflow', 'schema_name': 'test_schedule'}"
assert error_data['block'] == "opc_repository"
assert error_data['notification_id'] == f'OPC_WRITE_GET_NODE_ERROR_{opc_repository.id}'
assert (
error_data['message']
== "Failed to get node from OPC server: Test error | metadata: {'model_id': 'test_model', 'model_name': 'test_model', 'workflow_name': 'test_workflow', 'schema_name': 'test_schedule'}"
)
assert error_data['block'] == 'opc_repository'
assert error_data['level'] == NotificationLevel.ERROR
assert error_data['attachment_content'] is not None
@@ -318,16 +321,20 @@ async def test_write_data_invalid_data_type(opc_repository, mock_client):
mock_node = AsyncMock()
mock_client.get_node = MagicMock(return_value=mock_node)
is_success, error_data = await opc_repository.write_data("ns=2;s=TestNode", 42.0,
"invalid_type", opc_repository.logger, metadata['metadata'])
is_success, error_data = await opc_repository.write_data(
'ns=2;s=TestNode', 42.0, 'invalid_type', opc_repository.logger, metadata['metadata']
)
opc_repository.validate_connection.assert_called_once()
mock_client.get_node.assert_called_once_with("ns=2;s=TestNode")
mock_client.get_node.assert_called_once_with('ns=2;s=TestNode')
assert is_success is False
assert error_data['notification_id'] == f"OPC_WRITE_DATA_TYPE_ERROR_{opc_repository.id}"
assert error_data['message'] == "Unsupported data type: invalid_type | metadata: {'model_id': 'test_model', 'model_name': 'test_model', 'workflow_name': 'test_workflow', 'schema_name': 'test_schedule'}"
assert error_data['block'] == "opc_repository"
assert error_data['notification_id'] == f'OPC_WRITE_DATA_TYPE_ERROR_{opc_repository.id}'
assert (
error_data['message']
== "Unsupported data type: invalid_type | metadata: {'model_id': 'test_model', 'model_name': 'test_model', 'workflow_name': 'test_workflow', 'schema_name': 'test_schedule'}"
)
assert error_data['block'] == 'opc_repository'
assert error_data['level'] == NotificationLevel.ERROR
assert error_data.get('attachment_content') is None
@@ -340,10 +347,11 @@ async def test_write_data(mock_metrics, opc_repository, mock_client):
mock_node = AsyncMock()
mock_client.get_node = MagicMock(return_value=mock_node)
result = await opc_repository.write_data("ns=2;s=TestNode", 42.0,
"float", opc_repository.logger, metadata['metadata'])
result = await opc_repository.write_data(
'ns=2;s=TestNode', 42.0, 'float', opc_repository.logger, metadata['metadata']
)
mock_client.get_node.assert_called_once_with("ns=2;s=TestNode")
mock_client.get_node.assert_called_once_with('ns=2;s=TestNode')
mock_node.write_value.assert_called_once()
assert result == (True, {})
@@ -351,7 +359,7 @@ async def test_write_data(mock_metrics, opc_repository, mock_client):
pod_id=opc_repository.pod_id,
model_name=metadata['metadata']['model_name'],
pipeline_name=metadata['metadata']['workflow_name'],
opc_server_id=opc_repository.id
opc_server_id=opc_repository.id,
)
mock_metrics.PREDICTION_OPC_WRITING_COUNT.labels.return_value.inc.assert_called_once_with()
@@ -359,10 +367,11 @@ async def test_write_data(mock_metrics, opc_repository, mock_client):
pod_id=opc_repository.pod_id,
model_name=metadata['metadata']['model_name'],
pipeline_name=metadata['metadata']['workflow_name'],
opc_server_id=opc_repository.id
opc_server_id=opc_repository.id,
)
mock_metrics.PREDICTION_OPC_WRITING_RESPONSE_TIME_MONITOR.labels.return_value.observe.assert_called_once_with(
ANY)
ANY
)
@pytest.mark.asyncio
@@ -372,17 +381,21 @@ async def test_write_data_write_value_failed(opc_repository, mock_client):
mock_node = AsyncMock()
opc_repository.error_count = 0
mock_client.get_node = MagicMock(return_value=mock_node)
mock_node.write_value.side_effect = Exception("Test error")
mock_node.write_value.side_effect = Exception('Test error')
is_success, error_data = await opc_repository.write_data("ns=2;s=TestNode", 42.0,
"float", opc_repository.logger, metadata['metadata'])
is_success, error_data = await opc_repository.write_data(
'ns=2;s=TestNode', 42.0, 'float', opc_repository.logger, metadata['metadata']
)
opc_repository.validate_connection.assert_called_once()
mock_client.get_node.assert_called_once_with("ns=2;s=TestNode")
mock_client.get_node.assert_called_once_with('ns=2;s=TestNode')
mock_node.write_value.assert_called_once()
assert is_success is False
assert error_data['notification_id'] == f"OPC_WRITE_DATA_ERROR_{opc_repository.id}"
assert error_data['message'] == "Failed to write data to OPC server: Test error | metadata: {'model_id': 'test_model', 'model_name': 'test_model', 'workflow_name': 'test_workflow', 'schema_name': 'test_schedule'}"
assert error_data['block'] == "opc_repository"
assert error_data['notification_id'] == f'OPC_WRITE_DATA_ERROR_{opc_repository.id}'
assert (
error_data['message']
== "Failed to write data to OPC server: Test error | metadata: {'model_id': 'test_model', 'model_name': 'test_model', 'workflow_name': 'test_workflow', 'schema_name': 'test_schedule'}"
)
assert error_data['block'] == 'opc_repository'
assert error_data['level'] == NotificationLevel.ERROR
assert error_data['attachment_content'] is not None

11
tests/laborious/utils/test_connectors_config.py
View File

@@ -1,8 +1,11 @@
from os import environ
from laborious.utils.connectors_config import (build_mlflow_config,
from laborious.utils.connectors_config import (
build_mlflow_config,
build_mongodb_config,
build_opc_config,
build_postgres_config,
build_mongodb_config)
)
def test_build_mlflow_config_with_env_vars():
@@ -144,7 +147,7 @@ def test_build_mongo_db_config_with_env_vars():
assert build_mongodb_config() == {
'connection_string': 'mongodb://sientia1:sientia1@localhost:27018',
'database_name': 'test_db',
'ttl_index_seconds': 3600
'ttl_index_seconds': 3600,
}
@@ -157,5 +160,5 @@ def test_build_mongo_db_config_with_defaults():
assert build_mongodb_config() == {
'connection_string': 'mongodb://root:wKZDbMNU1c@localhost:27018',
'database_name': 'sientia',
'ttl_index_seconds': 3600
'ttl_index_seconds': 3600,
}

135
tests/laborious/workflows/subworkflows/test_format_and_export_prediction.py
View File

@@ -1,9 +1,10 @@
from unittest.mock import call, patch, AsyncMock, ANY
from pytest import mark, fixture
from unittest.mock import ANY, AsyncMock, call, patch
from pytest import fixture, mark
from sientia_do.temporal.constants import DATETIME_FORMAT_WITH_TZ
from laborious.activities.activities import Activities
from laborious.workflows.sub_workflows.format_and_export_prediction import FormatAndExportPrediction
from sientia_do.temporal.constants import DATETIME_FORMAT_WITH_TZ
@fixture
@@ -12,36 +13,39 @@ def format_and_export_prediction():
metadata = {
"metadata": {
"model_id": "test_model",
"model_name": "test_model",
"workflow_name": "test_workflow",
"schema_name": "test_schedule",
'metadata': {
'model_id': 'test_model',
'model_name': 'test_model',
'workflow_name': 'test_workflow',
'schema_name': 'test_schedule',
},
}
@mark.asyncio
@patch("laborious.workflows.sub_workflows.format_and_export_prediction.workflow", new_callable=AsyncMock)
@patch(
'laborious.workflows.sub_workflows.format_and_export_prediction.workflow',
new_callable=AsyncMock,
)
async def test_run_none_path_flag(workflow_mock, format_and_export_prediction):
input_data = {
'metadata': metadata,
"path_flag": None,
"data": {"test": "data"},
"timestamp": "2021-01-01",
"model_id": 1,
"prediction_confidence": 0,
"schema": "test_schema",
"table_name": "test_table",
"opc_servers": ["test_server"],
"opc_output_config": {"test": "config"},
"prediction_store_policy": "erl:1"
'path_flag': None,
'data': {'test': 'data'},
'timestamp': '2021-01-01',
'model_id': 1,
'prediction_confidence': 0,
'schema': 'test_schema',
'table_name': 'test_table',
'opc_servers': ['test_server'],
'opc_output_config': {'test': 'config'},
'prediction_store_policy': 'erl:1',
}
await format_and_export_prediction.run(input_data)
workflow_mock.execute_local_activity_method.assert_has_calls([
workflow_mock.execute_local_activity_method.assert_has_calls(
[
call(
Activities.format_prediction,
{
@@ -50,26 +54,31 @@ async def test_run_none_path_flag(workflow_mock, format_and_export_prediction):
'model_id': input_data['model_id'],
'prediction_confidence': input_data['prediction_confidence'],
'prediction_store_policy': input_data['prediction_store_policy'],
**metadata
**metadata,
},
retry_policy=ANY,
start_to_close_timeout=ANY
)])
start_to_close_timeout=ANY,
)
]
)
workflow_mock.execute_activity_method.assert_has_calls([
workflow_mock.execute_activity_method.assert_has_calls(
[
call(
Activities.write_opc_data,
{
'opc_output_config': input_data['opc_output_config'],
'data': workflow_mock.execute_local_activity_method.return_value,
**metadata
**metadata,
},
retry_policy=ANY,
start_to_close_timeout=ANY
start_to_close_timeout=ANY,
)
]
)
])
workflow_mock.execute_activity_method.assert_has_calls([
workflow_mock.execute_activity_method.assert_has_calls(
[
call(
Activities.export_data_to_postgres,
{
@@ -79,38 +88,43 @@ async def test_run_none_path_flag(workflow_mock, format_and_export_prediction):
**metadata,
'timestamp_conversion': {
'column': 'timestamp',
'format': DATETIME_FORMAT_WITH_TZ
}
'format': DATETIME_FORMAT_WITH_TZ,
},
},
retry_policy=ANY,
start_to_close_timeout=ANY
)])
start_to_close_timeout=ANY,
)
]
)
assert workflow_mock.execute_activity_method.call_count == 3
assert workflow_mock.execute_local_activity_method.call_count == 1
@mark.asyncio
@patch("laborious.workflows.sub_workflows.format_and_export_prediction.workflow", new_callable=AsyncMock)
@patch(
'laborious.workflows.sub_workflows.format_and_export_prediction.workflow',
new_callable=AsyncMock,
)
async def test_run_default_path_flag(workflow_mock, format_and_export_prediction):
input_data = {
'metadata': metadata,
"path_flag": "default",
"data": {"test": "data"},
"timestamp": "2021-01-01",
"model_id": 1,
"prediction_confidence": 0,
"schema": "test_schema",
"table_name": "test_table",
"opc_servers": ["test_server"],
"opc_output_config": {"test": "config"},
"comment": "test_comment"
'path_flag': 'default',
'data': {'test': 'data'},
'timestamp': '2021-01-01',
'model_id': 1,
'prediction_confidence': 0,
'schema': 'test_schema',
'table_name': 'test_table',
'opc_servers': ['test_server'],
'opc_output_config': {'test': 'config'},
'comment': 'test_comment',
}
await format_and_export_prediction.run(input_data)
workflow_mock.execute_local_activity_method.assert_has_calls([
workflow_mock.execute_local_activity_method.assert_has_calls(
[
call(
Activities.format_default_prediction,
{
@@ -118,27 +132,31 @@ async def test_run_default_path_flag(workflow_mock, format_and_export_prediction
'model_id': input_data['model_id'],
'prediction_confidence': input_data['prediction_confidence'],
'comment': input_data['comment'],
**metadata
**metadata,
},
retry_policy=ANY,
start_to_close_timeout=ANY
start_to_close_timeout=ANY,
)
]
)
])
workflow_mock.execute_activity_method.assert_has_calls([
workflow_mock.execute_activity_method.assert_has_calls(
[
call(
Activities.write_opc_data,
{
'opc_output_config': input_data['opc_output_config'],
'data': workflow_mock.execute_local_activity_method.return_value,
**metadata
**metadata,
},
retry_policy=ANY,
start_to_close_timeout=ANY
start_to_close_timeout=ANY,
)
]
)
])
workflow_mock.execute_activity_method.assert_has_calls([
workflow_mock.execute_activity_method.assert_has_calls(
[
call(
Activities.export_data_to_postgres,
{
@@ -148,13 +166,14 @@ async def test_run_default_path_flag(workflow_mock, format_and_export_prediction
**metadata,
'timestamp_conversion': {
'column': 'timestamp',
'format': DATETIME_FORMAT_WITH_TZ
}
'format': DATETIME_FORMAT_WITH_TZ,
},
},
retry_policy=ANY,
start_to_close_timeout=ANY
start_to_close_timeout=ANY,
)
]
)
])
assert workflow_mock.execute_activity_method.call_count == 3
assert workflow_mock.execute_local_activity_method.call_count == 1

516
tests/laborious/workflows/subworkflows/test_prediction_process.py
View File

@@ -1,5 +1,7 @@
from unittest.mock import AsyncMock, patch, call, ANY
from unittest.mock import ANY, AsyncMock, call, patch
from pytest import fixture, mark
from laborious.activities.activities import Activities
from laborious.workflows.sub_workflows.prediction_process import PredictionProcess
@@ -10,17 +12,17 @@ def prediction_process():
metadata = {
"metadata": {
"model_id": "test_model",
"model_name": "test_model",
"workflow_name": "test_workflow",
"schema_name": "test_schedule",
'metadata': {
'model_id': 'test_model',
'model_name': 'test_model',
'workflow_name': 'test_workflow',
'schema_name': 'test_schedule',
},
}
@mark.asyncio
@patch("laborious.workflows.sub_workflows.prediction_process.workflow", new_callable=AsyncMock)
@patch('laborious.workflows.sub_workflows.prediction_process.workflow', new_callable=AsyncMock)
async def test_run(workflow_mock, prediction_process):
prediction_process.path_flag_handler = AsyncMock(return_value=False)
# Arrange
@@ -34,26 +36,24 @@ async def test_run(workflow_mock, prediction_process):
'mlflow_transform_filters': {'test': 'filter'},
'mlflow_predict_filters': {'test': 'filter'},
'model_name': 'test_model_name',
'model_config': {
'retention': '30'
},
'model_config': {'retention': '30'},
'path_priority': ['continue', 'repeat', 'stop'],
'opc_output_config': {'test': 'config'},
'prediction_store_policy': 'lts:1'
'prediction_store_policy': 'lts:1',
}
# Mock the activity responses
workflow_mock.execute_local_activity_method.side_effect = [
'2024-01-01', # get_last_timestamp
('continue', 0.95, "Input data with bad quality"), # input_gate
('continue', 0.95, 'Input data with bad quality'), # input_gate
{'content': 'transformed_data', 'timestamp': '2024-01-01'}, # transform_data
# mlflow_response_gate (transform)
('continue', 0.95, "Error"),
('continue', 0.95, 'Error'),
# mlflow_content_gate (transform)
('continue', 0.95, "Transformed data not passed the content filter"),
('continue', 0.95, 'Transformed data not passed the content filter'),
{'content': 'predicted_data', 'timestamp': '2024-01-01'}, # request_predict
# mlflow_response_gate (predict)
('continue', 0.95, "Error"),
('continue', 0.95, 'Error'),
]
# Act
@@ -62,57 +62,112 @@ async def test_run(workflow_mock, prediction_process):
# Assert
assert workflow_mock.execute_local_activity_method.call_count == 7
workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.get_last_timestamp, {
workflow_mock.execute_local_activity_method.assert_has_calls(
[
call(
Activities.get_last_timestamp,
{
**metadata,
'data': input_data['data'],
},
retry_policy=ANY, start_to_close_timeout=ANY)])
workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.input_gate, {
retry_policy=ANY,
start_to_close_timeout=ANY,
)
]
)
workflow_mock.execute_local_activity_method.assert_has_calls(
[
call(
Activities.input_gate,
{
**metadata,
'filters': input_data['input_filters'],
'data': input_data['data'],
'path_priority': input_data['path_priority'],
}, retry_policy=ANY, start_to_close_timeout=ANY)])
workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.request_transform, {
},
retry_policy=ANY,
start_to_close_timeout=ANY,
)
]
)
workflow_mock.execute_local_activity_method.assert_has_calls(
[
call(
Activities.request_transform,
{
**metadata,
'data': input_data['data'],
'model_name': input_data['model_name'],
'model_config': input_data['model_config'],
}, retry_policy=ANY, start_to_close_timeout=ANY)])
workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.mlflow_response_gate, {
},
retry_policy=ANY,
start_to_close_timeout=ANY,
)
]
)
workflow_mock.execute_local_activity_method.assert_has_calls(
[
call(
Activities.mlflow_response_gate,
{
**metadata,
'filters': input_data['mlflow_transform_filters'],
'data': {'content': 'transformed_data', 'timestamp': '2024-01-01'},
'type': 'transform',
'path_priority': input_data['path_priority'],
}, retry_policy=ANY, start_to_close_timeout=ANY)])
workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.mlflow_content_gate, {
},
retry_policy=ANY,
start_to_close_timeout=ANY,
)
]
)
workflow_mock.execute_local_activity_method.assert_has_calls(
[
call(
Activities.mlflow_content_gate,
{
**metadata,
'filters': input_data['mlflow_transform_filters'],
'data': 'transformed_data',
'type': 'transform',
'path_priority': input_data['path_priority'],
}, retry_policy=ANY, start_to_close_timeout=ANY)])
workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.request_predict, {
},
retry_policy=ANY,
start_to_close_timeout=ANY,
)
]
)
workflow_mock.execute_local_activity_method.assert_has_calls(
[
call(
Activities.request_predict,
{
**metadata,
'data': 'transformed_data',
'model_name': input_data['model_name'],
'model_config': input_data['model_config'],
}, retry_policy=ANY, start_to_close_timeout=ANY)])
workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.mlflow_response_gate, {
},
retry_policy=ANY,
start_to_close_timeout=ANY,
)
]
)
workflow_mock.execute_local_activity_method.assert_has_calls(
[
call(
Activities.mlflow_response_gate,
{
**metadata,
'filters': input_data['mlflow_predict_filters'],
'data': {'content': 'predicted_data', 'timestamp': '2024-01-01'},
'type': 'predict',
'path_priority': input_data['path_priority'],
}, retry_policy=ANY, start_to_close_timeout=ANY)])
},
retry_policy=ANY,
start_to_close_timeout=ANY,
)
]
)
workflow_mock.execute_child_workflow.assert_called_once_with(
'format_and_export_prediction',
@@ -129,13 +184,13 @@ async def test_run(workflow_mock, prediction_process):
'schema': input_data['schema'],
'table_name': input_data['table_name'],
'comment': 'Error',
'prediction_store_policy': input_data['prediction_store_policy']
}
'prediction_store_policy': input_data['prediction_store_policy'],
},
)
@mark.asyncio
@patch("laborious.workflows.sub_workflows.prediction_process.workflow", new_callable=AsyncMock)
@patch('laborious.workflows.sub_workflows.prediction_process.workflow', new_callable=AsyncMock)
async def test_run_stop_at_input_gate(workflow_mock, prediction_process):
prediction_process.path_flag_handler = AsyncMock(return_value=True)
# Arrange
@@ -149,17 +204,15 @@ async def test_run_stop_at_input_gate(workflow_mock, prediction_process):
'mlflow_transform_filters': {'test': 'filter'},
'mlflow_predict_filters': {'test': 'filter'},
'model_name': 'test_model_name',
'model_config': {
'retention': '30'
},
'model_config': {'retention': '30'},
'path_priority': ['continue', 'repeat', 'stop'],
'opc_output_config': {'test': 'config'}
'opc_output_config': {'test': 'config'},
}
# Mock the activity responses
workflow_mock.execute_local_activity_method.side_effect = [
'2024-01-01', # get_last_timestamp
('stop', 0.95, "Input data with bad quality"), # input_gate
('stop', 0.95, 'Input data with bad quality'), # input_gate
]
# Act
@@ -167,23 +220,35 @@ async def test_run_stop_at_input_gate(workflow_mock, prediction_process):
# Assert
assert workflow_mock.execute_local_activity_method.call_count == 2
workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.get_last_timestamp, {
workflow_mock.execute_local_activity_method.assert_has_calls(
[
call(
Activities.get_last_timestamp,
{
'data': input_data['data'],
**metadata,
}, retry_policy=ANY, start_to_close_timeout=ANY),
call(Activities.input_gate, {
},
retry_policy=ANY,
start_to_close_timeout=ANY,
),
call(
Activities.input_gate,
{
'filters': input_data['input_filters'],
'data': input_data['data'],
'path_priority': input_data['path_priority'],
**metadata,
}, retry_policy=ANY, start_to_close_timeout=ANY)
])
},
retry_policy=ANY,
start_to_close_timeout=ANY,
),
]
)
workflow_mock.execute_child_workflow.assert_not_called()
@mark.asyncio
@patch("laborious.workflows.sub_workflows.prediction_process.workflow", new_callable=AsyncMock)
@patch('laborious.workflows.sub_workflows.prediction_process.workflow', new_callable=AsyncMock)
async def test_run_stop_at_first_mlflow_response_gate(workflow_mock, prediction_process):
prediction_process.path_flag_handler = AsyncMock(side_effect=[False, True])
# Arrange
@@ -197,19 +262,17 @@ async def test_run_stop_at_first_mlflow_response_gate(workflow_mock, prediction_
'mlflow_transform_filters': {'test': 'filter'},
'mlflow_predict_filters': {'test': 'filter'},
'model_name': 'test_model_name',
'model_config': {
'retention': '30'
},
'model_config': {'retention': '30'},
'path_priority': ['continue', 'repeat', 'stop'],
'opc_output_config': {'test': 'config'}
'opc_output_config': {'test': 'config'},
}
# Mock the activity responses
workflow_mock.execute_local_activity_method.side_effect = [
'2024-01-01', # get_last_timestamp
('repeat', 0.95, "Input data with bad quality"), # input_gate
('repeat', 0.95, 'Input data with bad quality'), # input_gate
{'content': 'transformed_data', 'timestamp': '2024-01-01'}, # transform_data
('continue', 0.95, "Error"), # mlflow_response_gate (transform)
('continue', 0.95, 'Error'), # mlflow_response_gate (transform)
]
# Act
@@ -217,46 +280,72 @@ async def test_run_stop_at_first_mlflow_response_gate(workflow_mock, prediction_
# Assert
assert workflow_mock.execute_local_activity_method.call_count == 4
workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.get_last_timestamp, {
workflow_mock.execute_local_activity_method.assert_has_calls(
[
call(
Activities.get_last_timestamp,
{
'data': input_data['data'],
**metadata,
},
retry_policy=ANY, start_to_close_timeout=ANY)])
workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.input_gate, {
retry_policy=ANY,
start_to_close_timeout=ANY,
)
]
)
workflow_mock.execute_local_activity_method.assert_has_calls(
[
call(
Activities.input_gate,
{
'filters': input_data['input_filters'],
'data': input_data['data'],
'path_priority': input_data['path_priority'],
**metadata,
},
retry_policy=ANY, start_to_close_timeout=ANY)])
workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.request_transform, {
retry_policy=ANY,
start_to_close_timeout=ANY,
)
]
)
workflow_mock.execute_local_activity_method.assert_has_calls(
[
call(
Activities.request_transform,
{
'data': input_data['data'],
'model_name': input_data['model_name'],
'model_config': input_data['model_config'],
**metadata
**metadata,
},
retry_policy=ANY, start_to_close_timeout=ANY)
])
workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.mlflow_response_gate, {
retry_policy=ANY,
start_to_close_timeout=ANY,
)
]
)
workflow_mock.execute_local_activity_method.assert_has_calls(
[
call(
Activities.mlflow_response_gate,
{
'filters': input_data['mlflow_transform_filters'],
'data': {'content': 'transformed_data', 'timestamp': '2024-01-01'},
'type': 'transform',
'path_priority': input_data['path_priority'],
**metadata
}, retry_policy=ANY, start_to_close_timeout=ANY)
])
**metadata,
},
retry_policy=ANY,
start_to_close_timeout=ANY,
)
]
)
workflow_mock.execute_child_workflow.assert_not_called()
@mark.asyncio
@patch("laborious.workflows.sub_workflows.prediction_process.workflow", new_callable=AsyncMock)
@patch('laborious.workflows.sub_workflows.prediction_process.workflow', new_callable=AsyncMock)
async def test_run_stop_at_mlflow_content_gate(workflow_mock, prediction_process):
prediction_process.path_flag_handler = AsyncMock(
side_effect=[False, False, True])
prediction_process.path_flag_handler = AsyncMock(side_effect=[False, False, True])
# Arrange
input_data = {
'metadata': metadata,
@@ -268,22 +357,20 @@ async def test_run_stop_at_mlflow_content_gate(workflow_mock, prediction_process
'mlflow_transform_filters': {'test': 'filter'},
'mlflow_predict_filters': {'test': 'filter'},
'model_name': 'test_model_name',
'model_config': {
'retention': '30'
},
'model_config': {'retention': '30'},
'path_priority': ['continue', 'repeat', 'stop'],
'opc_output_config': {'test': 'config'}
'opc_output_config': {'test': 'config'},
}
# Mock the activity responses
workflow_mock.execute_local_activity_method.side_effect = [
'2024-01-01', # get_last_timestamp
('continue', 0.95, "Input data with bad quality"), # input_gate
('continue', 0.95, 'Input data with bad quality'), # input_gate
{'content': 'transformed_data', 'timestamp': '2024-01-01'}, # transform_data
# mlflow_response_gate (transform)
('continue', 0.95, "Error"),
('continue', 0.95, 'Error'),
# mlflow_content_gate (transform)
('continue', 0.95, "Transformed data not passed the content filter"),
('continue', 0.95, 'Transformed data not passed the content filter'),
]
# Act
@@ -292,51 +379,88 @@ async def test_run_stop_at_mlflow_content_gate(workflow_mock, prediction_process
# Assert
assert workflow_mock.execute_local_activity_method.call_count == 5
workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.get_last_timestamp, {
workflow_mock.execute_local_activity_method.assert_has_calls(
[
call(
Activities.get_last_timestamp,
{
'data': input_data['data'],
**metadata,
},
retry_policy=ANY, start_to_close_timeout=ANY)])
workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.input_gate, {
retry_policy=ANY,
start_to_close_timeout=ANY,
)
]
)
workflow_mock.execute_local_activity_method.assert_has_calls(
[
call(
Activities.input_gate,
{
'filters': input_data['input_filters'],
'data': input_data['data'],
'path_priority': input_data['path_priority'],
**metadata,
},
retry_policy=ANY, start_to_close_timeout=ANY)])
workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.request_transform, {
retry_policy=ANY,
start_to_close_timeout=ANY,
)
]
)
workflow_mock.execute_local_activity_method.assert_has_calls(
[
call(
Activities.request_transform,
{
'data': input_data['data'],
'model_name': input_data['model_name'],
'model_config': input_data['model_config'],
**metadata
}, retry_policy=ANY, start_to_close_timeout=ANY)])
workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.mlflow_response_gate, {
**metadata,
},
retry_policy=ANY,
start_to_close_timeout=ANY,
)
]
)
workflow_mock.execute_local_activity_method.assert_has_calls(
[
call(
Activities.mlflow_response_gate,
{
'filters': input_data['mlflow_transform_filters'],
'data': {'content': 'transformed_data', 'timestamp': '2024-01-01'},
'type': 'transform',
'path_priority': input_data['path_priority'],
**metadata,
}, retry_policy=ANY, start_to_close_timeout=ANY)])
workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.mlflow_content_gate, {
},
retry_policy=ANY,
start_to_close_timeout=ANY,
)
]
)
workflow_mock.execute_local_activity_method.assert_has_calls(
[
call(
Activities.mlflow_content_gate,
{
'filters': input_data['mlflow_transform_filters'],
'data': 'transformed_data',
'type': 'transform',
'path_priority': input_data['path_priority'],
**metadata,
}, retry_policy=ANY, start_to_close_timeout=ANY)])
},
retry_policy=ANY,
start_to_close_timeout=ANY,
)
]
)
workflow_mock.execute_child_workflow.assert_not_called()
@mark.asyncio
@patch("laborious.workflows.sub_workflows.prediction_process.workflow", new_callable=AsyncMock)
@patch('laborious.workflows.sub_workflows.prediction_process.workflow', new_callable=AsyncMock)
async def test_run_stop_at_mlflow_last_response_gate(workflow_mock, prediction_process):
prediction_process.path_flag_handler = AsyncMock(
side_effect=[False, False, False, True])
prediction_process.path_flag_handler = AsyncMock(side_effect=[False, False, False, True])
# Arrange
input_data = {
'metadata': metadata,
@@ -348,24 +472,22 @@ async def test_run_stop_at_mlflow_last_response_gate(workflow_mock, prediction_p
'mlflow_transform_filters': {'test': 'filter'},
'mlflow_predict_filters': {'test': 'filter'},
'model_name': 'test_model_name',
'model_config': {
'retention': '30'
},
'model_config': {'retention': '30'},
'path_priority': ['continue', 'repeat', 'stop'],
'opc_output_config': {'test': 'config'}
'opc_output_config': {'test': 'config'},
}
# Mock the activity responses
workflow_mock.execute_local_activity_method.side_effect = [
'2024-01-01', # get_last_timestamp
('continue', 0.95, "Input data with bad quality"), # input_gate
('continue', 0.95, 'Input data with bad quality'), # input_gate
{'content': 'transformed_data', 'timestamp': '2024-01-01'}, # transform_data
# mlflow_response_gate (transform)
('continue', 0.95, "Error"),
('continue', 0.95, 'Error'),
# mlflow_content_gate (transform)
('continue', 0.95, "Transformed data not passed the content filter"),
('continue', 0.95, 'Transformed data not passed the content filter'),
{'content': 'predicted_data', 'timestamp': '2024-01-01'}, # request_predict
('continue', 0.95, "Error"), # mlflow_response_gate (predict)
('continue', 0.95, 'Error'), # mlflow_response_gate (predict)
]
# Act
@@ -373,63 +495,117 @@ async def test_run_stop_at_mlflow_last_response_gate(workflow_mock, prediction_p
# Assert
assert workflow_mock.execute_local_activity_method.call_count == 7
workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.get_last_timestamp, {
workflow_mock.execute_local_activity_method.assert_has_calls(
[
call(
Activities.get_last_timestamp,
{
'data': input_data['data'],
**metadata,
},
retry_policy=ANY, start_to_close_timeout=ANY)])
workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.input_gate, {
retry_policy=ANY,
start_to_close_timeout=ANY,
)
]
)
workflow_mock.execute_local_activity_method.assert_has_calls(
[
call(
Activities.input_gate,
{
'filters': input_data['input_filters'],
'data': input_data['data'],
'path_priority': input_data['path_priority'],
**metadata,
},
retry_policy=ANY, start_to_close_timeout=ANY)])
workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.request_transform, {
retry_policy=ANY,
start_to_close_timeout=ANY,
)
]
)
workflow_mock.execute_local_activity_method.assert_has_calls(
[
call(
Activities.request_transform,
{
'data': input_data['data'],
'model_name': input_data['model_name'],
'model_config': input_data['model_config'],
**metadata
}, retry_policy=ANY, start_to_close_timeout=ANY)])
workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.mlflow_response_gate, {
**metadata,
},
retry_policy=ANY,
start_to_close_timeout=ANY,
)
]
)
workflow_mock.execute_local_activity_method.assert_has_calls(
[
call(
Activities.mlflow_response_gate,
{
'filters': input_data['mlflow_transform_filters'],
'data': {'content': 'transformed_data', 'timestamp': '2024-01-01'},
'type': 'transform',
'path_priority': input_data['path_priority'],
**metadata,
}, retry_policy=ANY, start_to_close_timeout=ANY)])
workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.mlflow_content_gate, {
},
retry_policy=ANY,
start_to_close_timeout=ANY,
)
]
)
workflow_mock.execute_local_activity_method.assert_has_calls(
[
call(
Activities.mlflow_content_gate,
{
'filters': input_data['mlflow_transform_filters'],
'data': 'transformed_data',
'type': 'transform',
'path_priority': input_data['path_priority'],
**metadata,
}, retry_policy=ANY, start_to_close_timeout=ANY)])
workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.request_predict, {
},
retry_policy=ANY,
start_to_close_timeout=ANY,
)
]
)
workflow_mock.execute_local_activity_method.assert_has_calls(
[
call(
Activities.request_predict,
{
'data': 'transformed_data',
'model_name': input_data['model_name'],
'model_config': input_data['model_config'],
**metadata
}, retry_policy=ANY, start_to_close_timeout=ANY)])
workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.mlflow_response_gate, {
**metadata,
},
retry_policy=ANY,
start_to_close_timeout=ANY,
)
]
)
workflow_mock.execute_local_activity_method.assert_has_calls(
[
call(
Activities.mlflow_response_gate,
{
'filters': input_data['mlflow_predict_filters'],
'data': {'content': 'predicted_data', 'timestamp': '2024-01-01'},
'type': 'predict',
'path_priority': input_data['path_priority'],
**metadata,
}, retry_policy=ANY, start_to_close_timeout=ANY)])
},
retry_policy=ANY,
start_to_close_timeout=ANY,
)
]
)
workflow_mock.execute_child_workflow.assert_not_called()
@mark.asyncio
@patch("laborious.workflows.sub_workflows.prediction_process.workflow", new_callable=AsyncMock)
@patch('laborious.workflows.sub_workflows.prediction_process.workflow', new_callable=AsyncMock)
async def test_path_flag_handler_stop(workflow_mock, prediction_process):
# Arrange
data = {'test': 'data'}
@@ -440,21 +616,24 @@ async def test_path_flag_handler_stop(workflow_mock, prediction_process):
model = 'test_model'
last_timestamp = '2024-01-01'
model_name = 'test_model_name'
model_config = {
'retention': '30'
}
model_config = {'retention': '30'}
# Act
result = await prediction_process.path_flag_handler(
data, path_flag, {
data,
path_flag,
{
'metadata': metadata,
'schema': schema,
'table_name': table_name,
'model_id': model,
'last_timestamp': last_timestamp,
'model_name': model_name,
'model_config': model_config
}, confidence, last_timestamp, ""
'model_config': model_config,
},
confidence,
last_timestamp,
'',
)
# Assert
@@ -464,7 +643,7 @@ async def test_path_flag_handler_stop(workflow_mock, prediction_process):
@mark.asyncio
@patch("laborious.workflows.sub_workflows.prediction_process.workflow", new_callable=AsyncMock)
@patch('laborious.workflows.sub_workflows.prediction_process.workflow', new_callable=AsyncMock)
async def test_path_flag_handler_repeat(workflow_mock, prediction_process):
# Arrange
data = {'test': 'data'}
@@ -475,21 +654,24 @@ async def test_path_flag_handler_repeat(workflow_mock, prediction_process):
model = 'test_model'
last_timestamp = '2024-01-01'
model_name = 'test_model_name'
model_config = {
'retention': '30'
}
model_config = {'retention': '30'}
# Act
result = await prediction_process.path_flag_handler(
data, path_flag, {
data,
path_flag,
{
'metadata': metadata,
'schema': schema,
'table_name': table_name,
'model_id': model,
'last_timestamp': last_timestamp,
'model_name': model_name,
'model_config': model_config
}, confidence, last_timestamp, ""
'model_config': model_config,
},
confidence,
last_timestamp,
'',
)
# Assert
@@ -504,13 +686,13 @@ async def test_path_flag_handler_repeat(workflow_mock, prediction_process):
'last_timestamp': last_timestamp,
},
retry_policy=ANY,
start_to_close_timeout=ANY
start_to_close_timeout=ANY,
)
workflow_mock.execute_child_workflow.assert_not_called()
@mark.asyncio
@patch("laborious.workflows.sub_workflows.prediction_process.workflow", new_callable=AsyncMock)
@patch('laborious.workflows.sub_workflows.prediction_process.workflow', new_callable=AsyncMock)
async def test_path_flag_handler_continue(workflow_mock, prediction_process):
# Arrange
data = {'test': 'data'}
@@ -521,14 +703,14 @@ async def test_path_flag_handler_continue(workflow_mock, prediction_process):
model = 'test_model'
last_timestamp = '2024-01-01'
model_name = 'test_model_name'
model_config = {
'retention': '30'
}
model_config = {'retention': '30'}
prediction_store_policy = 'erl:1'
# Act
result = await prediction_process.path_flag_handler(
data, path_flag, {
data,
path_flag,
{
'metadata': metadata,
'schema': schema,
'table_name': table_name,
@@ -537,8 +719,11 @@ async def test_path_flag_handler_continue(workflow_mock, prediction_process):
'model_name': model_name,
'model_config': model_config,
'opc_output_config': {'test': 'config'},
'prediction_store_policy': prediction_store_policy
}, confidence, last_timestamp, 'Prediction Process'
'prediction_store_policy': prediction_store_policy,
},
confidence,
last_timestamp,
'Prediction Process',
)
# Assert
@@ -559,13 +744,13 @@ async def test_path_flag_handler_continue(workflow_mock, prediction_process):
'table_name': table_name,
'comment': 'Prediction Process',
'opc_output_config': {'test': 'config'},
'prediction_store_policy': prediction_store_policy
}
'prediction_store_policy': prediction_store_policy,
},
)
@mark.asyncio
@patch("laborious.workflows.sub_workflows.prediction_process.workflow", new_callable=AsyncMock)
@patch('laborious.workflows.sub_workflows.prediction_process.workflow', new_callable=AsyncMock)
async def test_path_flag_handler_unknown(workflow_mock, prediction_process):
# Arrange
data = {'test': 'data'}
@@ -576,13 +761,13 @@ async def test_path_flag_handler_unknown(workflow_mock, prediction_process):
model = 'test_model'
last_timestamp = '2024-01-01'
model_name = 'test_model_name'
model_config = {
'retention': '30'
}
model_config = {'retention': '30'}
prediction_store_policy = 'erl:1'
# Act
result = await prediction_process.path_flag_handler(
data, path_flag, {
data,
path_flag,
{
**metadata,
'schema': schema,
'table_name': table_name,
@@ -591,8 +776,11 @@ async def test_path_flag_handler_unknown(workflow_mock, prediction_process):
'model_name': model_name,
'model_config': model_config,
'opc_output_config': {'test': 'config'},
'prediction_store_policy': prediction_store_policy
}, confidence, last_timestamp, ""
'prediction_store_policy': prediction_store_policy,
},
confidence,
last_timestamp,
'',
)
# Assert

267
tests/laborious/workflows/test_minimal_retrain.py
View File

@@ -1,5 +1,7 @@
from unittest.mock import AsyncMock, MagicMock, call, patch, ANY
from unittest.mock import ANY, AsyncMock, call, patch
from pytest import fixture, mark
from laborious.activities.activities import Activities
from laborious.workflows.minimal_retrain import MinimalRetrain
@@ -10,11 +12,11 @@ def minimal_retrain() -> MinimalRetrain:
metadata = {
"metadata": {
"model_id": "test_model_id",
"model_name": "test_model",
"workflow_name": "minimal_retrain",
"schedule_name": "test_schedule",
'metadata': {
'model_id': 'test_model_id',
'model_name': 'test_model',
'workflow_name': 'minimal_retrain',
'schedule_name': 'test_schedule',
},
}
@@ -23,76 +25,271 @@ metadata = {
@patch('laborious.workflows.minimal_retrain.workflow', new_callable=AsyncMock)
async def test_run(workflow_mock: AsyncMock, minimal_retrain: MinimalRetrain):
input_data = {
"model_id": "test_model_id",
"model_name": "test_model",
"workflow_name": "minimal_retrain",
"schedule_name": "test_schedule",
"query": "test_query",
"schema": "test_schema",
"table_name": "test_table",
'model_id': 'test_model_id',
'model_name': 'test_model',
'workflow_name': 'minimal_retrain',
'schedule_name': 'test_schedule',
'query': 'test_query',
'schema': 'test_schema',
'table_name': 'test_table',
'model_config': {
'target': 'test_target',
'transform_flavor': 'test_transform_flavor',
'predict_flavor': 'test_predict_flavor',
},
}
workflow_mock.execute_activity_method = AsyncMock(
return_value={
"data1": "1",
"data2": "2",
}
side_effect=[
{'success': True, 'object_key': 'test_object_key'},
{'success': True, 'experiment': 'test_experiment'},
{
'success': True,
'version': 'test_version',
'mlflow_run_id': 'test_mlflow_run_id',
'mlflow_experiment_id': 'test_mlflow_experiment_id',
},
{'report': 'test_report'},
]
)
await minimal_retrain.run(input_data)
workflow_mock.execute_local_activity_method.assert_has_calls(
workflow_mock.execute_activity_method.assert_has_calls(
[
call(
Activities.load_custom_query,
Activities.query_to_minio,
{
**metadata,
"query": input_data["query"],
'datetime_columns': input_data.get('datetime_columns', [])
'query': input_data['query'],
'datetime_columns': input_data.get('datetime_columns', []),
'model_name': input_data['model_name'],
'object_prefix': f'retrain_datasets/{input_data["model_name"]}/data',
},
retry_policy=ANY,
start_to_close_timeout=ANY
start_to_close_timeout=ANY,
)
]
)
workflow_mock.execute_activity_method.assert_has_calls([
workflow_mock.execute_activity_method.assert_has_calls(
[
call(
Activities.retrain_model,
{
**metadata,
'data': workflow_mock.execute_local_activity_method.return_value,
'object_key': 'test_object_key',
'model_name': input_data['model_name'],
'model_config': input_data['model_config'],
},
retry_policy=ANY,
start_to_close_timeout=ANY
start_to_close_timeout=ANY,
)
]
)
])
workflow_mock.execute_activity_method.assert_has_calls([
workflow_mock.execute_activity_method.assert_has_calls(
[
call(
Activities.update_production_model,
{
**metadata,
'model_name': input_data['model_name'],
'model_id': input_data['model_id'],
**workflow_mock.execute_activity_method.return_value,
'success': True,
'experiment': 'test_experiment',
},
retry_policy=ANY,
start_to_close_timeout=ANY
start_to_close_timeout=ANY,
)
]
)
])
workflow_mock.execute_activity_method.assert_has_calls([
workflow_mock.execute_local_activity_method.assert_has_calls(
[
call(
Activities.format_retrain_report,
{
**metadata,
'experiment_response': {'success': True, 'experiment': 'test_experiment'},
'model_name': input_data['model_name'],
'update_report': {
'success': True,
'version': 'test_version',
'mlflow_run_id': 'test_mlflow_run_id',
'mlflow_experiment_id': 'test_mlflow_experiment_id',
},
},
retry_policy=ANY,
start_to_close_timeout=ANY,
)
]
)
workflow_mock.execute_activity_method.assert_has_calls(
[
call(
Activities.export_data_to_postgres,
{
**metadata,
'data': workflow_mock.execute_activity_method.return_value,
'data': workflow_mock.execute_local_activity_method.return_value,
'schema': input_data['schema'],
'table_name': input_data['table_name'],
},
retry_policy=ANY,
start_to_close_timeout=ANY
start_to_close_timeout=ANY,
)
])
]
)
@mark.asyncio
@patch('laborious.workflows.minimal_retrain.workflow', new_callable=AsyncMock)
async def test_run_storage_fail(workflow_mock: AsyncMock, minimal_retrain: MinimalRetrain):
input_data = {
'model_id': 'test_model_id',
'model_name': 'test_model',
'workflow_name': 'minimal_retrain',
'schedule_name': 'test_schedule',
'query': 'test_query',
'schema': 'test_schema',
'table_name': 'test_table',
'model_config': {
'target': 'test_target',
'transform_flavor': 'test_transform_flavor',
'predict_flavor': 'test_predict_flavor',
},
}
workflow_mock.execute_activity_method = AsyncMock(
side_effect=[
{'success': False, 'object_key': 'test_object_key'},
{'success': True, 'experiment': 'test_experiment'},
{
'success': True,
'version': 'test_version',
'mlflow_run_id': 'test_mlflow_run_id',
'mlflow_experiment_id': 'test_mlflow_experiment_id',
},
{'report': 'test_report'},
]
)
await minimal_retrain.run(input_data)
workflow_mock.execute_activity_method.assert_called_once_with(
Activities.query_to_minio,
{
**metadata,
'query': input_data['query'],
'datetime_columns': input_data.get('datetime_columns', []),
'model_name': input_data['model_name'],
'object_prefix': f'retrain_datasets/{input_data["model_name"]}/data',
},
retry_policy=ANY,
start_to_close_timeout=ANY,
)
workflow_mock.execute_local_activity_method.assert_not_called()
@mark.asyncio
@patch('laborious.workflows.minimal_retrain.workflow', new_callable=AsyncMock)
async def test_run_fail_retrain(workflow_mock: AsyncMock, minimal_retrain: MinimalRetrain):
input_data = {
'model_id': 'test_model_id',
'model_name': 'test_model',
'workflow_name': 'minimal_retrain',
'schedule_name': 'test_schedule',
'query': 'test_query',
'schema': 'test_schema',
'table_name': 'test_table',
'model_config': {
'target': 'test_target',
'transform_flavor': 'test_transform_flavor',
'predict_flavor': 'test_predict_flavor',
},
}
workflow_mock.execute_activity_method = AsyncMock(
side_effect=[
{'success': True, 'object_key': 'test_object_key'},
{'success': False, 'experiment': 'test_experiment'},
{
'success': True,
'version': 'test_version',
'mlflow_run_id': 'test_mlflow_run_id',
'mlflow_experiment_id': 'test_mlflow_experiment_id',
},
{'report': 'test_report'},
]
)
await minimal_retrain.run(input_data)
workflow_mock.execute_activity_method.assert_has_calls(
[
call(
Activities.query_to_minio,
{
**metadata,
'query': input_data['query'],
'datetime_columns': input_data.get('datetime_columns', []),
'model_name': input_data['model_name'],
'object_prefix': f'retrain_datasets/{input_data["model_name"]}/data',
},
retry_policy=ANY,
start_to_close_timeout=ANY,
)
]
)
workflow_mock.execute_activity_method.assert_has_calls(
[
call(
Activities.retrain_model,
{
**metadata,
'object_key': 'test_object_key',
'model_name': input_data['model_name'],
'model_config': input_data['model_config'],
},
retry_policy=ANY,
start_to_close_timeout=ANY,
)
]
)
workflow_mock.execute_local_activity_method.assert_has_calls(
[
call(
Activities.format_retrain_report,
{
**metadata,
'experiment_response': {'success': False, 'experiment': 'test_experiment'},
'model_name': input_data['model_name'],
'update_report': {},
},
retry_policy=ANY,
start_to_close_timeout=ANY,
)
]
)
workflow_mock.execute_activity_method.assert_has_calls(
[
call(
Activities.export_data_to_postgres,
{
**metadata,
'data': workflow_mock.execute_local_activity_method.return_value,
'schema': input_data['schema'],
'table_name': input_data['table_name'],
},
retry_policy=ANY,
start_to_close_timeout=ANY,
)
]
)
assert workflow_mock.execute_activity_method.call_count == 3
assert workflow_mock.execute_local_activity_method.call_count == 1

63
tests/laborious/workflows/test_predictions_batch.py
View File

@@ -1,5 +1,7 @@
from unittest.mock import AsyncMock, call, patch, ANY
from unittest.mock import ANY, AsyncMock, call, patch
from pytest import fixture, mark
from laborious.activities.activities import Activities
from laborious.workflows.predictions_batch import PredictionsBatch
@@ -10,11 +12,11 @@ def predictions_batch() -> PredictionsBatch:
metadata = {
"metadata": {
"model_id": "test_model_id",
"model_name": "test_model",
"workflow_name": "predictions_batch",
"schedule_name": "test_schedule",
'metadata': {
'model_id': 'test_model_id',
'model_name': 'test_model',
'workflow_name': 'predictions_batch',
'schedule_name': 'test_schedule',
},
}
@@ -22,9 +24,7 @@ metadata = {
@mark.asyncio
@patch('laborious.workflows.predictions_batch.workflow', new_callable=AsyncMock)
async def test_run(workflow_mock: AsyncMock, predictions_batch: PredictionsBatch):
workflow_mock.execute_local_activity_method.return_value = {
'data': 'test_data'
}
workflow_mock.execute_local_activity_method.return_value = {'data': 'test_data'}
input_data = {
'schedule_name': 'test_schedule',
'model_name': 'test_model',
@@ -35,25 +35,25 @@ async def test_run(workflow_mock: AsyncMock, predictions_batch: PredictionsBatch
'opc_output_config': 'test_opc_output_config',
'datetime_columns': ['timestamp', 'created_at'],
'prediction_store_policy': 'erl:1',
'model_config': {
'retention': '30'
}
'model_config': {'retention': '30'},
}
await predictions_batch.run(input_data)
workflow_mock.execute_local_activity_method.assert_has_calls([
workflow_mock.execute_local_activity_method.assert_has_calls(
[
call(
Activities.load_custom_query,
{
**metadata,
'query': input_data['query'],
'datetime_columns': input_data.get('datetime_columns', [])
'datetime_columns': input_data.get('datetime_columns', []),
},
retry_policy=ANY,
start_to_close_timeout=ANY
start_to_close_timeout=ANY,
)
]
)
])
prediction_input = {
'metadata': metadata,
'data': {'data': 'test_data'},
@@ -61,28 +61,19 @@ async def test_run(workflow_mock: AsyncMock, predictions_batch: PredictionsBatch
'table_name': input_data['table_name'],
'model_id': input_data['model_id'],
'model_name': input_data['model_name'],
'input_filters': input_data.get('input_filters', {
'EMPTY_DATA': {
'POLICY': 'STOP'
}
}),
'mlflow_transform_filters': input_data.get('mlflow_transform_filters', {
'API_ERROR': {
'POLICY': 'STOP'
}
}),
'mlflow_predict_filters': input_data.get('mlflow_predict_filters', {
'API_ERROR': {
'POLICY': 'STOP'
}
}),
'input_filters': input_data.get('input_filters', {'EMPTY_DATA': {'POLICY': 'STOP'}}),
'mlflow_transform_filters': input_data.get(
'mlflow_transform_filters', {'API_ERROR': {'POLICY': 'STOP'}}
),
'mlflow_predict_filters': input_data.get(
'mlflow_predict_filters', {'API_ERROR': {'POLICY': 'STOP'}}
),
'model_config': input_data.get('model_config', {}),
'path_priority': input_data.get('path_priority', ['STOP', 'CONTINUE', 'REPEAT']),
'opc_output_config': input_data.get('opc_output_config', {}),
'prediction_store_policy': input_data.get('prediction_store_policy', 'erl:1')
'prediction_store_policy': input_data.get('prediction_store_policy', 'erl:1'),
}
workflow_mock.execute_child_workflow.assert_has_calls([
call(
'prediction_process', prediction_input)
])
workflow_mock.execute_child_workflow.assert_has_calls(
[call('prediction_process', prediction_input)]
)

99
validate.sh Executable file
View File

@@ -0,0 +1,99 @@
#!/bin/bash
# Model Manager Code Validation Script
# This script runs all code quality checks before committing or deploying
set -e # Exit on any error
# Colors for output
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m' # No Color
echo -e "${BLUE}╔════════════════════════════════════════════════════════╗${NC}"
echo -e "${BLUE}║ Model Manager - Code Validation Suite ║${NC}"
echo -e "${BLUE}╚════════════════════════════════════════════════════════╝${NC}"
echo ""
# Check if virtual environment is activated
if [[ -z "${VIRTUAL_ENV}" ]] && [[ -z "${CONDA_DEFAULT_ENV}" ]]; then
echo -e "${YELLOW}⚠️ Warning: No virtual environment detected${NC}"
echo -e "${YELLOW} Consider activating your venv/conda environment${NC}"
echo ""
fi
# Function to run a validation step
run_step() {
local step_name=$1
local step_command=$2
echo -e "${BLUE}━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━${NC}"
echo -e "${BLUE}▶ ${step_name}${NC}"
echo -e "${BLUE}━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━${NC}"
if eval "$step_command"; then
echo -e "${GREEN}✅ ${step_name} - PASSED${NC}"
echo ""
return 0
else
echo -e "${RED}❌ ${step_name} - FAILED${NC}"
echo ""
return 1
fi
}
# Track failures
FAILED_STEPS=()
# Step 1: Code Formatting Check (Ruff)
if ! run_step "1. Code Formatting (Ruff)" "ruff format --check laborious/ tests/"; then
FAILED_STEPS+=("Code Formatting")
fi
# Step 2: Linting (Ruff)
if ! run_step "2. Code Linting (Ruff)" "ruff check laborious/ tests/"; then
FAILED_STEPS+=("Linting")
fi
# Step 3: Type Checking (mypy)
if ! run_step "3. Type Checking (mypy)" "mypy laborious/"; then
FAILED_STEPS+=("Type Checking")
fi
# Step 4: Security Analysis (Bandit)
if ! run_step "4. Security Analysis (Bandit)" "bandit -r laborious/ -ll -q"; then
FAILED_STEPS+=("Security Analysis")
fi
# Step 5: Unit Tests (pytest)
if ! run_step "5. Unit Tests (pytest)" "pytest tests/ --cov=laborious --cov-report=term-missing --cov-report=xml --cov-report=html --cov-fail-under=80 -q"; then
FAILED_STEPS+=("Unit Tests")
fi
# Summary
echo -e "${BLUE}╔════════════════════════════════════════════════════════╗${NC}"
echo -e "${BLUE}║ Validation Summary ║${NC}"
echo -e "${BLUE}╚════════════════════════════════════════════════════════╝${NC}"
echo ""
if [ ${#FAILED_STEPS[@]} -eq 0 ]; then
echo -e "${GREEN}✅ All validation checks passed!${NC}"
echo -e "${GREEN} Your code is ready for commit/deployment.${NC}"
echo ""
exit 0
else
echo -e "${RED}❌ Validation failed for the following steps:${NC}"
for step in "${FAILED_STEPS[@]}"; do
echo -e "${RED} • ${step}${NC}"
done
echo ""
echo -e "${YELLOW}💡 Tips:${NC}"
echo -e "${YELLOW} • Run 'ruff format laborious/ tests/' to auto-fix formatting${NC}"
echo -e "${YELLOW} • Run 'ruff check --fix laborious/ tests/' to auto-fix linting issues${NC}"
echo -e "${YELLOW} • Review mypy errors and add type hints where needed${NC}"
echo -e "${YELLOW} • Check bandit warnings for security issues${NC}"
echo -e "${YELLOW} • Fix failing tests or improve test coverage${NC}"
echo ""
exit 1
fi