SIENTIAPDE-1255: Refactor data quality gates to training focused metrics and repositories. This commit removes the data quality gates and filters, focusing on training-specific metrics and data repositories. It also updates the README to reflect these changes, including new training metrics and a streamlined data services section.

This commit is contained in:
Bruno Domingues
2025-10-17 00:39:06 -03:00
parent e191e7849f
commit 85a8ba1e68
11 changed files with 15 additions and 1407 deletions

View File

@@ -167,17 +167,18 @@ The Model Manager system uses a Temporal-based workflow architecture with clear
#### **Data Services (`model_manager/utils/`)**
- **Connectors Config**: Environment variable-based configuration management
- **Repository**: Data access layer for MLFlow operations
- `model_repository.py`: MLFlow model operations and retraining
- **Filters**: Data quality validation and MLFlow response filtering
- `conditional_filters.py`: Input data validation filters
- `mlflow_filters.py`: MLFlow API response validation filters
- **Repository**: Data access layer for training operations
- `training_repository.py`: Training business logic and operations
- **Models**: Data models and schemas
- `train_model_params.py`: Training parameters model
- `train_model_result.py`: Training result model
- `experiment_status.py`: Experiment status enum
- **Key Features**:
- Environment variable-based configuration with sensible defaults
- Connection pool management and optimization
- Security credential management
- Configuration validation and error handling
- Support for MLFlow model flavors
- Type-safe data models with validation
### Data Flow Architecture
@@ -814,10 +815,10 @@ The Model Manager system exposes comprehensive Prometheus metrics for operationa
- 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]
### Data Quality Metrics
- Filter pass/fail rates through notification system
- MLFlow API response validation metrics
- Data quality gate performance tracking
### Training Metrics
- Training success/failure rates through notification system
- Model save performance metrics
- Experiment status tracking
## Configuration
@@ -887,7 +888,7 @@ The project maintains **99%+ code coverage** with comprehensive unit and integra
pytest tests/ --cov=model_manager --cov-report=term-missing --cov-report=xml --cov-report=html
# Run specific test file
pytest tests/activities/test_gates.py -v
pytest tests/activities/test_training.py -v
# Run with coverage visualization
pytest tests/ --cov=model_manager --cov-report=xml
@@ -1019,7 +1020,7 @@ model_manager/
4. **Workflow Execution Failures**
- Review activity error logs and notifications
- Check data quality filter configurations
- Check training parameter validation errors
- Verify input data format and required fields
### Debug Mode

View File

@@ -7,13 +7,12 @@ with workflow.unsafe.imports_passed_through():
from sientia_do.observability.logger import Logger
from model_manager.activities.experiment_tracking import ExperimentTracking
from model_manager.activities.gates import Gates
from model_manager.activities.minio import MinIO
from model_manager.activities.mlflow import MLFlow
from model_manager.activities.training import Training
class Activities(ExperimentTracking, MLFlow, MinIO, Gates, Training):
class Activities(ExperimentTracking, MLFlow, MinIO, Training):
"""
Main activities orchestrator for the Model Manager system.
@@ -23,9 +22,8 @@ class Activities(ExperimentTracking, MLFlow, MinIO, Gates, Training):
The class implements multiple inheritance to combine specialized functionality:
- ExperimentTracking: ML experiment lifecycle tracking and database operations (extends Postgres)
- MLFlow: Model inference and transformation operations
- MLFlow: Model saving and artifact management operations
- MinIO: Object storage operations (file upload/download/delete)
- Gates: Data quality validation and filtering mechanisms
- Training: ML model training operations (extends BaseActivity)
Attributes:
@@ -102,8 +100,6 @@ class Activities(ExperimentTracking, MLFlow, MinIO, Gates, Training):
notification_handler=notification_handler,
)
Gates.__init__(self, logger=logger, notification_handler=notification_handler)
Training.__init__(self, logger=logger, notification_handler=notification_handler)
async def shutdown(self):

View File

@@ -1,583 +0,0 @@
from temporalio import activity, workflow
with workflow.unsafe.imports_passed_through():
import traceback
from typing import Any
from pandas import DataFrame
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_WITH_TZ, now
from model_manager import metrics
from model_manager.utils.filters.conditional_filters import (
filter_empty_data,
filter_specific_variables_null_values,
)
from model_manager.utils.filters.mlflow_filters import api_error_filter, nan_values_filter
# Input filter function mappings
input_filter_functions = {
'SPECIFIC_VARIABLES_NULL_VALUES': filter_specific_variables_null_values,
'EMPTY_DATA': filter_empty_data,
'path_confidence': {'STOP': -1, 'CONTINUE': 2, 'REPEAT': -1},
}
# MLFlow response filter function mappings
mlflow_response_filter_functions = {
'API_ERROR': api_error_filter,
'path_confidence': {'STOP': -1, 'CONTINUE': 10, 'REPEAT': -1},
}
# MLFlow content filter function mappings
mlflow_content_filter_functions = {
'NAN_VALUES': nan_values_filter,
'EMPTY_DATA': filter_empty_data,
'path_confidence': {'STOP': -1, 'CONTINUE': 18, 'REPEAT': -1},
}
class Gates(BaseActivity):
"""
Data quality gates and filtering activities for the Model Manager system.
This class implements comprehensive data quality validation and filtering
mechanisms that can be applied at different stages of the prediction pipeline.
It provides configurable filters with policy-based decision making to ensure
data integrity and quality throughout the ML workflow.
The class supports multiple filter types and implements a flexible policy
system that can be configured for different validation requirements. Each
filter returns a path decision (STOP, CONTINUE, REPEAT) along with confidence
scores and detailed comments for monitoring and debugging.
Attributes:
input_filter_functions (dict): Mapping of input filter names to functions
mlflow_response_filter_functions (dict): Mapping of MLFlow response filter names to functions
mlflow_content_filter_functions (dict): Mapping of MLFlow content filter names to functions
"""
def __init__(self, logger: Logger, notification_handler: NotificationHandler):
"""
Initialize data quality gates with logging and notification capabilities.
Args:
logger: Logger instance for observability and debugging
notification_handler: Notification handler for alerts and monitoring
Raises:
Exception: If BaseActivity initialization fails
"""
BaseActivity.__init__(self, logger, notification_handler, set_error_counter=True)
@activity.defn(name='input_gate')
async def input_gate(self, input_data: dict[str, Any]) -> tuple[str | None, int, str]:
"""
Apply input data quality filters and validation.
This activity validates input data quality using configurable filters
before proceeding with ML operations. It applies multiple filter types
and returns a path decision based on the filter results and configured
policies.
The method implements a comprehensive filtering system that:
1. Applies configured filters to input data
2. Evaluates filter results against policy configurations
3. Determines appropriate path decisions (STOP, CONTINUE, REPEAT)
4. Provides confidence scores and detailed comments
5. Handles errors gracefully with notification integration
Args:
input_data: Configuration and data for input validation
Required keys:
- metadata (dict): Workflow execution metadata
- filters (dict): Filter configuration and policies
- data (dict): Input data to validate
- path_priority (list[str]): Priority order for path decisions
Returns:
tuple: (path_flag, confidence, comment)
- path_flag (str | None): Decision path (STOP, CONTINUE, REPEAT, or None)
- confidence (int): Confidence score for the decision
- comment (str): Detailed explanation of the decision
Raises:
Exception: If filter execution fails or configuration is invalid
"""
metadata = input_data['metadata']
self.info('Performing input gate...', metadata)
filters = input_data['filters']
data = DataFrame(input_data['data'])
path_priority = input_data['path_priority']
filter_output = []
self.debug(f'Input data: {data.head(5).to_string()}', metadata)
self.debug(f'Filters: {filters}', metadata)
# Apply each configured filter
for fil, config in filters.items():
if fil not in input_filter_functions:
self.error(f'Filter {fil} not found', metadata)
continue
try:
if input_filter_functions[fil](data, config['config']): # type: ignore[operator]
self.debug(f'Data not passed the input filter {fil}:{config}', metadata)
filter_output.append(config['policy']) # type: ignore[index]
except Exception as e: # noqa: BLE001
trace = traceback.format_exc()
self.send_notification(
metadata=metadata,
notification_id=f'INTPUT_GATE_ERROR__{fil}',
message=f'Error in filter {fil}:{config}: \n {e}',
block='input_gate',
level=NotificationLevel.ERROR,
attachment_content=trace,
)
for path_flag in path_priority:
if path_flag in filter_output:
self.info(f'Input gate result: {path_flag}', metadata)
return (
path_flag,
input_filter_functions['path_confidence'][path_flag], # type: ignore[index]
'Input data with bad quality',
)
self.info('Nothing was filtered by the input gate', metadata)
return None, 0, ''
@activity.defn(name='mlflow_response_gate')
async def mlflow_response_gate(self, input_data: dict[str, Any]) -> tuple[str | None, int, str]:
"""
Validate MLFlow API response quality and integrity.
This activity validates MLFlow API responses to ensure they meet quality
standards before proceeding with further processing. It applies response-specific
filters and determines appropriate path decisions based on response quality.
The method implements response validation that:
1. Applies MLFlow response-specific filters
2. Evaluates API response quality and integrity
3. Determines path decisions based on response validation results
4. Provides confidence scores and detailed validation comments
5. Handles API errors and response validation failures
Args:
input_data: Configuration and data for response validation
Required keys:
- metadata (dict): Workflow execution metadata
- filters (dict): Response filter configuration and policies
- data (dict): MLFlow API response data to validate
- type (str): Type of MLFlow operation (transform, predict)
- path_priority (list[str]): Priority order for path decisions
Returns:
tuple: (path_flag, confidence, comment)
- path_flag (str | None): Decision path (STOP, CONTINUE, REPEAT, or None)
- confidence (int): Confidence score for the decision
- comment (str): Detailed explanation of the decision
Raises:
Exception: If response validation fails or configuration is invalid
"""
metadata = input_data['metadata']
self.info('Performing mlflow response gate...', metadata)
filters = input_data['filters']
data = input_data['data']
gate_type = input_data['type']
path_priority = input_data['path_priority']
filter_output = []
self.debug(f'Input data: \n {create_sample_dict(data, max_items=5, max_depth=2)}', metadata)
self.debug(f'Filters: {filters}', metadata)
comments = []
for fil, config in filters.items():
if fil not in mlflow_response_filter_functions:
self.error(f'Filter {fil} not found', metadata)
continue
try:
if mlflow_response_filter_functions[fil](data, config): # type: ignore[operator]
filter_output.append(config['policy']) # type: ignore[index]
comments.append(data['content']['message'])
self.send_notification(
metadata=metadata,
notification_id=f'{gate_type.upper()}_GATE_RESPONSE_FILTER__{fil}',
message=data['content']['message'],
block='mlflow_gate',
level=NotificationLevel.ERROR,
attachment_content=data['content']['traceback'],
)
except Exception as e: # noqa: BLE001
trace = traceback.format_exc()
self.send_notification(
metadata=metadata,
notification_id=f'MLFLOW_GATE_RESPONSE_FILTER__{fil}',
message=f'Error in filter {fil}:{config}: \n {e}',
block='mlflow_gate',
level=NotificationLevel.ERROR,
attachment_content=trace,
)
for path_flag in path_priority:
if path_flag in filter_output:
self.info(f'Mlflow response gate result: {path_flag}', metadata)
return (
path_flag,
mlflow_response_filter_functions['path_confidence'][path_flag], # type: ignore[index]
', '.join(comments),
)
self.info('Nothing was filtered by the mlflow response gate', metadata)
return None, 0, ''
@activity.defn(name='mlflow_content_gate')
async def mlflow_content_gate(self, input_data: dict[str, Any]) -> tuple[str | None, int, str]:
"""
Validate MLFlow prediction content quality and integrity.
This activity validates the content of MLFlow predictions to ensure they
meet quality standards before export and persistence. It applies content-specific
filters and determines appropriate path decisions based on content quality.
The method implements content validation that:
1. Applies MLFlow content-specific filters
2. Evaluates prediction content quality and integrity
3. Determines path decisions based on content validation results
4. Provides confidence scores and detailed validation comments
5. Handles content validation failures and quality issues
Args:
input_data: Configuration and data for content validation
Required keys:
- metadata (dict): Workflow execution metadata
- filters (dict): Content filter configuration and policies
- data (dict): MLFlow prediction content to validate
- type (str): Type of MLFlow operation (transform, predict)
- path_priority (list[str]): Priority order for path decisions
Returns:
tuple: (path_flag, confidence, comment)
- path_flag (str | None): Decision path (STOP, CONTINUE, REPEAT, or None)
- confidence (int): Confidence score for the decision
- comment (str): Detailed explanation of the decision
Raises:
Exception: If content validation fails or configuration is invalid
"""
metadata = input_data['metadata']
self.info('Performing mlflow content gate...', metadata)
filters = input_data['filters']
data = DataFrame(input_data['data'])
gate_type = input_data['type']
path_priority = input_data['path_priority']
filter_output = []
self.debug(f'Input data:\n {data.head(5).to_string()}', metadata)
self.debug(f'Filters: \n {create_sample_dict(filters)}', metadata)
for fil, config in filters.items():
if fil not in mlflow_content_filter_functions:
continue
try:
if mlflow_content_filter_functions[fil](data, config): # type: ignore[operator]
filter_output.append(config['policy']) # type: ignore[index]
self.send_notification(
metadata=metadata,
notification_id=f'{gate_type.upper()}_GATE_CONTENT_FILTER__{fil}',
message=f'Data not passed the content filter {fil}:{config}',
block='mlflow_gate',
level=NotificationLevel.WARNING,
attachment_content=data.to_string(),
)
except Exception as e: # noqa: BLE001
trace = traceback.format_exc()
self.send_notification(
metadata=metadata,
notification_id=f'MLFLOW_GATE_CONTENT_FILTER__{fil}',
message=f'Error in filter {fil}:{config}: \n {e}',
block='mlflow_gate',
level=NotificationLevel.ERROR,
attachment_content=trace,
)
for path_flag in path_priority:
if path_flag in filter_output:
self.info(f'Mlflow content gate result: {path_flag}', metadata)
return (
path_flag,
mlflow_content_filter_functions['path_confidence'][path_flag], # type: ignore[index]
'Transformed data not passed the content filter',
)
self.info('Nothing was filtered by the mlflow content gate', metadata)
return None, 0, ''
def get_prediction_store_policy(
self, prediction_store_policy: str, metadata: dict[str, Any]
) -> tuple[str, int]:
"""
Parse and validate prediction store policy configuration.
This method parses prediction store policy strings in the format 'type:value'
and validates them against allowed policy types and values. It provides
sensible defaults for invalid configurations and logs policy validation
failures for operational monitoring.
Supported Policy Types:
- 'lts': Latest timestamp - sorts data by timestamp descending
- 'erl': Earliest timestamp - sorts data by timestamp ascending
Args:
prediction_store_policy (str): Policy string in format 'type:value'
metadata (dict[str, Any]): Context metadata for logging and notifications
Returns:
tuple[str, int]: (policy_type, policy_value)
- policy_type (str): Validated policy type ('lts' or 'erl')
- policy_value (int): Number of rows to retain
"""
policy_elements = prediction_store_policy.split(':')
if len(policy_elements) < 2:
self.error(
f'Invalid prediction store policy: {prediction_store_policy}, using default policy',
metadata,
)
return 'lts', 1
policy_type = policy_elements[0]
policy_value = policy_elements[1]
# If the policy_type is not lts or erl, we use the default policy
# If the policty_value is not a number or 0, we use the default policy
if (
policy_type not in ['lts', 'erl']
or not policy_value.isdigit()
or int(policy_value) == 0
):
self.error(
f'Invalid prediction store policy: {prediction_store_policy}, using default policy',
metadata,
)
return 'lts', 1
return policy_type, int(policy_value)
@activity.defn(name='format_prediction')
async def format_prediction(self, input_data: dict[str, Any]) -> dict[Any, Any]:
"""
Format prediction data according to configured storage policies.
This method formats prediction data for storage and export operations.
It applies timestamp-based sorting policies, adds metadata fields,
and ensures data consistency before persistence. The method supports
multiple storage policies for flexible data retention strategies.
Storage Policies:
- 'lts:N': Latest timestamp - retains N most recent predictions
- 'erl:N': Earliest timestamp - retains N oldest predictions
Args:
input_data (dict): Input data containing:
- data (dict[str, Any]): Raw prediction data to format
- timestamp (str): Default timestamp if data lacks timestamp column
- model_id (str): Unique identifier for the ML model
- prediction_confidence (float): Confidence score for the prediction
- prediction_store_policy (str): Storage policy in format 'type:value'
Returns:
dict: Formatted prediction data ready for storage and export
"""
metadata = input_data['metadata']
prediction_store_policy = input_data['prediction_store_policy']
self.info('Formatting prediction...', metadata)
data = DataFrame(input_data['data'])
# Create timestamp column from index and reset index
data['timestamp'] = data.index
data = data.reset_index(drop=True)
self.debug(f'Prediction store policy: {prediction_store_policy}', metadata)
self.debug(f'Prediction data: {data.head(5).to_string()}', metadata)
policy_type, policy_value = self.get_prediction_store_policy(
prediction_store_policy, metadata
)
# If data has no timestamp, we use the default timestamp and not sort the data
self.info(
f'Sorting data by timestamp and applying policy: {policy_type}:{policy_value}', metadata
)
# If policy_type is lts, we need to sort the data by timestamp descending and take the first policy_value rows
if policy_type == 'lts':
self.debug('Sorting data by timestamp descending', metadata)
data = data.sort_values(by='timestamp', ascending=False)
# If policy_type is erl, we need to sort the data by timestamp ascending and take the first policy_value rows
elif policy_type == 'erl':
self.debug('Sorting data by timestamp ascending', metadata)
data = data.sort_values(by='timestamp', ascending=True)
else:
self.error(f'Invalid policy type: {policy_type}, using default policy', metadata)
raise ValueError(f'Invalid policy type: {policy_type}')
data = data.head(int(policy_value))
data['model_id'] = input_data['model_id']
data['prediction_confidence'] = input_data['prediction_confidence']
data['prediction_status'] = 'Good'
data['comments'] = ''
data = data.sort_values(by='timestamp', ascending=False)
data = data.reset_index(drop=True)
self.info(f'Prediction formatted: {len(data)} rows', metadata)
self.debug(f'Prediction data: {data.head(5).to_string()}', metadata)
return data.to_dict()
@activity.defn(name='format_default_prediction')
async def format_default_prediction(self, input_data: dict[str, Any]) -> dict[Any, Any]:
"""
Create and format default prediction data for error conditions.
This method generates default prediction data when the main prediction
pipeline encounters errors or quality issues. It creates a standardized
data structure with zero values for predictions and useful metadata
for operational monitoring and debugging.
The default prediction serves as a fallback mechanism to:
1. Maintain data pipeline continuity during failures
2. Provide operational visibility into prediction quality issues
3. Enable downstream systems to handle error conditions gracefully
4. Support debugging and troubleshooting efforts
Args:
input_data (dict): Input data containing:
- timestamp (str): Timestamp for the default prediction
- model_id (str): Unique identifier for the ML model
- prediction_confidence (float): Confidence score (typically low for errors)
- comment (str): Error description or operational comment
Returns:
dict: Formatted default prediction data with error indicators
"""
metadata = input_data['metadata']
self.debug('Formatting default prediction...', metadata)
data = DataFrame(
{
'prediction': [0],
'response_time': [0],
'timestamp': [input_data['timestamp']],
'model_id': [input_data['model_id']],
'prediction_confidence': [input_data['prediction_confidence']],
'prediction_status': ['Bad'],
'comments': [input_data['comment']],
}
)
self.info(f'Default prediction formatted: {data.size} rows', metadata)
return data.to_dict()
@activity.defn(name='get_last_timestamp')
async def get_last_timestamp(self, input_data: dict[str, Any]) -> str:
"""
Extract the most recent timestamp from prediction data.
This method analyzes prediction data to find the latest timestamp,
enabling incremental processing and data continuity tracking.
It handles empty datasets gracefully by returning the current time
as a fallback timestamp.
The method is essential for:
1. Incremental data processing workflows
2. Data continuity validation
3. Timestamp-based data loading optimization
4. Workflow execution tracking
Args:
input_data (dict): Input data containing:
- data (dict[str, Any]): Prediction data to analyze
Returns:
str: Formatted timestamp string in UTC with timezone
"""
metadata = input_data['metadata']
self.info('Getting last timestamp...', metadata)
data = DataFrame(input_data['data'])
self.debug(f'Input data: {data.head(5).to_string()}', metadata)
if data.empty:
return now().strftime(DATETIME_FORMAT_WITH_TZ)
max_timestamp = max(data['timestamp'].values.tolist())
self.info(f'Last timestamp: {max_timestamp}', metadata)
return max_timestamp
@activity.defn(name='write_metrics')
async def write_metrics(self, input_data: dict[str, Any]):
"""
Write prediction performance metrics to Prometheus monitoring system.
This method records comprehensive metrics for prediction operations,
enabling operational monitoring, performance analysis, and alerting.
It tracks prediction counts, confidence levels, and response times
for each model and pipeline combination.
Metrics Recorded:
1. Prediction Count: Incremental counter for successful predictions
2. Confidence Monitor: Current confidence level for predictions
3. Response Time Monitor: Histogram of prediction response times
Args:
input_data (dict): Input data containing:
- metadata (dict[str, Any]): Workflow execution metadata
- prediction (dict[str, Any]): Prediction data with metrics
Raises:
Exception: If metrics writing fails or configuration is invalid
"""
metadata = input_data['metadata']
prediction = DataFrame(input_data['prediction'])
prediction_confidence = prediction['prediction_confidence'].values[0]
response_time = prediction['response_time'].values[0]
self.info(f'Writing metrics for model {metadata["model_name"]}', metadata)
metrics.PREDICTIONS_WRITTEN_COUNT.labels(
pod_id=self.pod_id,
model_name=metadata['model_name'],
pipeline_name=metadata['workflow_name'],
).inc()
metrics.PREDICTION_CONFIDENCE_MONITOR.labels(
pod_id=self.pod_id,
model_name=metadata['model_name'],
pipeline_name=metadata['workflow_name'],
).set(prediction_confidence)
metrics.PREDICTION_RESPONSE_TIME_MONITOR.labels(
pod_id=self.pod_id,
model_name=metadata['model_name'],
pipeline_name=metadata['workflow_name'],
).observe(response_time)
self.info(f'Metrics written for model {metadata["model_name"]}', metadata)

View File

@@ -1,44 +0,0 @@
from pandas import DataFrame
def filter_specific_variables_null_values(data: DataFrame, config: dict) -> bool:
"""
Filter to check if specific variables contain null values.
This function examines a DataFrame to determine if any of the specified variables
contain null (NaN) values. It returns True if null values are found for any of
the specified variables, False otherwise.
Args:
data (DataFrame): The pandas DataFrame to be examined. Must contain columns
named 'variable' and 'value'.
config (dict): Configuration dictionary containing the following key:
- variables (list): List of variable names to check for null values
Returns:
bool: True if any of the specified variables contain null values,
False if none of the specified variables contain null values.
"""
return not data[data['variable'].isin(config['variables']) & data['value'].isna()].empty
def filter_empty_data(data: DataFrame, _config: dict) -> bool:
"""
Filter to check if the DataFrame is empty.
This function determines whether the provided DataFrame contains any data.
It's a simple utility function that can be used in conditional logic to
handle cases where no data is available.
Args:
data (DataFrame): The pandas DataFrame to be checked for emptiness.
_config (dict): Configuration dictionary (unused in this function).
The underscore prefix indicates this parameter is required for
interface consistency but not used in the implementation.
Returns:
bool: True if the DataFrame is empty (has no rows), False if it contains data.
"""
return data.empty

View File

@@ -1,64 +0,0 @@
import numpy as np
from pandas import DataFrame
def api_error_filter(response: dict, _config: dict) -> bool:
"""
Filter MLFlow API responses for error conditions.
This function analyzes MLFlow API responses to detect error conditions
and determine if the response should be filtered out due to quality
or reliability issues.
Args:
response: MLFlow API response data (dict)
_config: Filter configuration dictionary
Required keys:
- error_codes (list, optional): List of error codes to detect
- error_keywords (list, optional): List of error keywords to detect
- check_structure (bool, optional): Whether to validate response structure
Returns:
bool: True if data should be filtered (contains errors), False otherwise
"""
if not response:
return True
if not response['success']:
return True
return False
def nan_values_filter(predictions: DataFrame, _config: dict) -> bool:
"""
Filter data for NaN (Not a Number) values.
This function detects NaN values in MLFlow prediction results and
determines if the data quality is sufficient for further processing
or export operations.
Args:
predictions: DataFrame containing prediction data to check for NaN values
_config: Filter configuration dictionary
Required keys:
- max_nan_ratio (float, optional): Maximum allowed NaN value ratio (0.0 to 1.0)
- max_nan_count (int, optional): Maximum allowed NaN value count
- check_nested (bool, optional): Whether to check nested data structures
Returns:
bool: True if data should be filtered (too many NaN values), False otherwise
"""
data = (
predictions.replace({None: np.nan})
.infer_objects(copy=False)
.drop(columns=['timestamp'], errors='ignore')
)
if data.isna().all().all():
return True
return False

View File

@@ -4,7 +4,6 @@ from pytest import mark
from model_manager.activities.activities import Activities
from model_manager.activities.experiment_tracking import ExperimentTracking
from model_manager.activities.gates import Gates
from model_manager.activities.mlflow import MLFlow
from model_manager.activities.training import Training
@@ -12,11 +11,9 @@ from model_manager.activities.training import Training
@patch('model_manager.activities.activities.ExperimentTracking.__init__')
@patch('model_manager.activities.activities.MLFlow.__init__')
@patch('model_manager.activities.activities.MinIO.__init__')
@patch('model_manager.activities.activities.Gates.__init__')
@patch('model_manager.activities.activities.Training.__init__')
def test___init__(
mock_training_init,
mock_gates_init,
mock_minio_init,
mock_mlflow_init,
mock_experiment_tracking_init,
@@ -59,7 +56,6 @@ def test___init__(
assert isinstance(activities, Activities)
assert isinstance(activities, ExperimentTracking)
assert isinstance(activities, MLFlow)
assert isinstance(activities, Gates)
assert isinstance(activities, Training)
mock_experiment_tracking_init.assert_called_once_with(
@@ -100,10 +96,6 @@ def test___init__(
notification_handler=notification_handler,
)
mock_gates_init.assert_called_once_with(
ANY, logger=logger, notification_handler=notification_handler
)
mock_training_init.assert_called_once_with(
ANY, logger=logger, notification_handler=notification_handler
)

View File

@@ -1,630 +0,0 @@
from unittest.mock import ANY, MagicMock, patch
from pytest import fixture, mark
from sientia_do.notifications.models import NotificationLevel
from model_manager.activities.gates import Gates
@fixture
def gates_activity():
gates = Gates(
logger=MagicMock(),
notification_handler=MagicMock(),
)
gates.error = MagicMock()
gates.debug = MagicMock()
gates.info = MagicMock()
gates.warning = MagicMock()
gates.critical = MagicMock()
gates.send_notification = MagicMock()
return gates
metadata = {
'metadata': {
'model_id': 'test_model',
'model_name': 'test_model',
'workflow_name': 'test_workflow',
'schema_name': 'test_schedule',
},
}
@mark.asyncio
async def test_input_gate_invalid_filter(gates_activity):
# Arrange
input_data = {
**metadata,
'filters': {'INVALID_FILTER': {'POLICY': 'STOP'}},
'data': {'value': [1, 2, 3]},
'path_priority': ['STOP', 'CONTINUE', 'REPEAT'],
}
# Act
result = await gates_activity.input_gate(input_data)
# Assert
assert result == (None, 0, '')
gates_activity.error.assert_called_once_with(
'Filter INVALID_FILTER not found', metadata['metadata']
)
@mark.asyncio
@patch('model_manager.activities.gates.input_filter_functions')
async def test_input_gate_filter_exception(mock_input_filter_functions, gates_activity):
# Arrange
mock_input_filter_functions.__contains__.return_value = True
mock_input_filter_functions.__getitem__.return_value = MagicMock(
side_effect=Exception('Test error')
)
input_data = {
**metadata,
'filters': {'EMPTY_DATA': {'policy': 'STOP', 'config': {}}},
'data': {'value': []},
'path_priority': ['STOP', 'CONTINUE', 'REPEAT'],
}
# Act
result = await gates_activity.input_gate(input_data)
# Assert
assert result == (None, 0, '')
gates_activity.send_notification.assert_called_once_with(
metadata=metadata['metadata'],
notification_id='INTPUT_GATE_ERROR__EMPTY_DATA',
message="Error in filter EMPTY_DATA:{'policy': 'STOP', 'config': {}}: \n Test error",
block='input_gate',
level=NotificationLevel.ERROR,
attachment_content=ANY,
)
@mark.asyncio
async def test_input_gate_no_filters(gates_activity):
# Arrange
input_data = {
**metadata,
'filters': {},
'data': {'value': [1, 2, 3]},
'path_priority': ['CONTINUE', 'STOP', 'REPEAT'],
}
# Act
result = await gates_activity.input_gate(input_data)
# Assert
assert result == (None, 0, '')
gates_activity.debug.assert_called()
@mark.asyncio
async def test_input_gate_with_filter(gates_activity):
# Arrange
input_data = {
**metadata,
'filters': {'EMPTY_DATA': {'policy': 'STOP', 'config': {}}},
'data': {'value': []},
'path_priority': ['STOP', 'CONTINUE', 'REPEAT'],
}
# Act
result = await gates_activity.input_gate(input_data)
# Assert
assert result == ('STOP', -1, 'Input data with bad quality')
gates_activity.debug.assert_called()
@mark.asyncio
async def test_input_gate_filter_returns_false(gates_activity):
"""Test to cover line 129 branch when filter returns False (filter passes)."""
# Arrange - Use data that will NOT trigger EMPTY_DATA filter (has data)
input_data = {
**metadata,
'filters': {'EMPTY_DATA': {'policy': 'STOP', 'config': {}}},
'data': {'value': [1, 2, 3, 4, 5]}, # Has data, filter returns False
'path_priority': ['STOP', 'CONTINUE', 'REPEAT'],
}
# Act
result = await gates_activity.input_gate(input_data)
# Assert - Filter returns False, so no policy is added to filter_output
assert result == (None, 0, '') # No filter triggered
gates_activity.debug.assert_called()
@mark.asyncio
async def test_mlflow_response_gate_invalid_filter(gates_activity):
# Arrange
input_data = {
**metadata,
'filters': {'INVALID_FILTER': {'POLICY': 'STOP'}},
'data': {'content': {'message': 'success'}},
'type': 'test',
'path_priority': ['STOP', 'CONTINUE', 'REPEAT'],
}
# Act
result = await gates_activity.mlflow_response_gate(input_data)
# Assert
assert result == (None, 0, '')
@mark.asyncio
@patch('model_manager.activities.gates.mlflow_response_filter_functions')
async def test_mlflow_response_gate_filter_exception(
mock_mlflow_response_filter_functions, gates_activity
):
# Arrange
mock_mlflow_response_filter_functions.__contains__.return_value = True
mock_mlflow_response_filter_functions.__getitem__.return_value = MagicMock(
side_effect=Exception('Test error')
)
input_data = {
**metadata,
'filters': {'INVALID_FILTER': {'POLICY': 'STOP'}},
'data': {'content': {'message': 'success'}},
'type': 'test',
'path_priority': ['STOP', 'CONTINUE', 'REPEAT'],
}
# Act
result = await gates_activity.mlflow_response_gate(input_data)
# Assert
assert result == (None, 0, '')
gates_activity.send_notification.assert_called_once_with(
metadata=metadata['metadata'],
notification_id='MLFLOW_GATE_RESPONSE_FILTER__INVALID_FILTER',
message="Error in filter INVALID_FILTER:{'POLICY': 'STOP'}: \n Test error",
block='mlflow_gate',
level=NotificationLevel.ERROR,
attachment_content=ANY,
)
@mark.asyncio
async def test_mlflow_response_gate_no_filters(gates_activity):
# Arrange
input_data = {
**metadata,
'filters': {},
'data': {'content': {'message': 'success'}},
'type': 'test',
'path_priority': ['CONTINUE', 'STOP', 'REPEAT'],
}
# Act
result = await gates_activity.mlflow_response_gate(input_data)
# Assert
assert result == (None, 0, '')
gates_activity.debug.assert_called()
@mark.asyncio
async def test_mlflow_response_gate_with_filter(gates_activity):
# Arrange
input_data = {
**metadata,
'filters': {'API_ERROR': {'policy': 'STOP'}},
'data': {
'success': False,
'content': {'message': 'API error occurred', 'traceback': 'error trace'},
},
'type': 'test',
'path_priority': ['STOP', 'CONTINUE', 'REPEAT'],
}
# Act
result = await gates_activity.mlflow_response_gate(input_data)
# Assert
assert result == ('STOP', -1, 'API error occurred')
gates_activity.debug.assert_called()
gates_activity.send_notification.assert_called()
@mark.asyncio
async def test_mlflow_response_gate_filter_returns_false(gates_activity):
"""Test to cover line 208 branch when filter returns False (no API error)."""
# Arrange - Use data that will NOT trigger API_ERROR filter (success=True)
input_data = {
**metadata,
'filters': {'API_ERROR': {'policy': 'STOP'}},
'data': {
'success': True, # Success=True, filter returns False
'content': {'message': 'Operation successful', 'result': 'data'},
},
'type': 'transform',
'path_priority': ['STOP', 'CONTINUE', 'REPEAT'],
}
# Act
result = await gates_activity.mlflow_response_gate(input_data)
# Assert - Filter returns False, so no policy is added to filter_output
assert result == (None, 0, '') # No filter triggered
gates_activity.debug.assert_called()
@mark.asyncio
async def test_mlflow_content_gate_invalid_filter(gates_activity):
# Arrange
input_data = {
**metadata,
'filters': {'INVALID_FILTER': {'POLICY': 'STOP'}},
'data': {'value': [1, 2, 3]},
'type': 'test',
'path_priority': ['STOP', 'CONTINUE', 'REPEAT'],
}
# Act
result = await gates_activity.mlflow_content_gate(input_data)
# Assert
assert result == (None, 0, '')
@mark.asyncio
@patch('model_manager.activities.gates.mlflow_content_filter_functions')
async def test_mlflow_content_gate_filter_exception(
mock_mlflow_content_filter_functions, gates_activity
):
# Arrange
mock_mlflow_content_filter_functions.__contains__.return_value = True
mock_mlflow_content_filter_functions.__getitem__.return_value = MagicMock(
side_effect=Exception('Test error')
)
input_data = {
**metadata,
'filters': {'API_ERROR': {'POLICY': 'STOP'}},
'data': {
'success': False,
'content': {'message': 'API error occurred', 'traceback': 'error trace'},
},
'type': 'test',
'path_priority': ['STOP', 'CONTINUE', 'REPEAT'],
}
# Act
result = await gates_activity.mlflow_content_gate(input_data)
# Assert
assert result == (None, 0, '')
gates_activity.debug.assert_called()
gates_activity.send_notification.assert_called_once_with(
metadata=metadata['metadata'],
notification_id='MLFLOW_GATE_CONTENT_FILTER__API_ERROR',
message="Error in filter API_ERROR:{'POLICY': 'STOP'}: \n Test error",
block='mlflow_gate',
level=NotificationLevel.ERROR,
attachment_content=ANY,
)
@mark.asyncio
async def test_mlflow_content_gate_no_filters(gates_activity):
# Arrange
input_data = {
**metadata,
'filters': {},
'data': {'value': [1, 2, 3]},
'type': 'test',
'path_priority': ['CONTINUE', 'STOP', 'REPEAT'],
}
# Act
result = await gates_activity.mlflow_content_gate(input_data)
# Assert
assert result == (None, 0, '')
gates_activity.debug.assert_called()
@mark.asyncio
async def test_mlflow_content_gate_with_filter(gates_activity):
# Arrange
input_data = {
**metadata,
'filters': {'NAN_VALUES': {'policy': 'STOP', 'config': {}}},
'data': {'value': [None, None, None]},
'type': 'test',
'path_priority': ['STOP', 'CONTINUE', 'REPEAT'],
}
# Act
result = await gates_activity.mlflow_content_gate(input_data)
# Assert
assert result == ('STOP', -1, 'Transformed data not passed the content filter')
gates_activity.debug.assert_called()
gates_activity.send_notification.assert_called()
@mark.asyncio
async def test_mlflow_content_gate_filter_returns_false(gates_activity):
"""Test to cover line 293 branch when filter returns False (no NaN values)."""
# Arrange - Use data that will NOT trigger NAN_VALUES filter (no NaN)
input_data = {
**metadata,
'filters': {'NAN_VALUES': {'policy': 'STOP', 'config': {}}},
'data': {'value': [1.0, 2.0, 3.0, 4.0, 5.0]}, # All valid numbers, no NaN
'type': 'predict',
'path_priority': ['STOP', 'CONTINUE', 'REPEAT'],
}
# Act
result = await gates_activity.mlflow_content_gate(input_data)
# Assert - Filter returns False, so no policy is added to filter_output
assert result == (None, 0, '') # No filter triggered
gates_activity.debug.assert_called()
def test_get_prediction_store_policy_invalid_policy(gates_activity):
# Arrange
prediction_store_policy = 'INVALID_POLICY'
# Act
policy_type, policy_value = gates_activity.get_prediction_store_policy(
prediction_store_policy, metadata
)
# Assert
assert policy_type == 'lts'
assert policy_value == 1
def test_get_prediction_store_policy_invalid_policy_value(gates_activity):
# Arrange
prediction_store_policy = 'abc:INVALID_VALUE'
# Act
policy_type, policy_value = gates_activity.get_prediction_store_policy(
prediction_store_policy, metadata
)
# Assert
assert policy_type == 'lts'
assert policy_value == 1
def test_get_prediction_store_policy_valid_policy_type(gates_activity):
# Arrange
prediction_store_policy = 'abc:1'
# Act
policy_type, policy_value = gates_activity.get_prediction_store_policy(
prediction_store_policy, metadata
)
# Assert
assert policy_type == 'lts'
assert policy_value == 1
def test_get_prediction_store_policy_valid_policy(gates_activity):
# Arrange
prediction_store_policy = 'erl:1'
# Act
policy_type, policy_value = gates_activity.get_prediction_store_policy(
prediction_store_policy, metadata
)
# Assert
assert policy_type == 'erl'
assert policy_value == 1
@mark.asyncio
async def test_format_prediction_no_timestamp(gates_activity):
# Arrange
input_data = {
**metadata,
'data': {
'prediction': {'2023-05-26 11:12:27': 1},
'response_time': {'2023-05-26 11:12:27': 0.1},
},
'model_id': 'test_model',
'prediction_confidence': 0.9,
'prediction_store_policy': 'lts:1',
}
# Act
result = await gates_activity.format_prediction(input_data)
# Assert
assert result['prediction'] == {0: 1}
assert result['response_time'] == {0: ANY}
assert result['timestamp'] == {0: '2023-05-26 11:12:27'}
assert result['model_id'] == {0: 'test_model'}
assert result['prediction_confidence'] == {0: 0.9}
assert result['prediction_status'] == {0: 'Good'}
assert result['comments'] == {0: ''}
@mark.asyncio
async def test_format_prediction_with_timestamp_erl(gates_activity):
# Arrange
input_data = {
**metadata,
'data': {
'prediction': {
'2023-05-26 11:12:27': 1,
'2023-05-26 11:12:28': 2,
'2023-05-26 11:12:29': 3,
},
'response_time': {
'2023-05-26 11:12:27': 0.1,
'2023-05-26 11:12:28': 0.2,
'2023-05-26 11:12:29': 0.3,
},
},
'model_id': 'test_model',
'prediction_confidence': 0.9,
'prediction_store_policy': 'erl:2',
}
# Act
result = await gates_activity.format_prediction(input_data)
# Assert
assert result['prediction'] == {0: 2, 1: 1}
assert result['response_time'] == {0: 0.2, 1: 0.1}
assert result['timestamp'] == {0: '2023-05-26 11:12:28', 1: '2023-05-26 11:12:27'}
assert result['model_id'] == {0: 'test_model', 1: 'test_model'}
assert result['prediction_confidence'] == {0: 0.9, 1: 0.9}
assert result['prediction_status'] == {0: 'Good', 1: 'Good'}
assert result['comments'] == {0: '', 1: ''}
@mark.asyncio
async def test_format_prediction_with_timestamp_lts(gates_activity):
# Arrange
input_data = {
**metadata,
'data': {
'prediction': {
'2023-05-26 11:12:27': 1,
'2023-05-26 11:12:28': 2,
'2023-05-26 11:12:29': 3,
},
'response_time': {
'2023-05-26 11:12:27': 0.1,
'2023-05-26 11:12:28': 0.2,
'2023-05-26 11:12:29': 0.3,
},
},
'model_id': 'test_model',
'prediction_confidence': 0.9,
'prediction_store_policy': 'lts:2',
}
# Act
result = await gates_activity.format_prediction(input_data)
# Assert
assert result['prediction'] == {0: 3, 1: 2}
assert result['response_time'] == {0: 0.3, 1: 0.2}
assert result['timestamp'] == {0: '2023-05-26 11:12:29', 1: '2023-05-26 11:12:28'}
assert result['model_id'] == {0: 'test_model', 1: 'test_model'}
assert result['prediction_confidence'] == {0: 0.9, 1: 0.9}
assert result['prediction_status'] == {0: 'Good', 1: 'Good'}
assert result['comments'] == {0: '', 1: ''}
@mark.asyncio
async def test_format_prediction_with_timestamp_invalid_policy(gates_activity):
# Arrange
input_data = {
**metadata,
'data': {
'prediction': [1, 2, 3],
'response_time': [0.1, 0.2, 0.3],
'timestamp': ['2023-05-26 11:12:27', '2023-05-26 11:12:28', '2023-05-26 11:12:29'],
},
'model_id': 'test_model',
'prediction_confidence': 0.9,
'prediction_store_policy': 'lts:2',
}
gates_activity.get_prediction_store_policy = MagicMock(return_value=('invalid', 1))
try:
await gates_activity.format_prediction(input_data)
except ValueError as e:
assert str(e) == 'Invalid policy type: invalid'
else:
raise AssertionError('Expected ValueError')
@mark.asyncio
async def test_format_default_prediction(gates_activity):
# Arrange
input_data = {
**metadata,
'timestamp': '2023-05-26 11:12:27',
'model_id': 'test_model',
'prediction_confidence': 0.1,
'comment': 'Test comment',
}
# Act
result = await gates_activity.format_default_prediction(input_data)
# Assert
assert result['prediction'] == {0: 0}
assert result['response_time'] == {0: 0}
assert result['timestamp'] == {0: '2023-05-26 11:12:27'}
assert result['model_id'] == {0: 'test_model'}
assert result['prediction_confidence'] == {0: 0.1}
assert result['prediction_status'] == {0: 'Bad'}
assert result['comments'] == {0: 'Test comment'}
gates_activity.debug.assert_called()
@mark.asyncio
async def test_get_last_timestamp_with_data(gates_activity):
# Arrange
input_data = {**metadata, 'data': {'timestamp': ['2023-05-26 11:12:27', '2023-05-26 11:12:28']}}
# Act
result = await gates_activity.get_last_timestamp(input_data)
# Assert
assert result == '2023-05-26 11:12:28'
@mark.asyncio
async def test_get_last_timestamp_no_data(gates_activity):
# Arrange
input_data = {'data': {}, **metadata}
# Act
result = await gates_activity.get_last_timestamp(input_data)
# Assert
assert isinstance(result, str) # Should be a timestamp string
assert len(result) > 0
@mark.asyncio
@patch('model_manager.activities.gates.metrics')
async def test_write_metrics(mock_metrics, gates_activity):
"""Test write_metrics method."""
input_data = {
**metadata,
'prediction': {
'prediction': [1, 2, 3],
'prediction_confidence': [0.9, 0.8, 0.7],
'response_time': [0.1, 0.2, 0.3],
},
}
await gates_activity.write_metrics(input_data)
mock_metrics.PREDICTIONS_WRITTEN_COUNT.labels.assert_called_once_with(
pod_id=gates_activity.pod_id,
model_name=metadata['metadata']['model_name'],
pipeline_name=metadata['metadata']['workflow_name'],
)
mock_metrics.PREDICTIONS_WRITTEN_COUNT.labels.return_value.inc.assert_called_once_with()
mock_metrics.PREDICTION_CONFIDENCE_MONITOR.labels.assert_called_once_with(
pod_id=gates_activity.pod_id,
model_name=metadata['metadata']['model_name'],
pipeline_name=metadata['metadata']['workflow_name'],
)
mock_metrics.PREDICTION_CONFIDENCE_MONITOR.labels.return_value.set.assert_called_once_with(0.9)
mock_metrics.PREDICTION_RESPONSE_TIME_MONITOR.labels.assert_called_once_with(
pod_id=gates_activity.pod_id,
model_name=metadata['metadata']['model_name'],
pipeline_name=metadata['metadata']['workflow_name'],
)
mock_metrics.PREDICTION_RESPONSE_TIME_MONITOR.labels.return_value.observe.assert_called_once_with(
0.1
)

View File

@@ -1,37 +0,0 @@
from pandas import DataFrame
from model_manager.utils.filters.conditional_filters import (
filter_empty_data,
filter_specific_variables_null_values,
)
def test_filter_specific_variables_null_values():
assert (
filter_specific_variables_null_values(
DataFrame({'variable': ['variable1', 'variable2'], 'value': [1, 2]}),
config={'variables': ['variable2']},
)
is False
)
def test_filter_specific_variables_null_values_with_null_values():
assert (
filter_specific_variables_null_values(
DataFrame({'variable': ['variable1', 'variable2'], 'value': [1, None]}),
config={'variables': ['variable2']},
)
is True
)
def test_filter_empty_data():
assert filter_empty_data(DataFrame(), {}) is True
def test_filter_empty_data_with_data():
assert (
filter_empty_data(DataFrame({'variable': ['variable1', 'variable2'], 'value': [1, 2]}), {})
is False
)

View File

@@ -1,23 +0,0 @@
from pandas import DataFrame
from model_manager.utils.filters.mlflow_filters import api_error_filter, nan_values_filter
def test_api_error_filter_invalid_response():
assert api_error_filter(None, {})
def test_api_error_filter_valid_response_fail():
assert api_error_filter({'success': False}, {})
def test_api_error_filter_valid_response_success():
assert not api_error_filter({'success': True}, {})
def test_nan_values_filter_all_nan_values():
assert nan_values_filter(DataFrame({'variable': [None, None]}), {})
def test_nan_values_filter_no_nan_values():
assert not nan_values_filter(DataFrame({'variable': [1, 2]}), {})