vitor-aignosi
9cdbc0afda
Update dependencies and enhance datetime handling - Bump sientia-dataops-library version in requirements.txt to 1.4.4. - Update image tag in values.yaml to 0.4.4. - Add 'datetime_columns' to metadata in MinimalRetrain and PredictionsBatch workflows for improved data handling.
95 lines
3.1 KiB
Python
95 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'],
|
|
'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)
|
|
)
|