gitea_admin/sientia-dataops-model-manager
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

SIENTIAPDE-1243: Refactor and enhance model manager activities and workflows

Browse Source 
This commit includes several changes:

- Reorganized imports and class inheritance in activities.py, gates.py and mlflow.py for better readability and maintainability.
- Improved error handling and logging in gates.py and mlflow.py.
- Added input validation and filtering in gates.py to ensure data quality.
- Enhanced prediction formatting and storage policy management in gates.py.
- Updated metrics.py to use consistent naming conventions and labels.
- Refactored connectors_config.py to use type hints and improve code clarity.
- Updated conditional and MLFlow filters for better data quality checks.
- Improved model repository logic for retraining and updating models.
- Enhanced worker.py to include SDK metrics and improved error handling.
- Refactored workflows for better modularity and error handling.
- Updated tests to reflect the changes and improve test coverage.
This commit is contained in:
Bruno Domingues
2025-10-01 17:28:57 -03:00
parent b102f79087
commit dfc190c818
24 changed files with 1482 additions and 1399 deletions
Download Patch File Download Diff File Expand all files Collapse all files

34
model_manager/workflows/minimal_retrain.py
View File

@@ -1,14 +1,16 @@
from temporalio import workflow
with workflow.unsafe.imports_passed_through():
from model_manager.activities.activities import Activities
from typing import Any
from sientia_do.temporal.policies import retry_policy
from datetime import timedelta
from typing import Any
from sientia_do.temporal.policies import retry_policy
from model_manager.activities.activities import Activities
@workflow.defn(name="minimal_retrain")
class MinimalRetrain():
@workflow.defn(name='minimal_retrain')
class MinimalRetrain:
"""
Automated model retraining workflow for the Model Manager system.
@@ -63,7 +65,7 @@ class MinimalRetrain():
'schedule_name': input_data['schedule_name'],
'model_name': input_data['model_name'],
'model_id': input_data['model_id'],
'workflow_name': 'minimal_retrain'
'workflow_name': 'minimal_retrain',
}
}
@@ -74,21 +76,17 @@ class MinimalRetrain():
{
**metadata,
'query': input_data['query'],
'datetime_columns': input_data.get('datetime_columns', [])
'datetime_columns': input_data.get('datetime_columns', []),
},
retry_policy=retry_policy,
start_to_close_timeout=timedelta(seconds=60)
start_to_close_timeout=timedelta(seconds=60),
)
experiment_response = await workflow.execute_activity_method(
Activities.retrain_model,
{
**metadata,
'data': data,
'model_name': model_name
},
{**metadata, 'data': data, 'model_name': model_name},
retry_policy=retry_policy,
start_to_close_timeout=timedelta(seconds=60)
start_to_close_timeout=timedelta(seconds=60),
)
report = await workflow.execute_activity_method(
@@ -97,10 +95,10 @@ class MinimalRetrain():
**metadata,
'model_name': model_name,
'model_id': input_data['model_id'],
**experiment_response
**experiment_response,
},
retry_policy=retry_policy,
start_to_close_timeout=timedelta(seconds=60)
start_to_close_timeout=timedelta(seconds=60),
)
await workflow.execute_activity_method(
@@ -109,8 +107,8 @@ class MinimalRetrain():
**metadata,
'data': report,
'schema': input_data['schema'],
'table_name': input_data['table_name']
'table_name': input_data['table_name'],
},
retry_policy=retry_policy,
start_to_close_timeout=timedelta(seconds=60)
start_to_close_timeout=timedelta(seconds=60),
)

47
model_manager/workflows/predictions_batch.py
View File

@@ -1,14 +1,16 @@
from temporalio import workflow
with workflow.unsafe.imports_passed_through():
from model_manager.activities.activities import Activities
from typing import Any
from sientia_do.temporal.policies import retry_policy
from datetime import timedelta
from typing import Any
from sientia_do.temporal.policies import retry_policy
from model_manager.activities.activities import Activities
@workflow.defn(name="predictions_batch")
class PredictionsBatch():
@workflow.defn(name='predictions_batch')
class PredictionsBatch:
"""
Main batch prediction workflow for the Model Manager system.
@@ -74,7 +76,7 @@ class PredictionsBatch():
'schedule_name': input_data['schedule_name'],
'model_name': input_data['model_name'],
'model_id': input_data['model_id'],
'workflow_name': 'predictions_batch'
'workflow_name': 'predictions_batch',
}
}
@@ -84,10 +86,10 @@ class PredictionsBatch():
{
**metadata,
'query': input_data['query'],
'datetime_columns': input_data.get('datetime_columns', [])
'datetime_columns': input_data.get('datetime_columns', []),
},
retry_policy=retry_policy,
start_to_close_timeout=timedelta(seconds=300)
start_to_close_timeout=timedelta(seconds=300),
)
# Prepare input for prediction_process workflow
@@ -98,28 +100,17 @@ class PredictionsBatch():
'table_name': input_data['table_name'],
'model_id': input_data['model_id'],
'model_name': input_data['model_name'],
'input_filters': input_data.get('input_filters', {
'EMPTY_DATA': {
'POLICY': 'STOP'
}
}),
'mlflow_transform_filters': input_data.get('mlflow_transform_filters', {
'API_ERROR': {
'POLICY': 'STOP'
}
}),
'mlflow_predict_filters': input_data.get('mlflow_predict_filters', {
'API_ERROR': {
'POLICY': 'STOP'
}
}),
'input_filters': input_data.get('input_filters', {'EMPTY_DATA': {'POLICY': 'STOP'}}),
'mlflow_transform_filters': input_data.get(
'mlflow_transform_filters', {'API_ERROR': {'POLICY': 'STOP'}}
),
'mlflow_predict_filters': input_data.get(
'mlflow_predict_filters', {'API_ERROR': {'POLICY': 'STOP'}}
),
'model_config': input_data.get('model_config', {}),
'path_priority': input_data.get('path_priority', ['STOP', 'CONTINUE', 'REPEAT']),
'prediction_store_policy': input_data.get(
'prediction_store_policy', 'lts:1')
'prediction_store_policy': input_data.get('prediction_store_policy', 'lts:1'),
}
# Execute prediction process workflow
await workflow.execute_child_workflow(
'prediction_process', prediction_input)
await workflow.execute_child_workflow('prediction_process', prediction_input)

34
model_manager/workflows/sub_workflows/format_and_export_prediction.py
View File

@@ -1,15 +1,17 @@
from temporalio import workflow
with workflow.unsafe.imports_passed_through():
from model_manager.activities.activities import Activities
from typing import Any
from datetime import timedelta
from sientia_do.temporal.policies import retry_policy
from typing import Any
from sientia_do.temporal.constants import DATETIME_FORMAT_WITH_TZ
from sientia_do.temporal.policies import retry_policy
from model_manager.activities.activities import Activities
@workflow.defn(name="format_and_export_prediction")
class FormatAndExportPrediction():
@workflow.defn(name='format_and_export_prediction')
class FormatAndExportPrediction:
"""
Data formatting and export workflow for prediction results.
@@ -75,10 +77,10 @@ class FormatAndExportPrediction():
'timestamp': input_data['timestamp'],
'model_id': input_data['model_id'],
'prediction_confidence': prediction_confidence,
'prediction_store_policy': input_data['prediction_store_policy']
'prediction_store_policy': input_data['prediction_store_policy'],
},
retry_policy=retry_policy,
start_to_close_timeout=timedelta(seconds=60)
start_to_close_timeout=timedelta(seconds=60),
)
else:
@@ -90,10 +92,10 @@ class FormatAndExportPrediction():
'timestamp': input_data['timestamp'],
'model_id': input_data['model_id'],
'prediction_confidence': prediction_confidence,
'comment': input_data['comment']
'comment': input_data['comment'],
},
retry_policy=retry_policy,
start_to_close_timeout=timedelta(seconds=60)
start_to_close_timeout=timedelta(seconds=60),
)
# write to postgres
@@ -104,21 +106,15 @@ class FormatAndExportPrediction():
'schema': input_data['schema'],
'table_name': input_data['table_name'],
'data': prediction,
'timestamp_conversion': {
'column': 'timestamp',
'format': DATETIME_FORMAT_WITH_TZ
}
'timestamp_conversion': {'column': 'timestamp', 'format': DATETIME_FORMAT_WITH_TZ},
},
retry_policy=retry_policy,
start_to_close_timeout=timedelta(seconds=60)
start_to_close_timeout=timedelta(seconds=60),
)
await workflow.execute_activity_method(
Activities.write_metrics,
{
**metadata,
'prediction': prediction
},
{**metadata, 'prediction': prediction},
retry_policy=retry_policy,
start_to_close_timeout=timedelta(seconds=60)
start_to_close_timeout=timedelta(seconds=60),
)

55
model_manager/workflows/sub_workflows/prediction_process.py
View File

@@ -1,14 +1,16 @@
from temporalio import workflow
with workflow.unsafe.imports_passed_through():
from model_manager.activities.activities import Activities
from typing import Any
from sientia_do.temporal.policies import retry_policy
from datetime import timedelta
from typing import Any
from sientia_do.temporal.policies import retry_policy
from model_manager.activities.activities import Activities
@workflow.defn(name="prediction_process")
class PredictionProcess():
@workflow.defn(name='prediction_process')
class PredictionProcess:
"""
Core prediction processing workflow for the Model Manager system.
@@ -84,10 +86,7 @@ class PredictionProcess():
# Get last timestamp for incremental processing
last_timestamp = await workflow.execute_local_activity_method(
Activities.get_last_timestamp,
{
**metadata,
'data': data
},
{**metadata, 'data': data},
retry_policy=retry_policy,
start_to_close_timeout=timedelta(minutes=1),
)
@@ -97,7 +96,7 @@ class PredictionProcess():
**metadata,
'filters': input_data['input_filters'],
'data': data,
'path_priority': input_data['path_priority']
'path_priority': input_data['path_priority'],
}
path_flag, confidence, comment = await workflow.execute_local_activity_method(
@@ -116,12 +115,7 @@ class PredictionProcess():
# Request MLFlow model transformation
response_data = await workflow.execute_local_activity_method(
Activities.request_transform,
{
**metadata,
'data': data,
'model_name': model_name,
'model_config': model_config
},
{**metadata, 'data': data, 'model_name': model_name, 'model_config': model_config},
retry_policy=retry_policy,
start_to_close_timeout=timedelta(minutes=5),
)
@@ -134,7 +128,7 @@ class PredictionProcess():
'filters': input_data['mlflow_transform_filters'],
'data': response_data,
'type': 'transform',
'path_priority': input_data['path_priority']
'path_priority': input_data['path_priority'],
},
retry_policy=retry_policy,
start_to_close_timeout=timedelta(minutes=1),
@@ -155,7 +149,7 @@ class PredictionProcess():
'filters': input_data['mlflow_transform_filters'],
'data': transformed_data,
'type': 'transform',
'path_priority': input_data['path_priority']
'path_priority': input_data['path_priority'],
},
retry_policy=retry_policy,
start_to_close_timeout=timedelta(minutes=1),
@@ -172,7 +166,7 @@ class PredictionProcess():
**metadata,
'data': transformed_data,
'model_name': model_name,
'model_config': model_config
'model_config': model_config,
},
retry_policy=retry_policy,
start_to_close_timeout=timedelta(minutes=5),
@@ -186,7 +180,7 @@ class PredictionProcess():
'filters': input_data['mlflow_predict_filters'],
'data': response_data,
'type': 'predict',
'path_priority': input_data['path_priority']
'path_priority': input_data['path_priority'],
},
retry_policy=retry_policy,
start_to_close_timeout=timedelta(minutes=1),
@@ -213,12 +207,19 @@ class PredictionProcess():
'schema': input_data['schema'],
'table_name': input_data['table_name'],
'comment': comment,
'prediction_store_policy': input_data['prediction_store_policy']
}
'prediction_store_policy': input_data['prediction_store_policy'],
},
)
async def path_flag_handler(self, data: dict, path_flag: str, input_data: dict,
confidence: int, last_timestamp: str, comment: str) -> bool:
async def path_flag_handler(
self,
data: dict,
path_flag: str,
input_data: dict,
confidence: int,
last_timestamp: str,
comment: str,
) -> bool:
"""
Handle path decisions based on filter results and confidence levels.
@@ -264,7 +265,7 @@ class PredictionProcess():
'schema': schema,
'table_name': table_name,
'model': model_id,
'last_timestamp': last_timestamp
'last_timestamp': last_timestamp,
},
retry_policy=retry_policy,
start_to_close_timeout=timedelta(minutes=1),
@@ -286,8 +287,8 @@ class PredictionProcess():
'schema': schema,
'table_name': table_name,
'comment': comment,
'prediction_store_policy': input_data['prediction_store_policy']
}
'prediction_store_policy': input_data['prediction_store_policy'],
},
)
return True