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

61
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(): 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.notifications.handlers import CoreNotificationHandler as NotificationHandler
from sientia_do.observability.logger import Logger 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 model_manager.activities.gates import Gates
from typing import Any from model_manager.activities.mlflow import MLFlow
class Activities(Postgres, MLFlow, Gates): class Activities(Postgres, MLFlow, Gates):
@@ -29,11 +31,13 @@ class Activities(Postgres, MLFlow, Gates):
notification_handler (NotificationHandler): Notification management instance notification_handler (NotificationHandler): Notification management instance
""" """
def __init__(self, def __init__(
postgres_config: dict[str, Any], self,
mlflow_config: dict[str, Any], postgres_config: dict[str, Any],
logger: Logger, mlflow_config: dict[str, Any],
notification_handler: NotificationHandler): logger: Logger,
notification_handler: NotificationHandler,
):
""" """
Initialize the Activities orchestrator with all required configurations. Initialize the Activities orchestrator with all required configurations.
@@ -52,25 +56,30 @@ class Activities(Postgres, MLFlow, Gates):
Exception: If any parent class initialization fails Exception: If any parent class initialization fails
""" """
# Initialize parent classes # Initialize parent classes
Postgres.__init__(self, host=postgres_config['host'], Postgres.__init__(
port=postgres_config['port'], self,
user=postgres_config['user'], host=postgres_config['host'],
password=postgres_config['password'], port=postgres_config['port'],
dbname=postgres_config['dbname'], user=postgres_config['user'],
min_connections=postgres_config['min_connections'], password=postgres_config['password'],
max_connections=postgres_config['max_connections'], dbname=postgres_config['dbname'],
logger=logger, min_connections=postgres_config['min_connections'],
notification_handler=notification_handler) max_connections=postgres_config['max_connections'],
logger=logger,
notification_handler=notification_handler,
)
MLFlow.__init__(self, mlflow_host=mlflow_config['host'], MLFlow.__init__(
mlflow_port=mlflow_config['port'], self,
mlflow_username=mlflow_config['username'], mlflow_host=mlflow_config['host'],
mlflow_password=mlflow_config['password'], mlflow_port=mlflow_config['port'],
logger=logger, mlflow_username=mlflow_config['username'],
notification_handler=notification_handler) mlflow_password=mlflow_config['password'],
logger=logger,
notification_handler=notification_handler,
)
Gates.__init__(self, logger=logger, Gates.__init__(self, logger=logger, notification_handler=notification_handler)
notification_handler=notification_handler)
async def shutdown(self): async def shutdown(self):
""" """

258
model_manager/activities/gates.py
View File

@@ -1,53 +1,42 @@
from temporalio import activity, workflow from temporalio import activity, workflow
with workflow.unsafe.imports_passed_through(): with workflow.unsafe.imports_passed_through():
import traceback 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.handlers import CoreNotificationHandler as NotificationHandler
from sientia_do.notifications.models import NotificationLevel 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.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.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 ( from model_manager.utils.filters.conditional_filters import (
filter_empty_data, filter_empty_data,
filter_specific_variables_null_values filter_specific_variables_null_values,
) )
from pandas import DataFrame from model_manager.utils.filters.mlflow_filters import api_error_filter, nan_values_filter
from model_manager import metrics
# Input filter function mappings # Input filter function mappings
input_filter_functions = { input_filter_functions = {
'SPECIFIC_VARIABLES_NULL_VALUES': filter_specific_variables_null_values, 'SPECIFIC_VARIABLES_NULL_VALUES': filter_specific_variables_null_values,
'EMPTY_DATA': filter_empty_data, 'EMPTY_DATA': filter_empty_data,
'path_confidence': { 'path_confidence': {'STOP': -1, 'CONTINUE': 2, 'REPEAT': -1},
'STOP': -1,
'CONTINUE': 2,
'REPEAT': -1
}
} }
# MLFlow response filter function mappings # MLFlow response filter function mappings
mlflow_response_filter_functions = { mlflow_response_filter_functions = {
'API_ERROR': api_error_filter, 'API_ERROR': api_error_filter,
'path_confidence': { 'path_confidence': {'STOP': -1, 'CONTINUE': 10, 'REPEAT': -1},
'STOP': -1,
'CONTINUE': 10,
'REPEAT': -1
},
} }
# MLFlow content filter function mappings # MLFlow content filter function mappings
mlflow_content_filter_functions = { mlflow_content_filter_functions = {
'NAN_VALUES': nan_values_filter, 'NAN_VALUES': nan_values_filter,
'EMPTY_DATA': filter_empty_data, 'EMPTY_DATA': filter_empty_data,
'path_confidence': { 'path_confidence': {'STOP': -1, 'CONTINUE': 18, 'REPEAT': -1},
'STOP': -1,
'CONTINUE': 18,
'REPEAT': -1
}
} }
@@ -82,10 +71,9 @@ class Gates(BaseActivity):
Raises: Raises:
Exception: If BaseActivity initialization fails Exception: If BaseActivity initialization fails
""" """
BaseActivity.__init__( BaseActivity.__init__(self, logger, notification_handler, set_error_counter=True)
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]: async def input_gate(self, input_data: dict[str, Any]) -> tuple[str | None, int, str]:
""" """
Apply input data quality filters and validation. Apply input data quality filters and validation.
@@ -121,7 +109,7 @@ class Gates(BaseActivity):
""" """
metadata = input_data['metadata'] metadata = input_data['metadata']
self.info("Performing input gate...", metadata) self.info('Performing input gate...', metadata)
filters = input_data['filters'] filters = input_data['filters']
data = DataFrame(input_data['data']) data = DataFrame(input_data['data'])
@@ -129,40 +117,42 @@ class Gates(BaseActivity):
filter_output = [] filter_output = []
self.debug(f"Input data: {data.head(5).to_string()}", metadata) self.debug(f'Input data: {data.head(5).to_string()}', metadata)
self.debug(f"Filters: {filters}", metadata) self.debug(f'Filters: {filters}', metadata)
# Apply each configured filter # Apply each configured filter
for fil, config in filters.items(): for fil, config in filters.items():
if fil not in input_filter_functions: if fil not in input_filter_functions:
self.error(f"Filter {fil} not found", metadata) self.error(f'Filter {fil} not found', metadata)
continue continue
try: try:
if input_filter_functions[fil](data, config['config']): if input_filter_functions[fil](data, config['config']):
self.debug( self.debug(f'Data not passed the input filter {fil}:{config}', metadata)
f"Data not passed the input filter {fil}:{config}", metadata)
filter_output.append(config['policy']) filter_output.append(config['policy'])
except Exception as e: except Exception as e: # noqa: BLE001
trace = traceback.format_exc() trace = traceback.format_exc()
self.send_notification( self.send_notification(
metadata=metadata, metadata=metadata,
notification_id=f"INTPUT_GATE_ERROR__{fil}", notification_id=f'INTPUT_GATE_ERROR__{fil}',
message=f"Error in filter {fil}:{config}: \n {e}", message=f'Error in filter {fil}:{config}: \n {e}',
block="input_gate", block='input_gate',
level=NotificationLevel.ERROR, level=NotificationLevel.ERROR,
attachment_content=trace attachment_content=trace,
) )
for path_flag in path_priority: for path_flag in path_priority:
if path_flag in filter_output: if path_flag in filter_output:
self.info(f"Input gate result: {path_flag}", metadata) self.info(f'Input gate result: {path_flag}', metadata)
return path_flag, input_filter_functions['path_confidence'][path_flag], \ return (
"Input data with bad quality" path_flag,
input_filter_functions['path_confidence'][path_flag],
'Input data with bad quality',
)
self.info("Nothing was filtered by the input gate", metadata) self.info('Nothing was filtered by the input gate', metadata)
return None, 0, "" 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]: async def mlflow_response_gate(self, input_data: dict[str, Any]) -> tuple[str | None, int, str]:
""" """
Validate MLFlow API response quality and integrity. Validate MLFlow API response quality and integrity.
@@ -197,7 +187,7 @@ class Gates(BaseActivity):
Exception: If response validation fails or configuration is invalid Exception: If response validation fails or configuration is invalid
""" """
metadata = input_data['metadata'] metadata = input_data['metadata']
self.info("Performing mlflow response gate...", metadata) self.info('Performing mlflow response gate...', metadata)
filters = input_data['filters'] filters = input_data['filters']
data = input_data['data'] data = input_data['data']
@@ -206,14 +196,13 @@ class Gates(BaseActivity):
filter_output = [] filter_output = []
self.debug( self.debug(f'Input data: \n {create_sample_dict(data, max_items=5, max_depth=2)}', metadata)
f"Input data: \n {create_sample_dict(data, max_items=5, max_depth=2)}", metadata) self.debug(f'Filters: {filters}', metadata)
self.debug(f"Filters: {filters}", metadata)
comments = [] comments = []
for fil, config in filters.items(): for fil, config in filters.items():
if fil not in mlflow_response_filter_functions: 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 continue
try: try:
if mlflow_response_filter_functions[fil](data, config): if mlflow_response_filter_functions[fil](data, config):
@@ -221,34 +210,36 @@ class Gates(BaseActivity):
comments.append(data['content']['message']) comments.append(data['content']['message'])
self.send_notification( self.send_notification(
metadata=metadata, 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'], message=data['content']['message'],
block="mlflow_gate", block='mlflow_gate',
level=NotificationLevel.ERROR, 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() trace = traceback.format_exc()
self.send_notification( self.send_notification(
metadata=metadata, metadata=metadata,
notification_id=f"MLFLOW_GATE_RESPONSE_FILTER__{fil}", notification_id=f'MLFLOW_GATE_RESPONSE_FILTER__{fil}',
message=f"Error in filter {fil}:{config}: \n {e}", message=f'Error in filter {fil}:{config}: \n {e}',
block="mlflow_gate", block='mlflow_gate',
level=NotificationLevel.ERROR, level=NotificationLevel.ERROR,
attachment_content=trace attachment_content=trace,
) )
for path_flag in path_priority: for path_flag in path_priority:
if path_flag in filter_output: if path_flag in filter_output:
self.info( self.info(f'Mlflow response gate result: {path_flag}', metadata)
f"Mlflow response gate result: {path_flag}", metadata) return (
return path_flag, mlflow_response_filter_functions['path_confidence'][path_flag], \ path_flag,
", ".join(comments) mlflow_response_filter_functions['path_confidence'][path_flag],
', '.join(comments),
)
self.info("Nothing was filtered by the mlflow response gate", metadata) self.info('Nothing was filtered by the mlflow response gate', metadata)
return None, 0, "" 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]: async def mlflow_content_gate(self, input_data: dict[str, Any]) -> tuple[str | None, int, str]:
""" """
Validate MLFlow prediction content quality and integrity. Validate MLFlow prediction content quality and integrity.
@@ -283,7 +274,7 @@ class Gates(BaseActivity):
Exception: If content validation fails or configuration is invalid Exception: If content validation fails or configuration is invalid
""" """
metadata = input_data['metadata'] metadata = input_data['metadata']
self.info("Performing mlflow content gate...", metadata) self.info('Performing mlflow content gate...', metadata)
filters = input_data['filters'] filters = input_data['filters']
data = DataFrame(input_data['data']) data = DataFrame(input_data['data'])
@@ -292,8 +283,8 @@ class Gates(BaseActivity):
filter_output = [] filter_output = []
self.debug(f"Input data:\n {data.head(5).to_string()}", metadata) 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'Filters: \n {create_sample_dict(filters)}', metadata)
for fil, config in filters.items(): for fil, config in filters.items():
if fil not in mlflow_content_filter_functions: if fil not in mlflow_content_filter_functions:
@@ -303,36 +294,38 @@ class Gates(BaseActivity):
filter_output.append(config['policy']) filter_output.append(config['policy'])
self.send_notification( self.send_notification(
metadata=metadata, metadata=metadata,
notification_id=f"{gate_type.upper()}_GATE_CONTENT_FILTER__{fil}", notification_id=f'{gate_type.upper()}_GATE_CONTENT_FILTER__{fil}',
message=f"Data not passed the content filter {fil}:{config}", message=f'Data not passed the content filter {fil}:{config}',
block="mlflow_gate", block='mlflow_gate',
level=NotificationLevel.WARNING, 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() trace = traceback.format_exc()
self.send_notification( self.send_notification(
metadata=metadata, metadata=metadata,
notification_id=f"MLFLOW_GATE_CONTENT_FILTER__{fil}", notification_id=f'MLFLOW_GATE_CONTENT_FILTER__{fil}',
message=f"Error in filter {fil}:{config}: \n {e}", message=f'Error in filter {fil}:{config}: \n {e}',
block="mlflow_gate", block='mlflow_gate',
level=NotificationLevel.ERROR, level=NotificationLevel.ERROR,
attachment_content=trace attachment_content=trace,
) )
for path_flag in path_priority: for path_flag in path_priority:
if path_flag in filter_output: if path_flag in filter_output:
self.info( self.info(f'Mlflow content gate result: {path_flag}', metadata)
f"Mlflow content gate result: {path_flag}", metadata) return (
return path_flag, mlflow_content_filter_functions['path_confidence'][path_flag], \ path_flag,
"Transformed data not passed the content filter" 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) self.info('Nothing was filtered by the mlflow content gate', metadata)
return None, 0, "" return None, 0, ''
def get_prediction_store_policy(self, def get_prediction_store_policy(
prediction_store_policy: str, self, prediction_store_policy: str, metadata: dict[str, Any]
metadata: dict[str, Any]) -> tuple[str, int]: ) -> tuple[str, int]:
""" """
Parse and validate prediction store policy configuration. Parse and validate prediction store policy configuration.
@@ -358,7 +351,9 @@ class Gates(BaseActivity):
if len(policy_elements) < 2: if len(policy_elements) < 2:
self.error( 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 'lts', 1
policy_type = policy_elements[0] 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 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 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( 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 'lts', 1
return policy_type, int(policy_value) 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]: async def format_prediction(self, input_data: dict[str, Any]) -> dict[Any, Any]:
""" """
Format prediction data according to configured storage policies. Format prediction data according to configured storage policies.
@@ -400,7 +401,7 @@ class Gates(BaseActivity):
""" """
metadata = input_data['metadata'] metadata = input_data['metadata']
prediction_store_policy = input_data['prediction_store_policy'] prediction_store_policy = input_data['prediction_store_policy']
self.info("Formatting prediction...", metadata) self.info('Formatting prediction...', metadata)
data = DataFrame(input_data['data']) data = DataFrame(input_data['data'])
@@ -408,48 +409,45 @@ class Gates(BaseActivity):
data['timestamp'] = data.index data['timestamp'] = data.index
data = data.reset_index(drop=True) data = data.reset_index(drop=True)
self.debug( self.debug(f'Prediction store policy: {prediction_store_policy}', metadata)
f"Prediction store policy: {prediction_store_policy}", metadata) self.debug(f'Prediction data: {data.head(5).to_string()}', metadata)
self.debug(f"Prediction data: {data.head(5).to_string()}", metadata)
policy_type, policy_value = self.get_prediction_store_policy( 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 # If data has no timestamp, we use the default timestamp and not sort the data
self.info( 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 is lts, we need to sort the data by timestamp descending and take the first policy_value rows
if policy_type == 'lts': if policy_type == 'lts':
self.debug( self.debug('Sorting data by timestamp descending', metadata)
"Sorting data by timestamp descending", metadata)
data = data.sort_values(by='timestamp', ascending=False) 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 # 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': elif policy_type == 'erl':
self.debug( self.debug('Sorting data by timestamp ascending', metadata)
"Sorting data by timestamp ascending", metadata)
data = data.sort_values(by='timestamp', ascending=True) data = data.sort_values(by='timestamp', ascending=True)
else: else:
self.error( self.error(f'Invalid policy type: {policy_type}, using default policy', metadata)
f"Invalid policy type: {policy_type}, using default policy", metadata) raise ValueError(f'Invalid policy type: {policy_type}')
raise ValueError(
f"Invalid policy type: {policy_type}")
data = data.head(int(policy_value)) data = data.head(int(policy_value))
data['model_id'] = input_data['model_id'] data['model_id'] = input_data['model_id']
data['prediction_confidence'] = input_data['prediction_confidence'] data['prediction_confidence'] = input_data['prediction_confidence']
data['prediction_status'] = 'Good' data['prediction_status'] = 'Good'
data['comments'] = "" data['comments'] = ''
data = data.sort_values(by='timestamp', ascending=False) data = data.sort_values(by='timestamp', ascending=False)
data = data.reset_index(drop=True) data = data.reset_index(drop=True)
self.info(f"Prediction formatted: {len(data)} rows", metadata) self.info(f'Prediction formatted: {len(data)} rows', metadata)
self.debug(f"Prediction data: {data.head(5).to_string()}", metadata) self.debug(f'Prediction data: {data.head(5).to_string()}', metadata)
return data.to_dict() 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]: async def format_default_prediction(self, input_data: dict[str, Any]) -> dict[Any, Any]:
""" """
Create and format default prediction data for error conditions. Create and format default prediction data for error conditions.
@@ -477,22 +475,24 @@ class Gates(BaseActivity):
""" """
metadata = input_data['metadata'] 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], 'prediction': [0],
'timestamp': [input_data['timestamp']], 'response_time': [0],
'model_id': [input_data['model_id']], 'timestamp': [input_data['timestamp']],
'prediction_confidence': [input_data['prediction_confidence']], 'model_id': [input_data['model_id']],
'prediction_status': ['Bad'], 'prediction_confidence': [input_data['prediction_confidence']],
'comments': [input_data['comment']] 'prediction_status': ['Bad'],
}) 'comments': [input_data['comment']],
}
)
self.info(f"Default prediction formatted: {data.size} rows", metadata) self.info(f'Default prediction formatted: {data.size} rows', metadata)
return data.to_dict() 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: async def get_last_timestamp(self, input_data: dict[str, Any]) -> str:
""" """
Extract the most recent timestamp from prediction data. Extract the most recent timestamp from prediction data.
@@ -517,24 +517,22 @@ class Gates(BaseActivity):
""" """
metadata = input_data['metadata'] metadata = input_data['metadata']
self.info("Getting last timestamp...", metadata) self.info('Getting last timestamp...', metadata)
data = DataFrame(input_data['data']) 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: if data.empty:
return now().strftime(DATETIME_FORMAT_WITH_TZ) return now().strftime(DATETIME_FORMAT_WITH_TZ)
max_timestamp = max( max_timestamp = max(data['timestamp'].values.tolist())
data['timestamp'].values.tolist())
self.info( self.info(f'Last timestamp: {max_timestamp}', metadata)
f"Last timestamp: {max_timestamp}", metadata)
return max_timestamp return max_timestamp
@activity.defn(name="write_metrics") @activity.defn(name='write_metrics')
async def write_metrics(self, input_data: dict[str, Any]): async def write_metrics(self, input_data: dict[str, Any]):
""" """
Write prediction performance metrics to Prometheus monitoring system. Write prediction performance metrics to Prometheus monitoring system.
@@ -562,26 +560,24 @@ class Gates(BaseActivity):
prediction_confidence = prediction['prediction_confidence'].values[0] prediction_confidence = prediction['prediction_confidence'].values[0]
response_time = prediction['response_time'].values[0] response_time = prediction['response_time'].values[0]
self.info( self.info(f'Writing metrics for model {metadata["model_name"]}', metadata)
f"Writing metrics for model {metadata['model_name']}", metadata)
metrics.PREDICTIONS_WRITTEN_COUNT.labels( metrics.PREDICTIONS_WRITTEN_COUNT.labels(
pod_id=self.pod_id, pod_id=self.pod_id,
model_name=metadata['model_name'], model_name=metadata['model_name'],
pipeline_name=metadata['workflow_name'] pipeline_name=metadata['workflow_name'],
).inc() ).inc()
metrics.PREDICTION_CONFIDENCE_MONITOR.labels( metrics.PREDICTION_CONFIDENCE_MONITOR.labels(
pod_id=self.pod_id, pod_id=self.pod_id,
model_name=metadata['model_name'], model_name=metadata['model_name'],
pipeline_name=metadata['workflow_name'] pipeline_name=metadata['workflow_name'],
).set(prediction_confidence) ).set(prediction_confidence)
metrics.PREDICTION_RESPONSE_TIME_MONITOR.labels( metrics.PREDICTION_RESPONSE_TIME_MONITOR.labels(
pod_id=self.pod_id, pod_id=self.pod_id,
model_name=metadata['model_name'], model_name=metadata['model_name'],
pipeline_name=metadata['workflow_name'] pipeline_name=metadata['workflow_name'],
).observe(response_time) ).observe(response_time)
self.info( self.info(f'Metrics written for model {metadata["model_name"]}', metadata)
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 from temporalio import activity, workflow
with workflow.unsafe.imports_passed_through(): with workflow.unsafe.imports_passed_through():
from datetime import datetime import traceback
from pandas import Timestamp, to_datetime from typing import Any
from sientia_do.temporal.constants import DATETIME_FORMAT, DATETIME_FORMAT_WITH_TZ
from sientia_do.temporal.activities.base import BaseActivity 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.handlers import CoreNotificationHandler as NotificationHandler
from sientia_do.notifications.models import NotificationLevel from sientia_do.notifications.models import NotificationLevel
from sientia_do.observability.logger import Logger 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 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): class MLFlow(BaseActivity):
@@ -36,8 +35,15 @@ class MLFlow(BaseActivity):
model_monitoring_repository (MLFlowRepository): Repository for MLFlow operations model_monitoring_repository (MLFlowRepository): Repository for MLFlow operations
""" """
def __init__(self, mlflow_host: str, mlflow_port: int, mlflow_username: str, def __init__(
mlflow_password: str, logger: Logger, notification_handler: NotificationHandler): self,
mlflow_host: str,
mlflow_port: int,
mlflow_username: str,
mlflow_password: str,
logger: Logger,
notification_handler: NotificationHandler,
):
""" """
Initialize MLFlow activities with server configuration. Initialize MLFlow activities with server configuration.
@@ -52,18 +58,17 @@ class MLFlow(BaseActivity):
Raises: Raises:
Exception: If MLFlowRepository initialization fails Exception: If MLFlowRepository initialization fails
""" """
BaseActivity.__init__( BaseActivity.__init__(self, logger, notification_handler, set_error_counter=True)
self, logger, notification_handler, set_error_counter=True)
self.mlflow_host = mlflow_host self.mlflow_host = mlflow_host
self.mlflow_port = mlflow_port self.mlflow_port = mlflow_port
self.mlflow_username = mlflow_username self.mlflow_username = mlflow_username
self.mlflow_password = mlflow_password self.mlflow_password = mlflow_password
self.model_monitoring_repository = MLFlowRepository( 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]: async def request_transform(self, input_data: dict[str, Any]) -> dict[str, Any]:
""" """
Transform input data using MLFlow models. Transform input data using MLFlow models.
@@ -99,7 +104,7 @@ class MLFlow(BaseActivity):
model_name = input_data['model_name'] model_name = input_data['model_name']
model_config = input_data.get('model_config', {}) 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) self.debug(data.head(5).to_string(), metadata)
# Sort by created_at in descending order and keep first occurrence of each variable/timestamp pair # 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 # Pivot data for model input format
data = data.pivot( data = data.pivot(index='timestamp', columns='variable', values='value')
index='timestamp', columns='variable',
values='value')
data.fillna(np.nan, inplace=True) data.fillna(np.nan, inplace=True)
# data.reset_index(inplace=True) # data.reset_index(inplace=True)
data.columns.name = None data.columns.name = None
self.debug("Processed input data:", metadata) self.debug('Processed input data:', metadata)
self.debug(data.head(5).to_string(), metadata) self.debug(data.head(5).to_string(), metadata)
# Request transformation from MLFlow model # Request transformation from MLFlow model
@@ -124,16 +127,20 @@ class MLFlow(BaseActivity):
) )
self.debug( 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( 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 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]: async def request_predict(self, input_data: dict[str, Any]) -> dict[str, Any]:
""" """
Execute predictions using MLFlow models. Execute predictions using MLFlow models.
@@ -169,14 +176,15 @@ class MLFlow(BaseActivity):
model_name = input_data['model_name'] model_name = input_data['model_name']
model_config = input_data.get('model_config', {}) 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 # Convert numpy.nan to None for model compatibility
data.replace(np.nan, None, inplace=True) data.replace(np.nan, None, inplace=True)
data['timestamp'] = data.index data['timestamp'] = data.index
data['timestamp'] = to_datetime( 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 # Request prediction from MLFlow model
response_data = self.model_monitoring_repository.predict( response_data = self.model_monitoring_repository.predict(
@@ -184,13 +192,15 @@ class MLFlow(BaseActivity):
) )
self.debug( 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 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]: async def retrain_model(self, input_data: dict[str, Any]) -> dict[str, Any]:
""" """
Retrain MLFlow models with updated training data. 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=['model_id'], inplace=True, errors='ignore')
data.drop(columns=['created_at'], inplace=True, errors='ignore') data.drop(columns=['created_at'], inplace=True, errors='ignore')
data = data.pivot(index='timestamp', columns='variable', data = data.pivot(index='timestamp', columns='variable', values='value')
values='value')
data.sort_index(inplace=True) data.sort_index(inplace=True)
data.reset_index(inplace=True) data.reset_index(inplace=True)
@@ -244,15 +253,10 @@ class MLFlow(BaseActivity):
try: try:
retrain_output, experiment = self.model_monitoring_repository.retrain_model( retrain_output, experiment = self.model_monitoring_repository.retrain_model(
data=data, data=data, model_name=model_name
model_name=model_name
) )
return { return {'status': retrain_output, 'timestamp': timestamp, 'experiment': experiment}
'status': retrain_output,
'timestamp': timestamp,
'experiment': experiment
}
except Exception as e: except Exception as e:
trace = traceback.format_exc() trace = traceback.format_exc()
self.send_notification( self.send_notification(
@@ -261,12 +265,12 @@ class MLFlow(BaseActivity):
message=f'Error retraining model {model_name}: {e}', message=f'Error retraining model {model_name}: {e}',
block='retrain_model', block='retrain_model',
level=NotificationLevel.ERROR, level=NotificationLevel.ERROR,
attachment_content=trace attachment_content=trace,
) )
self.error(trace, metadata=metadata) self.error(trace, metadata=metadata)
raise e 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]: async def update_production_model(self, input_data: dict[str, Any]) -> dict[Any, Any]:
""" """
Update production model with newly trained model version. Update production model with newly trained model version.
@@ -311,12 +315,12 @@ class MLFlow(BaseActivity):
status = input_data['status'] status = input_data['status']
self.info( self.info(
f'Updating production model {model_name} from experiment {experiment}...', metadata) f'Updating production model {model_name} from experiment {experiment}...', metadata
)
try: try:
response = self.model_monitoring_repository.update_production_model( response = self.model_monitoring_repository.update_production_model(
experiment=experiment, experiment=experiment, model_name=model_name
model_name=model_name
) )
report = DataFrame([response]) report = DataFrame([response])
@@ -325,8 +329,7 @@ class MLFlow(BaseActivity):
report['timestamp'] = timestamp report['timestamp'] = timestamp
report['status'] = status report['status'] = status
self.info( self.info(f'Production model {model_name} updated successfully', metadata)
f'Production model {model_name} updated successfully', metadata)
return report.to_dict() return report.to_dict()
except Exception as e: except Exception as e:
@@ -337,7 +340,7 @@ class MLFlow(BaseActivity):
message=f'Error updating production model {model_name}: {e}', message=f'Error updating production model {model_name}: {e}',
block='update_production_model', block='update_production_model',
level=NotificationLevel.ERROR, level=NotificationLevel.ERROR,
attachment_content=trace attachment_content=trace,
) )
self.error(trace, metadata=metadata) self.error(trace, metadata=metadata)
raise e raise e

24
model_manager/metrics.py
View File

@@ -22,36 +22,36 @@ Metric Labels:
- pipeline_name: Name of the prediction pipeline - pipeline_name: Name of the prediction pipeline
""" """
from prometheus_client import Gauge, Counter, Histogram from prometheus_client import Counter, Gauge, Histogram
# Application health metric # Application health metric
APP_UP = Gauge( APP_UP = Gauge(
"app_up", 'app_up',
"Indicates if the application is running (1) or shutting down (0)", 'Indicates if the application is running (1) or shutting down (0)',
["pod_id"], ['pod_id'],
) )
# Core labels used across multiple metrics # 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 # Prediction operation metrics
PREDICTIONS_WRITTEN_COUNT = Counter( PREDICTIONS_WRITTEN_COUNT = Counter(
"model_manager_predictions_written_count", 'model_manager_predictions_written_count',
"Number of predictions written to the database table predictions", 'Number of predictions written to the database table predictions',
CORE_LABELS, CORE_LABELS,
) )
# Prediction quality metrics # Prediction quality metrics
PREDICTION_CONFIDENCE_MONITOR = Gauge( PREDICTION_CONFIDENCE_MONITOR = Gauge(
"model_manager_prediction_confidence_monitor", 'model_manager_prediction_confidence_monitor',
"Current confidence of each prediction", 'Current confidence of each prediction',
CORE_LABELS, CORE_LABELS,
) )
# Performance monitoring metrics # Performance monitoring metrics
PREDICTION_RESPONSE_TIME_MONITOR = Histogram( PREDICTION_RESPONSE_TIME_MONITOR = Histogram(
"model_manager_prediction_response_time_monitor", 'model_manager_prediction_response_time_monitor',
"Current response time of each prediction", 'Current response time of each prediction',
CORE_LABELS, 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 from os import getenv
import json from typing import Any
from typing import Dict, Any
def build_postgres_config() -> Dict[str, Any]: def build_postgres_config() -> dict[str, Any]:
""" """
Build PostgreSQL database configuration from environment variables. Build PostgreSQL database configuration from environment variables.
@@ -30,11 +29,11 @@ def build_postgres_config() -> Dict[str, Any]:
'password': getenv('POSTGRES_PASSWORD', 'sientia'), 'password': getenv('POSTGRES_PASSWORD', 'sientia'),
'dbname': getenv('POSTGRES_DBNAME', 'sientia'), 'dbname': getenv('POSTGRES_DBNAME', 'sientia'),
'min_connections': int(getenv('POSTGRES_MIN_CONNECTIONS', '5')), '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. Build MLFlow server configuration from environment variables.
@@ -55,11 +54,11 @@ def build_mlflow_config() -> Dict[str, Any]:
'host': getenv('MLFLOW_HOST', 'http://localhost'), 'host': getenv('MLFLOW_HOST', 'http://localhost'),
'port': int(getenv('MLFLOW_PORT', '5080')), 'port': int(getenv('MLFLOW_PORT', '5080')),
'username': getenv('MLFLOW_USERNAME', 'aignosi'), '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. Build MongoDB configuration from environment variables.
@@ -86,5 +85,5 @@ def build_mongodb_config() -> Dict[str, Any]:
return { return {
'connection_string': connection_string, 'connection_string': connection_string,
'database_name': getenv('MONGODB_DATABASE_NAME', 'sientia'), '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. False if none of the specified variables contain null values.
""" """
return not data[ return not data[data['variable'].isin(config['variables']) & data['value'].isna()].empty
data['variable'].isin(config['variables']) & data['value'].isna()].empty
def filter_empty_data(data: DataFrame, _config: dict) -> bool: 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 bool: True if data should be filtered (too many NaN values), False otherwise
""" """
data = predictions.replace({None: np.nan}).drop( data = (
columns=['timestamp'], errors='ignore').infer_objects() predictions.replace({None: np.nan})
.drop(columns=['timestamp'], errors='ignore')
.infer_objects()
)
if data.isna().all().all(): if data.isna().all().all():
return True 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. By Monitoring we mean the evaluation of the performance of models, the generation of reports.
""" """
from datetime import datetime
import traceback import traceback
import pandas as pd from datetime import datetime
import mlflow
from os import makedirs, path, remove from os import makedirs, path, remove
import mlflow
import pandas as pd
from sientia.ModelServing import ModelServing 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.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): def __init__(self, host, username, password, logger: Logger):
self.model_serving = ModelServing(
self.model_serving = ModelServing(tracking_uri=host, tracking_uri=host, username=username, password=password, logger=logger
username=username, password=password, )
logger=logger)
self.logger = logger self.logger = logger
def detect_and_parse_datetime_index(self, data: pd.DataFrame, metadata: dict) -> pd.DataFrame: 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 # Get type of first element of index
index_type = type(index[0]) 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 # Check if all in index are of the same type
if not all(isinstance(i, index_type) for i in index): if not all(isinstance(i, index_type) for i in index):
raise ValueError( raise ValueError(f'{message}')
f"{message}")
# Check type and converts to DATETIME_FORMAT_WITH_TZ # 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 # Validate format of string and return error if not valid
try: try:
pd.to_datetime(data.index, format=DATETIME_FORMAT_WITH_TZ) pd.to_datetime(data.index, format=DATETIME_FORMAT_WITH_TZ)
except ValueError: except ValueError as e:
raise ValueError( raise ValueError(f'{message}') from e
f"{message}")
elif index_type == datetime or index_type == pd.Timestamp: elif index_type == datetime or index_type == pd.Timestamp:
data.index = data.index.strftime(DATETIME_FORMAT_WITH_TZ) data.index = data.index.strftime(DATETIME_FORMAT_WITH_TZ)
else: else:
raise ValueError( raise ValueError(f'{message}')
f"{message}")
return data 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. Transform data using a model.
@@ -81,41 +81,42 @@ class MLFlowRepository():
try: try:
self.logger.custom_debug( 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) model_retention = model_config.get('retention_minutes', 0)
flavor = model_config.get('transform_flavor', 'sklearn') flavor = model_config.get('transform_flavor', 'sklearn')
compressed = model_config.get('is_compressed', False) compressed = model_config.get('is_compressed', False)
retention_target = model_config.get('retention_target', 'model') retention_target = model_config.get('retention_target', 'model')
transform_keyword = model_config.get( transform_keyword = model_config.get('transform_function_keyword', 'predict')
'transform_function_keyword', 'predict')
transformed_data = self.model_serving.get_cached_transform( transformed_data = self.model_serving.get_cached_transform(
model_name, data, model_retention, flavor, model_name,
compressed, retention_target, transform_keyword data,
model_retention,
flavor,
compressed,
retention_target,
transform_keyword,
) )
self.logger.custom_debug( 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 = self.detect_and_parse_datetime_index(transformed_data, metadata)
transformed_data, metadata)
return { return {'success': True, 'content': transformed_data.to_dict()}
'success': True,
'content': transformed_data.to_dict()
}
except Exception as e: except Exception as e: # noqa: BLE001
return { return {
'success': False, 'success': False,
'content': { 'content': {'message': str(e), 'traceback': traceback.format_exc()},
'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. Predict data using a model.
@@ -137,31 +138,26 @@ class MLFlowRepository():
start_time = datetime.now() start_time = datetime.now()
self.logger.custom_debug( 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( data = self.model_serving.get_cached_predict(
model_name, data, model_retention, flavor, model_name, data, model_retention, flavor, compressed, retention_target
compressed, retention_target
) )
end_time = datetime.now() end_time = datetime.now()
data = pd.DataFrame(data, columns=['prediction']) data = pd.DataFrame(data, columns=['prediction'])
self.logger.custom_debug( 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.index = input_index
data['response_time'] = (end_time - start_time).total_seconds() data['response_time'] = (end_time - start_time).total_seconds()
return { return {'success': True, 'content': data.to_dict()}
'success': True,
'content': data.to_dict()
}
except Exception as e: except Exception as e: # noqa: BLE001
return { return {
'success': False, 'success': False,
'content': { 'content': {'message': str(e), 'traceback': traceback.format_exc()},
'message': str(e),
'traceback': traceback.format_exc()
}
} }
def get_experiment_by_run_id(self, run_id: str) -> dict: def get_experiment_by_run_id(self, run_id: str) -> dict:
@@ -190,10 +186,9 @@ class MLFlowRepository():
Returns: Returns:
str: The next run name in format 'model_name-run_number' str: The next run name in format 'model_name-run_number'
""" """
runs = mlflow.search_runs( runs = mlflow.search_runs(experiment_names=[model_name], order_by=['start_time desc'])
experiment_names=[model_name], order_by=["start_time desc"])
next_run_number = len(runs) + 1 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: def create_model_experiment(self, model_name: str, data: pd.DataFrame) -> tuple:
""" """
@@ -217,14 +212,10 @@ class MLFlowRepository():
- experiment: MLFlow experiment name - experiment: MLFlow experiment name
""" """
# load predictor model # load predictor model
predictor_uri = f"models:/{model_name}/production" predictor_uri = f'models:/{model_name}/production'
# load transform model # load transform model
latest_production_id = self.model_serving.get_model_run_id( latest_production_id = self.model_serving.get_model_run_id(model_name, stage='Production')
model_name, stage="Production" transform_uri = self.model_serving.get_model_uri(latest_production_id, prediction=False)
)
transform_uri = self.model_serving.get_model_uri(
latest_production_id, prediction=False
)
# load # load
data_model = mlflow.sklearn.load_model(transform_uri) data_model = mlflow.sklearn.load_model(transform_uri)
prediction_model = mlflow.sklearn.load_model(predictor_uri) prediction_model = mlflow.sklearn.load_model(predictor_uri)
@@ -233,20 +224,16 @@ class MLFlowRepository():
target_name = data_model.target_variable target_name = data_model.target_variable
y = data[target_name] y = data[target_name]
treated_data = pd.merge( treated_data = pd.merge(treated_data, y, left_index=True, right_index=True)
treated_data, y, left_index=True, right_index=True)
prediction_model = prediction_model.fit(treated_data) prediction_model = prediction_model.fit(treated_data)
experiment = self.get_experiment_by_run_id(latest_production_id) experiment = self.get_experiment_by_run_id(latest_production_id)
mlflow.set_experiment(experiment) mlflow.set_experiment(experiment)
return prediction_model, data_model, experiment return prediction_model, data_model, experiment
def perform_model_retrain(self, def perform_model_retrain(
prediction_model, self, prediction_model, data_model, experiment: str, model_name: str, data: pd.DataFrame
data_model, ):
experiment: str,
model_name: str,
data: pd.DataFrame):
""" """
Execute the complete model retraining process in MLFlow. Execute the complete model retraining process in MLFlow.
@@ -271,7 +258,7 @@ class MLFlowRepository():
""" """
pred_model_atributes = vars(prediction_model) # load class attributes pred_model_atributes = vars(prediction_model) # load class attributes
data_model_atributes = vars(data_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) current_run_name = self.get_next_run_name(experiment)
with mlflow.start_run( with mlflow.start_run(
run_name=current_run_name, description=experiment_description run_name=current_run_name, description=experiment_description
@@ -279,32 +266,32 @@ class MLFlowRepository():
# update transfomation model # update transfomation model
# fixed parameters # fixed parameters
for name_atribute, val_atribute in pred_model_atributes.items(): 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) mlflow.log_param(name_atribute, val_atribute)
# update prediction model # update prediction model
for name_atribute, val_atribute in data_model_atributes.items(): 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) mlflow.log_param(name_atribute, val_atribute)
# dynamic parameters, including model itself # 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) data.to_csv(file_path, index=True)
# log the data raw # log the data raw
mlflow.log_artifact(file_path) mlflow.log_artifact(file_path)
# dynamic parameters, including model itself # dynamic parameters, including model itself
mlflow.sklearn.log_model(prediction_model, "prediction_model") mlflow.sklearn.log_model(prediction_model, 'prediction_model')
mlflow.log_param("retrain", True) mlflow.log_param('retrain', True)
# clear temp file # clear temp file
if path.exists(file_path): if path.exists(file_path):
remove(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: def retrain_model(self, data: pd.DataFrame, model_name: str) -> tuple:
""" """
@@ -325,10 +312,10 @@ class MLFlowRepository():
- status_message (str): Retraining operation status - status_message (str): Retraining operation status
- experiment_name (str): MLFlow experiment identifier - experiment_name (str): MLFlow experiment identifier
""" """
prediction_model, data_model, experiment = self.create_model_experiment( prediction_model, data_model, experiment = self.create_model_experiment(model_name, data)
model_name, data)
retrain_result = self.perform_model_retrain( 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 return retrain_result
def get_experiment(self, experiment_name: str) -> int: def get_experiment(self, experiment_name: str) -> int:
@@ -374,22 +361,21 @@ class MLFlowRepository():
""" """
runs = mlflow.search_runs( runs = mlflow.search_runs(
experiment_ids=[experiment_id], experiment_ids=[experiment_id],
filter_string="", # Sem filtro no MLflow ainda filter_string='', # Sem filtro no MLflow ainda
output_format="pandas" output_format='pandas',
) )
if not isinstance(runs, pd.DataFrame): if not isinstance(runs, pd.DataFrame):
raise ValueError('Runs is not a pandas DataFrame') raise ValueError('Runs is not a pandas DataFrame')
# Filtrar apenas as runs onde params.retrain == True # 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 # Converter a coluna 'end_time' para datetime
filtered_runs['end_time'] = pd.to_datetime(filtered_runs['end_time']) filtered_runs['end_time'] = pd.to_datetime(filtered_runs['end_time'])
# Ordenar o DataFrame de forma descendente pela coluna 'end_time' # Ordenar o DataFrame de forma descendente pela coluna 'end_time'
filtered_runs = filtered_runs.sort_values( filtered_runs = filtered_runs.sort_values(by='end_time', ascending=False)
by='end_time', ascending=False)
# Pegar a última run_id do DataFrame filtrado e ordenado # Pegar a última run_id do DataFrame filtrado e ordenado
latest_run_id = filtered_runs.iloc[0]['run_id'] latest_run_id = filtered_runs.iloc[0]['run_id']
@@ -423,16 +409,14 @@ class MLFlowRepository():
# Registrar o modelo # Registrar o modelo
# Aqui estamos assumindo que você já tem um modelo salvo, caso contrário você precisará treiná-lo e salvá-lo primeiro. # 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. # Se o modelo já está registrado, você pode usar o método register_model() ou pyfunc.load_model() para isso.
mlflow.register_model( mlflow.register_model(f'runs:/{run_id}/prediction_model', model_name)
f"runs:/{run_id}/prediction_model", model_name)
# Colocar a versão do modelo em produção # 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' # 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() client = mlflow.tracking.MlflowClient()
# Obter a versão mais recente registrada do modelo # Obter a versão mais recente registrada do modelo
model_versions = client.get_registered_model( model_versions = client.get_registered_model(model_name).latest_versions
model_name).latest_versions
if not isinstance(model_versions, list): if not isinstance(model_versions, list):
raise ValueError('Model versions is not a 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' # Mover a versão mais recente do modelo para o estágio de 'Production'
client.transition_model_version_stage( client.transition_model_version_stage(
name=model_name, name=model_name, version=max_version, stage='Production', archive_existing_versions=True
version=max_version,
stage="Production",
archive_existing_versions=True
) )
return { return {'model_name': model_name, 'version': max_version, 'mlflow_run_id': run_id}
'model_name': model_name,
'version': max_version,
'mlflow_run_id': run_id
}
def update_production_model(self, experiment: str, model_name: str) -> dict: 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) - PROJECT_NAME: Project name for notifications (default: model_manager)
""" """
from temporalio import workflow, client from temporalio import client, workflow
from temporalio.worker import Worker, PollerBehaviorAutoscaling from temporalio.runtime import PrometheusConfig, Runtime, TelemetryConfig
from temporalio.runtime import Runtime, TelemetryConfig, PrometheusConfig from temporalio.worker import PollerBehaviorAutoscaling, Worker
with workflow.unsafe.imports_passed_through(): with workflow.unsafe.imports_passed_through():
import asyncio
import os import os
import sys import sys
import asyncio
from model_manager.workflows.minimal_retrain import MinimalRetrain from prometheus_client import start_http_server
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 sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler
from sientia_do.observability.logger import get_logger from sientia_do.observability.logger import get_logger
from model_manager import metrics 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') 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(): 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(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() start_prometheus_server()
logger.custom_info('Starting Notification Handler...', metadata) logger.custom_info('Starting Notification Handler...', metadata)
@@ -95,7 +98,7 @@ async def main():
connection_string=mongo_config['connection_string'], connection_string=mongo_config['connection_string'],
database=mongo_config['database_name'], database=mongo_config['database_name'],
logger=logger, logger=logger,
project_name=os.getenv('PROJECT_NAME', 'model-manager') project_name=os.getenv('PROJECT_NAME', 'model-manager'),
) )
logger.custom_info('Starting Activities...', metadata) logger.custom_info('Starting Activities...', metadata)
@@ -104,16 +107,14 @@ async def main():
postgres_config=build_postgres_config(), postgres_config=build_postgres_config(),
mlflow_config=build_mlflow_config(), mlflow_config=build_mlflow_config(),
logger=logger, logger=logger,
notification_handler=notification_handler notification_handler=notification_handler,
) )
logger.custom_info( logger.custom_info(f'Starting SDK Metrics Server on port {SDK_METRICS_PORT}...', metadata)
f'Starting SDK Metrics Server on port {SDK_METRICS_PORT}...', metadata)
new_runtime = Runtime( new_runtime = Runtime(
telemetry=TelemetryConfig( telemetry=TelemetryConfig(
metrics=PrometheusConfig( metrics=PrometheusConfig(bind_address=f'0.0.0.0:{SDK_METRICS_PORT}')
bind_address=f"0.0.0.0:{SDK_METRICS_PORT}")
) )
) )
@@ -122,7 +123,7 @@ async def main():
temporal_client = await client.Client.connect( temporal_client = await client.Client.connect(
target_host=host, target_host=host,
namespace=os.getenv('TEMPORAL_NAMESPACE', 'model-manager'), namespace=os.getenv('TEMPORAL_NAMESPACE', 'model-manager'),
runtime=new_runtime runtime=new_runtime,
) )
logger.custom_info('Starting Workers...', metadata) logger.custom_info('Starting Workers...', metadata)
@@ -136,20 +137,19 @@ async def main():
activities.load_custom_query, activities.load_custom_query,
activities.retrain_model, activities.retrain_model,
activities.update_production_model, activities.update_production_model,
activities.export_data_to_postgres activities.export_data_to_postgres,
], ],
max_concurrent_workflow_tasks=50, max_concurrent_workflow_tasks=50,
max_concurrent_activities=50, max_concurrent_activities=50,
max_concurrent_local_activities=50, max_concurrent_local_activities=50,
max_cached_workflows=200, max_cached_workflows=200,
workflow_task_poller_behavior=PollerBehaviorAutoscaling(), workflow_task_poller_behavior=PollerBehaviorAutoscaling(),
activity_task_poller_behavior=PollerBehaviorAutoscaling() activity_task_poller_behavior=PollerBehaviorAutoscaling(),
), ),
Worker( Worker(
temporal_client, temporal_client,
task_queue='predictions_batch-queue', task_queue='predictions_batch-queue',
workflows=[PredictionsBatch, PredictionProcess, workflows=[PredictionsBatch, PredictionProcess, FormatAndExportPrediction],
FormatAndExportPrediction],
activities=[ activities=[
# MLFlow # MLFlow
activities.request_predict, activities.request_predict,
@@ -165,15 +165,15 @@ async def main():
activities.load_custom_query, activities.load_custom_query,
activities.repeat_last_prediction, activities.repeat_last_prediction,
activities.export_data_to_postgres, activities.export_data_to_postgres,
activities.write_metrics activities.write_metrics,
], ],
max_concurrent_workflow_tasks=50, max_concurrent_workflow_tasks=50,
max_concurrent_activities=50, max_concurrent_activities=50,
max_concurrent_local_activities=50, max_concurrent_local_activities=50,
max_cached_workflows=200, max_cached_workflows=200,
workflow_task_poller_behavior=PollerBehaviorAutoscaling(), workflow_task_poller_behavior=PollerBehaviorAutoscaling(),
activity_task_poller_behavior=PollerBehaviorAutoscaling() activity_task_poller_behavior=PollerBehaviorAutoscaling(),
) ),
] ]
handlers = [] handlers = []
@@ -186,8 +186,8 @@ async def main():
# This will run the workers and wait for them to complete. # 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. # If an exception occurs in any of the worker handlers, it will be propagated here.
await asyncio.gather(*handlers) await asyncio.gather(*handlers)
except BaseException as e: # NOSONAR except BaseException as e: # noqa: BLE001
logger.custom_error(f"An unhandled exception occurred: {e}", metadata) logger.custom_error(f'An unhandled exception occurred: {e}', metadata)
finally: finally:
if notification_handler: if notification_handler:
notification_handler.shutdown() notification_handler.shutdown()
@@ -216,12 +216,12 @@ def start_prometheus_server():
SystemExit: If the metrics server fails to start SystemExit: If the metrics server fails to start
""" """
try: try:
port = int(os.getenv("HTTP_METRICS_PORT", 9090)) port = int(os.getenv('HTTP_METRICS_PORT', 9090))
start_http_server(port) 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 metrics.APP_UP.labels(pod_id=POD_ID).set(1) # Mark app as UP
except Exception as e: except Exception as e: # noqa: BLE001
print(f"Failed to start Prometheus server: {e}") print(f'Failed to start Prometheus server: {e}')
os._exit(1) os._exit(1)

34
model_manager/workflows/minimal_retrain.py
View File

@@ -1,14 +1,16 @@
from temporalio import workflow from temporalio import workflow
with workflow.unsafe.imports_passed_through(): 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 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") @workflow.defn(name='minimal_retrain')
class MinimalRetrain(): class MinimalRetrain:
""" """
Automated model retraining workflow for the Model Manager system. Automated model retraining workflow for the Model Manager system.
@@ -63,7 +65,7 @@ class MinimalRetrain():
'schedule_name': input_data['schedule_name'], 'schedule_name': input_data['schedule_name'],
'model_name': input_data['model_name'], 'model_name': input_data['model_name'],
'model_id': input_data['model_id'], 'model_id': input_data['model_id'],
'workflow_name': 'minimal_retrain' 'workflow_name': 'minimal_retrain',
} }
} }
@@ -74,21 +76,17 @@ class MinimalRetrain():
{ {
**metadata, **metadata,
'query': input_data['query'], 'query': input_data['query'],
'datetime_columns': input_data.get('datetime_columns', []) 'datetime_columns': input_data.get('datetime_columns', []),
}, },
retry_policy=retry_policy, 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( experiment_response = await workflow.execute_activity_method(
Activities.retrain_model, Activities.retrain_model,
{ {**metadata, 'data': data, 'model_name': model_name},
**metadata,
'data': data,
'model_name': model_name
},
retry_policy=retry_policy, retry_policy=retry_policy,
start_to_close_timeout=timedelta(seconds=60) start_to_close_timeout=timedelta(seconds=60),
) )
report = await workflow.execute_activity_method( report = await workflow.execute_activity_method(
@@ -97,10 +95,10 @@ class MinimalRetrain():
**metadata, **metadata,
'model_name': model_name, 'model_name': model_name,
'model_id': input_data['model_id'], 'model_id': input_data['model_id'],
**experiment_response **experiment_response,
}, },
retry_policy=retry_policy, retry_policy=retry_policy,
start_to_close_timeout=timedelta(seconds=60) start_to_close_timeout=timedelta(seconds=60),
) )
await workflow.execute_activity_method( await workflow.execute_activity_method(
@@ -109,8 +107,8 @@ class MinimalRetrain():
**metadata, **metadata,
'data': report, 'data': report,
'schema': input_data['schema'], 'schema': input_data['schema'],
'table_name': input_data['table_name'] 'table_name': input_data['table_name'],
}, },
retry_policy=retry_policy, 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 from temporalio import workflow
with workflow.unsafe.imports_passed_through(): 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 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") @workflow.defn(name='predictions_batch')
class PredictionsBatch(): class PredictionsBatch:
""" """
Main batch prediction workflow for the Model Manager system. Main batch prediction workflow for the Model Manager system.
@@ -74,7 +76,7 @@ class PredictionsBatch():
'schedule_name': input_data['schedule_name'], 'schedule_name': input_data['schedule_name'],
'model_name': input_data['model_name'], 'model_name': input_data['model_name'],
'model_id': input_data['model_id'], 'model_id': input_data['model_id'],
'workflow_name': 'predictions_batch' 'workflow_name': 'predictions_batch',
} }
} }
@@ -84,10 +86,10 @@ class PredictionsBatch():
{ {
**metadata, **metadata,
'query': input_data['query'], 'query': input_data['query'],
'datetime_columns': input_data.get('datetime_columns', []) 'datetime_columns': input_data.get('datetime_columns', []),
}, },
retry_policy=retry_policy, retry_policy=retry_policy,
start_to_close_timeout=timedelta(seconds=300) start_to_close_timeout=timedelta(seconds=300),
) )
# Prepare input for prediction_process workflow # Prepare input for prediction_process workflow
@@ -98,28 +100,17 @@ class PredictionsBatch():
'table_name': input_data['table_name'], 'table_name': input_data['table_name'],
'model_id': input_data['model_id'], 'model_id': input_data['model_id'],
'model_name': input_data['model_name'], 'model_name': input_data['model_name'],
'input_filters': input_data.get('input_filters', { 'input_filters': input_data.get('input_filters', {'EMPTY_DATA': {'POLICY': 'STOP'}}),
'EMPTY_DATA': { 'mlflow_transform_filters': input_data.get(
'POLICY': 'STOP' 'mlflow_transform_filters', {'API_ERROR': {'POLICY': 'STOP'}}
} ),
}), 'mlflow_predict_filters': input_data.get(
'mlflow_transform_filters': input_data.get('mlflow_transform_filters', { 'mlflow_predict_filters', {'API_ERROR': {'POLICY': 'STOP'}}
'API_ERROR': { ),
'POLICY': 'STOP'
}
}),
'mlflow_predict_filters': input_data.get('mlflow_predict_filters', {
'API_ERROR': {
'POLICY': 'STOP'
}
}),
'model_config': input_data.get('model_config', {}), 'model_config': input_data.get('model_config', {}),
'path_priority': input_data.get('path_priority', ['STOP', 'CONTINUE', 'REPEAT']), '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 # Execute prediction process workflow
await workflow.execute_child_workflow( await workflow.execute_child_workflow('prediction_process', prediction_input)
'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 from temporalio import workflow
with workflow.unsafe.imports_passed_through(): with workflow.unsafe.imports_passed_through():
from model_manager.activities.activities import Activities
from typing import Any
from datetime import timedelta 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.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") @workflow.defn(name='format_and_export_prediction')
class FormatAndExportPrediction(): class FormatAndExportPrediction:
""" """
Data formatting and export workflow for prediction results. Data formatting and export workflow for prediction results.
@@ -75,10 +77,10 @@ class FormatAndExportPrediction():
'timestamp': input_data['timestamp'], 'timestamp': input_data['timestamp'],
'model_id': input_data['model_id'], 'model_id': input_data['model_id'],
'prediction_confidence': prediction_confidence, 'prediction_confidence': prediction_confidence,
'prediction_store_policy': input_data['prediction_store_policy'] 'prediction_store_policy': input_data['prediction_store_policy'],
}, },
retry_policy=retry_policy, retry_policy=retry_policy,
start_to_close_timeout=timedelta(seconds=60) start_to_close_timeout=timedelta(seconds=60),
) )
else: else:
@@ -90,10 +92,10 @@ class FormatAndExportPrediction():
'timestamp': input_data['timestamp'], 'timestamp': input_data['timestamp'],
'model_id': input_data['model_id'], 'model_id': input_data['model_id'],
'prediction_confidence': prediction_confidence, 'prediction_confidence': prediction_confidence,
'comment': input_data['comment'] 'comment': input_data['comment'],
}, },
retry_policy=retry_policy, retry_policy=retry_policy,
start_to_close_timeout=timedelta(seconds=60) start_to_close_timeout=timedelta(seconds=60),
) )
# write to postgres # write to postgres
@@ -104,21 +106,15 @@ class FormatAndExportPrediction():
'schema': input_data['schema'], 'schema': input_data['schema'],
'table_name': input_data['table_name'], 'table_name': input_data['table_name'],
'data': prediction, 'data': prediction,
'timestamp_conversion': { 'timestamp_conversion': {'column': 'timestamp', 'format': DATETIME_FORMAT_WITH_TZ},
'column': 'timestamp',
'format': DATETIME_FORMAT_WITH_TZ
}
}, },
retry_policy=retry_policy, retry_policy=retry_policy,
start_to_close_timeout=timedelta(seconds=60) start_to_close_timeout=timedelta(seconds=60),
) )
await workflow.execute_activity_method( await workflow.execute_activity_method(
Activities.write_metrics, Activities.write_metrics,
{ {**metadata, 'prediction': prediction},
**metadata,
'prediction': prediction
},
retry_policy=retry_policy, 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 from temporalio import workflow
with workflow.unsafe.imports_passed_through(): 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 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") @workflow.defn(name='prediction_process')
class PredictionProcess(): class PredictionProcess:
""" """
Core prediction processing workflow for the Model Manager system. Core prediction processing workflow for the Model Manager system.
@@ -84,10 +86,7 @@ class PredictionProcess():
# Get last timestamp for incremental processing # Get last timestamp for incremental processing
last_timestamp = await workflow.execute_local_activity_method( last_timestamp = await workflow.execute_local_activity_method(
Activities.get_last_timestamp, Activities.get_last_timestamp,
{ {**metadata, 'data': data},
**metadata,
'data': data
},
retry_policy=retry_policy, retry_policy=retry_policy,
start_to_close_timeout=timedelta(minutes=1), start_to_close_timeout=timedelta(minutes=1),
) )
@@ -97,7 +96,7 @@ class PredictionProcess():
**metadata, **metadata,
'filters': input_data['input_filters'], 'filters': input_data['input_filters'],
'data': data, 'data': data,
'path_priority': input_data['path_priority'] 'path_priority': input_data['path_priority'],
} }
path_flag, confidence, comment = await workflow.execute_local_activity_method( path_flag, confidence, comment = await workflow.execute_local_activity_method(
@@ -116,12 +115,7 @@ class PredictionProcess():
# Request MLFlow model transformation # Request MLFlow model transformation
response_data = await workflow.execute_local_activity_method( response_data = await workflow.execute_local_activity_method(
Activities.request_transform, 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, retry_policy=retry_policy,
start_to_close_timeout=timedelta(minutes=5), start_to_close_timeout=timedelta(minutes=5),
) )
@@ -134,7 +128,7 @@ class PredictionProcess():
'filters': input_data['mlflow_transform_filters'], 'filters': input_data['mlflow_transform_filters'],
'data': response_data, 'data': response_data,
'type': 'transform', 'type': 'transform',
'path_priority': input_data['path_priority'] 'path_priority': input_data['path_priority'],
}, },
retry_policy=retry_policy, retry_policy=retry_policy,
start_to_close_timeout=timedelta(minutes=1), start_to_close_timeout=timedelta(minutes=1),
@@ -155,7 +149,7 @@ class PredictionProcess():
'filters': input_data['mlflow_transform_filters'], 'filters': input_data['mlflow_transform_filters'],
'data': transformed_data, 'data': transformed_data,
'type': 'transform', 'type': 'transform',
'path_priority': input_data['path_priority'] 'path_priority': input_data['path_priority'],
}, },
retry_policy=retry_policy, retry_policy=retry_policy,
start_to_close_timeout=timedelta(minutes=1), start_to_close_timeout=timedelta(minutes=1),
@@ -172,7 +166,7 @@ class PredictionProcess():
**metadata, **metadata,
'data': transformed_data, 'data': transformed_data,
'model_name': model_name, 'model_name': model_name,
'model_config': model_config 'model_config': model_config,
}, },
retry_policy=retry_policy, retry_policy=retry_policy,
start_to_close_timeout=timedelta(minutes=5), start_to_close_timeout=timedelta(minutes=5),
@@ -186,7 +180,7 @@ class PredictionProcess():
'filters': input_data['mlflow_predict_filters'], 'filters': input_data['mlflow_predict_filters'],
'data': response_data, 'data': response_data,
'type': 'predict', 'type': 'predict',
'path_priority': input_data['path_priority'] 'path_priority': input_data['path_priority'],
}, },
retry_policy=retry_policy, retry_policy=retry_policy,
start_to_close_timeout=timedelta(minutes=1), start_to_close_timeout=timedelta(minutes=1),
@@ -213,12 +207,19 @@ class PredictionProcess():
'schema': input_data['schema'], 'schema': input_data['schema'],
'table_name': input_data['table_name'], 'table_name': input_data['table_name'],
'comment': comment, '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, async def path_flag_handler(
confidence: int, last_timestamp: str, comment: str) -> bool: 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. Handle path decisions based on filter results and confidence levels.
@@ -264,7 +265,7 @@ class PredictionProcess():
'schema': schema, 'schema': schema,
'table_name': table_name, 'table_name': table_name,
'model': model_id, 'model': model_id,
'last_timestamp': last_timestamp 'last_timestamp': last_timestamp,
}, },
retry_policy=retry_policy, retry_policy=retry_policy,
start_to_close_timeout=timedelta(minutes=1), start_to_close_timeout=timedelta(minutes=1),
@@ -286,8 +287,8 @@ class PredictionProcess():
'schema': schema, 'schema': schema,
'table_name': table_name, 'table_name': table_name,
'comment': comment, 'comment': comment,
'prediction_store_policy': input_data['prediction_store_policy'] 'prediction_store_policy': input_data['prediction_store_policy'],
} },
) )
return True 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 pytest import mark
from unittest.mock import patch, MagicMock, ANY
from sientia_do.temporal.activities.postgres import Postgres from sientia_do.temporal.activities.postgres import Postgres
from model_manager.activities.activities import Activities 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.gates import Gates
from model_manager.activities.mlflow import MLFlow
@patch('model_manager.activities.activities.Postgres.__init__') @patch('model_manager.activities.activities.Postgres.__init__')
@patch('model_manager.activities.activities.MLFlow.__init__') @patch('model_manager.activities.activities.MLFlow.__init__')
@patch('model_manager.activities.activities.Gates.__init__') @patch('model_manager.activities.activities.Gates.__init__')
def test___init__(mock_gates_init, mock_mlflow_init, mock_postgres_init): def test___init__(mock_gates_init, mock_mlflow_init, mock_postgres_init):
postgres_config = { postgres_config = {
'host': 'localhost', 'host': 'localhost',
'port': 5432, 'port': 5432,
@@ -18,15 +19,10 @@ def test___init__(mock_gates_init, mock_mlflow_init, mock_postgres_init):
'password': 'postgres', 'password': 'postgres',
'dbname': 'postgres', 'dbname': 'postgres',
'min_connections': 1, 'min_connections': 1,
'max_connections': 10 'max_connections': 10,
} }
mlflow_config = { mlflow_config = {'host': 'localhost', 'port': 5000, 'username': 'mlflow', 'password': 'mlflow'}
'host': 'localhost',
'port': 5000,
'username': 'mlflow',
'password': 'mlflow'
}
logger = MagicMock() logger = MagicMock()
notification_handler = MagicMock() notification_handler = MagicMock()
@@ -35,7 +31,7 @@ def test___init__(mock_gates_init, mock_mlflow_init, mock_postgres_init):
postgres_config=postgres_config, postgres_config=postgres_config,
mlflow_config=mlflow_config, mlflow_config=mlflow_config,
logger=logger, logger=logger,
notification_handler=notification_handler notification_handler=notification_handler,
) )
assert isinstance(activities, Activities) 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'], min_connections=postgres_config['min_connections'],
max_connections=postgres_config['max_connections'], max_connections=postgres_config['max_connections'],
logger=logger, logger=logger,
notification_handler=notification_handler notification_handler=notification_handler,
) )
mock_mlflow_init.assert_called_once_with( 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_username=mlflow_config['username'],
mlflow_password=mlflow_config['password'], mlflow_password=mlflow_config['password'],
logger=logger, logger=logger,
notification_handler=notification_handler notification_handler=notification_handler,
) )
mock_gates_init.assert_called_once_with( mock_gates_init.assert_called_once_with(
ANY, ANY, logger=logger, notification_handler=notification_handler
logger=logger,
notification_handler=notification_handler
) )
@@ -84,15 +78,10 @@ async def test_shutdown(_mock_mlflow_init, mock_postgres_init):
'password': 'postgres', 'password': 'postgres',
'dbname': 'postgres', 'dbname': 'postgres',
'min_connections': 1, 'min_connections': 1,
'max_connections': 10 'max_connections': 10,
} }
mlflow_config = { mlflow_config = {'host': 'localhost', 'port': 5000, 'username': 'mlflow', 'password': 'mlflow'}
'host': 'localhost',
'port': 5000,
'username': 'mlflow',
'password': 'mlflow'
}
logger = MagicMock() logger = MagicMock()
notification_handler = MagicMock() notification_handler = MagicMock()
@@ -101,7 +90,7 @@ async def test_shutdown(_mock_mlflow_init, mock_postgres_init):
postgres_config=postgres_config, postgres_config=postgres_config,
mlflow_config=mlflow_config, mlflow_config=mlflow_config,
logger=logger, logger=logger,
notification_handler=notification_handler notification_handler=notification_handler,
) )
await activities.shutdown() await activities.shutdown()

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

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

@@ -1,45 +1,42 @@
from datetime import datetime
from unittest.mock import ANY, MagicMock, patch from unittest.mock import ANY, MagicMock, patch
import numpy as np import numpy as np
from pandas import DataFrame, Timestamp from pytest import fixture, mark
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 sientia_do.notifications.models import NotificationLevel 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): def test___init__(mock_mlflow_repository):
mlflow = MLFlow( mlflow = MLFlow(
mlflow_host="http://localhost", mlflow_host='http://localhost',
mlflow_port=5000, mlflow_port=5000,
mlflow_username="admin", mlflow_username='admin',
mlflow_password="admin", mlflow_password='admin',
logger=MagicMock(), 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_port == 5000
assert mlflow.mlflow_username == "admin" assert mlflow.mlflow_username == 'admin'
assert mlflow.mlflow_password == "admin" assert mlflow.mlflow_password == 'admin'
mock_mlflow_repository.assert_called_once_with( mock_mlflow_repository.assert_called_once_with('http://localhost:5000', 'admin', 'admin', ANY)
"http://localhost:5000", "admin", "admin", ANY
)
@fixture @fixture
@patch("model_manager.activities.mlflow.MLFlowRepository") @patch('model_manager.activities.mlflow.MLFlowRepository')
def mlflow(mock_mlflow_repository): def mlflow(mock_mlflow_repository):
mlflow = MLFlow( mlflow = MLFlow(
mlflow_host="http://localhost:5000", mlflow_host='http://localhost:5000',
mlflow_port=5000, mlflow_port=5000,
mlflow_username="admin", mlflow_username='admin',
mlflow_password="admin", mlflow_password='admin',
logger=MagicMock(), logger=MagicMock(),
notification_handler=MagicMock() notification_handler=MagicMock(),
) )
mlflow.send_notification = MagicMock() mlflow.send_notification = MagicMock()
@@ -48,44 +45,67 @@ def mlflow(mock_mlflow_repository):
metadata = { metadata = {
"metadata": { 'metadata': {
"model_id": "test_model", 'model_id': 'test_model',
"model_name": "test_model", 'model_name': 'test_model',
"workflow_name": "test_workflow", 'workflow_name': 'test_workflow',
"schema_name": "test_schedule", 'schema_name': 'test_schedule',
}, },
} }
@mark.asyncio @mark.asyncio
@patch("model_manager.activities.mlflow.DataFrame") @patch('model_manager.activities.mlflow.DataFrame')
@patch("model_manager.activities.mlflow.max") @patch('model_manager.activities.mlflow.max')
async def test_request_transform_success(mock_max, mock_dataframe, mlflow): async def test_request_transform_success(mock_max, mock_dataframe, mlflow):
mock_max.return_value = '2024-01-02' mock_max.return_value = '2024-01-02'
# Mock input data # Mock input data
input_data = { input_data = {
**metadata, **metadata,
'data': [ 'data': [
{'timestamp': '2024-01-01', 'variable': 'var1', {
'value': 1.0, 'created_at': '2024-01-01 12:00:00'}, 'timestamp': '2024-01-01',
{'timestamp': '2024-01-01', 'variable': 'var2', 'variable': 'var1',
'value': 2.0, 'created_at': '2024-01-01 12:00:00'}, 'value': 1.0,
{'timestamp': '2024-01-02', 'variable': 'var1', 'created_at': '2024-01-01 12:00:00',
'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-01',
{'timestamp': '2024-01-02', 'variable': 'var1', 'variable': 'var2',
'value': 1.0, 'created_at': '2024-01-01 12:00:00'}, 'value': 2.0,
{'timestamp': '2024-01-02', 'variable': 'var2', 'created_at': '2024-01-01 12:00:00',
'value': 1.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_name': 'test_model',
'model_config': {} 'model_config': {},
} }
# Mock the transform response # Mock the transform response
expected_response = {'prediction': [0.5, 0.6], 'timestamp': [ expected_response = {'prediction': [0.5, 0.6], 'timestamp': ['2024-01-01', '2024-01-02']}
'2024-01-01', '2024-01-02']}
mlflow.model_monitoring_repository.transform.return_value = expected_response mlflow.model_monitoring_repository.transform.return_value = expected_response
mock_dataframe.return_value.sort_values.return_value = mock_dataframe.return_value mock_dataframe.return_value.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 @mark.asyncio
@patch("model_manager.activities.mlflow.DataFrame") @patch('model_manager.activities.mlflow.DataFrame')
@patch("model_manager.activities.mlflow.to_datetime") @patch('model_manager.activities.mlflow.to_datetime')
@patch("model_manager.activities.mlflow.max") @patch('model_manager.activities.mlflow.max')
async def test_request_predict(mock_max, mock_to_datetime, mock_dataframe, mlflow): async def test_request_predict(mock_max, mock_to_datetime, mock_dataframe, mlflow):
mock_max.return_value = '2024-01-02' mock_max.return_value = '2024-01-02'
# Mock input data # Mock input data
input_data = { input_data = {
**metadata, **metadata,
'data': { 'data': {
"variable": { 'variable': {
"2024-01-01": "var1", '2024-01-01': 'var1',
"2024-01-02": "var2", '2024-01-02': 'var2',
"2024-01-03": "var1", '2024-01-03': 'var1',
"2024-01-04": "var2" '2024-01-04': 'var2',
}, },
"value": { 'value': {'2024-01-01': 1.0, '2024-01-02': 2.0, '2024-01-03': 3.0, '2024-01-04': 4.0},
"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_name': 'test_model',
'model_config': {} 'model_config': {},
} }
# Mock the predict response # 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) response_data = await mlflow.request_predict(input_data)
mock_dataframe.assert_called_once_with(input_data['data']) mock_dataframe.assert_called_once_with(input_data['data'])
mock_dataframe.return_value.replace.assert_called_once_with( mock_dataframe.return_value.replace.assert_called_once_with(np.nan, None, inplace=True)
np.nan, None, inplace=True
)
mock_dataframe.return_value.__setitem__.assert_any_call( mock_dataframe.return_value.__setitem__.assert_any_call(
'timestamp', mock_to_datetime.return_value.dt.strftime.return_value '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_to_datetime.assert_called_once_with(
mock_dataframe.return_value.__getitem__.return_value, format=DATETIME_FORMAT_WITH_TZ mock_dataframe.return_value.__getitem__.return_value, format=DATETIME_FORMAT_WITH_TZ
) )
mock_to_datetime.return_value.dt.strftime.assert_called_once_with( mock_to_datetime.return_value.dt.strftime.assert_called_once_with(DATETIME_FORMAT)
DATETIME_FORMAT
)
# Verify the response # Verify the response
assert response_data == expected_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 @mark.asyncio
async def test_retrain_model(mlflow): async def test_retrain_model(mlflow):
data = { data = {
"model_id": [4, 5, 6, 7], 'model_id': [4, 5, 6, 7],
"created_at": [1, 2, 3, 4], 'created_at': [1, 2, 3, 4],
"timestamp": [1, 1, 2, 2], 'timestamp': [1, 1, 2, 2],
"variable": ["var1", "var2", "var1", "var2"], 'variable': ['var1', 'var2', 'var1', 'var2'],
"value": [1, 2, 3, 4] 'value': [1, 2, 3, 4],
} }
mlflow.model_monitoring_repository.retrain_model.return_value = ( mlflow.model_monitoring_repository.retrain_model.return_value = (
'Model retrained successfully', 'test') 'Model retrained successfully',
'test',
)
response = await mlflow.retrain_model({ response = await mlflow.retrain_model({**metadata, 'data': data, 'model_name': 'test_model'})
**metadata,
'data': data,
'model_name': 'test_model'
})
mlflow.model_monitoring_repository.retrain_model.assert_called_once() mlflow.model_monitoring_repository.retrain_model.assert_called_once()
assert response == { assert response == {
"status": 'Model retrained successfully', 'status': 'Model retrained successfully',
"timestamp": 2, 'timestamp': 2,
"experiment": 'test' 'experiment': 'test',
} }
@@ -206,20 +215,16 @@ async def test_retrain_model_error(mlflow):
) )
data = { data = {
"model_id": [4, 5, 6, 7], 'model_id': [4, 5, 6, 7],
"created_at": [1, 2, 3, 4], 'created_at': [1, 2, 3, 4],
"timestamp": [1, 1, 2, 2], 'timestamp': [1, 1, 2, 2],
"variable": ["var1", "var2", "var1", "var2"], 'variable': ['var1', 'var2', 'var1', 'var2'],
"value": [1, 2, 3, 4] 'value': [1, 2, 3, 4],
} }
try: try:
await mlflow.retrain_model({ await mlflow.retrain_model({**metadata, 'data': data, 'model_name': 'test_model'})
**metadata, except Exception as e: # noqa: BLE001
'data': data,
'model_name': 'test_model'
})
except Exception as e:
assert str(e) == 'Error retraining model' assert str(e) == 'Error retraining model'
mlflow.send_notification.assert_called_once_with( mlflow.send_notification.assert_called_once_with(
metadata=metadata['metadata'], metadata=metadata['metadata'],
@@ -227,20 +232,18 @@ async def test_retrain_model_error(mlflow):
message='Error retraining model test_model: Error retraining model', message='Error retraining model test_model: Error retraining model',
block='retrain_model', block='retrain_model',
level=NotificationLevel.ERROR, level=NotificationLevel.ERROR,
attachment_content=ANY attachment_content=ANY,
) )
else: else:
assert False, "No exception raised" raise AssertionError('No exception raised')
@mark.asyncio @mark.asyncio
async def test_update_production_model(mlflow): async def test_update_production_model(mlflow):
mlflow.model_monitoring_repository.update_production_model.return_value = ( mlflow.model_monitoring_repository.update_production_model.return_value = {
{ 'data1': 1,
"data1": 1, 'data2': 2,
"data2": 2 }
}
)
input_data = { input_data = {
**metadata, **metadata,
@@ -248,13 +251,14 @@ async def test_update_production_model(mlflow):
'model_id': 1, 'model_id': 1,
'experiment': 'test', 'experiment': 'test',
'timestamp': 2, 'timestamp': 2,
'status': 'success' 'status': 'success',
} }
response = await mlflow.update_production_model(input_data) response = await mlflow.update_production_model(input_data)
mlflow.model_monitoring_repository.update_production_model.assert_called_once_with( 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 == { assert response == {
'data1': {0: 1}, 'data1': {0: 1},
@@ -262,7 +266,7 @@ async def test_update_production_model(mlflow):
'model_id': {0: 1}, 'model_id': {0: 1},
'model_name': {0: 'test_model'}, 'model_name': {0: 'test_model'},
'timestamp': {0: 2}, 'timestamp': {0: 2},
'status': {0: 'success'} 'status': {0: 'success'},
} }
@@ -278,12 +282,12 @@ async def test_update_production_model_error(mlflow):
'model_id': 1, 'model_id': 1,
'experiment': 'test', 'experiment': 'test',
'timestamp': 2, 'timestamp': 2,
'status': 'success' 'status': 'success',
} }
try: try:
await mlflow.update_production_model(input_data) await mlflow.update_production_model(input_data)
except Exception as e: except Exception as e: # noqa: BLE001
assert str(e) == 'Error updating production model' assert str(e) == 'Error updating production model'
mlflow.send_notification.assert_called_once_with( mlflow.send_notification.assert_called_once_with(
metadata=metadata['metadata'], 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', message='Error updating production model test_model: Error updating production model',
block='update_production_model', block='update_production_model',
level=NotificationLevel.ERROR, level=NotificationLevel.ERROR,
attachment_content=ANY attachment_content=ANY,
) )
else: 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 pandas import DataFrame
from model_manager.utils.filters.conditional_filters import ( from model_manager.utils.filters.conditional_filters import (
filter_empty_data,
filter_specific_variables_null_values, filter_specific_variables_null_values,
filter_empty_data
) )
def test_filter_specific_variables_null_values(): def test_filter_specific_variables_null_values():
assert filter_specific_variables_null_values( assert (
DataFrame( filter_specific_variables_null_values(
{'variable': ['variable1', 'variable2'], 'value': [1, 2]}), DataFrame({'variable': ['variable1', 'variable2'], 'value': [1, 2]}),
config={'variables': ['variable2']}) is False config={'variables': ['variable2']},
)
is False
)
def test_filter_specific_variables_null_values_with_null_values(): def test_filter_specific_variables_null_values_with_null_values():
assert filter_specific_variables_null_values( assert (
DataFrame( filter_specific_variables_null_values(
{'variable': ['variable1', 'variable2'], 'value': [1, None]}), DataFrame({'variable': ['variable1', 'variable2'], 'value': [1, None]}),
config={'variables': ['variable2']}) is True config={'variables': ['variable2']},
)
is True
)
def test_filter_empty_data(): def test_filter_empty_data():
@@ -25,6 +31,7 @@ def test_filter_empty_data():
def test_filter_empty_data_with_data(): def test_filter_empty_data_with_data():
assert filter_empty_data( assert (
DataFrame({'variable': ['variable1', 'variable2'], 'value': [1, 2]}), filter_empty_data(DataFrame({'variable': ['variable1', 'variable2'], 'value': [1, 2]}), {})
{}) is False is False
)

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

@@ -1,22 +1,23 @@
from pandas import DataFrame from pandas import DataFrame
from model_manager.utils.filters.mlflow_filters import api_error_filter, nan_values_filter from model_manager.utils.filters.mlflow_filters import api_error_filter, nan_values_filter
def test_api_error_filter_invalid_response(): 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(): 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(): 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(): 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(): 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]}), {})

335
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 from unittest.mock import ANY, MagicMock, call, patch
import numpy as np import numpy as np
from pandas import DataFrame
import pytest import pytest
from datetime import datetime, timezone from pandas import DataFrame, Timestamp
from pandas import Timestamp
from model_manager.utils.repository.model_repository import MLFlowRepository from model_manager.utils.repository.model_repository import MLFlowRepository
@pytest.fixture @pytest.fixture
def mlflow_repository(): def mlflow_repository():
with patch('model_manager.utils.repository.model_repository.ModelServing', with patch(
autospec=True) as mock_model_serving: 'model_manager.utils.repository.model_repository.ModelServing', autospec=True
) as mock_model_serving:
mock_instance = mock_model_serving.return_value mock_instance = mock_model_serving.return_value
mock_instance.get_transformed_data = MagicMock() mock_instance.get_transformed_data = MagicMock()
repo = MLFlowRepository( repo = MLFlowRepository(
host='http://localhost:5000', host='http://localhost:5000', username='admin', password='admin', logger=MagicMock()
username='admin',
password='admin',
logger=MagicMock()
) )
return repo return repo
metadata = { metadata = {
"metadata": { 'metadata': {
"model_id": "test_model", 'model_id': 'test_model',
"model_name": "test_model", 'model_name': 'test_model',
"workflow_name": "test_workflow", 'workflow_name': 'test_workflow',
"schema_name": "test_schedule", 'schema_name': 'test_schedule',
}, },
} }
@@ -38,80 +37,56 @@ class Any:
invalid_cases = [ invalid_cases = [
( ({'value': {'2024-01-01 12:00:00': 1, 2024: 2}}),
{ ({'value': {'2024-01-01': 1, '2024-01-02': 2}}),
'value': { ({'value': {Any(): 1, Any(): 2}}),
'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): def test_detect_and_parse_datetime_index_error_cases(mlflow_repository, data):
input_data = DataFrame( input_data = DataFrame(data)
data
)
with pytest.raises(ValueError) as e: with pytest.raises(ValueError) as e:
mlflow_repository.detect_and_parse_datetime_index( mlflow_repository.detect_and_parse_datetime_index(input_data, metadata['metadata'])
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 = [ valid_cases = [
( (
{ {'value': {'2024-01-01 12:00:00+0000': 1, '2024-01-02 12:00:00+0000': 2}},
'value': { ['2024-01-01 12:00:00+0000', '2024-01-02 12:00:00+0000'],
'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': { 'value': {
datetime(2025, 1, 1, 12, 0, 0, tzinfo=timezone.utc): 1, datetime(2025, 1, 1, 12, 0, 0, tzinfo=UTC): 1,
datetime(2025, 1, 2, 12, 0, 0, tzinfo=timezone.utc): 2 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': { 'value': {
Timestamp(2026, 1, 1, 12, 0, 0, tzinfo=timezone.utc): 1, Timestamp(2026, 1, 1, 12, 0, 0, tzinfo=UTC): 1,
Timestamp(2026, 1, 2, 12, 0, 0, tzinfo=timezone.utc): 2 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): def test_detect_and_parse_datetime_index_valid_format(mlflow_repository, data, expected):
input_data = DataFrame(data) input_data = DataFrame(data)
response = mlflow_repository.detect_and_parse_datetime_index( response = mlflow_repository.detect_and_parse_datetime_index(input_data, metadata['metadata'])
input_data, metadata['metadata'])
assert response.index.tolist() == expected assert response.index.tolist() == expected
@@ -122,18 +97,19 @@ def test_transform_success(mlflow_repository):
mlflow_repository.detect_and_parse_datetime_index = MagicMock() mlflow_repository.detect_and_parse_datetime_index = MagicMock()
output = mlflow_repository.transform( output = mlflow_repository.transform(model_name, data, {}, metadata['metadata'])
model_name, data, {}, metadata['metadata'])
mlflow_repository.model_serving.get_cached_transform.assert_called_once_with( 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.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 == { assert output == {
'success': True, '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() data = MagicMock()
model_name = 'model' model_name = 'model'
mlflow_repository.model_serving.get_cached_transform.side_effect = Exception( mlflow_repository.model_serving.get_cached_transform.side_effect = Exception('error')
'error')
output = mlflow_repository.transform( output = mlflow_repository.transform(model_name, data, {}, metadata['metadata'])
model_name, data, {}, metadata['metadata'])
mlflow_repository.model_serving.get_cached_transform.assert_called_once_with( 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 == { assert output == {'success': False, 'content': {'message': 'error', 'traceback': ANY}}
'success': False,
'content': {
'message': 'error',
'traceback': ANY
}
}
def test_predict_success(mlflow_repository): def test_predict_success(mlflow_repository):
data = DataFrame({ data = DataFrame({'feat_1': {'index_1': 2, 'index_2': 3}})
'feat_1': {
'index_1': 2,
'index_2': 3
}
})
model_name = 'model' model_name = 'model'
mlflow_repository.model_serving.get_cached_predict.return_value = np.array( mlflow_repository.model_serving.get_cached_predict.return_value = np.array([2, 3])
[2, 3]
)
output = mlflow_repository.predict( output = mlflow_repository.predict(model_name, data, {}, metadata['metadata'])
model_name, data, {}, metadata['metadata'])
mlflow_repository.model_serving.get_cached_predict.assert_called_once_with( 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['success'] is True
assert output['content'] == { assert output['content'] == {
'prediction': { 'prediction': {'index_1': 2, 'index_2': 3},
'index_1': 2, 'response_time': {'index_1': ANY, 'index_2': ANY},
'index_2': 3
}, 'response_time': {
'index_1': ANY,
'index_2': ANY
}
} }
def test_predict_error(mlflow_repository): def test_predict_error(mlflow_repository):
data = DataFrame({ data = DataFrame({'feat_1': {'index_1': 2, 'index_2': 3}})
'feat_1': {
'index_1': 2,
'index_2': 3
}
})
model_name = 'model' model_name = 'model'
mlflow_repository.model_serving.get_cached_predict = MagicMock( mlflow_repository.model_serving.get_cached_predict = MagicMock(side_effect=Exception('error'))
side_effect=Exception('error')
)
output = mlflow_repository.predict( output = mlflow_repository.predict(model_name, data, {}, metadata['metadata'])
model_name, data, {}, metadata['metadata'])
mlflow_repository.model_serving.get_cached_predict.assert_called_once_with( 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 == { assert output == {'success': False, 'content': {'message': 'error', 'traceback': ANY}}
'success': False,
'content': {
'message': 'error',
'traceback': ANY
}
}
@patch('model_manager.utils.repository.model_repository.mlflow') @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') @patch('model_manager.utils.repository.model_repository.mlflow')
def test_get_experiment_success(mlflow, mlflow_repository): def test_get_experiment_success(mlflow, mlflow_repository):
mlflow.get_experiment_by_name.return_value = MagicMock( mlflow.get_experiment_by_name.return_value = MagicMock(experiment_id='0')
experiment_id='0')
output = mlflow_repository.get_experiment('test') output = mlflow_repository.get_experiment('test')
@@ -263,23 +206,25 @@ def test_get_experiment_error(mlflow, mlflow_repository):
except ValueError as e: except ValueError as e:
assert str(e) == 'Experiment test not found' assert str(e) == 'Experiment test not found'
else: else:
assert False raise AssertionError('Expected exception')
@patch('model_manager.utils.repository.model_repository.mlflow') @patch('model_manager.utils.repository.model_repository.mlflow')
def test_get_experiment_last_run(mlflow, mlflow_repository): 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'], 'params.retrain': ['True', 'False', 'True', 'False'],
'run_id': ['0', '1', '2', '3'], '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) output = mlflow_repository.get_experiment_last_run(0)
mlflow.search_runs.assert_called_once_with( mlflow.search_runs.assert_called_once_with(
experiment_ids=[0], experiment_ids=[0],
filter_string="", filter_string='',
output_format="pandas", output_format='pandas',
) )
assert output == '2' assert output == '2'
@@ -294,17 +239,14 @@ def test_get_experiment_last_run_error(mlflow, mlflow_repository):
except ValueError as e: except ValueError as e:
assert str(e) == 'Runs is not a pandas DataFrame' assert str(e) == 'Runs is not a pandas DataFrame'
else: else:
assert False raise AssertionError('Expected exception')
@patch('model_manager.utils.repository.model_repository.mlflow.sklearn') @patch('model_manager.utils.repository.model_repository.mlflow.sklearn')
@patch('model_manager.utils.repository.model_repository.mlflow.set_experiment') @patch('model_manager.utils.repository.model_repository.mlflow.set_experiment')
def test_create_model_experiment(set_experiment, sklearn, mlflow_repository): 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_run_id = MagicMock( mlflow_repository.model_serving.get_model_uri = MagicMock(return_value='test')
return_value='0')
mlflow_repository.model_serving.get_model_uri = MagicMock(
return_value='test')
mlflow_repository.get_experiment_by_run_id = MagicMock() mlflow_repository.get_experiment_by_run_id = MagicMock()
data_model_mock = 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] sklearn.load_model.side_effect = [data_model_mock, prediction_model_mock]
data_model_mock.fit.return_value = data_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], {
}) 'x': [10, 20, 30],
}
)
data_model_mock.target_variable = 'y' data_model_mock.target_variable = 'y'
prediction_model_mock.fit.return_value = prediction_model_mock prediction_model_mock.fit.return_value = prediction_model_mock
data = DataFrame({ data = DataFrame({'x': [1, 2, 3], 'y': [4, 5, 6]})
'x': [1, 2, 3],
'y': [4, 5, 6]
})
output = mlflow_repository.create_model_experiment('test', data) output = mlflow_repository.create_model_experiment('test', data)
mlflow_repository.model_serving.get_model_run_id.assert_called_once_with( mlflow_repository.model_serving.get_model_run_id.assert_called_once_with(
'test', stage='Production') 'test', stage='Production'
mlflow_repository.model_serving.get_model_uri.assert_called_once_with( )
'0', prediction=False) 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(mlflow_repository.model_serving.get_model_uri.return_value),
]) call('models:/test/production'),
]
)
assert sklearn.load_model.call_count == 2 assert sklearn.load_model.call_count == 2
data_model_mock.fit.assert_called_once_with(data) 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] fit_args = prediction_model_mock.fit.call_args[0][0]
assert fit_args.equals( assert fit_args.equals(
DataFrame({ DataFrame(
'x': [10, 20, 30], {
'y': [4, 5, 6], 'x': [10, 20, 30],
}) 'y': [4, 5, 6],
}
)
) )
mlflow_repository.get_experiment_by_run_id.assert_called_once_with('0') mlflow_repository.get_experiment_by_run_id.assert_called_once_with('0')
set_experiment.assert_called_once_with( set_experiment.assert_called_once_with(mlflow_repository.get_experiment_by_run_id.return_value)
mlflow_repository.get_experiment_by_run_id.return_value
)
assert output == (prediction_model_mock, assert output == (
data_model_mock, prediction_model_mock,
mlflow_repository.get_experiment_by_run_id.return_value) data_model_mock,
mlflow_repository.get_experiment_by_run_id.return_value,
)
@patch('model_manager.utils.repository.model_repository.mlflow.start_run') @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.sklearn.log_model')
@patch('model_manager.utils.repository.model_repository.mlflow.log_artifact') @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): def test_perform_model_retrain(log_artifact, log_model, log_param, start_run, mlflow_repository):
prediction_model_mock = MagicMock() prediction_model_mock = MagicMock()
data_model_mock = MagicMock() data_model_mock = MagicMock()
experiment = 'test' experiment = 'test'
model_name = 'test' model_name = 'test'
data = MagicMock() data = MagicMock()
mlflow_repository.get_next_run_name = MagicMock( mlflow_repository.get_next_run_name = MagicMock(return_value='test-1')
return_value='test-1')
run = MagicMock() run = MagicMock()
start_run.__enter__.return_value = run start_run.__enter__.return_value = run
output = mlflow_repository.perform_model_retrain( 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) mlflow_repository.get_next_run_name.assert_called_once_with(experiment)
start_run.assert_called_once_with( 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([ log_model.assert_has_calls(
call(data_model_mock, "data_model"), [
call(prediction_model_mock, "prediction_model"), call(data_model_mock, 'data_model'),
]) call(prediction_model_mock, 'prediction_model'),
]
)
data.to_csv.assert_called_once_with( data.to_csv.assert_called_once_with('temp/raw_data_test.csv', index=True)
"temp/raw_data_test.csv", index=True)
log_artifact.assert_called_once_with( log_artifact.assert_called_once_with('temp/raw_data_test.csv')
"temp/raw_data_test.csv")
log_param.assert_has_calls([ log_param.assert_has_calls(
call("retrain", True), [
]) call('retrain', True),
]
)
assert output == ("Model retrained successfully", experiment) assert output == ('Model retrained successfully', experiment)
def test_retrain_model(mlflow_repository): def test_retrain_model(mlflow_repository):
@@ -407,18 +354,18 @@ def test_retrain_model(mlflow_repository):
model_name = 'test' model_name = 'test'
mlflow_repository.create_model_experiment = MagicMock( 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( mlflow_repository.perform_model_retrain = MagicMock(return_value='Model retrained successfully')
return_value='Model retrained successfully')
output = mlflow_repository.retrain_model(data, model_name) output = mlflow_repository.retrain_model(data, model_name)
mlflow_repository.create_model_experiment.assert_called_once_with( mlflow_repository.create_model_experiment.assert_called_once_with(model_name, data)
model_name, data)
mlflow_repository.perform_model_retrain.assert_called_once_with( 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' 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') output = mlflow_repository.update_production_model_by_run_id('0', 'test')
mlflow.register_model.assert_called_once_with( mlflow.register_model.assert_called_once_with(
"runs:/0/prediction_model", 'runs:/0/prediction_model',
'test', 'test',
) )
@@ -461,38 +408,34 @@ def test_update_production_model_by_run_id(mlflow, mlflow_repository):
@patch('model_manager.utils.repository.model_repository.mlflow') @patch('model_manager.utils.repository.model_repository.mlflow')
def test_update_production_model_by_run_id_error(mlflow, mlflow_repository): def test_update_production_model_by_run_id_error(mlflow, mlflow_repository):
mlflow.tracking.MlflowClient.return_value = MagicMock( mlflow.tracking.MlflowClient.return_value = MagicMock(
get_registered_model=MagicMock( get_registered_model=MagicMock(return_value=MagicMock(latest_versions={}))
return_value=MagicMock(
latest_versions={}
)
)
) )
try: try:
mlflow_repository.update_production_model_by_run_id('0', 'test') 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' assert str(e) == 'Model versions is not a list'
else: else:
assert False raise AssertionError('Expected exception')
def test_update_production_model(mlflow_repository): def test_update_production_model(mlflow_repository):
connector = mlflow_repository connector = mlflow_repository
with patch.object(connector, 'get_experiment', with patch.object(connector, 'get_experiment', return_value='0') as get_experiment:
return_value='0') as get_experiment: with patch.object(
with patch.object(connector, 'get_experiment_last_run', connector, 'get_experiment_last_run', return_value='2'
return_value='2') as get_experiment_last_run: ) as get_experiment_last_run:
with patch.object(connector, 'update_production_model_by_run_id', with patch.object(
return_value={'model_name': 'test', 'version': '3', connector,
'mlflow_run_id': '0'}) as update_production_model_by_run_id: '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') output = connector.update_production_model('0', 'test')
get_experiment.assert_called_once_with('0') get_experiment.assert_called_once_with('0')
get_experiment_last_run.assert_called_once_with('0') get_experiment_last_run.assert_called_once_with('0')
update_production_model_by_run_id.assert_called_once_with( update_production_model_by_run_id.assert_called_once_with('2', 'test')
'2', 'test')
assert output == { assert output == {
'model_name': 'test', 'model_name': 'test',

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

@@ -1,7 +1,10 @@
from os import environ from os import environ
from model_manager.utils.connectors_config import (build_mlflow_config,
build_postgres_config, from model_manager.utils.connectors_config import (
build_mongodb_config) build_mlflow_config,
build_mongodb_config,
build_postgres_config,
)
def test_build_mlflow_config_with_env_vars(): 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() == { assert build_mongodb_config() == {
'connection_string': 'mongodb://sientia1:sientia1@localhost:27018', 'connection_string': 'mongodb://sientia1:sientia1@localhost:27018',
'database_name': 'test_db', '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() == { assert build_mongodb_config() == {
'connection_string': 'mongodb://root:wKZDbMNU1c@localhost:27018', 'connection_string': 'mongodb://root:wKZDbMNU1c@localhost:27018',
'database_name': 'sientia', 'database_name': 'sientia',
'ttl_index_seconds': 3600 'ttl_index_seconds': 3600,
} }

197
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 unittest.mock import ANY, AsyncMock, call, patch
from pytest import mark, fixture
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.activities.activities import Activities
from model_manager.workflows.sub_workflows.format_and_export_prediction import FormatAndExportPrediction from model_manager.workflows.sub_workflows.format_and_export_prediction import (
from sientia_do.temporal.constants import DATETIME_FORMAT_WITH_TZ FormatAndExportPrediction,
)
@fixture @fixture
@@ -12,119 +15,133 @@ def format_and_export_prediction():
metadata = { metadata = {
"metadata": { 'metadata': {
"model_id": "test_model", 'model_id': 'test_model',
"model_name": "test_model", 'model_name': 'test_model',
"workflow_name": "test_workflow", 'workflow_name': 'test_workflow',
"schema_name": "test_schedule", 'schema_name': 'test_schedule',
}, },
} }
@mark.asyncio @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): async def test_run_none_path_flag(workflow_mock, format_and_export_prediction):
input_data = { input_data = {
'metadata': metadata, 'metadata': metadata,
"path_flag": None, 'path_flag': None,
"data": {"test": "data"}, 'data': {'test': 'data'},
"timestamp": "2021-01-01", 'timestamp': '2021-01-01',
"model_id": 1, 'model_id': 1,
"prediction_confidence": 0, 'prediction_confidence': 0,
"schema": "test_schema", 'schema': 'test_schema',
"table_name": "test_table", 'table_name': 'test_table',
"prediction_store_policy": "erl:1" 'prediction_store_policy': 'erl:1',
} }
await format_and_export_prediction.run(input_data) 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, call(
{ Activities.format_prediction,
'data': input_data['data'], {
'timestamp': input_data['timestamp'], 'data': input_data['data'],
'model_id': input_data['model_id'], 'timestamp': input_data['timestamp'],
'prediction_confidence': input_data['prediction_confidence'], 'model_id': input_data['model_id'],
'prediction_store_policy': input_data['prediction_store_policy'], 'prediction_confidence': input_data['prediction_confidence'],
**metadata 'prediction_store_policy': input_data['prediction_store_policy'],
}, **metadata,
retry_policy=ANY, },
start_to_close_timeout=ANY retry_policy=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, call(
{ Activities.export_data_to_postgres,
'schema': input_data['schema'], {
'table_name': input_data['table_name'], 'schema': input_data['schema'],
'data': workflow_mock.execute_activity_method.return_value, 'table_name': input_data['table_name'],
**metadata, 'data': workflow_mock.execute_activity_method.return_value,
'timestamp_conversion': { **metadata,
'column': 'timestamp', 'timestamp_conversion': {
'format': DATETIME_FORMAT_WITH_TZ 'column': 'timestamp',
} 'format': DATETIME_FORMAT_WITH_TZ,
}, },
retry_policy=ANY, },
start_to_close_timeout=ANY retry_policy=ANY,
)]) start_to_close_timeout=ANY,
)
]
)
assert workflow_mock.execute_activity_method.call_count == 2 assert workflow_mock.execute_activity_method.call_count == 2
assert workflow_mock.execute_local_activity_method.call_count == 1 assert workflow_mock.execute_local_activity_method.call_count == 1
@mark.asyncio @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): async def test_run_default_path_flag(workflow_mock, format_and_export_prediction):
input_data = { input_data = {
'metadata': metadata, 'metadata': metadata,
"path_flag": "default", 'path_flag': 'default',
"data": {"test": "data"}, 'data': {'test': 'data'},
"timestamp": "2021-01-01", 'timestamp': '2021-01-01',
"model_id": 1, 'model_id': 1,
"prediction_confidence": 0, 'prediction_confidence': 0,
"schema": "test_schema", 'schema': 'test_schema',
"table_name": "test_table", 'table_name': 'test_table',
"comment": "test_comment" 'comment': 'test_comment',
} }
await format_and_export_prediction.run(input_data) 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, call(
{ Activities.format_default_prediction,
'timestamp': input_data['timestamp'], {
'model_id': input_data['model_id'], 'timestamp': input_data['timestamp'],
'prediction_confidence': input_data['prediction_confidence'], 'model_id': input_data['model_id'],
'comment': input_data['comment'], 'prediction_confidence': input_data['prediction_confidence'],
**metadata 'comment': input_data['comment'],
}, **metadata,
retry_policy=ANY, },
start_to_close_timeout=ANY retry_policy=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, call(
{ Activities.export_data_to_postgres,
'schema': input_data['schema'], {
'table_name': input_data['table_name'], 'schema': input_data['schema'],
'data': workflow_mock.execute_activity_method.return_value, 'table_name': input_data['table_name'],
**metadata, 'data': workflow_mock.execute_activity_method.return_value,
'timestamp_conversion': { **metadata,
'column': 'timestamp', 'timestamp_conversion': {
'format': DATETIME_FORMAT_WITH_TZ 'column': 'timestamp',
} 'format': DATETIME_FORMAT_WITH_TZ,
}, },
retry_policy=ANY, },
start_to_close_timeout=ANY retry_policy=ANY,
) start_to_close_timeout=ANY,
]) )
]
)
assert workflow_mock.execute_activity_method.call_count == 3 assert workflow_mock.execute_activity_method.call_count == 3
assert workflow_mock.execute_local_activity_method.call_count == 1 assert workflow_mock.execute_local_activity_method.call_count == 1

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

120
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 pytest import fixture, mark
from model_manager.activities.activities import Activities from model_manager.activities.activities import Activities
from model_manager.workflows.minimal_retrain import MinimalRetrain from model_manager.workflows.minimal_retrain import MinimalRetrain
@@ -10,11 +12,11 @@ def minimal_retrain() -> MinimalRetrain:
metadata = { metadata = {
"metadata": { 'metadata': {
"model_id": "test_model_id", 'model_id': 'test_model_id',
"model_name": "test_model", 'model_name': 'test_model',
"workflow_name": "minimal_retrain", 'workflow_name': 'minimal_retrain',
"schedule_name": "test_schedule", 'schedule_name': 'test_schedule',
}, },
} }
@@ -23,19 +25,19 @@ metadata = {
@patch('model_manager.workflows.minimal_retrain.workflow', new_callable=AsyncMock) @patch('model_manager.workflows.minimal_retrain.workflow', new_callable=AsyncMock)
async def test_run(workflow_mock: AsyncMock, minimal_retrain: MinimalRetrain): async def test_run(workflow_mock: AsyncMock, minimal_retrain: MinimalRetrain):
input_data = { input_data = {
"model_id": "test_model_id", 'model_id': 'test_model_id',
"model_name": "test_model", 'model_name': 'test_model',
"workflow_name": "minimal_retrain", 'workflow_name': 'minimal_retrain',
"schedule_name": "test_schedule", 'schedule_name': 'test_schedule',
"query": "test_query", 'query': 'test_query',
"schema": "test_schema", 'schema': 'test_schema',
"table_name": "test_table", 'table_name': 'test_table',
} }
workflow_mock.execute_activity_method = AsyncMock( workflow_mock.execute_activity_method = AsyncMock(
return_value={ return_value={
"data1": "1", 'data1': '1',
"data2": "2", 'data2': '2',
} }
) )
@@ -47,52 +49,58 @@ async def test_run(workflow_mock: AsyncMock, minimal_retrain: MinimalRetrain):
Activities.load_custom_query, Activities.load_custom_query,
{ {
**metadata, **metadata,
"query": input_data["query"], 'query': input_data['query'],
'datetime_columns': input_data.get('datetime_columns', []) 'datetime_columns': input_data.get('datetime_columns', []),
}, },
retry_policy=ANY, 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, call(
{ Activities.retrain_model,
**metadata, {
'data': workflow_mock.execute_local_activity_method.return_value, **metadata,
'model_name': input_data['model_name'], 'data': workflow_mock.execute_local_activity_method.return_value,
}, 'model_name': input_data['model_name'],
retry_policy=ANY, },
start_to_close_timeout=ANY retry_policy=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, call(
{ Activities.update_production_model,
**metadata, {
'model_name': input_data['model_name'], **metadata,
'model_id': input_data['model_id'], 'model_name': input_data['model_name'],
**workflow_mock.execute_activity_method.return_value, 'model_id': input_data['model_id'],
}, **workflow_mock.execute_activity_method.return_value,
retry_policy=ANY, },
start_to_close_timeout=ANY retry_policy=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, call(
{ Activities.export_data_to_postgres,
**metadata, {
'data': workflow_mock.execute_activity_method.return_value, **metadata,
'schema': input_data['schema'], 'data': workflow_mock.execute_activity_method.return_value,
'table_name': input_data['table_name'], 'schema': input_data['schema'],
}, 'table_name': input_data['table_name'],
retry_policy=ANY, },
start_to_close_timeout=ANY retry_policy=ANY,
) start_to_close_timeout=ANY,
]) )
]
)

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