vitor-aignosi
e95e65b72b
QTZPOC-20: Implement explicit joblib model flavor and enhance model import workflow with extensive E2E tests. This commit introduces a first-class joblib model flavor for direct loading and saving of Python pickle artifacts, including custom code within the artifact's code/ directory, bypassing MLflow's flavor loaders. The model import workflow's status database configuration is now streamlined to default to the POSTGRES_* server, with explicit IMPORT_STATUS_DB_* overrides. Comprehensive end-to-end tests have been added for the model import workflow, covering success paths, gate rejections, digest enforcement, provisioning failures, status record integrity, cleanup, and retry/recovery scenarios using a real-clock environment. New E2
241 lines
8.5 KiB
Python
241 lines
8.5 KiB
Python
"""
|
|
End-to-end tests for the explicit `joblib` flavor on the MLflow model load paths.
|
|
|
|
`joblib.load` is NOT mocked here: the artifact fixture holds a real object serialized with
|
|
`joblib.dump(..., compress=3)` whose class lives only inside the artifact's `code/` directory.
|
|
A passing run therefore proves the flavor decompresses the pickle **and** resolves the module,
|
|
without any `mlflow.*.load_model` being called for that model.
|
|
"""
|
|
|
|
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_FLAVOR_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}'
|
|
|
|
FLAVOR_LOADERS = ('sklearn', 'pyfunc', 'pytorch')
|
|
|
|
|
|
def build_input_data(
|
|
model_id: int,
|
|
model_name: str,
|
|
transform_flavor: str = 'sklearn',
|
|
predict_flavor: str = 'sklearn',
|
|
) -> 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': transform_flavor,
|
|
'predict_flavor': predict_flavor,
|
|
},
|
|
'datetime_columns': ['timestamp', 'created_at'],
|
|
}
|
|
|
|
|
|
def expected_joblib_prediction(model) -> float:
|
|
"""Prediction the artifact's model produces for the transform output of this suite."""
|
|
features = pd.DataFrame(
|
|
{
|
|
'feature_1': [0.234],
|
|
'feature_2': [0.783],
|
|
'timestamp': ['2024-01-01 12:00:00'],
|
|
}
|
|
)
|
|
return float(model.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)
|
|
|
|
|
|
def loader_uris(patch_mlflow) -> list[str]:
|
|
"""Every model URI handed to an MLflow flavor loader during the run."""
|
|
return [
|
|
str(loaded.args[0])
|
|
for flavor in FLAVOR_LOADERS
|
|
for loaded in getattr(patch_mlflow, flavor).load_model.call_args_list
|
|
]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.integration
|
|
async def test_prediction_workflow_with_joblib_flavor(
|
|
temporal_test_env: WorkflowEnvironment,
|
|
temporal_worker: Worker,
|
|
test_activities: Activities,
|
|
postgres_engine,
|
|
joblib_flavor_artifacts,
|
|
mlflow_client,
|
|
patch_mlflow,
|
|
):
|
|
"""`predict_flavor: 'joblib'` deserializes the real artifact, and the 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_FLAVOR_MODEL_NAME, predict_flavor='joblib')
|
|
|
|
await start_and_await_workflow(
|
|
client, PredictionsBatch.run, input_data, make_workflow_id('test-joblib-flavor')
|
|
)
|
|
|
|
assert_single_prediction_close(
|
|
postgres_engine,
|
|
model_id,
|
|
expected_joblib_prediction(joblib_flavor_artifacts['prediction_model']['model']),
|
|
)
|
|
|
|
assert mlflow_client.download_artifacts.call_count == 1
|
|
assert mlflow_client.download_artifacts.call_args.args[1] == 'prediction_model'
|
|
|
|
# The joblib flavor never builds a `models:/` URI, so no flavor loader saw this model.
|
|
assert all(JOBLIB_FLAVOR_MODEL_NAME not in uri for uri in loader_uris(patch_mlflow))
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.integration
|
|
async def test_prediction_workflow_with_joblib_transform_flavor(
|
|
temporal_test_env: WorkflowEnvironment,
|
|
temporal_worker: Worker,
|
|
test_activities: Activities,
|
|
postgres_engine,
|
|
joblib_flavor_artifacts,
|
|
mlflow_client,
|
|
patch_mlflow,
|
|
mock_mlflow_models,
|
|
):
|
|
"""`transform_flavor: 'joblib'` downloads and deserializes `data_model` on its own."""
|
|
client = temporal_test_env.client
|
|
model_id = 274
|
|
insert_sample_data(postgres_engine, model_id, [60.0, 78.2])
|
|
input_data = build_input_data(model_id, JOBLIB_FLAVOR_MODEL_NAME, transform_flavor='joblib')
|
|
|
|
await start_and_await_workflow(
|
|
client, PredictionsBatch.run, input_data, make_workflow_id('test-joblib-transform')
|
|
)
|
|
|
|
assert_prediction(postgres_engine, model_id, prediction=0.5)
|
|
mock_mlflow_models['predict_model'].predict.assert_called()
|
|
|
|
assert mlflow_client.download_artifacts.call_count == 1
|
|
assert mlflow_client.download_artifacts.call_args.args[1] == 'data_model'
|
|
|
|
# The joblib flavor resolves no artifact URI, so `get_run` is never reached and no flavor
|
|
# loader was given a `data_model` URI.
|
|
patch_mlflow.get_run.assert_not_called()
|
|
assert all(not uri.endswith('/data_model') for uri in loader_uris(patch_mlflow))
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.integration
|
|
async def test_prediction_workflow_with_sklearn_flavor_is_unchanged(
|
|
temporal_test_env: WorkflowEnvironment,
|
|
temporal_worker: Worker,
|
|
test_activities: Activities,
|
|
postgres_engine,
|
|
mlflow_client,
|
|
mock_mlflow_models,
|
|
):
|
|
"""A model on an MLflow flavor 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-sklearn-flavor')
|
|
)
|
|
|
|
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_joblib_flavor_failure_reports_single_cause(
|
|
temporal_test_env: WorkflowEnvironment,
|
|
temporal_worker: Worker,
|
|
test_activities: Activities,
|
|
postgres_engine,
|
|
joblib_flavor_artifact_without_pickle,
|
|
mlflow_client,
|
|
mock_logger,
|
|
):
|
|
"""The joblib flavor failing stops the workflow with one cause, not a chain of two."""
|
|
mlflow_client.download_artifacts.side_effect = lambda _run_id, _artifact_name: (
|
|
joblib_flavor_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_FLAVOR_MODEL_NAME, predict_flavor='joblib')
|
|
|
|
await start_and_await_workflow(
|
|
client, PredictionsBatch.run, input_data, make_workflow_id('test-joblib-flavor-failed')
|
|
)
|
|
|
|
assert_stop(postgres_engine, model_id)
|
|
|
|
logged = logged_messages(mock_logger, 'debug')
|
|
assert 'No .pkl file found in prediction_model artifact' in logged
|
|
assert 'Original error' not in logged
|