Repository: caikit/caikit (python)

Issue:
Retry failed local model loads
## Describe the bug

With the periodic load OR with an explicitly requested load based on an inference call, we can easily hit the case where the artifacts themselves have not yet finished writing to the shared disk. When this happens, a thread is spawned to load the model which will eventually raise for one reason or another (corrupted file, missing file, etc...). These exceptions won't tank the server, and the model future will be removed from the loaded_models dict, but any inference calls that are in flight will get 500s. Furthermore, this can happen really easily in tests!


Pull request:
Retry load failures 404
**What this PR does / why we need it**:

Closes #404 

This PR adds retry support for loading models from `local_models_dir`. This is needed to handle races with model artifact copying since an artifact will not be atomically copied to the mounted directory. If the periodic sync or a user inference request is triggered when a model is partially uploaded, the `load` function will begin, but will fail due to corrupted resources. Prior to this PR, this situation was handled by simply removing the model from the set of loaded models and letting the lazy loading restart on the next stimulus (period or inference request). The problem with this approach is that any requests received during the beginning of the load and the eventual failure/unload would receive 500 responses due to the failed load. With this PR, those requests will continue to wait until either the load succeeds or all retries are exhausted.

Patch:
diff --git a/caikit/config/config.yml b/caikit/config/config.yml
index 804b6381..c7954cf6 100644
--- a/caikit/config/config.yml
+++ b/caikit/config/config.yml
@@ -83,6 +83,10 @@ runtime:
     # unloaded.
     lazy_load_local_models: false
     lazy_load_poll_period_seconds: 10
+    # Number of times to retry on lazy load failure. This mitigates lazy loads
+    # that fail due to partially uploaded model artifacts when the load is
+    # initiated.
+    lazy_load_retries: 2
 
     # Number of threads to make available for load jobs (null => all)
     load_threads: null
diff --git a/caikit/runtime/model_management/loaded_model.py b/caikit/runtime/model_management/loaded_model.py
index 0b318a1b..b0affe6c 100644
--- a/caikit/runtime/model_management/loaded_model.py
+++ b/caikit/runtime/model_management/loaded_model.py
@@ -40,8 +40,10 @@ try:
 except TypeError:  # pragma: no cover
     CaikitModelFuture = Future
 
+CaikitModelFutureFactory = Callable[[], CaikitModelFuture]
 
-class LoadedModel:
+
+class LoadedModel:  # pylint: disable=too-many-instance-attributes
     __doc__ = __doc__
 
     class Builder:
@@ -52,10 +54,12 @@ class LoadedModel:
         def __init__(self):
             self._model_to_build = LoadedModel()
 
-        def model_future(
-            self, caikit_model_future: CaikitModelFuture
+        def model_future_factory(
+            self, caikit_model_future_factory: CaikitModelFutureFactory
         ) -> "LoadedModel.Builder":
-            self._model_to_build._caikit_model_future = caikit_model_future
+            self._model_to_build._caikit_model_future_factory = (
+                caikit_model_future_factory
+            )
             return self
 
         def fail_callback(self, callback: Callable) -> "LoadedModel.Builder":
@@ -74,10 +78,14 @@ class LoadedModel:
             self._model_to_build._model_id = model_id
             return self
 
+        def retries(self, retries: int) -> "LoadedModel.Builder":
+            self._model_to_build._retries = retries
+            return self
+
         def build(self) -> "LoadedModel":
             error.value_check(
                 "<RUN12786023E>",
-                self._model_to_build._caikit_model_future
+                self._model_to_build._caikit_model_future_factory
                 and self._model_to_build._model_id
                 and self._model_to_build._model_type,
                 "Cannot build LoadedModel with incomplete required fields."
@@ -86,16 +94,21 @@ class LoadedModel:
                 self._model_to_build._model_id,
                 self._model_to_build._model_type,
             )
+            self._model_to_build._caikit_model_future = (
+                self._model_to_build._caikit_model_future_factory()
+            )
             return self._model_to_build
 
     def __init__(self):
         # Use the builder ^^
+        self._caikit_model_future_factory: Optional[CaikitModelFutureFactory] = None
         self._caikit_model_future: Optional[CaikitModelFuture] = None
         self._model: Optional[ModuleBase] = None
         self._fail_callback: Optional[Callable] = None
         self._model_id: str = ""
         self._model_path: str = ""
         self._model_type: str = ""
+        self._retries: int = 0
         self._size: Optional[int] = None
 
     def id(self) -> str:
@@ -110,9 +123,25 @@ class LoadedModel:
             try:
                 self._model = self._caikit_model_future.result()
             except CaikitRuntimeException:
-                if self._fail_callback:
-                    self._fail_callback()
-                raise
+                if self._retries:
+                    self._retries -= 1
+                    log.debug(
+                        "Failed to load %s from %s. Retrying with %d retries left",
+                        self.id,
+                        self.path,
+                        self._retries,
+                    )
+                    self._caikit_model_future = self._caikit_model_future_factory()
+                    # Try waiting again with a fresh load future. This may open
+                    # a recursive retry. Once all retries are exhausted, if the
+                    # load still fails, the deepest call will invoke the fail
+                    # callback and the exception will percolate up from there to
+                    # here where it will be raised to the external waiter.
+                    self.wait()
+                else:
+                    if self._fail_callback:
+                        self._fail_callback()
+                    raise
 
     def type(self) -> str:
         return self._model_type
diff --git a/caikit/runtime/model_management/model_loader.py b/caikit/runtime/model_management/model_loader.py
index 692baf63..ab8807c8 100644
--- a/caikit/runtime/model_management/model_loader.py
+++ b/caikit/runtime/model_management/model_loader.py
@@ -13,6 +13,7 @@
 # limitations under the License.
 # Standard
 from concurrent.futures import ThreadPoolExecutor
+from functools import partial
 from typing import Callable, Optional, Union
 
 # Third Party
@@ -64,6 +65,7 @@ class ModelLoader:
         model_type: str,
         aborter: Optional[ActionAborter] = None,
         fail_callback: Optional[Callable] = None,
+        retries: int = 0,
     ) -> LoadedModel:
         """Start loading a model from disk and associate the ID/size with it
 
@@ -75,6 +77,7 @@ class ModelLoader:
                 the call's parent to abort the load
             fail_callback (Optional[Callable]): Optional no-arg callback to call
                 on load failure
+            retries (int): Number of times to retry loading
         Returns:
             model (LoadedModel) : The model that was loaded
         """
@@ -85,6 +88,7 @@ class ModelLoader:
             .type(model_type)
             .path(local_model_path)
             .fail_callback(fail_callback)
+            .retries(retries)
         )
 
         # Set up the async loading
@@ -93,10 +97,12 @@ class ModelLoader:
         if aborter is not None:
             log.debug3("Using abortable action to load %s", model_id)
             action = AbortableAction(aborter, self._load_module, *args)
-            future = self._load_thread_pool.submit(action.do)
+            future_factory = partial(self._load_thread_pool.submit, action.do)
         else:
-            future = self._load_thread_pool.submit(self._load_module, *args)
-        model_builder.model_future(future)
+            future_factory = partial(
+                self._load_thread_pool.submit, self._load_module, *args
+            )
+        model_builder.model_future_factory(future_factory)
 
         # Return the built model with the future handle
         return model_builder.build()
diff --git a/caikit/runtime/model_management/model_manager.py b/caikit/runtime/model_management/model_manager.py
index 6f84e615..468859ef 100644
--- a/caikit/runtime/model_management/model_manager.py
+++ b/caikit/runtime/model_management/model_manager.py
@@ -157,6 +157,7 @@ class ModelManager:  # pylint: disable=too-many-instance-attributes
         model_type: str,
         wait: bool = True,
         aborter: Optional[ActionAborter] = None,
+        retries: Optional[int] = None,
     ) -> int:
         """Load a model using model_path (in Cloud Object Storage) & give it a model ID
         Args:
@@ -164,6 +165,8 @@ class ModelManager:  # pylint: disable=too-many-instance-attributes
             local_model_path (str): Local path to load the model from.
             model_type (str): Type of the model to load.
             wait (bool): Wait for the model to finish loading
+            aborter (Optional[ActionAborter]): The aborter to use for this load
+            retries: Optional[int]: Number of times to retry on load failure
         Returns:
             Model_size (int) : Size of the loaded model in bytes
         """
@@ -191,6 +194,7 @@ class ModelManager:  # pylint: disable=too-many-instance-attributes
                             model_type,
                             aborter=aborter,
                             fail_callback=partial(self.unload_model, model_id),
+                            retries=retries,
                         )
                     except Exception as ex:
                         self.__increment_load_model_exception_count_metric(model_type)
@@ -382,6 +386,7 @@ class ModelManager:  # pylint: disable=too-many-instance-attributes
                 local_model_path=local_model_path,
                 model_type=self._LOCAL_MODEL_TYPE,
                 wait=True,
+                retries=get_config().runtime.lazy_load_retries,
             )
             model_loaded = True
 
@@ -440,7 +445,13 @@ 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)
-            self.load_model(model_id, model_path, self._LOCAL_MODEL_TYPE, wait=False)
+            self.load_model(
+                model_id,
+                model_path,
+                self._LOCAL_MODEL_TYPE,
+                wait=False,
+                retries=get_config().runtime.lazy_load_retries,
+            )
 
         # Unload old models
         # NOTE: No need for error handling here since unload_model will warn on


Test patch:
diff --git a/tests/conftest.py b/tests/conftest.py
index f1c6a180..f52a882d 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -4,7 +4,7 @@ This sets up global test configs when pytest starts
 
 # Standard
 from contextlib import contextmanager
-from typing import List, Union
+from typing import Callable, List, Union
 from unittest.mock import patch
 import copy
 import json
@@ -202,6 +202,26 @@ def backend_priority(backend_cfg: Union[List[dict], dict]):
         yield
 
 
+class TempFailWrapper:
+    """Helper that can wrap a callable with a sequence of failures"""
+
+    def __init__(
+        self,
+        func: Callable,
+        num_failures: int = 1,
+        exc: Exception = RuntimeError("Yikes"),
+    ):
+        self.func = func
+        self.num_failures = num_failures
+        self.exc = exc
+
+    def __call__(self, *args, **kwargs):
+        if self.num_failures:
+            self.num_failures -= 1
+            raise self.exc
+        return self.func(*args, **kwargs)
+
+
 # IMPLEMENTATION DETAILS ############################################################
 def _random_id():
     return str(uuid.uuid4())
diff --git a/tests/runtime/conftest.py b/tests/runtime/conftest.py
index a0029369..6fbda163 100644
--- a/tests/runtime/conftest.py
+++ b/tests/runtime/conftest.py
@@ -4,6 +4,7 @@ This sets up global test configs when pytest starts
 
 # Standard
 from contextlib import contextmanager
+from functools import partial
 from typing import Dict, Type, Union
 import os
 import shlex
@@ -303,13 +304,13 @@ def register_trained_model(
     """Helper to auto-load a model that has completed training. This replaces
     the old auto-load feature which was only needed for unit tests
     """
-    model_future = MODEL_MANAGER.get_model_future(training_id)
+    model_future_factory = partial(MODEL_MANAGER.get_model_future, training_id)
     loaded_model = (
         LoadedModel.Builder()
         .id(model_id)
         .type("trained")
         .path("")
-        .model_future(model_future)
+        .model_future_factory(model_future_factory)
         .build()
     )
     if isinstance(servicer, RuntimeGRPCServer):
diff --git a/tests/runtime/model_management/test_model_loader.py b/tests/runtime/model_management/test_model_loader.py
index 4bd675bb..812bdf3f 100644
--- a/tests/runtime/model_management/test_model_loader.py
+++ b/tests/runtime/model_management/test_model_loader.py
@@ -14,6 +14,7 @@
 # Standard
 from concurrent.futures import Future
 from contextlib import contextmanager
+from typing import Callable
 from unittest import mock
 import copy
 import tempfile
@@ -32,7 +33,7 @@ from caikit.runtime.model_management.model_loader import ModelLoader
 from caikit.runtime.types.caikit_runtime_exception import CaikitRuntimeException
 from sample_lib.data_model import SampleInputType, SampleOutputType
 from sample_lib.modules.sample_task import SampleModule
-from tests.conftest import random_test_id, temp_config
+from tests.conftest import TempFailWrapper, random_test_id, temp_config
 from tests.core.helpers import MockBackend
 from tests.fixtures import Fixtures
 import caikit.core
@@ -328,3 +329,51 @@ def test_load_model_without_waiting_deferred_error(model_loader):
         loaded_model.model()
     assert context.value.status_code == grpc.StatusCode.INTERNAL
     assert model_id in context.value.message
+
+
+def test_load_model_succeed_after_retry(model_loader):
+    """Make sure that a model which fails to load on the first try can load
+    successfully on a retry.
+    """
+    failures = 2
+    fail_wrapper = TempFailWrapper(
+        model_loader._load_module,
+        num_failures=failures,
+        exc=CaikitRuntimeException(grpc.StatusCode.INTERNAL, "Yikes!"),
+    )
+    with mock.patch.object(model_loader, "_load_module", fail_wrapper):
+        model_id = random_test_id()
+        loaded_model = model_loader.load_model(
+            model_id=model_id,
+            local_model_path=Fixtures.get_good_model_path(),
+            model_type=Fixtures.get_good_model_type(),
+            retries=failures + 1,
+        )
+        assert loaded_model.model() is not None
+        assert isinstance(loaded_model.model(), base.ModuleBase)
+        assert loaded_model._retries == 1
+
+
+def test_load_model_fail_callback_once(model_loader):
+    """Make sure that a model which fails all of its retries will only call the
+    fail callback once
+    """
+    failures = 3
+    fail_wrapper = TempFailWrapper(
+        model_loader._load_module,
+        num_failures=failures,
+        exc=CaikitRuntimeException(grpc.StatusCode.INTERNAL, "Yikes!"),
+    )
+    with mock.patch.object(model_loader, "_load_module", fail_wrapper):
+        model_id = random_test_id()
+        fail_cb = mock.MagicMock()
+        loaded_model = model_loader.load_model(
+            model_id=model_id,
+            local_model_path=Fixtures.get_good_model_path(),
+            model_type=Fixtures.get_good_model_type(),
+            fail_callback=fail_cb,
+            retries=failures - 1,
+        )
+        with pytest.raises(CaikitRuntimeException):
+            loaded_model.wait()
+        fail_cb.assert_called_once()
diff --git a/tests/runtime/model_management/test_model_manager.py b/tests/runtime/model_management/test_model_manager.py
index 63f1fd5e..eb325e56 100644
--- a/tests/runtime/model_management/test_model_manager.py
+++ b/tests/runtime/model_management/test_model_manager.py
@@ -14,6 +14,7 @@
 # Standard
 from concurrent.futures import Future, ThreadPoolExecutor
 from contextlib import contextmanager
+from functools import partial
 from tempfile import TemporaryDirectory
 from unittest.mock import MagicMock, patch
 import os
@@ -33,7 +34,7 @@ from caikit.runtime.model_management.loaded_model import LoadedModel
 from caikit.runtime.model_management.model_manager import ModelManager
 from caikit.runtime.types.caikit_runtime_exception import CaikitRuntimeException
 from caikit.runtime.utils.import_util import get_dynamic_module
-from tests.conftest import random_test_id, temp_config
+from tests.conftest import TempFailWrapper, random_test_id, temp_config
 from tests.core.helpers import TestFinder
 from tests.fixtures import Fixtures
 import caikit.runtime.model_management.model_loader
@@ -626,9 +627,10 @@ def test_retrieve_model_returns_the_module_from_the_model_loader():
             mock_sizer.get_model_size.return_value = 1
             model_future = Future()
             model_future.result = lambda *_, **__: expected_module
+            model_future_factory = lambda: model_future
             mock_loader.load_model.return_value = (
                 LoadedModel.Builder()
-                .model_future(model_future)
+                .model_future_factory(model_future_factory)
                 .id("foo")
                 .type("bar")
                 .build()
@@ -647,7 +649,7 @@ def test_unload_partially_loaded():
     # completed successfully
     slow_loader = SlowLoader()
     pool = ThreadPoolExecutor()
-    model_future = pool.submit(slow_loader.load)
+    model_future_factory = partial(pool.submit, slow_loader.load)
 
     model_id = random_test_id()
     mock_sizer = MagicMock()
@@ -657,7 +659,7 @@ def test_unload_partially_loaded():
             mock_sizer.get_model_size.return_value = 1
             mock_loader.load_model.return_value = (
                 LoadedModel.Builder()
-                .model_future(model_future)
+                .model_future_factory(model_future_factory)
                 .id("foo")
                 .type("bar")
                 .build()
@@ -686,7 +688,7 @@ def test_unload_unexpected_error_loaded():
     # completed successfully
     slow_loader = SlowLoader(RuntimeError("yikes"))
     pool = ThreadPoolExecutor()
-    model_future = pool.submit(slow_loader.load)
+    model_future_factory = partial(pool.submit, slow_loader.load)
 
     model_id = random_test_id()
     mock_sizer = MagicMock()
@@ -696,7 +698,7 @@ def test_unload_unexpected_error_loaded():
             mock_sizer.get_model_size.return_value = 1
             mock_loader.load_model.return_value = (
                 LoadedModel.Builder()
-                .model_future(model_future)
+                .model_future_factory(model_future_factory)
                 .id("foo")
                 .type("bar")
                 .build()
@@ -727,7 +729,7 @@ def test_reload_partially_loaded():
     # completed successfully
     slow_loader = SlowLoader()
     pool = ThreadPoolExecutor()
-    model_future = pool.submit(slow_loader.load)
+    model_future_factory = partial(pool.submit, slow_loader.load)
 
     model_id = random_test_id()
     mock_sizer = MagicMock()
@@ -736,13 +738,14 @@ def test_reload_partially_loaded():
     with patch.object(MODEL_MANAGER, "model_loader", mock_loader):
         with patch.object(MODEL_MANAGER, "model_sizer", mock_sizer):
             mock_sizer.get_model_size.return_value = special_model_size
-            mock_loader.load_model.return_value = (
+            loaded_model = (
                 LoadedModel.Builder()
-                .model_future(model_future)
+                .model_future_factory(model_future_factory)
                 .id("foo")
                 .type("bar")
                 .build()
             )
+            mock_loader.load_model.return_value = loaded_model
             model_size = MODEL_MANAGER.load_model(
                 model_id, ANY_MODEL_PATH, ANY_MODEL_TYPE, wait=False
             )
@@ -757,7 +760,7 @@ def test_reload_partially_loaded():
             # Unblock the load
             assert not slow_loader.done_loading()
             slow_loader.unblock_load()
-            model_future.result()
+            loaded_model.wait()
             assert slow_loader.done_loading()
 
 
@@ -835,3 +838,77 @@ def test_periodic_sync_handles_errors():
                 manager.sync_local_models(True)
                 mock_sync.assert_called_once()
                 assert manager._lazy_sync_timer is not None
+
+
+def test_periodic_sync_handles_temporary_errors():
+    """Test that models loaded with the periodic sync can retry if the initial
+    load operation fails
+    """
+    with TemporaryDirectory() as cache_dir:
+        with non_singleton_model_managers(
+            1,
+            {
+                "runtime": {
+                    "local_models_dir": cache_dir,
+                    "lazy_load_local_models": True,
+                    "lazy_load_retries": 1,
+                },
+            },
+            "merge",
+        ) as managers:
+            manager = managers[0]
+            flakey_loader = TempFailWrapper(
+                manager.model_loader._load_module,
+                num_failures=1,
+                exc=CaikitRuntimeException(grpc.StatusCode.INTERNAL, "Dang"),
+            )
+            with patch.object(
+                manager.model_loader,
+                "_load_module",
+                flakey_loader,
+            ):
+                assert manager._lazy_sync_timer is not None
+                model_path = Fixtures.get_good_model_path()
+                model_name = os.path.basename(model_path)
+                shutil.copytree(model_path, os.path.join(cache_dir, model_name))
+                manager.sync_local_models(True)
+                assert manager._lazy_sync_timer is not None
+                model = manager.retrieve_model(model_name)
+                assert model
+
+
+def test_lazy_load_handles_temporary_errors():
+    """Test that a lazy load without a periodic sync correctly retries failed
+    loads
+    """
+    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,
+                    "lazy_load_retries": 1,
+                },
+            },
+            "merge",
+        ) as managers:
+            manager = managers[0]
+            flakey_loader = TempFailWrapper(
+                manager.model_loader._load_module,
+                num_failures=1,
+                exc=CaikitRuntimeException(grpc.StatusCode.INTERNAL, "Dang"),
+            )
+            with patch.object(
+                manager.model_loader,
+                "_load_module",
+                flakey_loader,
+            ):
+                assert manager._lazy_sync_timer is None
+                model_path = Fixtures.get_good_model_path()
+                model_name = os.path.basename(model_path)
+                shutil.copytree(model_path, os.path.join(cache_dir, model_name))
+                assert manager._lazy_sync_timer is None
+                model = manager.retrieve_model(model_name)
+                assert model