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

SIENTIAPDE-1273

Browse Source 
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.
This commit is contained in:
vitor-aignosi
2025-11-13 16:40:23 -03:00
parent 6ac0f38d59
commit b68674fe64
15 changed files with 1639 additions and 68 deletions
Download Patch File Download Diff File Expand all files Collapse all files

12
laborious/activities/activities.py
View File

@@ -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)

34
laborious/activities/mlflow.py
View File

@@ -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()

266
laborious/activities/model_metrics.py Normal file
View File

@@ -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()

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

@@ -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 {}

20
laborious/worker/worker.py
View File

@@ -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,

99
laborious/workflows/drift.py Normal file
View File

@@ -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),
)

1
laborious/workflows/predictions_batch.py
View File

@@ -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'}}),

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

@@ -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,

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

@@ -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'],

101
tests.ipynb
View File

@@ -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": [
"<div>\n",
"<style scoped>\n",
" .dataframe tbody tr th:only-of-type {\n",
" vertical-align: middle;\n",
" }\n",
"\n",
" .dataframe tbody tr th {\n",
" vertical-align: top;\n",
" }\n",
"\n",
" .dataframe thead th {\n",
" text-align: right;\n",
" }\n",
"</style>\n",
"<table border=\"1\" class=\"dataframe\">\n",
" <thead>\n",
" <tr style=\"text-align: right;\">\n",
" <th></th>\n",
" <th>a</th>\n",
" <th>b</th>\n",
" <th>target</th>\n",
" <th>prediction</th>\n",
" <th>timestamp</th>\n",
" </tr>\n",
" </thead>\n",
" <tbody>\n",
" <tr>\n",
" <th>0</th>\n",
" <td>1</td>\n",
" <td>4</td>\n",
" <td>7</td>\n",
" <td>1</td>\n",
" <td>2025-01-01</td>\n",
" </tr>\n",
" <tr>\n",
" <th>1</th>\n",
" <td>2</td>\n",
" <td>5</td>\n",
" <td>8</td>\n",
" <td>2</td>\n",
" <td>2025-01-02</td>\n",
" </tr>\n",
" <tr>\n",
" <th>2</th>\n",
" <td>3</td>\n",
" <td>6</td>\n",
" <td>9</td>\n",
" <td>3</td>\n",
" <td>2025-01-03</td>\n",
" </tr>\n",
" </tbody>\n",
"</table>\n",
"</div>"
],
"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": {

641
tests/laborious/activities/test_model_metrics.py Normal file
View File

@@ -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')

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

@@ -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]

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

@@ -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,

247
tests/laborious/workflows/test_drift.py Normal file
View File

@@ -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,
)

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

@@ -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(