sientia-dataops-model-manager/model_manager/workflows/train_model.py

"""
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)