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

SIENTIAPDE-1273

Browse Source 
SIENTIAPDE-1273
Enhance security analysis and SQL injection handling

- Added skip for potential SQL injection false positives in Bandit configuration.
- Updated validate.sh to use the pyproject.toml configuration for Bandit security analysis.
- Refactored code to replace ensure_dataframe utility with direct DataFrame usage in multiple activities, improving clarity and reducing dependencies.
- Removed the deprecated dataframe_utils module to streamline the codebase.
This commit is contained in:
vitor-aignosi
2025-11-17 16:04:54 -03:00
parent 1014c33dd9
commit a88a15c60a
18 changed files with 516 additions and 439 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
laborious/activities/activities.py
View File

@@ -9,9 +9,10 @@ with workflow.unsafe.imports_passed_through():
from laborious.activities.gates import Gates
from laborious.activities.mlflow import MLFlow
from laborious.activities.model_metrics import ModelMetrics
from laborious.activities.opc import OPC
from laborious.activities.storage import Storage
from laborious.activities.model_metrics import ModelMetrics
class Activities(Storage, MLFlow, Gates, OPC, ModelMetrics):
"""
@@ -131,4 +132,4 @@ class Activities(Storage, MLFlow, Gates, OPC, ModelMetrics):
MLFlow.close(self)
Gates.close(self)
await OPC.close(self)
ModelMetrics.close(self)
ModelMetrics.close(self)

12
laborious/activities/gates.py
View File

@@ -15,7 +15,6 @@ with workflow.unsafe.imports_passed_through():
from sientia_do.temporal.constants import DATETIME_FORMAT_WITH_TZ, now
from laborious import metrics
from laborious.utils.dataframe_utils import ensure_dataframe
from laborious.utils.filters.conditional_filters import (
filter_empty_data,
filter_specific_variables_null_values,
@@ -150,7 +149,7 @@ class Gates(SientiaMonitoring):
self.info('Performing input gate...', metadata)
filters = input_data['filters']
data = ensure_dataframe(input_data['data'])
data = DataFrame(input_data['data'])
path_priority = input_data['path_priority']
filter_output = []
@@ -403,19 +402,18 @@ class Gates(SientiaMonitoring):
return policy_type, int(policy_value)
@activity.defn(name='format_transformed_data')
async def format_transformed_data(self, input_data: dict[str, Any]) -> dict:
"""
Format transformed data according to configured storage policies.
"""
metadata = input_data['metadata']
model_id = input_data['model_id']
self.info('Formatting transformed data...', metadata)
data = ensure_dataframe(input_data['data'])
data = DataFrame(input_data['data'])
data['timestamp'] = data.index
data = data.reset_index(drop=True)
@@ -454,7 +452,7 @@ class Gates(SientiaMonitoring):
prediction_store_policy = input_data['prediction_store_policy']
self.info('Formatting prediction...', metadata)
data = ensure_dataframe(input_data['data'])
data = DataFrame(input_data['data'])
# Create timestamp column from index and reset index
data['timestamp'] = data.index
@@ -602,7 +600,7 @@ class Gates(SientiaMonitoring):
self.info('Getting last timestamp...', metadata)
data = ensure_dataframe(input_data['data'])
data = DataFrame(input_data['data'])
self.debug(f'Input data: {data.head(5).to_string()}', metadata)

11
laborious/activities/mlflow.py
View File

@@ -1,4 +1,3 @@
from typing import Hashable
from temporalio import activity, workflow
with workflow.unsafe.imports_passed_through():
@@ -20,7 +19,6 @@ with workflow.unsafe.imports_passed_through():
now,
)
from laborious.utils.dataframe_utils import ensure_dataframe
from laborious.utils.repository.minio_repository import MinioRepository
from laborious.utils.repository.model_repository import MLFlowRepository
@@ -140,7 +138,7 @@ class MLFlow(SientiaMonitoring):
"""
metadata = input_data['metadata']
self.info('Transforming data...', metadata)
data = ensure_dataframe(input_data['data'])
data = DataFrame(input_data['data'])
model_name = input_data['model_name']
model_config = input_data.get('model_config', {})
@@ -212,7 +210,7 @@ class MLFlow(SientiaMonitoring):
"""
metadata = input_data['metadata']
self.info('Predicting data...', metadata)
data = ensure_dataframe(input_data['data'])
data = DataFrame(input_data['data'])
model_name = input_data['model_name']
model_config = input_data.get('model_config', {})
@@ -421,7 +419,6 @@ class MLFlow(SientiaMonitoring):
self.error(trace, metadata=metadata)
raise e
@activity.defn(name='get_reference_data')
async def get_reference_data(self, input_data: dict[str, Any]) -> list[dict] | None:
"""
@@ -439,7 +436,7 @@ class MLFlow(SientiaMonitoring):
metadata = input_data['metadata']
model_name = input_data['model_name']
artifact = "evaluation_data.csv"
artifact = 'evaluation_data.csv'
reference_data = await self.model_monitoring_repository.load_artifact_dataframe(
model_name=model_name, artifact_path=artifact, metadata=metadata
@@ -452,4 +449,4 @@ class MLFlow(SientiaMonitoring):
reference_data['timestamp'] = to_datetime(reference_data['timestamp'])
reference_data['timestamp'] = reference_data['timestamp'].dt.strftime(DATETIME_FORMAT)
return reference_data.to_dict(orient='records')
return reference_data.to_dict(orient='records')

139
laborious/activities/model_metrics.py
View File

@@ -1,25 +1,28 @@
from temporalio import activity, workflow
with workflow.unsafe.imports_passed_through():
import time
import traceback
import warnings
from typing import Any
import numpy as np
from pandas import DataFrame, Index, to_datetime
from sientia_do.observability.sientia_monitoring import SientiaMonitoring
from sientia_do.observability.metrics_controller import MetricsController
from sientia_do.observability.logger import Logger
from sientia.ModelAnalysis import ModelAnalysis
from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler
from sientia_do.notifications.models import NotificationLevel
from sientia.ModelAnalysis import ModelAnalysis
from laborious import metrics
import time
import numpy as np
from sientia_do.observability.logger import Logger
from sientia_do.observability.metrics_controller import MetricsController
from sientia_do.observability.sientia_monitoring import SientiaMonitoring
from sientia_do.temporal.constants import DATETIME_FORMAT, DATETIME_FORMAT_WITH_TZ
import warnings
import traceback
from laborious import metrics
warnings.filterwarnings('ignore', category=RuntimeWarning, message='Degrees of freedom <= 0')
warnings.filterwarnings('ignore', category=RuntimeWarning, message='invalid value encountered in scalar divide')
warnings.filterwarnings(
'ignore', category=RuntimeWarning, message='invalid value encountered in scalar divide'
)
class ModelMetrics(SientiaMonitoring):
"""
@@ -28,12 +31,12 @@ class ModelMetrics(SientiaMonitoring):
This class provides activities for writing metrics to the Prometheus monitoring system.
"""
def __init__(self,
def __init__(
self,
logger: Logger,
notification_handler: NotificationHandler,
metrics_controller: MetricsController,
):
SientiaMonitoring.__init__(self, logger, notification_handler, metrics_controller)
def close(self) -> None:
@@ -42,11 +45,11 @@ class ModelMetrics(SientiaMonitoring):
"""
SientiaMonitoring.shutdown(self)
def __del__(self):
self.close()
async def get_drift_metrics(self,
async def get_drift_metrics(
self,
reference_data: DataFrame,
target_data: DataFrame,
target_name: str,
@@ -73,34 +76,37 @@ class ModelMetrics(SientiaMonitoring):
'features': reference_columns,
}
model_analysis = ModelAnalysis(config=config)
self.debug(f'Reference data: Size {reference_data.shape} \n{reference_data.head(5).to_string()}', metadata)
self.debug(
f'Reference data: Size {reference_data.shape} \n{reference_data.head(5).to_string()}',
metadata,
)
self.debug(f'Target data: Size {target_data.shape} \n{target_data.head(5).to_string()}', metadata)
self.debug(
f'Target data: Size {target_data.shape} \n{target_data.head(5).to_string()}', metadata
)
core_labels = self.get_core_labels(metadata, operation_type='detect_univariate_drift')
start_time = time.time()
try:
univariate_drift = model_analysis.detect_univariate_drift(
reference_df=reference_data,
analysis_df=target_data,
features=reference_columns,
timestamp_col=config['timestamp'],
methods=drift_metrics,
chunk_period=chunk_period
chunk_period=chunk_period,
)
except Exception as e:
self.error(f'Error detecting univariate drift: {e}', metadata)
await self.emit_metric(
metric_object=metrics.MODEL_ANALYZE_ERROR_COUNT, tags=core_labels)
metric_object=metrics.MODEL_ANALYZE_ERROR_COUNT, tags=core_labels
)
raise e
await self.observe_lag(start_time, metrics.MODEL_ANALYZE_LAG, core_labels)
await self.emit_metric(metric_object=metrics.MODEL_ANALYZE_COUNT, tags=core_labels)
core_labels = self.get_core_labels(metadata, operation_type='detect_multivariate_drift')
start_time = time.time()
try:
@@ -109,11 +115,13 @@ class ModelMetrics(SientiaMonitoring):
analysis_df=target_data,
features=reference_columns,
timestamp_col=config['timestamp'],
chunk_period=chunk_period
chunk_period=chunk_period,
)
except Exception as e:
self.error(f'Error detecting multivariate drift: {e}', metadata)
await self.emit_metric(metric_object=metrics.MODEL_ANALYZE_ERROR_COUNT, tags=core_labels)
await self.emit_metric(
metric_object=metrics.MODEL_ANALYZE_ERROR_COUNT, tags=core_labels
)
raise e
await self.observe_lag(start_time, metrics.MODEL_ANALYZE_LAG, core_labels)
await self.emit_metric(metric_object=metrics.MODEL_ANALYZE_COUNT, tags=core_labels)
@@ -127,15 +135,18 @@ class ModelMetrics(SientiaMonitoring):
)
except Exception as e:
self.error(f'Error getting drift metrics: {e}', metadata)
await self.emit_metric(metric_object=metrics.MODEL_ANALYZE_ERROR_COUNT, tags=core_labels)
await self.emit_metric(
metric_object=metrics.MODEL_ANALYZE_ERROR_COUNT, tags=core_labels
)
raise e
await self.observe_lag(start_time, metrics.MODEL_ANALYZE_LAG, core_labels)
await self.emit_metric(metric_object=metrics.MODEL_ANALYZE_COUNT, tags=core_labels)
self.debug(f'Drift dataframe: Size {drift_df.shape} \n{drift_df.head(5).to_string()}', metadata)
return drift_df
self.debug(
f'Drift dataframe: Size {drift_df.shape} \n{drift_df.head(5).to_string()}', metadata
)
return drift_df
@activity.defn(name='calculate_drift')
async def calculate_drift(self, input_data: dict[str, Any]) -> list[dict]:
@@ -194,8 +205,8 @@ class ModelMetrics(SientiaMonitoring):
)
reference_columns = reference_data.drop(
columns=[target_name, 'timestamp', 'target', 'prediction'],
errors='ignore').columns
columns=[target_name, 'timestamp', 'target', 'prediction'], errors='ignore'
).columns
try:
drift_df = await self.get_drift_metrics(
@@ -224,12 +235,11 @@ class ModelMetrics(SientiaMonitoring):
return []
# Drop unnecessary columns
drift_df.drop(columns=['p_value', 'chunk_start_date', 'chunk_end_date'], inplace=True)
drift_df.drop(columns=['p_value'], inplace=True)
# Extract timestamps only until minutes
if chunk_period == 'min':
target_timestamps = target_data['timestamp'].apply(
lambda x: x[:16])
target_timestamps = target_data['timestamp'].apply(lambda x: x[:16])
else:
target_timestamps = target_data['timestamp']
@@ -237,32 +247,39 @@ class ModelMetrics(SientiaMonitoring):
drift_df = drift_df[drift_df['timestamp'].isin(target_timestamps)]
if drift_df.empty:
self.warning('No drift metrics found after dropping rows where timestamp is not in target data', metadata)
self.warning(
'No drift metrics found after dropping rows where timestamp is not in target data',
metadata,
)
return []
# Rename columns to match database columns
drift_df.rename(columns={
'metric': 'method',
'statistic': 'value',
}, inplace=True)
drift_df.rename(
columns={
'metric': 'method',
'statistic': 'value',
},
inplace=True,
)
# Drop duplicates
drift_df.drop_duplicates(
subset=['timestamp', 'method', 'feature'],
keep='first', inplace=True)
subset=['timestamp', 'method', 'feature'], keep='first', inplace=True
)
drift_df['model_id'] = model_id
drift_df['accurate'] = accurate
drift_df['timestamp'] = to_datetime(drift_df['timestamp'])
drift_df['timestamp'] = drift_df['timestamp'].dt.tz_localize('UTC')
drift_df['timestamp'] = drift_df['timestamp'].dt.strftime(DATETIME_FORMAT_WITH_TZ)
self.debug(f'Drift dataframe: Size {drift_df.shape} \n{drift_df.head(5).to_string()}', metadata)
self.debug(
f'Drift dataframe: Size {drift_df.shape} \n{drift_df.head(5).to_string()}', metadata
)
self.debug(f'Drift dataframe: {drift_df.head(5).to_string()}', metadata)
return drift_df.to_dict(orient='records')
@activity.defn(name='calculate_simple_metrics')
@@ -298,52 +315,40 @@ class ModelMetrics(SientiaMonitoring):
output_data = []
diff = target_data['target'] - target_data['prediction']
diff_squared = diff ** 2
diff_squared = diff**2
self.info(f'Calculating simple metrics for model {model_id}: {metrics}', metadata)
for metric in metrics:
if metric == 'rmse':
output_data.append({
'metric': 'rmse',
'value': np.sqrt(np.mean(diff_squared))
})
output_data.append({'metric': 'rmse', 'value': np.sqrt(np.mean(diff_squared))})
elif metric == 'mse':
output_data.append({
'metric': 'mse',
'value': np.mean(diff_squared)
})
output_data.append({'metric': 'mse', 'value': np.mean(diff_squared)})
elif metric == 'mae':
output_data.append({
'metric': 'mae',
'value': np.mean(np.abs(diff))
})
output_data.append({'metric': 'mae', 'value': np.mean(np.abs(diff))})
elif metric == 'r2':
y_true = target_data['target']
y_mean = np.mean(y_true)
ss_res = np.sum(diff_squared)
ss_tot = np.sum((y_true - y_mean) ** 2)
# Evita divisão por zero
if ss_tot == 0:
r2_score = 0.0
else:
r2_score = 1 - (ss_res / ss_tot)
output_data.append({
'metric': 'r2',
'value': r2_score
})
output_data.append({'metric': 'r2', 'value': r2_score})
data = DataFrame(output_data)
data['model_id'] = model_id
data['timestamp'] = target_data['timestamp'].max()
data['data_size'] = data_size
data['interval_minutes'] = interval_minutes
self.debug(f'Simple metrics dataframe: Size {data.shape} \n{data.head(5).to_string()}', metadata)
self.debug(
f'Simple metrics dataframe: Size {data.shape} \n{data.head(5).to_string()}', metadata
)
return data.to_dict(orient='records')

6
laborious/activities/opc.py
View File

@@ -1,4 +1,5 @@
from typing import Hashable
from collections.abc import Hashable
from temporalio import activity, workflow
with workflow.unsafe.imports_passed_through():
@@ -12,7 +13,6 @@ with workflow.unsafe.imports_passed_through():
from sientia_do.observability.metrics_controller import MetricsController
from sientia_do.observability.sientia_monitoring import SientiaMonitoring
from laborious.utils.dataframe_utils import ensure_dataframe
from laborious.utils.repository.opc_repository import OpcRepository
OPC_WRITTING_ERROR_CONFIDENCE = 12
@@ -297,7 +297,7 @@ class OPC(SientiaMonitoring):
"""
metadata = input_data['metadata']
self.info('Writing data to OPC servers...', metadata)
data = ensure_dataframe(input_data['data'])
data = DataFrame(input_data['data'])
opc_output_config = input_data['opc_output_config']
self.info(f'Data to write: {data.size} rows', metadata)

2
laborious/metrics.py
View File

@@ -187,4 +187,4 @@ MODEL_ANALYZE_ERROR_COUNT = Counter(
'laborious_model_analyze_error_count',
'Number of errors during analyze operations',
SIENTIA_CORE_LABELS,
)
)

31
laborious/utils/dataframe_utils.py
View File

@@ -1,31 +0,0 @@
"""
DataFrame utility functions for handling serialized DataFrames.
This module provides helper functions to work with DataFrames that may
come from Temporal serialization (already as DataFrame) or from legacy
code (as dict).
"""
from typing import Any
from pandas import DataFrame
def ensure_dataframe(data: Any) -> DataFrame:
"""
Ensure that data is a DataFrame, converting from dict if necessary.
This function handles both cases:
- Data already deserialized as DataFrame (from Temporal codec)
- Data as dict (legacy format or non-DataFrame serialization)
Args:
data: Data that should be a DataFrame (can be DataFrame or dict)
Returns:
DataFrame: The data as a pandas DataFrame
"""
if isinstance(data, DataFrame):
return data
return DataFrame(data)

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

@@ -16,11 +16,11 @@ Capabilities:
import ctypes
import gc
from io import StringIO
import threading
import time
import traceback
from datetime import datetime, timedelta
from io import StringIO
from os import environ, makedirs, path
from shutil import rmtree
from typing import Any, Literal, overload
@@ -34,7 +34,6 @@ from sientia_do.observability.logger import Logger
from sientia_do.observability.metrics_controller import MetricsController
from sientia_do.observability.sientia_monitoring import SientiaMonitoring
from sientia_do.temporal.constants import DATETIME_FORMAT_WITH_TZ
from sientia.ModelAnalysis import ModelAnalysis
from laborious import metrics
@@ -217,8 +216,9 @@ class MLFlowRepository(SientiaMonitoring):
run_info = mlflow.get_run(run_id)
return run_info.data.params
def check_artifact_exists(self, run_id: str,
artifact_path: str, metadata: dict[str, Any]) -> bool:
def check_artifact_exists(
self, run_id: str, artifact_path: str, metadata: dict[str, Any]
) -> bool:
"""
Check if an artifact exists in the MLflow Model Registry.
@@ -233,8 +233,9 @@ class MLFlowRepository(SientiaMonitoring):
self.debug(f'Artifacts of {run_id}: \n{artifacts}', metadata)
self.debug(f'Looking for artifact {artifact_path} in {run_id}', metadata)
return any(artifact.path == artifact_path for artifact in artifacts)
"""
Functions related to download and load models
"""
@@ -279,10 +280,9 @@ class MLFlowRepository(SientiaMonitoring):
return artifacts
async def load_artifact_dataframe(self, model_name: str, artifact_path: str,
metadata: dict[str, Any]) -> pd.DataFrame | None:
async def load_artifact_dataframe(
self, model_name: str, artifact_path: str, metadata: dict[str, Any]
) -> pd.DataFrame | None:
"""
Load the dataframe content of an artifact from the MLflow Model Registry.
@@ -300,7 +300,7 @@ class MLFlowRepository(SientiaMonitoring):
if not self.check_artifact_exists(run_id, artifact_path, metadata):
return None
artifact_path = path.join("runs:/", run_id, artifact_path)
artifact_path = path.join('runs:/', run_id, artifact_path)
start_time = time.time()
try:
@@ -704,8 +704,9 @@ class MLFlowRepository(SientiaMonitoring):
Functions related to model retraining
"""
def get_prediction_data(self, prediction_model: Any, retrain_dataset: pd.DataFrame,
target_name: str) -> pd.DataFrame:
def get_prediction_data(
self, prediction_model: Any, retrain_dataset: pd.DataFrame, target_name: str
) -> pd.DataFrame:
"""
Get prediction data from prediction model.
"""
@@ -714,7 +715,6 @@ class MLFlowRepository(SientiaMonitoring):
prediction_data = prediction_model.predict(retrain_dataset)
if isinstance(prediction_data, pd.DataFrame):
prediction_data.columns = pd.Index(['prediction'])
else:
@@ -724,7 +724,8 @@ class MLFlowRepository(SientiaMonitoring):
# Merge prediction data with retrain_dataset on index
prediction_data = pd.merge(
retrain_dataset, prediction_data, left_index=True, right_index=True, how='left')
retrain_dataset, prediction_data, left_index=True, right_index=True, how='left'
)
# Rename column "target_name" to "target"
prediction_data.rename(columns={target_name: 'target'}, inplace=True)
@@ -733,9 +734,7 @@ class MLFlowRepository(SientiaMonitoring):
prediction_data.reset_index(drop=True, inplace=True)
prediction_data.sort_values(
by='timestamp', ascending=True, inplace=True
)
prediction_data.sort_values(by='timestamp', ascending=True, inplace=True)
return prediction_data
@@ -859,8 +858,7 @@ class MLFlowRepository(SientiaMonitoring):
prediction_model.fit(retrain_dataset)
# get prediction data
prediction_data = self.get_prediction_data(
prediction_model, retrain_dataset, target_name)
prediction_data = self.get_prediction_data(prediction_model, retrain_dataset, target_name)
self.info(f'Model experiment creation completed successfully for {model_name}', metadata)
@@ -1439,4 +1437,3 @@ class MLFlowRepository(SientiaMonitoring):
metadata_result['mlflow_experiment_id'] = experiment_id
return metadata_result

4
laborious/worker/worker.py
View File

@@ -29,7 +29,6 @@ from temporalio import client, workflow
from temporalio.runtime import PrometheusConfig, Runtime, TelemetryConfig
from temporalio.worker import PollerBehaviorAutoscaling, Worker
with workflow.unsafe.imports_passed_through():
import asyncio
import os
@@ -48,11 +47,10 @@ with workflow.unsafe.imports_passed_through():
build_opc_config,
build_postgres_config,
)
from laborious.workflows.drift import Drift
from laborious.workflows.minimal_retrain import MinimalRetrain
from laborious.workflows.predictions_batch import PredictionsBatch
from laborious.workflows.drift import Drift
from laborious.workflows.simple_metrics import SimpleMetrics
from laborious.workflows.sub_workflows.format_and_export_prediction import (
FormatAndExportPrediction,
)

25
laborious/workflows/drift.py
View File

@@ -4,10 +4,10 @@ with workflow.unsafe.imports_passed_through():
from datetime import timedelta
from typing import Any
from sientia_do.temporal.constants import DATETIME_FORMAT_WITH_TZ
from sientia_do.temporal.policies import retry_policy
from laborious.activities.activities import Activities
from sientia_do.temporal.constants import DATETIME_FORMAT_WITH_TZ
@workflow.defn(name='drift')
@@ -39,10 +39,10 @@ class Drift:
gathering_query = f"""
SELECT *
FROM {input_data['schema']}.{input_data['source_table_name']}
FROM "{input_data['schema']}"."{input_data['source_table_name']}"
WHERE
model_id = {input_data['model_id']} AND
timestamp > NOW() - INTERVAL '{input_data['interval']} minutes'
model_id = '{input_data['model_id']}' AND
timestamp > NOW() - INTERVAL {input_data['interval']} minutes
ORDER BY timestamp ASC
"""
@@ -60,10 +60,7 @@ class Drift:
reference_data_handler = workflow.start_local_activity_method(
Activities.get_reference_data,
{
**metadata,
'model_name': input_data['model_name']
},
{**metadata, 'model_name': input_data['model_name']},
retry_policy=retry_policy,
start_to_close_timeout=timedelta(seconds=300),
)
@@ -83,8 +80,9 @@ class Drift:
'model_name': input_data['model_name'],
'model_id': input_data['model_id'],
'target_name': target_name,
'drift_metrics': input_data.get('drift_metrics',
['kolmogorov_smirnov', 'jensen_shannon', 'wasserstein']),
'drift_metrics': input_data.get(
'drift_metrics', ['kolmogorov_smirnov', 'jensen_shannon', 'wasserstein']
),
'chunk_period': input_data.get('chunk_period', 'min'),
},
retry_policy=retry_policy,
@@ -99,8 +97,11 @@ class Drift:
'data': drift_data,
'schema': input_data['schema'],
'table_name': input_data['target_table_name'],
'timestamp_conversion': {'column': 'timestamp', 'format': DATETIME_FORMAT_WITH_TZ},
'timestamp_conversion': {
'column': 'timestamp',
'format': DATETIME_FORMAT_WITH_TZ,
},
},
retry_policy=retry_policy,
start_to_close_timeout=timedelta(seconds=300),
)
)

17
laborious/workflows/simple_metrics.py
View File

@@ -4,10 +4,10 @@ with workflow.unsafe.imports_passed_through():
from datetime import timedelta
from typing import Any
from sientia_do.temporal.constants import DATETIME_FORMAT_WITH_TZ
from sientia_do.temporal.policies import retry_policy
from laborious.activities.activities import Activities
from sientia_do.temporal.constants import DATETIME_FORMAT_WITH_TZ
@workflow.defn(name='simple_metrics')
@@ -19,10 +19,10 @@ class SimpleMetrics:
"""
metadata = {
'metadata': {
'schedule_name': input_data['schedule_name'],
'model_name': input_data['model_name'],
'model_id': input_data['model_id'],
'model_name': input_data['model_name'],
'workflow_name': 'simple_metrics',
'schedule_name': input_data['schedule_name'],
}
}
@@ -34,15 +34,15 @@ class SimpleMetrics:
query = f"""
select p."timestamp", p.prediction, ld.value as "target"
from {input_data['schema']}.{input_data['predictions_table_name']} p
inner join {input_data['schema']}.{input_data['data_table_name']} ld
from "{input_data['schema']}"."{input_data['predictions_table_name']}" p
inner join "{input_data['schema']}"."{input_data['data_table_name']}" ld
on p."timestamp" = ld."timestamp"
where
p.model_id = {model_id} and
p.model_id = '{model_id}' and
p.prediction is not null and
ld.variable = '{target_name}' and
ld.value is not null and
p."timestamp" >= NOW() - INTERVAL '{interval_minutes} minutes'
p."timestamp" >= NOW() - INTERVAL {interval_minutes} minutes
order by
p."timestamp" desc;
"""
@@ -74,7 +74,6 @@ class SimpleMetrics:
retry_policy=retry_policy,
start_to_close_timeout=timedelta(seconds=300),
)
if not simple_metrics:
return
@@ -93,4 +92,4 @@ class SimpleMetrics:
},
retry_policy=retry_policy,
start_to_close_timeout=timedelta(seconds=300),
)
)

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

@@ -89,7 +89,6 @@ class FormatAndExportPrediction:
)
if transformed_data is not None:
transformed = await workflow.execute_local_activity_method(
Activities.format_transformed_data,
{

3
pyproject.toml
View File

@@ -48,6 +48,7 @@ ignore = [
"S101", # use of assert (needed for tests)
"S105", # possible hardcoded password (false positives)
"S106", # possible hardcoded password (false positives)
"S608", # potential sql injection (false positives)
"N802", # function name should be lowercase (temporal decorators)
"N806", # variable in function should be lowercase
]
@@ -152,4 +153,4 @@ directory = "htmlcov"
[tool.bandit]
exclude_dirs = ["tests", "venv", ".venv"]
skips = ["B101", "B601"] # Skip assert and shell injection in controlled environments
skips = ["B101", "B601", "B608"] # Skip assert, shell injection, and SQL injection (false positives)

538
tests/laborious/activities/test_model_metrics.py
View File

@@ -1,4 +1,4 @@
from unittest.mock import ANY, AsyncMock, MagicMock, call, patch
from unittest.mock import ANY, AsyncMock, MagicMock, patch
from pandas import DataFrame
from pytest import fixture, mark
@@ -22,7 +22,13 @@ def model_metrics_activity():
model_metrics.send_notification = MagicMock()
model_metrics.send_notification_async = AsyncMock()
model_metrics.emit_metric = AsyncMock()
model_metrics.get_core_labels = MagicMock(return_value={'pod_id': 'test_pod', 'model_name': 'test_model', 'workflow_name': 'test_workflow'})
model_metrics.get_core_labels = MagicMock(
return_value={
'pod_id': 'test_pod',
'model_name': 'test_model',
'workflow_name': 'test_workflow',
}
)
model_metrics.observe_lag = AsyncMock()
model_metrics.pod_id = 'test_pod'
return model_metrics
@@ -60,7 +66,7 @@ async def test_calculate_drift_invalid_chunk_period(model_metrics_activity):
try:
await model_metrics_activity.calculate_drift(input_data)
except ValueError as e:
assert str(e) == "Invalid chunk period: invalid, must be \"min\" or \"s\""
assert str(e) == 'Invalid chunk period: invalid, must be "min" or "s"'
model_metrics_activity.error.assert_called_once_with(
'Invalid chunk period: invalid', metadata['metadata']
)
@@ -76,26 +82,41 @@ async def test_calculate_drift_with_reference_data(
):
# Arrange
mock_to_datetime.return_value.dt.strftime.return_value = '2023-05-26 11:12:27'
mock_to_datetime.return_value.dt.tz_localize.return_value.dt.strftime.return_value = '2023-05-26 11:12:27+00:00'
mock_to_datetime.return_value.dt.tz_localize.return_value.dt.strftime.return_value = (
'2023-05-26 11:12:27+00:00'
)
mock_drift_df = MagicMock()
mock_drift_df.empty = False
mock_drift_df.drop.return_value = mock_drift_df
mock_drift_df.__getitem__.return_value.isin.return_value = [True]
mock_drift_df.__getitem__.return_value = mock_drift_df
mock_drift_df.__getitem__.return_value.dt.tz_localize.return_value.dt.strftime.return_value = '2023-05-26 11:12:27+00:00'
mock_drift_df.__getitem__.return_value.dt.tz_localize.return_value.dt.strftime.return_value = (
'2023-05-26 11:12:27+00:00'
)
mock_drift_df.rename.return_value = mock_drift_df
mock_drift_df.drop_duplicates.return_value = mock_drift_df
mock_drift_df.to_dict.return_value = {'method': ['ks_test'], 'value': [0.5], 'feature': ['feature1'], 'timestamp': ['2023-05-26 11:12:27+00:00'], 'model_id': ['test_model_id'], 'accurate': [True]}
mock_drift_df.to_dict.return_value = [
{
'method': 'ks_test',
'value': 0.5,
'feature': 'feature1',
'timestamp': '2023-05-26 11:12:27+00:00',
'model_id': 'test_model_id',
'accurate': True,
}
]
model_metrics_activity.get_drift_metrics = AsyncMock(return_value=mock_drift_df)
reference_data = DataFrame({
'timestamp': ['2023-05-26 11:12:27'],
'target': [1.0],
'feature1': [1.0],
})
reference_data = DataFrame(
{
'timestamp': ['2023-05-26 11:12:27'],
'target': [1.0],
'feature1': [1.0],
}
)
mock_target_df = MagicMock()
mock_target_df.pivot.return_value = mock_target_df
mock_target_df.index = ['2023-05-26 11:12:27']
@@ -124,16 +145,20 @@ async def test_calculate_drift_with_reference_data(
result = await model_metrics_activity.calculate_drift(input_data)
# Assert
assert isinstance(result, dict)
assert result == mock_drift_df.to_dict.return_value
assert isinstance(result, list)
assert result == mock_drift_df.to_dict.return_value # type: ignore[comparison-overlap]
model_metrics_activity.info.assert_called()
model_metrics_activity.get_drift_metrics.assert_called_once()
# Verify transformations were called
mock_drift_df.drop.assert_called_once_with(columns=['p_value', 'chunk_start_date', 'chunk_end_date'], inplace=True)
mock_drift_df.drop.assert_called_once_with(columns=['p_value'], inplace=True)
mock_drift_df.__getitem__.assert_called()
mock_drift_df.rename.assert_called_once_with(columns={'metric': 'method', 'statistic': 'value'}, inplace=True)
mock_drift_df.drop_duplicates.assert_called_once_with(subset=['timestamp', 'method', 'feature'], keep='first', inplace=True)
mock_drift_df.to_dict.assert_called_once()
mock_drift_df.rename.assert_called_once_with(
columns={'metric': 'method', 'statistic': 'value'}, inplace=True
)
mock_drift_df.drop_duplicates.assert_called_once_with(
subset=['timestamp', 'method', 'feature'], keep='first', inplace=True
)
mock_drift_df.to_dict.assert_called_once_with(orient='records')
@mark.asyncio
@@ -144,18 +169,31 @@ async def test_calculate_drift_without_reference_data(
):
# Arrange
mock_to_datetime.return_value.dt.strftime.return_value = '2023-05-26 11:12:27'
mock_to_datetime.return_value.dt.tz_localize.return_value.dt.strftime.return_value = '2023-05-26 11:12:27+00:00'
mock_to_datetime.return_value.dt.tz_localize.return_value.dt.strftime.return_value = (
'2023-05-26 11:12:27+00:00'
)
mock_drift_df = MagicMock()
mock_drift_df.empty = False
mock_drift_df.drop.return_value = mock_drift_df
mock_drift_df.__getitem__.return_value.isin.return_value = [True]
mock_drift_df.__getitem__.return_value = mock_drift_df
mock_drift_df.__getitem__.return_value.dt.tz_localize.return_value.dt.strftime.return_value = '2023-05-26 11:12:27+00:00'
mock_drift_df.__getitem__.return_value.dt.tz_localize.return_value.dt.strftime.return_value = (
'2023-05-26 11:12:27+00:00'
)
mock_drift_df.rename.return_value = mock_drift_df
mock_drift_df.drop_duplicates.return_value = mock_drift_df
mock_drift_df.to_dict.return_value = {'method': ['ks_test'], 'value': [0.5], 'feature': ['feature1'], 'timestamp': ['2023-05-26 11:12:27+00:00'], 'model_id': ['test_model_id'], 'accurate': [False]}
mock_drift_df.to_dict.return_value = [
{
'method': 'ks_test',
'value': 0.5,
'feature': 'feature1',
'timestamp': '2023-05-26 11:12:27+00:00',
'model_id': 'test_model_id',
'accurate': False,
}
]
model_metrics_activity.get_drift_metrics = AsyncMock(return_value=mock_drift_df)
target_data_dict = {
@@ -163,19 +201,25 @@ async def test_calculate_drift_without_reference_data(
'variable': ['feature1', 'feature1', 'feature1'],
'value': [1.0, 2.0, 3.0],
}
mock_target_df = MagicMock()
mock_target_df.pivot.return_value = mock_target_df
mock_target_df.index = ['2023-05-26 11:12:27', '2023-05-26 11:12:28', '2023-05-26 11:12:29']
mock_target_df.reset_index.return_value = mock_target_df
mock_target_df.dropna.return_value = mock_target_df
mock_target_df.sort_values.return_value = mock_target_df
mock_target_df.head.return_value = DataFrame({'timestamp': ['2023-05-26 11:12:27'], 'feature1': [1.0]})
mock_target_df.__getitem__.return_value.apply.return_value = ['2023-05-26 11:12:27', '2023-05-26 11:12:28', '2023-05-26 11:12:29']
mock_target_df.head.return_value = DataFrame(
{'timestamp': ['2023-05-26 11:12:27'], 'feature1': [1.0]}
)
mock_target_df.__getitem__.return_value.apply.return_value = [
'2023-05-26 11:12:27',
'2023-05-26 11:12:28',
'2023-05-26 11:12:29',
]
mock_target_df.drop.return_value.columns = ['feature1']
mock_dataframe.return_value = mock_target_df
mock_dataframe.side_effect = lambda x=None: mock_target_df if x is not None else mock_target_df
input_data = {
**metadata,
'model_name': 'test_model',
@@ -191,7 +235,7 @@ async def test_calculate_drift_without_reference_data(
result = await model_metrics_activity.calculate_drift(input_data)
# Assert
assert isinstance(result, dict)
assert isinstance(result, list)
assert result == mock_drift_df.to_dict.return_value
model_metrics_activity.warning.assert_called()
model_metrics_activity.send_notification_async.assert_called_once_with(
@@ -203,11 +247,15 @@ async def test_calculate_drift_without_reference_data(
attachment_content=ANY,
)
# Verify transformations were called
mock_drift_df.drop.assert_called_once_with(columns=['p_value', 'chunk_start_date', 'chunk_end_date'], inplace=True)
mock_drift_df.drop.assert_called_once_with(columns=['p_value'], inplace=True)
mock_drift_df.__getitem__.assert_called()
mock_drift_df.rename.assert_called_once_with(columns={'metric': 'method', 'statistic': 'value'}, inplace=True)
mock_drift_df.drop_duplicates.assert_called_once_with(subset=['timestamp', 'method', 'feature'], keep='first', inplace=True)
mock_drift_df.to_dict.assert_called_once()
mock_drift_df.rename.assert_called_once_with(
columns={'metric': 'method', 'statistic': 'value'}, inplace=True
)
mock_drift_df.drop_duplicates.assert_called_once_with(
subset=['timestamp', 'method', 'feature'], keep='first', inplace=True
)
mock_drift_df.to_dict.assert_called_once_with(orient='records')
@mark.asyncio
@@ -218,15 +266,17 @@ async def test_calculate_drift_empty_drift_df(
):
# Arrange
mock_to_datetime.return_value.dt.strftime.return_value = '2023-05-26 11:12:27'
model_metrics_activity.get_drift_metrics = AsyncMock(return_value=DataFrame())
reference_data = DataFrame({
'timestamp': ['2023-05-26 11:12:27'],
'target': [1.0],
'feature1': [1.0],
})
reference_data = DataFrame(
{
'timestamp': ['2023-05-26 11:12:27'],
'target': [1.0],
'feature1': [1.0],
}
)
mock_target_df = MagicMock()
mock_target_df.pivot.return_value = mock_target_df
mock_target_df.index = ['2023-05-26 11:12:27']
@@ -235,7 +285,7 @@ async def test_calculate_drift_empty_drift_df(
mock_target_df.__getitem__.return_value.apply.return_value = ['2023-05-26 11:12:27']
mock_target_df.drop.return_value.columns = ['feature1']
mock_dataframe.return_value = mock_target_df
input_data = {
**metadata,
'model_name': 'test_model',
@@ -255,8 +305,10 @@ async def test_calculate_drift_empty_drift_df(
result = await model_metrics_activity.calculate_drift(input_data)
# Assert
assert result == {}
model_metrics_activity.warning.assert_called_with('No drift metrics found', metadata['metadata'])
assert result == []
model_metrics_activity.warning.assert_called_with(
'No drift metrics found', metadata['metadata']
)
@mark.asyncio
@@ -267,35 +319,37 @@ async def test_calculate_drift_empty_after_timestamp_filter(
):
# Arrange
mock_to_datetime.return_value.dt.strftime.return_value = '2023-05-26 11:12:27'
mock_drift_df = MagicMock()
mock_drift_df.empty = False
mock_drift_df.drop.return_value = mock_drift_df
# Set up __getitem__ to handle filtering - timestamp access returns series with isin=False
# and filtering returns empty DataFrame
mock_timestamp_series = MagicMock()
mock_timestamp_series.isin.return_value = [False]
mock_empty_df = MagicMock()
mock_empty_df.empty = True
def getitem_side_effect(key):
if key == 'timestamp':
return mock_timestamp_series
else:
# This is the filtering operation - return empty DataFrame
return mock_empty_df
mock_drift_df.__getitem__.side_effect = getitem_side_effect
model_metrics_activity.get_drift_metrics = AsyncMock(return_value=mock_drift_df)
reference_data = DataFrame({
'timestamp': ['2023-05-26 11:12:27'],
'target': [1.0],
'feature1': [1.0],
})
reference_data = DataFrame(
{
'timestamp': ['2023-05-26 11:12:27'],
'target': [1.0],
'feature1': [1.0],
}
)
mock_target_df = MagicMock()
mock_target_df.pivot.return_value = mock_target_df
mock_target_df.index = ['2023-05-26 11:12:27']
@@ -324,13 +378,13 @@ async def test_calculate_drift_empty_after_timestamp_filter(
result = await model_metrics_activity.calculate_drift(input_data)
# Assert
assert result == {}
assert result == []
model_metrics_activity.warning.assert_called_with(
'No drift metrics found after dropping rows where timestamp is not in target data',
metadata['metadata']
metadata['metadata'],
)
# Verify transformations were called
mock_drift_df.drop.assert_called_once_with(columns=['p_value', 'chunk_start_date', 'chunk_end_date'], inplace=True)
mock_drift_df.drop.assert_called_once_with(columns=['p_value'], inplace=True)
mock_drift_df.__getitem__.assert_called()
@@ -342,26 +396,41 @@ async def test_calculate_drift_success_min(
):
# Arrange
mock_to_datetime.return_value.dt.strftime.return_value = '2023-05-26 11:12:27'
mock_to_datetime.return_value.dt.tz_localize.return_value.dt.strftime.return_value = '2023-05-26 11:12:27+00:00'
mock_to_datetime.return_value.dt.tz_localize.return_value.dt.strftime.return_value = (
'2023-05-26 11:12:27+00:00'
)
mock_drift_df = MagicMock()
mock_drift_df.empty = False
mock_drift_df.drop.return_value = mock_drift_df
mock_drift_df.__getitem__.return_value.isin.return_value = [True]
mock_drift_df.__getitem__.return_value = mock_drift_df
mock_drift_df.__getitem__.return_value.dt.tz_localize.return_value.dt.strftime.return_value = '2023-05-26 11:12:27+00:00'
mock_drift_df.__getitem__.return_value.dt.tz_localize.return_value.dt.strftime.return_value = (
'2023-05-26 11:12:27+00:00'
)
mock_drift_df.rename.return_value = mock_drift_df
mock_drift_df.drop_duplicates.return_value = mock_drift_df
mock_drift_df.to_dict.return_value = {'method': ['ks_test'], 'value': [0.5], 'feature': ['feature1'], 'timestamp': ['2023-05-26 11:12:27+00:00'], 'model_id': ['test_model_id'], 'accurate': [True]}
mock_drift_df.to_dict.return_value = [
{
'method': 'ks_test',
'value': 0.5,
'feature': 'feature1',
'timestamp': '2023-05-26 11:12:27+00:00',
'model_id': 'test_model_id',
'accurate': True,
}
]
model_metrics_activity.get_drift_metrics = AsyncMock(return_value=mock_drift_df)
reference_data = DataFrame({
'timestamp': ['2023-05-26 11:12:27'],
'target': [1.0],
'feature1': [1.0],
})
reference_data = DataFrame(
{
'timestamp': ['2023-05-26 11:12:27'],
'target': [1.0],
'feature1': [1.0],
}
)
mock_target_df = MagicMock()
mock_target_df.pivot.return_value = mock_target_df
mock_target_df.index = ['2023-05-26 11:12:27']
@@ -390,46 +459,63 @@ async def test_calculate_drift_success_min(
result = await model_metrics_activity.calculate_drift(input_data)
# Assert
assert isinstance(result, dict)
assert result == mock_drift_df.to_dict.return_value
assert isinstance(result, list)
assert result == mock_drift_df.to_dict.return_value # type: ignore[comparison-overlap]
model_metrics_activity.info.assert_called()
model_metrics_activity.get_drift_metrics.assert_called_once()
# Verify transformations were called
mock_drift_df.drop.assert_called_once_with(columns=['p_value', 'chunk_start_date', 'chunk_end_date'], inplace=True)
mock_drift_df.drop.assert_called_once_with(columns=['p_value'], inplace=True)
mock_drift_df.__getitem__.assert_called()
mock_drift_df.rename.assert_called_once_with(columns={'metric': 'method', 'statistic': 'value'}, inplace=True)
mock_drift_df.drop_duplicates.assert_called_once_with(subset=['timestamp', 'method', 'feature'], keep='first', inplace=True)
mock_drift_df.to_dict.assert_called_once()
mock_drift_df.rename.assert_called_once_with(
columns={'metric': 'method', 'statistic': 'value'}, inplace=True
)
mock_drift_df.drop_duplicates.assert_called_once_with(
subset=['timestamp', 'method', 'feature'], keep='first', inplace=True
)
mock_drift_df.to_dict.assert_called_once_with(orient='records')
@mark.asyncio
@patch('laborious.activities.model_metrics.DataFrame')
@patch('laborious.activities.model_metrics.to_datetime')
async def test_calculate_drift_success_s(
mock_to_datetime, mock_dataframe, model_metrics_activity
):
async def test_calculate_drift_success_s(mock_to_datetime, mock_dataframe, model_metrics_activity):
# Arrange
mock_to_datetime.return_value.dt.strftime.return_value = '2023-05-26 11:12:27'
mock_to_datetime.return_value.dt.tz_localize.return_value.dt.strftime.return_value = '2023-05-26 11:12:27+00:00'
mock_to_datetime.return_value.dt.tz_localize.return_value.dt.strftime.return_value = (
'2023-05-26 11:12:27+00:00'
)
mock_drift_df = MagicMock()
mock_drift_df.empty = False
mock_drift_df.drop.return_value = mock_drift_df
mock_drift_df.__getitem__.return_value.isin.return_value = [True]
mock_drift_df.__getitem__.return_value = mock_drift_df
mock_drift_df.__getitem__.return_value.dt.tz_localize.return_value.dt.strftime.return_value = '2023-05-26 11:12:27+00:00'
mock_drift_df.__getitem__.return_value.dt.tz_localize.return_value.dt.strftime.return_value = (
'2023-05-26 11:12:27+00:00'
)
mock_drift_df.rename.return_value = mock_drift_df
mock_drift_df.drop_duplicates.return_value = mock_drift_df
mock_drift_df.to_dict.return_value = {'method': ['ks_test'], 'value': [0.5], 'feature': ['feature1'], 'timestamp': ['2023-05-26 11:12:27+00:00'], 'model_id': ['test_model_id'], 'accurate': [True]}
mock_drift_df.to_dict.return_value = [
{
'method': 'ks_test',
'value': 0.5,
'feature': 'feature1',
'timestamp': '2023-05-26 11:12:27+00:00',
'model_id': 'test_model_id',
'accurate': True,
}
]
model_metrics_activity.get_drift_metrics = AsyncMock(return_value=mock_drift_df)
reference_data = DataFrame({
'timestamp': ['2023-05-26 11:12:27'],
'target': [1.0],
'feature1': [1.0],
})
reference_data = DataFrame(
{
'timestamp': ['2023-05-26 11:12:27'],
'target': [1.0],
'feature1': [1.0],
}
)
mock_target_df = MagicMock()
mock_target_df.pivot.return_value = mock_target_df
mock_target_df.index = ['2023-05-26 11:12:27']
@@ -458,16 +544,20 @@ async def test_calculate_drift_success_s(
result = await model_metrics_activity.calculate_drift(input_data)
# Assert
assert isinstance(result, dict)
assert result == mock_drift_df.to_dict.return_value
assert isinstance(result, list)
assert result == mock_drift_df.to_dict.return_value # type: ignore[comparison-overlap]
model_metrics_activity.info.assert_called()
model_metrics_activity.get_drift_metrics.assert_called_once()
# Verify transformations were called
mock_drift_df.drop.assert_called_once_with(columns=['p_value', 'chunk_start_date', 'chunk_end_date'], inplace=True)
mock_drift_df.drop.assert_called_once_with(columns=['p_value'], inplace=True)
mock_drift_df.__getitem__.assert_called()
mock_drift_df.rename.assert_called_once_with(columns={'metric': 'method', 'statistic': 'value'}, inplace=True)
mock_drift_df.drop_duplicates.assert_called_once_with(subset=['timestamp', 'method', 'feature'], keep='first', inplace=True)
mock_drift_df.to_dict.assert_called_once()
mock_drift_df.rename.assert_called_once_with(
columns={'metric': 'method', 'statistic': 'value'}, inplace=True
)
mock_drift_df.drop_duplicates.assert_called_once_with(
subset=['timestamp', 'method', 'feature'], keep='first', inplace=True
)
mock_drift_df.to_dict.assert_called_once_with(orient='records')
@mark.asyncio
@@ -478,15 +568,19 @@ async def test_calculate_drift_get_drift_metrics_error(
):
# Arrange
mock_to_datetime.return_value.dt.strftime.return_value = '2023-05-26 11:12:27'
model_metrics_activity.get_drift_metrics = AsyncMock(side_effect=Exception('Get drift metrics error'))
reference_data = DataFrame({
'timestamp': ['2023-05-26 11:12:27'],
'target': [1.0],
'feature1': [1.0],
})
model_metrics_activity.get_drift_metrics = AsyncMock(
side_effect=Exception('Get drift metrics error')
)
reference_data = DataFrame(
{
'timestamp': ['2023-05-26 11:12:27'],
'target': [1.0],
'feature1': [1.0],
}
)
mock_target_df = MagicMock()
mock_target_df.pivot.return_value = mock_target_df
mock_target_df.index = ['2023-05-26 11:12:27']
@@ -515,10 +609,9 @@ async def test_calculate_drift_get_drift_metrics_error(
result = await model_metrics_activity.calculate_drift(input_data)
# Assert
assert result == {}
assert result == []
model_metrics_activity.error.assert_called_once_with(
'Error getting drift metrics: Get drift metrics error',
metadata['metadata']
'Error getting drift metrics: Get drift metrics error', metadata['metadata']
)
model_metrics_activity.send_notification_async.assert_called_once_with(
metadata=metadata['metadata'],
@@ -540,30 +633,36 @@ async def test_get_drift_metrics_success(
):
# Arrange
mock_time.return_value = 1000.0
mock_drift_df = DataFrame({
'timestamp': ['2023-05-26 11:12:27'],
'metric': ['ks_test'],
'statistic': [0.5],
'feature': ['feature1'],
})
mock_drift_df = DataFrame(
{
'timestamp': ['2023-05-26 11:12:27'],
'metric': ['ks_test'],
'statistic': [0.5],
'feature': ['feature1'],
}
)
mock_model_analysis.return_value.detect_univariate_drift.return_value = MagicMock()
mock_model_analysis.return_value.detect_multivariate_drift.return_value = MagicMock()
mock_model_analysis.return_value.get_drift_metrics_dataframe.return_value = mock_drift_df
reference_data = DataFrame({
'timestamp': ['2023-05-26 11:12:27'],
'target': [1.0],
'feature1': [1.0],
})
target_data = DataFrame({
'timestamp': ['2023-05-26 11:12:27'],
'target': [1.0],
'feature1': [1.0],
})
reference_data = DataFrame(
{
'timestamp': ['2023-05-26 11:12:27'],
'target': [1.0],
'feature1': [1.0],
}
)
target_data = DataFrame(
{
'timestamp': ['2023-05-26 11:12:27'],
'target': [1.0],
'feature1': [1.0],
}
)
reference_columns = reference_data.drop(
columns=['target', 'timestamp'], errors='ignore'
).columns
@@ -596,21 +695,27 @@ async def test_get_drift_metrics_univariate_error(
):
# Arrange
mock_time.return_value = 1000.0
mock_model_analysis.return_value.detect_univariate_drift.side_effect = Exception('Univariate drift error')
reference_data = DataFrame({
'timestamp': ['2023-05-26 11:12:27'],
'target': [1.0],
'feature1': [1.0],
})
target_data = DataFrame({
'timestamp': ['2023-05-26 11:12:27'],
'target': [1.0],
'feature1': [1.0],
})
mock_model_analysis.return_value.detect_univariate_drift.side_effect = Exception(
'Univariate drift error'
)
reference_data = DataFrame(
{
'timestamp': ['2023-05-26 11:12:27'],
'target': [1.0],
'feature1': [1.0],
}
)
target_data = DataFrame(
{
'timestamp': ['2023-05-26 11:12:27'],
'target': [1.0],
'feature1': [1.0],
}
)
reference_columns = reference_data.drop(
columns=['target', 'timestamp'], errors='ignore'
).columns
@@ -629,28 +734,26 @@ async def test_get_drift_metrics_univariate_error(
except Exception as e:
assert str(e) == 'Univariate drift error'
model_metrics_activity.error.assert_called_once_with(
'Error detecting univariate drift: Univariate drift error',
metadata['metadata']
'Error detecting univariate drift: Univariate drift error', metadata['metadata']
)
model_metrics_activity.emit_metric.assert_called_with(
metric_object=mock_metrics.MODEL_ANALYZE_ERROR_COUNT,
tags=ANY
metric_object=mock_metrics.MODEL_ANALYZE_ERROR_COUNT, tags=ANY
)
else:
raise AssertionError('Expected Exception')
@mark.asyncio
async def test_calculate_simple_metrics_success_all_metrics(
model_metrics_activity
):
async def test_calculate_simple_metrics_success_all_metrics(model_metrics_activity):
# Arrange
target_data = DataFrame({
'timestamp': ['2023-05-26 11:12:27', '2023-05-26 11:12:28', '2023-05-26 11:12:29'],
'target': [1.0, 2.0, 3.0],
'prediction': [1.1, 2.1, 2.9],
})
target_data = DataFrame(
{
'timestamp': ['2023-05-26 11:12:27', '2023-05-26 11:12:28', '2023-05-26 11:12:29'],
'target': [1.0, 2.0, 3.0],
'prediction': [1.1, 2.1, 2.9],
}
)
input_data = {
**metadata,
'model_id': 'test_model_id',
@@ -673,23 +776,23 @@ async def test_calculate_simple_metrics_success_all_metrics(
assert all(data_size == 3 for data_size in result['data_size'].values)
assert all(interval_minutes == 5 for interval_minutes in result['interval_minutes'].values)
model_metrics_activity.info.assert_called_once_with(
'Calculating simple metrics for model test_model_id: [\'rmse\', \'mse\', \'mae\', \'r2\']',
metadata['metadata']
"Calculating simple metrics for model test_model_id: ['rmse', 'mse', 'mae', 'r2']",
metadata['metadata'],
)
model_metrics_activity.debug.assert_called_once()
@mark.asyncio
async def test_calculate_simple_metrics_success_rmse_only(
model_metrics_activity
):
async def test_calculate_simple_metrics_success_rmse_only(model_metrics_activity):
# Arrange
target_data = DataFrame({
'timestamp': ['2023-05-26 11:12:27', '2023-05-26 11:12:28'],
'target': [1.0, 2.0],
'prediction': [1.1, 2.1],
})
target_data = DataFrame(
{
'timestamp': ['2023-05-26 11:12:27', '2023-05-26 11:12:28'],
'target': [1.0, 2.0],
'prediction': [1.1, 2.1],
}
)
input_data = {
**metadata,
'model_id': 'test_model_id',
@@ -709,22 +812,21 @@ async def test_calculate_simple_metrics_success_rmse_only(
assert result['data_size'].values[0] == 2
assert result['interval_minutes'].values[0] == 5
model_metrics_activity.info.assert_called_once_with(
'Calculating simple metrics for model test_model_id: [\'rmse\']',
metadata['metadata']
"Calculating simple metrics for model test_model_id: ['rmse']", metadata['metadata']
)
@mark.asyncio
async def test_calculate_simple_metrics_success_mse_only(
model_metrics_activity
):
async def test_calculate_simple_metrics_success_mse_only(model_metrics_activity):
# Arrange
target_data = DataFrame({
'timestamp': ['2023-05-26 11:12:27', '2023-05-26 11:12:28'],
'target': [1.0, 2.0],
'prediction': [1.1, 2.1],
})
target_data = DataFrame(
{
'timestamp': ['2023-05-26 11:12:27', '2023-05-26 11:12:28'],
'target': [1.0, 2.0],
'prediction': [1.1, 2.1],
}
)
input_data = {
**metadata,
'model_id': 'test_model_id',
@@ -744,22 +846,21 @@ async def test_calculate_simple_metrics_success_mse_only(
assert result['data_size'].values[0] == 2
assert result['interval_minutes'].values[0] == 5
model_metrics_activity.info.assert_called_once_with(
'Calculating simple metrics for model test_model_id: [\'mse\']',
metadata['metadata']
"Calculating simple metrics for model test_model_id: ['mse']", metadata['metadata']
)
@mark.asyncio
async def test_calculate_simple_metrics_success_mae_only(
model_metrics_activity
):
async def test_calculate_simple_metrics_success_mae_only(model_metrics_activity):
# Arrange
target_data = DataFrame({
'timestamp': ['2023-05-26 11:12:27', '2023-05-26 11:12:28'],
'target': [1.0, 2.0],
'prediction': [1.1, 2.1],
})
target_data = DataFrame(
{
'timestamp': ['2023-05-26 11:12:27', '2023-05-26 11:12:28'],
'target': [1.0, 2.0],
'prediction': [1.1, 2.1],
}
)
input_data = {
**metadata,
'model_id': 'test_model_id',
@@ -779,22 +880,21 @@ async def test_calculate_simple_metrics_success_mae_only(
assert result['data_size'].values[0] == 2
assert result['interval_minutes'].values[0] == 5
model_metrics_activity.info.assert_called_once_with(
'Calculating simple metrics for model test_model_id: [\'mae\']',
metadata['metadata']
"Calculating simple metrics for model test_model_id: ['mae']", metadata['metadata']
)
@mark.asyncio
async def test_calculate_simple_metrics_success_r2_only(
model_metrics_activity
):
async def test_calculate_simple_metrics_success_r2_only(model_metrics_activity):
# Arrange
target_data = DataFrame({
'timestamp': ['2023-05-26 11:12:27', '2023-05-26 11:12:28'],
'target': [1.0, 2.0],
'prediction': [1.1, 2.1],
})
target_data = DataFrame(
{
'timestamp': ['2023-05-26 11:12:27', '2023-05-26 11:12:28'],
'target': [1.0, 2.0],
'prediction': [1.1, 2.1],
}
)
input_data = {
**metadata,
'model_id': 'test_model_id',
@@ -814,23 +914,22 @@ async def test_calculate_simple_metrics_success_r2_only(
assert result['data_size'].values[0] == 2
assert result['interval_minutes'].values[0] == 5
model_metrics_activity.info.assert_called_once_with(
'Calculating simple metrics for model test_model_id: [\'r2\']',
metadata['metadata']
"Calculating simple metrics for model test_model_id: ['r2']", metadata['metadata']
)
@mark.asyncio
async def test_calculate_simple_metrics_r2_zero_ss_tot(
model_metrics_activity
):
async def test_calculate_simple_metrics_r2_zero_ss_tot(model_metrics_activity):
# Arrange
# All target values are the same, so ss_tot will be 0
target_data = DataFrame({
'timestamp': ['2023-05-26 11:12:27', '2023-05-26 11:12:28'],
'target': [1.0, 1.0],
'prediction': [1.1, 1.1],
})
target_data = DataFrame(
{
'timestamp': ['2023-05-26 11:12:27', '2023-05-26 11:12:28'],
'target': [1.0, 1.0],
'prediction': [1.1, 1.1],
}
)
input_data = {
**metadata,
'model_id': 'test_model_id',
@@ -851,22 +950,21 @@ async def test_calculate_simple_metrics_r2_zero_ss_tot(
assert result['data_size'].values[0] == 2
assert result['interval_minutes'].values[0] == 5
model_metrics_activity.info.assert_called_once_with(
'Calculating simple metrics for model test_model_id: [\'r2\']',
metadata['metadata']
"Calculating simple metrics for model test_model_id: ['r2']", metadata['metadata']
)
@mark.asyncio
async def test_calculate_simple_metrics_success_multiple_metrics_subset(
model_metrics_activity
):
async def test_calculate_simple_metrics_success_multiple_metrics_subset(model_metrics_activity):
# Arrange
target_data = DataFrame({
'timestamp': ['2023-05-26 11:12:27', '2023-05-26 11:12:28', '2023-05-26 11:12:29'],
'target': [1.0, 2.0, 3.0],
'prediction': [1.1, 2.1, 2.9],
})
target_data = DataFrame(
{
'timestamp': ['2023-05-26 11:12:27', '2023-05-26 11:12:28', '2023-05-26 11:12:29'],
'target': [1.0, 2.0, 3.0],
'prediction': [1.1, 2.1, 2.9],
}
)
input_data = {
**metadata,
'model_id': 'test_model_id',
@@ -887,7 +985,5 @@ async def test_calculate_simple_metrics_success_multiple_metrics_subset(
assert all(data_size == 3 for data_size in result['data_size'].values)
assert all(interval_minutes == 5 for interval_minutes in result['interval_minutes'].values)
model_metrics_activity.info.assert_called_once_with(
'Calculating simple metrics for model test_model_id: [\'rmse\', \'mae\']',
metadata['metadata']
"Calculating simple metrics for model test_model_id: ['rmse', 'mae']", metadata['metadata']
)

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

@@ -193,9 +193,11 @@ def test_get_model_params(mlflow, mlflow_repository):
def test_check_artifact_exists_true(mlflow_repository):
artifact = MagicMock(path='test_artifact')
mlflow_repository.client.list_artifacts.return_value = [artifact]
result = mlflow_repository.check_artifact_exists('run_id', 'test_artifact', metadata['metadata'])
result = mlflow_repository.check_artifact_exists(
'run_id', 'test_artifact', metadata['metadata']
)
assert result is True
mlflow_repository.client.list_artifacts.assert_called_once_with('run_id')
@@ -203,9 +205,11 @@ def test_check_artifact_exists_true(mlflow_repository):
def test_check_artifact_exists_false(mlflow_repository):
artifact = MagicMock(path='other_artifact')
mlflow_repository.client.list_artifacts.return_value = [artifact]
result = mlflow_repository.check_artifact_exists('run_id', 'test_artifact', metadata['metadata'])
result = mlflow_repository.check_artifact_exists(
'run_id', 'test_artifact', metadata['metadata']
)
assert result is False
mlflow_repository.client.list_artifacts.assert_called_once_with('run_id')
@@ -299,16 +303,22 @@ async def test_download_artifacts_error(makedirs, rmtree, path, mlflow_repositor
@patch('laborious.utils.repository.model_repository.mlflow')
@patch('laborious.utils.repository.model_repository.pd')
@patch('laborious.utils.repository.model_repository.StringIO')
async def test_load_artifact_dataframe_success(StringIO, pd, mlflow, mlflow_repository):
async def test_load_artifact_dataframe_success(_stringio, pd, mlflow, mlflow_repository):
mlflow_repository.get_model_run_id = MagicMock(return_value='run_id')
mlflow_repository.check_artifact_exists = MagicMock(return_value=True)
mlflow.artifacts.load_text.return_value = 'col1,col2\n1,2\n3,4'
result = await mlflow_repository.load_artifact_dataframe('model_name', 'artifact_path', metadata['metadata'])
mlflow_repository.get_model_run_id.assert_called_once_with(model_name='model_name', stage='Production')
mlflow_repository.check_artifact_exists.assert_called_once_with('run_id', 'artifact_path', metadata['metadata'])
result = await mlflow_repository.load_artifact_dataframe(
'model_name', 'artifact_path', metadata['metadata']
)
mlflow_repository.get_model_run_id.assert_called_once_with(
model_name='model_name', stage='Production'
)
mlflow_repository.check_artifact_exists.assert_called_once_with(
'run_id', 'artifact_path', metadata['metadata']
)
mlflow.artifacts.load_text.assert_called_once_with('runs:/run_id/artifact_path')
assert result == pd.read_csv.return_value
mlflow_repository.observe_lag.assert_called_once_with(ANY, metrics.MODEL_READ_LAG, ANY)
@@ -321,12 +331,18 @@ async def test_load_artifact_dataframe_success(StringIO, pd, mlflow, mlflow_repo
async def test_load_artifact_dataframe_not_exists(mlflow_repository):
mlflow_repository.get_model_run_id = MagicMock(return_value='run_id')
mlflow_repository.check_artifact_exists = MagicMock(return_value=False)
result = await mlflow_repository.load_artifact_dataframe('model_name', 'artifact_path', metadata['metadata'])
result = await mlflow_repository.load_artifact_dataframe(
'model_name', 'artifact_path', metadata['metadata']
)
assert result is None
mlflow_repository.get_model_run_id.assert_called_once_with(model_name='model_name', stage='Production')
mlflow_repository.check_artifact_exists.assert_called_once_with('run_id', 'artifact_path', metadata['metadata'])
mlflow_repository.get_model_run_id.assert_called_once_with(
model_name='model_name', stage='Production'
)
mlflow_repository.check_artifact_exists.assert_called_once_with(
'run_id', 'artifact_path', metadata['metadata']
)
@pytest.mark.asyncio
@@ -335,10 +351,12 @@ async def test_load_artifact_dataframe_error(mlflow, mlflow_repository):
mlflow_repository.get_model_run_id = MagicMock(return_value='run_id')
mlflow_repository.check_artifact_exists = MagicMock(return_value=True)
mlflow.artifacts.load_text.side_effect = ValueError('error')
with pytest.raises(ValueError):
await mlflow_repository.load_artifact_dataframe('model_name', 'artifact_path', metadata['metadata'])
await mlflow_repository.load_artifact_dataframe(
'model_name', 'artifact_path', metadata['metadata']
)
mlflow_repository.emit_metric.assert_called_once_with(
metric_object=metrics.MODEL_READ_ERROR_COUNT, tags=ANY
)
@@ -1059,7 +1077,9 @@ async def test_create_new_experiment(
data.to_csv.assert_called_once_with('./tmp/artifacts/model_name/retrain_data.csv', index=False)
# Verify prediction_data.to_csv was called with correct arguments
prediction_data.to_csv.assert_called_once()
assert prediction_data.to_csv.call_args[0][0] == './tmp/artifacts/model_name/evaluation_data.csv'
assert (
prediction_data.to_csv.call_args[0][0] == './tmp/artifacts/model_name/evaluation_data.csv'
)
assert prediction_data.to_csv.call_args[1]['index'] is False
mlflow.start_run.assert_called_once_with(
@@ -1089,10 +1109,12 @@ async def test_create_new_experiment(
}
)
mlflow.log_artifact.assert_has_calls([
call('./tmp/artifacts/model_name/retrain_data.csv'),
call('./tmp/artifacts/model_name/evaluation_data.csv'),
])
mlflow.log_artifact.assert_has_calls(
[
call('./tmp/artifacts/model_name/retrain_data.csv'),
call('./tmp/artifacts/model_name/evaluation_data.csv'),
]
)
force_memory_release.assert_called_once_with(mlflow_repository.logger)
@@ -1472,9 +1494,9 @@ def test_get_prediction_data_dataframe(mlflow_repository):
retrain_dataset = DataFrame({'feat_1': [1, 2], 'target': [3, 4]}, index=['idx1', 'idx2'])
prediction_model.predict.return_value = DataFrame({'pred': [5, 6]}, index=['idx1', 'idx2'])
target_name = 'target'
result = mlflow_repository.get_prediction_data(prediction_model, retrain_dataset, target_name)
prediction_model.predict.assert_called_once_with(retrain_dataset)
assert 'prediction' in result.columns
assert 'target' in result.columns
@@ -1487,9 +1509,9 @@ def test_get_prediction_data_array(mlflow_repository):
retrain_dataset = DataFrame({'feat_1': [1, 2], 'target': [3, 4]}, index=['idx1', 'idx2'])
prediction_model.predict.return_value = [5, 6]
target_name = 'target'
result = mlflow_repository.get_prediction_data(prediction_model, retrain_dataset, target_name)
prediction_model.predict.assert_called_once_with(retrain_dataset)
assert 'prediction' in result.columns
assert 'target' in result.columns

15
tests/laborious/workflows/test_drift.py
View File

@@ -1,10 +1,10 @@
from unittest.mock import ANY, AsyncMock, call, patch
from pytest import fixture, mark
from sientia_do.temporal.constants import DATETIME_FORMAT_WITH_TZ
from laborious.activities.activities import Activities
from laborious.workflows.drift import Drift
from sientia_do.temporal.constants import DATETIME_FORMAT_WITH_TZ
@fixture
@@ -45,9 +45,7 @@ async def test_run(workflow_mock: AsyncMock, drift: Drift):
reference_data = {'data': 'test_reference_data'}
drift_data = {'drift': 'test_drift_data'}
workflow_mock.start_local_activity_method.side_effect = [
target_data, reference_data
]
workflow_mock.start_local_activity_method.side_effect = [target_data, reference_data]
workflow_mock.execute_local_activity_method.return_value = drift_data
workflow_mock.execute_activity_method = AsyncMock()
@@ -56,12 +54,13 @@ async def test_run(workflow_mock: AsyncMock, drift: Drift):
await drift.run(input_data)
# Assert - Check start_local_activity_method calls
# Query format matches psycopg2.sql output (identifiers with double quotes, literals with single quotes)
expected_gathering_query = f"""
SELECT *
FROM {input_data['schema']}.{input_data['source_table_name']}
FROM "{input_data['schema']}"."{input_data['source_table_name']}"
WHERE
model_id = {input_data['model_id']} AND
timestamp > NOW() - INTERVAL '{input_data['interval']} minutes'
model_id = '{input_data['model_id']}' AND
timestamp > NOW() - INTERVAL {input_data['interval']} minutes
ORDER BY timestamp ASC
"""
@@ -73,6 +72,7 @@ async def test_run(workflow_mock: AsyncMock, drift: Drift):
**metadata,
'query': expected_gathering_query,
'datetime_columns': ['timestamp', 'created_at'],
'orient': 'records',
},
retry_policy=ANY,
start_to_close_timeout=ANY,
@@ -250,4 +250,3 @@ async def test_run_default_chunk_period(workflow_mock: AsyncMock, drift: Drift):
retry_policy=ANY,
start_to_close_timeout=ANY,
)

25
tests/laborious/workflows/test_simple_metrics.py
View File

@@ -1,10 +1,10 @@
from unittest.mock import ANY, AsyncMock, call, patch
from pytest import fixture, mark
from sientia_do.temporal.constants import DATETIME_FORMAT_WITH_TZ
from laborious.activities.activities import Activities
from laborious.workflows.simple_metrics import SimpleMetrics
from sientia_do.temporal.constants import DATETIME_FORMAT_WITH_TZ
@fixture
@@ -42,9 +42,7 @@ async def test_run(workflow_mock: AsyncMock, simple_metrics: SimpleMetrics):
target_data = {'data': 'test_target_data'}
simple_metrics_data = {'metrics': 'test_simple_metrics_data'}
workflow_mock.execute_local_activity_method.side_effect = [
target_data, simple_metrics_data
]
workflow_mock.execute_local_activity_method.side_effect = [target_data, simple_metrics_data]
workflow_mock.execute_activity_method = AsyncMock()
@@ -52,17 +50,18 @@ async def test_run(workflow_mock: AsyncMock, simple_metrics: SimpleMetrics):
await simple_metrics.run(input_data)
# Assert - Check load_custom_query call
# Query format matches psycopg2.sql output (identifiers with double quotes, literals with single quotes)
expected_query = f"""
select p."timestamp", p.prediction, ld.value as "target"
from {input_data['schema']}.{input_data['predictions_table_name']} p
inner join {input_data['schema']}.{input_data['data_table_name']} ld
from "{input_data['schema']}"."{input_data['predictions_table_name']}" p
inner join "{input_data['schema']}"."{input_data['data_table_name']}" ld
on p."timestamp" = ld."timestamp"
where
p.model_id = {input_data['model_id']} and
p.model_id = '{input_data['model_id']}' and
p.prediction is not null and
ld.variable = '{input_data['model_config']['target']}' and
ld.value is not null and
p."timestamp" >= NOW() - INTERVAL '{input_data['interval_minutes']} minutes'
p."timestamp" >= NOW() - INTERVAL {input_data['interval_minutes']} minutes
order by
p."timestamp" desc;
"""
@@ -75,6 +74,7 @@ async def test_run(workflow_mock: AsyncMock, simple_metrics: SimpleMetrics):
**metadata,
'query': expected_query,
'datetime_columns': ['timestamp'],
'orient': 'records',
},
retry_policy=ANY,
start_to_close_timeout=ANY,
@@ -159,9 +159,7 @@ async def test_run_empty_simple_metrics(workflow_mock: AsyncMock, simple_metrics
target_data = {'data': 'test_target_data'}
simple_metrics_data = None
workflow_mock.execute_local_activity_method.side_effect = [
target_data, simple_metrics_data
]
workflow_mock.execute_local_activity_method.side_effect = [target_data, simple_metrics_data]
workflow_mock.execute_activity_method = AsyncMock()
@@ -193,9 +191,7 @@ async def test_run_default_metrics(workflow_mock: AsyncMock, simple_metrics: Sim
target_data = {'data': 'test_target_data'}
simple_metrics_data = {'metrics': 'test_simple_metrics_data'}
workflow_mock.execute_local_activity_method.side_effect = [
target_data, simple_metrics_data
]
workflow_mock.execute_local_activity_method.side_effect = [target_data, simple_metrics_data]
workflow_mock.execute_activity_method = AsyncMock()
@@ -225,4 +221,3 @@ async def test_run_default_metrics(workflow_mock: AsyncMock, simple_metrics: Sim
),
]
)

2
validate.sh
View File

@@ -87,7 +87,7 @@ if ! run_step "3. Type Checking (mypy)" "mypy laborious/"; then
fi
# Step 4: Security Analysis (Bandit)
if ! run_step "4. Security Analysis (Bandit)" "bandit -r laborious/ -ll -q"; then
if ! run_step "4. Security Analysis (Bandit)" "bandit -c pyproject.toml -r laborious/ -ll -q"; then
FAILED_STEPS+=("Security Analysis")
fi