7e4e048ca3
Update sientia-mlops-library dependency to version 0.40.5 and fix SQL interval formatting in Drift and SimpleMetrics workflows for improved query accuracy.
108 lines
3.7 KiB
Python
108 lines
3.7 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='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',
|
|
}
|
|
}
|
|
|
|
print(f'Input data: {input_data}', metadata)
|
|
|
|
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'],
|
|
'orient': 'records',
|
|
},
|
|
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),
|
|
)
|