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

SIENTIAPDE-1222

Browse Source 
Refactor model configuration handling in MLFlow and workflows

- Replaced 'model_retention' with 'model_config' to encapsulate retention settings and improve consistency across various components.
- Updated test cases to reflect changes in argument structure, ensuring compatibility with the new model configuration format.
- Added 'prediction_store_policy' to input data handling in workflows for enhanced configuration management.
This commit is contained in:
vitor-aignosi
2025-09-17 16:02:54 -03:00
parent d6fdbc58bb
commit 84c1371d6c
4 changed files with 111 additions and 69 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
## Problemas no Courier:.md
View File

File diff suppressed because one or more lines are too long

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

@@ -4,6 +4,7 @@ from unittest.mock import ANY, MagicMock, patch
import numpy as np
from pandas import DataFrame, Timestamp
from pytest import fixture, mark, raises
from sientia_do.temporal.constants import DATETIME_FORMAT, DATETIME_FORMAT_WITH_TZ
from laborious.activities.mlflow import MLFlow
from sientia_do.notifications.models import NotificationLevel
@@ -79,7 +80,7 @@ async def test_request_transform_success(mock_max, mock_dataframe, mlflow):
'value': 1.0, 'created_at': '2024-01-01 12:00:00'}
],
'model_name': 'test_model',
'model_retention': 30
'model_config': {}
}
# Mock the transform response
@@ -108,26 +109,35 @@ async def test_request_transform_success(mock_max, mock_dataframe, mlflow):
# Verify the repository was called with correct arguments
mlflow.model_monitoring_repository.transform.assert_called_once_with(
'test_model', mock_dataframe, 30
'test_model', mock_dataframe, {}, metadata['metadata']
)
@mark.asyncio
@patch("laborious.activities.mlflow.DataFrame")
@patch("laborious.activities.mlflow.to_datetime")
@patch("laborious.activities.mlflow.max")
async def test_request_predict(mock_max, mock_dataframe, mlflow):
async def test_request_predict(mock_max, mock_to_datetime, mock_dataframe, mlflow):
mock_max.return_value = '2024-01-02'
# Mock input data
input_data = {
**metadata,
'data': [
{'timestamp': '2024-01-01', 'variable': 'var1', 'value': 1.0},
{'timestamp': '2024-01-01', 'variable': 'var2', 'value': 2.0},
{'timestamp': '2024-01-02', 'variable': 'var1', 'value': 3.0},
{'timestamp': '2024-01-02', 'variable': 'var2', 'value': 4.0}
],
'data': {
"variable": {
"2024-01-01": "var1",
"2024-01-02": "var2",
"2024-01-03": "var1",
"2024-01-04": "var2"
},
"value": {
"2024-01-01": 1.0,
"2024-01-02": 2.0,
"2024-01-03": 3.0,
"2024-01-04": 4.0
}
},
'model_name': 'test_model',
'model_retention': 30
'model_config': {}
}
# Mock the predict response
@@ -141,13 +151,26 @@ async def test_request_predict(mock_max, mock_dataframe, mlflow):
mock_dataframe.return_value.replace.assert_called_once_with(
np.nan, None, inplace=True
)
mock_dataframe.return_value.__setitem__.assert_any_call(
'timestamp', mock_to_datetime.return_value.dt.strftime.return_value
)
mock_dataframe.return_value.__setitem__.assert_any_call(
'timestamp', mock_to_datetime.return_value.dt.strftime.return_value
)
mock_to_datetime.assert_called_once_with(
mock_dataframe.return_value.__getitem__.return_value, format=DATETIME_FORMAT_WITH_TZ
)
mock_to_datetime.return_value.dt.strftime.assert_called_once_with(
DATETIME_FORMAT
)
# Verify the response
assert response_data == expected_response
# Verify the repository was called with correct arguments
mlflow.model_monitoring_repository.predict.assert_called_once_with(
'test_model', mock_dataframe.return_value, 30
'test_model', mock_dataframe.return_value, {}, metadata['metadata']
)

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

@@ -34,7 +34,9 @@ async def test_run(workflow_mock, prediction_process):
'mlflow_transform_filters': {'test': 'filter'},
'mlflow_predict_filters': {'test': 'filter'},
'model_name': 'test_model_name',
'model_retention': '30',
'model_config': {
'retention': '30'
},
'path_priority': ['continue', 'repeat', 'stop'],
'opc_output_config': {'test': 'config'},
'prediction_store_policy': 'lts:1'
@@ -62,54 +64,54 @@ async def test_run(workflow_mock, prediction_process):
workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.get_last_timestamp, {
**metadata,
'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, {
**metadata,
'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, {
**metadata,
'data': input_data['data'],
'model_name': input_data['model_name'],
'model_retention': input_data['model_retention'],
**metadata
'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'],
**metadata
}, 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'],
**metadata
}, 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_retention': input_data['model_retention'],
**metadata
'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'],
**metadata
}, retry_policy=ANY, start_to_close_timeout=ANY)])
workflow_mock.execute_child_workflow.assert_called_once_with(
@@ -122,11 +124,12 @@ async def test_run(workflow_mock, prediction_process):
'timestamp': '2024-01-01',
'model_id': 1,
'model_name': 'test_model_name',
'model_retention': '30',
'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'
'comment': 'Error',
'prediction_store_policy': input_data['prediction_store_policy']
}
)
@@ -146,7 +149,9 @@ 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_retention': '30',
'model_config': {
'retention': '30'
},
'path_priority': ['continue', 'repeat', 'stop'],
'opc_output_config': {'test': 'config'}
}
@@ -165,13 +170,13 @@ async def test_run_stop_at_input_gate(workflow_mock, prediction_process):
workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.get_last_timestamp, {
'data': input_data['data'],
**metadata
**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
**metadata,
}, retry_policy=ANY, start_to_close_timeout=ANY)
])
workflow_mock.execute_child_workflow.assert_not_called()
@@ -192,7 +197,9 @@ 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_retention': '30',
'model_config': {
'retention': '30'
},
'path_priority': ['continue', 'repeat', 'stop'],
'opc_output_config': {'test': 'config'}
}
@@ -213,7 +220,7 @@ async def test_run_stop_at_first_mlflow_response_gate(workflow_mock, prediction_
workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.get_last_timestamp, {
'data': input_data['data'],
**metadata
**metadata,
},
retry_policy=ANY, start_to_close_timeout=ANY)])
workflow_mock.execute_local_activity_method.assert_has_calls([
@@ -221,14 +228,14 @@ async def test_run_stop_at_first_mlflow_response_gate(workflow_mock, prediction_
'filters': input_data['input_filters'],
'data': input_data['data'],
'path_priority': input_data['path_priority'],
**metadata
**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_retention': input_data['model_retention'],
'model_config': input_data['model_config'],
**metadata
},
retry_policy=ANY, start_to_close_timeout=ANY)
@@ -261,7 +268,9 @@ 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_retention': '30',
'model_config': {
'retention': '30'
},
'path_priority': ['continue', 'repeat', 'stop'],
'opc_output_config': {'test': 'config'}
}
@@ -286,7 +295,7 @@ async def test_run_stop_at_mlflow_content_gate(workflow_mock, prediction_process
workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.get_last_timestamp, {
'data': input_data['data'],
**metadata
**metadata,
},
retry_policy=ANY, start_to_close_timeout=ANY)])
workflow_mock.execute_local_activity_method.assert_has_calls([
@@ -294,14 +303,14 @@ async def test_run_stop_at_mlflow_content_gate(workflow_mock, prediction_process
'filters': input_data['input_filters'],
'data': input_data['data'],
'path_priority': input_data['path_priority'],
**metadata
**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_retention': input_data['model_retention'],
'model_config': input_data['model_config'],
**metadata
}, retry_policy=ANY, start_to_close_timeout=ANY)])
workflow_mock.execute_local_activity_method.assert_has_calls([
@@ -310,7 +319,7 @@ async def test_run_stop_at_mlflow_content_gate(workflow_mock, prediction_process
'data': {'content': 'transformed_data', 'timestamp': '2024-01-01'},
'type': 'transform',
'path_priority': input_data['path_priority'],
**metadata
**metadata,
}, retry_policy=ANY, start_to_close_timeout=ANY)])
workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.mlflow_content_gate, {
@@ -318,7 +327,7 @@ async def test_run_stop_at_mlflow_content_gate(workflow_mock, prediction_process
'data': 'transformed_data',
'type': 'transform',
'path_priority': input_data['path_priority'],
**metadata
**metadata,
}, retry_policy=ANY, start_to_close_timeout=ANY)])
workflow_mock.execute_child_workflow.assert_not_called()
@@ -339,7 +348,9 @@ 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_retention': '30',
'model_config': {
'retention': '30'
},
'path_priority': ['continue', 'repeat', 'stop'],
'opc_output_config': {'test': 'config'}
}
@@ -365,7 +376,7 @@ async def test_run_stop_at_mlflow_last_response_gate(workflow_mock, prediction_p
workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.get_last_timestamp, {
'data': input_data['data'],
**metadata
**metadata,
},
retry_policy=ANY, start_to_close_timeout=ANY)])
workflow_mock.execute_local_activity_method.assert_has_calls([
@@ -373,14 +384,14 @@ async def test_run_stop_at_mlflow_last_response_gate(workflow_mock, prediction_p
'filters': input_data['input_filters'],
'data': input_data['data'],
'path_priority': input_data['path_priority'],
**metadata
**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_retention': input_data['model_retention'],
'model_config': input_data['model_config'],
**metadata
}, retry_policy=ANY, start_to_close_timeout=ANY)])
workflow_mock.execute_local_activity_method.assert_has_calls([
@@ -389,7 +400,7 @@ async def test_run_stop_at_mlflow_last_response_gate(workflow_mock, prediction_p
'data': {'content': 'transformed_data', 'timestamp': '2024-01-01'},
'type': 'transform',
'path_priority': input_data['path_priority'],
**metadata
**metadata,
}, retry_policy=ANY, start_to_close_timeout=ANY)])
workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.mlflow_content_gate, {
@@ -397,13 +408,13 @@ async def test_run_stop_at_mlflow_last_response_gate(workflow_mock, prediction_p
'data': 'transformed_data',
'type': 'transform',
'path_priority': input_data['path_priority'],
**metadata
**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_retention': input_data['model_retention'],
'model_config': input_data['model_config'],
**metadata
}, retry_policy=ANY, start_to_close_timeout=ANY)])
workflow_mock.execute_local_activity_method.assert_has_calls([
@@ -412,7 +423,7 @@ async def test_run_stop_at_mlflow_last_response_gate(workflow_mock, prediction_p
'data': {'content': 'predicted_data', 'timestamp': '2024-01-01'},
'type': 'predict',
'path_priority': input_data['path_priority'],
**metadata
**metadata,
}, retry_policy=ANY, start_to_close_timeout=ANY)])
workflow_mock.execute_child_workflow.assert_not_called()
@@ -429,7 +440,9 @@ 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_retention = '30'
model_config = {
'retention': '30'
}
# Act
result = await prediction_process.path_flag_handler(
@@ -440,7 +453,7 @@ async def test_path_flag_handler_stop(workflow_mock, prediction_process):
'model_id': model,
'last_timestamp': last_timestamp,
'model_name': model_name,
'model_retention': model_retention
'model_config': model_config
}, confidence, last_timestamp, ""
)
@@ -462,7 +475,9 @@ 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_retention = '30'
model_config = {
'retention': '30'
}
# Act
result = await prediction_process.path_flag_handler(
@@ -473,7 +488,7 @@ async def test_path_flag_handler_repeat(workflow_mock, prediction_process):
'model_id': model,
'last_timestamp': last_timestamp,
'model_name': model_name,
'model_retention': model_retention
'model_config': model_config
}, confidence, last_timestamp, ""
)
@@ -506,7 +521,10 @@ 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_retention = '30'
model_config = {
'retention': '30'
}
prediction_store_policy = 'erl:1'
# Act
result = await prediction_process.path_flag_handler(
@@ -517,8 +535,9 @@ async def test_path_flag_handler_continue(workflow_mock, prediction_process):
'model_id': model,
'last_timestamp': last_timestamp,
'model_name': model_name,
'model_retention': model_retention,
'opc_output_config': {'test': 'config'}
'model_config': model_config,
'opc_output_config': {'test': 'config'},
'prediction_store_policy': prediction_store_policy
}, confidence, last_timestamp, 'Prediction Process'
)
@@ -535,11 +554,12 @@ async def test_path_flag_handler_continue(workflow_mock, prediction_process):
'timestamp': last_timestamp,
'model_id': model,
'model_name': model_name,
'model_retention': model_retention,
'model_config': model_config,
'schema': schema,
'table_name': table_name,
'comment': 'Prediction Process',
'opc_output_config': {'test': 'config'}
'opc_output_config': {'test': 'config'},
'prediction_store_policy': prediction_store_policy
}
)
@@ -556,8 +576,10 @@ 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_retention = '30'
model_config = {
'retention': '30'
}
prediction_store_policy = 'erl:1'
# Act
result = await prediction_process.path_flag_handler(
data, path_flag, {
@@ -567,8 +589,9 @@ async def test_path_flag_handler_unknown(workflow_mock, prediction_process):
'model_id': model,
'last_timestamp': last_timestamp,
'model_name': model_name,
'model_retention': model_retention,
'opc_output_config': {'test': 'config'}
'model_config': model_config,
'opc_output_config': {'test': 'config'},
'prediction_store_policy': prediction_store_policy
}, confidence, last_timestamp, ""
)

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

@@ -33,7 +33,11 @@ async def test_run(workflow_mock: AsyncMock, predictions_batch: PredictionsBatch
'schema': 'test_schema',
'table_name': 'test_table',
'opc_output_config': 'test_opc_output_config',
'datetime_columns': ['timestamp', 'created_at']
'datetime_columns': ['timestamp', 'created_at'],
'prediction_store_policy': 'erl:1',
'model_config': {
'retention': '30'
}
}
await predictions_batch.run(input_data)
@@ -72,9 +76,10 @@ async def test_run(workflow_mock: AsyncMock, predictions_batch: PredictionsBatch
'POLICY': 'STOP'
}
}),
'model_retention': input_data.get('model_retention', 60),
'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', {})
'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([