gitea_admin/sientia-dataops-laborious-legacy
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

QTZPOC-20

Browse Source 
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
This commit is contained in:
Vitor Pimentel dos Santos
2026-08-26 15:45:33 -03:00
parent dac4f1718f
commit e95e65b72b
32 changed files with 4247 additions and 1040 deletions
Download Patch File Download Diff File Expand all files Collapse all files

647
tests/laborious/utils/repository/test_model_repository.py
View File

@@ -1,4 +1,8 @@
import ast
import importlib
import sys
from datetime import UTC, datetime
from pathlib import Path
from unittest.mock import ANY, AsyncMock, MagicMock, call, patch
import joblib as joblib_lib
@@ -10,6 +14,7 @@ from sientia_do.temporal.constants import DATETIME_FORMAT_WITH_TZ
from sklearn.linear_model import LinearRegression
from laborious import metrics
from laborious.utils.repository import model_repository
from laborious.utils.repository.model_repository import (
INVALID_FLAVOR_MESSAGE,
MLFlowRepository,
@@ -465,7 +470,7 @@ async def test_load_predict_model_pytorch(mlflow, mlflow_repository):
async def test_load_predict_model_error(mlflow_repository):
with pytest.raises(ValueError) as e:
await mlflow_repository.load_predict_model('test_model', {}, 'invalid')
assert str(e) == "Invalid flavor. Use 'sklearn' or 'pyfunc' or 'pytorch'."
assert str(e) == INVALID_FLAVOR_MESSAGE
mlflow_repository.emit_metric.assert_called_once_with(
metric_object=metrics.MODEL_READ_ERROR_COUNT, tags=ANY
@@ -542,7 +547,7 @@ async def test_load_transform_model_error(mlflow_repository):
with pytest.raises(ValueError) as e:
await mlflow_repository.load_transform_model('test_model', {}, 'invalid')
assert str(e) == "Invalid flavor. Use 'sklearn' or 'pyfunc' or 'pytorch'."
assert str(e) == INVALID_FLAVOR_MESSAGE
mlflow_repository.emit_metric.assert_called_once_with(
metric_object=metrics.MODEL_READ_ERROR_COUNT, tags=ANY
@@ -1124,7 +1129,7 @@ async def test_log_model_error(mlflow_repository):
await mlflow_repository.log_model(
model_data, 'invalid', 'prediction_model', metadata['metadata']
)
assert str(e) == "Invalid flavor. Use 'sklearn' or 'pyfunc' or 'pytorch'."
assert str(e) == INVALID_FLAVOR_MESSAGE
mlflow_repository.emit_metric.assert_called_once_with(
metric_object=metrics.MODEL_WRITE_ERROR_COUNT, tags=ANY
)
@@ -1768,10 +1773,19 @@ async def test_predict_success_ndarray(mlflow_repository):
"""
Tests for the joblib fallback on the model load paths
Tests for the joblib flavor on the model load and write paths
"""
JOBLIB_PATH = 'laborious.utils.repository.model_repository.joblib'
FLAVOR_LOADERS = ('sklearn', 'pyfunc', 'pytorch')
@pytest.fixture
def preserved_sys_path():
"""Restore `sys.path`: the joblib flavor inserts entries and never removes them."""
original = list(sys.path)
yield
sys.path[:] = original
def build_artifact_dir(tmp_path, *filenames):
@@ -1781,6 +1795,21 @@ def build_artifact_dir(tmp_path, *filenames):
return str(tmp_path)
def build_code_dir(tmp_path, module_name='probe_flavor_model'):
"""Add a `code/` directory carrying the model's class module to an artifact dir."""
code_dir = tmp_path / 'code'
code_dir.mkdir()
(code_dir / f'{module_name}.py').write_text(
'class ProbeModel:\n'
' def __init__(self, offset):\n'
' self.offset = offset\n'
'\n'
' def predict(self, data):\n'
' return [value + self.offset for value in data]\n'
)
return str(code_dir)
def emitted_tags(emit_metric_mock, metric_object):
"""Collect the `tags` kwarg of every emit_metric call for a given metric."""
return [
@@ -1790,86 +1819,22 @@ def emitted_tags(emit_metric_mock, 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()
def assert_no_flavor_loader_called(mlflow):
for flavor in FLAVOR_LOADERS:
getattr(mlflow, flavor).load_model.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(
async def test_load_predict_model_joblib_flavor_loads_pickle_directly(
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')
result = await mlflow_repository.load_predict_model('test_model', {}, 'joblib')
assert result is joblib_mock.load.return_value
mlflow_repository.get_model_run_id.assert_called_once_with(
@@ -1879,46 +1844,43 @@ async def test_load_predict_model_fallback_recovers_load(
'run-123', 'prediction_model'
)
joblib_mock.load.assert_called_once_with(str(tmp_path / 'model.pkl'))
assert_no_flavor_loader_called(mlflow)
@patch(JOBLIB_PATH)
@pytest.mark.asyncio
async def test_load_transform_model_fallback_recovers_load(
async def test_load_transform_model_joblib_flavor_loads_pickle_directly(
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')
result = await mlflow_repository.load_transform_model('test_model', {}, 'joblib')
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'))
mlflow_repository.get_model_uri.assert_not_called()
assert_no_flavor_loader_called(mlflow)
@patch(JOBLIB_PATH)
@pytest.mark.asyncio
async def test_load_transform_model_fallback_reuses_resolved_run(
async def test_load_transform_model_joblib_resolves_run_once(
joblib_mock, mlflow, mlflow_repository, tmp_path
):
"""A promotion between both steps must not change the run served."""
mlflow.sklearn.load_model.side_effect = Exception('flavor boom')
"""A promotion mid-load must not change the run served: one resolution per load."""
mlflow_repository.get_model_run_id = MagicMock(side_effect=['run-A', 'run-B'])
mlflow_repository.get_model_uri = MagicMock()
mlflow_repository.client.download_artifacts.return_value = build_artifact_dir(
tmp_path, 'MLmodel', 'model.pkl'
tmp_path, 'model.pkl'
)
result = await mlflow_repository.load_transform_model('test_model', {}, 'sklearn')
await mlflow_repository.load_transform_model('test_model', {}, 'joblib')
assert result is joblib_mock.load.return_value
mlflow_repository.get_model_run_id.assert_called_once_with(
model_name='test_model', stage='Production'
)
@@ -1927,57 +1889,91 @@ async def test_load_transform_model_fallback_reuses_resolved_run(
@patch(JOBLIB_PATH)
@pytest.mark.asyncio
async def test_download_model_wrapper_fallback_reuses_downloaded_dir(
joblib_mock, mlflow, mlflow_repository, tmp_path
@pytest.mark.parametrize('flavor', ['sklearn', 'pyfunc', 'pytorch'])
async def test_load_predict_model_flavor_failure_propagates(
joblib_mock, mlflow, mlflow_repository, flavor
):
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')
loader_error = mlflow_lib.exceptions.MlflowException('flavor boom')
getattr(mlflow, flavor).load_model.side_effect = loader_error
model, returned_path = await mlflow_repository.download_model(
'test_model', {}, 'predict', 'pyfunc', True
)
with pytest.raises(mlflow_lib.exceptions.MlflowException) as exc_info:
await mlflow_repository.load_predict_model('test_model', {}, flavor)
mlflow_repository.dowload_artifacts.assert_called_once_with(
'test_model', {}, 'prediction_model'
)
assert exc_info.value is loader_error
joblib_mock.load.assert_not_called()
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
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'
@patch(JOBLIB_PATH)
@pytest.mark.asyncio
async def test_load_predict_model_both_attempts_fail_chains_causes(
@pytest.mark.parametrize('flavor', ['sklearn', 'pyfunc', 'pytorch'])
async def test_load_transform_model_flavor_failure_propagates(
joblib_mock, mlflow, mlflow_repository, flavor
):
mlflow_repository.get_model_run_id = MagicMock()
mlflow_repository.get_model_uri = MagicMock()
loader_error = mlflow_lib.exceptions.MlflowException('flavor boom')
getattr(mlflow, flavor).load_model.side_effect = loader_error
with pytest.raises(mlflow_lib.exceptions.MlflowException) as exc_info:
await mlflow_repository.load_transform_model('test_model', {}, flavor)
assert exc_info.value is loader_error
joblib_mock.load.assert_not_called()
mlflow_repository.client.download_artifacts.assert_not_called()
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'
@patch(JOBLIB_PATH)
@pytest.mark.asyncio
async def test_download_model_wrapper_failure_propagates(joblib_mock, mlflow, mlflow_repository):
mlflow_repository.dowload_artifacts = AsyncMock()
loader_error = mlflow_lib.exceptions.MlflowException('wrapper boom')
mlflow.pyfunc.load_model.side_effect = loader_error
with pytest.raises(mlflow_lib.exceptions.MlflowException) as exc_info:
await mlflow_repository.download_model('test_model', {}, 'predict', 'pyfunc', True)
assert exc_info.value is loader_error
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_joblib_success_metrics(
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')
await mlflow_repository.load_predict_model('test_model', {}, 'joblib')
assert str(exc_info.value) == 'joblib boom'
assert exc_info.value.__cause__ is flavor_error
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'
assert emitted_tags(mlflow_repository.emit_metric, metrics.MODEL_READ_ERROR_COUNT) == []
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'
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'
@patch(JOBLIB_PATH)
@pytest.mark.asyncio
async def test_load_transform_model_both_attempts_fail_emits_error_metric(
async def test_load_transform_model_joblib_failure_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.get_model_uri = MagicMock()
mlflow_repository.client.download_artifacts.return_value = build_artifact_dir(
@@ -1985,65 +1981,41 @@ async def test_load_transform_model_both_attempts_fail_emits_error_metric(
)
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_download_model_wrapper_both_attempts_fail_chains_causes(
joblib_mock, mlflow, mlflow_repository, tmp_path
):
flavor_error = Exception('flavor boom')
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 = flavor_error
joblib_mock.load.side_effect = ValueError('joblib boom')
with pytest.raises(ValueError) as exc_info:
await mlflow_repository.download_model('test_model', {}, 'predict', 'pyfunc', True)
await mlflow_repository.load_transform_model('test_model', {}, 'joblib')
assert str(exc_info.value) == 'joblib boom'
assert exc_info.value.__cause__ is flavor_error
assert exc_info.value.__cause__ is None
error_tags = emitted_tags(mlflow_repository.emit_metric, metrics.MODEL_READ_ERROR_COUNT)
assert len(error_tags) == 1
assert error_tags[0]['operation_type'] == 'download_model_wrapper_joblib_fallback'
assert error_tags[0]['operation_type'] == 'load_transform_model'
assert emitted_tags(mlflow_repository.emit_metric, metrics.MODEL_READ_COUNT) == []
mlflow_repository.observe_lag.assert_not_called()
@patch(JOBLIB_PATH)
@pytest.mark.asyncio
async def test_load_predict_model_fallback_success_metrics(
async def test_no_load_emits_a_fallback_operation_type(
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', {}, 'joblib')
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) == []
emitted = [
emitted.kwargs['tags']['operation_type']
for emitted in mlflow_repository.emit_metric.call_args_list
]
observed = [
lag.args[2]['operation_type'] for lag in mlflow_repository.observe_lag.call_args_list
]
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]
assert emitted
assert all(not operation.endswith('_joblib_fallback') for operation in emitted + observed)
@pytest.mark.parametrize(
@@ -2057,9 +2029,7 @@ async def test_load_predict_model_fallback_success_metrics(
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')
)
result = mlflow_repository._find_model_pickle(local_dir, 'prediction_model')
assert result == str(tmp_path / expected)
@@ -2070,16 +2040,14 @@ def test_find_model_pickle_ignores_nested_pickles(mlflow_repository, tmp_path):
(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')
)
mlflow_repository._find_model_pickle(str(tmp_path), 'prediction_model')
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'))
mlflow_repository._find_model_pickle(str(tmp_path), 'data_model')
assert 'data_model' in str(exc_info.value)
@@ -2088,80 +2056,373 @@ def test_find_model_pickle_error_message_carries_context(mlflow_repository, tmp_
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')
)
mlflow_repository._find_model_pickle(local_dir, 'prediction_model')
message = str(exc_info.value)
assert 'prediction_model' in message
assert 'flavor boom' in message
assert 'MLmodel' in message
assert 'Original error' not 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')
)
mlflow_repository._find_model_pickle(str(tmp_path), 'prediction_model')
assert 'algo.pkl' in str(exc_info.value)
def test_resolve_artifact_code_dir_present(tmp_path, mlflow_repository):
code_dir = build_code_dir(tmp_path)
assert mlflow_repository.resolve_artifact_code_dir(str(tmp_path)) == code_dir
def test_resolve_artifact_code_dir_absent(tmp_path, mlflow_repository):
assert mlflow_repository.resolve_artifact_code_dir(str(tmp_path)) is None
def test_resolve_artifact_code_dir_without_artifact_path(mlflow_repository):
assert mlflow_repository.resolve_artifact_code_dir(None) is None
@patch(JOBLIB_PATH)
def test_load_model_pickle_adds_code_dir_before_load(
joblib_mock, mlflow_repository, tmp_path, preserved_sys_path
):
code_dir = build_code_dir(tmp_path)
local_dir = build_artifact_dir(tmp_path, 'model.pkl')
sys_path_at_load = []
joblib_mock.load.side_effect = lambda _path: sys_path_at_load.extend(sys.path)
mlflow_repository._load_model_pickle(local_dir, 'prediction_model')
joblib_mock.load.assert_called_once_with(str(tmp_path / 'model.pkl'))
assert code_dir in sys_path_at_load
assert sys.path[0] == code_dir
@patch(JOBLIB_PATH)
def test_load_model_pickle_without_code_dir_leaves_sys_path_alone(
joblib_mock, mlflow_repository, tmp_path, preserved_sys_path
):
local_dir = build_artifact_dir(tmp_path, 'model.pkl')
before = list(sys.path)
mlflow_repository._load_model_pickle(local_dir, 'data_model')
joblib_mock.load.assert_called_once_with(str(tmp_path / 'model.pkl'))
assert sys.path == before
@patch(JOBLIB_PATH)
def test_load_model_pickle_adds_code_dir_only_once(
joblib_mock, mlflow_repository, tmp_path, preserved_sys_path
):
code_dir = build_code_dir(tmp_path)
local_dir = build_artifact_dir(tmp_path, 'model.pkl')
mlflow_repository._load_model_pickle(local_dir, 'prediction_model')
mlflow_repository._load_model_pickle(local_dir, 'prediction_model')
assert sys.path.count(code_dir) == 1
@patch(JOBLIB_PATH)
def test_load_model_pickle_keeps_code_dir_after_load(
joblib_mock, mlflow_repository, tmp_path, preserved_sys_path
):
code_dir = build_code_dir(tmp_path)
local_dir = build_artifact_dir(tmp_path, 'model.pkl')
mlflow_repository._load_model_pickle(local_dir, 'prediction_model')
assert code_dir in sys.path
@patch(JOBLIB_PATH)
@pytest.mark.asyncio
async def test_get_model_caches_model_loaded_by_fallback(
async def test_log_model_joblib_writes_compressed_pickle_and_code(
joblib_mock, mlflow, mlflow_repository, tmp_path
):
code_dir = build_code_dir(tmp_path)
model_data = {'model': MagicMock(), 'artifact_path': str(tmp_path)}
await mlflow_repository.log_model(
model_data, 'joblib', 'prediction_model', metadata['metadata']
)
joblib_mock.dump.assert_called_once()
dumped_model, dumped_path = joblib_mock.dump.call_args.args
assert dumped_model is model_data['model']
assert dumped_path.endswith('/model.pkl')
assert joblib_mock.dump.call_args.kwargs == {'compress': 3}
mlflow.log_artifact.assert_called_once_with(dumped_path, artifact_path='prediction_model')
mlflow.log_artifacts.assert_called_once_with(code_dir, artifact_path='prediction_model/code')
for flavor in FLAVOR_LOADERS:
getattr(mlflow, flavor).log_model.assert_not_called()
mlflow_repository.observe_lag.assert_called_once_with(ANY, metrics.MODEL_WRITE_LAG, ANY)
mlflow_repository.emit_metric.assert_called_once_with(
metric_object=metrics.MODEL_WRITE_COUNT, tags=ANY
)
@patch(JOBLIB_PATH)
@pytest.mark.asyncio
async def test_log_model_joblib_without_code_dir_logs_only_pickle(
joblib_mock, mlflow, mlflow_repository, tmp_path
):
model_data = {'model': MagicMock(), 'artifact_path': str(tmp_path)}
await mlflow_repository.log_model(model_data, 'joblib', 'data_model', metadata['metadata'])
mlflow.log_artifact.assert_called_once_with(ANY, artifact_path='data_model')
mlflow.log_artifacts.assert_not_called()
@patch(JOBLIB_PATH)
@pytest.mark.asyncio
async def test_log_model_sklearn_does_not_dump_joblib(joblib_mock, mlflow, mlflow_repository):
model_data = {'model': MagicMock(), 'artifact_path': 'artifact_path'}
await mlflow_repository.log_model(
model_data, 'sklearn', 'prediction_model', metadata['metadata']
)
mlflow.sklearn.log_model.assert_called_once_with(model_data['model'], 'prediction_model')
joblib_mock.dump.assert_not_called()
mlflow.log_artifact.assert_not_called()
@patch(JOBLIB_PATH)
@pytest.mark.asyncio
async def test_download_model_joblib_returns_download_dir(
joblib_mock, mlflow, mlflow_repository, tmp_path
):
mlflow_repository.get_model_run_id = MagicMock(return_value='run-123')
local_dir = build_artifact_dir(tmp_path, 'model.pkl')
mlflow_repository.client.download_artifacts.return_value = local_dir
mlflow_repository.dowload_artifacts = AsyncMock()
model, artifact_path = await mlflow_repository.download_model(
'test_model', {}, 'predict', 'joblib'
)
assert model is joblib_mock.load.return_value
assert artifact_path == local_dir
mlflow_repository.dowload_artifacts.assert_not_called()
assert_no_flavor_loader_called(mlflow)
@patch(JOBLIB_PATH)
@pytest.mark.asyncio
async def test_download_model_sklearn_returns_no_artifact_path(
joblib_mock, mlflow, mlflow_repository
):
model, artifact_path = await mlflow_repository.download_model(
'test_model', {}, 'predict', 'sklearn'
)
assert model is mlflow.sklearn.load_model.return_value
assert artifact_path is None
joblib_mock.load.assert_not_called()
@patch(JOBLIB_PATH)
@pytest.mark.asyncio
async def test_download_model_transform_joblib_downloads_data_model(
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')
await mlflow_repository.download_model('test_model', {}, 'transform', 'joblib')
mlflow_repository.client.download_artifacts.assert_called_once_with('run-123', 'data_model')
@patch('laborious.utils.repository.model_repository.force_memory_release')
@patch('laborious.utils.repository.model_repository.rmtree')
@patch(JOBLIB_PATH)
@pytest.mark.asyncio
async def test_retrain_model_joblib_reaches_joblib_write_branch(
joblib_mock, _rmtree, _force_memory_release, mlflow, mlflow_repository, tmp_path
):
code_dir = build_code_dir(tmp_path)
model_config = {'target': 'target', 'transform_flavor': 'sklearn', 'predict_flavor': 'joblib'}
mlflow_repository.get_model_run_id = MagicMock(return_value='run-123')
mlflow_repository.get_model_params = MagicMock(return_value={})
mlflow_repository.get_experiment = MagicMock()
mlflow_repository.get_next_run_name = MagicMock()
mlflow_repository.fit_models = AsyncMock(
return_value={
'prediction_model': {'model': MagicMock(), 'artifact_path': str(tmp_path)},
'data_model': {'model': MagicMock(), 'artifact_path': None},
'prediction_data': MagicMock(spec=DataFrame),
}
)
output = await mlflow_repository.retrain_model(
MagicMock(), 'test_model', model_config, metadata['metadata']
)
assert output['success'] is True
mlflow_repository.fit_models.assert_called_once()
assert mlflow_repository.fit_models.call_args.kwargs['predict_flavor'] == 'joblib'
joblib_mock.dump.assert_called_once()
assert joblib_mock.dump.call_args.kwargs == {'compress': 3}
mlflow.log_artifact.assert_any_call(ANY, artifact_path='prediction_model')
mlflow.log_artifacts.assert_called_once_with(code_dir, artifact_path='prediction_model/code')
mlflow.sklearn.log_model.assert_called_once_with(ANY, 'data_model')
def test_module_defines_no_joblib_fallback_routine():
"""Guard against reintroducing 'try the flavor loader, then try pickle'."""
source = Path(model_repository.__file__).read_text()
assert 'joblib_fallback' not in source
# The fallback threaded the flavor loader's exception into the pickle helpers as
# `flavor_error`; nothing in the module may take that parameter again.
assert 'flavor_error' not in source
pickle_helpers = {
'_find_model_pickle',
'_load_model_pickle',
'_download_and_load_pickle',
'_load_joblib_artifact',
}
for handler in ast.walk(ast.parse(source)):
if not isinstance(handler, ast.ExceptHandler):
continue
for node in ast.walk(handler):
if not isinstance(node, ast.Call):
continue
func = node.func
if isinstance(func, ast.Attribute):
assert func.attr not in pickle_helpers, (
f'{func.attr} is called from an except handler'
)
assert not (isinstance(func.value, ast.Name) and func.value.id == 'joblib'), (
'joblib is called from an except handler'
)
@patch(JOBLIB_PATH)
@pytest.mark.asyncio
async def test_load_predict_model_invalid_flavor_before_any_io(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_before_any_io(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_log_model_invalid_flavor(joblib_mock, mlflow, mlflow_repository):
model_data = {'model': MagicMock(), 'artifact_path': 'artifact_path'}
with pytest.raises(ValueError) as exc_info:
await mlflow_repository.log_model(
model_data, 'invalid', 'prediction_model', metadata['metadata']
)
assert str(exc_info.value) == INVALID_FLAVOR_MESSAGE
joblib_mock.dump.assert_not_called()
mlflow_repository.emit_metric.assert_called_once_with(
metric_object=metrics.MODEL_WRITE_ERROR_COUNT, tags=ANY
)
mlflow_repository.observe_lag.assert_not_called()
def test_invalid_flavor_message_lists_every_accepted_flavor():
for flavor in ('sklearn', 'pyfunc', 'pytorch', 'joblib'):
assert f"'{flavor}'" in INVALID_FLAVOR_MESSAGE
@patch(JOBLIB_PATH)
@pytest.mark.asyncio
async def test_get_model_caches_joblib_loaded_model(
joblib_mock, mlflow, mlflow_repository, tmp_path
):
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', 'joblib')
second = await mlflow_repository.get_model('test_model', {}, 5, 'predict', 'joblib')
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
async def test_joblib_flavor_real_round_trip_with_code_dir(
mlflow_repository, tmp_path, preserved_sys_path
):
mlflow.sklearn.load_model.side_effect = Exception('flavor boom')
"""Real joblib, no mock: proves decompression and module resolution together."""
code_dir = build_code_dir(tmp_path)
sys.path.insert(0, code_dir)
try:
module = importlib.import_module('probe_flavor_model')
joblib_lib.dump(module.ProbeModel(2), tmp_path / 'model.pkl', compress=3)
finally:
sys.path.remove(code_dir)
del sys.modules['probe_flavor_model']
importlib.invalidate_caches()
with pytest.raises(ModuleNotFoundError):
importlib.import_module('probe_flavor_model')
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')
mlflow_repository.client.download_artifacts.return_value = str(tmp_path)
data = DataFrame({'feat_1': {'index_1': 2, 'index_2': 3}})
try:
loaded = await mlflow_repository.load_predict_model('test_model', {}, 'joblib')
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']
assert loaded.predict([1.0, 2.0]) == [3.0, 4.0]
assert code_dir in sys.path
finally:
sys.modules.pop('probe_flavor_model', None)
importlib.invalidate_caches()
def test_load_model_pickle_real_joblib_round_trip(mlflow_repository, tmp_path):
def test_log_model_joblib_real_round_trip(tmp_path, mlflow, mlflow_repository):
"""The write flavor dumps what the read flavor loads: same dialect by construction."""
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)
written = tmp_path / 'model.pkl'
joblib_lib.dump(estimator, written, compress=3)
loaded = mlflow_repository._load_model_pickle(
str(tmp_path), 'prediction_model', Exception('flavor boom')
)
loaded = mlflow_repository._load_model_pickle(str(tmp_path), 'prediction_model')
assert loaded is not estimator
assert loaded.predict(np.array([[3.0]]))[0] == pytest.approx(3.0)

50
tests/laborious/utils/test_connectors_config.py
View File

@@ -29,11 +29,20 @@ IMPORT_ENV_VARS = (
'IMPORT_STATUS_DB_PASSWORD',
)
# The import log connection falls back to the POSTGRES_* server, so a developer's own POSTGRES_*
# would otherwise decide what `build_import_status_config` returns under test.
POSTGRES_ENV_VARS = (
'POSTGRES_HOST',
'POSTGRES_PORT',
'POSTGRES_USER',
'POSTGRES_PASSWORD',
)
@pytest.fixture
def clean_import_env(monkeypatch):
"""Import configuration is read from a known-empty environment."""
for name in IMPORT_ENV_VARS:
for name in IMPORT_ENV_VARS + POSTGRES_ENV_VARS:
monkeypatch.delenv(name, raising=False)
return monkeypatch
@@ -208,12 +217,49 @@ def test_build_import_status_config_with_defaults(clean_import_env):
'host': 'localhost',
'port': 5432,
'dbname': 'sientia-core-mlops-bff',
'user': 'mlops_bff_user',
'user': 'sientia',
'password': '',
'schema': 'public',
}
def test_build_import_status_config_uses_the_postgres_server(clean_import_env):
"""The log is another database on the POSTGRES_* server: only its name is import configuration."""
clean_import_env.setenv('POSTGRES_HOST', 'paradedb-rw')
clean_import_env.setenv('POSTGRES_PORT', '5433')
clean_import_env.setenv('POSTGRES_USER', 'sientia-user')
clean_import_env.setenv('POSTGRES_PASSWORD', 'sientia-password')
clean_import_env.setenv('IMPORT_STATUS_DB_NAME', 'bff-db')
assert build_import_status_config() == {
'host': 'paradedb-rw',
'port': 5433,
'dbname': 'bff-db',
'user': 'sientia-user',
'password': 'sientia-password',
'schema': 'public',
}
def test_build_import_status_config_overrides_beat_the_postgres_server(clean_import_env):
"""A log on a server of its own: the four IMPORT_STATUS_DB_* variables win over POSTGRES_*."""
clean_import_env.setenv('POSTGRES_HOST', 'paradedb-rw')
clean_import_env.setenv('POSTGRES_PORT', '5433')
clean_import_env.setenv('POSTGRES_USER', 'sientia-user')
clean_import_env.setenv('POSTGRES_PASSWORD', 'sientia-password')
clean_import_env.setenv('IMPORT_STATUS_DB_HOST', 'bff-postgres')
clean_import_env.setenv('IMPORT_STATUS_DB_PORT', '6543')
clean_import_env.setenv('IMPORT_STATUS_DB_USER', 'bff-user')
clean_import_env.setenv('IMPORT_STATUS_DB_PASSWORD', 'bff-password')
config = build_import_status_config()
assert config['host'] == 'bff-postgres'
assert config['port'] == 6543
assert config['user'] == 'bff-user'
assert config['password'] == 'bff-password'
def test_build_import_status_config_reads_the_environment(clean_import_env):
clean_import_env.setenv('IMPORT_STATUS_DB_HOST', 'bff-postgres')
clean_import_env.setenv('IMPORT_STATUS_DB_PORT', '6543')