Files
sientia-dataops-model-manager/model-manager/workflows/minimal_retrain.py
Bruno Domingues 93d0849c80 SIENTIAPDE-1243: Initial commit of the model manager project, adding core files and configurations.
This commit introduces the initial project structure, including:
- .env.example: Example environment configuration.
- .github/workflows/quality-gate.yml: CI workflow for quality checks.
- .gitignore: Specifies intentionally untracked files that Git should ignore.
- Makefile: Automation of tasks like docker builds.
- README.md: Project documentation.
- Source code for model management, activities, utils, worker and workflows.
- Test suite.
- Dockerfile for the simulator.
- sonar-project.properties: SonarQube configuration file.
- values.yaml: Helm chart values for deployment.
2025-09-30 14:55:38 -03:00

117 lines
4.4 KiB
Python

from temporalio import workflow
with workflow.unsafe.imports_passed_through():
from laborious.activities.activities import Activities
from typing import Any
from sientia_do.temporal.policies import retry_policy
from datetime import timedelta
@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']
data = await workflow.execute_local_activity_method(
Activities.load_custom_query,
{
**metadata,
'query': input_data['query'],
'datetime_columns': input_data.get('datetime_columns', [])
},
retry_policy=retry_policy,
start_to_close_timeout=timedelta(seconds=60)
)
experiment_response = await workflow.execute_activity_method(
Activities.retrain_model,
{
**metadata,
'data': data,
'model_name': model_name
},
retry_policy=retry_policy,
start_to_close_timeout=timedelta(seconds=60)
)
report = await workflow.execute_activity_method(
Activities.update_production_model,
{
**metadata,
'model_name': model_name,
'model_id': input_data['model_id'],
**experiment_response
},
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=60)
)