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

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.models.train_model_result import TrainModelResult
    from model_manager.utils.repository.model_repository import MLFlowRepository


class MLFlow(BaseActivity):
    """
    MLFlow integration activities for model training operations.

    This class provides activities for saving trained models and managing
    artifacts in MLFlow. It handles model persistence, artifact generation,
    and cleanup operations with comprehensive error handling.

    The class implements robust error handling and logging for all
    MLFlow operations, ensuring reliable model management in production environments.

    Attributes:
        mlflow_host (str): MLFlow server hostname
        mlflow_port (int): MLFlow server port
        mlflow_username (str): MLFlow authentication username
        mlflow_password (str): MLFlow authentication password
        model_monitoring_repository (MLFlowRepository): Repository for MLFlow operations
    """

    def __init__(
        self,
        mlflow_host: str,
        mlflow_port: int,
        mlflow_username: str,
        mlflow_password: str,
        logger: Logger,
        notification_handler: NotificationHandler,
    ):
        """
        Initialize MLFlow activities with server configuration.

        Args:
            mlflow_host: MLFlow server hostname or IP address
            mlflow_port: MLFlow server port number
            mlflow_username: Username for MLFlow authentication
            mlflow_password: Password for MLFlow authentication
            logger: Logger instance for observability and debugging
            notification_handler: Notification handler for alerts and monitoring

        Raises:
            Exception: If MLFlowRepository initialization fails
        """
        BaseActivity.__init__(self, logger, notification_handler, set_error_counter=True)
        self.mlflow_host = mlflow_host
        self.mlflow_port = mlflow_port
        self.mlflow_username = mlflow_username
        self.mlflow_password = mlflow_password

        self.model_monitoring_repository = MLFlowRepository(
            f'{mlflow_host}:{mlflow_port}', mlflow_username, mlflow_password, logger
        )

    @activity.defn(name='save_model')
    async def save_model(self, input_data: dict[str, Any]) -> TrainModelResult:
        """
        Save a trained ML model and its artifacts to MLflow.

        This activity orchestrates the complete model saving pipeline:
        1. Generates the next run name for the experiment
        2. Creates and organizes artifacts (reports, data files)
        3. Logs model, parameters, metrics, and artifacts to MLflow

        Args:
            input_data: Configuration for model saving operation
                Required keys:
                    - metadata (dict): Workflow execution metadata
                    - train_result (TrainModelResult): Training result with model and metrics

        Returns:
            TrainModelResult: Updated training result with run_name and artifacts

        Raises:
            Exception: If model saving fails (after sending notification)

        Example:
            result = await save_model({
                'metadata': {'workflow_id': 'save-123', 'experiment_run_id': 456},
                'train_result': TrainModelResult(...)
            })
            # Returns: TrainModelResult with run_name and artifacts
        """
        metadata = input_data.get('metadata', {})
        train_result = input_data['train_result']

        try:
            experiment_name = train_result.params.experiment_name

            self.info(
                f'Starting model save for experiment: {experiment_name}',
                metadata,
            )

            # Step 1: Generate next run name
            self.info('Generating run name', metadata)
            train_result.run_name = self.model_monitoring_repository.get_next_run_name(
                experiment_name
            )
            self.info(f'Generated run name: {train_result.run_name}', metadata)

            # Step 2: Generate artifacts (reports, CSV files)
            self.info('Generating artifacts', metadata)
            train_result = self.model_monitoring_repository.generate_artifacts(train_result)
            self.info('Artifacts generated successfully', metadata)

            # Step 3: Save run to MLflow
            self.info('Saving run to MLflow', metadata)
            self.model_monitoring_repository.save_run(train_result)

            self.info(
                f'Model saved successfully - Run: {train_result.run_name}, '
                f'Experiment: {experiment_name}',
                metadata,
            )

            return train_result

        except Exception as e:  # noqa: BLE001
            error_msg = f'Error saving model - Experiment: {train_result.params.experiment_name if train_result and train_result.params else "unknown"}, Error: {str(e)}'
            trace = traceback.format_exc()

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

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

            # Re-raise exception to stop workflow
            raise

    @activity.defn(name='cleanup_run_directory')
    async def cleanup_run_directory(self, input_data: dict[str, Any]) -> None:
        """
        Clean up temporary run directory after model training.

        This activity deletes the temporary directory created during model training
        and artifact generation. It implements idempotent cleanup to handle cases
        where the directory may have already been deleted.

        Args:
            input_data: Configuration for cleanup operation
                Required keys:
                    - metadata (dict): Workflow execution metadata
                    - run_dir (str): Path to the run directory to delete

        Raises:
            Exception: If cleanup fails for reasons other than directory not existing

        Example:
            await cleanup_run_directory({
                'metadata': {'workflow_id': 'cleanup-123'},
                'run_dir': '/path/to/run_dir'
            })
        """
        import os
        import shutil

        metadata = input_data.get('metadata', {})
        run_dir = input_data.get('run_dir')

        try:
            if not run_dir:
                self.info('No run directory specified, skipping cleanup', metadata)
                return

            self.info(f'Cleaning up run directory: {run_dir}', metadata)

            # Idempotent cleanup: check if directory exists before deleting
            if os.path.exists(run_dir):
                shutil.rmtree(run_dir)
                self.info(f'Run directory deleted successfully: {run_dir}', metadata)
            else:
                self.info(f'Run directory already deleted: {run_dir}', metadata)

        except Exception as e:
            error_msg = f'Error cleaning up run directory {run_dir}: {str(e)}'
            trace = traceback.format_exc()

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

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

            # Re-raise exception
            raise