vitor-aignosi
c12f2996c2
QTZPOC-13: Introduce joblib fallback for MLflow model loading Enhances model loading resilience by attempting to deserialize models directly with joblib if MLflow's native flavor loaders fail. This mechanism handles cases where .pkl files are present in the artifact. Adds joblib as a core dependency and includes new end-to-end and unit tests to validate fallback behavior, covering success, no-fallback, and double-failure scenarios. Also updates .gitignore for .claude/ files and incorporates minor refactorings and linting adjustments across E2E tests and utility functions.
386 lines
13 KiB
Python
386 lines
13 KiB
Python
"""
|
|
End-to-end tests for PredictionsBatch workflow - Prediction Process scenarios.
|
|
"""
|
|
|
|
from decimal import Decimal
|
|
from unittest.mock import MagicMock
|
|
|
|
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_repeat,
|
|
assert_stop,
|
|
insert_sample_data,
|
|
make_workflow_id,
|
|
start_and_await_workflow,
|
|
)
|
|
from laborious.activities.activities import Activities
|
|
from laborious.workflows.predictions_batch import PredictionsBatch
|
|
|
|
base_input_data = {
|
|
'schedule_name': 'test-schedule',
|
|
'model_name': 'test_model',
|
|
'model_id': 201,
|
|
'query': 'SELECT timestamp, variable, value, created_at FROM predictions_schema.laborious_data WHERE model_id = 201',
|
|
'schema': 'predictions_schema',
|
|
'table_name': 'predictions',
|
|
'transform_table_name': 'transformed_data',
|
|
'input_filters': {
|
|
'SPECIFIC_VARIABLES_NULL_VALUES': {
|
|
'POLICY': 'CONTINUE',
|
|
'CONFIG': {'variables': ['sensor_1']},
|
|
},
|
|
},
|
|
'mlflow_transform_filters': {
|
|
'API_ERROR': {'POLICY': 'STOP', 'CONFIG': {}},
|
|
},
|
|
'mlflow_predict_filters': {
|
|
'API_ERROR': {'POLICY': 'STOP', 'CONFIG': {}},
|
|
},
|
|
'path_priority': ['STOP', 'CONTINUE', 'REPEAT'],
|
|
'opc_output_config': {},
|
|
'pi_web_api_output_config': {},
|
|
'save_transform': True,
|
|
'prediction_store_policy': 'lts:1',
|
|
'model_config': {
|
|
'retention_minutes': 0,
|
|
'transform_flavor': 'sklearn',
|
|
'predict_flavor': 'sklearn',
|
|
},
|
|
'datetime_columns': ['timestamp', 'created_at'],
|
|
}
|
|
|
|
base_query = 'SELECT timestamp, variable, value, created_at FROM predictions_schema.laborious_data WHERE model_id = {model_id}'
|
|
|
|
|
|
def get_base_input_data(model_id):
|
|
return {
|
|
**base_input_data,
|
|
'model_id': model_id,
|
|
'query': base_query.format(model_id=model_id),
|
|
}
|
|
|
|
|
|
def insert_sample_prediction(postgres_engine, model_id):
|
|
with postgres_engine.begin() as conn:
|
|
conn.execute(
|
|
text(f'DELETE FROM predictions_schema.predictions WHERE model_id = {model_id}')
|
|
)
|
|
insert_sql = f"""
|
|
INSERT INTO predictions_schema.predictions (model_id, timestamp, prediction, prediction_confidence, prediction_status, comments, response_time)
|
|
VALUES
|
|
({model_id}, '2024-01-01 12:00:00+00:00', 10, 0, 'Good', '', 0.1)
|
|
"""
|
|
conn.execute(text(insert_sql))
|
|
return (model_id, Decimal(10), Decimal(0), 'Good')
|
|
|
|
|
|
@pytest.fixture
|
|
def bad_data_model(patch_mlflow):
|
|
model = MagicMock(predict=MagicMock(side_effect=Exception('Bad data model')))
|
|
patch_mlflow.sklearn.load_model = MagicMock(return_value=model)
|
|
return model
|
|
|
|
|
|
@pytest.fixture
|
|
def bad_predict_model(patch_mlflow, mock_mlflow_models):
|
|
model = MagicMock(predict=MagicMock(side_effect=Exception('Bad predict model')))
|
|
|
|
def mock_sklearn_load_model(model_uri):
|
|
if 'data_model' in model_uri or 'transform' in model_uri.lower():
|
|
return mock_mlflow_models['transform_model']
|
|
return model
|
|
|
|
patch_mlflow.sklearn = MagicMock()
|
|
patch_mlflow.sklearn.load_model = MagicMock(side_effect=mock_sklearn_load_model)
|
|
return model
|
|
|
|
|
|
@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,
|
|
mock_mlflow_models,
|
|
):
|
|
"""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)
|
|
mock_mlflow_models['transform_model'].predict.assert_not_called()
|
|
mock_mlflow_models['predict_model'].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,
|
|
mock_mlflow_models,
|
|
):
|
|
"""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)
|
|
mock_mlflow_models['transform_model'].predict.assert_not_called()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.integration
|
|
async def test_scenario_2_1_3_input_gate_triggers_repeat(
|
|
temporal_test_env: WorkflowEnvironment,
|
|
temporal_worker: Worker,
|
|
test_activities: Activities,
|
|
postgres_engine,
|
|
mock_mlflow_models,
|
|
):
|
|
"""Input gate REPEAT with existing history."""
|
|
client = temporal_test_env.client
|
|
model_id = 213
|
|
insert_sample_data(postgres_engine, model_id, ['NULL', 78.2])
|
|
data = insert_sample_prediction(postgres_engine, 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')
|
|
)
|
|
assert_repeat(postgres_engine, model_id, data)
|
|
mock_mlflow_models['transform_model'].predict.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 predictions_schema.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,
|
|
mock_mlflow_models,
|
|
):
|
|
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)
|
|
mock_mlflow_models['predict_model'].predict.assert_not_called()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.integration
|
|
async def test_scenario_2_2_3_transform_gate_triggers_repeat(
|
|
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 = insert_sample_prediction(postgres_engine, model_id)
|
|
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')
|
|
)
|
|
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,
|
|
mock_mlflow_models,
|
|
):
|
|
"""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):
|
|
num_rows = max(len(data), 1) if hasattr(data, '__len__') else 1
|
|
result = pd.DataFrame({'feature_1': [np.nan] * num_rows, 'feature_2': [np.nan] * num_rows})
|
|
result.index = data.index
|
|
return result
|
|
|
|
mock_mlflow_models['transform_model'].predict = 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)
|
|
mock_mlflow_models['predict_model'].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_triggers_repeat(
|
|
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 = insert_sample_prediction(postgres_engine, model_id)
|
|
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')
|
|
)
|
|
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 predictions_schema.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)
|