Repository: caikit/caikit
Pull request #106 (2023-05-12 18:43:19)

🐛  Fix name in shared backends
## Description

This PR walks back an accidental API breaking change in `0.4.0` where we added `name` as a required initializer arg to `BackendBase`. We had decided to remove this since it's not actually necessary and it adds additional awkwardness, but forgot to make the change before merging and cutting `0.4.0`.


Patch:
diff --git a/caikit/core/model_manager.py b/caikit/core/model_manager.py
index c3650af6..e1214158 100644
--- a/caikit/core/model_manager.py
+++ b/caikit/core/model_manager.py
@@ -188,16 +188,17 @@ class ModelManager:
             # a version of the module available for the given backend
             loaded_model = None
             log.debug("Available load backends: %s", configured_load_backends)
-            for load_backend in configured_load_backends:
+            for i, load_backend in enumerate(configured_load_backends):
                 # If this is a shared loader, try loading the model directly
                 if isinstance(load_backend, SharedLoadBackendBase):
                     log.debug("Trying shared backend loader")
                     model = load_backend.load(module_path, *args, **kwargs)
                     if model is not None:
                         log.debug2(
-                            "Successfully loaded %s with loader %s",
+                            "Successfully loaded %s with loader (%d)%s",
                             module_path,
-                            load_backend.name,
+                            i,
+                            load_backend.backend_type,
                         )
                         error.type_check(
                             "<COR76726077E>",
@@ -209,9 +210,10 @@ class ModelManager:
                         model.set_load_backend(load_backend)
                         break
                     log.debug3(
-                        "Could not load %s with loader %s",
+                        "Could not load %s with loader (%d)%s",
                         module_path,
-                        load_backend.name,
+                        i,
+                        load_backend.backend_type,
                     )
 
                 # If this is not a shared loader, look for an implementation of the
diff --git a/caikit/core/module_backend_config.py b/caikit/core/module_backend_config.py
index b8623dae..61672042 100644
--- a/caikit/core/module_backend_config.py
+++ b/caikit/core/module_backend_config.py
@@ -15,7 +15,6 @@
 # Standard
 from typing import List
 import copy
-import threading
 
 # First Party
 import aconfig
@@ -41,38 +40,22 @@ def start_backends() -> None:
         None
     """
     for backend in _CONFIGURED_LOAD_BACKENDS:
-        with _BACKEND_START_LOCKS[backend.name]:
+        with backend.start_lock():
             if not backend.is_started:
                 backend.start()
     for backend in _CONFIGURED_TRAIN_BACKENDS:
-        with _BACKEND_START_LOCKS[backend.name]:
+        with backend.start_lock():
             if not backend.is_started:
                 backend.start()
 
 
-def get_load_backend(backend_name: str) -> BackendBase:
-    """Get the configured instance of the given backend type. If not configured,
-    a ValueError is raised
-    """
-    return _get_registry_backend(backend_name, "load", _CONFIGURED_LOAD_BACKENDS)
-
-
-def get_train_backend(backend_name: str) -> BackendBase:
-    """Get the configured instance of the given backend type. If not configured,
-    a ValueError is raised
-    """
-    return _get_registry_backend(backend_name, "train", _CONFIGURED_TRAIN_BACKENDS)
-
-
 def configured_load_backends() -> List[BackendBase]:
-    """This function returns the mapping of named"""
+    """This function returns the list of configured load backends"""
     return copy.copy(_CONFIGURED_LOAD_BACKENDS)
 
 
 def configured_train_backends() -> List[BackendBase]:
-    """This function exposes the list of configured train backends for downstream
-    checks
-    """
+    """This function returns the list of configured train backends"""
     return copy.copy(_CONFIGURED_TRAIN_BACKENDS)
 
 
@@ -133,21 +116,18 @@ def configure():
                 i,
             )
 
-            backend_name = backend_config.get("name", backend_type)
-            log.debug("Configuring backend [%s]", backend_name)
+            log.debug("Configuring backend (%d)[%s]", i, backend_type)
             backend_instance_config = backend_config.get("config", {})
-            log.debug3("Backend [%s] config: %s", backend_name, backend_instance_config)
+            log.debug3(
+                "Backend (%d)[%s] config: %s", i, backend_type, backend_instance_config
+            )
 
             backend_class = MODULE_BACKEND_CONFIG_FUNCTIONS.get(backend_type)
             error.value_check(
                 "<COR64618509E>",
-                not any(
-                    backend.name == backend_name
-                    and backend.backend_type == backend_type
-                    for backend in registry
-                ),
-                "{}/{} already configured",
-                backend_name,
+                len(registry) == i,
+                "({})[{}] already configured",
+                i,
                 backend_type,
             )
             error.value_check(
@@ -166,17 +146,14 @@ def configure():
                     ),
                 )
 
-            log.debug2("Performing config for [%s]", backend_name)
-            backend_instance = backend_class(backend_name, backend_instance_config)
+            log.debug2("Performing config for (%d)[%s]", i, backend_type)
+            backend_instance = backend_class(backend_instance_config)
 
             # Add configuration to backends as per individual module requirements
-            _configure_backend_overrides(backend_name, backend_instance)
+            _configure_backend_overrides(backend_type, backend_instance)
 
-            # NOTE: Configured backends holds the object of backend classes that are based
-            # on BaseBackend
-            # The start operation of the backend needs to be performed separately
+            # Add the instance to the registry
             registry.append(backend_instance)
-            _BACKEND_START_LOCKS[backend_name] = threading.Lock()
 
         log.debug2("All configured %s backends: %s", registry_name, registry)
 
@@ -187,13 +164,6 @@ def configure():
 _CONFIGURED_LOAD_BACKENDS = []
 _CONFIGURED_TRAIN_BACKENDS = []
 
-# Locks for starting backends
-# NOTE: A single lock is held created for a backend name, even if it is repeated
-#   between train and load. Configuration is a load-time operation, so while
-#   this might be slightly overly locking, it protects against the case where
-#   a train and load backend have overlapping global state.
-_BACKEND_START_LOCKS = {}
-
 
 def _configure_backend_overrides(backend: str, backend_instance: object):
     """Function to go over all the modules registered in the MODULE_BACKEND_REGISTRY
@@ -220,30 +190,3 @@ def _configure_backend_overrides(backend: str, backend_instance: object):
                 log.debug2(
                     f"No backend overrides configured for {module_id} module and {backend} backend"
                 )
-
-
-def _get_registry_backend(
-    backend_name: str, registry_name: str, registry: list
-) -> BackendBase:
-    """Get the configured instance of the given backend type. If not configured,
-    a ValueError is raised
-    """
-    matching_backends = [
-        backend for backend in registry if backend.name == backend_name
-    ]
-    error.value_check(
-        "<COR82987857E>",
-        matching_backends,
-        "Cannot fetch unconfigured {} backend [{}]",
-        registry_name,
-        backend_name,
-    )
-    assert (
-        len(matching_backends) == 1
-    ), "PROGRAMMING ERROR: Duplicate names should be prohibited"
-    backend = matching_backends[0]
-    if not backend.is_started:
-        with _BACKEND_START_LOCKS[backend_name]:
-            if not backend.is_started:
-                backend.start()
-    return backend
diff --git a/caikit/core/module_backends/base.py b/caikit/core/module_backends/base.py
index 0b8b6ad4..d99f41b4 100644
--- a/caikit/core/module_backends/base.py
+++ b/caikit/core/module_backends/base.py
@@ -16,9 +16,14 @@
 """
 
 # Standard
+from contextlib import contextmanager
+from threading import Lock
 from typing import Optional, Type
 import abc
 
+# First Party
+import aconfig
+
 # Local
 from ..module import ModuleBase
 
@@ -26,10 +31,10 @@ from ..module import ModuleBase
 class BackendBase(abc.ABC):
     """Interface for creating configuration setup for backends"""
 
-    def __init__(self, name: str, config=None) -> None:
-        self.name = name
+    def __init__(self, config: Optional[aconfig.Config] = None) -> None:
         self.config = config or {}
         self._started = False
+        self._start_lock = Lock()
 
     @property
     @classmethod
@@ -60,6 +65,11 @@ class BackendBase(abc.ABC):
         """Function to stop a distributed backend. This function
         should set self._started variable to False"""
 
+    @contextmanager
+    def start_lock(self):
+        with self._start_lock:
+            yield
+
 
 class SharedTrainBackendBase(BackendBase, abc.ABC):
     """Interface for a backend that can perform train on any given module


Test patch:
diff --git a/tests/core/helpers.py b/tests/core/helpers.py
index c95bfd02..6ede7fdb 100644
--- a/tests/core/helpers.py
+++ b/tests/core/helpers.py
@@ -36,8 +36,8 @@ from caikit.core.module_backends.base import SharedLoadBackendBase
 class MockBackend(BackendBase):
     backend_type = "MOCK"
 
-    def __init__(self, name, config=...) -> None:
-        super().__init__(name, config)
+    def __init__(self, config=...) -> None:
+        super().__init__(config)
         self._started = False
 
     def start(self):
@@ -64,7 +64,7 @@ class TestLoader(SharedLoadBackendBase):
                 # Don't load in this loader
                 return None
         # use the "Local" loader to actually load the model
-        model = LocalBackend(name="test").load(model_path)
+        model = LocalBackend().load(model_path)
         return model
 
     def register_config(self, config):
diff --git a/tests/core/test_model_manager.py b/tests/core/test_model_manager.py
index 55919eff..7cc65a20 100644
--- a/tests/core/test_model_manager.py
+++ b/tests/core/test_model_manager.py
@@ -17,18 +17,16 @@ and download and load them.
 
 # Standard
 from contextlib import contextmanager
-from typing import Optional
 import os
 import tempfile
 import uuid
 
 # Local
-from caikit.core import ModuleBase
-from caikit.core.module_backend_config import configure
+from caikit.core.module_backend_config import configure, configured_load_backends
 from caikit.core.module_backends import LocalBackend
+from caikit.core.module_backends.base import SharedLoadBackendBase
 
 # Unit Test Infrastructure
-from caikit.core.module_backends.base import SharedLoadBackendBase
 from sample_lib.blocks.sample_task import SampleBlock
 from sample_lib.data_model import SampleTask
 from tests.base import TestCaseBase
@@ -534,7 +532,7 @@ def test_load_with_new_shared_backend(good_model_path, reset_globals):
         {
             "module_backends": {
                 "load_priority": [
-                    {"type": TestLoader.backend_type, "name": loader_name},
+                    {"type": TestLoader.backend_type},
                     {"type": backend_types.LOCAL},
                 ],
             }
@@ -542,7 +540,7 @@ def test_load_with_new_shared_backend(good_model_path, reset_globals):
     ):
         configure()
         model = caikit.core.load(good_model_path)
-        assert model.load_backend.name == loader_name
+        assert model.load_backend.backend_type == TestLoader.backend_type
 
 
 def test_load_with_two_shared_loaders_of_the_same_type(good_model_path, reset_globals):
@@ -553,12 +551,10 @@ def test_load_with_two_shared_loaders_of_the_same_type(good_model_path, reset_gl
                 "load_priority": [
                     {
                         "type": TestLoader.backend_type,
-                        "name": "loader one",
                         "config": {"model_type": "model one"},
                     },
                     {
                         "type": TestLoader.backend_type,
-                        "name": "loader two",
                         "config": {"model_type": "model two"},
                     },
                     {"type": backend_types.LOCAL},
@@ -567,13 +563,15 @@ def test_load_with_two_shared_loaders_of_the_same_type(good_model_path, reset_gl
         }
     ):
         configure()
+        backends = configured_load_backends()
+        assert len(backends) == 3
         # plain model load should use first loader
         model = caikit.core.load(good_model_path)
-        assert model.load_backend.name == "loader one"
+        assert model.load_backend is backends[0]
 
         # model load that fails in the first loader will use the second
         model = caikit.core.load(good_model_path, model_type="model two")
-        assert model.load_backend.name == "loader two"
+        assert model.load_backend is backends[1]
 
 
 def test_load_does_not_read_config_yml_if_loader_does_not_require_it(
diff --git a/tests/core/test_module.py b/tests/core/test_module.py
index d81d8333..2e2f31ee 100644
--- a/tests/core/test_module.py
+++ b/tests/core/test_module.py
@@ -27,7 +27,7 @@ import aconfig
 
 # Local
 from caikit.core import ModuleConfig, module
-from caikit.core.module_backend_config import get_load_backend
+from caikit.core.module_backend_config import configured_load_backends
 from caikit.core.module_backends import backend_types
 from caikit.core.module_type import SUPPORTED_LOAD_BACKENDS_VAR_NAME
 
@@ -78,7 +78,11 @@ def configure_alternate_backend_impl():
     @caikit.core.blocks.block(base_module=DummyFoo, backend_type=backend_types.MOCK)
     class DummyBar(caikit.core.blocks.base.BlockBase):
         def test_fetching_backend(self):
-            return get_load_backend(backend_types.MOCK)
+            return [
+                backend
+                for backend in configured_load_backends()
+                if backend.backend_type == backend_types.MOCK
+            ][0]
 
     return DummyFoo, DummyBar
 
@@ -331,9 +335,8 @@ def test_class_attributes(reset_globals):
     assert DummyBar.BLOCK_CLASS
     assert DummyBar.PRODUCER_ID
 
-    # Fetch without config raises
-    with pytest.raises(ValueError):
-        assert get_load_backend(DummyBar.BACKEND_TYPE)
+    # No configured modules before configure
+    assert not configured_load_backends()
 
     # Configure and make sure it can be fetched by the class
     with temp_config(
diff --git a/tests/core/test_module_backend_config.py b/tests/core/test_module_backend_config.py
index 568330d5..dcad0b25 100644
--- a/tests/core/test_module_backend_config.py
+++ b/tests/core/test_module_backend_config.py
@@ -22,8 +22,6 @@ from caikit.core.module_backend_config import (
     configure,
     configured_load_backends,
     configured_train_backends,
-    get_load_backend,
-    get_train_backend,
     start_backends,
 )
 from caikit.core.module_backends import backend_types
@@ -35,6 +33,26 @@ from tests.core.helpers import *
 
 foo_cfg = {"mock": 1}
 
+
+def _get_backend(backend_type, backend_list):
+    """Get the single configured backend of the given type. An error is raised
+    if the number of matches != 1
+    """
+    matches = [
+        backend for backend in backend_list if backend.backend_type == backend_type
+    ]
+    assert len(matches) == 1
+    return matches[0]
+
+
+def get_load_backend(backend_type):
+    return _get_backend(backend_type, configured_load_backends())
+
+
+def get_train_backend(backend_type):
+    return _get_backend(backend_type, configured_train_backends())
+
+
 ## Tests #######################################################################
 
 
@@ -96,7 +114,7 @@ def test_configure_load_only(reset_globals):
         assert foo_cfg == mock_load_backend.config
 
         # Test train backend config
-        with pytest.raises(ValueError):
+        with pytest.raises(AssertionError):
             get_train_backend(backend_types.MOCK)
 
 
@@ -128,9 +146,9 @@ def test_disabling_local_backend(reset_globals):
         configure()
         assert get_load_backend(backend_types.MOCK)
         assert get_train_backend(backend_types.MOCK)
-        with pytest.raises(ValueError):
+        with pytest.raises(AssertionError):
             get_load_backend(backend_types.LOCAL)
-        with pytest.raises(ValueError):
+        with pytest.raises(AssertionError):
             get_train_backend(backend_types.LOCAL)
 
 
@@ -152,8 +170,8 @@ def test_duplicate_config_raises(reset_globals):
             configure()
 
 
-def test_duplicate_implied_names_raise(reset_globals):
-    """Test that duplicate entries with the same name implied from type raises"""
+def test_duplicate_explicit_instances_allowed(reset_globals):
+    """Test that duplicate entries are allowed"""
     with temp_config(
         {
             "module_backends": {
@@ -163,42 +181,6 @@ def test_duplicate_implied_names_raise(reset_globals):
                 ],
             }
         }
-    ):
-        with pytest.raises(ValueError):
-            configure()
-
-
-def test_duplicate_explicit_names_raise(reset_globals):
-    """Test that duplicate entries with the same name given explicitly from
-    raises
-    """
-    with temp_config(
-        {
-            "module_backends": {
-                "load_priority": [
-                    {"type": backend_types.MOCK, "name": "foo"},
-                    {"type": backend_types.MOCK, "name": "foo"},
-                ],
-            }
-        }
-    ):
-        with pytest.raises(ValueError):
-            configure()
-
-
-def test_duplicate_type_name_disambig(reset_globals):
-    """Test that multiple instances of the same type can be configured with
-    different names
-    """
-    with temp_config(
-        {
-            "module_backends": {
-                "load_priority": [
-                    {"type": backend_types.MOCK},
-                    {"type": backend_types.MOCK, "name": "foo"},
-                ],
-            }
-        }
     ):
         configure()
 
@@ -272,15 +254,3 @@ def test_multiple_module_same_backend_configures(reset_globals):
         assert "bar1" in mock_load_backend.config
         assert "bar2" in mock_load_backend.config
         assert mock_load_backend.config["bar1"] == 1
-
-
-def test_get_backend_starts_backend(reset_globals):
-    """Test that fetching a handle to a backend with get_backend ensures that it
-    is started
-    """
-    with temp_config(
-        {"module_backends": {"train_priority": [{"type": backend_types.MOCK}]}}
-    ):
-        configure()
-        mock_train_backend = get_train_backend(backend_types.MOCK)
-        assert mock_train_backend.is_started