Repository: caikit/caikit
Pull request #129 (2023-05-22 21:08:36)

🐛  Load backend arg
This PR is a little bugfix to supply the concrete load backend as a kwarg to `.load` classmethods on alternative backend module implementations.

This is to support users who need to have access to the load backend at initialization time. Previously they used `caikit.get_backend(type)` but this was removed in favor of explicitly setting the load backend. For shared load backends this is fine, but for backend-specific module implementations this deferred knowledge of the backend until post initialization. 

We have some design work to do to fix this gap, but in the meantime a quick extra kwarg to load will unblock us.


Patch:
diff --git a/caikit/core/model_manager.py b/caikit/core/model_manager.py
index c038a662..8ab38208 100644
--- a/caikit/core/model_manager.py
+++ b/caikit/core/model_manager.py
@@ -262,6 +262,7 @@ class ModelManager:
                         loaded_model = module_backend_impl.load(
                             module_path,
                             *args,
+                            load_backend=load_backend,
                             **kwargs,
                         )
                         if loaded_model is not None:


Test patch:
diff --git a/tests/core/test_model_manager.py b/tests/core/test_model_manager.py
index b5bcd516..ab6b6f1a 100644
--- a/tests/core/test_model_manager.py
+++ b/tests/core/test_model_manager.py
@@ -17,11 +17,14 @@ and download and load them.
 
 # Standard
 from contextlib import contextmanager
+from unittest import mock
+from unittest.mock import MagicMock
 import os
 import tempfile
 import uuid
 
 # Local
+from caikit.core.module_backends import module_backend_config
 from caikit.core.module_backends.module_backend_config import (
     configure,
     configured_load_backends,
@@ -405,6 +408,34 @@ def test_preferred_backend_enabled(reset_globals):
         assert isinstance(model, DummyBar)
 
 
+def test_module_backend_instance_is_passed_to_load_classmethod(reset_globals):
+    """When an alternate module implementation is loaded via a backend, the concrete
+    instance of the module backend is passed to .load via the load_backend kwarg.
+    """
+    _, DummyBar = setup_saved_model(MockBackend)
+    with temp_config(
+        {
+            "module_backends": {
+                "load_priority": [{"type": backend_types.MOCK}],
+            }
+        }
+    ):
+        configure()
+        with mock.patch.object(DummyBar, "load", MagicMock()) as mock_load:
+            mock_load.return_value = DummyBar()
+            dummy_model_path = os.path.join(TEST_DATA_PATH, DUMMY_LOCAL_MODEL_NAME)
+            model = caikit.core.load(dummy_model_path)
+
+            load_backends = module_backend_config.configured_load_backends()
+            expected_load_backend = [
+                be for be in load_backends if be.backend_type == backend_types.MOCK
+            ][0]
+
+            mock_load.assert_called_with(
+                dummy_model_path, **{"load_backend": expected_load_backend}
+            )
+
+
 def test_preferred_backend_disabled(reset_globals):
     """Make sure that for a model artifact saved with a local backend loads as
     local even with a preferred_backend when the preferred backend is disabled.
diff --git a/tests/runtime/model_management/test_model_loader.py b/tests/runtime/model_management/test_model_loader.py
index 2a6d995c..4a38fd43 100644
--- a/tests/runtime/model_management/test_model_loader.py
+++ b/tests/runtime/model_management/test_model_loader.py
@@ -265,7 +265,8 @@ def test_load_distributed_impl():
                     )
 
                 @classmethod
-                def load(cls, model_load_path) -> "DistributedGadget":
+                def load(cls, model_load_path, **kwargs) -> "DistributedGadget":
+                    # NOTE: kwargs needed here for load_backend
                     config = ModuleConfig.load(model_load_path)
                     return cls(bar=config.bar)