SIENTIAPDE-1243: Refactor and enhance model manager activities and workflows

This commit includes several changes:

- Reorganized imports and class inheritance in activities.py, gates.py and mlflow.py for better readability and maintainability.
- Improved error handling and logging in gates.py and mlflow.py.
- Added input validation and filtering in gates.py to ensure data quality.
- Enhanced prediction formatting and storage policy management in gates.py.
- Updated metrics.py to use consistent naming conventions and labels.
- Refactored connectors_config.py to use type hints and improve code clarity.
- Updated conditional and MLFlow filters for better data quality checks.
- Improved model repository logic for retraining and updating models.
- Enhanced worker.py to include SDK metrics and improved error handling.
- Refactored workflows for better modularity and error handling.
- Updated tests to reflect the changes and improve test coverage.
This commit is contained in:
Bruno Domingues
2025-10-01 17:28:57 -03:00
parent b102f79087
commit dfc190c818
24 changed files with 1482 additions and 1399 deletions

View File

@@ -1,53 +1,42 @@
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.temporal.activities.base import BaseActivity
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 sientia_do.formatters import create_sample_dict
from model_manager.utils.filters.mlflow_filters import nan_values_filter, api_error_filter
from model_manager import metrics
from model_manager.utils.filters.conditional_filters import (
filter_empty_data,
filter_specific_variables_null_values
filter_specific_variables_null_values,
)
from pandas import DataFrame
from model_manager import metrics
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
}
'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
},
'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
}
'path_confidence': {'STOP': -1, 'CONTINUE': 18, 'REPEAT': -1},
}
@@ -82,10 +71,9 @@ class Gates(BaseActivity):
Raises:
Exception: If BaseActivity initialization fails
"""
BaseActivity.__init__(
self, logger, notification_handler, set_error_counter=True)
BaseActivity.__init__(self, logger, notification_handler, set_error_counter=True)
@activity.defn(name="input_gate")
@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.
@@ -121,7 +109,7 @@ class Gates(BaseActivity):
"""
metadata = input_data['metadata']
self.info("Performing input gate...", metadata)
self.info('Performing input gate...', metadata)
filters = input_data['filters']
data = DataFrame(input_data['data'])
@@ -129,40 +117,42 @@ class Gates(BaseActivity):
filter_output = []
self.debug(f"Input data: {data.head(5).to_string()}", metadata)
self.debug(f"Filters: {filters}", metadata)
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)
self.error(f'Filter {fil} not found', metadata)
continue
try:
if input_filter_functions[fil](data, config['config']):
self.debug(
f"Data not passed the input filter {fil}:{config}", metadata)
self.debug(f'Data not passed the input filter {fil}:{config}', metadata)
filter_output.append(config['policy'])
except Exception as e:
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",
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
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], \
"Input data with bad quality"
self.info(f'Input gate result: {path_flag}', metadata)
return (
path_flag,
input_filter_functions['path_confidence'][path_flag],
'Input data with bad quality',
)
self.info("Nothing was filtered by the input gate", metadata)
return None, 0, ""
self.info('Nothing was filtered by the input gate', metadata)
return None, 0, ''
@activity.defn(name="mlflow_response_gate")
@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.
@@ -197,7 +187,7 @@ class Gates(BaseActivity):
Exception: If response validation fails or configuration is invalid
"""
metadata = input_data['metadata']
self.info("Performing mlflow response gate...", metadata)
self.info('Performing mlflow response gate...', metadata)
filters = input_data['filters']
data = input_data['data']
@@ -206,14 +196,13 @@ class Gates(BaseActivity):
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)
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)
self.error(f'Filter {fil} not found', metadata)
continue
try:
if mlflow_response_filter_functions[fil](data, config):
@@ -221,34 +210,36 @@ class Gates(BaseActivity):
comments.append(data['content']['message'])
self.send_notification(
metadata=metadata,
notification_id=f"{gate_type.upper()}_GATE_RESPONSE_FILTER__{fil}",
notification_id=f'{gate_type.upper()}_GATE_RESPONSE_FILTER__{fil}',
message=data['content']['message'],
block="mlflow_gate",
block='mlflow_gate',
level=NotificationLevel.ERROR,
attachment_content=data['content']['traceback']
attachment_content=data['content']['traceback'],
)
except Exception as e:
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",
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
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], \
", ".join(comments)
self.info(f'Mlflow response gate result: {path_flag}', metadata)
return (
path_flag,
mlflow_response_filter_functions['path_confidence'][path_flag],
', '.join(comments),
)
self.info("Nothing was filtered by the mlflow response gate", metadata)
return None, 0, ""
self.info('Nothing was filtered by the mlflow response gate', metadata)
return None, 0, ''
@activity.defn(name="mlflow_content_gate")
@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.
@@ -283,7 +274,7 @@ class Gates(BaseActivity):
Exception: If content validation fails or configuration is invalid
"""
metadata = input_data['metadata']
self.info("Performing mlflow content gate...", metadata)
self.info('Performing mlflow content gate...', metadata)
filters = input_data['filters']
data = DataFrame(input_data['data'])
@@ -292,8 +283,8 @@ class Gates(BaseActivity):
filter_output = []
self.debug(f"Input data:\n {data.head(5).to_string()}", metadata)
self.debug(f"Filters: \n {create_sample_dict(filters)}", metadata)
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:
@@ -303,36 +294,38 @@ class Gates(BaseActivity):
filter_output.append(config['policy'])
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",
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()
attachment_content=data.to_string(),
)
except Exception as e:
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",
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
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], \
"Transformed data not passed the content filter"
self.info(f'Mlflow content gate result: {path_flag}', metadata)
return (
path_flag,
mlflow_content_filter_functions['path_confidence'][path_flag],
'Transformed data not passed the content filter',
)
self.info("Nothing was filtered by the mlflow content gate", metadata)
return None, 0, ""
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]:
def get_prediction_store_policy(
self, prediction_store_policy: str, metadata: dict[str, Any]
) -> tuple[str, int]:
"""
Parse and validate prediction store policy configuration.
@@ -358,7 +351,9 @@ class Gates(BaseActivity):
if len(policy_elements) < 2:
self.error(
f"Invalid prediction store policy: {prediction_store_policy}, using default policy", metadata)
f'Invalid prediction store policy: {prediction_store_policy}, using default policy',
metadata,
)
return 'lts', 1
policy_type = policy_elements[0]
@@ -366,14 +361,20 @@ class Gates(BaseActivity):
# 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:
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)
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")
@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.
@@ -400,7 +401,7 @@ class Gates(BaseActivity):
"""
metadata = input_data['metadata']
prediction_store_policy = input_data['prediction_store_policy']
self.info("Formatting prediction...", metadata)
self.info('Formatting prediction...', metadata)
data = DataFrame(input_data['data'])
@@ -408,48 +409,45 @@ class Gates(BaseActivity):
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)
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)
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)
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)
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)
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}")
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['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)
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")
@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.
@@ -477,22 +475,24 @@ class Gates(BaseActivity):
"""
metadata = input_data['metadata']
self.debug("Formatting default prediction...", 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']]
})
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)
self.info(f'Default prediction formatted: {data.size} rows', metadata)
return data.to_dict()
@activity.defn(name="get_last_timestamp")
@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.
@@ -517,24 +517,22 @@ class Gates(BaseActivity):
"""
metadata = input_data['metadata']
self.info("Getting last timestamp...", metadata)
self.info('Getting last timestamp...', metadata)
data = DataFrame(input_data['data'])
self.debug(f"Input data: {data.head(5).to_string()}", metadata)
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())
max_timestamp = max(data['timestamp'].values.tolist())
self.info(
f"Last timestamp: {max_timestamp}", metadata)
self.info(f'Last timestamp: {max_timestamp}', metadata)
return max_timestamp
@activity.defn(name="write_metrics")
@activity.defn(name='write_metrics')
async def write_metrics(self, input_data: dict[str, Any]):
"""
Write prediction performance metrics to Prometheus monitoring system.
@@ -562,26 +560,24 @@ class Gates(BaseActivity):
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)
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']
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']
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']
pipeline_name=metadata['workflow_name'],
).observe(response_time)
self.info(
f"Metrics written for model {metadata['model_name']}", metadata)
self.info(f'Metrics written for model {metadata["model_name"]}', metadata)