SIENTIAPDE-1222

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.
2025-09-17 16:02:54 -03:00
parent d6fdbc58bb
commit 84c1371d6c
4 changed files with 111 additions and 69 deletions
--- a/tests/laborious/activities/test_mlflow.py
+++ b/tests/laborious/activities/test_mlflow.py
@@ -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']
    )


--- a/tests/laborious/workflows/subworkflows/test_prediction_process.py
+++ b/tests/laborious/workflows/subworkflows/test_prediction_process.py
@@ -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, ""
    )

--- a/tests/laborious/workflows/test_predictions_batch.py
+++ b/tests/laborious/workflows/test_predictions_batch.py
@@ -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([