Files
sientia-dataops-laborious-l…/laborious/workflows/minimal_retrain.py
vitor-aignosi 3a06f0ca39 SIENTIAPDE-1199
Update dependencies and refactor logging imports for observability

- Updated the sientia-dataops-library dependency version to 1.4.0 in requirements.txt.
- Changed image tag in values.yaml from 0.3.2 to 0.4.1.
- Refactored logging imports across multiple files to use the new observability module instead of the temporal.utils.logger.
- Updated retry policy imports in workflow files to reflect the new module structure.
2025-08-19 10:42:27 -03:00

94 lines
3.1 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():
@workflow.run
async def run(self, input_data: dict[str, Any]):
"""
This workflow runs a minimal retrain of a model.
The workflow executes in four steps:
1. Loads the data from the database
2. Formats the data and perform the retrain
3. Updates the production model
4. Saves a model
Args:
- input_data (dict[str, Any]): The input data for the workflow.
- schedule_name (str): The name of the schedule.
- model_name (str): The name of the model.
- model_id (int): The id of the model.
- query (str): The SQL query to be executed to load data.
- schema (dict, optional): The schema to store the report.
- table_name (str, optional): The name of the table to store report.
Returns:
None
Raises:
Exception: If any of the required parameters are missing or if the workflow fails.
"""
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'],
},
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)
)