vitor-aignosi
e95e65b72b
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
127 lines
4.6 KiB
Python
127 lines
4.6 KiB
Python
"""
|
|
End-to-end test for retraining a model served through the explicit `joblib` flavor.
|
|
|
|
Retraining is not part of the `predictions_batch` workflow, so it gets its own module. The test
|
|
is a write→read round trip in one run: `retrain_model` loads the production artifact through the
|
|
joblib flavor, `log_model` writes the retrained model back as `model.pkl` + `code/`, and the
|
|
written artifact is then loaded again through the same flavor. Neither `joblib.dump` nor
|
|
`joblib.load` is mocked; only MLflow's own I/O is.
|
|
"""
|
|
|
|
import importlib
|
|
import shutil
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
import pandas as pd
|
|
import pytest
|
|
|
|
from e2e.conftest import JOBLIB_FLAVOR_CODE_MODULE, JOBLIB_FLAVOR_MODEL_NAME
|
|
from laborious.activities.activities import Activities
|
|
|
|
RETRAIN_INDEX = [
|
|
'2024-01-01 12:00:00+0000',
|
|
'2024-01-01 12:01:00+0000',
|
|
'2024-01-01 12:02:00+0000',
|
|
]
|
|
|
|
METADATA = {
|
|
'model_id': 'test_model',
|
|
'model_name': JOBLIB_FLAVOR_MODEL_NAME,
|
|
'workflow_name': 'test_workflow',
|
|
'schema_name': 'test_schedule',
|
|
}
|
|
|
|
|
|
def build_retrain_data() -> pd.DataFrame:
|
|
"""Training frame indexed by timestamp strings, carrying the target column."""
|
|
return pd.DataFrame(
|
|
{
|
|
'sensor_1': [60.0, 61.0, 62.0],
|
|
'sensor_2': [78.2, 78.4, 78.6],
|
|
'target': [1.0, 2.0, 3.0],
|
|
},
|
|
index=RETRAIN_INDEX,
|
|
)
|
|
|
|
|
|
class FakeArtifactStore:
|
|
"""Minimal stand-in for the run's artifact tree, so the written layout is inspectable."""
|
|
|
|
def __init__(self, root: Path):
|
|
self.root = root
|
|
|
|
def log_artifact(self, local_path, artifact_path=None):
|
|
if artifact_path is None:
|
|
# Retraining also logs the CSVs at the run root; irrelevant to this test.
|
|
return
|
|
target = self.root / artifact_path
|
|
target.mkdir(parents=True, exist_ok=True)
|
|
shutil.copy(local_path, target / Path(local_path).name)
|
|
|
|
def log_artifacts(self, local_dir, artifact_path=None):
|
|
target = self.root / artifact_path if artifact_path else self.root
|
|
shutil.copytree(local_dir, target, dirs_exist_ok=True)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.integration
|
|
async def test_retrain_joblib_model_writes_pickle_and_code_and_reloads(
|
|
test_activities: Activities,
|
|
joblib_flavor_artifacts,
|
|
mlflow_client,
|
|
patch_mlflow,
|
|
tmp_path,
|
|
):
|
|
"""A joblib retrain writes `model.pkl` + `code/`, and the joblib flavor reads it back."""
|
|
repository = test_activities.model_monitoring_repository
|
|
|
|
store = FakeArtifactStore(tmp_path / 'run-retrained')
|
|
patch_mlflow.log_artifact.side_effect = store.log_artifact
|
|
patch_mlflow.log_artifacts.side_effect = store.log_artifacts
|
|
patch_mlflow.get_run.return_value.data.params = {}
|
|
patch_mlflow.start_run.return_value.__enter__.return_value.info.run_id = 'run-retrained'
|
|
|
|
output = await repository.retrain_model(
|
|
data=build_retrain_data(),
|
|
model_name=JOBLIB_FLAVOR_MODEL_NAME,
|
|
model_config={
|
|
'target': 'target',
|
|
'transform_flavor': 'joblib',
|
|
'predict_flavor': 'joblib',
|
|
},
|
|
metadata=METADATA,
|
|
)
|
|
|
|
assert output['success'] is True, output.get('traceback')
|
|
|
|
# Both models were loaded through the joblib flavor, not through a flavor loader.
|
|
downloaded = [loaded.args[1] for loaded in mlflow_client.download_artifacts.call_args_list]
|
|
assert sorted(downloaded) == ['data_model', 'prediction_model']
|
|
assert patch_mlflow.sklearn.load_model.call_count == 0
|
|
assert patch_mlflow.pyfunc.load_model.call_count == 0
|
|
|
|
written = store.root / 'prediction_model'
|
|
assert (written / 'model.pkl').is_file()
|
|
assert (written / 'code' / f'{JOBLIB_FLAVOR_CODE_MODULE}.py').is_file()
|
|
assert not (written / 'MLmodel').exists()
|
|
assert (store.root / 'data_model' / 'model.pkl').is_file()
|
|
|
|
# Drop the module the first load put in place: the reload must resolve it from the `code/`
|
|
# the retrain wrote, not from what is already imported.
|
|
for code_dir in (artifact['code_dir'] for artifact in joblib_flavor_artifacts.values()):
|
|
if code_dir in sys.path:
|
|
sys.path.remove(code_dir)
|
|
sys.modules.pop(JOBLIB_FLAVOR_CODE_MODULE, None)
|
|
importlib.invalidate_caches()
|
|
|
|
mlflow_client.download_artifacts.side_effect = lambda _run_id, _artifact_name: str(written)
|
|
|
|
reloaded = await repository.load_predict_model(JOBLIB_FLAVOR_MODEL_NAME, METADATA, 'joblib')
|
|
|
|
features = pd.DataFrame({'feature_1': [0.234], 'feature_2': [0.783]})
|
|
expected = joblib_flavor_artifacts['prediction_model']['model'].predict(features)
|
|
|
|
assert reloaded.predict(features) == pytest.approx(expected)
|
|
assert str(written / 'code') in sys.path
|