""" 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 # 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( run_dir=train_result.get('run_dir'), 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.TRAINING_SUCCESS, run_name=train_result.get('run_name'), ) return train_result 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.TRAINING_ERROR, error_message=self._extract_error_message(e), ) except Exception: pass raise async def _cleanup_resources( self, run_dir: str, metadata: dict[str, Any], ) -> None: """ Cleanup resources. This method removes the temporary run directory via activity. Args: run_dir: Temporary directory to remove metadata: Workflow execution metadata """ await workflow.execute_activity_method( Activities.cleanup_resources, { **metadata, 'run_dir': run_dir, }, retry_policy=network_retry_policy, start_to_close_timeout=timedelta(seconds=TIMEOUT_DELETE_FILE), ) 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)