sientia-dataops-model-manager/model_manager/activities/training.py

"""
Training activities for ML model training operations.

This module provides activities for training machine learning models.
The activity extends BaseActivity and receives pre-downloaded files
to return success/failure status without raising exceptions.
"""

from temporalio import activity, workflow

with workflow.unsafe.imports_passed_through():
    import traceback
    from typing import Any

    from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler
    from sientia_do.notifications.models import NotificationLevel
    from sientia_do.observability.logger import Logger
    from sientia_do.temporal.activities.base import BaseActivity

    from model_manager.utils.exceptions import ModelTrainingError
    from model_manager.utils.models.train_model_params import TrainModelParams
    from model_manager.utils.repository.model_repository import ModelRepository
    from model_manager.utils.repository.storage_repository import StorageRepository
    from model_manager.utils.repository.training_repository import TrainingRepository


class Training(BaseActivity):
    """
    Activity for ML model training operations.

    This activity extends BaseActivity and handles machine learning model
    training with comprehensive error handling. It receives pre-downloaded
    files from the workflow and returns success/failure status without
    raising exceptions.

    Attributes:
        logger (Logger): Logger instance for observability (inherited from BaseActivity)
        notification_handler (NotificationHandler): Handler for sending notifications (inherited)
    """

    def __init__(
        self,
        model_repository: ModelRepository,
        storage_repository: StorageRepository,
        logger: Logger,
        notification_handler: NotificationHandler,
    ):
        """
        Initialize Training activity.

        Args:
            logger: Logger instance for observability
            notification_handler: Handler for sending notifications
        """
        super().__init__(logger, notification_handler, set_error_counter=True)
        self.training_repository = TrainingRepository(logger)
        self.model_repository = model_repository
        self.storage_repository = storage_repository

    @activity.defn(name='validate_train_params')
    async def validate_train_params(self, input_data: dict[str, Any]) -> TrainModelParams:
        """
        Validate and convert training parameters from dict to TrainModelParams.

        This activity validates the input training parameters and converts them
        to a TrainModelParams object.

        Args:
            input_data: Training parameters and metadata at the same level
                Required keys:
                    - metadata (dict): Workflow execution metadata
                    - All TrainModelParams fields (experiment_run_id, target_variable, etc.)

        Returns:
            TrainModelParams: Validated and converted training parameters

        Raises:
            ValueError, TypeError, KeyError: If validation fails (after sending notification)
        """
        metadata = input_data.get('metadata', {})

        try:
            train_params = TrainModelParams.from_dict(input_data)
            train_params.validate_business_rules()

            self.info(
                f'Training parameters validated successfully - '
                f'Target: {train_params.target_variable}, '
                f'Experiment: {train_params.experiment_name}',
                metadata,
            )

            return train_params
        except (ValueError, TypeError, KeyError) as e:
            error_msg = f'Error validating training parameters: {str(e)}'
            trace = traceback.format_exc()

            self.send_notification(
                metadata=metadata,
                notification_id='VALIDATE_TRAIN_PARAMS_ERROR',
                message=error_msg,
                block='validate_train_params',
                level=NotificationLevel.ERROR,
                attachment_content=trace,
            )
            raise

    @activity.defn(name='train_model')
    async def train_model(self, input_data: dict[str, Any]) -> dict[str, str | None]:
        """
        Train a machine learning model.

        This activity orchestrates the ML training pipeline:
        1. Validate input parameters.
        2. Train the model via TrainingRepository.
        3. Perform post-training calculations.

        Args:
            input_data: Training configuration containing:
                - metadata (dict): Workflow execution metadata.
                - uploaded_file (BytesIO): Training data already downloaded from MinIO.
                - train_params (TrainModelParams | dict): Training parameters.

        Returns:
            dict: Keys `run_name` and `run_dir` when training and saving succeed.

        Raises:
            ValueError: If input validation fails.
            Exception: If training fails (after sending notification).
        """
        metadata = input_data.get('metadata', {})
        train_params = input_data['train_params']

        if isinstance(train_params, dict):
            train_params = TrainModelParams.from_dict(train_params)
        # type: ignore[assignment]

        model_trained = False
        model_saved = False

        try:
            with self.storage_repository.fetch_file(
                train_params.bucket_name, train_params.file_name
            ) as uploaded_file:
                train_result = self.training_repository.train(uploaded_file, train_params)

                train_result = self.training_repository.after_train_calculation(
                    train_params, train_result
                )

                model_trained = True
                train_result = self.model_repository.save_model(train_result)
                model_saved = True

                return {
                    'run_name': train_result.run_name,
                    'run_dir': train_result.run_dir,
                }
        except Exception as e:  # noqa: BLE001
            error_msg = (
                'Error training model - '
                f'model_trained={model_trained}, model_saved={model_saved}, '
                f'error: {str(e)}'
            )

            trace = traceback.format_exc()

            self.send_notification(
                metadata=metadata,
                notification_id='TRAIN_MODEL_ERROR',
                message=error_msg,
                block='train_model',
                level=NotificationLevel.ERROR,
                attachment_content=trace,
            )

            raise ModelTrainingError(
                model_trained=model_trained,
                model_saved=model_saved,
            ) from e

    @activity.defn(name='cleanup_resources')
    async def cleanup_resources(self, input_data: dict[str, Any]) -> None:
        """
        Cleanup temporary resources created during training.

        Args:
            input_data: Cleanup configuration containing:
                - metadata (dict): Workflow execution metadata.
                - run_dir (str): Temporary directory to remove.
                - bucket_name (str): MinIO bucket of the uploaded file.
                - file_name (str): MinIO object key to delete.

        Raises:
            Exception: If cleanup fails (after sending notification).
        """
        metadata = input_data.get('metadata', {})
        run_dir = input_data.get('run_dir', '')
        bucket_name = input_data.get('bucket_name', '')
        file_name = input_data.get('file_name', '')

        try:
            self.model_repository.cleanup_run_directory(run_dir)
            self.storage_repository.delete_file(bucket_name, file_name)
        except Exception as e:  # noqa: BLE001
            error_msg = (
                f'Error cleaning up resources - Run directory: {run_dir}, '
                f'File: {bucket_name}/{file_name}, Error: {str(e)}'
            )

            trace = traceback.format_exc()

            self.send_notification(
                metadata=metadata,
                notification_id='CLEANUP_RESOURCES_ERROR',
                message=error_msg,
                block='cleanup_resources',
                level=NotificationLevel.ERROR,
                attachment_content=trace,
            )
            raise