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

SIENTIAPDE-1273

Browse Source 
Enhance MLFlowRepository and Activities classes with new methods and metrics

- Added `check_artifact_exists` method to MLFlowRepository for verifying artifact presence in the MLflow Model Registry.
- Implemented `get_prediction_data` method in MLFlowRepository to retrieve prediction data from models.
- Updated Activities class to integrate ModelMetrics for improved metrics handling.
- Enhanced tests for artifact existence checks and prediction data retrieval, ensuring robust coverage for new functionalities.
- Updated various workflows to include `transform_table_name` in input data for better data handling.
This commit is contained in:
Vitor Pimentel dos Santos
2025-11-13 16:40:23 -03:00
parent 6ac0f38d59
commit b68674fe64
15 changed files with 1639 additions and 68 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -31,6 +31,7 @@ async def test_run(workflow_mock, prediction_process):
'data': {'test': 'data'},
'schema': 'test_schema',
'table_name': 'test_table',
'transform_table_name': 'test_transform_table',
'model_id': 1,
'input_filters': {'test': 'filter'},
'mlflow_transform_filters': {'test': 'filter'},
@@ -175,6 +176,7 @@ async def test_run(workflow_mock, prediction_process):
'metadata': metadata,
'path_flag': 'continue',
'data': 'predicted_data',
'transformed_data': 'transformed_data',
'prediction_confidence': 0.95,
'timestamp': '2024-01-01',
'model_id': 1,
@@ -183,6 +185,7 @@ async def test_run(workflow_mock, prediction_process):
'opc_output_config': input_data['opc_output_config'],
'schema': input_data['schema'],
'table_name': input_data['table_name'],
'transform_table_name': input_data['transform_table_name'],
'comment': 'Error',
'prediction_store_policy': input_data['prediction_store_policy'],
},
@@ -199,6 +202,7 @@ async def test_run_stop_at_input_gate(workflow_mock, prediction_process):
'data': {'test': 'data'},
'schema': 'test_schema',
'table_name': 'test_table',
'transform_table_name': 'test_transform_table',
'model_id': 1,
'input_filters': {'test': 'filter'},
'mlflow_transform_filters': {'test': 'filter'},
@@ -257,6 +261,7 @@ async def test_run_stop_at_first_mlflow_response_gate(workflow_mock, prediction_
'data': {'test': 'data'},
'schema': 'test_schema',
'table_name': 'test_table',
'transform_table_name': 'test_transform_table',
'model_id': 1,
'input_filters': {'test': 'filter'},
'mlflow_transform_filters': {'test': 'filter'},
@@ -352,6 +357,7 @@ async def test_run_stop_at_mlflow_content_gate(workflow_mock, prediction_process
'data': {'test': 'data'},
'schema': 'test_schema',
'table_name': 'test_table',
'transform_table_name': 'test_transform_table',
'model_id': 1,
'input_filters': {'test': 'filter'},
'mlflow_transform_filters': {'test': 'filter'},
@@ -467,6 +473,7 @@ async def test_run_stop_at_mlflow_last_response_gate(workflow_mock, prediction_p
'data': {'test': 'data'},
'schema': 'test_schema',
'table_name': 'test_table',
'transform_table_name': 'test_transform_table',
'model_id': 1,
'input_filters': {'test': 'filter'},
'mlflow_transform_filters': {'test': 'filter'},
@@ -626,6 +633,7 @@ async def test_path_flag_handler_stop(workflow_mock, prediction_process):
'metadata': metadata,
'schema': schema,
'table_name': table_name,
'transform_table_name': 'test_transform_table',
'model_id': model,
'last_timestamp': last_timestamp,
'model_name': model_name,
@@ -664,6 +672,7 @@ async def test_path_flag_handler_repeat(workflow_mock, prediction_process):
'metadata': metadata,
'schema': schema,
'table_name': table_name,
'transform_table_name': 'test_transform_table',
'model_id': model,
'last_timestamp': last_timestamp,
'model_name': model_name,
@@ -714,6 +723,7 @@ async def test_path_flag_handler_continue(workflow_mock, prediction_process):
'metadata': metadata,
'schema': schema,
'table_name': table_name,
'transform_table_name': 'test_transform_table',
'model_id': model,
'last_timestamp': last_timestamp,
'model_name': model_name,
@@ -742,6 +752,7 @@ async def test_path_flag_handler_continue(workflow_mock, prediction_process):
'model_config': model_config,
'schema': schema,
'table_name': table_name,
'transform_table_name': 'test_transform_table',
'comment': 'Prediction Process',
'opc_output_config': {'test': 'config'},
'prediction_store_policy': prediction_store_policy,
@@ -771,6 +782,7 @@ async def test_path_flag_handler_unknown(workflow_mock, prediction_process):
**metadata,
'schema': schema,
'table_name': table_name,
'transform_table_name': 'test_transform_table',
'model_id': model,
'last_timestamp': last_timestamp,
'model_name': model_name,

247
tests/laborious/workflows/test_drift.py Normal file
View File

@@ -0,0 +1,247 @@
from unittest.mock import ANY, AsyncMock, call, patch
from pytest import fixture, mark
from laborious.activities.activities import Activities
from laborious.workflows.drift import Drift
from sientia_do.temporal.constants import DATETIME_FORMAT_WITH_TZ
@fixture
def drift() -> Drift:
return Drift()
metadata = {
'metadata': {
'model_id': 'test_model_id',
'model_name': 'test_model',
'workflow_name': 'drift',
'schedule_name': 'test_schedule',
},
}
@mark.asyncio
@patch('laborious.workflows.drift.workflow', new_callable=AsyncMock)
async def test_run(workflow_mock: AsyncMock, drift: Drift):
# Arrange
input_data = {
'schedule_name': 'test_schedule',
'model_name': 'test_model',
'model_id': 'test_model_id',
'schema': 'test_schema',
'source_table_name': 'test_source_table',
'target_table_name': 'test_target_table',
'interval': 60,
'target_name': 'test_target',
'drift_metrics': ['psi', 'ks'],
'chunk_period': 'hour',
}
target_data = {'data': 'test_target_data'}
reference_data = {'data': 'test_reference_data'}
drift_data = {'drift': 'test_drift_data'}
workflow_mock.start_local_activity_method.side_effect = [
target_data, reference_data
]
workflow_mock.execute_local_activity_method.return_value = drift_data
workflow_mock.execute_activity_method = AsyncMock()
# Act
await drift.run(input_data)
# Assert - Check start_local_activity_method calls
expected_gathering_query = f"""
SELECT *
FROM {input_data['schema']}.{input_data['source_table_name']}
WHERE
model_id = {input_data['model_id']} AND
timestamp > NOW() - INTERVAL '{input_data['interval']} minutes'
ORDER BY timestamp ASC
"""
workflow_mock.start_local_activity_method.assert_has_calls(
[
call(
Activities.load_custom_query,
{
**metadata,
'query': expected_gathering_query,
'datetime_columns': ['timestamp', 'created_at'],
},
retry_policy=ANY,
start_to_close_timeout=ANY,
),
call(
Activities.get_reference_data,
{
**metadata,
'model_name': input_data['model_name'],
},
retry_policy=ANY,
start_to_close_timeout=ANY,
),
]
)
# Assert - Check calculate_drift call
workflow_mock.execute_local_activity_method.assert_called_once_with(
Activities.calculate_drift,
{
**metadata,
'target_data': target_data,
'reference_data': reference_data,
'model_name': input_data['model_name'],
'model_id': input_data['model_id'],
'target_name': input_data['target_name'],
'drift_metrics': input_data['drift_metrics'],
'chunk_period': input_data['chunk_period'],
},
retry_policy=ANY,
start_to_close_timeout=ANY,
)
# Assert - Check export_data_to_postgres call
workflow_mock.execute_activity_method.assert_called_once_with(
Activities.export_data_to_postgres,
{
**metadata,
'data': drift_data,
'schema': input_data['schema'],
'table_name': input_data['target_table_name'],
'timestamp_conversion': {'column': 'timestamp', 'format': DATETIME_FORMAT_WITH_TZ},
},
retry_policy=ANY,
start_to_close_timeout=ANY,
)
@mark.asyncio
@patch('laborious.workflows.drift.workflow', new_callable=AsyncMock)
async def test_run_empty_target_data(workflow_mock: AsyncMock, drift: Drift):
# Arrange
input_data = {
'schedule_name': 'test_schedule',
'model_name': 'test_model',
'model_id': 'test_model_id',
'schema': 'test_schema',
'source_table_name': 'test_source_table',
'target_table_name': 'test_target_table',
'interval': 60,
'target_name': 'test_target',
'drift_metrics': ['psi', 'ks'],
}
target_data = None
reference_data = {'data': 'test_reference_data'}
workflow_mock.start_local_activity_method.side_effect = [target_data, reference_data]
workflow_mock.execute_local_activity_method = AsyncMock()
workflow_mock.execute_activity_method = AsyncMock()
# Act
await drift.run(input_data)
# Assert - Should not call calculate_drift or export
workflow_mock.execute_local_activity_method.assert_not_called()
workflow_mock.execute_activity_method.assert_not_called()
@mark.asyncio
@patch('laborious.workflows.drift.workflow', new_callable=AsyncMock)
async def test_run_empty_drift_data(workflow_mock: AsyncMock, drift: Drift):
# Arrange
input_data = {
'schedule_name': 'test_schedule',
'model_name': 'test_model',
'model_id': 'test_model_id',
'schema': 'test_schema',
'source_table_name': 'test_source_table',
'target_table_name': 'test_target_table',
'interval': 60,
'target_name': 'test_target',
'drift_metrics': ['psi', 'ks'],
}
target_data = {'data': 'test_target_data'}
reference_data = {'data': 'test_reference_data'}
drift_data = None
workflow_mock.start_local_activity_method.side_effect = [target_data, reference_data]
workflow_mock.execute_local_activity_method.return_value = drift_data
workflow_mock.execute_activity_method = AsyncMock()
# Act
await drift.run(input_data)
# Assert - Should call calculate_drift but not export
workflow_mock.execute_local_activity_method.assert_called_once_with(
Activities.calculate_drift,
{
**metadata,
'target_data': target_data,
'reference_data': reference_data,
'model_name': input_data['model_name'],
'model_id': input_data['model_id'],
'target_name': input_data['target_name'],
'drift_metrics': input_data['drift_metrics'],
'chunk_period': input_data.get('chunk_period', 'min'),
},
retry_policy=ANY,
start_to_close_timeout=ANY,
)
workflow_mock.execute_activity_method.assert_not_called()
@mark.asyncio
@patch('laborious.workflows.drift.workflow', new_callable=AsyncMock)
async def test_run_default_chunk_period(workflow_mock: AsyncMock, drift: Drift):
# Arrange
input_data = {
'schedule_name': 'test_schedule',
'model_name': 'test_model',
'model_id': 'test_model_id',
'schema': 'test_schema',
'source_table_name': 'test_source_table',
'target_table_name': 'test_target_table',
'interval': 60,
'target_name': 'test_target',
'drift_metrics': ['psi', 'ks'],
# chunk_period not provided, should default to 'min'
}
target_data = {'data': 'test_target_data'}
reference_data = {'data': 'test_reference_data'}
drift_data = {'drift': 'test_drift_data'}
workflow_mock.start_local_activity_method.side_effect = [target_data, reference_data]
workflow_mock.execute_local_activity_method.return_value = drift_data
workflow_mock.execute_activity_method = AsyncMock()
# Act
await drift.run(input_data)
# Assert - Check calculate_drift call with default chunk_period
workflow_mock.execute_local_activity_method.assert_called_once_with(
Activities.calculate_drift,
{
**metadata,
'target_data': target_data,
'reference_data': reference_data,
'model_name': input_data['model_name'],
'model_id': input_data['model_id'],
'target_name': input_data['target_name'],
'drift_metrics': input_data['drift_metrics'],
'chunk_period': 'min', # Default value
},
retry_policy=ANY,
start_to_close_timeout=ANY,
)

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

@@ -32,6 +32,7 @@ async def test_run(workflow_mock: AsyncMock, predictions_batch: PredictionsBatch
'query': 'SELECT * FROM test',
'schema': 'test_schema',
'table_name': 'test_table',
'transform_table_name': 'test_transform_table',
'opc_output_config': 'test_opc_output_config',
'datetime_columns': ['timestamp', 'created_at'],
'prediction_store_policy': 'erl:1',
@@ -59,6 +60,7 @@ async def test_run(workflow_mock: AsyncMock, predictions_batch: PredictionsBatch
'data': {'data': 'test_data'},
'schema': input_data['schema'],
'table_name': input_data['table_name'],
'transform_table_name': input_data['transform_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'}}),
@@ -71,7 +73,7 @@ async def test_run(workflow_mock: AsyncMock, predictions_batch: PredictionsBatch
'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'),
'prediction_store_policy': input_data.get('prediction_store_policy', 'lts:1'),
}
workflow_mock.execute_child_workflow.assert_has_calls(