diff --git a/laborious/utils/repository/model_repository.py b/laborious/utils/repository/model_repository.py index 4d6fcc8..8d3d9f6 100644 --- a/laborious/utils/repository/model_repository.py +++ b/laborious/utils/repository/model_repository.py @@ -1330,7 +1330,7 @@ class MLFlowRepository(SientiaMonitoring): return code_dir if path.isdir(code_dir) else None def resolve_code_paths(self, artifact_path: str | None, metadata: dict) -> list[str]: - """Resolve the code dependencies to forward when re-logging a pyfunc model. + """Resolve the code dependencies to forward when re-logging a model. MLflow flattens whatever was given as `code_paths` into `/code/`, so the layout differs per model: some ship a `utils` package, others a @@ -1390,15 +1390,21 @@ class MLFlowRepository(SientiaMonitoring): core_labels = self.get_core_labels(metadata, operation_type='log_model') start_time = time.time() + code_paths = ( + self.resolve_code_paths(model_data.get('artifact_path'), metadata) + if flavor in ('sklearn', 'pyfunc') + else None + ) + try: if flavor == 'sklearn': - mlflow.sklearn.log_model(model, model_type) + self.debug(f'Code paths: {code_paths}', metadata) + + mlflow.sklearn.log_model(model, model_type, code_paths=code_paths or None) elif flavor == 'pyfunc': - code_path = self.resolve_code_paths(model_data.get('artifact_path'), metadata) + self.debug(f'Code paths: {code_paths}', metadata) - self.debug(f'Code path: {code_path}', metadata) - - model.store_model(artifact_path=model_type, code_path=code_path, to_disk=False) + model.store_model(artifact_path=model_type, code_path=code_paths, to_disk=False) self.debug('Model uploaded successfully', metadata) elif flavor == 'joblib': diff --git a/laborious/worker/worker.py b/laborious/worker/worker.py index 231cc8f..e585a2d 100644 --- a/laborious/worker/worker.py +++ b/laborious/worker/worker.py @@ -305,7 +305,14 @@ async def main(): # This will run the workers and wait for them to complete. # If an exception occurs in any of the worker handlers, it will be propagated here. await asyncio.gather(*handlers) - except BaseException as e: + except (SystemExit, KeyboardInterrupt) as e: + # Deliberate shutdown (Ctrl+C, a handler's own sys.exit) — not a failure to log as one. + # `finally` below calls sys.exit(exit_code) unconditionally, which supersedes this raise + # and becomes the exception that actually propagates — the raise is here so this branch + # still honors the reraise-on-SystemExit contract, not because it changes what exits. + exit_code = e.code if isinstance(e, SystemExit) and isinstance(e.code, int) else 0 + raise + except Exception as e: logger.custom_error(f'An unhandled exception occurred: {e}', metadata) exit_code = 1 finally: diff --git a/tests/laborious/utils/repository/test_model_repository.py b/tests/laborious/utils/repository/test_model_repository.py index 2014bdd..0311ef1 100644 --- a/tests/laborious/utils/repository/test_model_repository.py +++ b/tests/laborious/utils/repository/test_model_repository.py @@ -1131,7 +1131,9 @@ async def test_log_model_sklearn(mlflow, mlflow_repository): 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') + mlflow.sklearn.log_model.assert_called_once_with( + model_data['model'], 'prediction_model', code_paths=None + ) mlflow_repository.observe_lag.assert_called_once_with(ANY, metrics.MODEL_WRITE_LAG, ANY) mlflow_repository.emit_metric.assert_called_once_with( @@ -1139,6 +1141,22 @@ async def test_log_model_sklearn(mlflow, mlflow_repository): ) +@pytest.mark.asyncio +async def test_log_model_sklearn_forwards_code_paths(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, 'sklearn', 'prediction_model', metadata['metadata'] + ) + + mlflow.sklearn.log_model.assert_called_once_with( + model_data['model'], + 'prediction_model', + code_paths=[str(Path(code_dir) / 'probe_flavor_model.py')], + ) + + def test_resolve_code_paths_package_layout(tmp_path, mlflow_repository): code_dir = tmp_path / 'code' (code_dir / 'utils').mkdir(parents=True) @@ -2489,7 +2507,9 @@ async def test_log_model_sklearn_does_not_dump_joblib(joblib_mock, mlflow, mlflo model_data, 'sklearn', 'prediction_model', metadata['metadata'] ) - mlflow.sklearn.log_model.assert_called_once_with(model_data['model'], 'prediction_model') + mlflow.sklearn.log_model.assert_called_once_with( + model_data['model'], 'prediction_model', code_paths=None + ) joblib_mock.dump.assert_not_called() mlflow.log_artifact.assert_not_called() @@ -2617,7 +2637,7 @@ async def test_retrain_model_joblib_reaches_joblib_write_branch( 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') + mlflow.sklearn.log_model.assert_called_once_with(ANY, 'data_model', code_paths=None) def test_module_defines_no_joblib_fallback_routine(): diff --git a/tests/laborious/utils/repository/test_sklearn_flavor_retrain_e2e.py b/tests/laborious/utils/repository/test_sklearn_flavor_retrain_e2e.py new file mode 100644 index 0000000..f1b77fe --- /dev/null +++ b/tests/laborious/utils/repository/test_sklearn_flavor_retrain_e2e.py @@ -0,0 +1,104 @@ +""" +Integration test: a `sklearn`-flavor model with a custom class survives a retrain round trip. + +Unlike the rest of `test_model_repository.py`, this test does not mock the `mlflow` module — +`resolve_code_paths` forwarding `code_paths=` to `mlflow.sklearn.log_model` is only proven by +letting the real MLflow file-based `code_paths` validation/copy machinery run, and by really +dropping and reloading the custom class module afterwards. A fully mocked `mlflow.sklearn` (as +used elsewhere in this file, and in `e2e/test_retrain_joblib_flavor.py`) would never exercise +that machinery at all. + +The Model Registry itself is not exercised here — MLflow's registry requires a database-backed +store, not a bare `file:` tracking URI — so this drives the write (`log_model`) and read +(`client.download_artifacts` + `mlflow.sklearn.load_model`) by run id directly, the same +building blocks `download_model` composes for the `sklearn` flavor once a Production run is +resolved. +""" + +import importlib +import sys +from pathlib import Path +from unittest.mock import AsyncMock, MagicMock + +import mlflow +import pytest + +from laborious.utils.repository.model_repository import MLFlowRepository + +CODE_MODULE_NAME = 'sklearn_flavor_retrain_probe' + +CODE_MODULE_SOURCE = '''"""Custom estimator that exists only inside the artifact's code/ directory.""" + + +class SklearnFlavorProbePredictor: + def __init__(self, offset): + self.offset = offset + + def fit(self, X, y=None): + return self + + def predict(self, X): + return [row[0] + self.offset for row in X] +''' + + +def _write_code_module(code_dir: Path) -> Path: + code_dir.mkdir(parents=True) + (code_dir / f'{CODE_MODULE_NAME}.py').write_text(CODE_MODULE_SOURCE) + return code_dir + + +def _import_fresh_probe_model(code_dir: Path, offset: int): + """Instantiate the custom class with its module on `sys.path`, then undo both.""" + sys.path.insert(0, str(code_dir)) + try: + module = importlib.import_module(CODE_MODULE_NAME) + return module.SklearnFlavorProbePredictor(offset) + finally: + sys.path.remove(str(code_dir)) + sys.modules.pop(CODE_MODULE_NAME, None) + importlib.invalidate_caches() + + +@pytest.mark.asyncio +@pytest.mark.integration +async def test_retrain_sklearn_model_preserves_code_dir_and_reloads(tmp_path): + """A sklearn retrain with a custom `code/` class keeps it and can reload it afterwards.""" + tracking_uri = f'file://{tmp_path / "mlruns"}' + mlflow.set_tracking_uri(tracking_uri) + mlflow.set_experiment('sklearn_flavor_retrain_e2e') + + repository = MLFlowRepository( + host=tracking_uri, + username='test', + password='test', + logger=MagicMock(), + notification_handler=MagicMock(), + metrics_controller=AsyncMock(), + ) + repository.emit_metric = AsyncMock() + repository.observe_lag = AsyncMock() + + source_code_dir = _write_code_module(tmp_path / 'source_artifact' / 'code') + model = _import_fresh_probe_model(source_code_dir, offset=10) + model_data = {'model': model, 'artifact_path': str(tmp_path / 'source_artifact')} + + with mlflow.start_run() as run: + await repository.log_model(model_data, 'sklearn', 'prediction_model', {'metadata': {}}) + run_id = run.info.run_id + + written_code_dir = Path( + repository.client.download_artifacts(run_id, 'prediction_model/code') + ) + assert (written_code_dir / f'{CODE_MODULE_NAME}.py').is_file() + + # The class must resolve from the artifact's written code/, not from anything left in memory. + sys.modules.pop(CODE_MODULE_NAME, None) + importlib.invalidate_caches() + + local_model_dir = repository.client.download_artifacts(run_id, 'prediction_model') + repository.prepend_artifact_code_dir(local_model_dir, 'sklearn', {}) + + reloaded = mlflow.sklearn.load_model(local_model_dir) + + assert reloaded.predict([[1], [2]]) == [11, 12]