QTZPOC-13
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.
This commit is contained in:
@@ -1,13 +1,20 @@
|
||||
from datetime import UTC, datetime
|
||||
from unittest.mock import ANY, AsyncMock, MagicMock, call, patch
|
||||
|
||||
import joblib as joblib_lib
|
||||
import mlflow as mlflow_lib
|
||||
import numpy as np
|
||||
import pytest
|
||||
from pandas import DataFrame, Timestamp
|
||||
from sientia_do.temporal.constants import DATETIME_FORMAT_WITH_TZ
|
||||
from sklearn.linear_model import LinearRegression
|
||||
|
||||
from laborious import metrics
|
||||
from laborious.utils.repository.model_repository import MLFlowRepository, force_memory_release
|
||||
from laborious.utils.repository.model_repository import (
|
||||
INVALID_FLAVOR_MESSAGE,
|
||||
MLFlowRepository,
|
||||
force_memory_release,
|
||||
)
|
||||
|
||||
|
||||
@patch('laborious.utils.repository.model_repository.ctypes')
|
||||
@@ -608,32 +615,46 @@ def test_detect_and_parse_datetime_index_empty(mlflow_repository):
|
||||
assert response.equals(input_data)
|
||||
|
||||
|
||||
# Each case is (index, fragments expected in the error message).
|
||||
#
|
||||
# The index is passed explicitly instead of through a dict-of-dict: building a
|
||||
# DataFrame from nested dicts makes pandas sort the keys to infer the index, which a
|
||||
# mixed-type index cannot do.
|
||||
#
|
||||
# Fragments, not whole messages: the parse branch appends pandas' own multi-line
|
||||
# error, and the type branch appends the fully qualified class name, so neither is
|
||||
# stable enough to compare literally.
|
||||
invalid_cases = [
|
||||
(
|
||||
{'value': {'2024-01-01 12:00:00': 1, 2024: 2}},
|
||||
"Index must be all timestamp like column. Valid formats are: pandas Timestamp, datetime, string in format %Y-%m-%d %H:%M:%S. Elements are <class 'str'>, <class 'int'>.",
|
||||
['2024-01-01 12:00:00', 2024],
|
||||
["Elements are <class 'str'>,<class 'int'>"],
|
||||
),
|
||||
(
|
||||
{'value': {'2024-01-01': 1, '2024-01-02': 2}},
|
||||
'Index must be all timestamp like column. Valid formats are: pandas Timestamp, datetime, string in format %Y-%m-%d %H:%M:%S. Unable to parse given date format',
|
||||
['2024-01-01', '2024-01-02'],
|
||||
['Unable to parse given date format'],
|
||||
),
|
||||
(
|
||||
{'value': {Any(): 1, Any(): 2}},
|
||||
'Index must be all timestamp like column. Valid formats are: pandas Timestamp, datetime, string in format %Y-%m-%d %H:%M:%S. Got <class Any>.',
|
||||
[Any(), Any()],
|
||||
['Got <class', "Any'>."],
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize('data', invalid_cases)
|
||||
def test_detect_and_parse_datetime_index_error_cases(mlflow_repository, data):
|
||||
input_data = DataFrame(data[0])
|
||||
input_data = DataFrame({'value': [1, 2]}, index=data[0])
|
||||
|
||||
message = data[1]
|
||||
expected_fragments = data[1]
|
||||
|
||||
with pytest.raises(ValueError) as e:
|
||||
mlflow_repository.detect_and_parse_datetime_index(input_data, metadata['metadata'])
|
||||
|
||||
assert str(e) == message
|
||||
message = str(e.value)
|
||||
|
||||
assert 'Index must be all timestamp like column' in message
|
||||
assert f'string in format {DATETIME_FORMAT_WITH_TZ}' in message
|
||||
for fragment in expected_fragments:
|
||||
assert fragment in message, f'Expected {fragment!r} in {message!r}'
|
||||
|
||||
|
||||
valid_cases = [
|
||||
@@ -1713,3 +1734,357 @@ async def test_predict_success_ndarray(mlflow_repository):
|
||||
assert 'prediction' in content.columns
|
||||
assert 'response_time' in content.columns
|
||||
assert content.index.tolist() == data.index.tolist()
|
||||
|
||||
|
||||
"""
|
||||
Tests for the joblib fallback on the model load paths
|
||||
"""
|
||||
|
||||
JOBLIB_PATH = 'laborious.utils.repository.model_repository.joblib'
|
||||
|
||||
|
||||
def build_artifact_dir(tmp_path, *filenames):
|
||||
"""Create a fake downloaded artifact directory with the given top-level files."""
|
||||
for filename in filenames:
|
||||
(tmp_path / filename).write_bytes(b'pickle-bytes')
|
||||
return str(tmp_path)
|
||||
|
||||
|
||||
def emitted_tags(emit_metric_mock, metric_object):
|
||||
"""Collect the `tags` kwarg of every emit_metric call for a given metric."""
|
||||
return [
|
||||
emitted.kwargs['tags']
|
||||
for emitted in emit_metric_mock.call_args_list
|
||||
if emitted.kwargs['metric_object'] is metric_object
|
||||
]
|
||||
|
||||
|
||||
@patch(JOBLIB_PATH)
|
||||
@pytest.mark.asyncio
|
||||
async def test_load_predict_model_flavor_success_skips_fallback(
|
||||
joblib_mock, mlflow, mlflow_repository
|
||||
):
|
||||
result = await mlflow_repository.load_predict_model('test_model', {}, 'sklearn')
|
||||
|
||||
assert result == mlflow.sklearn.load_model.return_value
|
||||
joblib_mock.load.assert_not_called()
|
||||
mlflow_repository.client.download_artifacts.assert_not_called()
|
||||
|
||||
|
||||
@patch(JOBLIB_PATH)
|
||||
@pytest.mark.asyncio
|
||||
async def test_load_transform_model_flavor_success_skips_fallback(
|
||||
joblib_mock, mlflow, mlflow_repository
|
||||
):
|
||||
mlflow_repository.get_model_run_id = MagicMock()
|
||||
mlflow_repository.get_model_uri = MagicMock()
|
||||
|
||||
result = await mlflow_repository.load_transform_model('test_model', {}, 'sklearn')
|
||||
|
||||
assert result == mlflow.sklearn.load_model.return_value
|
||||
joblib_mock.load.assert_not_called()
|
||||
mlflow_repository.client.download_artifacts.assert_not_called()
|
||||
|
||||
|
||||
@patch(JOBLIB_PATH)
|
||||
@pytest.mark.asyncio
|
||||
async def test_download_model_wrapper_success_skips_fallback(
|
||||
joblib_mock, mlflow, mlflow_repository
|
||||
):
|
||||
mlflow_repository.dowload_artifacts = AsyncMock()
|
||||
|
||||
result = await mlflow_repository.download_model('test_model', {}, 'predict', 'pyfunc', True)
|
||||
|
||||
assert result == (
|
||||
mlflow.pyfunc.load_model.return_value._model_impl.python_model,
|
||||
mlflow_repository.dowload_artifacts.return_value,
|
||||
)
|
||||
joblib_mock.load.assert_not_called()
|
||||
mlflow_repository.client.download_artifacts.assert_not_called()
|
||||
|
||||
|
||||
@patch(JOBLIB_PATH)
|
||||
@pytest.mark.asyncio
|
||||
async def test_load_predict_model_invalid_flavor_skips_fallback(joblib_mock, mlflow_repository):
|
||||
with pytest.raises(ValueError) as exc_info:
|
||||
await mlflow_repository.load_predict_model('test_model', {}, 'invalid')
|
||||
|
||||
assert str(exc_info.value) == INVALID_FLAVOR_MESSAGE
|
||||
joblib_mock.load.assert_not_called()
|
||||
mlflow_repository.client.download_artifacts.assert_not_called()
|
||||
mlflow_repository.emit_metric.assert_not_called()
|
||||
|
||||
|
||||
@patch(JOBLIB_PATH)
|
||||
@pytest.mark.asyncio
|
||||
async def test_load_transform_model_invalid_flavor_skips_fallback(joblib_mock, mlflow_repository):
|
||||
with pytest.raises(ValueError) as exc_info:
|
||||
await mlflow_repository.load_transform_model('test_model', {}, 'invalid')
|
||||
|
||||
assert str(exc_info.value) == INVALID_FLAVOR_MESSAGE
|
||||
joblib_mock.load.assert_not_called()
|
||||
mlflow_repository.client.download_artifacts.assert_not_called()
|
||||
mlflow_repository.emit_metric.assert_not_called()
|
||||
|
||||
|
||||
@patch(JOBLIB_PATH)
|
||||
@pytest.mark.asyncio
|
||||
async def test_load_predict_model_fallback_recovers_load(
|
||||
joblib_mock, mlflow, mlflow_repository, tmp_path
|
||||
):
|
||||
mlflow.sklearn.load_model.side_effect = Exception('flavor boom')
|
||||
mlflow_repository.get_model_run_id = MagicMock(return_value='run-123')
|
||||
mlflow_repository.client.download_artifacts.return_value = build_artifact_dir(
|
||||
tmp_path, 'MLmodel', 'model.pkl'
|
||||
)
|
||||
|
||||
result = await mlflow_repository.load_predict_model('test_model', {}, 'sklearn')
|
||||
|
||||
assert result is joblib_mock.load.return_value
|
||||
mlflow_repository.get_model_run_id.assert_called_once_with(
|
||||
model_name='test_model', stage='Production'
|
||||
)
|
||||
mlflow_repository.client.download_artifacts.assert_called_once_with(
|
||||
'run-123', 'prediction_model'
|
||||
)
|
||||
joblib_mock.load.assert_called_once_with(str(tmp_path / 'model.pkl'))
|
||||
|
||||
|
||||
@patch(JOBLIB_PATH)
|
||||
@pytest.mark.asyncio
|
||||
async def test_load_transform_model_fallback_recovers_load(
|
||||
joblib_mock, mlflow, mlflow_repository, tmp_path
|
||||
):
|
||||
mlflow.sklearn.load_model.side_effect = Exception('flavor boom')
|
||||
mlflow_repository.get_model_run_id = MagicMock(return_value='run-123')
|
||||
mlflow_repository.get_model_uri = MagicMock()
|
||||
mlflow_repository.client.download_artifacts.return_value = build_artifact_dir(
|
||||
tmp_path, 'MLmodel', 'model.pkl'
|
||||
)
|
||||
|
||||
result = await mlflow_repository.load_transform_model('test_model', {}, 'sklearn')
|
||||
|
||||
assert result is joblib_mock.load.return_value
|
||||
mlflow_repository.get_model_run_id.assert_called_with(
|
||||
model_name='test_model', stage='Production'
|
||||
)
|
||||
mlflow_repository.client.download_artifacts.assert_called_once_with('run-123', 'data_model')
|
||||
joblib_mock.load.assert_called_once_with(str(tmp_path / 'model.pkl'))
|
||||
|
||||
|
||||
@patch(JOBLIB_PATH)
|
||||
@pytest.mark.asyncio
|
||||
async def test_download_model_wrapper_fallback_reuses_downloaded_dir(
|
||||
joblib_mock, mlflow, mlflow_repository, tmp_path
|
||||
):
|
||||
artifact_dir = build_artifact_dir(tmp_path, 'MLmodel', 'python_model.pkl')
|
||||
mlflow_repository.dowload_artifacts = AsyncMock(return_value=artifact_dir)
|
||||
mlflow.pyfunc.load_model.side_effect = Exception('flavor boom')
|
||||
|
||||
model, returned_path = await mlflow_repository.download_model(
|
||||
'test_model', {}, 'predict', 'pyfunc', True
|
||||
)
|
||||
|
||||
mlflow_repository.dowload_artifacts.assert_called_once_with(
|
||||
'test_model', {}, 'prediction_model'
|
||||
)
|
||||
mlflow_repository.client.download_artifacts.assert_not_called()
|
||||
joblib_mock.load.assert_called_once_with(str(tmp_path / 'python_model.pkl'))
|
||||
|
||||
assert model is joblib_mock.load.return_value
|
||||
assert returned_path == artifact_dir
|
||||
|
||||
|
||||
@patch(JOBLIB_PATH)
|
||||
@pytest.mark.asyncio
|
||||
async def test_load_predict_model_both_attempts_fail_chains_causes(
|
||||
joblib_mock, mlflow, mlflow_repository, tmp_path
|
||||
):
|
||||
flavor_error = Exception('flavor boom')
|
||||
mlflow.sklearn.load_model.side_effect = flavor_error
|
||||
mlflow_repository.get_model_run_id = MagicMock(return_value='run-123')
|
||||
mlflow_repository.client.download_artifacts.return_value = build_artifact_dir(
|
||||
tmp_path, 'model.pkl'
|
||||
)
|
||||
joblib_mock.load.side_effect = ValueError('joblib boom')
|
||||
|
||||
with pytest.raises(ValueError) as exc_info:
|
||||
await mlflow_repository.load_predict_model('test_model', {}, 'sklearn')
|
||||
|
||||
assert str(exc_info.value) == 'joblib boom'
|
||||
assert exc_info.value.__cause__ is flavor_error
|
||||
|
||||
error_tags = emitted_tags(mlflow_repository.emit_metric, metrics.MODEL_READ_ERROR_COUNT)
|
||||
assert len(error_tags) == 1
|
||||
assert error_tags[0]['operation_type'] == 'load_predict_model_joblib_fallback'
|
||||
|
||||
|
||||
@patch(JOBLIB_PATH)
|
||||
@pytest.mark.asyncio
|
||||
async def test_load_transform_model_both_attempts_fail_emits_error_metric(
|
||||
joblib_mock, mlflow, mlflow_repository, tmp_path
|
||||
):
|
||||
mlflow.sklearn.load_model.side_effect = Exception('flavor boom')
|
||||
mlflow_repository.get_model_run_id = MagicMock(return_value='run-123')
|
||||
mlflow_repository.get_model_uri = MagicMock()
|
||||
mlflow_repository.client.download_artifacts.return_value = build_artifact_dir(
|
||||
tmp_path, 'model.pkl'
|
||||
)
|
||||
joblib_mock.load.side_effect = ValueError('joblib boom')
|
||||
|
||||
with pytest.raises(ValueError):
|
||||
await mlflow_repository.load_transform_model('test_model', {}, 'sklearn')
|
||||
|
||||
error_tags = emitted_tags(mlflow_repository.emit_metric, metrics.MODEL_READ_ERROR_COUNT)
|
||||
assert len(error_tags) == 1
|
||||
assert error_tags[0]['operation_type'] == 'load_transform_model_joblib_fallback'
|
||||
assert emitted_tags(mlflow_repository.emit_metric, metrics.MODEL_READ_COUNT) == []
|
||||
|
||||
|
||||
@patch(JOBLIB_PATH)
|
||||
@pytest.mark.asyncio
|
||||
async def test_load_predict_model_fallback_success_metrics(
|
||||
joblib_mock, mlflow, mlflow_repository, tmp_path
|
||||
):
|
||||
mlflow.sklearn.load_model.side_effect = Exception('flavor boom')
|
||||
mlflow_repository.get_model_run_id = MagicMock(return_value='run-123')
|
||||
mlflow_repository.client.download_artifacts.return_value = build_artifact_dir(
|
||||
tmp_path, 'model.pkl'
|
||||
)
|
||||
mlflow_repository.warning = MagicMock()
|
||||
|
||||
await mlflow_repository.load_predict_model('test_model', {}, 'sklearn')
|
||||
|
||||
read_tags = emitted_tags(mlflow_repository.emit_metric, metrics.MODEL_READ_COUNT)
|
||||
assert len(read_tags) == 1
|
||||
assert read_tags[0]['operation_type'] == 'load_predict_model_joblib_fallback'
|
||||
assert emitted_tags(mlflow_repository.emit_metric, metrics.MODEL_READ_ERROR_COUNT) == []
|
||||
|
||||
mlflow_repository.observe_lag.assert_called_once()
|
||||
lag_call = mlflow_repository.observe_lag.call_args
|
||||
assert lag_call.args[1] is metrics.MODEL_READ_LAG
|
||||
assert lag_call.args[2]['operation_type'] == 'load_predict_model_joblib_fallback'
|
||||
|
||||
mlflow_repository.warning.assert_called_once()
|
||||
assert 'flavor boom' in mlflow_repository.warning.call_args.args[0]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
('filenames', 'expected'),
|
||||
[
|
||||
(('MLmodel', 'model.pkl'), 'model.pkl'),
|
||||
(('MLmodel', 'python_model.pkl'), 'python_model.pkl'),
|
||||
(('b_model.pkl', 'a_model.pkl'), 'a_model.pkl'),
|
||||
],
|
||||
)
|
||||
def test_find_model_pickle_selection(mlflow_repository, tmp_path, filenames, expected):
|
||||
local_dir = build_artifact_dir(tmp_path, *filenames)
|
||||
|
||||
result = mlflow_repository._find_model_pickle(
|
||||
local_dir, 'prediction_model', Exception('flavor boom')
|
||||
)
|
||||
|
||||
assert result == str(tmp_path / expected)
|
||||
|
||||
|
||||
def test_find_model_pickle_ignores_nested_pickles(mlflow_repository, tmp_path):
|
||||
nested = tmp_path / 'artifacts'
|
||||
nested.mkdir()
|
||||
(nested / 'stacking_model.pkl').write_bytes(b'pickle-bytes')
|
||||
|
||||
with pytest.raises(ValueError) as exc_info:
|
||||
mlflow_repository._find_model_pickle(
|
||||
str(tmp_path), 'prediction_model', Exception('flavor boom')
|
||||
)
|
||||
|
||||
assert 'No .pkl file found' in str(exc_info.value)
|
||||
|
||||
|
||||
def test_find_model_pickle_empty_directory(mlflow_repository, tmp_path):
|
||||
with pytest.raises(ValueError) as exc_info:
|
||||
mlflow_repository._find_model_pickle(str(tmp_path), 'data_model', Exception('flavor boom'))
|
||||
|
||||
assert 'data_model' in str(exc_info.value)
|
||||
|
||||
|
||||
def test_find_model_pickle_error_message_carries_context(mlflow_repository, tmp_path):
|
||||
local_dir = build_artifact_dir(tmp_path, 'MLmodel')
|
||||
|
||||
with pytest.raises(ValueError) as exc_info:
|
||||
mlflow_repository._find_model_pickle(
|
||||
local_dir, 'prediction_model', Exception('flavor boom')
|
||||
)
|
||||
|
||||
message = str(exc_info.value)
|
||||
assert 'prediction_model' in message
|
||||
assert 'flavor boom' in message
|
||||
assert 'MLmodel' in message
|
||||
|
||||
|
||||
def test_find_model_pickle_ignores_directory_named_pkl(mlflow_repository, tmp_path):
|
||||
(tmp_path / 'algo.pkl').mkdir()
|
||||
|
||||
with pytest.raises(ValueError) as exc_info:
|
||||
mlflow_repository._find_model_pickle(
|
||||
str(tmp_path), 'prediction_model', Exception('flavor boom')
|
||||
)
|
||||
|
||||
assert 'algo.pkl' in str(exc_info.value)
|
||||
|
||||
|
||||
@patch(JOBLIB_PATH)
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_model_caches_model_loaded_by_fallback(
|
||||
joblib_mock, mlflow, mlflow_repository, tmp_path
|
||||
):
|
||||
mlflow.sklearn.load_model.side_effect = Exception('flavor boom')
|
||||
mlflow_repository.get_model_run_id = MagicMock(return_value='run-123')
|
||||
mlflow_repository.client.download_artifacts.return_value = build_artifact_dir(
|
||||
tmp_path, 'model.pkl'
|
||||
)
|
||||
|
||||
first = await mlflow_repository.get_model('test_model', {}, 5, 'predict', 'sklearn')
|
||||
second = await mlflow_repository.get_model('test_model', {}, 5, 'predict', 'sklearn')
|
||||
|
||||
assert first is joblib_mock.load.return_value
|
||||
assert second is first
|
||||
assert mlflow.sklearn.load_model.call_count == 1
|
||||
assert joblib_mock.load.call_count == 1
|
||||
|
||||
|
||||
@patch(JOBLIB_PATH)
|
||||
@pytest.mark.asyncio
|
||||
async def test_predict_reports_both_causes_when_fallback_fails(
|
||||
joblib_mock, mlflow, mlflow_repository, tmp_path
|
||||
):
|
||||
mlflow.sklearn.load_model.side_effect = Exception('flavor boom')
|
||||
mlflow_repository.get_model_run_id = MagicMock(return_value='run-123')
|
||||
mlflow_repository.client.download_artifacts.return_value = build_artifact_dir(
|
||||
tmp_path, 'model.pkl'
|
||||
)
|
||||
joblib_mock.load.side_effect = ValueError('joblib boom')
|
||||
|
||||
data = DataFrame({'feat_1': {'index_1': 2, 'index_2': 3}})
|
||||
|
||||
output = await mlflow_repository.predict(
|
||||
'test_model',
|
||||
data,
|
||||
{'retention_minutes': 0, 'predict_flavor': 'sklearn'},
|
||||
metadata['metadata'],
|
||||
)
|
||||
|
||||
assert output['success'] is False
|
||||
assert 'joblib boom' in output['content']['traceback']
|
||||
assert 'flavor boom' in output['content']['traceback']
|
||||
|
||||
|
||||
def test_load_model_pickle_real_joblib_round_trip(mlflow_repository, tmp_path):
|
||||
estimator = LinearRegression().fit(np.array([[0.0], [1.0], [2.0]]), np.array([0.0, 1.0, 2.0]))
|
||||
joblib_lib.dump(estimator, tmp_path / 'model.pkl', compress=3)
|
||||
|
||||
loaded = mlflow_repository._load_model_pickle(
|
||||
str(tmp_path), 'prediction_model', Exception('flavor boom')
|
||||
)
|
||||
|
||||
assert loaded is not estimator
|
||||
assert loaded.predict(np.array([[3.0]]))[0] == pytest.approx(3.0)
|
||||
|
||||
Reference in New Issue
Block a user