Repository: caikit/caikit
Pull request #32 (2023-04-13 19:41:48)

Allow ServicePackageFactory to read config's lists of exclusions and inclusions
Issue https://github.com/caikit/caikit/issues/8

This PR allows for users to specify an `included` and `excluded` fields in `service_generation` that ServicePackageFactory will respect.
Example:

```
service_generation:
    enabled: true
    modules:
        included: ["00110203-baad-beef-0809-0a0b02dd0e0f"],
        excluded: ["00af2203-0405-0607-0263-0a0b02dd0c2f"]
    task_types:
        included: ["other_task"]
        excluded: ["sample_task"]
```

- If a module happens to be on both `included` and `excluded` lists, it's excluded
- If `included` and `excluded` field is `[]` or `null`, it equates to no-op


Patch:
diff --git a/.gitignore b/.gitignore
index 08f2dc05..a5a1e139 100644
--- a/.gitignore
+++ b/.gitignore
@@ -11,6 +11,9 @@ test
 # IDEs
 .vscode/
 .idea/
+
+# Virtual Env
+venv/
 # Mac personalization files
 *.DS_Store
 
diff --git a/caikit/runtime/config/config.yml b/caikit/runtime/config/config.yml
index 3c7150d5..15ae8a34 100644
--- a/caikit/runtime/config/config.yml
+++ b/caikit/runtime/config/config.yml
@@ -7,10 +7,12 @@ service_generation:
     client_package_name: ""
     module_types:
         included: ["blocks", "workflows"]
-    modules:
+    module_guids:
+        included: []
         excluded: []
     primitive_data_model_types: []
     task_types:
+        included: []
         excluded: []
 
 use_abortable_threads: true
@@ -143,11 +145,13 @@ test:
         client_package_name: ""
         module_types:
             included: ["blocks", "workflows"]
-        modules:
+        module_guids:
+            included: []
             excluded: []
         primitive_data_model_types:
             - "sample_lib.data_model.SampleInputType"
         task_types:
+            included: []
             excluded: []
     # training configs for tests
     training:
diff --git a/caikit/runtime/service_factory.py b/caikit/runtime/service_factory.py
index 1e5e1863..8342ca6c 100644
--- a/caikit/runtime/service_factory.py
+++ b/caikit/runtime/service_factory.py
@@ -15,7 +15,7 @@
 # Standard
 from enum import Enum
 from types import ModuleType
-from typing import Callable, Dict, List, Type
+from typing import Callable, Dict, List, Set, Type
 import dataclasses
 import inspect
 
@@ -36,11 +36,13 @@ import alog
 # Local
 from caikit.core import dataobject
 from caikit.core.data_model.base import DataBase
+from caikit.core.module import ModuleBase
 from caikit.interfaces.runtime.data_model import (
     TrainingInfoRequest,
     TrainingInfoResponse,
 )
 from caikit.runtime import service_generation
+from caikit.runtime.service_generation.core_module_helpers import get_module_info
 from caikit.runtime.service_generation.serializers import (
     RPCSerializerBase,
     snake_to_upper_camel,
@@ -187,21 +189,21 @@ class ServicePackageFactory:
             # !!!! This will use the `caikit_library` config
             _ = import_util.get_data_model()
 
-            lib = ConfigParser.get_instance().caikit_library
+            config_parser = ConfigParser.get_instance()
+            lib = config_parser.caikit_library
             ai_domain_name = snake_to_upper_camel(lib.replace("caikit_", ""))
             package_name = f"caikit.runtime.{ai_domain_name}"
 
             # Then do API introspection to come up with all the API definitions to support
-            modules = [
-                module_class
-                for module_class in caikit.core.MODULE_REGISTRY.values()
-                if module_class.__module__.partition(".")[0] == lib
-            ]
+            clean_modules = ServicePackageFactory._get_and_filter_modules(
+                config_parser, lib
+            )
+
             if service_type == cls.ServiceType.INFERENCE:
-                task_rpc_list = service_generation.create_inference_rpcs(modules)
+                task_rpc_list = service_generation.create_inference_rpcs(clean_modules)
                 service_name = f"{ai_domain_name}Service"
             else:  # service_type == cls.ServiceType.TRAINING
-                task_rpc_list = service_generation.create_training_rpcs(modules)
+                task_rpc_list = service_generation.create_training_rpcs(clean_modules)
                 service_name = f"{ai_domain_name}TrainingService"
 
             for rpc in task_rpc_list:
@@ -242,6 +244,75 @@ class ServicePackageFactory:
             )
 
     # Implementation details for pure python service packages #
+    @staticmethod
+    def _get_and_filter_modules(
+        config_parser: ConfigParser, lib: str
+    ) -> Set[Type[ModuleBase]]:
+        clean_modules = set()
+        modules = [
+            module_class
+            for module_class in caikit.core.MODULE_REGISTRY.values()
+            if module_class.__module__.partition(".")[0] == lib
+        ]
+        log.debug("Found all modules %s for library %s.", modules, lib)
+
+        # Check config for any explicit inclusions
+        included_task_types = (
+            config_parser.service_generation
+            and config_parser.service_generation.task_types
+            and config_parser.service_generation.task_types.included
+        )
+        included_modules = (
+            config_parser.service_generation
+            and config_parser.service_generation.module_guids
+            and config_parser.service_generation.module_guids.included
+        )
+
+        # Check config for any exclusions
+        excluded_task_types = (
+            config_parser.service_generation
+            and config_parser.service_generation.task_types
+            and config_parser.service_generation.task_types.excluded
+        )
+        excluded_modules = (
+            config_parser.service_generation
+            and config_parser.service_generation.module_guids
+            and config_parser.service_generation.module_guids.excluded
+        )
+
+        for ck_module in modules:
+            # Only create for modules from defined included and exclusion list
+            module_info = get_module_info(ck_module)
+            if excluded_task_types and module_info.type in excluded_task_types:
+                log.debug("Skipping module %s of type %s", ck_module, module_info.type)
+                continue
+
+            if excluded_modules and ck_module.MODULE_ID in excluded_modules:
+                log.debug("Skipping module %s of id %s", ck_module, ck_module.MODULE_ID)
+                continue
+
+            # no inclusions specified means include everything
+            if (included_task_types is None or included_task_types == []) and (
+                included_modules is None or included_modules == []
+            ):
+                clean_modules.add(ck_module)
+
+            # if inclusion is specified, use that
+            else:
+                if (included_modules and ck_module.MODULE_ID in included_modules) or (
+                    included_task_types and module_info.type in included_task_types
+                ):
+                    clean_modules.add(ck_module)
+
+        log.debug(
+            "Filtered list of modules %s after excluding task types: %s and modules ids: %s. \
+                Exclusions are defined in config",
+            clean_modules,
+            excluded_task_types,
+            excluded_modules,
+        )
+        return clean_modules
+
     @staticmethod
     def _create_request_message_types(
         rpcs_list: List[RPCSerializerBase],


Test patch:
diff --git a/tests/runtime/test_service_factory.py b/tests/runtime/test_service_factory.py
index 342d18c7..61c6e642 100644
--- a/tests/runtime/test_service_factory.py
+++ b/tests/runtime/test_service_factory.py
@@ -22,6 +22,8 @@ import pytest
 # Local
 from caikit.runtime.service_factory import ServicePackage, ServicePackageFactory
 from caikit.runtime.types.caikit_runtime_exception import CaikitRuntimeException
+from tests.conftest import temp_config_parser
+import caikit
 
 
 # 🌶️🌶️🌶️🌶️🌶️
@@ -165,3 +167,130 @@ def test_get_service_proto_module():
         "caikit_runtime_pb2", mock_config
     )
     assert service_proto_module == caikit_runtime_pb2
+
+
+# Local
+import sample_lib
+
+MODULE_LIST = [
+    module_class
+    for module_class in caikit.core.MODULE_REGISTRY.values()
+    if module_class.__module__.partition(".")[0] == "sample_lib"
+]
+
+### Test ServicePackageFactory._get_and_filter_modules function
+def test_get_and_filter_modules_respects_excluded_task_type():
+    assert len(MODULE_LIST) == 6  # there are 6 modules in sample_lib
+    with temp_config_parser(
+        {"service_generation": {"task_types": {"excluded": ["sample_task"]}}}
+    ) as cfg:
+        clean_modules = ServicePackageFactory._get_and_filter_modules(cfg, "sample_lib")
+        assert len(clean_modules) == 1
+        assert "sample_task" not in str(clean_modules)
+
+
+def test_get_and_filter_modules_respects_excluded_modules():
+    assert "InnerBlock" in str(MODULE_LIST)
+    with temp_config_parser(
+        {
+            "service_generation": {
+                "module_guids": {"excluded": ["00110203-baad-beef-0809-0a0b02dd0e0f"]}
+            }
+        }  # excluding InnerBlock
+    ) as cfg:
+
+        clean_modules = ServicePackageFactory._get_and_filter_modules(cfg, "sample_lib")
+        assert len(clean_modules) == 5
+        assert "InnerBlock" not in str(clean_modules)
+
+
+def test_get_and_filter_modules_respects_excluded_modules_and_excluded_task_type():
+    assert "InnerBlock" in str(MODULE_LIST)
+    assert len(MODULE_LIST) == 6  # there are 6 modules in sample_lib
+    with temp_config_parser(
+        {
+            "service_generation": {
+                "module_guids": {"excluded": ["00110203-baad-beef-0809-0a0b02dd0e0f"]},
+                "task_types": {"excluded": ["other_task"]},
+            }
+        }  # excluding InnerBlock and OtherBlock
+    ) as cfg:
+
+        clean_modules = ServicePackageFactory._get_and_filter_modules(cfg, "sample_lib")
+        assert len(clean_modules) == 4
+        assert "InnerBlock" not in str(clean_modules)
+        assert "OtherBlock" not in str(clean_modules)
+        assert "other_task" not in str(clean_modules)
+
+
+def test_get_and_filter_modules_respects_included_modules_and_included_task_types():
+    assert len(MODULE_LIST) == 6  # there are 6 modules in sample_lib
+    with temp_config_parser(
+        {
+            "service_generation": {
+                "module_guids": {"included": ["00110203-baad-beef-0809-0a0b02dd0e0f"]},
+                "task_types": {"included": ["other_task"]},
+            }
+        }  # only want InnerBlock and OtherBlock
+    ) as cfg:
+
+        clean_modules = ServicePackageFactory._get_and_filter_modules(cfg, "sample_lib")
+        assert len(clean_modules) == 2
+        assert "InnerBlock" in str(clean_modules)
+        assert "OtherBlock" in str(clean_modules)
+        assert "ListBlock" not in str(clean_modules)
+
+
+def test_get_and_filter_modules_respects_included_modules():
+    assert len(MODULE_LIST) == 6  # there are 6 modules in sample_lib
+    with temp_config_parser(
+        {
+            "service_generation": {
+                "module_guids": {
+                    "included": [
+                        "00110203-baad-beef-0809-0a0b02dd0e0f",
+                        "00af2203-0405-0607-0263-0a0b02dd0c2f",
+                    ]
+                },
+            }
+        }  # only want InnerBlock and ListBlock
+    ) as cfg:
+
+        clean_modules = ServicePackageFactory._get_and_filter_modules(cfg, "sample_lib")
+        assert len(clean_modules) == 2
+        assert "InnerBlock" in str(clean_modules)
+        assert "ListBlock" in str(clean_modules)
+        assert "OtherBlock" not in str(clean_modules)
+
+
+def test_get_and_filter_modules_respects_included_task_types():
+    assert len(MODULE_LIST) == 6  # there are 6 modules in sample_lib
+    with temp_config_parser(
+        {
+            "service_generation": {
+                "task_types": {"included": ["sample_task"]},
+            }
+        }  # only want sample_task which has 5 modules
+    ) as cfg:
+
+        clean_modules = ServicePackageFactory._get_and_filter_modules(cfg, "sample_lib")
+        assert len(clean_modules) == 5
+        assert "InnerBlock" in str(clean_modules)
+        assert "OtherBlock" not in str(clean_modules)
+
+
+def test_get_and_filter_modules_respects_included_task_types_and_excluded_modules():
+    assert len(MODULE_LIST) == 6  # there are 6 modules in sample_lib
+    with temp_config_parser(
+        {
+            "service_generation": {
+                "task_types": {"included": ["sample_task"]},
+                "module_guids": {"excluded": ["00af2203-0405-0607-0263-0a0b02dd0c2f"]},
+            }
+        }  # only want sample_task but not ListBlock
+    ) as cfg:
+
+        clean_modules = ServicePackageFactory._get_and_filter_modules(cfg, "sample_lib")
+        assert len(clean_modules) == 4
+        assert "InnerBlock" in str(clean_modules)
+        assert "ListBlock" not in str(clean_modules)