Repository: caikit/caikit
Pull request #588 (2023-11-27 23:44:25)

Don't lazy load models that are actively being written
<!--  Thanks for sending a pull request!  Here are some tips for you:
1. Make sure to read the Contributing Guide before submitting your PR: https://github.com/caikit/caikit/blob/main/CONTRIBUTING.md
2. If this PR closes another issue, add 'closes #<issue number>' somewhere in the PR summary. GitHub will automatically close that issue when this PR gets merged. Alternatively, adding 'refs #<issue number>' will not close the issue, but help provide the reviewer more context.-->

**What this PR does / why we need it**:

This fixes an issue where the lazy load loop can load a large model that is in the progress of being written, where the module itself may gracefully handle partial model artifacts.

This is done by running two back-to-back size checks on the model directory before loading, and only loading if the sizes match. This could have performance implications for models with a very large number of file artifacts. The amount of time that needs to elapse between these two checks to correctly catch models mid-write could vary depending on the specifics about how the bytes are making it to disk.

Possibly related to #467 

**Special notes for your reviewer**:

This is a very simple fix, it may be _too_ simple

**If applicable**:
- [ ] this PR contains documentation
- [ ] this PR contains unit tests
- [ ] this PR has been tested for backwards compatibility



Patch:
diff --git a/caikit/config/config.yml b/caikit/config/config.yml
index af3f1dd0..4be5c480 100644
--- a/caikit/config/config.yml
+++ b/caikit/config/config.yml
@@ -103,6 +103,10 @@ runtime:
     # that fail due to partially uploaded model artifacts when the load is
     # initiated.
     lazy_load_retries: 2
+    # Amount of time to watch for file updates to detect if a model is being written.
+    # Models that are being written will not be prematurely loaded
+    # Use zero or a negative value or None to disable
+    lazy_load_write_detection_period_seconds: 0.05
 
     # Number of threads to make available for load jobs (null => all)
     load_threads: null
diff --git a/caikit/runtime/model_management/model_manager.py b/caikit/runtime/model_management/model_manager.py
index ba94cda9..45f6da5c 100644
--- a/caikit/runtime/model_management/model_manager.py
+++ b/caikit/runtime/model_management/model_manager.py
@@ -14,11 +14,13 @@
 # Standard
 from collections import Counter as DictCounter
 from functools import partial
+from pathlib import Path
 from typing import Dict, Optional
 import atexit
 import gc
 import os
 import threading
+import time
 
 # Third Party
 from grpc import StatusCode
@@ -131,6 +133,16 @@ class ModelManager:  # pylint: disable=too-many-instance-attributes
             and self._lazy_load_local_models
             and self._lazy_load_poll_period_seconds
         )
+        self._lazy_load_write_detection_period_seconds = (
+            runtime_cfg.lazy_load_write_detection_period_seconds
+        )
+        error.type_check(
+            "<RUN58138047E>",
+            int,
+            float,
+            allow_none=True,
+            lazy_load_write_detection_period_seconds=self._lazy_load_write_detection_period_seconds,
+        )
         if self._enable_lazy_load_poll:
             atexit.register(self.shut_down)
 
@@ -452,6 +464,11 @@ class ModelManager:  # pylint: disable=too-many-instance-attributes
         # Load new models
         for model_id in new_models:
             model_path = os.path.join(self._local_models_dir, model_id)
+
+            if self._model_write_in_progress(model_path):
+                log.debug("Model %s is still being written", model_id)
+                continue
+
             self.load_model(
                 model_id,
                 model_path,
@@ -490,6 +507,35 @@ class ModelManager:  # pylint: disable=too-many-instance-attributes
                         exc_info=True,
                     )
 
+    def _model_write_in_progress(self, model_dir: str) -> bool:
+        """Returns true if model_dir is currently being written to. Uses the
+        runtime.lazy_load_write_detection_period_seconds configuration to sleep between
+        consecutive size checks of the directory.
+
+        Always returns false if runtime.lazy_load_write_detection_period_seconds is zero,
+        negative, or None.
+        """
+        if (
+            self._lazy_load_write_detection_period_seconds is None
+            or self._lazy_load_write_detection_period_seconds <= 0
+        ):
+            return False
+
+        # Get the current directory size
+        size = self._get_total_disk_size(model_dir)
+        # Sleep a bit to wait out another write
+        time.sleep(self._lazy_load_write_detection_period_seconds)
+        # Get the size again. If it has changed, then a write is currently  in progress
+        return self._get_total_disk_size(model_dir) != size
+
+    @staticmethod
+    def _get_total_disk_size(model_dir: str) -> int:
+        """Returns the sum of st_size of all files contained within the directory structure rooted
+        at model_dir.
+        """
+        dir_path = Path(model_dir)
+        return sum([f.stat().st_size for f in dir_path.rglob("*") if f.is_file()])
+
     def __report_total_model_size_metric(self):
         # Just a happy little lock to ensure that with concurrent loading and unloading,
         # the last metric reported will be correct.


Test patch:
diff --git a/tests/runtime/model_management/test_model_manager.py b/tests/runtime/model_management/test_model_manager.py
index 265b76a8..9688d41a 100644
--- a/tests/runtime/model_management/test_model_manager.py
+++ b/tests/runtime/model_management/test_model_manager.py
@@ -485,6 +485,63 @@ def test_model_manager_disk_caching_periodic_sync(good_model_path):
             assert mgr_one_unloaded and mgr_two_unloaded
 
 
+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
+    """
+    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]
+
+            # Start with a valid model
+            model_name = 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)
+
+            # Then kick off a thread that will start writing a large file inside this model dir.
+            # This simulates uploading a large model artifact
+            def write_big_file(path: str, stop_event: threading.Event):
+                big_file = os.path.join(path, "big_model_artifact.txt")
+                with open(big_file, "w") as bf:
+                    while not stop_event.is_set():
+                        bf.write("This is a big file\n" * 1000)
+
+            stop_write_event = threading.Event()
+            writer_thread = threading.Thread(
+                target=write_big_file, args=(model_cache_path, stop_write_event)
+            )
+            writer_thread.start()
+
+            try:
+                # 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
+
+                # Stop the model writing thread (Finish the model upload)
+                stop_write_event.set()
+                writer_thread.join()
+
+                # Re-trigger the sync and make sure the model is loaded this time
+                manager.sync_local_models(wait=True)
+                assert model_name in manager.loaded_models
+
+            finally:
+                stop_write_event.set()
+                writer_thread.join()
+
+
 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.