""" 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 # Activity timeouts (seconds). Tune per environment (large uploads, long training). # Training uses no_retry_policy: extend TIMEOUT_TRAIN_MODEL instead of adding retries # to avoid duplicate MLflow side effects. Cleanup/delete uses network_retry_policy. 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, 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 """ experiment_run_id = self._validate_experiment_run_id(input_data) input_data = {**input_data, 'experiment_run_id': experiment_run_id} 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, Any] | None = 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: if train_result is not None: await self._cleanup_resources( run_dir=train_result.get('run_dir'), metadata=metadata, ) else: pass except Exception: # noqa: BLE001 # If cleanup fails after training failed, there is nothing extra to log (DB not committed). if training_succeeded: # pragma: no branch workflow.logger.warning( 'cleanup_resources failed after successful training; model and DB status ' 'are already committed. Temp files may remain until scheduled cleanup.', ) 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], ) -> dict[str, Any]: """ 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: dict[str, Any]: Validated training parameters 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 as secondary: # noqa: BLE001 workflow.logger.warning( 'Failed to persist ORCHESTRATOR_VALIDATION_ERROR to experiment_run: %s', secondary, ) raise async def _train_model( self, train_params: dict[str, Any], experiment_run_id: int, metadata: dict[str, Any], ) -> dict[str, Any]: """ 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: dict[str, Any]: Serializable training summary from the 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 as secondary: # noqa: BLE001 workflow.logger.warning( 'Failed to persist TRAINING_ERROR status to experiment_run: %s', secondary, ) raise async def _cleanup_resources( self, run_dir: str | None, 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 """ if run_dir is None: return 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)