diff --git a/orchestrator/utils/orchestrator_functions.py b/orchestrator/utils/orchestrator_functions.py index b49a771..90bfdfb 100644 --- a/orchestrator/utils/orchestrator_functions.py +++ b/orchestrator/utils/orchestrator_functions.py @@ -211,7 +211,8 @@ def predictions_batch(config: dict[str, Any]): "config": {} } }, config.get('mlflow_predict_filters', [])), - "path_priority": path_priority + "path_priority": path_priority, + "predictions_storage_policy": config.get('predictions_storage_policy', 'lts:1') } diff --git a/tests/orchestrator/utils/test_orchestrator_functions.py b/tests/orchestrator/utils/test_orchestrator_functions.py index 09fae10..cd59478 100644 --- a/tests/orchestrator/utils/test_orchestrator_functions.py +++ b/tests/orchestrator/utils/test_orchestrator_functions.py @@ -197,7 +197,8 @@ def test_predictions_batch(mock_process_path_priority, } ], "path_priority": ["STOP", "CONTINUE", "REPEAT"], - "datetime_columns": ["timestamp"] + "datetime_columns": ["timestamp"], + "predictions_storage_policy": "erl:1" } result = predictions_batch(config) @@ -272,7 +273,8 @@ def test_predictions_batch(mock_process_path_priority, } }, "path_priority": ["STOP", "CONTINUE", "REPEAT"], - "datetime_columns": ["timestamp"] + "datetime_columns": ["timestamp"], + "predictions_storage_policy": "erl:1" } assert result == expected