vitor-aignosi
d2b365a34d
Refactor data handling in various modules to ensure DataFrame consistency - Replaced direct DataFrame instantiation with `ensure_dataframe` utility in Gates, MLFlow, OPC, and ModelMetrics classes to standardize data handling. - Updated return types in several asynchronous methods to return DataFrames instead of dictionaries for improved usability. - Adjusted data export processes in workflows to convert DataFrames to dictionaries with `to_dict(orient='records')` for compatibility with downstream systems.
137 lines
5.2 KiB
Python
137 lines
5.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.policies import retry_policy
|
|
|
|
from laborious.activities.activities import Activities
|
|
|
|
|
|
@workflow.defn(name='minimal_retrain')
|
|
class MinimalRetrain:
|
|
"""
|
|
Automated model retraining workflow for the Laborious system.
|
|
|
|
This workflow implements a complete model retraining pipeline that loads
|
|
training data, executes model retraining, updates production models,
|
|
and maintains comprehensive audit trails. It's designed for automated
|
|
model lifecycle management with minimal manual intervention.
|
|
|
|
The workflow provides a robust retraining process with:
|
|
- Automated data loading from configured data sources
|
|
- MLFlow model retraining with quality validation
|
|
- Production model updates with version control
|
|
- Comprehensive reporting and audit trail maintenance
|
|
- Error handling and notification integration
|
|
"""
|
|
|
|
@workflow.run
|
|
async def run(self, input_data: dict[str, Any]):
|
|
"""
|
|
Execute the automated model retraining workflow.
|
|
|
|
This method orchestrates the complete model retraining process by:
|
|
1. Loading training data using the provided custom SQL query
|
|
2. Executing MLFlow model retraining with the loaded data
|
|
3. Updating production models with newly trained versions
|
|
4. Persisting comprehensive retraining reports to database
|
|
|
|
The method implements comprehensive error handling and ensures all
|
|
required parameters are properly configured before proceeding.
|
|
|
|
Args:
|
|
input_data: Complete configuration for the retraining workflow
|
|
Required keys:
|
|
- schedule_name (str): Schedule identifier for the retraining
|
|
- model_name (str): Name of the ML model to retrain
|
|
- model_id (int): Unique identifier for the model version
|
|
- query (str): SQL query for training data loading
|
|
- schema (str, optional): Database schema for report storage
|
|
- table_name (str, optional): Target table for retraining reports
|
|
- datetime_columns (list[str], optional): Columns to treat as datetime
|
|
|
|
Returns:
|
|
None: The workflow completes successfully when all steps finish
|
|
|
|
Raises:
|
|
Exception: If any required parameters are missing or if the workflow fails
|
|
during data loading, retraining, or model update operations
|
|
"""
|
|
|
|
metadata = {
|
|
'metadata': {
|
|
'schedule_name': input_data['schedule_name'],
|
|
'model_name': input_data['model_name'],
|
|
'model_id': input_data['model_id'],
|
|
'workflow_name': 'minimal_retrain',
|
|
}
|
|
}
|
|
|
|
model_name = input_data['model_name']
|
|
model_config = input_data.get('model_config', {})
|
|
|
|
storage_result = await workflow.execute_activity_method(
|
|
Activities.query_to_minio,
|
|
{
|
|
**metadata,
|
|
'query': input_data['query'],
|
|
'datetime_columns': input_data.get('datetime_columns', []),
|
|
'model_name': model_name,
|
|
'object_prefix': f'retrain_datasets/{model_name}/data',
|
|
},
|
|
retry_policy=retry_policy,
|
|
start_to_close_timeout=timedelta(seconds=600),
|
|
)
|
|
|
|
if not storage_result['success']:
|
|
return
|
|
|
|
experiment_response = await workflow.execute_activity_method(
|
|
Activities.retrain_model,
|
|
{
|
|
**metadata,
|
|
'object_key': storage_result['object_key'],
|
|
'model_name': model_name,
|
|
'model_config': model_config,
|
|
},
|
|
retry_policy=retry_policy,
|
|
start_to_close_timeout=timedelta(hours=1),
|
|
)
|
|
|
|
if experiment_response['success']:
|
|
update_report = await workflow.execute_activity_method(
|
|
Activities.update_production_model,
|
|
{**metadata, 'model_name': model_name, **experiment_response},
|
|
retry_policy=retry_policy,
|
|
start_to_close_timeout=timedelta(seconds=60),
|
|
)
|
|
else:
|
|
update_report = {}
|
|
|
|
report = await workflow.execute_local_activity_method(
|
|
Activities.format_retrain_report,
|
|
{
|
|
**metadata,
|
|
'experiment_response': experiment_response,
|
|
'model_name': model_name,
|
|
'model_id': input_data['model_id'],
|
|
'update_report': update_report,
|
|
},
|
|
retry_policy=retry_policy,
|
|
start_to_close_timeout=timedelta(seconds=60),
|
|
)
|
|
|
|
await workflow.execute_activity_method(
|
|
Activities.export_data_to_postgres,
|
|
{
|
|
**metadata,
|
|
'data': report.to_dict(orient='records'),
|
|
'schema': input_data['schema'],
|
|
'table_name': input_data['table_name'],
|
|
},
|
|
retry_policy=retry_policy,
|
|
start_to_close_timeout=timedelta(seconds=600),
|
|
)
|