vitor-aignosi
f0fb9b854e
Enhance validation script and refactor code in various modules - Updated the validation script to include automatic code formatting and linting fixes using Ruff. - Removed the `clean_tmp_files` method from the Gates class to streamline functionality. - Simplified conditional checks in the OpcRepository for better clarity and error handling. - Added model ID to the minimal retrain workflow for improved tracking. - Introduced new test cases for error handling in MLFlow and storage operations, ensuring robustness in repository interactions.
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,
|
|
'schema': input_data['schema'],
|
|
'table_name': input_data['table_name'],
|
|
},
|
|
retry_policy=retry_policy,
|
|
start_to_close_timeout=timedelta(seconds=600),
|
|
)
|