Files
sientia-dataops-laborious_t…/laborious/workflows/drift.py
vitor-aignosi 64f0747e8e SIENTIAPDE-1273
Update version and enhance metrics calculation in Laborious system

- Updated image tag in values.yaml from 1.1.0 to 1.1.1.
- Modified GITHUB_BRANCH environment variable for consistency.
- Added a new method `calculate_simple_metrics` in model_metrics.py to compute various model performance metrics including RMSE, MSE, MAE, and R2.
- Integrated the new metrics calculation into the worker setup, allowing for concurrent processing of simple metrics.
- Updated tests to cover the new metrics calculation functionality, ensuring comprehensive validation of the implementation.
2025-11-14 15:43:28 -03:00

103 lines
3.6 KiB
Python

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',
}
}
model_config = input_data['model_config']
target_name = model_config['target']
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': target_name,
'drift_metrics': input_data.get('drift_metrics',
['kolmogorov_smirnov', 'jensen_shannon', 'wasserstein']),
'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),
)