SIENTIAPDE-1243: Initial commit of the model manager project, adding core files and configurations.
This commit introduces the initial project structure, including: - .env.example: Example environment configuration. - .github/workflows/quality-gate.yml: CI workflow for quality checks. - .gitignore: Specifies intentionally untracked files that Git should ignore. - Makefile: Automation of tasks like docker builds. - README.md: Project documentation. - Source code for model management, activities, utils, worker and workflows. - Test suite. - Dockerfile for the simulator. - sonar-project.properties: SonarQube configuration file. - values.yaml: Helm chart values for deployment.
This commit is contained in:
Bruno Domingues
@@ -0,0 +1,160 @@
|
||||
from unittest.mock import call, patch, AsyncMock, ANY
|
||||
from pytest import mark, fixture
|
||||
|
||||
from laborious.activities.activities import Activities
|
||||
from laborious.workflows.sub_workflows.format_and_export_prediction import FormatAndExportPrediction
|
||||
from sientia_do.temporal.constants import DATETIME_FORMAT_WITH_TZ
|
||||
|
||||
|
||||
@fixture
|
||||
def format_and_export_prediction():
|
||||
return FormatAndExportPrediction()
|
||||
|
||||
|
||||
metadata = {
|
||||
"metadata": {
|
||||
"model_id": "test_model",
|
||||
"model_name": "test_model",
|
||||
"workflow_name": "test_workflow",
|
||||
"schema_name": "test_schedule",
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
@patch("laborious.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",
|
||||
"opc_servers": ["test_server"],
|
||||
"opc_output_config": {"test": "config"},
|
||||
"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_activity_method.assert_has_calls([
|
||||
call(
|
||||
Activities.write_opc_data,
|
||||
{
|
||||
'opc_output_config': input_data['opc_output_config'],
|
||||
'data': workflow_mock.execute_local_activity_method.return_value,
|
||||
**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
|
||||
)])
|
||||
|
||||
assert workflow_mock.execute_activity_method.call_count == 3
|
||||
assert workflow_mock.execute_local_activity_method.call_count == 1
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
@patch("laborious.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",
|
||||
"opc_servers": ["test_server"],
|
||||
"opc_output_config": {"test": "config"},
|
||||
"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_activity_method.assert_has_calls([
|
||||
call(
|
||||
Activities.write_opc_data,
|
||||
{
|
||||
'opc_output_config': input_data['opc_output_config'],
|
||||
'data': workflow_mock.execute_local_activity_method.return_value,
|
||||
**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
|
||||
)
|
||||
])
|
||||
|
||||
assert workflow_mock.execute_activity_method.call_count == 3
|
||||
assert workflow_mock.execute_local_activity_method.call_count == 1
|
||||
@@ -0,0 +1,601 @@
|
||||
from unittest.mock import AsyncMock, patch, call, ANY
|
||||
from pytest import fixture, mark
|
||||
from laborious.activities.activities import Activities
|
||||
from laborious.workflows.sub_workflows.prediction_process import PredictionProcess
|
||||
|
||||
|
||||
@fixture
|
||||
def prediction_process():
|
||||
return PredictionProcess()
|
||||
|
||||
|
||||
metadata = {
|
||||
"metadata": {
|
||||
"model_id": "test_model",
|
||||
"model_name": "test_model",
|
||||
"workflow_name": "test_workflow",
|
||||
"schema_name": "test_schedule",
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
@patch("laborious.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
|
||||
input_data = {
|
||||
'metadata': metadata,
|
||||
'data': {'test': 'data'},
|
||||
'schema': 'test_schema',
|
||||
'table_name': 'test_table',
|
||||
'model_id': 1,
|
||||
'input_filters': {'test': 'filter'},
|
||||
'mlflow_transform_filters': {'test': 'filter'},
|
||||
'mlflow_predict_filters': {'test': 'filter'},
|
||||
'model_name': 'test_model_name',
|
||||
'model_config': {
|
||||
'retention': '30'
|
||||
},
|
||||
'path_priority': ['continue', 'repeat', 'stop'],
|
||||
'opc_output_config': {'test': 'config'},
|
||||
'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
|
||||
{'content': 'transformed_data', 'timestamp': '2024-01-01'}, # transform_data
|
||||
# mlflow_response_gate (transform)
|
||||
('continue', 0.95, "Error"),
|
||||
# mlflow_content_gate (transform)
|
||||
('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"),
|
||||
]
|
||||
|
||||
# Act
|
||||
await prediction_process.run(input_data)
|
||||
|
||||
# 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_child_workflow.assert_called_once_with(
|
||||
'format_and_export_prediction',
|
||||
{
|
||||
'metadata': metadata,
|
||||
'path_flag': 'continue',
|
||||
'data': 'predicted_data',
|
||||
'prediction_confidence': 0.95,
|
||||
'timestamp': '2024-01-01',
|
||||
'model_id': 1,
|
||||
'model_name': 'test_model_name',
|
||||
'model_config': input_data['model_config'],
|
||||
'opc_output_config': input_data['opc_output_config'],
|
||||
'schema': input_data['schema'],
|
||||
'table_name': input_data['table_name'],
|
||||
'comment': 'Error',
|
||||
'prediction_store_policy': input_data['prediction_store_policy']
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
@patch("laborious.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
|
||||
input_data = {
|
||||
'metadata': metadata,
|
||||
'data': {'test': 'data'},
|
||||
'schema': 'test_schema',
|
||||
'table_name': 'test_table',
|
||||
'model_id': 1,
|
||||
'input_filters': {'test': 'filter'},
|
||||
'mlflow_transform_filters': {'test': 'filter'},
|
||||
'mlflow_predict_filters': {'test': 'filter'},
|
||||
'model_name': 'test_model_name',
|
||||
'model_config': {
|
||||
'retention': '30'
|
||||
},
|
||||
'path_priority': ['continue', 'repeat', 'stop'],
|
||||
'opc_output_config': {'test': 'config'}
|
||||
}
|
||||
|
||||
# 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
|
||||
]
|
||||
|
||||
# Act
|
||||
await prediction_process.run(input_data)
|
||||
|
||||
# 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_child_workflow.assert_not_called()
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
@patch("laborious.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
|
||||
input_data = {
|
||||
'metadata': metadata,
|
||||
'data': {'test': 'data'},
|
||||
'schema': 'test_schema',
|
||||
'table_name': 'test_table',
|
||||
'model_id': 1,
|
||||
'input_filters': {'test': 'filter'},
|
||||
'mlflow_transform_filters': {'test': 'filter'},
|
||||
'mlflow_predict_filters': {'test': 'filter'},
|
||||
'model_name': 'test_model_name',
|
||||
'model_config': {
|
||||
'retention': '30'
|
||||
},
|
||||
'path_priority': ['continue', 'repeat', 'stop'],
|
||||
'opc_output_config': {'test': 'config'}
|
||||
}
|
||||
|
||||
# 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
|
||||
{'content': 'transformed_data', 'timestamp': '2024-01-01'}, # transform_data
|
||||
('continue', 0.95, "Error"), # mlflow_response_gate (transform)
|
||||
]
|
||||
|
||||
# Act
|
||||
await prediction_process.run(input_data)
|
||||
|
||||
# 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_child_workflow.assert_not_called()
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
@patch("laborious.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])
|
||||
# Arrange
|
||||
input_data = {
|
||||
'metadata': metadata,
|
||||
'data': {'test': 'data'},
|
||||
'schema': 'test_schema',
|
||||
'table_name': 'test_table',
|
||||
'model_id': 1,
|
||||
'input_filters': {'test': 'filter'},
|
||||
'mlflow_transform_filters': {'test': 'filter'},
|
||||
'mlflow_predict_filters': {'test': 'filter'},
|
||||
'model_name': 'test_model_name',
|
||||
'model_config': {
|
||||
'retention': '30'
|
||||
},
|
||||
'path_priority': ['continue', 'repeat', 'stop'],
|
||||
'opc_output_config': {'test': 'config'}
|
||||
}
|
||||
|
||||
# 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
|
||||
{'content': 'transformed_data', 'timestamp': '2024-01-01'}, # transform_data
|
||||
# mlflow_response_gate (transform)
|
||||
('continue', 0.95, "Error"),
|
||||
# mlflow_content_gate (transform)
|
||||
('continue', 0.95, "Transformed data not passed the content filter"),
|
||||
]
|
||||
|
||||
# Act
|
||||
await prediction_process.run(input_data)
|
||||
|
||||
# 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_child_workflow.assert_not_called()
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
@patch("laborious.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])
|
||||
# Arrange
|
||||
input_data = {
|
||||
'metadata': metadata,
|
||||
'data': {'test': 'data'},
|
||||
'schema': 'test_schema',
|
||||
'table_name': 'test_table',
|
||||
'model_id': 1,
|
||||
'input_filters': {'test': 'filter'},
|
||||
'mlflow_transform_filters': {'test': 'filter'},
|
||||
'mlflow_predict_filters': {'test': 'filter'},
|
||||
'model_name': 'test_model_name',
|
||||
'model_config': {
|
||||
'retention': '30'
|
||||
},
|
||||
'path_priority': ['continue', 'repeat', 'stop'],
|
||||
'opc_output_config': {'test': 'config'}
|
||||
}
|
||||
|
||||
# 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
|
||||
{'content': 'transformed_data', 'timestamp': '2024-01-01'}, # transform_data
|
||||
# mlflow_response_gate (transform)
|
||||
('continue', 0.95, "Error"),
|
||||
# mlflow_content_gate (transform)
|
||||
('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)
|
||||
]
|
||||
|
||||
# Act
|
||||
await prediction_process.run(input_data)
|
||||
|
||||
# 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_child_workflow.assert_not_called()
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
@patch("laborious.workflows.sub_workflows.prediction_process.workflow", new_callable=AsyncMock)
|
||||
async def test_path_flag_handler_stop(workflow_mock, prediction_process):
|
||||
# Arrange
|
||||
data = {'test': 'data'}
|
||||
path_flag = 'STOP'
|
||||
confidence = 0.95
|
||||
schema = 'test_schema'
|
||||
table_name = 'test_table'
|
||||
model = 'test_model'
|
||||
last_timestamp = '2024-01-01'
|
||||
model_name = 'test_model_name'
|
||||
model_config = {
|
||||
'retention': '30'
|
||||
}
|
||||
|
||||
# Act
|
||||
result = await prediction_process.path_flag_handler(
|
||||
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, ""
|
||||
)
|
||||
|
||||
# Assert
|
||||
assert result is True
|
||||
workflow_mock.execute_local_activity_method.assert_not_called()
|
||||
workflow_mock.execute_child_workflow.assert_not_called()
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
@patch("laborious.workflows.sub_workflows.prediction_process.workflow", new_callable=AsyncMock)
|
||||
async def test_path_flag_handler_repeat(workflow_mock, prediction_process):
|
||||
# Arrange
|
||||
data = {'test': 'data'}
|
||||
path_flag = 'repeat'
|
||||
confidence = 0.95
|
||||
schema = 'test_schema'
|
||||
table_name = 'test_table'
|
||||
model = 'test_model'
|
||||
last_timestamp = '2024-01-01'
|
||||
model_name = 'test_model_name'
|
||||
model_config = {
|
||||
'retention': '30'
|
||||
}
|
||||
|
||||
# Act
|
||||
result = await prediction_process.path_flag_handler(
|
||||
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, ""
|
||||
)
|
||||
|
||||
# Assert
|
||||
assert result is True
|
||||
workflow_mock.execute_activity_method.assert_called_once_with(
|
||||
Activities.repeat_last_prediction,
|
||||
{
|
||||
**metadata,
|
||||
'schema': schema,
|
||||
'table_name': table_name,
|
||||
'model': model,
|
||||
'last_timestamp': last_timestamp,
|
||||
},
|
||||
retry_policy=ANY,
|
||||
start_to_close_timeout=ANY
|
||||
)
|
||||
workflow_mock.execute_child_workflow.assert_not_called()
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
@patch("laborious.workflows.sub_workflows.prediction_process.workflow", new_callable=AsyncMock)
|
||||
async def test_path_flag_handler_continue(workflow_mock, prediction_process):
|
||||
# Arrange
|
||||
data = {'test': 'data'}
|
||||
path_flag = 'CONTINUE'
|
||||
confidence = 0.95
|
||||
schema = 'test_schema'
|
||||
table_name = 'test_table'
|
||||
model = 'test_model'
|
||||
last_timestamp = '2024-01-01'
|
||||
model_name = 'test_model_name'
|
||||
model_config = {
|
||||
'retention': '30'
|
||||
}
|
||||
prediction_store_policy = 'erl:1'
|
||||
|
||||
# Act
|
||||
result = await prediction_process.path_flag_handler(
|
||||
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,
|
||||
'opc_output_config': {'test': 'config'},
|
||||
'prediction_store_policy': prediction_store_policy
|
||||
}, confidence, last_timestamp, 'Prediction Process'
|
||||
)
|
||||
|
||||
# Assert
|
||||
assert result is True
|
||||
workflow_mock.execute_activity_method.assert_not_called()
|
||||
workflow_mock.execute_child_workflow.assert_called_once_with(
|
||||
'format_and_export_prediction',
|
||||
{
|
||||
'metadata': metadata,
|
||||
'path_flag': path_flag,
|
||||
'data': data,
|
||||
'prediction_confidence': confidence,
|
||||
'timestamp': last_timestamp,
|
||||
'model_id': model,
|
||||
'model_name': model_name,
|
||||
'model_config': model_config,
|
||||
'schema': schema,
|
||||
'table_name': table_name,
|
||||
'comment': 'Prediction Process',
|
||||
'opc_output_config': {'test': 'config'},
|
||||
'prediction_store_policy': prediction_store_policy
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
@patch("laborious.workflows.sub_workflows.prediction_process.workflow", new_callable=AsyncMock)
|
||||
async def test_path_flag_handler_unknown(workflow_mock, prediction_process):
|
||||
# Arrange
|
||||
data = {'test': 'data'}
|
||||
path_flag = 'unknown'
|
||||
confidence = 0.95
|
||||
schema = 'test_schema'
|
||||
table_name = 'test_table'
|
||||
model = 'test_model'
|
||||
last_timestamp = '2024-01-01'
|
||||
model_name = 'test_model_name'
|
||||
model_config = {
|
||||
'retention': '30'
|
||||
}
|
||||
prediction_store_policy = 'erl:1'
|
||||
# Act
|
||||
result = await prediction_process.path_flag_handler(
|
||||
data, path_flag, {
|
||||
**metadata,
|
||||
'schema': schema,
|
||||
'table_name': table_name,
|
||||
'model_id': model,
|
||||
'last_timestamp': last_timestamp,
|
||||
'model_name': model_name,
|
||||
'model_config': model_config,
|
||||
'opc_output_config': {'test': 'config'},
|
||||
'prediction_store_policy': prediction_store_policy
|
||||
}, confidence, last_timestamp, ""
|
||||
)
|
||||
|
||||
# Assert
|
||||
assert result is False
|
||||
workflow_mock.execute_activity_method.assert_not_called()
|
||||
workflow_mock.execute_child_workflow.assert_not_called()
|
||||
98
tests/laborious/workflows/test_minimal_retrain.py
Normal file
98
tests/laborious/workflows/test_minimal_retrain.py
Normal file
@@ -0,0 +1,98 @@
|
||||
from unittest.mock import AsyncMock, MagicMock, call, patch, ANY
|
||||
from pytest import fixture, mark
|
||||
from laborious.activities.activities import Activities
|
||||
from laborious.workflows.minimal_retrain import MinimalRetrain
|
||||
|
||||
|
||||
@fixture
|
||||
def minimal_retrain() -> MinimalRetrain:
|
||||
return MinimalRetrain()
|
||||
|
||||
|
||||
metadata = {
|
||||
"metadata": {
|
||||
"model_id": "test_model_id",
|
||||
"model_name": "test_model",
|
||||
"workflow_name": "minimal_retrain",
|
||||
"schedule_name": "test_schedule",
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
@patch('laborious.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",
|
||||
}
|
||||
|
||||
workflow_mock.execute_activity_method = AsyncMock(
|
||||
return_value={
|
||||
"data1": "1",
|
||||
"data2": "2",
|
||||
}
|
||||
)
|
||||
|
||||
await minimal_retrain.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_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.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
|
||||
)
|
||||
])
|
||||
88
tests/laborious/workflows/test_predictions_batch.py
Normal file
88
tests/laborious/workflows/test_predictions_batch.py
Normal file
@@ -0,0 +1,88 @@
|
||||
from unittest.mock import AsyncMock, call, patch, ANY
|
||||
from pytest import fixture, mark
|
||||
from laborious.activities.activities import Activities
|
||||
from laborious.workflows.predictions_batch import PredictionsBatch
|
||||
|
||||
|
||||
@fixture
|
||||
def predictions_batch() -> PredictionsBatch:
|
||||
return PredictionsBatch()
|
||||
|
||||
|
||||
metadata = {
|
||||
"metadata": {
|
||||
"model_id": "test_model_id",
|
||||
"model_name": "test_model",
|
||||
"workflow_name": "predictions_batch",
|
||||
"schedule_name": "test_schedule",
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
@patch('laborious.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'
|
||||
}
|
||||
input_data = {
|
||||
'schedule_name': 'test_schedule',
|
||||
'model_name': 'test_model',
|
||||
'model_id': 'test_model_id',
|
||||
'query': 'SELECT * FROM test',
|
||||
'schema': 'test_schema',
|
||||
'table_name': 'test_table',
|
||||
'opc_output_config': 'test_opc_output_config',
|
||||
'datetime_columns': ['timestamp', 'created_at'],
|
||||
'prediction_store_policy': 'erl:1',
|
||||
'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
|
||||
)
|
||||
])
|
||||
prediction_input = {
|
||||
'metadata': metadata,
|
||||
'data': {'data': 'test_data'},
|
||||
'schema': input_data['schema'],
|
||||
'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'
|
||||
}
|
||||
}),
|
||||
'model_config': input_data.get('model_config', {}),
|
||||
'path_priority': input_data.get('path_priority', ['STOP', 'CONTINUE', 'REPEAT']),
|
||||
'opc_output_config': input_data.get('opc_output_config', {}),
|
||||
'prediction_store_policy': input_data.get('prediction_store_policy', 'erl:1')
|
||||
}
|
||||
|
||||
workflow_mock.execute_child_workflow.assert_has_calls([
|
||||
call(
|
||||
'prediction_process', prediction_input)
|
||||
])
|
||||
Reference in New Issue
Block a user