SIENTIAPDE-1255: Refactor MLFlow activities for training operations and update metrics

This commit refactors the MLFlow activities to focus on model training rather than prediction operations. It removes prediction-related activities and metrics, and updates the MLFlow activity descriptions to reflect the change in focus. The README is also updated to reflect these changes.

README.md

@@ -40,8 +40,7 @@ An enterprise-grade ML model training orchestration platform built on Temporal.
   - [Test Execution](#test-execution)
 - [Monitoring and Metrics](#monitoring-and-metrics)
   - [Application Health Metrics](#application-health-metrics)
-  - [Prediction Operation Metrics](#prediction-operation-metrics)
+  - [Training Metrics](#training-metrics)
-  - [Data Quality Metrics](#data-quality-metrics)
 - [Configuration](#configuration-1)
   - [Environment Variables](#environment-variables)
   - [Workflow Configuration](#workflow-configuration)
@@ -806,15 +805,6 @@ The Model Manager system exposes comprehensive Prometheus metrics for operationa
 - `app_up`: Application health status (1=healthy, 0=unhealthy)
   - Labels: `pod_id`
 
-### Prediction Operation Metrics
-- `model_manager_predictions_written_count`: Counter for successful prediction exports
-  - Labels: `pod_id`, `model_name`, `pipeline_name`
-- `model_manager_prediction_confidence_monitor`: Gauge for current prediction confidence levels
-  - Labels: `pod_id`, `model_name`, `pipeline_name`
-- `model_manager_prediction_response_time_monitor`: Histogram for prediction response times
-  - Labels: `pod_id`, `model_name`, `pipeline_name`
-  - Buckets: [0.01, 0.05, 0.1, 0.2, 0.5, 1.0, 2.0, 5.0, 10.0]
-
 ### Training Metrics
 - Training success/failure rates through notification system
 - Model save performance metrics

model_manager/activities/mlflow.py

@@ -4,14 +4,10 @@ with workflow.unsafe.imports_passed_through():
     import traceback
     from typing import Any
 
-    import numpy as np
-    from pandas import DataFrame, to_datetime
-    from sientia_do.formatters import create_sample_dict
     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 sientia_do.temporal.constants import DATETIME_FORMAT, DATETIME_FORMAT_WITH_TZ
 
     from model_manager.utils.models.train_model_result import TrainModelResult
     from model_manager.utils.repository.model_repository import MLFlowRepository
@@ -19,14 +15,14 @@ with workflow.unsafe.imports_passed_through():
 
 class MLFlow(BaseActivity):
     """
-    MLFlow integration activities for model inference operations.
+    MLFlow integration activities for model training operations.
 
-    This class provides activities for interacting with MLFlow models, including
+    This class provides activities for saving trained models and managing
-    data transformation and prediction operations. It handles authentication,
+    artifacts in MLFlow. It handles model persistence, artifact generation,
-    data preprocessing, and model management with configurable retention policies.
+    and cleanup operations with comprehensive error handling.
 
-    The class implements comprehensive error handling and logging for all
+    The class implements robust error handling and logging for all
-    MLFlow operations, ensuring reliable model inference in production environments.
+    MLFlow operations, ensuring reliable model management in production environments.
 
     Attributes:
         mlflow_host (str): MLFlow server hostname
@@ -69,283 +65,6 @@ class MLFlow(BaseActivity):
             f'{mlflow_host}:{mlflow_port}', mlflow_username, mlflow_password, logger
         )
 
-    @activity.defn(name='request_transform')
-    async def request_transform(self, input_data: dict[str, Any]) -> dict[str, Any]:
-        """
-        Transform input data using MLFlow models.
-
-        This activity processes input data through MLFlow model transformation,
-        including data preprocessing, format conversion, and validation. It handles
-        data deduplication, pivoting, and cleanup to ensure optimal model performance.
-
-        The transformation process includes:
-        1. Data deduplication based on variable and timestamp
-        2. Data pivoting for model input format
-        3. Null value handling and cleanup
-        4. MLFlow model transformation request
-        5. Response validation and logging
-
-        Args:
-            input_data: Configuration and data for transformation
-                Required keys:
-                    - metadata (dict): Workflow execution metadata
-                    - data (dict): Input data for transformation
-                    - model_name (str): Name of the MLFlow model to use
-                    - model_retention (int): Model retention period in minutes
-
-        Returns:
-            dict: Transformed data from MLFlow model
-
-        Raises:
-            Exception: If transformation fails or MLFlow model is unavailable
-        """
-        metadata = input_data['metadata']
-        self.info('Transforming data...', metadata)
-        data = DataFrame(input_data['data'])
-        model_name = input_data['model_name']
-        model_config = input_data.get('model_config', {})
-
-        self.debug('Raw input data:', metadata)
-        self.debug(data.head(5).to_string(), metadata)
-
-        # Sort by created_at in descending order and keep first occurrence of each variable/timestamp pair
-        data = data.sort_values('created_at', ascending=False).drop_duplicates(
-            subset=['variable', 'timestamp'], keep='first'
-        )
-
-        # Pivot data for model input format
-        data = data.pivot(index='timestamp', columns='variable', values='value')
-        data.fillna(np.nan, inplace=True)
-        # data.reset_index(inplace=True)
-        data.columns.name = None
-
-        self.debug('Processed input data:', metadata)
-        self.debug(data.head(5).to_string(), metadata)
-
-        # Request transformation from MLFlow model
-        response_data = self.model_monitoring_repository.transform(
-            model_name, data, model_config, metadata
-        )
-
-        self.debug(
-            f'Transform raw response data: \n {create_sample_dict(response_data, max_items=5, max_depth=5)}',
-            metadata,
-        )
-
-        self.debug(
-            f'Transform response data: \n {create_sample_dict(response_data, max_items=5, max_depth=5)}',
-            metadata,
-        )
-
-        self.info('Data transformed successfully', metadata)
-
-        return response_data
-
-    @activity.defn(name='request_predict')
-    async def request_predict(self, input_data: dict[str, Any]) -> dict[str, Any]:
-        """
-        Execute predictions using MLFlow models.
-
-        This activity performs ML model inference using MLFlow models with the
-        transformed data. It handles data format conversion, null value processing,
-        and model prediction requests with comprehensive error handling.
-
-        The prediction process includes:
-        1. Data format validation and cleanup
-        2. Null value handling for model compatibility
-        3. MLFlow model prediction request
-        4. Response validation and logging
-        5. Performance monitoring and metrics
-
-        Args:
-            input_data: Configuration and data for prediction
-                Required keys:
-                    - metadata (dict): Workflow execution metadata
-                    - data (dict): Transformed data for prediction
-                    - model_name (str): Name of the MLFlow model to use
-                    - model_retention (int): Model retention period in minutes
-
-        Returns:
-            dict: Prediction results from MLFlow model
-
-        Raises:
-            Exception: If prediction fails or MLFlow model is unavailable
-        """
-        metadata = input_data['metadata']
-        self.info('Predicting data...', metadata)
-        data = DataFrame(input_data['data'])
-        model_name = input_data['model_name']
-        model_config = input_data.get('model_config', {})
-
-        self.debug(f'Input data for: \n {data.head(5).to_string()}', metadata)
-
-        # Convert numpy.nan to None for model compatibility
-        data.replace(np.nan, None, inplace=True)
-
-        data['timestamp'] = data.index
-        data['timestamp'] = to_datetime(
-            data['timestamp'], format=DATETIME_FORMAT_WITH_TZ
-        ).dt.strftime(DATETIME_FORMAT)
-
-        # Request prediction from MLFlow model
-        response_data = self.model_monitoring_repository.predict(
-            model_name, data, model_config, metadata
-        )
-
-        self.debug(
-            f'Prediction response data: \n {create_sample_dict(response_data, max_items=5, max_depth=5)}',
-            metadata,
-        )
-
-        self.info('Data predicted successfully', metadata)
-
-        return response_data
-
-    @activity.defn(name='retrain_model')
-    async def retrain_model(self, input_data: dict[str, Any]) -> dict[str, Any]:
-        """
-        Retrain MLFlow models with updated training data.
-
-        This activity orchestrates the complete model retraining process,
-        including data preparation, model retraining execution, and result
-        validation. It handles data preprocessing, column cleanup, and
-        comprehensive error handling for production model management.
-
-        The retraining process includes:
-        1. Data timestamp extraction and validation
-        2. Column cleanup and data preparation
-        3. Data pivoting for model input format
-        4. MLFlow model retraining execution
-        5. Result validation and error handling
-
-        Args:
-            input_data (dict): Input data containing:
-                - metadata (dict): Workflow execution metadata
-                - data (dict[str, Any]): Training data for model retraining
-                - model_name (str): Name of the MLFlow model to retrain
-
-        Returns:
-            dict: Retraining results containing:
-                - status (str): Retraining operation status
-                - timestamp (str): Timestamp of the retraining operation
-                - experiment (str): MLFlow experiment identifier
-
-        Raises:
-            Exception: If retraining fails or encounters critical errors
-        """
-        metadata = input_data['metadata']
-        data = DataFrame(input_data['data'])
-        model_name = input_data['model_name']
-
-        self.info(f'Retraining model {model_name}...', metadata)
-
-        timestamp = data['timestamp'].max()
-        self.debug(f'Timestamp: {timestamp}', metadata)
-
-        data.drop(columns=['model_id'], inplace=True, errors='ignore')
-        data.drop(columns=['created_at'], inplace=True, errors='ignore')
-
-        data = data.pivot(index='timestamp', columns='variable', values='value')
-        data.sort_index(inplace=True)
-        data.reset_index(inplace=True)
-
-        data = data.dropna()
-        data.columns.name = None
-
-        try:
-            retrain_output, experiment = self.model_monitoring_repository.retrain_model(
-                data=data, model_name=model_name
-            )
-
-            return {'status': retrain_output, 'timestamp': timestamp, 'experiment': experiment}
-        except Exception as e:
-            trace = traceback.format_exc()
-            self.send_notification(
-                metadata=metadata,
-                notification_id='RETRAIN_MODEL_ERROR',
-                message=f'Error retraining model {model_name}: {e}',
-                block='retrain_model',
-                level=NotificationLevel.ERROR,
-                attachment_content=trace,
-            )
-            self.error(trace, metadata=metadata)
-            raise e
-
-    @activity.defn(name='update_production_model')
-    async def update_production_model(self, input_data: dict[str, Any]) -> dict[Any, Any]:
-        """
-        Update production model with newly trained model version.
-
-        This activity manages the critical process of updating production
-        models with newly trained versions. It handles model deployment,
-        status tracking, and comprehensive reporting for operational
-        visibility and audit trails.
-
-        The update process includes:
-        1. Production model update execution
-        2. Status and metadata tracking
-        3. Comprehensive reporting and logging
-        4. Error handling and notification
-        5. Audit trail maintenance
-
-        Args:
-            input_data (dict): Input data containing:
-                - metadata (dict): Workflow execution metadata
-                - model_name (str): Name of the MLFlow model to update
-                - experiment (str): MLFlow experiment identifier
-                - model_id (str): Unique identifier for the model version
-                - timestamp (str): Timestamp of the update operation
-                - status (str): Current status of the model update
-
-        Returns:
-            dict[Any, Any]: Comprehensive update report containing:
-                - model_id (str): Model version identifier
-                - model_name (str): Name of the updated model
-                - timestamp (str): Update operation timestamp
-                - status (str): Update operation status
-                - Additional MLFlow response metadata
-
-        Raises:
-            Exception: If production model update fails
-        """
-        metadata = input_data['metadata']
-        model_name = input_data['model_name']
-        model_id = input_data['model_id']
-        experiment = input_data['experiment']
-        timestamp = input_data['timestamp']
-        status = input_data['status']
-
-        self.info(
-            f'Updating production model {model_name} from experiment {experiment}...', metadata
-        )
-
-        try:
-            response = self.model_monitoring_repository.update_production_model(
-                experiment=experiment, model_name=model_name
-            )
-
-            report = DataFrame([response])
-            report['model_id'] = model_id
-            report['model_name'] = model_name
-            report['timestamp'] = timestamp
-            report['status'] = status
-
-            self.info(f'Production model {model_name} updated successfully', metadata)
-            return report.to_dict()  # type: ignore[no-any-return]
-
-        except Exception as e:
-            trace = traceback.format_exc()
-            self.send_notification(
-                metadata=metadata,
-                notification_id='UPDATE_PRODUCTION_MODEL_ERROR',
-                message=f'Error updating production model {model_name}: {e}',
-                block='update_production_model',
-                level=NotificationLevel.ERROR,
-                attachment_content=trace,
-            )
-            self.error(trace, metadata=metadata)
-            raise e
-
     @activity.defn(name='save_model')
     async def save_model(self, input_data: dict[str, Any]) -> TrainModelResult:
         """

model_manager/metrics.py

@@ -3,7 +3,7 @@ Model Manager Metrics Module
 
 This module defines all Prometheus metrics used by the Sientia DataOps Model Manager system
 for monitoring and observability. The metrics provide insights into system performance,
-prediction quality, and operational health.
+training operations, and operational health.
 
 The metrics are designed to be scraped by Prometheus and can be visualized in
 Grafana or other monitoring dashboards to provide real-time visibility into
@@ -11,18 +11,12 @@ the system's operation.
 
 Key Metric Categories:
 - Application Health: Overall system status and availability
-- Prediction Operations: Count and performance of prediction operations
-- Data Quality: Confidence levels and validation results
-- Export Operations: Database export performance
-- Response Times: Performance monitoring for various operations
 
 Metric Labels:
 - pod_id: Kubernetes pod identifier for multi-instance deployments
-- model_name: Name of the ML model being used
-- pipeline_name: Name of the prediction pipeline
 """
 
-from prometheus_client import Counter, Gauge, Histogram
+from prometheus_client import Gauge
 
 # Application health metric
 APP_UP = Gauge(
@@ -30,28 +24,3 @@ APP_UP = Gauge(
     'Indicates if the application is running (1) or shutting down (0)',
     ['pod_id'],
 )
-
-# Core labels used across multiple metrics
-CORE_LABELS = ['pod_id', 'model_name', 'pipeline_name']
-
-# Prediction operation metrics
-PREDICTIONS_WRITTEN_COUNT = Counter(
-    'model_manager_predictions_written_count',
-    'Number of predictions written to the database table predictions',
-    CORE_LABELS,
-)
-
-# Prediction quality metrics
-PREDICTION_CONFIDENCE_MONITOR = Gauge(
-    'model_manager_prediction_confidence_monitor',
-    'Current confidence of each prediction',
-    CORE_LABELS,
-)
-
-# Performance monitoring metrics
-PREDICTION_RESPONSE_TIME_MONITOR = Histogram(
-    'model_manager_prediction_response_time_monitor',
-    'Current response time of each prediction',
-    CORE_LABELS,
-    buckets=[0.01, 0.05, 0.1, 0.2, 0.5, 1.0, 2.0, 5.0, 10.0],
-)

tests/activities/test_mlflow.py

@@ -1,10 +1,8 @@
 from unittest.mock import ANY, MagicMock, patch
 
-import numpy as np
 import pytest
 from pytest import fixture, mark
 from sientia_do.notifications.models import NotificationLevel
-from sientia_do.temporal.constants import DATETIME_FORMAT, DATETIME_FORMAT_WITH_TZ
 
 from model_manager.activities.mlflow import MLFlow
 
@@ -55,253 +53,6 @@ metadata = {
 }
 
 
-@mark.asyncio
-@patch('model_manager.activities.mlflow.DataFrame')
-@patch('model_manager.activities.mlflow.max')
-async def test_request_transform_success(mock_max, mock_dataframe, mlflow):
-    mock_max.return_value = '2024-01-02'
-    # Mock input data
-    input_data = {
-        **metadata,
-        'data': [
-            {
-                'timestamp': '2024-01-01',
-                'variable': 'var1',
-                'value': 1.0,
-                'created_at': '2024-01-01 12:00:00',
-            },
-            {
-                'timestamp': '2024-01-01',
-                'variable': 'var2',
-                'value': 2.0,
-                'created_at': '2024-01-01 12:00:00',
-            },
-            {
-                'timestamp': '2024-01-02',
-                'variable': 'var1',
-                'value': 3.0,
-                'created_at': '2024-01-02 12:00:00',
-            },
-            {
-                'timestamp': '2024-01-02',
-                'variable': 'var2',
-                'value': 4.0,
-                'created_at': '2024-01-02 12:00:00',
-            },
-            {
-                'timestamp': '2024-01-02',
-                'variable': 'var1',
-                'value': 1.0,
-                'created_at': '2024-01-01 12:00:00',
-            },
-            {
-                'timestamp': '2024-01-02',
-                'variable': 'var2',
-                'value': 1.0,
-                'created_at': '2024-01-01 12:00:00',
-            },
-        ],
-        'model_name': 'test_model',
-        'model_config': {},
-    }
-
-    # Mock the transform response
-    expected_response = {'prediction': [0.5, 0.6], 'timestamp': ['2024-01-01', '2024-01-02']}
-    mlflow.model_monitoring_repository.transform.return_value = expected_response
-
-    mock_dataframe.return_value.sort_values.return_value = mock_dataframe.return_value
-    mock_dataframe.return_value.drop_duplicates.return_value = mock_dataframe.return_value
-
-    # Call the method
-    response_data = await mlflow.request_transform(input_data)
-
-    # Verify the data was correctly transformed
-    mock_dataframe.assert_called_once_with(input_data['data'])
-    mock_dataframe.return_value.pivot.assert_called_once_with(
-        index='timestamp', columns='variable', values='value'
-    )
-    mock_dataframe = mock_dataframe.return_value.pivot.return_value
-    mock_dataframe.fillna.assert_called_once_with(np.nan, inplace=True)
-    # mock_dataframe.reset_index.assert_called_once()
-    mock_dataframe.columns.name = None
-
-    # Verify the response
-    assert response_data == expected_response
-
-    # Verify the repository was called with correct arguments
-    mlflow.model_monitoring_repository.transform.assert_called_once_with(
-        'test_model', mock_dataframe, {}, metadata['metadata']
-    )
-
-
-@mark.asyncio
-@patch('model_manager.activities.mlflow.DataFrame')
-@patch('model_manager.activities.mlflow.to_datetime')
-@patch('model_manager.activities.mlflow.max')
-async def test_request_predict(mock_max, mock_to_datetime, mock_dataframe, mlflow):
-    mock_max.return_value = '2024-01-02'
-    # Mock input data
-    input_data = {
-        **metadata,
-        'data': {
-            'variable': {
-                '2024-01-01': 'var1',
-                '2024-01-02': 'var2',
-                '2024-01-03': 'var1',
-                '2024-01-04': 'var2',
-            },
-            'value': {'2024-01-01': 1.0, '2024-01-02': 2.0, '2024-01-03': 3.0, '2024-01-04': 4.0},
-        },
-        'model_name': 'test_model',
-        'model_config': {},
-    }
-
-    # Mock the predict response
-    expected_response = {'prediction': [0.5, 0.6]}
-    mlflow.model_monitoring_repository.predict.return_value = expected_response
-
-    # Call the method
-    response_data = await mlflow.request_predict(input_data)
-
-    mock_dataframe.assert_called_once_with(input_data['data'])
-    mock_dataframe.return_value.replace.assert_called_once_with(np.nan, None, inplace=True)
-    mock_dataframe.return_value.__setitem__.assert_any_call(
-        'timestamp', mock_to_datetime.return_value.dt.strftime.return_value
-    )
-
-    mock_to_datetime.assert_called_once_with(
-        mock_dataframe.return_value.__getitem__.return_value, format=DATETIME_FORMAT_WITH_TZ
-    )
-    mock_to_datetime.return_value.dt.strftime.assert_called_once_with(DATETIME_FORMAT)
-
-    # Verify the response
-    assert response_data == expected_response
-
-    # Verify the repository was called with correct arguments
-    mlflow.model_monitoring_repository.predict.assert_called_once_with(
-        'test_model', mock_dataframe.return_value, {}, metadata['metadata']
-    )
-
-
-@mark.asyncio
-async def test_retrain_model(mlflow):
-    data = {
-        'model_id': [4, 5, 6, 7],
-        'created_at': [1, 2, 3, 4],
-        'timestamp': [1, 1, 2, 2],
-        'variable': ['var1', 'var2', 'var1', 'var2'],
-        'value': [1, 2, 3, 4],
-    }
-
-    mlflow.model_monitoring_repository.retrain_model.return_value = (
-        'Model retrained successfully',
-        'test',
-    )
-
-    response = await mlflow.retrain_model({**metadata, 'data': data, 'model_name': 'test_model'})
-
-    mlflow.model_monitoring_repository.retrain_model.assert_called_once()
-
-    assert response == {
-        'status': 'Model retrained successfully',
-        'timestamp': 2,
-        'experiment': 'test',
-    }
-
-
-@mark.asyncio
-async def test_retrain_model_error(mlflow):
-    mlflow.model_monitoring_repository.retrain_model.side_effect = Exception(
-        'Error retraining model'
-    )
-
-    data = {
-        'model_id': [4, 5, 6, 7],
-        'created_at': [1, 2, 3, 4],
-        'timestamp': [1, 1, 2, 2],
-        'variable': ['var1', 'var2', 'var1', 'var2'],
-        'value': [1, 2, 3, 4],
-    }
-
-    try:
-        await mlflow.retrain_model({**metadata, 'data': data, 'model_name': 'test_model'})
-    except Exception as e:  # noqa: BLE001
-        assert str(e) == 'Error retraining model'
-        mlflow.send_notification.assert_called_once_with(
-            metadata=metadata['metadata'],
-            notification_id='RETRAIN_MODEL_ERROR',
-            message='Error retraining model test_model: Error retraining model',
-            block='retrain_model',
-            level=NotificationLevel.ERROR,
-            attachment_content=ANY,
-        )
-    else:
-        raise AssertionError('No exception raised')
-
-
-@mark.asyncio
-async def test_update_production_model(mlflow):
-    mlflow.model_monitoring_repository.update_production_model.return_value = {
-        'data1': 1,
-        'data2': 2,
-    }
-
-    input_data = {
-        **metadata,
-        'model_name': 'test_model',
-        'model_id': 1,
-        'experiment': 'test',
-        'timestamp': 2,
-        'status': 'success',
-    }
-
-    response = await mlflow.update_production_model(input_data)
-
-    mlflow.model_monitoring_repository.update_production_model.assert_called_once_with(
-        experiment='test', model_name='test_model'
-    )
-
-    assert response == {
-        'data1': {0: 1},
-        'data2': {0: 2},
-        'model_id': {0: 1},
-        'model_name': {0: 'test_model'},
-        'timestamp': {0: 2},
-        'status': {0: 'success'},
-    }
-
-
-@mark.asyncio
-async def test_update_production_model_error(mlflow):
-    mlflow.model_monitoring_repository.update_production_model.side_effect = Exception(
-        'Error updating production model'
-    )
-
-    input_data = {
-        **metadata,
-        'model_name': 'test_model',
-        'model_id': 1,
-        'experiment': 'test',
-        'timestamp': 2,
-        'status': 'success',
-    }
-
-    try:
-        await mlflow.update_production_model(input_data)
-    except Exception as e:  # noqa: BLE001
-        assert str(e) == 'Error updating production model'
-        mlflow.send_notification.assert_called_once_with(
-            metadata=metadata['metadata'],
-            notification_id='UPDATE_PRODUCTION_MODEL_ERROR',
-            message='Error updating production model test_model: Error updating production model',
-            block='update_production_model',
-            level=NotificationLevel.ERROR,
-            attachment_content=ANY,
-        )
-    else:
-        raise AssertionError('No exception raised')
-
-
 @mark.asyncio
 async def test_save_model_success(mlflow):
     """Test save_model successfully saves model and artifacts to MLflow."""