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 io import BytesIO
    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.models.train_model_params import TrainModelParams
    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,
        logger: Logger,
        notification_handler: NotificationHandler,
    ):
        """
        Initialize Training activity.

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

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

        This activity orchestrates the complete ML training pipeline:
        1. Validates input parameters
        2. Trains the model using TrainingRepository
        3. Performs post-training calculations
        4. Returns success/failure status with results or error message

        The activity does NOT raise exceptions on failure - it catches all errors,
        sends notifications, and returns a failure status. This allows the workflow
        to handle the error gracefully and update the database accordingly.

        Args:
            input_data: Configuration for model training operation
                Required keys:
                    - metadata (dict): Workflow execution metadata
                    - uploaded_file (BytesIO): Training data file (already downloaded from MinIO)
                    - train_params (TrainModelParams): Training parameters object

        Returns:
            dict: Training result with the following structure:
                {
                    'success': bool,  # True if training succeeded, False otherwise
                    'result': TrainModelResult | None,  # Training result if success=True
                    'error_message': str | None  # Error message if success=False
                }

        Example:
            # Successful training
            result = await train_model({
                'metadata': {'workflow_id': 'train-123', 'experiment_run_id': 456},
                'uploaded_file': BytesIO(csv_data),
                'train_params': TrainModelParams(...)  # Already converted object
            })
            # Returns: {'success': True, 'result': TrainModelResult(...), 'error_message': None}

            # Failed training
            # Returns: {'success': False, 'result': None, 'error_message': 'Error details...'}
        """
        metadata = input_data.get('metadata', {})
        uploaded_file = input_data['uploaded_file']
        train_params = input_data['train_params']

        try:
            # Validate uploaded_file is BytesIO
            if not isinstance(uploaded_file, BytesIO):
                raise ValueError(f'uploaded_file must be BytesIO, got {type(uploaded_file)}')

            # Validate train_params is TrainModelParams
            if not isinstance(train_params, TrainModelParams):
                raise ValueError(f'train_params must be TrainModelParams, got {type(train_params)}')

            self.info(
                f'Starting model training for target: {train_params.target_variable}',
                metadata,
            )

            # Step 1: Train the model
            self.info('Training model with TrainingRepository', metadata)
            train_result = self.training_repository.train(uploaded_file, train_params)

            # Step 2: Perform post-training calculations
            self.info('Performing post-training calculations', metadata)
            final_result = self.training_repository.after_train_calculation(
                train_params, train_result
            )

            self.info(
                f'Model training completed successfully - '
                f'MSE: {final_result.mse_val}, MAE: {final_result.mae_val}, R²: {final_result.r2_val}',
                metadata,
            )

            return {
                'success': True,
                'result': final_result,
                'error_message': None,
            }

        except Exception as e:  # noqa: BLE001
            target = (
                train_params.target_variable
                if hasattr(train_params, 'target_variable')
                else 'unknown'
            )
            error_msg = f'Error training model - Target: {target}, Error: {str(e)}'
            trace = traceback.format_exc()

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

            # Log error with metadata
            self.error(trace, metadata=metadata)

            # Return failure result (do NOT raise exception)
            # This allows workflow to update database with error status
            return {
                'success': False,
                'result': None,
                'error_message': str(e),
            }