Repository: caikit/caikit
Pull request #413 (2023-08-09 22:28:10)

Multi mm component 412
**What this PR does / why we need it**:

Supports #412 

This PR adds a `MultiModelFinder` that is capable of proxying to an ordered list of multiple finders in the `model_management.finders` section.

**Special notes for your reviewer**:

This will be needed in order to fully take advantage of the `TGISAutoFinder` in `caikit-nlp` ([PR](https://github.com/caikit/caikit-nlp/pull/123))


Patch:
diff --git a/caikit/core/exceptions/caikit_core_exception.py b/caikit/core/exceptions/caikit_core_exception.py
index 1af8d705..9e662b7a 100644
--- a/caikit/core/exceptions/caikit_core_exception.py
+++ b/caikit/core/exceptions/caikit_core_exception.py
@@ -25,14 +25,12 @@ class CaikitCoreStatusCode(Enum):
     CONNECTION_ERROR = 3
     UNAUTHORIZED = 4
     FORBIDDEN = 5
-    UNKNOWN = 6
-    FATAL = 7
 
 
 class CaikitCoreException(Exception):
-    status_code: CaikitCoreStatusCode
-    message: str
+    StatusCode: CaikitCoreStatusCode
+    Message: str
 
-    def __init__(self, status_code: CaikitCoreStatusCode, message: str) -> None:
-        self.status_code = status_code
-        self.message = message
+    def __init__(self, status_code: CaikitCoreStatusCode, msg: str) -> None:
+        self.StatusCode = status_code
+        self.Message = msg
diff --git a/caikit/core/model_management/factories.py b/caikit/core/model_management/factories.py
index 7b4d0908..c0aaccdb 100644
--- a/caikit/core/model_management/factories.py
+++ b/caikit/core/model_management/factories.py
@@ -27,6 +27,7 @@ from ..toolkit.factory import Factory, FactoryConstructible
 from .local_model_finder import LocalModelFinder
 from .local_model_initializer import LocalModelInitializer
 from .local_model_trainer import LocalModelTrainer
+from .multi_model_finder import MultiModelFinder
 
 log = alog.use_channel("MMFCTRY")
 error = error_handler.get(log)
@@ -94,6 +95,7 @@ model_trainer_factory.register(LocalModelTrainer)
 # configuration for a model based on a unique path or id.
 model_finder_factory = ImportableFactory("ModelFinder")
 model_finder_factory.register(LocalModelFinder)
+model_finder_factory.register(MultiModelFinder)
 
 # Model initializer factory. An initializer is responsible for taking a model
 # configuration and preparing the model to be run in a configured runtime
diff --git a/caikit/core/model_management/local_model_trainer.py b/caikit/core/model_management/local_model_trainer.py
index e1348c6b..c215e47d 100644
--- a/caikit/core/model_management/local_model_trainer.py
+++ b/caikit/core/model_management/local_model_trainer.py
@@ -324,7 +324,7 @@ class LocalModelTrainer(ModelTrainerBase):
             return model_future
         raise CaikitCoreException(
             status_code=CaikitCoreStatusCode.NOT_FOUND,
-            message=f"Unknown training_id: {training_id}",
+            msg=f"Unknown training_id: {training_id}",
         )
 
     ## Impl ##
diff --git a/caikit/core/model_management/multi_model_finder.py b/caikit/core/model_management/multi_model_finder.py
new file mode 100644
index 00000000..4722702c
--- /dev/null
+++ b/caikit/core/model_management/multi_model_finder.py
@@ -0,0 +1,137 @@
+# Copyright The Caikit Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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.
+"""
+The MultiModelFinder configures a set of other model finders that will be used
+in sequence to try loading models.
+
+Configuration for MultiModelFinder lives under the config as follows:
+
+model_management:
+    finders:
+        <finder name>:
+            type: MULTI
+            config:
+                # Sequence of other finder names to use in priority order
+                finder_sequence:
+                    - other_finder1
+                    - other_finder2
+"""
+# Standard
+from typing import Optional
+
+# First Party
+import aconfig
+import alog
+
+# Local
+from ...config import get_config
+from ..exceptions import error_handler
+from ..modules import ModuleConfig
+from .model_finder_base import ModelFinderBase
+
+# NOTE: Top-level import done so that global MODEL_MANAGER can be used at
+#   construction time without incurring a circular dependency
+import caikit.core
+
+log = alog.use_channel("MFIND")
+error = error_handler.get(log)
+
+
+class MultiModelFinder(ModelFinderBase):
+    __doc__ = __doc__
+
+    name = "MULTI"
+
+    def __init__(self, config: aconfig.Config, instance_name: str):
+        """Initialize with the sequence of finders to use"""
+        self._instance_name = instance_name
+        finder_priority = config.finder_priority
+        error.type_check(
+            "<COR47518221E>",
+            list,
+            finder_priority=finder_priority,
+        )
+        error.type_check_all(
+            "<COR47518222E>",
+            str,
+            finder_priority=finder_priority,
+        )
+        error.value_check(
+            "<COR05042343E>",
+            finder_priority,
+            "Must provide at least one valid finder",
+        )
+        config_finders = get_config().model_management.finders
+        invalid_finders = [
+            finder for finder in finder_priority if finder not in config_finders
+        ]
+        error.value_check(
+            "<COR68252034E>",
+            not invalid_finders,
+            "Invalid finders given in finder_priority: {}",
+            invalid_finders,
+        )
+        error.value_check(
+            "<COR54613971E>",
+            self._instance_name not in finder_priority,
+            "Cannot include self in multi finder priority",
+        )
+        model_manager = config.model_manager or caikit.core.MODEL_MANAGER
+        log.debug2("Setting up %s with finder priority: %s", self.name, finder_priority)
+        self._finders = [model_manager.get_finder(finder) for finder in finder_priority]
+
+    def find_model(
+        self,
+        model_path: str,
+        **kwargs,
+    ) -> Optional[ModuleConfig]:
+        """Iterate through the sequence of finders and return the first one that
+        succeeds
+        """
+        for idx, finder in enumerate(self._finders):
+            log.debug2(
+                "Trying to find %s with finder %d of type %s",
+                model_path,
+                idx,
+                finder.name,
+            )
+            try:
+                module_config = finder.find_model(model_path, **kwargs)
+                if module_config:
+                    log.debug(
+                        "Found %s with finder %d of type %s",
+                        model_path,
+                        idx,
+                        finder.name,
+                    )
+                    return module_config
+                log.debug2(
+                    "Finder %d of type %s unable to find %s",
+                    idx,
+                    finder.name,
+                    model_path,
+                )
+            except Exception as err:  # pylint: disable=broad-exception-caught
+                log.debug2(
+                    "Finder %d of type %s failed to load %s: %s",
+                    idx,
+                    finder.name,
+                    model_path,
+                    err,
+                )
+                log.debug4("Finder error", exc_info=True)
+
+        # No finder succeeded
+        log.warning("Unable to find %s with any finder", model_path)
+        return None
diff --git a/caikit/core/model_manager.py b/caikit/core/model_manager.py
index e4fc50c2..7a00172a 100644
--- a/caikit/core/model_manager.py
+++ b/caikit/core/model_manager.py
@@ -82,11 +82,11 @@ class ModelManager:
         """
         # Initialize all configured components
         mm_config = get_config().model_management
-        for trainer in mm_config.trainers:
+        for trainer in mm_config.get("trainers", {}):
             self.get_trainer(trainer)
-        for finder in mm_config.finders:
+        for finder in mm_config.get("finders", {}):
             self.get_finder(finder)
-        for initializer in mm_config.initializers:
+        for initializer in mm_config.get("initializers", {}):
             self.get_initializer(initializer)
 
     ## Public ##################################################################
@@ -454,6 +454,7 @@ class ModelManager:
             # NOTE: This will lazily construct named finders if needed
             log.debug("Attempting to find [%s] with finder %s", module_path, finder)
             finder = self.get_finder(finder)
+            log.debug2("Finder type: %s", finder.name)
             model_config = finder.find_model(module_path, **kwargs)
             error.value_check(
                 "<COR92173495E>",
@@ -465,7 +466,13 @@ class ModelManager:
             # Use the given initializer to try to load the model
             #
             # NOTE: This will lazily construct named initializers if needed
+            log.debug(
+                "Attempting to initialize [%s] with initializer %s",
+                module_path,
+                initializer,
+            )
             initializer = self.get_initializer(initializer)
+            log.debug2("Initializer type: %s", initializer.name)
             loaded_model = initializer.init(model_config, **kwargs)
             error.value_check(
                 "<COR50207494E>",
diff --git a/caikit/runtime/servicers/training_management_servicer.py b/caikit/runtime/servicers/training_management_servicer.py
index 412bb56b..2f33d3d1 100644
--- a/caikit/runtime/servicers/training_management_servicer.py
+++ b/caikit/runtime/servicers/training_management_servicer.py
@@ -95,7 +95,7 @@ class TrainingManagementServicerImpl:
             # This is to handle stateful trainers that implement `cancel` by fully deleting
             # the training. NB: Future `GetTrainingStatus` calls for these canceled trainings
             # would raise a not found error to the user.
-            if err.status_code == CaikitCoreStatusCode.NOT_FOUND:
+            if err.StatusCode == CaikitCoreStatusCode.NOT_FOUND:
                 return TrainingStatusResponse(
                     training_id=training_info_request.training_id,
                     state=TrainingStatus.CANCELED,
diff --git a/caikit/runtime/utils/servicer_util.py b/caikit/runtime/utils/servicer_util.py
index f12a66e1..bf3deb30 100644
--- a/caikit/runtime/utils/servicer_util.py
+++ b/caikit/runtime/utils/servicer_util.py
@@ -458,58 +458,44 @@ def build_caikit_library_request_dict(
 
 
 def raise_caikit_runtime_exception(exception: CaikitCoreException):
-    if exception.status_code == CaikitCoreStatusCode.NOT_FOUND:
-        log.debug2("Not Found Error: [%s]", exception.message, exc_info=True)
+    if exception.StatusCode == CaikitCoreStatusCode.NOT_FOUND:
+        log.debug2("Not Found Error: [%s]", exception.Message, exc_info=True)
         raise CaikitRuntimeException(
             grpc.StatusCode.NOT_FOUND,
-            exception.message,
+            exception.Message,
         ) from exception
 
-    if exception.status_code == CaikitCoreStatusCode.INVALID_ARGUMENT:
-        log.debug2("Invalid Argument Error: [%s]", exception.message, exc_info=True)
+    if exception.StatusCode == CaikitCoreStatusCode.INVALID_ARGUMENT:
+        log.debug2("Invalid Argument Error: [%s]", exception.Message, exc_info=True)
         raise CaikitRuntimeException(
             grpc.StatusCode.INVALID_ARGUMENT,
-            exception.message,
+            exception.Message,
         ) from exception
 
-    if exception.status_code == CaikitCoreStatusCode.FORBIDDEN:
-        log.debug2("Forbidden Error: [%s]", exception.message, exc_info=True)
+    if exception.StatusCode == CaikitCoreStatusCode.FORBIDDEN:
+        log.debug2("Forbidden Error: [%s]", exception.Message, exc_info=True)
         raise CaikitRuntimeException(
             grpc.StatusCode.PERMISSION_DENIED,
-            exception.message,
+            exception.Message,
         ) from exception
 
-    if exception.status_code == CaikitCoreStatusCode.CONNECTION_ERROR:
-        log.debug2("Connection Error: [%s]", exception.message, exc_info=True)
+    if exception.StatusCode == CaikitCoreStatusCode.CONNECTION_ERROR:
+        log.debug2("Connection Error: [%s]", exception.Message, exc_info=True)
         raise CaikitRuntimeException(
             grpc.StatusCode.UNAVAILABLE,
-            exception.message,
+            exception.Message,
         ) from exception
 
-    if exception.status_code == CaikitCoreStatusCode.UNAUTHORIZED:
+    if exception.StatusCode == CaikitCoreStatusCode.UNAUTHORIZED:
         log.debug2("Unauthorized Error: [%s]", exception.Message, exc_info=True)
         raise CaikitRuntimeException(
             grpc.StatusCode.UNAUTHENTICATED,
-            exception.message,
-        ) from exception
-
-    if exception.status_code == CaikitCoreStatusCode.FATAL:
-        log.debug2("Fatal Error: [%s]", exception.message, exc_info=True)
-        raise CaikitRuntimeException(
-            grpc.StatusCode.INTERNAL,
-            exception.message,
-        ) from exception
-
-    if exception.status_code == CaikitCoreStatusCode.UNKNOWN:
-        log.debug2("Unknown Error: [%s]", exception.message, exc_info=True)
-        raise CaikitRuntimeException(
-            grpc.StatusCode.UNKNOWN,
-            exception.message,
+            exception.Message,
         ) from exception
 
     log.debug2(
-        "Unhandled CaikitCoreException Error: [%s]", exception.message, exc_info=True
+        "Unknown CaikitCoreException Error: [%s]", exception.Message, exc_info=True
     )
     raise CaikitRuntimeException(
-        grpc.StatusCode.INTERNAL, exception.message
+        grpc.StatusCode.UNKNOWN, exception.Message
     ) from exception


Test patch:
diff --git a/tests/core/helpers.py b/tests/core/helpers.py
index 41499c17..df6b2020 100644
--- a/tests/core/helpers.py
+++ b/tests/core/helpers.py
@@ -75,11 +75,17 @@ class TestFinder(ModelFinderBase):
         config.setdefault("module_id", SampleModule.MODULE_ID)
         config.setdefault("train", {}).setdefault("batch_size", 1)
         config.setdefault("train", {}).setdefault("learning_rate", 0.1)
+        self._fail_to_find = config.fail_to_find
+        self._raise_on_find = config.raise_on_find
         self._config = config
         self._local_finder = model_finder_factory.construct({"type": "LOCAL"})
         self._instance_name = instance_name
 
     def find_model(self, model_path, *args, **kwargs):
+        if self._raise_on_find:
+            raise RuntimeError("You told me to")
+        if self._fail_to_find:
+            return None
         if os.path.exists(model_path):
             return self._local_finder.find_model(model_path, *args, **kwargs)
         return ModuleConfig(self._config)
diff --git a/tests/core/model_management/test_multi_model_finder.py b/tests/core/model_management/test_multi_model_finder.py
new file mode 100644
index 00000000..0d3ddf31
--- /dev/null
+++ b/tests/core/model_management/test_multi_model_finder.py
@@ -0,0 +1,138 @@
+# Copyright The Caikit Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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.
+"""
+Tests for the MultiModelFinder
+"""
+# Standard
+from contextlib import contextmanager
+
+# Third Party
+import pytest
+
+# First Party
+import aconfig
+
+# Local
+from caikit.core import ModelManager
+from caikit.core.model_management.multi_model_finder import MultiModelFinder
+from tests.conftest import temp_config
+from tests.core.helpers import TestFinder
+
+## Helpers #####################################################################
+
+
+@contextmanager
+def temp_config_setup(
+    multi_finder_name="default",
+    multi_finder_cfg=None,
+    config_overrides=None,
+):
+    config_overrides = config_overrides or {}
+    finders = config_overrides.setdefault("model_management", {},).setdefault(
+        "finders",
+        {},
+    )
+    # Add a local one to proxy to
+    if not finders:
+        finders["local"] = {"type": "LOCAL"}
+    multi_finder_cfg = multi_finder_cfg or {"finder_priority": ["local"]}
+
+    multi_model_section = {
+        "type": MultiModelFinder.name,
+        "config": multi_finder_cfg,
+    }
+    finders.setdefault(multi_finder_name, multi_model_section)
+    with temp_config(config_overrides, "merge"):
+        config_overrides = aconfig.Config(config_overrides, override_env_vars=False)
+        yield config_overrides.model_management.finders[multi_finder_name].config
+
+
+@contextmanager
+def temp_config_finder(
+    multi_finder_name="default",
+    multi_finder_cfg=None,
+    config_overrides=None,
+):
+    with temp_config_setup(
+        multi_finder_name=multi_finder_name,
+        multi_finder_cfg=multi_finder_cfg,
+        config_overrides=config_overrides,
+    ) as mf_config:
+        mmgr = ModelManager()
+        mf_config.model_manager = mmgr
+        yield MultiModelFinder(mf_config, multi_finder_name)
+
+
+## Tests #######################################################################
+
+
+def test_multi_model_finder_model_found(good_model_path):
+    """Make sure that a simple proxy to local works"""
+    with temp_config_finder() as finder:
+        assert isinstance(finder, MultiModelFinder)
+        assert finder.find_model(good_model_path)
+
+
+@pytest.mark.parametrize(
+    "test_finder_config",
+    [{"fail_to_find": True}, {"raise_on_find": True}],
+)
+def test_multi_model_finder_first_not_found(test_finder_config, good_model_path):
+    """Make sure that a model can be found if the first finder fails"""
+    with temp_config_finder(
+        config_overrides={
+            "model_management": {
+                "finders": {
+                    "default": {
+                        "type": MultiModelFinder.name,
+                        "config": {
+                            "finder_priority": ["test-not-found", "local"],
+                        },
+                    },
+                    "test-not-found": {
+                        "type": TestFinder.name,
+                        "config": test_finder_config,
+                    },
+                    "local": {"type": "LOCAL"},
+                }
+            }
+        }
+    ) as finder:
+        assert isinstance(finder, MultiModelFinder)
+        assert finder.find_model(good_model_path)
+
+
+def test_multi_model_finder_not_found():
+    """Make sure that a simple proxy to local works"""
+    with temp_config_finder() as finder:
+        assert isinstance(finder, MultiModelFinder)
+        assert finder.find_model("not/a/valid/path") is None
+
+
+@pytest.mark.parametrize(
+    "params",
+    [
+        ({"finder_priority": "local"}, TypeError),
+        ({"finder_priority": [0, 1]}, TypeError),
+        ({"finder_priority": []}, ValueError),
+        ({"finder_priority": ["invalid-name"]}, ValueError),
+        ({"finder_priority": ["default"]}, ValueError),
+    ],
+)
+def test_multi_model_finder_invalid_config(params):
+    """Validate all forms of bad config"""
+    multi_finder_cfg, error_type = params
+    with temp_config_setup(multi_finder_cfg=multi_finder_cfg) as mf_config:
+        with pytest.raises(error_type):
+            MultiModelFinder(mf_config, "default")
diff --git a/tests/runtime/servicers/test_training_management_servicer.py b/tests/runtime/servicers/test_training_management_servicer.py
index 7f36510b..9b67d3a1 100644
--- a/tests/runtime/servicers/test_training_management_servicer.py
+++ b/tests/runtime/servicers/test_training_management_servicer.py
@@ -64,7 +64,7 @@ class MockModelFuture:
             return TrainingInfo(status=TrainingStatus.RUNNING)
         else:
             raise CaikitCoreException(
-                status_code=CaikitCoreStatusCode.NOT_FOUND, message="Training not found"
+                status_code=CaikitCoreStatusCode.NOT_FOUND, msg="Training not found"
             )
 
     def cancel(self):