Repository: caikit/caikit
Pull request #400 (2023-08-04 21:19:44)

:goal_net: add more compatibility for toolkit.error_handler
<!--  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**:

Whoops, adds back the ability to
```
from caikit.core.toolkit import error_handler
```

with warnings

**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/local_model_initializer.py b/caikit/core/model_management/local_model_initializer.py
index d3d42680..6057458e 100644
--- a/caikit/core/model_management/local_model_initializer.py
+++ b/caikit/core/model_management/local_model_initializer.py
@@ -28,8 +28,7 @@ model_management:
                       config: {}
 """
 # Standard
-from typing import Callable, List, Optional
-import copy
+from typing import Callable, Optional
 import inspect
 
 # First Party
@@ -178,29 +177,13 @@ class LocalModelInitializer(ModelInitializerBase):
                     module_backend_impl.__name__,
                 )
                 extra_kwargs = {}
-                if self._supports_arg(module_backend_impl.load, "load_backend"):
+                if self._supports_load_backend_kwarg(module_backend_impl.load):
                     extra_kwargs["load_backend"] = load_backend
-                # Try loading with the ModuleConfig directly and fall back to
-                # loading with the model_path for compatibility
-                try:
-                    loaded_model = module_backend_impl.load(
-                        model_config,
-                        **extra_kwargs,
-                        **kwargs,
-                    )
-                except TypeError as err:
-                    log.warning(
-                        "<COR98539580W>",
-                        "DEPRECATION: Loading %s failed with ModuleConfig. Using model_path. %s",
-                        module_backend_impl.MODULE_ID,
-                        err,
-                    )
-                    loaded_model = module_backend_impl.load(
-                        model_path,
-                        **extra_kwargs,
-                        **kwargs,
-                    )
-
+                loaded_model = module_backend_impl.load(
+                    model_path,
+                    **extra_kwargs,
+                    **kwargs,
+                )
                 error.type_check("<COR40080753E>", ModuleBase, model=loaded_model)
                 if loaded_model is not None:
                     log.debug2(
@@ -214,20 +197,15 @@ class LocalModelInitializer(ModelInitializerBase):
         # Return the loaded model if it was able to load
         return loaded_model
 
-    @property
-    def backends(self) -> List[BackendBase]:
-        """Return an immutable view of the priority sequence of backends"""
-        return copy.copy(self._backends)
-
     ## Implementation Details ##################################################
 
     @staticmethod
-    def _supports_arg(load_fn: Callable, arg_name: str) -> bool:
+    def _supports_load_backend_kwarg(load_fn: Callable) -> bool:
         """A load function supports the load_backend kwarg IFF it has an arg
         explicitly named load_backend or it has a ** kwarg capture
         """
         sig = inspect.signature(load_fn)
-        return arg_name in sig.parameters
+        return "load_backend" in sig.parameters
 
     def _get_supported_load_backends(self, backend_impl: ModuleBase):
         """Function to get a list of supported load backends
diff --git a/caikit/core/model_manager.py b/caikit/core/model_manager.py
index e4fc50c2..80d18151 100644
--- a/caikit/core/model_manager.py
+++ b/caikit/core/model_manager.py
@@ -83,11 +83,11 @@ class ModelManager:
         # Initialize all configured components
         mm_config = get_config().model_management
         for trainer in mm_config.trainers:
-            self.get_trainer(trainer)
+            self._get_trainer(trainer)
         for finder in mm_config.finders:
-            self.get_finder(finder)
+            self._get_finder(finder)
         for initializer in mm_config.initializers:
-            self.get_initializer(initializer)
+            self._get_initializer(initializer)
 
     ## Public ##################################################################
 
@@ -149,7 +149,7 @@ class ModelManager:
         error.subclass_check("<COR05418775E>", module, ModuleBase)
 
         # Get the trainer to use
-        trainer: ModelTrainerBase = self.get_trainer(trainer)
+        trainer: ModelTrainerBase = self._get_trainer(trainer)
 
         # Start the training
         with alog.ContextTimer(log.debug, "Started training in: "):
@@ -193,11 +193,11 @@ class ModelManager:
                 to the model which holds the status of the in-flight training.
         """
         try:
-            trainer = self.get_trainer(ModelTrainerBase.get_trainer_name(training_id))
+            trainer = self._get_trainer(ModelTrainerBase.get_trainer_name(training_id))
 
         # Fall back to the default trainer to try to find this ID
         except ValueError:
-            trainer = self.get_trainer("default")
+            trainer = self._get_trainer("default")
 
         return trainer.get_model_future(training_id)
 
@@ -383,41 +383,6 @@ class ModelManager:
         with self.__singleton_lock:
             self._singleton_module_cache.clear()
 
-    def get_trainer(self, trainer: Union[str, ModelTrainerBase]) -> ModelTrainerBase:
-        """Get the configured model trainer or the one passed by value"""
-        return self._get_component(
-            component=trainer,
-            component_dict=self._trainers,
-            component_factory=model_trainer_factory,
-            component_name="trainer",
-            component_cfg=get_config().model_management.trainers,
-            component_type=ModelTrainerBase,
-        )
-
-    def get_finder(self, finder: Union[str, ModelFinderBase]) -> ModelFinderBase:
-        """Get the configured model finder or the one passed by value"""
-        return self._get_component(
-            component=finder,
-            component_dict=self._finders,
-            component_factory=model_finder_factory,
-            component_name="finder",
-            component_cfg=get_config().model_management.finders,
-            component_type=ModelFinderBase,
-        )
-
-    def get_initializer(
-        self, initializer: Union[str, ModelInitializerBase]
-    ) -> ModelInitializerBase:
-        """Get the configured model initializer or the one passed by value"""
-        return self._get_component(
-            component=initializer,
-            component_dict=self._initializers,
-            component_factory=model_initializer_factory,
-            component_name="initializer",
-            component_cfg=get_config().model_management.initializers,
-            component_type=ModelInitializerBase,
-        )
-
     ## Implementation Details ##################################################
 
     def _load_from_dir(
@@ -453,7 +418,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)
+            finder = self._get_finder(finder)
             model_config = finder.find_model(module_path, **kwargs)
             error.value_check(
                 "<COR92173495E>",
@@ -465,7 +430,7 @@ class ModelManager:
             # Use the given initializer to try to load the model
             #
             # NOTE: This will lazily construct named initializers if needed
-            initializer = self.get_initializer(initializer)
+            initializer = self._get_initializer(initializer)
             loaded_model = initializer.init(model_config, **kwargs)
             error.value_check(
                 "<COR50207494E>",
@@ -595,3 +560,38 @@ class ModelManager:
             )
             component_dict[component] = component_factory.construct(cfg, component)
         return component_dict[component]
+
+    def _get_trainer(self, trainer: Union[str, ModelTrainerBase]) -> ModelTrainerBase:
+        """Get the configured model trainer or the one passed by value"""
+        return self._get_component(
+            component=trainer,
+            component_dict=self._trainers,
+            component_factory=model_trainer_factory,
+            component_name="trainer",
+            component_cfg=get_config().model_management.trainers,
+            component_type=ModelTrainerBase,
+        )
+
+    def _get_finder(self, finder: Union[str, ModelFinderBase]) -> ModelFinderBase:
+        """Get the configured model finder or the one passed by value"""
+        return self._get_component(
+            component=finder,
+            component_dict=self._finders,
+            component_factory=model_finder_factory,
+            component_name="finder",
+            component_cfg=get_config().model_management.finders,
+            component_type=ModelFinderBase,
+        )
+
+    def _get_initializer(
+        self, initializer: Union[str, ModelInitializerBase]
+    ) -> ModelInitializerBase:
+        """Get the configured model initializer or the one passed by value"""
+        return self._get_component(
+            component=initializer,
+            component_dict=self._initializers,
+            component_factory=model_initializer_factory,
+            component_name="initializer",
+            component_cfg=get_config().model_management.initializers,
+            component_type=ModelInitializerBase,
+        )
diff --git a/caikit/core/modules/config.py b/caikit/core/modules/config.py
index f93dbd84..262197b8 100644
--- a/caikit/core/modules/config.py
+++ b/caikit/core/modules/config.py
@@ -13,7 +13,6 @@
 # limitations under the License.
 
 # Standard
-from typing import Union
 import os
 
 # First Party
@@ -83,23 +82,20 @@ class ModuleConfig(aconfig.Config):
         )
 
     @classmethod
-    def load(cls, model_path: Union[str, "ModuleConfig"]) -> "ModuleConfig":
+    def load(cls, model_path):
         """Load a new module configuration from a directory on disk.
 
         Args:
-            model_path (Union[str, ModuleConfig]): Path to model directory. At
-                the top level of directory is `config.yml` which holds info
-                about the model. Note that the model_path here is assumed to be
-                operating system correct as a consequence of the way this method
-                is invoked by the model manager.
+            model_path (str): Path to model directory. At the top level of
+                directory is `config.yml` which holds info about the model. Note
+                that the model_path here is assumed to be operating system
+                correct as a consequence of the way this method is invoked by
+                the model manager.
 
         Returns:
-            model_config (ModuleConfig): Instantiated ModuleConfig for model
-                given model_path.
+            BlockConfig: Instantiated BlockConfig for model given model_path.
         """
-        error.type_check("<COR71170339E>", str, cls, model_path=model_path)
-        if isinstance(model_path, cls):
-            return model_path
+        error.type_check("<COR71170339E>", str, model_path=model_path)
 
         # Validate config.yml
         config_path = os.path.join(model_path, "config.yml")
diff --git a/caikit/core/modules/loader.py b/caikit/core/modules/loader.py
index e6837591..94936ad8 100644
--- a/caikit/core/modules/loader.py
+++ b/caikit/core/modules/loader.py
@@ -17,7 +17,6 @@ Contains recursive functions for loading modules saved inside modules.
 """
 
 # Standard
-from typing import Union
 import os
 
 # First Party
@@ -35,15 +34,16 @@ error = error_handler.get(log)
 class ModuleLoader:
     MODULE_PATHS_KEY = "module_paths"
 
-    def __init__(self, model_path: Union[str, ModuleConfig]):
+    def __init__(self, model_path):
         """Construct a new module loader.
 
         Args:
-            model_path (Union[str, ModuleConfig]): The path to the directory
-                where the model is to be loaded from, or a preloaded config.
+            model_path (str): The path to the directory where the model is to be
+                loaded from.
         """
+        self.model_path = os.path.normpath(model_path)
+        error.dir_check("<COR43014802E>", model_path)
         self.config = ModuleConfig.load(model_path)
-        self.model_path = self.config.model_path
 
     def load_arg(self, arg):
         """Extract arg value from the loaded model's config"""
diff --git a/caikit/core/toolkit/error_handler.py b/caikit/core/toolkit/error_handler.py
new file mode 100644
index 00000000..64d93d0f
--- /dev/null
+++ b/caikit/core/toolkit/error_handler.py
@@ -0,0 +1,32 @@
+# 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.
+
+
+"""This file exists for backwards compatibility.
+It imports from the caikit.core.exceptions package where things moved"""
+
+# Standard
+import warnings as _warnings
+
+# Local
+from caikit.core.exceptions.error_handler import *  # pylint: disable=wildcard-import,unused-wildcard-import
+
+# 🌶️🌶️🌶️ Warn if this package is ever imported
+# Allow DeprecationWarnings through if anything tries to import from `toolkit.errors`
+_warnings.filterwarnings("default", category=DeprecationWarning)
+# And actually warn them
+_warnings.warn(
+    "The caikit.toolkit.error_handler package has moved to caikit.core.exceptions",
+    DeprecationWarning,
+)
diff --git a/caikit/runtime/model_management/model_loader.py b/caikit/runtime/model_management/model_loader.py
index 692baf63..2d42e295 100644
--- a/caikit/runtime/model_management/model_loader.py
+++ b/caikit/runtime/model_management/model_loader.py
@@ -24,7 +24,7 @@ import alog
 
 # Local
 from caikit.config import get_config
-from caikit.core import MODEL_MANAGER, ModuleBase
+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
@@ -32,6 +32,7 @@ from caikit.runtime.work_management.abortable_action import (
     AbortableAction,
     ActionAborter,
 )
+import caikit.core
 
 log = alog.use_channel("MODEL-LOADER")
 
@@ -109,7 +110,7 @@ class ModelLoader:
 
             # Load using the caikit.core
             with CAIKIT_CORE_LOAD_DURATION_SUMMARY.labels(model_type=model_type).time():
-                model = MODEL_MANAGER.load(model_path)
+                model = caikit.core.load(model_path)
 
             # If this model needs batching, configure a Batcher to wrap it
             model = self._wrap_in_batcher_if_configured(
@@ -156,10 +157,10 @@ class ModelLoader:
 
     def _wrap_in_batcher_if_configured(
         self,
-        caikit_core_model: ModuleBase,
+        caikit_core_model: caikit.core.ModuleBase,
         model_type: str,
         model_id: str,
-    ) -> Union[Batcher, ModuleBase]:
+    ) -> Union[Batcher, caikit.core.ModuleBase]:
         """Perform Batcher wrapping on the given module if configured, otherwise
         return the model as is
         """
diff --git a/caikit/runtime/model_management/model_manager.py b/caikit/runtime/model_management/model_manager.py
index 6f84e615..1637f9d9 100644
--- a/caikit/runtime/model_management/model_manager.py
+++ b/caikit/runtime/model_management/model_manager.py
@@ -197,16 +197,9 @@ class ModelManager:  # pylint: disable=too-many-instance-attributes
                         raise ex
 
                     # Estimate the model's size and update the LoadedModel
-                    try:
-                        model_size = self.model_sizer.get_model_size(
-                            model_id, local_model_path, model_type
-                        )
-                    except CaikitRuntimeException:
-                        log.debug(
-                            "Unable to estimate model size for non-disk model: %s",
-                            model_id,
-                        )
-                        model_size = 0
+                    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
@@ -374,16 +367,17 @@ class ModelManager:  # pylint: disable=too-many-instance-attributes
         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)
-            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 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:


Test patch:
diff --git a/tests/core/helpers.py b/tests/core/helpers.py
index 41499c17..d1864694 100644
--- a/tests/core/helpers.py
+++ b/tests/core/helpers.py
@@ -14,7 +14,6 @@
 
 # Standard
 from typing import Optional, Type
-import os
 import uuid
 
 # Third Party
@@ -24,11 +23,9 @@ import pytest
 from caikit.core import MODEL_MANAGER
 from caikit.core.data_model import TrainingStatus
 from caikit.core.model_management import (
-    ModelFinderBase,
     ModelInitializerBase,
     ModelTrainerBase,
     TrainingInfo,
-    model_finder_factory,
     model_initializer_factory,
     model_trainer_factory,
 )
@@ -41,7 +38,6 @@ from caikit.core.registries import (
     module_backend_types,
     module_registry,
 )
-from sample_lib.modules import SampleModule
 
 
 # Add mock backend
@@ -65,30 +61,8 @@ class MockBackend(BackendBase):
 
 backend_types.register_backend_type(MockBackend)
 
-# Add a new model finder that tests can use
-class TestFinder(ModelFinderBase):
-    name = "TestFinder"
-    __test__ = False
-
-    def __init__(self, config, instance_name):
-        # Fill in expected fields for SampleModule
-        config.setdefault("module_id", SampleModule.MODULE_ID)
-        config.setdefault("train", {}).setdefault("batch_size", 1)
-        config.setdefault("train", {}).setdefault("learning_rate", 0.1)
-        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 os.path.exists(model_path):
-            return self._local_finder.find_model(model_path, *args, **kwargs)
-        return ModuleConfig(self._config)
-
-
-model_finder_factory.register(TestFinder)
-
 
-# Add a new model initializer that tests can use
+# Add a new shared load backend that tests can use
 class TestInitializer(ModelInitializerBase):
     name = "TestInitializer"
     __test__ = False
diff --git a/tests/core/modules/test_module.py b/tests/core/modules/test_module.py
index f86d61c0..d7189c9a 100644
--- a/tests/core/modules/test_module.py
+++ b/tests/core/modules/test_module.py
@@ -350,7 +350,7 @@ def test_class_attributes(reset_globals):
         # Normally non-local backend modules are loaded through caikit.load, so
         # we need to make sure the local loader has been properly configured to
         # enable local instance instantiation.
-        assert MODEL_MANAGER.get_initializer("default")
+        assert MODEL_MANAGER._get_initializer("default")
 
         # Make sure an instance can fetch via get_backend()
         inst = DummyBar()
diff --git a/tests/core/test_model_manager.py b/tests/core/test_model_manager.py
index 4bc63eb8..f156a22a 100644
--- a/tests/core/test_model_manager.py
+++ b/tests/core/test_model_manager.py
@@ -489,8 +489,8 @@ def test_load_with_two_shared_initializers_of_the_same_type(
             }
         }
     ):
-        model_one_loader = MODEL_MANAGER.get_initializer("model-one")
-        model_two_loader = MODEL_MANAGER.get_initializer("model-two")
+        model_one_loader = MODEL_MANAGER._get_initializer("model-one")
+        model_two_loader = MODEL_MANAGER._get_initializer("model-two")
 
         # plain model load should use first loader
         model = caikit.core.load(good_model_path, initializer="model-one")
@@ -600,22 +600,6 @@ def test_initialize_all_components(reset_globals):
         assert set(MODEL_MANAGER._initializers.keys()) == {"default", "foobar"}
 
 
-def test_find_without_on_disk_model(good_model_path, reset_globals):
-    """Make sure that ephemeral models can be found without existing on disk"""
-
-    with temp_config(
-        {
-            "model_management": {
-                "finders": {"default": {"type": TestFinder.name}},
-                "initializers": {"default": {"type": "LOCAL"}},
-            }
-        }
-    ):
-        model = caikit.core.load("some_pretend_model")
-        assert model
-        assert not os.path.exists("some_pretend_model")
-
-
 def test_train_by_module_class(reset_globals):
     """Make sure training can be accessed through the central train function
     with the class directly
diff --git a/tests/runtime/model_management/test_model_manager.py b/tests/runtime/model_management/test_model_manager.py
index 63f1fd5e..2efc9a87 100644
--- a/tests/runtime/model_management/test_model_manager.py
+++ b/tests/runtime/model_management/test_model_manager.py
@@ -27,16 +27,13 @@ import pytest
 
 # Local
 from caikit import get_config
-from caikit.core.model_manager import ModelManager as CoreModelManager
 from caikit.core.modules import ModuleBase
 from caikit.runtime.model_management.loaded_model import LoadedModel
 from caikit.runtime.model_management.model_manager import ModelManager
 from caikit.runtime.types.caikit_runtime_exception import CaikitRuntimeException
 from caikit.runtime.utils.import_util import get_dynamic_module
 from tests.conftest import random_test_id, temp_config
-from tests.core.helpers import TestFinder
 from tests.fixtures import Fixtures
-import caikit.runtime.model_management.model_loader
 
 get_dynamic_module("caikit.core")
 ANY_MODEL_TYPE = "test-any-model-type"
@@ -61,20 +58,16 @@ def temp_local_models_dir(workdir, model_manager=MODEL_MANAGER):
 @contextmanager
 def non_singleton_model_managers(num_mgrs=1, *args, **kwargs):
     with temp_config(*args, **kwargs):
-        with patch(
-            "caikit.runtime.model_management.model_loader.MODEL_MANAGER",
-            new_callable=CoreModelManager,
-        ):
-            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
+        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:
@@ -540,27 +533,6 @@ def test_load_local_model_deleted_dir():
             assert not manager.loaded_models
 
 
-def test_load_local_model_not_on_disk():
-    """Make sure bad models in local_models_dir at boot don't cause exceptions"""
-    with TemporaryDirectory() as cache_dir:
-        with non_singleton_model_managers(
-            1,
-            {
-                "model_management": {"finders": {"default": {"type": TestFinder.name}}},
-                "runtime": {
-                    "local_models_dir": cache_dir,
-                    "lazy_load_local_models": True,
-                },
-            },
-            "merge",
-        ) as managers:
-            manager = managers[0]
-            model_id = random_test_id()
-            manager.load_model(model_id, "some_non_disk_model", "test")
-            model = manager.retrieve_model(model_id)
-            assert model
-
-
 # ****************************** Unit Tests ****************************** #
 # These tests patch in mocks for the manager's dependencies, to test its code in isolation