SIENTIAPDE-1243: Refactor and enhance model manager activities and workflows
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:
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user