Repository: caikit/caikit
Pull request #341 (2023-07-18 17:45:38)

Lazy local model loading 324
## Description

This is a better implementation of https://github.com/caikit/caikit/pull/336

This PR adds a new feature to `caikit.runtime` to use a local directory on disk (which can be mounted) to act as a shared-state cache. This will allow replicas of `caikit.runtime` to host an eventually-consistent set of models that changes dynamically. This is a significantly weaker value proposition than `model-mesh`, but for many model types, especially those that use a `module_backend` to delegate the work to another engine (e.g. `caikit-tgis-backend`), it's fine to require that all models be loaded on all replicas.


Patch:
diff --git a/caikit/config/config.yml b/caikit/config/config.yml
index 46159b10..d5b26983 100644
--- a/caikit/config/config.yml
+++ b/caikit/config/config.yml
@@ -73,6 +73,17 @@ runtime:
     library: sample_lib # TODO: replace with libraries below when runtime can support multiple libraries
     local_models_dir: models
 
+    # 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
+    # previously loaded from local_models_dir, but are no longer present will be
+    # unloaded.
+    lazy_load_local_models: false
+    lazy_load_poll_period_seconds: 10
+
+    # Number of threads to make available for load jobs (null => all)
+    load_threads: null
+
     # TLS configs
     tls:
         server:
diff --git a/caikit/core/model_management/model_trainer_base.py b/caikit/core/model_management/model_trainer_base.py
index e46a9bdb..3047c887 100644
--- a/caikit/core/model_management/model_trainer_base.py
+++ b/caikit/core/model_management/model_trainer_base.py
@@ -97,6 +97,10 @@ class ModelTrainerBase(FactoryConstructible):
 
         ## Common Impl ##
 
+        def result(self) -> ModuleBase:
+            """Support result() to match concurrent.futures.Future"""
+            return self.load()
+
         @classmethod
         def _save_path_with_id(
             cls,
diff --git a/caikit/runtime/http_server.py b/caikit/runtime/http_server.py
index 309b8ff4..c2e2514c 100644
--- a/caikit/runtime/http_server.py
+++ b/caikit/runtime/http_server.py
@@ -472,14 +472,14 @@ class RuntimeHTTPServer(RuntimeServerBase):
         if get_origin(field_type) is Annotated:
             return cls._get_pydantic_type(get_args(field_type)[0])
         if get_origin(field_type) is Union:
-            return Union.__getitem__(
+            return Union[  # type: ignore
                 tuple(
                     (
                         cls._get_pydantic_type(arg_type)
                         for arg_type in get_args(field_type)
                     )
                 )
-            )
+            ]
         if get_origin(field_type) is list:
             return List[cls._get_pydantic_type(get_args(field_type)[0])]
 
diff --git a/caikit/runtime/model_management/loaded_model.py b/caikit/runtime/model_management/loaded_model.py
index dcc55075..de3ca51d 100644
--- a/caikit/runtime/model_management/loaded_model.py
+++ b/caikit/runtime/model_management/loaded_model.py
@@ -11,23 +11,55 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+"""
+A LoadedModel is a metadata wrapper around an instance of a core.ModuleBase
+class that contains the additional information needed to manage that model in
+the runtime.
+"""
+
+# Standard
+from concurrent.futures import Future
+from typing import Callable, Optional
 
 # First Party
 import alog
 
 # Local
 from caikit.core import ModuleBase
+from caikit.core.toolkit.errors import error_handler
+from caikit.runtime.types.caikit_runtime_exception import CaikitRuntimeException
 
 log = alog.use_channel("LOADED-MODEL")
+error = error_handler.get(log)
+
+
+# A future object that will yield an instance of a caikit module (a model)
+# NOTE: 3.9 introduced subscript typing for Futures
+try:
+    CaikitModelFuture = Future[ModuleBase]  # pylint: disable=unsubscriptable-object
+except TypeError:  # pragma: no cover
+    CaikitModelFuture = Future
 
 
 class LoadedModel:
+    __doc__ = __doc__
+
     class Builder:
+        """The LoadedModel.Builder allows the LoadedModel instance to be
+        constructed in pieces with chained '.' getattr semantics.
+        """
+
         def __init__(self):
             self._model_to_build = LoadedModel()
 
-        def module(self, caikit_module: ModuleBase) -> "LoadedModel.Builder":
-            self._model_to_build._caikit_module = caikit_module
+        def model_future(
+            self, caikit_model_future: CaikitModelFuture
+        ) -> "LoadedModel.Builder":
+            self._model_to_build._caikit_model_future = caikit_model_future
+            return self
+
+        def fail_callback(self, callback: Callable) -> "LoadedModel.Builder":
+            self._model_to_build._fail_callback = callback
             return self
 
         def path(self, model_path: str) -> "LoadedModel.Builder":
@@ -43,21 +75,44 @@ class LoadedModel:
             return self
 
         def build(self) -> "LoadedModel":
+            error.value_check(
+                "<RUN12786023E>",
+                self._model_to_build._caikit_model_future
+                and self._model_to_build._model_id
+                and self._model_to_build._model_type,
+                "Cannot build LoadedModel with incomplete required fields."
+                + " Future: {}, ID: {}, Type: {}",
+                self._model_to_build._caikit_model_future,
+                self._model_to_build._model_id,
+                self._model_to_build._model_type,
+            )
             return self._model_to_build
 
     def __init__(self):
         # Use the builder ^^
-        self._caikit_module = None
-        self._model_id = ""
-        self._model_path = ""
-        self._model_type = ""
-        self._size = 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._size: Optional[int] = None
 
     def id(self) -> str:
         return self._model_id
 
-    def module(self) -> ModuleBase:
-        return self._caikit_module
+    def model(self) -> ModuleBase:
+        self.wait()
+        return self._model
+
+    def wait(self):
+        if self._model is None:
+            try:
+                self._model = self._caikit_model_future.result()
+            except CaikitRuntimeException:
+                if self._fail_callback:
+                    self._fail_callback()
+                raise
 
     def type(self) -> str:
         return self._model_type
@@ -79,8 +134,8 @@ class LoadedModel:
             self._size = model_size
         elif self._size != model_size:
             log.warning(
-                "COM62206343W",
-                "Attempted to set size of model %s to %v, but it was already %v",
+                "<RUN46815705W>",
+                "Attempted to set size of model %s to %s, but it was already %s",
                 self.id(),
                 model_size,
                 self.size(),
diff --git a/caikit/runtime/model_management/model_loader.py b/caikit/runtime/model_management/model_loader.py
index 820685f6..2d42e295 100644
--- a/caikit/runtime/model_management/model_loader.py
+++ b/caikit/runtime/model_management/model_loader.py
@@ -12,7 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # Standard
-from typing import Union
+from concurrent.futures import ThreadPoolExecutor
+from typing import Callable, Optional, Union
 
 # Third Party
 from grpc import StatusCode
@@ -27,6 +28,10 @@ from caikit.core import MODEL_MANAGER
 from caikit.runtime.model_management.batcher import Batcher
 from caikit.runtime.model_management.loaded_model import LoadedModel
 from caikit.runtime.types.caikit_runtime_exception import CaikitRuntimeException
+from caikit.runtime.work_management.abortable_action import (
+    AbortableAction,
+    ActionAborter,
+)
 import caikit.core
 
 log = alog.use_channel("MODEL-LOADER")
@@ -48,43 +53,64 @@ class ModelLoader:
         # Re-instantiating this is a programming error
         assert self.__class__.__instance is None, "This class is a singleton!"
         ModelLoader.__instance = self
-        # Instead of storing self.config = get_config(), we will call get_config() everywhere
-        # since this class is a singleton and configuration may be updated during its
-        # lifetime
+        self._load_thread_pool = ThreadPoolExecutor(get_config().runtime.load_threads)
+        # Instead of storing config-based batching information here, we call
+        # get_config() when needed to support dynamic config changes for
+        # batching
 
-    def load_model(self, model_id, local_model_path, model_type) -> LoadedModel:
-        """Load a model using model_path (in Cloud Object Storage) & give it a model ID.
-        This always cleans up files on disk, no matter if the load succeeds or fails
+    def load_model(
+        self,
+        model_id: str,
+        local_model_path: str,
+        model_type: str,
+        aborter: Optional[ActionAborter] = None,
+        fail_callback: Optional[Callable] = None,
+    ) -> LoadedModel:
+        """Start loading a model from disk and associate the ID/size with it
 
         Args:
-            model_id (string):  Model ID string for the model to load.
-            local_model_path (string): Local filesystem path to load the model from.
-            model_type (string): Type of the model to load.
+            model_id (str): Model ID string for the model to load.
+            local_model_path (str): Local filesystem path to load the model from.
+            model_type (str): Type of the model to load.
+            aborter (Optional[ActionAborter]): An aborter to use that will allow
+                the call's parent to abort the load
+            fail_callback (Optional[Callable]): Optional no-arg callback to call
+                on load failure
         Returns:
             model (LoadedModel) : The model that was loaded
         """
-        with CAIKIT_CORE_LOAD_DURATION_SUMMARY.labels(model_type=model_type).time():
-            caikit_core_model = self._load_module(
-                local_model_path, model_id, model_type
-            )
-
-        model = (
+        # Set up the basics of the model's metadata
+        model_builder = (
             LoadedModel.Builder()
             .id(model_id)
             .type(model_type)
             .path(local_model_path)
-            .module(caikit_core_model)
-            .build()
+            .fail_callback(fail_callback)
         )
 
-        return model
-
-    def _load_module(self, model_path, model_id, model_type):
+        # Set up the async loading
+        args = (local_model_path, model_id, model_type)
+        log.debug2("Loading model %s async", model_id)
+        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)
+        else:
+            future = self._load_thread_pool.submit(self._load_module, *args)
+        model_builder.model_future(future)
+
+        # Return the built model with the future handle
+        return model_builder.build()
+
+    def _load_module(
+        self, model_path: str, model_id: str, model_type: str
+    ) -> LoadedModel:
         try:
             log.info("<RUN89711114I>", "Loading model '%s'", model_id)
 
             # Load using the caikit.core
-            model = caikit.core.load(model_path)
+            with CAIKIT_CORE_LOAD_DURATION_SUMMARY.labels(model_type=model_type).time():
+                model = caikit.core.load(model_path)
 
             # If this model needs batching, configure a Batcher to wrap it
             model = self._wrap_in_batcher_if_configured(
diff --git a/caikit/runtime/model_management/model_manager.py b/caikit/runtime/model_management/model_manager.py
index 6ece3d53..24f097f5 100644
--- a/caikit/runtime/model_management/model_manager.py
+++ b/caikit/runtime/model_management/model_manager.py
@@ -13,7 +13,9 @@
 # limitations under the License.
 # Standard
 from collections import Counter as DictCounter
-from typing import Dict
+from functools import partial
+from typing import Dict, Optional
+import atexit
 import gc
 import os
 import threading
@@ -28,12 +30,15 @@ import alog
 # Local
 from caikit import get_config
 from caikit.core import ModuleBase
+from caikit.core.toolkit.errors import error_handler
 from caikit.runtime.model_management.loaded_model import LoadedModel
 from caikit.runtime.model_management.model_loader import ModelLoader
 from caikit.runtime.model_management.model_sizer import ModelSizer
 from caikit.runtime.types.caikit_runtime_exception import CaikitRuntimeException
+from caikit.runtime.work_management.abortable_action import ActionAborter
 
 log = alog.use_channel("MODEL-MANAGR")
+error = error_handler.get(log)
 
 MODEL_SIZE_GAUGE = Gauge(
     "total_loaded_models_size",
@@ -56,7 +61,7 @@ LOAD_MODEL_DURATION_SUMMARY = Summary(
 LOCAL_MODEL_TYPE = "LOCAL"
 
 
-class ModelManager:
+class ModelManager:  # pylint: disable=too-many-instance-attributes
     """Model Manager class. The singleton class contains the core implementational details
     for the Model Runtime (i.e., load/unload functionality, etc). It does not provide the core
     details for predict calls."""
@@ -65,92 +70,189 @@ class ModelManager:
 
     __model_size_gauge_lock = threading.Lock()
 
+    _LOCAL_MODEL_TYPE = "standalone-model"
+
+    ## Construction ##
+
+    @classmethod
+    def get_instance(cls) -> "ModelManager":
+        """This method returns the instance of Model Manager"""
+        if not cls.__instance:
+            cls.__instance = ModelManager()
+        return cls.__instance
+
     def __init__(self):
         """Initialize a ModelManager instance."""
         # Re-instantiating this is a programming error
         assert self.__class__.__instance is None, "This class is a singleton!"
         ModelManager.__instance = self
 
-        # Pull in a ModelLoader
+        # Pull in a ModelLoader and ModelSizer
         self.model_loader = ModelLoader.get_instance()
-        # And a ModelSizer
         self.model_sizer = ModelSizer.get_instance()
 
-        # set up the local map for model ID to loaded model
+        # In-memory mapping of model_id to LoadedModel instance
         self.loaded_models: Dict[str, LoadedModel] = {}
 
+        # Lock for mutating operations on loaded_models
+        self._loaded_models_lock = threading.Lock()
+
         # Optionally load models mounted into a local directory
-        local_models_path = get_config().runtime.local_models_dir
-        if os.path.exists(local_models_path) and len(os.listdir(local_models_path)) > 0:
+        runtime_cfg = get_config().runtime
+        self._local_models_dir = runtime_cfg.local_models_dir or ""
+        if self._local_models_dir and not os.path.exists(self._local_models_dir):
+            log.warning(
+                "<RUN53709826W>",
+                "Invalid runtime.local_models_dir %s. Does not exist",
+                self._local_models_dir,
+            )
+            self._local_models_dir = ""
+
+        # Keep track of whether lazy loading is enabled
+        self._lazy_load_local_models = runtime_cfg.lazy_load_local_models
+        error.value_check(
+            "<RUN44773514E>",
+            not self._lazy_load_local_models or self._local_models_dir,
+            "Must set runtime.local_models_dir with runtime.lazy_load_local_models",
+        )
+
+        # Set up local model periodic sync
+        self._lazy_load_poll_period_seconds = runtime_cfg.lazy_load_poll_period_seconds
+        error.type_check(
+            "<RUN59138047E>",
+            int,
+            float,
+            allow_none=True,
+            lazy_load_poll_period_seconds=self._lazy_load_poll_period_seconds,
+        )
+        self._lazy_sync_timer = None
+        self._enable_lazy_load_poll = (
+            self._local_models_dir
+            and self._lazy_load_local_models
+            and self._lazy_load_poll_period_seconds
+        )
+        if self._enable_lazy_load_poll:
+            atexit.register(self.shut_down)
+
+        # Do the initial local models load
+        if self._local_models_dir:
             log.info("<RUN44739400I>", "Loading local models into Caikit Runtime...")
-            self.load_local_models(local_models_path)
-
-    def load_model(self, model_id, local_model_path, model_type) -> int:
+            self.sync_local_models(wait=True)
+
+    def shut_down(self):
+        """Shut down cache purging"""
+        timer = getattr(self, "_lazy_sync_timer", None)
+        if timer is not None:
+            timer.cancel()
+
+    ## Model Management ##
+
+    def load_model(
+        self,
+        model_id: str,
+        local_model_path: str,
+        model_type: str,
+        wait: bool = True,
+        aborter: Optional[ActionAborter] = None,
+    ) -> int:
         """Load a model using model_path (in Cloud Object Storage) & give it a model ID
         Args:
-            model_id (string):  Model ID string for the model to load.
-            local_model_path (string): Local path to load the model from.
-            model_type (string): Type of the model to load.
+            model_id (str):  Model ID string for the model to load.
+            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
         Returns:
             Model_size (int) : Size of the loaded model in bytes
         """
         with LOAD_MODEL_DURATION_SUMMARY.labels(model_type=model_type).time():
-            if model_id in self.loaded_models:
-                log.debug("Model '%s' is already loaded", model_id)
-                return self.loaded_models[model_id].size()
 
-            try:
-                model = self.model_loader.load_model(
-                    model_id, local_model_path, model_type
-                )
-            except Exception as ex:
-                self.__increment_load_model_exception_count_metric(model_type)
-                raise ex
-
-            model_size = self.model_sizer.get_model_size(
-                model_id, local_model_path, model_type
-            )
-            model.set_size(model_size)
-
-            # Add model + helpful metadata to our loaded models map
-            self.loaded_models[model_id] = model
-
-            # Update Prometheus metrics
-            self.__increment_model_count_metric(model_type, model_id)
-            self.__report_total_model_size_metric()
-
-            return self.loaded_models[model_id].size()
+            # If already loaded, just return the size
+            # NOTE: We make the dict access atomic here to avoid the race where
+            #   we check if model_id in the map, then re-look it up to get the
+            #   size which could fail if it is unloaded between the two.
+            model = self.loaded_models.get(model_id)
+            if model is not None:
+                log.debug("Model '%s' is already loaded", model_id)
+                return model.size()
+
+            # Grab the mutation lock and load the model if needed
+            with self._loaded_models_lock:
+                # Re-check now that the mutation lock is held
+                model = self.loaded_models.get(model_id)
+                if model is None:
+                    log.debug("Loading %s from %s", model_id, local_model_path)
+                    try:
+                        model = self.model_loader.load_model(
+                            model_id,
+                            local_model_path,
+                            model_type,
+                            aborter=aborter,
+                            fail_callback=partial(self.unload_model, model_id),
+                        )
+                    except Exception as ex:
+                        self.__increment_load_model_exception_count_metric(model_type)
+                        raise ex
+
+                    # Estimate the model's size and update the LoadedModel
+                    model_size = self.model_sizer.get_model_size(
+                        model_id, local_model_path, model_type
+                    )
+                    model.set_size(model_size)
+
+                    # Add model + helpful metadata to our loaded models map
+                    self.loaded_models[model_id] = model
+
+                    # Update Prometheus metrics
+                    self.__increment_model_count_metric(model_type, model_id)
+                    self.__report_total_model_size_metric()
+
+            # If waiting, do so outside of the mutation lock
+            if wait:
+                model.wait()
+
+            # Return the model's size
+            return model.size()
+
+    def sync_local_models(self, wait: bool = False):
+        """Sync in-memory models with models in the configured local_model_dir
+
+        New models will be loaded and models previously loaded from local will
+        be unloaded.
 
-    def load_local_models(self, local_model_dir):
-        """Load models mounted into a local directory
         Args:
-            local_model_dir (string):  Path to the local model directory
-        Returns:
-            None
+            wait (bool): Wait for loading to complete
         """
-        for model_base_path in os.listdir(local_model_dir):
-            try:
-                # Use the file name as the model id
-                model_id = model_base_path
-                model_path = os.path.join(local_model_dir, model_base_path)
-                self.load_model(model_id, model_path, "standalone-model")
-            except CaikitRuntimeException as e:
-                log.warning(
-                    "<RUN56627484W>",
-                    "Failed to load model %s: %s",
-                    model_base_path,
-                    repr(e),
-                    exc_info=True,
-                )
-                continue
-
-        if len(self.loaded_models) == 0:
-            log.error(
-                "<RUN56336804E>", "No models loaded in directory: %s", local_model_dir
+        try:
+            self._local_models_dir_sync(wait)
+        except StopIteration:
+            log.warning(
+                "<RUN56519883W>",
+                "local_models_dir %s unreachable. Terminating synchronization",
+                self._local_models_dir,
             )
-            raise CaikitRuntimeException(
-                StatusCode.INTERNAL, "No standalone models loaded"
+            self._enable_lazy_load_poll = False
+        except Exception as err:  # pylint: disable=broad-exception-caught
+            log.warning(
+                "<RUN44524933W>",
+                "Exception raised during local_models_dir sync: %s",
+                str(err),
+                exc_info=True,
+            )
+
+        # If running periodically, kick off the next iteration
+        if self._enable_lazy_load_poll:
+            if self._lazy_sync_timer is None:
+                log.info(
+                    "Initializing local_models_dir sync with period %s",
+                    self._lazy_load_poll_period_seconds,
+                )
+            if self._lazy_sync_timer is not None and self._lazy_sync_timer.is_alive():
+                log.debug2("Canceling live timer")
+                self._lazy_sync_timer.cancel()
+            self._lazy_sync_timer = threading.Timer(
+                self._lazy_load_poll_period_seconds, self.sync_local_models
             )
+            self._lazy_sync_timer.start()
 
     def unload_model(self, model_id) -> int:
         """Unload a model by ID model.
@@ -173,7 +275,13 @@ class ModelManager:
 
         # Delete the model and remove it from the model map
         try:
-            del self.loaded_models[model_id]
+            model = self.loaded_models.pop(model_id)
+            # If the model is still loading, we need to wait for it to finish so
+            # that we can do our best to fully free it
+            model.wait()
+            del model
+        except CaikitRuntimeException:
+            raise
         except Exception as ex:
             log.debug("Model '%s' failed to unload with error: %s", model_id, repr(ex))
             raise CaikitRuntimeException(
@@ -192,16 +300,11 @@ class ModelManager:
 
         return model_size
 
-    def unload_all_models(self) -> None:
-        """Unload all loaded models. This will also remove any lingering artifacts from strangely
-        packaged zip files, which initially resulted in load failures."""
-        try:
-            self.loaded_models.clear()
-        except Exception as ex:
-            log.debug("Unload all models failed with error: %s", repr(ex))
-            raise CaikitRuntimeException(
-                StatusCode.INTERNAL, "All models could not be unloaded!!"
-            ) from ex
+    def unload_all_models(self):
+        """Unload all loaded models"""
+        all_model_ids = list(self.loaded_models.keys())
+        for model_id in all_model_ids:
+            self.unload_model(model_id)
 
     def get_model_size(self, model_id) -> int:
         """Look up size of a model by model ID.
@@ -222,22 +325,6 @@ class ModelManager:
             )
 
         loaded_model = self.loaded_models[model_id]
-
-        # Model sizes should all be cached
-        if not loaded_model.has_size():
-            # This really shouldn't happen, because we size the models on load
-            log.warning(
-                "<RUN61106343W>",
-                "Loaded model %s unexpectedly did not have a size. Sizing it again now.",
-                model_id,
-            )
-            loaded_model.set_size(
-                self.model_sizer.get_model_size(
-                    model_id, loaded_model.path(), loaded_model.type()
-                )
-            )
-            self.loaded_models[model_id] = loaded_model
-
         self.__report_total_model_size_metric()
         return loaded_model.size()
 
@@ -252,33 +339,122 @@ class ModelManager:
         """
         return self.model_sizer.get_model_size(model_id, local_model_path, model_type)
 
-    def retrieve_model(self, model_id) -> ModuleBase:
+    def retrieve_model(self, model_id: str) -> ModuleBase:
         """Retrieve a model from the loaded model map by model ID.
 
         Args:
-            model_id(string): Model ID of the model to retrieve
+            model_id (str): Model ID of the model to retrieve
         Returns:
             response (caikit.core.module.ModuleBase):
                 A loaded Caikit model
         """
-        if not model_id:
+        if not model_id or not isinstance(model_id, str):
             raise CaikitRuntimeException(
                 StatusCode.INVALID_ARGUMENT, "Missing required model ID"
             )
 
-        # Now retrieve the model
-        if model_id not in self.loaded_models:
-            # We should not encounter this scenario, so if it happens, log
-            # it as an error-level log
+        # Now retrieve the model and fall back to lazy loading
+        model_loaded = model_id in self.loaded_models
+        if not model_loaded and self._lazy_load_local_models:
+            local_model_path = os.path.join(self._local_models_dir, model_id)
+            if os.path.exists(local_model_path):
+                log.debug2(
+                    "Lazy loading local model %s from %s", model_id, local_model_path
+                )
+                self.load_model(
+                    model_id=model_id,
+                    local_model_path=local_model_path,
+                    model_type=self._LOCAL_MODEL_TYPE,
+                    wait=True,
+                )
+                model_loaded = True
+
+        # If still not loaded, there's nothing to find, so raise NOT_FOUND
+        if not model_loaded:
             msg = "Model '%s' not loaded" % model_id
-            log.error(
-                {"log_code": "<RUN61105243E>", "message": msg, "model_id": model_id}
+            log.debug(
+                {"log_code": "<RUN61105243D>", "message": msg, "model_id": model_id}
             )
             raise CaikitRuntimeException(
                 StatusCode.NOT_FOUND, msg, {"model_id": model_id}
             )
 
-        return self.loaded_models[model_id].module()
+        # NOTE: If the model is partially loaded, this call will wait on the
+        #   model future in the LoadedModel
+        return self.loaded_models[model_id].model()
+
+    ## Implementation Details ##
+
+    def _local_models_dir_sync(self, wait: bool = False):
+        """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.
+
+        NOTE: In the case that the local_models_dir becomes unreadable, it will
+            raise StopIteration to indicate that any periodic synchronization
+            should terminate.
+        """
+        # Get the list of models on disk
+        # NOTE: If the local_models_dir has disappeared, this is likely a unit
+        #   test with a temp dir, but in any event, we should stop trying to
+        #   sync going forward
+        try:
+            disk_models = os.listdir(self._local_models_dir)
+        except FileNotFoundError as err:
+            raise StopIteration() from err
+
+        # 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)
+
+        # Find all models that are currently loaded from the local models dir
+        # that no longer exist
+        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,
+            )
+        ]
+        log.debug("Unloaded local models: %s", unload_models)
+
+        # 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)
+
+        # Unload old models
+        # NOTE: No need for error handling here since unload_model will warn on
+        #   errors and move on
+        for model_id in unload_models:
+            log.debug2("Unloading local model %s", model_id)
+            self.unload_model(model_id)
+
+        # Wait for models to load
+        if wait:
+            for model_id in new_models:
+                loaded_model = self.loaded_models.get(model_id)
+                # If somehow already purged, there's nothing to wait on. This is
+                # extremely unlikely since it would require another thread to
+                # explicitly call unload on the model AND have the model finish
+                # loading between then and now. Better to be safe than sorry,
+                # though!
+                if loaded_model is None:  # pragma: no cover
+                    continue
+                # Wait for it and make sure it didn't fail
+                try:
+                    loaded_model.wait()
+                except CaikitRuntimeException as err:
+                    log.warning(
+                        "<RUN56627485W>",
+                        "Failed to load model %s: %s",
+                        model_id,
+                        repr(err),
+                        exc_info=True,
+                    )
 
     def __report_total_model_size_metric(self):
         # Just a happy little lock to ensure that with concurrent loading and unloading,
@@ -302,10 +478,3 @@ class ModelManager:
     @staticmethod
     def __increment_load_model_exception_count_metric(model_type):
         LOAD_MODEL_EXCEPTION_COUNTER.labels(model_type=model_type).inc()
-
-    @classmethod
-    def get_instance(cls) -> "ModelManager":
-        """This method returns the instance of Model Manager"""
-        if not cls.__instance:
-            cls.__instance = ModelManager()
-        return cls.__instance
diff --git a/caikit/runtime/service_generation/protoable.py b/caikit/runtime/service_generation/protoable.py
index 24acf3cc..2ce842c6 100644
--- a/caikit/runtime/service_generation/protoable.py
+++ b/caikit/runtime/service_generation/protoable.py
@@ -125,7 +125,7 @@ def get_union_list_type(field_name: str, union_protoables: List) -> Type[DataBas
             )
         else:
             param_list.append(arg)
-    return Union.__getitem__(tuple((param_list)))
+    return Union[tuple(param_list)]  # type: ignore
 
 
 def get_protoable_return_type(arg_type: Type) -> Type:
diff --git a/caikit/runtime/servicers/global_predict_servicer.py b/caikit/runtime/servicers/global_predict_servicer.py
index a8a2af69..13a94ce5 100644
--- a/caikit/runtime/servicers/global_predict_servicer.py
+++ b/caikit/runtime/servicers/global_predict_servicer.py
@@ -46,7 +46,7 @@ from caikit.runtime.utils.servicer_util import (
     validate_data_model,
 )
 from caikit.runtime.work_management.abortable_action import AbortableAction
-from caikit.runtime.work_management.call_aborter import CallAborter
+from caikit.runtime.work_management.rpc_aborter import RpcAborter
 
 PREDICT_RPC_COUNTER = Counter(
     "predict_rpc_count",
@@ -196,9 +196,7 @@ class GlobalPredictServicer:
                     request_name,
                     model_id,
                     inference_func_name=inference_signature.method_name,
-                    aborter=CallAborter(context)
-                    if self.use_abortable_threads
-                    else None,
+                    aborter=RpcAborter(context) if self.use_abortable_threads else None,
                     **caikit_library_request,
                 )
 
@@ -217,7 +215,7 @@ class GlobalPredictServicer:
         request_name: str,
         model_id: str,
         inference_func_name: str = "run",
-        aborter: Optional[CallAborter] = None,
+        aborter: Optional[RpcAborter] = None,
         **kwargs,
     ) -> Union[DataBase, Iterable[DataBase]]:
         """Run a prediction against the given model using the raw arguments to
@@ -230,7 +228,7 @@ class GlobalPredictServicer:
                 The ID of the loaded model
             inference_func_name (str):
                 The name of the inference function to run
-            aborter (Optional[CallAborter]):
+            aborter (Optional[RpcAborter]):
                 If using abortable calls, this is the aborter to use
             **kwargs: Keyword arguments to pass to the model's run function
         Returns:
diff --git a/caikit/runtime/servicers/model_runtime_servicer.py b/caikit/runtime/servicers/model_runtime_servicer.py
index 55a35d7d..ac846024 100644
--- a/caikit/runtime/servicers/model_runtime_servicer.py
+++ b/caikit/runtime/servicers/model_runtime_servicer.py
@@ -25,7 +25,7 @@ from caikit.runtime.protobufs import model_runtime_pb2, model_runtime_pb2_grpc
 from caikit.runtime.types.aborted_exception import AbortedException
 from caikit.runtime.types.caikit_runtime_exception import CaikitRuntimeException
 from caikit.runtime.work_management.abortable_action import AbortableAction
-from caikit.runtime.work_management.call_aborter import CallAborter
+from caikit.runtime.work_management.rpc_aborter import RpcAborter
 
 log = alog.use_channel("MR-SERVICR-I")
 
@@ -57,12 +57,14 @@ class ModelRuntimeServicerImpl(model_runtime_pb2_grpc.ModelRuntimeServicer):
             )
             caikit_config = get_config()
             if caikit_config.runtime.use_abortable_threads:
+                aborter = RpcAborter(context)
                 work = AbortableAction(
-                    CallAborter(context),
+                    aborter,
                     self.model_manager.load_model,
                     request.modelId,
                     request.modelPath,
                     request.modelType,
+                    aborter=aborter,
                 )
                 model_size = work.do()
             else:
diff --git a/caikit/runtime/work_management/abortable_action.py b/caikit/runtime/work_management/abortable_action.py
index de2eae3f..b561d27d 100644
--- a/caikit/runtime/work_management/abortable_action.py
+++ b/caikit/runtime/work_management/abortable_action.py
@@ -13,6 +13,8 @@
 # limitations under the License.
 
 # Standard
+from typing import Callable
+import abc
 import threading
 
 # First Party
@@ -25,6 +27,21 @@ from caikit.runtime.types.aborted_exception import AbortedException
 log = alog.use_channel("ABORT-ACTION")
 
 
+class ActionAborter(abc.ABC):
+    """Simple interface to wrap up a notification that an action must abort.
+
+    Children of this class can bind to any notification tool (e.g. grpc context)
+    """
+
+    @abc.abstractmethod
+    def must_abort(self) -> bool:
+        """Indicate whether or not the action must be aborted"""
+
+    @abc.abstractmethod
+    def add_event(self, event: threading.Event):
+        """Add an event to notify when abort happens"""
+
+
 class AbortableAction:
     """A class for Abortable Actions. We want actions that are computationally heavy to be
     abortable by Model Mesh! Currently, we use this for the following operations.
@@ -40,17 +57,23 @@ class AbortableAction:
         Instances of this class create a threading.Event, which will be used to signal that either:
         - The RPC was terminated
         - The heavy work that we wanted to complete is done
-        This is done by using a CallAborter and a DestroyableThread.
-        Registering the event with the CallAborter will cause it to set when the RPC is
+        This is done by using a RpcAborter and a DestroyableThread.
+        Registering the event with the RpcAborter will cause it to set when the RPC is
         terminated, and creating a DestroyableThread with the event will cause it to set when
         the thread terminates.
 
         The action will start the DestroyableThread and then wait on the event. When it wakes, it
-        will check the reason and destroy the thread if it was woken by the CallAborter or return
+        will check the reason and destroy the thread if it was woken by the RpcAborter or return
         the result if it was woken by the thread completing.
     """
 
-    def __init__(self, call_aborter, runnable_func, *args, **kwargs):
+    def __init__(
+        self,
+        call_aborter: ActionAborter,
+        runnable_func: Callable,
+        *args,
+        **kwargs,
+    ):
         """
         Args:
             call_aborter - call aborter capable of aborting the runnable_func
@@ -71,7 +94,7 @@ class AbortableAction:
             self.__runnable_func,
             *args,
             work_done_event=self.__done_or_aborted_event,
-            **kwargs
+            **kwargs,
         )
 
     def do(self):
diff --git a/caikit/runtime/work_management/call_aborter.py b/caikit/runtime/work_management/rpc_aborter.py
similarity index 96%
rename from caikit/runtime/work_management/call_aborter.py
rename to caikit/runtime/work_management/rpc_aborter.py
index 196ffd60..39f8df00 100644
--- a/caikit/runtime/work_management/call_aborter.py
+++ b/caikit/runtime/work_management/rpc_aborter.py
@@ -26,11 +26,12 @@ import alog
 
 # Local
 from caikit.runtime.types.caikit_runtime_exception import CaikitRuntimeException
+from caikit.runtime.work_management.abortable_action import ActionAborter
 
 log = alog.use_channel("CALL-ABORTER")
 
 
-class CallAborter:
+class RpcAborter(ActionAborter):
     """
     This class registers a callback with a grpc context, to be called in the event of rpc
     termination. Termination could be nominal (we returned a response) but we should have


Test patch:
diff --git a/tests/runtime/conftest.py b/tests/runtime/conftest.py
index 7bcaa789..f672ba5e 100644
--- a/tests/runtime/conftest.py
+++ b/tests/runtime/conftest.py
@@ -3,6 +3,7 @@ This sets up global test configs when pytest starts
 """
 
 # Standard
+from concurrent.futures import Future
 from contextlib import contextmanager
 from typing import Type, Union
 import os
@@ -246,13 +247,12 @@ def register_trained_model(
     the old auto-load feature which was only needed for unit tests
     """
     model_future = MODEL_MANAGER.get_model_future(training_id)
-    model = model_future.load()
     loaded_model = (
         LoadedModel.Builder()
         .id(model_id)
         .type("trained")
         .path("")
-        .module(model)
+        .model_future(model_future)
         .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 4a38fd43..4bd675bb 100644
--- a/tests/runtime/model_management/test_model_loader.py
+++ b/tests/runtime/model_management/test_model_loader.py
@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # Standard
+from concurrent.futures import Future
 from contextlib import contextmanager
 from unittest import mock
 import copy
@@ -53,6 +54,12 @@ def model_loader():
     return ModelLoader.get_instance()
 
 
+def make_model_future(model_instance):
+    fake_future = Future()
+    fake_future.result = lambda *_, **__: model_instance
+    return fake_future
+
+
 ## Tests #######################################################################
 
 
@@ -64,8 +71,8 @@ def test_load_model_ok_response(model_loader):
         local_model_path=Fixtures.get_good_model_path(),
         model_type=Fixtures.get_good_model_type(),
     )
-    assert loaded_model.module() is not None
-    assert isinstance(loaded_model.module(), base.ModuleBase)
+    assert loaded_model.model() is not None
+    assert isinstance(loaded_model.model(), base.ModuleBase)
     assert model_id == loaded_model.id()
     assert Fixtures.get_good_model_type() == loaded_model.type()
     assert Fixtures.get_good_model_path() == loaded_model.path()
@@ -82,8 +89,8 @@ def test_load_model_archive(model_loader):
         local_model_path=Fixtures.get_good_model_archive_path(),
         model_type=Fixtures.get_good_model_type(),
     )
-    assert loaded_model.module() is not None
-    assert isinstance(loaded_model.module(), base.ModuleBase)
+    assert loaded_model.model() is not None
+    assert isinstance(loaded_model.model(), base.ModuleBase)
 
 
 def test_load_model_error_not_found_response(model_loader):
@@ -94,7 +101,7 @@ def test_load_model_error_not_found_response(model_loader):
             model_id=model_id,
             local_model_path="test/this/does/not/exist.zip",
             model_type="categories_esa",
-        )
+        ).wait()
     assert context.value.status_code == grpc.StatusCode.NOT_FOUND
     assert model_id in context.value.message
 
@@ -107,7 +114,7 @@ def test_load_invalid_model_error_response(model_loader):
             model_id=model_id,
             local_model_path=Fixtures.get_bad_model_archive_path(),
             model_type="not_real",
-        )
+        ).wait()
     assert context.value.status_code == grpc.StatusCode.INTERNAL
     assert model_id in context.value.message
 
@@ -144,7 +151,7 @@ def test_nonzip_extract_fails(model_loader):
             model_id,
             Fixtures.get_invalid_model_archive_path(),
             Fixtures.get_good_model_type(),
-        )
+        ).wait()
     # This ends up returning a FileNotFoundError from caikit core.
     # maybe not the best? But it does include an error message at least
     assert (
@@ -168,7 +175,7 @@ def test_with_batching(model_loader):
         "load_with_batch",
         Fixtures.get_good_model_path(),
         model_type="fake_batch_module",
-    ).module()
+    ).model()
     assert isinstance(model, Batcher)
     assert model._batch_size == get_config().runtime.batching.fake_batch_module.size
 
@@ -177,7 +184,7 @@ def test_with_batching(model_loader):
         "load_without_batch",
         Fixtures.get_good_model_path(),
         model_type=Fixtures.get_good_model_type(),
-    ).module()
+    ).model()
     assert not isinstance(model, Batcher)
 
 
@@ -190,7 +197,7 @@ def test_with_batching_by_default(model_loader):
             "load_with_batch_default",
             Fixtures.get_good_model_path(),
             model_type=Fixtures.get_good_model_type(),
-        ).module()
+        ).model()
         assert isinstance(model, Batcher)
         assert model._batch_size == cfg.runtime.batching.default.size
 
@@ -214,7 +221,7 @@ def test_with_batching_collect_delay(model_loader):
             "load_with_batch_default",
             Fixtures.get_good_model_path(),
             model_type=model_type,
-        ).module()
+        ).model()
         assert isinstance(model, Batcher)
         assert model._batch_size == getattr(cfg.runtime.batching, model_type).size
         assert (
@@ -282,5 +289,42 @@ def test_load_distributed_impl():
                         random_test_id(),
                         model_path,
                         model_type=model_type,
-                    ).module()
+                    ).model()
                     assert isinstance(model, DistributedGadget)
+
+
+def test_load_model_without_waiting_success(model_loader):
+    """Make sure that loading a model can defer the model to a future and access
+    the loaded model when complete
+    """
+    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(),
+    )
+    # Block to get the model
+    assert loaded_model.model() is not None
+    assert isinstance(loaded_model.model(), base.ModuleBase)
+    assert model_id == loaded_model.id()
+    assert Fixtures.get_good_model_type() == loaded_model.type()
+    assert Fixtures.get_good_model_path() == loaded_model.path()
+
+    # Models are not sized by the loader
+    assert loaded_model.size() == 0
+
+
+def test_load_model_without_waiting_deferred_error(model_loader):
+    """Make sure that loading a model can defer the model to a future and raise
+    when the future is used if the loading failed
+    """
+    model_id = random_test_id()
+    loaded_model = model_loader.load_model(
+        model_id=model_id,
+        local_model_path=Fixtures.get_bad_model_archive_path(),
+        model_type="not_real",
+    )
+    with pytest.raises(CaikitRuntimeException) as context:
+        loaded_model.model()
+    assert context.value.status_code == grpc.StatusCode.INTERNAL
+    assert model_id in context.value.message
diff --git a/tests/runtime/model_management/test_model_manager.py b/tests/runtime/model_management/test_model_manager.py
index 864f75fa..2efc9a87 100644
--- a/tests/runtime/model_management/test_model_manager.py
+++ b/tests/runtime/model_management/test_model_manager.py
@@ -12,14 +12,18 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # Standard
+from concurrent.futures import Future, ThreadPoolExecutor
+from contextlib import contextmanager
 from tempfile import TemporaryDirectory
 from unittest.mock import MagicMock, patch
 import os
 import shutil
-import unittest
+import threading
+import time
 
 # Third Party
 import grpc
+import pytest
 
 # Local
 from caikit import get_config
@@ -34,370 +38,772 @@ from tests.fixtures import Fixtures
 get_dynamic_module("caikit.core")
 ANY_MODEL_TYPE = "test-any-model-type"
 ANY_MODEL_PATH = "test-any-model-path"
+MODEL_MANAGER = ModelManager.get_instance()
 
 
-class TestModelManager(unittest.TestCase):
-    """This test suite tests the modelmanager class"""
+@pytest.fixture(autouse=True)
+def tear_down():
+    yield
+    MODEL_MANAGER.unload_all_models()
 
-    def setUp(self):
-        """This method runs before each test begins to run"""
-        self.model_manager = ModelManager.get_instance()
 
-    def tearDown(self):
-        """This method runs after each test executes"""
-        self.model_manager.unload_all_models()
+@contextmanager
+def temp_local_models_dir(workdir, model_manager=MODEL_MANAGER):
+    prev_local_models_dir = model_manager._local_models_dir
+    model_manager._local_models_dir = workdir
+    yield
+    model_manager._local_models_dir = prev_local_models_dir
 
-    # ****************************** Integration Tests ****************************** #
-    # These tests do not patch in mocks, the manager will use real instances of its dependencies
 
-    def test_load_model_ok_response(self):
-        model_id = "happy_load_test"
-        model_size = self.model_manager.load_model(
-            model_id=model_id,
-            local_model_path=Fixtures.get_good_model_path(),
-            model_type=Fixtures.get_good_model_type(),
+@contextmanager
+def non_singleton_model_managers(num_mgrs=1, *args, **kwargs):
+    with temp_config(*args, **kwargs):
+        instances = []
+        try:
+            for _ in range(num_mgrs):
+                ModelManager._ModelManager__instance = None
+                instances.append(ModelManager.get_instance())
+            yield instances
+        finally:
+            for inst in instances:
+                inst.shut_down()
+            ModelManager._ModelManager__instance = MODEL_MANAGER
+
+
+class SlowLoader:
+    """Helper class to simulate slow loading"""
+
+    def __init__(self, load_result="STUB"):
+        self._load_result = load_result
+        self._load_start_event = threading.Event()
+        self._load_end_event = threading.Event()
+
+    def load(self, *_, **__):
+        self._load_start_event.wait()
+        self._load_end_event.set()
+        if isinstance(self._load_result, Exception):
+            raise self._load_result
+        return self._load_result
+
+    def unblock_load(self):
+        self._load_start_event.set()
+
+    def done_loading(self):
+        return self._load_end_event.is_set()
+
+
+# ****************************** Integration Tests ****************************** #
+# These tests do not patch in mocks, the manager will use real instances of its dependencies
+
+
+def test_load_model_ok_response():
+    model_id = "happy_load_test"
+    model_size = MODEL_MANAGER.load_model(
+        model_id=model_id,
+        local_model_path=Fixtures.get_good_model_path(),
+        model_type=Fixtures.get_good_model_type(),
+    )
+    assert model_size > 0
+
+
+@pytest.mark.parametrize("model_id", ["", b"bytes-id", 123])
+def test_retrieve_invalid_model_ids(model_id):
+    with pytest.raises(CaikitRuntimeException) as context:
+        MODEL_MANAGER.retrieve_model(model_id)
+    assert context.value.status_code == grpc.StatusCode.INVALID_ARGUMENT
+
+
+def test_load_model_no_size_update():
+    model_id = random_test_id()
+    model_size = MODEL_MANAGER.load_model(
+        model_id=model_id,
+        local_model_path=Fixtures.get_good_model_path(),
+        model_type=Fixtures.get_good_model_type(),
+    )
+    assert model_size > 0
+    loaded_model = MODEL_MANAGER.loaded_models[model_id]
+    assert loaded_model.size() == model_size
+    loaded_model.set_size(model_size * 2)
+    assert loaded_model.size() == model_size
+
+
+def test_load_local_models():
+    with TemporaryDirectory() as tempdir:
+        shutil.copytree(Fixtures.get_good_model_path(), os.path.join(tempdir, "model1"))
+        shutil.copy(
+            Fixtures.get_good_model_archive_path(),
+            os.path.join(tempdir, "model2.zip"),
         )
-        self.assertGreater(model_size, 0)
-
-    def test_load_local_models(self):
-        with TemporaryDirectory() as tempdir:
-            shutil.copytree(
-                Fixtures.get_good_model_path(), os.path.join(tempdir, "model1")
-            )
-            shutil.copy(
-                Fixtures.get_good_model_archive_path(),
-                os.path.join(tempdir, "model2.zip"),
-            )
-
-            self.model_manager.load_local_models(tempdir)
-            self.assertEqual(len(self.model_manager.loaded_models), 2)
-            self.assertIn("model1", self.model_manager.loaded_models.keys())
-            self.assertIn("model2.zip", self.model_manager.loaded_models.keys())
-            self.assertNotIn(
-                "model-does-not-exist.zip", self.model_manager.loaded_models.keys()
-            )
-
-    def test_model_manager_loads_local_models_on_init(self):
-        with TemporaryDirectory() as tempdir:
-            shutil.copytree(
-                Fixtures.get_good_model_path(), os.path.join(tempdir, "model1")
-            )
-            shutil.copy(
-                Fixtures.get_good_model_archive_path(),
-                os.path.join(tempdir, "model2.zip"),
-            )
-            ModelManager._ModelManager__instance = None
-            with temp_config(
-                {"runtime": {"local_models_dir": tempdir}}, merge_strategy="merge"
-            ):
-                self.model_manager = ModelManager()
-
-                self.assertEqual(len(self.model_manager.loaded_models), 2)
-                self.assertIn("model1", self.model_manager.loaded_models.keys())
-                self.assertIn("model2.zip", self.model_manager.loaded_models.keys())
-                self.assertNotIn(
-                    "model-does-not-exist.zip", self.model_manager.loaded_models.keys()
-                )
 
-    def test_model_manager_raises_if_all_local_models_fail_to_load(self):
-        with TemporaryDirectory() as tempdir:
-            shutil.copy(
-                Fixtures.get_bad_model_archive_path(), os.path.join(tempdir, "model1")
-            )
-            shutil.copy(
-                Fixtures.get_invalid_model_archive_path(),
-                os.path.join(tempdir, "model2.zip"),
-            )
-            with self.assertRaises(CaikitRuntimeException) as ctx:
-                self.model_manager.load_local_models(tempdir)
-            self.assertEqual(grpc.StatusCode.INTERNAL, ctx.exception.status_code)
-
-    def test_load_model_error_response(self):
-        """Test load model's model does not exist when the loader throws"""
-        with self.assertRaises(CaikitRuntimeException) as context:
-            self.model_manager.load_model(
-                model_id=random_test_id(),
-                local_model_path=Fixtures().get_invalid_model_archive_path(),
-                model_type="categories_esa",
-            )
+        with temp_local_models_dir(tempdir):
+            MODEL_MANAGER.sync_local_models(wait=True)
+        assert len(MODEL_MANAGER.loaded_models) == 2
+        assert "model1" in MODEL_MANAGER.loaded_models.keys()
+        assert "model2.zip" in MODEL_MANAGER.loaded_models.keys()
+        assert "model-does-not-exist.zip" not in MODEL_MANAGER.loaded_models.keys()
 
-        self.assertEqual(context.exception.status_code, grpc.StatusCode.NOT_FOUND)
-        self.assertEqual(len(self.model_manager.loaded_models), 0)
 
-    def test_load_model_map_insertion(self):
-        """Test if loaded model is correctly added to map storing model data"""
-        model = random_test_id()
-        self.model_manager.load_model(
-            model_id=model,
-            local_model_path=Fixtures.get_good_model_path(),
-            model_type=Fixtures.get_good_model_type(),
+def test_model_manager_loads_local_models_on_init():
+    with TemporaryDirectory() as tempdir:
+        shutil.copytree(Fixtures.get_good_model_path(), os.path.join(tempdir, "model1"))
+        shutil.copy(
+            Fixtures.get_good_model_archive_path(),
+            os.path.join(tempdir, "model2.zip"),
         )
-        self.assertGreater(self.model_manager.loaded_models[model].size(), 0)
-
-    def test_load_model_count(self):
-        """Test if multiple loaded models are added to map storing model data"""
-        self.model_manager.load_model(
+        ModelManager._ModelManager__instance = None
+        with temp_config(
+            {"runtime": {"local_models_dir": tempdir}}, merge_strategy="merge"
+        ):
+            MODEL_MANAGER = ModelManager()
+
+            assert len(MODEL_MANAGER.loaded_models) == 2
+            assert "model1" in MODEL_MANAGER.loaded_models.keys()
+            assert "model2.zip" in MODEL_MANAGER.loaded_models.keys()
+            assert "model-does-not-exist.zip" not in MODEL_MANAGER.loaded_models.keys()
+
+
+def test_load_model_error_response():
+    """Test load model's model does not exist when the loader throws"""
+    with pytest.raises(CaikitRuntimeException) as context:
+        MODEL_MANAGER.load_model(
             model_id=random_test_id(),
-            local_model_path=Fixtures.get_good_model_path(),
-            model_type=Fixtures.get_good_model_type(),
-        )
-        self.model_manager.load_model(
-            model_id=random_test_id(),
-            local_model_path=Fixtures.get_good_model_path(),
-            model_type=Fixtures.get_good_model_type(),
-        )
-        self.assertEqual(len(self.model_manager.loaded_models), 2)
-
-    def test_unload_model_ok_response(self):
-        """Test to make sure that given a loaded model ID, the model manager is able to correctly
-        unload a model, giving a nonzero model size back."""
-        model_id = "happy_unload_test"
-        self.model_manager.load_model(
-            model_id=model_id,
-            local_model_path=Fixtures.get_good_model_path(),
-            model_type=Fixtures.get_good_model_type(),
-        )
-        model_size = self.model_manager.unload_model(model_id=model_id)
-        self.assertTrue(model_size >= 0)
-
-    def test_unload_model_count(self):
-        """Test if unloaded models are deleted from loaded models map"""
-        id_1 = "test"
-        id_2 = "test2"
-        # Load models from COS
-        self.model_manager.load_model(
-            model_id=id_1,
-            local_model_path=Fixtures.get_good_model_path(),
-            model_type=Fixtures.get_good_model_type(),
+            local_model_path=Fixtures().get_invalid_model_archive_path(),
+            model_type="categories_esa",
+            wait=True,
         )
-        self.model_manager.load_model(
-            model_id=id_2,
-            local_model_path=Fixtures.get_good_model_path(),
-            model_type=Fixtures.get_good_model_type(),
-        )
-        # Unload one of the models, and make sure it was properly removed
-        self.model_manager.unload_model(model_id=id_2)
-        self.assertEqual(len(self.model_manager.loaded_models), 1)
-        self.assertIn(id_1, self.model_manager.loaded_models.keys())
-
-    def test_unload_all_models_count(self):
-        """Test if unload all models deletes every model from loaded models map"""
-        id_1 = "test"
-        id_2 = "test2"
-        self.model_manager.load_model(
-            model_id=id_1,
-            local_model_path=Fixtures.get_good_model_path(),
-            model_type=Fixtures.get_good_model_type(),
-        )
-        self.model_manager.load_model(
-            model_id=id_2,
-            local_model_path=Fixtures.get_good_model_path(),
-            model_type=Fixtures.get_good_model_type(),
+
+    assert context.value.status_code == grpc.StatusCode.NOT_FOUND
+    assert len(MODEL_MANAGER.loaded_models) == 0
+
+
+def test_load_model_map_insertion():
+    """Test if loaded model is correctly added to map storing model data"""
+    model = random_test_id()
+    MODEL_MANAGER.load_model(
+        model_id=model,
+        local_model_path=Fixtures.get_good_model_path(),
+        model_type=Fixtures.get_good_model_type(),
+    )
+    assert MODEL_MANAGER.loaded_models[model].size() > 0
+
+
+def test_load_model_count():
+    """Test if multiple loaded models are added to map storing model data"""
+    MODEL_MANAGER.load_model(
+        model_id=random_test_id(),
+        local_model_path=Fixtures.get_good_model_path(),
+        model_type=Fixtures.get_good_model_type(),
+    )
+    MODEL_MANAGER.load_model(
+        model_id=random_test_id(),
+        local_model_path=Fixtures.get_good_model_path(),
+        model_type=Fixtures.get_good_model_type(),
+    )
+    assert len(MODEL_MANAGER.loaded_models) == 2
+
+
+def test_unload_model_ok_response():
+    """Test to make sure that given a loaded model ID, the model manager is able to correctly
+    unload a model, giving a nonzero model size back."""
+    model_id = "happy_unload_test"
+    MODEL_MANAGER.load_model(
+        model_id=model_id,
+        local_model_path=Fixtures.get_good_model_path(),
+        model_type=Fixtures.get_good_model_type(),
+    )
+    model_size = MODEL_MANAGER.unload_model(model_id=model_id)
+    assert model_size >= 0
+
+
+def test_unload_model_count():
+    """Test if unloaded models are deleted from loaded models map"""
+    id_1 = "test"
+    id_2 = "test2"
+    # Load models from COS
+    MODEL_MANAGER.load_model(
+        model_id=id_1,
+        local_model_path=Fixtures.get_good_model_path(),
+        model_type=Fixtures.get_good_model_type(),
+    )
+    MODEL_MANAGER.load_model(
+        model_id=id_2,
+        local_model_path=Fixtures.get_good_model_path(),
+        model_type=Fixtures.get_good_model_type(),
+    )
+    # Unload one of the models, and make sure it was properly removed
+    MODEL_MANAGER.unload_model(model_id=id_2)
+    assert len(MODEL_MANAGER.loaded_models) == 1
+    assert id_1 in MODEL_MANAGER.loaded_models.keys()
+
+
+def test_unload_all_models_count():
+    """Test if unload all models deletes every model from loaded models map"""
+    id_1 = "test"
+    id_2 = "test2"
+    MODEL_MANAGER.load_model(
+        model_id=id_1,
+        local_model_path=Fixtures.get_good_model_path(),
+        model_type=Fixtures.get_good_model_type(),
+    )
+    MODEL_MANAGER.load_model(
+        model_id=id_2,
+        local_model_path=Fixtures.get_good_model_path(),
+        model_type=Fixtures.get_good_model_type(),
+    )
+    MODEL_MANAGER.unload_all_models()
+    assert len(MODEL_MANAGER.loaded_models) == 0
+
+
+def test_unload_model_not_loaded_response():
+    """Test unload model for model not loaded does NOT throw an error"""
+    MODEL_MANAGER.unload_model(model_id=random_test_id())
+
+
+def test_retrieve_model_returns_loaded_model():
+    """Test that a loaded model can be retrieved"""
+    model_id = random_test_id()
+    MODEL_MANAGER.load_model(
+        model_id,
+        local_model_path=Fixtures.get_good_model_path(),
+        model_type=Fixtures.get_good_model_type(),  # eventually we'd like to be determining the type from the model itself...
+    )
+    model = MODEL_MANAGER.retrieve_model(model_id)
+    assert isinstance(model, ModuleBase)
+    assert len(MODEL_MANAGER.loaded_models) == 1
+
+
+def test_retrieve_model_raises_error_for_not_found_model():
+    """Test that gRPC NOT_FOUND exception raised when non-existent model retrieved"""
+    with pytest.raises(CaikitRuntimeException) as context:
+        MODEL_MANAGER.retrieve_model("not-found")
+    assert context.value.status_code == grpc.StatusCode.NOT_FOUND
+
+
+def test_model_size_ok_response():
+    """Test if loaded model correctly returns model size"""
+    model = random_test_id()
+    MODEL_MANAGER.load_model(
+        model_id=model,
+        local_model_path=Fixtures.get_good_model_path(),
+        model_type=Fixtures.get_good_model_type(),
+    )
+    assert MODEL_MANAGER.get_model_size(model) > 0
+
+
+def test_model_size_error_not_found_response():
+    """Test model size's model does not exist error response"""
+    with pytest.raises(CaikitRuntimeException) as context:
+        MODEL_MANAGER.get_model_size("no_exist_model")
+    assert context.value.status_code == grpc.StatusCode.NOT_FOUND
+
+
+def test_model_size_error_none_not_found_response():
+    """Test model size's model is None error response"""
+    with pytest.raises(CaikitRuntimeException) as context:
+        MODEL_MANAGER.get_model_size(None)
+    assert context.value.status_code == grpc.StatusCode.NOT_FOUND
+
+
+def test_estimate_model_size_ok_response_on_loaded_model():
+    """Test if loaded model correctly returns model size"""
+    MODEL_MANAGER.load_model(
+        model_id=random_test_id(),
+        local_model_path=Fixtures.get_good_model_path(),
+        model_type=Fixtures.get_good_model_type(),
+    )
+    assert (
+        MODEL_MANAGER.estimate_model_size(
+            "test", Fixtures.get_good_model_path(), Fixtures.get_good_model_type()
         )
-        self.model_manager.unload_all_models()
-        self.assertEqual(len(self.model_manager.loaded_models), 0)
+        > 0
+    )
 
-    def test_unload_model_not_loaded_response(self):
-        """Test unload model for model not loaded does NOT throw an error"""
-        try:
-            self.model_manager.unload_model(model_id=random_test_id())
-        except Exception:
-            self.fail("Unload for a model that does not exist threw an error!")
 
-    # TODO: If this is refactored to pytest, the loaded_model_id fixture can be used
-    def test_retrieve_model_returns_loaded_model(self):
-        """Test that a loaded model can be retrieved"""
-        model_id = random_test_id()
-        Fixtures.load_model(
-            model_id=model_id,
-            local_model_path=Fixtures.get_good_model_path(),
-            model_type=Fixtures.get_good_model_type(),
-        )
-        model = self.model_manager.retrieve_model(model_id)
-        self.assertIsInstance(model, ModuleBase)
-        self.assertEqual(len(self.model_manager.loaded_models), 1)
-
-    def test_retrieve_model_raises_error_for_not_found_model(self):
-        """Test that gRPC NOT_FOUND exception raised when non-existent model retrieved"""
-        with self.assertRaises(CaikitRuntimeException) as context:
-            self.model_manager.retrieve_model("not-found")
-        self.assertEqual(context.exception.status_code, grpc.StatusCode.NOT_FOUND)
-
-    def test_model_size_ok_response(self):
-        """Test if loaded model correctly returns model size"""
-        model = random_test_id()
-        self.model_manager.load_model(
-            model_id=model,
-            local_model_path=Fixtures.get_good_model_path(),
-            model_type=Fixtures.get_good_model_type(),
-        )
-        self.assertTrue(self.model_manager.get_model_size(model) > 0)
-
-    def test_model_size_error_not_found_response(self):
-        """Test model size's model does not exist error response"""
-        with self.assertRaises(CaikitRuntimeException) as context:
-            self.model_manager.get_model_size("no_exist_model")
-        self.assertEqual(context.exception.status_code, grpc.StatusCode.NOT_FOUND)
-
-    def test_model_size_error_none_not_found_response(self):
-        """Test model size's model is None error response"""
-        with self.assertRaises(CaikitRuntimeException) as context:
-            self.model_manager.get_model_size(None)
-        self.assertEqual(context.exception.status_code, grpc.StatusCode.NOT_FOUND)
-
-    def test_estimate_model_size_ok_response_on_loaded_model(self):
-        """Test if loaded model correctly returns model size"""
-        self.model_manager.load_model(
-            model_id=random_test_id(),
-            local_model_path=Fixtures.get_good_model_path(),
-            model_type=Fixtures.get_good_model_type(),
-        )
-        self.assertTrue(
-            self.model_manager.estimate_model_size(
-                "test", Fixtures.get_good_model_path(), Fixtures.get_good_model_type()
-            )
-            > 0
+def test_estimate_model_size_ok_response_on_nonloaded_model():
+    """Test if a model that's not loaded, correctly returns predicted model size"""
+    assert (
+        MODEL_MANAGER.estimate_model_size(
+            "test", Fixtures.get_good_model_path(), Fixtures.get_good_model_type()
         )
+        > 0
+    )
 
-    def test_estimate_model_size_ok_response_on_nonloaded_model(self):
-        """Test if a model that's not loaded, correctly returns predicted model size"""
-        self.assertTrue(
-            self.model_manager.estimate_model_size(
-                "test", Fixtures.get_good_model_path(), Fixtures.get_good_model_type()
-            )
-            > 0
-        )
 
-    def test_estimate_model_size_by_type(self):
-        """Test that a model's size is estimated differently based on its type"""
-        config = get_config().inference_plugin.model_mesh
-        self.assertTrue(Fixtures.get_good_model_type() in config.model_size_multipliers)
+def test_estimate_model_size_by_type():
+    """Test that a model's size is estimated differently based on its type"""
+    config = get_config().inference_plugin.model_mesh
+    assert Fixtures.get_good_model_type() in config.model_size_multipliers
 
-        typed_model_size = self.model_manager.estimate_model_size(
-            "test", Fixtures.get_good_model_path(), Fixtures.get_good_model_type()
-        )
-        untyped_model_size = self.model_manager.estimate_model_size(
-            "test", Fixtures.get_good_model_path(), "test-not-a-model-type"
-        )
+    typed_model_size = MODEL_MANAGER.estimate_model_size(
+        "test", Fixtures.get_good_model_path(), Fixtures.get_good_model_type()
+    )
+    untyped_model_size = MODEL_MANAGER.estimate_model_size(
+        "test", Fixtures.get_good_model_path(), "test-not-a-model-type"
+    )
 
-        self.assertTrue(typed_model_size > 0)
-        self.assertTrue(untyped_model_size > 0)
+    assert typed_model_size > 0
+    assert untyped_model_size > 0
+    assert typed_model_size != untyped_model_size
 
-        self.assertNotAlmostEqual(typed_model_size, untyped_model_size)
 
-    def test_estimate_model_size_error_not_found_response(self):
-        """Test if error in predict model size on unknown model path"""
-        with self.assertRaises(CaikitRuntimeException) as context:
-            self.model_manager.estimate_model_size(
-                model_id=random_test_id(),
-                local_model_path="no_exist.zip",
-                model_type="categories_esa",
+def test_estimate_model_size_error_not_found_response():
+    """Test if error in predict model size on unknown model path"""
+    with pytest.raises(CaikitRuntimeException) as context:
+        MODEL_MANAGER.estimate_model_size(
+            model_id=random_test_id(),
+            local_model_path="no_exist.zip",
+            model_type="categories_esa",
+        )
+    assert context.value.status_code == grpc.StatusCode.NOT_FOUND
+
+
+def test_model_manager_replicas_with_disk_caching(good_model_path):
+    """Test that multiple model manager instances can co-exist and share a set
+    of models when local_models_cache is used.
+
+    This test simulates running concurrent replicas of caikit.runtime that are
+    logically independent from one another.
+    """
+    # NOTE: This test requires that the ModelManager class not be a singleton.
+    #   To accomplish this, the singleton instance is temporarily removed.
+    with TemporaryDirectory() as cache_dir:
+        with non_singleton_model_managers(
+            2,
+            {
+                "runtime": {
+                    "local_models_dir": cache_dir,
+                    "lazy_load_local_models": True,
+                    "lazy_load_poll_period_seconds": 0,
+                },
+            },
+            "merge",
+        ) as managers:
+            # Make two standalone manager instances
+            manager_one, manager_two = managers
+            assert manager_two is not manager_one
+            assert manager_one._local_models_dir == cache_dir
+            assert manager_two._local_models_dir == cache_dir
+
+            # Trying to retrieve a model that is not loaded in either and is
+            # not saved in the cache dir should yield NOT_FOUND
+            model_id = random_test_id()
+            with pytest.raises(CaikitRuntimeException) as context:
+                manager_one.retrieve_model(model_id)
+            assert context.value.status_code == grpc.StatusCode.NOT_FOUND
+            with pytest.raises(CaikitRuntimeException) as context:
+                manager_two.retrieve_model(model_id)
+            assert context.value.status_code == grpc.StatusCode.NOT_FOUND
+
+            # Place a saved model into the cache dir manually
+            model_one_path = os.path.join(cache_dir, model_id)
+            shutil.copytree(good_model_path, model_one_path)
+
+            # Retrieve the model by ID and ensure it's available now
+            model_one = manager_one.retrieve_model(model_id)
+            assert isinstance(model_one, ModuleBase)
+            assert (
+                manager_one.loaded_models.get(model_id)
+                and manager_one.loaded_models[model_id].model() is model_one
             )
-        self.assertEqual(context.exception.status_code, grpc.StatusCode.NOT_FOUND)
-
-    # ****************************** Unit Tests ****************************** #
-    # These tests patch in mocks for the manager's dependencies, to test its code in isolation
 
-    def test_load_model(self):
-        """Test to make sure that given valid input, the model manager gives a happy response
-        when we tried to load in a model (model size > 0 or 0 if the model size will be computed
-        at a later time)."""
-        mock_loader = MagicMock()
-        mock_sizer = MagicMock()
+            # Make sure the same model is available on the second manager
+            assert manager_two.loaded_models.get(model_id) is None
+            assert manager_two.retrieve_model(model_id)
+            assert manager_two.loaded_models.get(model_id)
+
+            # Remove the model from the local dir and trigger both managers to
+            # sync their local model dirs to make sure it gets removed
+            shutil.rmtree(model_one_path)
+            manager_one.sync_local_models(wait=True)
+            manager_two.sync_local_models(wait=True)
+            with pytest.raises(CaikitRuntimeException) as context:
+                manager_one.retrieve_model(model_id)
+            with pytest.raises(CaikitRuntimeException) as context:
+                manager_two.retrieve_model(model_id)
+
+
+def test_model_manager_disk_caching_periodic_sync(good_model_path):
+    """Make sure that when using disk caching, the manager periodically syncs
+    its loaded models based on their presence in the cache
+    """
+    purge_period = 0.001
+    with TemporaryDirectory() as cache_dir:
+        with non_singleton_model_managers(
+            2,
+            {
+                "runtime": {
+                    "local_models_dir": cache_dir,
+                    "lazy_load_local_models": True,
+                    "lazy_load_poll_period_seconds": purge_period,
+                },
+            },
+            "merge",
+        ) as managers:
+            manager_one, manager_two = managers
+
+            # Load the model on the first manager and get it loaded into the
+            # the first by fetching it
+            model_id = random_test_id()
+            model_cache_path = os.path.join(cache_dir, model_id)
+            assert not os.path.exists(model_cache_path)
+            shutil.copytree(good_model_path, model_cache_path)
+            assert manager_one.retrieve_model(model_id)
+            assert model_id in manager_one.loaded_models
+
+            # Wait for the sync period and make sure it's available in the
+            # second manager
+            start = time.time()
+            mgr_two_loaded = False
+            while (time.time() - start) < (purge_period * 1000):
+                time.sleep(purge_period)
+                if model_id in manager_two.loaded_models:
+                    mgr_two_loaded = True
+                    break
+            assert mgr_two_loaded
+
+            # Remove the model from disk and wait to ensure that the model gets
+            # unloaded
+            shutil.rmtree(model_cache_path)
+            start = time.time()
+            mgr_one_unloaded = False
+            mgr_two_unloaded = False
+            while (time.time() - start) < (purge_period * 1000):
+                if model_id not in manager_one.loaded_models:
+                    mgr_one_unloaded = True
+                if model_id not in manager_two.loaded_models:
+                    mgr_two_unloaded = True
+                if mgr_one_unloaded and mgr_two_unloaded:
+                    break
+            assert mgr_one_unloaded and mgr_two_unloaded
+
+
+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
+    """
+    with TemporaryDirectory() as tempdir:
+        cache_dir = os.path.join(tempdir, "cache_dir")
+        os.makedirs(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.001,
+                },
+            },
+            "merge",
+        ) as managers:
+            manager = managers[0]
+
+            # Make sure the timer started
+            assert manager._lazy_sync_timer is not None
+
+            # Delete the cache dir and force a sync
+            shutil.rmtree(cache_dir)
+            while True:
+                try:
+                    os.listdir(cache_dir)
+                except FileNotFoundError:
+                    break
+            manager.sync_local_models(wait=True)
+
+            # Make sure the timer is removed
+            assert manager._lazy_sync_timer is None
+
+
+def test_load_local_model_deleted_dir():
+    """Make sure bad models in local_models_dir at boot don't cause exceptions"""
+    with TemporaryDirectory() as cache_dir:
         model_id = random_test_id()
-        expected_model_size = 1234
+        model_cache_path = os.path.join(cache_dir, model_id)
+        os.makedirs(model_cache_path)
+        with non_singleton_model_managers(
+            1,
+            {
+                "runtime": {
+                    "local_models_dir": cache_dir,
+                    "lazy_load_local_models": False,
+                },
+            },
+            "merge",
+        ) as managers:
+            manager = managers[0]
+            assert not manager.loaded_models
+
+
+# ****************************** Unit Tests ****************************** #
+# These tests patch in mocks for the manager's dependencies, to test its code in isolation
+
+
+def test_load_model():
+    """Test to make sure that given valid input, the model manager gives a happy response
+    when we tried to load in a model (model size > 0 or 0 if the model size will be computed
+    at a later time)."""
+    mock_loader = MagicMock()
+    mock_sizer = MagicMock()
+    model_id = random_test_id()
+    expected_model_size = 1234
+
+    with patch.object(MODEL_MANAGER, "model_loader", mock_loader):
+        with patch.object(MODEL_MANAGER, "model_sizer", mock_sizer):
+            loaded_model = LoadedModel()
+            loaded_model._model = "something"
+            mock_loader.load_model.return_value = loaded_model
+            mock_sizer.get_model_size.return_value = expected_model_size
+
+            model_size = MODEL_MANAGER.load_model(
+                model_id, ANY_MODEL_PATH, ANY_MODEL_TYPE
+            )
+            assert expected_model_size == model_size
+            mock_loader.load_model.assert_called_once()
+            call_args = mock_loader.load_model.call_args
+            assert call_args.args == (
+                model_id,
+                ANY_MODEL_PATH,
+                ANY_MODEL_TYPE,
+            )
+            assert call_args.kwargs["aborter"] is None
+            assert "fail_callback" in call_args.kwargs
+            mock_sizer.get_model_size.assert_called_once_with(
+                model_id, ANY_MODEL_PATH, ANY_MODEL_TYPE
+            )
 
-        with patch.object(self.model_manager, "model_loader", mock_loader):
-            with patch.object(self.model_manager, "model_sizer", mock_sizer):
-                mock_loader.load_model.return_value = LoadedModel.Builder().build()
-                mock_sizer.get_model_size.return_value = expected_model_size
 
-                model_size = self.model_manager.load_model(
-                    model_id, ANY_MODEL_PATH, ANY_MODEL_TYPE
-                )
-                self.assertEqual(expected_model_size, model_size)
-                mock_loader.load_model.assert_called_once_with(
-                    model_id, ANY_MODEL_PATH, ANY_MODEL_TYPE
-                )
-                mock_sizer.get_model_size.assert_called_once_with(
-                    model_id, ANY_MODEL_PATH, ANY_MODEL_TYPE
-                )
+def test_load_model_throws_if_the_model_loader_throws():
+    """Test load model's model does not exist when the loader throws"""
+    mock_loader = MagicMock()
+    with patch.object(MODEL_MANAGER, "model_loader", mock_loader):
+        mock_loader.load_model.side_effect = CaikitRuntimeException(
+            grpc.StatusCode.NOT_FOUND, "test any not found exception"
+        )
 
-    def test_load_model_throws_if_the_model_loader_throws(self):
-        """Test load model's model does not exist when the loader throws"""
-        mock_loader = MagicMock()
-        with patch.object(self.model_manager, "model_loader", mock_loader):
-            mock_loader.load_model.side_effect = CaikitRuntimeException(
-                grpc.StatusCode.NOT_FOUND, "test any not found exception"
+        with pytest.raises(CaikitRuntimeException) as context:
+            MODEL_MANAGER.load_model(random_test_id(), ANY_MODEL_PATH, ANY_MODEL_TYPE)
+
+        assert context.value.status_code == grpc.StatusCode.NOT_FOUND
+        assert len(MODEL_MANAGER.loaded_models) == 0
+
+
+def test_retrieve_model_returns_the_module_from_the_model_loader():
+    """Test that a loaded model can be retrieved"""
+    model_id = random_test_id()
+    expected_module = "this is definitely a stub module"
+    mock_sizer = MagicMock()
+    mock_loader = MagicMock()
+
+    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 = 1
+            model_future = Future()
+            model_future.result = lambda *_, **__: expected_module
+            mock_loader.load_model.return_value = (
+                LoadedModel.Builder()
+                .model_future(model_future)
+                .id("foo")
+                .type("bar")
+                .build()
+            )
+            MODEL_MANAGER.load_model(model_id, ANY_MODEL_PATH, ANY_MODEL_TYPE)
+
+            model = MODEL_MANAGER.retrieve_model(model_id)
+            assert expected_module == model
+
+
+def test_unload_partially_loaded():
+    """Make sure that when unloading a model that is still loading, the load
+    completes correctly
+    """
+    # Set up a "slow load" that we can use to ensure that the loading has
+    # completed successfully
+    slow_loader = SlowLoader()
+    pool = ThreadPoolExecutor()
+    model_future = pool.submit(slow_loader.load)
+
+    model_id = random_test_id()
+    mock_sizer = MagicMock()
+    mock_loader = MagicMock()
+    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 = 1
+            mock_loader.load_model.return_value = (
+                LoadedModel.Builder()
+                .model_future(model_future)
+                .id("foo")
+                .type("bar")
+                .build()
+            )
+            MODEL_MANAGER.load_model(
+                model_id, ANY_MODEL_PATH, ANY_MODEL_TYPE, wait=False
             )
 
-            with self.assertRaises(CaikitRuntimeException) as context:
-                self.model_manager.load_model(
-                    random_test_id(), ANY_MODEL_PATH, ANY_MODEL_TYPE
-                )
-
-            self.assertEqual(context.exception.status_code, grpc.StatusCode.NOT_FOUND)
-            self.assertEqual(len(self.model_manager.loaded_models), 0)
+            # Start the unload assertion that will block on loading
+            unload_future = pool.submit(MODEL_MANAGER.unload_model, model_id)
+
+            # Unblock the load
+            assert not slow_loader.done_loading()
+            slow_loader.unblock_load()
+
+            # Make sure unload completes and the model finished loading
+            unload_future.result()
+            assert slow_loader.done_loading()
+
+
+def test_unload_unexpected_error_loaded():
+    """Make sure that when unloading a model that is still loading, errors when
+    waiting for the load to complete are handled
+    """
+    # Set up a "slow load" that we can use to ensure that the loading has
+    # completed successfully
+    slow_loader = SlowLoader(RuntimeError("yikes"))
+    pool = ThreadPoolExecutor()
+    model_future = pool.submit(slow_loader.load)
+
+    model_id = random_test_id()
+    mock_sizer = MagicMock()
+    mock_loader = MagicMock()
+    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 = 1
+            mock_loader.load_model.return_value = (
+                LoadedModel.Builder()
+                .model_future(model_future)
+                .id("foo")
+                .type("bar")
+                .build()
+            )
+            MODEL_MANAGER.load_model(
+                model_id, ANY_MODEL_PATH, ANY_MODEL_TYPE, wait=False
+            )
 
-    def test_retrieve_model_returns_the_module_from_the_model_loader(self):
-        """Test that a loaded model can be retrieved"""
-        model_id = random_test_id()
-        expected_module = "this is definitely a stub module"
-        mock_sizer = MagicMock()
-        mock_loader = MagicMock()
-
-        with patch.object(self.model_manager, "model_loader", mock_loader):
-            with patch.object(self.model_manager, "model_sizer", mock_sizer):
-                mock_sizer.get_model_size.return_value = 1
-                mock_loader.load_model.return_value = (
-                    LoadedModel.Builder().module(expected_module).build()
+            # Start the unload assertion that will block on loading
+            unload_future = pool.submit(MODEL_MANAGER.unload_model, model_id)
+
+            # Unblock the load
+            assert not slow_loader.done_loading()
+            slow_loader.unblock_load()
+
+            # Make sure unload completes and the model finished loading
+            with pytest.raises(CaikitRuntimeException) as context:
+                unload_future.result()
+            assert context.value.status_code == grpc.StatusCode.INTERNAL
+            assert slow_loader.done_loading()
+
+
+def test_reload_partially_loaded():
+    """Make sure that attempting to reload a model that has already started to
+    load simply returns the existing LoadedModel
+    """
+    # Set up a "slow load" that we can use to ensure that the loading has
+    # completed successfully
+    slow_loader = SlowLoader()
+    pool = ThreadPoolExecutor()
+    model_future = pool.submit(slow_loader.load)
+
+    model_id = random_test_id()
+    mock_sizer = MagicMock()
+    mock_loader = MagicMock()
+    special_model_size = 123321
+    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 = (
+                LoadedModel.Builder()
+                .model_future(model_future)
+                .id("foo")
+                .type("bar")
+                .build()
+            )
+            model_size = MODEL_MANAGER.load_model(
+                model_id, ANY_MODEL_PATH, ANY_MODEL_TYPE, wait=False
+            )
+            assert model_size == special_model_size
+            assert (
+                MODEL_MANAGER.load_model(
+                    model_id, ANY_MODEL_PATH, ANY_MODEL_TYPE, wait=False
                 )
-                self.model_manager.load_model(model_id, ANY_MODEL_PATH, ANY_MODEL_TYPE)
+                == special_model_size
+            )
 
-                model = self.model_manager.retrieve_model(model_id)
-                self.assertEqual(expected_module, model)
+            # Unblock the load
+            assert not slow_loader.done_loading()
+            slow_loader.unblock_load()
+            model_future.result()
+            assert slow_loader.done_loading()
+
+
+def test_get_model_size_returns_size_from_model_sizer():
+    """Test that loading a model stores the size from the ModelSizer"""
+    mock_loader = MagicMock()
+    mock_sizer = MagicMock()
+    expected_model_size = 1234
+    model_id = random_test_id()
+
+    with patch.object(MODEL_MANAGER, "model_loader", mock_loader):
+        with patch.object(MODEL_MANAGER, "model_sizer", mock_sizer):
+            loaded_model = LoadedModel()
+            loaded_model._model = "something"
+            mock_loader.load_model.return_value = loaded_model
+            mock_sizer.get_model_size.return_value = expected_model_size
 
-    def test_get_model_size_returns_size_from_model_sizer(self):
-        """Test that loading a model stores the size from the ModelSizer"""
-        mock_loader = MagicMock()
-        mock_sizer = MagicMock()
-        expected_model_size = 1234
-        model_id = random_test_id()
+            MODEL_MANAGER.load_model(model_id, ANY_MODEL_PATH, ANY_MODEL_TYPE)
 
-        with patch.object(self.model_manager, "model_loader", mock_loader):
-            with patch.object(self.model_manager, "model_sizer", mock_sizer):
-                mock_loader.load_model.return_value = LoadedModel.Builder().build()
-                mock_sizer.get_model_size.return_value = expected_model_size
+            model_size = MODEL_MANAGER.get_model_size(model_id)
+            assert expected_model_size == model_size
 
-                self.model_manager.load_model(model_id, ANY_MODEL_PATH, ANY_MODEL_TYPE)
 
-                model_size = self.model_manager.get_model_size(model_id)
-                self.assertEqual(expected_model_size, model_size)
+def test_estimate_model_size_returns_size_from_model_sizer():
+    """Test that estimating a model size uses the ModelSizer"""
+    mock_sizer = MagicMock()
+    expected_model_size = 5678
+    model_id = random_test_id()
 
-    def test_estimate_model_size_returns_size_from_model_sizer(self):
-        """Test that estimating a model size uses the ModelSizer"""
-        mock_sizer = MagicMock()
-        expected_model_size = 5678
-        model_id = random_test_id()
+    with patch.object(MODEL_MANAGER, "model_sizer", mock_sizer):
+        mock_sizer.get_model_size.return_value = expected_model_size
+        model_size = MODEL_MANAGER.estimate_model_size(
+            model_id, ANY_MODEL_PATH, ANY_MODEL_TYPE
+        )
+        assert expected_model_size == model_size
 
-        with patch.object(self.model_manager, "model_sizer", mock_sizer):
-            mock_sizer.get_model_size.return_value = expected_model_size
-            model_size = self.model_manager.estimate_model_size(
-                model_id, ANY_MODEL_PATH, ANY_MODEL_TYPE
-            )
-            self.assertEqual(expected_model_size, model_size)
 
-    def test_estimate_model_size_throws_if_model_sizer_throws(self):
-        """Test that estimating a model size uses the ModelSizer"""
-        mock_sizer = MagicMock()
-        model_id = random_test_id()
+def test_estimate_model_size_throws_if_model_sizer_throws():
+    """Test that estimating a model size uses the ModelSizer"""
+    mock_sizer = MagicMock()
+    model_id = random_test_id()
 
-        with patch.object(self.model_manager, "model_sizer", mock_sizer):
-            mock_sizer.get_model_size.side_effect = CaikitRuntimeException(
-                grpc.StatusCode.UNAVAILABLE, "test-any-exception"
-            )
-            with self.assertRaises(CaikitRuntimeException) as context:
-                self.model_manager.estimate_model_size(
-                    model_id, ANY_MODEL_PATH, ANY_MODEL_TYPE
-                )
-            self.assertEqual(context.exception.status_code, grpc.StatusCode.UNAVAILABLE)
+    with patch.object(MODEL_MANAGER, "model_sizer", mock_sizer):
+        mock_sizer.get_model_size.side_effect = CaikitRuntimeException(
+            grpc.StatusCode.UNAVAILABLE, "test-any-exception"
+        )
+        with pytest.raises(CaikitRuntimeException) as context:
+            MODEL_MANAGER.estimate_model_size(model_id, ANY_MODEL_PATH, ANY_MODEL_TYPE)
+        assert context.value.status_code == grpc.StatusCode.UNAVAILABLE
+
+
+def test_periodic_sync_handles_errors():
+    """Test that any exception raised during syncing local models is handled
+    without terminating the polling loop
+    """
+
+    class SecretException(Exception):
+        pass
+
+    with TemporaryDirectory() as cache_dir:
+        with non_singleton_model_managers(
+            1,
+            {
+                "runtime": {
+                    "local_models_dir": cache_dir,
+                    "lazy_load_local_models": True,
+                },
+            },
+            "merge",
+        ) as managers:
+            manager = managers[0]
+            with patch.object(manager, "_local_models_dir_sync") as mock_sync:
+                mock_sync.side_effect = SecretException()
+                assert manager._lazy_sync_timer is not None
+                manager.sync_local_models(True)
+                mock_sync.assert_called_once()
+                assert manager._lazy_sync_timer is not None
diff --git a/tests/runtime/servicers/test_model_runtime_servicer_impl.py b/tests/runtime/servicers/test_model_runtime_servicer_impl.py
index f94bf5a8..d7f7b8d5 100644
--- a/tests/runtime/servicers/test_model_runtime_servicer_impl.py
+++ b/tests/runtime/servicers/test_model_runtime_servicer_impl.py
@@ -90,7 +90,7 @@ class TestModelRuntimeServicerImpl(unittest.TestCase):
         mock_manager = MagicMock()
         started = Event()
 
-        def never_return(*args):
+        def never_return(*args, **kwargs):
             started.set()
             while True:
                 time.sleep(0.01)
diff --git a/tests/runtime/test_grpc_server.py b/tests/runtime/test_grpc_server.py
index a275f93a..66ea83a6 100644
--- a/tests/runtime/test_grpc_server.py
+++ b/tests/runtime/test_grpc_server.py
@@ -828,7 +828,7 @@ def test_canceling_model_loads_causes_exceptions(runtime_grpc_server):
         modelKey="bar",
     )
 
-    def never_return(*args):
+    def never_return(*args, **kwargs):
         request_received.set()
         try:
             while True:
diff --git a/tests/runtime/work_management/test_abortable_action.py b/tests/runtime/work_management/test_abortable_action.py
index eec05785..744e9857 100644
--- a/tests/runtime/work_management/test_abortable_action.py
+++ b/tests/runtime/work_management/test_abortable_action.py
@@ -19,7 +19,7 @@ import unittest
 # Local
 from caikit.runtime.types.aborted_exception import AbortedException
 from caikit.runtime.work_management.abortable_action import AbortableAction
-from caikit.runtime.work_management.call_aborter import CallAborter
+from caikit.runtime.work_management.rpc_aborter import RpcAborter
 from tests.fixtures import Fixtures
 
 
@@ -29,7 +29,7 @@ class TestAbortableAction(unittest.TestCase):
     def setUp(self):
         """This method runs before each test begins to run"""
         self.rpc_context = Fixtures.build_context()
-        self.aborter = CallAborter(self.rpc_context)
+        self.aborter = RpcAborter(self.rpc_context)
 
     def test_it_can_run_a_function(self):
         expected_result = "test-any-result"
diff --git a/tests/runtime/work_management/test_call_aborter.py b/tests/runtime/work_management/test_call_aborter.py
index 3af78402..f1e002cd 100644
--- a/tests/runtime/work_management/test_call_aborter.py
+++ b/tests/runtime/work_management/test_call_aborter.py
@@ -16,11 +16,11 @@ import threading
 import unittest
 
 # Local
-from caikit.runtime.work_management.call_aborter import CallAborter
+from caikit.runtime.work_management.rpc_aborter import RpcAborter
 from tests.fixtures import Fixtures
 
 
-class TestCallAborter(unittest.TestCase):
+class TestRpcAborter(unittest.TestCase):
     """This test suite tests the call aborter utility"""
 
     CHANNEL = None
@@ -29,7 +29,7 @@ class TestCallAborter(unittest.TestCase):
     def test_call_aborter_sets_event(self):
         ctx = Fixtures.build_context("call_aborter_event_party")
         # Create a new Call aborter
-        aborter = CallAborter(ctx)
+        aborter = RpcAborter(ctx)
         # Create a new threading event and add it to the call aborter
         event = threading.Event()
         aborter.add_event(event)
@@ -41,7 +41,7 @@ class TestCallAborter(unittest.TestCase):
     def test_call_aborter_sets_event_added_after_termination(self):
         ctx = Fixtures.build_context("call_aborter_event_party")
         # Create a new call aborter
-        aborter = CallAborter(ctx)
+        aborter = RpcAborter(ctx)
         # Cancel the call before creating the threading event and adding to the aborter
         ctx.cancel()
         event = threading.Event()