vitor-aignosi
c8b809f189
Update dependencies and refactor data handling in various modules - Updated sientia-dataops-library dependency version from 1.5.3 to 1.5.4 in requirements files. - Updated sientia-mlops-library dependency version from 0.39.0 to 0.40.2 in requirements files. - Refactored return types in Gates, MLFlow, and ModelMetrics classes to return dictionaries instead of DataFrames for improved compatibility with downstream systems. - Removed the temporal_codec module as it is no longer needed for DataFrame serialization. - Adjusted data handling in the Drift workflow to ensure proper data structure is maintained.
105 lines
3.6 KiB
Python
105 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',
|
|
}
|
|
}
|
|
|
|
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'],
|
|
},
|
|
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),
|
|
) |