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')

0
tests/scripts/__init__.py
View File

217
tests/scripts/test_import_model_manual_run.py
View File

@@ -1,217 +0,0 @@
"""The anti-rot guard for `scripts/import_model_manual_run.py`.
The script itself is only ever smoke-run by a person (task 13.6 of `import-model-e2e-suite`), so
nothing else would notice the three things most likely to be quietly broken in it:
1. the `# %%` cell layout — the whole contract is "cell N does step N", and a lost or reordered
marker turns a ten-step procedure into an unreadable one;
2. the boundary of design E15 — the script stands in for the frontend, so borrowing
`laborious.activities` or `laborious.workflows` would make it exercise the implementation instead
of the contract, and that is exactly the shortcut somebody reaches for when a cell needs one more
value;
3. the configuration rule — no host, user, password or key literal in the file, every target from
the environment, and every variable it reads documented in `.env.example`.
The wire shapes the script produces (the digest, the envelope, the object key, the row) are asserted
by the e2e suite against the real workflow, so they are deliberately not re-asserted here. This file
is about the script staying the *kind of thing* it is.
"""
from __future__ import annotations
import ast
import re
from pathlib import Path
import pytest
REPO_ROOT = Path(__file__).resolve().parents[2]
SCRIPT_PATH = REPO_ROOT / 'scripts' / 'import_model_manual_run.py'
ENV_EXAMPLE_PATH = REPO_ROOT / '.env.example'
CELL_MARKER = '# %%'
EXPECTED_CELLS = 10
FORBIDDEN_IMPORT_ROOTS = ('laborious.activities', 'laborious.workflows')
ALLOWED_LABORIOUS_IMPORT = 'laborious.utils.bundle.format'
# Substrings that would mean a target or a credential was written into the file instead of read
# from the environment. `mongodb://` and `postgresql` are schemes, not hosts, so they are not here.
FORBIDDEN_LITERAL_SUBSTRINGS = (
'localhost',
'127.0.0.1',
'0.0.0.0', # noqa: S104 - a denylist entry, not an interface this test binds
'::1',
'.svc.cluster.local',
'minioadmin',
'mlops_bff_user',
)
# A dotted host with a port (`db.example.com:5432`), an IPv4 address, or inline URL credentials.
HOSTLIKE_PATTERNS = (
re.compile(r'[A-Za-z0-9][A-Za-z0-9.-]*\.[A-Za-z]{2,}:\d{2,5}'),
re.compile(r'\b\d{1,3}(?:\.\d{1,3}){3}\b'),
re.compile(r'//[^/\s]+:[^/\s]+@'),
)
@pytest.fixture(scope='module')
def source() -> str:
"""The script's source text."""
return SCRIPT_PATH.read_text(encoding='utf-8')
@pytest.fixture(scope='module')
def tree(source: str) -> ast.Module:
"""The script parsed. A syntax error here is itself the finding — the module must be runnable."""
return ast.parse(source, filename=str(SCRIPT_PATH))
def _string_literals(tree: ast.Module) -> list[str]:
"""Every string constant in the module, comments and docstrings excluded."""
docstrings = set()
for node in ast.walk(tree):
if isinstance(node, ast.Module | ast.FunctionDef | ast.AsyncFunctionDef | ast.ClassDef):
body = getattr(node, 'body', [])
if (
body
and isinstance(body[0], ast.Expr)
and isinstance(body[0].value, ast.Constant)
and isinstance(body[0].value.value, str)
):
docstrings.add(id(body[0].value))
return [
node.value
for node in ast.walk(tree)
if isinstance(node, ast.Constant)
and isinstance(node.value, str)
and id(node) not in docstrings
]
def _env_names(tree: ast.Module) -> set[str]:
"""Every environment variable the script reads, by name."""
names: set[str] = set()
readers = {'env', 'env_required', 'getenv'}
for node in ast.walk(tree):
if not isinstance(node, ast.Call) or not node.args:
continue
target = node.func
called = (
target.attr
if isinstance(target, ast.Attribute)
else target.id
if isinstance(target, ast.Name)
else ''
)
first = node.args[0]
if called in readers and isinstance(first, ast.Constant) and isinstance(first.value, str):
names.add(first.value)
return names
class TestCellLayout:
"""The `# %%` contract of design E15."""
def test_script_exists(self) -> None:
assert SCRIPT_PATH.is_file(), f'{SCRIPT_PATH} is missing'
def test_every_cell_marker_is_present(self, source: str) -> None:
markers = [line for line in source.splitlines() if line.strip() == CELL_MARKER]
assert len(markers) == EXPECTED_CELLS, (
f'expected {EXPECTED_CELLS} `{CELL_MARKER}` cell markers, found {len(markers)}'
)
def test_cells_are_numbered_in_order(self, source: str) -> None:
"""Each cell opens with `# Cell N — ...`, and N counts 1..10 without a gap."""
lines = source.splitlines()
numbers: list[int] = []
for index, line in enumerate(lines):
if line.strip() != CELL_MARKER:
continue
following = lines[index + 1] if index + 1 < len(lines) else ''
match = re.match(r'#\s*Cell\s+(\d+)\b', following.strip())
assert match, f'the line after the marker on line {index + 1} does not name its cell'
numbers.append(int(match.group(1)))
assert numbers == list(range(1, EXPECTED_CELLS + 1))
def test_the_marker_precedes_the_first_statement(self, source: str) -> None:
"""The module docstring comes first, then cell 1 — nothing executable before a marker."""
body = source.split('"""', 2)[-1]
assert body.lstrip().startswith(CELL_MARKER)
class TestBoundary:
"""Design E15: the script stands in for the frontend, so it borrows no consumer code."""
def test_no_import_from_activities_or_workflows(self, tree: ast.Module) -> None:
for node in ast.walk(tree):
if isinstance(node, ast.Import):
for alias in node.names:
assert not alias.name.startswith(FORBIDDEN_IMPORT_ROOTS), (
f'the script imports {alias.name}'
)
elif isinstance(node, ast.ImportFrom) and node.module:
assert not node.module.startswith(FORBIDDEN_IMPORT_ROOTS), (
f'the script imports from {node.module}'
)
def test_the_only_laborious_import_is_the_format_helper(self, tree: ast.Module) -> None:
modules: set[str] = set()
for node in ast.walk(tree):
if isinstance(node, ast.Import):
modules.update(a.name for a in node.names if a.name.startswith('laborious'))
elif (
isinstance(node, ast.ImportFrom)
and node.module
and node.module.startswith('laborious')
):
modules.add(node.module)
assert modules <= {ALLOWED_LABORIOUS_IMPORT}, (
f'unexpected laborious imports: {sorted(modules - {ALLOWED_LABORIOUS_IMPORT})}'
)
class TestConfiguration:
"""No credential and no target literal in the file; every variable documented."""
@pytest.mark.parametrize('forbidden', FORBIDDEN_LITERAL_SUBSTRINGS)
def test_no_forbidden_substring(self, tree: ast.Module, forbidden: str) -> None:
offenders = [text for text in _string_literals(tree) if forbidden in text]
assert not offenders, f'{forbidden!r} appears in a string literal: {offenders}'
def test_no_hostlike_literal(self, tree: ast.Module) -> None:
for text in _string_literals(tree):
for pattern in HOSTLIKE_PATTERNS:
assert not pattern.search(text), (
f'{text!r} looks like a host, an address or inline credentials'
)
def test_every_variable_it_reads_is_documented(self, tree: ast.Module) -> None:
documented = {
line.split('=', 1)[0].strip()
for line in ENV_EXAMPLE_PATH.read_text(encoding='utf-8').splitlines()
if '=' in line and not line.lstrip().startswith('#')
}
undocumented = sorted(
_env_names(tree) - documented - {'MLFLOW_TRACKING_USERNAME', 'MLFLOW_TRACKING_PASSWORD'}
)
assert not undocumented, f'.env.example does not document: {undocumented}'
def test_it_reads_the_variables_the_contract_names(self, tree: ast.Module) -> None:
names = _env_names(tree)
required = {
'IMPORT_PASSWORD_KEY',
'IMPORT_BUNDLE_BUCKET',
'IMPORT_BUNDLE_PREFIX',
'IMPORT_STATUS_DB_HOST',
'IMPORT_STATUS_DB_PORT',
'IMPORT_STATUS_DB_NAME',
'IMPORT_STATUS_DB_USER',
'IMPORT_STATUS_DB_PASSWORD',
'IMPORT_SCRIPT_BUNDLE_PATH',
'IMPORT_SCRIPT_BUNDLE_PASSWORD',
'IMPORT_SCRIPT_USERNAME',
'RUNTIME',
}
assert required <= names, f'the script no longer reads: {sorted(required - names)}'