Merge pull request #57 from Aignosi/fix/SAM0123-320-evict-stale-artifact-modules

SAM0123 320 evict stale artifact modules

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,70 @@ 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)
+
+        # Snapshot: another thread, or a lazy module's `__getattr__`, may import mid-scan.
+        evicted = [
+            name
+            for name, module in sys.modules.copy().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,
@@ -1257,7 +1342,10 @@ class MLFlowRepository(SientiaMonitoring):
 
         if not isinstance(treated_data_candidate, pd.DataFrame):
             data_model = treated_data_candidate
-            treated_data = data_model.predict(data)
+            if transform_flavor == 'pyfunc':
+                treated_data = data_model.predict({}, data)
+            else:
+                treated_data = data_model.predict(data)
         else:
             treated_data = treated_data_candidate
 

requirements-light.txt

@@ -1,6 +1,6 @@
 temporalio
 psycopg2-binary
-sqlalchemy
+sqlalchemy<2.1
 asyncua==1.0.6
 redis
 sientia_do>=1.12.2

requirements-local.txt

@@ -1,6 +1,6 @@
 temporalio
 psycopg2-binary
-sqlalchemy
+sqlalchemy<2.1
 asyncua==1.0.6
 redis
 git+ssh://git@github.com/Aignosi/sientia-dataops-library.git@1.12.2

requirements.txt

@@ -1,6 +1,6 @@
 temporalio
 psycopg2-binary
-sqlalchemy
+sqlalchemy<2.1
 asyncua==1.0.6
 redis
 sientia_do>=1.12.2

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
@@ -1042,6 +1043,38 @@ async def test_fit_models_not_df_target_name_none_and_not_in_model(
     }
 
 
+@patch('laborious.utils.repository.model_repository.pd.merge')
+@patch('laborious.utils.repository.model_repository.isinstance')
+@pytest.mark.asyncio
+async def test_fit_models_pyfunc_transform_predicts_with_wrapper_signature(
+    isinstance_mock, pd_merge, mlflow_repository
+):
+    """A pyfunc data model refit returns the wrapper, whose `predict` takes `(context, model_input)`."""
+    isinstance_mock.return_value = False
+
+    data_model = MagicMock()
+    mlflow_repository.download_model = AsyncMock(
+        side_effect=[(data_model, 'artifact_path'), (MagicMock(), 'artifact_path')],
+    )
+    mlflow_repository.detect_and_parse_datetime_index = MagicMock()
+    mlflow_repository.get_prediction_data = MagicMock(return_value=DataFrame())
+
+    data = MagicMock()
+
+    await mlflow_repository.fit_models(
+        'model_name',
+        data,
+        'latest_production_id',
+        metadata['metadata'],
+        'pyfunc',
+        False,
+        'pyfunc',
+        'target',
+    )
+
+    data_model.fit.return_value.predict.assert_called_once_with({}, data)
+
+
 @patch('laborious.utils.repository.model_repository.pd.merge')
 @patch('laborious.utils.repository.model_repository.isinstance')
 @pytest.mark.asyncio
@@ -2336,6 +2369,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`.