diff --git a/laborious/utils/repository/model_repository.py b/laborious/utils/repository/model_repository.py index 285957a..bddfc60 100644 --- a/laborious/utils/repository/model_repository.py +++ b/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 `/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, diff --git a/tests/laborious/utils/repository/test_model_repository.py b/tests/laborious/utils/repository/test_model_repository.py index 454605a..296ba6c 100644 --- a/tests/laborious/utils/repository/test_model_repository.py +++ b/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`.