Repository: caikit/caikit
Pull request #594 (2023-12-04 21:11:30)

Add MultiModelInitializer Implementation
<!--  Thanks for sending a pull request!  Here are some tips for you:
1. Make sure to read the Contributing Guide before submitting your PR: https://github.com/caikit/caikit/blob/main/CONTRIBUTING.md
2. If this PR closes another issue, add 'closes #<issue number>' somewhere in the PR summary. GitHub will automatically close that issue when this PR gets merged. Alternatively, adding 'refs #<issue number>' will not close the issue, but help provide the reviewer more context.-->

**What this PR does / why we need it**:
This PR adds a `MultiModelInitializer` implementation that is very similar to the existing MultiModelFinder except that it works with initializers. This will support future initializers that might need a backup/default override. 

**Special notes for your reviewer**:

**If applicable**:
- [ ] this PR contains documentation
- [ ] this PR contains unit tests
- [ ] this PR has been tested for backwards compatibility



Patch:
diff --git a/caikit/core/model_management/factories.py b/caikit/core/model_management/factories.py
index c39cc5db..24fd26a3 100644
--- a/caikit/core/model_management/factories.py
+++ b/caikit/core/model_management/factories.py
@@ -21,6 +21,7 @@ from .local_model_finder import LocalModelFinder
 from .local_model_initializer import LocalModelInitializer
 from .local_model_trainer import LocalModelTrainer
 from .multi_model_finder import MultiModelFinder
+from .multi_model_initializer import MultiModelInitializer
 
 # Model trainer factory. A trainer is responsible for performing the train
 # operation against a configured framework connection.
@@ -38,3 +39,4 @@ model_finder_factory.register(MultiModelFinder)
 # location.
 model_initializer_factory = ImportableFactory("ModelInitializer")
 model_initializer_factory.register(LocalModelInitializer)
+model_initializer_factory.register(MultiModelInitializer)
diff --git a/caikit/core/model_management/multi_model_initializer.py b/caikit/core/model_management/multi_model_initializer.py
new file mode 100644
index 00000000..f5b5a0bb
--- /dev/null
+++ b/caikit/core/model_management/multi_model_initializer.py
@@ -0,0 +1,146 @@
+# 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 MultiModelInitializer configures a set of other model initializers that will be used
+in sequence to try loading models.
+
+Configuration for MultiModelInitializer lives under the config as follows:
+
+model_management:
+    initializers:
+        <initializer name>:
+            type: MULTI
+            config:
+                # Sequence of other initializer names to use in priority order
+                initializer_priority:
+                    - other_initializer1
+                    - other_initializer2
+"""
+# 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 ModuleBase, ModuleConfig
+from .model_initializer_base import ModelInitializerBase
+
+# 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("MINIT")
+error = error_handler.get(log)
+
+
+class MultiModelInitializer(ModelInitializerBase):
+    __doc__ = __doc__
+
+    name = "MULTI"
+
+    def __init__(self, config: aconfig.Config, instance_name: str):
+        """Initialize with the sequence of initializers to use"""
+        self._instance_name = instance_name
+        initializer_priority = config.initializer_priority
+        error.type_check(
+            "<COR47518221E>",
+            list,
+            initializer_priority=initializer_priority,
+        )
+        error.type_check_all(
+            "<COR47518222E>",
+            str,
+            initializer_priority=initializer_priority,
+        )
+        error.value_check(
+            "<COR05042343E>",
+            initializer_priority,
+            "Must provide at least one valid initializer",
+        )
+        config_initializers = get_config().model_management.initializers
+        invalid_initializers = [
+            initializer
+            for initializer in initializer_priority
+            if initializer not in config_initializers
+        ]
+        error.value_check(
+            "<COR68252034E>",
+            not invalid_initializers,
+            "Invalid initializers given in initializer_priority: {}",
+            invalid_initializers,
+        )
+        error.value_check(
+            "<COR54613971E>",
+            self._instance_name not in config_initializers,
+            "Cannot include self in multi initializer priority",
+        )
+        model_manager = config.model_manager or caikit.core.MODEL_MANAGER
+        log.debug2(
+            "Setting up %s with initializer priority: %s",
+            self.name,
+            initializer_priority,
+        )
+        self._initializers = [
+            model_manager.get_initializer(initializer)
+            for initializer in initializer_priority
+        ]
+
+    def init(
+        self,
+        model_config: ModuleConfig,
+        **kwargs,
+    ) -> Optional[ModuleBase]:
+        """Iterate through the sequence of initializers and return the first one that
+        succeeds
+        """
+        for idx, initializer in enumerate(self._initializers):
+            log.debug2(
+                "Trying to init %s with initializer %d of type %s",
+                model_config.module_id,
+                idx,
+                initializer.name,
+            )
+            try:
+                module = initializer.init(model_config, **kwargs)
+                if module:
+                    log.debug(
+                        "Init model %s with initializer %d of type %s",
+                        model_config.module_id,
+                        idx,
+                        initializer.name,
+                    )
+                    return module
+                log.debug2(
+                    "Initializer %d of type %s unable to init %s",
+                    idx,
+                    initializer.name,
+                    model_config.module_id,
+                )
+            except Exception as err:  # pylint: disable=broad-exception-caught
+                log.debug2(
+                    "Initializer %d of type %s failed to load %s: %s",
+                    idx,
+                    initializer.name,
+                    model_config.module_id,
+                    err,
+                )
+                log.debug4("Initializer error", exc_info=True)
+
+        # No initializer succeeded
+        log.warning("Unable to init %s with any initializer", model_config.module_id)
+        return None
diff --git a/pyproject.toml b/pyproject.toml
index 1a1a6e66..54428d91 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -65,9 +65,7 @@ all = [
 ## Dev Extra Sets ##
 
 dev-test = [
-    # NOTE: pytest-asyncio>=0.22 breaks importing with an error about multiple
-    #   imports of sample modules
-    "pytest-asyncio>=0.21.0,<0.22",
+    "pytest-asyncio>=0.21.0,<1",
     "pytest-cov>=2.10.1,<5.0",
     "pytest-html>=3.1.1,<5.0",
     "pytest>=6.2.5,<8.0",


Test patch:
diff --git a/tests/core/model_management/test_multi_model_initializer.py b/tests/core/model_management/test_multi_model_initializer.py
new file mode 100644
index 00000000..0d83a270
--- /dev/null
+++ b/tests/core/model_management/test_multi_model_initializer.py
@@ -0,0 +1,101 @@
+# 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.model_management.factories import model_initializer_factory
+from caikit.core.model_management.local_model_finder import LocalModelFinder
+from caikit.core.model_management.model_initializer_base import ModelInitializerBase
+from caikit.core.modules import ModuleConfig
+from tests.conftest import temp_config
+
+## Helpers #####################################################################
+
+# Add bad initializer to model factory
+class BadModelInitializer(ModelInitializerBase):
+    name = "BAD"
+
+    def __init__(self, config: aconfig.Config, instance_name: str):
+        """A FactoryConstructible object must be constructed with a config
+        object that it uses to pull in all configuration
+        """
+        pass
+
+    def init(
+        self,
+        model_config,
+        **kwargs,
+    ):
+        raise ValueError("Bad Model Initializer")
+
+
+model_initializer_factory.register(BadModelInitializer)
+
+
+@contextmanager
+def construct_mm_initializer(multi_model_config, config_override={}):
+    config_override = config_override or {
+        "model_management": {
+            "initializers": {
+                "local": {
+                    "type": "LOCAL",
+                },
+                "bad": {"type": "BAD"},
+            }
+        }
+    }
+
+    with temp_config(config_override, "merge"):
+        model_config = {
+            "type": "MULTI",
+            "config": multi_model_config,
+        }
+        yield model_initializer_factory.construct(model_config, "instance_name")
+
+
+## Tests #######################################################################
+
+
+@pytest.mark.parametrize(
+    ["initializers", "load_successful"],
+    [[["local"], True], [["bad", "local"], True], [["bad"], False]],
+)
+def test_multi_model_initializer(good_model_path, initializers, load_successful):
+    finder = LocalModelFinder(aconfig.Config({}), "local")
+    config = finder.find_model(good_model_path)
+    with construct_mm_initializer(
+        {"initializer_priority": initializers}
+    ) as initializer:
+        if load_successful:
+            assert initializer.init(config)
+        else:
+            assert not initializer.init(config)
+
+
+def test_multi_model_initializer_bad_config():
+    config = ModuleConfig({"module_id": "bad"})
+    with construct_mm_initializer(
+        {"initializer_priority": ["bad", "local"]}
+    ) as initializer:
+        assert not initializer.init(config)