Files
sientia-dataops-laborious-l…/laborious/workflows/simple_metrics.py
vitor-aignosi 7e4e048ca3 SIENTIAPDE-1273
Update sientia-mlops-library dependency to version 0.40.5 and fix SQL interval formatting in Drift and SimpleMetrics workflows for improved query accuracy.
2025-11-19 08:47:48 -03:00

96 lines
3.2 KiB
Python

from temporalio import workflow
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
@workflow.defn(name='simple_metrics')
class SimpleMetrics:
@workflow.run
async def run(self, input_data: dict[str, Any]):
"""
Execute the simple metrics workflow.
"""
metadata = {
'metadata': {
'model_id': input_data['model_id'],
'model_name': input_data['model_name'],
'workflow_name': 'simple_metrics',
'schedule_name': input_data['schedule_name'],
}
}
model_id = input_data['model_id']
interval_minutes = input_data['interval_minutes']
model_config = input_data['model_config']
target_name = model_config['target']
query = f"""
select p."timestamp", p.prediction, ld.value as "target"
from "{input_data['schema']}"."{input_data['predictions_table_name']}" p
inner join "{input_data['schema']}"."{input_data['data_table_name']}" ld
on p."timestamp" = ld."timestamp"
where
p.model_id = '{model_id}' and
p.prediction is not null and
ld.variable = '{target_name}' and
ld.value is not null and
p."timestamp" >= NOW() - INTERVAL '{interval_minutes} minutes'
order by
p."timestamp" desc;
"""
target_data = await workflow.execute_local_activity_method(
Activities.load_custom_query,
{
**metadata,
'query': query,
'datetime_columns': ['timestamp'],
'orient': 'records',
},
retry_policy=retry_policy,
start_to_close_timeout=timedelta(seconds=300),
)
if not target_data:
return
simple_metrics = await workflow.execute_local_activity_method(
Activities.calculate_simple_metrics,
{
**metadata,
'model_id': model_id,
'target_data': target_data,
'metrics': input_data.get('metrics', ['rmse', 'mse', 'mae', 'r2']),
'interval_minutes': interval_minutes,
},
retry_policy=retry_policy,
start_to_close_timeout=timedelta(seconds=300),
)
if not simple_metrics:
return
await workflow.execute_activity_method(
Activities.export_data_to_postgres,
{
**metadata,
'data': simple_metrics,
'schema': input_data['schema'],
'table_name': input_data['target_table_name'],
'timestamp_conversion': {
'column': 'timestamp',
'format': DATETIME_FORMAT_WITH_TZ,
},
},
retry_policy=retry_policy,
start_to_close_timeout=timedelta(seconds=300),
)