bf2b6b3888
- Updated the `Training` class to raise `ModelTrainingError` on training failures for better error management. - Enhanced the `run` method in `TrainModel` to return training results and ensure proper resource cleanup, including validation files. - Refactored exception handling to prevent silent failures during resource cleanup and experiment run updates. - Adjusted type hints for improved clarity and consistency in method signatures.
421 lines
15 KiB
Python
421 lines
15 KiB
Python
"""
|
|
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 os
|
|
from datetime import timedelta
|
|
from typing import Any
|
|
|
|
from temporalio.common import RetryPolicy
|
|
|
|
from model_manager.activities.activities import Activities
|
|
from model_manager.activities.experiment_tracking import UpdateType
|
|
from model_manager.utils.exceptions import ModelTrainingError
|
|
from model_manager.utils.models.experiment_status import ExperimentStatus
|
|
from model_manager.utils.models.train_model_params import TrainModelParams
|
|
|
|
# Activity Timeouts (in seconds) - Configurable via environment variables
|
|
# Defaults are designed to handle large files (up to 200MB)
|
|
TIMEOUT_VALIDATE_PARAMS = int(os.getenv('TIMEOUT_VALIDATE_PARAMS', '30'))
|
|
TIMEOUT_TRAIN_MODEL = int(os.getenv('TIMEOUT_TRAIN_MODEL', '2700'))
|
|
TIMEOUT_DELETE_FILE = int(os.getenv('TIMEOUT_DELETE_FILE', '120'))
|
|
TIMEOUT_UPDATE_DATABASE = int(os.getenv('TIMEOUT_UPDATE_DATABASE', '30'))
|
|
|
|
# Retry Policies - Granular strategies for different operation types
|
|
# Fast retry for transient network errors (MinIO operations)
|
|
network_retry_policy = RetryPolicy(
|
|
initial_interval=timedelta(seconds=1),
|
|
maximum_interval=timedelta(seconds=10),
|
|
backoff_coefficient=2.0,
|
|
maximum_attempts=5,
|
|
)
|
|
|
|
# No retry for training - data errors are permanent
|
|
no_retry_policy = RetryPolicy(
|
|
maximum_attempts=1,
|
|
)
|
|
|
|
# Database retry with exponential backoff
|
|
database_retry_policy = RetryPolicy(
|
|
initial_interval=timedelta(seconds=2),
|
|
maximum_interval=timedelta(seconds=20),
|
|
backoff_coefficient=2.0,
|
|
maximum_attempts=5,
|
|
)
|
|
|
|
|
|
|
|
@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]) -> dict[str, str | 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
|
|
"""
|
|
experiment_run_id = self._validate_experiment_run_id(input_data)
|
|
model_name = input_data.get('model_name')
|
|
model_id = input_data.get('model_id')
|
|
|
|
metadata = {
|
|
'metadata': {
|
|
'experiment_run_id': experiment_run_id,
|
|
'workflow_name': 'train_model',
|
|
'model_name': model_name,
|
|
'model_id': model_id,
|
|
}
|
|
}
|
|
|
|
train_params = await self._validate_training_parameters(
|
|
input_data, experiment_run_id, metadata
|
|
)
|
|
|
|
training_succeeded = False
|
|
train_result: dict[str, str | None]
|
|
|
|
try:
|
|
train_result = await self._train_model(
|
|
train_params=train_params,
|
|
experiment_run_id=experiment_run_id,
|
|
metadata=metadata,
|
|
)
|
|
training_succeeded = True
|
|
finally:
|
|
try:
|
|
await self._cleanup_resources(
|
|
experiment_run_id=experiment_run_id,
|
|
bucket_name=train_params.bucket_name,
|
|
file_name=train_params.file_name,
|
|
val_file_name=train_params.val_file_name,
|
|
metadata=metadata,
|
|
)
|
|
except Exception:
|
|
if training_succeeded:
|
|
raise
|
|
return train_result
|
|
|
|
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:
|
|
raise ValueError('experiment_run_id is required but was not provided')
|
|
|
|
if isinstance(experiment_run_id, int):
|
|
return experiment_run_id
|
|
|
|
if isinstance(experiment_run_id, str) and experiment_run_id.strip().isdigit():
|
|
return int(experiment_run_id.strip())
|
|
|
|
raise ValueError(
|
|
f'experiment_run_id must be an integer or numeric string, got {type(experiment_run_id).__name__}'
|
|
)
|
|
|
|
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 ORCHESTRATOR_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)
|
|
"""
|
|
try:
|
|
input_data = await workflow.execute_activity_method(
|
|
Activities.load_model_metadata,
|
|
{
|
|
**input_data,
|
|
**metadata,
|
|
},
|
|
retry_policy=no_retry_policy,
|
|
start_to_close_timeout=timedelta(seconds=TIMEOUT_VALIDATE_PARAMS),
|
|
)
|
|
|
|
train_params = await workflow.execute_activity_method(
|
|
Activities.validate_train_params,
|
|
{
|
|
**input_data,
|
|
**metadata,
|
|
},
|
|
retry_policy=no_retry_policy,
|
|
start_to_close_timeout=timedelta(seconds=TIMEOUT_VALIDATE_PARAMS),
|
|
)
|
|
|
|
await self._update_experiment_run(
|
|
metadata=metadata,
|
|
experiment_run_id=experiment_run_id,
|
|
update_type=UpdateType.STATUS,
|
|
status=ExperimentStatus.ORCHESTRATOR_WAITING_PROC,
|
|
)
|
|
|
|
return train_params
|
|
except Exception as e:
|
|
try:
|
|
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=self._extract_error_message(e),
|
|
)
|
|
except Exception:
|
|
pass
|
|
raise
|
|
|
|
async def _train_model(
|
|
self,
|
|
train_params: TrainModelParams,
|
|
experiment_run_id: int,
|
|
metadata: dict[str, Any],
|
|
) -> dict[str, str | None]:
|
|
"""
|
|
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)
|
|
"""
|
|
try:
|
|
train_result = await workflow.execute_activity_method(
|
|
Activities.train_model,
|
|
{
|
|
**metadata,
|
|
'train_params': train_params,
|
|
},
|
|
retry_policy=no_retry_policy,
|
|
start_to_close_timeout=timedelta(seconds=TIMEOUT_TRAIN_MODEL),
|
|
)
|
|
|
|
await self._update_experiment_run(
|
|
metadata=metadata,
|
|
experiment_run_id=experiment_run_id,
|
|
update_type=UpdateType.MODEL_SAVED,
|
|
status=ExperimentStatus.TRACKING_SENT,
|
|
run_name=train_result.get('run_name'),
|
|
)
|
|
|
|
return train_result
|
|
except Exception as e:
|
|
# Mapear flags -> status
|
|
# False/False: erro no treino
|
|
# True/False: erro ao salvar (MLflow)
|
|
# False/True: estado inconsistente, tratar como erro de treino
|
|
# True/True: não deveria cair aqui; tratar como erro genérico de treino
|
|
status = ExperimentStatus.TRAINING_ERROR
|
|
|
|
if isinstance(e, ModelTrainingError) and (e.model_trained and not e.model_saved):
|
|
status = ExperimentStatus.TRACKING_SEND_ERROR
|
|
|
|
try:
|
|
await self._update_experiment_run(
|
|
metadata=metadata,
|
|
experiment_run_id=experiment_run_id,
|
|
update_type=UpdateType.STATUS_WITH_ERROR,
|
|
status=status,
|
|
error_message=self._extract_error_message(e),
|
|
)
|
|
except Exception:
|
|
pass
|
|
raise
|
|
|
|
async def _cleanup_resources(
|
|
self,
|
|
experiment_run_id: int,
|
|
bucket_name: str,
|
|
file_name: str,
|
|
val_file_name: str | None,
|
|
metadata: dict[str, Any],
|
|
) -> None:
|
|
"""
|
|
Cleanup resources and delete file from MinIO.
|
|
|
|
This method deletes the training file from MinIO. On success, updates
|
|
DB status to FILE_DELETED.
|
|
On error, updates DB status to FILE_DELETE_ERROR.
|
|
|
|
Args:
|
|
experiment_run_id: Validated experiment run ID
|
|
metadata: Workflow execution metadata
|
|
|
|
Raises:
|
|
Exception: If cleanup fails (after updating DB status)
|
|
"""
|
|
try:
|
|
await workflow.execute_activity_method(
|
|
Activities.cleanup_resources,
|
|
{
|
|
**metadata,
|
|
'bucket_name': bucket_name,
|
|
'file_name': file_name,
|
|
'val_file_name': val_file_name,
|
|
},
|
|
retry_policy=network_retry_policy,
|
|
start_to_close_timeout=timedelta(seconds=TIMEOUT_DELETE_FILE),
|
|
)
|
|
|
|
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:
|
|
try:
|
|
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=self._extract_error_message(e),
|
|
)
|
|
except Exception:
|
|
pass
|
|
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,
|
|
}
|
|
|
|
if error_message is not None:
|
|
update_input['error_message'] = error_message
|
|
|
|
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=database_retry_policy,
|
|
start_to_close_timeout=timedelta(seconds=TIMEOUT_UPDATE_DATABASE),
|
|
)
|
|
|
|
def _extract_error_message(self, exc: Exception) -> str:
|
|
message_parts: list[str] = []
|
|
seen: set[int] = set()
|
|
current: Exception | None = exc
|
|
|
|
while current and id(current) not in seen:
|
|
seen.add(id(current))
|
|
text = str(current).strip()
|
|
|
|
if text and text not in message_parts:
|
|
message_parts.append(text)
|
|
|
|
cause = getattr(current, '__cause__', None)
|
|
context = getattr(current, '__context__', None)
|
|
current = cause if isinstance(cause, Exception) else context
|
|
|
|
if not message_parts:
|
|
return repr(exc)
|
|
|
|
return ' | '.join(message_parts)
|