From ac795c7c5344f5f96e4191c41ece1843685cf0dd Mon Sep 17 00:00:00 2001 From: vitor-aignosi Date: Wed, 15 Oct 2025 16:00:18 -0300 Subject: [PATCH] SIENTIAPDE-1231 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. --- .gitignore | 1 + laborious/activities/activities.py | 76 +- laborious/activities/gates.py | 320 +++-- laborious/activities/mlflow.py | 144 +- laborious/activities/opc.py | 119 +- laborious/activities/storage.py | 121 +- laborious/metrics.py | 38 +- laborious/utils/connectors_config.py | 24 +- .../utils/filters/conditional_filters.py | 3 +- laborious/utils/filters/mlflow_filters.py | 7 +- .../utils/repository/minio_repository.py | 90 +- .../utils/repository/model_repository.py | 765 ++++------ laborious/utils/repository/opc_repository.py | 182 ++- laborious/worker/worker.py | 77 +- laborious/workflows/minimal_retrain.py | 44 +- laborious/workflows/predictions_batch.py | 46 +- .../format_and_export_prediction.py | 42 +- .../sub_workflows/prediction_process.py | 55 +- mlruns/0/meta.yaml | 6 + mlruns/586524947870967910/meta.yaml | 6 + mlruns/models/test/meta.yaml | 5 + pyproject.toml | 159 +++ requirements-dev.txt | 19 + tests/laborious/activities/test_activities.py | 79 +- tests/laborious/activities/test_gates.py | 237 ++-- tests/laborious/activities/test_mlflow.py | 414 ++++-- tests/laborious/activities/test_opc.py | 339 ++--- tests/laborious/activities/test_storage.py | 203 +++ .../utils/filters/test_conditional_filters.py | 31 +- .../utils/filters/test_mlflow_filters.py | 11 +- .../utils/repository/test_minio_repository.py | 134 ++ .../utils/repository/test_model_repository.py | 1245 ++++++++++++----- .../utils/repository/test_opc_repository.py | 203 +-- .../laborious/utils/test_connectors_config.py | 15 +- .../test_format_and_export_prediction.py | 253 ++-- .../subworkflows/test_prediction_process.py | 720 ++++++---- .../workflows/test_minimal_retrain.py | 313 ++++- .../workflows/test_predictions_batch.py | 79 +- validate.sh | 99 ++ 39 files changed, 4122 insertions(+), 2602 deletions(-) create mode 100644 mlruns/0/meta.yaml create mode 100644 mlruns/586524947870967910/meta.yaml create mode 100644 mlruns/models/test/meta.yaml create mode 100644 pyproject.toml create mode 100644 requirements-dev.txt create mode 100644 tests/laborious/activities/test_storage.py create mode 100644 tests/laborious/utils/repository/test_minio_repository.py create mode 100755 validate.sh diff --git a/.gitignore b/.gitignore index 2c8240d..00fc074 100644 --- a/.gitignore +++ b/.gitignore @@ -37,6 +37,7 @@ __pycache__/ # Ignorar coverage htmlcov/ .coverage +coverage.xml # git keys git_key* diff --git a/laborious/activities/activities.py b/laborious/activities/activities.py index d028064..5487d2a 100644 --- a/laborious/activities/activities.py +++ b/laborious/activities/activities.py @@ -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, - 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): + 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, + ): """ Initialize the Activities orchestrator with all required configurations. @@ -59,32 +63,36 @@ class Activities(Storage, MLFlow, Gates, OPC): Exception: If any parent class initialization fails """ # Initialize parent classes - Storage.__init__(self, host=postgres_config['host'], - port=postgres_config['port'], - user=postgres_config['user'], - password=postgres_config['password'], - 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) + Storage.__init__( + self, + host=postgres_config['host'], + port=postgres_config['port'], + user=postgres_config['user'], + password=postgres_config['password'], + 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, + ) - 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) + 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, + ) - 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): """ diff --git a/laborious/activities/gates.py b/laborious/activities/gates.py index 7e6f785..0653161 100644 --- a/laborious/activities/gates.py +++ b/laborious/activities/gates.py @@ -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': { - 'STOP': -1, - 'CONTINUE': 2, - 'REPEAT': -1 - } +} + +# Confidence mappings kept separate from function maps to avoid Union types +input_path_confidence: Mapping[str, int] = { + 'STOP': -1, + 'CONTINUE': 2, + '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': { - 'STOP': -1, - 'CONTINUE': 10, - 'REPEAT': -1 - }, +} + +mlflow_response_path_confidence: Mapping[str, int] = { + 'STOP': -1, + 'CONTINUE': 10, + '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': { - 'STOP': -1, - 'CONTINUE': 18, - 'REPEAT': -1 - } +} + +mlflow_content_path_confidence: Mapping[str, int] = { + 'STOP': -1, + 'CONTINUE': 18, + '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({ - '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']] - }) + 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']], + } + ) - 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({ - 'model_id': [model_id], - 'model_name': [model_name], - 'timestamp': [experiment_response['timestamp']], - 'status': [experiment_response['message']] - }) + report = DataFrame( + { + 'model_id': [model_id], + 'model_name': [model_name], + 'timestamp': [experiment_response['timestamp']], + '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) diff --git a/laborious/activities/mlflow.py b/laborious/activities/mlflow.py index 822cd69..a289c9f 100644 --- a/laborious/activities/mlflow.py +++ b/laborious/activities/mlflow.py @@ -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 diff --git a/laborious/activities/opc.py b/laborious/activities/opc.py index 2c427e0..b4604d8 100644 --- a/laborious/activities/opc.py +++ b/laborious/activities/opc.py @@ -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() diff --git a/laborious/activities/storage.py b/laborious/activities/storage.py index ef61fb4..8aec756 100644 --- a/laborious/activities/storage.py +++ b/laborious/activities/storage.py @@ -1,20 +1,21 @@ from temporalio import activity, workflow -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 + from typing import Any -DATETIME_FILENAME_FORMAT = "%Y-%m-%d_%H-%M-%S" + 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 + + +DATETIME_FILENAME_FORMAT = '%Y-%m-%d_%H-%M-%S' class Storage(Postgres): @@ -23,36 +24,33 @@ class Storage(Postgres): directly to MinIO as Parquet and return the object name. """ - def __init__(self, - host: str, - port: int, - user: str, - password: str, - dbname: str, - min_connections: int, - max_connections: int, - minio_config: dict[str, Any], - logger: Logger, - notification_handler: NotificationHandler): - super().__init__(host=host, - port=port, - user=user, - password=password, - dbname=dbname, - min_connections=min_connections, - max_connections=max_connections, - logger=logger, - notification_handler=notification_handler) + def __init__( + self, + host: str, + port: int, + user: str, + password: str, + dbname: str, + min_connections: int, + max_connections: int, + minio_config: dict[str, Any], + logger: Logger, + notification_handler: NotificationHandler, + ): + super().__init__( + host=host, + port=port, + user=user, + password=password, + dbname=dbname, + min_connections=min_connections, + max_connections=max_connections, + logger=logger, + 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 + self.close() diff --git a/laborious/metrics.py b/laborious/metrics.py index 97f7bb9..7826fac 100644 --- a/laborious/metrics.py +++ b/laborious/metrics.py @@ -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], ) diff --git a/laborious/utils/connectors_config.py b/laborious/utils/connectors_config.py index 969b95c..7501a07 100644 --- a/laborious/utils/connectors_config.py +++ b/laborious/utils/connectors_config.py @@ -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'), } diff --git a/laborious/utils/filters/conditional_filters.py b/laborious/utils/filters/conditional_filters.py index 6caf1e5..967e52c 100644 --- a/laborious/utils/filters/conditional_filters.py +++ b/laborious/utils/filters/conditional_filters.py @@ -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: diff --git a/laborious/utils/filters/mlflow_filters.py b/laborious/utils/filters/mlflow_filters.py index f792a0e..82970e8 100644 --- a/laborious/utils/filters/mlflow_filters.py +++ b/laborious/utils/filters/mlflow_filters.py @@ -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 diff --git a/laborious/utils/repository/minio_repository.py b/laborious/utils/repository/minio_repository.py index d957e65..df20a06 100644 --- a/laborious/utils/repository/minio_repository.py +++ b/laborious/utils/repository/minio_repository.py @@ -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, - 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): - +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, + ): # 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.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 + self.logger.custom_info(f"Creating bucket '{self.minio_bucket}'", metadata) + self.s3_client.create_bucket(Bucket=self.minio_bucket) - def store_dataframe_as_parquet(self, dataframe: DataFrame, uri: str, - object_name: str, metadata: dict[str, Any]): + return True + 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()) diff --git a/laborious/utils/repository/model_repository.py b/laborious/utils/repository/model_repository.py index 691857a..d84f527 100644 --- a/laborious/utils/repository/model_repository.py +++ b/laborious/utils/repository/model_repository.py @@ -13,53 +13,48 @@ The repository provides comprehensive functionality for: - Model retraining workflows - Production model updates and versioning """ -from datetime import datetime, timedelta -import traceback -from mlflow.entities import Experiment, experiment -import pandas as pd -import mlflow -from os import makedirs, path, remove, environ -from shutil import rmtree -from sys import path as sys_path -from sientia_do.observability.logger import Logger -import lzma -import gzip -import pickle -from numpy import ndarray -from typing import Any -import gc -from sientia_do.temporal.constants import DATETIME_FORMAT_WITH_TZ -import ctypes -ARTIFACTS_PATH = "./tmp/artifacts" -TRANSFORMED_COMPRESSED_PATH = "artifacts/training_transformer.pkl" -PREDICTION_COMPRESSED_PATH = "artifacts/stacking_model.pkl" +import ctypes +import gc +import traceback +from datetime import datetime, timedelta +from os import environ, makedirs, path +from shutil import rmtree +from typing import Any, Literal, overload + +import mlflow +import pandas as pd +from mlflow.entities import Experiment +from numpy import ndarray +from sientia_do.observability.logger import Logger +from sientia_do.temporal.constants import DATETIME_FORMAT_WITH_TZ + +ARTIFACTS_PATH = './tmp/artifacts' +TRANSFORMED_COMPRESSED_PATH = 'artifacts/training_transformer.pkl' +PREDICTION_COMPRESSED_PATH = 'artifacts/stacking_model.pkl' def force_memory_release(logger: Logger): gc.collect() try: - ctypes.CDLL("libc.so.6").malloc_trim(0) - logger.info( - f"Memory released") + ctypes.CDLL('libc.so.6').malloc_trim(0) + logger.info('Memory released') except Exception as e: - logger.info( - f"Memory release failed: {e}") + logger.info(f'Memory release failed: {e}') -class MLFlowRepository(): +class MLFlowRepository: def __init__(self, host: str, username: str, password: str, logger: Logger): - # set tracking uri mlflow.set_tracking_uri(host) - environ["MLFLOW_TRACKING_USERNAME"] = username - environ["MLFLOW_TRACKING_PASSWORD"] = password + environ['MLFLOW_TRACKING_USERNAME'] = username + environ['MLFLOW_TRACKING_PASSWORD'] = password # Create an MLflow client self.client = mlflow.tracking.MlflowClient() - self.model_cache = {} + self.model_cache: dict[str, Any] = {} self.logger = logger """ @@ -79,12 +74,12 @@ class MLFlowRepository(): """ run_info = mlflow.get_run(run_id) if prediction: - model_uri = run_info.info.artifact_uri + "/prediction_model" + model_uri = run_info.info.artifact_uri + '/prediction_model' else: - model_uri = run_info.info.artifact_uri + "/data_model" + model_uri = run_info.info.artifact_uri + '/data_model' return model_uri - def get_model_run_id(self, model_name: str, stage: str = "Production"): + def get_model_run_id(self, model_name: str, stage: str = 'Production') -> str: """ Get the run_id of a model based on its name and stage. @@ -106,13 +101,10 @@ class MLFlowRepository(): ) # Get all versions of the model and filter by stage - model_versions = self.client.search_model_versions( - filter_string=f"name='{model_name}'" - ) + model_versions = self.client.search_model_versions(filter_string=f"name='{model_name}'") # Filter versions by the desired stage using current_stage attribute - stage_versions = [ - mv for mv in model_versions if mv.current_stage == stage] + stage_versions = [mv for mv in model_versions if mv.current_stage == stage] if not stage_versions: raise mlflow.exceptions.MlflowException( @@ -121,28 +113,9 @@ class MLFlowRepository(): # Sort by version number to get the latest latest_version = max(stage_versions, key=lambda v: int(v.version)) - run_id = latest_version.source.split("/") + run_id = latest_version.source.split('/') return run_id[2] - def get_experiment_by_run_id(self, run_id: str) -> Experiment: - """ - Get experiment name by run ID. - - Args: - run_id (str): The MLFlow run ID - - Returns: - str: The experiment name - """ - # Get the run information using the run_id - run = mlflow.get_run(run_id) - - # Extract the experiment ID from the run - experiment_id = run.info.experiment_id - - # Get the experiment details using the experiment ID - return mlflow.get_experiment(experiment_id) - def get_next_run_name(self, model_name: str) -> str: """ Generate the next run name for a specific MLFlow model. @@ -157,12 +130,13 @@ class MLFlowRepository(): Returns: str: The next run name in format 'model_name-run_number' """ - runs = mlflow.search_runs( - experiment_names=[model_name], order_by=["start_time desc"]) + runs = mlflow.search_runs(experiment_names=[model_name], order_by=['start_time desc']) next_run_number = len(runs) + 1 - return f"{model_name}-{next_run_number}" + return f'{model_name}-{next_run_number}' - def get_experiment(self, experiment_name: str, create_if_not_exists: bool = False) -> Experiment: + def get_experiment( + self, experiment_name: str, create_if_not_exists: bool = False + ) -> Experiment: """ Retrieve MLFlow experiment ID by experiment name. @@ -189,47 +163,6 @@ class MLFlowRepository(): return experiment - def get_experiment_last_run(self, experiment_id: int) -> str: - """ - Retrieve the most recent retraining run ID for an experiment. - - This method searches for the latest run in an MLFlow experiment - that has been marked as a retraining run. It filters runs by - the 'retrain' parameter and orders them by completion time. - - Args: - experiment_id (int): MLFlow experiment ID - - Returns: - str: MLFlow run ID of the most recent retraining run - - Raises: - ValueError: If runs data is not in expected DataFrame format - """ - runs = mlflow.search_runs( - experiment_ids=[experiment_id], - filter_string="", # Sem filtro no MLflow ainda - output_format="pandas" - ) - - if not isinstance(runs, pd.DataFrame): - raise ValueError('Runs is not a pandas DataFrame') - - # Filtrar apenas as runs onde params.retrain == True - filtered_runs = runs[runs["params.retrain"] == 'True'] - - # Converter a coluna 'end_time' para datetime - filtered_runs['end_time'] = pd.to_datetime(filtered_runs['end_time']) - - # Ordenar o DataFrame de forma descendente pela coluna 'end_time' - filtered_runs = filtered_runs.sort_values( - by='end_time', ascending=False) - - # Pegar a última run_id do DataFrame filtrado e ordenado - latest_run_id = filtered_runs.iloc[0]['run_id'] - - return latest_run_id - def get_model_params(self, run_id: str): """Obtém os parâmetros de uma run""" run_info = mlflow.get_run(run_id) @@ -239,7 +172,7 @@ class MLFlowRepository(): Functions related to download and load models """ - def dowload_artifacts(self, model_name: str, artifact_path: str = "data_model") -> str: + def dowload_artifacts(self, model_name: str, artifact_path: str = 'data_model') -> str: """ Downloads artifacts from a specific MLFlow run. @@ -250,10 +183,8 @@ class MLFlowRepository(): Returns: str: Path to the downloaded artifacts """ - run_id = self.get_model_run_id( - model_name=model_name, stage="Production" - ) - output_dir = f"{ARTIFACTS_PATH}/{model_name}" + run_id = self.get_model_run_id(model_name=model_name, stage='Production') + output_dir = f'{ARTIFACTS_PATH}/{model_name}' full_path = path.join(output_dir, artifact_path) @@ -262,14 +193,9 @@ class MLFlowRepository(): rmtree(full_path) makedirs(output_dir, exist_ok=True) - self.logger.info( - f"Downloading artifacts from {run_id} to {output_dir}") + self.logger.info(f'Downloading artifacts from {run_id} to {output_dir}') - return self.client.download_artifacts( - run_id, - artifact_path, - output_dir - ) + return self.client.download_artifacts(run_id, artifact_path, output_dir) def load_predict_model(self, model_name: str, flavor: str = 'sklearn') -> Any: """ @@ -287,9 +213,8 @@ class MLFlowRepository(): - The model is fetched from the "production" stage of the MLflow Model Registry. - Warnings during the model loading process are suppressed. """ - model_uri = f"models:/{model_name}/production" - self.logger.info( - f"Loading prediction model {model_name} from {model_uri}") + model_uri = f'models:/{model_name}/production' + self.logger.info(f'Loading prediction model {model_name} from {model_uri}') if flavor == 'pyfunc': model = mlflow.pyfunc.load_model(model_uri) elif flavor == 'sklearn': @@ -297,8 +222,7 @@ class MLFlowRepository(): elif flavor == 'pytorch': model = mlflow.pytorch.load_model(model_uri) else: - raise ValueError( - "Invalid flavor. Use 'sklearn' or 'pyfunc' or 'pytorch'.") + raise ValueError("Invalid flavor. Use 'sklearn' or 'pyfunc' or 'pytorch'.") return model @@ -322,14 +246,10 @@ class MLFlowRepository(): model cannot be loaded. """ - latest_production_id = self.get_model_run_id( - model_name=model_name, stage="Production" - ) - model_uri = self.get_model_uri( - latest_production_id, prediction=False) + latest_production_id = self.get_model_run_id(model_name=model_name, stage='Production') + model_uri = self.get_model_uri(latest_production_id, prediction=False) - self.logger.info( - f"Loading data model {model_name} from {model_uri}") + self.logger.info(f'Loading data model {model_name} from {model_uri}') if flavor == 'sklearn': model = mlflow.sklearn.load_model(model_uri) elif flavor == 'pyfunc': @@ -337,64 +257,12 @@ class MLFlowRepository(): elif flavor == 'pytorch': model = mlflow.pytorch.load_model(model_uri) else: - raise ValueError( - "Invalid flavor. Use 'sklearn' or 'pyfunc' or 'pytorch'.") + raise ValueError("Invalid flavor. Use 'sklearn' or 'pyfunc' or 'pytorch'.") return model - def load_model_with_compression(self, artifact_path: str, model_type: str) -> Any: - """ - Load model from pickle file trying different compression methods. - - Args: - artifact_path (str): Path to the artifact directory - type (str): Type of model ('transformer' or 'prediction') - - Returns: - Any: Loaded model object - - Raises: - ValueError: If model cannot be loaded with any compression method - """ - code_path = path.join( - artifact_path, "code") - - pickle_file = TRANSFORMED_COMPRESSED_PATH if model_type == "transform" else PREDICTION_COMPRESSED_PATH - - pickle_path = path.join(artifact_path, pickle_file) - - self.logger.info( - f"Loading model with type {model_type} from {pickle_path}") - - if code_path not in sys_path: - sys_path.insert(0, code_path) - self.logger.info( - f"Added {code_path} to Python path") - - loading_methods = [ - ("lzma", lambda p: lzma.open(p, "rb")), - ("gzip", lambda p: gzip.open(p, "rb")), - ("pickle", lambda p: open(p, "rb")), - ] - - for format_name, open_func in loading_methods: - try: - self.logger.inf - (f"Trying to load with {format_name}...") - with open_func(pickle_path) as f: - model = pickle.load(f) - self.logger.info( - f"Successfully loaded with {format_name}!") - return model - except (lzma.LZMAError, gzip.BadGzipFile, OSError, pickle.UnpicklingError, ValueError) as e: - self.logger.info( - f"Failed with {format_name}: {e.__class__.__name__}:{e}") - continue - - raise ValueError( - f"Could not load model from {pickle_path} - unknown or corrupted format") - - def download_model(self, model_name: str, model_type: str, flavor: str, - load_wrapper: bool = False) -> tuple[Any, str]: + def download_model( + self, model_name: str, model_type: str, flavor: str, load_wrapper: bool = False + ) -> tuple[Any, str | None]: """ Download model based on type (predict or transform). @@ -409,37 +277,35 @@ class MLFlowRepository(): """ self.logger.info( - f"Downloading {model_type} model {model_name} with flavor {flavor} and load_wrapper {load_wrapper}") + f'Downloading {model_type} model {model_name} with flavor {flavor} and load_wrapper {load_wrapper}' + ) - if model_type not in ["predict", "transform"]: - raise ValueError( - "Invalid model_type. Use 'predict' or 'transform'.") + if model_type not in ['predict', 'transform']: + raise ValueError("Invalid model_type. Use 'predict' or 'transform'.") artifact_path = None if load_wrapper: self.logger.info( - f"Loading wrapper for {model_type} model {model_name} with flavor {flavor}") + f'Loading wrapper for {model_type} model {model_name} with flavor {flavor}' + ) - target = "prediction_model" if model_type == "predict" else "data_model" + target = 'prediction_model' if model_type == 'predict' else 'data_model' - artifact_path = self.dowload_artifacts( - model_name, target) + artifact_path = self.dowload_artifacts(model_name, target) self.logger.info( - f"Model with type {model_type} and name {model_name} is compressed, loading from {artifact_path}") + f'Model with type {model_type} and name {model_name} is compressed, loading from {artifact_path}' + ) raw_model = mlflow.pyfunc.load_model(artifact_path) model = raw_model._model_impl.python_model else: + if model_type == 'predict': + model = self.load_predict_model(model_name, flavor) - if model_type == "predict": - model = self.load_predict_model( - model_name, flavor) - - elif model_type == "transform": - model = self.load_transform_model( - model_name, flavor) + elif model_type == 'transform': + model = self.load_transform_model(model_name, flavor) return model, artifact_path @@ -466,32 +332,31 @@ class MLFlowRepository(): # Get type of first element of index index_type = type(index[0]) - self.logger.custom_info(f"Index type: {index_type}", metadata) + self.logger.custom_info(f'Index type: {index_type}', metadata) - message = f"Index must be all timestamp like column. Valid formats are: pandas Timestamp, datetime, string in format {DATETIME_FORMAT_WITH_TZ}. Got {index_type}." + message = f'Index must be all timestamp like column. Valid formats are: pandas Timestamp, datetime, string in format {DATETIME_FORMAT_WITH_TZ}.' # Check if all in index are of the same type if not all(isinstance(i, index_type) for i in index): - raise ValueError( - f"{message}. Elements are {','.join(map(type, index))}") + types = map(str, map(type, index)) + raise ValueError(f'{message}. Elements are {",".join(types)}') # Check type and converts to DATETIME_FORMAT_WITH_TZ - if index_type == str: + if index_type is str: # Validate format of string and return error if not valid try: pd.to_datetime(data.index, format=DATETIME_FORMAT_WITH_TZ) - except ValueError: - raise ValueError( - f"{message}") + except ValueError as e: + raise ValueError(f'{message}. Unable to parse given date format: {e}') from e elif index_type == datetime or index_type == pd.Timestamp: - if data.index.tz is None: - data.index = data.index.tz_localize('UTC') + index = data.index + if hasattr(index, 'tz') and index.tz is None: + data.index = index.tz_localize('UTC') # type: ignore[attr-defined] - data.index = data.index.strftime(DATETIME_FORMAT_WITH_TZ) + data.index = data.index.strftime(DATETIME_FORMAT_WITH_TZ) # type: ignore[attr-defined] else: - raise ValueError( - f"{message}") + raise ValueError(f'{message}. Got {index_type}.') return data @@ -516,8 +381,7 @@ class MLFlowRepository(): return False return True - def handle_valid_model(self, model_name: str, model_type: str, - cache: dict) -> dict: + def handle_valid_model(self, model_name: str, cache: dict) -> dict: """ Handle valid cached model by returning appropriate model configuration. @@ -531,9 +395,7 @@ class MLFlowRepository(): Returns: dict: Model configuration with model and artifact path """ - if self.logger: - self.logger.debug( - f"Model {model_name} is still valid, using cached version") + self.logger.debug(f'Model {model_name} is still valid, using cached version') return cache['target'] @@ -548,9 +410,7 @@ class MLFlowRepository(): Returns: None """ - if self.logger: - self.logger.debug( - f"Model {model_name} is outdated, downloading a new one") + self.logger.debug(f'Model {model_name} is outdated, downloading a new one') del self.model_cache[model_key]['target']['model'] del self.model_cache[model_key] @@ -571,7 +431,8 @@ class MLFlowRepository(): # Retention is 0, download a new model if retention <= 0: model, _artifact_path = self.download_model( - model_name=model_name, model_type=model_type, flavor=flavor, load_wrapper=False) + model_name=model_name, model_type=model_type, flavor=flavor, load_wrapper=False + ) return model model_key = f'{model_name}_{model_type}' @@ -581,32 +442,49 @@ class MLFlowRepository(): # Check if config has changed or is outdated if self.check_cache_retention(cache, retention): - return self.handle_valid_model( - model_name=model_name, model_type=model_type, cache=cache) + return self.handle_valid_model(model_name=model_name, cache=cache) else: # Model is outdated, delete old model files - self.handle_outdated_model(model_name, model_key) + self.handle_outdated_model(model_name=model_name, model_key=model_key) else: - if self.logger: - self.logger.debug( - f"Model {model_name} is not in {model_type} cache, downloading a new one") + self.logger.debug( + f'Model {model_name} is not in {model_type} cache, downloading a new one' + ) # Donwload new model model, _artifact_path = self.download_model( - model_name=model_name, model_type=model_type, flavor=flavor, - load_wrapper=False) + model_name=model_name, model_type=model_type, flavor=flavor, load_wrapper=False + ) - cache = { - 'target': model, - 'timestamp': datetime.now() - } + cache = {'target': model, 'timestamp': datetime.now()} self.model_cache[model_key] = cache return model - def get_cached_transform(self, model_name: str, data: pd.DataFrame, - retention: int, flavor: str) -> pd.DataFrame: + @overload + def get_cached_operation( + self, + model_name: str, + data: pd.DataFrame, + operation: Literal['transform'], + retention: int, + flavor: str, + ) -> pd.DataFrame: ... + + @overload + def get_cached_operation( + self, + model_name: str, + data: pd.DataFrame, + operation: Literal['predict'], + retention: int, + flavor: str, + ) -> pd.DataFrame | ndarray: ... + + def get_cached_operation( + self, model_name: str, data: pd.DataFrame, operation: str, retention: int, flavor: str + ) -> pd.DataFrame | ndarray: """ Get transformed data using cached transform model. @@ -619,44 +497,17 @@ class MLFlowRepository(): Returns: pd.DataFrame: Transformed data """ + if operation not in ['transform', 'predict']: + raise ValueError("Invalid operation. Use 'transform' or 'predict'.") + model = self.get_model( - model_name=model_name, retention=retention, - model_type="transform", flavor=flavor) + model_name=model_name, retention=retention, model_type=operation, flavor=flavor + ) prediction = model.predict(data) if retention == 0: - self.logger.info( - f"Deleting model {model_name} from memory") - del model - - force_memory_release(self.logger) - - return prediction - - def get_cached_predict(self, model_name: str, data: pd.DataFrame, retention: int, - flavor: str) -> ndarray: - """ - Get predictions using cached prediction model. - - Args: - model_name (str): Name of the prediction model - data (pd.DataFrame): Data to make predictions on - retention (int): Cache retention time in minutes - flavor (str): Model flavor ('sklearn', 'pyfunc', 'pytorch') - - Returns: - pd.DataFrame: Model predictions - """ - model = self.get_model( - model_name=model_name, retention=retention, - model_type="predict", flavor=flavor) - - prediction = model.predict(data) - - if retention == 0: - self.logger.info( - f"Deleting model {model_name} from memory") + self.logger.info(f'Deleting model {model_name}:{operation} from memory') del model force_memory_release(self.logger) @@ -667,10 +518,16 @@ class MLFlowRepository(): Functions related to model retraining """ - def create_model_experiment(self, model_name: str, data: pd.DataFrame, latest_production_id: str, - transform_flavor: str = 'sklearn', predict_flavor: str = 'sklearn', - fit_config: dict = {}, target_name: str = None, - metadata: dict = {}) -> tuple: + def fit_models( + self, + model_name: str, + data: pd.DataFrame, + latest_production_id: str, + metadata: dict, + transform_flavor: str = 'sklearn', + predict_flavor: str = 'sklearn', + target_name: str | None = None, + ) -> dict[str, dict[str, Any]]: """ Create a new MLFlow experiment for model retraining. @@ -686,7 +543,6 @@ class MLFlowRepository(): data (pd.DataFrame): Training data for model retraining transform_flavor (str): Flavor for transformation model predict_flavor (str): Flavor for prediction model - fit_config (dict): Fit configuration target_name (str): Target name metadata (dict): Metadata for logging @@ -696,78 +552,90 @@ class MLFlowRepository(): - data_model: Fitted transformation model - experiment: MLFlow experiment name """ - self.logger.custom_info( - f"Starting model experiment creation for {model_name}", metadata) + + self.logger.custom_info(f'Starting model experiment creation for {model_name}', metadata) self.logger.custom_debug( - f"Model configuration - transform_flavor: {transform_flavor}, predict_flavor: {predict_flavor}, fit_config: {fit_config}, target_name: {target_name}", metadata) + f'Model configuration - transform_flavor: {transform_flavor}, predict_flavor: {predict_flavor}, target_name: {target_name}', + metadata, + ) # data.to_csv( # f"tmp/retrain_data_{model_name}.csv", index=True) self.logger.custom_info( - f"Retrieved latest production run ID: {latest_production_id}", metadata) - self.logger.custom_info( - f"Loading transformation model for {model_name}", metadata) + f'Retrieved latest production run ID: {latest_production_id}', metadata + ) + self.logger.custom_info(f'Loading transformation model for {model_name}', metadata) load_transform_wrapper = transform_flavor == 'pyfunc' data_model, data_artifact_path = self.download_model( - model_name=model_name, model_type="transform", flavor=transform_flavor, - load_wrapper=load_transform_wrapper + model_name=model_name, + model_type='transform', + flavor=transform_flavor, + load_wrapper=load_transform_wrapper, ) - self.logger.custom_info( - f"Loading prediction model for {model_name}", metadata) + self.logger.custom_info(f'Loading prediction model for {model_name}', metadata) load_predict_wrapper = predict_flavor == 'pyfunc' prediction_model, prediction_artifact_path = self.download_model( - model_name=model_name, model_type="predict", flavor=predict_flavor, - load_wrapper=load_predict_wrapper + model_name=model_name, + model_type='predict', + flavor=predict_flavor, + load_wrapper=load_predict_wrapper, ) - treated_data = data_model.fit(data) + treated_data_candidate = data_model.fit(data) + + if not isinstance(treated_data_candidate, pd.DataFrame): + data_model = treated_data_candidate + treated_data = data_model.predict(data) + else: + treated_data = treated_data_candidate # Stores current index as timestamp, Courier model expects a timestamp column # with specific format treated_data['timestamp'] = treated_data.index # Parses timestamp column to datetime format to align with data - treated_data = self.detect_and_parse_datetime_index( - treated_data, metadata) + treated_data = self.detect_and_parse_datetime_index(treated_data, metadata) - treated_data = treated_data.drop_duplicates( - subset=['timestamp'], keep='first') + treated_data = treated_data.drop_duplicates(subset=['timestamp'], keep='first') - self.logger.custom_debug( - f"Treated data index: {treated_data.index}", metadata) + self.logger.custom_debug(f'Treated data index: {treated_data.index}', metadata) # treated_data.to_csv( # f"tmp/retrain_treated_data_{model_name}.csv", index=True) - self.logger.custom_debug( - f"Transformed data shape: {treated_data.shape}", metadata) + self.logger.custom_debug(f'Transformed data shape: {treated_data.shape}', metadata) if target_name is None: target_name = data_model.target_variable self.logger.custom_debug( - f"Using target variable from data model: {target_name}", metadata) + f'Using target variable from data model: {target_name}', metadata + ) else: - self.logger.custom_debug( - f"Using provided target variable: {target_name}", metadata) + self.logger.custom_debug(f'Using provided target variable: {target_name}', metadata) # Check if treated_data contains target variable if target_name not in treated_data.columns: self.logger.custom_debug( - f"Target variable {target_name} not found in treated data, aligning data with treated data indexes", metadata) + f'Target variable {target_name} not found in treated data, aligning data with treated data indexes', + metadata, + ) # Aligns data with treated data indexes to get target variable aligned_data = data.loc[treated_data.index] + aligned_series = aligned_data[target_name] retrain_dataset = pd.merge( - treated_data, aligned_data, left_index=True, right_index=True) + treated_data, aligned_series, left_index=True, right_index=True + ) else: # Uses target variable from treated data self.logger.custom_debug( - f"Target variable {target_name} found in treated data, using it", metadata) + f'Target variable {target_name} found in treated data, using it', metadata + ) retrain_dataset = treated_data # retrain_dataset.to_csv( @@ -776,91 +644,47 @@ class MLFlowRepository(): prediction_model.fit(retrain_dataset) self.logger.custom_info( - f"Model experiment creation completed successfully for {model_name}", metadata) + f'Model experiment creation completed successfully for {model_name}', metadata + ) retrain_data = { 'prediction_model': { 'model': prediction_model, - 'artifact_path': prediction_artifact_path + 'artifact_path': prediction_artifact_path, }, - 'data_model': { - 'model': data_model, - 'artifact_path': data_artifact_path - } + 'data_model': {'model': data_model, 'artifact_path': data_artifact_path}, } return retrain_data - def export_model_to_pkl(self, model_data: dict, model_type: str, metadata: dict = {}): - artifact_local_path = model_data['artifact_path'] + def log_model(self, model_data: dict, flavor: str, model_type: str, metadata: dict): model = model_data['model'] - if artifact_local_path: - pickle_file = TRANSFORMED_COMPRESSED_PATH if model_type == "data_model" else PREDICTION_COMPRESSED_PATH - pickle_path = path.join(artifact_local_path, pickle_file) - - self.logger.custom_info( - f"Saving model to {pickle_path}", metadata) - - with open(pickle_path, "wb") as f: - pickle.dump(model, f) - - def export_model_to_lzma(self, model_data: dict, model_type: str, metadata: dict = {}): - artifact_local_path = model_data['artifact_path'] - model = model_data['model'] - - if artifact_local_path: - pickle_file = TRANSFORMED_COMPRESSED_PATH if model_type == "data_model" else PREDICTION_COMPRESSED_PATH - pickle_path = path.join(artifact_local_path, pickle_file) - - self.logger.custom_info( - f"Saving model to {pickle_path}", metadata) - - if path.exists(pickle_path): - remove(pickle_path) - - with lzma.open(pickle_path + ".xz", "wb") as f: - self.logger.custom_info( - "Compressing model", metadata) - pickle.dump(model, f) - - def log_model(self, model_data: dict, flavor: str, model_type: str, - metadata: dict = {}): - - model = model_data['model'] - - self.logger.custom_debug( - f"Logging {model_type} model to {model_type}", metadata) + self.logger.custom_debug(f'Logging {model_type} model to {model_type}', metadata) if flavor == 'sklearn': mlflow.sklearn.log_model(model, model_type) elif flavor == 'pyfunc': - code_path = [path.join( - model_data['artifact_path'], 'code', "utils")] + code_path = [path.join(model_data['artifact_path'], 'code', 'utils')] - self.logger.custom_debug( - f"Code path: {code_path}", metadata) + self.logger.custom_debug(f'Code path: {code_path}', metadata) - model.store_model( - artifact_path=model_type, - code_path=code_path, - to_disk=False - ) + model.store_model(artifact_path=model_type, code_path=code_path, to_disk=False) - self.logger.custom_debug( - f"Model uploaded successfully", metadata) + self.logger.custom_debug('Model uploaded successfully', metadata) elif flavor == 'pytorch': mlflow.pytorch.log_model(model, model_type) else: - raise ValueError( - "Invalid flavor. Use 'sklearn' or 'pyfunc' or 'pytorch'.") + raise ValueError("Invalid flavor. Use 'sklearn' or 'pyfunc' or 'pytorch'.") - def perform_model_retrain(self, - model_name: str, - data: pd.DataFrame, - retrain_data: dict, - transform_flavor: str = 'sklearn', - predict_flavor: str = 'sklearn', - metadata: dict = {}, - latest_production_id: str = None) -> dict: + def create_new_experiment( + self, + model_name: str, + data: pd.DataFrame, + retrain_data: dict, + latest_production_id: str, + metadata: dict, + transform_flavor: str = 'sklearn', + predict_flavor: str = 'sklearn', + ) -> dict: """ Execute the complete model retraining process in MLFlow. @@ -890,61 +714,53 @@ class MLFlowRepository(): prediction_model = retrain_data['prediction_model'] data_model = retrain_data['data_model'] - model_temp_path = path.join( - ARTIFACTS_PATH, model_name) + model_temp_path = path.join(ARTIFACTS_PATH, model_name) - self.logger.custom_info( - f"Starting model retraining process for {model_name}", metadata) + self.logger.custom_info(f'Starting model retraining process for {model_name}', metadata) original_params = self.get_model_params(latest_production_id) retrain_params = { **original_params, - "retrain": True, + 'retrain': True, 'retrain_date': datetime.now().isoformat(), 'source_run_id': latest_production_id, 'retrain_samples': str(data.shape), } - experiment_description = f"Retrain model {model_name} with new data" + experiment_description = f'Retrain model {model_name} with new data' - experiment = self.get_experiment(model_name, - create_if_not_exists=True) + experiment = self.get_experiment(model_name, create_if_not_exists=True) experiment_name = experiment.name current_run_name = self.get_next_run_name(experiment_name) - self.logger.custom_debug( - f"Attributes: {retrain_params}", metadata) + self.logger.custom_debug(f'Attributes: {retrain_params}', metadata) - data_path = f"{model_temp_path}/retrain_data.csv" + data_path = f'{model_temp_path}/retrain_data.csv' data.to_csv(data_path, index=True) self.logger.custom_info( - f"Starting model upload for {experiment_name} with run name {current_run_name}", metadata) + f'Starting model upload for {experiment_name} with run name {current_run_name}', + metadata, + ) with mlflow.start_run( experiment_id=experiment.experiment_id, run_name=current_run_name, - description=experiment_description + description=experiment_description, ) as _run: run_id = _run.info.run_id - self.logger.custom_info( - f"Logging data model", metadata) + self.logger.custom_info('Logging data model', metadata) # dynamic parameters, including model itself - self.log_model(data_model, transform_flavor, - "data_model", metadata) + self.log_model(data_model, transform_flavor, 'data_model', metadata) # dynamic parameters, including model itself - self.logger.custom_info( - f"Logging prediction model", metadata) - self.log_model(prediction_model, predict_flavor, - "prediction_model", metadata) + self.logger.custom_info('Logging prediction model', metadata) + self.log_model(prediction_model, predict_flavor, 'prediction_model', metadata) - self.logger.custom_info( - f"Model logged successfully for {model_name}", metadata) + self.logger.custom_info(f'Model logged successfully for {model_name}', metadata) - self.logger.custom_info( - f"Logging remaining parameters for {model_name}", metadata) + self.logger.custom_info(f'Logging remaining parameters for {model_name}', metadata) # update transfomation model # fixed parameters @@ -953,28 +769,29 @@ class MLFlowRepository(): # log the data raw mlflow.log_artifact(data_path) - self.logger.custom_info( - f"Deleting model from filesystem", metadata) + self.logger.custom_info('Deleting model from filesystem', metadata) if path.exists(model_temp_path): rmtree(model_temp_path) - self.logger.custom_info( - f"Deleting prediction model from memory", metadata) + self.logger.custom_info('Deleting prediction model from memory', metadata) del prediction_model['model'] del prediction_model - self.logger.custom_info( - f"Deleting data model from memory", metadata) + self.logger.custom_info('Deleting data model from memory', metadata) del data_model['model'] del data_model + force_memory_release(self.logger) + return { 'run_id': run_id, 'experiment_id': experiment.experiment_id, - 'experiment_name': experiment.name + 'experiment_name': experiment.name, } - def update_production_model_by_run_id(self, run_id: str, model_name: str, metadata: dict = {}) -> dict: + def update_production_model_by_run_id( + self, run_id: str, model_name: str, metadata: dict + ) -> dict: """ Update production model with a specific MLFlow run. @@ -999,18 +816,18 @@ class MLFlowRepository(): 3. Transitions the model to 'Production' stage 4. Archives existing production versions """ + self.logger.custom_info( - f"Starting production model update for {model_name} with run ID: {run_id}", metadata) + f'Starting production model update for {model_name} with run ID: {run_id}', metadata + ) # Registrar o modelo # Aqui estamos assumindo que você já tem um modelo salvo, caso contrário você precisará treiná-lo e salvá-lo primeiro. # Se o modelo já está registrado, você pode usar o método register_model() ou pyfunc.load_model() para isso. - mlflow.register_model( - f"runs:/{run_id}/prediction_model", model_name) + mlflow.register_model(f'runs:/{run_id}/prediction_model', model_name) # Obter a versão mais recente registrada do modelo - model_versions = self.client.get_registered_model( - model_name).latest_versions + model_versions = self.client.get_registered_model(model_name).latest_versions if not isinstance(model_versions, list): raise ValueError('Model versions is not a list') @@ -1019,24 +836,16 @@ class MLFlowRepository(): # Mover a versão mais recente do modelo para o estágio de 'Production' self.client.transition_model_version_stage( - name=model_name, - version=max_version, - stage="Production", - archive_existing_versions=True + name=model_name, version=max_version, stage='Production', archive_existing_versions=True ) - return { - 'model_name': model_name, - 'version': max_version, - 'mlflow_run_id': run_id - } + return {'model_name': model_name, 'version': max_version, 'mlflow_run_id': run_id} """ Functions that provide the interface to model operations """ - def transform(self, model_name: str, data: pd.DataFrame, - model_config: dict, metadata: dict): + def transform(self, model_name: str, data: pd.DataFrame, model_config: dict, metadata: dict): """ Transform data using a cached transformation model. @@ -1069,8 +878,10 @@ class MLFlowRepository(): Exception: Any exception during model loading or transformation is caught and returned in the response structure rather than propagated. """ + self.logger.custom_debug( - f"Data received for model transformation: {data.head(5).to_csv()}", metadata) + f'Data received for model transformation: {data.head(5).to_csv()}', metadata + ) # data.to_csv( # f"tmp/data_{model_name}.csv", index=True) @@ -1079,35 +890,29 @@ class MLFlowRepository(): flavor = model_config.get('transform_flavor', 'sklearn') try: - transformed_data = self.get_cached_transform( - model_name, data, model_retention, flavor + transformed_data: pd.DataFrame = self.get_cached_operation( + model_name, data, 'transform', model_retention, flavor ) self.logger.custom_debug( - f"Data received from model transformation: {transformed_data.head(5).to_csv()}", metadata) + f'Data received from model transformation: {transformed_data.head(5).to_csv()}', + metadata, + ) # transformed_data.to_csv( # f"tmp/transformed_data_{model_name}.csv", index=True) - transformed_data = self.detect_and_parse_datetime_index( - transformed_data, metadata) + transformed_data = self.detect_and_parse_datetime_index(transformed_data, metadata) - return { - 'success': True, - 'content': transformed_data.to_dict() - } + return {'success': True, 'content': transformed_data.to_dict()} except Exception as e: return { 'success': False, - 'content': { - 'message': str(e), - 'traceback': traceback.format_exc() - } + 'content': {'message': str(e), 'traceback': traceback.format_exc()}, } - def predict(self, model_name: str, data: pd.DataFrame, - model_config: dict, metadata: dict): + def predict(self, model_name: str, data: pd.DataFrame, model_config: dict, metadata: dict): """ Generate predictions using a cached prediction model. @@ -1152,55 +957,50 @@ class MLFlowRepository(): flavor = model_config.get('predict_flavor', 'sklearn') try: - input_index = data.index start_time = datetime.now() self.logger.custom_debug( - f"Data received for model prediction: {data.head(5).to_csv()}", metadata) + f'Data received for model prediction: {data.head(5).to_csv()}', metadata + ) # data.to_csv( # f"tmp/treated_data_{model_name}.csv", index=True) - predict_data = self.get_cached_predict( - model_name, data, model_retention, flavor) + predict_data = self.get_cached_operation( + model_name, data, 'predict', model_retention, flavor + ) end_time = datetime.now() if isinstance(predict_data, pd.DataFrame): self.logger.custom_debug( - f"Data received from model prediction: {data.head(5).to_csv()}", metadata) + f'Data received from model prediction: {data.head(5).to_csv()}', metadata + ) # predict_data.to_csv( # f"tmp/predicted_data_{model_name}.csv", index=True) - data.columns = ['prediction'] + predict_data.columns = ['prediction'] else: - predict_data = pd.DataFrame( - predict_data, columns=['prediction']) + predict_data = pd.DataFrame(predict_data, columns=['prediction']) # predict_data.to_csv( # f"tmp/predicted_data_{model_name}.csv", index=True) predict_data.index = input_index - predict_data['response_time'] = ( - end_time - start_time).total_seconds() + predict_data['response_time'] = (end_time - start_time).total_seconds() - return { - 'success': True, - 'content': predict_data.to_dict() - } + return {'success': True, 'content': predict_data.to_dict()} except Exception as e: return { 'success': False, - 'content': { - 'message': str(e), - 'traceback': traceback.format_exc() - } + 'content': {'message': str(e), 'traceback': traceback.format_exc()}, } - def retrain_model(self, data: pd.DataFrame, model_name: str, - model_config: dict, metadata: dict) -> tuple: + def retrain_model( + self, data: pd.DataFrame, model_name: str, model_config: dict, metadata: dict + ) -> dict[str, Any]: """ Orchestrate the complete model retraining workflow. @@ -1245,49 +1045,53 @@ class MLFlowRepository(): Exception: Any other exception during the retraining process """ - self.logger.custom_info( - f"Starting model retraining workflow for {model_name}", metadata) - self.logger.custom_debug( - f"Data received for model retraining: {data.to_csv()}", metadata) + self.logger.custom_info(f'Starting model retraining workflow for {model_name}', metadata) + self.logger.custom_debug(f'Data received for model retraining: {data.to_csv()}', metadata) target_name = model_config.get('target', None) transform_flavor = model_config.get('transform_flavor', 'sklearn') predict_flavor = model_config.get('predict_flavor', 'sklearn') - fit_config = { - 'split_fit_data': model_config.get('split_fit_data', False), - 'split_fit_first': model_config.get('split_fit_first', 'x').lower(), - 'y_type': model_config.get('y_type', 'series').lower() - } - self.logger.custom_debug( - f"Model configuration - transform_flavor: {transform_flavor}, predict_flavor: {predict_flavor}, fit_config: {fit_config}, target_name: {target_name}", metadata) + f'Model configuration - transform_flavor: {transform_flavor}, predict_flavor: {predict_flavor}, target_name: {target_name}', + metadata, + ) try: - latest_production_id = self.get_model_run_id( - model_name, stage="Production" + latest_production_id = self.get_model_run_id(model_name, stage='Production') + self.logger.custom_info('Creating model experiment environment', metadata) + retrain_data = self.fit_models( + model_name=model_name, + data=data, + transform_flavor=transform_flavor, + predict_flavor=predict_flavor, + target_name=target_name, + metadata=metadata, + latest_production_id=latest_production_id, ) self.logger.custom_info( - "Creating model experiment environment", metadata) - retrain_data = self.create_model_experiment( - model_name=model_name, data=data, transform_flavor=transform_flavor, - predict_flavor=predict_flavor, fit_config=fit_config, target_name=target_name, - metadata=metadata, latest_production_id=latest_production_id) - self.logger.custom_info( - f"Model experiment created successfully: {retrain_data}", metadata) + f'Model experiment created successfully: {retrain_data}', metadata + ) - self.logger.custom_info("Saving model retrain", metadata) - experiment = self.perform_model_retrain( - model_name=model_name, data=data, retrain_data=retrain_data, - transform_flavor=transform_flavor, predict_flavor=predict_flavor, metadata=metadata, latest_production_id=latest_production_id) + self.logger.custom_info('Saving model retrain', metadata) + experiment = self.create_new_experiment( + model_name=model_name, + data=data, + retrain_data=retrain_data, + transform_flavor=transform_flavor, + predict_flavor=predict_flavor, + metadata=metadata, + latest_production_id=latest_production_id, + ) self.logger.custom_info( - f"Model retraining completed successfully for experiment: {experiment}", metadata) + f'Model retraining completed successfully for experiment: {experiment}', metadata + ) return { 'success': True, 'experiment': experiment, - 'message': 'Model retrained successfully.' + 'message': 'Model retrained successfully.', } except Exception as e: error_msg = f'Error retraining model {model_name}: {e}' @@ -1296,10 +1100,12 @@ class MLFlowRepository(): 'success': False, 'experiment': None, 'message': error_msg, - 'traceback': traceback.format_exc() + 'traceback': traceback.format_exc(), } - def update_production_model(self, experiment: str, model_name: str, metadata: dict = {}) -> dict: + def update_production_model( + self, experiment: dict[str, Any], model_name: str, metadata: dict + ) -> dict: """ Update production model using the latest retraining run. @@ -1346,8 +1152,7 @@ class MLFlowRepository(): """ run_id = experiment['run_id'] experiment_id = experiment['experiment_id'] - metadata_result = self.update_production_model_by_run_id( - run_id, model_name, metadata) + metadata_result = self.update_production_model_by_run_id(run_id, model_name, metadata) metadata_result['mlflow_experiment_id'] = experiment_id diff --git a/laborious/utils/repository/opc_repository.py b/laborious/utils/repository/opc_repository.py index 96ecfef..072274a 100644 --- a/laborious/utils/repository/opc_repository.py +++ b/laborious/utils/repository/opc_repository.py @@ -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, - 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): +class OpcRepository: + def __init__( + self, + opc_id: str, + url: str, + logger: Logger, + notification_handler: NotificationHandler, + 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 diff --git a/laborious/worker/worker.py b/laborious/worker/worker.py index 4658524..781ad55 100644 --- a/laborious/worker/worker.py +++ b/laborious/worker/worker.py @@ -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) diff --git a/laborious/workflows/minimal_retrain.py b/laborious/workflows/minimal_retrain.py index 71f72ee..5497432 100644 --- a/laborious/workflows/minimal_retrain.py +++ b/laborious/workflows/minimal_retrain.py @@ -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), ) diff --git a/laborious/workflows/predictions_batch.py b/laborious/workflows/predictions_batch.py index e522eaa..f79206f 100644 --- a/laborious/workflows/predictions_batch.py +++ b/laborious/workflows/predictions_batch.py @@ -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) diff --git a/laborious/workflows/sub_workflows/format_and_export_prediction.py b/laborious/workflows/sub_workflows/format_and_export_prediction.py index 8e7df07..72ff12d 100644 --- a/laborious/workflows/sub_workflows/format_and_export_prediction.py +++ b/laborious/workflows/sub_workflows/format_and_export_prediction.py @@ -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), ) diff --git a/laborious/workflows/sub_workflows/prediction_process.py b/laborious/workflows/sub_workflows/prediction_process.py index 777fa1c..7bd36e2 100644 --- a/laborious/workflows/sub_workflows/prediction_process.py +++ b/laborious/workflows/sub_workflows/prediction_process.py @@ -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 diff --git a/mlruns/0/meta.yaml b/mlruns/0/meta.yaml new file mode 100644 index 0000000..2a5def3 --- /dev/null +++ b/mlruns/0/meta.yaml @@ -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 diff --git a/mlruns/586524947870967910/meta.yaml b/mlruns/586524947870967910/meta.yaml new file mode 100644 index 0000000..7221134 --- /dev/null +++ b/mlruns/586524947870967910/meta.yaml @@ -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 diff --git a/mlruns/models/test/meta.yaml b/mlruns/models/test/meta.yaml new file mode 100644 index 0000000..9d225cc --- /dev/null +++ b/mlruns/models/test/meta.yaml @@ -0,0 +1,5 @@ +aliases: {} +creation_timestamp: 1760447068191 +description: null +last_updated_timestamp: 1760447068191 +name: test diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..5b153a3 --- /dev/null +++ b/pyproject.toml @@ -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 \ No newline at end of file diff --git a/requirements-dev.txt b/requirements-dev.txt new file mode 100644 index 0000000..56ab376 --- /dev/null +++ b/requirements-dev.txt @@ -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 \ No newline at end of file diff --git a/tests/laborious/activities/test_activities.py b/tests/laborious/activities/test_activities.py index de418be..1999e91 100644 --- a/tests/laborious/activities/test_activities.py +++ b/tests/laborious/activities/test_activities.py @@ -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() diff --git a/tests/laborious/activities/test_gates.py b/tests/laborious/activities/test_gates.py index a65c6d4..1a376e2 100644 --- a/tests/laborious/activities/test_gates.py +++ b/tests/laborious/activities/test_gates.py @@ -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], - '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']}, + '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'], + }, '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 diff --git a/tests/laborious/activities/test_mlflow.py b/tests/laborious/activities/test_mlflow.py index 85cb7f1..cfb66a1 100644 --- a/tests/laborious/activities/test_mlflow.py +++ b/tests/laborious/activities/test_mlflow.py @@ -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( + { + **metadata, + 'object_key': 'test_object_key', + 'model_name': 'test_model', + 'model_config': { + 'target': 'target', + 'transform_flavor': 'sklearn', + 'predict_flavor': 'pyfunc', + }, + } + ) - response = await mlflow.retrain_model({ - **metadata, - 'data': data, - 'model_name': 'test_model' - }) + raw_data = mlflow.minio_repository.get_parquet_as_dataframe.return_value - mlflow.model_monitoring_repository.retrain_model.assert_called_once() + 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' - mlflow.send_notification.assert_called_once_with( - metadata=metadata['metadata'], - notification_id='RETRAIN_MODEL_ERROR', - message='Error retraining model test_model: Error retraining model', - block='retrain_model', - level=NotificationLevel.ERROR, - attachment_content=ANY - ) - else: - assert False, "No exception raised" + '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: Model retrained failed.', + block='retrain_model', + level=NotificationLevel.ERROR, + attachment_content=ANY, + ) + + 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') diff --git a/tests/laborious/activities/test_opc.py b/tests/laborious/activities/test_opc.py index 07f00b8..815df27 100644 --- a/tests/laborious/activities/test_opc.py +++ b/tests/laborious/activities/test_opc.py @@ -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, { - '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, {})) + 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, {})), ) 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([ - call( - id="server1", - url="http://localhost:8080", - logger=mock_logger, - 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' - ), - ]) - mock_opc_repository.assert_has_calls([ - call( - id="server2", - url="http://localhost:8080", - logger=mock_logger, - 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' - ) - ]) + mock_opc_repository.assert_has_calls( + [ + call( + opc_id='server1', + url='http://localhost:8080', + logger=mock_logger, + 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', + ), + ] + ) + mock_opc_repository.assert_has_calls( + [ + call( + opc_id='server2', + url='http://localhost:8080', + logger=mock_logger, + 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', + ) + ] + ) server1.connect.assert_called_once() server2.connect.assert_called_once() - mock_send_notification.assert_has_calls([ - call( - metadata={ - 'model_id': '-', - 'model_name': '-', - 'workflow_name': '-', - 'schedule_name': 'INITIALIZATION' - }, - notification_id="OPC_CONNECTION_ERROR_server3", - message="Failed to connect to OPC server: Test error", - block="opc_repository", - level=NotificationLevel.ERROR, - attachment_content=ANY - ) - ]) + mock_send_notification.assert_has_calls( + [ + call( + metadata={ + 'model_id': '-', + 'model_name': '-', + 'workflow_name': '-', + 'schedule_name': 'INITIALIZATION', + }, + notification_id='OPC_CONNECTION_ERROR_server3', + message='Failed to connect to OPC server: Test error', + block='opc_repository', + level=NotificationLevel.ERROR, + 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, { - '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' - }) + 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', + }, + ) - 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'} - }, - 'confidence_tags': { - 'tag2': {'data_type': 'float'} - } + 'prediction_tags': {'tag1': {'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([ - call( - server_id='server1', - tag='tag1', - data=0.75, - data_type='float', - tag_type='prediction', - 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'] - ) - ]) + 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( + [ + call( + server_id='server1', + tag='tag2', + data=0.95, + data_type='float', + tag_type='confidence', + 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'} - }, - 'confidence_tags': { - 'tag2': {'data_type': 'float'} - } + 'prediction_tags': {'tag1': {'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', [ - (DataFrame({'prediction_confidence': [0]}), True, 0), - (DataFrame({'prediction_confidence': [0]}), False, 12), -]) +@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) diff --git a/tests/laborious/activities/test_storage.py b/tests/laborious/activities/test_storage.py new file mode 100644 index 0000000..87c28f0 --- /dev/null +++ b/tests/laborious/activities/test_storage.py @@ -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() diff --git a/tests/laborious/utils/filters/test_conditional_filters.py b/tests/laborious/utils/filters/test_conditional_filters.py index 405bc9b..de25b7b 100644 --- a/tests/laborious/utils/filters/test_conditional_filters.py +++ b/tests/laborious/utils/filters/test_conditional_filters.py @@ -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 + ) diff --git a/tests/laborious/utils/filters/test_mlflow_filters.py b/tests/laborious/utils/filters/test_mlflow_filters.py index f9c61e9..5353a6f 100644 --- a/tests/laborious/utils/filters/test_mlflow_filters.py +++ b/tests/laborious/utils/filters/test_mlflow_filters.py @@ -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 diff --git a/tests/laborious/utils/repository/test_minio_repository.py b/tests/laborious/utils/repository/test_minio_repository.py new file mode 100644 index 0000000..ec279d2 --- /dev/null +++ b/tests/laborious/utils/repository/test_minio_repository.py @@ -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 diff --git a/tests/laborious/utils/repository/test_model_repository.py b/tests/laborious/utils/repository/test_model_repository.py index cdc59b4..0e77bdf 100644 --- a/tests/laborious/utils/repository/test_model_repository.py +++ b/tests/laborious/utils/repository/test_model_repository.py @@ -1,34 +1,60 @@ +from datetime import UTC, datetime from unittest.mock import ANY, MagicMock, call, patch + +import mlflow as mlflow_lib import numpy as np -from pandas import DataFrame import pytest -from datetime import datetime, timezone -from pandas import Timestamp -from laborious.utils.repository.model_repository import MLFlowRepository +from pandas import DataFrame, Timestamp + +from laborious.utils.repository.model_repository import MLFlowRepository, force_memory_release + + +@patch('laborious.utils.repository.model_repository.ctypes') +@patch('laborious.utils.repository.model_repository.gc') +def test_force_memory_release_success(gc, ctypes): + logger = MagicMock() + + force_memory_release(logger) + + gc.collect.assert_called_once() + ctypes.CDLL.return_value.malloc_trim.assert_called_once_with(0) + logger.info.assert_called_once_with('Memory released') + + +@patch('laborious.utils.repository.model_repository.ctypes') +@patch('laborious.utils.repository.model_repository.gc') +def test_force_memory_release_error(gc, ctypes): + logger = MagicMock() + ctypes.CDLL.return_value.malloc_trim.side_effect = Exception('error') + force_memory_release(logger) + + gc.collect.assert_called_once() + ctypes.CDLL.return_value.malloc_trim.assert_called_once_with(0) + + logger.info.assert_called_once_with('Memory release failed: error') @pytest.fixture def mlflow_repository(): - with patch('laborious.utils.repository.model_repository.ModelServing', - autospec=True) as mock_model_serving: - mock_instance = mock_model_serving.return_value - mock_instance.get_transformed_data = MagicMock() - + with patch('laborious.utils.repository.model_repository.mlflow'): repo = MLFlowRepository( - host='http://localhost:5000', - username='admin', - password='admin', - logger=MagicMock() + host='http://localhost:5000', username='admin', password='admin', logger=MagicMock() ) return repo +@pytest.fixture +def mlflow(): + with patch('laborious.utils.repository.model_repository.mlflow') as mlflow: + yield mlflow + + 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', }, } @@ -37,203 +63,75 @@ class Any: pass -invalid_cases = [ - ( - { - 'value': { - '2024-01-01 12:00:00': 1, - 2024: 2 - } - } - ), - ( - { - 'value': { - '2024-01-01': 1, - '2024-01-02': 2 - } - } - ), - ( - { - 'value': { - Any(): 1, - Any(): 2 - } - } - ) -] - - -@pytest.mark.parametrize("data", invalid_cases) -def test_detect_and_parse_datetime_index_error_cases(mlflow_repository, data): - input_data = DataFrame( - data - ) - - with pytest.raises(ValueError) as e: - mlflow_repository.detect_and_parse_datetime_index( - input_data, metadata['metadata']) - - assert str(e) == "Index must be all timestamp like column. Valid formats are: pandas Timestamp, datetime, string in format %Y-%m-%d %H:%M:%S" - - -valid_cases = [ - ( - { - 'value': { - '2024-01-01 12:00:00+0000': 1, - '2024-01-02 12:00:00+0000': 2 - } - }, ['2024-01-01 12:00:00+0000', '2024-01-02 12:00:00+0000'] - ), - ( - { - 'value': { - datetime(2025, 1, 1, 12, 0, 0, tzinfo=timezone.utc): 1, - datetime(2025, 1, 2, 12, 0, 0, tzinfo=timezone.utc): 2 - } - }, ['2025-01-01 12:00:00+0000', '2025-01-02 12:00:00+0000'] - ), - ( - { - 'value': { - Timestamp(2026, 1, 1, 12, 0, 0, tzinfo=timezone.utc): 1, - Timestamp(2026, 1, 2, 12, 0, 0, tzinfo=timezone.utc): 2 - } - }, ['2026-01-01 12:00:00+0000', '2026-01-02 12:00:00+0000'] - ), -] - - -@pytest.mark.parametrize("data,expected", valid_cases) -def test_detect_and_parse_datetime_index_valid_format(mlflow_repository, data, expected): - input_data = DataFrame(data) - - response = mlflow_repository.detect_and_parse_datetime_index( - input_data, metadata['metadata']) - - assert response.index.tolist() == expected - - -def test_transform_success(mlflow_repository): - data = MagicMock() - model_name = 'model' - - mlflow_repository.detect_and_parse_datetime_index = MagicMock() - - output = mlflow_repository.transform( - model_name, data, {}, metadata['metadata']) - - mlflow_repository.model_serving.get_cached_transform.assert_called_once_with( - model_name, data, 0, 'sklearn', False, 'model', 'predict') - - mlflow_repository.detect_and_parse_datetime_index.assert_called_once_with( - mlflow_repository.model_serving.get_cached_transform.return_value, metadata['metadata']) - - assert output == { - 'success': True, - 'content': mlflow_repository.detect_and_parse_datetime_index.return_value.to_dict.return_value - } - - -def test_transform_error(mlflow_repository): - data = MagicMock() - model_name = 'model' - - mlflow_repository.model_serving.get_cached_transform.side_effect = Exception( - 'error') - - output = mlflow_repository.transform( - model_name, data, {}, metadata['metadata']) - - mlflow_repository.model_serving.get_cached_transform.assert_called_once_with( - model_name, data, 0, 'sklearn', False, 'model', 'predict') - - assert output == { - 'success': False, - 'content': { - 'message': 'error', - 'traceback': ANY - } - } - - -def test_predict_success(mlflow_repository): - data = DataFrame({ - 'feat_1': { - 'index_1': 2, - 'index_2': 3 - } - }) - model_name = 'model' - mlflow_repository.model_serving.get_cached_predict.return_value = np.array( - [2, 3] - ) - - output = mlflow_repository.predict( - model_name, data, {}, metadata['metadata']) - - mlflow_repository.model_serving.get_cached_predict.assert_called_once_with( - model_name, data, 0, 'pyfunc', False, 'model') - - assert output['success'] is True - assert output['content'] == { - 'prediction': { - 'index_1': 2, - 'index_2': 3 - }, 'response_time': { - 'index_1': ANY, - 'index_2': ANY - } - } - - -def test_predict_error(mlflow_repository): - data = DataFrame({ - 'feat_1': { - 'index_1': 2, - 'index_2': 3 - } - }) - model_name = 'model' - - mlflow_repository.model_serving.get_cached_predict = MagicMock( - side_effect=Exception('error') - ) - - output = mlflow_repository.predict( - model_name, data, {}, metadata['metadata']) - - mlflow_repository.model_serving.get_cached_predict.assert_called_once_with( - model_name, data, 0, 'pyfunc', False, 'model') - - assert output == { - 'success': False, - 'content': { - 'message': 'error', - 'traceback': ANY - } - } - - -@patch('laborious.utils.repository.model_repository.mlflow') -def test_get_experiment_by_run_id(mlflow, mlflow_repository): - mlflow.get_run.return_value = MagicMock( - info=MagicMock( - experiment_id='0', - ) - ) - mlflow.get_experiment.return_value = MagicMock() - mlflow.get_experiment.return_value.name = 'test' - - output = mlflow_repository.get_experiment_by_run_id('0') - assert output == 'test' +def test_get_model_uri_prediction(mlflow, mlflow_repository): + mlflow.get_run.return_value = MagicMock(info=MagicMock(artifact_uri='test')) + output = mlflow_repository.get_model_uri('0', prediction=True) + assert output == 'test/prediction_model' mlflow.get_run.assert_called_once_with('0') - mlflow.get_experiment.assert_called_once_with('0') -@patch('laborious.utils.repository.model_repository.mlflow') +def test_get_model_uri_transform(mlflow, mlflow_repository): + mlflow.get_run.return_value = MagicMock(info=MagicMock(artifact_uri='test')) + output = mlflow_repository.get_model_uri('0', prediction=False) + assert output == 'test/data_model' + mlflow.get_run.assert_called_once_with('0') + + +def test_get_model_run_id_not_registered_models(mlflow_repository): + mlflow_repository.client.search_registered_models.return_value = [] + with pytest.raises(mlflow_lib.exceptions.MlflowException) as e: + mlflow_repository.get_model_run_id('test') + + mlflow_repository.client.search_registered_models.assert_called_once_with( + filter_string="name='test'" + ) + + assert str(e.value) == "Model 'test' not found in the Model Registry." + + +def test_get_model_run_id_not_stage_versions(mlflow_repository): + mlflow_repository.client.search_registered_models.return_value = [MagicMock(name='test')] + + mlflow_repository.client.search_model_versions.return_value = [ + MagicMock(current_stage='Staging'), + MagicMock(current_stage='Staging'), + MagicMock(current_stage='Archived'), + ] + + with pytest.raises(mlflow_lib.exceptions.MlflowException) as e: + mlflow_repository.get_model_run_id('test') + + mlflow_repository.client.search_registered_models.assert_called_once_with( + filter_string="name='test'" + ) + mlflow_repository.client.search_model_versions.assert_called_once_with( + filter_string="name='test'" + ) + + assert str(e.value) == "Model 'test' in stage 'Production' not found in the Model Registry." + + +def test_get_model_run_id_success(mlflow_repository): + mlflow_repository.client.search_registered_models.return_value = [MagicMock(name='test')] + + mlflow_repository.client.search_model_versions.return_value = [ + MagicMock(current_stage='Production', version='1'), + MagicMock(current_stage='Production', version='2', source='runs/test/1'), + MagicMock(current_stage='Archived', version='3'), + ] + + output = mlflow_repository.get_model_run_id('test') + + mlflow_repository.client.search_registered_models.assert_called_once_with( + filter_string="name='test'" + ) + mlflow_repository.client.search_model_versions.assert_called_once_with( + filter_string="name='test'" + ) + + assert output == '1' + + def test_get_next_run_name(mlflow, mlflow_repository): mlflow.search_runs.return_value = [1, 2, 3] output = mlflow_repository.get_next_run_name('run') @@ -244,17 +142,66 @@ def test_get_next_run_name(mlflow, mlflow_repository): ) -@patch('laborious.utils.repository.model_repository.mlflow') -def test_get_experiment_success(mlflow, mlflow_repository): - mlflow.get_experiment_by_name.return_value = MagicMock( - experiment_id='0') +def test_get_experiment_experiment_exists(mlflow, mlflow_repository): + experiment = MagicMock(experiment_id='0') + + mlflow.get_experiment_by_name.return_value = experiment output = mlflow_repository.get_experiment('test') - assert output == 0 + assert output == experiment + + +def test_get_experiment_none_create(mlflow, mlflow_repository): + experiment = MagicMock(experiment_id='0') + + mlflow.get_experiment_by_name.return_value = None + + mlflow.create_experiment.return_value = experiment + + output = mlflow_repository.get_experiment('test', create_if_not_exists=True) + + assert output == experiment + + +def test_get_experiment_none_not_create(mlflow, mlflow_repository): + mlflow.get_experiment_by_name.return_value = None + + with pytest.raises(ValueError) as e: + mlflow_repository.get_experiment('test', create_if_not_exists=False) + + assert str(e) == 'Experiment test not found' + + +@patch('laborious.utils.repository.model_repository.path') +@patch('laborious.utils.repository.model_repository.rmtree') +@patch('laborious.utils.repository.model_repository.makedirs') +def test_download_artifacts_success(makedirs, rmtree, path, mlflow_repository): + mlflow_repository.get_model_run_id = MagicMock(return_value='test') + + path.exists.return_value = True + + output = mlflow_repository.dowload_artifacts('test', 'path') + + mlflow_repository.get_model_run_id.assert_called_once_with( + model_name='test', stage='Production' + ) + + path.join.assert_called_once_with('./tmp/artifacts/test', 'path') + + path.exists.assert_called_once_with(path.join.return_value) + + rmtree.assert_called_once_with(path.join.return_value) + + makedirs.assert_called_once_with('./tmp/artifacts/test', exist_ok=True) + + mlflow_repository.client.download_artifacts.assert_called_once_with( + mlflow_repository.get_model_run_id.return_value, 'path', './tmp/artifacts/test' + ) + + assert output == mlflow_repository.client.download_artifacts.return_value -@patch('laborious.utils.repository.model_repository.mlflow') def test_get_experiment_error(mlflow, mlflow_repository): mlflow.get_experiment_by_name.return_value = None @@ -263,188 +210,611 @@ def test_get_experiment_error(mlflow, mlflow_repository): except ValueError as e: assert str(e) == 'Experiment test not found' else: - assert False + raise AssertionError('Expected ValueError') -@patch('laborious.utils.repository.model_repository.mlflow') -def test_get_experiment_last_run(mlflow, mlflow_repository): - mlflow.search_runs.return_value = DataFrame({ - 'params.retrain': ['True', 'False', 'True', 'False'], - 'end_time': ['2021-01-01', '2021-01-02', '2021-01-03', '2021-01-04'], - 'run_id': ['0', '1', '2', '3'], - }) +def test_load_predict_model_sklearn(mlflow, mlflow_repository): + result = mlflow_repository.load_predict_model('test_model', 'sklearn') - output = mlflow_repository.get_experiment_last_run(0) + assert result == mlflow.sklearn.load_model.return_value + mlflow.sklearn.load_model.assert_called_once_with('models:/test_model/production') - mlflow.search_runs.assert_called_once_with( - experiment_ids=[0], - filter_string="", - output_format="pandas", + +def test_load_predict_model_pyfunc(mlflow, mlflow_repository): + result = mlflow_repository.load_predict_model('test_model', 'pyfunc') + assert result == mlflow.pyfunc.load_model.return_value + mlflow.pyfunc.load_model.assert_called_once_with('models:/test_model/production') + + +def test_load_predict_model_pytorch(mlflow, mlflow_repository): + result = mlflow_repository.load_predict_model('test_model', 'pytorch') + assert result == mlflow.pytorch.load_model.return_value + mlflow.pytorch.load_model.assert_called_once_with('models:/test_model/production') + + +def test_load_predict_model_error(mlflow_repository): + with pytest.raises(ValueError) as e: + mlflow_repository.load_predict_model('test_model', 'invalid') + assert str(e) == "Invalid flavor. Use 'sklearn' or 'pyfunc' or 'pytorch'." + + +def validate_common_load_transform_model_mocks(mlflow_repository, model_name): + mlflow_repository.get_model_run_id.assert_called_once_with( + model_name=model_name, stage='Production' + ) + mlflow_repository.get_model_uri.assert_called_once_with( + mlflow_repository.get_model_run_id.return_value, prediction=False ) - assert output == '2' + +def test_load_transform_model_sklearn(mlflow, mlflow_repository): + mlflow_repository.get_model_run_id = MagicMock() + mlflow_repository.get_model_uri = MagicMock() + + result = mlflow_repository.load_transform_model('test_model', 'sklearn') + + validate_common_load_transform_model_mocks(mlflow_repository, 'test_model') + + assert result == mlflow.sklearn.load_model.return_value + mlflow.sklearn.load_model.assert_called_once_with(mlflow_repository.get_model_uri.return_value) -@patch('laborious.utils.repository.model_repository.mlflow') -def test_get_experiment_last_run_error(mlflow, mlflow_repository): - mlflow.search_runs.return_value = [] +def test_load_transform_model_pyfunc(mlflow, mlflow_repository): + mlflow_repository.get_model_run_id = MagicMock() + mlflow_repository.get_model_uri = MagicMock() - try: - mlflow_repository.get_experiment_last_run(0) - except ValueError as e: - assert str(e) == 'Runs is not a pandas DataFrame' - else: - assert False + result = mlflow_repository.load_transform_model('test_model', 'pyfunc') + + validate_common_load_transform_model_mocks(mlflow_repository, 'test_model') + + assert result == mlflow.pyfunc.load_model.return_value + mlflow.pyfunc.load_model.assert_called_once_with(mlflow_repository.get_model_uri.return_value) -@patch('laborious.utils.repository.model_repository.mlflow.sklearn') -@patch('laborious.utils.repository.model_repository.mlflow.set_experiment') -def test_create_model_experiment(set_experiment, sklearn, mlflow_repository): +def test_load_transform_model_pytorch(mlflow, mlflow_repository): + mlflow_repository.get_model_run_id = MagicMock() + mlflow_repository.get_model_uri = MagicMock() - mlflow_repository.model_serving.get_model_run_id = MagicMock( - return_value='0') - mlflow_repository.model_serving.get_model_uri = MagicMock( - return_value='test') - mlflow_repository.get_experiment_by_run_id = MagicMock() + result = mlflow_repository.load_transform_model('test_model', 'pytorch') + validate_common_load_transform_model_mocks(mlflow_repository, 'test_model') - data_model_mock = MagicMock() - prediction_model_mock = MagicMock() + assert result == mlflow.pytorch.load_model.return_value + mlflow.pytorch.load_model.assert_called_once_with(mlflow_repository.get_model_uri.return_value) - sklearn.load_model.side_effect = [data_model_mock, prediction_model_mock] - data_model_mock.fit.return_value = data_model_mock - data_model_mock.predict.return_value = DataFrame({ - 'x': [10, 20, 30], - }) - data_model_mock.target_variable = 'y' +def test_load_transform_model_error(mlflow_repository): + mlflow_repository.get_model_run_id = MagicMock() + mlflow_repository.get_model_uri = MagicMock() - prediction_model_mock.fit.return_value = prediction_model_mock + with pytest.raises(ValueError) as e: + mlflow_repository.load_transform_model('test_model', 'invalid') + assert str(e) == "Invalid flavor. Use 'sklearn' or 'pyfunc' or 'pytorch'." - data = DataFrame({ - 'x': [1, 2, 3], - 'y': [4, 5, 6] - }) - output = mlflow_repository.create_model_experiment('test', data) +def test_download_model_invalid_model_type(mlflow_repository): + with pytest.raises(ValueError) as e: + mlflow_repository.download_model('test_model', 'invalid', 'sklearn') + assert str(e) == "Invalid model_type. Use 'predict' or 'transform'." - mlflow_repository.model_serving.get_model_run_id.assert_called_once_with( - 'test', stage='Production') - mlflow_repository.model_serving.get_model_uri.assert_called_once_with( - '0', prediction=False) - sklearn.load_model.assert_has_calls([ - call(mlflow_repository.model_serving.get_model_uri.return_value), - call("models:/test/production"), - ]) - assert sklearn.load_model.call_count == 2 +@pytest.mark.parametrize( + 'model_type', [('predict', 'prediction_model'), ('transform', 'data_model')] +) +def test_download_model_load_wrapper(mlflow, mlflow_repository, model_type): + mlflow_repository.dowload_artifacts = MagicMock() - data_model_mock.fit.assert_called_once_with(data) - data_model_mock.predict.assert_called_once_with(data) + result = mlflow_repository.download_model('test_model', model_type[0], 'pyfunc', True) - fit_args = prediction_model_mock.fit.call_args[0][0] - assert fit_args.equals( - DataFrame({ - 'x': [10, 20, 30], - 'y': [4, 5, 6], - }) + mlflow_repository.dowload_artifacts.assert_called_once_with('test_model', model_type[1]) + + mlflow.pyfunc.load_model.assert_called_once_with( + mlflow_repository.dowload_artifacts.return_value ) - mlflow_repository.get_experiment_by_run_id.assert_called_once_with('0') - - set_experiment.assert_called_once_with( - mlflow_repository.get_experiment_by_run_id.return_value + assert result == ( + mlflow.pyfunc.load_model.return_value._model_impl.python_model, + mlflow_repository.dowload_artifacts.return_value, ) - assert output == (prediction_model_mock, - data_model_mock, - mlflow_repository.get_experiment_by_run_id.return_value) + +def test_download_model_predict(mlflow_repository): + mlflow_repository.load_predict_model = MagicMock() + mlflow_repository.load_transform_model = MagicMock() + + result = mlflow_repository.download_model('test_model', 'predict', 'pyfunc', False) + + mlflow_repository.load_predict_model.assert_called_once_with('test_model', 'pyfunc') + mlflow_repository.load_transform_model.assert_not_called() + + assert result == (mlflow_repository.load_predict_model.return_value, None) -@patch('laborious.utils.repository.model_repository.mlflow.start_run') -@patch('laborious.utils.repository.model_repository.mlflow.log_param') -@patch('laborious.utils.repository.model_repository.mlflow.sklearn.log_model') -@patch('laborious.utils.repository.model_repository.mlflow.log_artifact') -def test_perform_model_retrain(log_artifact, log_model, log_param, start_run, mlflow_repository): +def test_download_model_transform(mlflow_repository): + mlflow_repository.load_predict_model = MagicMock() + mlflow_repository.load_transform_model = MagicMock() + + result = mlflow_repository.download_model('test_model', 'transform', 'pyfunc', False) + + mlflow_repository.load_predict_model.assert_not_called() + mlflow_repository.load_transform_model.assert_called_once_with('test_model', 'pyfunc') + + assert result == (mlflow_repository.load_transform_model.return_value, None) + + +invalid_cases = [ + ( + {'value': {'2024-01-01 12:00:00': 1, 2024: 2}}, + "Index must be all timestamp like column. Valid formats are: pandas Timestamp, datetime, string in format %Y-%m-%d %H:%M:%S. Elements are , .", + ), + ( + {'value': {'2024-01-01': 1, '2024-01-02': 2}}, + 'Index must be all timestamp like column. Valid formats are: pandas Timestamp, datetime, string in format %Y-%m-%d %H:%M:%S. Unable to parse given date format', + ), + ( + {'value': {Any(): 1, Any(): 2}}, + 'Index must be all timestamp like column. Valid formats are: pandas Timestamp, datetime, string in format %Y-%m-%d %H:%M:%S. Got .', + ), +] + + +@pytest.mark.parametrize('data', invalid_cases) +def test_detect_and_parse_datetime_index_error_cases(mlflow_repository, data): + input_data = DataFrame(data[0]) + + message = data[1] + + with pytest.raises(ValueError) as e: + mlflow_repository.detect_and_parse_datetime_index(input_data, metadata['metadata']) + + assert str(e) == message + + +valid_cases = [ + ( + {'value': {'2024-01-01 12:00:00+0000': 1, '2024-01-02 12:00:00+0000': 2}}, + ['2024-01-01 12:00:00+0000', '2024-01-02 12:00:00+0000'], + ), + ( + { + 'value': { + datetime(2025, 1, 1, 12, 0, 0, tzinfo=UTC): 1, + datetime(2025, 1, 2, 12, 0, 0, tzinfo=UTC): 2, + } + }, + ['2025-01-01 12:00:00+0000', '2025-01-02 12:00:00+0000'], + ), + ( + { + 'value': { + Timestamp(2026, 1, 1, 12, 0, 0, tzinfo=UTC): 1, + Timestamp(2026, 1, 2, 12, 0, 0, tzinfo=UTC): 2, + } + }, + ['2026-01-01 12:00:00+0000', '2026-01-02 12:00:00+0000'], + ), +] + + +@pytest.mark.parametrize('data,expected', valid_cases) +def test_detect_and_parse_datetime_index_valid_format(mlflow_repository, data, expected): + input_data = DataFrame(data) + + response = mlflow_repository.detect_and_parse_datetime_index(input_data, metadata['metadata']) + + assert response.index.tolist() == expected + + +@patch('laborious.utils.repository.model_repository.datetime') +def test_check_cache_retention_false(datetime_mock, mlflow_repository): + datetime_mock.now = MagicMock(return_value=datetime.strptime('2025-01-02', '%Y-%m-%d')) + cache = {'timestamp': datetime.strptime('2025-01-01', '%Y-%m-%d')} + + assert mlflow_repository.check_cache_retention(cache, 1) is False + + +@patch('laborious.utils.repository.model_repository.datetime') +def test_check_cache_retention_true(datetime_mock, mlflow_repository): + datetime_mock.now = MagicMock(return_value=datetime.strptime('2025-01-01', '%Y-%m-%d')) + cache = {'timestamp': datetime.strptime('2025-01-01', '%Y-%m-%d')} + assert mlflow_repository.check_cache_retention(cache, 1) is True + + +def test_handle_valid_model(mlflow_repository): + cache = {'target': {'model': 'model', 'artifact_path': 'test_artifact_path'}} + output = mlflow_repository.handle_valid_model('model_name', cache) + assert output == {'model': 'model', 'artifact_path': 'test_artifact_path'} + + +def test_handle_outdated_model(mlflow_repository): + mlflow_repository.model_cache = { + 'model_name_transform': { + 'target': {'model': 'model', 'artifact_path': 'test_artifact_path'} + } + } + mlflow_repository.handle_outdated_model('model_name', 'model_name_transform') + assert mlflow_repository.model_cache == {} + + +def test_get_model_retention_0(mlflow_repository): + model = MagicMock() + + mlflow_repository.download_model = MagicMock(return_value=(model, 'artifact_path')) + output = mlflow_repository.get_model('model_name', 0, 'predict', 'pyfunc') + + assert output == model + mlflow_repository.download_model.assert_called_once_with( + model_name='model_name', model_type='predict', flavor='pyfunc', load_wrapper=False + ) + + +def test_get_model_cached_valid(mlflow_repository): + mlflow_repository.check_cache_retention = MagicMock(return_value=True) + mlflow_repository.handle_valid_model = MagicMock() + mlflow_repository.handle_outdated_model = MagicMock() + mlflow_repository.model_cache = { + 'model_name_predict': { + 'target': 'cached_model', + } + } + + output = mlflow_repository.get_model('model_name', 1, 'predict', 'pyfunc') + + assert output == mlflow_repository.handle_valid_model.return_value + mlflow_repository.check_cache_retention.assert_called_once_with( + mlflow_repository.model_cache['model_name_predict'], 1 + ) + + mlflow_repository.handle_valid_model.assert_called_once_with( + model_name='model_name', cache=mlflow_repository.model_cache['model_name_predict'] + ) + + mlflow_repository.handle_outdated_model.assert_not_called() + + +def test_get_model_cached_outdated(mlflow_repository): + mlflow_repository.check_cache_retention = MagicMock(return_value=False) + mlflow_repository.handle_valid_model = MagicMock() + mlflow_repository.handle_outdated_model = MagicMock() + model = MagicMock() + mlflow_repository.download_model = MagicMock(return_value=(model, 'artifact_path')) + cache = { + 'model_name_predict': { + 'target': 'cached_model', + } + } + mlflow_repository.model_cache = cache + + output = mlflow_repository.get_model('model_name', 1, 'predict', 'pyfunc') + assert output == model + mlflow_repository.check_cache_retention.assert_called_once_with( + { + 'target': 'cached_model', + }, + 1, + ) + mlflow_repository.handle_valid_model.assert_not_called() + mlflow_repository.handle_outdated_model.assert_called_once_with( + model_name='model_name', model_key='model_name_predict' + ) + + +def test_get_model_cached_not_found(mlflow_repository): + mlflow_repository.check_cache_retention = MagicMock(return_value=False) + mlflow_repository.handle_valid_model = MagicMock() + mlflow_repository.handle_outdated_model = MagicMock() + mlflow_repository.model_cache = {} + model = MagicMock() + mlflow_repository.download_model = MagicMock(return_value=(model, 'artifact_path')) + output = mlflow_repository.get_model('model_name', 1, 'predict', 'pyfunc') + assert output == model + mlflow_repository.check_cache_retention.assert_not_called() + mlflow_repository.handle_valid_model.assert_not_called() + mlflow_repository.handle_outdated_model.assert_not_called() + + +@patch('laborious.utils.repository.model_repository.force_memory_release') +def test_get_cached_operation_retention_0(force_memory_release, mlflow_repository): + model = MagicMock() + data = MagicMock() + mlflow_repository.get_model = MagicMock(return_value=model) + output = mlflow_repository.get_cached_operation('model_name', data, 'transform', 0, 'sklearn') + assert output == model.predict.return_value + force_memory_release.assert_called_once_with(mlflow_repository.logger) + + +@patch('laborious.utils.repository.model_repository.force_memory_release') +def test_get_cached_predict_retention_not_0(force_memory_release, mlflow_repository): + model = MagicMock() + data = MagicMock() + mlflow_repository.get_model = MagicMock(return_value=model) + output = mlflow_repository.get_cached_operation('model_name', data, 'predict', 1, 'sklearn') + assert output == model.predict.return_value + force_memory_release.assert_not_called() + + +@patch('laborious.utils.repository.model_repository.force_memory_release') +def test_get_cached_operation_invalid_operation(force_memory_release, mlflow_repository): + data = MagicMock() + with pytest.raises(ValueError) as e: + mlflow_repository.get_cached_operation('model_name', data, 'invalid', 0, 'sklearn') + assert str(e) == "Invalid operation. Use 'transform' or 'predict'." + + +@patch('laborious.utils.repository.model_repository.pd.merge') +@patch('laborious.utils.repository.model_repository.isinstance') +def test_fit_models_not_df_target_name_none_and_not_in_model( + isinstance_mock, pd_merge, mlflow_repository +): + isinstance_mock.return_value = False + + data_model = MagicMock() + prediction_model = MagicMock() + mlflow_repository.download_model = MagicMock( + side_effect=[(data_model, 'artifact_path'), (prediction_model, 'artifact_path')], + ) + mlflow_repository.detect_and_parse_datetime_index = MagicMock( + return_value=MagicMock(drop_duplicates=MagicMock(return_value=MagicMock(columns=[]))) + ) - prediction_model_mock = MagicMock() - data_model_mock = MagicMock() - experiment = 'test' - model_name = 'test' data = MagicMock() - mlflow_repository.get_next_run_name = MagicMock( - return_value='test-1') - run = MagicMock() - start_run.__enter__.return_value = run + output = mlflow_repository.fit_models( + 'model_name', data, 'latest_production_id', metadata['metadata'], 'sklearn', 'pyfunc', None + ) - output = mlflow_repository.perform_model_retrain( - prediction_model_mock, data_model_mock, experiment, model_name, data) + mlflow_repository.download_model.assert_has_calls( + [ + call( + model_name='model_name', + model_type='transform', + flavor='sklearn', + load_wrapper=False, + ), + call(model_name='model_name', model_type='predict', flavor='pyfunc', load_wrapper=True), + ] + ) - mlflow_repository.get_next_run_name.assert_called_once_with(experiment) - start_run.assert_called_once_with( - run_name='test-1', description='Retrain model test with new data') + data_model.fit.assert_called_once_with(data) - log_model.assert_has_calls([ - call(data_model_mock, "data_model"), - call(prediction_model_mock, "prediction_model"), - ]) + data_model.fit.return_value.predict.assert_called_once_with(data) - data.to_csv.assert_called_once_with( - "temp/raw_data_test.csv", index=True) + transformed_data = data_model.fit.return_value.predict.return_value - log_artifact.assert_called_once_with( - "temp/raw_data_test.csv") + transformed_data.__setitem__.assert_called_once_with('timestamp', transformed_data.index) - log_param.assert_has_calls([ - call("retrain", True), - ]) + mlflow_repository.detect_and_parse_datetime_index.assert_called_once_with( + transformed_data, metadata['metadata'] + ) - assert output == ("Model retrained successfully", experiment) + transformed_data = mlflow_repository.detect_and_parse_datetime_index.return_value + + transformed_data.drop_duplicates.assert_called_once_with(subset=['timestamp'], keep='first') + + transformed_data = transformed_data.drop_duplicates.return_value + + data.loc.__getitem__.assert_called_once_with(transformed_data.index) + + aligned_data = data.loc.__getitem__.return_value + + aligned_data.__getitem__.assert_called_once_with(data_model.fit.return_value.target_variable) + + pd_merge.assert_called_once_with( + transformed_data, aligned_data.__getitem__.return_value, left_index=True, right_index=True + ) + + prediction_model.fit.assert_called_once_with(pd_merge.return_value) + + assert output == { + 'prediction_model': {'model': prediction_model, 'artifact_path': 'artifact_path'}, + 'data_model': {'model': data_model.fit.return_value, 'artifact_path': 'artifact_path'}, + } -def test_retrain_model(mlflow_repository): +@patch('laborious.utils.repository.model_repository.pd.merge') +@patch('laborious.utils.repository.model_repository.isinstance') +def test_fit_models_df_target_name_not_none_and_in_model( + isinstance_mock, pd_merge, mlflow_repository +): + isinstance_mock.return_value = True + + data_model = MagicMock( + target_variable='feat_2', + ) + prediction_model = MagicMock() + mlflow_repository.download_model = MagicMock( + side_effect=[(data_model, 'artifact_path'), (prediction_model, 'artifact_path')], + ) + mlflow_repository.detect_and_parse_datetime_index = MagicMock( + return_value=MagicMock( + drop_duplicates=MagicMock(return_value=MagicMock(columns=['feat_1'])) + ) + ) + data = MagicMock() - model_name = 'test' - mlflow_repository.create_model_experiment = MagicMock( - return_value=('data_model', 'prediction_model', '0')) + output = mlflow_repository.fit_models( + 'model_name', + data, + 'latest_production_id', + metadata['metadata'], + 'sklearn', + 'pyfunc', + 'feat_1', + ) - mlflow_repository.perform_model_retrain = MagicMock( - return_value='Model retrained successfully') + mlflow_repository.download_model.assert_has_calls( + [ + call( + model_name='model_name', + model_type='transform', + flavor='sklearn', + load_wrapper=False, + ), + call(model_name='model_name', model_type='predict', flavor='pyfunc', load_wrapper=True), + ] + ) - output = mlflow_repository.retrain_model(data, model_name) + data_model.fit.assert_called_once_with(data) - mlflow_repository.create_model_experiment.assert_called_once_with( - model_name, data) + transformed_data = data_model.fit.return_value - mlflow_repository.perform_model_retrain.assert_called_once_with( - 'data_model', 'prediction_model', '0', model_name, data) + transformed_data.__setitem__.assert_called_once_with('timestamp', transformed_data.index) - assert output == 'Model retrained successfully' + mlflow_repository.detect_and_parse_datetime_index.assert_called_once_with( + transformed_data, metadata['metadata'] + ) + + transformed_data = mlflow_repository.detect_and_parse_datetime_index.return_value + + transformed_data.drop_duplicates.assert_called_once_with(subset=['timestamp'], keep='first') + + transformed_data = transformed_data.drop_duplicates.return_value + + data.loc.__getitem__.assert_not_called() + + pd_merge.assert_not_called() + + prediction_model.fit.assert_called_once_with(transformed_data) + + assert output == { + 'prediction_model': {'model': prediction_model, 'artifact_path': 'artifact_path'}, + 'data_model': {'model': data_model, 'artifact_path': 'artifact_path'}, + } + + +def test_log_model_sklearn(mlflow, mlflow_repository): + model_data = {'model': MagicMock(), 'artifact_path': 'artifact_path'} + mlflow_repository.log_model(model_data, 'sklearn', 'prediction_model', metadata['metadata']) + mlflow.sklearn.log_model.assert_called_once_with(model_data['model'], 'prediction_model') + + +@patch('laborious.utils.repository.model_repository.path') +def test_log_model_pyfunc(path, mlflow, mlflow_repository): + model_data = {'model': MagicMock(), 'artifact_path': 'artifact_path'} + mlflow_repository.log_model(model_data, 'pyfunc', 'prediction_model', metadata['metadata']) + + mlflow.pyfunc.log_model.assert_not_called() + + path.join.assert_called_once_with('artifact_path', 'code', 'utils') + + model_data['model'].store_model.assert_called_once_with( + artifact_path='prediction_model', code_path=[path.join.return_value], to_disk=False + ) + + +def test_log_model_pytorch(mlflow, mlflow_repository): + model_data = {'model': MagicMock(), 'artifact_path': 'artifact_path'} + mlflow_repository.log_model(model_data, 'pytorch', 'prediction_model', metadata['metadata']) + mlflow.pytorch.log_model.assert_called_once_with(model_data['model'], 'prediction_model') + + +def test_log_model_error(mlflow_repository): + model_data = {'model': MagicMock(), 'artifact_path': 'artifact_path'} + with pytest.raises(ValueError) as e: + mlflow_repository.log_model(model_data, 'invalid', 'prediction_model', metadata['metadata']) + assert str(e) == "Invalid flavor. Use 'sklearn' or 'pyfunc' or 'pytorch'." + + +@patch('laborious.utils.repository.model_repository.force_memory_release') +@patch('laborious.utils.repository.model_repository.path') +@patch('laborious.utils.repository.model_repository.rmtree') +def test_create_new_experiment(_rmtree, path, force_memory_release, mlflow, mlflow_repository): + model_name = 'model_name' + data = MagicMock() + retrain_data = { + 'prediction_model': {'model': MagicMock(), 'artifact_path': 'artifact_path'}, + 'data_model': {'model': MagicMock(), 'artifact_path': 'artifact_path'}, + } + + mlflow_repository.get_model_params = MagicMock( + return_value={ + 'transform_flavor': 'sklearn', + 'predict_flavor': 'pyfunc', + 'target_name': 'target_name', + } + ) + + mlflow_repository.get_experiment = MagicMock() + mlflow_repository.get_next_run_name = MagicMock() + mlflow_repository.log_model = MagicMock() + path.exists.return_value = True + path.join.return_value = './tmp/artifacts/model_name' + + report = mlflow_repository.create_new_experiment( + model_name, + data, + retrain_data, + 'latest_production_id', + metadata['metadata'], + 'sklearn', + 'pyfunc', + ) + + path.join.assert_called_once_with('./tmp/artifacts', 'model_name') + + mlflow_repository.get_model_params.assert_called_once_with('latest_production_id') + mlflow_repository.get_experiment.assert_called_once_with(model_name, create_if_not_exists=True) + mlflow_repository.get_next_run_name.assert_called_once_with( + mlflow_repository.get_experiment.return_value.name + ) + + data.to_csv.assert_called_once_with('./tmp/artifacts/model_name/retrain_data.csv', index=True) + + mlflow.start_run.assert_called_once_with( + experiment_id=mlflow_repository.get_experiment.return_value.experiment_id, + run_name=mlflow_repository.get_next_run_name.return_value, + description='Retrain model model_name with new data', + ) + + mlflow_repository.log_model.assert_has_calls( + [ + call(retrain_data['data_model'], 'sklearn', 'data_model', metadata['metadata']), + call( + retrain_data['prediction_model'], 'pyfunc', 'prediction_model', metadata['metadata'] + ), + ] + ) + + mlflow.log_params.assert_called_once_with( + { + 'transform_flavor': 'sklearn', + 'predict_flavor': 'pyfunc', + 'target_name': 'target_name', + 'retrain': True, + 'retrain_date': ANY, + 'source_run_id': 'latest_production_id', + 'retrain_samples': data.shape.__str__.return_value, + } + ) + + mlflow.log_artifact.assert_called_once_with('./tmp/artifacts/model_name/retrain_data.csv') + + force_memory_release.assert_called_once_with(mlflow_repository.logger) + + assert report == { + 'run_id': mlflow.start_run.return_value.__enter__.return_value.info.run_id, + 'experiment_id': mlflow_repository.get_experiment.return_value.experiment_id, + 'experiment_name': mlflow_repository.get_experiment.return_value.name, + } -@patch('laborious.utils.repository.model_repository.mlflow') def test_update_production_model_by_run_id(mlflow, mlflow_repository): - client_mock = MagicMock() - mlflow.tracking.MlflowClient.return_value = client_mock - - client_mock.get_registered_model.return_value = MagicMock( + mlflow_repository.client.get_registered_model.return_value = MagicMock( latest_versions=[ MagicMock(version='1'), MagicMock(version='2'), MagicMock(version='3'), ] ) - output = mlflow_repository.update_production_model_by_run_id('0', 'test') + output = mlflow_repository.update_production_model_by_run_id('0', 'test', metadata['metadata']) mlflow.register_model.assert_called_once_with( - "runs:/0/prediction_model", + 'runs:/0/prediction_model', 'test', ) - mlflow.tracking.MlflowClient.assert_called_once() - client_mock.get_registered_model.assert_called_once_with('test') - client_mock.transition_model_version_stage.assert_called_once_with( + mlflow_repository.client.get_registered_model.assert_called_once_with('test') + mlflow_repository.client.transition_model_version_stage.assert_called_once_with( name='test', version='3', stage='Production', @@ -458,45 +828,184 @@ def test_update_production_model_by_run_id(mlflow, mlflow_repository): } -@patch('laborious.utils.repository.model_repository.mlflow') def test_update_production_model_by_run_id_error(mlflow, mlflow_repository): - mlflow.tracking.MlflowClient.return_value = MagicMock( - get_registered_model=MagicMock( - return_value=MagicMock( - latest_versions={} - ) - ) + mlflow_repository.client.get_registered_model.return_value = MagicMock( + get_registered_model=MagicMock(return_value=MagicMock(latest_versions={})) ) try: - mlflow_repository.update_production_model_by_run_id('0', 'test') + mlflow_repository.update_production_model_by_run_id('0', 'test', metadata['metadata']) except Exception as e: assert str(e) == 'Model versions is not a list' else: - assert False + raise AssertionError('Expected Exception') + + +def test_transform_success(mlflow_repository): + data = MagicMock() + model_name = 'model' + model_config = { + 'retention_minutes': 60, + 'transform_flavor': 'sklearn', + 'predict_flavor': 'pyfunc', + } + + mlflow_repository.get_cached_operation = MagicMock() + + mlflow_repository.detect_and_parse_datetime_index = MagicMock() + + output = mlflow_repository.transform(model_name, data, model_config, metadata['metadata']) + + mlflow_repository.get_cached_operation.assert_called_once_with( + model_name, data, 'transform', 60, 'sklearn' + ) + + mlflow_repository.detect_and_parse_datetime_index.assert_called_once_with( + mlflow_repository.get_cached_operation.return_value, metadata['metadata'] + ) + + assert output == { + 'success': True, + 'content': mlflow_repository.detect_and_parse_datetime_index.return_value.to_dict.return_value, + } + + +def test_transform_error(mlflow_repository): + data = MagicMock() + model_name = 'model' + model_config = { + 'retention_minutes': 60, + 'transform_flavor': 'sklearn', + 'predict_flavor': 'pyfunc', + } + + mlflow_repository.get_cached_operation = MagicMock(side_effect=Exception('error')) + + output = mlflow_repository.transform(model_name, data, model_config, metadata['metadata']) + + mlflow_repository.get_cached_operation.assert_called_once_with( + model_name, data, 'transform', 60, 'sklearn' + ) + + assert output == {'success': False, 'content': {'message': 'error', 'traceback': ANY}} + + +def test_predict_success_array(mlflow_repository): + data = DataFrame({'feat_1': {'index_1': 2, 'index_2': 3}}) + model_config = {'retention_minutes': 60, 'predict_flavor': 'pyfunc'} + model_name = 'model' + mlflow_repository.get_cached_operation = MagicMock(return_value=np.array([2, 3])) + + output = mlflow_repository.predict(model_name, data, model_config, metadata['metadata']) + + mlflow_repository.get_cached_operation.assert_called_once_with( + model_name, data, 'predict', 60, 'pyfunc' + ) + + assert output['success'] is True + assert output['content'] == { + 'prediction': {'index_1': 2, 'index_2': 3}, + 'response_time': {'index_1': ANY, 'index_2': ANY}, + } + + +def test_predict_success_df(mlflow_repository): + data = DataFrame({'feat_1': {'index_1': 2, 'index_2': 3}}) + model_config = {'retention_minutes': 60, 'predict_flavor': 'pyfunc'} + model_name = 'model' + + mlflow_repository.get_cached_operation = MagicMock( + return_value=DataFrame({'feat_1': {'index_3': 2, 'index_4': 3}}) + ) + + output = mlflow_repository.predict(model_name, data, model_config, metadata['metadata']) + + mlflow_repository.get_cached_operation.assert_called_once_with( + model_name, data, 'predict', 60, 'pyfunc' + ) + + assert output['success'] is True + assert output['content'] == { + 'prediction': {'index_1': 2, 'index_2': 3}, + 'response_time': {'index_1': ANY, 'index_2': ANY}, + } + + +def test_predict_error(mlflow_repository): + data = DataFrame({'feat_1': {'index_1': 2, 'index_2': 3}}) + model_name = 'model' + model_config = {'retention_minutes': 60, 'predict_flavor': 'pyfunc'} + + mlflow_repository.get_cached_operation = MagicMock(side_effect=Exception('error')) + + output = mlflow_repository.predict(model_name, data, model_config, metadata['metadata']) + + mlflow_repository.get_cached_operation.assert_called_once_with( + model_name, data, 'predict', 60, 'pyfunc' + ) + + assert output == {'success': False, 'content': {'message': 'error', 'traceback': ANY}} + + +def test_retrain_model(mlflow_repository): + data = MagicMock() + model_name = 'test' + model_config = {'target': 'target', 'transform_flavor': 'sklearn', 'predict_flavor': 'pyfunc'} + + mlflow_repository.get_model_run_id = MagicMock() + mlflow_repository.fit_models = MagicMock() + mlflow_repository.create_new_experiment = MagicMock() + + output = mlflow_repository.retrain_model(data, model_name, model_config, metadata['metadata']) + + mlflow_repository.get_model_run_id.assert_called_once_with(model_name, stage='Production') + + mlflow_repository.fit_models.assert_called_once_with( + model_name=model_name, + data=data, + transform_flavor='sklearn', + predict_flavor='pyfunc', + target_name='target', + metadata=metadata['metadata'], + latest_production_id=mlflow_repository.get_model_run_id.return_value, + ) + + mlflow_repository.create_new_experiment.assert_called_once_with( + model_name=model_name, + data=data, + retrain_data=mlflow_repository.fit_models.return_value, + transform_flavor='sklearn', + predict_flavor='pyfunc', + metadata=metadata['metadata'], + latest_production_id=mlflow_repository.get_model_run_id.return_value, + ) + + assert output == { + 'success': True, + 'experiment': mlflow_repository.create_new_experiment.return_value, + 'message': 'Model retrained successfully.', + } def test_update_production_model(mlflow_repository): - connector = mlflow_repository + experiment = {'run_id': '0', 'experiment_id': '0'} + model_name = 'test' + mlflow_repository.update_production_model_by_run_id = MagicMock() + mlflow_repository.update_production_model_by_run_id.return_value = { + 'model_name': 'test', + 'version': '3', + 'mlflow_run_id': '0', + } - with patch.object(connector, 'get_experiment', - return_value='0') as get_experiment: - with patch.object(connector, 'get_experiment_last_run', - return_value='2') as get_experiment_last_run: - with patch.object(connector, 'update_production_model_by_run_id', - return_value={'model_name': 'test', 'version': '3', - 'mlflow_run_id': '0'}) as update_production_model_by_run_id: + output = mlflow_repository.update_production_model(experiment, model_name, metadata['metadata']) - output = connector.update_production_model('0', 'test') + mlflow_repository.update_production_model_by_run_id.assert_called_once_with( + '0', 'test', metadata['metadata'] + ) - get_experiment.assert_called_once_with('0') - get_experiment_last_run.assert_called_once_with('0') - update_production_model_by_run_id.assert_called_once_with( - '2', 'test') - - assert output == { - 'model_name': 'test', - 'version': '3', - 'mlflow_run_id': '0', - 'mlflow_experiment_id': '0', - } + assert output == { + 'model_name': 'test', + 'version': '3', + 'mlflow_run_id': '0', + 'mlflow_experiment_id': '0', + } diff --git a/tests/laborious/utils/repository/test_opc_repository.py b/tests/laborious/utils/repository/test_opc_repository.py index bc84db8..a0048e8 100644 --- a/tests/laborious/utils/repository/test_opc_repository.py +++ b/tests/laborious/utils/repository/test_opc_repository.py @@ -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 diff --git a/tests/laborious/utils/test_connectors_config.py b/tests/laborious/utils/test_connectors_config.py index 910439c..cf2a6b8 100644 --- a/tests/laborious/utils/test_connectors_config.py +++ b/tests/laborious/utils/test_connectors_config.py @@ -1,8 +1,11 @@ from os import environ -from laborious.utils.connectors_config import (build_mlflow_config, - build_opc_config, - build_postgres_config, - build_mongodb_config) + +from laborious.utils.connectors_config import ( + build_mlflow_config, + build_mongodb_config, + build_opc_config, + build_postgres_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, } diff --git a/tests/laborious/workflows/subworkflows/test_format_and_export_prediction.py b/tests/laborious/workflows/subworkflows/test_format_and_export_prediction.py index a8e6e20..006ef4e 100644 --- a/tests/laborious/workflows/subworkflows/test_format_and_export_prediction.py +++ b/tests/laborious/workflows/subworkflows/test_format_and_export_prediction.py @@ -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,149 +13,167 @@ 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([ - call( - Activities.format_prediction, - { - 'data': input_data['data'], - 'timestamp': input_data['timestamp'], - 'model_id': input_data['model_id'], - 'prediction_confidence': input_data['prediction_confidence'], - 'prediction_store_policy': input_data['prediction_store_policy'], - **metadata - }, - retry_policy=ANY, - start_to_close_timeout=ANY - )]) + workflow_mock.execute_local_activity_method.assert_has_calls( + [ + call( + Activities.format_prediction, + { + 'data': input_data['data'], + 'timestamp': input_data['timestamp'], + 'model_id': input_data['model_id'], + 'prediction_confidence': input_data['prediction_confidence'], + 'prediction_store_policy': input_data['prediction_store_policy'], + **metadata, + }, + retry_policy=ANY, + start_to_close_timeout=ANY, + ) + ] + ) - 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 - }, - retry_policy=ANY, - start_to_close_timeout=ANY - ) - ]) + 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, + }, + retry_policy=ANY, + start_to_close_timeout=ANY, + ) + ] + ) - workflow_mock.execute_activity_method.assert_has_calls([ - call( - Activities.export_data_to_postgres, - { - 'schema': input_data['schema'], - 'table_name': input_data['table_name'], - 'data': workflow_mock.execute_activity_method.return_value, - **metadata, - 'timestamp_conversion': { - 'column': 'timestamp', - 'format': DATETIME_FORMAT_WITH_TZ - } - }, - retry_policy=ANY, - start_to_close_timeout=ANY - )]) + workflow_mock.execute_activity_method.assert_has_calls( + [ + call( + Activities.export_data_to_postgres, + { + 'schema': input_data['schema'], + 'table_name': input_data['table_name'], + 'data': workflow_mock.execute_activity_method.return_value, + **metadata, + 'timestamp_conversion': { + 'column': 'timestamp', + 'format': DATETIME_FORMAT_WITH_TZ, + }, + }, + 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 @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([ - call( - Activities.format_default_prediction, - { - 'timestamp': input_data['timestamp'], - 'model_id': input_data['model_id'], - 'prediction_confidence': input_data['prediction_confidence'], - 'comment': input_data['comment'], - **metadata - }, - retry_policy=ANY, - start_to_close_timeout=ANY - ) - ]) + workflow_mock.execute_local_activity_method.assert_has_calls( + [ + call( + Activities.format_default_prediction, + { + 'timestamp': input_data['timestamp'], + 'model_id': input_data['model_id'], + 'prediction_confidence': input_data['prediction_confidence'], + 'comment': input_data['comment'], + **metadata, + }, + retry_policy=ANY, + start_to_close_timeout=ANY, + ) + ] + ) - 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 - }, - retry_policy=ANY, - start_to_close_timeout=ANY - ) - ]) + 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, + }, + retry_policy=ANY, + start_to_close_timeout=ANY, + ) + ] + ) - workflow_mock.execute_activity_method.assert_has_calls([ - call( - Activities.export_data_to_postgres, - { - 'schema': input_data['schema'], - 'table_name': input_data['table_name'], - 'data': workflow_mock.execute_activity_method.return_value, - **metadata, - 'timestamp_conversion': { - 'column': 'timestamp', - 'format': DATETIME_FORMAT_WITH_TZ - } - }, - retry_policy=ANY, - start_to_close_timeout=ANY - ) - ]) + workflow_mock.execute_activity_method.assert_has_calls( + [ + call( + Activities.export_data_to_postgres, + { + 'schema': input_data['schema'], + 'table_name': input_data['table_name'], + 'data': workflow_mock.execute_activity_method.return_value, + **metadata, + 'timestamp_conversion': { + 'column': 'timestamp', + 'format': DATETIME_FORMAT_WITH_TZ, + }, + }, + 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 diff --git a/tests/laborious/workflows/subworkflows/test_prediction_process.py b/tests/laborious/workflows/subworkflows/test_prediction_process.py index df60ada..6baec30 100644 --- a/tests/laborious/workflows/subworkflows/test_prediction_process.py +++ b/tests/laborious/workflows/subworkflows/test_prediction_process.py @@ -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, { - **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, { - **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, { - **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, { - **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, { - **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, { - **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, { - **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)]) + 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, + { + **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, + { + **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, + { + **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, + { + **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, + { + **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, + { + **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, + ) + ] + ) 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, { - 'data': input_data['data'], - **metadata, - }, 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) - ]) + 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, + { + '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_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, { - '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, { - '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, { - '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, { - '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.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, + { + '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, + { + '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, + { + '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_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, { - '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, { - '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, { - '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, { - '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, { - '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.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, + { + '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, + { + '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, + { + '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, + { + '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_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, { - '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, { - '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, { - '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, { - '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, { - '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, { - '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, { - '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)]) + 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, + { + '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, + { + '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, + { + '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, + { + '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, + { + '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, + { + '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, + ) + ] + ) 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 diff --git a/tests/laborious/workflows/test_minimal_retrain.py b/tests/laborious/workflows/test_minimal_retrain.py index b3b03b5..c088702 100644 --- a/tests/laborious/workflows/test_minimal_retrain.py +++ b/tests/laborious/workflows/test_minimal_retrain.py @@ -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([ - call( - Activities.retrain_model, - { - **metadata, - 'data': workflow_mock.execute_local_activity_method.return_value, - 'model_name': input_data['model_name'], - }, - 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_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, - }, - retry_policy=ANY, - start_to_close_timeout=ANY - ) - ]) + workflow_mock.execute_activity_method.assert_has_calls( + [ + call( + Activities.update_production_model, + { + **metadata, + 'model_name': input_data['model_name'], + 'success': True, + 'experiment': 'test_experiment', + }, + retry_policy=ANY, + start_to_close_timeout=ANY, + ) + ] + ) - workflow_mock.execute_activity_method.assert_has_calls([ - call( - Activities.export_data_to_postgres, + 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_local_activity_method.return_value, + 'schema': input_data['schema'], + 'table_name': input_data['table_name'], + }, + retry_policy=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'}, { - **metadata, - 'data': workflow_mock.execute_activity_method.return_value, - 'schema': input_data['schema'], - 'table_name': input_data['table_name'], + '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 - ) - ]) + {'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 diff --git a/tests/laborious/workflows/test_predictions_batch.py b/tests/laborious/workflows/test_predictions_batch.py index 90d7d21..62cc43c 100644 --- a/tests/laborious/workflows/test_predictions_batch.py +++ b/tests/laborious/workflows/test_predictions_batch.py @@ -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([ - call( - Activities.load_custom_query, - { - **metadata, - 'query': input_data['query'], - 'datetime_columns': input_data.get('datetime_columns', []) - }, - retry_policy=ANY, - start_to_close_timeout=ANY - ) - ]) + 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', []), + }, + retry_policy=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)] + ) diff --git a/validate.sh b/validate.sh new file mode 100755 index 0000000..8728312 --- /dev/null +++ b/validate.sh @@ -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 \ No newline at end of file