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

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