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

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

Browse Source 
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
Download Patch File Download Diff File Expand all files Collapse all files

33
model_manager/activities/activities.py
View File

@@ -1,12 +1,14 @@
from temporalio import activity, workflow
from temporalio import workflow
with workflow.unsafe.imports_passed_through():
from sientia_do.temporal.activities.postgres import Postgres
from typing import Any
from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler
from sientia_do.observability.logger import Logger
from model_manager.activities.mlflow import MLFlow
from sientia_do.temporal.activities.postgres import Postgres
from model_manager.activities.gates import Gates
from typing import Any
from model_manager.activities.mlflow import MLFlow
class Activities(Postgres, MLFlow, Gates):
@@ -29,11 +31,13 @@ class Activities(Postgres, MLFlow, Gates):
notification_handler (NotificationHandler): Notification management instance
"""
def __init__(self,
def __init__(
self,
postgres_config: dict[str, Any],
mlflow_config: dict[str, Any],
logger: Logger,
notification_handler: NotificationHandler):
notification_handler: NotificationHandler,
):
"""
Initialize the Activities orchestrator with all required configurations.
@@ -52,7 +56,9 @@ class Activities(Postgres, MLFlow, Gates):
Exception: If any parent class initialization fails
"""
# Initialize parent classes
Postgres.__init__(self, host=postgres_config['host'],
Postgres.__init__(
self,
host=postgres_config['host'],
port=postgres_config['port'],
user=postgres_config['user'],
password=postgres_config['password'],
@@ -60,17 +66,20 @@ class Activities(Postgres, MLFlow, Gates):
min_connections=postgres_config['min_connections'],
max_connections=postgres_config['max_connections'],
logger=logger,
notification_handler=notification_handler)
notification_handler=notification_handler,
)
MLFlow.__init__(self, mlflow_host=mlflow_config['host'],
MLFlow.__init__(
self,
mlflow_host=mlflow_config['host'],
mlflow_port=mlflow_config['port'],
mlflow_username=mlflow_config['username'],
mlflow_password=mlflow_config['password'],
logger=logger,
notification_handler=notification_handler)
notification_handler=notification_handler,
)
Gates.__init__(self, logger=logger,
notification_handler=notification_handler)
Gates.__init__(self, logger=logger, notification_handler=notification_handler)
async def shutdown(self):
"""

246
model_manager/activities/gates.py
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({
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']]
})
'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)

97
model_manager/activities/mlflow.py
View File

@@ -1,20 +1,19 @@
from temporalio import activity, workflow
with workflow.unsafe.imports_passed_through():
from datetime import datetime
from pandas import Timestamp, to_datetime
from sientia_do.temporal.constants import DATETIME_FORMAT, DATETIME_FORMAT_WITH_TZ
from sientia_do.temporal.activities.base import BaseActivity
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.formatters import create_sample_dict
from sientia_do.temporal.activities.base import BaseActivity
from sientia_do.temporal.constants import DATETIME_FORMAT, DATETIME_FORMAT_WITH_TZ
from model_manager.utils.repository.model_repository import MLFlowRepository
from typing import Any
import numpy as np
from pandas import DataFrame
import traceback
class MLFlow(BaseActivity):
@@ -36,8 +35,15 @@ class MLFlow(BaseActivity):
model_monitoring_repository (MLFlowRepository): Repository for MLFlow operations
"""
def __init__(self, mlflow_host: str, mlflow_port: int, mlflow_username: str,
mlflow_password: str, logger: Logger, notification_handler: NotificationHandler):
def __init__(
self,
mlflow_host: str,
mlflow_port: int,
mlflow_username: str,
mlflow_password: str,
logger: Logger,
notification_handler: NotificationHandler,
):
"""
Initialize MLFlow activities with server configuration.
@@ -52,18 +58,17 @@ class MLFlow(BaseActivity):
Raises:
Exception: If MLFlowRepository initialization fails
"""
BaseActivity.__init__(
self, logger, notification_handler, set_error_counter=True)
BaseActivity.__init__(self, logger, notification_handler, set_error_counter=True)
self.mlflow_host = mlflow_host
self.mlflow_port = mlflow_port
self.mlflow_username = mlflow_username
self.mlflow_password = mlflow_password
self.model_monitoring_repository = MLFlowRepository(
f"{mlflow_host}:{mlflow_port}", mlflow_username, mlflow_password, logger
f'{mlflow_host}:{mlflow_port}', mlflow_username, mlflow_password, logger
)
@activity.defn(name="request_transform")
@activity.defn(name='request_transform')
async def request_transform(self, input_data: dict[str, Any]) -> dict[str, Any]:
"""
Transform input data using MLFlow models.
@@ -99,7 +104,7 @@ class MLFlow(BaseActivity):
model_name = input_data['model_name']
model_config = input_data.get('model_config', {})
self.debug("Raw input data:", metadata)
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
@@ -108,14 +113,12 @@ class MLFlow(BaseActivity):
)
# Pivot data for model input format
data = data.pivot(
index='timestamp', columns='variable',
values='value')
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('Processed input data:', metadata)
self.debug(data.head(5).to_string(), metadata)
# Request transformation from MLFlow model
@@ -124,16 +127,20 @@ class MLFlow(BaseActivity):
)
self.debug(
f"Transform raw response data: \n {create_sample_dict(response_data, max_items=5, max_depth=5)}", metadata)
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)
f'Transform response data: \n {create_sample_dict(response_data, max_items=5, max_depth=5)}',
metadata,
)
self.info("Data transformed successfully", metadata)
self.info('Data transformed successfully', metadata)
return response_data
@activity.defn(name="request_predict")
@activity.defn(name='request_predict')
async def request_predict(self, input_data: dict[str, Any]) -> dict[str, Any]:
"""
Execute predictions using MLFlow models.
@@ -169,14 +176,15 @@ class MLFlow(BaseActivity):
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)
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)
data['timestamp'], format=DATETIME_FORMAT_WITH_TZ
).dt.strftime(DATETIME_FORMAT)
# Request prediction from MLFlow model
response_data = self.model_monitoring_repository.predict(
@@ -184,13 +192,15 @@ class MLFlow(BaseActivity):
)
self.debug(
f"Prediction response data: \n {create_sample_dict(response_data, max_items=5, max_depth=5)}", metadata)
f'Prediction response data: \n {create_sample_dict(response_data, max_items=5, max_depth=5)}',
metadata,
)
self.info("Data predicted successfully", metadata)
self.info('Data predicted successfully', metadata)
return response_data
@activity.defn(name="retrain_model")
@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.
@@ -234,8 +244,7 @@ class MLFlow(BaseActivity):
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 = data.pivot(index='timestamp', columns='variable', values='value')
data.sort_index(inplace=True)
data.reset_index(inplace=True)
@@ -244,15 +253,10 @@ class MLFlow(BaseActivity):
try:
retrain_output, experiment = self.model_monitoring_repository.retrain_model(
data=data,
model_name=model_name
data=data, model_name=model_name
)
return {
'status': retrain_output,
'timestamp': timestamp,
'experiment': experiment
}
return {'status': retrain_output, 'timestamp': timestamp, 'experiment': experiment}
except Exception as e:
trace = traceback.format_exc()
self.send_notification(
@@ -261,12 +265,12 @@ class MLFlow(BaseActivity):
message=f'Error retraining model {model_name}: {e}',
block='retrain_model',
level=NotificationLevel.ERROR,
attachment_content=trace
attachment_content=trace,
)
self.error(trace, metadata=metadata)
raise e
@activity.defn(name="update_production_model")
@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.
@@ -311,12 +315,12 @@ class MLFlow(BaseActivity):
status = input_data['status']
self.info(
f'Updating production model {model_name} from experiment {experiment}...', metadata)
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
experiment=experiment, model_name=model_name
)
report = DataFrame([response])
@@ -325,8 +329,7 @@ class MLFlow(BaseActivity):
report['timestamp'] = timestamp
report['status'] = status
self.info(
f'Production model {model_name} updated successfully', metadata)
self.info(f'Production model {model_name} updated successfully', metadata)
return report.to_dict()
except Exception as e:
@@ -337,7 +340,7 @@ class MLFlow(BaseActivity):
message=f'Error updating production model {model_name}: {e}',
block='update_production_model',
level=NotificationLevel.ERROR,
attachment_content=trace
attachment_content=trace,
)
self.error(trace, metadata=metadata)
raise e

24
model_manager/metrics.py
View File

@@ -22,36 +22,36 @@ Metric Labels:
- pipeline_name: Name of the prediction pipeline
"""
from prometheus_client import Gauge, Counter, Histogram
from prometheus_client import Counter, Gauge, Histogram
# Application health metric
APP_UP = Gauge(
"app_up",
"Indicates if the application is running (1) or shutting down (0)",
["pod_id"],
'app_up',
'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"]
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",
'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",
'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",
'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]
buckets=[0.01, 0.05, 0.1, 0.2, 0.5, 1.0, 2.0, 5.0, 10.0],
)

15
model_manager/utils/connectors_config.py
View File

@@ -1,9 +1,8 @@
from os import getenv
import json
from typing import Dict, Any
from typing import Any
def build_postgres_config() -> Dict[str, Any]:
def build_postgres_config() -> dict[str, Any]:
"""
Build PostgreSQL database configuration from environment variables.
@@ -30,11 +29,11 @@ def build_postgres_config() -> Dict[str, Any]:
'password': getenv('POSTGRES_PASSWORD', 'sientia'),
'dbname': getenv('POSTGRES_DBNAME', 'sientia'),
'min_connections': int(getenv('POSTGRES_MIN_CONNECTIONS', '5')),
'max_connections': int(getenv('POSTGRES_MAX_CONNECTIONS', '20'))
'max_connections': int(getenv('POSTGRES_MAX_CONNECTIONS', '20')),
}
def build_mlflow_config() -> Dict[str, Any]:
def build_mlflow_config() -> dict[str, Any]:
"""
Build MLFlow server configuration from environment variables.
@@ -55,11 +54,11 @@ def build_mlflow_config() -> Dict[str, Any]:
'host': getenv('MLFLOW_HOST', 'http://localhost'),
'port': int(getenv('MLFLOW_PORT', '5080')),
'username': getenv('MLFLOW_USERNAME', 'aignosi'),
'password': getenv('MLFLOW_PASSWORD', 'aignosi')
'password': getenv('MLFLOW_PASSWORD', 'aignosi'),
}
def build_mongodb_config() -> Dict[str, Any]:
def build_mongodb_config() -> dict[str, Any]:
"""
Build MongoDB configuration from environment variables.
@@ -86,5 +85,5 @@ def build_mongodb_config() -> Dict[str, Any]:
return {
'connection_string': connection_string,
'database_name': getenv('MONGODB_DATABASE_NAME', 'sientia'),
'ttl_index_seconds': int(getenv('MONGODB_TTL_INDEX_HOURS', '1')) * 3600
'ttl_index_seconds': int(getenv('MONGODB_TTL_INDEX_HOURS', '1')) * 3600,
}

3
model_manager/utils/filters/conditional_filters.py
View File

@@ -20,8 +20,7 @@ def filter_specific_variables_null_values(data: DataFrame, config: dict) -> bool
False if none of the specified variables contain null values.
"""
return not data[
data['variable'].isin(config['variables']) & data['value'].isna()].empty
return not data[data['variable'].isin(config['variables']) & data['value'].isna()].empty
def filter_empty_data(data: DataFrame, _config: dict) -> bool:

7
model_manager/utils/filters/mlflow_filters.py
View File

@@ -52,8 +52,11 @@ def nan_values_filter(predictions: DataFrame, _config: dict) -> bool:
bool: True if data should be filtered (too many NaN values), False otherwise
"""
data = predictions.replace({None: np.nan}).drop(
columns=['timestamp'], errors='ignore').infer_objects()
data = (
predictions.replace({None: np.nan})
.drop(columns=['timestamp'], errors='ignore')
.infer_objects()
)
if data.isna().all().all():
return True

175
model_manager/utils/repository/model_repository.py
View File

@@ -10,22 +10,23 @@ requests using the Model Monitoring API functions.
By Monitoring we mean the evaluation of the performance of models, the generation of reports.
"""
from datetime import datetime
import traceback
import pandas as pd
import mlflow
from datetime import datetime
from os import makedirs, path, remove
import mlflow
import pandas as pd
from sientia.ModelServing import ModelServing
from sientia_do.temporal.constants import DATETIME_FORMAT_WITH_TZ
from sientia_do.observability.logger import Logger
from sientia_do.temporal.constants import DATETIME_FORMAT_WITH_TZ
class MLFlowRepository():
class MLFlowRepository:
def __init__(self, host, username, password, logger: Logger):
self.model_serving = ModelServing(tracking_uri=host,
username=username, password=password,
logger=logger)
self.model_serving = ModelServing(
tracking_uri=host, username=username, password=password, logger=logger
)
self.logger = logger
def detect_and_parse_datetime_index(self, data: pd.DataFrame, metadata: dict) -> pd.DataFrame:
@@ -40,33 +41,32 @@ class MLFlowRepository():
# Get type of first element of index
index_type = type(index[0])
self.logger.custom_info(f"Index type: {index_type}", metadata)
self.logger.custom_info(f'Index type: {index_type}', metadata)
message = f"Index must be all timestamp like column. Valid formats are: pandas Timestamp, datetime, string in format {DATETIME_FORMAT_WITH_TZ}"
message = f'Index must be all timestamp like column. Valid formats are: pandas Timestamp, datetime, string in format {DATETIME_FORMAT_WITH_TZ}'
# Check if all in index are of the same type
if not all(isinstance(i, index_type) for i in index):
raise ValueError(
f"{message}")
raise ValueError(f'{message}')
# Check type and converts to DATETIME_FORMAT_WITH_TZ
if index_type == str:
if index_type is str:
# Validate format of string and return error if not valid
try:
pd.to_datetime(data.index, format=DATETIME_FORMAT_WITH_TZ)
except ValueError:
raise ValueError(
f"{message}")
except ValueError as e:
raise ValueError(f'{message}') from e
elif index_type == datetime or index_type == pd.Timestamp:
data.index = data.index.strftime(DATETIME_FORMAT_WITH_TZ)
else:
raise ValueError(
f"{message}")
raise ValueError(f'{message}')
return data
def transform(self, model_name: str, data: pd.DataFrame, model_config: dict, metadata: dict) -> dict:
def transform(
self, model_name: str, data: pd.DataFrame, model_config: dict, metadata: dict
) -> dict:
"""
Transform data using a model.
@@ -81,41 +81,42 @@ class MLFlowRepository():
try:
self.logger.custom_debug(
f"Data received for model transformation: {data.to_csv()}", metadata)
f'Data received for model transformation: {data.to_csv()}', metadata
)
model_retention = model_config.get('retention_minutes', 0)
flavor = model_config.get('transform_flavor', 'sklearn')
compressed = model_config.get('is_compressed', False)
retention_target = model_config.get('retention_target', 'model')
transform_keyword = model_config.get(
'transform_function_keyword', 'predict')
transform_keyword = model_config.get('transform_function_keyword', 'predict')
transformed_data = self.model_serving.get_cached_transform(
model_name, data, model_retention, flavor,
compressed, retention_target, transform_keyword
model_name,
data,
model_retention,
flavor,
compressed,
retention_target,
transform_keyword,
)
self.logger.custom_debug(
f"Data received from model transformation: {transformed_data.to_csv()}", metadata)
f'Data received from model transformation: {transformed_data.to_csv()}', metadata
)
transformed_data = self.detect_and_parse_datetime_index(
transformed_data, metadata)
transformed_data = self.detect_and_parse_datetime_index(transformed_data, metadata)
return {
'success': True,
'content': transformed_data.to_dict()
}
return {'success': True, 'content': transformed_data.to_dict()}
except Exception as e:
except Exception as e: # noqa: BLE001
return {
'success': False,
'content': {
'message': str(e),
'traceback': traceback.format_exc()
}
'content': {'message': str(e), 'traceback': traceback.format_exc()},
}
def predict(self, model_name: str, data: pd.DataFrame, model_config: dict, metadata: dict) -> dict:
def predict(
self, model_name: str, data: pd.DataFrame, model_config: dict, metadata: dict
) -> dict:
"""
Predict data using a model.
@@ -137,31 +138,26 @@ class MLFlowRepository():
start_time = datetime.now()
self.logger.custom_debug(
f"Data received for model prediction: {data.to_csv()}", metadata)
f'Data received for model prediction: {data.to_csv()}', metadata
)
data = self.model_serving.get_cached_predict(
model_name, data, model_retention, flavor,
compressed, retention_target
model_name, data, model_retention, flavor, compressed, retention_target
)
end_time = datetime.now()
data = pd.DataFrame(data, columns=['prediction'])
self.logger.custom_debug(
f"Data received from model prediction: {data.to_csv()}", metadata)
f'Data received from model prediction: {data.to_csv()}', metadata
)
data.index = input_index
data['response_time'] = (end_time - start_time).total_seconds()
return {
'success': True,
'content': data.to_dict()
}
return {'success': True, 'content': data.to_dict()}
except Exception as e:
except Exception as e: # noqa: BLE001
return {
'success': False,
'content': {
'message': str(e),
'traceback': traceback.format_exc()
}
'content': {'message': str(e), 'traceback': traceback.format_exc()},
}
def get_experiment_by_run_id(self, run_id: str) -> dict:
@@ -190,10 +186,9 @@ class MLFlowRepository():
Returns:
str: The next run name in format 'model_name-run_number'
"""
runs = mlflow.search_runs(
experiment_names=[model_name], order_by=["start_time desc"])
runs = mlflow.search_runs(experiment_names=[model_name], order_by=['start_time desc'])
next_run_number = len(runs) + 1
return f"{model_name}-{next_run_number}"
return f'{model_name}-{next_run_number}'
def create_model_experiment(self, model_name: str, data: pd.DataFrame) -> tuple:
"""
@@ -217,14 +212,10 @@ class MLFlowRepository():
- experiment: MLFlow experiment name
"""
# load predictor model
predictor_uri = f"models:/{model_name}/production"
predictor_uri = f'models:/{model_name}/production'
# load transform model
latest_production_id = self.model_serving.get_model_run_id(
model_name, stage="Production"
)
transform_uri = self.model_serving.get_model_uri(
latest_production_id, prediction=False
)
latest_production_id = self.model_serving.get_model_run_id(model_name, stage='Production')
transform_uri = self.model_serving.get_model_uri(latest_production_id, prediction=False)
# load
data_model = mlflow.sklearn.load_model(transform_uri)
prediction_model = mlflow.sklearn.load_model(predictor_uri)
@@ -233,20 +224,16 @@ class MLFlowRepository():
target_name = data_model.target_variable
y = data[target_name]
treated_data = pd.merge(
treated_data, y, left_index=True, right_index=True)
treated_data = pd.merge(treated_data, y, left_index=True, right_index=True)
prediction_model = prediction_model.fit(treated_data)
experiment = self.get_experiment_by_run_id(latest_production_id)
mlflow.set_experiment(experiment)
return prediction_model, data_model, experiment
def perform_model_retrain(self,
prediction_model,
data_model,
experiment: str,
model_name: str,
data: pd.DataFrame):
def perform_model_retrain(
self, prediction_model, data_model, experiment: str, model_name: str, data: pd.DataFrame
):
"""
Execute the complete model retraining process in MLFlow.
@@ -271,7 +258,7 @@ class MLFlowRepository():
"""
pred_model_atributes = vars(prediction_model) # load class attributes
data_model_atributes = vars(data_model) # load class attributes
experiment_description = f"Retrain model {model_name} with new data"
experiment_description = f'Retrain model {model_name} with new data'
current_run_name = self.get_next_run_name(experiment)
with mlflow.start_run(
run_name=current_run_name, description=experiment_description
@@ -279,32 +266,32 @@ class MLFlowRepository():
# update transfomation model
# fixed parameters
for name_atribute, val_atribute in pred_model_atributes.items():
if name_atribute != "model":
if name_atribute != 'model':
mlflow.log_param(name_atribute, val_atribute)
# update prediction model
for name_atribute, val_atribute in data_model_atributes.items():
if name_atribute != "model":
if name_atribute != 'model':
mlflow.log_param(name_atribute, val_atribute)
# dynamic parameters, including model itself
mlflow.sklearn.log_model(data_model, "data_model")
mlflow.sklearn.log_model(data_model, 'data_model')
makedirs("temp", exist_ok=True)
makedirs('temp', exist_ok=True)
file_path = f"temp/raw_data_{model_name}.csv"
file_path = f'temp/raw_data_{model_name}.csv'
data.to_csv(file_path, index=True)
# log the data raw
mlflow.log_artifact(file_path)
# dynamic parameters, including model itself
mlflow.sklearn.log_model(prediction_model, "prediction_model")
mlflow.log_param("retrain", True)
mlflow.sklearn.log_model(prediction_model, 'prediction_model')
mlflow.log_param('retrain', True)
# clear temp file
if path.exists(file_path):
remove(file_path)
return "Model retrained successfully", experiment
return 'Model retrained successfully', experiment
def retrain_model(self, data: pd.DataFrame, model_name: str) -> tuple:
"""
@@ -325,10 +312,10 @@ class MLFlowRepository():
- status_message (str): Retraining operation status
- experiment_name (str): MLFlow experiment identifier
"""
prediction_model, data_model, experiment = self.create_model_experiment(
model_name, data)
prediction_model, data_model, experiment = self.create_model_experiment(model_name, data)
retrain_result = self.perform_model_retrain(
prediction_model, data_model, experiment, model_name, data)
prediction_model, data_model, experiment, model_name, data
)
return retrain_result
def get_experiment(self, experiment_name: str) -> int:
@@ -374,22 +361,21 @@ class MLFlowRepository():
"""
runs = mlflow.search_runs(
experiment_ids=[experiment_id],
filter_string="", # Sem filtro no MLflow ainda
output_format="pandas"
filter_string='', # Sem filtro no MLflow ainda
output_format='pandas',
)
if not isinstance(runs, pd.DataFrame):
raise ValueError('Runs is not a pandas DataFrame')
# Filtrar apenas as runs onde params.retrain == True
filtered_runs = runs[runs["params.retrain"] == 'True']
filtered_runs = runs[runs['params.retrain'] == 'True']
# Converter a coluna 'end_time' para datetime
filtered_runs['end_time'] = pd.to_datetime(filtered_runs['end_time'])
# Ordenar o DataFrame de forma descendente pela coluna 'end_time'
filtered_runs = filtered_runs.sort_values(
by='end_time', ascending=False)
filtered_runs = filtered_runs.sort_values(by='end_time', ascending=False)
# Pegar a última run_id do DataFrame filtrado e ordenado
latest_run_id = filtered_runs.iloc[0]['run_id']
@@ -423,16 +409,14 @@ class MLFlowRepository():
# Registrar o modelo
# Aqui estamos assumindo que você já tem um modelo salvo, caso contrário você precisará treiná-lo e salvá-lo primeiro.
# Se o modelo já está registrado, você pode usar o método register_model() ou pyfunc.load_model() para isso.
mlflow.register_model(
f"runs:/{run_id}/prediction_model", model_name)
mlflow.register_model(f'runs:/{run_id}/prediction_model', model_name)
# Colocar a versão do modelo em produção
# Depois de registrar o modelo, precisamos pegar a versão mais recente do modelo e movê-lo para o estágio 'Production'
client = mlflow.tracking.MlflowClient()
# Obter a versão mais recente registrada do modelo
model_versions = client.get_registered_model(
model_name).latest_versions
model_versions = client.get_registered_model(model_name).latest_versions
if not isinstance(model_versions, list):
raise ValueError('Model versions is not a list')
@@ -441,17 +425,10 @@ class MLFlowRepository():
# Mover a versão mais recente do modelo para o estágio de 'Production'
client.transition_model_version_stage(
name=model_name,
version=max_version,
stage="Production",
archive_existing_versions=True
name=model_name, version=max_version, stage='Production', archive_existing_versions=True
)
return {
'model_name': model_name,
'version': max_version,
'mlflow_run_id': run_id
}
return {'model_name': model_name, 'version': max_version, 'mlflow_run_id': run_id}
def update_production_model(self, experiment: str, model_name: str) -> dict:
"""

76
model_manager/worker/worker.py
View File

@@ -25,32 +25,35 @@ Environment Variables:
- PROJECT_NAME: Project name for notifications (default: model_manager)
"""
from temporalio import workflow, client
from temporalio.worker import Worker, PollerBehaviorAutoscaling
from temporalio.runtime import Runtime, TelemetryConfig, PrometheusConfig
from temporalio import client, workflow
from temporalio.runtime import PrometheusConfig, Runtime, TelemetryConfig
from temporalio.worker import PollerBehaviorAutoscaling, Worker
with workflow.unsafe.imports_passed_through():
import asyncio
import os
import sys
import asyncio
from model_manager.workflows.minimal_retrain import MinimalRetrain
from model_manager.workflows.predictions_batch import PredictionsBatch
from model_manager.workflows.sub_workflows.prediction_process import PredictionProcess
from model_manager.workflows.sub_workflows.format_and_export_prediction import \
FormatAndExportPrediction
from model_manager.activities.activities import Activities
from model_manager.utils.connectors_config import (
build_postgres_config,
build_mlflow_config,
build_mongodb_config
)
from prometheus_client import start_http_server
from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler
from sientia_do.observability.logger import get_logger
from model_manager import metrics
from prometheus_client import start_http_server
from model_manager.activities.activities import Activities
from model_manager.utils.connectors_config import (
build_mlflow_config,
build_mongodb_config,
build_postgres_config,
)
from model_manager.workflows.minimal_retrain import MinimalRetrain
from model_manager.workflows.predictions_batch import PredictionsBatch
from model_manager.workflows.sub_workflows.format_and_export_prediction import (
FormatAndExportPrediction,
)
from model_manager.workflows.sub_workflows.prediction_process import PredictionProcess
POD_ID = os.getenv('POD_ID')
SDK_METRICS_PORT = int(os.getenv('HTTP_SDK_METRICS_PORT', "9091"))
SDK_METRICS_PORT = int(os.getenv('HTTP_SDK_METRICS_PORT', '9091'))
async def main():
@@ -85,7 +88,7 @@ async def main():
logger.custom_info(f'Starting Worker with POD_ID: {POD_ID}', metadata)
logger.custom_info("Starting prometheus client...", metadata)
logger.custom_info('Starting prometheus client...', metadata)
start_prometheus_server()
logger.custom_info('Starting Notification Handler...', metadata)
@@ -95,7 +98,7 @@ async def main():
connection_string=mongo_config['connection_string'],
database=mongo_config['database_name'],
logger=logger,
project_name=os.getenv('PROJECT_NAME', 'model-manager')
project_name=os.getenv('PROJECT_NAME', 'model-manager'),
)
logger.custom_info('Starting Activities...', metadata)
@@ -104,16 +107,14 @@ async def main():
postgres_config=build_postgres_config(),
mlflow_config=build_mlflow_config(),
logger=logger,
notification_handler=notification_handler
notification_handler=notification_handler,
)
logger.custom_info(
f'Starting SDK Metrics Server on port {SDK_METRICS_PORT}...', metadata)
logger.custom_info(f'Starting SDK Metrics Server on port {SDK_METRICS_PORT}...', metadata)
new_runtime = Runtime(
telemetry=TelemetryConfig(
metrics=PrometheusConfig(
bind_address=f"0.0.0.0:{SDK_METRICS_PORT}")
metrics=PrometheusConfig(bind_address=f'0.0.0.0:{SDK_METRICS_PORT}')
)
)
@@ -122,7 +123,7 @@ async def main():
temporal_client = await client.Client.connect(
target_host=host,
namespace=os.getenv('TEMPORAL_NAMESPACE', 'model-manager'),
runtime=new_runtime
runtime=new_runtime,
)
logger.custom_info('Starting Workers...', metadata)
@@ -136,20 +137,19 @@ async def main():
activities.load_custom_query,
activities.retrain_model,
activities.update_production_model,
activities.export_data_to_postgres
activities.export_data_to_postgres,
],
max_concurrent_workflow_tasks=50,
max_concurrent_activities=50,
max_concurrent_local_activities=50,
max_cached_workflows=200,
workflow_task_poller_behavior=PollerBehaviorAutoscaling(),
activity_task_poller_behavior=PollerBehaviorAutoscaling()
activity_task_poller_behavior=PollerBehaviorAutoscaling(),
),
Worker(
temporal_client,
task_queue='predictions_batch-queue',
workflows=[PredictionsBatch, PredictionProcess,
FormatAndExportPrediction],
workflows=[PredictionsBatch, PredictionProcess, FormatAndExportPrediction],
activities=[
# MLFlow
activities.request_predict,
@@ -165,15 +165,15 @@ async def main():
activities.load_custom_query,
activities.repeat_last_prediction,
activities.export_data_to_postgres,
activities.write_metrics
activities.write_metrics,
],
max_concurrent_workflow_tasks=50,
max_concurrent_activities=50,
max_concurrent_local_activities=50,
max_cached_workflows=200,
workflow_task_poller_behavior=PollerBehaviorAutoscaling(),
activity_task_poller_behavior=PollerBehaviorAutoscaling()
)
activity_task_poller_behavior=PollerBehaviorAutoscaling(),
),
]
handlers = []
@@ -186,8 +186,8 @@ async def main():
# This will run the workers and wait for them to complete.
# If an exception occurs in any of the worker handlers, it will be propagated here.
await asyncio.gather(*handlers)
except BaseException as e: # NOSONAR
logger.custom_error(f"An unhandled exception occurred: {e}", metadata)
except BaseException as e: # noqa: BLE001
logger.custom_error(f'An unhandled exception occurred: {e}', metadata)
finally:
if notification_handler:
notification_handler.shutdown()
@@ -216,12 +216,12 @@ def start_prometheus_server():
SystemExit: If the metrics server fails to start
"""
try:
port = int(os.getenv("HTTP_METRICS_PORT", 9090))
port = int(os.getenv('HTTP_METRICS_PORT', 9090))
start_http_server(port)
print(f"Prometheus server started on port {port}.")
print(f'Prometheus server started on port {port}.')
metrics.APP_UP.labels(pod_id=POD_ID).set(1) # Mark app as UP
except Exception as e:
print(f"Failed to start Prometheus server: {e}")
except Exception as e: # noqa: BLE001
print(f'Failed to start Prometheus server: {e}')
os._exit(1)

34
model_manager/workflows/minimal_retrain.py
View File

@@ -1,14 +1,16 @@
from temporalio import workflow
with workflow.unsafe.imports_passed_through():
from model_manager.activities.activities import Activities
from typing import Any
from sientia_do.temporal.policies import retry_policy
from datetime import timedelta
from typing import Any
from sientia_do.temporal.policies import retry_policy
from model_manager.activities.activities import Activities
@workflow.defn(name="minimal_retrain")
class MinimalRetrain():
@workflow.defn(name='minimal_retrain')
class MinimalRetrain:
"""
Automated model retraining workflow for the Model Manager system.
@@ -63,7 +65,7 @@ class MinimalRetrain():
'schedule_name': input_data['schedule_name'],
'model_name': input_data['model_name'],
'model_id': input_data['model_id'],
'workflow_name': 'minimal_retrain'
'workflow_name': 'minimal_retrain',
}
}
@@ -74,21 +76,17 @@ class MinimalRetrain():
{
**metadata,
'query': input_data['query'],
'datetime_columns': input_data.get('datetime_columns', [])
'datetime_columns': input_data.get('datetime_columns', []),
},
retry_policy=retry_policy,
start_to_close_timeout=timedelta(seconds=60)
start_to_close_timeout=timedelta(seconds=60),
)
experiment_response = await workflow.execute_activity_method(
Activities.retrain_model,
{
**metadata,
'data': data,
'model_name': model_name
},
{**metadata, 'data': data, 'model_name': model_name},
retry_policy=retry_policy,
start_to_close_timeout=timedelta(seconds=60)
start_to_close_timeout=timedelta(seconds=60),
)
report = await workflow.execute_activity_method(
@@ -97,10 +95,10 @@ class MinimalRetrain():
**metadata,
'model_name': model_name,
'model_id': input_data['model_id'],
**experiment_response
**experiment_response,
},
retry_policy=retry_policy,
start_to_close_timeout=timedelta(seconds=60)
start_to_close_timeout=timedelta(seconds=60),
)
await workflow.execute_activity_method(
@@ -109,8 +107,8 @@ class MinimalRetrain():
**metadata,
'data': report,
'schema': input_data['schema'],
'table_name': input_data['table_name']
'table_name': input_data['table_name'],
},
retry_policy=retry_policy,
start_to_close_timeout=timedelta(seconds=60)
start_to_close_timeout=timedelta(seconds=60),
)

47
model_manager/workflows/predictions_batch.py
View File

@@ -1,14 +1,16 @@
from temporalio import workflow
with workflow.unsafe.imports_passed_through():
from model_manager.activities.activities import Activities
from typing import Any
from sientia_do.temporal.policies import retry_policy
from datetime import timedelta
from typing import Any
from sientia_do.temporal.policies import retry_policy
from model_manager.activities.activities import Activities
@workflow.defn(name="predictions_batch")
class PredictionsBatch():
@workflow.defn(name='predictions_batch')
class PredictionsBatch:
"""
Main batch prediction workflow for the Model Manager system.
@@ -74,7 +76,7 @@ class PredictionsBatch():
'schedule_name': input_data['schedule_name'],
'model_name': input_data['model_name'],
'model_id': input_data['model_id'],
'workflow_name': 'predictions_batch'
'workflow_name': 'predictions_batch',
}
}
@@ -84,10 +86,10 @@ class PredictionsBatch():
{
**metadata,
'query': input_data['query'],
'datetime_columns': input_data.get('datetime_columns', [])
'datetime_columns': input_data.get('datetime_columns', []),
},
retry_policy=retry_policy,
start_to_close_timeout=timedelta(seconds=300)
start_to_close_timeout=timedelta(seconds=300),
)
# Prepare input for prediction_process workflow
@@ -98,28 +100,17 @@ class PredictionsBatch():
'table_name': input_data['table_name'],
'model_id': input_data['model_id'],
'model_name': input_data['model_name'],
'input_filters': input_data.get('input_filters', {
'EMPTY_DATA': {
'POLICY': 'STOP'
}
}),
'mlflow_transform_filters': input_data.get('mlflow_transform_filters', {
'API_ERROR': {
'POLICY': 'STOP'
}
}),
'mlflow_predict_filters': input_data.get('mlflow_predict_filters', {
'API_ERROR': {
'POLICY': 'STOP'
}
}),
'input_filters': input_data.get('input_filters', {'EMPTY_DATA': {'POLICY': 'STOP'}}),
'mlflow_transform_filters': input_data.get(
'mlflow_transform_filters', {'API_ERROR': {'POLICY': 'STOP'}}
),
'mlflow_predict_filters': input_data.get(
'mlflow_predict_filters', {'API_ERROR': {'POLICY': 'STOP'}}
),
'model_config': input_data.get('model_config', {}),
'path_priority': input_data.get('path_priority', ['STOP', 'CONTINUE', 'REPEAT']),
'prediction_store_policy': input_data.get(
'prediction_store_policy', 'lts:1')
'prediction_store_policy': input_data.get('prediction_store_policy', 'lts:1'),
}
# Execute prediction process workflow
await workflow.execute_child_workflow(
'prediction_process', prediction_input)
await workflow.execute_child_workflow('prediction_process', prediction_input)

34
model_manager/workflows/sub_workflows/format_and_export_prediction.py
View File

@@ -1,15 +1,17 @@
from temporalio import workflow
with workflow.unsafe.imports_passed_through():
from model_manager.activities.activities import Activities
from typing import Any
from datetime import timedelta
from sientia_do.temporal.policies import retry_policy
from typing import Any
from sientia_do.temporal.constants import DATETIME_FORMAT_WITH_TZ
from sientia_do.temporal.policies import retry_policy
from model_manager.activities.activities import Activities
@workflow.defn(name="format_and_export_prediction")
class FormatAndExportPrediction():
@workflow.defn(name='format_and_export_prediction')
class FormatAndExportPrediction:
"""
Data formatting and export workflow for prediction results.
@@ -75,10 +77,10 @@ class FormatAndExportPrediction():
'timestamp': input_data['timestamp'],
'model_id': input_data['model_id'],
'prediction_confidence': prediction_confidence,
'prediction_store_policy': input_data['prediction_store_policy']
'prediction_store_policy': input_data['prediction_store_policy'],
},
retry_policy=retry_policy,
start_to_close_timeout=timedelta(seconds=60)
start_to_close_timeout=timedelta(seconds=60),
)
else:
@@ -90,10 +92,10 @@ class FormatAndExportPrediction():
'timestamp': input_data['timestamp'],
'model_id': input_data['model_id'],
'prediction_confidence': prediction_confidence,
'comment': input_data['comment']
'comment': input_data['comment'],
},
retry_policy=retry_policy,
start_to_close_timeout=timedelta(seconds=60)
start_to_close_timeout=timedelta(seconds=60),
)
# write to postgres
@@ -104,21 +106,15 @@ class FormatAndExportPrediction():
'schema': input_data['schema'],
'table_name': input_data['table_name'],
'data': prediction,
'timestamp_conversion': {
'column': 'timestamp',
'format': DATETIME_FORMAT_WITH_TZ
}
'timestamp_conversion': {'column': 'timestamp', 'format': DATETIME_FORMAT_WITH_TZ},
},
retry_policy=retry_policy,
start_to_close_timeout=timedelta(seconds=60)
start_to_close_timeout=timedelta(seconds=60),
)
await workflow.execute_activity_method(
Activities.write_metrics,
{
**metadata,
'prediction': prediction
},
{**metadata, 'prediction': prediction},
retry_policy=retry_policy,
start_to_close_timeout=timedelta(seconds=60)
start_to_close_timeout=timedelta(seconds=60),
)

55
model_manager/workflows/sub_workflows/prediction_process.py
View File

@@ -1,14 +1,16 @@
from temporalio import workflow
with workflow.unsafe.imports_passed_through():
from model_manager.activities.activities import Activities
from typing import Any
from sientia_do.temporal.policies import retry_policy
from datetime import timedelta
from typing import Any
from sientia_do.temporal.policies import retry_policy
from model_manager.activities.activities import Activities
@workflow.defn(name="prediction_process")
class PredictionProcess():
@workflow.defn(name='prediction_process')
class PredictionProcess:
"""
Core prediction processing workflow for the Model Manager system.
@@ -84,10 +86,7 @@ class PredictionProcess():
# Get last timestamp for incremental processing
last_timestamp = await workflow.execute_local_activity_method(
Activities.get_last_timestamp,
{
**metadata,
'data': data
},
{**metadata, 'data': data},
retry_policy=retry_policy,
start_to_close_timeout=timedelta(minutes=1),
)
@@ -97,7 +96,7 @@ class PredictionProcess():
**metadata,
'filters': input_data['input_filters'],
'data': data,
'path_priority': input_data['path_priority']
'path_priority': input_data['path_priority'],
}
path_flag, confidence, comment = await workflow.execute_local_activity_method(
@@ -116,12 +115,7 @@ class PredictionProcess():
# Request MLFlow model transformation
response_data = await workflow.execute_local_activity_method(
Activities.request_transform,
{
**metadata,
'data': data,
'model_name': model_name,
'model_config': model_config
},
{**metadata, 'data': data, 'model_name': model_name, 'model_config': model_config},
retry_policy=retry_policy,
start_to_close_timeout=timedelta(minutes=5),
)
@@ -134,7 +128,7 @@ class PredictionProcess():
'filters': input_data['mlflow_transform_filters'],
'data': response_data,
'type': 'transform',
'path_priority': input_data['path_priority']
'path_priority': input_data['path_priority'],
},
retry_policy=retry_policy,
start_to_close_timeout=timedelta(minutes=1),
@@ -155,7 +149,7 @@ class PredictionProcess():
'filters': input_data['mlflow_transform_filters'],
'data': transformed_data,
'type': 'transform',
'path_priority': input_data['path_priority']
'path_priority': input_data['path_priority'],
},
retry_policy=retry_policy,
start_to_close_timeout=timedelta(minutes=1),
@@ -172,7 +166,7 @@ class PredictionProcess():
**metadata,
'data': transformed_data,
'model_name': model_name,
'model_config': model_config
'model_config': model_config,
},
retry_policy=retry_policy,
start_to_close_timeout=timedelta(minutes=5),
@@ -186,7 +180,7 @@ class PredictionProcess():
'filters': input_data['mlflow_predict_filters'],
'data': response_data,
'type': 'predict',
'path_priority': input_data['path_priority']
'path_priority': input_data['path_priority'],
},
retry_policy=retry_policy,
start_to_close_timeout=timedelta(minutes=1),
@@ -213,12 +207,19 @@ class PredictionProcess():
'schema': input_data['schema'],
'table_name': input_data['table_name'],
'comment': comment,
'prediction_store_policy': input_data['prediction_store_policy']
}
'prediction_store_policy': input_data['prediction_store_policy'],
},
)
async def path_flag_handler(self, data: dict, path_flag: str, input_data: dict,
confidence: int, last_timestamp: str, comment: str) -> bool:
async def path_flag_handler(
self,
data: dict,
path_flag: str,
input_data: dict,
confidence: int,
last_timestamp: str,
comment: str,
) -> bool:
"""
Handle path decisions based on filter results and confidence levels.
@@ -264,7 +265,7 @@ class PredictionProcess():
'schema': schema,
'table_name': table_name,
'model': model_id,
'last_timestamp': last_timestamp
'last_timestamp': last_timestamp,
},
retry_policy=retry_policy,
start_to_close_timeout=timedelta(minutes=1),
@@ -286,8 +287,8 @@ class PredictionProcess():
'schema': schema,
'table_name': table_name,
'comment': comment,
'prediction_store_policy': input_data['prediction_store_policy']
}
'prediction_store_policy': input_data['prediction_store_policy'],
},
)
return True

37
tests/laborious/activities/test_activities.py
View File

@@ -1,16 +1,17 @@
from unittest.mock import ANY, MagicMock, patch
from pytest import mark
from unittest.mock import patch, MagicMock, ANY
from sientia_do.temporal.activities.postgres import Postgres
from model_manager.activities.activities import Activities
from model_manager.activities.mlflow import MLFlow
from model_manager.activities.gates import Gates
from model_manager.activities.mlflow import MLFlow
@patch('model_manager.activities.activities.Postgres.__init__')
@patch('model_manager.activities.activities.MLFlow.__init__')
@patch('model_manager.activities.activities.Gates.__init__')
def test___init__(mock_gates_init, mock_mlflow_init, mock_postgres_init):
postgres_config = {
'host': 'localhost',
'port': 5432,
@@ -18,15 +19,10 @@ def test___init__(mock_gates_init, mock_mlflow_init, mock_postgres_init):
'password': 'postgres',
'dbname': 'postgres',
'min_connections': 1,
'max_connections': 10
'max_connections': 10,
}
mlflow_config = {
'host': 'localhost',
'port': 5000,
'username': 'mlflow',
'password': 'mlflow'
}
mlflow_config = {'host': 'localhost', 'port': 5000, 'username': 'mlflow', 'password': 'mlflow'}
logger = MagicMock()
notification_handler = MagicMock()
@@ -35,7 +31,7 @@ def test___init__(mock_gates_init, mock_mlflow_init, mock_postgres_init):
postgres_config=postgres_config,
mlflow_config=mlflow_config,
logger=logger,
notification_handler=notification_handler
notification_handler=notification_handler,
)
assert isinstance(activities, Activities)
@@ -53,7 +49,7 @@ def test___init__(mock_gates_init, mock_mlflow_init, mock_postgres_init):
min_connections=postgres_config['min_connections'],
max_connections=postgres_config['max_connections'],
logger=logger,
notification_handler=notification_handler
notification_handler=notification_handler,
)
mock_mlflow_init.assert_called_once_with(
@@ -63,13 +59,11 @@ def test___init__(mock_gates_init, mock_mlflow_init, mock_postgres_init):
mlflow_username=mlflow_config['username'],
mlflow_password=mlflow_config['password'],
logger=logger,
notification_handler=notification_handler
notification_handler=notification_handler,
)
mock_gates_init.assert_called_once_with(
ANY,
logger=logger,
notification_handler=notification_handler
ANY, logger=logger, notification_handler=notification_handler
)
@@ -84,15 +78,10 @@ async def test_shutdown(_mock_mlflow_init, mock_postgres_init):
'password': 'postgres',
'dbname': 'postgres',
'min_connections': 1,
'max_connections': 10
'max_connections': 10,
}
mlflow_config = {
'host': 'localhost',
'port': 5000,
'username': 'mlflow',
'password': 'mlflow'
}
mlflow_config = {'host': 'localhost', 'port': 5000, 'username': 'mlflow', 'password': 'mlflow'}
logger = MagicMock()
notification_handler = MagicMock()
@@ -101,7 +90,7 @@ async def test_shutdown(_mock_mlflow_init, mock_postgres_init):
postgres_config=postgres_config,
mlflow_config=mlflow_config,
logger=logger,
notification_handler=notification_handler
notification_handler=notification_handler,
)
await activities.shutdown()

235
tests/laborious/activities/test_gates.py
View File

@@ -1,6 +1,8 @@
from unittest.mock import MagicMock, ANY, patch
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
@@ -20,11 +22,11 @@ def gates_activity():
metadata = {
"metadata": {
"model_id": "test_model",
"model_name": "test_model",
"workflow_name": "test_workflow",
"schema_name": "test_schedule",
'metadata': {
'model_id': 'test_model',
'model_name': 'test_model',
'workflow_name': 'test_workflow',
'schema_name': 'test_schedule',
},
}
@@ -34,20 +36,18 @@ async def test_input_gate_invalid_filter(gates_activity):
# Arrange
input_data = {
**metadata,
'filters': {
'INVALID_FILTER': {'POLICY': 'STOP'}
},
'filters': {'INVALID_FILTER': {'POLICY': 'STOP'}},
'data': {'value': [1, 2, 3]},
'path_priority': ['STOP', 'CONTINUE', 'REPEAT']
'path_priority': ['STOP', 'CONTINUE', 'REPEAT'],
}
# Act
result = await gates_activity.input_gate(input_data)
# Assert
assert result == (None, 0, "")
assert result == (None, 0, '')
gates_activity.error.assert_called_once_with(
"Filter INVALID_FILTER not found", metadata['metadata']
'Filter INVALID_FILTER not found', metadata['metadata']
)
@@ -57,28 +57,27 @@ async def test_input_gate_filter_exception(mock_input_filter_functions, gates_ac
# Arrange
mock_input_filter_functions.__contains__.return_value = True
mock_input_filter_functions.__getitem__.return_value = MagicMock(
side_effect=Exception("Test error"))
side_effect=Exception('Test error')
)
input_data = {
**metadata,
'filters': {
'EMPTY_DATA': {'policy': 'STOP', 'config': {}}
},
'filters': {'EMPTY_DATA': {'policy': 'STOP', 'config': {}}},
'data': {'value': []},
'path_priority': ['STOP', 'CONTINUE', 'REPEAT']
'path_priority': ['STOP', 'CONTINUE', 'REPEAT'],
}
# Act
result = await gates_activity.input_gate(input_data)
# Assert
assert result == (None, 0, "")
assert result == (None, 0, '')
gates_activity.send_notification.assert_called_once_with(
metadata=metadata['metadata'],
notification_id="INTPUT_GATE_ERROR__EMPTY_DATA",
notification_id='INTPUT_GATE_ERROR__EMPTY_DATA',
message="Error in filter EMPTY_DATA:{'policy': 'STOP', 'config': {}}: \n Test error",
block="input_gate",
block='input_gate',
level=NotificationLevel.ERROR,
attachment_content=ANY
attachment_content=ANY,
)
@@ -89,14 +88,14 @@ async def test_input_gate_no_filters(gates_activity):
**metadata,
'filters': {},
'data': {'value': [1, 2, 3]},
'path_priority': ['CONTINUE', 'STOP', 'REPEAT']
'path_priority': ['CONTINUE', 'STOP', 'REPEAT'],
}
# Act
result = await gates_activity.input_gate(input_data)
# Assert
assert result == (None, 0, "")
assert result == (None, 0, '')
gates_activity.debug.assert_called()
@@ -105,18 +104,16 @@ async def test_input_gate_with_filter(gates_activity):
# Arrange
input_data = {
**metadata,
'filters': {
'EMPTY_DATA': {'policy': 'STOP', 'config': {}}
},
'filters': {'EMPTY_DATA': {'policy': 'STOP', 'config': {}}},
'data': {'value': []},
'path_priority': ['STOP', 'CONTINUE', 'REPEAT']
'path_priority': ['STOP', 'CONTINUE', 'REPEAT'],
}
# Act
result = await gates_activity.input_gate(input_data)
# Assert
assert result == ('STOP', -1, "Input data with bad quality")
assert result == ('STOP', -1, 'Input data with bad quality')
gates_activity.debug.assert_called()
@@ -125,51 +122,49 @@ async def test_mlflow_response_gate_invalid_filter(gates_activity):
# Arrange
input_data = {
**metadata,
'filters': {
'INVALID_FILTER': {'POLICY': 'STOP'}
},
'filters': {'INVALID_FILTER': {'POLICY': 'STOP'}},
'data': {'content': {'message': 'success'}},
'type': 'test',
'path_priority': ['STOP', 'CONTINUE', 'REPEAT']
'path_priority': ['STOP', 'CONTINUE', 'REPEAT'],
}
# Act
result = await gates_activity.mlflow_response_gate(input_data)
# Assert
assert result == (None, 0, "")
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):
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"))
side_effect=Exception('Test error')
)
input_data = {
**metadata,
'filters': {
'INVALID_FILTER': {'POLICY': 'STOP'}
},
'filters': {'INVALID_FILTER': {'POLICY': 'STOP'}},
'data': {'content': {'message': 'success'}},
'type': 'test',
'path_priority': ['STOP', 'CONTINUE', 'REPEAT']
'path_priority': ['STOP', 'CONTINUE', 'REPEAT'],
}
# Act
result = await gates_activity.mlflow_response_gate(input_data)
# Assert
assert result == (None, 0, "")
assert result == (None, 0, '')
gates_activity.send_notification.assert_called_once_with(
metadata=metadata['metadata'],
notification_id="MLFLOW_GATE_RESPONSE_FILTER__INVALID_FILTER",
notification_id='MLFLOW_GATE_RESPONSE_FILTER__INVALID_FILTER',
message="Error in filter INVALID_FILTER:{'POLICY': 'STOP'}: \n Test error",
block="mlflow_gate",
block='mlflow_gate',
level=NotificationLevel.ERROR,
attachment_content=ANY
attachment_content=ANY,
)
@@ -181,14 +176,14 @@ async def test_mlflow_response_gate_no_filters(gates_activity):
'filters': {},
'data': {'content': {'message': 'success'}},
'type': 'test',
'path_priority': ['CONTINUE', 'STOP', 'REPEAT']
'path_priority': ['CONTINUE', 'STOP', 'REPEAT'],
}
# Act
result = await gates_activity.mlflow_response_gate(input_data)
# Assert
assert result == (None, 0, "")
assert result == (None, 0, '')
gates_activity.debug.assert_called()
@@ -197,25 +192,20 @@ async def test_mlflow_response_gate_with_filter(gates_activity):
# Arrange
input_data = {
**metadata,
'filters': {
'API_ERROR': {'policy': 'STOP'}
},
'filters': {'API_ERROR': {'policy': 'STOP'}},
'data': {
'success': False,
'content': {
'message': 'API error occurred',
'traceback': 'error trace'
}
'content': {'message': 'API error occurred', 'traceback': 'error trace'},
},
'type': 'test',
'path_priority': ['STOP', 'CONTINUE', 'REPEAT']
'path_priority': ['STOP', 'CONTINUE', 'REPEAT'],
}
# Act
result = await gates_activity.mlflow_response_gate(input_data)
# Assert
assert result == ('STOP', -1, "API error occurred")
assert result == ('STOP', -1, 'API error occurred')
gates_activity.debug.assert_called()
gates_activity.send_notification.assert_called()
@@ -225,58 +215,53 @@ async def test_mlflow_content_gate_invalid_filter(gates_activity):
# Arrange
input_data = {
**metadata,
'filters': {
'INVALID_FILTER': {'POLICY': 'STOP'}
},
'filters': {'INVALID_FILTER': {'POLICY': 'STOP'}},
'data': {'value': [1, 2, 3]},
'type': 'test',
'path_priority': ['STOP', 'CONTINUE', 'REPEAT']
'path_priority': ['STOP', 'CONTINUE', 'REPEAT'],
}
# Act
result = await gates_activity.mlflow_content_gate(input_data)
# Assert
assert result == (None, 0, "")
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):
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"))
side_effect=Exception('Test error')
)
input_data = {
**metadata,
'filters': {
'API_ERROR': {'POLICY': 'STOP'}
},
'filters': {'API_ERROR': {'POLICY': 'STOP'}},
'data': {
'success': False,
'content': {
'message': 'API error occurred',
'traceback': 'error trace'
}
'content': {'message': 'API error occurred', 'traceback': 'error trace'},
},
'type': 'test',
'path_priority': ['STOP', 'CONTINUE', 'REPEAT']
'path_priority': ['STOP', 'CONTINUE', 'REPEAT'],
}
# Act
result = await gates_activity.mlflow_content_gate(input_data)
# Assert
assert result == (None, 0, "")
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",
notification_id='MLFLOW_GATE_CONTENT_FILTER__API_ERROR',
message="Error in filter API_ERROR:{'POLICY': 'STOP'}: \n Test error",
block="mlflow_gate",
block='mlflow_gate',
level=NotificationLevel.ERROR,
attachment_content=ANY
attachment_content=ANY,
)
@@ -288,14 +273,14 @@ async def test_mlflow_content_gate_no_filters(gates_activity):
'filters': {},
'data': {'value': [1, 2, 3]},
'type': 'test',
'path_priority': ['CONTINUE', 'STOP', 'REPEAT']
'path_priority': ['CONTINUE', 'STOP', 'REPEAT'],
}
# Act
result = await gates_activity.mlflow_content_gate(input_data)
# Assert
assert result == (None, 0, "")
assert result == (None, 0, '')
gates_activity.debug.assert_called()
@@ -304,20 +289,17 @@ async def test_mlflow_content_gate_with_filter(gates_activity):
# Arrange
input_data = {
**metadata,
'filters': {
'NAN_VALUES': {'policy': 'STOP', 'config': {}}
},
'filters': {'NAN_VALUES': {'policy': 'STOP', 'config': {}}},
'data': {'value': [None, None, None]},
'type': 'test',
'path_priority': ['STOP', 'CONTINUE', 'REPEAT']
'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")
assert result == ('STOP', -1, 'Transformed data not passed the content filter')
gates_activity.debug.assert_called()
gates_activity.send_notification.assert_called()
@@ -328,7 +310,8 @@ def test_get_prediction_store_policy_invalid_policy(gates_activity):
# Act
policy_type, policy_value = gates_activity.get_prediction_store_policy(
prediction_store_policy, metadata)
prediction_store_policy, metadata
)
# Assert
assert policy_type == 'lts'
@@ -341,7 +324,8 @@ def test_get_prediction_store_policy_invalid_policy_value(gates_activity):
# Act
policy_type, policy_value = gates_activity.get_prediction_store_policy(
prediction_store_policy, metadata)
prediction_store_policy, metadata
)
# Assert
assert policy_type == 'lts'
@@ -354,7 +338,8 @@ def test_get_prediction_store_policy_valid_policy_type(gates_activity):
# Act
policy_type, policy_value = gates_activity.get_prediction_store_policy(
prediction_store_policy, metadata)
prediction_store_policy, metadata
)
# Assert
assert policy_type == 'lts'
@@ -367,7 +352,8 @@ def test_get_prediction_store_policy_valid_policy(gates_activity):
# Act
policy_type, policy_value = gates_activity.get_prediction_store_policy(
prediction_store_policy, metadata)
prediction_store_policy, metadata
)
# Assert
assert policy_type == 'erl'
@@ -380,16 +366,12 @@ async def test_format_prediction_no_timestamp(gates_activity):
input_data = {
**metadata,
'data': {
'prediction': {
'2023-05-26 11:12:27': 1
},
'response_time': {
'2023-05-26 11:12:27': 0.1
}
'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'
'prediction_store_policy': 'lts:1',
}
# Act
@@ -402,7 +384,7 @@ async def test_format_prediction_no_timestamp(gates_activity):
assert result['model_id'] == {0: 'test_model'}
assert result['prediction_confidence'] == {0: 0.9}
assert result['prediction_status'] == {0: 'Good'}
assert result['comments'] == {0: ""}
assert result['comments'] == {0: ''}
@mark.asyncio
@@ -420,11 +402,11 @@ async def test_format_prediction_with_timestamp_erl(gates_activity):
'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'
'prediction_store_policy': 'erl:2',
}
# Act
@@ -433,12 +415,11 @@ async def test_format_prediction_with_timestamp_erl(gates_activity):
# 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['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: ""}
assert result['comments'] == {0: '', 1: ''}
@mark.asyncio
@@ -456,11 +437,11 @@ async def test_format_prediction_with_timestamp_lts(gates_activity):
'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'
'prediction_store_policy': 'lts:2',
}
# Act
@@ -469,12 +450,11 @@ async def test_format_prediction_with_timestamp_lts(gates_activity):
# 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['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: ""}
assert result['comments'] == {0: '', 1: ''}
@mark.asyncio
@@ -482,22 +462,23 @@ async def test_format_prediction_with_timestamp_invalid_policy(gates_activity):
# Arrange
input_data = {
**metadata,
'data': {'prediction': [1, 2, 3],
'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']},
'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'
'prediction_store_policy': 'lts:2',
}
gates_activity.get_prediction_store_policy = MagicMock(
return_value=('invalid', 1))
gates_activity.get_prediction_store_policy = MagicMock(return_value=('invalid', 1))
try:
result = await gates_activity.format_prediction(input_data)
await gates_activity.format_prediction(input_data)
except ValueError as e:
assert str(e) == "Invalid policy type: invalid"
assert str(e) == 'Invalid policy type: invalid'
else:
assert False, "Expected ValueError"
raise AssertionError('Expected ValueError')
@mark.asyncio
@@ -508,7 +489,7 @@ async def test_format_default_prediction(gates_activity):
'timestamp': '2023-05-26 11:12:27',
'model_id': 'test_model',
'prediction_confidence': 0.1,
'comment': 'Test comment'
'comment': 'Test comment',
}
# Act
@@ -528,12 +509,7 @@ async def test_format_default_prediction(gates_activity):
@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']
}
}
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)
@@ -545,10 +521,7 @@ async def test_get_last_timestamp_with_data(gates_activity):
@mark.asyncio
async def test_get_last_timestamp_no_data(gates_activity):
# Arrange
input_data = {
'data': {},
**metadata
}
input_data = {'data': {}, **metadata}
# Act
result = await gates_activity.get_last_timestamp(input_data)
@@ -567,30 +540,28 @@ async def test_write_metrics(mock_metrics, gates_activity):
'prediction': {
'prediction': [1, 2, 3],
'prediction_confidence': [0.9, 0.8, 0.7],
'response_time': [0.1, 0.2, 0.3]
}
'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']
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
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']
pipeline_name=metadata['metadata']['workflow_name'],
)
mock_metrics.PREDICTION_RESPONSE_TIME_MONITOR.labels.return_value.observe.assert_called_once_with(
0.1

210
tests/laborious/activities/test_mlflow.py
View File

@@ -1,45 +1,42 @@
from datetime import datetime
from unittest.mock import ANY, MagicMock, patch
import numpy as np
from pandas import DataFrame, Timestamp
from pytest import fixture, mark, raises
from sientia_do.temporal.constants import DATETIME_FORMAT, DATETIME_FORMAT_WITH_TZ
from model_manager.activities.mlflow import MLFlow
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
@patch("model_manager.activities.mlflow.MLFlowRepository")
@patch('model_manager.activities.mlflow.MLFlowRepository')
def test___init__(mock_mlflow_repository):
mlflow = MLFlow(
mlflow_host="http://localhost",
mlflow_host='http://localhost',
mlflow_port=5000,
mlflow_username="admin",
mlflow_password="admin",
mlflow_username='admin',
mlflow_password='admin',
logger=MagicMock(),
notification_handler=MagicMock()
notification_handler=MagicMock(),
)
assert mlflow.mlflow_host == "http://localhost"
assert mlflow.mlflow_host == 'http://localhost'
assert mlflow.mlflow_port == 5000
assert mlflow.mlflow_username == "admin"
assert mlflow.mlflow_password == "admin"
assert mlflow.mlflow_username == 'admin'
assert mlflow.mlflow_password == 'admin'
mock_mlflow_repository.assert_called_once_with(
"http://localhost:5000", "admin", "admin", ANY
)
mock_mlflow_repository.assert_called_once_with('http://localhost:5000', 'admin', 'admin', ANY)
@fixture
@patch("model_manager.activities.mlflow.MLFlowRepository")
@patch('model_manager.activities.mlflow.MLFlowRepository')
def mlflow(mock_mlflow_repository):
mlflow = MLFlow(
mlflow_host="http://localhost:5000",
mlflow_host='http://localhost:5000',
mlflow_port=5000,
mlflow_username="admin",
mlflow_password="admin",
mlflow_username='admin',
mlflow_password='admin',
logger=MagicMock(),
notification_handler=MagicMock()
notification_handler=MagicMock(),
)
mlflow.send_notification = MagicMock()
@@ -48,44 +45,67 @@ def mlflow(mock_mlflow_repository):
metadata = {
"metadata": {
"model_id": "test_model",
"model_name": "test_model",
"workflow_name": "test_workflow",
"schema_name": "test_schedule",
'metadata': {
'model_id': 'test_model',
'model_name': 'test_model',
'workflow_name': 'test_workflow',
'schema_name': 'test_schedule',
},
}
@mark.asyncio
@patch("model_manager.activities.mlflow.DataFrame")
@patch("model_manager.activities.mlflow.max")
@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'}
{
'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': {}
'model_config': {},
}
# Mock the transform response
expected_response = {'prediction': [0.5, 0.6], 'timestamp': [
'2024-01-01', '2024-01-02']}
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
@@ -114,30 +134,25 @@ async def test_request_transform_success(mock_max, mock_dataframe, mlflow):
@mark.asyncio
@patch("model_manager.activities.mlflow.DataFrame")
@patch("model_manager.activities.mlflow.to_datetime")
@patch("model_manager.activities.mlflow.max")
@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"
'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
}
'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': {}
'model_config': {},
}
# Mock the predict response
@@ -148,9 +163,7 @@ async def test_request_predict(mock_max, mock_to_datetime, mock_dataframe, mlflo
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.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
)
@@ -158,9 +171,7 @@ async def test_request_predict(mock_max, mock_to_datetime, mock_dataframe, mlflo
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
)
mock_to_datetime.return_value.dt.strftime.assert_called_once_with(DATETIME_FORMAT)
# Verify the response
assert response_data == expected_response
@@ -174,28 +185,26 @@ async def test_request_predict(mock_max, mock_to_datetime, mock_dataframe, mlflo
@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]
'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')
'Model retrained successfully',
'test',
)
response = await mlflow.retrain_model({
**metadata,
'data': data,
'model_name': 'test_model'
})
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'
'status': 'Model retrained successfully',
'timestamp': 2,
'experiment': 'test',
}
@@ -206,20 +215,16 @@ async def test_retrain_model_error(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]
'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:
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'],
@@ -227,20 +232,18 @@ async def test_retrain_model_error(mlflow):
message='Error retraining model test_model: Error retraining model',
block='retrain_model',
level=NotificationLevel.ERROR,
attachment_content=ANY
attachment_content=ANY,
)
else:
assert False, "No exception raised"
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
mlflow.model_monitoring_repository.update_production_model.return_value = {
'data1': 1,
'data2': 2,
}
)
input_data = {
**metadata,
@@ -248,13 +251,14 @@ async def test_update_production_model(mlflow):
'model_id': 1,
'experiment': 'test',
'timestamp': 2,
'status': 'success'
'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')
experiment='test', model_name='test_model'
)
assert response == {
'data1': {0: 1},
@@ -262,7 +266,7 @@ async def test_update_production_model(mlflow):
'model_id': {0: 1},
'model_name': {0: 'test_model'},
'timestamp': {0: 2},
'status': {0: 'success'}
'status': {0: 'success'},
}
@@ -278,12 +282,12 @@ async def test_update_production_model_error(mlflow):
'model_id': 1,
'experiment': 'test',
'timestamp': 2,
'status': 'success'
'status': 'success',
}
try:
await mlflow.update_production_model(input_data)
except Exception as e:
except Exception as e: # noqa: BLE001
assert str(e) == 'Error updating production model'
mlflow.send_notification.assert_called_once_with(
metadata=metadata['metadata'],
@@ -291,7 +295,7 @@ async def test_update_production_model_error(mlflow):
message='Error updating production model test_model: Error updating production model',
block='update_production_model',
level=NotificationLevel.ERROR,
attachment_content=ANY
attachment_content=ANY,
)
else:
assert False, "No exception raised"
raise AssertionError('No exception raised')

31
tests/laborious/utils/filters/test_conditional_filters.py
View File

@@ -1,23 +1,29 @@
from pandas import DataFrame
from model_manager.utils.filters.conditional_filters import (
filter_empty_data,
filter_specific_variables_null_values,
filter_empty_data
)
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
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
assert (
filter_specific_variables_null_values(
DataFrame({'variable': ['variable1', 'variable2'], 'value': [1, None]}),
config={'variables': ['variable2']},
)
is True
)
def test_filter_empty_data():
@@ -25,6 +31,7 @@ def test_filter_empty_data():
def test_filter_empty_data_with_data():
assert filter_empty_data(
DataFrame({'variable': ['variable1', 'variable2'], 'value': [1, 2]}),
{}) is False
assert (
filter_empty_data(DataFrame({'variable': ['variable1', 'variable2'], 'value': [1, 2]}), {})
is False
)

11
tests/laborious/utils/filters/test_mlflow_filters.py
View File

@@ -1,22 +1,23 @@
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, {}) == True # NOSONAR
assert api_error_filter(None, {})
def test_api_error_filter_valid_response_fail():
assert api_error_filter({'success': False}, {}) == True
assert api_error_filter({'success': False}, {})
def test_api_error_filter_valid_response_success():
assert api_error_filter({'success': True}, {}) == False
assert not api_error_filter({'success': True}, {})
def test_nan_values_filter_all_nan_values():
assert nan_values_filter(DataFrame({'variable': [None, None]}), {}) == True
assert nan_values_filter(DataFrame({'variable': [None, None]}), {})
def test_nan_values_filter_no_nan_values():
assert nan_values_filter(DataFrame({'variable': [1, 2]}), {}) == False
assert not nan_values_filter(DataFrame({'variable': [1, 2]}), {})

319
tests/laborious/utils/repository/test_model_repository.py
View File

@@ -1,34 +1,33 @@
from datetime import UTC, datetime
from unittest.mock import ANY, MagicMock, call, patch
import numpy as np
from pandas import DataFrame
import pytest
from datetime import datetime, timezone
from pandas import Timestamp
from pandas import DataFrame, Timestamp
from model_manager.utils.repository.model_repository import MLFlowRepository
@pytest.fixture
def mlflow_repository():
with patch('model_manager.utils.repository.model_repository.ModelServing',
autospec=True) as mock_model_serving:
with patch(
'model_manager.utils.repository.model_repository.ModelServing', autospec=True
) as mock_model_serving:
mock_instance = mock_model_serving.return_value
mock_instance.get_transformed_data = MagicMock()
repo = MLFlowRepository(
host='http://localhost:5000',
username='admin',
password='admin',
logger=MagicMock()
host='http://localhost:5000', username='admin', password='admin', logger=MagicMock()
)
return repo
metadata = {
"metadata": {
"model_id": "test_model",
"model_name": "test_model",
"workflow_name": "test_workflow",
"schema_name": "test_schedule",
'metadata': {
'model_id': 'test_model',
'model_name': 'test_model',
'workflow_name': 'test_workflow',
'schema_name': 'test_schedule',
},
}
@@ -38,80 +37,56 @@ class Any:
invalid_cases = [
(
{
'value': {
'2024-01-01 12:00:00': 1,
2024: 2
}
}
),
(
{
'value': {
'2024-01-01': 1,
'2024-01-02': 2
}
}
),
(
{
'value': {
Any(): 1,
Any(): 2
}
}
)
({'value': {'2024-01-01 12:00:00': 1, 2024: 2}}),
({'value': {'2024-01-01': 1, '2024-01-02': 2}}),
({'value': {Any(): 1, Any(): 2}}),
]
@pytest.mark.parametrize("data", invalid_cases)
@pytest.mark.parametrize('data', invalid_cases)
def test_detect_and_parse_datetime_index_error_cases(mlflow_repository, data):
input_data = DataFrame(
data
)
input_data = DataFrame(data)
with pytest.raises(ValueError) as e:
mlflow_repository.detect_and_parse_datetime_index(
input_data, metadata['metadata'])
mlflow_repository.detect_and_parse_datetime_index(input_data, metadata['metadata'])
assert str(e) == "Index must be all timestamp like column. Valid formats are: pandas Timestamp, datetime, string in format %Y-%m-%d %H:%M:%S"
assert (
str(e)
== 'Index must be all timestamp like column. Valid formats are: pandas Timestamp, datetime, string in format %Y-%m-%d %H:%M:%S'
)
valid_cases = [
(
{
'value': {
'2024-01-01 12:00:00+0000': 1,
'2024-01-02 12:00:00+0000': 2
}
}, ['2024-01-01 12:00:00+0000', '2024-01-02 12:00:00+0000']
{'value': {'2024-01-01 12:00:00+0000': 1, '2024-01-02 12:00:00+0000': 2}},
['2024-01-01 12:00:00+0000', '2024-01-02 12:00:00+0000'],
),
(
{
'value': {
datetime(2025, 1, 1, 12, 0, 0, tzinfo=timezone.utc): 1,
datetime(2025, 1, 2, 12, 0, 0, tzinfo=timezone.utc): 2
datetime(2025, 1, 1, 12, 0, 0, tzinfo=UTC): 1,
datetime(2025, 1, 2, 12, 0, 0, tzinfo=UTC): 2,
}
}, ['2025-01-01 12:00:00+0000', '2025-01-02 12:00:00+0000']
},
['2025-01-01 12:00:00+0000', '2025-01-02 12:00:00+0000'],
),
(
{
'value': {
Timestamp(2026, 1, 1, 12, 0, 0, tzinfo=timezone.utc): 1,
Timestamp(2026, 1, 2, 12, 0, 0, tzinfo=timezone.utc): 2
Timestamp(2026, 1, 1, 12, 0, 0, tzinfo=UTC): 1,
Timestamp(2026, 1, 2, 12, 0, 0, tzinfo=UTC): 2,
}
}, ['2026-01-01 12:00:00+0000', '2026-01-02 12:00:00+0000']
},
['2026-01-01 12:00:00+0000', '2026-01-02 12:00:00+0000'],
),
]
@pytest.mark.parametrize("data,expected", valid_cases)
@pytest.mark.parametrize('data,expected', valid_cases)
def test_detect_and_parse_datetime_index_valid_format(mlflow_repository, data, expected):
input_data = DataFrame(data)
response = mlflow_repository.detect_and_parse_datetime_index(
input_data, metadata['metadata'])
response = mlflow_repository.detect_and_parse_datetime_index(input_data, metadata['metadata'])
assert response.index.tolist() == expected
@@ -122,18 +97,19 @@ def test_transform_success(mlflow_repository):
mlflow_repository.detect_and_parse_datetime_index = MagicMock()
output = mlflow_repository.transform(
model_name, data, {}, metadata['metadata'])
output = mlflow_repository.transform(model_name, data, {}, metadata['metadata'])
mlflow_repository.model_serving.get_cached_transform.assert_called_once_with(
model_name, data, 0, 'sklearn', False, 'model', 'predict')
model_name, data, 0, 'sklearn', False, 'model', 'predict'
)
mlflow_repository.detect_and_parse_datetime_index.assert_called_once_with(
mlflow_repository.model_serving.get_cached_transform.return_value, metadata['metadata'])
mlflow_repository.model_serving.get_cached_transform.return_value, metadata['metadata']
)
assert output == {
'success': True,
'content': mlflow_repository.detect_and_parse_datetime_index.return_value.to_dict.return_value
'content': mlflow_repository.detect_and_parse_datetime_index.return_value.to_dict.return_value,
}
@@ -141,80 +117,48 @@ def test_transform_error(mlflow_repository):
data = MagicMock()
model_name = 'model'
mlflow_repository.model_serving.get_cached_transform.side_effect = Exception(
'error')
mlflow_repository.model_serving.get_cached_transform.side_effect = Exception('error')
output = mlflow_repository.transform(
model_name, data, {}, metadata['metadata'])
output = mlflow_repository.transform(model_name, data, {}, metadata['metadata'])
mlflow_repository.model_serving.get_cached_transform.assert_called_once_with(
model_name, data, 0, 'sklearn', False, 'model', 'predict')
model_name, data, 0, 'sklearn', False, 'model', 'predict'
)
assert output == {
'success': False,
'content': {
'message': 'error',
'traceback': ANY
}
}
assert output == {'success': False, 'content': {'message': 'error', 'traceback': ANY}}
def test_predict_success(mlflow_repository):
data = DataFrame({
'feat_1': {
'index_1': 2,
'index_2': 3
}
})
data = DataFrame({'feat_1': {'index_1': 2, 'index_2': 3}})
model_name = 'model'
mlflow_repository.model_serving.get_cached_predict.return_value = np.array(
[2, 3]
)
mlflow_repository.model_serving.get_cached_predict.return_value = np.array([2, 3])
output = mlflow_repository.predict(
model_name, data, {}, metadata['metadata'])
output = mlflow_repository.predict(model_name, data, {}, metadata['metadata'])
mlflow_repository.model_serving.get_cached_predict.assert_called_once_with(
model_name, data, 0, 'pyfunc', False, 'model')
model_name, data, 0, 'pyfunc', False, 'model'
)
assert output['success'] is True
assert output['content'] == {
'prediction': {
'index_1': 2,
'index_2': 3
}, 'response_time': {
'index_1': ANY,
'index_2': ANY
}
'prediction': {'index_1': 2, 'index_2': 3},
'response_time': {'index_1': ANY, 'index_2': ANY},
}
def test_predict_error(mlflow_repository):
data = DataFrame({
'feat_1': {
'index_1': 2,
'index_2': 3
}
})
data = DataFrame({'feat_1': {'index_1': 2, 'index_2': 3}})
model_name = 'model'
mlflow_repository.model_serving.get_cached_predict = MagicMock(
side_effect=Exception('error')
)
mlflow_repository.model_serving.get_cached_predict = MagicMock(side_effect=Exception('error'))
output = mlflow_repository.predict(
model_name, data, {}, metadata['metadata'])
output = mlflow_repository.predict(model_name, data, {}, metadata['metadata'])
mlflow_repository.model_serving.get_cached_predict.assert_called_once_with(
model_name, data, 0, 'pyfunc', False, 'model')
model_name, data, 0, 'pyfunc', False, 'model'
)
assert output == {
'success': False,
'content': {
'message': 'error',
'traceback': ANY
}
}
assert output == {'success': False, 'content': {'message': 'error', 'traceback': ANY}}
@patch('model_manager.utils.repository.model_repository.mlflow')
@@ -246,8 +190,7 @@ def test_get_next_run_name(mlflow, mlflow_repository):
@patch('model_manager.utils.repository.model_repository.mlflow')
def test_get_experiment_success(mlflow, mlflow_repository):
mlflow.get_experiment_by_name.return_value = MagicMock(
experiment_id='0')
mlflow.get_experiment_by_name.return_value = MagicMock(experiment_id='0')
output = mlflow_repository.get_experiment('test')
@@ -263,23 +206,25 @@ def test_get_experiment_error(mlflow, mlflow_repository):
except ValueError as e:
assert str(e) == 'Experiment test not found'
else:
assert False
raise AssertionError('Expected exception')
@patch('model_manager.utils.repository.model_repository.mlflow')
def test_get_experiment_last_run(mlflow, mlflow_repository):
mlflow.search_runs.return_value = DataFrame({
mlflow.search_runs.return_value = DataFrame(
{
'params.retrain': ['True', 'False', 'True', 'False'],
'end_time': ['2021-01-01', '2021-01-02', '2021-01-03', '2021-01-04'],
'run_id': ['0', '1', '2', '3'],
})
}
)
output = mlflow_repository.get_experiment_last_run(0)
mlflow.search_runs.assert_called_once_with(
experiment_ids=[0],
filter_string="",
output_format="pandas",
filter_string='',
output_format='pandas',
)
assert output == '2'
@@ -294,17 +239,14 @@ def test_get_experiment_last_run_error(mlflow, mlflow_repository):
except ValueError as e:
assert str(e) == 'Runs is not a pandas DataFrame'
else:
assert False
raise AssertionError('Expected exception')
@patch('model_manager.utils.repository.model_repository.mlflow.sklearn')
@patch('model_manager.utils.repository.model_repository.mlflow.set_experiment')
def test_create_model_experiment(set_experiment, sklearn, mlflow_repository):
mlflow_repository.model_serving.get_model_run_id = MagicMock(
return_value='0')
mlflow_repository.model_serving.get_model_uri = MagicMock(
return_value='test')
mlflow_repository.model_serving.get_model_run_id = MagicMock(return_value='0')
mlflow_repository.model_serving.get_model_uri = MagicMock(return_value='test')
mlflow_repository.get_experiment_by_run_id = MagicMock()
data_model_mock = MagicMock()
@@ -313,29 +255,30 @@ def test_create_model_experiment(set_experiment, sklearn, mlflow_repository):
sklearn.load_model.side_effect = [data_model_mock, prediction_model_mock]
data_model_mock.fit.return_value = data_model_mock
data_model_mock.predict.return_value = DataFrame({
data_model_mock.predict.return_value = DataFrame(
{
'x': [10, 20, 30],
})
}
)
data_model_mock.target_variable = 'y'
prediction_model_mock.fit.return_value = prediction_model_mock
data = DataFrame({
'x': [1, 2, 3],
'y': [4, 5, 6]
})
data = DataFrame({'x': [1, 2, 3], 'y': [4, 5, 6]})
output = mlflow_repository.create_model_experiment('test', data)
mlflow_repository.model_serving.get_model_run_id.assert_called_once_with(
'test', stage='Production')
mlflow_repository.model_serving.get_model_uri.assert_called_once_with(
'0', prediction=False)
'test', stage='Production'
)
mlflow_repository.model_serving.get_model_uri.assert_called_once_with('0', prediction=False)
sklearn.load_model.assert_has_calls([
sklearn.load_model.assert_has_calls(
[
call(mlflow_repository.model_serving.get_model_uri.return_value),
call("models:/test/production"),
])
call('models:/test/production'),
]
)
assert sklearn.load_model.call_count == 2
data_model_mock.fit.assert_called_once_with(data)
@@ -343,21 +286,23 @@ def test_create_model_experiment(set_experiment, sklearn, mlflow_repository):
fit_args = prediction_model_mock.fit.call_args[0][0]
assert fit_args.equals(
DataFrame({
DataFrame(
{
'x': [10, 20, 30],
'y': [4, 5, 6],
})
}
)
)
mlflow_repository.get_experiment_by_run_id.assert_called_once_with('0')
set_experiment.assert_called_once_with(
mlflow_repository.get_experiment_by_run_id.return_value
)
set_experiment.assert_called_once_with(mlflow_repository.get_experiment_by_run_id.return_value)
assert output == (prediction_model_mock,
assert output == (
prediction_model_mock,
data_model_mock,
mlflow_repository.get_experiment_by_run_id.return_value)
mlflow_repository.get_experiment_by_run_id.return_value,
)
@patch('model_manager.utils.repository.model_repository.mlflow.start_run')
@@ -365,41 +310,43 @@ def test_create_model_experiment(set_experiment, sklearn, mlflow_repository):
@patch('model_manager.utils.repository.model_repository.mlflow.sklearn.log_model')
@patch('model_manager.utils.repository.model_repository.mlflow.log_artifact')
def test_perform_model_retrain(log_artifact, log_model, log_param, start_run, mlflow_repository):
prediction_model_mock = MagicMock()
data_model_mock = MagicMock()
experiment = 'test'
model_name = 'test'
data = MagicMock()
mlflow_repository.get_next_run_name = MagicMock(
return_value='test-1')
mlflow_repository.get_next_run_name = MagicMock(return_value='test-1')
run = MagicMock()
start_run.__enter__.return_value = run
output = mlflow_repository.perform_model_retrain(
prediction_model_mock, data_model_mock, experiment, model_name, data)
prediction_model_mock, data_model_mock, experiment, model_name, data
)
mlflow_repository.get_next_run_name.assert_called_once_with(experiment)
start_run.assert_called_once_with(
run_name='test-1', description='Retrain model test with new data')
run_name='test-1', description='Retrain model test with new data'
)
log_model.assert_has_calls([
call(data_model_mock, "data_model"),
call(prediction_model_mock, "prediction_model"),
])
log_model.assert_has_calls(
[
call(data_model_mock, 'data_model'),
call(prediction_model_mock, 'prediction_model'),
]
)
data.to_csv.assert_called_once_with(
"temp/raw_data_test.csv", index=True)
data.to_csv.assert_called_once_with('temp/raw_data_test.csv', index=True)
log_artifact.assert_called_once_with(
"temp/raw_data_test.csv")
log_artifact.assert_called_once_with('temp/raw_data_test.csv')
log_param.assert_has_calls([
call("retrain", True),
])
log_param.assert_has_calls(
[
call('retrain', True),
]
)
assert output == ("Model retrained successfully", experiment)
assert output == ('Model retrained successfully', experiment)
def test_retrain_model(mlflow_repository):
@@ -407,18 +354,18 @@ def test_retrain_model(mlflow_repository):
model_name = 'test'
mlflow_repository.create_model_experiment = MagicMock(
return_value=('data_model', 'prediction_model', '0'))
return_value=('data_model', 'prediction_model', '0')
)
mlflow_repository.perform_model_retrain = MagicMock(
return_value='Model retrained successfully')
mlflow_repository.perform_model_retrain = MagicMock(return_value='Model retrained successfully')
output = mlflow_repository.retrain_model(data, model_name)
mlflow_repository.create_model_experiment.assert_called_once_with(
model_name, data)
mlflow_repository.create_model_experiment.assert_called_once_with(model_name, data)
mlflow_repository.perform_model_retrain.assert_called_once_with(
'data_model', 'prediction_model', '0', model_name, data)
'data_model', 'prediction_model', '0', model_name, data
)
assert output == 'Model retrained successfully'
@@ -438,7 +385,7 @@ def test_update_production_model_by_run_id(mlflow, mlflow_repository):
output = mlflow_repository.update_production_model_by_run_id('0', 'test')
mlflow.register_model.assert_called_once_with(
"runs:/0/prediction_model",
'runs:/0/prediction_model',
'test',
)
@@ -461,38 +408,34 @@ def test_update_production_model_by_run_id(mlflow, mlflow_repository):
@patch('model_manager.utils.repository.model_repository.mlflow')
def test_update_production_model_by_run_id_error(mlflow, mlflow_repository):
mlflow.tracking.MlflowClient.return_value = MagicMock(
get_registered_model=MagicMock(
return_value=MagicMock(
latest_versions={}
)
)
get_registered_model=MagicMock(return_value=MagicMock(latest_versions={}))
)
try:
mlflow_repository.update_production_model_by_run_id('0', 'test')
except Exception as e:
except Exception as e: # noqa: BLE001
assert str(e) == 'Model versions is not a list'
else:
assert False
raise AssertionError('Expected exception')
def test_update_production_model(mlflow_repository):
connector = mlflow_repository
with patch.object(connector, 'get_experiment',
return_value='0') as get_experiment:
with patch.object(connector, 'get_experiment_last_run',
return_value='2') as get_experiment_last_run:
with patch.object(connector, 'update_production_model_by_run_id',
return_value={'model_name': 'test', 'version': '3',
'mlflow_run_id': '0'}) as update_production_model_by_run_id:
with patch.object(connector, 'get_experiment', return_value='0') as get_experiment:
with patch.object(
connector, 'get_experiment_last_run', return_value='2'
) as get_experiment_last_run:
with patch.object(
connector,
'update_production_model_by_run_id',
return_value={'model_name': 'test', 'version': '3', 'mlflow_run_id': '0'},
) as update_production_model_by_run_id:
output = connector.update_production_model('0', 'test')
get_experiment.assert_called_once_with('0')
get_experiment_last_run.assert_called_once_with('0')
update_production_model_by_run_id.assert_called_once_with(
'2', 'test')
update_production_model_by_run_id.assert_called_once_with('2', 'test')
assert output == {
'model_name': 'test',

11
tests/laborious/utils/test_connectors_config.py
View File

@@ -1,7 +1,10 @@
from os import environ
from model_manager.utils.connectors_config import (build_mlflow_config,
from model_manager.utils.connectors_config import (
build_mlflow_config,
build_mongodb_config,
build_postgres_config,
build_mongodb_config)
)
def test_build_mlflow_config_with_env_vars():
@@ -95,7 +98,7 @@ def test_build_mongo_db_config_with_env_vars():
assert build_mongodb_config() == {
'connection_string': 'mongodb://sientia1:sientia1@localhost:27018',
'database_name': 'test_db',
'ttl_index_seconds': 3600
'ttl_index_seconds': 3600,
}
@@ -108,5 +111,5 @@ def test_build_mongo_db_config_with_defaults():
assert build_mongodb_config() == {
'connection_string': 'mongodb://root:wKZDbMNU1c@localhost:27018',
'database_name': 'sientia',
'ttl_index_seconds': 3600
'ttl_index_seconds': 3600,
}

111
tests/laborious/workflows/subworkflows/test_format_and_export_prediction.py
View File

@@ -1,9 +1,12 @@
from unittest.mock import call, patch, AsyncMock, ANY
from pytest import mark, fixture
from unittest.mock import ANY, AsyncMock, call, patch
from pytest import fixture, mark
from sientia_do.temporal.constants import DATETIME_FORMAT_WITH_TZ
from model_manager.activities.activities import Activities
from model_manager.workflows.sub_workflows.format_and_export_prediction import FormatAndExportPrediction
from sientia_do.temporal.constants import DATETIME_FORMAT_WITH_TZ
from model_manager.workflows.sub_workflows.format_and_export_prediction import (
FormatAndExportPrediction,
)
@fixture
@@ -12,34 +15,37 @@ def format_and_export_prediction():
metadata = {
"metadata": {
"model_id": "test_model",
"model_name": "test_model",
"workflow_name": "test_workflow",
"schema_name": "test_schedule",
'metadata': {
'model_id': 'test_model',
'model_name': 'test_model',
'workflow_name': 'test_workflow',
'schema_name': 'test_schedule',
},
}
@mark.asyncio
@patch("model_manager.workflows.sub_workflows.format_and_export_prediction.workflow", new_callable=AsyncMock)
@patch(
'model_manager.workflows.sub_workflows.format_and_export_prediction.workflow',
new_callable=AsyncMock,
)
async def test_run_none_path_flag(workflow_mock, format_and_export_prediction):
input_data = {
'metadata': metadata,
"path_flag": None,
"data": {"test": "data"},
"timestamp": "2021-01-01",
"model_id": 1,
"prediction_confidence": 0,
"schema": "test_schema",
"table_name": "test_table",
"prediction_store_policy": "erl:1"
'path_flag': None,
'data': {'test': 'data'},
'timestamp': '2021-01-01',
'model_id': 1,
'prediction_confidence': 0,
'schema': 'test_schema',
'table_name': 'test_table',
'prediction_store_policy': 'erl:1',
}
await format_and_export_prediction.run(input_data)
workflow_mock.execute_local_activity_method.assert_has_calls([
workflow_mock.execute_local_activity_method.assert_has_calls(
[
call(
Activities.format_prediction,
{
@@ -48,13 +54,16 @@ async def test_run_none_path_flag(workflow_mock, format_and_export_prediction):
'model_id': input_data['model_id'],
'prediction_confidence': input_data['prediction_confidence'],
'prediction_store_policy': input_data['prediction_store_policy'],
**metadata
**metadata,
},
retry_policy=ANY,
start_to_close_timeout=ANY
)])
start_to_close_timeout=ANY,
)
]
)
workflow_mock.execute_activity_method.assert_has_calls([
workflow_mock.execute_activity_method.assert_has_calls(
[
call(
Activities.export_data_to_postgres,
{
@@ -64,36 +73,41 @@ async def test_run_none_path_flag(workflow_mock, format_and_export_prediction):
**metadata,
'timestamp_conversion': {
'column': 'timestamp',
'format': DATETIME_FORMAT_WITH_TZ
}
'format': DATETIME_FORMAT_WITH_TZ,
},
},
retry_policy=ANY,
start_to_close_timeout=ANY
)])
start_to_close_timeout=ANY,
)
]
)
assert workflow_mock.execute_activity_method.call_count == 2
assert workflow_mock.execute_local_activity_method.call_count == 1
@mark.asyncio
@patch("model_manager.workflows.sub_workflows.format_and_export_prediction.workflow", new_callable=AsyncMock)
@patch(
'model_manager.workflows.sub_workflows.format_and_export_prediction.workflow',
new_callable=AsyncMock,
)
async def test_run_default_path_flag(workflow_mock, format_and_export_prediction):
input_data = {
'metadata': metadata,
"path_flag": "default",
"data": {"test": "data"},
"timestamp": "2021-01-01",
"model_id": 1,
"prediction_confidence": 0,
"schema": "test_schema",
"table_name": "test_table",
"comment": "test_comment"
'path_flag': 'default',
'data': {'test': 'data'},
'timestamp': '2021-01-01',
'model_id': 1,
'prediction_confidence': 0,
'schema': 'test_schema',
'table_name': 'test_table',
'comment': 'test_comment',
}
await format_and_export_prediction.run(input_data)
workflow_mock.execute_local_activity_method.assert_has_calls([
workflow_mock.execute_local_activity_method.assert_has_calls(
[
call(
Activities.format_default_prediction,
{
@@ -101,14 +115,16 @@ async def test_run_default_path_flag(workflow_mock, format_and_export_prediction
'model_id': input_data['model_id'],
'prediction_confidence': input_data['prediction_confidence'],
'comment': input_data['comment'],
**metadata
**metadata,
},
retry_policy=ANY,
start_to_close_timeout=ANY
start_to_close_timeout=ANY,
)
]
)
])
workflow_mock.execute_activity_method.assert_has_calls([
workflow_mock.execute_activity_method.assert_has_calls(
[
call(
Activities.export_data_to_postgres,
{
@@ -118,13 +134,14 @@ async def test_run_default_path_flag(workflow_mock, format_and_export_prediction
**metadata,
'timestamp_conversion': {
'column': 'timestamp',
'format': DATETIME_FORMAT_WITH_TZ
}
'format': DATETIME_FORMAT_WITH_TZ,
},
},
retry_policy=ANY,
start_to_close_timeout=ANY
start_to_close_timeout=ANY,
)
]
)
])
assert workflow_mock.execute_activity_method.call_count == 3
assert workflow_mock.execute_local_activity_method.call_count == 1

517
tests/laborious/workflows/subworkflows/test_prediction_process.py
View File

@@ -1,5 +1,7 @@
from unittest.mock import AsyncMock, patch, call, ANY
from unittest.mock import ANY, AsyncMock, call, patch
from pytest import fixture, mark
from model_manager.activities.activities import Activities
from model_manager.workflows.sub_workflows.prediction_process import PredictionProcess
@@ -10,17 +12,17 @@ def prediction_process():
metadata = {
"metadata": {
"model_id": "test_model",
"model_name": "test_model",
"workflow_name": "test_workflow",
"schema_name": "test_schedule",
'metadata': {
'model_id': 'test_model',
'model_name': 'test_model',
'workflow_name': 'test_workflow',
'schema_name': 'test_schedule',
},
}
@mark.asyncio
@patch("model_manager.workflows.sub_workflows.prediction_process.workflow", new_callable=AsyncMock)
@patch('model_manager.workflows.sub_workflows.prediction_process.workflow', new_callable=AsyncMock)
async def test_run(workflow_mock, prediction_process):
prediction_process.path_flag_handler = AsyncMock(return_value=False)
# Arrange
@@ -34,26 +36,23 @@ async def test_run(workflow_mock, prediction_process):
'mlflow_transform_filters': {'test': 'filter'},
'mlflow_predict_filters': {'test': 'filter'},
'model_name': 'test_model_name',
'model_config': {
'retention': '30'
},
'model_config': {'retention': '30'},
'path_priority': ['continue', 'repeat', 'stop'],
'prediction_store_policy': 'lts:1'
'prediction_store_policy': 'lts:1',
}
# Mock the activity responses
workflow_mock.execute_local_activity_method.side_effect = [
'2024-01-01', # get_last_timestamp
('continue', 0.95, "Input data with bad quality"), # input_gate
('continue', 0.95, 'Input data with bad quality'), # input_gate
{'content': 'transformed_data', 'timestamp': '2024-01-01'}, # transform_data
# mlflow_response_gate (transform)
('continue', 0.95, "Error"),
('continue', 0.95, 'Error'),
# mlflow_content_gate (transform)
('continue', 0.95, "Transformed data not passed the content filter"),
('continue', 0.95, 'Transformed data not passed the content filter'),
{'content': 'predicted_data', 'timestamp': '2024-01-01'}, # request_predict
# mlflow_response_gate (predict)
('continue', 0.95, "Error"),
('continue', 0.95, 'Error'),
]
# Act
@@ -62,57 +61,112 @@ async def test_run(workflow_mock, prediction_process):
# Assert
assert workflow_mock.execute_local_activity_method.call_count == 7
workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.get_last_timestamp, {
workflow_mock.execute_local_activity_method.assert_has_calls(
[
call(
Activities.get_last_timestamp,
{
**metadata,
'data': input_data['data'],
},
retry_policy=ANY, start_to_close_timeout=ANY)])
workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.input_gate, {
retry_policy=ANY,
start_to_close_timeout=ANY,
)
]
)
workflow_mock.execute_local_activity_method.assert_has_calls(
[
call(
Activities.input_gate,
{
**metadata,
'filters': input_data['input_filters'],
'data': input_data['data'],
'path_priority': input_data['path_priority'],
}, retry_policy=ANY, start_to_close_timeout=ANY)])
workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.request_transform, {
},
retry_policy=ANY,
start_to_close_timeout=ANY,
)
]
)
workflow_mock.execute_local_activity_method.assert_has_calls(
[
call(
Activities.request_transform,
{
**metadata,
'data': input_data['data'],
'model_name': input_data['model_name'],
'model_config': input_data['model_config'],
}, retry_policy=ANY, start_to_close_timeout=ANY)])
workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.mlflow_response_gate, {
},
retry_policy=ANY,
start_to_close_timeout=ANY,
)
]
)
workflow_mock.execute_local_activity_method.assert_has_calls(
[
call(
Activities.mlflow_response_gate,
{
**metadata,
'filters': input_data['mlflow_transform_filters'],
'data': {'content': 'transformed_data', 'timestamp': '2024-01-01'},
'type': 'transform',
'path_priority': input_data['path_priority'],
}, retry_policy=ANY, start_to_close_timeout=ANY)])
workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.mlflow_content_gate, {
},
retry_policy=ANY,
start_to_close_timeout=ANY,
)
]
)
workflow_mock.execute_local_activity_method.assert_has_calls(
[
call(
Activities.mlflow_content_gate,
{
**metadata,
'filters': input_data['mlflow_transform_filters'],
'data': 'transformed_data',
'type': 'transform',
'path_priority': input_data['path_priority'],
}, retry_policy=ANY, start_to_close_timeout=ANY)])
workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.request_predict, {
},
retry_policy=ANY,
start_to_close_timeout=ANY,
)
]
)
workflow_mock.execute_local_activity_method.assert_has_calls(
[
call(
Activities.request_predict,
{
**metadata,
'data': 'transformed_data',
'model_name': input_data['model_name'],
'model_config': input_data['model_config'],
}, retry_policy=ANY, start_to_close_timeout=ANY)])
workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.mlflow_response_gate, {
},
retry_policy=ANY,
start_to_close_timeout=ANY,
)
]
)
workflow_mock.execute_local_activity_method.assert_has_calls(
[
call(
Activities.mlflow_response_gate,
{
**metadata,
'filters': input_data['mlflow_predict_filters'],
'data': {'content': 'predicted_data', 'timestamp': '2024-01-01'},
'type': 'predict',
'path_priority': input_data['path_priority'],
}, retry_policy=ANY, start_to_close_timeout=ANY)])
},
retry_policy=ANY,
start_to_close_timeout=ANY,
)
]
)
workflow_mock.execute_child_workflow.assert_called_once_with(
'format_and_export_prediction',
@@ -125,17 +179,16 @@ async def test_run(workflow_mock, prediction_process):
'model_id': 1,
'model_name': 'test_model_name',
'model_config': input_data['model_config'],
'schema': input_data['schema'],
'table_name': input_data['table_name'],
'comment': 'Error',
'prediction_store_policy': input_data['prediction_store_policy']
}
'prediction_store_policy': input_data['prediction_store_policy'],
},
)
@mark.asyncio
@patch("model_manager.workflows.sub_workflows.prediction_process.workflow", new_callable=AsyncMock)
@patch('model_manager.workflows.sub_workflows.prediction_process.workflow', new_callable=AsyncMock)
async def test_run_stop_at_input_gate(workflow_mock, prediction_process):
prediction_process.path_flag_handler = AsyncMock(return_value=True)
# Arrange
@@ -149,17 +202,14 @@ async def test_run_stop_at_input_gate(workflow_mock, prediction_process):
'mlflow_transform_filters': {'test': 'filter'},
'mlflow_predict_filters': {'test': 'filter'},
'model_name': 'test_model_name',
'model_config': {
'retention': '30'
},
'model_config': {'retention': '30'},
'path_priority': ['continue', 'repeat', 'stop'],
}
# Mock the activity responses
workflow_mock.execute_local_activity_method.side_effect = [
'2024-01-01', # get_last_timestamp
('stop', 0.95, "Input data with bad quality"), # input_gate
('stop', 0.95, 'Input data with bad quality'), # input_gate
]
# Act
@@ -167,23 +217,35 @@ async def test_run_stop_at_input_gate(workflow_mock, prediction_process):
# Assert
assert workflow_mock.execute_local_activity_method.call_count == 2
workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.get_last_timestamp, {
workflow_mock.execute_local_activity_method.assert_has_calls(
[
call(
Activities.get_last_timestamp,
{
'data': input_data['data'],
**metadata,
}, retry_policy=ANY, start_to_close_timeout=ANY),
call(Activities.input_gate, {
},
retry_policy=ANY,
start_to_close_timeout=ANY,
),
call(
Activities.input_gate,
{
'filters': input_data['input_filters'],
'data': input_data['data'],
'path_priority': input_data['path_priority'],
**metadata,
}, retry_policy=ANY, start_to_close_timeout=ANY)
])
},
retry_policy=ANY,
start_to_close_timeout=ANY,
),
]
)
workflow_mock.execute_child_workflow.assert_not_called()
@mark.asyncio
@patch("model_manager.workflows.sub_workflows.prediction_process.workflow", new_callable=AsyncMock)
@patch('model_manager.workflows.sub_workflows.prediction_process.workflow', new_callable=AsyncMock)
async def test_run_stop_at_first_mlflow_response_gate(workflow_mock, prediction_process):
prediction_process.path_flag_handler = AsyncMock(side_effect=[False, True])
# Arrange
@@ -197,19 +259,16 @@ async def test_run_stop_at_first_mlflow_response_gate(workflow_mock, prediction_
'mlflow_transform_filters': {'test': 'filter'},
'mlflow_predict_filters': {'test': 'filter'},
'model_name': 'test_model_name',
'model_config': {
'retention': '30'
},
'model_config': {'retention': '30'},
'path_priority': ['continue', 'repeat', 'stop'],
}
# Mock the activity responses
workflow_mock.execute_local_activity_method.side_effect = [
'2024-01-01', # get_last_timestamp
('repeat', 0.95, "Input data with bad quality"), # input_gate
('repeat', 0.95, 'Input data with bad quality'), # input_gate
{'content': 'transformed_data', 'timestamp': '2024-01-01'}, # transform_data
('continue', 0.95, "Error"), # mlflow_response_gate (transform)
('continue', 0.95, 'Error'), # mlflow_response_gate (transform)
]
# Act
@@ -217,46 +276,72 @@ async def test_run_stop_at_first_mlflow_response_gate(workflow_mock, prediction_
# Assert
assert workflow_mock.execute_local_activity_method.call_count == 4
workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.get_last_timestamp, {
workflow_mock.execute_local_activity_method.assert_has_calls(
[
call(
Activities.get_last_timestamp,
{
'data': input_data['data'],
**metadata,
},
retry_policy=ANY, start_to_close_timeout=ANY)])
workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.input_gate, {
retry_policy=ANY,
start_to_close_timeout=ANY,
)
]
)
workflow_mock.execute_local_activity_method.assert_has_calls(
[
call(
Activities.input_gate,
{
'filters': input_data['input_filters'],
'data': input_data['data'],
'path_priority': input_data['path_priority'],
**metadata,
},
retry_policy=ANY, start_to_close_timeout=ANY)])
workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.request_transform, {
retry_policy=ANY,
start_to_close_timeout=ANY,
)
]
)
workflow_mock.execute_local_activity_method.assert_has_calls(
[
call(
Activities.request_transform,
{
'data': input_data['data'],
'model_name': input_data['model_name'],
'model_config': input_data['model_config'],
**metadata
**metadata,
},
retry_policy=ANY, start_to_close_timeout=ANY)
])
workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.mlflow_response_gate, {
retry_policy=ANY,
start_to_close_timeout=ANY,
)
]
)
workflow_mock.execute_local_activity_method.assert_has_calls(
[
call(
Activities.mlflow_response_gate,
{
'filters': input_data['mlflow_transform_filters'],
'data': {'content': 'transformed_data', 'timestamp': '2024-01-01'},
'type': 'transform',
'path_priority': input_data['path_priority'],
**metadata
}, retry_policy=ANY, start_to_close_timeout=ANY)
])
**metadata,
},
retry_policy=ANY,
start_to_close_timeout=ANY,
)
]
)
workflow_mock.execute_child_workflow.assert_not_called()
@mark.asyncio
@patch("model_manager.workflows.sub_workflows.prediction_process.workflow", new_callable=AsyncMock)
@patch('model_manager.workflows.sub_workflows.prediction_process.workflow', new_callable=AsyncMock)
async def test_run_stop_at_mlflow_content_gate(workflow_mock, prediction_process):
prediction_process.path_flag_handler = AsyncMock(
side_effect=[False, False, True])
prediction_process.path_flag_handler = AsyncMock(side_effect=[False, False, True])
# Arrange
input_data = {
'metadata': metadata,
@@ -268,22 +353,19 @@ async def test_run_stop_at_mlflow_content_gate(workflow_mock, prediction_process
'mlflow_transform_filters': {'test': 'filter'},
'mlflow_predict_filters': {'test': 'filter'},
'model_name': 'test_model_name',
'model_config': {
'retention': '30'
},
'model_config': {'retention': '30'},
'path_priority': ['continue', 'repeat', 'stop'],
}
# Mock the activity responses
workflow_mock.execute_local_activity_method.side_effect = [
'2024-01-01', # get_last_timestamp
('continue', 0.95, "Input data with bad quality"), # input_gate
('continue', 0.95, 'Input data with bad quality'), # input_gate
{'content': 'transformed_data', 'timestamp': '2024-01-01'}, # transform_data
# mlflow_response_gate (transform)
('continue', 0.95, "Error"),
('continue', 0.95, 'Error'),
# mlflow_content_gate (transform)
('continue', 0.95, "Transformed data not passed the content filter"),
('continue', 0.95, 'Transformed data not passed the content filter'),
]
# Act
@@ -292,51 +374,88 @@ async def test_run_stop_at_mlflow_content_gate(workflow_mock, prediction_process
# Assert
assert workflow_mock.execute_local_activity_method.call_count == 5
workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.get_last_timestamp, {
workflow_mock.execute_local_activity_method.assert_has_calls(
[
call(
Activities.get_last_timestamp,
{
'data': input_data['data'],
**metadata,
},
retry_policy=ANY, start_to_close_timeout=ANY)])
workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.input_gate, {
retry_policy=ANY,
start_to_close_timeout=ANY,
)
]
)
workflow_mock.execute_local_activity_method.assert_has_calls(
[
call(
Activities.input_gate,
{
'filters': input_data['input_filters'],
'data': input_data['data'],
'path_priority': input_data['path_priority'],
**metadata,
},
retry_policy=ANY, start_to_close_timeout=ANY)])
workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.request_transform, {
retry_policy=ANY,
start_to_close_timeout=ANY,
)
]
)
workflow_mock.execute_local_activity_method.assert_has_calls(
[
call(
Activities.request_transform,
{
'data': input_data['data'],
'model_name': input_data['model_name'],
'model_config': input_data['model_config'],
**metadata
}, retry_policy=ANY, start_to_close_timeout=ANY)])
workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.mlflow_response_gate, {
**metadata,
},
retry_policy=ANY,
start_to_close_timeout=ANY,
)
]
)
workflow_mock.execute_local_activity_method.assert_has_calls(
[
call(
Activities.mlflow_response_gate,
{
'filters': input_data['mlflow_transform_filters'],
'data': {'content': 'transformed_data', 'timestamp': '2024-01-01'},
'type': 'transform',
'path_priority': input_data['path_priority'],
**metadata,
}, retry_policy=ANY, start_to_close_timeout=ANY)])
workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.mlflow_content_gate, {
},
retry_policy=ANY,
start_to_close_timeout=ANY,
)
]
)
workflow_mock.execute_local_activity_method.assert_has_calls(
[
call(
Activities.mlflow_content_gate,
{
'filters': input_data['mlflow_transform_filters'],
'data': 'transformed_data',
'type': 'transform',
'path_priority': input_data['path_priority'],
**metadata,
}, retry_policy=ANY, start_to_close_timeout=ANY)])
},
retry_policy=ANY,
start_to_close_timeout=ANY,
)
]
)
workflow_mock.execute_child_workflow.assert_not_called()
@mark.asyncio
@patch("model_manager.workflows.sub_workflows.prediction_process.workflow", new_callable=AsyncMock)
@patch('model_manager.workflows.sub_workflows.prediction_process.workflow', new_callable=AsyncMock)
async def test_run_stop_at_mlflow_last_response_gate(workflow_mock, prediction_process):
prediction_process.path_flag_handler = AsyncMock(
side_effect=[False, False, False, True])
prediction_process.path_flag_handler = AsyncMock(side_effect=[False, False, False, True])
# Arrange
input_data = {
'metadata': metadata,
@@ -348,24 +467,21 @@ async def test_run_stop_at_mlflow_last_response_gate(workflow_mock, prediction_p
'mlflow_transform_filters': {'test': 'filter'},
'mlflow_predict_filters': {'test': 'filter'},
'model_name': 'test_model_name',
'model_config': {
'retention': '30'
},
'model_config': {'retention': '30'},
'path_priority': ['continue', 'repeat', 'stop'],
}
# Mock the activity responses
workflow_mock.execute_local_activity_method.side_effect = [
'2024-01-01', # get_last_timestamp
('continue', 0.95, "Input data with bad quality"), # input_gate
('continue', 0.95, 'Input data with bad quality'), # input_gate
{'content': 'transformed_data', 'timestamp': '2024-01-01'}, # transform_data
# mlflow_response_gate (transform)
('continue', 0.95, "Error"),
('continue', 0.95, 'Error'),
# mlflow_content_gate (transform)
('continue', 0.95, "Transformed data not passed the content filter"),
('continue', 0.95, 'Transformed data not passed the content filter'),
{'content': 'predicted_data', 'timestamp': '2024-01-01'}, # request_predict
('continue', 0.95, "Error"), # mlflow_response_gate (predict)
('continue', 0.95, 'Error'), # mlflow_response_gate (predict)
]
# Act
@@ -373,63 +489,117 @@ async def test_run_stop_at_mlflow_last_response_gate(workflow_mock, prediction_p
# Assert
assert workflow_mock.execute_local_activity_method.call_count == 7
workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.get_last_timestamp, {
workflow_mock.execute_local_activity_method.assert_has_calls(
[
call(
Activities.get_last_timestamp,
{
'data': input_data['data'],
**metadata,
},
retry_policy=ANY, start_to_close_timeout=ANY)])
workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.input_gate, {
retry_policy=ANY,
start_to_close_timeout=ANY,
)
]
)
workflow_mock.execute_local_activity_method.assert_has_calls(
[
call(
Activities.input_gate,
{
'filters': input_data['input_filters'],
'data': input_data['data'],
'path_priority': input_data['path_priority'],
**metadata,
},
retry_policy=ANY, start_to_close_timeout=ANY)])
workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.request_transform, {
retry_policy=ANY,
start_to_close_timeout=ANY,
)
]
)
workflow_mock.execute_local_activity_method.assert_has_calls(
[
call(
Activities.request_transform,
{
'data': input_data['data'],
'model_name': input_data['model_name'],
'model_config': input_data['model_config'],
**metadata
}, retry_policy=ANY, start_to_close_timeout=ANY)])
workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.mlflow_response_gate, {
**metadata,
},
retry_policy=ANY,
start_to_close_timeout=ANY,
)
]
)
workflow_mock.execute_local_activity_method.assert_has_calls(
[
call(
Activities.mlflow_response_gate,
{
'filters': input_data['mlflow_transform_filters'],
'data': {'content': 'transformed_data', 'timestamp': '2024-01-01'},
'type': 'transform',
'path_priority': input_data['path_priority'],
**metadata,
}, retry_policy=ANY, start_to_close_timeout=ANY)])
workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.mlflow_content_gate, {
},
retry_policy=ANY,
start_to_close_timeout=ANY,
)
]
)
workflow_mock.execute_local_activity_method.assert_has_calls(
[
call(
Activities.mlflow_content_gate,
{
'filters': input_data['mlflow_transform_filters'],
'data': 'transformed_data',
'type': 'transform',
'path_priority': input_data['path_priority'],
**metadata,
}, retry_policy=ANY, start_to_close_timeout=ANY)])
workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.request_predict, {
},
retry_policy=ANY,
start_to_close_timeout=ANY,
)
]
)
workflow_mock.execute_local_activity_method.assert_has_calls(
[
call(
Activities.request_predict,
{
'data': 'transformed_data',
'model_name': input_data['model_name'],
'model_config': input_data['model_config'],
**metadata
}, retry_policy=ANY, start_to_close_timeout=ANY)])
workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.mlflow_response_gate, {
**metadata,
},
retry_policy=ANY,
start_to_close_timeout=ANY,
)
]
)
workflow_mock.execute_local_activity_method.assert_has_calls(
[
call(
Activities.mlflow_response_gate,
{
'filters': input_data['mlflow_predict_filters'],
'data': {'content': 'predicted_data', 'timestamp': '2024-01-01'},
'type': 'predict',
'path_priority': input_data['path_priority'],
**metadata,
}, retry_policy=ANY, start_to_close_timeout=ANY)])
},
retry_policy=ANY,
start_to_close_timeout=ANY,
)
]
)
workflow_mock.execute_child_workflow.assert_not_called()
@mark.asyncio
@patch("model_manager.workflows.sub_workflows.prediction_process.workflow", new_callable=AsyncMock)
@patch('model_manager.workflows.sub_workflows.prediction_process.workflow', new_callable=AsyncMock)
async def test_path_flag_handler_stop(workflow_mock, prediction_process):
# Arrange
data = {'test': 'data'}
@@ -440,21 +610,24 @@ async def test_path_flag_handler_stop(workflow_mock, prediction_process):
model = 'test_model'
last_timestamp = '2024-01-01'
model_name = 'test_model_name'
model_config = {
'retention': '30'
}
model_config = {'retention': '30'}
# Act
result = await prediction_process.path_flag_handler(
data, path_flag, {
data,
path_flag,
{
'metadata': metadata,
'schema': schema,
'table_name': table_name,
'model_id': model,
'last_timestamp': last_timestamp,
'model_name': model_name,
'model_config': model_config
}, confidence, last_timestamp, ""
'model_config': model_config,
},
confidence,
last_timestamp,
'',
)
# Assert
@@ -464,7 +637,7 @@ async def test_path_flag_handler_stop(workflow_mock, prediction_process):
@mark.asyncio
@patch("model_manager.workflows.sub_workflows.prediction_process.workflow", new_callable=AsyncMock)
@patch('model_manager.workflows.sub_workflows.prediction_process.workflow', new_callable=AsyncMock)
async def test_path_flag_handler_repeat(workflow_mock, prediction_process):
# Arrange
data = {'test': 'data'}
@@ -475,21 +648,24 @@ async def test_path_flag_handler_repeat(workflow_mock, prediction_process):
model = 'test_model'
last_timestamp = '2024-01-01'
model_name = 'test_model_name'
model_config = {
'retention': '30'
}
model_config = {'retention': '30'}
# Act
result = await prediction_process.path_flag_handler(
data, path_flag, {
data,
path_flag,
{
'metadata': metadata,
'schema': schema,
'table_name': table_name,
'model_id': model,
'last_timestamp': last_timestamp,
'model_name': model_name,
'model_config': model_config
}, confidence, last_timestamp, ""
'model_config': model_config,
},
confidence,
last_timestamp,
'',
)
# Assert
@@ -504,13 +680,13 @@ async def test_path_flag_handler_repeat(workflow_mock, prediction_process):
'last_timestamp': last_timestamp,
},
retry_policy=ANY,
start_to_close_timeout=ANY
start_to_close_timeout=ANY,
)
workflow_mock.execute_child_workflow.assert_not_called()
@mark.asyncio
@patch("model_manager.workflows.sub_workflows.prediction_process.workflow", new_callable=AsyncMock)
@patch('model_manager.workflows.sub_workflows.prediction_process.workflow', new_callable=AsyncMock)
async def test_path_flag_handler_continue(workflow_mock, prediction_process):
# Arrange
data = {'test': 'data'}
@@ -521,14 +697,14 @@ async def test_path_flag_handler_continue(workflow_mock, prediction_process):
model = 'test_model'
last_timestamp = '2024-01-01'
model_name = 'test_model_name'
model_config = {
'retention': '30'
}
model_config = {'retention': '30'}
prediction_store_policy = 'erl:1'
# Act
result = await prediction_process.path_flag_handler(
data, path_flag, {
data,
path_flag,
{
'metadata': metadata,
'schema': schema,
'table_name': table_name,
@@ -536,9 +712,11 @@ async def test_path_flag_handler_continue(workflow_mock, prediction_process):
'last_timestamp': last_timestamp,
'model_name': model_name,
'model_config': model_config,
'prediction_store_policy': prediction_store_policy
}, confidence, last_timestamp, 'Prediction Process'
'prediction_store_policy': prediction_store_policy,
},
confidence,
last_timestamp,
'Prediction Process',
)
# Assert
@@ -558,14 +736,13 @@ async def test_path_flag_handler_continue(workflow_mock, prediction_process):
'schema': schema,
'table_name': table_name,
'comment': 'Prediction Process',
'prediction_store_policy': prediction_store_policy
}
'prediction_store_policy': prediction_store_policy,
},
)
@mark.asyncio
@patch("model_manager.workflows.sub_workflows.prediction_process.workflow", new_callable=AsyncMock)
@patch('model_manager.workflows.sub_workflows.prediction_process.workflow', new_callable=AsyncMock)
async def test_path_flag_handler_unknown(workflow_mock, prediction_process):
# Arrange
data = {'test': 'data'}
@@ -576,13 +753,13 @@ async def test_path_flag_handler_unknown(workflow_mock, prediction_process):
model = 'test_model'
last_timestamp = '2024-01-01'
model_name = 'test_model_name'
model_config = {
'retention': '30'
}
model_config = {'retention': '30'}
prediction_store_policy = 'erl:1'
# Act
result = await prediction_process.path_flag_handler(
data, path_flag, {
data,
path_flag,
{
**metadata,
'schema': schema,
'table_name': table_name,
@@ -590,9 +767,11 @@ async def test_path_flag_handler_unknown(workflow_mock, prediction_process):
'last_timestamp': last_timestamp,
'model_name': model_name,
'model_config': model_config,
'prediction_store_policy': prediction_store_policy
}, confidence, last_timestamp, ""
'prediction_store_policy': prediction_store_policy,
},
confidence,
last_timestamp,
'',
)
# Assert

62
tests/laborious/workflows/test_minimal_retrain.py
View File

@@ -1,5 +1,7 @@
from unittest.mock import AsyncMock, MagicMock, call, patch, ANY
from unittest.mock import ANY, AsyncMock, call, patch
from pytest import fixture, mark
from model_manager.activities.activities import Activities
from model_manager.workflows.minimal_retrain import MinimalRetrain
@@ -10,11 +12,11 @@ def minimal_retrain() -> MinimalRetrain:
metadata = {
"metadata": {
"model_id": "test_model_id",
"model_name": "test_model",
"workflow_name": "minimal_retrain",
"schedule_name": "test_schedule",
'metadata': {
'model_id': 'test_model_id',
'model_name': 'test_model',
'workflow_name': 'minimal_retrain',
'schedule_name': 'test_schedule',
},
}
@@ -23,19 +25,19 @@ metadata = {
@patch('model_manager.workflows.minimal_retrain.workflow', new_callable=AsyncMock)
async def test_run(workflow_mock: AsyncMock, minimal_retrain: MinimalRetrain):
input_data = {
"model_id": "test_model_id",
"model_name": "test_model",
"workflow_name": "minimal_retrain",
"schedule_name": "test_schedule",
"query": "test_query",
"schema": "test_schema",
"table_name": "test_table",
'model_id': 'test_model_id',
'model_name': 'test_model',
'workflow_name': 'minimal_retrain',
'schedule_name': 'test_schedule',
'query': 'test_query',
'schema': 'test_schema',
'table_name': 'test_table',
}
workflow_mock.execute_activity_method = AsyncMock(
return_value={
"data1": "1",
"data2": "2",
'data1': '1',
'data2': '2',
}
)
@@ -47,16 +49,17 @@ async def test_run(workflow_mock: AsyncMock, minimal_retrain: MinimalRetrain):
Activities.load_custom_query,
{
**metadata,
"query": input_data["query"],
'datetime_columns': input_data.get('datetime_columns', [])
'query': input_data['query'],
'datetime_columns': input_data.get('datetime_columns', []),
},
retry_policy=ANY,
start_to_close_timeout=ANY
start_to_close_timeout=ANY,
)
]
)
workflow_mock.execute_activity_method.assert_has_calls([
workflow_mock.execute_activity_method.assert_has_calls(
[
call(
Activities.retrain_model,
{
@@ -65,11 +68,13 @@ async def test_run(workflow_mock: AsyncMock, minimal_retrain: MinimalRetrain):
'model_name': input_data['model_name'],
},
retry_policy=ANY,
start_to_close_timeout=ANY
start_to_close_timeout=ANY,
)
]
)
])
workflow_mock.execute_activity_method.assert_has_calls([
workflow_mock.execute_activity_method.assert_has_calls(
[
call(
Activities.update_production_model,
{
@@ -79,11 +84,13 @@ async def test_run(workflow_mock: AsyncMock, minimal_retrain: MinimalRetrain):
**workflow_mock.execute_activity_method.return_value,
},
retry_policy=ANY,
start_to_close_timeout=ANY
start_to_close_timeout=ANY,
)
]
)
])
workflow_mock.execute_activity_method.assert_has_calls([
workflow_mock.execute_activity_method.assert_has_calls(
[
call(
Activities.export_data_to_postgres,
{
@@ -93,6 +100,7 @@ async def test_run(workflow_mock: AsyncMock, minimal_retrain: MinimalRetrain):
'table_name': input_data['table_name'],
},
retry_policy=ANY,
start_to_close_timeout=ANY
start_to_close_timeout=ANY,
)
]
)
])

65
tests/laborious/workflows/test_predictions_batch.py
View File

@@ -1,5 +1,7 @@
from unittest.mock import AsyncMock, call, patch, ANY
from unittest.mock import ANY, AsyncMock, call, patch
from pytest import fixture, mark
from model_manager.activities.activities import Activities
from model_manager.workflows.predictions_batch import PredictionsBatch
@@ -10,11 +12,11 @@ def predictions_batch() -> PredictionsBatch:
metadata = {
"metadata": {
"model_id": "test_model_id",
"model_name": "test_model",
"workflow_name": "predictions_batch",
"schedule_name": "test_schedule",
'metadata': {
'model_id': 'test_model_id',
'model_name': 'test_model',
'workflow_name': 'predictions_batch',
'schedule_name': 'test_schedule',
},
}
@@ -22,9 +24,7 @@ metadata = {
@mark.asyncio
@patch('model_manager.workflows.predictions_batch.workflow', new_callable=AsyncMock)
async def test_run(workflow_mock: AsyncMock, predictions_batch: PredictionsBatch):
workflow_mock.execute_local_activity_method.return_value = {
'data': 'test_data'
}
workflow_mock.execute_local_activity_method.return_value = {'data': 'test_data'}
input_data = {
'schedule_name': 'test_schedule',
'model_name': 'test_model',
@@ -32,28 +32,27 @@ async def test_run(workflow_mock: AsyncMock, predictions_batch: PredictionsBatch
'query': 'SELECT * FROM test',
'schema': 'test_schema',
'table_name': 'test_table',
'datetime_columns': ['timestamp', 'created_at'],
'prediction_store_policy': 'erl:1',
'model_config': {
'retention': '30'
}
'model_config': {'retention': '30'},
}
await predictions_batch.run(input_data)
workflow_mock.execute_local_activity_method.assert_has_calls([
workflow_mock.execute_local_activity_method.assert_has_calls(
[
call(
Activities.load_custom_query,
{
**metadata,
'query': input_data['query'],
'datetime_columns': input_data.get('datetime_columns', [])
'datetime_columns': input_data.get('datetime_columns', []),
},
retry_policy=ANY,
start_to_close_timeout=ANY
start_to_close_timeout=ANY,
)
]
)
])
prediction_input = {
'metadata': metadata,
'data': {'data': 'test_data'},
@@ -61,28 +60,18 @@ async def test_run(workflow_mock: AsyncMock, predictions_batch: PredictionsBatch
'table_name': input_data['table_name'],
'model_id': input_data['model_id'],
'model_name': input_data['model_name'],
'input_filters': input_data.get('input_filters', {
'EMPTY_DATA': {
'POLICY': 'STOP'
}
}),
'mlflow_transform_filters': input_data.get('mlflow_transform_filters', {
'API_ERROR': {
'POLICY': 'STOP'
}
}),
'mlflow_predict_filters': input_data.get('mlflow_predict_filters', {
'API_ERROR': {
'POLICY': 'STOP'
}
}),
'input_filters': input_data.get('input_filters', {'EMPTY_DATA': {'POLICY': 'STOP'}}),
'mlflow_transform_filters': input_data.get(
'mlflow_transform_filters', {'API_ERROR': {'POLICY': 'STOP'}}
),
'mlflow_predict_filters': input_data.get(
'mlflow_predict_filters', {'API_ERROR': {'POLICY': 'STOP'}}
),
'model_config': input_data.get('model_config', {}),
'path_priority': input_data.get('path_priority', ['STOP', 'CONTINUE', 'REPEAT']),
'prediction_store_policy': input_data.get('prediction_store_policy', 'erl:1')
'prediction_store_policy': input_data.get('prediction_store_policy', 'erl:1'),
}
workflow_mock.execute_child_workflow.assert_has_calls([
call(
'prediction_process', prediction_input)
])
workflow_mock.execute_child_workflow.assert_has_calls(
[call('prediction_process', prediction_input)]
)