Repository: caikit/caikit (python)

Issue:
Backend modules don't participate in service generation
## Describe the bug

A backend module (e.g. https://github.com/caikit/caikit-nlp/blob/0f14b486ef7840cef5a39863523c75ab06e7a650/caikit_nlp/modules/text_generation/text_generation_tgis.py#L45) may have a `base_module` with a `task`, and may define `@taskmethods`, but the parameteres from those task methods don't make it into the generated runtime services. 

The root of the problem is that service generation only looks at plain modules, and not backend modules.

## Platform

Please provide details about the environment you are using, including the following:

- Interpreter version:
- Library version:

## Sample Code

Dump the services out for `caikit_nlp` and see that there is no `preserve_input_text` parameter anywhere.

## Expected behavior

The text gen interfaces should include `preserve_input_text`

## Observed behavior


## Additional context



Pull request description:
🐛  Backends fix
<!--  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**:

Fixes #446 

This includes backend modules in the service generation.
It also fixes our `reset_globals` fixture since that wasn't properly cleaning up backend modules.
It also moves the `reset_*` stuff into `conftest.py` so that
- We don't need squirrely imports in our tests, and
- The runtime tests can use them too

**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/registries.py b/caikit/core/registries.py
index 53305a30..97ff6b27 100644
--- a/caikit/core/registries.py
+++ b/caikit/core/registries.py
@@ -65,7 +65,7 @@ def module_backend_registry() -> Dict[
     Dict[ module_id, Dict[ backend_type, Tuple[ backend_impl_class, backend_config_dict ] ] ]
 
     Returns:
-        Dict[str, Dict[str, Tuple["caikit.core.BackendBase", Dict]]]: The module
+        Dict[str, Dict[str, Tuple["caikit.core.ModuleBase", Dict]]]: The module
             backend registry
     """
     # TODO: put a real data structure here instead of nested dicts
diff --git a/caikit/runtime/service_factory.py b/caikit/runtime/service_factory.py
index 9713b3ec..8e0307de 100644
--- a/caikit/runtime/service_factory.py
+++ b/caikit/runtime/service_factory.py
@@ -30,7 +30,7 @@ import alog
 
 # Local
 from caikit import get_config
-from caikit.core import ModuleBase
+from caikit.core import LocalBackend, ModuleBase, registries
 from caikit.interfaces.runtime.data_model import (
     TrainingInfoRequest,
     TrainingStatusResponse,
@@ -38,7 +38,6 @@ from caikit.interfaces.runtime.data_model import (
 from caikit.runtime import service_generation
 from caikit.runtime.service_generation.rpcs import CaikitRPCBase
 from caikit.runtime.utils import import_util
-import caikit.core
 
 log = alog.use_channel("SVC-FACTORY")
 
@@ -178,9 +177,18 @@ class ServicePackageFactory:
         clean_modules = set()
         modules = [
             module_class
-            for module_class in caikit.core.registries.module_registry().values()
+            for module_class in registries.module_registry().values()
             if module_class.__module__.partition(".")[0] == lib
         ]
+        # NB: The `module_registry` only includes the `LOCAL` backend modules.
+        # Implementations of the same module for different backends need to be fetched from the
+        # backend registry
+        backend_modules = []
+        for backend_dict in registries.module_backend_registry().values():
+            for backend, config in backend_dict.items():
+                if backend != LocalBackend.backend_type:
+                    backend_modules.append(config.impl_class)
+        modules.extend(backend_modules)
         log.debug("Found all modules %s for library %s.", modules, lib)
 
         # Check config for any explicit inclusions


Test patch:
diff --git a/tests/conftest.py b/tests/conftest.py
index f52a882d..3e9896a7 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -21,6 +21,13 @@ import alog
 
 # Local
 from caikit import get_config
+from caikit.core import MODEL_MANAGER
+from caikit.core.registries import (
+    module_backend_classes,
+    module_backend_registry,
+    module_backend_types,
+    module_registry,
+)
 from caikit.core.toolkit import logging
 import caikit
 
@@ -225,3 +232,58 @@ class TempFailWrapper:
 # IMPLEMENTATION DETAILS ############################################################
 def _random_id():
     return str(uuid.uuid4())
+
+
+@pytest.fixture
+def reset_backend_types():
+    """Fixture that will reset the backend types if a test modifies them"""
+    base_backend_types = {key: val for key, val in module_backend_types().items()}
+    base_backend_classes = {key: val for key, val in module_backend_classes().items()}
+    yield
+    module_backend_types().clear()
+    module_backend_types().update(base_backend_types)
+    module_backend_classes().clear()
+    module_backend_classes().update(base_backend_classes)
+
+
+@pytest.fixture
+def reset_module_backend_registry():
+    """Fixture that will reset the module distribution registry if a test modifies them"""
+    # NB: Deepcopy here because the values are nested dicts that will be persisted by reference
+    orig_module_backend_registry = copy.deepcopy(module_backend_registry())
+    yield
+    module_backend_registry().clear()
+    module_backend_registry().update(orig_module_backend_registry)
+
+
+@pytest.fixture
+def reset_module_registry():
+    """Fixture that will reset caikit.core module registry if a test modifies it"""
+    orig_module_registry = {key: val for key, val in module_registry().items()}
+    yield
+    module_registry().clear()
+    module_registry().update(orig_module_registry)
+
+
+@pytest.fixture
+def reset_model_manager():
+    prev_finders = MODEL_MANAGER._finders
+    prev_initializers = MODEL_MANAGER._initializers
+    prev_trainers = MODEL_MANAGER._trainers
+    MODEL_MANAGER._finders = {}
+    MODEL_MANAGER._initializers = {}
+    MODEL_MANAGER._trainers = {}
+    yield
+    MODEL_MANAGER._finders = prev_finders
+    MODEL_MANAGER._initializers = prev_initializers
+    MODEL_MANAGER._trainers = prev_trainers
+
+
+@pytest.fixture
+def reset_globals(
+    reset_backend_types,
+    reset_model_manager,
+    reset_module_backend_registry,
+    reset_module_registry,
+):
+    """Fixture that will reset the backend types and module registries if a test modifies them"""
diff --git a/tests/core/helpers.py b/tests/core/helpers.py
index df6b2020..ea3e83cf 100644
--- a/tests/core/helpers.py
+++ b/tests/core/helpers.py
@@ -17,9 +17,6 @@ from typing import Optional, Type
 import os
 import uuid
 
-# Third Party
-import pytest
-
 # Local
 from caikit.core import MODEL_MANAGER
 from caikit.core.data_model import TrainingStatus
@@ -35,12 +32,6 @@ from caikit.core.model_management import (
 from caikit.core.model_management.local_model_initializer import LocalModelInitializer
 from caikit.core.module_backends import BackendBase, backend_types
 from caikit.core.modules import ModuleBase, ModuleConfig
-from caikit.core.registries import (
-    module_backend_classes,
-    module_backend_registry,
-    module_backend_types,
-    module_registry,
-)
 from sample_lib.modules import SampleModule
 
 
@@ -192,59 +183,3 @@ def configured_backends():
         if isinstance(loader, LocalModelInitializer)
     ]
     return [backend for loader in local_initializers for backend in loader._backends]
-
-
-@pytest.fixture
-def reset_backend_types():
-    """Fixture that will reset the backend types if a test modifies them"""
-    base_backend_types = {key: val for key, val in module_backend_types().items()}
-    base_backend_classes = {key: val for key, val in module_backend_classes().items()}
-    yield
-    module_backend_types().clear()
-    module_backend_types().update(base_backend_types)
-    module_backend_classes().clear()
-    module_backend_classes().update(base_backend_classes)
-
-
-@pytest.fixture
-def reset_module_backend_registry():
-    """Fixture that will reset the module distribution registry if a test modifies them"""
-    orig_module_backend_registry = {
-        key: val for key, val in module_backend_registry().items()
-    }
-    yield
-    module_backend_registry().clear()
-    module_backend_registry().update(orig_module_backend_registry)
-
-
-@pytest.fixture
-def reset_module_registry():
-    """Fixture that will reset caikit.core module registry if a test modifies it"""
-    orig_module_registry = {key: val for key, val in module_registry().items()}
-    yield
-    module_registry().clear()
-    module_registry().update(orig_module_registry)
-
-
-@pytest.fixture
-def reset_model_manager():
-    prev_finders = MODEL_MANAGER._finders
-    prev_initializers = MODEL_MANAGER._initializers
-    prev_trainers = MODEL_MANAGER._trainers
-    MODEL_MANAGER._finders = {}
-    MODEL_MANAGER._initializers = {}
-    MODEL_MANAGER._trainers = {}
-    yield
-    MODEL_MANAGER._finders = prev_finders
-    MODEL_MANAGER._initializers = prev_initializers
-    MODEL_MANAGER._trainers = prev_trainers
-
-
-@pytest.fixture
-def reset_globals(
-    reset_backend_types,
-    reset_model_manager,
-    reset_module_backend_registry,
-    reset_module_registry,
-):
-    """Fixture that will reset the backend types and module registries if a test modifies them"""
diff --git a/tests/core/module_backends/test_backend_types.py b/tests/core/module_backends/test_backend_types.py
index 66ef0fb4..966394f7 100644
--- a/tests/core/module_backends/test_backend_types.py
+++ b/tests/core/module_backends/test_backend_types.py
@@ -24,7 +24,8 @@ import pytest
 # Local
 from caikit.core import backend_types
 from caikit.core.module_backends.base import BackendBase
-from tests.core.helpers import MockBackend, reset_backend_types
+from tests.conftest import reset_backend_types
+from tests.core.helpers import MockBackend
 
 ## Tests #######################################################################
 
diff --git a/tests/core/modules/test_module.py b/tests/core/modules/test_module.py
index 5a9048da..eaab139f 100644
--- a/tests/core/modules/test_module.py
+++ b/tests/core/modules/test_module.py
@@ -17,12 +17,16 @@ import io
 import os
 import tempfile
 
+# Third Party
+import pytest
+
 # First Party
 import aconfig
 
 # Local
 from caikit.core import ModuleConfig, ModuleLoader, ModuleSaver
 from caikit.core.modules.decorator import SUPPORTED_LOAD_BACKENDS_VAR_NAME
+from caikit.core.registries import module_backend_registry, module_registry
 
 # pylint: disable=import-error
 from sample_lib.data_model.sample import SampleInputType
diff --git a/tests/core/test_model_manager.py b/tests/core/test_model_manager.py
index b9a003d1..7a8e869e 100644
--- a/tests/core/test_model_manager.py
+++ b/tests/core/test_model_manager.py
@@ -20,6 +20,9 @@ import os
 import tempfile
 import uuid
 
+# Third Party
+import pytest
+
 # Local
 from caikit.core import LocalBackend
 from caikit.core.data_model import DataStream, TrainingStatus
diff --git a/tests/runtime/model_management/test_model_loader.py b/tests/runtime/model_management/test_model_loader.py
index 812bdf3f..d17da8af 100644
--- a/tests/runtime/model_management/test_model_loader.py
+++ b/tests/runtime/model_management/test_model_loader.py
@@ -14,9 +14,7 @@
 # Standard
 from concurrent.futures import Future
 from contextlib import contextmanager
-from typing import Callable
 from unittest import mock
-import copy
 import tempfile
 
 # Third Party
@@ -231,67 +229,50 @@ def test_with_batching_collect_delay(model_loader):
         )
 
 
-def test_load_distributed_impl():
+def test_load_distributed_impl(reset_globals):
     """Make sure that when configured, an alternate distributed
     implementation of a module can be loaded
     """
 
-    reg_copy = copy.deepcopy(caikit.core.registries.module_registry())
-    backend_registry_copy = copy.deepcopy(
-        caikit.core.registries.module_backend_registry()
+    @module(
+        base_module=SampleModule,
+        backend_type=backend_types.MOCK,
+        backend_config_override={"bar1": 1},
     )
-    # 🌶️🌶️🌶️: the MODULE_BACKEND_REGISTRY can't be easily patched since two separate modules hold
-    # an imported reference to it and one edits it (decorator.py) while the other reads it
-    # (model_manager.py)
-
-    with mock.patch.object(caikit.core.registries, "MODULE_REGISTRY", reg_copy):
-        with mock.patch.object(
-            caikit.core.registries,
-            "MODULE_BACKEND_REGISTRY",
-            backend_registry_copy,
-        ):
-
-            @module(
-                base_module=SampleModule,
-                backend_type=backend_types.MOCK,
-                backend_config_override={"bar1": 1},
-            )
-            class DistributedGadget(caikit.core.ModuleBase):
-                """An alternate implementation of a Gadget"""
-
-                SUPPORTED_LOAD_BACKENDS = [
-                    MockBackend.backend_type,
-                    backend_types.LOCAL,
-                ]
-
-                def __init__(self, bar):
-                    self.bar = bar
-
-                def run(self, sample_input: SampleInputType) -> SampleOutputType:
-                    return SampleOutputType(
-                        greeting=f"hello distributed {sample_input.name}"
-                    )
-
-                @classmethod
-                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)
-
-            with tempfile.TemporaryDirectory() as model_path:
-                # Create and save the model directly with the local impl
-                SampleModule().save(model_path)
-
-                model_type = "gadget"
-
-                with temp_model_loader() as model_loader:
-                    # Load the distributed version
-                    model = model_loader.load_model(
-                        random_test_id(),
-                        model_path,
-                        model_type=model_type,
-                    ).model()
-                    assert isinstance(model, DistributedGadget)
+    class DistributedGadget(caikit.core.ModuleBase):
+        """An alternate implementation of a Gadget"""
+
+        SUPPORTED_LOAD_BACKENDS = [
+            MockBackend.backend_type,
+            backend_types.LOCAL,
+        ]
+
+        def __init__(self, bar):
+            self.bar = bar
+
+        def run(self, sample_input: SampleInputType) -> SampleOutputType:
+            return SampleOutputType(greeting=f"hello distributed {sample_input.name}")
+
+        @classmethod
+        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)
+
+    with tempfile.TemporaryDirectory() as model_path:
+        # Create and save the model directly with the local impl
+        SampleModule().save(model_path)
+
+        model_type = "gadget"
+
+        with temp_model_loader() as model_loader:
+            # Load the distributed version
+            model = model_loader.load_model(
+                random_test_id(),
+                model_path,
+                model_type=model_type,
+            ).model()
+            assert isinstance(model, DistributedGadget)
 
 
 def test_load_model_without_waiting_success(model_loader):
diff --git a/tests/runtime/service_generation/test_create_service.py b/tests/runtime/service_generation/test_create_service.py
index 4bf37880..22a838d0 100644
--- a/tests/runtime/service_generation/test_create_service.py
+++ b/tests/runtime/service_generation/test_create_service.py
@@ -16,6 +16,7 @@ from typing import Iterable
 import uuid
 
 # Local
+from caikit.core import LocalBackend
 from caikit.runtime.service_generation.create_service import (
     create_inference_rpcs,
     create_training_rpcs,
@@ -38,6 +39,9 @@ widget_class = sample_lib.modules.sample_task.SampleModule
 
 ### create_inference_rpcs tests #################################################
 
+# Local
+from tests.core.helpers import MockBackend
+
 
 def test_create_inference_rpcs_uses_task_from_module_decorator():
     # make a new module with SampleTask
@@ -55,6 +59,30 @@ def test_create_inference_rpcs_uses_task_from_module_decorator():
     assert SampleModule in rpcs[0].module_list
 
 
+def test_create_inference_rpcs_includes_backend_modules():
+    @caikit.module(
+        id=str(uuid.uuid4()), name="something", version="0.0.0", task=SampleTask
+    )
+    class NewModule(caikit.core.ModuleBase):
+        def run(self, sample_input: SampleInputType) -> SampleOutputType:
+            pass
+
+    @caikit.module(backend_type=MockBackend.backend_type, base_module=NewModule)
+    class NewBackendModule(caikit.core.ModuleBase):
+        def run(
+            self, sample_input: SampleInputType, backend_param: str
+        ) -> SampleOutputType:
+            pass
+
+    rpcs = create_inference_rpcs([NewModule, NewBackendModule])
+    assert len(rpcs) == 1
+    assert NewModule in rpcs[0].module_list
+    # Make sure the backend module and its parameters are added as well
+    assert NewBackendModule in rpcs[0].module_list
+    assert len(rpcs[0].request.triples) == 2
+    assert rpcs[0].request.triples[1][1] == "backend_param"
+
+
 def test_create_inference_rpcs_uses_task_from_module_decorator_with_streaming():
     @caikit.module(
         id=str(uuid.uuid4()),
diff --git a/tests/runtime/servicers/test_global_train_servicer_impl.py b/tests/runtime/servicers/test_global_train_servicer_impl.py
index 6f2756ea..4dc15f1e 100644
--- a/tests/runtime/servicers/test_global_train_servicer_impl.py
+++ b/tests/runtime/servicers/test_global_train_servicer_impl.py
@@ -38,8 +38,7 @@ from sample_lib.data_model.sample import (
     SampleTrainingType,
 )
 from sample_lib.modules.sample_task.sample_implementation import SampleModule
-from tests.conftest import random_test_id, temp_config
-from tests.core.helpers import reset_model_manager
+from tests.conftest import random_test_id, reset_model_manager, temp_config
 from tests.fixtures import Fixtures
 from tests.runtime.conftest import register_trained_model
 import caikit.core
diff --git a/tests/runtime/test_service_factory.py b/tests/runtime/test_service_factory.py
index 0b84c80d..1ecefaf8 100644
--- a/tests/runtime/test_service_factory.py
+++ b/tests/runtime/test_service_factory.py
@@ -24,8 +24,11 @@ import pytest
 # Local
 from caikit.core.data_model import render_dataobject_protos
 from caikit.runtime.service_factory import ServicePackage, ServicePackageFactory
+from sample_lib import SampleModule
+from sample_lib.data_model import SampleInputType, SampleOutputType
 from sample_lib.modules.sample_task import ListModule
 from tests.conftest import temp_config
+from tests.core.helpers import MockBackend
 from tests.data_model_helpers import reset_global_protobuf_registry, temp_dpool
 from tests.runtime.conftest import sample_inference_service, sample_train_service
 import caikit
@@ -340,3 +343,23 @@ def test_override_package_and_domain_with_proto_gen(clean_data_model):
                                     )
                                 else:
                                     assert service_name == f"{domain_override}Service"
+
+
+def test_backend_modules_included_in_service_generation(
+    clean_data_model, reset_globals
+):
+    # Add a new backend module for the good ol' `SampleModule`
+    @caikit.module(backend_type=MockBackend.backend_type, base_module=SampleModule)
+    class NewBackendModule(caikit.core.ModuleBase):
+        def run(
+            self, sample_input: SampleInputType, backend_param: str
+        ) -> SampleOutputType:
+            pass
+
+    inference_service = ServicePackageFactory.get_service_package(
+        ServicePackageFactory.ServiceType.INFERENCE
+    )
+    sample_task_request = inference_service.messages.SampleTaskRequest
+
+    # Check that the new parameter defined in this backend module exists in the service
+    assert "backend_param" in sample_task_request.DESCRIPTOR.fields_by_name.keys()