Repository: caikit/caikit
Pull request #578 (2023-11-13 20:14:27)

Nested model names
**What this PR does / why we need it**:

This PR makes it possible to load nested models inside of `local_models_dir` without them being immediately unloaded by the periodic sync logic.

**Special notes for your reviewer**:

This _does_ change a log code from `W` to `D` because we're changing the semantics of finding non-model directories in `local_models_dir` from being "bad" to being expected. It does NOT silence `COR80419785E` which is raised [here](https://github.com/caikit/caikit/blob/main/caikit/core/model_manager.py#L280). From the perspective of `load`, this _is_ an error, but when used by `runtime` it is a valid outcome of attempting to load something that isn't actually a model.

The fix I attempted for this was to [filter `disk_models`](https://github.com/caikit/caikit/blob/main/caikit/runtime/model_management/model_manager.py#L431) to only those containing a `config.yml` file, but that breaks files with `.zip` suffixes. I could certainly adjust the filter to allow `.zip` files, but this gets into reimplementing the logic of the core `load` function and could make for difficult maintenance if more archive formats are added.

**If applicable**:
- [ ] this PR contains documentation
- [ ] this PR contains unit tests
- [ ] this PR has been tested for backwards compatibility



Patch:
diff --git a/caikit/runtime/model_management/model_manager.py b/caikit/runtime/model_management/model_manager.py
index c3d0c46b..77b0872b 100644
--- a/caikit/runtime/model_management/model_manager.py
+++ b/caikit/runtime/model_management/model_manager.py
@@ -443,10 +443,8 @@ class ModelManager:  # pylint: disable=too-many-instance-attributes
         unload_models = [
             model_id
             for model_id, loaded_model in self.loaded_models.items()
-            if model_id not in disk_models
-            and loaded_model.path().startswith(
-                self._local_models_dir,
-            )
+            if loaded_model.path().startswith(self._local_models_dir)
+            and not os.path.exists(loaded_model.path())
         ]
         log.debug("Unloaded local models: %s", unload_models)
 
@@ -483,8 +481,8 @@ class ModelManager:  # pylint: disable=too-many-instance-attributes
                 try:
                     loaded_model.wait()
                 except CaikitRuntimeException as err:
-                    log.warning(
-                        "<RUN56627485W>",
+                    log.debug(
+                        "<RUN56627485D>",
                         "Failed to load model %s: %s",
                         model_id,
                         repr(err),


Test patch:
diff --git a/tests/runtime/model_management/test_model_manager.py b/tests/runtime/model_management/test_model_manager.py
index 43096725..341daf7e 100644
--- a/tests/runtime/model_management/test_model_manager.py
+++ b/tests/runtime/model_management/test_model_manager.py
@@ -485,6 +485,46 @@ def test_model_manager_disk_caching_periodic_sync(good_model_path):
             assert mgr_one_unloaded and mgr_two_unloaded
 
 
+def test_nested_local_model_load_unload(good_model_path):
+    """Test that a model can be loaded in a subdirectory of the local_models_dir
+    and that the periodic sync does not unload the model.
+    """
+    with TemporaryDirectory() as cache_dir:
+        with non_singleton_model_managers(
+            1,
+            {
+                "runtime": {
+                    "local_models_dir": cache_dir,
+                    "lazy_load_local_models": True,
+                    "lazy_load_poll_period_seconds": 0,
+                },
+            },
+            "merge",
+        ) as managers:
+            manager = managers[0]
+
+            # Copy the model into a nested model directory
+            model_name = os.path.join("parent", os.path.basename(good_model_path))
+            model_cache_path = os.path.join(cache_dir, model_name)
+            assert not os.path.exists(model_cache_path)
+            shutil.copytree(good_model_path, model_cache_path)
+
+            # Trigger the periodic sync and make sure the model is NOT loaded
+            assert model_name not in manager.loaded_models
+            manager.sync_local_models(wait=True)
+            assert model_name not in manager.loaded_models
+
+            # Explicitly ask to load the nested model name to trigger the lazy
+            # load
+            model = manager.retrieve_model(model_name)
+            assert model
+            assert model_name in manager.loaded_models
+
+            # Re-trigger the sync and make sure the model does not get unloaded
+            manager.sync_local_models(wait=True)
+            assert model_name in manager.loaded_models
+
+
 def test_load_local_model_deleted_dir():
     """Make sure losing the local_models_dir out from under a running manager
     doesn't kill the whole thing