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.
187 lines
6.5 KiB
Python
187 lines
6.5 KiB
Python
"""
|
|
End-to-end tests for the joblib fallback on the MLflow model load paths.
|
|
|
|
`joblib.load` is NOT mocked here: the artifact fixture holds a real sklearn
|
|
estimator serialized with `joblib.dump(..., compress=3)`, so a successful run
|
|
proves the fallback really deserializes a compressed pickle and serves it.
|
|
"""
|
|
|
|
import pandas as pd
|
|
import pytest
|
|
from sqlalchemy import text
|
|
from temporalio.testing import WorkflowEnvironment
|
|
from temporalio.worker import Worker
|
|
|
|
from e2e.conftest import JOBLIB_FALLBACK_FLAVOR_ERROR, JOBLIB_FALLBACK_MODEL_NAME
|
|
from e2e.helpers import (
|
|
assert_prediction,
|
|
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_query = 'SELECT timestamp, variable, value, created_at FROM predictions_schema.laborious_data WHERE model_id = {model_id}'
|
|
|
|
|
|
def build_input_data(model_id: int, model_name: str) -> dict:
|
|
"""Workflow input for a plain transform + predict run of a given model."""
|
|
return {
|
|
'schedule_name': 'test-schedule',
|
|
'model_name': model_name,
|
|
'model_id': model_id,
|
|
'query': base_query.format(model_id=model_id),
|
|
'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'],
|
|
}
|
|
|
|
|
|
def expected_fallback_prediction(estimator) -> float:
|
|
"""Prediction the fixture estimator produces for the mocked transform output."""
|
|
features = pd.DataFrame(
|
|
{
|
|
'feature_1': [0.234],
|
|
'feature_2': [0.783],
|
|
'timestamp': ['2024-01-01 12:00:00'],
|
|
}
|
|
)
|
|
return float(estimator.predict(features)[0])
|
|
|
|
|
|
def assert_single_prediction_close(postgres_engine, model_id, expected, tolerance=1e-9):
|
|
"""Assert exactly one Good prediction row whose value matches `expected`."""
|
|
with postgres_engine.connect() as conn:
|
|
rows = conn.execute(
|
|
text(
|
|
'SELECT prediction, prediction_status '
|
|
'FROM predictions_schema.predictions '
|
|
f'WHERE model_id = {model_id}'
|
|
)
|
|
).fetchall()
|
|
|
|
assert len(rows) == 1, f'Expected one prediction record, got {len(rows)}'
|
|
assert abs(float(rows[0][0]) - expected) < tolerance, (
|
|
f'Expected prediction close to {expected}, got {rows[0][0]}'
|
|
)
|
|
assert rows[0][1] == 'Good', f"Expected prediction_status='Good', got {rows[0][1]}"
|
|
|
|
|
|
def logged_messages(mock_logger, level: str) -> str:
|
|
"""Concatenate every message logged at a given custom level."""
|
|
calls = getattr(mock_logger, f'custom_{level}').call_args_list
|
|
return ' '.join(str(logged.args[0]) for logged in calls)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.integration
|
|
async def test_prediction_workflow_recovered_by_joblib_fallback(
|
|
temporal_test_env: WorkflowEnvironment,
|
|
temporal_worker: Worker,
|
|
test_activities: Activities,
|
|
postgres_engine,
|
|
joblib_fallback_artifact,
|
|
mlflow_client,
|
|
mock_logger,
|
|
):
|
|
"""Flavor loader fails, joblib fallback loads the real compressed pickle, prediction lands."""
|
|
client = temporal_test_env.client
|
|
model_id = 271
|
|
insert_sample_data(postgres_engine, model_id, [60.0, 78.2])
|
|
input_data = build_input_data(model_id, JOBLIB_FALLBACK_MODEL_NAME)
|
|
|
|
await start_and_await_workflow(
|
|
client, PredictionsBatch.run, input_data, make_workflow_id('test-joblib-fallback')
|
|
)
|
|
|
|
assert_single_prediction_close(
|
|
postgres_engine,
|
|
model_id,
|
|
expected_fallback_prediction(joblib_fallback_artifact['estimator']),
|
|
)
|
|
|
|
assert mlflow_client.download_artifacts.call_count == 1
|
|
assert mlflow_client.download_artifacts.call_args.args[1] == 'prediction_model'
|
|
|
|
assert JOBLIB_FALLBACK_FLAVOR_ERROR in logged_messages(mock_logger, 'warning')
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.integration
|
|
async def test_prediction_workflow_without_fallback_is_unchanged(
|
|
temporal_test_env: WorkflowEnvironment,
|
|
temporal_worker: Worker,
|
|
test_activities: Activities,
|
|
postgres_engine,
|
|
mlflow_client,
|
|
mock_mlflow_models,
|
|
):
|
|
"""A model whose flavor loader works keeps the current result and downloads nothing."""
|
|
client = temporal_test_env.client
|
|
model_id = 272
|
|
insert_sample_data(postgres_engine, model_id, [60.0, 78.2])
|
|
input_data = build_input_data(model_id, 'test_model')
|
|
|
|
await start_and_await_workflow(
|
|
client, PredictionsBatch.run, input_data, make_workflow_id('test-joblib-no-fallback')
|
|
)
|
|
|
|
assert_prediction(postgres_engine, model_id, prediction=0.5)
|
|
mock_mlflow_models['predict_model'].predict.assert_called()
|
|
mlflow_client.download_artifacts.assert_not_called()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.integration
|
|
async def test_prediction_workflow_reports_both_causes_when_fallback_fails(
|
|
temporal_test_env: WorkflowEnvironment,
|
|
temporal_worker: Worker,
|
|
test_activities: Activities,
|
|
postgres_engine,
|
|
joblib_fallback_artifact_without_pickle,
|
|
mlflow_client,
|
|
mock_logger,
|
|
):
|
|
"""Both attempts fail: workflow stops and the logged traceback carries both causes."""
|
|
mlflow_client.download_artifacts.return_value = joblib_fallback_artifact_without_pickle
|
|
|
|
client = temporal_test_env.client
|
|
model_id = 273
|
|
insert_sample_data(postgres_engine, model_id, [60.0, 78.2])
|
|
input_data = build_input_data(model_id, JOBLIB_FALLBACK_MODEL_NAME)
|
|
|
|
await start_and_await_workflow(
|
|
client, PredictionsBatch.run, input_data, make_workflow_id('test-joblib-fallback-failed')
|
|
)
|
|
|
|
assert_stop(postgres_engine, model_id)
|
|
|
|
logged = logged_messages(mock_logger, 'debug')
|
|
assert JOBLIB_FALLBACK_FLAVOR_ERROR in logged
|
|
assert 'No .pkl file found in prediction_model artifact' in logged
|