vitor-aignosi
30ff378cb7
QTZPOC-20: Prioritize code/ directory download and sys.path injection before model loading, enabling MLflow flavors to stream models directly from URIs and simplifying joblib artifact handling.
129 lines
4.6 KiB
Python
129 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, _local_dir = 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
|