Repository: caikit/caikit (python)

Issue:
Allow ephemeral models to lazy load in runtime
## Is your feature request related to a problem? Please describe.

The current implementation of lazy loading in `caikit.runtime.model_management.model_manager` is centered around `local_models_dir`. This makes sense for physical models, but with the presence of the `ModelFinder` abstraction, we also want to support "auto-finders" that are able to infer a model's configuration without needing the model on disk physically.

An "auto-finder" may be configured to receive the `model_path` and find a model based on it in any arbitrary way. Given this, an auto-finder is free to play tricks with the on-disk path such as `os.path.basename` to strip off path prefixes, but this also interferes with the lazy-load sync's unloading logic which assumes that a model whose path begins with `local_models_dir` which is not present on disk has been deleted and therefore should be unloaded.

## Describe the solution you'd like

To solve this, we want the runtime `ModelManager` to explicitly permit lazy loading of non-disk models by using the `model_id` as the `model_path` for models where `os.path.join(local_models_dir, model_id)` does not exist.

## Describe alternatives you've considered

We could solve this by requiring auto-finders to be smart enough to handle this themselves, but this places several limitations on what finders can do and decentralizes the logic for this kind of ephemeral model.

## Additional context

This is currently limiting the use of `TGISAutoFinder` ([here](https://github.com/caikit/caikit-nlp/blob/main/caikit_nlp/model_management/tgis_auto_finder.py)) with `runtime`


Pull request:
LazyLoadEphemeralModels
**What this PR does / why we need it**:

Closes #429 

**Special notes for your reviewer**:

**If applicable**:
- [X] this PR contains unit tests


Patch:
diff --git a/caikit/runtime/model_management/model_manager.py b/caikit/runtime/model_management/model_manager.py
index 468859ef..1b01b3ba 100644
--- a/caikit/runtime/model_management/model_manager.py
+++ b/caikit/runtime/model_management/model_manager.py
@@ -381,6 +381,13 @@ class ModelManager:  # pylint: disable=too-many-instance-attributes
             log.debug2(
                 "Lazy loading local model %s from %s", model_id, local_model_path
             )
+            # If the model is not present on disk, attempt to lazy load it
+            # anyway using the model_id as the "model_path". This allows
+            # auto-finders that can infer the model's config to load based on
+            # the ID.
+            if not os.path.exists(local_model_path):
+                log.debug2("Attempting to load ephemeral model %s", model_id)
+                local_model_path = model_id
             self.load_model(
                 model_id=model_id,
                 local_model_path=local_model_path,
@@ -392,7 +399,7 @@ class ModelManager:  # pylint: disable=too-many-instance-attributes
 
         # If still not loaded, there's nothing to find, so raise NOT_FOUND
         if not model_loaded:
-            msg = "Model '%s' not loaded" % model_id
+            msg = f"Model '{model_id}' not loaded"
             log.debug(
                 {"log_code": "<RUN61105243D>", "message": msg, "model_id": model_id}
             )


Test patch:
diff --git a/tests/runtime/model_management/test_model_manager.py b/tests/runtime/model_management/test_model_manager.py
index eb325e56..43096725 100644
--- a/tests/runtime/model_management/test_model_manager.py
+++ b/tests/runtime/model_management/test_model_manager.py
@@ -541,8 +541,10 @@ def test_load_local_model_deleted_dir():
             assert not manager.loaded_models
 
 
-def test_load_local_model_not_on_disk():
-    """Make sure bad models in local_models_dir at boot don't cause exceptions"""
+def test_lazy_load_ephemeral_model():
+    """Make sure an ephemeral model (not on disk) can be lazy loaded if the
+    right finder configuration is present to load it without hitting disk.
+    """
     with TemporaryDirectory() as cache_dir:
         with non_singleton_model_managers(
             1,
@@ -557,9 +559,13 @@ def test_load_local_model_not_on_disk():
         ) as managers:
             manager = managers[0]
             model_id = random_test_id()
-            manager.load_model(model_id, "some_non_disk_model", "test")
             model = manager.retrieve_model(model_id)
             assert model
+            assert cache_dir not in manager.loaded_models[model_id].path()
+
+            # Make sure the model does not get unloaded on sync
+            manager._local_models_dir_sync(wait=True)
+            assert model_id in manager.loaded_models
 
 
 # ****************************** Unit Tests ****************************** #