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

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

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

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