diff --git a/model_manager/activities/activities.py b/model_manager/activities/activities.py index 81407a7..9fedd57 100644 --- a/model_manager/activities/activities.py +++ b/model_manager/activities/activities.py @@ -1,12 +1,14 @@ -from temporalio import activity, workflow +from temporalio import workflow with workflow.unsafe.imports_passed_through(): - from sientia_do.temporal.activities.postgres import Postgres + from typing import Any + from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler from sientia_do.observability.logger import Logger - from model_manager.activities.mlflow import MLFlow + from sientia_do.temporal.activities.postgres import Postgres + from model_manager.activities.gates import Gates - from typing import Any + from model_manager.activities.mlflow import MLFlow class Activities(Postgres, MLFlow, Gates): @@ -29,11 +31,13 @@ class Activities(Postgres, MLFlow, Gates): notification_handler (NotificationHandler): Notification management instance """ - def __init__(self, - postgres_config: dict[str, Any], - mlflow_config: dict[str, Any], - logger: Logger, - notification_handler: NotificationHandler): + def __init__( + self, + postgres_config: dict[str, Any], + mlflow_config: dict[str, Any], + logger: Logger, + notification_handler: NotificationHandler, + ): """ Initialize the Activities orchestrator with all required configurations. @@ -52,25 +56,30 @@ class Activities(Postgres, MLFlow, Gates): Exception: If any parent class initialization fails """ # Initialize parent classes - Postgres.__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'], - logger=logger, - notification_handler=notification_handler) + Postgres.__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'], + 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'], - 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'], + logger=logger, + notification_handler=notification_handler, + ) - Gates.__init__(self, logger=logger, - notification_handler=notification_handler) + Gates.__init__(self, logger=logger, notification_handler=notification_handler) async def shutdown(self): """ diff --git a/model_manager/activities/gates.py b/model_manager/activities/gates.py index 338e212..091977c 100644 --- a/model_manager/activities/gates.py +++ b/model_manager/activities/gates.py @@ -1,53 +1,42 @@ from temporalio import activity, workflow - with workflow.unsafe.imports_passed_through(): import traceback + 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.temporal.activities.base import BaseActivity 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 sientia_do.formatters import create_sample_dict - from model_manager.utils.filters.mlflow_filters import nan_values_filter, api_error_filter + from model_manager import metrics from model_manager.utils.filters.conditional_filters import ( filter_empty_data, - filter_specific_variables_null_values + filter_specific_variables_null_values, ) - from pandas import DataFrame - from model_manager import metrics + from model_manager.utils.filters.mlflow_filters import api_error_filter, nan_values_filter # Input filter function mappings input_filter_functions = { 'SPECIFIC_VARIABLES_NULL_VALUES': filter_specific_variables_null_values, 'EMPTY_DATA': filter_empty_data, - 'path_confidence': { - 'STOP': -1, - 'CONTINUE': 2, - 'REPEAT': -1 - } + 'path_confidence': {'STOP': -1, 'CONTINUE': 2, 'REPEAT': -1}, } # MLFlow response filter function mappings mlflow_response_filter_functions = { 'API_ERROR': api_error_filter, - 'path_confidence': { - 'STOP': -1, - 'CONTINUE': 10, - 'REPEAT': -1 - }, + 'path_confidence': {'STOP': -1, 'CONTINUE': 10, 'REPEAT': -1}, } # MLFlow content filter function mappings mlflow_content_filter_functions = { 'NAN_VALUES': nan_values_filter, 'EMPTY_DATA': filter_empty_data, - 'path_confidence': { - 'STOP': -1, - 'CONTINUE': 18, - 'REPEAT': -1 - } + 'path_confidence': {'STOP': -1, 'CONTINUE': 18, 'REPEAT': -1}, } @@ -82,10 +71,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. @@ -121,7 +109,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']) @@ -129,40 +117,42 @@ 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: + except Exception as e: # noqa: BLE001 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_filter_functions['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. @@ -197,7 +187,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'] @@ -206,14 +196,13 @@ class Gates(BaseActivity): filter_output = [] - self.debug( - f"Input data: \n {create_sample_dict(data, max_items=5, max_depth=2)}", metadata) - self.debug(f"Filters: {filters}", metadata) + self.debug(f'Input data: \n {create_sample_dict(data, max_items=5, max_depth=2)}', metadata) + self.debug(f'Filters: {filters}', metadata) comments = [] for fil, config in filters.items(): if fil not in mlflow_response_filter_functions: - self.error(f"Filter {fil} not found", metadata) + self.error(f'Filter {fil} not found', metadata) continue try: if mlflow_response_filter_functions[fil](data, config): @@ -221,34 +210,36 @@ 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: + except Exception as e: # noqa: BLE001 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_filter_functions['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. @@ -283,7 +274,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']) @@ -292,8 +283,8 @@ class Gates(BaseActivity): filter_output = [] - self.debug(f"Input data:\n {data.head(5).to_string()}", metadata) - self.debug(f"Filters: \n {create_sample_dict(filters)}", metadata) + self.debug(f'Input data:\n {data.head(5).to_string()}', metadata) + self.debug(f'Filters: \n {create_sample_dict(filters)}', metadata) for fil, config in filters.items(): if fil not in mlflow_content_filter_functions: @@ -303,36 +294,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: + except Exception as e: # noqa: BLE001 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_filter_functions['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. @@ -358,7 +351,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] @@ -366,14 +361,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. @@ -400,7 +401,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']) @@ -408,48 +409,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. @@ -477,22 +475,24 @@ 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="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. @@ -517,24 +517,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. @@ -562,26 +560,24 @@ 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) diff --git a/model_manager/activities/mlflow.py b/model_manager/activities/mlflow.py index e55a55e..4de863b 100644 --- a/model_manager/activities/mlflow.py +++ b/model_manager/activities/mlflow.py @@ -1,20 +1,19 @@ from temporalio import activity, workflow - with workflow.unsafe.imports_passed_through(): - from datetime import datetime - from pandas import Timestamp, 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 sientia_do.temporal.activities.base import BaseActivity + from sientia_do.temporal.constants import DATETIME_FORMAT, DATETIME_FORMAT_WITH_TZ + from model_manager.utils.repository.model_repository import MLFlowRepository - from typing import Any - import numpy as np - from pandas import DataFrame - import traceback class MLFlow(BaseActivity): @@ -36,8 +35,15 @@ class MLFlow(BaseActivity): model_monitoring_repository (MLFlowRepository): Repository for MLFlow operations """ - def __init__(self, mlflow_host: str, mlflow_port: int, mlflow_username: str, - mlflow_password: str, logger: Logger, notification_handler: NotificationHandler): + def __init__( + self, + mlflow_host: str, + mlflow_port: int, + mlflow_username: str, + mlflow_password: str, + logger: Logger, + notification_handler: NotificationHandler, + ): """ Initialize MLFlow activities with server configuration. @@ -52,18 +58,17 @@ 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 ) - @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. @@ -99,7 +104,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 @@ -108,14 +113,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 @@ -124,16 +127,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. @@ -169,14 +176,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( @@ -184,13 +192,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. @@ -234,8 +244,7 @@ class MLFlow(BaseActivity): data.drop(columns=['model_id'], inplace=True, errors='ignore') data.drop(columns=['created_at'], inplace=True, errors='ignore') - data = data.pivot(index='timestamp', columns='variable', - values='value') + data = data.pivot(index='timestamp', columns='variable', values='value') data.sort_index(inplace=True) data.reset_index(inplace=True) @@ -244,15 +253,10 @@ class MLFlow(BaseActivity): try: retrain_output, experiment = self.model_monitoring_repository.retrain_model( - data=data, - model_name=model_name + data=data, model_name=model_name ) - return { - 'status': retrain_output, - 'timestamp': timestamp, - 'experiment': experiment - } + return {'status': retrain_output, 'timestamp': timestamp, 'experiment': experiment} except Exception as e: trace = traceback.format_exc() self.send_notification( @@ -261,12 +265,12 @@ class MLFlow(BaseActivity): message=f'Error retraining model {model_name}: {e}', block='retrain_model', level=NotificationLevel.ERROR, - attachment_content=trace + attachment_content=trace, ) self.error(trace, metadata=metadata) raise e - @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. @@ -311,12 +315,12 @@ class MLFlow(BaseActivity): status = input_data['status'] 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 ) report = DataFrame([response]) @@ -325,8 +329,7 @@ class MLFlow(BaseActivity): report['timestamp'] = timestamp report['status'] = status - self.info( - f'Production model {model_name} updated successfully', metadata) + self.info(f'Production model {model_name} updated successfully', metadata) return report.to_dict() except Exception as e: @@ -337,7 +340,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/model_manager/metrics.py b/model_manager/metrics.py index a19b03c..311f80f 100644 --- a/model_manager/metrics.py +++ b/model_manager/metrics.py @@ -22,36 +22,36 @@ Metric Labels: - pipeline_name: Name of the prediction pipeline """ -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( - "model_manager_predictions_written_count", - "Number of predictions written to the database table predictions", + 'model_manager_predictions_written_count', + 'Number of predictions written to the database table predictions', CORE_LABELS, ) # Prediction quality metrics PREDICTION_CONFIDENCE_MONITOR = Gauge( - "model_manager_prediction_confidence_monitor", - "Current confidence of each prediction", + 'model_manager_prediction_confidence_monitor', + 'Current confidence of each prediction', CORE_LABELS, ) # Performance monitoring metrics PREDICTION_RESPONSE_TIME_MONITOR = Histogram( - "model_manager_prediction_response_time_monitor", - "Current response time of each prediction", + 'model_manager_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], ) diff --git a/model_manager/utils/connectors_config.py b/model_manager/utils/connectors_config.py index 40e7c07..aec7625 100644 --- a/model_manager/utils/connectors_config.py +++ b/model_manager/utils/connectors_config.py @@ -1,9 +1,8 @@ from os import getenv -import json -from typing import Dict, Any +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 +29,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 +54,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_mongodb_config() -> Dict[str, Any]: +def build_mongodb_config() -> dict[str, Any]: """ Build MongoDB configuration from environment variables. @@ -86,5 +85,5 @@ 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, } diff --git a/model_manager/utils/filters/conditional_filters.py b/model_manager/utils/filters/conditional_filters.py index 2c51805..717a6f8 100644 --- a/model_manager/utils/filters/conditional_filters.py +++ b/model_manager/utils/filters/conditional_filters.py @@ -20,8 +20,7 @@ def filter_specific_variables_null_values(data: DataFrame, config: dict) -> bool False if none of the specified variables contain null values. """ - 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/model_manager/utils/filters/mlflow_filters.py b/model_manager/utils/filters/mlflow_filters.py index f792a0e..82970e8 100644 --- a/model_manager/utils/filters/mlflow_filters.py +++ b/model_manager/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/model_manager/utils/repository/model_repository.py b/model_manager/utils/repository/model_repository.py index 3107d1c..e2465a9 100644 --- a/model_manager/utils/repository/model_repository.py +++ b/model_manager/utils/repository/model_repository.py @@ -10,22 +10,23 @@ requests using the Model Monitoring API functions. By Monitoring we mean the evaluation of the performance of models, the generation of reports. """ -from datetime import datetime + import traceback -import pandas as pd -import mlflow +from datetime import datetime from os import makedirs, path, remove + +import mlflow +import pandas as pd from sientia.ModelServing import ModelServing -from sientia_do.temporal.constants import DATETIME_FORMAT_WITH_TZ from sientia_do.observability.logger import Logger +from sientia_do.temporal.constants import DATETIME_FORMAT_WITH_TZ -class MLFlowRepository(): +class MLFlowRepository: def __init__(self, host, username, password, logger: Logger): - - self.model_serving = ModelServing(tracking_uri=host, - username=username, password=password, - logger=logger) + self.model_serving = ModelServing( + tracking_uri=host, username=username, password=password, logger=logger + ) self.logger = logger def detect_and_parse_datetime_index(self, data: pd.DataFrame, metadata: dict) -> pd.DataFrame: @@ -40,33 +41,32 @@ 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}" + 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}") + raise ValueError(f'{message}') # 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}') from e elif index_type == datetime or index_type == pd.Timestamp: data.index = data.index.strftime(DATETIME_FORMAT_WITH_TZ) else: - raise ValueError( - f"{message}") + raise ValueError(f'{message}') return data - def transform(self, model_name: str, data: pd.DataFrame, model_config: dict, metadata: dict) -> dict: + def transform( + self, model_name: str, data: pd.DataFrame, model_config: dict, metadata: dict + ) -> dict: """ Transform data using a model. @@ -81,41 +81,42 @@ class MLFlowRepository(): try: self.logger.custom_debug( - f"Data received for model transformation: {data.to_csv()}", metadata) + f'Data received for model transformation: {data.to_csv()}', metadata + ) model_retention = model_config.get('retention_minutes', 0) flavor = model_config.get('transform_flavor', 'sklearn') compressed = model_config.get('is_compressed', False) retention_target = model_config.get('retention_target', 'model') - transform_keyword = model_config.get( - 'transform_function_keyword', 'predict') + transform_keyword = model_config.get('transform_function_keyword', 'predict') transformed_data = self.model_serving.get_cached_transform( - model_name, data, model_retention, flavor, - compressed, retention_target, transform_keyword + model_name, + data, + model_retention, + flavor, + compressed, + retention_target, + transform_keyword, ) self.logger.custom_debug( - f"Data received from model transformation: {transformed_data.to_csv()}", metadata) + f'Data received from model transformation: {transformed_data.to_csv()}', metadata + ) - 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: + except Exception as e: # noqa: BLE001 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) -> dict: + def predict( + self, model_name: str, data: pd.DataFrame, model_config: dict, metadata: dict + ) -> dict: """ Predict data using a model. @@ -137,31 +138,26 @@ class MLFlowRepository(): start_time = datetime.now() self.logger.custom_debug( - f"Data received for model prediction: {data.to_csv()}", metadata) + f'Data received for model prediction: {data.to_csv()}', metadata + ) data = self.model_serving.get_cached_predict( - model_name, data, model_retention, flavor, - compressed, retention_target + model_name, data, model_retention, flavor, compressed, retention_target ) end_time = datetime.now() data = pd.DataFrame(data, columns=['prediction']) self.logger.custom_debug( - f"Data received from model prediction: {data.to_csv()}", metadata) + f'Data received from model prediction: {data.to_csv()}', metadata + ) data.index = input_index data['response_time'] = (end_time - start_time).total_seconds() - return { - 'success': True, - 'content': data.to_dict() - } + return {'success': True, 'content': data.to_dict()} - except Exception as e: + except Exception as e: # noqa: BLE001 return { 'success': False, - 'content': { - 'message': str(e), - 'traceback': traceback.format_exc() - } + 'content': {'message': str(e), 'traceback': traceback.format_exc()}, } def get_experiment_by_run_id(self, run_id: str) -> dict: @@ -190,10 +186,9 @@ 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 create_model_experiment(self, model_name: str, data: pd.DataFrame) -> tuple: """ @@ -217,14 +212,10 @@ class MLFlowRepository(): - experiment: MLFlow experiment name """ # load predictor model - predictor_uri = f"models:/{model_name}/production" + predictor_uri = f'models:/{model_name}/production' # load transform model - latest_production_id = self.model_serving.get_model_run_id( - model_name, stage="Production" - ) - transform_uri = self.model_serving.get_model_uri( - latest_production_id, prediction=False - ) + latest_production_id = self.model_serving.get_model_run_id(model_name, stage='Production') + transform_uri = self.model_serving.get_model_uri(latest_production_id, prediction=False) # load data_model = mlflow.sklearn.load_model(transform_uri) prediction_model = mlflow.sklearn.load_model(predictor_uri) @@ -233,20 +224,16 @@ class MLFlowRepository(): target_name = data_model.target_variable y = data[target_name] - treated_data = pd.merge( - treated_data, y, left_index=True, right_index=True) + treated_data = pd.merge(treated_data, y, left_index=True, right_index=True) prediction_model = prediction_model.fit(treated_data) experiment = self.get_experiment_by_run_id(latest_production_id) mlflow.set_experiment(experiment) return prediction_model, data_model, experiment - def perform_model_retrain(self, - prediction_model, - data_model, - experiment: str, - model_name: str, - data: pd.DataFrame): + def perform_model_retrain( + self, prediction_model, data_model, experiment: str, model_name: str, data: pd.DataFrame + ): """ Execute the complete model retraining process in MLFlow. @@ -271,7 +258,7 @@ class MLFlowRepository(): """ pred_model_atributes = vars(prediction_model) # load class attributes data_model_atributes = vars(data_model) # load class attributes - experiment_description = f"Retrain model {model_name} with new data" + experiment_description = f'Retrain model {model_name} with new data' current_run_name = self.get_next_run_name(experiment) with mlflow.start_run( run_name=current_run_name, description=experiment_description @@ -279,32 +266,32 @@ class MLFlowRepository(): # update transfomation model # fixed parameters for name_atribute, val_atribute in pred_model_atributes.items(): - if name_atribute != "model": + if name_atribute != 'model': mlflow.log_param(name_atribute, val_atribute) # update prediction model for name_atribute, val_atribute in data_model_atributes.items(): - if name_atribute != "model": + if name_atribute != 'model': mlflow.log_param(name_atribute, val_atribute) # dynamic parameters, including model itself - mlflow.sklearn.log_model(data_model, "data_model") + mlflow.sklearn.log_model(data_model, 'data_model') - makedirs("temp", exist_ok=True) + makedirs('temp', exist_ok=True) - file_path = f"temp/raw_data_{model_name}.csv" + file_path = f'temp/raw_data_{model_name}.csv' data.to_csv(file_path, index=True) # log the data raw mlflow.log_artifact(file_path) # dynamic parameters, including model itself - mlflow.sklearn.log_model(prediction_model, "prediction_model") - mlflow.log_param("retrain", True) + mlflow.sklearn.log_model(prediction_model, 'prediction_model') + mlflow.log_param('retrain', True) # clear temp file if path.exists(file_path): remove(file_path) - return "Model retrained successfully", experiment + return 'Model retrained successfully', experiment def retrain_model(self, data: pd.DataFrame, model_name: str) -> tuple: """ @@ -325,10 +312,10 @@ class MLFlowRepository(): - status_message (str): Retraining operation status - experiment_name (str): MLFlow experiment identifier """ - prediction_model, data_model, experiment = self.create_model_experiment( - model_name, data) + prediction_model, data_model, experiment = self.create_model_experiment(model_name, data) retrain_result = self.perform_model_retrain( - prediction_model, data_model, experiment, model_name, data) + prediction_model, data_model, experiment, model_name, data + ) return retrain_result def get_experiment(self, experiment_name: str) -> int: @@ -374,22 +361,21 @@ class MLFlowRepository(): """ runs = mlflow.search_runs( experiment_ids=[experiment_id], - filter_string="", # Sem filtro no MLflow ainda - output_format="pandas" + 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'] + 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) + 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'] @@ -423,16 +409,14 @@ class MLFlowRepository(): # 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) # Colocar a versão do modelo em produção # Depois de registrar o modelo, precisamos pegar a versão mais recente do modelo e movê-lo para o estágio 'Production' client = mlflow.tracking.MlflowClient() # Obter a versão mais recente registrada do modelo - model_versions = client.get_registered_model( - model_name).latest_versions + model_versions = client.get_registered_model(model_name).latest_versions if not isinstance(model_versions, list): raise ValueError('Model versions is not a list') @@ -441,17 +425,10 @@ class MLFlowRepository(): # Mover a versão mais recente do modelo para o estágio de 'Production' 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} def update_production_model(self, experiment: str, model_name: str) -> dict: """ diff --git a/model_manager/worker/worker.py b/model_manager/worker/worker.py index 92631ab..76fc37e 100644 --- a/model_manager/worker/worker.py +++ b/model_manager/worker/worker.py @@ -25,32 +25,35 @@ Environment Variables: - PROJECT_NAME: Project name for notifications (default: model_manager) """ -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 model_manager.workflows.minimal_retrain import MinimalRetrain - from model_manager.workflows.predictions_batch import PredictionsBatch - from model_manager.workflows.sub_workflows.prediction_process import PredictionProcess - from model_manager.workflows.sub_workflows.format_and_export_prediction import \ - FormatAndExportPrediction - from model_manager.activities.activities import Activities - from model_manager.utils.connectors_config import ( - build_postgres_config, - build_mlflow_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 model_manager import metrics - from prometheus_client import start_http_server + from model_manager.activities.activities import Activities + from model_manager.utils.connectors_config import ( + build_mlflow_config, + build_mongodb_config, + build_postgres_config, + ) + from model_manager.workflows.minimal_retrain import MinimalRetrain + from model_manager.workflows.predictions_batch import PredictionsBatch + from model_manager.workflows.sub_workflows.format_and_export_prediction import ( + FormatAndExportPrediction, + ) + from model_manager.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(): @@ -85,7 +88,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) @@ -95,7 +98,7 @@ async def main(): connection_string=mongo_config['connection_string'], database=mongo_config['database_name'], logger=logger, - project_name=os.getenv('PROJECT_NAME', 'model-manager') + project_name=os.getenv('PROJECT_NAME', 'model-manager'), ) logger.custom_info('Starting Activities...', metadata) @@ -104,16 +107,14 @@ async def main(): postgres_config=build_postgres_config(), mlflow_config=build_mlflow_config(), logger=logger, - notification_handler=notification_handler + notification_handler=notification_handler, ) - 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}') ) ) @@ -122,7 +123,7 @@ async def main(): temporal_client = await client.Client.connect( target_host=host, namespace=os.getenv('TEMPORAL_NAMESPACE', 'model-manager'), - runtime=new_runtime + runtime=new_runtime, ) logger.custom_info('Starting Workers...', metadata) @@ -136,20 +137,19 @@ async def main(): activities.load_custom_query, activities.retrain_model, activities.update_production_model, - activities.export_data_to_postgres + activities.export_data_to_postgres, ], 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(), ), Worker( temporal_client, task_queue='predictions_batch-queue', - workflows=[PredictionsBatch, PredictionProcess, - FormatAndExportPrediction], + workflows=[PredictionsBatch, PredictionProcess, FormatAndExportPrediction], activities=[ # MLFlow activities.request_predict, @@ -165,15 +165,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 = [] @@ -186,8 +186,8 @@ async def main(): # This will run the workers and wait for them to complete. # 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) + except BaseException as e: # noqa: BLE001 + logger.custom_error(f'An unhandled exception occurred: {e}', metadata) finally: if notification_handler: notification_handler.shutdown() @@ -216,12 +216,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}") + except Exception as e: # noqa: BLE001 + print(f'Failed to start Prometheus server: {e}') os._exit(1) diff --git a/model_manager/workflows/minimal_retrain.py b/model_manager/workflows/minimal_retrain.py index f6e2c80..03618bf 100644 --- a/model_manager/workflows/minimal_retrain.py +++ b/model_manager/workflows/minimal_retrain.py @@ -1,14 +1,16 @@ from temporalio import workflow with workflow.unsafe.imports_passed_through(): - from model_manager.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 model_manager.activities.activities import Activities -@workflow.defn(name="minimal_retrain") -class MinimalRetrain(): +@workflow.defn(name='minimal_retrain') +class MinimalRetrain: """ Automated model retraining workflow for the Model Manager system. @@ -63,7 +65,7 @@ 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', } } @@ -74,21 +76,17 @@ class MinimalRetrain(): { **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=60) + start_to_close_timeout=timedelta(seconds=60), ) experiment_response = await workflow.execute_activity_method( Activities.retrain_model, - { - **metadata, - 'data': data, - 'model_name': model_name - }, + {**metadata, 'data': data, 'model_name': model_name}, retry_policy=retry_policy, - start_to_close_timeout=timedelta(seconds=60) + start_to_close_timeout=timedelta(seconds=60), ) report = await workflow.execute_activity_method( @@ -97,10 +95,10 @@ class MinimalRetrain(): **metadata, 'model_name': model_name, 'model_id': input_data['model_id'], - **experiment_response + **experiment_response, }, retry_policy=retry_policy, - start_to_close_timeout=timedelta(seconds=60) + start_to_close_timeout=timedelta(seconds=60), ) await workflow.execute_activity_method( @@ -109,8 +107,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=60) + start_to_close_timeout=timedelta(seconds=60), ) diff --git a/model_manager/workflows/predictions_batch.py b/model_manager/workflows/predictions_batch.py index dec0fb7..b221e9f 100644 --- a/model_manager/workflows/predictions_batch.py +++ b/model_manager/workflows/predictions_batch.py @@ -1,14 +1,16 @@ from temporalio import workflow with workflow.unsafe.imports_passed_through(): - from model_manager.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 model_manager.activities.activities import Activities -@workflow.defn(name="predictions_batch") -class PredictionsBatch(): +@workflow.defn(name='predictions_batch') +class PredictionsBatch: """ Main batch prediction workflow for the Model Manager 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,17 @@ 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']), - - '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/model_manager/workflows/sub_workflows/format_and_export_prediction.py b/model_manager/workflows/sub_workflows/format_and_export_prediction.py index 60fad76..2a4266a 100644 --- a/model_manager/workflows/sub_workflows/format_and_export_prediction.py +++ b/model_manager/workflows/sub_workflows/format_and_export_prediction.py @@ -1,15 +1,17 @@ from temporalio import workflow with workflow.unsafe.imports_passed_through(): - from model_manager.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 model_manager.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. @@ -75,10 +77,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: @@ -90,10 +92,10 @@ 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 postgres @@ -104,21 +106,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/model_manager/workflows/sub_workflows/prediction_process.py b/model_manager/workflows/sub_workflows/prediction_process.py index c5407ec..bac0490 100644 --- a/model_manager/workflows/sub_workflows/prediction_process.py +++ b/model_manager/workflows/sub_workflows/prediction_process.py @@ -1,14 +1,16 @@ from temporalio import workflow with workflow.unsafe.imports_passed_through(): - from model_manager.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 model_manager.activities.activities import Activities -@workflow.defn(name="prediction_process") -class PredictionProcess(): +@workflow.defn(name='prediction_process') +class PredictionProcess: """ Core prediction processing workflow for the Model Manager 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), @@ -213,12 +207,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. @@ -264,7 +265,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), @@ -286,8 +287,8 @@ class PredictionProcess(): 'schema': schema, 'table_name': table_name, 'comment': comment, - 'prediction_store_policy': input_data['prediction_store_policy'] - } + 'prediction_store_policy': input_data['prediction_store_policy'], + }, ) return True diff --git a/tests/laborious/activities/test_activities.py b/tests/laborious/activities/test_activities.py index b6e53df..b0cd4e6 100644 --- a/tests/laborious/activities/test_activities.py +++ b/tests/laborious/activities/test_activities.py @@ -1,16 +1,17 @@ +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 model_manager.activities.activities import Activities -from model_manager.activities.mlflow import MLFlow from model_manager.activities.gates import Gates +from model_manager.activities.mlflow import MLFlow @patch('model_manager.activities.activities.Postgres.__init__') @patch('model_manager.activities.activities.MLFlow.__init__') @patch('model_manager.activities.activities.Gates.__init__') def test___init__(mock_gates_init, mock_mlflow_init, mock_postgres_init): - postgres_config = { 'host': 'localhost', 'port': 5432, @@ -18,15 +19,10 @@ def test___init__(mock_gates_init, mock_mlflow_init, mock_postgres_init): 'password': 'postgres', 'dbname': 'postgres', 'min_connections': 1, - 'max_connections': 10 + 'max_connections': 10, } - mlflow_config = { - 'host': 'localhost', - 'port': 5000, - 'username': 'mlflow', - 'password': 'mlflow' - } + mlflow_config = {'host': 'localhost', 'port': 5000, 'username': 'mlflow', 'password': 'mlflow'} logger = MagicMock() notification_handler = MagicMock() @@ -35,7 +31,7 @@ def test___init__(mock_gates_init, mock_mlflow_init, mock_postgres_init): postgres_config=postgres_config, mlflow_config=mlflow_config, logger=logger, - notification_handler=notification_handler + notification_handler=notification_handler, ) assert isinstance(activities, Activities) @@ -53,7 +49,7 @@ def test___init__(mock_gates_init, mock_mlflow_init, mock_postgres_init): min_connections=postgres_config['min_connections'], max_connections=postgres_config['max_connections'], logger=logger, - notification_handler=notification_handler + notification_handler=notification_handler, ) mock_mlflow_init.assert_called_once_with( @@ -63,13 +59,11 @@ def test___init__(mock_gates_init, mock_mlflow_init, mock_postgres_init): mlflow_username=mlflow_config['username'], mlflow_password=mlflow_config['password'], logger=logger, - notification_handler=notification_handler + 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 ) @@ -84,15 +78,10 @@ async def test_shutdown(_mock_mlflow_init, mock_postgres_init): 'password': 'postgres', 'dbname': 'postgres', 'min_connections': 1, - 'max_connections': 10 + 'max_connections': 10, } - mlflow_config = { - 'host': 'localhost', - 'port': 5000, - 'username': 'mlflow', - 'password': 'mlflow' - } + mlflow_config = {'host': 'localhost', 'port': 5000, 'username': 'mlflow', 'password': 'mlflow'} logger = MagicMock() notification_handler = MagicMock() @@ -101,7 +90,7 @@ async def test_shutdown(_mock_mlflow_init, mock_postgres_init): postgres_config=postgres_config, mlflow_config=mlflow_config, logger=logger, - notification_handler=notification_handler + notification_handler=notification_handler, ) await activities.shutdown() diff --git a/tests/laborious/activities/test_gates.py b/tests/laborious/activities/test_gates.py index b587b93..ee68007 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 model_manager.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('model_manager.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('model_manager.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 243982f..80b7e85 100644 --- a/tests/laborious/activities/test_mlflow.py +++ b/tests/laborious/activities/test_mlflow.py @@ -1,45 +1,42 @@ -from datetime import datetime from unittest.mock import ANY, MagicMock, 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 model_manager.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 model_manager.activities.mlflow import MLFlow -@patch("model_manager.activities.mlflow.MLFlowRepository") +@patch('model_manager.activities.mlflow.MLFlowRepository') def test___init__(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', 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) @fixture -@patch("model_manager.activities.mlflow.MLFlowRepository") +@patch('model_manager.activities.mlflow.MLFlowRepository') def mlflow(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', logger=MagicMock(), - notification_handler=MagicMock() + notification_handler=MagicMock(), ) mlflow.send_notification = MagicMock() @@ -48,44 +45,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("model_manager.activities.mlflow.DataFrame") -@patch("model_manager.activities.mlflow.max") +@patch('model_manager.activities.mlflow.DataFrame') +@patch('model_manager.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 +134,25 @@ async def test_request_transform_success(mock_max, mock_dataframe, mlflow): @mark.asyncio -@patch("model_manager.activities.mlflow.DataFrame") -@patch("model_manager.activities.mlflow.to_datetime") -@patch("model_manager.activities.mlflow.max") +@patch('model_manager.activities.mlflow.DataFrame') +@patch('model_manager.activities.mlflow.to_datetime') +@patch('model_manager.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 +163,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 +171,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 @@ -174,28 +185,26 @@ 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] + '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], } mlflow.model_monitoring_repository.retrain_model.return_value = ( - 'Model retrained successfully', 'test') + 'Model retrained successfully', + 'test', + ) - response = await mlflow.retrain_model({ - **metadata, - 'data': data, - 'model_name': 'test_model' - }) + response = await mlflow.retrain_model({**metadata, 'data': data, 'model_name': 'test_model'}) mlflow.model_monitoring_repository.retrain_model.assert_called_once() assert response == { - "status": 'Model retrained successfully', - "timestamp": 2, - "experiment": 'test' + 'status': 'Model retrained successfully', + 'timestamp': 2, + 'experiment': 'test', } @@ -206,20 +215,16 @@ async def test_retrain_model_error(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] + '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], } try: - await mlflow.retrain_model({ - **metadata, - 'data': data, - 'model_name': 'test_model' - }) - except Exception as e: + await mlflow.retrain_model({**metadata, 'data': data, 'model_name': 'test_model'}) + except Exception as e: # noqa: BLE001 assert str(e) == 'Error retraining model' mlflow.send_notification.assert_called_once_with( metadata=metadata['metadata'], @@ -227,20 +232,18 @@ async def test_retrain_model_error(mlflow): message='Error retraining model test_model: Error retraining model', block='retrain_model', level=NotificationLevel.ERROR, - attachment_content=ANY + attachment_content=ANY, ) else: - assert False, "No exception raised" + raise AssertionError('No exception raised') @mark.asyncio async def test_update_production_model(mlflow): - mlflow.model_monitoring_repository.update_production_model.return_value = ( - { - "data1": 1, - "data2": 2 - } - ) + mlflow.model_monitoring_repository.update_production_model.return_value = { + 'data1': 1, + 'data2': 2, + } input_data = { **metadata, @@ -248,13 +251,14 @@ async def test_update_production_model(mlflow): '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' + ) assert response == { 'data1': {0: 1}, @@ -262,7 +266,7 @@ async def test_update_production_model(mlflow): 'model_id': {0: 1}, 'model_name': {0: 'test_model'}, 'timestamp': {0: 2}, - 'status': {0: 'success'} + 'status': {0: 'success'}, } @@ -278,12 +282,12 @@ async def test_update_production_model_error(mlflow): 'model_id': 1, 'experiment': 'test', 'timestamp': 2, - 'status': 'success' + 'status': 'success', } try: await mlflow.update_production_model(input_data) - except Exception as e: + except Exception as e: # noqa: BLE001 assert str(e) == 'Error updating production model' mlflow.send_notification.assert_called_once_with( metadata=metadata['metadata'], @@ -291,7 +295,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/utils/filters/test_conditional_filters.py b/tests/laborious/utils/filters/test_conditional_filters.py index a7f3efd..e23105e 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 model_manager.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 6fddf51..43e96a5 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 model_manager.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, {}) def test_api_error_filter_valid_response_fail(): - assert api_error_filter({'success': False}, {}) == True + assert api_error_filter({'success': False}, {}) def test_api_error_filter_valid_response_success(): - assert api_error_filter({'success': True}, {}) == False + assert not api_error_filter({'success': True}, {}) 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]}), {}) def test_nan_values_filter_no_nan_values(): - assert nan_values_filter(DataFrame({'variable': [1, 2]}), {}) == False + assert not nan_values_filter(DataFrame({'variable': [1, 2]}), {}) diff --git a/tests/laborious/utils/repository/test_model_repository.py b/tests/laborious/utils/repository/test_model_repository.py index ab65913..58fcf7e 100644 --- a/tests/laborious/utils/repository/test_model_repository.py +++ b/tests/laborious/utils/repository/test_model_repository.py @@ -1,34 +1,33 @@ +from datetime import UTC, datetime from unittest.mock import ANY, MagicMock, call, patch + import numpy as np -from pandas import DataFrame import pytest -from datetime import datetime, timezone -from pandas import Timestamp +from pandas import DataFrame, Timestamp + from model_manager.utils.repository.model_repository import MLFlowRepository @pytest.fixture def mlflow_repository(): - with patch('model_manager.utils.repository.model_repository.ModelServing', - autospec=True) as mock_model_serving: + with patch( + 'model_manager.utils.repository.model_repository.ModelServing', autospec=True + ) as mock_model_serving: mock_instance = mock_model_serving.return_value mock_instance.get_transformed_data = MagicMock() repo = MLFlowRepository( - host='http://localhost:5000', - username='admin', - password='admin', - logger=MagicMock() + host='http://localhost:5000', username='admin', password='admin', logger=MagicMock() ) return repo 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', }, } @@ -38,80 +37,56 @@ class Any: 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 - } - } - ) + ({'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) +@pytest.mark.parametrize('data', invalid_cases) def test_detect_and_parse_datetime_index_error_cases(mlflow_repository, data): - input_data = DataFrame( - data - ) + input_data = DataFrame(data) with pytest.raises(ValueError) as e: - mlflow_repository.detect_and_parse_datetime_index( - input_data, metadata['metadata']) + 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" + 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': {'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 + 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'] + }, + ['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 + 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'] + }, + ['2026-01-01 12:00:00+0000', '2026-01-02 12:00:00+0000'], ), ] -@pytest.mark.parametrize("data,expected", valid_cases) +@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']) + response = mlflow_repository.detect_and_parse_datetime_index(input_data, metadata['metadata']) assert response.index.tolist() == expected @@ -122,18 +97,19 @@ def test_transform_success(mlflow_repository): mlflow_repository.detect_and_parse_datetime_index = MagicMock() - output = mlflow_repository.transform( - model_name, data, {}, metadata['metadata']) + 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') + 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']) + 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 + 'content': mlflow_repository.detect_and_parse_datetime_index.return_value.to_dict.return_value, } @@ -141,80 +117,48 @@ def test_transform_error(mlflow_repository): data = MagicMock() model_name = 'model' - mlflow_repository.model_serving.get_cached_transform.side_effect = Exception( - 'error') + mlflow_repository.model_serving.get_cached_transform.side_effect = Exception('error') - output = mlflow_repository.transform( - model_name, data, {}, metadata['metadata']) + 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') + model_name, data, 0, 'sklearn', False, 'model', 'predict' + ) - assert output == { - 'success': False, - 'content': { - 'message': 'error', - 'traceback': ANY - } - } + 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 - } - }) + 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] - ) + mlflow_repository.model_serving.get_cached_predict.return_value = np.array([2, 3]) - output = mlflow_repository.predict( - model_name, data, {}, metadata['metadata']) + 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') + 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 - } + '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 - } - }) + 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') - ) + mlflow_repository.model_serving.get_cached_predict = MagicMock(side_effect=Exception('error')) - output = mlflow_repository.predict( - model_name, data, {}, metadata['metadata']) + 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') + model_name, data, 0, 'pyfunc', False, 'model' + ) - assert output == { - 'success': False, - 'content': { - 'message': 'error', - 'traceback': ANY - } - } + assert output == {'success': False, 'content': {'message': 'error', 'traceback': ANY}} @patch('model_manager.utils.repository.model_repository.mlflow') @@ -246,8 +190,7 @@ def test_get_next_run_name(mlflow, mlflow_repository): @patch('model_manager.utils.repository.model_repository.mlflow') def test_get_experiment_success(mlflow, mlflow_repository): - mlflow.get_experiment_by_name.return_value = MagicMock( - experiment_id='0') + mlflow.get_experiment_by_name.return_value = MagicMock(experiment_id='0') output = mlflow_repository.get_experiment('test') @@ -263,23 +206,25 @@ 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 exception') @patch('model_manager.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'], - }) + 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'], + } + ) output = mlflow_repository.get_experiment_last_run(0) mlflow.search_runs.assert_called_once_with( experiment_ids=[0], - filter_string="", - output_format="pandas", + filter_string='', + output_format='pandas', ) assert output == '2' @@ -294,17 +239,14 @@ def test_get_experiment_last_run_error(mlflow, mlflow_repository): except ValueError as e: assert str(e) == 'Runs is not a pandas DataFrame' else: - assert False + raise AssertionError('Expected exception') @patch('model_manager.utils.repository.model_repository.mlflow.sklearn') @patch('model_manager.utils.repository.model_repository.mlflow.set_experiment') def test_create_model_experiment(set_experiment, sklearn, mlflow_repository): - - 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.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() data_model_mock = MagicMock() @@ -313,29 +255,30 @@ def test_create_model_experiment(set_experiment, sklearn, mlflow_repository): 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.predict.return_value = DataFrame( + { + 'x': [10, 20, 30], + } + ) data_model_mock.target_variable = 'y' prediction_model_mock.fit.return_value = prediction_model_mock - data = DataFrame({ - 'x': [1, 2, 3], - 'y': [4, 5, 6] - }) + data = DataFrame({'x': [1, 2, 3], 'y': [4, 5, 6]}) output = mlflow_repository.create_model_experiment('test', data) 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) + '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"), - ]) + 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 data_model_mock.fit.assert_called_once_with(data) @@ -343,21 +286,23 @@ def test_create_model_experiment(set_experiment, sklearn, mlflow_repository): fit_args = prediction_model_mock.fit.call_args[0][0] assert fit_args.equals( - DataFrame({ - 'x': [10, 20, 30], - 'y': [4, 5, 6], - }) + DataFrame( + { + 'x': [10, 20, 30], + 'y': [4, 5, 6], + } + ) ) 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 - ) + set_experiment.assert_called_once_with(mlflow_repository.get_experiment_by_run_id.return_value) - assert output == (prediction_model_mock, - data_model_mock, - mlflow_repository.get_experiment_by_run_id.return_value) + assert output == ( + prediction_model_mock, + data_model_mock, + mlflow_repository.get_experiment_by_run_id.return_value, + ) @patch('model_manager.utils.repository.model_repository.mlflow.start_run') @@ -365,41 +310,43 @@ def test_create_model_experiment(set_experiment, sklearn, mlflow_repository): @patch('model_manager.utils.repository.model_repository.mlflow.sklearn.log_model') @patch('model_manager.utils.repository.model_repository.mlflow.log_artifact') def test_perform_model_retrain(log_artifact, log_model, log_param, start_run, mlflow_repository): - 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') + mlflow_repository.get_next_run_name = MagicMock(return_value='test-1') run = MagicMock() start_run.__enter__.return_value = run output = mlflow_repository.perform_model_retrain( - prediction_model_mock, data_model_mock, experiment, model_name, data) + prediction_model_mock, data_model_mock, experiment, model_name, data + ) 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') + run_name='test-1', description='Retrain model test with new data' + ) - log_model.assert_has_calls([ - call(data_model_mock, "data_model"), - call(prediction_model_mock, "prediction_model"), - ]) + log_model.assert_has_calls( + [ + call(data_model_mock, 'data_model'), + call(prediction_model_mock, 'prediction_model'), + ] + ) - data.to_csv.assert_called_once_with( - "temp/raw_data_test.csv", index=True) + data.to_csv.assert_called_once_with('temp/raw_data_test.csv', index=True) - log_artifact.assert_called_once_with( - "temp/raw_data_test.csv") + log_artifact.assert_called_once_with('temp/raw_data_test.csv') - log_param.assert_has_calls([ - call("retrain", True), - ]) + log_param.assert_has_calls( + [ + call('retrain', True), + ] + ) - assert output == ("Model retrained successfully", experiment) + assert output == ('Model retrained successfully', experiment) def test_retrain_model(mlflow_repository): @@ -407,18 +354,18 @@ def test_retrain_model(mlflow_repository): model_name = 'test' mlflow_repository.create_model_experiment = MagicMock( - return_value=('data_model', 'prediction_model', '0')) + return_value=('data_model', 'prediction_model', '0') + ) - mlflow_repository.perform_model_retrain = MagicMock( - return_value='Model retrained successfully') + mlflow_repository.perform_model_retrain = MagicMock(return_value='Model retrained successfully') output = mlflow_repository.retrain_model(data, model_name) - mlflow_repository.create_model_experiment.assert_called_once_with( - model_name, data) + mlflow_repository.create_model_experiment.assert_called_once_with(model_name, data) mlflow_repository.perform_model_retrain.assert_called_once_with( - 'data_model', 'prediction_model', '0', model_name, data) + 'data_model', 'prediction_model', '0', model_name, data + ) assert output == 'Model retrained successfully' @@ -438,7 +385,7 @@ def test_update_production_model_by_run_id(mlflow, mlflow_repository): output = mlflow_repository.update_production_model_by_run_id('0', 'test') mlflow.register_model.assert_called_once_with( - "runs:/0/prediction_model", + 'runs:/0/prediction_model', 'test', ) @@ -461,38 +408,34 @@ def test_update_production_model_by_run_id(mlflow, mlflow_repository): @patch('model_manager.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={} - ) - ) + get_registered_model=MagicMock(return_value=MagicMock(latest_versions={})) ) try: mlflow_repository.update_production_model_by_run_id('0', 'test') - except Exception as e: + except Exception as e: # noqa: BLE001 assert str(e) == 'Model versions is not a list' else: - assert False + raise AssertionError('Expected exception') def test_update_production_model(mlflow_repository): connector = mlflow_repository - 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: - + 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 = connector.update_production_model('0', 'test') 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') + update_production_model_by_run_id.assert_called_once_with('2', 'test') assert output == { 'model_name': 'test', diff --git a/tests/laborious/utils/test_connectors_config.py b/tests/laborious/utils/test_connectors_config.py index 6179b33..0f5e005 100644 --- a/tests/laborious/utils/test_connectors_config.py +++ b/tests/laborious/utils/test_connectors_config.py @@ -1,7 +1,10 @@ from os import environ -from model_manager.utils.connectors_config import (build_mlflow_config, - build_postgres_config, - build_mongodb_config) + +from model_manager.utils.connectors_config import ( + build_mlflow_config, + build_mongodb_config, + build_postgres_config, +) def test_build_mlflow_config_with_env_vars(): @@ -95,7 +98,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, } @@ -108,5 +111,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 3606d5d..fb83fac 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,12 @@ -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 model_manager.activities.activities import Activities -from model_manager.workflows.sub_workflows.format_and_export_prediction import FormatAndExportPrediction -from sientia_do.temporal.constants import DATETIME_FORMAT_WITH_TZ +from model_manager.workflows.sub_workflows.format_and_export_prediction import ( + FormatAndExportPrediction, +) @fixture @@ -12,119 +15,133 @@ 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("model_manager.workflows.sub_workflows.format_and_export_prediction.workflow", new_callable=AsyncMock) +@patch( + 'model_manager.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", - "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', + '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.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 == 2 assert workflow_mock.execute_local_activity_method.call_count == 1 @mark.asyncio -@patch("model_manager.workflows.sub_workflows.format_and_export_prediction.workflow", new_callable=AsyncMock) +@patch( + 'model_manager.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", - "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', + '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.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 6922ef1..d45585a 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 model_manager.activities.activities import Activities from model_manager.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("model_manager.workflows.sub_workflows.prediction_process.workflow", new_callable=AsyncMock) +@patch('model_manager.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,23 @@ 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'], - - '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 +61,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', @@ -125,17 +179,16 @@ async def test_run(workflow_mock, prediction_process): 'model_id': 1, 'model_name': 'test_model_name', 'model_config': input_data['model_config'], - '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("model_manager.workflows.sub_workflows.prediction_process.workflow", new_callable=AsyncMock) +@patch('model_manager.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 +202,14 @@ 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'], - } # 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 +217,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("model_manager.workflows.sub_workflows.prediction_process.workflow", new_callable=AsyncMock) +@patch('model_manager.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 +259,16 @@ 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'], - } # 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 +276,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("model_manager.workflows.sub_workflows.prediction_process.workflow", new_callable=AsyncMock) +@patch('model_manager.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 +353,19 @@ 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'], - } # 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 +374,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("model_manager.workflows.sub_workflows.prediction_process.workflow", new_callable=AsyncMock) +@patch('model_manager.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 +467,21 @@ 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'], - } # 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 +489,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("model_manager.workflows.sub_workflows.prediction_process.workflow", new_callable=AsyncMock) +@patch('model_manager.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 +610,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 +637,7 @@ async def test_path_flag_handler_stop(workflow_mock, prediction_process): @mark.asyncio -@patch("model_manager.workflows.sub_workflows.prediction_process.workflow", new_callable=AsyncMock) +@patch('model_manager.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 +648,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 +680,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("model_manager.workflows.sub_workflows.prediction_process.workflow", new_callable=AsyncMock) +@patch('model_manager.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 +697,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, @@ -536,9 +712,11 @@ async def test_path_flag_handler_continue(workflow_mock, prediction_process): 'last_timestamp': last_timestamp, 'model_name': model_name, 'model_config': model_config, - - 'prediction_store_policy': prediction_store_policy - }, confidence, last_timestamp, 'Prediction Process' + 'prediction_store_policy': prediction_store_policy, + }, + confidence, + last_timestamp, + 'Prediction Process', ) # Assert @@ -558,14 +736,13 @@ async def test_path_flag_handler_continue(workflow_mock, prediction_process): 'schema': schema, 'table_name': table_name, 'comment': 'Prediction Process', - - 'prediction_store_policy': prediction_store_policy - } + 'prediction_store_policy': prediction_store_policy, + }, ) @mark.asyncio -@patch("model_manager.workflows.sub_workflows.prediction_process.workflow", new_callable=AsyncMock) +@patch('model_manager.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 +753,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, @@ -590,9 +767,11 @@ async def test_path_flag_handler_unknown(workflow_mock, prediction_process): 'last_timestamp': last_timestamp, 'model_name': model_name, 'model_config': model_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 1d8e6d7..5296e80 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 model_manager.activities.activities import Activities from model_manager.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,19 +25,19 @@ metadata = { @patch('model_manager.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', } workflow_mock.execute_activity_method = AsyncMock( return_value={ - "data1": "1", - "data2": "2", + 'data1': '1', + 'data2': '2', } ) @@ -47,52 +49,58 @@ async def test_run(workflow_mock: AsyncMock, minimal_retrain: MinimalRetrain): Activities.load_custom_query, { **metadata, - "query": input_data["query"], - 'datetime_columns': input_data.get('datetime_columns', []) + 'query': input_data['query'], + 'datetime_columns': input_data.get('datetime_columns', []), }, 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, + '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.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'], + '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.export_data_to_postgres, - { - **metadata, - 'data': workflow_mock.execute_activity_method.return_value, - 'schema': input_data['schema'], - 'table_name': input_data['table_name'], - }, - retry_policy=ANY, - start_to_close_timeout=ANY - ) - ]) + workflow_mock.execute_activity_method.assert_has_calls( + [ + call( + Activities.export_data_to_postgres, + { + **metadata, + 'data': workflow_mock.execute_activity_method.return_value, + 'schema': input_data['schema'], + 'table_name': input_data['table_name'], + }, + retry_policy=ANY, + start_to_close_timeout=ANY, + ) + ] + ) diff --git a/tests/laborious/workflows/test_predictions_batch.py b/tests/laborious/workflows/test_predictions_batch.py index e71a462..c49ffea 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 model_manager.activities.activities import Activities from model_manager.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('model_manager.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', @@ -32,28 +32,27 @@ async def test_run(workflow_mock: AsyncMock, predictions_batch: PredictionsBatch 'query': 'SELECT * FROM test', 'schema': 'test_schema', 'table_name': 'test_table', - '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 +60,18 @@ 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']), - - '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)] + )