From a66b996cc1a46ad435c35346bf20dbac819d795d Mon Sep 17 00:00:00 2001 From: Bruno Domingues Date: Fri, 17 Oct 2025 00:51:29 -0300 Subject: [PATCH] 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 | 12 +- model_manager/activities/mlflow.py | 293 +---------------------------- model_manager/metrics.py | 35 +--- tests/activities/test_mlflow.py | 249 ------------------------ 4 files changed, 9 insertions(+), 580 deletions(-) diff --git a/README.md b/README.md index d53e1d6..ee62cfe 100644 --- a/README.md +++ b/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) - - [Data Quality Metrics](#data-quality-metrics) + - [Training Metrics](#training-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 diff --git a/model_manager/activities/mlflow.py b/model_manager/activities/mlflow.py index b01b48a..c158a46 100644 --- a/model_manager/activities/mlflow.py +++ b/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 - data transformation and prediction operations. It handles authentication, - data preprocessing, and model management with configurable retention policies. + 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 comprehensive error handling and logging for all - MLFlow operations, ensuring reliable model inference in production environments. + 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 @@ -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: """ diff --git a/model_manager/metrics.py b/model_manager/metrics.py index 311f80f..0372ebb 100644 --- a/model_manager/metrics.py +++ b/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], -) diff --git a/tests/activities/test_mlflow.py b/tests/activities/test_mlflow.py index 0c965ec..0406083 100644 --- a/tests/activities/test_mlflow.py +++ b/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."""