502 lines
18 KiB
Python
502 lines
18 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 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
|
|
|
|
# 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_DOWNLOAD_FILE = int(os.getenv('TIMEOUT_DOWNLOAD_FILE', '600'))
|
|
TIMEOUT_TRAIN_MODEL = int(os.getenv('TIMEOUT_TRAIN_MODEL', '1800'))
|
|
TIMEOUT_SAVE_MODEL = int(os.getenv('TIMEOUT_SAVE_MODEL', '300'))
|
|
TIMEOUT_CLEANUP_DIRECTORY = int(os.getenv('TIMEOUT_CLEANUP_DIRECTORY', '60'))
|
|
TIMEOUT_DELETE_FILE = int(os.getenv('TIMEOUT_DELETE_FILE', '60'))
|
|
TIMEOUT_UPDATE_DATABASE = int(os.getenv('TIMEOUT_UPDATE_DATABASE', '30'))
|
|
|
|
|
|
@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=TIMEOUT_VALIDATE_PARAMS),
|
|
)
|
|
|
|
# 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=TIMEOUT_DOWNLOAD_FILE),
|
|
)
|
|
|
|
# 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=TIMEOUT_TRAIN_MODEL),
|
|
)
|
|
|
|
# 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=TIMEOUT_SAVE_MODEL),
|
|
)
|
|
|
|
# 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 via activity 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 via activity (deterministic)
|
|
if hasattr(saved_result, 'run_dir') and saved_result.run_dir:
|
|
cleanup_input = {
|
|
**metadata,
|
|
'run_dir': saved_result.run_dir,
|
|
}
|
|
|
|
await workflow.execute_activity_method(
|
|
Activities.cleanup_run_directory,
|
|
cleanup_input,
|
|
retry_policy=retry_policy,
|
|
start_to_close_timeout=timedelta(seconds=TIMEOUT_CLEANUP_DIRECTORY),
|
|
)
|
|
|
|
workflow.logger.info(
|
|
f'Run directory cleanup completed 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=TIMEOUT_DELETE_FILE),
|
|
)
|
|
|
|
# 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=TIMEOUT_UPDATE_DATABASE),
|
|
)
|