Repository: caikit/caikit (python)

Issue:
Runtime lazy load only mode
## Is your feature request related to a problem? Please describe.

When managing large numbers of models, it can be very expensive to perform the proactive loading in [sync_local_models](https://github.com/caikit/caikit/blob/main/caikit/runtime/model_management/model_manager.py#L263). This can result in slow boots and starved server thread pools.

The request here is to add a feature to enable the periodic sync to only handle deletion and not perform proactive loading. This would defer all loading to lazy-loading on inference while maintaining the ability to unload models that have been removed from `local_models_dir`.

## Describe the solution you'd like

* Add a new config argument `runtime.proactive_load_local_models`
* In [_local_models_dir_sync](https://github.com/caikit/caikit/blob/main/caikit/runtime/model_management/model_manager.py#L490), only perform the loading if `proactive_load_local_models` is `true`

## Describe alternatives you've considered

This issue is related to [an issue in `caikit-tgis-runtime`](https://github.com/caikit/caikit-tgis-backend/issues/51) that is coming up due to massive repetition of copying prompt artifacts from `local_models_dir` to TGIS' prompt cache. We've considered much more sophisticated solutions and may implement them in the future, but this will act as a simple mitigation to avoid ongoing problems.

## Additional context

The main downside of this proposal is that the `info` service loses the ability to guarantee eventual consistency between replicas of a given server.


Pull request:
Proactive load toggle 661
**What this PR does / why we need it**:

Closes #661 

This PR adds the `runtime.load_new_local_models` config option which can disable loading new models during synchronization with `local_models_dir`.

**Special notes for your reviewer**:

When loading new local models is disabled, the model info endpoints are not guaranteed to be consistent across replicas of `caikit.runtime` running behind a loadbalancer (e.g. k8s `Deployment` / `Service`) since loading only happens on inference calls which are distributed across the replicas.

**If applicable**:
- [x] this PR contains documentation (comment in `config.yml`)
- [x] this PR contains unit tests
- [x] this PR has been tested for backwards compatibility


Patch:
diff --git a/caikit/config/config.yml b/caikit/config/config.yml
index 434e1f03..0659949c 100644
--- a/caikit/config/config.yml
+++ b/caikit/config/config.yml
@@ -97,6 +97,11 @@ runtime:
     # the server
     wait_for_initial_model_loads: true
 
+    # If true, on each local_models_dir sync (include the initial one), any new
+    # models will start loading. If false, new models will be ignored and will
+    # only lazy load on inference.
+    load_new_local_models: true
+
     # If enabled, the models in local_models_dir will be periodically sync'ed
     # with the in-memory models. New models that are not in-memory that are
     # found in local_models_dir will be loaded and existing models that were
diff --git a/caikit/runtime/model_management/model_manager.py b/caikit/runtime/model_management/model_manager.py
index a29b20f7..0d9cedc2 100644
--- a/caikit/runtime/model_management/model_manager.py
+++ b/caikit/runtime/model_management/model_manager.py
@@ -164,12 +164,15 @@ class ModelManager:  # pylint: disable=too-many-instance-attributes
         # Do the initial local models load
         if self._local_models_dir:
             wait = runtime_cfg.wait_for_initial_model_loads
+            load = runtime_cfg.load_new_local_models
             log.info(
                 "<RUN44739400I>",
-                "Loading local models into Caikit Runtime. Wait: %s",
+                "Initializing local_models_dir %s. Wait: %s. Load: %s",
+                self._local_models_dir,
                 wait,
+                load,
             )
-            self.sync_local_models(wait=wait)
+            self.sync_local_models(wait=wait, load=load)
 
     def shut_down(self):
         """Shut down cache purging"""
@@ -260,7 +263,7 @@ class ModelManager:  # pylint: disable=too-many-instance-attributes
             # Return the loaded model handle
             return model
 
-    def sync_local_models(self, wait: bool = False):
+    def sync_local_models(self, wait: bool = False, load: bool = True):
         """Sync in-memory models with models in the configured local_model_dir
 
         New models will be loaded and models previously loaded from local will
@@ -268,9 +271,10 @@ class ModelManager:  # pylint: disable=too-many-instance-attributes
 
         Args:
             wait (bool): After starting all loads, wait for them to complete
+            load (bool): Perform loading during sync
         """
         try:
-            self._local_models_dir_sync(wait)
+            self._local_models_dir_sync(wait, load)
         except StopIteration:
             log.warning(
                 "<RUN56519883W>",
@@ -304,7 +308,9 @@ class ModelManager:  # pylint: disable=too-many-instance-attributes
                 [thread.name for thread in threading.enumerate()],
             )
             self._lazy_sync_timer = threading.Timer(
-                self._lazy_load_poll_period_seconds, self.sync_local_models
+                self._lazy_load_poll_period_seconds,
+                self.sync_local_models,
+                kwargs={"load": load},
             )
             self._lazy_sync_timer.daemon = True
             self._lazy_sync_timer.start()
@@ -446,7 +452,7 @@ class ModelManager:  # pylint: disable=too-many-instance-attributes
 
     ## Implementation Details ##
 
-    def _local_models_dir_sync(self, wait: bool = False):
+    def _local_models_dir_sync(self, wait: bool = False, load: bool = True):
         """This function implements the mechanics of synchronizing the
         local_models_dir and the in-memory loaded_models map. It may raise and
         therefore errors should be handled by the wrapper function.
@@ -471,10 +477,16 @@ class ModelManager:  # pylint: disable=too-many-instance-attributes
         log.debug3("Currently loaded models: %s", list(self.loaded_models.keys()))
 
         # Find all models that aren't currently loaded
-        new_models = [
-            model_id for model_id in disk_models if model_id not in self.loaded_models
-        ]
-        log.debug("New local models: %s", new_models)
+        if load:
+            new_models = [
+                model_id
+                for model_id in disk_models
+                if model_id not in self.loaded_models
+            ]
+            log.debug("New local models: %s", new_models)
+        else:
+            log.debug("Skipping new model loading")
+            new_models = []
 
         # Find all models that are currently loaded from the local models dir
         # that no longer exist


Test patch:
diff --git a/tests/runtime/model_management/test_model_manager.py b/tests/runtime/model_management/test_model_manager.py
index 27825da8..022e4e32 100644
--- a/tests/runtime/model_management/test_model_manager.py
+++ b/tests/runtime/model_management/test_model_manager.py
@@ -505,6 +505,55 @@ def test_model_manager_disk_caching_periodic_sync(good_model_path):
             assert mgr_one_unloaded and mgr_two_unloaded
 
 
+def test_periodic_sync_without_loading(good_model_path):
+    """Test that periodic synchronization of local_models_dir can proceed
+    without loading new models found there (unload only with lazy loading)
+    """
+    purge_period = 0.001
+    with TemporaryDirectory() as cache_dir:
+        # Copy the good model to the cache dir before starting the manager
+        model_id = random_test_id()
+        model_cache_path = os.path.join(cache_dir, model_id)
+        shutil.copytree(good_model_path, model_cache_path)
+
+        # Start the manager without loading new local models
+        with non_singleton_model_managers(
+            1,
+            {
+                "runtime": {
+                    "load_new_local_models": False,
+                    "local_models_dir": cache_dir,
+                    "lazy_load_local_models": True,
+                    "lazy_load_poll_period_seconds": purge_period,
+                    # NOTE: There won't be any initial model loads, but this
+                    #   ensures that if there were, they would happen
+                    #   synchronously during __init__
+                    "wait_for_initial_model_loads": True,
+                },
+            },
+            "merge",
+        ) as managers:
+            manager = managers[0]
+
+            # The model doesn't load at boot
+            assert model_id not in manager.loaded_models
+
+            # Wait for the purge period to run and make sure it's still not
+            # loaded
+            manager._lazy_sync_timer.join()
+            assert model_id not in manager.loaded_models
+
+            # Explicitly retrieve the model and make sure it _does_ lazy load
+            model = manager.retrieve_model(model_id)
+            assert model
+            assert model_id in manager.loaded_models
+
+            # Remove the file from local_models_dir and make sure it gets purged
+            shutil.rmtree(model_cache_path)
+            manager._lazy_sync_timer.join()
+            assert model_id not in manager.loaded_models
+
+
 def test_lazy_load_of_large_model(good_model_path):
     """Test that a large model that is actively being written to disk is not incorrectly loaded
     too soon by the lazy loading poll