gitea_admin/sientia-dataops-laborious_temporal
1
0
Fork 0
Code Issues Pull Requests Actions 1 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 b279244160
commit 673fc79df3
4 changed files with 111 additions and 69 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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([