Repository: caikit/caikit
Pull request #651 (2024-01-31 16:30:44)

Single module server
**What this PR does / why we need it**:

This PR enables the usage of `__main__` as the `runtime.library` so that an end-to-end example of `caikit.runtime` can be contained in a single python script.

**Special notes for your reviewer**:

While not strictly necessary for standard operation, this PR is worth merging because (a) it helps with concise examples and (b) users _may_ want to write this kind of single-module server in the future.

**If applicable**:
- [ ] this PR contains documentation
- [x] this PR contains unit tests
- [x] this PR has been tested for backwards compatibility



Patch:
diff --git a/caikit/core/toolkit/name_tools.py b/caikit/core/toolkit/name_tools.py
index 999a5eb7..71ca2a1d 100644
--- a/caikit/core/toolkit/name_tools.py
+++ b/caikit/core/toolkit/name_tools.py
@@ -19,4 +19,4 @@ and other Protobuf names
 
 def snake_to_upper_camel(string: str) -> str:
     """Simple snake -> upper camel conversion for descriptors"""
-    return "".join([part[0].upper() + part[1:] for part in string.split("_")])
+    return "".join([part[0].upper() + part[1:] for part in string.split("_") if part])
diff --git a/caikit/runtime/grpc_server.py b/caikit/runtime/grpc_server.py
index d061eb02..8f8745ae 100644
--- a/caikit/runtime/grpc_server.py
+++ b/caikit/runtime/grpc_server.py
@@ -120,14 +120,14 @@ class RuntimeGRPCServer(RuntimeServerBase):
             )
 
             # Add training management servicer to the gRPC server
-            training_management_service: ServicePackage = (
+            self.training_management_service: ServicePackage = (
                 ServicePackageFactory.get_service_package(
                     ServicePackageFactory.ServiceType.TRAINING_MANAGEMENT,
                 )
             )
-            service_names.append(training_management_service.descriptor.full_name)
+            service_names.append(self.training_management_service.descriptor.full_name)
 
-            training_management_service.registration_function(
+            self.training_management_service.registration_function(
                 TrainingManagementServicerImpl(), self.server
             )
 
diff --git a/caikit/runtime/servicers/global_train_servicer.py b/caikit/runtime/servicers/global_train_servicer.py
index 3ebd61a3..2bb15b80 100644
--- a/caikit/runtime/servicers/global_train_servicer.py
+++ b/caikit/runtime/servicers/global_train_servicer.py
@@ -73,7 +73,7 @@ class GlobalTrainServicer:
             lib_version = "unknown"
 
         log.info(
-            "<RUN76884779I>",
+            "<RUN76884778I>",
             "Constructed train service for library: %s, version: %s",
             self.library,
             lib_version,
diff --git a/caikit/runtime/utils/import_util.py b/caikit/runtime/utils/import_util.py
index fb3f74b2..23e1914a 100644
--- a/caikit/runtime/utils/import_util.py
+++ b/caikit/runtime/utils/import_util.py
@@ -101,6 +101,7 @@ def get_data_model(config: aconfig.Config = None) -> UnifiedDataModel:
 
     cdm = UnifiedDataModel()
     for lib_name in lib_names:
+        log.debug2("Importing library %s", lib_name)
         cdm = _get_cdm_from_lib(lib_name, cdm)
 
     # Check module registry to get base modules
@@ -161,6 +162,8 @@ def get_dynamic_module(module_name: str, module_dir: str = None) -> ModuleType:
     Returns:
         (module): Handle to the module after dynamic import
     """
+    if module := sys.modules.get(module_name):
+        return module
     module_path = f"{module_dir}.{module_name}" if module_dir else module_name
     log.info("<RUN11997772I>", "Loading service module: %s", module_path)
     # Try to find the spec for the module that we're interested in.


Test patch:
diff --git a/tests/runtime/utils/test_import_util.py b/tests/runtime/utils/test_import_util.py
index 6a58d36d..e672aeda 100644
--- a/tests/runtime/utils/test_import_util.py
+++ b/tests/runtime/utils/test_import_util.py
@@ -89,6 +89,15 @@ def test_get_caikit_library_loads_caikit_core():
     assert sample_module == caikit.core
 
 
+def test_get_caikit_library_loads_main():
+    """Make sure __main__ works"""
+    # Standard
+    import sys
+
+    sample_module = get_dynamic_module("__main__")
+    assert sample_module is sys.modules["__main__"]
+
+
 ### get_data_model #############################################################