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(