""" 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.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')) # 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, ) # Moderate retry with backoff for MLFlow operations mlflow_retry_policy = RetryPolicy( initial_interval=timedelta(seconds=5), maximum_interval=timedelta(seconds=30), backoff_coefficient=2.0, maximum_attempts=3, ) # 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, ) # Filesystem retry for cleanup operations filesystem_retry_policy = RetryPolicy( initial_interval=timedelta(seconds=2), maximum_interval=timedelta(seconds=10), backoff_coefficient=1.5, maximum_attempts=3, ) @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=no_retry_policy, # Validation errors are permanent 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=network_retry_policy, # Fast retry for network issues 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=no_retry_policy, # Training errors are permanent (bad data) 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=mlflow_retry_policy, # Retry MLFlow with backoff 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=filesystem_retry_policy, # Retry filesystem operations 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=network_retry_policy, # Fast retry for network issues 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=database_retry_policy, # Retry DB with exponential backoff start_to_close_timeout=timedelta(seconds=TIMEOUT_UPDATE_DATABASE), )