vitor-aignosi
30ff378cb7
QTZPOC-20: Prioritize code/ directory download and sys.path injection before model loading, enabling MLflow flavors to stream models directly from URIs and simplifying joblib artifact handling.
278 lines
9.8 KiB
Python
278 lines
9.8 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.
|
|
"""
|
|
|
|
from unittest.mock import MagicMock
|
|
|
|
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')
|
|
|
|
|
|
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_downloads_only_the_artifact_code_directory(
|
|
temporal_test_env: WorkflowEnvironment,
|
|
temporal_worker: Worker,
|
|
test_activities: Activities,
|
|
postgres_engine,
|
|
joblib_flavor_artifacts,
|
|
mlflow_client,
|
|
mock_mlflow_models,
|
|
):
|
|
"""A flavor model whose artifact carries `code` downloads that directory — and only it.
|
|
|
|
The flavor loaders read the model URI themselves, so nothing else is fetched through the
|
|
client: the `code` directory is the whole of what this workflow downloads.
|
|
"""
|
|
mlflow_client.list_artifacts.side_effect = lambda _run_id, artifact_path=None: (
|
|
[MagicMock(path=f'{artifact_path}/code')] if artifact_path else []
|
|
)
|
|
|
|
client = temporal_test_env.client
|
|
model_id = 275
|
|
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-code-only')
|
|
)
|
|
|
|
assert_prediction(postgres_engine, model_id, prediction=0.5)
|
|
|
|
downloaded = [loaded.args[1] for loaded in mlflow_client.download_artifacts.call_args_list]
|
|
assert sorted(downloaded) == ['data_model/code', 'prediction_model/code']
|
|
|
|
|
|
@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
|