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

1
model_manager/utils/models/experiment_status.py
View File

@@ -23,6 +23,7 @@ class ExperimentStatus(str, Enum):
FILE_DELETE_ERROR: Cleanup failed due to file system or MinIO errors.
"""
ORCHESTRATOR_VALIDATION_ERROR = 'ORCHESTRATOR_VALIDATION_ERROR'
MAGE_WAITING_PROC = 'MAGE_WAITING_PROC'
TRAINING_SUCCESS = 'TRAINING_SUCCESS'
TRAINING_ERROR = 'TRAINING_ERROR'

480
model_manager/workflows/train_model.py Normal file
View File

@@ -0,0 +1,480 @@
"""
Train Model Workflow for ML model training pipeline.
This workflow orchestrates the complete ML model training process, including:
- Parameter validation and conversion
- Data download from MinIO
- Model training
- Model saving to MLFlow
- Experiment tracking and status updates
"""
from temporalio import workflow
with workflow.unsafe.imports_passed_through():
import shutil
from datetime import timedelta
from typing import Any
from sientia_do.temporal.policies import retry_policy
from model_manager.activities.activities import Activities
from model_manager.activities.experiment_tracking import UpdateType
from model_manager.utils.models.experiment_status import ExperimentStatus
from model_manager.utils.models.train_model_params import TrainModelParams
from model_manager.utils.models.train_model_result import TrainModelResult
@workflow.defn(name='train_model')
class TrainModel:
"""
Complete ML model training workflow.
This workflow implements the full training pipeline from parameter validation
through model training and saving to MLFlow. It provides comprehensive error
handling with database status updates at each stage.
The workflow ensures:
- Proper parameter validation before training starts
- Status tracking in database for monitoring
- Error handling with detailed error messages
- Cleanup and proper resource management
"""
@workflow.run
async def run(self, input_data: dict[str, Any]) -> None:
"""
Execute the complete model training workflow.
This method orchestrates all steps of the training pipeline:
1. Validate and convert training parameters
2. Download training data from MinIO
3. Train the model
4. Save model to MLFlow
5. Update experiment tracking status
Args:
input_data: Complete configuration for the training workflow
All TrainModelParams fields at the same level:
- experiment_run_id (int): Unique identifier for the experiment run (REQUIRED)
- target_variable (str): Target variable to predict
- variable_columns (list[str]): Feature columns
- train_size (int): Training data percentage
- ... (all other TrainModelParams fields)
Raises:
ValueError: If experiment_run_id is missing or invalid
"""
# CRITICAL: Validate experiment_run_id first
# Without it, we cannot update database status, so fail immediately
experiment_run_id = self._validate_experiment_run_id(input_data)
metadata = {
'metadata': {
'experiment_run_id': experiment_run_id,
'workflow_name': 'train_model',
}
}
# Step 1: Validate and convert training parameters
train_params = await self._validate_training_parameters(
input_data, experiment_run_id, metadata
)
# Step 2 & 3: Download from MinIO and Train model
# The _download_and_train_model method handles both steps:
# - Downloads file from MinIO (returns BytesIO)
# - Trains model with the downloaded file
# Any error (download OR training) = TRAINING_ERROR
train_result = await self._download_and_train_model(
train_params=train_params,
experiment_run_id=experiment_run_id,
metadata=metadata,
)
# Step 4: Save model to MLFlow
saved_result = await self._save_model_to_mlflow(
train_result=train_result,
experiment_run_id=experiment_run_id,
metadata=metadata,
)
# Step 5: Cleanup resources and delete file from MinIO
await self._cleanup_resources(
saved_result=saved_result,
experiment_run_id=experiment_run_id,
metadata=metadata,
)
def _validate_experiment_run_id(self, input_data: dict[str, Any]) -> int:
"""
Validate experiment_run_id from input data.
This method ensures that experiment_run_id is present and valid.
Without a valid experiment_run_id, we cannot update database status,
so this validation must happen before any other operation.
Args:
input_data: Input data dictionary containing experiment_run_id
Returns:
int: Validated experiment_run_id
Raises:
ValueError: If experiment_run_id is missing or not an integer
"""
experiment_run_id = input_data.get('experiment_run_id')
if experiment_run_id is None:
error_msg = 'experiment_run_id is required but was not provided'
workflow.logger.error(error_msg)
raise ValueError(error_msg)
if not isinstance(experiment_run_id, int):
error_msg = (
f'experiment_run_id must be an integer, got {type(experiment_run_id).__name__}'
)
workflow.logger.error(error_msg)
raise ValueError(error_msg)
return experiment_run_id
async def _validate_training_parameters(
self,
input_data: dict[str, Any],
experiment_run_id: int,
metadata: dict[str, Any],
) -> TrainModelParams:
"""
Validate and convert training parameters from dict to TrainModelParams.
This method calls the validate_train_params activity to convert and validate
the input parameters. On success, updates DB status to MAGE_WAITING_PROC.
On error, updates DB status to ORCHESTRATOR_VALIDATION_ERROR.
Args:
input_data: Input data dictionary containing all training parameters
experiment_run_id: Validated experiment run ID
metadata: Workflow execution metadata
Returns:
TrainModelParams: Validated training parameters object
Raises:
Exception: If validation fails (after updating DB status)
"""
validation_input = {
**metadata,
**input_data, # All training params at same level
}
try:
# Execute validation activity
train_params = await workflow.execute_activity_method(
Activities.validate_train_params,
validation_input,
retry_policy=retry_policy,
start_to_close_timeout=timedelta(seconds=30),
)
# Validation succeeded: Update status
await self._update_experiment_run(
metadata=metadata,
experiment_run_id=experiment_run_id,
update_type=UpdateType.STATUS,
status=ExperimentStatus.MAGE_WAITING_PROC,
)
return train_params
except Exception as e:
# Validation failed: Update status with error
error_message = str(e)
workflow.logger.error(
f'Validation failed for experiment {experiment_run_id}: {error_message}'
)
await self._update_experiment_run(
metadata=metadata,
experiment_run_id=experiment_run_id,
update_type=UpdateType.STATUS_WITH_ERROR,
status=ExperimentStatus.ORCHESTRATOR_VALIDATION_ERROR,
error_message=error_message,
)
# Re-raise exception to stop workflow
raise
async def _download_and_train_model(
self,
train_params: TrainModelParams,
experiment_run_id: int,
metadata: dict[str, Any],
) -> TrainModelResult:
"""
Download file from MinIO and train model.
This method orchestrates the download and training steps using proper
resource management with try/catch/finally. On success, updates DB status
to TRAINING_SUCCESS. On error, updates DB status to TRAINING_ERROR.
Args:
train_params: TrainModelParams object with training configuration
experiment_run_id: Validated experiment run ID
metadata: Workflow execution metadata
Returns:
TrainModelResult: Training result from train_model activity
Raises:
Exception: If download or training fails (after updating DB status)
"""
uploaded_file = None
try:
# Step 1: Download file from MinIO
download_input = {
**metadata,
'bucket_name': train_params.bucket_name,
'file_name': train_params.file_name,
}
uploaded_file = await workflow.execute_activity_method(
Activities.fetch_file_from_minio,
download_input,
retry_policy=retry_policy,
start_to_close_timeout=timedelta(seconds=60),
)
# Step 2: Train model with downloaded file
train_input = {
**metadata,
'uploaded_file': uploaded_file,
'train_params': train_params,
}
train_result = await workflow.execute_activity_method(
Activities.train_model,
train_input,
retry_policy=retry_policy,
start_to_close_timeout=timedelta(seconds=300),
)
# Training succeeded: Update status
await self._update_experiment_run(
metadata=metadata,
experiment_run_id=experiment_run_id,
update_type=UpdateType.STATUS,
status=ExperimentStatus.TRAINING_SUCCESS,
)
return train_result
except Exception as e:
# Download or training failed: Update status with error
error_message = str(e)
workflow.logger.error(
f'Download or training failed for experiment {experiment_run_id}: {error_message}'
)
await self._update_experiment_run(
metadata=metadata,
experiment_run_id=experiment_run_id,
update_type=UpdateType.STATUS_WITH_ERROR,
status=ExperimentStatus.TRAINING_ERROR,
error_message=error_message,
)
# Re-raise exception to stop workflow
raise
finally:
# Ensure BytesIO is closed
if uploaded_file and hasattr(uploaded_file, 'close'):
uploaded_file.close()
async def _save_model_to_mlflow(
self,
train_result: TrainModelResult,
experiment_run_id: int,
metadata: dict[str, Any],
) -> TrainModelResult:
"""
Save trained model to MLFlow.
This method calls the save_model activity to save the trained model and
its artifacts to MLFlow. On success, updates DB status to MLFLOW_SENT
with MODEL_SAVED type and run_name. On error, updates DB status to
MLFLOW_SEND_ERROR.
Args:
train_result: TrainModelResult from training step
experiment_run_id: Validated experiment run ID
metadata: Workflow execution metadata
Returns:
TrainModelResult: Updated training result with MLFlow run name
Raises:
Exception: If model saving fails (after updating DB status)
"""
try:
# Save model to MLFlow
save_input = {
**metadata,
'train_result': train_result,
}
saved_result = await workflow.execute_activity_method(
Activities.save_model,
save_input,
retry_policy=retry_policy,
start_to_close_timeout=timedelta(seconds=120),
)
# Model saved successfully: Update status to MLFLOW_SENT with run_name
await self._update_experiment_run(
metadata=metadata,
experiment_run_id=experiment_run_id,
update_type=UpdateType.MODEL_SAVED,
status=ExperimentStatus.MLFLOW_SENT,
run_name=saved_result.run_name,
)
return saved_result
except Exception as e:
# Model saving failed: Update status with error
error_message = str(e)
workflow.logger.error(
f'Model saving failed for experiment {experiment_run_id}: {error_message}'
)
await self._update_experiment_run(
metadata=metadata,
experiment_run_id=experiment_run_id,
update_type=UpdateType.STATUS_WITH_ERROR,
status=ExperimentStatus.MLFLOW_SEND_ERROR,
error_message=error_message,
)
# Re-raise exception to stop workflow
raise
async def _cleanup_resources(
self,
saved_result: TrainModelResult,
experiment_run_id: int,
metadata: dict[str, Any],
) -> None:
"""
Cleanup resources and delete file from MinIO.
This method removes the temporary run directory and deletes the training
file from MinIO. On success, updates DB status to FILE_DELETED.
On error, updates DB status to FILE_DELETE_ERROR.
Args:
saved_result: TrainModelResult with run_dir and params information
experiment_run_id: Validated experiment run ID
metadata: Workflow execution metadata
Raises:
Exception: If cleanup fails (after updating DB status)
"""
try:
# Step 1: Remove temporary run directory
if hasattr(saved_result, 'run_dir') and saved_result.run_dir:
shutil.rmtree(saved_result.run_dir)
workflow.logger.info(
f'Removed run directory: {saved_result.run_dir} for experiment {experiment_run_id}'
)
# Step 2: Delete file from MinIO
delete_input = {
**metadata,
'bucket_name': saved_result.params.bucket_name,
'file_name': saved_result.params.file_name,
}
await workflow.execute_activity_method(
Activities.delete_file_from_minio,
delete_input,
retry_policy=retry_policy,
start_to_close_timeout=timedelta(seconds=30),
)
# Cleanup succeeded: Update status to FILE_DELETED
await self._update_experiment_run(
metadata=metadata,
experiment_run_id=experiment_run_id,
update_type=UpdateType.STATUS,
status=ExperimentStatus.FILE_DELETED,
)
except Exception as e:
# Cleanup failed: Update status with error
error_message = str(e)
workflow.logger.error(
f'Cleanup failed for experiment {experiment_run_id}: {error_message}'
)
await self._update_experiment_run(
metadata=metadata,
experiment_run_id=experiment_run_id,
update_type=UpdateType.STATUS_WITH_ERROR,
status=ExperimentStatus.FILE_DELETE_ERROR,
error_message=error_message,
)
# Re-raise exception to stop workflow
raise
async def _update_experiment_run(
self,
metadata: dict[str, Any],
experiment_run_id: int,
update_type: UpdateType,
status: ExperimentStatus,
error_message: str | None = None,
run_name: str | None = None,
) -> None:
"""
Update experiment run status in the database.
This is a helper method to simplify calls to the update_experiment_run activity.
It handles both success and error status updates.
Args:
metadata: Workflow execution metadata
experiment_run_id: Unique identifier for the experiment run
update_type: Type of update (STATUS, STATUS_WITH_ERROR, or MODEL_SAVED)
status: Status to set in the database
error_message: Error message (required if update_type is STATUS_WITH_ERROR)
run_name: MLFlow run name (required if update_type is MODEL_SAVED)
"""
update_input = {
**metadata,
'experiment_run_id': experiment_run_id,
'update_type': update_type,
'status': status,
}
# Add error_message only if provided
if error_message is not None:
update_input['error_message'] = error_message
# Add run_name only if provided
if run_name is not None:
update_input['run_name'] = run_name
await workflow.execute_activity_method(
Activities.update_experiment_run,
update_input,
retry_policy=retry_policy,
start_to_close_timeout=timedelta(seconds=30),
)