diff --git a/laborious/utils/repository/model_repository.py b/laborious/utils/repository/model_repository.py index 285957a..e4e70e1 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,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 diff --git a/requirements-light.txt b/requirements-light.txt index 6220a72..fe03725 100644 --- a/requirements-light.txt +++ b/requirements-light.txt @@ -1,6 +1,6 @@ temporalio psycopg2-binary -sqlalchemy +sqlalchemy<2.1 asyncua==1.0.6 redis sientia_do>=1.12.2 diff --git a/requirements-local.txt b/requirements-local.txt index 4588a28..a8c4365 100644 --- a/requirements-local.txt +++ b/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 diff --git a/requirements.txt b/requirements.txt index 6564587..7f17536 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,6 +1,6 @@ temporalio psycopg2-binary -sqlalchemy +sqlalchemy<2.1 asyncua==1.0.6 redis sientia_do>=1.12.2 diff --git a/tests/laborious/utils/repository/test_model_repository.py b/tests/laborious/utils/repository/test_model_repository.py index 454605a..cf5e0be 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 @@ -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`.