gitea_admin/sientia-dataops-model-manager
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

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

Browse Source 
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.
This commit is contained in:
Bruno Domingues
2025-10-17 00:51:29 -03:00
parent 85a8ba1e68
commit a66b996cc1
4 changed files with 9 additions and 580 deletions
Download Patch File Download Diff File Expand all files Collapse all files

12
README.md
View File

@@ -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

293
model_manager/activities/mlflow.py
View File

@@ -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:
"""

35
model_manager/metrics.py
View File

@@ -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],
)

249
tests/activities/test_mlflow.py
View File

@@ -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."""