From 6ac0f38d59dad6da31a7911940f3a2d540926776 Mon Sep 17 00:00:00 2001 From: vitor-aignosi Date: Tue, 11 Nov 2025 16:50:22 -0300 Subject: [PATCH 01/27] SIENTIAPDE-1273 Update requirements and enhance metrics and data handling - Updated the sientia-dataops-library dependency version in requirements.txt to 1.5.3. - Added new metrics for model analysis, including lag, count, and error count in metrics.py. - Implemented a new method for formatting transformed data in gates.py. - Enhanced MLFlowRepository with methods to load artifact dataframes and calculate model metrics, including drift and performance metrics. - Updated the prediction process to handle transformed data and ensure proper execution of related activities in format_and_export_prediction.py and prediction_process.py. --- laborious/activities/gates.py | 22 +++++ laborious/metrics.py | 19 ++++ .../utils/repository/model_repository.py | 93 +++++++++++++++++++ .../format_and_export_prediction.py | 41 +++++++- .../sub_workflows/prediction_process.py | 2 + requirements.txt | 2 +- 6 files changed, 177 insertions(+), 2 deletions(-) diff --git a/laborious/activities/gates.py b/laborious/activities/gates.py index 9dc9eb1..56a1569 100644 --- a/laborious/activities/gates.py +++ b/laborious/activities/gates.py @@ -402,6 +402,28 @@ class Gates(SientiaMonitoring): return policy_type, int(policy_value) + + @activity.defn(name='format_transformed_data') + async def format_transformed_data(self, input_data: dict[str, Any]) -> dict[Any, Any]: + """ + Format transformed data according to configured storage policies. + """ + metadata = input_data['metadata'] + + model_id = input_data['model_id'] + + self.info('Formatting transformed data...', metadata) + + data = DataFrame(input_data['data']) + + data['timestamp'] = data.index + data = data.reset_index(drop=True) + + data = data.melt(id_vars='timestamp', var_name='variable', value_name='value') + data['model_id'] = model_id + + return data.to_dict() + @activity.defn(name='format_prediction') async def format_prediction(self, input_data: dict[str, Any]) -> dict[Any, Any]: """ diff --git a/laborious/metrics.py b/laborious/metrics.py index ec87867..59e4a81 100644 --- a/laborious/metrics.py +++ b/laborious/metrics.py @@ -169,3 +169,22 @@ MODEL_WRITE_ERROR_COUNT = Counter( 'Number of errors writing to the model', SIENTIA_CORE_LABELS, ) + +MODEL_ANALYZE_LAG = Histogram( + 'laborious_model_analyze_lag', + 'Lag between the start and end of analyze operations', + SIENTIA_CORE_LABELS, + buckets=[0.01, 0.05, 0.1, 0.2, 0.5, 1.0, 2.0, 5.0, 10.0], +) + +MODEL_ANALYZE_COUNT = Counter( + 'laborious_model_analyze_count', + 'Number of analyze operations', + SIENTIA_CORE_LABELS, +) + +MODEL_ANALYZE_ERROR_COUNT = Counter( + 'laborious_model_analyze_error_count', + 'Number of errors during analyze operations', + SIENTIA_CORE_LABELS, +) \ No newline at end of file diff --git a/laborious/utils/repository/model_repository.py b/laborious/utils/repository/model_repository.py index 1414713..edd9671 100644 --- a/laborious/utils/repository/model_repository.py +++ b/laborious/utils/repository/model_repository.py @@ -16,6 +16,7 @@ Capabilities: import ctypes import gc +from io import StringIO import threading import time import traceback @@ -33,6 +34,7 @@ from sientia_do.observability.logger import Logger from sientia_do.observability.metrics_controller import MetricsController from sientia_do.observability.sientia_monitoring import SientiaMonitoring from sientia_do.temporal.constants import DATETIME_FORMAT_WITH_TZ +from sientia.ModelAnalysis import ModelAnalysis from laborious import metrics @@ -259,6 +261,40 @@ class MLFlowRepository(SientiaMonitoring): return artifacts + + async def load_artifact_dataframe(self, model_name: str, artifact_path: str, + metadata: dict[str, Any]) -> pd.DataFrame: + + """ + Load the dataframe content of an artifact from the MLflow Model Registry. + + Args: + model_name (str): The name of the model to download from the registry. + artifact_path (str): The path to the artifact to load. + metadata (dict[str, Any]): Metadata used for structured logging. + + Returns: + pd.DataFrame: The dataframe content of the artifact. + """ + run_id = self.get_model_run_id(model_name=model_name, stage='Production') + artifact_path = path.join("runs:/", run_id, artifact_path) + core_labels = self.get_core_labels(metadata, operation_type='load_text') + + start_time = time.time() + try: + content = mlflow.artifacts.load_text(artifact_path) + except Exception as e: + await self.emit_metric(metric_object=metrics.MODEL_READ_ERROR_COUNT, tags=core_labels) + raise e + + await self.observe_lag(start_time, metrics.MODEL_READ_LAG, core_labels) + await self.emit_metric(metric_object=metrics.MODEL_READ_COUNT, tags=core_labels) + + dataframe = pd.read_csv(StringIO(content)) + + self.info(f'Loaded dataframe from {model_name}:{artifact_path}', metadata) + return dataframe + async def load_predict_model( self, model_name: str, metadata: dict[str, Any], flavor: str = 'sklearn' ) -> Any: @@ -1335,3 +1371,60 @@ class MLFlowRepository(SientiaMonitoring): metadata_result['mlflow_experiment_id'] = experiment_id return metadata_result + + async def get_model_metrics(self, + analyse_data: pd.DataFrame, + train_reference_data: pd.DataFrame, + test_reference_data: pd.DataFrame, + target_name: str, + drift_metrics: list[str], + performance_metrics: list[str], + metadata: dict) -> dict[str, Any]: + + """ + Calculates drift and preformance metrics for a model. + + Args: + analyse_data (pd.DataFrame): The data to analyse. + train_reference_data (pd.DataFrame): The train reference data. + test_reference_data (pd.DataFrame): The test reference data. + target_name (str): The target name. + drift_metrics (list[str]): The drift metrics. + performance_metrics (list[str]): The performance metrics. + metadata (dict): The metadata. + + Returns: + dict: The metrics. + """ + columns = train_reference_data.drop( + columns=[target_name, 'timestamp', 'target', 'prediction'], errors='ignore').columns + + config = { + "target": target_name, + "prediction": "prediction", + "timestamp": "timestamp", + "columns": columns, + } + + model_analysis = ModelAnalysis(config=config) + + core_labels = self.get_core_labels(metadata, operation_type='detect_univariate_drift') + + start_time = time.time() + try: + model_analysis.detect_univariate_drift( + reference_df=train_reference_data, + analyse_df=analyse_data, + features=columns, + timestamp_column=config['timestamp'], + metrics=drift_metrics, + chunk_period="s" + ) + except Exception as e: + await self.emit_metric(metric_object=metrics.MODEL_ANALYZE_ERROR_COUNT, tags=core_labels) + raise e + + await self.observe_lag(start_time, metrics.MODEL_ANALYZE_LAG, core_labels) + await self.emit_metric(metric_object=metrics.MODEL_ANALYZE_COUNT, tags=core_labels) + + return {} diff --git a/laborious/workflows/sub_workflows/format_and_export_prediction.py b/laborious/workflows/sub_workflows/format_and_export_prediction.py index 8d8c9da..bd28966 100644 --- a/laborious/workflows/sub_workflows/format_and_export_prediction.py +++ b/laborious/workflows/sub_workflows/format_and_export_prediction.py @@ -69,6 +69,7 @@ class FormatAndExportPrediction: metadata = input_data['metadata'] path_flag = input_data['path_flag'] data = input_data['data'] + transformed_data = input_data.get('transformed_data', None) prediction_confidence = input_data['prediction_confidence'] if path_flag is None: @@ -87,6 +88,37 @@ class FormatAndExportPrediction: start_to_close_timeout=timedelta(seconds=60), ) + if transformed_data is not None: + + transformed = await workflow.execute_local_activity_method( + Activities.format_transformed_data, + { + **metadata, + 'data': transformed_data, + }, + retry_policy=retry_policy, + start_to_close_timeout=timedelta(seconds=60), + ) + + write_transformed_handler = workflow.execute_activity_method( + Activities.export_data_to_postgres, + { + **metadata, + 'schema': input_data['schema'], + 'table_name': input_data['transform_table_name'], + 'data': transformed, + 'timestamp_conversion': { + 'column': 'timestamp', + 'format': DATETIME_FORMAT_WITH_TZ, + }, + }, + retry_policy=retry_policy, + start_to_close_timeout=timedelta(seconds=60), + ) + + else: + write_transformed_handler = None + else: # create default prediction prediction = await workflow.execute_local_activity_method( @@ -102,6 +134,8 @@ class FormatAndExportPrediction: start_to_close_timeout=timedelta(seconds=60), ) + write_transformed_handler = None + # write to opc prediction, opc_metrics = await workflow.execute_activity_method( Activities.write_opc_data, @@ -115,7 +149,7 @@ class FormatAndExportPrediction: ) # write to postgres - await workflow.execute_activity_method( + prediction_handler = workflow.execute_activity_method( Activities.export_data_to_postgres, { **metadata, @@ -128,6 +162,11 @@ class FormatAndExportPrediction: start_to_close_timeout=timedelta(seconds=180), ) + await prediction_handler + + if write_transformed_handler is not None: + await write_transformed_handler + await workflow.execute_activity_method( Activities.write_metrics, { diff --git a/laborious/workflows/sub_workflows/prediction_process.py b/laborious/workflows/sub_workflows/prediction_process.py index ab78fb1..fa07504 100644 --- a/laborious/workflows/sub_workflows/prediction_process.py +++ b/laborious/workflows/sub_workflows/prediction_process.py @@ -82,6 +82,7 @@ class PredictionProcess: model_id = input_data['model_id'] model_name = input_data['model_name'] model_config = input_data.get('model_config', {}) + save_transform = input_data.get('save_transform', True) # Get last timestamp for incremental processing last_timestamp = await workflow.execute_local_activity_method( @@ -199,6 +200,7 @@ class PredictionProcess: 'metadata': metadata, 'path_flag': path_flag, 'data': response_data['content'], + 'transformed_data': transformed_data if save_transform else None, 'prediction_confidence': confidence, 'timestamp': last_timestamp, 'model_id': model_id, diff --git a/requirements.txt b/requirements.txt index 12feb2e..750ef66 100644 --- a/requirements.txt +++ b/requirements.txt @@ -3,7 +3,7 @@ psycopg2-binary sqlalchemy asyncua redis -git+ssh://git@github.com/Aignosi/sientia-dataops-library.git@1.5.2 +git+ssh://git@github.com/Aignosi/sientia-dataops-library.git@1.5.3 git+ssh://git@github.com/Aignosi/sientia-mlops-library.git@0.39.0 prometheus-client botocore From b68674fe64293cf8835557588517b9779cb36be4 Mon Sep 17 00:00:00 2001 From: vitor-aignosi Date: Thu, 13 Nov 2025 16:40:23 -0300 Subject: [PATCH 02/27] SIENTIAPDE-1273 Enhance MLFlowRepository and Activities classes with new methods and metrics - Added `check_artifact_exists` method to MLFlowRepository for verifying artifact presence in the MLflow Model Registry. - Implemented `get_prediction_data` method in MLFlowRepository to retrieve prediction data from models. - Updated Activities class to integrate ModelMetrics for improved metrics handling. - Enhanced tests for artifact existence checks and prediction data retrieval, ensuring robust coverage for new functionalities. - Updated various workflows to include `transform_table_name` in input data for better data handling. --- laborious/activities/activities.py | 12 +- laborious/activities/mlflow.py | 34 + laborious/activities/model_metrics.py | 266 ++++++++ .../utils/repository/model_repository.py | 132 ++-- laborious/worker/worker.py | 20 + laborious/workflows/drift.py | 99 +++ laborious/workflows/predictions_batch.py | 1 + .../format_and_export_prediction.py | 3 +- .../sub_workflows/prediction_process.py | 3 + tests.ipynb | 101 +++ .../activities/test_model_metrics.py | 641 ++++++++++++++++++ .../utils/repository/test_model_repository.py | 132 +++- .../subworkflows/test_prediction_process.py | 12 + tests/laborious/workflows/test_drift.py | 247 +++++++ .../workflows/test_predictions_batch.py | 4 +- 15 files changed, 1639 insertions(+), 68 deletions(-) create mode 100644 laborious/activities/model_metrics.py create mode 100644 laborious/workflows/drift.py create mode 100644 tests/laborious/activities/test_model_metrics.py create mode 100644 tests/laborious/workflows/test_drift.py diff --git a/laborious/activities/activities.py b/laborious/activities/activities.py index e09d95f..d9c07ac 100644 --- a/laborious/activities/activities.py +++ b/laborious/activities/activities.py @@ -11,9 +11,9 @@ with workflow.unsafe.imports_passed_through(): from laborious.activities.mlflow import MLFlow from laborious.activities.opc import OPC from laborious.activities.storage import Storage + from laborious.activities.model_metrics import ModelMetrics - -class Activities(Storage, MLFlow, Gates, OPC): +class Activities(Storage, MLFlow, Gates, OPC, ModelMetrics): """ Main activities orchestrator for the Laborious system. @@ -108,6 +108,13 @@ class Activities(Storage, MLFlow, Gates, OPC): metrics_controller=metrics_controller, ) + ModelMetrics.__init__( + self, + logger=logger, + notification_handler=notification_handler, + metrics_controller=metrics_controller, + ) + async def shutdown(self): """ Gracefully shutdown all activities and clean up resources. @@ -124,3 +131,4 @@ class Activities(Storage, MLFlow, Gates, OPC): MLFlow.close(self) Gates.close(self) await OPC.close(self) + ModelMetrics.close(self) \ No newline at end of file diff --git a/laborious/activities/mlflow.py b/laborious/activities/mlflow.py index 98be4f2..2c50abe 100644 --- a/laborious/activities/mlflow.py +++ b/laborious/activities/mlflow.py @@ -1,3 +1,4 @@ +from typing import Hashable from temporalio import activity, workflow with workflow.unsafe.imports_passed_through(): @@ -418,3 +419,36 @@ class MLFlow(SientiaMonitoring): ) self.error(trace, metadata=metadata) raise e + + + @activity.defn(name='get_reference_data') + async def get_reference_data(self, input_data: dict[str, Any]) -> dict[Hashable, Any] | None: + """ + Get reference data from the MLflow Model Registry. + + Args: + input_data (dict): Input data containing: + - metadata (dict): Workflow execution metadata + - model_name (str): Name of the MLFlow model to get reference data from + + Returns: + dict[Hashable, Any] | None: Reference data from the MLflow Model Registry. + + """ + + metadata = input_data['metadata'] + model_name = input_data['model_name'] + artifact = "evaluation_data.csv" + + reference_data = await self.model_monitoring_repository.load_artifact_dataframe( + model_name=model_name, artifact_path=artifact, metadata=metadata + ) + + if reference_data is None: + self.warning(f'Reference data not found for model {model_name}', metadata) + return None + + reference_data['timestamp'] = to_datetime(reference_data['timestamp']) + reference_data['timestamp'] = reference_data['timestamp'].dt.strftime(DATETIME_FORMAT) + + return reference_data.to_dict() \ No newline at end of file diff --git a/laborious/activities/model_metrics.py b/laborious/activities/model_metrics.py new file mode 100644 index 0000000..4e24ec7 --- /dev/null +++ b/laborious/activities/model_metrics.py @@ -0,0 +1,266 @@ +from temporalio import activity, workflow + +with workflow.unsafe.imports_passed_through(): + from typing import Any, Hashable + + from pandas import DataFrame, Index, to_datetime + + from sientia_do.observability.sientia_monitoring import SientiaMonitoring + from sientia_do.observability.metrics_controller import MetricsController + from sientia_do.observability.logger import Logger + from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler + from sientia_do.notifications.models import NotificationLevel + from sientia.ModelAnalysis import ModelAnalysis + from laborious import metrics + import time + from sientia_do.temporal.constants import DATETIME_FORMAT, DATETIME_FORMAT_WITH_TZ + import warnings + import traceback + +warnings.filterwarnings('ignore', category=RuntimeWarning, message='Degrees of freedom <= 0') +warnings.filterwarnings('ignore', category=RuntimeWarning, message='invalid value encountered in scalar divide') + + +class ModelMetrics(SientiaMonitoring): + """ + Metrics activities for the Laborious system. + + This class provides activities for writing metrics to the Prometheus monitoring system. + """ + + def __init__(self, + logger: Logger, + notification_handler: NotificationHandler, + metrics_controller: MetricsController, + ): + + SientiaMonitoring.__init__(self, logger, notification_handler, metrics_controller) + + def close(self) -> None: + """ + Close the model metrics activity and clean up resources. + """ + SientiaMonitoring.shutdown(self) + + + def __del__(self): + self.close() + + async def get_drift_metrics(self, + reference_data: DataFrame, + target_data: DataFrame, + target_name: str, + reference_columns: Index, + drift_metrics: list[str], + chunk_period: str, + metadata: dict[str, Any], + ) -> DataFrame: + """ + Calculate univariate drift metrics for a model. + Args: + model_analysis (ModelAnalysis): Model analysis object + reference_data (DataFrame): Reference data + target_data (DataFrame): Target data + reference_columns (list[str]): Reference columns + drift_metrics (list[str]): Drift metrics + metadata (dict[str, Any]): Workflow execution metadata + """ + + config = { + 'target': target_name, + 'prediction': 'prediction', + 'timestamp': 'timestamp', + 'features': reference_columns, + } + + + model_analysis = ModelAnalysis(config=config) + + self.debug(f'Reference data: Size {reference_data.shape} \n{reference_data.head(5).to_string()}', metadata) + + self.debug(f'Target data: Size {target_data.shape} \n{target_data.head(5).to_string()}', metadata) + + core_labels = self.get_core_labels(metadata, operation_type='detect_univariate_drift') + start_time = time.time() + try: + + univariate_drift = model_analysis.detect_univariate_drift( + reference_df=reference_data, + analysis_df=target_data, + features=reference_columns, + timestamp_col=config['timestamp'], + methods=drift_metrics, + chunk_period=chunk_period + ) + except Exception as e: + self.error(f'Error detecting univariate drift: {e}', metadata) + await self.emit_metric( + metric_object=metrics.MODEL_ANALYZE_ERROR_COUNT, tags=core_labels) + raise e + await self.observe_lag(start_time, metrics.MODEL_ANALYZE_LAG, core_labels) + await self.emit_metric(metric_object=metrics.MODEL_ANALYZE_COUNT, tags=core_labels) + + + core_labels = self.get_core_labels(metadata, operation_type='detect_multivariate_drift') + start_time = time.time() + try: + multivariate_drift = model_analysis.detect_multivariate_drift( + reference_df=reference_data, + analysis_df=target_data, + features=reference_columns, + timestamp_col=config['timestamp'], + chunk_period=chunk_period + ) + except Exception as e: + self.error(f'Error detecting multivariate drift: {e}', metadata) + await self.emit_metric(metric_object=metrics.MODEL_ANALYZE_ERROR_COUNT, tags=core_labels) + raise e + await self.observe_lag(start_time, metrics.MODEL_ANALYZE_LAG, core_labels) + await self.emit_metric(metric_object=metrics.MODEL_ANALYZE_COUNT, tags=core_labels) + + start_time = time.time() + core_labels = self.get_core_labels(metadata, operation_type='get_drift_metrics_dataframe') + try: + drift_df = model_analysis.get_drift_metrics_dataframe( + univariate_drift=univariate_drift, + multivariate_drift=multivariate_drift, + ) + except Exception as e: + self.error(f'Error getting drift metrics: {e}', metadata) + await self.emit_metric(metric_object=metrics.MODEL_ANALYZE_ERROR_COUNT, tags=core_labels) + raise e + await self.observe_lag(start_time, metrics.MODEL_ANALYZE_LAG, core_labels) + await self.emit_metric(metric_object=metrics.MODEL_ANALYZE_COUNT, tags=core_labels) + + self.debug(f'Drift dataframe: Size {drift_df.shape} \n{drift_df.head(5).to_string()}', metadata) + + return drift_df + + + @activity.defn(name='calculate_drift') + async def calculate_drift(self, input_data: dict[str, Any]) -> dict[Hashable, Any]: + """ + Calculate drift metrics for a model. + + Args: + input_data (dict[str, Any]): Input data containing: + - metadata (dict): Workflow execution metadata + - model_name (str): Name of the MLFlow model to calculate drift for + - reference_data (pd.DataFrame): Reference data for the model + - target_data (pd.DataFrame): Target data for calculating drift + - target_name (str): Name of the target column + - drift_metrics (list[str]): List of drift metrics to calculate + """ + metadata = input_data['metadata'] + model_name = input_data['model_name'] + model_id = input_data['model_id'] + reference_raw_data = input_data['reference_data'] + target_data = DataFrame(input_data['target_data']) + target_name = input_data['target_name'] + drift_metrics = input_data['drift_metrics'] + chunk_period = input_data['chunk_period'] + + if chunk_period not in ['min', 's']: + self.error(f'Invalid chunk period: {chunk_period}', metadata) + raise ValueError(f'Invalid chunk period: {chunk_period}, must be "min" or "s"') + + self.info(f'Calculating drift for model {model_name}', metadata) + + target_data = target_data.pivot(index='timestamp', columns='variable', values='value') + target_data['timestamp'] = target_data.index + target_data['timestamp'] = to_datetime(target_data['timestamp']) + target_data['timestamp'] = target_data['timestamp'].dt.strftime(DATETIME_FORMAT) + target_data = target_data.reset_index(drop=True) + target_data.dropna(inplace=True) + + if reference_raw_data is not None: + self.info('Using reference data', metadata) + reference_data = DataFrame(reference_raw_data) + accurate = True + else: + # Get 30% first rows of target_data + self.warning('Using 30% first rows of target data as reference data', metadata) + target_data.sort_values(by='timestamp', ascending=True, inplace=True) + reference_data = target_data.head(int(len(target_data) * 0.3)) + accurate = False + + await self.send_notification_async( + metadata=metadata, + notification_id='MODEL_METRICS_REFERENCE_DATA_WARNING', + message='Using 30% first rows of target data as reference data', + block='model_metrics', + level=NotificationLevel.WARNING, + attachment_content=reference_data.to_csv(), + ) + + reference_columns = reference_data.drop( + columns=[target_name, 'timestamp', 'target', 'prediction'], + errors='ignore').columns + + try: + drift_df = await self.get_drift_metrics( + reference_data=reference_data, + target_data=target_data, + target_name=target_name, + reference_columns=reference_columns, + drift_metrics=drift_metrics, + chunk_period=chunk_period, + metadata=metadata, + ) + except Exception as e: + self.error(f'Error getting drift metrics: {e}', metadata) + await self.send_notification_async( + metadata=metadata, + notification_id='MODEL_METRICS_GET_DRIFT_METRICS_ERROR', + message=f'Error getting drift metrics: {e}', + block='model_metrics', + level=NotificationLevel.ERROR, + attachment_content=traceback.format_exc(), + ) + return {} + + if drift_df.empty: + self.warning('No drift metrics found', metadata) + return {} + + # Drop unnecessary columns + drift_df.drop(columns=['p_value', 'chunk_start_date', 'chunk_end_date'], inplace=True) + + # Extract timestamps only until minutes + if chunk_period == 'min': + target_timestamps = target_data['timestamp'].apply( + lambda x: x[:16]) + else: + target_timestamps = target_data['timestamp'] + + # Drop rows where timestamp is not in target data, to avoid save drift from reference + drift_df = drift_df[drift_df['timestamp'].isin(target_timestamps)] + + if drift_df.empty: + self.warning('No drift metrics found after dropping rows where timestamp is not in target data', metadata) + return {} + + # Rename columns to match database columns + drift_df.rename(columns={ + 'metric': 'method', + 'statistic': 'value', + }, inplace=True) + + # Drop duplicates + drift_df.drop_duplicates( + subset=['timestamp', 'method', 'feature'], + keep='first', inplace=True) + + drift_df['model_id'] = model_id + drift_df['accurate'] = accurate + + + drift_df['timestamp'] = to_datetime(drift_df['timestamp']) + drift_df['timestamp'] = drift_df['timestamp'].dt.tz_localize('UTC') + drift_df['timestamp'] = drift_df['timestamp'].dt.strftime(DATETIME_FORMAT_WITH_TZ) + + self.debug(f'Drift dataframe: Size {drift_df.shape} \n{drift_df.head(5).to_string()}', metadata) + + + return drift_df.to_dict() + diff --git a/laborious/utils/repository/model_repository.py b/laborious/utils/repository/model_repository.py index edd9671..e80d095 100644 --- a/laborious/utils/repository/model_repository.py +++ b/laborious/utils/repository/model_repository.py @@ -217,6 +217,24 @@ class MLFlowRepository(SientiaMonitoring): run_info = mlflow.get_run(run_id) return run_info.data.params + def check_artifact_exists(self, run_id: str, + artifact_path: str, metadata: dict[str, Any]) -> bool: + """ + Check if an artifact exists in the MLflow Model Registry. + + Args: + run_id (str): Run identifier to inspect. + artifact_path (str): Path to the artifact to check. + + Returns: + bool: True if the artifact exists, False otherwise. + """ + artifacts = self.client.list_artifacts(run_id) + + self.debug(f'Artifacts of {run_id}: \n{artifacts}', metadata) + self.debug(f'Looking for artifact {artifact_path} in {run_id}', metadata) + + return any(artifact.path == artifact_path for artifact in artifacts) """ Functions related to download and load models """ @@ -263,7 +281,7 @@ class MLFlowRepository(SientiaMonitoring): async def load_artifact_dataframe(self, model_name: str, artifact_path: str, - metadata: dict[str, Any]) -> pd.DataFrame: + metadata: dict[str, Any]) -> pd.DataFrame | None: """ Load the dataframe content of an artifact from the MLflow Model Registry. @@ -277,9 +295,13 @@ class MLFlowRepository(SientiaMonitoring): pd.DataFrame: The dataframe content of the artifact. """ run_id = self.get_model_run_id(model_name=model_name, stage='Production') - artifact_path = path.join("runs:/", run_id, artifact_path) core_labels = self.get_core_labels(metadata, operation_type='load_text') + if not self.check_artifact_exists(run_id, artifact_path, metadata): + return None + + artifact_path = path.join("runs:/", run_id, artifact_path) + start_time = time.time() try: content = mlflow.artifacts.load_text(artifact_path) @@ -290,6 +312,8 @@ class MLFlowRepository(SientiaMonitoring): await self.observe_lag(start_time, metrics.MODEL_READ_LAG, core_labels) await self.emit_metric(metric_object=metrics.MODEL_READ_COUNT, tags=core_labels) + self.debug(f'Content of {run_id}/{artifact_path}: \n{content}', metadata) + dataframe = pd.read_csv(StringIO(content)) self.info(f'Loaded dataframe from {model_name}:{artifact_path}', metadata) @@ -680,6 +704,41 @@ class MLFlowRepository(SientiaMonitoring): Functions related to model retraining """ + def get_prediction_data(self, prediction_model: Any, retrain_dataset: pd.DataFrame, + target_name: str) -> pd.DataFrame: + """ + Get prediction data from prediction model. + """ + input_index = retrain_dataset.index + + prediction_data = prediction_model.predict(retrain_dataset) + + if isinstance(prediction_data, pd.DataFrame): + + prediction_data.columns = pd.Index(['prediction']) + + else: + prediction_data = pd.DataFrame(prediction_data, columns=['prediction']) + + prediction_data.index = input_index + + # Merge prediction data with retrain_dataset on index + prediction_data = pd.merge( + retrain_dataset, prediction_data, left_index=True, right_index=True, how='left') + + # Rename column "target_name" to "target" + prediction_data.rename(columns={target_name: 'target'}, inplace=True) + + prediction_data['timestamp'] = prediction_data.index + + prediction_data.reset_index(drop=True, inplace=True) + + prediction_data.sort_values( + by='timestamp', ascending=True, inplace=True + ) + + return prediction_data + async def fit_models( self, model_name: str, @@ -689,7 +748,7 @@ class MLFlowRepository(SientiaMonitoring): transform_flavor: str = 'sklearn', predict_flavor: str = 'sklearn', target_name: str | None = None, - ) -> dict[str, dict[str, Any]]: + ) -> dict[str, Any]: """ Prepare models and data for a retraining run. @@ -799,6 +858,10 @@ class MLFlowRepository(SientiaMonitoring): prediction_model.fit(retrain_dataset) + # get prediction data + prediction_data = self.get_prediction_data( + prediction_model, retrain_dataset, target_name) + self.info(f'Model experiment creation completed successfully for {model_name}', metadata) retrain_data = { @@ -807,6 +870,7 @@ class MLFlowRepository(SientiaMonitoring): 'artifact_path': prediction_artifact_path, }, 'data_model': {'model': data_model, 'artifact_path': data_artifact_path}, + 'prediction_data': prediction_data, } return retrain_data @@ -885,6 +949,7 @@ class MLFlowRepository(SientiaMonitoring): prediction_model = retrain_data['prediction_model'] data_model = retrain_data['data_model'] + prediction_data = retrain_data['prediction_data'] model_temp_path = path.join(ARTIFACTS_PATH, model_name) @@ -908,10 +973,12 @@ class MLFlowRepository(SientiaMonitoring): self.debug(f'Attributes: {retrain_params}', metadata) data_path = f'{model_temp_path}/retrain_data.csv' + prediction_data_path = f'{model_temp_path}/evaluation_data.csv' makedirs(model_temp_path, exist_ok=True) - data.to_csv(data_path, index=True) + data.to_csv(data_path, index=False) + prediction_data.to_csv(prediction_data_path, index=False) self.info( f'Starting model upload for {experiment_name} with run name {current_run_name}', @@ -945,6 +1012,7 @@ class MLFlowRepository(SientiaMonitoring): # log the data raw mlflow.log_artifact(data_path) + mlflow.log_artifact(prediction_data_path) except Exception as e: await self.emit_metric(metric_object=metrics.MODEL_WRITE_ERROR_COUNT, tags=core_labels) @@ -1372,59 +1440,3 @@ class MLFlowRepository(SientiaMonitoring): return metadata_result - async def get_model_metrics(self, - analyse_data: pd.DataFrame, - train_reference_data: pd.DataFrame, - test_reference_data: pd.DataFrame, - target_name: str, - drift_metrics: list[str], - performance_metrics: list[str], - metadata: dict) -> dict[str, Any]: - - """ - Calculates drift and preformance metrics for a model. - - Args: - analyse_data (pd.DataFrame): The data to analyse. - train_reference_data (pd.DataFrame): The train reference data. - test_reference_data (pd.DataFrame): The test reference data. - target_name (str): The target name. - drift_metrics (list[str]): The drift metrics. - performance_metrics (list[str]): The performance metrics. - metadata (dict): The metadata. - - Returns: - dict: The metrics. - """ - columns = train_reference_data.drop( - columns=[target_name, 'timestamp', 'target', 'prediction'], errors='ignore').columns - - config = { - "target": target_name, - "prediction": "prediction", - "timestamp": "timestamp", - "columns": columns, - } - - model_analysis = ModelAnalysis(config=config) - - core_labels = self.get_core_labels(metadata, operation_type='detect_univariate_drift') - - start_time = time.time() - try: - model_analysis.detect_univariate_drift( - reference_df=train_reference_data, - analyse_df=analyse_data, - features=columns, - timestamp_column=config['timestamp'], - metrics=drift_metrics, - chunk_period="s" - ) - except Exception as e: - await self.emit_metric(metric_object=metrics.MODEL_ANALYZE_ERROR_COUNT, tags=core_labels) - raise e - - await self.observe_lag(start_time, metrics.MODEL_ANALYZE_LAG, core_labels) - await self.emit_metric(metric_object=metrics.MODEL_ANALYZE_COUNT, tags=core_labels) - - return {} diff --git a/laborious/worker/worker.py b/laborious/worker/worker.py index 5faf004..3d65de4 100644 --- a/laborious/worker/worker.py +++ b/laborious/worker/worker.py @@ -29,6 +29,7 @@ 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 @@ -49,6 +50,7 @@ with workflow.unsafe.imports_passed_through(): ) from laborious.workflows.minimal_retrain import MinimalRetrain from laborious.workflows.predictions_batch import PredictionsBatch + from laborious.workflows.drift import Drift from laborious.workflows.sub_workflows.format_and_export_prediction import ( FormatAndExportPrediction, ) @@ -156,6 +158,23 @@ async def main(): workflow_task_poller_behavior=PollerBehaviorAutoscaling(), activity_task_poller_behavior=PollerBehaviorAutoscaling(), ), + Worker( + temporal_client, + task_queue='drift-queue', + workflows=[Drift], + activities=[ + activities.load_custom_query, + activities.get_reference_data, + activities.calculate_drift, + activities.export_data_to_postgres, + ], + max_concurrent_workflow_tasks=50, + max_concurrent_activities=50, + max_concurrent_local_activities=50, + max_cached_workflows=2, + workflow_task_poller_behavior=PollerBehaviorAutoscaling(), + activity_task_poller_behavior=PollerBehaviorAutoscaling(), + ), Worker( temporal_client, task_queue='predictions_batch-queue', @@ -168,6 +187,7 @@ async def main(): activities.input_gate, activities.mlflow_response_gate, activities.mlflow_content_gate, + activities.format_transformed_data, activities.format_prediction, activities.format_default_prediction, activities.get_last_timestamp, diff --git a/laborious/workflows/drift.py b/laborious/workflows/drift.py new file mode 100644 index 0000000..5677b37 --- /dev/null +++ b/laborious/workflows/drift.py @@ -0,0 +1,99 @@ +from temporalio import workflow + +with workflow.unsafe.imports_passed_through(): + from datetime import timedelta + from typing import Any + + from sientia_do.temporal.policies import retry_policy + + from laborious.activities.activities import Activities + from sientia_do.temporal.constants import DATETIME_FORMAT_WITH_TZ + + +@workflow.defn(name='drift') +class Drift: + @workflow.run + async def run(self, input_data: dict[str, Any]): + """ + Execute the drift workflow. + + This method orchestrates the complete drift process by: + 1. Loading data using the provided custom SQL query + 2. Preparing prediction configuration and filters + 3. Delegating to the PredictionProcess workflow for ML operations + """ + + metadata = { + 'metadata': { + 'schedule_name': input_data['schedule_name'], + 'model_name': input_data['model_name'], + 'model_id': input_data['model_id'], + 'workflow_name': 'drift', + } + } + + gathering_query = f""" + SELECT * + FROM {input_data['schema']}.{input_data['source_table_name']} + WHERE + model_id = {input_data['model_id']} AND + timestamp > NOW() - INTERVAL '{input_data['interval']} minutes' + ORDER BY timestamp ASC + """ + + target_data_handler = workflow.start_local_activity_method( + Activities.load_custom_query, + { + **metadata, + 'query': gathering_query, + 'datetime_columns': ['timestamp', 'created_at'], + }, + retry_policy=retry_policy, + start_to_close_timeout=timedelta(seconds=300), + ) + + reference_data_handler = workflow.start_local_activity_method( + Activities.get_reference_data, + { + **metadata, + 'model_name': input_data['model_name'] + }, + retry_policy=retry_policy, + start_to_close_timeout=timedelta(seconds=300), + ) + + target_data = await target_data_handler + reference_data = await reference_data_handler + + if not target_data: + return + + drift_data = await workflow.execute_local_activity_method( + Activities.calculate_drift, + { + **metadata, + 'target_data': target_data, + 'reference_data': reference_data, + 'model_name': input_data['model_name'], + 'model_id': input_data['model_id'], + 'target_name': input_data['target_name'], + 'drift_metrics': input_data['drift_metrics'], + 'chunk_period': input_data.get('chunk_period', 'min'), + }, + retry_policy=retry_policy, + start_to_close_timeout=timedelta(seconds=300), + ) + + if drift_data: + await workflow.execute_activity_method( + Activities.export_data_to_postgres, + { + **metadata, + 'data': drift_data, + 'schema': input_data['schema'], + 'table_name': input_data['target_table_name'], + 'timestamp_conversion': {'column': 'timestamp', 'format': DATETIME_FORMAT_WITH_TZ}, + }, + retry_policy=retry_policy, + start_to_close_timeout=timedelta(seconds=300), + ) \ No newline at end of file diff --git a/laborious/workflows/predictions_batch.py b/laborious/workflows/predictions_batch.py index afd7c86..6c5f770 100644 --- a/laborious/workflows/predictions_batch.py +++ b/laborious/workflows/predictions_batch.py @@ -98,6 +98,7 @@ class PredictionsBatch: 'data': data, 'schema': input_data['schema'], 'table_name': input_data['table_name'], + 'transform_table_name': input_data['transform_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'}}), diff --git a/laborious/workflows/sub_workflows/format_and_export_prediction.py b/laborious/workflows/sub_workflows/format_and_export_prediction.py index bd28966..4280ccf 100644 --- a/laborious/workflows/sub_workflows/format_and_export_prediction.py +++ b/laborious/workflows/sub_workflows/format_and_export_prediction.py @@ -95,12 +95,13 @@ class FormatAndExportPrediction: { **metadata, 'data': transformed_data, + 'model_id': input_data['model_id'], }, retry_policy=retry_policy, start_to_close_timeout=timedelta(seconds=60), ) - write_transformed_handler = workflow.execute_activity_method( + write_transformed_handler = workflow.start_activity_method( Activities.export_data_to_postgres, { **metadata, diff --git a/laborious/workflows/sub_workflows/prediction_process.py b/laborious/workflows/sub_workflows/prediction_process.py index fa07504..73f0227 100644 --- a/laborious/workflows/sub_workflows/prediction_process.py +++ b/laborious/workflows/sub_workflows/prediction_process.py @@ -209,6 +209,7 @@ class PredictionProcess: 'opc_output_config': input_data['opc_output_config'], 'schema': input_data['schema'], 'table_name': input_data['table_name'], + 'transform_table_name': input_data['transform_table_name'], 'comment': comment, 'prediction_store_policy': input_data['prediction_store_policy'], }, @@ -250,6 +251,7 @@ class PredictionProcess: schema = input_data['schema'] table_name = input_data['table_name'] + transform_table_name = input_data['transform_table_name'] model_id = input_data['model_id'] model_name = input_data['model_name'] model_config = input_data.get('model_config', {}) @@ -289,6 +291,7 @@ class PredictionProcess: 'model_config': model_config, 'schema': schema, 'table_name': table_name, + 'transform_table_name': transform_table_name, 'comment': comment, 'opc_output_config': input_data['opc_output_config'], 'prediction_store_policy': input_data['prediction_store_policy'], diff --git a/tests.ipynb b/tests.ipynb index 3e4b61a..808277c 100644 --- a/tests.ipynb +++ b/tests.ipynb @@ -487,6 +487,107 @@ "except Exception as e:\n", " print(e)\n" ] + }, + { + "cell_type": "code", + "execution_count": 7, + "id": "771ab4ee", + "metadata": {}, + "outputs": [ + { + "data": { + "text/html": [ + "
\n", + "\n", + "\n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + "
abtargetpredictiontimestamp
014712025-01-01
125822025-01-02
236932025-01-03
\n", + "
" + ], + "text/plain": [ + " a b target prediction timestamp\n", + "0 1 4 7 1 2025-01-01\n", + "1 2 5 8 2 2025-01-02\n", + "2 3 6 9 3 2025-01-03" + ] + }, + "metadata": {}, + "output_type": "display_data" + } + ], + "source": [ + "from pandas import DataFrame, merge\n", + "\n", + "retrain_dataset = DataFrame({\n", + " 'a': {'2025-01-01': 1, '2025-01-02': 2, '2025-01-03': 3},\n", + " 'b': {'2025-01-01': 4, '2025-01-02': 5, '2025-01-03': 6},\n", + " 'c': {'2025-01-01': 7, '2025-01-02': 8, '2025-01-03': 9},\n", + "})\n", + "\n", + "prediction_data = DataFrame({\n", + " 'prediction': {'1': 1, '2': 2, '3': 3},\n", + "})\n", + "\n", + "prediction_data.index = retrain_dataset.index\n", + "\n", + "prediction_data = merge(\n", + " retrain_dataset, prediction_data, left_index=True, right_index=True, how='left')\n", + "\n", + "prediction_data.rename(columns={'c': 'target'}, inplace=True)\n", + "\n", + "prediction_data['timestamp'] = prediction_data.index\n", + "\n", + "prediction_data.reset_index(drop=True, inplace=True)\n", + "\n", + "display(prediction_data)" + ] } ], "metadata": { diff --git a/tests/laborious/activities/test_model_metrics.py b/tests/laborious/activities/test_model_metrics.py new file mode 100644 index 0000000..8b55df8 --- /dev/null +++ b/tests/laborious/activities/test_model_metrics.py @@ -0,0 +1,641 @@ +from unittest.mock import ANY, AsyncMock, MagicMock, call, patch + +from pandas import DataFrame +from pytest import fixture, mark +from sientia_do.notifications.models import NotificationLevel + +from laborious.activities.model_metrics import ModelMetrics + + +@fixture +def model_metrics_activity(): + model_metrics = ModelMetrics( + logger=MagicMock(), + notification_handler=MagicMock(), + metrics_controller=AsyncMock(), + ) + model_metrics.error = MagicMock() + model_metrics.debug = MagicMock() + model_metrics.info = MagicMock() + model_metrics.warning = MagicMock() + model_metrics.critical = MagicMock() + model_metrics.send_notification = MagicMock() + model_metrics.send_notification_async = AsyncMock() + model_metrics.emit_metric = AsyncMock() + model_metrics.get_core_labels = MagicMock(return_value={'pod_id': 'test_pod', 'model_name': 'test_model', 'workflow_name': 'test_workflow'}) + model_metrics.observe_lag = AsyncMock() + model_metrics.pod_id = 'test_pod' + return model_metrics + + +metadata = { + 'metadata': { + 'model_id': 'test_model', + 'model_name': 'test_model', + 'workflow_name': 'test_workflow', + 'schema_name': 'test_schedule', + }, +} + + +@mark.asyncio +async def test_calculate_drift_invalid_chunk_period(model_metrics_activity): + # Arrange + input_data = { + **metadata, + 'model_name': 'test_model', + 'model_id': 'test_model_id', + 'reference_data': None, + 'target_data': { + 'timestamp': ['2023-05-26 11:12:27'], + 'variable': ['feature1'], + 'value': [1.0], + }, + 'target_name': 'target', + 'drift_metrics': ['ks_test'], + 'chunk_period': 'invalid', + } + + # Act & Assert + try: + await model_metrics_activity.calculate_drift(input_data) + except ValueError as e: + assert str(e) == "Invalid chunk period: invalid, must be \"min\" or \"s\"" + model_metrics_activity.error.assert_called_once_with( + 'Invalid chunk period: invalid', metadata['metadata'] + ) + else: + raise AssertionError('Expected ValueError') + + +@mark.asyncio +@patch('laborious.activities.model_metrics.DataFrame') +@patch('laborious.activities.model_metrics.to_datetime') +async def test_calculate_drift_with_reference_data( + mock_to_datetime, mock_dataframe, model_metrics_activity +): + # Arrange + mock_to_datetime.return_value.dt.strftime.return_value = '2023-05-26 11:12:27' + mock_to_datetime.return_value.dt.tz_localize.return_value.dt.strftime.return_value = '2023-05-26 11:12:27+00:00' + + mock_drift_df = MagicMock() + mock_drift_df.empty = False + mock_drift_df.drop.return_value = mock_drift_df + mock_drift_df.__getitem__.return_value.isin.return_value = [True] + mock_drift_df.__getitem__.return_value = mock_drift_df + mock_drift_df.__getitem__.return_value.dt.tz_localize.return_value.dt.strftime.return_value = '2023-05-26 11:12:27+00:00' + mock_drift_df.rename.return_value = mock_drift_df + mock_drift_df.drop_duplicates.return_value = mock_drift_df + mock_drift_df.to_dict.return_value = {'method': ['ks_test'], 'value': [0.5], 'feature': ['feature1'], 'timestamp': ['2023-05-26 11:12:27+00:00'], 'model_id': ['test_model_id'], 'accurate': [True]} + + model_metrics_activity.get_drift_metrics = AsyncMock(return_value=mock_drift_df) + + reference_data = DataFrame({ + 'timestamp': ['2023-05-26 11:12:27'], + 'target': [1.0], + 'feature1': [1.0], + }) + + mock_target_df = MagicMock() + mock_target_df.pivot.return_value = mock_target_df + mock_target_df.index = ['2023-05-26 11:12:27'] + mock_target_df.reset_index.return_value = mock_target_df + mock_target_df.dropna.return_value = mock_target_df + mock_target_df.__getitem__.return_value.apply.return_value = ['2023-05-26 11:12:27'] + mock_target_df.drop.return_value.columns = ['feature1'] + mock_dataframe.return_value = mock_target_df + + input_data = { + **metadata, + 'model_name': 'test_model', + 'model_id': 'test_model_id', + 'reference_data': reference_data.to_dict(), + 'target_data': { + 'timestamp': ['2023-05-26 11:12:27'], + 'variable': ['feature1'], + 'value': [1.0], + }, + 'target_name': 'target', + 'drift_metrics': ['ks_test'], + 'chunk_period': 'min', + } + + # Act + result = await model_metrics_activity.calculate_drift(input_data) + + # Assert + assert isinstance(result, dict) + assert result == mock_drift_df.to_dict.return_value + model_metrics_activity.info.assert_called() + model_metrics_activity.get_drift_metrics.assert_called_once() + # Verify transformations were called + mock_drift_df.drop.assert_called_once_with(columns=['p_value', 'chunk_start_date', 'chunk_end_date'], inplace=True) + mock_drift_df.__getitem__.assert_called() + mock_drift_df.rename.assert_called_once_with(columns={'metric': 'method', 'statistic': 'value'}, inplace=True) + mock_drift_df.drop_duplicates.assert_called_once_with(subset=['timestamp', 'method', 'feature'], keep='first', inplace=True) + mock_drift_df.to_dict.assert_called_once() + + +@mark.asyncio +@patch('laborious.activities.model_metrics.DataFrame') +@patch('laborious.activities.model_metrics.to_datetime') +async def test_calculate_drift_without_reference_data( + mock_to_datetime, mock_dataframe, model_metrics_activity +): + # Arrange + mock_to_datetime.return_value.dt.strftime.return_value = '2023-05-26 11:12:27' + mock_to_datetime.return_value.dt.tz_localize.return_value.dt.strftime.return_value = '2023-05-26 11:12:27+00:00' + + mock_drift_df = MagicMock() + mock_drift_df.empty = False + mock_drift_df.drop.return_value = mock_drift_df + mock_drift_df.__getitem__.return_value.isin.return_value = [True] + mock_drift_df.__getitem__.return_value = mock_drift_df + mock_drift_df.__getitem__.return_value.dt.tz_localize.return_value.dt.strftime.return_value = '2023-05-26 11:12:27+00:00' + mock_drift_df.rename.return_value = mock_drift_df + mock_drift_df.drop_duplicates.return_value = mock_drift_df + mock_drift_df.to_dict.return_value = {'method': ['ks_test'], 'value': [0.5], 'feature': ['feature1'], 'timestamp': ['2023-05-26 11:12:27+00:00'], 'model_id': ['test_model_id'], 'accurate': [False]} + + model_metrics_activity.get_drift_metrics = AsyncMock(return_value=mock_drift_df) + + target_data_dict = { + 'timestamp': ['2023-05-26 11:12:27', '2023-05-26 11:12:28', '2023-05-26 11:12:29'], + 'variable': ['feature1', 'feature1', 'feature1'], + 'value': [1.0, 2.0, 3.0], + } + + mock_target_df = MagicMock() + mock_target_df.pivot.return_value = mock_target_df + mock_target_df.index = ['2023-05-26 11:12:27', '2023-05-26 11:12:28', '2023-05-26 11:12:29'] + mock_target_df.reset_index.return_value = mock_target_df + mock_target_df.dropna.return_value = mock_target_df + mock_target_df.sort_values.return_value = mock_target_df + mock_target_df.head.return_value = DataFrame({'timestamp': ['2023-05-26 11:12:27'], 'feature1': [1.0]}) + mock_target_df.__getitem__.return_value.apply.return_value = ['2023-05-26 11:12:27', '2023-05-26 11:12:28', '2023-05-26 11:12:29'] + mock_target_df.drop.return_value.columns = ['feature1'] + mock_dataframe.return_value = mock_target_df + mock_dataframe.side_effect = lambda x=None: mock_target_df if x is not None else mock_target_df + + input_data = { + **metadata, + 'model_name': 'test_model', + 'model_id': 'test_model_id', + 'reference_data': None, + 'target_data': target_data_dict, + 'target_name': 'target', + 'drift_metrics': ['ks_test'], + 'chunk_period': 's', + } + + # Act + result = await model_metrics_activity.calculate_drift(input_data) + + # Assert + assert isinstance(result, dict) + assert result == mock_drift_df.to_dict.return_value + model_metrics_activity.warning.assert_called() + model_metrics_activity.send_notification_async.assert_called_once_with( + metadata=metadata['metadata'], + notification_id='MODEL_METRICS_REFERENCE_DATA_WARNING', + message='Using 30% first rows of target data as reference data', + block='model_metrics', + level=NotificationLevel.WARNING, + attachment_content=ANY, + ) + # Verify transformations were called + mock_drift_df.drop.assert_called_once_with(columns=['p_value', 'chunk_start_date', 'chunk_end_date'], inplace=True) + mock_drift_df.__getitem__.assert_called() + mock_drift_df.rename.assert_called_once_with(columns={'metric': 'method', 'statistic': 'value'}, inplace=True) + mock_drift_df.drop_duplicates.assert_called_once_with(subset=['timestamp', 'method', 'feature'], keep='first', inplace=True) + mock_drift_df.to_dict.assert_called_once() + + +@mark.asyncio +@patch('laborious.activities.model_metrics.DataFrame') +@patch('laborious.activities.model_metrics.to_datetime') +async def test_calculate_drift_empty_drift_df( + mock_to_datetime, mock_dataframe, model_metrics_activity +): + # Arrange + mock_to_datetime.return_value.dt.strftime.return_value = '2023-05-26 11:12:27' + + model_metrics_activity.get_drift_metrics = AsyncMock(return_value=DataFrame()) + + reference_data = DataFrame({ + 'timestamp': ['2023-05-26 11:12:27'], + 'target': [1.0], + 'feature1': [1.0], + }) + + mock_target_df = MagicMock() + mock_target_df.pivot.return_value = mock_target_df + mock_target_df.index = ['2023-05-26 11:12:27'] + mock_target_df.reset_index.return_value = mock_target_df + mock_target_df.dropna.return_value = mock_target_df + mock_target_df.__getitem__.return_value.apply.return_value = ['2023-05-26 11:12:27'] + mock_target_df.drop.return_value.columns = ['feature1'] + mock_dataframe.return_value = mock_target_df + + input_data = { + **metadata, + 'model_name': 'test_model', + 'model_id': 'test_model_id', + 'reference_data': reference_data.to_dict(), + 'target_data': { + 'timestamp': ['2023-05-26 11:12:27'], + 'variable': ['feature1'], + 'value': [1.0], + }, + 'target_name': 'target', + 'drift_metrics': ['ks_test'], + 'chunk_period': 'min', + } + + # Act + result = await model_metrics_activity.calculate_drift(input_data) + + # Assert + assert result == {} + model_metrics_activity.warning.assert_called_with('No drift metrics found', metadata['metadata']) + + +@mark.asyncio +@patch('laborious.activities.model_metrics.DataFrame') +@patch('laborious.activities.model_metrics.to_datetime') +async def test_calculate_drift_empty_after_timestamp_filter( + mock_to_datetime, mock_dataframe, model_metrics_activity +): + # Arrange + mock_to_datetime.return_value.dt.strftime.return_value = '2023-05-26 11:12:27' + + mock_drift_df = MagicMock() + mock_drift_df.empty = False + mock_drift_df.drop.return_value = mock_drift_df + + # Set up __getitem__ to handle filtering - timestamp access returns series with isin=False + # and filtering returns empty DataFrame + mock_timestamp_series = MagicMock() + mock_timestamp_series.isin.return_value = [False] + mock_empty_df = MagicMock() + mock_empty_df.empty = True + + def getitem_side_effect(key): + if key == 'timestamp': + return mock_timestamp_series + else: + # This is the filtering operation - return empty DataFrame + return mock_empty_df + + mock_drift_df.__getitem__.side_effect = getitem_side_effect + + model_metrics_activity.get_drift_metrics = AsyncMock(return_value=mock_drift_df) + + reference_data = DataFrame({ + 'timestamp': ['2023-05-26 11:12:27'], + 'target': [1.0], + 'feature1': [1.0], + }) + + mock_target_df = MagicMock() + mock_target_df.pivot.return_value = mock_target_df + mock_target_df.index = ['2023-05-26 11:12:27'] + mock_target_df.reset_index.return_value = mock_target_df + mock_target_df.dropna.return_value = mock_target_df + mock_target_df.__getitem__.return_value.apply.return_value = ['2023-05-26 11:12:27'] + mock_target_df.drop.return_value.columns = ['feature1'] + mock_dataframe.return_value = mock_target_df + + input_data = { + **metadata, + 'model_name': 'test_model', + 'model_id': 'test_model_id', + 'reference_data': reference_data.to_dict(), + 'target_data': { + 'timestamp': ['2023-05-26 11:12:27'], + 'variable': ['feature1'], + 'value': [1.0], + }, + 'target_name': 'target', + 'drift_metrics': ['ks_test'], + 'chunk_period': 'min', + } + + # Act + result = await model_metrics_activity.calculate_drift(input_data) + + # Assert + assert result == {} + model_metrics_activity.warning.assert_called_with( + 'No drift metrics found after dropping rows where timestamp is not in target data', + metadata['metadata'] + ) + # Verify transformations were called + mock_drift_df.drop.assert_called_once_with(columns=['p_value', 'chunk_start_date', 'chunk_end_date'], inplace=True) + mock_drift_df.__getitem__.assert_called() + + +@mark.asyncio +@patch('laborious.activities.model_metrics.DataFrame') +@patch('laborious.activities.model_metrics.to_datetime') +async def test_calculate_drift_success_min( + mock_to_datetime, mock_dataframe, model_metrics_activity +): + # Arrange + mock_to_datetime.return_value.dt.strftime.return_value = '2023-05-26 11:12:27' + mock_to_datetime.return_value.dt.tz_localize.return_value.dt.strftime.return_value = '2023-05-26 11:12:27+00:00' + + mock_drift_df = MagicMock() + mock_drift_df.empty = False + mock_drift_df.drop.return_value = mock_drift_df + mock_drift_df.__getitem__.return_value.isin.return_value = [True] + mock_drift_df.__getitem__.return_value = mock_drift_df + mock_drift_df.__getitem__.return_value.dt.tz_localize.return_value.dt.strftime.return_value = '2023-05-26 11:12:27+00:00' + mock_drift_df.rename.return_value = mock_drift_df + mock_drift_df.drop_duplicates.return_value = mock_drift_df + mock_drift_df.to_dict.return_value = {'method': ['ks_test'], 'value': [0.5], 'feature': ['feature1'], 'timestamp': ['2023-05-26 11:12:27+00:00'], 'model_id': ['test_model_id'], 'accurate': [True]} + + model_metrics_activity.get_drift_metrics = AsyncMock(return_value=mock_drift_df) + + reference_data = DataFrame({ + 'timestamp': ['2023-05-26 11:12:27'], + 'target': [1.0], + 'feature1': [1.0], + }) + + mock_target_df = MagicMock() + mock_target_df.pivot.return_value = mock_target_df + mock_target_df.index = ['2023-05-26 11:12:27'] + mock_target_df.reset_index.return_value = mock_target_df + mock_target_df.dropna.return_value = mock_target_df + mock_target_df.__getitem__.return_value.isin.return_value = [True] + mock_target_df.drop.return_value.columns = ['feature1'] + mock_dataframe.return_value = mock_target_df + + input_data = { + **metadata, + 'model_name': 'test_model', + 'model_id': 'test_model_id', + 'reference_data': reference_data.to_dict(), + 'target_data': { + 'timestamp': ['2023-05-26 11:12:27'], + 'variable': ['feature1'], + 'value': [1.0], + }, + 'target_name': 'target', + 'drift_metrics': ['ks_test'], + 'chunk_period': 'min', + } + + # Act + result = await model_metrics_activity.calculate_drift(input_data) + + # Assert + assert isinstance(result, dict) + assert result == mock_drift_df.to_dict.return_value + model_metrics_activity.info.assert_called() + model_metrics_activity.get_drift_metrics.assert_called_once() + # Verify transformations were called + mock_drift_df.drop.assert_called_once_with(columns=['p_value', 'chunk_start_date', 'chunk_end_date'], inplace=True) + mock_drift_df.__getitem__.assert_called() + mock_drift_df.rename.assert_called_once_with(columns={'metric': 'method', 'statistic': 'value'}, inplace=True) + mock_drift_df.drop_duplicates.assert_called_once_with(subset=['timestamp', 'method', 'feature'], keep='first', inplace=True) + mock_drift_df.to_dict.assert_called_once() + + +@mark.asyncio +@patch('laborious.activities.model_metrics.DataFrame') +@patch('laborious.activities.model_metrics.to_datetime') +async def test_calculate_drift_success_s( + mock_to_datetime, mock_dataframe, model_metrics_activity +): + # Arrange + mock_to_datetime.return_value.dt.strftime.return_value = '2023-05-26 11:12:27' + mock_to_datetime.return_value.dt.tz_localize.return_value.dt.strftime.return_value = '2023-05-26 11:12:27+00:00' + + mock_drift_df = MagicMock() + mock_drift_df.empty = False + mock_drift_df.drop.return_value = mock_drift_df + mock_drift_df.__getitem__.return_value.isin.return_value = [True] + mock_drift_df.__getitem__.return_value = mock_drift_df + mock_drift_df.__getitem__.return_value.dt.tz_localize.return_value.dt.strftime.return_value = '2023-05-26 11:12:27+00:00' + mock_drift_df.rename.return_value = mock_drift_df + mock_drift_df.drop_duplicates.return_value = mock_drift_df + mock_drift_df.to_dict.return_value = {'method': ['ks_test'], 'value': [0.5], 'feature': ['feature1'], 'timestamp': ['2023-05-26 11:12:27+00:00'], 'model_id': ['test_model_id'], 'accurate': [True]} + + model_metrics_activity.get_drift_metrics = AsyncMock(return_value=mock_drift_df) + + reference_data = DataFrame({ + 'timestamp': ['2023-05-26 11:12:27'], + 'target': [1.0], + 'feature1': [1.0], + }) + + mock_target_df = MagicMock() + mock_target_df.pivot.return_value = mock_target_df + mock_target_df.index = ['2023-05-26 11:12:27'] + mock_target_df.reset_index.return_value = mock_target_df + mock_target_df.dropna.return_value = mock_target_df + mock_target_df.__getitem__.return_value.isin.return_value = [True] + mock_target_df.drop.return_value.columns = ['feature1'] + mock_dataframe.return_value = mock_target_df + + input_data = { + **metadata, + 'model_name': 'test_model', + 'model_id': 'test_model_id', + 'reference_data': reference_data.to_dict(), + 'target_data': { + 'timestamp': ['2023-05-26 11:12:27'], + 'variable': ['feature1'], + 'value': [1.0], + }, + 'target_name': 'target', + 'drift_metrics': ['ks_test'], + 'chunk_period': 's', + } + + # Act + result = await model_metrics_activity.calculate_drift(input_data) + + # Assert + assert isinstance(result, dict) + assert result == mock_drift_df.to_dict.return_value + model_metrics_activity.info.assert_called() + model_metrics_activity.get_drift_metrics.assert_called_once() + # Verify transformations were called + mock_drift_df.drop.assert_called_once_with(columns=['p_value', 'chunk_start_date', 'chunk_end_date'], inplace=True) + mock_drift_df.__getitem__.assert_called() + mock_drift_df.rename.assert_called_once_with(columns={'metric': 'method', 'statistic': 'value'}, inplace=True) + mock_drift_df.drop_duplicates.assert_called_once_with(subset=['timestamp', 'method', 'feature'], keep='first', inplace=True) + mock_drift_df.to_dict.assert_called_once() + + +@mark.asyncio +@patch('laborious.activities.model_metrics.DataFrame') +@patch('laborious.activities.model_metrics.to_datetime') +async def test_calculate_drift_get_drift_metrics_error( + mock_to_datetime, mock_dataframe, model_metrics_activity +): + # Arrange + mock_to_datetime.return_value.dt.strftime.return_value = '2023-05-26 11:12:27' + + model_metrics_activity.get_drift_metrics = AsyncMock(side_effect=Exception('Get drift metrics error')) + + reference_data = DataFrame({ + 'timestamp': ['2023-05-26 11:12:27'], + 'target': [1.0], + 'feature1': [1.0], + }) + + mock_target_df = MagicMock() + mock_target_df.pivot.return_value = mock_target_df + mock_target_df.index = ['2023-05-26 11:12:27'] + mock_target_df.reset_index.return_value = mock_target_df + mock_target_df.dropna.return_value = mock_target_df + mock_target_df.__getitem__.return_value.apply.return_value = ['2023-05-26 11:12:27'] + mock_target_df.drop.return_value.columns = ['feature1'] + mock_dataframe.return_value = mock_target_df + + input_data = { + **metadata, + 'model_name': 'test_model', + 'model_id': 'test_model_id', + 'reference_data': reference_data.to_dict(), + 'target_data': { + 'timestamp': ['2023-05-26 11:12:27'], + 'variable': ['feature1'], + 'value': [1.0], + }, + 'target_name': 'target', + 'drift_metrics': ['ks_test'], + 'chunk_period': 'min', + } + + # Act + result = await model_metrics_activity.calculate_drift(input_data) + + # Assert + assert result == {} + model_metrics_activity.error.assert_called_once_with( + 'Error getting drift metrics: Get drift metrics error', + metadata['metadata'] + ) + model_metrics_activity.send_notification_async.assert_called_once_with( + metadata=metadata['metadata'], + notification_id='MODEL_METRICS_GET_DRIFT_METRICS_ERROR', + message='Error getting drift metrics: Get drift metrics error', + block='model_metrics', + level=NotificationLevel.ERROR, + attachment_content=ANY, + ) + + +@mark.asyncio +@patch('laborious.activities.model_metrics.to_datetime') +@patch('laborious.activities.model_metrics.time.time') +@patch('laborious.activities.model_metrics.ModelAnalysis') +@patch('laborious.activities.model_metrics.metrics') +async def test_get_drift_metrics_success( + mock_metrics, mock_model_analysis, mock_time, mock_to_datetime, model_metrics_activity +): + # Arrange + mock_time.return_value = 1000.0 + + mock_drift_df = DataFrame({ + 'timestamp': ['2023-05-26 11:12:27'], + 'metric': ['ks_test'], + 'statistic': [0.5], + 'feature': ['feature1'], + }) + + mock_model_analysis.return_value.detect_univariate_drift.return_value = MagicMock() + mock_model_analysis.return_value.detect_multivariate_drift.return_value = MagicMock() + mock_model_analysis.return_value.get_drift_metrics_dataframe.return_value = mock_drift_df + + reference_data = DataFrame({ + 'timestamp': ['2023-05-26 11:12:27'], + 'target': [1.0], + 'feature1': [1.0], + }) + + target_data = DataFrame({ + 'timestamp': ['2023-05-26 11:12:27'], + 'target': [1.0], + 'feature1': [1.0], + }) + + reference_columns = reference_data.drop( + columns=['target', 'timestamp'], errors='ignore' + ).columns + + # Act + result = await model_metrics_activity.get_drift_metrics( + reference_data=reference_data, + target_data=target_data, + target_name='target', + reference_columns=reference_columns, + drift_metrics=['ks_test'], + chunk_period='min', + metadata=metadata['metadata'], + ) + + # Assert + assert isinstance(result, DataFrame) + model_metrics_activity.debug.assert_called() + model_metrics_activity.observe_lag.assert_called() + model_metrics_activity.emit_metric.assert_called() + + +@mark.asyncio +@patch('laborious.activities.model_metrics.to_datetime') +@patch('laborious.activities.model_metrics.time.time') +@patch('laborious.activities.model_metrics.ModelAnalysis') +@patch('laborious.activities.model_metrics.metrics') +async def test_get_drift_metrics_univariate_error( + mock_metrics, mock_model_analysis, mock_time, mock_to_datetime, model_metrics_activity +): + # Arrange + mock_time.return_value = 1000.0 + + mock_model_analysis.return_value.detect_univariate_drift.side_effect = Exception('Univariate drift error') + + reference_data = DataFrame({ + 'timestamp': ['2023-05-26 11:12:27'], + 'target': [1.0], + 'feature1': [1.0], + }) + + target_data = DataFrame({ + 'timestamp': ['2023-05-26 11:12:27'], + 'target': [1.0], + 'feature1': [1.0], + }) + + reference_columns = reference_data.drop( + columns=['target', 'timestamp'], errors='ignore' + ).columns + + # Act & Assert + try: + await model_metrics_activity.get_drift_metrics( + reference_data=reference_data, + target_data=target_data, + target_name='target', + reference_columns=reference_columns, + drift_metrics=['ks_test'], + chunk_period='min', + metadata=metadata['metadata'], + ) + except Exception as e: + assert str(e) == 'Univariate drift error' + model_metrics_activity.error.assert_called_once_with( + 'Error detecting univariate drift: Univariate drift error', + metadata['metadata'] + ) + model_metrics_activity.emit_metric.assert_called_with( + metric_object=mock_metrics.MODEL_ANALYZE_ERROR_COUNT, + tags=ANY + ) + else: + raise AssertionError('Expected Exception') + diff --git a/tests/laborious/utils/repository/test_model_repository.py b/tests/laborious/utils/repository/test_model_repository.py index 9380853..9d79f70 100644 --- a/tests/laborious/utils/repository/test_model_repository.py +++ b/tests/laborious/utils/repository/test_model_repository.py @@ -190,6 +190,26 @@ def test_get_model_params(mlflow, mlflow_repository): assert output == mlflow.get_run.return_value.data.params +def test_check_artifact_exists_true(mlflow_repository): + artifact = MagicMock(path='test_artifact') + mlflow_repository.client.list_artifacts.return_value = [artifact] + + result = mlflow_repository.check_artifact_exists('run_id', 'test_artifact', metadata['metadata']) + + assert result is True + mlflow_repository.client.list_artifacts.assert_called_once_with('run_id') + + +def test_check_artifact_exists_false(mlflow_repository): + artifact = MagicMock(path='other_artifact') + mlflow_repository.client.list_artifacts.return_value = [artifact] + + result = mlflow_repository.check_artifact_exists('run_id', 'test_artifact', metadata['metadata']) + + assert result is False + mlflow_repository.client.list_artifacts.assert_called_once_with('run_id') + + @pytest.mark.asyncio @patch('laborious.utils.repository.model_repository.path') @patch('laborious.utils.repository.model_repository.rmtree') @@ -275,6 +295,56 @@ async def test_download_artifacts_error(makedirs, rmtree, path, mlflow_repositor mlflow_repository.observe_lag.assert_not_called() +@pytest.mark.asyncio +@patch('laborious.utils.repository.model_repository.mlflow') +@patch('laborious.utils.repository.model_repository.pd') +@patch('laborious.utils.repository.model_repository.StringIO') +async def test_load_artifact_dataframe_success(StringIO, pd, mlflow, mlflow_repository): + mlflow_repository.get_model_run_id = MagicMock(return_value='run_id') + mlflow_repository.check_artifact_exists = MagicMock(return_value=True) + + mlflow.artifacts.load_text.return_value = 'col1,col2\n1,2\n3,4' + + result = await mlflow_repository.load_artifact_dataframe('model_name', 'artifact_path', metadata['metadata']) + + mlflow_repository.get_model_run_id.assert_called_once_with(model_name='model_name', stage='Production') + mlflow_repository.check_artifact_exists.assert_called_once_with('run_id', 'artifact_path', metadata['metadata']) + mlflow.artifacts.load_text.assert_called_once_with('runs:/run_id/artifact_path') + assert result == pd.read_csv.return_value + mlflow_repository.observe_lag.assert_called_once_with(ANY, metrics.MODEL_READ_LAG, ANY) + mlflow_repository.emit_metric.assert_called_once_with( + metric_object=metrics.MODEL_READ_COUNT, tags=ANY + ) + + +@pytest.mark.asyncio +async def test_load_artifact_dataframe_not_exists(mlflow_repository): + mlflow_repository.get_model_run_id = MagicMock(return_value='run_id') + mlflow_repository.check_artifact_exists = MagicMock(return_value=False) + + result = await mlflow_repository.load_artifact_dataframe('model_name', 'artifact_path', metadata['metadata']) + + assert result is None + mlflow_repository.get_model_run_id.assert_called_once_with(model_name='model_name', stage='Production') + mlflow_repository.check_artifact_exists.assert_called_once_with('run_id', 'artifact_path', metadata['metadata']) + + +@pytest.mark.asyncio +@patch('laborious.utils.repository.model_repository.mlflow') +async def test_load_artifact_dataframe_error(mlflow, mlflow_repository): + mlflow_repository.get_model_run_id = MagicMock(return_value='run_id') + mlflow_repository.check_artifact_exists = MagicMock(return_value=True) + mlflow.artifacts.load_text.side_effect = ValueError('error') + + with pytest.raises(ValueError): + await mlflow_repository.load_artifact_dataframe('model_name', 'artifact_path', metadata['metadata']) + + mlflow_repository.emit_metric.assert_called_once_with( + metric_object=metrics.MODEL_READ_ERROR_COUNT, tags=ANY + ) + mlflow_repository.observe_lag.assert_not_called() + + def test_get_experiment_error(mlflow, mlflow_repository): mlflow.get_experiment_by_name.return_value = None @@ -719,6 +789,7 @@ async def test_fit_models_not_df_target_name_none_and_not_in_model( mlflow_repository.detect_and_parse_datetime_index = MagicMock( return_value=MagicMock(drop_duplicates=MagicMock(return_value=MagicMock(columns=[]))) ) + mlflow_repository.get_prediction_data = MagicMock(return_value=DataFrame()) data = MagicMock() @@ -775,9 +846,14 @@ async def test_fit_models_not_df_target_name_none_and_not_in_model( prediction_model.fit.assert_called_once_with(pd_merge.return_value) + mlflow_repository.get_prediction_data.assert_called_once_with( + prediction_model, pd_merge.return_value, data_model.fit.return_value.target_variable + ) + assert output == { 'prediction_model': {'model': prediction_model, 'artifact_path': 'artifact_path'}, 'data_model': {'model': data_model.fit.return_value, 'artifact_path': 'artifact_path'}, + 'prediction_data': mlflow_repository.get_prediction_data.return_value, } @@ -801,6 +877,7 @@ async def test_fit_models_df_target_name_not_none_and_in_model( drop_duplicates=MagicMock(return_value=MagicMock(columns=['feat_1'])) ) ) + mlflow_repository.get_prediction_data = MagicMock(return_value=DataFrame()) data = MagicMock() @@ -855,9 +932,14 @@ async def test_fit_models_df_target_name_not_none_and_in_model( prediction_model.fit.assert_called_once_with(transformed_data) + mlflow_repository.get_prediction_data.assert_called_once_with( + prediction_model, transformed_data, 'feat_1' + ) + assert output == { 'prediction_model': {'model': prediction_model, 'artifact_path': 'artifact_path'}, 'data_model': {'model': data_model, 'artifact_path': 'artifact_path'}, + 'prediction_data': mlflow_repository.get_prediction_data.return_value, } @@ -878,7 +960,8 @@ async def test_log_model_sklearn(mlflow, mlflow_repository): @patch('laborious.utils.repository.model_repository.path') @pytest.mark.asyncio async def test_log_model_pyfunc(path, mlflow, mlflow_repository): - model_data = {'model': MagicMock(), 'artifact_path': 'artifact_path'} + model_mock = MagicMock() + model_data = {'model': model_mock, 'artifact_path': 'artifact_path'} await mlflow_repository.log_model( model_data, 'pyfunc', 'prediction_model', metadata['metadata'] ) @@ -887,7 +970,7 @@ async def test_log_model_pyfunc(path, mlflow, mlflow_repository): path.join.assert_called_once_with('artifact_path', 'code', 'utils') - model_data['model'].store_model.assert_called_once_with( + model_mock.store_model.assert_called_once_with( artifact_path='prediction_model', code_path=[path.join.return_value], to_disk=False ) @@ -934,9 +1017,11 @@ async def test_create_new_experiment( ): model_name = 'model_name' data = MagicMock() + prediction_data = MagicMock(spec=DataFrame) retrain_data = { 'prediction_model': {'model': MagicMock(), 'artifact_path': 'artifact_path'}, 'data_model': {'model': MagicMock(), 'artifact_path': 'artifact_path'}, + 'prediction_data': prediction_data, } mlflow_repository.get_model_params = MagicMock( @@ -971,7 +1056,11 @@ async def test_create_new_experiment( mlflow_repository.get_experiment.return_value.name ) - data.to_csv.assert_called_once_with('./tmp/artifacts/model_name/retrain_data.csv', index=True) + data.to_csv.assert_called_once_with('./tmp/artifacts/model_name/retrain_data.csv', index=False) + # Verify prediction_data.to_csv was called with correct arguments + prediction_data.to_csv.assert_called_once() + assert prediction_data.to_csv.call_args[0][0] == './tmp/artifacts/model_name/evaluation_data.csv' + assert prediction_data.to_csv.call_args[1]['index'] is False mlflow.start_run.assert_called_once_with( experiment_id=mlflow_repository.get_experiment.return_value.experiment_id, @@ -1000,7 +1089,10 @@ async def test_create_new_experiment( } ) - mlflow.log_artifact.assert_called_once_with('./tmp/artifacts/model_name/retrain_data.csv') + mlflow.log_artifact.assert_has_calls([ + call('./tmp/artifacts/model_name/retrain_data.csv'), + call('./tmp/artifacts/model_name/evaluation_data.csv'), + ]) force_memory_release.assert_called_once_with(mlflow_repository.logger) @@ -1026,9 +1118,11 @@ async def test_create_new_experiment_error( mlflow.start_run.side_effect = ValueError('error') model_name = 'model_name' data = MagicMock() + prediction_data = MagicMock(spec=DataFrame) retrain_data = { 'prediction_model': {'model': MagicMock(), 'artifact_path': 'artifact_path'}, 'data_model': {'model': MagicMock(), 'artifact_path': 'artifact_path'}, + 'prediction_data': prediction_data, } mlflow_repository.get_model_params = MagicMock( @@ -1371,3 +1465,33 @@ async def test_update_production_model(mlflow_repository): 'mlflow_run_id': '0', 'mlflow_experiment_id': '0', } + + +def test_get_prediction_data_dataframe(mlflow_repository): + prediction_model = MagicMock() + retrain_dataset = DataFrame({'feat_1': [1, 2], 'target': [3, 4]}, index=['idx1', 'idx2']) + prediction_model.predict.return_value = DataFrame({'pred': [5, 6]}, index=['idx1', 'idx2']) + target_name = 'target' + + result = mlflow_repository.get_prediction_data(prediction_model, retrain_dataset, target_name) + + prediction_model.predict.assert_called_once_with(retrain_dataset) + assert 'prediction' in result.columns + assert 'target' in result.columns + assert 'timestamp' in result.columns + assert result.index.tolist() == [0, 1] + + +def test_get_prediction_data_array(mlflow_repository): + prediction_model = MagicMock() + retrain_dataset = DataFrame({'feat_1': [1, 2], 'target': [3, 4]}, index=['idx1', 'idx2']) + prediction_model.predict.return_value = [5, 6] + target_name = 'target' + + result = mlflow_repository.get_prediction_data(prediction_model, retrain_dataset, target_name) + + prediction_model.predict.assert_called_once_with(retrain_dataset) + assert 'prediction' in result.columns + assert 'target' in result.columns + assert 'timestamp' in result.columns + assert result.index.tolist() == [0, 1] diff --git a/tests/laborious/workflows/subworkflows/test_prediction_process.py b/tests/laborious/workflows/subworkflows/test_prediction_process.py index 0d8a0c1..6581033 100644 --- a/tests/laborious/workflows/subworkflows/test_prediction_process.py +++ b/tests/laborious/workflows/subworkflows/test_prediction_process.py @@ -31,6 +31,7 @@ async def test_run(workflow_mock, prediction_process): 'data': {'test': 'data'}, 'schema': 'test_schema', 'table_name': 'test_table', + 'transform_table_name': 'test_transform_table', 'model_id': 1, 'input_filters': {'test': 'filter'}, 'mlflow_transform_filters': {'test': 'filter'}, @@ -175,6 +176,7 @@ async def test_run(workflow_mock, prediction_process): 'metadata': metadata, 'path_flag': 'continue', 'data': 'predicted_data', + 'transformed_data': 'transformed_data', 'prediction_confidence': 0.95, 'timestamp': '2024-01-01', 'model_id': 1, @@ -183,6 +185,7 @@ async def test_run(workflow_mock, prediction_process): 'opc_output_config': input_data['opc_output_config'], 'schema': input_data['schema'], 'table_name': input_data['table_name'], + 'transform_table_name': input_data['transform_table_name'], 'comment': 'Error', 'prediction_store_policy': input_data['prediction_store_policy'], }, @@ -199,6 +202,7 @@ async def test_run_stop_at_input_gate(workflow_mock, prediction_process): 'data': {'test': 'data'}, 'schema': 'test_schema', 'table_name': 'test_table', + 'transform_table_name': 'test_transform_table', 'model_id': 1, 'input_filters': {'test': 'filter'}, 'mlflow_transform_filters': {'test': 'filter'}, @@ -257,6 +261,7 @@ async def test_run_stop_at_first_mlflow_response_gate(workflow_mock, prediction_ 'data': {'test': 'data'}, 'schema': 'test_schema', 'table_name': 'test_table', + 'transform_table_name': 'test_transform_table', 'model_id': 1, 'input_filters': {'test': 'filter'}, 'mlflow_transform_filters': {'test': 'filter'}, @@ -352,6 +357,7 @@ async def test_run_stop_at_mlflow_content_gate(workflow_mock, prediction_process 'data': {'test': 'data'}, 'schema': 'test_schema', 'table_name': 'test_table', + 'transform_table_name': 'test_transform_table', 'model_id': 1, 'input_filters': {'test': 'filter'}, 'mlflow_transform_filters': {'test': 'filter'}, @@ -467,6 +473,7 @@ async def test_run_stop_at_mlflow_last_response_gate(workflow_mock, prediction_p 'data': {'test': 'data'}, 'schema': 'test_schema', 'table_name': 'test_table', + 'transform_table_name': 'test_transform_table', 'model_id': 1, 'input_filters': {'test': 'filter'}, 'mlflow_transform_filters': {'test': 'filter'}, @@ -626,6 +633,7 @@ async def test_path_flag_handler_stop(workflow_mock, prediction_process): 'metadata': metadata, 'schema': schema, 'table_name': table_name, + 'transform_table_name': 'test_transform_table', 'model_id': model, 'last_timestamp': last_timestamp, 'model_name': model_name, @@ -664,6 +672,7 @@ async def test_path_flag_handler_repeat(workflow_mock, prediction_process): 'metadata': metadata, 'schema': schema, 'table_name': table_name, + 'transform_table_name': 'test_transform_table', 'model_id': model, 'last_timestamp': last_timestamp, 'model_name': model_name, @@ -714,6 +723,7 @@ async def test_path_flag_handler_continue(workflow_mock, prediction_process): 'metadata': metadata, 'schema': schema, 'table_name': table_name, + 'transform_table_name': 'test_transform_table', 'model_id': model, 'last_timestamp': last_timestamp, 'model_name': model_name, @@ -742,6 +752,7 @@ async def test_path_flag_handler_continue(workflow_mock, prediction_process): 'model_config': model_config, 'schema': schema, 'table_name': table_name, + 'transform_table_name': 'test_transform_table', 'comment': 'Prediction Process', 'opc_output_config': {'test': 'config'}, 'prediction_store_policy': prediction_store_policy, @@ -771,6 +782,7 @@ async def test_path_flag_handler_unknown(workflow_mock, prediction_process): **metadata, 'schema': schema, 'table_name': table_name, + 'transform_table_name': 'test_transform_table', 'model_id': model, 'last_timestamp': last_timestamp, 'model_name': model_name, diff --git a/tests/laborious/workflows/test_drift.py b/tests/laborious/workflows/test_drift.py new file mode 100644 index 0000000..09411aa --- /dev/null +++ b/tests/laborious/workflows/test_drift.py @@ -0,0 +1,247 @@ +from unittest.mock import ANY, AsyncMock, call, patch + +from pytest import fixture, mark + +from laborious.activities.activities import Activities +from laborious.workflows.drift import Drift +from sientia_do.temporal.constants import DATETIME_FORMAT_WITH_TZ + + +@fixture +def drift() -> Drift: + return Drift() + + +metadata = { + 'metadata': { + 'model_id': 'test_model_id', + 'model_name': 'test_model', + 'workflow_name': 'drift', + 'schedule_name': 'test_schedule', + }, +} + + +@mark.asyncio +@patch('laborious.workflows.drift.workflow', new_callable=AsyncMock) +async def test_run(workflow_mock: AsyncMock, drift: Drift): + # Arrange + input_data = { + 'schedule_name': 'test_schedule', + 'model_name': 'test_model', + 'model_id': 'test_model_id', + 'schema': 'test_schema', + 'source_table_name': 'test_source_table', + 'target_table_name': 'test_target_table', + 'interval': 60, + 'target_name': 'test_target', + 'drift_metrics': ['psi', 'ks'], + 'chunk_period': 'hour', + } + + target_data = {'data': 'test_target_data'} + reference_data = {'data': 'test_reference_data'} + drift_data = {'drift': 'test_drift_data'} + + workflow_mock.start_local_activity_method.side_effect = [ + target_data, reference_data + ] + + workflow_mock.execute_local_activity_method.return_value = drift_data + workflow_mock.execute_activity_method = AsyncMock() + + # Act + await drift.run(input_data) + + # Assert - Check start_local_activity_method calls + expected_gathering_query = f""" + SELECT * + FROM {input_data['schema']}.{input_data['source_table_name']} + WHERE + model_id = {input_data['model_id']} AND + timestamp > NOW() - INTERVAL '{input_data['interval']} minutes' + ORDER BY timestamp ASC + """ + + workflow_mock.start_local_activity_method.assert_has_calls( + [ + call( + Activities.load_custom_query, + { + **metadata, + 'query': expected_gathering_query, + 'datetime_columns': ['timestamp', 'created_at'], + }, + retry_policy=ANY, + start_to_close_timeout=ANY, + ), + call( + Activities.get_reference_data, + { + **metadata, + 'model_name': input_data['model_name'], + }, + retry_policy=ANY, + start_to_close_timeout=ANY, + ), + ] + ) + + # Assert - Check calculate_drift call + workflow_mock.execute_local_activity_method.assert_called_once_with( + Activities.calculate_drift, + { + **metadata, + 'target_data': target_data, + 'reference_data': reference_data, + 'model_name': input_data['model_name'], + 'model_id': input_data['model_id'], + 'target_name': input_data['target_name'], + 'drift_metrics': input_data['drift_metrics'], + 'chunk_period': input_data['chunk_period'], + }, + retry_policy=ANY, + start_to_close_timeout=ANY, + ) + + # Assert - Check export_data_to_postgres call + workflow_mock.execute_activity_method.assert_called_once_with( + Activities.export_data_to_postgres, + { + **metadata, + 'data': drift_data, + 'schema': input_data['schema'], + 'table_name': input_data['target_table_name'], + 'timestamp_conversion': {'column': 'timestamp', 'format': DATETIME_FORMAT_WITH_TZ}, + }, + retry_policy=ANY, + start_to_close_timeout=ANY, + ) + + +@mark.asyncio +@patch('laborious.workflows.drift.workflow', new_callable=AsyncMock) +async def test_run_empty_target_data(workflow_mock: AsyncMock, drift: Drift): + # Arrange + input_data = { + 'schedule_name': 'test_schedule', + 'model_name': 'test_model', + 'model_id': 'test_model_id', + 'schema': 'test_schema', + 'source_table_name': 'test_source_table', + 'target_table_name': 'test_target_table', + 'interval': 60, + 'target_name': 'test_target', + 'drift_metrics': ['psi', 'ks'], + } + + target_data = None + reference_data = {'data': 'test_reference_data'} + + workflow_mock.start_local_activity_method.side_effect = [target_data, reference_data] + + workflow_mock.execute_local_activity_method = AsyncMock() + workflow_mock.execute_activity_method = AsyncMock() + + # Act + await drift.run(input_data) + + # Assert - Should not call calculate_drift or export + workflow_mock.execute_local_activity_method.assert_not_called() + workflow_mock.execute_activity_method.assert_not_called() + + +@mark.asyncio +@patch('laborious.workflows.drift.workflow', new_callable=AsyncMock) +async def test_run_empty_drift_data(workflow_mock: AsyncMock, drift: Drift): + # Arrange + input_data = { + 'schedule_name': 'test_schedule', + 'model_name': 'test_model', + 'model_id': 'test_model_id', + 'schema': 'test_schema', + 'source_table_name': 'test_source_table', + 'target_table_name': 'test_target_table', + 'interval': 60, + 'target_name': 'test_target', + 'drift_metrics': ['psi', 'ks'], + } + + target_data = {'data': 'test_target_data'} + reference_data = {'data': 'test_reference_data'} + drift_data = None + + workflow_mock.start_local_activity_method.side_effect = [target_data, reference_data] + + workflow_mock.execute_local_activity_method.return_value = drift_data + workflow_mock.execute_activity_method = AsyncMock() + + # Act + await drift.run(input_data) + + # Assert - Should call calculate_drift but not export + workflow_mock.execute_local_activity_method.assert_called_once_with( + Activities.calculate_drift, + { + **metadata, + 'target_data': target_data, + 'reference_data': reference_data, + 'model_name': input_data['model_name'], + 'model_id': input_data['model_id'], + 'target_name': input_data['target_name'], + 'drift_metrics': input_data['drift_metrics'], + 'chunk_period': input_data.get('chunk_period', 'min'), + }, + retry_policy=ANY, + start_to_close_timeout=ANY, + ) + + workflow_mock.execute_activity_method.assert_not_called() + + +@mark.asyncio +@patch('laborious.workflows.drift.workflow', new_callable=AsyncMock) +async def test_run_default_chunk_period(workflow_mock: AsyncMock, drift: Drift): + # Arrange + input_data = { + 'schedule_name': 'test_schedule', + 'model_name': 'test_model', + 'model_id': 'test_model_id', + 'schema': 'test_schema', + 'source_table_name': 'test_source_table', + 'target_table_name': 'test_target_table', + 'interval': 60, + 'target_name': 'test_target', + 'drift_metrics': ['psi', 'ks'], + # chunk_period not provided, should default to 'min' + } + + target_data = {'data': 'test_target_data'} + reference_data = {'data': 'test_reference_data'} + drift_data = {'drift': 'test_drift_data'} + + workflow_mock.start_local_activity_method.side_effect = [target_data, reference_data] + + workflow_mock.execute_local_activity_method.return_value = drift_data + workflow_mock.execute_activity_method = AsyncMock() + + # Act + await drift.run(input_data) + + # Assert - Check calculate_drift call with default chunk_period + workflow_mock.execute_local_activity_method.assert_called_once_with( + Activities.calculate_drift, + { + **metadata, + 'target_data': target_data, + 'reference_data': reference_data, + 'model_name': input_data['model_name'], + 'model_id': input_data['model_id'], + 'target_name': input_data['target_name'], + 'drift_metrics': input_data['drift_metrics'], + 'chunk_period': 'min', # Default value + }, + 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 7979c47..6513da5 100644 --- a/tests/laborious/workflows/test_predictions_batch.py +++ b/tests/laborious/workflows/test_predictions_batch.py @@ -32,6 +32,7 @@ async def test_run(workflow_mock: AsyncMock, predictions_batch: PredictionsBatch 'query': 'SELECT * FROM test', 'schema': 'test_schema', 'table_name': 'test_table', + 'transform_table_name': 'test_transform_table', 'opc_output_config': 'test_opc_output_config', 'datetime_columns': ['timestamp', 'created_at'], 'prediction_store_policy': 'erl:1', @@ -59,6 +60,7 @@ async def test_run(workflow_mock: AsyncMock, predictions_batch: PredictionsBatch 'data': {'data': 'test_data'}, 'schema': input_data['schema'], 'table_name': input_data['table_name'], + 'transform_table_name': input_data['transform_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'}}), @@ -71,7 +73,7 @@ async def test_run(workflow_mock: AsyncMock, predictions_batch: PredictionsBatch 'model_config': input_data.get('model_config', {}), 'path_priority': input_data.get('path_priority', ['STOP', 'CONTINUE', 'REPEAT']), 'opc_output_config': input_data.get('opc_output_config', {}), - 'prediction_store_policy': input_data.get('prediction_store_policy', 'erl:1'), + 'prediction_store_policy': input_data.get('prediction_store_policy', 'lts:1'), } workflow_mock.execute_child_workflow.assert_has_calls( From 66193cea15343a53d606e62e9c9378463c82ed80 Mon Sep 17 00:00:00 2001 From: vitor-aignosi Date: Fri, 14 Nov 2025 08:47:39 -0300 Subject: [PATCH 03/27] SIENTIAPDE-1273 Refactor Drift class to improve target name handling - Extracted target name from model configuration in the Drift class for better clarity and maintainability. - Updated test cases to reflect changes in input data structure, ensuring consistency in target name retrieval. --- laborious/workflows/drift.py | 5 ++++- tests/laborious/workflows/test_drift.py | 2 +- 2 files changed, 5 insertions(+), 2 deletions(-) diff --git a/laborious/workflows/drift.py b/laborious/workflows/drift.py index 5677b37..d7c213c 100644 --- a/laborious/workflows/drift.py +++ b/laborious/workflows/drift.py @@ -32,6 +32,9 @@ class Drift: } } + model_config = input_data['model_config'] + target_name = model_config['target'] + gathering_query = f""" SELECT * FROM {input_data['schema']}.{input_data['source_table_name']} @@ -76,7 +79,7 @@ class Drift: 'reference_data': reference_data, 'model_name': input_data['model_name'], 'model_id': input_data['model_id'], - 'target_name': input_data['target_name'], + 'target_name': target_name, 'drift_metrics': input_data['drift_metrics'], 'chunk_period': input_data.get('chunk_period', 'min'), }, diff --git a/tests/laborious/workflows/test_drift.py b/tests/laborious/workflows/test_drift.py index 09411aa..82c9eb3 100644 --- a/tests/laborious/workflows/test_drift.py +++ b/tests/laborious/workflows/test_drift.py @@ -34,7 +34,7 @@ async def test_run(workflow_mock: AsyncMock, drift: Drift): 'source_table_name': 'test_source_table', 'target_table_name': 'test_target_table', 'interval': 60, - 'target_name': 'test_target', + 'model_config': {'target': 'test_target'}, 'drift_metrics': ['psi', 'ks'], 'chunk_period': 'hour', } From 64f0747e8e83887d88730dfec399d0a6bf670a17 Mon Sep 17 00:00:00 2001 From: vitor-aignosi Date: Fri, 14 Nov 2025 15:43:28 -0300 Subject: [PATCH 04/27] SIENTIAPDE-1273 Update version and enhance metrics calculation in Laborious system - Updated image tag in values.yaml from 1.1.0 to 1.1.1. - Modified GITHUB_BRANCH environment variable for consistency. - Added a new method `calculate_simple_metrics` in model_metrics.py to compute various model performance metrics including RMSE, MSE, MAE, and R2. - Integrated the new metrics calculation into the worker setup, allowing for concurrent processing of simple metrics. - Updated tests to cover the new metrics calculation functionality, ensuring comprehensive validation of the implementation. --- laborious/activities/model_metrics.py | 83 +++++- laborious/worker/worker.py | 18 ++ laborious/workflows/drift.py | 3 +- laborious/workflows/simple_metrics.py | 95 +++++++ .../activities/test_model_metrics.py | 252 ++++++++++++++++++ tests/laborious/workflows/test_drift.py | 18 +- .../workflows/test_simple_metrics.py | 228 ++++++++++++++++ values.yaml | 6 +- 8 files changed, 692 insertions(+), 11 deletions(-) create mode 100644 laborious/workflows/simple_metrics.py create mode 100644 tests/laborious/workflows/test_simple_metrics.py diff --git a/laborious/activities/model_metrics.py b/laborious/activities/model_metrics.py index 4e24ec7..c70cc81 100644 --- a/laborious/activities/model_metrics.py +++ b/laborious/activities/model_metrics.py @@ -13,6 +13,7 @@ with workflow.unsafe.imports_passed_through(): from sientia.ModelAnalysis import ModelAnalysis from laborious import metrics import time + import numpy as np from sientia_do.temporal.constants import DATETIME_FORMAT, DATETIME_FORMAT_WITH_TZ import warnings import traceback @@ -20,7 +21,6 @@ with workflow.unsafe.imports_passed_through(): warnings.filterwarnings('ignore', category=RuntimeWarning, message='Degrees of freedom <= 0') warnings.filterwarnings('ignore', category=RuntimeWarning, message='invalid value encountered in scalar divide') - class ModelMetrics(SientiaMonitoring): """ Metrics activities for the Laborious system. @@ -264,3 +264,84 @@ class ModelMetrics(SientiaMonitoring): return drift_df.to_dict() + async def calculate_simple_metrics(self, input_data: dict[str, Any]) -> dict[Hashable, Any]: + """ + Calculate simple metrics for a model. Metrics available are: + - rmse + - mse + - mae + - r2 + - accuracy + - precision + - recall + - f1 + Args: + input_data (dict[str, Any]): Input data containing: + - metadata (dict): Workflow execution metadata + - model_id (str): ID of the MLFlow model + - target_data (pd.DataFrame): Target data for calculating metrics, containing target and prediction columns + - metrics (list[str]): List of metrics to calculate + Returns: + dict[Hashable, Any]: Dictionary containing the calculated metrics + """ + + metadata = input_data['metadata'] + model_id = input_data['model_id'] + target_data = DataFrame(input_data['target_data']) + metrics = input_data['metrics'] + interval_minutes = input_data['interval_minutes'] + + data_size = len(target_data) + + output_data = [] + + diff = target_data['target'] - target_data['prediction'] + diff_squared = diff ** 2 + + self.info(f'Calculating simple metrics for model {model_id}: {metrics}', metadata) + + for metric in metrics: + if metric == 'rmse': + output_data.append({ + 'metric': 'rmse', + 'value': np.sqrt(np.mean(diff_squared)) + }) + elif metric == 'mse': + output_data.append({ + 'metric': 'mse', + 'value': np.mean(diff_squared) + }) + elif metric == 'mae': + output_data.append({ + 'metric': 'mae', + 'value': np.mean(np.abs(diff)) + }) + elif metric == 'r2': + y_true = target_data['target'] + y_mean = np.mean(y_true) + + ss_res = np.sum(diff_squared) + ss_tot = np.sum((y_true - y_mean) ** 2) + + # Evita divisão por zero + if ss_tot == 0: + r2_score = 0.0 + else: + r2_score = 1 - (ss_res / ss_tot) + + output_data.append({ + 'metric': 'r2', + 'value': r2_score + }) + + data = DataFrame(output_data) + data['model_id'] = model_id + data['timestamp'] = target_data['timestamp'].max() + data['data_size'] = data_size + data['interval_minutes'] = interval_minutes + + self.debug(f'Simple metrics dataframe: Size {data.shape} \n{data.head(5).to_string()}', metadata) + + return data.to_dict() + + diff --git a/laborious/worker/worker.py b/laborious/worker/worker.py index 3d65de4..6bfd779 100644 --- a/laborious/worker/worker.py +++ b/laborious/worker/worker.py @@ -51,6 +51,8 @@ with workflow.unsafe.imports_passed_through(): from laborious.workflows.minimal_retrain import MinimalRetrain from laborious.workflows.predictions_batch import PredictionsBatch from laborious.workflows.drift import Drift + from laborious.workflows.simple_metrics import SientiaMetrics + from laborious.workflows.sub_workflows.format_and_export_prediction import ( FormatAndExportPrediction, ) @@ -175,6 +177,22 @@ async def main(): workflow_task_poller_behavior=PollerBehaviorAutoscaling(), activity_task_poller_behavior=PollerBehaviorAutoscaling(), ), + Worker( + temporal_client, + task_queue='simple_metrics-queue', + workflows=[SimpleMetrics], + activities=[ + activities.load_custom_query, + activities.calculate_simple_metrics, + activities.export_data_to_postgres, + ], + max_concurrent_workflow_tasks=50, + max_concurrent_activities=50, + max_concurrent_local_activities=50, + max_cached_workflows=2, + workflow_task_poller_behavior=PollerBehaviorAutoscaling(), + activity_task_poller_behavior=PollerBehaviorAutoscaling(), + ), Worker( temporal_client, task_queue='predictions_batch-queue', diff --git a/laborious/workflows/drift.py b/laborious/workflows/drift.py index d7c213c..7afaef7 100644 --- a/laborious/workflows/drift.py +++ b/laborious/workflows/drift.py @@ -80,7 +80,8 @@ class Drift: 'model_name': input_data['model_name'], 'model_id': input_data['model_id'], 'target_name': target_name, - 'drift_metrics': input_data['drift_metrics'], + 'drift_metrics': input_data.get('drift_metrics', + ['kolmogorov_smirnov', 'jensen_shannon', 'wasserstein']), 'chunk_period': input_data.get('chunk_period', 'min'), }, retry_policy=retry_policy, diff --git a/laborious/workflows/simple_metrics.py b/laborious/workflows/simple_metrics.py new file mode 100644 index 0000000..f5f02d2 --- /dev/null +++ b/laborious/workflows/simple_metrics.py @@ -0,0 +1,95 @@ +from temporalio import workflow + +with workflow.unsafe.imports_passed_through(): + from datetime import timedelta + from typing import Any + + from sientia_do.temporal.policies import retry_policy + + from laborious.activities.activities import Activities + from sientia_do.temporal.constants import DATETIME_FORMAT_WITH_TZ + + +@workflow.defn(name='simple_metrics') +class SimpleMetrics: + @workflow.run + async def run(self, input_data: dict[str, Any]): + """ + Execute the simple metrics workflow. + """ + metadata = { + 'metadata': { + 'schedule_name': input_data['schedule_name'], + 'model_name': input_data['model_name'], + 'model_id': input_data['model_id'], + 'workflow_name': 'simple_metrics', + } + } + + model_id = input_data['model_id'] + interval_minutes = input_data['interval_minutes'] + + model_config = input_data['model_config'] + target_name = model_config['target'] + + query = f""" + select p."timestamp", p.prediction, ld.value as "target" + from {input_data['schema']}.{input_data['predictions_table_name']} p + inner join {input_data['schema']}.{input_data['data_table_name']} ld + on p."timestamp" = ld."timestamp" + where + p.model_id = {model_id} and + p.prediction is not null and + ld.variable = '{target_name}' and + ld.value is not null and + p."timestamp" >= NOW() - INTERVAL '{interval_minutes} minutes' + order by + p."timestamp" desc; + """ + + target_data = await workflow.execute_local_activity_method( + Activities.load_custom_query, + { + **metadata, + 'query': query, + 'datetime_columns': ['timestamp'], + }, + retry_policy=retry_policy, + start_to_close_timeout=timedelta(seconds=300), + ) + + if not target_data: + return + + simple_metrics = await workflow.execute_local_activity_method( + Activities.calculate_simple_metrics, + { + **metadata, + 'model_id': model_id, + 'target_data': target_data, + 'metrics': input_data.get('metrics', ['rmse', 'mse', 'mae', 'r2']), + 'interval_minutes': interval_minutes, + }, + retry_policy=retry_policy, + start_to_close_timeout=timedelta(seconds=300), + ) + + + if not simple_metrics: + return + + await workflow.execute_activity_method( + Activities.export_data_to_postgres, + { + **metadata, + 'data': simple_metrics, + 'schema': input_data['schema'], + 'table_name': input_data['target_table_name'], + 'timestamp_conversion': { + 'column': 'timestamp', + 'format': DATETIME_FORMAT_WITH_TZ, + }, + }, + retry_policy=retry_policy, + start_to_close_timeout=timedelta(seconds=300), + ) \ No newline at end of file diff --git a/tests/laborious/activities/test_model_metrics.py b/tests/laborious/activities/test_model_metrics.py index 8b55df8..87c3a68 100644 --- a/tests/laborious/activities/test_model_metrics.py +++ b/tests/laborious/activities/test_model_metrics.py @@ -639,3 +639,255 @@ async def test_get_drift_metrics_univariate_error( else: raise AssertionError('Expected Exception') + +@mark.asyncio +async def test_calculate_simple_metrics_success_all_metrics( + model_metrics_activity +): + # Arrange + target_data = DataFrame({ + 'timestamp': ['2023-05-26 11:12:27', '2023-05-26 11:12:28', '2023-05-26 11:12:29'], + 'target': [1.0, 2.0, 3.0], + 'prediction': [1.1, 2.1, 2.9], + }) + + input_data = { + **metadata, + 'model_id': 'test_model_id', + 'target_data': target_data.to_dict(), + 'metrics': ['rmse', 'mse', 'mae', 'r2'], + 'interval_minutes': 5, + } + + # Act + result = DataFrame(await model_metrics_activity.calculate_simple_metrics(input_data)) + + # Assert + assert len(result['metric']) == 4 + assert 'rmse' in result['metric'].values + assert 'mse' in result['metric'].values + assert 'mae' in result['metric'].values + assert 'r2' in result['metric'].values + assert all(model_id == 'test_model_id' for model_id in result['model_id'].values) + assert all(timestamp == '2023-05-26 11:12:29' for timestamp in result['timestamp'].values) + assert all(data_size == 3 for data_size in result['data_size'].values) + assert all(interval_minutes == 5 for interval_minutes in result['interval_minutes'].values) + model_metrics_activity.info.assert_called_once_with( + 'Calculating simple metrics for model test_model_id: [\'rmse\', \'mse\', \'mae\', \'r2\']', + metadata['metadata'] + ) + model_metrics_activity.debug.assert_called_once() + + +@mark.asyncio +async def test_calculate_simple_metrics_success_rmse_only( + model_metrics_activity +): + # Arrange + target_data = DataFrame({ + 'timestamp': ['2023-05-26 11:12:27', '2023-05-26 11:12:28'], + 'target': [1.0, 2.0], + 'prediction': [1.1, 2.1], + }) + + input_data = { + **metadata, + 'model_id': 'test_model_id', + 'target_data': target_data.to_dict(), + 'metrics': ['rmse'], + 'interval_minutes': 5, + } + + # Act + result = DataFrame(await model_metrics_activity.calculate_simple_metrics(input_data)) + + # Assert + assert len(result['metric']) == 1 + assert result['metric'].values[0] == 'rmse' + assert result['model_id'].values[0] == 'test_model_id' + assert result['timestamp'].values[0] == '2023-05-26 11:12:28' + assert result['data_size'].values[0] == 2 + assert result['interval_minutes'].values[0] == 5 + model_metrics_activity.info.assert_called_once_with( + 'Calculating simple metrics for model test_model_id: [\'rmse\']', + metadata['metadata'] + ) + + +@mark.asyncio +async def test_calculate_simple_metrics_success_mse_only( + model_metrics_activity +): + # Arrange + target_data = DataFrame({ + 'timestamp': ['2023-05-26 11:12:27', '2023-05-26 11:12:28'], + 'target': [1.0, 2.0], + 'prediction': [1.1, 2.1], + }) + + input_data = { + **metadata, + 'model_id': 'test_model_id', + 'target_data': target_data.to_dict(), + 'metrics': ['mse'], + 'interval_minutes': 5, + } + + # Act + result = DataFrame(await model_metrics_activity.calculate_simple_metrics(input_data)) + + # Assert + assert len(result['metric']) == 1 + assert result['metric'].values[0] == 'mse' + assert result['model_id'].values[0] == 'test_model_id' + assert result['timestamp'].values[0] == '2023-05-26 11:12:28' + assert result['data_size'].values[0] == 2 + assert result['interval_minutes'].values[0] == 5 + model_metrics_activity.info.assert_called_once_with( + 'Calculating simple metrics for model test_model_id: [\'mse\']', + metadata['metadata'] + ) + + +@mark.asyncio +async def test_calculate_simple_metrics_success_mae_only( + model_metrics_activity +): + # Arrange + target_data = DataFrame({ + 'timestamp': ['2023-05-26 11:12:27', '2023-05-26 11:12:28'], + 'target': [1.0, 2.0], + 'prediction': [1.1, 2.1], + }) + + input_data = { + **metadata, + 'model_id': 'test_model_id', + 'target_data': target_data.to_dict(), + 'metrics': ['mae'], + 'interval_minutes': 5, + } + + # Act + result = DataFrame(await model_metrics_activity.calculate_simple_metrics(input_data)) + + # Assert + assert len(result['metric']) == 1 + assert result['metric'].values[0] == 'mae' + assert result['model_id'].values[0] == 'test_model_id' + assert result['timestamp'].values[0] == '2023-05-26 11:12:28' + assert result['data_size'].values[0] == 2 + assert result['interval_minutes'].values[0] == 5 + model_metrics_activity.info.assert_called_once_with( + 'Calculating simple metrics for model test_model_id: [\'mae\']', + metadata['metadata'] + ) + + +@mark.asyncio +async def test_calculate_simple_metrics_success_r2_only( + model_metrics_activity +): + # Arrange + target_data = DataFrame({ + 'timestamp': ['2023-05-26 11:12:27', '2023-05-26 11:12:28'], + 'target': [1.0, 2.0], + 'prediction': [1.1, 2.1], + }) + + input_data = { + **metadata, + 'model_id': 'test_model_id', + 'target_data': target_data.to_dict(), + 'metrics': ['r2'], + 'interval_minutes': 5, + } + + # Act + result = DataFrame(await model_metrics_activity.calculate_simple_metrics(input_data)) + + # Assert + assert len(result['metric']) == 1 + assert result['metric'].values[0] == 'r2' + assert result['model_id'].values[0] == 'test_model_id' + assert result['timestamp'].values[0] == '2023-05-26 11:12:28' + assert result['data_size'].values[0] == 2 + assert result['interval_minutes'].values[0] == 5 + model_metrics_activity.info.assert_called_once_with( + 'Calculating simple metrics for model test_model_id: [\'r2\']', + metadata['metadata'] + ) + + +@mark.asyncio +async def test_calculate_simple_metrics_r2_zero_ss_tot( + model_metrics_activity +): + # Arrange + # All target values are the same, so ss_tot will be 0 + target_data = DataFrame({ + 'timestamp': ['2023-05-26 11:12:27', '2023-05-26 11:12:28'], + 'target': [1.0, 1.0], + 'prediction': [1.1, 1.1], + }) + + input_data = { + **metadata, + 'model_id': 'test_model_id', + 'target_data': target_data.to_dict(), + 'metrics': ['r2'], + 'interval_minutes': 5, + } + + # Act + result = DataFrame(await model_metrics_activity.calculate_simple_metrics(input_data)) + + # Assert + assert len(result['metric']) == 1 + assert result['metric'].values[0] == 'r2' + assert result['value'].values[0] == 0.0 # Should return 0.0 when ss_tot == 0 + assert result['model_id'].values[0] == 'test_model_id' + assert result['timestamp'].values[0] == '2023-05-26 11:12:28' + assert result['data_size'].values[0] == 2 + assert result['interval_minutes'].values[0] == 5 + model_metrics_activity.info.assert_called_once_with( + 'Calculating simple metrics for model test_model_id: [\'r2\']', + metadata['metadata'] + ) + + +@mark.asyncio +async def test_calculate_simple_metrics_success_multiple_metrics_subset( + model_metrics_activity +): + # Arrange + target_data = DataFrame({ + 'timestamp': ['2023-05-26 11:12:27', '2023-05-26 11:12:28', '2023-05-26 11:12:29'], + 'target': [1.0, 2.0, 3.0], + 'prediction': [1.1, 2.1, 2.9], + }) + + input_data = { + **metadata, + 'model_id': 'test_model_id', + 'target_data': target_data.to_dict(), + 'metrics': ['rmse', 'mae'], + 'interval_minutes': 5, + } + + # Act + result = DataFrame(await model_metrics_activity.calculate_simple_metrics(input_data)) + + # Assert + assert len(result['metric']) == 2 + assert 'rmse' in result['metric'].values + assert 'mae' in result['metric'].values + assert all(model_id == 'test_model_id' for model_id in result['model_id'].values) + assert all(timestamp == '2023-05-26 11:12:29' for timestamp in result['timestamp'].values) + assert all(data_size == 3 for data_size in result['data_size'].values) + assert all(interval_minutes == 5 for interval_minutes in result['interval_minutes'].values) + model_metrics_activity.info.assert_called_once_with( + 'Calculating simple metrics for model test_model_id: [\'rmse\', \'mae\']', + metadata['metadata'] + ) + diff --git a/tests/laborious/workflows/test_drift.py b/tests/laborious/workflows/test_drift.py index 82c9eb3..2f3f1dc 100644 --- a/tests/laborious/workflows/test_drift.py +++ b/tests/laborious/workflows/test_drift.py @@ -39,6 +39,8 @@ async def test_run(workflow_mock: AsyncMock, drift: Drift): 'chunk_period': 'hour', } + target_name = input_data['model_config']['target'] + target_data = {'data': 'test_target_data'} reference_data = {'data': 'test_reference_data'} drift_data = {'drift': 'test_drift_data'} @@ -96,7 +98,7 @@ async def test_run(workflow_mock: AsyncMock, drift: Drift): 'reference_data': reference_data, 'model_name': input_data['model_name'], 'model_id': input_data['model_id'], - 'target_name': input_data['target_name'], + 'target_name': target_name, 'drift_metrics': input_data['drift_metrics'], 'chunk_period': input_data['chunk_period'], }, @@ -131,7 +133,7 @@ async def test_run_empty_target_data(workflow_mock: AsyncMock, drift: Drift): 'source_table_name': 'test_source_table', 'target_table_name': 'test_target_table', 'interval': 60, - 'target_name': 'test_target', + 'model_config': {'target': 'test_target'}, 'drift_metrics': ['psi', 'ks'], } @@ -163,10 +165,12 @@ async def test_run_empty_drift_data(workflow_mock: AsyncMock, drift: Drift): 'source_table_name': 'test_source_table', 'target_table_name': 'test_target_table', 'interval': 60, - 'target_name': 'test_target', + 'model_config': {'target': 'test_target'}, 'drift_metrics': ['psi', 'ks'], } + target_name = input_data['model_config']['target'] + target_data = {'data': 'test_target_data'} reference_data = {'data': 'test_reference_data'} drift_data = None @@ -188,7 +192,7 @@ async def test_run_empty_drift_data(workflow_mock: AsyncMock, drift: Drift): 'reference_data': reference_data, 'model_name': input_data['model_name'], 'model_id': input_data['model_id'], - 'target_name': input_data['target_name'], + 'target_name': target_name, 'drift_metrics': input_data['drift_metrics'], 'chunk_period': input_data.get('chunk_period', 'min'), }, @@ -211,11 +215,13 @@ async def test_run_default_chunk_period(workflow_mock: AsyncMock, drift: Drift): 'source_table_name': 'test_source_table', 'target_table_name': 'test_target_table', 'interval': 60, - 'target_name': 'test_target', + 'model_config': {'target': 'test_target'}, 'drift_metrics': ['psi', 'ks'], # chunk_period not provided, should default to 'min' } + target_name = input_data['model_config']['target'] + target_data = {'data': 'test_target_data'} reference_data = {'data': 'test_reference_data'} drift_data = {'drift': 'test_drift_data'} @@ -237,7 +243,7 @@ async def test_run_default_chunk_period(workflow_mock: AsyncMock, drift: Drift): 'reference_data': reference_data, 'model_name': input_data['model_name'], 'model_id': input_data['model_id'], - 'target_name': input_data['target_name'], + 'target_name': target_name, 'drift_metrics': input_data['drift_metrics'], 'chunk_period': 'min', # Default value }, diff --git a/tests/laborious/workflows/test_simple_metrics.py b/tests/laborious/workflows/test_simple_metrics.py new file mode 100644 index 0000000..b185b1a --- /dev/null +++ b/tests/laborious/workflows/test_simple_metrics.py @@ -0,0 +1,228 @@ +from unittest.mock import ANY, AsyncMock, call, patch + +from pytest import fixture, mark + +from laborious.activities.activities import Activities +from laborious.workflows.simple_metrics import SimpleMetrics +from sientia_do.temporal.constants import DATETIME_FORMAT_WITH_TZ + + +@fixture +def simple_metrics() -> SimpleMetrics: + return SimpleMetrics() + + +metadata = { + 'metadata': { + 'model_id': 'test_model_id', + 'model_name': 'test_model', + 'workflow_name': 'simple_metrics', + 'schedule_name': 'test_schedule', + }, +} + + +@mark.asyncio +@patch('laborious.workflows.simple_metrics.workflow', new_callable=AsyncMock) +async def test_run(workflow_mock: AsyncMock, simple_metrics: SimpleMetrics): + # Arrange + input_data = { + 'schedule_name': 'test_schedule', + 'model_name': 'test_model', + 'model_id': 'test_model_id', + 'interval_minutes': 60, + 'model_config': {'target': 'test_target'}, + 'schema': 'test_schema', + 'predictions_table_name': 'test_predictions_table', + 'data_table_name': 'test_data_table', + 'target_table_name': 'test_target_table', + 'metrics': ['rmse', 'mse', 'mae', 'r2'], + } + + target_data = {'data': 'test_target_data'} + simple_metrics_data = {'metrics': 'test_simple_metrics_data'} + + workflow_mock.execute_local_activity_method.side_effect = [ + target_data, simple_metrics_data + ] + + workflow_mock.execute_activity_method = AsyncMock() + + # Act + await simple_metrics.run(input_data) + + # Assert - Check load_custom_query call + expected_query = f""" + select p."timestamp", p.prediction, ld.value as "target" + from {input_data['schema']}.{input_data['predictions_table_name']} p + inner join {input_data['schema']}.{input_data['data_table_name']} ld + on p."timestamp" = ld."timestamp" + where + p.model_id = {input_data['model_id']} and + p.prediction is not null and + ld.variable = '{input_data['model_config']['target']}' and + ld.value is not null and + p."timestamp" >= NOW() - INTERVAL '{input_data['interval_minutes']} minutes' + order by + p."timestamp" desc; + """ + + workflow_mock.execute_local_activity_method.assert_has_calls( + [ + call( + Activities.load_custom_query, + { + **metadata, + 'query': expected_query, + 'datetime_columns': ['timestamp'], + }, + retry_policy=ANY, + start_to_close_timeout=ANY, + ), + call( + Activities.calculate_simple_metrics, + { + **metadata, + 'model_id': input_data['model_id'], + 'target_data': target_data, + 'metrics': input_data['metrics'], + 'interval_minutes': input_data['interval_minutes'], + }, + retry_policy=ANY, + start_to_close_timeout=ANY, + ), + ] + ) + + # Assert - Check export_data_to_postgres call + workflow_mock.execute_activity_method.assert_called_once_with( + Activities.export_data_to_postgres, + { + **metadata, + 'data': simple_metrics_data, + 'schema': input_data['schema'], + 'table_name': input_data['target_table_name'], + 'timestamp_conversion': {'column': 'timestamp', 'format': DATETIME_FORMAT_WITH_TZ}, + }, + retry_policy=ANY, + start_to_close_timeout=ANY, + ) + + +@mark.asyncio +@patch('laborious.workflows.simple_metrics.workflow', new_callable=AsyncMock) +async def test_run_empty_target_data(workflow_mock: AsyncMock, simple_metrics: SimpleMetrics): + # Arrange + input_data = { + 'schedule_name': 'test_schedule', + 'model_name': 'test_model', + 'model_id': 'test_model_id', + 'interval_minutes': 60, + 'model_config': {'target': 'test_target'}, + 'schema': 'test_schema', + 'predictions_table_name': 'test_predictions_table', + 'data_table_name': 'test_data_table', + 'target_table_name': 'test_target_table', + 'metrics': ['rmse', 'mse'], + } + + target_data = None + + workflow_mock.execute_local_activity_method.return_value = target_data + workflow_mock.execute_activity_method = AsyncMock() + + # Act + await simple_metrics.run(input_data) + + # Assert - Should not call calculate_simple_metrics or export + assert workflow_mock.execute_local_activity_method.call_count == 1 + workflow_mock.execute_activity_method.assert_not_called() + + +@mark.asyncio +@patch('laborious.workflows.simple_metrics.workflow', new_callable=AsyncMock) +async def test_run_empty_simple_metrics(workflow_mock: AsyncMock, simple_metrics: SimpleMetrics): + # Arrange + input_data = { + 'schedule_name': 'test_schedule', + 'model_name': 'test_model', + 'model_id': 'test_model_id', + 'interval_minutes': 60, + 'model_config': {'target': 'test_target'}, + 'schema': 'test_schema', + 'predictions_table_name': 'test_predictions_table', + 'data_table_name': 'test_data_table', + 'target_table_name': 'test_target_table', + 'metrics': ['rmse', 'mse'], + } + + target_data = {'data': 'test_target_data'} + simple_metrics_data = None + + workflow_mock.execute_local_activity_method.side_effect = [ + target_data, simple_metrics_data + ] + + workflow_mock.execute_activity_method = AsyncMock() + + # Act + await simple_metrics.run(input_data) + + # Assert - Should call calculate_simple_metrics but not export + assert workflow_mock.execute_local_activity_method.call_count == 2 + workflow_mock.execute_activity_method.assert_not_called() + + +@mark.asyncio +@patch('laborious.workflows.simple_metrics.workflow', new_callable=AsyncMock) +async def test_run_default_metrics(workflow_mock: AsyncMock, simple_metrics: SimpleMetrics): + # Arrange + input_data = { + 'schedule_name': 'test_schedule', + 'model_name': 'test_model', + 'model_id': 'test_model_id', + 'interval_minutes': 60, + 'model_config': {'target': 'test_target'}, + 'schema': 'test_schema', + 'predictions_table_name': 'test_predictions_table', + 'data_table_name': 'test_data_table', + 'target_table_name': 'test_target_table', + # metrics not provided, should default to ['rmse', 'mse', 'mae', 'r2'] + } + + target_data = {'data': 'test_target_data'} + simple_metrics_data = {'metrics': 'test_simple_metrics_data'} + + workflow_mock.execute_local_activity_method.side_effect = [ + target_data, simple_metrics_data + ] + + workflow_mock.execute_activity_method = AsyncMock() + + # Act + await simple_metrics.run(input_data) + + # Assert - Check calculate_simple_metrics call with default metrics + workflow_mock.execute_local_activity_method.assert_has_calls( + [ + call( + Activities.load_custom_query, + ANY, + retry_policy=ANY, + start_to_close_timeout=ANY, + ), + call( + Activities.calculate_simple_metrics, + { + **metadata, + 'model_id': input_data['model_id'], + 'target_data': target_data, + 'metrics': ['rmse', 'mse', 'mae', 'r2'], # Default value + 'interval_minutes': input_data['interval_minutes'], + }, + retry_policy=ANY, + start_to_close_timeout=ANY, + ), + ] + ) + diff --git a/values.yaml b/values.yaml index f8ad89e..5034e12 100644 --- a/values.yaml +++ b/values.yaml @@ -11,7 +11,7 @@ image: # This sets the pull policy for images. pullPolicy: Always # Overrides the image tag whose default is the chart appVersion. - tag: "1.1.0" + tag: "1.1.1" 0# This is for the secrets for pulling an image from a private repository more information can be found here: https://kubernetes.io/docs/tasks/configure-pod-container/pull-image-private-registry/ imagePullSecrets: @@ -151,7 +151,7 @@ env: - name: GITHUB_REPO_URL value: "git@github.com:Aignosi/sientia-dataops-laborious_temporal.git" - name: GITHUB_BRANCH - value: "feature/SIENTIAPDE-1325-adicionar-metricas-especificas-de-operacoes-externas" + value: "feature/SIENTIAPDE-1273" - name: PYTHON_APP value: "laborious.worker.worker" @@ -234,7 +234,7 @@ ssh: # kubectl create secret docker-registry docker-hub-secret --namespace sientia --docker-server=http://aignosi.azurecr.io --docker-username=aignosi --docker-password=5I5zpQ6sRaHqX1hD3dr+2mo647yO3FRc359/wu6gsP+ACRDRz5mp -# helm upgrade --install sientia-laborious-worker sientia/sientia-module -n sientia --create-namespace -f ./values.yaml --version 0.5.0 +# helm upgrade --install sientia-laborious-worker sientia/sientia-module -n sientia --create-namespace -f ./values.yaml --version 0.6.0 # kubectl create secret generic git-ssh-key-sientia-laborious-worker \ # --namespace sientia \ From 4191f68d3b78aed76be357b20f940e11428a18ea Mon Sep 17 00:00:00 2001 From: vitor-aignosi Date: Fri, 14 Nov 2025 16:06:51 -0300 Subject: [PATCH 05/27] SIENTIAPDE-1273 Refactor ModelMetrics return format and correct import name in worker module - Changed the return format of the metrics data in ModelMetrics from a dictionary to a list for improved usability. - Corrected the import statement for SimpleMetrics in the worker module to ensure consistency and clarity. --- laborious/worker/worker.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/laborious/worker/worker.py b/laborious/worker/worker.py index 6bfd779..1d28016 100644 --- a/laborious/worker/worker.py +++ b/laborious/worker/worker.py @@ -51,7 +51,7 @@ with workflow.unsafe.imports_passed_through(): from laborious.workflows.minimal_retrain import MinimalRetrain from laborious.workflows.predictions_batch import PredictionsBatch from laborious.workflows.drift import Drift - from laborious.workflows.simple_metrics import SientiaMetrics + from laborious.workflows.simple_metrics import SimpleMetrics from laborious.workflows.sub_workflows.format_and_export_prediction import ( FormatAndExportPrediction, From d2b365a34d3fa65a6b9dce0c4ad5c322780e8563 Mon Sep 17 00:00:00 2001 From: vitor-aignosi Date: Fri, 14 Nov 2025 16:55:00 -0300 Subject: [PATCH 06/27] SIENTIAPDE-1273 Refactor data handling in various modules to ensure DataFrame consistency - Replaced direct DataFrame instantiation with `ensure_dataframe` utility in Gates, MLFlow, OPC, and ModelMetrics classes to standardize data handling. - Updated return types in several asynchronous methods to return DataFrames instead of dictionaries for improved usability. - Adjusted data export processes in workflows to convert DataFrames to dictionaries with `to_dict(orient='records')` for compatibility with downstream systems. --- laborious/activities/gates.py | 25 +-- laborious/activities/mlflow.py | 9 +- laborious/activities/model_metrics.py | 8 +- laborious/activities/opc.py | 8 +- laborious/utils/dataframe_utils.py | 31 ++++ laborious/utils/temporal_codec.py | 149 +++++++++++++++++ laborious/worker/worker.py | 5 + laborious/workflows/drift.py | 2 +- laborious/workflows/minimal_retrain.py | 2 +- laborious/workflows/simple_metrics.py | 2 +- .../format_and_export_prediction.py | 4 +- tests.ipynb | 151 ++++++++++++++++++ 12 files changed, 368 insertions(+), 28 deletions(-) create mode 100644 laborious/utils/dataframe_utils.py create mode 100644 laborious/utils/temporal_codec.py diff --git a/laborious/activities/gates.py b/laborious/activities/gates.py index 56a1569..7c8ae76 100644 --- a/laborious/activities/gates.py +++ b/laborious/activities/gates.py @@ -15,6 +15,7 @@ with workflow.unsafe.imports_passed_through(): from sientia_do.temporal.constants import DATETIME_FORMAT_WITH_TZ, now from laborious import metrics + from laborious.utils.dataframe_utils import ensure_dataframe from laborious.utils.filters.conditional_filters import ( filter_empty_data, filter_specific_variables_null_values, @@ -149,7 +150,7 @@ class Gates(SientiaMonitoring): self.info('Performing input gate...', metadata) filters = input_data['filters'] - data = DataFrame(input_data['data']) + data = ensure_dataframe(input_data['data']) path_priority = input_data['path_priority'] filter_output = [] @@ -404,7 +405,7 @@ class Gates(SientiaMonitoring): @activity.defn(name='format_transformed_data') - async def format_transformed_data(self, input_data: dict[str, Any]) -> dict[Any, Any]: + async def format_transformed_data(self, input_data: dict[str, Any]) -> DataFrame: """ Format transformed data according to configured storage policies. """ @@ -414,7 +415,7 @@ class Gates(SientiaMonitoring): self.info('Formatting transformed data...', metadata) - data = DataFrame(input_data['data']) + data = ensure_dataframe(input_data['data']) data['timestamp'] = data.index data = data.reset_index(drop=True) @@ -422,10 +423,10 @@ class Gates(SientiaMonitoring): data = data.melt(id_vars='timestamp', var_name='variable', value_name='value') data['model_id'] = model_id - return data.to_dict() + return data @activity.defn(name='format_prediction') - async def format_prediction(self, input_data: dict[str, Any]) -> dict[Any, Any]: + async def format_prediction(self, input_data: dict[str, Any]) -> DataFrame: """ Format prediction data according to configured storage policies. @@ -453,7 +454,7 @@ class Gates(SientiaMonitoring): prediction_store_policy = input_data['prediction_store_policy'] self.info('Formatting prediction...', metadata) - data = DataFrame(input_data['data']) + data = ensure_dataframe(input_data['data']) # Create timestamp column from index and reset index data['timestamp'] = data.index @@ -495,10 +496,10 @@ class Gates(SientiaMonitoring): self.info(f'Prediction formatted: {len(data)} rows', metadata) self.debug(f'Prediction data: {data.head(5).to_string()}', metadata) - return data.to_dict() + return data @activity.defn(name='format_default_prediction') - async def format_default_prediction(self, input_data: dict[str, Any]) -> dict[Any, Any]: + async def format_default_prediction(self, input_data: dict[str, Any]) -> DataFrame: """ Create and format default prediction data for error conditions. @@ -540,10 +541,10 @@ class Gates(SientiaMonitoring): ) self.info(f'Default prediction formatted: {data.size} rows', metadata) - return data.to_dict() + return data @activity.defn(name='format_retrain_report') - async def format_retrain_report(self, input_data: dict[str, Any]) -> dict[Any, Any]: + async def format_retrain_report(self, input_data: dict[str, Any]) -> DataFrame: """ Format retrain report data according to configured storage policies. """ @@ -572,7 +573,7 @@ class Gates(SientiaMonitoring): self.debug(f'Retrain report: {report.to_csv()}', metadata) - return report.to_dict() + return report @activity.defn(name='get_last_timestamp') async def get_last_timestamp(self, input_data: dict[str, Any]) -> str: @@ -601,7 +602,7 @@ class Gates(SientiaMonitoring): self.info('Getting last timestamp...', metadata) - data = DataFrame(input_data['data']) + data = ensure_dataframe(input_data['data']) self.debug(f'Input data: {data.head(5).to_string()}', metadata) diff --git a/laborious/activities/mlflow.py b/laborious/activities/mlflow.py index 2c50abe..0447215 100644 --- a/laborious/activities/mlflow.py +++ b/laborious/activities/mlflow.py @@ -20,6 +20,7 @@ with workflow.unsafe.imports_passed_through(): now, ) + from laborious.utils.dataframe_utils import ensure_dataframe from laborious.utils.repository.minio_repository import MinioRepository from laborious.utils.repository.model_repository import MLFlowRepository @@ -139,7 +140,7 @@ class MLFlow(SientiaMonitoring): """ metadata = input_data['metadata'] self.info('Transforming data...', metadata) - data = DataFrame(input_data['data']) + data = ensure_dataframe(input_data['data']) model_name = input_data['model_name'] model_config = input_data.get('model_config', {}) @@ -211,7 +212,7 @@ class MLFlow(SientiaMonitoring): """ metadata = input_data['metadata'] self.info('Predicting data...', metadata) - data = DataFrame(input_data['data']) + data = ensure_dataframe(input_data['data']) model_name = input_data['model_name'] model_config = input_data.get('model_config', {}) @@ -422,7 +423,7 @@ class MLFlow(SientiaMonitoring): @activity.defn(name='get_reference_data') - async def get_reference_data(self, input_data: dict[str, Any]) -> dict[Hashable, Any] | None: + async def get_reference_data(self, input_data: dict[str, Any]) -> DataFrame | None: """ Get reference data from the MLflow Model Registry. @@ -451,4 +452,4 @@ class MLFlow(SientiaMonitoring): reference_data['timestamp'] = to_datetime(reference_data['timestamp']) reference_data['timestamp'] = reference_data['timestamp'].dt.strftime(DATETIME_FORMAT) - return reference_data.to_dict() \ No newline at end of file + return reference_data \ No newline at end of file diff --git a/laborious/activities/model_metrics.py b/laborious/activities/model_metrics.py index c70cc81..9896fa7 100644 --- a/laborious/activities/model_metrics.py +++ b/laborious/activities/model_metrics.py @@ -138,7 +138,7 @@ class ModelMetrics(SientiaMonitoring): @activity.defn(name='calculate_drift') - async def calculate_drift(self, input_data: dict[str, Any]) -> dict[Hashable, Any]: + async def calculate_drift(self, input_data: dict[str, Any]) -> DataFrame | dict: """ Calculate drift metrics for a model. @@ -262,9 +262,9 @@ class ModelMetrics(SientiaMonitoring): self.debug(f'Drift dataframe: Size {drift_df.shape} \n{drift_df.head(5).to_string()}', metadata) - return drift_df.to_dict() + return drift_df - async def calculate_simple_metrics(self, input_data: dict[str, Any]) -> dict[Hashable, Any]: + async def calculate_simple_metrics(self, input_data: dict[str, Any]) -> DataFrame: """ Calculate simple metrics for a model. Metrics available are: - rmse @@ -342,6 +342,6 @@ class ModelMetrics(SientiaMonitoring): self.debug(f'Simple metrics dataframe: Size {data.shape} \n{data.head(5).to_string()}', metadata) - return data.to_dict() + return data diff --git a/laborious/activities/opc.py b/laborious/activities/opc.py index aa93ca5..404840e 100644 --- a/laborious/activities/opc.py +++ b/laborious/activities/opc.py @@ -1,3 +1,4 @@ +from typing import Hashable from temporalio import activity, workflow with workflow.unsafe.imports_passed_through(): @@ -11,6 +12,7 @@ with workflow.unsafe.imports_passed_through(): from sientia_do.observability.metrics_controller import MetricsController from sientia_do.observability.sientia_monitoring import SientiaMonitoring + from laborious.utils.dataframe_utils import ensure_dataframe from laborious.utils.repository.opc_repository import OpcRepository OPC_WRITTING_ERROR_CONFIDENCE = 12 @@ -275,7 +277,7 @@ class OPC(SientiaMonitoring): @activity.defn(name='write_opc_data') async def write_opc_data( self, input_data: dict[str, Any] - ) -> tuple[dict[Any, Any], dict[str, dict[str, float | None]]]: + ) -> tuple[dict[Hashable, Any], dict[str, dict[str, float | None]]]: """ Write prediction and confidence data to OPC servers. The two writing operations are optional and independent of each other. @@ -295,7 +297,7 @@ class OPC(SientiaMonitoring): """ metadata = input_data['metadata'] self.info('Writing data to OPC servers...', metadata) - data = DataFrame(input_data['data']) + data = ensure_dataframe(input_data['data']) opc_output_config = input_data['opc_output_config'] self.info(f'Data to write: {data.size} rows', metadata) @@ -324,7 +326,7 @@ class OPC(SientiaMonitoring): def process_confidence( self, data: DataFrame, success: bool, metadata: dict[str, Any] - ) -> dict[Any, Any]: + ) -> dict[Hashable, Any]: """ Process prediction confidence based on OPC write operation success. diff --git a/laborious/utils/dataframe_utils.py b/laborious/utils/dataframe_utils.py new file mode 100644 index 0000000..c87764d --- /dev/null +++ b/laborious/utils/dataframe_utils.py @@ -0,0 +1,31 @@ +""" +DataFrame utility functions for handling serialized DataFrames. + +This module provides helper functions to work with DataFrames that may +come from Temporal serialization (already as DataFrame) or from legacy +code (as dict). +""" + +from typing import Any + +from pandas import DataFrame + + +def ensure_dataframe(data: Any) -> DataFrame: + """ + Ensure that data is a DataFrame, converting from dict if necessary. + + This function handles both cases: + - Data already deserialized as DataFrame (from Temporal codec) + - Data as dict (legacy format or non-DataFrame serialization) + + Args: + data: Data that should be a DataFrame (can be DataFrame or dict) + + Returns: + DataFrame: The data as a pandas DataFrame + """ + if isinstance(data, DataFrame): + return data + return DataFrame(data) + diff --git a/laborious/utils/temporal_codec.py b/laborious/utils/temporal_codec.py new file mode 100644 index 0000000..2c6d70a --- /dev/null +++ b/laborious/utils/temporal_codec.py @@ -0,0 +1,149 @@ +""" +Temporal Codec for DataFrame Serialization + +This module provides a custom Temporal DataConverter that automatically +serializes pandas DataFrames to Parquet format and deserializes them back. + +The codec only handles DataFrames, leaving all other types to the default +Temporal serialization mechanism. +""" + +import io +from collections.abc import Sequence +from typing import Any, List, Optional, Type + +from temporalio.api.common.v1 import Payload +from temporalio.converter import DataConverter, PayloadConverter +from pandas import DataFrame, read_parquet + + +class DataFramePayloadConverter(PayloadConverter): + """ + Custom PayloadConverter that serializes pandas DataFrames to Parquet format. + + Only DataFrames are handled by this converter. All other types are passed + to the default Temporal serialization mechanism. + """ + + def __init__(self, default_payload_converter: PayloadConverter): + """ + Initialize the DataFrame payload converter. + + Args: + default_payload_converter: The default Temporal payload converter + to use for non-DataFrame types + """ + self._default = default_payload_converter + + def to_payloads(self, values: Sequence[Any]) -> List[Payload]: + """ + Convert values to Temporal Payloads. + + If a value is a pandas DataFrame, it is serialized to Parquet format. + Otherwise, the default converter is used. + + Args: + values: The values to serialize + + Returns: + List[Payload]: The serialized payloads + """ + payloads = [] + for value in values: + # Check if value is a DataFrame + try: + if isinstance(value, DataFrame): + buffer = io.BytesIO() + value.to_parquet(buffer, engine='pyarrow', index=True) + payloads.append(Payload( + metadata={"encoding": b"parquet-dataframe"}, + data=buffer.getvalue() + )) + continue + except ImportError: + # pandas not available, fall through to default + pass + except Exception: + # Error serializing DataFrame, fall through to default + pass + + # Use default converter for all other types + default_payloads = self._default.to_payloads([value]) + payloads.extend(default_payloads) + + return payloads + + def from_payloads( + self, + payloads: Sequence[Payload], + type_hints: Optional[List[Type]] = None, + ) -> List[Any]: + """ + Convert Temporal Payloads back to Python values. + + If a payload metadata indicates it's a Parquet-serialized DataFrame, + it is deserialized. Otherwise, the default converter is used. + + Args: + payloads: The payloads to deserialize + type_hints: Optional type hints for the expected return types + + Returns: + List[Any]: The deserialized values + """ + values = [] + for i, payload in enumerate(payloads): + # Check if this is a Parquet-serialized DataFrame + if payload.metadata.get("encoding") == b"parquet-dataframe": + try: + buffer = io.BytesIO(payload.data) + values.append(read_parquet(buffer)) + continue + except ImportError: + # pandas not available, fall through to default + pass + except Exception: + # Error deserializing DataFrame, fall through to default + pass + + # Use default converter for all other types + # type_hints must have same length as payloads if provided + type_hint = type_hints[i] if type_hints and i < len(type_hints) else None + default_values = self._default.from_payloads([payload], [type_hint] if type_hint is not None else None) + values.extend(default_values) + + return values + + +def create_dataframe_data_converter() -> DataConverter: + """ + Create a DataConverter with DataFrame serialization support. + + Returns: + DataConverter: A DataConverter that handles DataFrames automatically + """ + # Get default converter to use as fallback + # DataConverter.default is an attribute, not a method + default_converter = DataConverter.default + + # Create a factory class that extends PayloadConverter + class DataFramePayloadConverterFactory(PayloadConverter): + def __init__(self): + super().__init__() + # Create instance of default payload converter to use as fallback + self._default_converter = default_converter.payload_converter_class() + + def to_payloads(self, values: Sequence[Any]) -> List[Payload]: + return DataFramePayloadConverter(self._default_converter).to_payloads(values) + + def from_payloads( + self, + payloads: Sequence[Payload], + type_hints: Optional[List[Type]] = None, + ) -> List[Any]: + return DataFramePayloadConverter(self._default_converter).from_payloads(payloads, type_hints) + + return DataConverter( + payload_converter_class=DataFramePayloadConverterFactory + ) + diff --git a/laborious/worker/worker.py b/laborious/worker/worker.py index 1d28016..4700c48 100644 --- a/laborious/worker/worker.py +++ b/laborious/worker/worker.py @@ -48,6 +48,7 @@ with workflow.unsafe.imports_passed_through(): build_opc_config, build_postgres_config, ) + from laborious.utils.temporal_codec import create_dataframe_data_converter from laborious.workflows.minimal_retrain import MinimalRetrain from laborious.workflows.predictions_batch import PredictionsBatch from laborious.workflows.drift import Drift @@ -132,10 +133,14 @@ async def main(): logger.custom_info(f'Starting Temporal Client at {host}...', metadata) + # Create custom data converter with DataFrame support + data_converter = create_dataframe_data_converter() + temporal_client = await client.Client.connect( target_host=host, namespace=os.getenv('TEMPORAL_NAMESPACE', 'laborious'), runtime=new_runtime, + data_converter=data_converter, ) logger.custom_info('Starting Workers...', metadata) diff --git a/laborious/workflows/drift.py b/laborious/workflows/drift.py index 7afaef7..043ef7e 100644 --- a/laborious/workflows/drift.py +++ b/laborious/workflows/drift.py @@ -93,7 +93,7 @@ class Drift: Activities.export_data_to_postgres, { **metadata, - 'data': drift_data, + 'data': drift_data.to_dict(orient='records'), 'schema': input_data['schema'], 'table_name': input_data['target_table_name'], 'timestamp_conversion': {'column': 'timestamp', 'format': DATETIME_FORMAT_WITH_TZ}, diff --git a/laborious/workflows/minimal_retrain.py b/laborious/workflows/minimal_retrain.py index 02878cf..6ba88f9 100644 --- a/laborious/workflows/minimal_retrain.py +++ b/laborious/workflows/minimal_retrain.py @@ -127,7 +127,7 @@ class MinimalRetrain: Activities.export_data_to_postgres, { **metadata, - 'data': report, + 'data': report.to_dict(orient='records'), 'schema': input_data['schema'], 'table_name': input_data['table_name'], }, diff --git a/laborious/workflows/simple_metrics.py b/laborious/workflows/simple_metrics.py index f5f02d2..f001962 100644 --- a/laborious/workflows/simple_metrics.py +++ b/laborious/workflows/simple_metrics.py @@ -82,7 +82,7 @@ class SimpleMetrics: Activities.export_data_to_postgres, { **metadata, - 'data': simple_metrics, + 'data': simple_metrics.to_dict(orient='records'), 'schema': input_data['schema'], 'table_name': input_data['target_table_name'], 'timestamp_conversion': { diff --git a/laborious/workflows/sub_workflows/format_and_export_prediction.py b/laborious/workflows/sub_workflows/format_and_export_prediction.py index 4280ccf..1d59e7a 100644 --- a/laborious/workflows/sub_workflows/format_and_export_prediction.py +++ b/laborious/workflows/sub_workflows/format_and_export_prediction.py @@ -107,7 +107,7 @@ class FormatAndExportPrediction: **metadata, 'schema': input_data['schema'], 'table_name': input_data['transform_table_name'], - 'data': transformed, + 'data': transformed.to_dict(orient='records'), 'timestamp_conversion': { 'column': 'timestamp', 'format': DATETIME_FORMAT_WITH_TZ, @@ -156,7 +156,7 @@ class FormatAndExportPrediction: **metadata, 'schema': input_data['schema'], 'table_name': input_data['table_name'], - 'data': prediction, + 'data': prediction.to_dict(orient='records'), 'timestamp_conversion': {'column': 'timestamp', 'format': DATETIME_FORMAT_WITH_TZ}, }, retry_policy=retry_policy, diff --git a/tests.ipynb b/tests.ipynb index 808277c..25f1cf0 100644 --- a/tests.ipynb +++ b/tests.ipynb @@ -588,6 +588,157 @@ "\n", "display(prediction_data)" ] + }, + { + "cell_type": "code", + "execution_count": 10, + "id": "486b95b3", + "metadata": {}, + "outputs": [ + { + "data": { + "text/html": [ + "
\n", + "\n", + "\n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + "
ab
2025-01-0114
2025-01-0225
2025-01-0336
\n", + "
" + ], + "text/plain": [ + " a b\n", + "2025-01-01 1 4\n", + "2025-01-02 2 5\n", + "2025-01-03 3 6" + ] + }, + "metadata": {}, + "output_type": "display_data" + }, + { + "data": { + "text/plain": [ + "{'index': ['2025-01-01', '2025-01-02', '2025-01-03'],\n", + " 'columns': ['a', 'b'],\n", + " 'data': [[1, 4], [2, 5], [3, 6]]}" + ] + }, + "metadata": {}, + "output_type": "display_data" + }, + { + "data": { + "text/html": [ + "
\n", + "\n", + "\n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + "
012
index2025-01-012025-01-022025-01-03
columnsabNone
data[1, 4][2, 5][3, 6]
\n", + "
" + ], + "text/plain": [ + " 0 1 2\n", + "index 2025-01-01 2025-01-02 2025-01-03\n", + "columns a b None\n", + "data [1, 4] [2, 5] [3, 6]" + ] + }, + "metadata": {}, + "output_type": "display_data" + } + ], + "source": [ + "from pandas import DataFrame\n", + "\n", + "data = DataFrame({\n", + " \"a\": {\"2025-01-01\": 1, \"2025-01-02\": 2, \"2025-01-03\": 3},\n", + " \"b\": {\"2025-01-01\": 4, \"2025-01-02\": 5, \"2025-01-03\": 6},\n", + "})\n", + "\n", + "display(data)\n", + "\n", + "data_list = data.to_dict('split')\n", + "\n", + "display(data_list)\n", + "\n", + "data_rec = DataFrame.from_dict(data_list, orient='index')\n", + "\n", + "display(data_rec)" + ] } ], "metadata": { From c8b809f18994cccaa5a7f6c9461b639abd51eff2 Mon Sep 17 00:00:00 2001 From: vitor-aignosi Date: Mon, 17 Nov 2025 09:58:55 -0300 Subject: [PATCH 07/27] SIENTIAPDE-1273 Update dependencies and refactor data handling in various modules - Updated sientia-dataops-library dependency version from 1.5.3 to 1.5.4 in requirements files. - Updated sientia-mlops-library dependency version from 0.39.0 to 0.40.2 in requirements files. - Refactored return types in Gates, MLFlow, and ModelMetrics classes to return dictionaries instead of DataFrames for improved compatibility with downstream systems. - Removed the temporal_codec module as it is no longer needed for DataFrame serialization. - Adjusted data handling in the Drift workflow to ensure proper data structure is maintained. --- laborious/activities/gates.py | 16 +-- laborious/activities/mlflow.py | 6 +- laborious/activities/model_metrics.py | 16 +-- laborious/utils/temporal_codec.py | 149 -------------------------- laborious/worker/worker.py | 5 - laborious/workflows/drift.py | 4 +- requirements-light.txt | 2 +- requirements.txt | 4 +- tests.ipynb | 15 ++- 9 files changed, 40 insertions(+), 177 deletions(-) delete mode 100644 laborious/utils/temporal_codec.py diff --git a/laborious/activities/gates.py b/laborious/activities/gates.py index 7c8ae76..2ac33f1 100644 --- a/laborious/activities/gates.py +++ b/laborious/activities/gates.py @@ -405,7 +405,7 @@ class Gates(SientiaMonitoring): @activity.defn(name='format_transformed_data') - async def format_transformed_data(self, input_data: dict[str, Any]) -> DataFrame: + async def format_transformed_data(self, input_data: dict[str, Any]) -> dict[str, Any]: """ Format transformed data according to configured storage policies. """ @@ -423,10 +423,10 @@ class Gates(SientiaMonitoring): data = data.melt(id_vars='timestamp', var_name='variable', value_name='value') data['model_id'] = model_id - return data + return data.to_dict() @activity.defn(name='format_prediction') - async def format_prediction(self, input_data: dict[str, Any]) -> DataFrame: + async def format_prediction(self, input_data: dict[str, Any]) -> dict[str, Any]: """ Format prediction data according to configured storage policies. @@ -496,10 +496,10 @@ class Gates(SientiaMonitoring): self.info(f'Prediction formatted: {len(data)} rows', metadata) self.debug(f'Prediction data: {data.head(5).to_string()}', metadata) - return data + return data.to_dict() @activity.defn(name='format_default_prediction') - async def format_default_prediction(self, input_data: dict[str, Any]) -> DataFrame: + async def format_default_prediction(self, input_data: dict[str, Any]) -> dict[str, Any]: """ Create and format default prediction data for error conditions. @@ -541,10 +541,10 @@ class Gates(SientiaMonitoring): ) self.info(f'Default prediction formatted: {data.size} rows', metadata) - return data + return data.to_dict() @activity.defn(name='format_retrain_report') - async def format_retrain_report(self, input_data: dict[str, Any]) -> DataFrame: + async def format_retrain_report(self, input_data: dict[str, Any]) -> dict[str, Any]: """ Format retrain report data according to configured storage policies. """ @@ -573,7 +573,7 @@ class Gates(SientiaMonitoring): self.debug(f'Retrain report: {report.to_csv()}', metadata) - return report + return report.to_dict() @activity.defn(name='get_last_timestamp') async def get_last_timestamp(self, input_data: dict[str, Any]) -> str: diff --git a/laborious/activities/mlflow.py b/laborious/activities/mlflow.py index 0447215..129fb75 100644 --- a/laborious/activities/mlflow.py +++ b/laborious/activities/mlflow.py @@ -423,7 +423,7 @@ class MLFlow(SientiaMonitoring): @activity.defn(name='get_reference_data') - async def get_reference_data(self, input_data: dict[str, Any]) -> DataFrame | None: + async def get_reference_data(self, input_data: dict[str, Any]) -> list[dict[Hashable, Any]] | None: """ Get reference data from the MLflow Model Registry. @@ -433,7 +433,7 @@ class MLFlow(SientiaMonitoring): - model_name (str): Name of the MLFlow model to get reference data from Returns: - dict[Hashable, Any] | None: Reference data from the MLflow Model Registry. + list[dict[Hashable, Any]] | None: Reference data from the MLflow Model Registry. """ @@ -452,4 +452,4 @@ class MLFlow(SientiaMonitoring): reference_data['timestamp'] = to_datetime(reference_data['timestamp']) reference_data['timestamp'] = reference_data['timestamp'].dt.strftime(DATETIME_FORMAT) - return reference_data \ No newline at end of file + return reference_data.to_dict(orient='records') \ No newline at end of file diff --git a/laborious/activities/model_metrics.py b/laborious/activities/model_metrics.py index 9896fa7..78e7aad 100644 --- a/laborious/activities/model_metrics.py +++ b/laborious/activities/model_metrics.py @@ -138,7 +138,7 @@ class ModelMetrics(SientiaMonitoring): @activity.defn(name='calculate_drift') - async def calculate_drift(self, input_data: dict[str, Any]) -> DataFrame | dict: + async def calculate_drift(self, input_data: dict[str, Any]) -> list[dict[Hashable, Any]]: """ Calculate drift metrics for a model. @@ -217,11 +217,11 @@ class ModelMetrics(SientiaMonitoring): level=NotificationLevel.ERROR, attachment_content=traceback.format_exc(), ) - return {} + return [] if drift_df.empty: self.warning('No drift metrics found', metadata) - return {} + return [] # Drop unnecessary columns drift_df.drop(columns=['p_value', 'chunk_start_date', 'chunk_end_date'], inplace=True) @@ -238,7 +238,7 @@ class ModelMetrics(SientiaMonitoring): if drift_df.empty: self.warning('No drift metrics found after dropping rows where timestamp is not in target data', metadata) - return {} + return [] # Rename columns to match database columns drift_df.rename(columns={ @@ -261,10 +261,12 @@ class ModelMetrics(SientiaMonitoring): self.debug(f'Drift dataframe: Size {drift_df.shape} \n{drift_df.head(5).to_string()}', metadata) + self.debug(f'Drift dataframe: {drift_df.head(5).to_string()}', metadata) - return drift_df + return drift_df.to_dict(orient='records') - async def calculate_simple_metrics(self, input_data: dict[str, Any]) -> DataFrame: + @activity.defn(name='calculate_simple_metrics') + async def calculate_simple_metrics(self, input_data: dict[str, Any]) -> list[dict[Hashable, Any]]: """ Calculate simple metrics for a model. Metrics available are: - rmse @@ -342,6 +344,6 @@ class ModelMetrics(SientiaMonitoring): self.debug(f'Simple metrics dataframe: Size {data.shape} \n{data.head(5).to_string()}', metadata) - return data + return data.to_dict(orient='records') diff --git a/laborious/utils/temporal_codec.py b/laborious/utils/temporal_codec.py deleted file mode 100644 index 2c6d70a..0000000 --- a/laborious/utils/temporal_codec.py +++ /dev/null @@ -1,149 +0,0 @@ -""" -Temporal Codec for DataFrame Serialization - -This module provides a custom Temporal DataConverter that automatically -serializes pandas DataFrames to Parquet format and deserializes them back. - -The codec only handles DataFrames, leaving all other types to the default -Temporal serialization mechanism. -""" - -import io -from collections.abc import Sequence -from typing import Any, List, Optional, Type - -from temporalio.api.common.v1 import Payload -from temporalio.converter import DataConverter, PayloadConverter -from pandas import DataFrame, read_parquet - - -class DataFramePayloadConverter(PayloadConverter): - """ - Custom PayloadConverter that serializes pandas DataFrames to Parquet format. - - Only DataFrames are handled by this converter. All other types are passed - to the default Temporal serialization mechanism. - """ - - def __init__(self, default_payload_converter: PayloadConverter): - """ - Initialize the DataFrame payload converter. - - Args: - default_payload_converter: The default Temporal payload converter - to use for non-DataFrame types - """ - self._default = default_payload_converter - - def to_payloads(self, values: Sequence[Any]) -> List[Payload]: - """ - Convert values to Temporal Payloads. - - If a value is a pandas DataFrame, it is serialized to Parquet format. - Otherwise, the default converter is used. - - Args: - values: The values to serialize - - Returns: - List[Payload]: The serialized payloads - """ - payloads = [] - for value in values: - # Check if value is a DataFrame - try: - if isinstance(value, DataFrame): - buffer = io.BytesIO() - value.to_parquet(buffer, engine='pyarrow', index=True) - payloads.append(Payload( - metadata={"encoding": b"parquet-dataframe"}, - data=buffer.getvalue() - )) - continue - except ImportError: - # pandas not available, fall through to default - pass - except Exception: - # Error serializing DataFrame, fall through to default - pass - - # Use default converter for all other types - default_payloads = self._default.to_payloads([value]) - payloads.extend(default_payloads) - - return payloads - - def from_payloads( - self, - payloads: Sequence[Payload], - type_hints: Optional[List[Type]] = None, - ) -> List[Any]: - """ - Convert Temporal Payloads back to Python values. - - If a payload metadata indicates it's a Parquet-serialized DataFrame, - it is deserialized. Otherwise, the default converter is used. - - Args: - payloads: The payloads to deserialize - type_hints: Optional type hints for the expected return types - - Returns: - List[Any]: The deserialized values - """ - values = [] - for i, payload in enumerate(payloads): - # Check if this is a Parquet-serialized DataFrame - if payload.metadata.get("encoding") == b"parquet-dataframe": - try: - buffer = io.BytesIO(payload.data) - values.append(read_parquet(buffer)) - continue - except ImportError: - # pandas not available, fall through to default - pass - except Exception: - # Error deserializing DataFrame, fall through to default - pass - - # Use default converter for all other types - # type_hints must have same length as payloads if provided - type_hint = type_hints[i] if type_hints and i < len(type_hints) else None - default_values = self._default.from_payloads([payload], [type_hint] if type_hint is not None else None) - values.extend(default_values) - - return values - - -def create_dataframe_data_converter() -> DataConverter: - """ - Create a DataConverter with DataFrame serialization support. - - Returns: - DataConverter: A DataConverter that handles DataFrames automatically - """ - # Get default converter to use as fallback - # DataConverter.default is an attribute, not a method - default_converter = DataConverter.default - - # Create a factory class that extends PayloadConverter - class DataFramePayloadConverterFactory(PayloadConverter): - def __init__(self): - super().__init__() - # Create instance of default payload converter to use as fallback - self._default_converter = default_converter.payload_converter_class() - - def to_payloads(self, values: Sequence[Any]) -> List[Payload]: - return DataFramePayloadConverter(self._default_converter).to_payloads(values) - - def from_payloads( - self, - payloads: Sequence[Payload], - type_hints: Optional[List[Type]] = None, - ) -> List[Any]: - return DataFramePayloadConverter(self._default_converter).from_payloads(payloads, type_hints) - - return DataConverter( - payload_converter_class=DataFramePayloadConverterFactory - ) - diff --git a/laborious/worker/worker.py b/laborious/worker/worker.py index 4700c48..1d28016 100644 --- a/laborious/worker/worker.py +++ b/laborious/worker/worker.py @@ -48,7 +48,6 @@ with workflow.unsafe.imports_passed_through(): build_opc_config, build_postgres_config, ) - from laborious.utils.temporal_codec import create_dataframe_data_converter from laborious.workflows.minimal_retrain import MinimalRetrain from laborious.workflows.predictions_batch import PredictionsBatch from laborious.workflows.drift import Drift @@ -133,14 +132,10 @@ async def main(): logger.custom_info(f'Starting Temporal Client at {host}...', metadata) - # Create custom data converter with DataFrame support - data_converter = create_dataframe_data_converter() - temporal_client = await client.Client.connect( target_host=host, namespace=os.getenv('TEMPORAL_NAMESPACE', 'laborious'), runtime=new_runtime, - data_converter=data_converter, ) logger.custom_info('Starting Workers...', metadata) diff --git a/laborious/workflows/drift.py b/laborious/workflows/drift.py index 043ef7e..60bb274 100644 --- a/laborious/workflows/drift.py +++ b/laborious/workflows/drift.py @@ -32,6 +32,8 @@ class Drift: } } + print(f'Input data: {input_data}', metadata) + model_config = input_data['model_config'] target_name = model_config['target'] @@ -93,7 +95,7 @@ class Drift: Activities.export_data_to_postgres, { **metadata, - 'data': drift_data.to_dict(orient='records'), + 'data': drift_data, 'schema': input_data['schema'], 'table_name': input_data['target_table_name'], 'timestamp_conversion': {'column': 'timestamp', 'format': DATETIME_FORMAT_WITH_TZ}, diff --git a/requirements-light.txt b/requirements-light.txt index e41ba53..26f2290 100644 --- a/requirements-light.txt +++ b/requirements-light.txt @@ -3,7 +3,7 @@ psycopg2-binary sqlalchemy asyncua redis -git+ssh://git@github.com/Aignosi/sientia-dataops-library.git@1.5.2 +git+ssh://git@github.com/Aignosi/sientia-dataops-library.git@1.5.4 prometheus-client botocore boto3 diff --git a/requirements.txt b/requirements.txt index 750ef66..214e910 100644 --- a/requirements.txt +++ b/requirements.txt @@ -3,8 +3,8 @@ psycopg2-binary sqlalchemy asyncua redis -git+ssh://git@github.com/Aignosi/sientia-dataops-library.git@1.5.3 -git+ssh://git@github.com/Aignosi/sientia-mlops-library.git@0.39.0 +git+ssh://git@github.com/Aignosi/sientia-dataops-library.git@1.5.4 +git+ssh://git@github.com/Aignosi/sientia-mlops-library.git@0.40.2 prometheus-client botocore boto3 diff --git a/tests.ipynb b/tests.ipynb index 25f1cf0..6933169 100644 --- a/tests.ipynb +++ b/tests.ipynb @@ -591,10 +591,19 @@ }, { "cell_type": "code", - "execution_count": 10, + "execution_count": 1, "id": "486b95b3", "metadata": {}, "outputs": [ + { + "data": { + "text/plain": [ + "[]" + ] + }, + "metadata": {}, + "output_type": "display_data" + }, { "data": { "text/html": [ @@ -724,6 +733,10 @@ "source": [ "from pandas import DataFrame\n", "\n", + "data = DataFrame()\n", + "\n", + "display(data.to_dict(orient='records'))\n", + "\n", "data = DataFrame({\n", " \"a\": {\"2025-01-01\": 1, \"2025-01-02\": 2, \"2025-01-03\": 3},\n", " \"b\": {\"2025-01-01\": 4, \"2025-01-02\": 5, \"2025-01-03\": 6},\n", From a7ab8ffe4916b6dadb512a0eb65d7de8b87225ad Mon Sep 17 00:00:00 2001 From: vitor-aignosi Date: Mon, 17 Nov 2025 10:05:57 -0300 Subject: [PATCH 08/27] SIENTIAPDE-1273 Refactor return type of calculate_drift method in ModelMetrics class to improve compatibility. Changed from a list of dictionaries to a generic list for enhanced flexibility in data handling. --- laborious/activities/model_metrics.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/laborious/activities/model_metrics.py b/laborious/activities/model_metrics.py index 78e7aad..a3deda8 100644 --- a/laborious/activities/model_metrics.py +++ b/laborious/activities/model_metrics.py @@ -138,7 +138,7 @@ class ModelMetrics(SientiaMonitoring): @activity.defn(name='calculate_drift') - async def calculate_drift(self, input_data: dict[str, Any]) -> list[dict[Hashable, Any]]: + async def calculate_drift(self, input_data: dict[str, Any]) -> list: """ Calculate drift metrics for a model. From 46fe88f69e27ccd24c042bf2f77f0dea9492cfe9 Mon Sep 17 00:00:00 2001 From: vitor-aignosi Date: Mon, 17 Nov 2025 10:14:47 -0300 Subject: [PATCH 09/27] SIENTIAPDE-1273 Refactor return types in ModelMetrics class methods to enhance type clarity. Updated calculate_drift and calculate_simple_metrics methods to return lists of dictionaries instead of generic lists, improving type specificity for better data handling. --- laborious/activities/model_metrics.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/laborious/activities/model_metrics.py b/laborious/activities/model_metrics.py index a3deda8..dce8523 100644 --- a/laborious/activities/model_metrics.py +++ b/laborious/activities/model_metrics.py @@ -1,7 +1,7 @@ from temporalio import activity, workflow with workflow.unsafe.imports_passed_through(): - from typing import Any, Hashable + from typing import Any from pandas import DataFrame, Index, to_datetime @@ -138,7 +138,7 @@ class ModelMetrics(SientiaMonitoring): @activity.defn(name='calculate_drift') - async def calculate_drift(self, input_data: dict[str, Any]) -> list: + async def calculate_drift(self, input_data: dict[str, Any]) -> list[dict]: """ Calculate drift metrics for a model. @@ -266,7 +266,7 @@ class ModelMetrics(SientiaMonitoring): return drift_df.to_dict(orient='records') @activity.defn(name='calculate_simple_metrics') - async def calculate_simple_metrics(self, input_data: dict[str, Any]) -> list[dict[Hashable, Any]]: + async def calculate_simple_metrics(self, input_data: dict[str, Any]) -> list[dict]: """ Calculate simple metrics for a model. Metrics available are: - rmse From 91580f82551237c5ddf7398744de6b9b187b7b8c Mon Sep 17 00:00:00 2001 From: vitor-aignosi Date: Mon, 17 Nov 2025 10:25:30 -0300 Subject: [PATCH 10/27] SIENTIAPDE-1273 Refactor return type of get_reference_data method in MLFlow class to improve type specificity. Changed from a list of dictionaries with Hashable keys to a generic list of dictionaries for enhanced clarity in data handling. --- laborious/activities/mlflow.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/laborious/activities/mlflow.py b/laborious/activities/mlflow.py index 129fb75..4d542c5 100644 --- a/laborious/activities/mlflow.py +++ b/laborious/activities/mlflow.py @@ -423,7 +423,7 @@ class MLFlow(SientiaMonitoring): @activity.defn(name='get_reference_data') - async def get_reference_data(self, input_data: dict[str, Any]) -> list[dict[Hashable, Any]] | None: + async def get_reference_data(self, input_data: dict[str, Any]) -> list[dict] | None: """ Get reference data from the MLflow Model Registry. From 08f9522695ecb4859b8da5ed3ca2a80451476cd7 Mon Sep 17 00:00:00 2001 From: vitor-aignosi Date: Mon, 17 Nov 2025 10:36:17 -0300 Subject: [PATCH 11/27] SIENTIAPDE-1273 Refactor data export in MinimalRetrain workflow to improve structure. Changed the export format from a list of dictionaries to a direct report object for enhanced clarity and compatibility with downstream systems. --- laborious/workflows/minimal_retrain.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/laborious/workflows/minimal_retrain.py b/laborious/workflows/minimal_retrain.py index 6ba88f9..02878cf 100644 --- a/laborious/workflows/minimal_retrain.py +++ b/laborious/workflows/minimal_retrain.py @@ -127,7 +127,7 @@ class MinimalRetrain: Activities.export_data_to_postgres, { **metadata, - 'data': report.to_dict(orient='records'), + 'data': report, 'schema': input_data['schema'], 'table_name': input_data['table_name'], }, From c9b1ddfd1514adc3f6c9a6ee34aac2be9c50bbf5 Mon Sep 17 00:00:00 2001 From: vitor-aignosi Date: Mon, 17 Nov 2025 10:38:44 -0300 Subject: [PATCH 12/27] SIENTIAPDE-1273 Refactor data export in Drift, SimpleMetrics, and FormatAndExportPrediction workflows to improve data handling. Changed the export format from a list of dictionaries to direct objects for enhanced clarity and compatibility with downstream systems. --- laborious/workflows/drift.py | 1 + laborious/workflows/simple_metrics.py | 3 ++- .../workflows/sub_workflows/format_and_export_prediction.py | 4 ++-- 3 files changed, 5 insertions(+), 3 deletions(-) diff --git a/laborious/workflows/drift.py b/laborious/workflows/drift.py index 60bb274..d556dac 100644 --- a/laborious/workflows/drift.py +++ b/laborious/workflows/drift.py @@ -52,6 +52,7 @@ class Drift: **metadata, 'query': gathering_query, 'datetime_columns': ['timestamp', 'created_at'], + 'orient': 'records', }, retry_policy=retry_policy, start_to_close_timeout=timedelta(seconds=300), diff --git a/laborious/workflows/simple_metrics.py b/laborious/workflows/simple_metrics.py index f001962..ce950af 100644 --- a/laborious/workflows/simple_metrics.py +++ b/laborious/workflows/simple_metrics.py @@ -53,6 +53,7 @@ class SimpleMetrics: **metadata, 'query': query, 'datetime_columns': ['timestamp'], + 'orient': 'records', }, retry_policy=retry_policy, start_to_close_timeout=timedelta(seconds=300), @@ -82,7 +83,7 @@ class SimpleMetrics: Activities.export_data_to_postgres, { **metadata, - 'data': simple_metrics.to_dict(orient='records'), + 'data': simple_metrics, 'schema': input_data['schema'], 'table_name': input_data['target_table_name'], 'timestamp_conversion': { diff --git a/laborious/workflows/sub_workflows/format_and_export_prediction.py b/laborious/workflows/sub_workflows/format_and_export_prediction.py index 1d59e7a..4280ccf 100644 --- a/laborious/workflows/sub_workflows/format_and_export_prediction.py +++ b/laborious/workflows/sub_workflows/format_and_export_prediction.py @@ -107,7 +107,7 @@ class FormatAndExportPrediction: **metadata, 'schema': input_data['schema'], 'table_name': input_data['transform_table_name'], - 'data': transformed.to_dict(orient='records'), + 'data': transformed, 'timestamp_conversion': { 'column': 'timestamp', 'format': DATETIME_FORMAT_WITH_TZ, @@ -156,7 +156,7 @@ class FormatAndExportPrediction: **metadata, 'schema': input_data['schema'], 'table_name': input_data['table_name'], - 'data': prediction.to_dict(orient='records'), + 'data': prediction, 'timestamp_conversion': {'column': 'timestamp', 'format': DATETIME_FORMAT_WITH_TZ}, }, retry_policy=retry_policy, From 4563ffecac8d6d5eac9432e39644a46388c2a08c Mon Sep 17 00:00:00 2001 From: vitor-aignosi Date: Mon, 17 Nov 2025 10:45:13 -0300 Subject: [PATCH 13/27] SIENTIAPDE-1273 Update sientia-dataops-library dependency version to 1.6.1 and refactor return types in Gates class methods for improved type clarity. Changed return types from dict[str, Any] to dict for better compatibility with downstream systems. --- laborious/activities/gates.py | 8 ++++---- requirements-light.txt | 2 +- requirements.txt | 2 +- 3 files changed, 6 insertions(+), 6 deletions(-) diff --git a/laborious/activities/gates.py b/laborious/activities/gates.py index 2ac33f1..f3f7ad9 100644 --- a/laborious/activities/gates.py +++ b/laborious/activities/gates.py @@ -405,7 +405,7 @@ class Gates(SientiaMonitoring): @activity.defn(name='format_transformed_data') - async def format_transformed_data(self, input_data: dict[str, Any]) -> dict[str, Any]: + async def format_transformed_data(self, input_data: dict[str, Any]) -> dict: """ Format transformed data according to configured storage policies. """ @@ -426,7 +426,7 @@ class Gates(SientiaMonitoring): return data.to_dict() @activity.defn(name='format_prediction') - async def format_prediction(self, input_data: dict[str, Any]) -> dict[str, Any]: + async def format_prediction(self, input_data: dict[str, Any]) -> dict: """ Format prediction data according to configured storage policies. @@ -499,7 +499,7 @@ class Gates(SientiaMonitoring): return data.to_dict() @activity.defn(name='format_default_prediction') - async def format_default_prediction(self, input_data: dict[str, Any]) -> dict[str, Any]: + async def format_default_prediction(self, input_data: dict[str, Any]) -> dict: """ Create and format default prediction data for error conditions. @@ -544,7 +544,7 @@ class Gates(SientiaMonitoring): return data.to_dict() @activity.defn(name='format_retrain_report') - async def format_retrain_report(self, input_data: dict[str, Any]) -> dict[str, Any]: + async def format_retrain_report(self, input_data: dict[str, Any]) -> dict: """ Format retrain report data according to configured storage policies. """ diff --git a/requirements-light.txt b/requirements-light.txt index 26f2290..8831822 100644 --- a/requirements-light.txt +++ b/requirements-light.txt @@ -3,7 +3,7 @@ psycopg2-binary sqlalchemy asyncua redis -git+ssh://git@github.com/Aignosi/sientia-dataops-library.git@1.5.4 +git+ssh://git@github.com/Aignosi/sientia-dataops-library.git@1.6.1 prometheus-client botocore boto3 diff --git a/requirements.txt b/requirements.txt index 214e910..d475d8a 100644 --- a/requirements.txt +++ b/requirements.txt @@ -3,7 +3,7 @@ psycopg2-binary sqlalchemy asyncua redis -git+ssh://git@github.com/Aignosi/sientia-dataops-library.git@1.5.4 +git+ssh://git@github.com/Aignosi/sientia-dataops-library.git@1.6.1 git+ssh://git@github.com/Aignosi/sientia-mlops-library.git@0.40.2 prometheus-client botocore From 019b46009a0315cc1fa3f6e68ce95d3dee59556a Mon Sep 17 00:00:00 2001 From: vitor-aignosi Date: Mon, 17 Nov 2025 10:56:58 -0300 Subject: [PATCH 14/27] SIENTIAPDE-1273 Refactor data size calculation in ModelMetrics class to use shape method for improved accuracy. Changed from using len(target_data) to data.shape[0] for better clarity in data handling. --- laborious/activities/model_metrics.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/laborious/activities/model_metrics.py b/laborious/activities/model_metrics.py index dce8523..3d55997 100644 --- a/laborious/activities/model_metrics.py +++ b/laborious/activities/model_metrics.py @@ -293,7 +293,7 @@ class ModelMetrics(SientiaMonitoring): metrics = input_data['metrics'] interval_minutes = input_data['interval_minutes'] - data_size = len(target_data) + data_size = data.shape[0] output_data = [] From 59d3c26c71f51bd71c1f49f0ce808f02a7eac4db Mon Sep 17 00:00:00 2001 From: vitor-aignosi Date: Mon, 17 Nov 2025 10:57:55 -0300 Subject: [PATCH 15/27] SIENTIAPDE-1273 Update execution count in tests notebook and add data shape output for improved clarity in data handling. --- tests.ipynb | 16 ++++++++++++++-- 1 file changed, 14 insertions(+), 2 deletions(-) diff --git a/tests.ipynb b/tests.ipynb index 6933169..b8c1b5c 100644 --- a/tests.ipynb +++ b/tests.ipynb @@ -591,7 +591,7 @@ }, { "cell_type": "code", - "execution_count": 1, + "execution_count": 4, "id": "486b95b3", "metadata": {}, "outputs": [ @@ -728,6 +728,16 @@ }, "metadata": {}, "output_type": "display_data" + }, + { + "data": { + "text/plain": [ + "3" + ] + }, + "execution_count": 4, + "metadata": {}, + "output_type": "execute_result" } ], "source": [ @@ -750,7 +760,9 @@ "\n", "data_rec = DataFrame.from_dict(data_list, orient='index')\n", "\n", - "display(data_rec)" + "display(data_rec)\n", + "\n", + "data.shape[0]" ] } ], From 1014c33dd9f889b4dbb268f8308cf5584020ef40 Mon Sep 17 00:00:00 2001 From: vitor-aignosi Date: Mon, 17 Nov 2025 11:03:27 -0300 Subject: [PATCH 16/27] SIENTIAPDE-1273 Fix data size calculation in ModelMetrics class to use target_data.shape[0] for improved accuracy in metrics processing. --- laborious/activities/model_metrics.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/laborious/activities/model_metrics.py b/laborious/activities/model_metrics.py index 3d55997..ae0f231 100644 --- a/laborious/activities/model_metrics.py +++ b/laborious/activities/model_metrics.py @@ -293,7 +293,7 @@ class ModelMetrics(SientiaMonitoring): metrics = input_data['metrics'] interval_minutes = input_data['interval_minutes'] - data_size = data.shape[0] + data_size = target_data.shape[0] output_data = [] From a88a15c60aacd8bc68740336b6238f9e4e354799 Mon Sep 17 00:00:00 2001 From: vitor-aignosi Date: Mon, 17 Nov 2025 16:04:54 -0300 Subject: [PATCH 17/27] SIENTIAPDE-1273 SIENTIAPDE-1273 Enhance security analysis and SQL injection handling - Added skip for potential SQL injection false positives in Bandit configuration. - Updated validate.sh to use the pyproject.toml configuration for Bandit security analysis. - Refactored code to replace ensure_dataframe utility with direct DataFrame usage in multiple activities, improving clarity and reducing dependencies. - Removed the deprecated dataframe_utils module to streamline the codebase. --- laborious/activities/activities.py | 5 +- laborious/activities/gates.py | 12 +- laborious/activities/mlflow.py | 11 +- laborious/activities/model_metrics.py | 139 ++--- laborious/activities/opc.py | 6 +- laborious/metrics.py | 2 +- laborious/utils/dataframe_utils.py | 31 - .../utils/repository/model_repository.py | 37 +- laborious/worker/worker.py | 4 +- laborious/workflows/drift.py | 25 +- laborious/workflows/simple_metrics.py | 17 +- .../format_and_export_prediction.py | 1 - pyproject.toml | 3 +- .../activities/test_model_metrics.py | 538 +++++++++++------- .../utils/repository/test_model_repository.py | 82 ++- tests/laborious/workflows/test_drift.py | 15 +- .../workflows/test_simple_metrics.py | 25 +- validate.sh | 2 +- 18 files changed, 516 insertions(+), 439 deletions(-) delete mode 100644 laborious/utils/dataframe_utils.py diff --git a/laborious/activities/activities.py b/laborious/activities/activities.py index d9c07ac..077a433 100644 --- a/laborious/activities/activities.py +++ b/laborious/activities/activities.py @@ -9,9 +9,10 @@ with workflow.unsafe.imports_passed_through(): from laborious.activities.gates import Gates from laborious.activities.mlflow import MLFlow + from laborious.activities.model_metrics import ModelMetrics from laborious.activities.opc import OPC from laborious.activities.storage import Storage - from laborious.activities.model_metrics import ModelMetrics + class Activities(Storage, MLFlow, Gates, OPC, ModelMetrics): """ @@ -131,4 +132,4 @@ class Activities(Storage, MLFlow, Gates, OPC, ModelMetrics): MLFlow.close(self) Gates.close(self) await OPC.close(self) - ModelMetrics.close(self) \ No newline at end of file + ModelMetrics.close(self) diff --git a/laborious/activities/gates.py b/laborious/activities/gates.py index f3f7ad9..13072cf 100644 --- a/laborious/activities/gates.py +++ b/laborious/activities/gates.py @@ -15,7 +15,6 @@ with workflow.unsafe.imports_passed_through(): from sientia_do.temporal.constants import DATETIME_FORMAT_WITH_TZ, now from laborious import metrics - from laborious.utils.dataframe_utils import ensure_dataframe from laborious.utils.filters.conditional_filters import ( filter_empty_data, filter_specific_variables_null_values, @@ -150,7 +149,7 @@ class Gates(SientiaMonitoring): self.info('Performing input gate...', metadata) filters = input_data['filters'] - data = ensure_dataframe(input_data['data']) + data = DataFrame(input_data['data']) path_priority = input_data['path_priority'] filter_output = [] @@ -403,19 +402,18 @@ class Gates(SientiaMonitoring): return policy_type, int(policy_value) - @activity.defn(name='format_transformed_data') async def format_transformed_data(self, input_data: dict[str, Any]) -> dict: """ Format transformed data according to configured storage policies. """ metadata = input_data['metadata'] - + model_id = input_data['model_id'] self.info('Formatting transformed data...', metadata) - data = ensure_dataframe(input_data['data']) + data = DataFrame(input_data['data']) data['timestamp'] = data.index data = data.reset_index(drop=True) @@ -454,7 +452,7 @@ class Gates(SientiaMonitoring): prediction_store_policy = input_data['prediction_store_policy'] self.info('Formatting prediction...', metadata) - data = ensure_dataframe(input_data['data']) + data = DataFrame(input_data['data']) # Create timestamp column from index and reset index data['timestamp'] = data.index @@ -602,7 +600,7 @@ class Gates(SientiaMonitoring): self.info('Getting last timestamp...', metadata) - data = ensure_dataframe(input_data['data']) + data = DataFrame(input_data['data']) self.debug(f'Input data: {data.head(5).to_string()}', metadata) diff --git a/laborious/activities/mlflow.py b/laborious/activities/mlflow.py index 4d542c5..6eadf60 100644 --- a/laborious/activities/mlflow.py +++ b/laborious/activities/mlflow.py @@ -1,4 +1,3 @@ -from typing import Hashable from temporalio import activity, workflow with workflow.unsafe.imports_passed_through(): @@ -20,7 +19,6 @@ with workflow.unsafe.imports_passed_through(): now, ) - from laborious.utils.dataframe_utils import ensure_dataframe from laborious.utils.repository.minio_repository import MinioRepository from laborious.utils.repository.model_repository import MLFlowRepository @@ -140,7 +138,7 @@ class MLFlow(SientiaMonitoring): """ metadata = input_data['metadata'] self.info('Transforming data...', metadata) - data = ensure_dataframe(input_data['data']) + data = DataFrame(input_data['data']) model_name = input_data['model_name'] model_config = input_data.get('model_config', {}) @@ -212,7 +210,7 @@ class MLFlow(SientiaMonitoring): """ metadata = input_data['metadata'] self.info('Predicting data...', metadata) - data = ensure_dataframe(input_data['data']) + data = DataFrame(input_data['data']) model_name = input_data['model_name'] model_config = input_data.get('model_config', {}) @@ -421,7 +419,6 @@ class MLFlow(SientiaMonitoring): self.error(trace, metadata=metadata) raise e - @activity.defn(name='get_reference_data') async def get_reference_data(self, input_data: dict[str, Any]) -> list[dict] | None: """ @@ -439,7 +436,7 @@ class MLFlow(SientiaMonitoring): metadata = input_data['metadata'] model_name = input_data['model_name'] - artifact = "evaluation_data.csv" + artifact = 'evaluation_data.csv' reference_data = await self.model_monitoring_repository.load_artifact_dataframe( model_name=model_name, artifact_path=artifact, metadata=metadata @@ -452,4 +449,4 @@ class MLFlow(SientiaMonitoring): reference_data['timestamp'] = to_datetime(reference_data['timestamp']) reference_data['timestamp'] = reference_data['timestamp'].dt.strftime(DATETIME_FORMAT) - return reference_data.to_dict(orient='records') \ No newline at end of file + return reference_data.to_dict(orient='records') diff --git a/laborious/activities/model_metrics.py b/laborious/activities/model_metrics.py index ae0f231..29b2848 100644 --- a/laborious/activities/model_metrics.py +++ b/laborious/activities/model_metrics.py @@ -1,25 +1,28 @@ from temporalio import activity, workflow with workflow.unsafe.imports_passed_through(): + import time + import traceback + import warnings from typing import Any + import numpy as np from pandas import DataFrame, Index, to_datetime - - from sientia_do.observability.sientia_monitoring import SientiaMonitoring - from sientia_do.observability.metrics_controller import MetricsController - from sientia_do.observability.logger import Logger + from sientia.ModelAnalysis import ModelAnalysis from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler from sientia_do.notifications.models import NotificationLevel - from sientia.ModelAnalysis import ModelAnalysis - from laborious import metrics - import time - import numpy as np + from sientia_do.observability.logger import Logger + from sientia_do.observability.metrics_controller import MetricsController + from sientia_do.observability.sientia_monitoring import SientiaMonitoring from sientia_do.temporal.constants import DATETIME_FORMAT, DATETIME_FORMAT_WITH_TZ - import warnings - import traceback + + from laborious import metrics warnings.filterwarnings('ignore', category=RuntimeWarning, message='Degrees of freedom <= 0') -warnings.filterwarnings('ignore', category=RuntimeWarning, message='invalid value encountered in scalar divide') +warnings.filterwarnings( + 'ignore', category=RuntimeWarning, message='invalid value encountered in scalar divide' +) + class ModelMetrics(SientiaMonitoring): """ @@ -28,12 +31,12 @@ class ModelMetrics(SientiaMonitoring): This class provides activities for writing metrics to the Prometheus monitoring system. """ - def __init__(self, + def __init__( + self, logger: Logger, notification_handler: NotificationHandler, metrics_controller: MetricsController, ): - SientiaMonitoring.__init__(self, logger, notification_handler, metrics_controller) def close(self) -> None: @@ -42,11 +45,11 @@ class ModelMetrics(SientiaMonitoring): """ SientiaMonitoring.shutdown(self) - def __del__(self): self.close() - async def get_drift_metrics(self, + async def get_drift_metrics( + self, reference_data: DataFrame, target_data: DataFrame, target_name: str, @@ -73,34 +76,37 @@ class ModelMetrics(SientiaMonitoring): 'features': reference_columns, } - model_analysis = ModelAnalysis(config=config) - self.debug(f'Reference data: Size {reference_data.shape} \n{reference_data.head(5).to_string()}', metadata) + self.debug( + f'Reference data: Size {reference_data.shape} \n{reference_data.head(5).to_string()}', + metadata, + ) - self.debug(f'Target data: Size {target_data.shape} \n{target_data.head(5).to_string()}', metadata) + self.debug( + f'Target data: Size {target_data.shape} \n{target_data.head(5).to_string()}', metadata + ) core_labels = self.get_core_labels(metadata, operation_type='detect_univariate_drift') start_time = time.time() try: - univariate_drift = model_analysis.detect_univariate_drift( reference_df=reference_data, analysis_df=target_data, features=reference_columns, timestamp_col=config['timestamp'], methods=drift_metrics, - chunk_period=chunk_period + chunk_period=chunk_period, ) except Exception as e: self.error(f'Error detecting univariate drift: {e}', metadata) await self.emit_metric( - metric_object=metrics.MODEL_ANALYZE_ERROR_COUNT, tags=core_labels) + metric_object=metrics.MODEL_ANALYZE_ERROR_COUNT, tags=core_labels + ) raise e await self.observe_lag(start_time, metrics.MODEL_ANALYZE_LAG, core_labels) await self.emit_metric(metric_object=metrics.MODEL_ANALYZE_COUNT, tags=core_labels) - core_labels = self.get_core_labels(metadata, operation_type='detect_multivariate_drift') start_time = time.time() try: @@ -109,11 +115,13 @@ class ModelMetrics(SientiaMonitoring): analysis_df=target_data, features=reference_columns, timestamp_col=config['timestamp'], - chunk_period=chunk_period + chunk_period=chunk_period, ) except Exception as e: self.error(f'Error detecting multivariate drift: {e}', metadata) - await self.emit_metric(metric_object=metrics.MODEL_ANALYZE_ERROR_COUNT, tags=core_labels) + await self.emit_metric( + metric_object=metrics.MODEL_ANALYZE_ERROR_COUNT, tags=core_labels + ) raise e await self.observe_lag(start_time, metrics.MODEL_ANALYZE_LAG, core_labels) await self.emit_metric(metric_object=metrics.MODEL_ANALYZE_COUNT, tags=core_labels) @@ -127,15 +135,18 @@ class ModelMetrics(SientiaMonitoring): ) except Exception as e: self.error(f'Error getting drift metrics: {e}', metadata) - await self.emit_metric(metric_object=metrics.MODEL_ANALYZE_ERROR_COUNT, tags=core_labels) + await self.emit_metric( + metric_object=metrics.MODEL_ANALYZE_ERROR_COUNT, tags=core_labels + ) raise e await self.observe_lag(start_time, metrics.MODEL_ANALYZE_LAG, core_labels) await self.emit_metric(metric_object=metrics.MODEL_ANALYZE_COUNT, tags=core_labels) - self.debug(f'Drift dataframe: Size {drift_df.shape} \n{drift_df.head(5).to_string()}', metadata) - - return drift_df + self.debug( + f'Drift dataframe: Size {drift_df.shape} \n{drift_df.head(5).to_string()}', metadata + ) + return drift_df @activity.defn(name='calculate_drift') async def calculate_drift(self, input_data: dict[str, Any]) -> list[dict]: @@ -194,8 +205,8 @@ class ModelMetrics(SientiaMonitoring): ) reference_columns = reference_data.drop( - columns=[target_name, 'timestamp', 'target', 'prediction'], - errors='ignore').columns + columns=[target_name, 'timestamp', 'target', 'prediction'], errors='ignore' + ).columns try: drift_df = await self.get_drift_metrics( @@ -224,12 +235,11 @@ class ModelMetrics(SientiaMonitoring): return [] # Drop unnecessary columns - drift_df.drop(columns=['p_value', 'chunk_start_date', 'chunk_end_date'], inplace=True) + drift_df.drop(columns=['p_value'], inplace=True) # Extract timestamps only until minutes if chunk_period == 'min': - target_timestamps = target_data['timestamp'].apply( - lambda x: x[:16]) + target_timestamps = target_data['timestamp'].apply(lambda x: x[:16]) else: target_timestamps = target_data['timestamp'] @@ -237,32 +247,39 @@ class ModelMetrics(SientiaMonitoring): drift_df = drift_df[drift_df['timestamp'].isin(target_timestamps)] if drift_df.empty: - self.warning('No drift metrics found after dropping rows where timestamp is not in target data', metadata) + self.warning( + 'No drift metrics found after dropping rows where timestamp is not in target data', + metadata, + ) return [] # Rename columns to match database columns - drift_df.rename(columns={ - 'metric': 'method', - 'statistic': 'value', - }, inplace=True) + drift_df.rename( + columns={ + 'metric': 'method', + 'statistic': 'value', + }, + inplace=True, + ) # Drop duplicates drift_df.drop_duplicates( - subset=['timestamp', 'method', 'feature'], - keep='first', inplace=True) + subset=['timestamp', 'method', 'feature'], keep='first', inplace=True + ) drift_df['model_id'] = model_id drift_df['accurate'] = accurate - drift_df['timestamp'] = to_datetime(drift_df['timestamp']) drift_df['timestamp'] = drift_df['timestamp'].dt.tz_localize('UTC') drift_df['timestamp'] = drift_df['timestamp'].dt.strftime(DATETIME_FORMAT_WITH_TZ) - self.debug(f'Drift dataframe: Size {drift_df.shape} \n{drift_df.head(5).to_string()}', metadata) + self.debug( + f'Drift dataframe: Size {drift_df.shape} \n{drift_df.head(5).to_string()}', metadata + ) self.debug(f'Drift dataframe: {drift_df.head(5).to_string()}', metadata) - + return drift_df.to_dict(orient='records') @activity.defn(name='calculate_simple_metrics') @@ -298,52 +315,40 @@ class ModelMetrics(SientiaMonitoring): output_data = [] diff = target_data['target'] - target_data['prediction'] - diff_squared = diff ** 2 + diff_squared = diff**2 self.info(f'Calculating simple metrics for model {model_id}: {metrics}', metadata) for metric in metrics: if metric == 'rmse': - output_data.append({ - 'metric': 'rmse', - 'value': np.sqrt(np.mean(diff_squared)) - }) + output_data.append({'metric': 'rmse', 'value': np.sqrt(np.mean(diff_squared))}) elif metric == 'mse': - output_data.append({ - 'metric': 'mse', - 'value': np.mean(diff_squared) - }) + output_data.append({'metric': 'mse', 'value': np.mean(diff_squared)}) elif metric == 'mae': - output_data.append({ - 'metric': 'mae', - 'value': np.mean(np.abs(diff)) - }) + output_data.append({'metric': 'mae', 'value': np.mean(np.abs(diff))}) elif metric == 'r2': y_true = target_data['target'] y_mean = np.mean(y_true) - + ss_res = np.sum(diff_squared) ss_tot = np.sum((y_true - y_mean) ** 2) - + # Evita divisão por zero if ss_tot == 0: r2_score = 0.0 else: r2_score = 1 - (ss_res / ss_tot) - - output_data.append({ - 'metric': 'r2', - 'value': r2_score - }) - + + output_data.append({'metric': 'r2', 'value': r2_score}) + data = DataFrame(output_data) data['model_id'] = model_id data['timestamp'] = target_data['timestamp'].max() data['data_size'] = data_size data['interval_minutes'] = interval_minutes - self.debug(f'Simple metrics dataframe: Size {data.shape} \n{data.head(5).to_string()}', metadata) + self.debug( + f'Simple metrics dataframe: Size {data.shape} \n{data.head(5).to_string()}', metadata + ) return data.to_dict(orient='records') - - diff --git a/laborious/activities/opc.py b/laborious/activities/opc.py index 404840e..cfbf450 100644 --- a/laborious/activities/opc.py +++ b/laborious/activities/opc.py @@ -1,4 +1,5 @@ -from typing import Hashable +from collections.abc import Hashable + from temporalio import activity, workflow with workflow.unsafe.imports_passed_through(): @@ -12,7 +13,6 @@ with workflow.unsafe.imports_passed_through(): from sientia_do.observability.metrics_controller import MetricsController from sientia_do.observability.sientia_monitoring import SientiaMonitoring - from laborious.utils.dataframe_utils import ensure_dataframe from laborious.utils.repository.opc_repository import OpcRepository OPC_WRITTING_ERROR_CONFIDENCE = 12 @@ -297,7 +297,7 @@ class OPC(SientiaMonitoring): """ metadata = input_data['metadata'] self.info('Writing data to OPC servers...', metadata) - data = ensure_dataframe(input_data['data']) + data = DataFrame(input_data['data']) opc_output_config = input_data['opc_output_config'] self.info(f'Data to write: {data.size} rows', metadata) diff --git a/laborious/metrics.py b/laborious/metrics.py index 59e4a81..be5604d 100644 --- a/laborious/metrics.py +++ b/laborious/metrics.py @@ -187,4 +187,4 @@ MODEL_ANALYZE_ERROR_COUNT = Counter( 'laborious_model_analyze_error_count', 'Number of errors during analyze operations', SIENTIA_CORE_LABELS, -) \ No newline at end of file +) diff --git a/laborious/utils/dataframe_utils.py b/laborious/utils/dataframe_utils.py deleted file mode 100644 index c87764d..0000000 --- a/laborious/utils/dataframe_utils.py +++ /dev/null @@ -1,31 +0,0 @@ -""" -DataFrame utility functions for handling serialized DataFrames. - -This module provides helper functions to work with DataFrames that may -come from Temporal serialization (already as DataFrame) or from legacy -code (as dict). -""" - -from typing import Any - -from pandas import DataFrame - - -def ensure_dataframe(data: Any) -> DataFrame: - """ - Ensure that data is a DataFrame, converting from dict if necessary. - - This function handles both cases: - - Data already deserialized as DataFrame (from Temporal codec) - - Data as dict (legacy format or non-DataFrame serialization) - - Args: - data: Data that should be a DataFrame (can be DataFrame or dict) - - Returns: - DataFrame: The data as a pandas DataFrame - """ - if isinstance(data, DataFrame): - return data - return DataFrame(data) - diff --git a/laborious/utils/repository/model_repository.py b/laborious/utils/repository/model_repository.py index e80d095..33366e6 100644 --- a/laborious/utils/repository/model_repository.py +++ b/laborious/utils/repository/model_repository.py @@ -16,11 +16,11 @@ Capabilities: import ctypes import gc -from io import StringIO import threading import time import traceback from datetime import datetime, timedelta +from io import StringIO from os import environ, makedirs, path from shutil import rmtree from typing import Any, Literal, overload @@ -34,7 +34,6 @@ from sientia_do.observability.logger import Logger from sientia_do.observability.metrics_controller import MetricsController from sientia_do.observability.sientia_monitoring import SientiaMonitoring from sientia_do.temporal.constants import DATETIME_FORMAT_WITH_TZ -from sientia.ModelAnalysis import ModelAnalysis from laborious import metrics @@ -217,8 +216,9 @@ class MLFlowRepository(SientiaMonitoring): run_info = mlflow.get_run(run_id) return run_info.data.params - def check_artifact_exists(self, run_id: str, - artifact_path: str, metadata: dict[str, Any]) -> bool: + def check_artifact_exists( + self, run_id: str, artifact_path: str, metadata: dict[str, Any] + ) -> bool: """ Check if an artifact exists in the MLflow Model Registry. @@ -233,8 +233,9 @@ class MLFlowRepository(SientiaMonitoring): self.debug(f'Artifacts of {run_id}: \n{artifacts}', metadata) self.debug(f'Looking for artifact {artifact_path} in {run_id}', metadata) - + return any(artifact.path == artifact_path for artifact in artifacts) + """ Functions related to download and load models """ @@ -279,10 +280,9 @@ class MLFlowRepository(SientiaMonitoring): return artifacts - - async def load_artifact_dataframe(self, model_name: str, artifact_path: str, - metadata: dict[str, Any]) -> pd.DataFrame | None: - + async def load_artifact_dataframe( + self, model_name: str, artifact_path: str, metadata: dict[str, Any] + ) -> pd.DataFrame | None: """ Load the dataframe content of an artifact from the MLflow Model Registry. @@ -300,7 +300,7 @@ class MLFlowRepository(SientiaMonitoring): if not self.check_artifact_exists(run_id, artifact_path, metadata): return None - artifact_path = path.join("runs:/", run_id, artifact_path) + artifact_path = path.join('runs:/', run_id, artifact_path) start_time = time.time() try: @@ -704,8 +704,9 @@ class MLFlowRepository(SientiaMonitoring): Functions related to model retraining """ - def get_prediction_data(self, prediction_model: Any, retrain_dataset: pd.DataFrame, - target_name: str) -> pd.DataFrame: + def get_prediction_data( + self, prediction_model: Any, retrain_dataset: pd.DataFrame, target_name: str + ) -> pd.DataFrame: """ Get prediction data from prediction model. """ @@ -714,7 +715,6 @@ class MLFlowRepository(SientiaMonitoring): prediction_data = prediction_model.predict(retrain_dataset) if isinstance(prediction_data, pd.DataFrame): - prediction_data.columns = pd.Index(['prediction']) else: @@ -724,7 +724,8 @@ class MLFlowRepository(SientiaMonitoring): # Merge prediction data with retrain_dataset on index prediction_data = pd.merge( - retrain_dataset, prediction_data, left_index=True, right_index=True, how='left') + retrain_dataset, prediction_data, left_index=True, right_index=True, how='left' + ) # Rename column "target_name" to "target" prediction_data.rename(columns={target_name: 'target'}, inplace=True) @@ -733,9 +734,7 @@ class MLFlowRepository(SientiaMonitoring): prediction_data.reset_index(drop=True, inplace=True) - prediction_data.sort_values( - by='timestamp', ascending=True, inplace=True - ) + prediction_data.sort_values(by='timestamp', ascending=True, inplace=True) return prediction_data @@ -859,8 +858,7 @@ class MLFlowRepository(SientiaMonitoring): prediction_model.fit(retrain_dataset) # get prediction data - prediction_data = self.get_prediction_data( - prediction_model, retrain_dataset, target_name) + prediction_data = self.get_prediction_data(prediction_model, retrain_dataset, target_name) self.info(f'Model experiment creation completed successfully for {model_name}', metadata) @@ -1439,4 +1437,3 @@ class MLFlowRepository(SientiaMonitoring): metadata_result['mlflow_experiment_id'] = experiment_id return metadata_result - diff --git a/laborious/worker/worker.py b/laborious/worker/worker.py index 1d28016..0e9d3ca 100644 --- a/laborious/worker/worker.py +++ b/laborious/worker/worker.py @@ -29,7 +29,6 @@ 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 @@ -48,11 +47,10 @@ with workflow.unsafe.imports_passed_through(): build_opc_config, build_postgres_config, ) + from laborious.workflows.drift import Drift from laborious.workflows.minimal_retrain import MinimalRetrain from laborious.workflows.predictions_batch import PredictionsBatch - from laborious.workflows.drift import Drift from laborious.workflows.simple_metrics import SimpleMetrics - from laborious.workflows.sub_workflows.format_and_export_prediction import ( FormatAndExportPrediction, ) diff --git a/laborious/workflows/drift.py b/laborious/workflows/drift.py index d556dac..85780f7 100644 --- a/laborious/workflows/drift.py +++ b/laborious/workflows/drift.py @@ -4,10 +4,10 @@ with workflow.unsafe.imports_passed_through(): from datetime import timedelta from typing import Any + from sientia_do.temporal.constants import DATETIME_FORMAT_WITH_TZ from sientia_do.temporal.policies import retry_policy from laborious.activities.activities import Activities - from sientia_do.temporal.constants import DATETIME_FORMAT_WITH_TZ @workflow.defn(name='drift') @@ -39,10 +39,10 @@ class Drift: gathering_query = f""" SELECT * - FROM {input_data['schema']}.{input_data['source_table_name']} + FROM "{input_data['schema']}"."{input_data['source_table_name']}" WHERE - model_id = {input_data['model_id']} AND - timestamp > NOW() - INTERVAL '{input_data['interval']} minutes' + model_id = '{input_data['model_id']}' AND + timestamp > NOW() - INTERVAL {input_data['interval']} minutes ORDER BY timestamp ASC """ @@ -60,10 +60,7 @@ class Drift: reference_data_handler = workflow.start_local_activity_method( Activities.get_reference_data, - { - **metadata, - 'model_name': input_data['model_name'] - }, + {**metadata, 'model_name': input_data['model_name']}, retry_policy=retry_policy, start_to_close_timeout=timedelta(seconds=300), ) @@ -83,8 +80,9 @@ class Drift: 'model_name': input_data['model_name'], 'model_id': input_data['model_id'], 'target_name': target_name, - 'drift_metrics': input_data.get('drift_metrics', - ['kolmogorov_smirnov', 'jensen_shannon', 'wasserstein']), + 'drift_metrics': input_data.get( + 'drift_metrics', ['kolmogorov_smirnov', 'jensen_shannon', 'wasserstein'] + ), 'chunk_period': input_data.get('chunk_period', 'min'), }, retry_policy=retry_policy, @@ -99,8 +97,11 @@ class Drift: 'data': drift_data, 'schema': input_data['schema'], 'table_name': input_data['target_table_name'], - '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=300), - ) \ No newline at end of file + ) diff --git a/laborious/workflows/simple_metrics.py b/laborious/workflows/simple_metrics.py index ce950af..7ab7153 100644 --- a/laborious/workflows/simple_metrics.py +++ b/laborious/workflows/simple_metrics.py @@ -4,10 +4,10 @@ with workflow.unsafe.imports_passed_through(): from datetime import timedelta from typing import Any + from sientia_do.temporal.constants import DATETIME_FORMAT_WITH_TZ from sientia_do.temporal.policies import retry_policy from laborious.activities.activities import Activities - from sientia_do.temporal.constants import DATETIME_FORMAT_WITH_TZ @workflow.defn(name='simple_metrics') @@ -19,10 +19,10 @@ class SimpleMetrics: """ metadata = { 'metadata': { - 'schedule_name': input_data['schedule_name'], - 'model_name': input_data['model_name'], 'model_id': input_data['model_id'], + 'model_name': input_data['model_name'], 'workflow_name': 'simple_metrics', + 'schedule_name': input_data['schedule_name'], } } @@ -34,15 +34,15 @@ class SimpleMetrics: query = f""" select p."timestamp", p.prediction, ld.value as "target" - from {input_data['schema']}.{input_data['predictions_table_name']} p - inner join {input_data['schema']}.{input_data['data_table_name']} ld + from "{input_data['schema']}"."{input_data['predictions_table_name']}" p + inner join "{input_data['schema']}"."{input_data['data_table_name']}" ld on p."timestamp" = ld."timestamp" where - p.model_id = {model_id} and + p.model_id = '{model_id}' and p.prediction is not null and ld.variable = '{target_name}' and ld.value is not null and - p."timestamp" >= NOW() - INTERVAL '{interval_minutes} minutes' + p."timestamp" >= NOW() - INTERVAL {interval_minutes} minutes order by p."timestamp" desc; """ @@ -74,7 +74,6 @@ class SimpleMetrics: retry_policy=retry_policy, start_to_close_timeout=timedelta(seconds=300), ) - if not simple_metrics: return @@ -93,4 +92,4 @@ class SimpleMetrics: }, retry_policy=retry_policy, start_to_close_timeout=timedelta(seconds=300), - ) \ No newline at end of file + ) diff --git a/laborious/workflows/sub_workflows/format_and_export_prediction.py b/laborious/workflows/sub_workflows/format_and_export_prediction.py index 4280ccf..1851c61 100644 --- a/laborious/workflows/sub_workflows/format_and_export_prediction.py +++ b/laborious/workflows/sub_workflows/format_and_export_prediction.py @@ -89,7 +89,6 @@ class FormatAndExportPrediction: ) if transformed_data is not None: - transformed = await workflow.execute_local_activity_method( Activities.format_transformed_data, { diff --git a/pyproject.toml b/pyproject.toml index 12e06d4..0f912c2 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -48,6 +48,7 @@ ignore = [ "S101", # use of assert (needed for tests) "S105", # possible hardcoded password (false positives) "S106", # possible hardcoded password (false positives) + "S608", # potential sql injection (false positives) "N802", # function name should be lowercase (temporal decorators) "N806", # variable in function should be lowercase ] @@ -152,4 +153,4 @@ directory = "htmlcov" [tool.bandit] exclude_dirs = ["tests", "venv", ".venv"] -skips = ["B101", "B601"] # Skip assert and shell injection in controlled environments \ No newline at end of file +skips = ["B101", "B601", "B608"] # Skip assert, shell injection, and SQL injection (false positives) \ No newline at end of file diff --git a/tests/laborious/activities/test_model_metrics.py b/tests/laborious/activities/test_model_metrics.py index 87c3a68..f20687d 100644 --- a/tests/laborious/activities/test_model_metrics.py +++ b/tests/laborious/activities/test_model_metrics.py @@ -1,4 +1,4 @@ -from unittest.mock import ANY, AsyncMock, MagicMock, call, patch +from unittest.mock import ANY, AsyncMock, MagicMock, patch from pandas import DataFrame from pytest import fixture, mark @@ -22,7 +22,13 @@ def model_metrics_activity(): model_metrics.send_notification = MagicMock() model_metrics.send_notification_async = AsyncMock() model_metrics.emit_metric = AsyncMock() - model_metrics.get_core_labels = MagicMock(return_value={'pod_id': 'test_pod', 'model_name': 'test_model', 'workflow_name': 'test_workflow'}) + model_metrics.get_core_labels = MagicMock( + return_value={ + 'pod_id': 'test_pod', + 'model_name': 'test_model', + 'workflow_name': 'test_workflow', + } + ) model_metrics.observe_lag = AsyncMock() model_metrics.pod_id = 'test_pod' return model_metrics @@ -60,7 +66,7 @@ async def test_calculate_drift_invalid_chunk_period(model_metrics_activity): try: await model_metrics_activity.calculate_drift(input_data) except ValueError as e: - assert str(e) == "Invalid chunk period: invalid, must be \"min\" or \"s\"" + assert str(e) == 'Invalid chunk period: invalid, must be "min" or "s"' model_metrics_activity.error.assert_called_once_with( 'Invalid chunk period: invalid', metadata['metadata'] ) @@ -76,26 +82,41 @@ async def test_calculate_drift_with_reference_data( ): # Arrange mock_to_datetime.return_value.dt.strftime.return_value = '2023-05-26 11:12:27' - mock_to_datetime.return_value.dt.tz_localize.return_value.dt.strftime.return_value = '2023-05-26 11:12:27+00:00' - + mock_to_datetime.return_value.dt.tz_localize.return_value.dt.strftime.return_value = ( + '2023-05-26 11:12:27+00:00' + ) + mock_drift_df = MagicMock() mock_drift_df.empty = False mock_drift_df.drop.return_value = mock_drift_df mock_drift_df.__getitem__.return_value.isin.return_value = [True] mock_drift_df.__getitem__.return_value = mock_drift_df - mock_drift_df.__getitem__.return_value.dt.tz_localize.return_value.dt.strftime.return_value = '2023-05-26 11:12:27+00:00' + mock_drift_df.__getitem__.return_value.dt.tz_localize.return_value.dt.strftime.return_value = ( + '2023-05-26 11:12:27+00:00' + ) mock_drift_df.rename.return_value = mock_drift_df mock_drift_df.drop_duplicates.return_value = mock_drift_df - mock_drift_df.to_dict.return_value = {'method': ['ks_test'], 'value': [0.5], 'feature': ['feature1'], 'timestamp': ['2023-05-26 11:12:27+00:00'], 'model_id': ['test_model_id'], 'accurate': [True]} - + mock_drift_df.to_dict.return_value = [ + { + 'method': 'ks_test', + 'value': 0.5, + 'feature': 'feature1', + 'timestamp': '2023-05-26 11:12:27+00:00', + 'model_id': 'test_model_id', + 'accurate': True, + } + ] + model_metrics_activity.get_drift_metrics = AsyncMock(return_value=mock_drift_df) - reference_data = DataFrame({ - 'timestamp': ['2023-05-26 11:12:27'], - 'target': [1.0], - 'feature1': [1.0], - }) - + reference_data = DataFrame( + { + 'timestamp': ['2023-05-26 11:12:27'], + 'target': [1.0], + 'feature1': [1.0], + } + ) + mock_target_df = MagicMock() mock_target_df.pivot.return_value = mock_target_df mock_target_df.index = ['2023-05-26 11:12:27'] @@ -124,16 +145,20 @@ async def test_calculate_drift_with_reference_data( result = await model_metrics_activity.calculate_drift(input_data) # Assert - assert isinstance(result, dict) - assert result == mock_drift_df.to_dict.return_value + assert isinstance(result, list) + assert result == mock_drift_df.to_dict.return_value # type: ignore[comparison-overlap] model_metrics_activity.info.assert_called() model_metrics_activity.get_drift_metrics.assert_called_once() # Verify transformations were called - mock_drift_df.drop.assert_called_once_with(columns=['p_value', 'chunk_start_date', 'chunk_end_date'], inplace=True) + mock_drift_df.drop.assert_called_once_with(columns=['p_value'], inplace=True) mock_drift_df.__getitem__.assert_called() - mock_drift_df.rename.assert_called_once_with(columns={'metric': 'method', 'statistic': 'value'}, inplace=True) - mock_drift_df.drop_duplicates.assert_called_once_with(subset=['timestamp', 'method', 'feature'], keep='first', inplace=True) - mock_drift_df.to_dict.assert_called_once() + mock_drift_df.rename.assert_called_once_with( + columns={'metric': 'method', 'statistic': 'value'}, inplace=True + ) + mock_drift_df.drop_duplicates.assert_called_once_with( + subset=['timestamp', 'method', 'feature'], keep='first', inplace=True + ) + mock_drift_df.to_dict.assert_called_once_with(orient='records') @mark.asyncio @@ -144,18 +169,31 @@ async def test_calculate_drift_without_reference_data( ): # Arrange mock_to_datetime.return_value.dt.strftime.return_value = '2023-05-26 11:12:27' - mock_to_datetime.return_value.dt.tz_localize.return_value.dt.strftime.return_value = '2023-05-26 11:12:27+00:00' - + mock_to_datetime.return_value.dt.tz_localize.return_value.dt.strftime.return_value = ( + '2023-05-26 11:12:27+00:00' + ) + mock_drift_df = MagicMock() mock_drift_df.empty = False mock_drift_df.drop.return_value = mock_drift_df mock_drift_df.__getitem__.return_value.isin.return_value = [True] mock_drift_df.__getitem__.return_value = mock_drift_df - mock_drift_df.__getitem__.return_value.dt.tz_localize.return_value.dt.strftime.return_value = '2023-05-26 11:12:27+00:00' + mock_drift_df.__getitem__.return_value.dt.tz_localize.return_value.dt.strftime.return_value = ( + '2023-05-26 11:12:27+00:00' + ) mock_drift_df.rename.return_value = mock_drift_df mock_drift_df.drop_duplicates.return_value = mock_drift_df - mock_drift_df.to_dict.return_value = {'method': ['ks_test'], 'value': [0.5], 'feature': ['feature1'], 'timestamp': ['2023-05-26 11:12:27+00:00'], 'model_id': ['test_model_id'], 'accurate': [False]} - + mock_drift_df.to_dict.return_value = [ + { + 'method': 'ks_test', + 'value': 0.5, + 'feature': 'feature1', + 'timestamp': '2023-05-26 11:12:27+00:00', + 'model_id': 'test_model_id', + 'accurate': False, + } + ] + model_metrics_activity.get_drift_metrics = AsyncMock(return_value=mock_drift_df) target_data_dict = { @@ -163,19 +201,25 @@ async def test_calculate_drift_without_reference_data( 'variable': ['feature1', 'feature1', 'feature1'], 'value': [1.0, 2.0, 3.0], } - + mock_target_df = MagicMock() mock_target_df.pivot.return_value = mock_target_df mock_target_df.index = ['2023-05-26 11:12:27', '2023-05-26 11:12:28', '2023-05-26 11:12:29'] mock_target_df.reset_index.return_value = mock_target_df mock_target_df.dropna.return_value = mock_target_df mock_target_df.sort_values.return_value = mock_target_df - mock_target_df.head.return_value = DataFrame({'timestamp': ['2023-05-26 11:12:27'], 'feature1': [1.0]}) - mock_target_df.__getitem__.return_value.apply.return_value = ['2023-05-26 11:12:27', '2023-05-26 11:12:28', '2023-05-26 11:12:29'] + mock_target_df.head.return_value = DataFrame( + {'timestamp': ['2023-05-26 11:12:27'], 'feature1': [1.0]} + ) + mock_target_df.__getitem__.return_value.apply.return_value = [ + '2023-05-26 11:12:27', + '2023-05-26 11:12:28', + '2023-05-26 11:12:29', + ] mock_target_df.drop.return_value.columns = ['feature1'] mock_dataframe.return_value = mock_target_df mock_dataframe.side_effect = lambda x=None: mock_target_df if x is not None else mock_target_df - + input_data = { **metadata, 'model_name': 'test_model', @@ -191,7 +235,7 @@ async def test_calculate_drift_without_reference_data( result = await model_metrics_activity.calculate_drift(input_data) # Assert - assert isinstance(result, dict) + assert isinstance(result, list) assert result == mock_drift_df.to_dict.return_value model_metrics_activity.warning.assert_called() model_metrics_activity.send_notification_async.assert_called_once_with( @@ -203,11 +247,15 @@ async def test_calculate_drift_without_reference_data( attachment_content=ANY, ) # Verify transformations were called - mock_drift_df.drop.assert_called_once_with(columns=['p_value', 'chunk_start_date', 'chunk_end_date'], inplace=True) + mock_drift_df.drop.assert_called_once_with(columns=['p_value'], inplace=True) mock_drift_df.__getitem__.assert_called() - mock_drift_df.rename.assert_called_once_with(columns={'metric': 'method', 'statistic': 'value'}, inplace=True) - mock_drift_df.drop_duplicates.assert_called_once_with(subset=['timestamp', 'method', 'feature'], keep='first', inplace=True) - mock_drift_df.to_dict.assert_called_once() + mock_drift_df.rename.assert_called_once_with( + columns={'metric': 'method', 'statistic': 'value'}, inplace=True + ) + mock_drift_df.drop_duplicates.assert_called_once_with( + subset=['timestamp', 'method', 'feature'], keep='first', inplace=True + ) + mock_drift_df.to_dict.assert_called_once_with(orient='records') @mark.asyncio @@ -218,15 +266,17 @@ async def test_calculate_drift_empty_drift_df( ): # Arrange mock_to_datetime.return_value.dt.strftime.return_value = '2023-05-26 11:12:27' - + model_metrics_activity.get_drift_metrics = AsyncMock(return_value=DataFrame()) - reference_data = DataFrame({ - 'timestamp': ['2023-05-26 11:12:27'], - 'target': [1.0], - 'feature1': [1.0], - }) - + reference_data = DataFrame( + { + 'timestamp': ['2023-05-26 11:12:27'], + 'target': [1.0], + 'feature1': [1.0], + } + ) + mock_target_df = MagicMock() mock_target_df.pivot.return_value = mock_target_df mock_target_df.index = ['2023-05-26 11:12:27'] @@ -235,7 +285,7 @@ async def test_calculate_drift_empty_drift_df( mock_target_df.__getitem__.return_value.apply.return_value = ['2023-05-26 11:12:27'] mock_target_df.drop.return_value.columns = ['feature1'] mock_dataframe.return_value = mock_target_df - + input_data = { **metadata, 'model_name': 'test_model', @@ -255,8 +305,10 @@ async def test_calculate_drift_empty_drift_df( result = await model_metrics_activity.calculate_drift(input_data) # Assert - assert result == {} - model_metrics_activity.warning.assert_called_with('No drift metrics found', metadata['metadata']) + assert result == [] + model_metrics_activity.warning.assert_called_with( + 'No drift metrics found', metadata['metadata'] + ) @mark.asyncio @@ -267,35 +319,37 @@ async def test_calculate_drift_empty_after_timestamp_filter( ): # Arrange mock_to_datetime.return_value.dt.strftime.return_value = '2023-05-26 11:12:27' - + mock_drift_df = MagicMock() mock_drift_df.empty = False mock_drift_df.drop.return_value = mock_drift_df - + # Set up __getitem__ to handle filtering - timestamp access returns series with isin=False # and filtering returns empty DataFrame mock_timestamp_series = MagicMock() mock_timestamp_series.isin.return_value = [False] mock_empty_df = MagicMock() mock_empty_df.empty = True - + def getitem_side_effect(key): if key == 'timestamp': return mock_timestamp_series else: # This is the filtering operation - return empty DataFrame return mock_empty_df - + mock_drift_df.__getitem__.side_effect = getitem_side_effect model_metrics_activity.get_drift_metrics = AsyncMock(return_value=mock_drift_df) - reference_data = DataFrame({ - 'timestamp': ['2023-05-26 11:12:27'], - 'target': [1.0], - 'feature1': [1.0], - }) - + reference_data = DataFrame( + { + 'timestamp': ['2023-05-26 11:12:27'], + 'target': [1.0], + 'feature1': [1.0], + } + ) + mock_target_df = MagicMock() mock_target_df.pivot.return_value = mock_target_df mock_target_df.index = ['2023-05-26 11:12:27'] @@ -324,13 +378,13 @@ async def test_calculate_drift_empty_after_timestamp_filter( result = await model_metrics_activity.calculate_drift(input_data) # Assert - assert result == {} + assert result == [] model_metrics_activity.warning.assert_called_with( 'No drift metrics found after dropping rows where timestamp is not in target data', - metadata['metadata'] + metadata['metadata'], ) # Verify transformations were called - mock_drift_df.drop.assert_called_once_with(columns=['p_value', 'chunk_start_date', 'chunk_end_date'], inplace=True) + mock_drift_df.drop.assert_called_once_with(columns=['p_value'], inplace=True) mock_drift_df.__getitem__.assert_called() @@ -342,26 +396,41 @@ async def test_calculate_drift_success_min( ): # Arrange mock_to_datetime.return_value.dt.strftime.return_value = '2023-05-26 11:12:27' - mock_to_datetime.return_value.dt.tz_localize.return_value.dt.strftime.return_value = '2023-05-26 11:12:27+00:00' - + mock_to_datetime.return_value.dt.tz_localize.return_value.dt.strftime.return_value = ( + '2023-05-26 11:12:27+00:00' + ) + mock_drift_df = MagicMock() mock_drift_df.empty = False mock_drift_df.drop.return_value = mock_drift_df mock_drift_df.__getitem__.return_value.isin.return_value = [True] mock_drift_df.__getitem__.return_value = mock_drift_df - mock_drift_df.__getitem__.return_value.dt.tz_localize.return_value.dt.strftime.return_value = '2023-05-26 11:12:27+00:00' + mock_drift_df.__getitem__.return_value.dt.tz_localize.return_value.dt.strftime.return_value = ( + '2023-05-26 11:12:27+00:00' + ) mock_drift_df.rename.return_value = mock_drift_df mock_drift_df.drop_duplicates.return_value = mock_drift_df - mock_drift_df.to_dict.return_value = {'method': ['ks_test'], 'value': [0.5], 'feature': ['feature1'], 'timestamp': ['2023-05-26 11:12:27+00:00'], 'model_id': ['test_model_id'], 'accurate': [True]} - + mock_drift_df.to_dict.return_value = [ + { + 'method': 'ks_test', + 'value': 0.5, + 'feature': 'feature1', + 'timestamp': '2023-05-26 11:12:27+00:00', + 'model_id': 'test_model_id', + 'accurate': True, + } + ] + model_metrics_activity.get_drift_metrics = AsyncMock(return_value=mock_drift_df) - reference_data = DataFrame({ - 'timestamp': ['2023-05-26 11:12:27'], - 'target': [1.0], - 'feature1': [1.0], - }) - + reference_data = DataFrame( + { + 'timestamp': ['2023-05-26 11:12:27'], + 'target': [1.0], + 'feature1': [1.0], + } + ) + mock_target_df = MagicMock() mock_target_df.pivot.return_value = mock_target_df mock_target_df.index = ['2023-05-26 11:12:27'] @@ -390,46 +459,63 @@ async def test_calculate_drift_success_min( result = await model_metrics_activity.calculate_drift(input_data) # Assert - assert isinstance(result, dict) - assert result == mock_drift_df.to_dict.return_value + assert isinstance(result, list) + assert result == mock_drift_df.to_dict.return_value # type: ignore[comparison-overlap] model_metrics_activity.info.assert_called() model_metrics_activity.get_drift_metrics.assert_called_once() # Verify transformations were called - mock_drift_df.drop.assert_called_once_with(columns=['p_value', 'chunk_start_date', 'chunk_end_date'], inplace=True) + mock_drift_df.drop.assert_called_once_with(columns=['p_value'], inplace=True) mock_drift_df.__getitem__.assert_called() - mock_drift_df.rename.assert_called_once_with(columns={'metric': 'method', 'statistic': 'value'}, inplace=True) - mock_drift_df.drop_duplicates.assert_called_once_with(subset=['timestamp', 'method', 'feature'], keep='first', inplace=True) - mock_drift_df.to_dict.assert_called_once() + mock_drift_df.rename.assert_called_once_with( + columns={'metric': 'method', 'statistic': 'value'}, inplace=True + ) + mock_drift_df.drop_duplicates.assert_called_once_with( + subset=['timestamp', 'method', 'feature'], keep='first', inplace=True + ) + mock_drift_df.to_dict.assert_called_once_with(orient='records') @mark.asyncio @patch('laborious.activities.model_metrics.DataFrame') @patch('laborious.activities.model_metrics.to_datetime') -async def test_calculate_drift_success_s( - mock_to_datetime, mock_dataframe, model_metrics_activity -): +async def test_calculate_drift_success_s(mock_to_datetime, mock_dataframe, model_metrics_activity): # Arrange mock_to_datetime.return_value.dt.strftime.return_value = '2023-05-26 11:12:27' - mock_to_datetime.return_value.dt.tz_localize.return_value.dt.strftime.return_value = '2023-05-26 11:12:27+00:00' - + mock_to_datetime.return_value.dt.tz_localize.return_value.dt.strftime.return_value = ( + '2023-05-26 11:12:27+00:00' + ) + mock_drift_df = MagicMock() mock_drift_df.empty = False mock_drift_df.drop.return_value = mock_drift_df mock_drift_df.__getitem__.return_value.isin.return_value = [True] mock_drift_df.__getitem__.return_value = mock_drift_df - mock_drift_df.__getitem__.return_value.dt.tz_localize.return_value.dt.strftime.return_value = '2023-05-26 11:12:27+00:00' + mock_drift_df.__getitem__.return_value.dt.tz_localize.return_value.dt.strftime.return_value = ( + '2023-05-26 11:12:27+00:00' + ) mock_drift_df.rename.return_value = mock_drift_df mock_drift_df.drop_duplicates.return_value = mock_drift_df - mock_drift_df.to_dict.return_value = {'method': ['ks_test'], 'value': [0.5], 'feature': ['feature1'], 'timestamp': ['2023-05-26 11:12:27+00:00'], 'model_id': ['test_model_id'], 'accurate': [True]} - + mock_drift_df.to_dict.return_value = [ + { + 'method': 'ks_test', + 'value': 0.5, + 'feature': 'feature1', + 'timestamp': '2023-05-26 11:12:27+00:00', + 'model_id': 'test_model_id', + 'accurate': True, + } + ] + model_metrics_activity.get_drift_metrics = AsyncMock(return_value=mock_drift_df) - reference_data = DataFrame({ - 'timestamp': ['2023-05-26 11:12:27'], - 'target': [1.0], - 'feature1': [1.0], - }) - + reference_data = DataFrame( + { + 'timestamp': ['2023-05-26 11:12:27'], + 'target': [1.0], + 'feature1': [1.0], + } + ) + mock_target_df = MagicMock() mock_target_df.pivot.return_value = mock_target_df mock_target_df.index = ['2023-05-26 11:12:27'] @@ -458,16 +544,20 @@ async def test_calculate_drift_success_s( result = await model_metrics_activity.calculate_drift(input_data) # Assert - assert isinstance(result, dict) - assert result == mock_drift_df.to_dict.return_value + assert isinstance(result, list) + assert result == mock_drift_df.to_dict.return_value # type: ignore[comparison-overlap] model_metrics_activity.info.assert_called() model_metrics_activity.get_drift_metrics.assert_called_once() # Verify transformations were called - mock_drift_df.drop.assert_called_once_with(columns=['p_value', 'chunk_start_date', 'chunk_end_date'], inplace=True) + mock_drift_df.drop.assert_called_once_with(columns=['p_value'], inplace=True) mock_drift_df.__getitem__.assert_called() - mock_drift_df.rename.assert_called_once_with(columns={'metric': 'method', 'statistic': 'value'}, inplace=True) - mock_drift_df.drop_duplicates.assert_called_once_with(subset=['timestamp', 'method', 'feature'], keep='first', inplace=True) - mock_drift_df.to_dict.assert_called_once() + mock_drift_df.rename.assert_called_once_with( + columns={'metric': 'method', 'statistic': 'value'}, inplace=True + ) + mock_drift_df.drop_duplicates.assert_called_once_with( + subset=['timestamp', 'method', 'feature'], keep='first', inplace=True + ) + mock_drift_df.to_dict.assert_called_once_with(orient='records') @mark.asyncio @@ -478,15 +568,19 @@ async def test_calculate_drift_get_drift_metrics_error( ): # Arrange mock_to_datetime.return_value.dt.strftime.return_value = '2023-05-26 11:12:27' - - model_metrics_activity.get_drift_metrics = AsyncMock(side_effect=Exception('Get drift metrics error')) - reference_data = DataFrame({ - 'timestamp': ['2023-05-26 11:12:27'], - 'target': [1.0], - 'feature1': [1.0], - }) - + model_metrics_activity.get_drift_metrics = AsyncMock( + side_effect=Exception('Get drift metrics error') + ) + + reference_data = DataFrame( + { + 'timestamp': ['2023-05-26 11:12:27'], + 'target': [1.0], + 'feature1': [1.0], + } + ) + mock_target_df = MagicMock() mock_target_df.pivot.return_value = mock_target_df mock_target_df.index = ['2023-05-26 11:12:27'] @@ -515,10 +609,9 @@ async def test_calculate_drift_get_drift_metrics_error( result = await model_metrics_activity.calculate_drift(input_data) # Assert - assert result == {} + assert result == [] model_metrics_activity.error.assert_called_once_with( - 'Error getting drift metrics: Get drift metrics error', - metadata['metadata'] + 'Error getting drift metrics: Get drift metrics error', metadata['metadata'] ) model_metrics_activity.send_notification_async.assert_called_once_with( metadata=metadata['metadata'], @@ -540,30 +633,36 @@ async def test_get_drift_metrics_success( ): # Arrange mock_time.return_value = 1000.0 - - mock_drift_df = DataFrame({ - 'timestamp': ['2023-05-26 11:12:27'], - 'metric': ['ks_test'], - 'statistic': [0.5], - 'feature': ['feature1'], - }) - + + mock_drift_df = DataFrame( + { + 'timestamp': ['2023-05-26 11:12:27'], + 'metric': ['ks_test'], + 'statistic': [0.5], + 'feature': ['feature1'], + } + ) + mock_model_analysis.return_value.detect_univariate_drift.return_value = MagicMock() mock_model_analysis.return_value.detect_multivariate_drift.return_value = MagicMock() mock_model_analysis.return_value.get_drift_metrics_dataframe.return_value = mock_drift_df - reference_data = DataFrame({ - 'timestamp': ['2023-05-26 11:12:27'], - 'target': [1.0], - 'feature1': [1.0], - }) - - target_data = DataFrame({ - 'timestamp': ['2023-05-26 11:12:27'], - 'target': [1.0], - 'feature1': [1.0], - }) - + reference_data = DataFrame( + { + 'timestamp': ['2023-05-26 11:12:27'], + 'target': [1.0], + 'feature1': [1.0], + } + ) + + target_data = DataFrame( + { + 'timestamp': ['2023-05-26 11:12:27'], + 'target': [1.0], + 'feature1': [1.0], + } + ) + reference_columns = reference_data.drop( columns=['target', 'timestamp'], errors='ignore' ).columns @@ -596,21 +695,27 @@ async def test_get_drift_metrics_univariate_error( ): # Arrange mock_time.return_value = 1000.0 - - mock_model_analysis.return_value.detect_univariate_drift.side_effect = Exception('Univariate drift error') - reference_data = DataFrame({ - 'timestamp': ['2023-05-26 11:12:27'], - 'target': [1.0], - 'feature1': [1.0], - }) - - target_data = DataFrame({ - 'timestamp': ['2023-05-26 11:12:27'], - 'target': [1.0], - 'feature1': [1.0], - }) - + mock_model_analysis.return_value.detect_univariate_drift.side_effect = Exception( + 'Univariate drift error' + ) + + reference_data = DataFrame( + { + 'timestamp': ['2023-05-26 11:12:27'], + 'target': [1.0], + 'feature1': [1.0], + } + ) + + target_data = DataFrame( + { + 'timestamp': ['2023-05-26 11:12:27'], + 'target': [1.0], + 'feature1': [1.0], + } + ) + reference_columns = reference_data.drop( columns=['target', 'timestamp'], errors='ignore' ).columns @@ -629,28 +734,26 @@ async def test_get_drift_metrics_univariate_error( except Exception as e: assert str(e) == 'Univariate drift error' model_metrics_activity.error.assert_called_once_with( - 'Error detecting univariate drift: Univariate drift error', - metadata['metadata'] + 'Error detecting univariate drift: Univariate drift error', metadata['metadata'] ) model_metrics_activity.emit_metric.assert_called_with( - metric_object=mock_metrics.MODEL_ANALYZE_ERROR_COUNT, - tags=ANY + metric_object=mock_metrics.MODEL_ANALYZE_ERROR_COUNT, tags=ANY ) else: raise AssertionError('Expected Exception') @mark.asyncio -async def test_calculate_simple_metrics_success_all_metrics( - model_metrics_activity -): +async def test_calculate_simple_metrics_success_all_metrics(model_metrics_activity): # Arrange - target_data = DataFrame({ - 'timestamp': ['2023-05-26 11:12:27', '2023-05-26 11:12:28', '2023-05-26 11:12:29'], - 'target': [1.0, 2.0, 3.0], - 'prediction': [1.1, 2.1, 2.9], - }) - + target_data = DataFrame( + { + 'timestamp': ['2023-05-26 11:12:27', '2023-05-26 11:12:28', '2023-05-26 11:12:29'], + 'target': [1.0, 2.0, 3.0], + 'prediction': [1.1, 2.1, 2.9], + } + ) + input_data = { **metadata, 'model_id': 'test_model_id', @@ -673,23 +776,23 @@ async def test_calculate_simple_metrics_success_all_metrics( assert all(data_size == 3 for data_size in result['data_size'].values) assert all(interval_minutes == 5 for interval_minutes in result['interval_minutes'].values) model_metrics_activity.info.assert_called_once_with( - 'Calculating simple metrics for model test_model_id: [\'rmse\', \'mse\', \'mae\', \'r2\']', - metadata['metadata'] + "Calculating simple metrics for model test_model_id: ['rmse', 'mse', 'mae', 'r2']", + metadata['metadata'], ) model_metrics_activity.debug.assert_called_once() @mark.asyncio -async def test_calculate_simple_metrics_success_rmse_only( - model_metrics_activity -): +async def test_calculate_simple_metrics_success_rmse_only(model_metrics_activity): # Arrange - target_data = DataFrame({ - 'timestamp': ['2023-05-26 11:12:27', '2023-05-26 11:12:28'], - 'target': [1.0, 2.0], - 'prediction': [1.1, 2.1], - }) - + target_data = DataFrame( + { + 'timestamp': ['2023-05-26 11:12:27', '2023-05-26 11:12:28'], + 'target': [1.0, 2.0], + 'prediction': [1.1, 2.1], + } + ) + input_data = { **metadata, 'model_id': 'test_model_id', @@ -709,22 +812,21 @@ async def test_calculate_simple_metrics_success_rmse_only( assert result['data_size'].values[0] == 2 assert result['interval_minutes'].values[0] == 5 model_metrics_activity.info.assert_called_once_with( - 'Calculating simple metrics for model test_model_id: [\'rmse\']', - metadata['metadata'] + "Calculating simple metrics for model test_model_id: ['rmse']", metadata['metadata'] ) @mark.asyncio -async def test_calculate_simple_metrics_success_mse_only( - model_metrics_activity -): +async def test_calculate_simple_metrics_success_mse_only(model_metrics_activity): # Arrange - target_data = DataFrame({ - 'timestamp': ['2023-05-26 11:12:27', '2023-05-26 11:12:28'], - 'target': [1.0, 2.0], - 'prediction': [1.1, 2.1], - }) - + target_data = DataFrame( + { + 'timestamp': ['2023-05-26 11:12:27', '2023-05-26 11:12:28'], + 'target': [1.0, 2.0], + 'prediction': [1.1, 2.1], + } + ) + input_data = { **metadata, 'model_id': 'test_model_id', @@ -744,22 +846,21 @@ async def test_calculate_simple_metrics_success_mse_only( assert result['data_size'].values[0] == 2 assert result['interval_minutes'].values[0] == 5 model_metrics_activity.info.assert_called_once_with( - 'Calculating simple metrics for model test_model_id: [\'mse\']', - metadata['metadata'] + "Calculating simple metrics for model test_model_id: ['mse']", metadata['metadata'] ) @mark.asyncio -async def test_calculate_simple_metrics_success_mae_only( - model_metrics_activity -): +async def test_calculate_simple_metrics_success_mae_only(model_metrics_activity): # Arrange - target_data = DataFrame({ - 'timestamp': ['2023-05-26 11:12:27', '2023-05-26 11:12:28'], - 'target': [1.0, 2.0], - 'prediction': [1.1, 2.1], - }) - + target_data = DataFrame( + { + 'timestamp': ['2023-05-26 11:12:27', '2023-05-26 11:12:28'], + 'target': [1.0, 2.0], + 'prediction': [1.1, 2.1], + } + ) + input_data = { **metadata, 'model_id': 'test_model_id', @@ -779,22 +880,21 @@ async def test_calculate_simple_metrics_success_mae_only( assert result['data_size'].values[0] == 2 assert result['interval_minutes'].values[0] == 5 model_metrics_activity.info.assert_called_once_with( - 'Calculating simple metrics for model test_model_id: [\'mae\']', - metadata['metadata'] + "Calculating simple metrics for model test_model_id: ['mae']", metadata['metadata'] ) @mark.asyncio -async def test_calculate_simple_metrics_success_r2_only( - model_metrics_activity -): +async def test_calculate_simple_metrics_success_r2_only(model_metrics_activity): # Arrange - target_data = DataFrame({ - 'timestamp': ['2023-05-26 11:12:27', '2023-05-26 11:12:28'], - 'target': [1.0, 2.0], - 'prediction': [1.1, 2.1], - }) - + target_data = DataFrame( + { + 'timestamp': ['2023-05-26 11:12:27', '2023-05-26 11:12:28'], + 'target': [1.0, 2.0], + 'prediction': [1.1, 2.1], + } + ) + input_data = { **metadata, 'model_id': 'test_model_id', @@ -814,23 +914,22 @@ async def test_calculate_simple_metrics_success_r2_only( assert result['data_size'].values[0] == 2 assert result['interval_minutes'].values[0] == 5 model_metrics_activity.info.assert_called_once_with( - 'Calculating simple metrics for model test_model_id: [\'r2\']', - metadata['metadata'] + "Calculating simple metrics for model test_model_id: ['r2']", metadata['metadata'] ) @mark.asyncio -async def test_calculate_simple_metrics_r2_zero_ss_tot( - model_metrics_activity -): +async def test_calculate_simple_metrics_r2_zero_ss_tot(model_metrics_activity): # Arrange # All target values are the same, so ss_tot will be 0 - target_data = DataFrame({ - 'timestamp': ['2023-05-26 11:12:27', '2023-05-26 11:12:28'], - 'target': [1.0, 1.0], - 'prediction': [1.1, 1.1], - }) - + target_data = DataFrame( + { + 'timestamp': ['2023-05-26 11:12:27', '2023-05-26 11:12:28'], + 'target': [1.0, 1.0], + 'prediction': [1.1, 1.1], + } + ) + input_data = { **metadata, 'model_id': 'test_model_id', @@ -851,22 +950,21 @@ async def test_calculate_simple_metrics_r2_zero_ss_tot( assert result['data_size'].values[0] == 2 assert result['interval_minutes'].values[0] == 5 model_metrics_activity.info.assert_called_once_with( - 'Calculating simple metrics for model test_model_id: [\'r2\']', - metadata['metadata'] + "Calculating simple metrics for model test_model_id: ['r2']", metadata['metadata'] ) @mark.asyncio -async def test_calculate_simple_metrics_success_multiple_metrics_subset( - model_metrics_activity -): +async def test_calculate_simple_metrics_success_multiple_metrics_subset(model_metrics_activity): # Arrange - target_data = DataFrame({ - 'timestamp': ['2023-05-26 11:12:27', '2023-05-26 11:12:28', '2023-05-26 11:12:29'], - 'target': [1.0, 2.0, 3.0], - 'prediction': [1.1, 2.1, 2.9], - }) - + target_data = DataFrame( + { + 'timestamp': ['2023-05-26 11:12:27', '2023-05-26 11:12:28', '2023-05-26 11:12:29'], + 'target': [1.0, 2.0, 3.0], + 'prediction': [1.1, 2.1, 2.9], + } + ) + input_data = { **metadata, 'model_id': 'test_model_id', @@ -887,7 +985,5 @@ async def test_calculate_simple_metrics_success_multiple_metrics_subset( assert all(data_size == 3 for data_size in result['data_size'].values) assert all(interval_minutes == 5 for interval_minutes in result['interval_minutes'].values) model_metrics_activity.info.assert_called_once_with( - 'Calculating simple metrics for model test_model_id: [\'rmse\', \'mae\']', - metadata['metadata'] + "Calculating simple metrics for model test_model_id: ['rmse', 'mae']", metadata['metadata'] ) - diff --git a/tests/laborious/utils/repository/test_model_repository.py b/tests/laborious/utils/repository/test_model_repository.py index 9d79f70..40782fe 100644 --- a/tests/laborious/utils/repository/test_model_repository.py +++ b/tests/laborious/utils/repository/test_model_repository.py @@ -193,9 +193,11 @@ def test_get_model_params(mlflow, mlflow_repository): def test_check_artifact_exists_true(mlflow_repository): artifact = MagicMock(path='test_artifact') mlflow_repository.client.list_artifacts.return_value = [artifact] - - result = mlflow_repository.check_artifact_exists('run_id', 'test_artifact', metadata['metadata']) - + + result = mlflow_repository.check_artifact_exists( + 'run_id', 'test_artifact', metadata['metadata'] + ) + assert result is True mlflow_repository.client.list_artifacts.assert_called_once_with('run_id') @@ -203,9 +205,11 @@ def test_check_artifact_exists_true(mlflow_repository): def test_check_artifact_exists_false(mlflow_repository): artifact = MagicMock(path='other_artifact') mlflow_repository.client.list_artifacts.return_value = [artifact] - - result = mlflow_repository.check_artifact_exists('run_id', 'test_artifact', metadata['metadata']) - + + result = mlflow_repository.check_artifact_exists( + 'run_id', 'test_artifact', metadata['metadata'] + ) + assert result is False mlflow_repository.client.list_artifacts.assert_called_once_with('run_id') @@ -299,16 +303,22 @@ async def test_download_artifacts_error(makedirs, rmtree, path, mlflow_repositor @patch('laborious.utils.repository.model_repository.mlflow') @patch('laborious.utils.repository.model_repository.pd') @patch('laborious.utils.repository.model_repository.StringIO') -async def test_load_artifact_dataframe_success(StringIO, pd, mlflow, mlflow_repository): +async def test_load_artifact_dataframe_success(_stringio, pd, mlflow, mlflow_repository): mlflow_repository.get_model_run_id = MagicMock(return_value='run_id') mlflow_repository.check_artifact_exists = MagicMock(return_value=True) - + mlflow.artifacts.load_text.return_value = 'col1,col2\n1,2\n3,4' - - result = await mlflow_repository.load_artifact_dataframe('model_name', 'artifact_path', metadata['metadata']) - - mlflow_repository.get_model_run_id.assert_called_once_with(model_name='model_name', stage='Production') - mlflow_repository.check_artifact_exists.assert_called_once_with('run_id', 'artifact_path', metadata['metadata']) + + result = await mlflow_repository.load_artifact_dataframe( + 'model_name', 'artifact_path', metadata['metadata'] + ) + + mlflow_repository.get_model_run_id.assert_called_once_with( + model_name='model_name', stage='Production' + ) + mlflow_repository.check_artifact_exists.assert_called_once_with( + 'run_id', 'artifact_path', metadata['metadata'] + ) mlflow.artifacts.load_text.assert_called_once_with('runs:/run_id/artifact_path') assert result == pd.read_csv.return_value mlflow_repository.observe_lag.assert_called_once_with(ANY, metrics.MODEL_READ_LAG, ANY) @@ -321,12 +331,18 @@ async def test_load_artifact_dataframe_success(StringIO, pd, mlflow, mlflow_repo async def test_load_artifact_dataframe_not_exists(mlflow_repository): mlflow_repository.get_model_run_id = MagicMock(return_value='run_id') mlflow_repository.check_artifact_exists = MagicMock(return_value=False) - - result = await mlflow_repository.load_artifact_dataframe('model_name', 'artifact_path', metadata['metadata']) - + + result = await mlflow_repository.load_artifact_dataframe( + 'model_name', 'artifact_path', metadata['metadata'] + ) + assert result is None - mlflow_repository.get_model_run_id.assert_called_once_with(model_name='model_name', stage='Production') - mlflow_repository.check_artifact_exists.assert_called_once_with('run_id', 'artifact_path', metadata['metadata']) + mlflow_repository.get_model_run_id.assert_called_once_with( + model_name='model_name', stage='Production' + ) + mlflow_repository.check_artifact_exists.assert_called_once_with( + 'run_id', 'artifact_path', metadata['metadata'] + ) @pytest.mark.asyncio @@ -335,10 +351,12 @@ async def test_load_artifact_dataframe_error(mlflow, mlflow_repository): mlflow_repository.get_model_run_id = MagicMock(return_value='run_id') mlflow_repository.check_artifact_exists = MagicMock(return_value=True) mlflow.artifacts.load_text.side_effect = ValueError('error') - + with pytest.raises(ValueError): - await mlflow_repository.load_artifact_dataframe('model_name', 'artifact_path', metadata['metadata']) - + await mlflow_repository.load_artifact_dataframe( + 'model_name', 'artifact_path', metadata['metadata'] + ) + mlflow_repository.emit_metric.assert_called_once_with( metric_object=metrics.MODEL_READ_ERROR_COUNT, tags=ANY ) @@ -1059,7 +1077,9 @@ async def test_create_new_experiment( data.to_csv.assert_called_once_with('./tmp/artifacts/model_name/retrain_data.csv', index=False) # Verify prediction_data.to_csv was called with correct arguments prediction_data.to_csv.assert_called_once() - assert prediction_data.to_csv.call_args[0][0] == './tmp/artifacts/model_name/evaluation_data.csv' + assert ( + prediction_data.to_csv.call_args[0][0] == './tmp/artifacts/model_name/evaluation_data.csv' + ) assert prediction_data.to_csv.call_args[1]['index'] is False mlflow.start_run.assert_called_once_with( @@ -1089,10 +1109,12 @@ async def test_create_new_experiment( } ) - mlflow.log_artifact.assert_has_calls([ - call('./tmp/artifacts/model_name/retrain_data.csv'), - call('./tmp/artifacts/model_name/evaluation_data.csv'), - ]) + mlflow.log_artifact.assert_has_calls( + [ + call('./tmp/artifacts/model_name/retrain_data.csv'), + call('./tmp/artifacts/model_name/evaluation_data.csv'), + ] + ) force_memory_release.assert_called_once_with(mlflow_repository.logger) @@ -1472,9 +1494,9 @@ def test_get_prediction_data_dataframe(mlflow_repository): retrain_dataset = DataFrame({'feat_1': [1, 2], 'target': [3, 4]}, index=['idx1', 'idx2']) prediction_model.predict.return_value = DataFrame({'pred': [5, 6]}, index=['idx1', 'idx2']) target_name = 'target' - + result = mlflow_repository.get_prediction_data(prediction_model, retrain_dataset, target_name) - + prediction_model.predict.assert_called_once_with(retrain_dataset) assert 'prediction' in result.columns assert 'target' in result.columns @@ -1487,9 +1509,9 @@ def test_get_prediction_data_array(mlflow_repository): retrain_dataset = DataFrame({'feat_1': [1, 2], 'target': [3, 4]}, index=['idx1', 'idx2']) prediction_model.predict.return_value = [5, 6] target_name = 'target' - + result = mlflow_repository.get_prediction_data(prediction_model, retrain_dataset, target_name) - + prediction_model.predict.assert_called_once_with(retrain_dataset) assert 'prediction' in result.columns assert 'target' in result.columns diff --git a/tests/laborious/workflows/test_drift.py b/tests/laborious/workflows/test_drift.py index 2f3f1dc..25870ae 100644 --- a/tests/laborious/workflows/test_drift.py +++ b/tests/laborious/workflows/test_drift.py @@ -1,10 +1,10 @@ from unittest.mock import ANY, AsyncMock, call, patch from pytest import fixture, mark +from sientia_do.temporal.constants import DATETIME_FORMAT_WITH_TZ from laborious.activities.activities import Activities from laborious.workflows.drift import Drift -from sientia_do.temporal.constants import DATETIME_FORMAT_WITH_TZ @fixture @@ -45,9 +45,7 @@ async def test_run(workflow_mock: AsyncMock, drift: Drift): reference_data = {'data': 'test_reference_data'} drift_data = {'drift': 'test_drift_data'} - workflow_mock.start_local_activity_method.side_effect = [ - target_data, reference_data - ] + workflow_mock.start_local_activity_method.side_effect = [target_data, reference_data] workflow_mock.execute_local_activity_method.return_value = drift_data workflow_mock.execute_activity_method = AsyncMock() @@ -56,12 +54,13 @@ async def test_run(workflow_mock: AsyncMock, drift: Drift): await drift.run(input_data) # Assert - Check start_local_activity_method calls + # Query format matches psycopg2.sql output (identifiers with double quotes, literals with single quotes) expected_gathering_query = f""" SELECT * - FROM {input_data['schema']}.{input_data['source_table_name']} + FROM "{input_data['schema']}"."{input_data['source_table_name']}" WHERE - model_id = {input_data['model_id']} AND - timestamp > NOW() - INTERVAL '{input_data['interval']} minutes' + model_id = '{input_data['model_id']}' AND + timestamp > NOW() - INTERVAL {input_data['interval']} minutes ORDER BY timestamp ASC """ @@ -73,6 +72,7 @@ async def test_run(workflow_mock: AsyncMock, drift: Drift): **metadata, 'query': expected_gathering_query, 'datetime_columns': ['timestamp', 'created_at'], + 'orient': 'records', }, retry_policy=ANY, start_to_close_timeout=ANY, @@ -250,4 +250,3 @@ async def test_run_default_chunk_period(workflow_mock: AsyncMock, drift: Drift): retry_policy=ANY, start_to_close_timeout=ANY, ) - diff --git a/tests/laborious/workflows/test_simple_metrics.py b/tests/laborious/workflows/test_simple_metrics.py index b185b1a..12c2cac 100644 --- a/tests/laborious/workflows/test_simple_metrics.py +++ b/tests/laborious/workflows/test_simple_metrics.py @@ -1,10 +1,10 @@ from unittest.mock import ANY, AsyncMock, call, patch from pytest import fixture, mark +from sientia_do.temporal.constants import DATETIME_FORMAT_WITH_TZ from laborious.activities.activities import Activities from laborious.workflows.simple_metrics import SimpleMetrics -from sientia_do.temporal.constants import DATETIME_FORMAT_WITH_TZ @fixture @@ -42,9 +42,7 @@ async def test_run(workflow_mock: AsyncMock, simple_metrics: SimpleMetrics): target_data = {'data': 'test_target_data'} simple_metrics_data = {'metrics': 'test_simple_metrics_data'} - workflow_mock.execute_local_activity_method.side_effect = [ - target_data, simple_metrics_data - ] + workflow_mock.execute_local_activity_method.side_effect = [target_data, simple_metrics_data] workflow_mock.execute_activity_method = AsyncMock() @@ -52,17 +50,18 @@ async def test_run(workflow_mock: AsyncMock, simple_metrics: SimpleMetrics): await simple_metrics.run(input_data) # Assert - Check load_custom_query call + # Query format matches psycopg2.sql output (identifiers with double quotes, literals with single quotes) expected_query = f""" select p."timestamp", p.prediction, ld.value as "target" - from {input_data['schema']}.{input_data['predictions_table_name']} p - inner join {input_data['schema']}.{input_data['data_table_name']} ld + from "{input_data['schema']}"."{input_data['predictions_table_name']}" p + inner join "{input_data['schema']}"."{input_data['data_table_name']}" ld on p."timestamp" = ld."timestamp" where - p.model_id = {input_data['model_id']} and + p.model_id = '{input_data['model_id']}' and p.prediction is not null and ld.variable = '{input_data['model_config']['target']}' and ld.value is not null and - p."timestamp" >= NOW() - INTERVAL '{input_data['interval_minutes']} minutes' + p."timestamp" >= NOW() - INTERVAL {input_data['interval_minutes']} minutes order by p."timestamp" desc; """ @@ -75,6 +74,7 @@ async def test_run(workflow_mock: AsyncMock, simple_metrics: SimpleMetrics): **metadata, 'query': expected_query, 'datetime_columns': ['timestamp'], + 'orient': 'records', }, retry_policy=ANY, start_to_close_timeout=ANY, @@ -159,9 +159,7 @@ async def test_run_empty_simple_metrics(workflow_mock: AsyncMock, simple_metrics target_data = {'data': 'test_target_data'} simple_metrics_data = None - workflow_mock.execute_local_activity_method.side_effect = [ - target_data, simple_metrics_data - ] + workflow_mock.execute_local_activity_method.side_effect = [target_data, simple_metrics_data] workflow_mock.execute_activity_method = AsyncMock() @@ -193,9 +191,7 @@ async def test_run_default_metrics(workflow_mock: AsyncMock, simple_metrics: Sim target_data = {'data': 'test_target_data'} simple_metrics_data = {'metrics': 'test_simple_metrics_data'} - workflow_mock.execute_local_activity_method.side_effect = [ - target_data, simple_metrics_data - ] + workflow_mock.execute_local_activity_method.side_effect = [target_data, simple_metrics_data] workflow_mock.execute_activity_method = AsyncMock() @@ -225,4 +221,3 @@ async def test_run_default_metrics(workflow_mock: AsyncMock, simple_metrics: Sim ), ] ) - diff --git a/validate.sh b/validate.sh index c7c44fe..87a8df8 100755 --- a/validate.sh +++ b/validate.sh @@ -87,7 +87,7 @@ if ! run_step "3. Type Checking (mypy)" "mypy laborious/"; then fi # Step 4: Security Analysis (Bandit) -if ! run_step "4. Security Analysis (Bandit)" "bandit -r laborious/ -ll -q"; then +if ! run_step "4. Security Analysis (Bandit)" "bandit -c pyproject.toml -r laborious/ -ll -q"; then FAILED_STEPS+=("Security Analysis") fi From 07dc61211613f46ce660e4a4af3e9392612d2695 Mon Sep 17 00:00:00 2001 From: vitor-aignosi Date: Mon, 17 Nov 2025 16:52:33 -0300 Subject: [PATCH 18/27] SIENTIAPDE-1273 Enhance data handling and export processes in Laborious workflows - Updated `gates.py` to improve data quality validation, filtering, and formatting operations, including enhanced metrics recording. - Refined `mlflow.py` to better manage model transformations and reference data retrieval from MLflow Model Registry. - Enhanced `format_and_export_prediction.py` to support separate export of transformed data, improving flexibility in data handling. - Added comprehensive test coverage for new functionalities, including transformed data formatting and retrain report generation. - Improved documentation in `README.md` to reflect changes in activities and workflows, ensuring clarity on data processing and export paths. --- README.md | 53 +++++-- laborious/activities/gates.py | 67 +++++++- laborious/activities/mlflow.py | 18 ++- .../format_and_export_prediction.py | 28 +++- tests/laborious/activities/test_gates.py | 146 +++++++++++++++++ tests/laborious/activities/test_mlflow.py | 89 +++++++++++ .../test_format_and_export_prediction.py | 150 ++++++++++++++++++ 7 files changed, 531 insertions(+), 20 deletions(-) diff --git a/README.md b/README.md index 0143cbd..a34b939 100644 --- a/README.md +++ b/README.md @@ -130,8 +130,14 @@ Laborious uses a Temporal-based architecture with strong separation of concerns - `minimal_retrain.py`: Automated model retraining and production update #### **Activities (`laborious/activities/`)** -- `gates.py`: Data quality validation and filtering -- `mlflow.py`: Transform and predict operations +- `gates.py`: Data quality validation, filtering, and data formatting operations + - Input/response/content gates for quality validation + - Prediction and transformed data formatting + - Retrain report formatting and metrics recording +- `mlflow.py`: Transform, predict, and model management operations + - MLFlow model transformation and prediction + - Model retraining and production updates + - Reference data retrieval from MLflow Model Registry - `opc.py`: OPC UA export to industrial systems (optional) - `activities.py`: Aggregates activity interfaces @@ -146,7 +152,9 @@ Laborious uses a Temporal-based architecture with strong separation of concerns #### **1. Batch Prediction Pipeline** ``` Input Data (PostgreSQL) → Data Quality Gates → MLFlow Transform → -MLFlow Prediction → Response Validation → Export (PostgreSQL [+ OPC]) +MLFlow Prediction → Response Validation → Format & Export + ├─→ Predictions → PostgreSQL [+ OPC] + └─→ Transformed Data → PostgreSQL (optional) ``` #### **2. Model Retraining Pipeline** @@ -318,27 +326,37 @@ The **FormatAndExportPrediction** workflow handles prediction data formatting an #### Execution Flow 1. **Path Decision**: Determines formatting path based on configuration 2. **Data Formatting**: Formats prediction data for specific output requirements -3. **PostgreSQL Export**: Writes formatted predictions to database -4. **OPC Export**: Writes predictions to OPC servers -5. **Metrics Recording**: Records export performance and success metrics +3. **Transformed Data Processing**: Optionally formats and exports transformed data separately +4. **PostgreSQL Export**: Writes formatted predictions to database +5. **OPC Export**: Writes predictions to OPC servers +6. **Metrics Recording**: Records export performance and success metrics #### Key Features - **Flexible Formatting**: Configurable output formats for different destinations - **Multi-Destination Export**: PostgreSQL and OPC server integration +- **Transformed Data Export**: Optional separate export of MLFlow transformed data - **Performance Monitoring**: Comprehensive metrics for export operations - **Error Handling**: Robust error handling with notification integration #### Architecture Diagram ```mermaid flowchart LR - A[1. format_prediction/format_default_prediction] --> B[2. write_opc_data] --> C[3. export_data_to_postgres] --> D[4. write_metrics] + A[1. format_prediction/format_default_prediction] --> B[2. format_transformed_data] --> C[3. write_opc_data] --> D[4. export_data_to_postgres] --> E[5. write_metrics] A -.-> Format[Data Formatting] - B -.-> OPC[OPC Servers] - C -.-> PostgreSQL[(PostgreSQL)] - D -.-> Prometheus[Prometheus] + B -.-> Transform[Transformed Data] + C -.-> OPC[OPC Servers] + D -.-> PostgreSQL[(PostgreSQL)] + E -.-> Prometheus[Prometheus] ``` +#### Transformed Data Export +When `transformed_data` is provided in the input, the workflow will: +- Format the transformed data using `format_transformed_data` activity +- Export it to a separate table (`transform_table_name`) asynchronously +- Wait for both prediction and transformed data exports to complete +- This enables separate tracking of model transformations for analysis and debugging + ### 4. Minimal Retrain Workflow (`minimal_retrain.py`) The **MinimalRetrain** workflow handles automated model retraining and production model updates. @@ -579,11 +597,24 @@ The workflow at `.github/workflows/quality-gate.yml` executes validations on eac ``` tests/ ├── activities/ # Activity implementation tests -├── workflow/ # Workflow orchestration tests +│ ├── test_gates.py # Data quality gates and formatting tests +│ ├── test_mlflow.py # MLFlow operations and reference data tests +│ └── ... # Other activity tests +├── workflows/ # Workflow orchestration tests +│ └── subworkflows/ # Sub-workflow tests +│ └── test_format_and_export_prediction.py # Export workflow tests ├── utils/ # Utility function tests └── integration/ # End-to-end workflow tests ``` +### Test Coverage +The test suite provides comprehensive coverage for: +- **Data Quality Gates**: Input, response, and content validation filters +- **Data Formatting**: Prediction, transformed data, and retrain report formatting +- **MLFlow Operations**: Transform, predict, retrain, and reference data retrieval +- **Workflow Orchestration**: Complete workflow execution paths and error handling +- **Metrics Recording**: Performance monitoring and OPC export metrics + ### Test Execution ```bash # Install test dependencies diff --git a/laborious/activities/gates.py b/laborious/activities/gates.py index 13072cf..f0f8474 100644 --- a/laborious/activities/gates.py +++ b/laborious/activities/gates.py @@ -405,7 +405,32 @@ class Gates(SientiaMonitoring): @activity.defn(name='format_transformed_data') async def format_transformed_data(self, input_data: dict[str, Any]) -> dict: """ - Format transformed data according to configured storage policies. + Format transformed data for storage and export operations. + + This method formats transformed data from MLFlow model transformations + into a standardized format suitable for database storage. It converts + wide-format data (columns as variables) into long-format (melted) + with proper timestamp handling and model identification. + + The formatting process includes: + 1. Converting input data dictionary to DataFrame + 2. Extracting timestamps from DataFrame index + 3. Resetting index to create sequential row numbers + 4. Melting data from wide format to long format (variable-value pairs) + 5. Adding model_id for data lineage tracking + + Args: + input_data (dict): Input data containing: + - metadata (dict): Workflow execution metadata + - data (dict[str, Any]): Transformed data to format (DataFrame-compatible dict) + - model_id (str): Unique identifier for the ML model + + Returns: + dict: Formatted data dictionary with keys: + - timestamp (dict): Timestamp values indexed by row number + - variable (dict): Variable names indexed by row number + - value (dict): Variable values indexed by row number + - model_id (dict): Model identifiers indexed by row number """ metadata = input_data['metadata'] @@ -544,7 +569,45 @@ class Gates(SientiaMonitoring): @activity.defn(name='format_retrain_report') async def format_retrain_report(self, input_data: dict[str, Any]) -> dict: """ - Format retrain report data according to configured storage policies. + Format retrain report data for storage and audit trail maintenance. + + This method formats model retraining operation results into a standardized + report format suitable for database storage and operational monitoring. + It captures retraining status, timestamps, and model version information + for comprehensive audit trails and operational visibility. + + The formatting process includes: + 1. Extracting retraining experiment response data + 2. Capturing model update report information (version, MLflow IDs) + 3. Formatting timestamps and status information + 4. Conditionally including version information for successful retrains + + Args: + input_data (dict): Input data containing: + - metadata (dict): Workflow execution metadata + - experiment_response (dict): Retraining experiment response containing: + - success (bool): Retraining operation success status + - timestamp (str): Timestamp of the retraining operation + - message (str): Status message or error description + - update_report (dict): Model update report containing: + - version (str): New model version identifier + - mlflow_run_id (str): MLflow run identifier + - mlflow_experiment_id (str): MLflow experiment identifier + - model_id (str): Unique identifier for the ML model + - model_name (str): Name of the ML model + + Returns: + dict: Formatted retrain report dictionary with keys: + - model_id (dict): Model identifiers indexed by row number + - model_name (dict): Model names indexed by row number + - timestamp (dict): Retraining timestamps indexed by row number + - status (dict): Retraining status messages indexed by row number + - version (dict, optional): Model versions indexed by row number + Only included if experiment_response['success'] is True + - mlflow_run_id (dict, optional): MLflow run IDs indexed by row number + Only included if experiment_response['success'] is True + - mlflow_experiment_id (dict, optional): MLflow experiment IDs indexed by row number + Only included if experiment_response['success'] is True """ metadata = input_data['metadata'] self.info('Formatting retrain report...', metadata) diff --git a/laborious/activities/mlflow.py b/laborious/activities/mlflow.py index 6eadf60..3970bb7 100644 --- a/laborious/activities/mlflow.py +++ b/laborious/activities/mlflow.py @@ -424,14 +424,30 @@ class MLFlow(SientiaMonitoring): """ Get reference data from the MLflow Model Registry. + This method retrieves evaluation reference data stored as artifacts in the + MLflow Model Registry. The reference data is typically used for model + drift detection, performance comparison, and quality validation. The method + loads the data from a CSV artifact file and formats timestamps for + consistent processing. + + The method handles: + 1. Loading evaluation data artifact from MLflow Model Registry + 2. Timestamp parsing and formatting for consistency + 3. Data conversion to dictionary format for workflow consumption + 4. Graceful handling of missing reference data + Args: input_data (dict): Input data containing: - metadata (dict): Workflow execution metadata - model_name (str): Name of the MLFlow model to get reference data from Returns: - list[dict[Hashable, Any]] | None: Reference data from the MLflow Model Registry. + list[dict[Hashable, Any]] | None: Reference data from the MLflow Model Registry + as a list of dictionaries. Returns None if reference data is not found + or if the artifact does not exist. + Raises: + Exception: If artifact loading fails or encounters errors during processing """ metadata = input_data['metadata'] diff --git a/laborious/workflows/sub_workflows/format_and_export_prediction.py b/laborious/workflows/sub_workflows/format_and_export_prediction.py index 1851c61..e50c2cf 100644 --- a/laborious/workflows/sub_workflows/format_and_export_prediction.py +++ b/laborious/workflows/sub_workflows/format_and_export_prediction.py @@ -50,21 +50,36 @@ class FormatAndExportPrediction: Args: input_data: Complete configuration for the export workflow Required keys: + - metadata (dict): Workflow execution metadata - path_flag (str | None): Decision path flag for formatting strategy + - None: Normal prediction path with full formatting + - Any other value: Default prediction path for error conditions - data (dict[str, Any]): Prediction data to format and export - prediction_confidence (float): Confidence score for the prediction - timestamp (str): ISO-formatted timestamp for the prediction - model_id (int): Unique identifier for the ML model - model_name (str): Name of the ML model - - model_retention (str): Model retention policy configuration - - comment (str): Operational comment or error description - schema (str): Database schema for data storage - table_name (str): Target table for data persistence - opc_output_config (dict[str, Any]): OPC server export configuration - - prediction_store_policy (str, optional): Data retention policy + Optional keys: + - transformed_data (dict[str, Any]): Transformed data to export separately + Only processed when path_flag is None + - transform_table_name (str): Target table for transformed data export + Required if transformed_data is provided + - prediction_store_policy (str): Data retention policy (e.g., 'lts:1', 'erl:2') + Required when path_flag is None + - comment (str): Operational comment or error description + Required when path_flag is not None Returns: - bool: True if the workflow completes successfully, False otherwise + None: The workflow completes successfully when all export operations finish + + Note: + When transformed_data is provided and path_flag is None, the workflow will: + 1. Format the transformed data using format_transformed_data + 2. Export it to a separate table (transform_table_name) asynchronously + 3. Wait for both prediction and transformed data exports to complete """ metadata = input_data['metadata'] path_flag = input_data['path_flag'] @@ -73,7 +88,7 @@ class FormatAndExportPrediction: prediction_confidence = input_data['prediction_confidence'] if path_flag is None: - # proceed with formatting and exporting + # Normal prediction path: format prediction data with full metadata prediction = await workflow.execute_local_activity_method( Activities.format_prediction, { @@ -88,6 +103,7 @@ class FormatAndExportPrediction: start_to_close_timeout=timedelta(seconds=60), ) + # Optionally format and export transformed data to separate table if transformed_data is not None: transformed = await workflow.execute_local_activity_method( Activities.format_transformed_data, @@ -120,7 +136,7 @@ class FormatAndExportPrediction: write_transformed_handler = None else: - # create default prediction + # Error path: create default prediction with error indicators prediction = await workflow.execute_local_activity_method( Activities.format_default_prediction, { diff --git a/tests/laborious/activities/test_gates.py b/tests/laborious/activities/test_gates.py index 8399b43..e33ea4f 100644 --- a/tests/laborious/activities/test_gates.py +++ b/tests/laborious/activities/test_gates.py @@ -546,6 +546,80 @@ async def test_format_prediction_with_timestamp_invalid_policy(gates_activity): raise AssertionError('Expected ValueError') +@mark.asyncio +async def test_format_transformed_data_single_row(gates_activity): + # Arrange + input_data = { + **metadata, + 'data': { + 'var1': {'2023-05-26 11:12:27': 1.0}, + 'var2': {'2023-05-26 11:12:27': 2.0}, + }, + 'model_id': 'test_model', + } + + # Act + result = await gates_activity.format_transformed_data(input_data) + + # Assert + assert result['timestamp'] == {0: '2023-05-26 11:12:27', 1: '2023-05-26 11:12:27'} + assert result['variable'] == {0: 'var1', 1: 'var2'} + assert result['value'] == {0: 1.0, 1: 2.0} + assert result['model_id'] == {0: 'test_model', 1: 'test_model'} + gates_activity.info.assert_called() + + +@mark.asyncio +async def test_format_transformed_data_multiple_rows(gates_activity): + # Arrange + input_data = { + **metadata, + 'data': { + 'var1': { + '2023-05-26 11:12:27': 1.0, + '2023-05-26 11:12:28': 2.0, + }, + 'var2': { + '2023-05-26 11:12:27': 3.0, + '2023-05-26 11:12:28': 4.0, + }, + }, + 'model_id': 'test_model', + } + + # Act + result = await gates_activity.format_transformed_data(input_data) + + # Assert + assert len(result['timestamp']) == 4 + assert len(result['variable']) == 4 + assert len(result['value']) == 4 + assert len(result['model_id']) == 4 + assert all(v == 'test_model' for v in result['model_id'].values()) + assert set(result['variable'].values()) == {'var1', 'var2'} + gates_activity.info.assert_called() + + +@mark.asyncio +async def test_format_transformed_data_empty_data(gates_activity): + # Arrange + input_data = { + **metadata, + 'data': {}, + 'model_id': 'test_model', + } + + # Act + result = await gates_activity.format_transformed_data(input_data) + + # Assert + assert result['timestamp'] == {} + assert result['variable'] == {} + assert result['value'] == {} + assert result['model_id'] == {} + gates_activity.info.assert_called() + + @mark.asyncio async def test_format_default_prediction(gates_activity): # Arrange @@ -603,6 +677,40 @@ async def test_format_retrain_report(gates_activity): assert result['mlflow_experiment_id'] == {0: 'test_mlflow_experiment_id'} +@mark.asyncio +async def test_format_retrain_report_failure(gates_activity): + # Arrange + input_data = { + **metadata, + 'experiment_response': { + 'success': False, + 'timestamp': '2023-05-26 11:12:27', + 'message': 'failure', + }, + 'update_report': { + 'version': '1.0.0', + 'mlflow_run_id': 'test_mlflow_run_id', + 'mlflow_experiment_id': 'test_mlflow_experiment_id', + }, + 'model_id': 'test_model', + 'model_name': 'test_model', + } + + # Act + result = await gates_activity.format_retrain_report(input_data) + + # Assert + assert result['model_id'] == {0: 'test_model'} + assert result['model_name'] == {0: 'test_model'} + assert result['timestamp'] == {0: '2023-05-26 11:12:27'} + assert result['status'] == {0: 'failure'} + assert 'version' not in result + assert 'mlflow_run_id' not in result + assert 'mlflow_experiment_id' not in result + gates_activity.info.assert_called() + gates_activity.debug.assert_called() + + @mark.asyncio async def test_get_last_timestamp_with_data(gates_activity): # Arrange @@ -742,3 +850,41 @@ async def test_write_metrics(mock_metrics, gates_activity): ), ] ) + + +@mark.asyncio +@patch('laborious.activities.gates.metrics') +async def test_write_metrics_with_none_opc_response_time(mock_metrics, gates_activity): + """Test write_metrics method with None response_time in opc_metrics.""" + input_data = { + **metadata, + 'prediction': { + 'prediction': [1], + 'prediction_confidence': [0.9], + 'response_time': [0.1], + }, + 'opc_metrics': {'server1': {'tag1': 0.1, 'tag2': None}}, + } + await gates_activity.write_metrics(input_data) + + # Verify that metrics for tag1 are emitted + gates_activity.emit_metric.assert_any_call( + metric_object=mock_metrics.PREDICTION_OPC_WRITING_RESPONSE_TIME_MONITOR, + method='observe', + tags={ + 'pod_id': gates_activity.pod_id, + 'model_name': metadata['metadata']['model_name'], + 'workflow_name': metadata['metadata']['workflow_name'], + 'opc_server_id': 'server1', + 'tag': 'tag1', + }, + value=0.1, + ) + + # Verify that metrics for tag2 (with None response_time) are NOT emitted + calls = [ + c + for c in gates_activity.emit_metric.call_args_list + if len(c[1].get('tags', {})) > 0 and c[1]['tags'].get('tag') == 'tag2' + ] + assert len(calls) == 0, 'Metrics should not be emitted for None response_time' diff --git a/tests/laborious/activities/test_mlflow.py b/tests/laborious/activities/test_mlflow.py index 93d71ee..dd7a0c5 100644 --- a/tests/laborious/activities/test_mlflow.py +++ b/tests/laborious/activities/test_mlflow.py @@ -76,6 +76,11 @@ def mlflow(mock_minio_repository, mock_mlflow_repository): mlflow.send_notification = MagicMock() mlflow.emit_metric = AsyncMock() mlflow.send_notification_async = AsyncMock() + mlflow.error = MagicMock() + mlflow.debug = MagicMock() + mlflow.info = MagicMock() + mlflow.warning = MagicMock() + mlflow.critical = MagicMock() return mlflow @@ -488,3 +493,87 @@ async def test_update_production_model_error(mlflow): ) else: raise AssertionError('No exception raised') + + +@mark.asyncio +@patch('laborious.activities.mlflow.to_datetime') +async def test_get_reference_data_success(mock_to_datetime, mlflow): + # Arrange + input_data = { + **metadata, + 'model_name': 'test_model', + } + + # Mock reference data DataFrame + mock_reference_data = MagicMock() + mock_reference_data.__getitem__.return_value = MagicMock() + mock_to_datetime.return_value.dt.strftime.return_value = MagicMock() + mock_reference_data.to_dict.return_value = [ + {'timestamp': '2023-05-26 11:12:27', 'value': 1.0}, + {'timestamp': '2023-05-26 11:12:28', 'value': 2.0}, + ] + + mlflow.model_monitoring_repository.load_artifact_dataframe.return_value = mock_reference_data + + # Act + result = await mlflow.get_reference_data(input_data) + + # Assert + mlflow.model_monitoring_repository.load_artifact_dataframe.assert_called_once_with( + model_name='test_model', + artifact_path='evaluation_data.csv', + metadata=metadata['metadata'], + ) + mock_to_datetime.assert_called_once_with(mock_reference_data.__getitem__.return_value) + + mock_reference_data.to_dict.assert_called_once_with(orient='records') + assert result == mock_reference_data.to_dict.return_value + + +@mark.asyncio +async def test_get_reference_data_not_found(mlflow): + # Arrange + input_data = { + **metadata, + 'model_name': 'test_model', + } + + mlflow.model_monitoring_repository.load_artifact_dataframe.return_value = None + + # Act + result = await mlflow.get_reference_data(input_data) + + # Assert + mlflow.model_monitoring_repository.load_artifact_dataframe.assert_called_once_with( + model_name='test_model', + artifact_path='evaluation_data.csv', + metadata=metadata['metadata'], + ) + mlflow.warning.assert_called_once_with( + 'Reference data not found for model test_model', metadata['metadata'] + ) + assert result is None + + +@mark.asyncio +async def test_get_reference_data_exception(mlflow): + # Arrange + input_data = { + **metadata, + 'model_name': 'test_model', + } + + mlflow.model_monitoring_repository.load_artifact_dataframe.side_effect = Exception( + 'Error loading artifact' + ) + + # Act & Assert + with raises(Exception) as e: + await mlflow.get_reference_data(input_data) + + assert str(e.value) == 'Error loading artifact' + mlflow.model_monitoring_repository.load_artifact_dataframe.assert_called_once_with( + model_name='test_model', + artifact_path='evaluation_data.csv', + metadata=metadata['metadata'], + ) 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 ecd3bc2..f5619f3 100644 --- a/tests/laborious/workflows/subworkflows/test_format_and_export_prediction.py +++ b/tests/laborious/workflows/subworkflows/test_format_and_export_prediction.py @@ -125,6 +125,156 @@ async def test_run_none_path_flag(workflow_mock, format_and_export_prediction): assert workflow_mock.execute_local_activity_method.call_count == 1 +@mark.asyncio +@patch( + 'laborious.workflows.sub_workflows.format_and_export_prediction.workflow', + new_callable=AsyncMock, +) +async def test_run_none_path_flag_with_transformed_data( + workflow_mock, format_and_export_prediction +): + # Arrange + input_data = { + 'metadata': metadata, + 'path_flag': None, + 'data': {'test': 'data'}, + 'transformed_data': {'transformed': 'data'}, + 'timestamp': '2021-01-01', + 'model_id': 1, + 'prediction_confidence': 0.9, + 'schema': 'test_schema', + 'table_name': 'test_table', + 'transform_table_name': 'test_transform_table', + 'opc_servers': ['test_server'], + 'opc_output_config': {'test': 'config'}, + 'prediction_store_policy': 'lts:1', + } + + prediction_data = MagicMock() + opc_metrics = MagicMock() + transformed_data = MagicMock() + + workflow_mock.execute_local_activity_method.side_effect = [ + prediction_data, # format_prediction + transformed_data, # format_transformed_data + ] + + write_transformed_handler = AsyncMock() + workflow_mock.start_activity_method.return_value = write_transformed_handler + workflow_mock.execute_activity_method.side_effect = [ + (prediction_data, opc_metrics), # write_opc_data + MagicMock(), # export_data_to_postgres (prediction) + MagicMock(), # write_metrics + ] + + # Act + await format_and_export_prediction.run(input_data) + + # Assert - format_prediction call + 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, + ), + call( + Activities.format_transformed_data, + { + 'data': input_data['transformed_data'], + 'model_id': input_data['model_id'], + **metadata, + }, + retry_policy=ANY, + start_to_close_timeout=ANY, + ), + ] + ) + + # Assert - start_activity_method for transformed data export + workflow_mock.start_activity_method.assert_called_once_with( + Activities.export_data_to_postgres, + { + 'schema': input_data['schema'], + 'table_name': input_data['transform_table_name'], + 'data': transformed_data, + 'timestamp_conversion': { + 'column': 'timestamp', + 'format': DATETIME_FORMAT_WITH_TZ, + }, + **metadata, + }, + retry_policy=ANY, + start_to_close_timeout=ANY, + ) + + # Assert - write_opc_data call + workflow_mock.execute_activity_method.assert_has_calls( + [ + call( + Activities.write_opc_data, + { + 'opc_output_config': input_data['opc_output_config'], + 'data': prediction_data, + **metadata, + }, + retry_policy=ANY, + start_to_close_timeout=ANY, + ) + ] + ) + + # Assert - export_data_to_postgres for prediction call + 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': prediction_data, + 'timestamp_conversion': { + 'column': 'timestamp', + 'format': DATETIME_FORMAT_WITH_TZ, + }, + **metadata, + }, + retry_policy=ANY, + start_to_close_timeout=ANY, + ) + ] + ) + + # Assert - write_metrics call + workflow_mock.execute_activity_method.assert_has_calls( + [ + call( + Activities.write_metrics, + { + **metadata, + 'prediction': prediction_data, + 'opc_metrics': opc_metrics, + }, + retry_policy=ANY, + start_to_close_timeout=ANY, + ) + ] + ) + + # Assert - verify counts + assert workflow_mock.execute_activity_method.call_count == 3 + assert workflow_mock.execute_local_activity_method.call_count == 2 + assert workflow_mock.start_activity_method.call_count == 1 + + @mark.asyncio @patch( 'laborious.workflows.sub_workflows.format_and_export_prediction.workflow', From ffb829e0cb89183dee3db8cc7c8b4fac3552e329 Mon Sep 17 00:00:00 2001 From: vitor-aignosi Date: Tue, 18 Nov 2025 08:49:40 -0300 Subject: [PATCH 19/27] SIENTIAPDE-1273 Refactor data handling in Laborious workflows to enhance clarity and compatibility - Improved data validation and formatting in `gates.py`. - Streamlined model transformation management in `mlflow.py`. - Enhanced export functionality in `format_and_export_prediction.py` for better data handling flexibility. - Expanded test coverage for new features and improved documentation in `README.md`. --- tests/conftest.py | 15 +++++++++++++++ 1 file changed, 15 insertions(+) create mode 100644 tests/conftest.py diff --git a/tests/conftest.py b/tests/conftest.py new file mode 100644 index 0000000..78aec93 --- /dev/null +++ b/tests/conftest.py @@ -0,0 +1,15 @@ +""" +Pytest configuration file with global mocks for external dependencies. + +This module mocks the 'sientia' module to avoid requiring its installation +during unit tests. The mock is registered in sys.modules before any test +imports are executed. +""" +import sys +from unittest.mock import MagicMock + +# Mock sientia module +sientia_mock = MagicMock() +sientia_mock.ModelAnalysis = MagicMock +sys.modules['sientia'] = sientia_mock +sys.modules['sientia.ModelAnalysis'] = MagicMock() From 6ae83d90bee331e9c33d852e448a8fc3c9907b65 Mon Sep 17 00:00:00 2001 From: vitor-aignosi Date: Tue, 18 Nov 2025 08:59:43 -0300 Subject: [PATCH 20/27] SIENTIAPDE-1273 Update conftest.py to add a blank line for improved readability and code style consistency. --- tests/conftest.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/conftest.py b/tests/conftest.py index 78aec93..c0fad01 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -5,6 +5,7 @@ This module mocks the 'sientia' module to avoid requiring its installation during unit tests. The mock is registered in sys.modules before any test imports are executed. """ + import sys from unittest.mock import MagicMock From 388685df6e97b3abafd71f08ca6ff0b41a124e47 Mon Sep 17 00:00:00 2001 From: vitor-aignosi Date: Tue, 18 Nov 2025 09:16:42 -0300 Subject: [PATCH 21/27] SIENTIAPDE-1273 Update image tag to 1.1.2 in values.yaml and refactor prediction data sorting in model_repository.py for improved clarity and consistency. --- laborious/utils/repository/model_repository.py | 4 ++-- values.yaml | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/laborious/utils/repository/model_repository.py b/laborious/utils/repository/model_repository.py index 33366e6..395865e 100644 --- a/laborious/utils/repository/model_repository.py +++ b/laborious/utils/repository/model_repository.py @@ -723,7 +723,7 @@ class MLFlowRepository(SientiaMonitoring): prediction_data.index = input_index # Merge prediction data with retrain_dataset on index - prediction_data = pd.merge( + prediction_data = pd.merge( # NOSONAR retrain_dataset, prediction_data, left_index=True, right_index=True, how='left' ) @@ -734,7 +734,7 @@ class MLFlowRepository(SientiaMonitoring): prediction_data.reset_index(drop=True, inplace=True) - prediction_data.sort_values(by='timestamp', ascending=True, inplace=True) + prediction_data = prediction_data.sort_values(by='timestamp', ascending=True) return prediction_data diff --git a/values.yaml b/values.yaml index 5034e12..64988ab 100644 --- a/values.yaml +++ b/values.yaml @@ -11,7 +11,7 @@ image: # This sets the pull policy for images. pullPolicy: Always # Overrides the image tag whose default is the chart appVersion. - tag: "1.1.1" + tag: "1.1.2" 0# This is for the secrets for pulling an image from a private repository more information can be found here: https://kubernetes.io/docs/tasks/configure-pod-container/pull-image-private-registry/ imagePullSecrets: From 7e4e048ca3756374caca045aa5f0dafa240dba59 Mon Sep 17 00:00:00 2001 From: vitor-aignosi Date: Wed, 19 Nov 2025 08:47:48 -0300 Subject: [PATCH 22/27] SIENTIAPDE-1273 Update sientia-mlops-library dependency to version 0.40.5 and fix SQL interval formatting in Drift and SimpleMetrics workflows for improved query accuracy. --- laborious/workflows/drift.py | 2 +- laborious/workflows/simple_metrics.py | 2 +- requirements.txt | 2 +- tests/laborious/workflows/test_drift.py | 2 +- tests/laborious/workflows/test_simple_metrics.py | 2 +- 5 files changed, 5 insertions(+), 5 deletions(-) diff --git a/laborious/workflows/drift.py b/laborious/workflows/drift.py index 85780f7..769c613 100644 --- a/laborious/workflows/drift.py +++ b/laborious/workflows/drift.py @@ -42,7 +42,7 @@ class Drift: FROM "{input_data['schema']}"."{input_data['source_table_name']}" WHERE model_id = '{input_data['model_id']}' AND - timestamp > NOW() - INTERVAL {input_data['interval']} minutes + timestamp > NOW() - INTERVAL '{input_data['interval']} minutes' ORDER BY timestamp ASC """ diff --git a/laborious/workflows/simple_metrics.py b/laborious/workflows/simple_metrics.py index 7ab7153..438b98d 100644 --- a/laborious/workflows/simple_metrics.py +++ b/laborious/workflows/simple_metrics.py @@ -42,7 +42,7 @@ class SimpleMetrics: p.prediction is not null and ld.variable = '{target_name}' and ld.value is not null and - p."timestamp" >= NOW() - INTERVAL {interval_minutes} minutes + p."timestamp" >= NOW() - INTERVAL '{interval_minutes} minutes' order by p."timestamp" desc; """ diff --git a/requirements.txt b/requirements.txt index d475d8a..abc1753 100644 --- a/requirements.txt +++ b/requirements.txt @@ -4,7 +4,7 @@ sqlalchemy asyncua redis git+ssh://git@github.com/Aignosi/sientia-dataops-library.git@1.6.1 -git+ssh://git@github.com/Aignosi/sientia-mlops-library.git@0.40.2 +git+ssh://git@github.com/Aignosi/sientia-mlops-library.git@0.40.5 prometheus-client botocore boto3 diff --git a/tests/laborious/workflows/test_drift.py b/tests/laborious/workflows/test_drift.py index 25870ae..8f013d3 100644 --- a/tests/laborious/workflows/test_drift.py +++ b/tests/laborious/workflows/test_drift.py @@ -60,7 +60,7 @@ async def test_run(workflow_mock: AsyncMock, drift: Drift): FROM "{input_data['schema']}"."{input_data['source_table_name']}" WHERE model_id = '{input_data['model_id']}' AND - timestamp > NOW() - INTERVAL {input_data['interval']} minutes + timestamp > NOW() - INTERVAL '{input_data['interval']} minutes' ORDER BY timestamp ASC """ diff --git a/tests/laborious/workflows/test_simple_metrics.py b/tests/laborious/workflows/test_simple_metrics.py index 12c2cac..0d96bc5 100644 --- a/tests/laborious/workflows/test_simple_metrics.py +++ b/tests/laborious/workflows/test_simple_metrics.py @@ -61,7 +61,7 @@ async def test_run(workflow_mock: AsyncMock, simple_metrics: SimpleMetrics): p.prediction is not null and ld.variable = '{input_data['model_config']['target']}' and ld.value is not null and - p."timestamp" >= NOW() - INTERVAL {input_data['interval_minutes']} minutes + p."timestamp" >= NOW() - INTERVAL '{input_data['interval_minutes']} minutes' order by p."timestamp" desc; """ From 8f6a3cd93d9940fbe2601cf0269456c1b1423733 Mon Sep 17 00:00:00 2001 From: vitor-aignosi Date: Wed, 19 Nov 2025 11:53:50 -0300 Subject: [PATCH 23/27] SIENTIAPDE-1273 Enhance resource management and configuration in Laborious worker - Updated `values.yaml` to define resource limits and requests for better performance tuning. - Modified environment variables in `worker.py` to support resource-based scaling and improved task queue management. - Introduced new functions for creating resource tuners and poller behaviors, enhancing scalability and efficiency in handling workloads. --- laborious/worker/worker.py | 113 +++++++++++++++++++++++++++++-------- values.yaml | 26 ++++----- 2 files changed, 101 insertions(+), 38 deletions(-) diff --git a/laborious/worker/worker.py b/laborious/worker/worker.py index 0e9d3ca..50d00cf 100644 --- a/laborious/worker/worker.py +++ b/laborious/worker/worker.py @@ -5,12 +5,15 @@ This module provides the main worker implementation for the Sientia DataOps Labo It orchestrates Temporal workers, manages task queues, and handles the lifecycle of prediction and retraining workflows. -The worker supports two main task queues: -- predictions_batch-queue: Handles batch prediction workflows +The worker supports multiple task queues: +- predictions_batch-queue: Handles batch prediction workflows (heavy workload) - minimal_retrain-queue: Handles model retraining workflows +- drift-queue: Handles drift detection workflows +- simple_metrics-queue: Handles simple metrics calculation workflows Key Features: -- Automatic scaling with PollerBehaviorAutoscaling +- Resource-based scaling with WorkerTuner (CPU and memory aware) +- Automatic polling scaling with PollerBehaviorAutoscaling - Prometheus metrics integration - Comprehensive error handling and logging - Graceful shutdown with cleanup @@ -23,16 +26,38 @@ Environment Variables: - HTTP_METRICS_PORT: Prometheus metrics server port (default: 9090) - HTTP_SDK_METRICS_PORT: Temporal SDK metrics port (default: 9091) - PROJECT_NAME: Project name for notifications (default: laborious) + +Tuner Configuration (Resource-based scaling): +- TUNER_TARGET_MEMORY_USAGE: Target memory usage (0.0-1.0, default: 0.75) +- TUNER_TARGET_CPU_USAGE: Target CPU usage (0.0-1.0, default: 0.80) +- TUNER_WORKFLOW_MIN_SLOTS: Minimum workflow slots (default: 5) +- TUNER_WORKFLOW_MAX_SLOTS: Maximum workflow slots (default: 50) +- TUNER_ACTIVITY_MIN_SLOTS: Minimum activity slots (default: 5) +- TUNER_ACTIVITY_MAX_SLOTS: Maximum activity slots (default: 50) +- TUNER_WORKFLOW_RAMP_THROTTLE_MS: Workflow ramp throttle in ms (default: 100) +- TUNER_ACTIVITY_RAMP_THROTTLE_MS: Activity ramp throttle in ms (default: 50) + +Poller Configuration: +- POLLER_MINIMUM: Minimum number of pollers (default: 1) +- POLLER_MAXIMUM: Maximum number of pollers (default: 10) +- POLLER_INITIAL: Initial number of pollers (default: 2) """ +import os + from temporalio import client, workflow from temporalio.runtime import PrometheusConfig, Runtime, TelemetryConfig -from temporalio.worker import PollerBehaviorAutoscaling, Worker +from temporalio.worker import ( + PollerBehaviorAutoscaling, + ResourceBasedSlotConfig, + Worker, + WorkerTuner, +) with workflow.unsafe.imports_passed_through(): import asyncio - import os import sys + from datetime import timedelta from prometheus_client import start_http_server from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler @@ -60,6 +85,49 @@ POD_ID = os.getenv('POD_ID') SDK_METRICS_PORT = int(os.getenv('HTTP_SDK_METRICS_PORT', '9091')) +def create_resource_tuner() -> WorkerTuner: + """Create a resource-based tuner from environment variables.""" + target_memory = float(os.getenv('TUNER_TARGET_MEMORY_USAGE', '0.75')) + target_cpu = float(os.getenv('TUNER_TARGET_CPU_USAGE', '0.50')) + workflow_min = int(os.getenv('TUNER_WORKFLOW_MIN_SLOTS', '5')) + workflow_max = int(os.getenv('TUNER_WORKFLOW_MAX_SLOTS', '50')) + activity_min = int(os.getenv('TUNER_ACTIVITY_MIN_SLOTS', '5')) + activity_max = int(os.getenv('TUNER_ACTIVITY_MAX_SLOTS', '50')) + local_activity_min = int(os.getenv('TUNER_LOCAL_ACTIVITY_MIN_SLOTS', '1')) + local_activity_max = int(os.getenv('TUNER_LOCAL_ACTIVITY_MAX_SLOTS', '30')) + workflow_ramp = int(os.getenv('TUNER_WORKFLOW_RAMP_THROTTLE_MS', '100')) + activity_ramp = int(os.getenv('TUNER_ACTIVITY_RAMP_THROTTLE_MS', '50')) + local_activity_ramp = int(os.getenv('TUNER_LOCAL_ACTIVITY_RAMP_THROTTLE_MS', '50')) + + return WorkerTuner.create_resource_based( + target_memory_usage=target_memory, + target_cpu_usage=target_cpu, + workflow_config=ResourceBasedSlotConfig( + minimum_slots=workflow_min, + maximum_slots=workflow_max, + ramp_throttle=timedelta(milliseconds=workflow_ramp), + ), + activity_config=ResourceBasedSlotConfig( + minimum_slots=activity_min, + maximum_slots=activity_max, + ramp_throttle=timedelta(milliseconds=activity_ramp), + ), + local_activity_config=ResourceBasedSlotConfig( + minimum_slots=local_activity_min, + maximum_slots=local_activity_max, + ramp_throttle=timedelta(milliseconds=local_activity_ramp), + ), + ) + + +def create_poller_behavior() -> PollerBehaviorAutoscaling: + """Create poller behavior from environment variables.""" + minimum = int(os.getenv('POLLER_MINIMUM', '1')) + maximum = int(os.getenv('POLLER_MAXIMUM', '10')) + initial = int(os.getenv('POLLER_INITIAL', '2')) + return PollerBehaviorAutoscaling(minimum=minimum, maximum=maximum, initial=initial) + + async def main(): """ Main entry point for the Laborious worker application. @@ -138,6 +206,9 @@ async def main(): logger.custom_info('Starting Workers...', metadata) + tuner = create_resource_tuner() + poller = create_poller_behavior() + workers = [ Worker( temporal_client, @@ -151,12 +222,10 @@ async def main(): activities.format_retrain_report, activities.export_data_to_postgres, ], - max_concurrent_workflow_tasks=50, - max_concurrent_activities=50, - max_concurrent_local_activities=50, + tuner=tuner, max_cached_workflows=2, - workflow_task_poller_behavior=PollerBehaviorAutoscaling(), - activity_task_poller_behavior=PollerBehaviorAutoscaling(), + workflow_task_poller_behavior=poller, + activity_task_poller_behavior=poller, ), Worker( temporal_client, @@ -168,12 +237,10 @@ async def main(): activities.calculate_drift, activities.export_data_to_postgres, ], - max_concurrent_workflow_tasks=50, - max_concurrent_activities=50, - max_concurrent_local_activities=50, + tuner=tuner, max_cached_workflows=2, - workflow_task_poller_behavior=PollerBehaviorAutoscaling(), - activity_task_poller_behavior=PollerBehaviorAutoscaling(), + workflow_task_poller_behavior=poller, + activity_task_poller_behavior=poller, ), Worker( temporal_client, @@ -184,12 +251,10 @@ async def main(): activities.calculate_simple_metrics, activities.export_data_to_postgres, ], - max_concurrent_workflow_tasks=50, - max_concurrent_activities=50, - max_concurrent_local_activities=50, + tuner=tuner, max_cached_workflows=2, - workflow_task_poller_behavior=PollerBehaviorAutoscaling(), - activity_task_poller_behavior=PollerBehaviorAutoscaling(), + workflow_task_poller_behavior=poller, + activity_task_poller_behavior=poller, ), Worker( temporal_client, @@ -215,12 +280,10 @@ async def main(): activities.export_data_to_postgres, activities.write_metrics, ], - max_concurrent_workflow_tasks=50, - max_concurrent_activities=50, - max_concurrent_local_activities=50, + tuner=tuner, max_cached_workflows=200, - workflow_task_poller_behavior=PollerBehaviorAutoscaling(), - activity_task_poller_behavior=PollerBehaviorAutoscaling(), + workflow_task_poller_behavior=poller, + activity_task_poller_behavior=poller, ), ] diff --git a/values.yaml b/values.yaml index 64988ab..a41347c 100644 --- a/values.yaml +++ b/values.yaml @@ -52,17 +52,15 @@ securityContext: {} # runAsUser: 1000 -resources: {} - # We usually recommend not to specify default resources and to leave this as a conscious - # choice for the user. This also increases chances charts run on environments with little - # resources, such as Minikube. If you do want to specify resources, uncomment the following - # lines, adjust them as necessary, and remove the curly braces after 'resources:'. - # limits: - # cpu: 100m - # memory: 128Mi - # requests: - # cpu: 100m - # memory: 128Mi +resources: + # Resource limits and requests are important for ResourceBasedTuner to work correctly. + # The tuner monitors system CPU and memory usage, so proper resource limits must be set. + limits: + cpu: 2000m # 2 CPU cores + memory: 20Gi # 20 GB memory + requests: + cpu: 1000m # 1 CPU core + memory: 2Gi # 2 GB memory # This is to setup the liveness and readiness probes more information can be found here: https://kubernetes.io/docs/tasks/configure-pod-container/configure-liveness-readiness-startup-probes/ livenessProbe: @@ -167,9 +165,11 @@ env: - name: POSTGRES_DBNAME value: "sientia" - name: POSTGRES_MIN_CONNECTIONS - value: "10" + value: "20" + # max_connections = number_of_workers * max_concurrent_activities * safety_factor + # Example: 4 workers * 50 activities * 0.5 = 100 connections - name: POSTGRES_MAX_CONNECTIONS - value: "30" + value: "100" - name: MLFLOW_HOST value: "http://sientia-tracker-mlflow-tracking.sientia-tracker.svc.cluster.local" From 684bc9fb31b3e81587e1c9d67ac40f16ca7b699f Mon Sep 17 00:00:00 2001 From: vitor-aignosi Date: Wed, 19 Nov 2025 13:26:08 -0300 Subject: [PATCH 24/27] SIENTIAPDE-1273 Enhance prediction data retrieval in MLFlowRepository by adding predict_flavor parameter - Updated get_prediction_data method to accept a predict_flavor argument, allowing for different prediction model handling. - Adjusted calls to get_prediction_data throughout the codebase to include the new parameter. - Added new test cases to validate behavior for different predict_flavor values, ensuring robust functionality. --- .../utils/repository/model_repository.py | 9 ++++-- .../utils/repository/test_model_repository.py | 32 ++++++++++++++++--- 2 files changed, 34 insertions(+), 7 deletions(-) diff --git a/laborious/utils/repository/model_repository.py b/laborious/utils/repository/model_repository.py index 395865e..9543ddd 100644 --- a/laborious/utils/repository/model_repository.py +++ b/laborious/utils/repository/model_repository.py @@ -705,14 +705,17 @@ class MLFlowRepository(SientiaMonitoring): """ def get_prediction_data( - self, prediction_model: Any, retrain_dataset: pd.DataFrame, target_name: str + self, prediction_model: Any, retrain_dataset: pd.DataFrame, target_name: str, predict_flavor: str ) -> pd.DataFrame: """ Get prediction data from prediction model. """ input_index = retrain_dataset.index - prediction_data = prediction_model.predict(retrain_dataset) + if predict_flavor == 'pyfunc': + prediction_data = prediction_model.predict({}, retrain_dataset) + else: + prediction_data = prediction_model.predict(retrain_dataset) if isinstance(prediction_data, pd.DataFrame): prediction_data.columns = pd.Index(['prediction']) @@ -858,7 +861,7 @@ class MLFlowRepository(SientiaMonitoring): prediction_model.fit(retrain_dataset) # get prediction data - prediction_data = self.get_prediction_data(prediction_model, retrain_dataset, target_name) + prediction_data = self.get_prediction_data(prediction_model, retrain_dataset, target_name, predict_flavor) self.info(f'Model experiment creation completed successfully for {model_name}', metadata) diff --git a/tests/laborious/utils/repository/test_model_repository.py b/tests/laborious/utils/repository/test_model_repository.py index 40782fe..a86c7ec 100644 --- a/tests/laborious/utils/repository/test_model_repository.py +++ b/tests/laborious/utils/repository/test_model_repository.py @@ -865,7 +865,7 @@ async def test_fit_models_not_df_target_name_none_and_not_in_model( prediction_model.fit.assert_called_once_with(pd_merge.return_value) mlflow_repository.get_prediction_data.assert_called_once_with( - prediction_model, pd_merge.return_value, data_model.fit.return_value.target_variable + prediction_model, pd_merge.return_value, data_model.fit.return_value.target_variable, 'pyfunc' ) assert output == { @@ -951,7 +951,7 @@ async def test_fit_models_df_target_name_not_none_and_in_model( prediction_model.fit.assert_called_once_with(transformed_data) mlflow_repository.get_prediction_data.assert_called_once_with( - prediction_model, transformed_data, 'feat_1' + prediction_model, transformed_data, 'feat_1', 'pyfunc' ) assert output == { @@ -1494,8 +1494,11 @@ def test_get_prediction_data_dataframe(mlflow_repository): retrain_dataset = DataFrame({'feat_1': [1, 2], 'target': [3, 4]}, index=['idx1', 'idx2']) prediction_model.predict.return_value = DataFrame({'pred': [5, 6]}, index=['idx1', 'idx2']) target_name = 'target' + predict_flavor = 'sklearn' - result = mlflow_repository.get_prediction_data(prediction_model, retrain_dataset, target_name) + result = mlflow_repository.get_prediction_data( + prediction_model, retrain_dataset, target_name, predict_flavor + ) prediction_model.predict.assert_called_once_with(retrain_dataset) assert 'prediction' in result.columns @@ -1509,11 +1512,32 @@ def test_get_prediction_data_array(mlflow_repository): retrain_dataset = DataFrame({'feat_1': [1, 2], 'target': [3, 4]}, index=['idx1', 'idx2']) prediction_model.predict.return_value = [5, 6] target_name = 'target' + predict_flavor = 'sklearn' - result = mlflow_repository.get_prediction_data(prediction_model, retrain_dataset, target_name) + result = mlflow_repository.get_prediction_data( + prediction_model, retrain_dataset, target_name, predict_flavor + ) prediction_model.predict.assert_called_once_with(retrain_dataset) assert 'prediction' in result.columns assert 'target' in result.columns assert 'timestamp' in result.columns assert result.index.tolist() == [0, 1] + + +def test_get_prediction_data_pyfunc(mlflow_repository): + prediction_model = MagicMock() + retrain_dataset = DataFrame({'feat_1': [1, 2], 'target': [3, 4]}, index=['idx1', 'idx2']) + prediction_model.predict.return_value = DataFrame({'pred': [5, 6]}, index=['idx1', 'idx2']) + target_name = 'target' + predict_flavor = 'pyfunc' + + result = mlflow_repository.get_prediction_data( + prediction_model, retrain_dataset, target_name, predict_flavor + ) + + prediction_model.predict.assert_called_once_with({}, retrain_dataset) + assert 'prediction' in result.columns + assert 'target' in result.columns + assert 'timestamp' in result.columns + assert result.index.tolist() == [0, 1] From 3311409885a226e601b21661cf968c3e72372d77 Mon Sep 17 00:00:00 2001 From: vitor-aignosi Date: Wed, 19 Nov 2025 14:54:50 -0300 Subject: [PATCH 25/27] SIENTIAPDE-1273 Update sientia-mlops-library dependency to version 0.40.6 and refactor get_prediction_data method calls for improved readability in model_repository.py and test_model_repository.py. --- laborious/utils/repository/model_repository.py | 10 ++++++++-- requirements.txt | 2 +- .../utils/repository/test_model_repository.py | 5 ++++- 3 files changed, 13 insertions(+), 4 deletions(-) diff --git a/laborious/utils/repository/model_repository.py b/laborious/utils/repository/model_repository.py index 9543ddd..eaab16c 100644 --- a/laborious/utils/repository/model_repository.py +++ b/laborious/utils/repository/model_repository.py @@ -705,7 +705,11 @@ class MLFlowRepository(SientiaMonitoring): """ def get_prediction_data( - self, prediction_model: Any, retrain_dataset: pd.DataFrame, target_name: str, predict_flavor: str + self, + prediction_model: Any, + retrain_dataset: pd.DataFrame, + target_name: str, + predict_flavor: str, ) -> pd.DataFrame: """ Get prediction data from prediction model. @@ -861,7 +865,9 @@ class MLFlowRepository(SientiaMonitoring): prediction_model.fit(retrain_dataset) # get prediction data - prediction_data = self.get_prediction_data(prediction_model, retrain_dataset, target_name, predict_flavor) + prediction_data = self.get_prediction_data( + prediction_model, retrain_dataset, target_name, predict_flavor + ) self.info(f'Model experiment creation completed successfully for {model_name}', metadata) diff --git a/requirements.txt b/requirements.txt index abc1753..3921162 100644 --- a/requirements.txt +++ b/requirements.txt @@ -4,7 +4,7 @@ sqlalchemy asyncua redis git+ssh://git@github.com/Aignosi/sientia-dataops-library.git@1.6.1 -git+ssh://git@github.com/Aignosi/sientia-mlops-library.git@0.40.5 +git+ssh://git@github.com/Aignosi/sientia-mlops-library.git@0.40.6 prometheus-client botocore boto3 diff --git a/tests/laborious/utils/repository/test_model_repository.py b/tests/laborious/utils/repository/test_model_repository.py index a86c7ec..466fae5 100644 --- a/tests/laborious/utils/repository/test_model_repository.py +++ b/tests/laborious/utils/repository/test_model_repository.py @@ -865,7 +865,10 @@ async def test_fit_models_not_df_target_name_none_and_not_in_model( prediction_model.fit.assert_called_once_with(pd_merge.return_value) mlflow_repository.get_prediction_data.assert_called_once_with( - prediction_model, pd_merge.return_value, data_model.fit.return_value.target_variable, 'pyfunc' + prediction_model, + pd_merge.return_value, + data_model.fit.return_value.target_variable, + 'pyfunc', ) assert output == { From 2a451da976e895beb02ef68c13b01278f876325e Mon Sep 17 00:00:00 2001 From: vitor-aignosi Date: Wed, 19 Nov 2025 18:31:52 -0300 Subject: [PATCH 26/27] SIENTIAPDE-1273 Refactor imports in opc.py to improve code clarity - Moved the import of Hashable from collections.abc to the appropriate section, enhancing readability and organization of imports. --- laborious/activities/opc.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/laborious/activities/opc.py b/laborious/activities/opc.py index cfbf450..c4abe95 100644 --- a/laborious/activities/opc.py +++ b/laborious/activities/opc.py @@ -1,9 +1,8 @@ -from collections.abc import Hashable - from temporalio import activity, workflow with workflow.unsafe.imports_passed_through(): import traceback + from collections.abc import Hashable from typing import Any from pandas import DataFrame From 894d39faa4cd2bd5646dfb8e85790a5d03bb56d4 Mon Sep 17 00:00:00 2001 From: vitor-aignosi Date: Mon, 24 Nov 2025 09:06:00 -0300 Subject: [PATCH 27/27] SIENTIAPDE-1273 Refactor import statements in worker.py for improved organization - Moved the import of the os module to the appropriate section, enhancing clarity and consistency in the import structure. --- laborious/worker/worker.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/laborious/worker/worker.py b/laborious/worker/worker.py index 50d00cf..33384b0 100644 --- a/laborious/worker/worker.py +++ b/laborious/worker/worker.py @@ -43,7 +43,6 @@ Poller Configuration: - POLLER_INITIAL: Initial number of pollers (default: 2) """ -import os from temporalio import client, workflow from temporalio.runtime import PrometheusConfig, Runtime, TelemetryConfig @@ -80,6 +79,7 @@ with workflow.unsafe.imports_passed_through(): FormatAndExportPrediction, ) from laborious.workflows.sub_workflows.prediction_process import PredictionProcess + import os POD_ID = os.getenv('POD_ID') SDK_METRICS_PORT = int(os.getenv('HTTP_SDK_METRICS_PORT', '9091'))