QTZPOC-20

QTZPOC-20: Enable code_paths for sklearn models; improve worker shutdown handling

The MLFlowRepository now resolves and forwards code_paths to mlflow.sklearn.log_model, ensuring custom code dependencies for sklearn models are preserved and correctly reloaded. A new integration test validates this end-to-end functionality. Additionally, worker exception handling is refined to distinguish between deliberate shutdowns (SystemExit, KeyboardInterrupt) and unhandled exceptions, preventing unnecessary error logging for expected exits.

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 `<artifact>/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':

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:

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():

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]