SAM0123-320: fix: evict stale artifact modules before loading a model

A long-lived worker kept the first `utils_slr` it ever imported in
`sys.modules`, so every later load unpickled against that stale class and
ignored the freshly downloaded `code/`. Courier Stratified retrain kept
failing on an old `StratifiedTransformerWrapper.fit` from a months-old
`/tmp/tmpvutfjrti` download.

`prepend_artifact_code_dir` now evicts cached modules shipped by the
artifact's `code/` that were imported from an artifact code directory,
and moves the directory to the front of `sys.path` when already present.
Modules installed in the image are never touched.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

laborious/utils/repository/model_repository.py

@@ -16,6 +16,7 @@ Capabilities:
 
 import ctypes
 import gc
+import importlib
 import sys
 import tempfile
 import threading
@@ -51,6 +52,10 @@ MODEL_PICKLE_NAME = 'model.pkl'
 
 INVALID_FLAVOR_MESSAGE = "Invalid flavor. Use 'sklearn' or 'pyfunc' or 'joblib'."
 
+# Every artifact `code` directory this process has put on `sys.path`. A cached module is evicted
+# only when it was imported from one of these, so modules installed in the image are never touched.
+ARTIFACT_CODE_DIRS: set[str] = set()
+
 
 def force_memory_release(logger: Logger, metadata: dict[str, Any] | None = None):
     """Attempt to release memory from the Python process.
@@ -451,6 +456,15 @@ class MLFlowRepository(SientiaMonitoring):
         The entry is never removed. A loaded model may import lazily, long after
         the load returned.
 
+        `sys.path` only matters on the first import of a module: once
+        `utils_slr.mlflow.pyfunc_wrappers` sits in `sys.modules`, unpickling
+        resolves the class from it and the freshly downloaded `code` is never
+        read. So before prepending, the modules a previous artifact's `code`
+        left in `sys.modules` are evicted (see `evict_artifact_modules`), and
+        the directory is moved to the front even when it is already on
+        `sys.path`. Models already loaded keep the classes they were built
+        with: eviction only drops the cache entry, not the module objects.
+
         Args:
             local_dir (str | None): Local directory of the downloaded model
                 artifact, e.g. the download of `<run>/artifacts/prediction_model`.
@@ -467,8 +481,15 @@ class MLFlowRepository(SientiaMonitoring):
         if not code_dir:
             return None
 
+        code_dir = path.abspath(code_dir)
+
+        ARTIFACT_CODE_DIRS.add(code_dir)
+        self.evict_artifact_modules(code_dir, flavor, metadata)
+
         if code_dir in sys.path:
-            self.debug(f'{flavor} flavor: {code_dir} already on sys.path', metadata)
+            sys.path.remove(code_dir)
+            sys.path.insert(0, code_dir)
+            self.debug(f'{flavor} flavor: {code_dir} already on sys.path, moved to front', metadata)
             return None
 
         sys.path.insert(0, code_dir)
@@ -476,6 +497,69 @@ class MLFlowRepository(SientiaMonitoring):
 
         return code_dir
 
+    def evict_artifact_modules(
+        self,
+        code_dir: str,
+        flavor: str,
+        metadata: dict[str, Any] | None = None,
+    ) -> list[str]:
+        """
+        Drop from `sys.modules` what an artifact `code` directory imported before.
+
+        A long-lived worker loads many artifacts that ship the same top-level
+        package (`utils_slr`, `utils`) at different versions. Without eviction,
+        the first one imported wins for the life of the process, and every later
+        load unpickles against that stale class.
+
+        Only modules whose top-level name is shipped by `code_dir` *and* whose
+        file lives under an artifact `code` directory (a previous one, or this
+        very one re-downloaded in place) are evicted. A module of the same name
+        installed in the image is left alone.
+
+        Args:
+            code_dir (str): Absolute `code` directory about to be prepended,
+                already registered in `ARTIFACT_CODE_DIRS`.
+            flavor (str): Flavor doing the load, for the log line only.
+            metadata (dict[str, Any] | None): Metadata used for structured
+                logging.
+
+        Returns:
+            list[str]: Names of the evicted modules.
+        """
+        shipped = {
+            path.splitext(entry)[0]
+            for entry in listdir(code_dir)
+            if entry != '__pycache__'
+            and (path.isdir(path.join(code_dir, entry)) or entry.endswith('.py'))
+        }
+
+        artifact_prefixes = tuple(path.join(directory, '') for directory in ARTIFACT_CODE_DIRS)
+
+        evicted = [
+            name
+            for name, module in list(sys.modules.items())
+            if name.partition('.')[0] in shipped
+            # `__path__` covers namespace packages, which have no `__file__`.
+            and any(
+                path.abspath(origin).startswith(artifact_prefixes)
+                for origin in [getattr(module, '__file__', None), *getattr(module, '__path__', [])]
+                if origin
+            )
+        ]
+
+        for name in evicted:
+            sys.modules.pop(name, None)
+
+        if evicted:
+            importlib.invalidate_caches()
+            self.info(
+                f'{flavor} flavor: evicted {len(evicted)} cached artifact module(s) '
+                f'before loading from {code_dir}: {", ".join(sorted(evicted))}',
+                metadata,
+            )
+
+        return evicted
+
     async def download_artifact_code_dir(
         self,
         model_name: str,

tests/laborious/utils/repository/test_model_repository.py

@@ -1,6 +1,7 @@
 import ast
 import importlib
 import sys
+import types
 from datetime import UTC, datetime
 from pathlib import Path
 from unittest.mock import ANY, AsyncMock, MagicMock, call, patch
@@ -2336,6 +2337,82 @@ def test_prepend_artifact_code_dir_does_not_duplicate(
     assert sys.path.count(code_dir) == 1
 
 
+def build_versioned_package_dir(artifact_dir, version):
+    """Add a `code/` directory shipping a `probe_pkg` package at a given version."""
+    package_dir = artifact_dir / 'code' / 'probe_pkg'
+    package_dir.mkdir(parents=True)
+    (package_dir / '__init__.py').write_text('')
+    (package_dir / 'wrappers.py').write_text(f'VERSION = {version!r}\n')
+    return str(artifact_dir / 'code')
+
+
+@pytest.fixture
+def clean_probe_pkg():
+    """Keep `probe_pkg` out of `sys.modules` before and after each test."""
+
+    def purge():
+        for name in [n for n in sys.modules if n.partition('.')[0] == 'probe_pkg']:
+            del sys.modules[name]
+        importlib.invalidate_caches()
+
+    purge()
+    yield
+    purge()
+
+
+def test_prepend_artifact_code_dir_evicts_module_cached_from_previous_artifact(
+    mlflow_repository, tmp_path, preserved_sys_path, clean_probe_pkg
+):
+    """A later artifact's `code` wins over what an earlier artifact left in `sys.modules`."""
+    old_artifact, new_artifact = tmp_path / 'old', tmp_path / 'new'
+    build_versioned_package_dir(old_artifact, 'old')
+    build_versioned_package_dir(new_artifact, 'new')
+
+    mlflow_repository.prepend_artifact_code_dir(str(old_artifact), 'pyfunc')
+    stale = importlib.import_module('probe_pkg.wrappers')
+    assert stale.VERSION == 'old'
+
+    mlflow_repository.prepend_artifact_code_dir(str(new_artifact), 'pyfunc')
+
+    assert importlib.import_module('probe_pkg.wrappers').VERSION == 'new'
+    # Already-loaded objects keep the module they were built against.
+    assert stale.VERSION == 'old'
+
+
+def test_prepend_artifact_code_dir_reloads_code_dir_re_downloaded_in_place(
+    mlflow_repository, tmp_path, preserved_sys_path, clean_probe_pkg
+):
+    """The fixed `joblib` download dir is reused: re-prepending it must re-read the files."""
+    code_dir = build_versioned_package_dir(tmp_path, 'old')
+
+    mlflow_repository.prepend_artifact_code_dir(str(tmp_path), 'joblib')
+    assert importlib.import_module('probe_pkg.wrappers').VERSION == 'old'
+
+    # A different size, so the `.pyc` written a moment ago is not taken as current.
+    (tmp_path / 'code' / 'probe_pkg' / 'wrappers.py').write_text("VERSION = 'renewed'\n")
+    sys.path.insert(0, str(tmp_path / 'elsewhere'))
+
+    mlflow_repository.prepend_artifact_code_dir(str(tmp_path), 'joblib')
+
+    assert sys.path[0] == code_dir
+    assert importlib.import_module('probe_pkg.wrappers').VERSION == 'renewed'
+
+
+def test_evict_artifact_modules_leaves_modules_not_from_an_artifact_alone(
+    mlflow_repository, tmp_path, preserved_sys_path, clean_probe_pkg
+):
+    """A same-named module installed in the image is never evicted."""
+    code_dir = build_versioned_package_dir(tmp_path / 'artifact', 'new')
+    installed = types.ModuleType('probe_pkg')
+    installed.__file__ = str(tmp_path / 'site-packages' / 'probe_pkg' / '__init__.py')
+    sys.modules['probe_pkg'] = installed
+
+    evicted = mlflow_repository.evict_artifact_modules(code_dir, 'pyfunc')
+
+    assert evicted == []
+    assert sys.modules['probe_pkg'] is installed
+
+
 def stub_remote_code_dir(mlflow_repository, tmp_path, artifact_name, exists=True):
     """Point the client at a run whose artifact does — or does not — carry `code`.