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

SIENTIAPDE-1253: Refactor training workflow and activities to raise exceptions on failure

Browse Source 
This commit refactors the training workflow and associated activities to raise exceptions on failure instead of returning success/failure dictionaries. This allows the Temporal workflow to handle errors more effectively and ensures that the workflow stops when a critical error occurs.

Key changes:

- The train_model workflow is introduced to orchestrate the entire training process, including parameter validation, data download, model training, and model saving.
- The validate_train_params activity is added to validate and convert training parameters.
- The train_model and save_model activities are updated to raise exceptions on failure.
- The ExperimentStatus enum is updated to include a new status for orchestrator validation errors.
- The tests are updated to reflect the new exception-based error handling.
- The activities now return the TrainModelResult directly instead of a dictionary.
This commit is contained in:
Bruno Domingues
2025-10-15 15:19:29 -03:00
parent 61267ec49d
commit 8ea98360c3
8 changed files with 1472 additions and 129 deletions
Download Patch File Download Diff File Expand all files Collapse all files

41
model_manager/activities/mlflow.py
View File

@@ -13,6 +13,7 @@ with workflow.unsafe.imports_passed_through():
from sientia_do.temporal.activities.base import BaseActivity
from sientia_do.temporal.constants import DATETIME_FORMAT, DATETIME_FORMAT_WITH_TZ
from model_manager.utils.models.train_model_result import TrainModelResult
from model_manager.utils.repository.model_repository import MLFlowRepository
@@ -346,19 +347,14 @@ class MLFlow(BaseActivity):
raise e
@activity.defn(name='save_model')
async def save_model(self, input_data: dict[str, Any]) -> dict[str, Any]:
async def save_model(self, input_data: dict[str, Any]) -> TrainModelResult:
"""
Save a trained ML model and its artifacts to MLflow with comprehensive error handling.
Save a trained ML model and its artifacts to MLflow.
This activity orchestrates the complete model saving pipeline:
1. Generates the next run name for the experiment
2. Creates and organizes artifacts (reports, data files)
3. Logs model, parameters, metrics, and artifacts to MLflow
4. Returns success/failure status with results or error message
The activity does NOT raise exceptions on failure - it catches all errors,
sends notifications, and returns a failure status. This allows the workflow
to handle the error gracefully and update the database accordingly.
Args:
input_data: Configuration for model saving operation
@@ -367,23 +363,17 @@ class MLFlow(BaseActivity):
- train_result (TrainModelResult): Training result with model and metrics
Returns:
dict: Save result with the following structure:
{
'success': bool, # True if saving succeeded, False otherwise
'result': TrainModelResult | None, # Updated result if success=True
'error_message': str | None # Error message if success=False
}
TrainModelResult: Updated training result with run_name and artifacts
Raises:
Exception: If model saving fails (after sending notification)
Example:
# Successful save
result = await save_model({
'metadata': {'workflow_id': 'save-123', 'experiment_run_id': 456},
'train_result': TrainModelResult(...)
})
# Returns: {'success': True, 'result': TrainModelResult(...), 'error_message': None}
# Failed save
# Returns: {'success': False, 'result': None, 'error_message': 'Error details...'}
# Returns: TrainModelResult with run_name and artifacts
"""
metadata = input_data.get('metadata', {})
train_result = input_data['train_result']
@@ -418,11 +408,7 @@ class MLFlow(BaseActivity):
metadata,
)
return {
'success': True,
'result': train_result,
'error_message': None,
}
return train_result
except Exception as e: # noqa: BLE001
error_msg = f'Error saving model - Experiment: {train_result.params.experiment_name if train_result and train_result.params else "unknown"}, Error: {str(e)}'
@@ -441,10 +427,5 @@ class MLFlow(BaseActivity):
# Log error with metadata
self.error(trace, metadata=metadata)
# Return failure result (do NOT raise exception)
# This allows workflow to update database with error status
return {
'success': False,
'result': None,
'error_message': str(e),
}
# Re-raise exception to stop workflow
raise

115
model_manager/activities/training.py
View File

@@ -19,6 +19,7 @@ with workflow.unsafe.imports_passed_through():
from sientia_do.temporal.activities.base import BaseActivity
from model_manager.utils.models.train_model_params import TrainModelParams
from model_manager.utils.models.train_model_result import TrainModelResult
from model_manager.utils.repository.training_repository import TrainingRepository
@@ -51,20 +52,84 @@ class Training(BaseActivity):
BaseActivity.__init__(self, logger, notification_handler, set_error_counter=True)
self.training_repository = TrainingRepository(logger)
@activity.defn(name='train_model')
async def train_model(self, input_data: dict[str, Any]) -> dict[str, Any]:
@activity.defn(name='validate_train_params')
async def validate_train_params(self, input_data: dict[str, Any]) -> TrainModelParams:
"""
Train a machine learning model with comprehensive error handling.
Validate and convert training parameters from dict to TrainModelParams.
This activity validates the input training parameters and converts them
to a TrainModelParams object.
Args:
input_data: Training parameters and metadata at the same level
Required keys:
- metadata (dict): Workflow execution metadata
- All TrainModelParams fields (experiment_run_id, target_variable, etc.)
Returns:
TrainModelParams: Validated and converted training parameters
Raises:
ValueError, TypeError, KeyError: If validation fails (after sending notification)
Example:
result = await validate_train_params({
'metadata': {'workflow_id': 'train-123'},
'experiment_run_id': 456,
'target_variable': 'price',
'variable_columns': ['feature1', 'feature2'],
'train_size': 80,
# ... other required fields at same level
})
# Returns: TrainModelParams(...)
"""
metadata = input_data.get('metadata', {})
try:
self.info('Validating training parameters', metadata)
# Convert input_data directly to TrainModelParams (this validates all fields)
# The from_dict method will extract only the fields it needs
train_params = TrainModelParams.from_dict(input_data)
self.info(
f'Training parameters validated successfully - '
f'Target: {train_params.target_variable}, '
f'Experiment: {train_params.experiment_name}',
metadata,
)
return train_params
except (ValueError, TypeError, KeyError) as e:
error_msg = f'Error validating training parameters: {str(e)}'
trace = traceback.format_exc()
# Send notification (MongoDB)
self.send_notification(
metadata=metadata,
notification_id='VALIDATE_TRAIN_PARAMS_ERROR',
message=error_msg,
block='validate_train_params',
level=NotificationLevel.ERROR,
attachment_content=trace,
)
# Log error with metadata
self.error(trace, metadata=metadata)
# Re-raise exception to stop workflow
raise
@activity.defn(name='train_model')
async def train_model(self, input_data: dict[str, Any]) -> TrainModelResult:
"""
Train a machine learning model.
This activity orchestrates the complete ML training pipeline:
1. Validates input parameters
2. Trains the model using TrainingRepository
3. Performs post-training calculations
4. Returns success/failure status with results or error message
The activity does NOT raise exceptions on failure - it catches all errors,
sends notifications, and returns a failure status. This allows the workflow
to handle the error gracefully and update the database accordingly.
Args:
input_data: Configuration for model training operation
@@ -74,24 +139,19 @@ class Training(BaseActivity):
- train_params (TrainModelParams): Training parameters object
Returns:
dict: Training result with the following structure:
{
'success': bool, # True if training succeeded, False otherwise
'result': TrainModelResult | None, # Training result if success=True
'error_message': str | None # Error message if success=False
}
TrainModelResult: Training result with model, metrics, and data
Raises:
ValueError: If input validation fails
Exception: If training fails (after sending notification)
Example:
# Successful training
result = await train_model({
'metadata': {'workflow_id': 'train-123', 'experiment_run_id': 456},
'uploaded_file': BytesIO(csv_data),
'train_params': TrainModelParams(...) # Already converted object
'train_params': TrainModelParams(...)
})
# Returns: {'success': True, 'result': TrainModelResult(...), 'error_message': None}
# Failed training
# Returns: {'success': False, 'result': None, 'error_message': 'Error details...'}
# Returns: TrainModelResult(...)
"""
metadata = input_data.get('metadata', {})
uploaded_file = input_data['uploaded_file']
@@ -127,11 +187,7 @@ class Training(BaseActivity):
metadata,
)
return {
'success': True,
'result': final_result,
'error_message': None,
}
return final_result
except Exception as e: # noqa: BLE001
target = (
@@ -155,10 +211,5 @@ class Training(BaseActivity):
# Log error with metadata
self.error(trace, metadata=metadata)
# Return failure result (do NOT raise exception)
# This allows workflow to update database with error status
return {
'success': False,
'result': None,
'error_message': str(e),
}
# Re-raise exception to stop workflow
raise