Repository: caikit/caikit
Pull request #472 (2023-09-14 19:58:00)

Find regressed modules for backwards compatibility check
<!--  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**:
We need a check to compare a set of modules supported in a prior service generation to this current generation.
Raise an error if at least a module becomes unsupported, hence breaking backwards compatibility.

This PR:
- adds a function `assert_compatible()` in create_service.py to assert that no module was priorly supported is no longer.
- adds a flag in `get_service_package()` to write out a `modules.json` file. Should this just be an always behavior/ not a flag?
- assumes that we have a way to give `assert_compatible()` a prev module list to do the comparison with. This PR is a first round of putting it in the config.yml as a file path, but can be defined in a better way for libraries to input this.
- 
**Special notes for your reviewer**:

**If applicable**:
- [ ] this PR contains documentation
- [x] this PR contains unit tests
- [ ] this PR has been tested for backwards compatibility



Patch:
diff --git a/caikit/config/config.yml b/caikit/config/config.yml
index fdf72014..8dbe055c 100644
--- a/caikit/config/config.yml
+++ b/caikit/config/config.yml
@@ -162,6 +162,10 @@ runtime:
         task_types:
             included: []
             excluded: []
+        backwards_compatibility:
+            enabled: false
+            current_modules_path: modules.json
+            prev_modules_path: prev_modules.json
 
     # Configuration for batch inference. This dict contains entries for individual
     # model types and can be extended using additional overlay config files to add
diff --git a/caikit/runtime/dump_services.py b/caikit/runtime/dump_services.py
index f2946073..bec7f9f4 100644
--- a/caikit/runtime/dump_services.py
+++ b/caikit/runtime/dump_services.py
@@ -32,7 +32,7 @@ log = alog.use_channel("RUNTIME-DUMP-SVC")
 def dump_grpc_services(output_dir: str):
     """Utility for rendering the all generated interfaces to proto files"""
     inf_svc = ServicePackageFactory.get_service_package(
-        ServicePackageFactory.ServiceType.INFERENCE,
+        ServicePackageFactory.ServiceType.INFERENCE, write_modules_file=True
     )
     train_svc = ServicePackageFactory.get_service_package(
         ServicePackageFactory.ServiceType.TRAINING,
diff --git a/caikit/runtime/service_factory.py b/caikit/runtime/service_factory.py
index 144feddb..7c630d8f 100644
--- a/caikit/runtime/service_factory.py
+++ b/caikit/runtime/service_factory.py
@@ -17,6 +17,8 @@ from enum import Enum
 from types import ModuleType
 from typing import Callable, Dict, Set, Type, Union
 import dataclasses
+import json
+import os
 
 # Third Party
 import google.protobuf.descriptor
@@ -94,7 +96,9 @@ class ServicePackageFactory:
         TRAINING_MANAGEMENT = 3
 
     @classmethod
-    def get_service_package(cls, service_type: ServiceType) -> ServicePackage:
+    def get_service_package(
+        cls, service_type: ServiceType, write_modules_file: bool = False
+    ) -> ServicePackage:
         """Public factory API. Returns a service package of the requested type, from the
         configured source.
 
@@ -102,6 +106,9 @@ class ServicePackageFactory:
             service_type (ServicePackageFactory.ServiceType): The type of service to build,
                 to match the servicer implementation that will handle it. e.g. the
                 GlobalPredictServicer expects an "INFERENCE" service
+            write_modules_file (bool): if set, write out a modules.json file to list the
+                included modules in this service generation. See config to customize file name at
+                runtime.service_generation.backwards_compatibility.current_modules_path
 
         Returns:
             ServicePackage: A container with properties referencing everything you need to bind a
@@ -134,10 +141,15 @@ class ServicePackageFactory:
         # Then do API introspection to come up with all the API definitions to support
         caikit_config = get_config()
         clean_modules = ServicePackageFactory._get_and_filter_modules(
-            caikit_config, caikit_config.runtime.library
+            caikit_config, caikit_config.runtime.library, write_modules_file
         )
 
         if service_type == cls.ServiceType.INFERENCE:
+            # Assert for backwards compatibility, if enabled, when service type is INFERENCE
+            ServicePackageFactory._check_backwards_compatibility(
+                caikit_config, clean_modules
+            )
+
             rpc_list = service_generation.create_inference_rpcs(
                 clean_modules, caikit_config
             )
@@ -176,9 +188,36 @@ class ServicePackageFactory:
         )
 
     # Implementation details for pure python service packages #
+    @staticmethod
+    def _check_backwards_compatibility(
+        caikit_config: aconfig.Config, clean_modules: Set[Type[ModuleBase]]
+    ):
+        backwards_compat_conf = (
+            caikit_config.runtime.service_generation.backwards_compatibility
+        )
+        if backwards_compat_conf and backwards_compat_conf.enabled:
+            previous_included_modules = set()
+            prev_modules_path = backwards_compat_conf.prev_modules_path
+            error.value_check(
+                "<SVC98335764E>",
+                os.path.exists(prev_modules_path) and os.path.isfile(prev_modules_path),
+                "prev_modules_path {} is not a valid file path or is missing permissions",
+                prev_modules_path,
+            )
+            with open(prev_modules_path, "r", encoding="utf-8") as f:
+                previous_modules = json.load(f)
+                previous_included_task_map = previous_modules["included_modules"]
+                for task_module in previous_included_task_map.values():
+                    previous_included_modules.update(task_module.keys())
+
+            service_generation.assert_compatible(
+                [mod.MODULE_ID for mod in clean_modules],
+                previous_included_modules,
+            )
+
     @staticmethod
     def _get_and_filter_modules(
-        caikit_config: aconfig.Config, lib: str
+        caikit_config: aconfig.Config, lib: str, write_modules_file: bool
     ) -> Set[Type[ModuleBase]]:
         clean_modules = set()
         modules = [
@@ -244,6 +283,26 @@ class ServicePackageFactory:
             clean_modules,
             excluded_modules,
         )
+
+        # if enabled, write out the inclusions to modules.json
+        backwards_compat_conf = (
+            caikit_config.runtime.service_generation.backwards_compatibility
+        )
+        if write_modules_file:
+            modules_json_path = (
+                backwards_compat_conf and backwards_compat_conf.current_modules_path
+            ) or "modules.json"
+            included_dict = {}
+            for module in clean_modules:
+                for task_type in module.tasks:
+                    included_dict.setdefault(task_type.__name__, {})[
+                        module.MODULE_ID
+                    ] = str(module)
+            modules_dict = {
+                "included_modules": included_dict,
+            }
+            with open(modules_json_path, "w", encoding="utf-8") as f:
+                json.dump(modules_dict, f, indent=4, sort_keys=True)
         return clean_modules
 
 
diff --git a/caikit/runtime/service_generation/__init__.py b/caikit/runtime/service_generation/__init__.py
index 1fb7635e..ebbe6280 100644
--- a/caikit/runtime/service_generation/__init__.py
+++ b/caikit/runtime/service_generation/__init__.py
@@ -1,4 +1,8 @@
 # Local
 from . import create_service
-from .create_service import create_inference_rpcs, create_training_rpcs
+from .create_service import (
+    assert_compatible,
+    create_inference_rpcs,
+    create_training_rpcs,
+)
 from .proto_package import get_ai_domain, get_runtime_service_package
diff --git a/caikit/runtime/service_generation/create_service.py b/caikit/runtime/service_generation/create_service.py
index 3ef0ee21..2bf369b4 100644
--- a/caikit/runtime/service_generation/create_service.py
+++ b/caikit/runtime/service_generation/create_service.py
@@ -26,9 +26,11 @@ import alog
 # Local
 from .rpcs import CaikitRPCBase, ModuleClassTrainRPC, TaskPredictRPC
 from caikit.core import ModuleBase, TaskBase
+from caikit.core.exceptions import error_handler
 from caikit.core.signature_parsing.module_signature import CaikitMethodSignature
 
 log = alog.use_channel("CREATE-RPCS")
+error = error_handler.get(log)
 
 ## Globals #####################################################################
 
@@ -37,6 +39,33 @@ TRAIN_FUNCTION_NAME = "train"
 ## Utilities ###################################################################
 
 
+def assert_compatible(modules: List[str], previous_modules: List[str]):
+    """Logic about whether it's okay to include this set of modules in service generation
+
+    Args:
+        modules: list of module IDs that we are considering in service generation
+        previous_modules: list of module IDs that were supported in the previous service version
+
+    Raises:
+        If a new service should not be built with this set of modules
+    """
+    regressed_modules = set(previous_modules) - set(modules)
+    if len(regressed_modules) > 0:
+        log.error(
+            "BREAKING CHANGE FOUND! These modules became unsupported. These models were "
+            "on the supported list in previous version, but now are no longer supported."
+        )
+        for mod in regressed_modules:
+            log.error("Regressed module: %s", mod)
+
+    error.value_check(
+        "<SVC68235724E>",
+        len(regressed_modules) == 0,
+        "BREAKING CHANGE! Found unsupported module(s) that were previously supported: {}",
+        regressed_modules,
+    )
+
+
 def create_inference_rpcs(
     modules: List[Type[ModuleBase]], caikit_config: aconfig.Config = None
 ) -> List[CaikitRPCBase]:
diff --git a/examples/sample_lib/start_runtime_with_sample_lib.py b/examples/sample_lib/start_runtime_with_sample_lib.py
index 5fd3b022..4cda5176 100644
--- a/examples/sample_lib/start_runtime_with_sample_lib.py
+++ b/examples/sample_lib/start_runtime_with_sample_lib.py
@@ -61,3 +61,4 @@ if __name__ == "__main__":
     finally:
         # remove generated protos
         shutil.rmtree("protos", ignore_errors=True)
+        os.remove("modules.json")


Test patch:
diff --git a/tests/runtime/service_generation/test_create_service.py b/tests/runtime/service_generation/test_create_service.py
index c4261702..8248a36a 100644
--- a/tests/runtime/service_generation/test_create_service.py
+++ b/tests/runtime/service_generation/test_create_service.py
@@ -15,9 +15,13 @@
 from typing import Iterable
 import uuid
 
+# Third Party
+import pytest
+
 # Local
 from caikit.core import LocalBackend
 from caikit.runtime.service_generation.create_service import (
+    assert_compatible,
     create_inference_rpcs,
     create_training_rpcs,
 )
@@ -335,3 +339,38 @@ def test_create_training_rpcs():
     rpcs = create_training_rpcs([widget_class])
     assert len(rpcs) == 1
     assert widget_class in rpcs[0].module_list
+
+
+### assert_compatible tests #################################################
+def test_assert_compatible_does_not_raise_if_modules_continue_to_be_supported():
+    previous_module_list = [
+        sample_lib.modules.sample_task.SampleModule.MODULE_ID,
+        sample_lib.modules.sample_task.InnerModule.MODULE_ID,
+    ]
+
+    current_module_list = [
+        sample_lib.modules.sample_task.SampleModule.MODULE_ID,
+        sample_lib.modules.sample_task.InnerModule.MODULE_ID,
+        sample_lib.modules.other_task.OtherModule.MODULE_ID,
+    ]
+    assert_compatible(current_module_list, previous_module_list)
+
+
+def test_assert_compatible_raises_if_modules_are_no_longer_supported():
+    previous_module_list = [
+        sample_lib.modules.sample_task.SampleModule.MODULE_ID,
+        sample_lib.modules.sample_task.InnerModule.MODULE_ID,
+    ]
+
+    current_module_list = [
+        sample_lib.modules.sample_task.SampleModule.MODULE_ID,
+        sample_lib.modules.other_task.OtherModule.MODULE_ID,
+    ]  # missing InnerModule from prev
+
+    with pytest.raises(ValueError) as context:
+        assert_compatible(current_module_list, previous_module_list)
+    assert (
+        "BREAKING CHANGE! Found unsupported module(s) that were previously supported: "
+        in str(context.value)
+        and sample_lib.modules.sample_task.InnerModule.MODULE_ID in str(context.value)
+    )
diff --git a/tests/runtime/test_service_factory.py b/tests/runtime/test_service_factory.py
index f2fca653..d36d12af 100644
--- a/tests/runtime/test_service_factory.py
+++ b/tests/runtime/test_service_factory.py
@@ -15,7 +15,10 @@
 """Unit tests for the service factory"""
 # Standard
 from pathlib import Path
+import json
+import os
 import tempfile
+import uuid
 
 # Third Party
 from google.protobuf.message import Message
@@ -120,7 +123,9 @@ def test_get_and_filter_modules_respects_excluded_modules():
             }
         }
     ) as cfg:
-        clean_modules = ServicePackageFactory._get_and_filter_modules(cfg, "sample_lib")
+        clean_modules = ServicePackageFactory._get_and_filter_modules(
+            cfg, "sample_lib", False
+        )
         assert "ListModule" not in str(clean_modules)
         assert "SampleModule" in str(clean_modules)
         assert "OtherModule" in str(clean_modules)
@@ -136,12 +141,117 @@ def test_get_and_filter_modules_respects_included_modules():
             }
         }
     ) as cfg:
-        clean_modules = ServicePackageFactory._get_and_filter_modules(cfg, "sample_lib")
+        clean_modules = ServicePackageFactory._get_and_filter_modules(
+            cfg, "sample_lib", False
+        )
         assert len(clean_modules) == 1
         assert "ListModule" in str(clean_modules)
         assert "SampleModule" not in str(clean_modules)
 
 
+def test_assert_compatible_raises_if_prev_modules_path_is_not_valid():
+    with temp_config(
+        {
+            "runtime": {
+                "service_generation": {
+                    "backwards_compatibility": {
+                        "enabled": True,
+                        "prev_modules_path": "foobar.json",  # a file that's not valid
+                    }
+                },
+            }
+        },
+        "merge",
+    ):
+        with pytest.raises(ValueError) as e:
+            ServicePackageFactory.get_service_package(
+                ServicePackageFactory.ServiceType.INFERENCE
+            )
+
+
+def test_assert_compatible_does_not_raise_if_not_enabled():
+    with temp_config(
+        {
+            "runtime": {
+                "service_generation": {
+                    "backwards_compatibility": {
+                        "enabled": False,
+                    }
+                },
+            }
+        },
+        "merge",
+    ):
+        ServicePackageFactory.get_service_package(
+            ServicePackageFactory.ServiceType.INFERENCE
+        )
+
+
+def test_assert_compatible_raises_if_a_module_becomes_unsupported():
+    with tempfile.TemporaryDirectory() as workdir:
+        prev_module_file = os.path.join(workdir, "prev_modules.json")
+        random_uuid = str(uuid.uuid4())
+        with open(prev_module_file, "w", encoding="utf-8") as file:
+            json_content = {
+                "included_modules": {
+                    "SampleTask": {
+                        random_uuid: "<class 'sample_lib.modules.sample_task.sample_implementation.PrevSampleModule'>",
+                    }
+                },
+            }
+            file.write(json.dumps(json_content, indent=4))
+        with temp_config(
+            {
+                "runtime": {
+                    "service_generation": {
+                        "backwards_compatibility": {
+                            "enabled": True,
+                            "prev_modules_path": prev_module_file,
+                        }
+                    },
+                }
+            },
+            "merge",
+        ):
+            with pytest.raises(ValueError):
+                # this raises because PrevSampleModule will not be included in this service generation
+                ServicePackageFactory.get_service_package(
+                    ServicePackageFactory.ServiceType.INFERENCE
+                )
+
+
+def test_assert_compatible_does_not_raise_if_supported_modules_continue_to_be_supported():
+    with tempfile.TemporaryDirectory() as workdir:
+        prev_module_file = os.path.join(workdir, "prev_modules.json")
+
+        with open(prev_module_file, "w", encoding="utf-8") as file:
+            json_content = {
+                "included_modules": {
+                    "SampleTask": {
+                        "00110203-0405-0607-0809-0a0b02dd0e0f": "<class 'sample_lib.modules.sample_task.sample_implementation.SampleModule'>",
+                    }
+                },
+            }
+            file.write(json.dumps(json_content, indent=4))
+        with temp_config(
+            {
+                "runtime": {
+                    "service_generation": {
+                        "backwards_compatibility": {
+                            "enabled": True,
+                            "prev_modules_path": prev_module_file,
+                        }
+                    },
+                }
+            },
+            "merge",
+        ):
+            # this does not raise because SampleModule will be included in this service generation
+            ServicePackageFactory.get_service_package(
+                ServicePackageFactory.ServiceType.INFERENCE
+            )
+
+
 def test_override_domain(clean_data_model):
     """
     Test override of gRPC domain generation from config.
@@ -173,7 +283,9 @@ def test_override_domain(clean_data_model):
         )
 
         # Just double-check that basics are good.
-        clean_modules = ServicePackageFactory._get_and_filter_modules(cfg, "sample_lib")
+        clean_modules = ServicePackageFactory._get_and_filter_modules(
+            cfg, "sample_lib", False
+        )
         assert "SampleModule" in str(clean_modules)
 
 
@@ -202,7 +314,9 @@ def test_override_package(clean_data_model):
         )
 
         # Just double-check that basics are good.
-        clean_modules = ServicePackageFactory._get_and_filter_modules(cfg, "sample_lib")
+        clean_modules = ServicePackageFactory._get_and_filter_modules(
+            cfg, "sample_lib", False
+        )
         assert "SampleModule" in str(clean_modules)
 
 
@@ -238,7 +352,9 @@ def test_override_package_and_domain_with_proto_gen(clean_data_model):
         )
 
         # Just double-check that basics are good.
-        clean_modules = ServicePackageFactory._get_and_filter_modules(cfg, "sample_lib")
+        clean_modules = ServicePackageFactory._get_and_filter_modules(
+            cfg, "sample_lib", False
+        )
         assert "SampleModule" in str(clean_modules)
 
         # Full check on proto generation