vitor-aignosi
84e6501063
Implement new scheduling configurations for model retraining and drift analysis in input_sample.json - Added three new schedule configurations: `minimal-retrain-test-runtime`, `drift-test-runtime`, and `simple-metrics-test-runtime`. - Each configuration includes parameters such as model ID, workflow type, frequency, and specific queries for data retrieval. - Enhanced the structure to support active status and updated timestamps for better tracking of schedule states.
489 lines
18 KiB
Python
489 lines
18 KiB
Python
"""
|
|
End-to-end tests for PredictionsBatch workflow - Prediction Process scenarios.
|
|
"""
|
|
|
|
from decimal import Decimal
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
import numpy as np
|
|
import pandas as pd
|
|
import pytest
|
|
from sqlalchemy import text
|
|
from temporalio.testing import WorkflowEnvironment
|
|
from temporalio.worker import Worker
|
|
|
|
from e2e.helpers import (
|
|
assert_continue,
|
|
assert_postgres_unique_violation_in_chain,
|
|
assert_prediction,
|
|
assert_prediction_row_count,
|
|
assert_repeat,
|
|
assert_stop,
|
|
insert_sample_data,
|
|
insert_sample_prediction,
|
|
load_scenario_input,
|
|
make_workflow_id,
|
|
start_and_await_workflow,
|
|
)
|
|
from laborious.activities.activities import Activities
|
|
from laborious.utils.models import minio_dataframe_payload as minio_payload_module
|
|
from laborious.workflows.predictions_batch import PredictionsBatch
|
|
|
|
DISTINCT_BATCH_TIMESTAMP = '2024-01-01 13:00:00+00:00'
|
|
HISTORY_TIMESTAMP = '2024-01-01 12:00:00+00:00'
|
|
|
|
|
|
def get_base_input_data(model_id):
|
|
return load_scenario_input('prediction_process_base.json', model_id=model_id)
|
|
|
|
|
|
@pytest.fixture
|
|
def bad_data_model(mlflow_repository_stub):
|
|
mlflow_repository_stub.stub_wrapper.transform = MagicMock(
|
|
side_effect=Exception('Bad data model')
|
|
)
|
|
return mlflow_repository_stub.stub_wrapper
|
|
|
|
|
|
@pytest.fixture
|
|
def bad_predict_model(mlflow_repository_stub):
|
|
wrapper = mlflow_repository_stub.stub_wrapper
|
|
|
|
def _good_transform(data):
|
|
result = pd.DataFrame(
|
|
{
|
|
'feature_1': [0.234] * len(data),
|
|
'feature_2': [0.783] * len(data),
|
|
}
|
|
)
|
|
result.index = data.index
|
|
return result, {}
|
|
|
|
wrapper.transform.side_effect = _good_transform
|
|
wrapper.predict = MagicMock(side_effect=Exception('Bad predict model'))
|
|
return wrapper
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.integration
|
|
async def test_scenario_2_1_1_input_gate_triggers_continue(
|
|
temporal_test_env: WorkflowEnvironment,
|
|
temporal_worker: Worker,
|
|
test_activities: Activities,
|
|
postgres_engine,
|
|
mlflow_repository_stub,
|
|
):
|
|
"""Input gate CONTINUE: export default prediction; MLflow transform/predict not used."""
|
|
client = temporal_test_env.client
|
|
model_id = 211
|
|
insert_sample_data(postgres_engine, model_id, ['NULL', 78.2])
|
|
input_data = get_base_input_data(model_id)
|
|
await start_and_await_workflow(
|
|
client, PredictionsBatch.run, input_data, make_workflow_id('test-continue-policy')
|
|
)
|
|
assert_continue(postgres_engine, model_id)
|
|
mlflow_repository_stub.stub_wrapper.transform.assert_not_called()
|
|
mlflow_repository_stub.stub_wrapper.predict.assert_not_called()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.integration
|
|
async def test_scenario_2_1_2_input_gate_triggers_stop(
|
|
temporal_test_env: WorkflowEnvironment,
|
|
temporal_worker: Worker,
|
|
test_activities: Activities,
|
|
postgres_engine,
|
|
mlflow_repository_stub,
|
|
):
|
|
"""Input gate STOP: no export, no MLflow."""
|
|
client = temporal_test_env.client
|
|
model_id = 212
|
|
insert_sample_data(postgres_engine, model_id, ['NULL', 78.2])
|
|
input_data = get_base_input_data(model_id)
|
|
input_data['input_filters']['SPECIFIC_VARIABLES_NULL_VALUES']['POLICY'] = 'STOP'
|
|
await start_and_await_workflow(
|
|
client, PredictionsBatch.run, input_data, make_workflow_id('test-input-stop')
|
|
)
|
|
assert_stop(postgres_engine, model_id)
|
|
mlflow_repository_stub.stub_wrapper.transform.assert_not_called()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.integration
|
|
async def test_scenario_2_1_3_input_gate_repeat_batch_timestamp_equals_history_fails(
|
|
temporal_test_env: WorkflowEnvironment,
|
|
temporal_worker: Worker,
|
|
test_activities: Activities,
|
|
postgres_engine,
|
|
mlflow_repository_stub,
|
|
):
|
|
"""
|
|
REPEAT uses ``last_timestamp`` from the batch payload as the new row's ``timestamp``.
|
|
When it equals the only historical prediction row, Postgres rejects the duplicate key.
|
|
"""
|
|
client = temporal_test_env.client
|
|
model_id = 213
|
|
insert_sample_data(postgres_engine, model_id, ['NULL', 78.2], data_timestamp=HISTORY_TIMESTAMP)
|
|
insert_sample_prediction(postgres_engine, model_id, prediction_timestamp=HISTORY_TIMESTAMP)
|
|
input_data = get_base_input_data(model_id)
|
|
input_data['input_filters']['SPECIFIC_VARIABLES_NULL_VALUES']['POLICY'] = 'REPEAT'
|
|
with pytest.raises(Exception) as excinfo:
|
|
await start_and_await_workflow(
|
|
client, PredictionsBatch.run, input_data, make_workflow_id('test-input-repeat-collision')
|
|
)
|
|
assert_postgres_unique_violation_in_chain(excinfo.value)
|
|
assert_prediction_row_count(postgres_engine, model_id, 1)
|
|
mlflow_repository_stub.stub_wrapper.transform.assert_not_called()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.integration
|
|
async def test_scenario_2_1_3_input_gate_repeat_distinct_batch_timestamp_inserts_second_row(
|
|
temporal_test_env: WorkflowEnvironment,
|
|
temporal_worker: Worker,
|
|
test_activities: Activities,
|
|
postgres_engine,
|
|
mlflow_repository_stub,
|
|
):
|
|
"""REPEAT succeeds when batch ``last_timestamp`` differs from the historical prediction row."""
|
|
client = temporal_test_env.client
|
|
model_id = 2131
|
|
insert_sample_data(
|
|
postgres_engine, model_id, ['NULL', 78.2], data_timestamp=DISTINCT_BATCH_TIMESTAMP
|
|
)
|
|
data = insert_sample_prediction(postgres_engine, model_id, prediction_timestamp=HISTORY_TIMESTAMP)
|
|
input_data = get_base_input_data(model_id)
|
|
input_data['input_filters']['SPECIFIC_VARIABLES_NULL_VALUES']['POLICY'] = 'REPEAT'
|
|
await start_and_await_workflow(
|
|
client, PredictionsBatch.run, input_data, make_workflow_id('test-input-repeat-ok')
|
|
)
|
|
assert_repeat(postgres_engine, model_id, data)
|
|
mlflow_repository_stub.stub_wrapper.transform.assert_not_called()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.integration
|
|
async def test_scenario_2_1_4_input_gate_repeat_without_prior_prediction(
|
|
temporal_test_env: WorkflowEnvironment,
|
|
temporal_worker: Worker,
|
|
test_activities: Activities,
|
|
postgres_engine,
|
|
):
|
|
"""REPEAT when no prior row in predictions: repeat_last_prediction runs; still no new duplicate export path."""
|
|
client = temporal_test_env.client
|
|
model_id = 214
|
|
insert_sample_data(postgres_engine, model_id, ['NULL', 78.2])
|
|
with postgres_engine.begin() as conn:
|
|
conn.execute(text(f'DELETE FROM sientia_data.predictions WHERE model_id = {model_id}'))
|
|
input_data = get_base_input_data(model_id)
|
|
input_data['input_filters']['SPECIFIC_VARIABLES_NULL_VALUES']['POLICY'] = 'REPEAT'
|
|
await start_and_await_workflow(
|
|
client, PredictionsBatch.run, input_data, make_workflow_id('test-input-repeat-no-history')
|
|
)
|
|
assert_stop(postgres_engine, model_id)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.integration
|
|
async def test_scenario_2_2_1_transform_gate_triggers_continue(
|
|
temporal_test_env: WorkflowEnvironment,
|
|
temporal_worker: Worker,
|
|
test_activities: Activities,
|
|
postgres_engine,
|
|
bad_data_model,
|
|
):
|
|
client = temporal_test_env.client
|
|
model_id = 221
|
|
insert_sample_data(postgres_engine, model_id, [60.0, 78.2])
|
|
input_data = get_base_input_data(model_id)
|
|
input_data['mlflow_transform_filters']['API_ERROR']['POLICY'] = 'CONTINUE'
|
|
await start_and_await_workflow(
|
|
client, PredictionsBatch.run, input_data, make_workflow_id('test-transform-continue')
|
|
)
|
|
assert_continue(
|
|
postgres_engine=postgres_engine,
|
|
model_id=model_id,
|
|
prediction_confidence=Decimal(10),
|
|
comments='Unknown MLFlow API error',
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.integration
|
|
async def test_scenario_2_2_2_transform_gate_triggers_stop(
|
|
temporal_test_env: WorkflowEnvironment,
|
|
temporal_worker: Worker,
|
|
test_activities: Activities,
|
|
postgres_engine,
|
|
bad_data_model,
|
|
mlflow_repository_stub,
|
|
):
|
|
client = temporal_test_env.client
|
|
model_id = 222
|
|
insert_sample_data(postgres_engine, model_id, [60.0, 78.2])
|
|
input_data = get_base_input_data(model_id)
|
|
input_data['mlflow_transform_filters']['API_ERROR']['POLICY'] = 'STOP'
|
|
await start_and_await_workflow(
|
|
client, PredictionsBatch.run, input_data, make_workflow_id('test-transform-stop')
|
|
)
|
|
assert_stop(postgres_engine, model_id)
|
|
mlflow_repository_stub.stub_wrapper.predict.assert_not_called()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.integration
|
|
async def test_scenario_2_2_3_transform_gate_repeat_batch_timestamp_equals_history_fails(
|
|
temporal_test_env: WorkflowEnvironment,
|
|
temporal_worker: Worker,
|
|
test_activities: Activities,
|
|
postgres_engine,
|
|
bad_data_model,
|
|
):
|
|
client = temporal_test_env.client
|
|
model_id = 223
|
|
insert_sample_data(postgres_engine, model_id, [60.0, 78.2], data_timestamp=HISTORY_TIMESTAMP)
|
|
insert_sample_prediction(postgres_engine, model_id, prediction_timestamp=HISTORY_TIMESTAMP)
|
|
input_data = get_base_input_data(model_id)
|
|
input_data['mlflow_transform_filters']['API_ERROR']['POLICY'] = 'REPEAT'
|
|
with pytest.raises(Exception) as excinfo:
|
|
await start_and_await_workflow(
|
|
client, PredictionsBatch.run, input_data, make_workflow_id('test-transform-repeat-collision')
|
|
)
|
|
assert_postgres_unique_violation_in_chain(excinfo.value)
|
|
assert_prediction_row_count(postgres_engine, model_id, 1)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.integration
|
|
async def test_scenario_2_2_3_transform_gate_repeat_distinct_batch_timestamp_inserts_second_row(
|
|
temporal_test_env: WorkflowEnvironment,
|
|
temporal_worker: Worker,
|
|
test_activities: Activities,
|
|
postgres_engine,
|
|
bad_data_model,
|
|
):
|
|
client = temporal_test_env.client
|
|
model_id = 2231
|
|
insert_sample_data(
|
|
postgres_engine, model_id, [60.0, 78.2], data_timestamp=DISTINCT_BATCH_TIMESTAMP
|
|
)
|
|
data = insert_sample_prediction(postgres_engine, model_id, prediction_timestamp=HISTORY_TIMESTAMP)
|
|
input_data = get_base_input_data(model_id)
|
|
input_data['mlflow_transform_filters']['API_ERROR']['POLICY'] = 'REPEAT'
|
|
await start_and_await_workflow(
|
|
client, PredictionsBatch.run, input_data, make_workflow_id('test-transform-repeat-ok')
|
|
)
|
|
assert_repeat(postgres_engine, model_id, data)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.integration
|
|
async def test_scenario_2_2_4_transform_content_gate_nan_values_stop(
|
|
temporal_test_env: WorkflowEnvironment,
|
|
temporal_worker: Worker,
|
|
test_activities: Activities,
|
|
postgres_engine,
|
|
mlflow_repository_stub,
|
|
):
|
|
"""mlflow_content_gate triggers STOP when transform output is all NaN (NAN_VALUES filter)."""
|
|
client = temporal_test_env.client
|
|
model_id = 224
|
|
|
|
def all_nan_transform(data):
|
|
result = pd.DataFrame(
|
|
{'feature_1': [np.nan] * len(data), 'feature_2': [np.nan] * len(data)}
|
|
)
|
|
result.index = data.index
|
|
return result, {}
|
|
|
|
mlflow_repository_stub.stub_wrapper.transform = MagicMock(side_effect=all_nan_transform)
|
|
|
|
insert_sample_data(postgres_engine, model_id, [60.0, 78.2])
|
|
input_data = get_base_input_data(model_id)
|
|
input_data['mlflow_transform_filters'] = {
|
|
'API_ERROR': {'POLICY': 'STOP', 'CONFIG': {}},
|
|
'NAN_VALUES': {'POLICY': 'STOP', 'CONFIG': {}},
|
|
}
|
|
await start_and_await_workflow(
|
|
client, PredictionsBatch.run, input_data, make_workflow_id('test-transform-content-stop')
|
|
)
|
|
assert_stop(postgres_engine, model_id)
|
|
mlflow_repository_stub.stub_wrapper.predict.assert_not_called()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.integration
|
|
async def test_scenario_2_3_1_predict_gate_triggers_continue(
|
|
temporal_test_env: WorkflowEnvironment,
|
|
temporal_worker: Worker,
|
|
test_activities: Activities,
|
|
postgres_engine,
|
|
bad_predict_model,
|
|
):
|
|
client = temporal_test_env.client
|
|
model_id = 231
|
|
insert_sample_data(postgres_engine, model_id, [60.0, 78.2])
|
|
input_data = get_base_input_data(model_id)
|
|
input_data['mlflow_predict_filters']['API_ERROR']['POLICY'] = 'CONTINUE'
|
|
await start_and_await_workflow(
|
|
client, PredictionsBatch.run, input_data, make_workflow_id('test-predict-continue')
|
|
)
|
|
assert_continue(
|
|
postgres_engine=postgres_engine,
|
|
model_id=model_id,
|
|
prediction_confidence=Decimal(10),
|
|
comments='Unknown MLFlow API error',
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.integration
|
|
async def test_scenario_2_3_2_predict_gate_triggers_stop(
|
|
temporal_test_env: WorkflowEnvironment,
|
|
temporal_worker: Worker,
|
|
test_activities: Activities,
|
|
postgres_engine,
|
|
bad_predict_model,
|
|
):
|
|
client = temporal_test_env.client
|
|
model_id = 232
|
|
insert_sample_data(postgres_engine, model_id, [23.5, 78.2])
|
|
input_data = get_base_input_data(model_id)
|
|
input_data['mlflow_predict_filters']['API_ERROR']['POLICY'] = 'STOP'
|
|
await start_and_await_workflow(
|
|
client, PredictionsBatch.run, input_data, make_workflow_id('test-predict-stop')
|
|
)
|
|
assert_stop(postgres_engine, model_id)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.integration
|
|
async def test_scenario_2_3_3_predict_gate_repeat_batch_timestamp_equals_history_fails(
|
|
temporal_test_env: WorkflowEnvironment,
|
|
temporal_worker: Worker,
|
|
test_activities: Activities,
|
|
postgres_engine,
|
|
bad_predict_model,
|
|
):
|
|
client = temporal_test_env.client
|
|
model_id = 233
|
|
insert_sample_data(postgres_engine, model_id, [23.5, 78.2], data_timestamp=HISTORY_TIMESTAMP)
|
|
insert_sample_prediction(postgres_engine, model_id, prediction_timestamp=HISTORY_TIMESTAMP)
|
|
input_data = get_base_input_data(model_id)
|
|
input_data['mlflow_predict_filters']['API_ERROR']['POLICY'] = 'REPEAT'
|
|
input_data['path_priority'] = ['REPEAT', 'STOP', 'CONTINUE']
|
|
with pytest.raises(Exception) as excinfo:
|
|
await start_and_await_workflow(
|
|
client, PredictionsBatch.run, input_data, make_workflow_id('test-predict-repeat-collision')
|
|
)
|
|
assert_postgres_unique_violation_in_chain(excinfo.value)
|
|
assert_prediction_row_count(postgres_engine, model_id, 1)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.integration
|
|
async def test_scenario_2_3_3_predict_gate_repeat_distinct_batch_timestamp_inserts_second_row(
|
|
temporal_test_env: WorkflowEnvironment,
|
|
temporal_worker: Worker,
|
|
test_activities: Activities,
|
|
postgres_engine,
|
|
bad_predict_model,
|
|
):
|
|
client = temporal_test_env.client
|
|
model_id = 2331
|
|
insert_sample_data(
|
|
postgres_engine, model_id, [23.5, 78.2], data_timestamp=DISTINCT_BATCH_TIMESTAMP
|
|
)
|
|
data = insert_sample_prediction(postgres_engine, model_id, prediction_timestamp=HISTORY_TIMESTAMP)
|
|
input_data = get_base_input_data(model_id)
|
|
input_data['mlflow_predict_filters']['API_ERROR']['POLICY'] = 'REPEAT'
|
|
input_data['path_priority'] = ['REPEAT', 'STOP', 'CONTINUE']
|
|
await start_and_await_workflow(
|
|
client, PredictionsBatch.run, input_data, make_workflow_id('test-predict-repeat-ok')
|
|
)
|
|
assert_repeat(postgres_engine, model_id, data)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.integration
|
|
async def test_scenario_2_4_1_input_empty_data_stop(
|
|
temporal_test_env: WorkflowEnvironment,
|
|
temporal_worker: Worker,
|
|
test_activities: Activities,
|
|
postgres_engine,
|
|
):
|
|
"""EMPTY_DATA filter with STOP when query returns no rows (offload payload empty)."""
|
|
client = temporal_test_env.client
|
|
model_id = 241
|
|
with postgres_engine.begin() as conn:
|
|
conn.execute(text(f'DELETE FROM sientia_data.laborious_data WHERE model_id = {model_id}'))
|
|
input_data = get_base_input_data(model_id)
|
|
input_data['input_filters'] = {'EMPTY_DATA': {'POLICY': 'STOP', 'CONFIG': {}}}
|
|
await start_and_await_workflow(
|
|
client, PredictionsBatch.run, input_data, make_workflow_id('test-empty-data-stop')
|
|
)
|
|
assert_stop(postgres_engine, model_id)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.integration
|
|
async def test_scenario_2_4_1_priority_conflict_resolution(
|
|
temporal_test_env: WorkflowEnvironment,
|
|
temporal_worker: Worker,
|
|
test_activities: Activities,
|
|
postgres_engine,
|
|
bad_data_model,
|
|
):
|
|
"""Conflicting filter outputs must honor configured path_priority order."""
|
|
client = temporal_test_env.client
|
|
model_id = 242
|
|
insert_sample_data(postgres_engine, model_id, [60.0, 78.2])
|
|
input_data = get_base_input_data(model_id)
|
|
input_data['mlflow_transform_filters'] = {'API_ERROR': {'POLICY': 'CONTINUE', 'CONFIG': {}}}
|
|
input_data['mlflow_predict_filters'] = {'API_ERROR': {'POLICY': 'STOP', 'CONFIG': {}}}
|
|
input_data['path_priority'] = ['STOP', 'CONTINUE', 'REPEAT']
|
|
|
|
await start_and_await_workflow(
|
|
client, PredictionsBatch.run, input_data, make_workflow_id('test-priority-conflict')
|
|
)
|
|
assert_continue(
|
|
postgres_engine=postgres_engine,
|
|
model_id=model_id,
|
|
prediction_confidence=Decimal(10),
|
|
comments='Unknown MLFlow API error',
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.integration
|
|
async def test_e2e_request_predict_inline_minio_payload_with_datetimeindex(
|
|
temporal_test_env: WorkflowEnvironment,
|
|
temporal_worker: Worker,
|
|
test_activities: Activities,
|
|
postgres_engine,
|
|
mlflow_repository_stub,
|
|
):
|
|
"""
|
|
High offload threshold forces inline tabular dicts; ``DatetimeIndex`` must serialize as JSON
|
|
(string index keys via ``MinioDataFramePayload.from_dataframe``) so ``request_predict`` completes.
|
|
"""
|
|
client = temporal_test_env.client
|
|
model_id = 252
|
|
with postgres_engine.begin() as conn:
|
|
conn.execute(text(f'DELETE FROM sientia_data.laborious_data WHERE model_id = {model_id}'))
|
|
conn.execute(text(f'DELETE FROM sientia_data.predictions WHERE model_id = {model_id}'))
|
|
conn.execute(text(f'DELETE FROM sientia_data.transformed_data WHERE model_id = {model_id}'))
|
|
insert_sample_data(postgres_engine, model_id, [60.0, 78.2])
|
|
input_data = get_base_input_data(model_id)
|
|
|
|
with patch.object(minio_payload_module, 'OFFLOAD_THRESHOLD_BYTES', 10**9):
|
|
await start_and_await_workflow(
|
|
client,
|
|
PredictionsBatch.run,
|
|
input_data,
|
|
make_workflow_id('test-predict-inline-json-datetimeindex'),
|
|
)
|
|
|
|
assert_prediction(postgres_engine, model_id, prediction=0.5, prediction_confidence=0)
|
|
mlflow_repository_stub.stub_wrapper.predict.assert_called()
|