QTZPOC-20
QTZPOC-20: Unify artifact code directory handling for all model flavors and standardize joblib pickle path Implemented a robust mechanism to ensure an artifact's code/ directory is prepended to sys.path before loading for sklearn, pyfunc, and joblib flavors. sklearn and pyfunc now resolve MLflow URIs to local directories first to facilitate this sys.path modification. Standardized the joblib flavor to explicitly read and write model.pkl within the artifact directory. Updated documentation and added comprehensive tests to reflect these changes.
This commit is contained in:
@@ -17,6 +17,7 @@ from laborious import metrics
|
||||
from laborious.utils.repository import model_repository
|
||||
from laborious.utils.repository.model_repository import (
|
||||
INVALID_FLAVOR_MESSAGE,
|
||||
MODEL_PICKLE_NAME,
|
||||
MLFlowRepository,
|
||||
force_memory_release,
|
||||
)
|
||||
@@ -68,6 +69,11 @@ def mlflow_repository():
|
||||
@pytest.fixture
|
||||
def mlflow():
|
||||
with patch('laborious.utils.repository.model_repository.mlflow') as mlflow:
|
||||
# Every URI load resolves the artifact to a local directory before calling the flavor
|
||||
# loader, and `os.path` refuses a MagicMock. Echoing the URI back keeps the loader's
|
||||
# argument readable and points the code-directory lookup at a path that does not exist,
|
||||
# so no test touches `sys.path` unless it builds a real artifact directory.
|
||||
mlflow.artifacts.download_artifacts.side_effect = lambda artifact_uri: str(artifact_uri)
|
||||
yield mlflow
|
||||
|
||||
|
||||
@@ -467,6 +473,11 @@ async def test_load_predict_model_error(mlflow_repository):
|
||||
mlflow_repository.observe_lag.assert_not_called()
|
||||
|
||||
|
||||
# A real string, because the transform load now resolves its URI to a local directory before
|
||||
# calling the loader, and that resolution goes through `os.path`.
|
||||
TRANSFORM_MODEL_URI = 'runs:/run-123/data_model'
|
||||
|
||||
|
||||
def validate_common_load_transform_model_mocks(mlflow_repository, model_name):
|
||||
mlflow_repository.get_model_run_id.assert_called_once_with(
|
||||
model_name=model_name, stage='Production'
|
||||
@@ -479,14 +490,14 @@ def validate_common_load_transform_model_mocks(mlflow_repository, model_name):
|
||||
@pytest.mark.asyncio
|
||||
async def test_load_transform_model_sklearn(mlflow, mlflow_repository):
|
||||
mlflow_repository.get_model_run_id = MagicMock()
|
||||
mlflow_repository.get_model_uri = MagicMock()
|
||||
mlflow_repository.get_model_uri = MagicMock(return_value=TRANSFORM_MODEL_URI)
|
||||
|
||||
result = await mlflow_repository.load_transform_model('test_model', {}, 'sklearn')
|
||||
|
||||
validate_common_load_transform_model_mocks(mlflow_repository, 'test_model')
|
||||
|
||||
assert result == mlflow.sklearn.load_model.return_value
|
||||
mlflow.sklearn.load_model.assert_called_once_with(mlflow_repository.get_model_uri.return_value)
|
||||
mlflow.sklearn.load_model.assert_called_once_with(TRANSFORM_MODEL_URI)
|
||||
|
||||
mlflow_repository.observe_lag.assert_called_once_with(ANY, metrics.MODEL_READ_LAG, ANY)
|
||||
mlflow_repository.emit_metric.assert_called_once_with(
|
||||
@@ -497,14 +508,14 @@ async def test_load_transform_model_sklearn(mlflow, mlflow_repository):
|
||||
@pytest.mark.asyncio
|
||||
async def test_load_transform_model_pyfunc(mlflow, mlflow_repository):
|
||||
mlflow_repository.get_model_run_id = MagicMock()
|
||||
mlflow_repository.get_model_uri = MagicMock()
|
||||
mlflow_repository.get_model_uri = MagicMock(return_value=TRANSFORM_MODEL_URI)
|
||||
|
||||
result = await mlflow_repository.load_transform_model('test_model', {}, 'pyfunc')
|
||||
|
||||
validate_common_load_transform_model_mocks(mlflow_repository, 'test_model')
|
||||
|
||||
assert result == mlflow.pyfunc.load_model.return_value
|
||||
mlflow.pyfunc.load_model.assert_called_once_with(mlflow_repository.get_model_uri.return_value)
|
||||
mlflow.pyfunc.load_model.assert_called_once_with(TRANSFORM_MODEL_URI)
|
||||
|
||||
mlflow_repository.observe_lag.assert_called_once_with(ANY, metrics.MODEL_READ_LAG, ANY)
|
||||
mlflow_repository.emit_metric.assert_called_once_with(
|
||||
@@ -543,8 +554,8 @@ async def test_download_model_invalid_model_type(mlflow_repository):
|
||||
@pytest.mark.parametrize(
|
||||
'model_type', [('predict', 'prediction_model'), ('transform', 'data_model')]
|
||||
)
|
||||
async def test_download_model_load_wrapper(mlflow, mlflow_repository, model_type):
|
||||
mlflow_repository.dowload_artifacts = AsyncMock()
|
||||
async def test_download_model_load_wrapper(mlflow, mlflow_repository, model_type, tmp_path):
|
||||
mlflow_repository.dowload_artifacts = AsyncMock(return_value=str(tmp_path))
|
||||
|
||||
result = await mlflow_repository.download_model('test_model', {}, model_type[0], 'pyfunc', True)
|
||||
|
||||
@@ -1740,14 +1751,14 @@ FLAVOR_LOADERS = ('sklearn', 'pyfunc')
|
||||
|
||||
@pytest.fixture
|
||||
def preserved_sys_path():
|
||||
"""Restore `sys.path`: the joblib flavor inserts entries and never removes them."""
|
||||
"""Restore `sys.path`: every flavor inserts artifact code entries and never removes them."""
|
||||
original = list(sys.path)
|
||||
yield
|
||||
sys.path[:] = original
|
||||
|
||||
|
||||
def build_artifact_dir(tmp_path, *filenames):
|
||||
"""Create a fake downloaded artifact directory with the given top-level files."""
|
||||
"""Create a fake download of one model artifact directory, holding the given files."""
|
||||
for filename in filenames:
|
||||
(tmp_path / filename).write_bytes(b'pickle-bytes')
|
||||
return str(tmp_path)
|
||||
@@ -1979,7 +1990,12 @@ async def test_no_load_emits_a_fallback_operation_type(
|
||||
@pytest.mark.parametrize(
|
||||
('filenames', 'expected'),
|
||||
[
|
||||
# The canonical name `log_model` writes, alongside a real MLflow manifest.
|
||||
(('MLmodel', 'model.pkl'), 'model.pkl'),
|
||||
# `model.pkl` wins over any other candidate, whatever the name order.
|
||||
(('a_model.pkl', 'model.pkl'), 'model.pkl'),
|
||||
(('model.pkl', 'z_model.pkl'), 'model.pkl'),
|
||||
# Fallback for hand-built artifacts: no `model.pkl`, so name order decides.
|
||||
(('MLmodel', 'python_model.pkl'), 'python_model.pkl'),
|
||||
(('b_model.pkl', 'a_model.pkl'), 'a_model.pkl'),
|
||||
],
|
||||
@@ -1993,6 +2009,7 @@ def test_find_model_pickle_selection(mlflow_repository, tmp_path, filenames, exp
|
||||
|
||||
|
||||
def test_find_model_pickle_ignores_nested_pickles(mlflow_repository, tmp_path):
|
||||
"""A pyfunc wrapper's own `artifacts/*.pkl` is never taken for the model."""
|
||||
nested = tmp_path / 'artifacts'
|
||||
nested.mkdir()
|
||||
(nested / 'stacking_model.pkl').write_bytes(b'pickle-bytes')
|
||||
@@ -2003,6 +2020,15 @@ def test_find_model_pickle_ignores_nested_pickles(mlflow_repository, tmp_path):
|
||||
assert 'No .pkl file found' in str(exc_info.value)
|
||||
|
||||
|
||||
def test_find_model_pickle_reads_the_name_log_model_writes(mlflow_repository, tmp_path):
|
||||
"""Read and write agree: both go through `MODEL_PICKLE_NAME`."""
|
||||
local_dir = build_artifact_dir(tmp_path, MODEL_PICKLE_NAME)
|
||||
|
||||
result = mlflow_repository._find_model_pickle(local_dir, 'prediction_model')
|
||||
|
||||
assert result == str(tmp_path / MODEL_PICKLE_NAME)
|
||||
|
||||
|
||||
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')
|
||||
@@ -2099,6 +2125,125 @@ def test_load_model_pickle_keeps_code_dir_after_load(
|
||||
assert code_dir in sys.path
|
||||
|
||||
|
||||
"""
|
||||
Tests for the artifact `code/` directory on every flavor's load path
|
||||
"""
|
||||
|
||||
|
||||
def observe_sys_path(recorder):
|
||||
"""A loader stand-in that records `sys.path` as it was when the load was called."""
|
||||
|
||||
def loader(_path):
|
||||
recorder.extend(sys.path)
|
||||
return MagicMock()
|
||||
|
||||
return loader
|
||||
|
||||
|
||||
def test_prepend_artifact_code_dir_inserts_at_the_front(
|
||||
mlflow_repository, tmp_path, preserved_sys_path
|
||||
):
|
||||
code_dir = build_code_dir(tmp_path)
|
||||
|
||||
result = mlflow_repository.prepend_artifact_code_dir(str(tmp_path), 'sklearn')
|
||||
|
||||
assert result == code_dir
|
||||
assert sys.path[0] == code_dir
|
||||
|
||||
|
||||
def test_prepend_artifact_code_dir_without_code_dir_leaves_sys_path_alone(
|
||||
mlflow_repository, tmp_path, preserved_sys_path
|
||||
):
|
||||
before = list(sys.path)
|
||||
|
||||
assert mlflow_repository.prepend_artifact_code_dir(str(tmp_path), 'pyfunc') is None
|
||||
assert sys.path == before
|
||||
|
||||
|
||||
def test_prepend_artifact_code_dir_does_not_duplicate(
|
||||
mlflow_repository, tmp_path, preserved_sys_path
|
||||
):
|
||||
code_dir = build_code_dir(tmp_path)
|
||||
|
||||
mlflow_repository.prepend_artifact_code_dir(str(tmp_path), 'sklearn')
|
||||
|
||||
assert mlflow_repository.prepend_artifact_code_dir(str(tmp_path), 'sklearn') is None
|
||||
assert sys.path.count(code_dir) == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize('flavor', FLAVOR_LOADERS)
|
||||
async def test_load_predict_model_puts_code_dir_on_sys_path_before_the_loader(
|
||||
flavor, mlflow, mlflow_repository, tmp_path, preserved_sys_path
|
||||
):
|
||||
"""MLflow extends `sys.path` only when `MLmodel` declares `code`; the directory decides here."""
|
||||
code_dir = build_code_dir(tmp_path)
|
||||
local_dir = build_artifact_dir(tmp_path, 'MLmodel', 'model.pkl')
|
||||
mlflow.artifacts.download_artifacts.side_effect = None
|
||||
mlflow.artifacts.download_artifacts.return_value = local_dir
|
||||
sys_path_at_load = []
|
||||
getattr(mlflow, flavor).load_model.side_effect = observe_sys_path(sys_path_at_load)
|
||||
|
||||
await mlflow_repository.load_predict_model('test_model', {}, flavor)
|
||||
|
||||
mlflow.artifacts.download_artifacts.assert_called_once_with(
|
||||
artifact_uri='models:/test_model/production'
|
||||
)
|
||||
getattr(mlflow, flavor).load_model.assert_called_once_with(local_dir)
|
||||
assert sys_path_at_load[0] == code_dir
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize('flavor', FLAVOR_LOADERS)
|
||||
async def test_load_transform_model_puts_code_dir_on_sys_path_before_the_loader(
|
||||
flavor, mlflow, mlflow_repository, tmp_path, preserved_sys_path
|
||||
):
|
||||
code_dir = build_code_dir(tmp_path)
|
||||
local_dir = build_artifact_dir(tmp_path, 'MLmodel', 'model.pkl')
|
||||
mlflow_repository.get_model_run_id = MagicMock(return_value='run-123')
|
||||
mlflow_repository.get_model_uri = MagicMock(return_value=TRANSFORM_MODEL_URI)
|
||||
mlflow.artifacts.download_artifacts.side_effect = None
|
||||
mlflow.artifacts.download_artifacts.return_value = local_dir
|
||||
sys_path_at_load = []
|
||||
getattr(mlflow, flavor).load_model.side_effect = observe_sys_path(sys_path_at_load)
|
||||
|
||||
await mlflow_repository.load_transform_model('test_model', {}, flavor)
|
||||
|
||||
mlflow.artifacts.download_artifacts.assert_called_once_with(artifact_uri=TRANSFORM_MODEL_URI)
|
||||
getattr(mlflow, flavor).load_model.assert_called_once_with(local_dir)
|
||||
assert sys_path_at_load[0] == code_dir
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize('flavor', FLAVOR_LOADERS)
|
||||
async def test_load_predict_model_without_code_dir_leaves_sys_path_alone(
|
||||
flavor, mlflow, mlflow_repository, tmp_path, preserved_sys_path
|
||||
):
|
||||
mlflow.artifacts.download_artifacts.side_effect = None
|
||||
mlflow.artifacts.download_artifacts.return_value = build_artifact_dir(tmp_path, 'MLmodel')
|
||||
before = list(sys.path)
|
||||
|
||||
await mlflow_repository.load_predict_model('test_model', {}, flavor)
|
||||
|
||||
assert sys.path == before
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_download_model_load_wrapper_puts_code_dir_on_sys_path(
|
||||
mlflow, mlflow_repository, tmp_path, preserved_sys_path
|
||||
):
|
||||
"""The wrapper branch downloads itself, so it is the one that has to add the directory."""
|
||||
code_dir = build_code_dir(tmp_path)
|
||||
mlflow_repository.dowload_artifacts = AsyncMock(return_value=str(tmp_path))
|
||||
sys_path_at_load = []
|
||||
mlflow.pyfunc.load_model.side_effect = observe_sys_path(sys_path_at_load)
|
||||
|
||||
await mlflow_repository.download_model('test_model', {}, 'predict', 'pyfunc', True)
|
||||
|
||||
assert sys_path_at_load[0] == code_dir
|
||||
assert sys.path[0] == code_dir
|
||||
|
||||
|
||||
@patch(JOBLIB_PATH)
|
||||
@pytest.mark.asyncio
|
||||
async def test_log_model_joblib_writes_compressed_pickle_and_code(
|
||||
|
||||
Reference in New Issue
Block a user