Repository: caikit/caikit
Pull request #634 (2024-01-04 17:53:09)

Move Service/Server Constant Generation to Interfaces
<!--  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**:
This PR moves all constant/string formatting required for Service/RuntimeServer from `caikit.runtime` to `caikit.interfaces`. This will allow clients to dynamically generate GRPC Packages/Request endpoints without needing to install all of `caikit[runtime]`   


**Special notes for your reviewer**:
Please let me know if I should rebase off/wait for [PR584](https://github.com/caikit/caikit/pull/584)

**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/toolkit/name_tools.py b/caikit/core/toolkit/name_tools.py
new file mode 100644
index 00000000..999a5eb7
--- /dev/null
+++ b/caikit/core/toolkit/name_tools.py
@@ -0,0 +1,22 @@
+# 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.
+
+"""Common string functions that are generally helpful for generating runtime RPC names
+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("_")])
diff --git a/caikit/runtime/http_server/http_server.py b/caikit/runtime/http_server/http_server.py
index 0418e80f..b008d915 100644
--- a/caikit/runtime/http_server/http_server.py
+++ b/caikit/runtime/http_server/http_server.py
@@ -19,7 +19,6 @@ API based on the task definitions available at boot.
 # Standard
 from contextlib import contextmanager, nullcontext
 from dataclasses import dataclass
-from enum import Enum
 from functools import partial
 from typing import Any, Dict, Iterable, List, Optional, Type, Union, get_args
 import asyncio
@@ -27,7 +26,6 @@ import inspect
 import io
 import json
 import os
-import re
 import signal
 import ssl
 import tempfile
@@ -71,6 +69,16 @@ from caikit.core.exceptions.caikit_core_exception import (
     CaikitCoreStatusCode,
 )
 from caikit.core.toolkit.sync_to_async import async_wrap_iter
+from caikit.runtime.names import (
+    HEALTH_ENDPOINT,
+    MODEL_ID,
+    MODELS_INFO_ENDPOINT,
+    OPTIONAL_INPUTS_KEY,
+    REQUIRED_INPUTS_KEY,
+    RUNTIME_INFO_ENDPOINT,
+    StreamEventTypes,
+    get_http_route_name,
+)
 from caikit.runtime.server_base import RuntimeServerBase
 from caikit.runtime.service_factory import ServicePackage
 from caikit.runtime.service_generation.rpcs import (
@@ -120,26 +128,6 @@ STATUS_CODE_TO_HTTP = {
 }
 
 
-# These keys are used to define the logical sections of the request and response
-# data structures.
-REQUIRED_INPUTS_KEY = "inputs"
-OPTIONAL_INPUTS_KEY = "parameters"
-MODEL_ID = "model_id"
-
-# Endpoint to use for health checks
-HEALTH_ENDPOINT = "/health"
-
-# Endpoint to use for server info
-RUNTIME_INFO_ENDPOINT = "/info/version"
-MODELS_INFO_ENDPOINT = "/info/models"
-
-
-# Stream event types enum
-class StreamEventTypes(Enum):
-    MESSAGE = "message"
-    ERROR = "error"
-
-
 # Small dataclass for consolidating TLS files
 @dataclass
 class _TlsFiles:
@@ -429,7 +417,7 @@ class RuntimeHTTPServer(RuntimeServerBase):
         pydantic_response = dataobject_to_pydantic(response_data_object)
 
         @self.app.post(
-            self._get_route(rpc),
+            get_http_route_name(rpc.name),
             responses=self._get_response_openapi(
                 response_data_object, pydantic_response
             ),
@@ -490,7 +478,7 @@ class RuntimeHTTPServer(RuntimeServerBase):
         pydantic_response = dataobject_to_pydantic(response_data_object)
 
         @self.app.post(
-            self._get_route(rpc),
+            get_http_route_name(rpc.name),
             responses=self._get_response_openapi(
                 response_data_object, pydantic_response
             ),
@@ -573,7 +561,7 @@ class RuntimeHTTPServer(RuntimeServerBase):
 
         # pylint: disable=unused-argument
         @self.app.post(
-            self._get_route(rpc),
+            get_http_route_name(rpc.name),
             response_model=pydantic_response,
             openapi_extra=self._get_request_openapi(pydantic_request),
         )
@@ -656,25 +644,6 @@ class RuntimeHTTPServer(RuntimeServerBase):
 
             return EventSourceResponse(_generator())
 
-    def _get_route(self, rpc: CaikitRPCBase) -> str:
-        """Get the REST route for this rpc"""
-        if rpc.name.endswith("Predict"):
-            task_name = re.sub(
-                r"(?<!^)(?=[A-Z])",
-                "-",
-                re.sub("Task$", "", re.sub("Predict$", "", rpc.name)),
-            ).lower()
-            route = "/".join([self.config.runtime.http.route_prefix, "task", task_name])
-            if route[0] != "/":
-                route = "/" + route
-            return route
-        if rpc.name.endswith("Train"):
-            route = "/".join([self.config.runtime.http.route_prefix, rpc.name])
-            if route[0] != "/":
-                route = "/" + route
-            return route
-        raise NotImplementedError("No support for train rpcs yet!")
-
     def _get_request_dataobject(self, rpc: CaikitRPCBase) -> Type[DataBase]:
         """Get the dataobject request for the given rpc"""
         is_inference_rpc = hasattr(rpc, "task")
diff --git a/caikit/runtime/model_management/model_manager.py b/caikit/runtime/model_management/model_manager.py
index a29b20f7..ac6a5ea5 100644
--- a/caikit/runtime/model_management/model_manager.py
+++ b/caikit/runtime/model_management/model_manager.py
@@ -112,26 +112,11 @@ class ModelManager:  # pylint: disable=too-many-instance-attributes
 
         # Keep track of whether lazy loading is enabled
         self._lazy_load_local_models = runtime_cfg.lazy_load_local_models
-
-        if self._lazy_load_local_models:
-            error.value_check(
-                "<RUN44773525E>",
-                runtime_cfg.local_models_dir is not None,
-                (
-                    "runtime.local_models_dir must be set"
-                    " if using runtime.lazy_load_local_models. "
-                ),
-            )
-
-            error.value_check(
-                "<RUN44773514E>",
-                self._local_models_dir,
-                (
-                    "runtime.local_models_dir must be a valid path"
-                    " if set with runtime.lazy_load_local_models. "
-                    f"Provided path: {runtime_cfg.local_models_dir}"
-                ),
-            )
+        error.value_check(
+            "<RUN44773514E>",
+            not self._lazy_load_local_models or self._local_models_dir,
+            "Must set runtime.local_models_dir with runtime.lazy_load_local_models",
+        )
 
         # Set up local model periodic sync
         self._lazy_load_poll_period_seconds = runtime_cfg.lazy_load_poll_period_seconds
diff --git a/caikit/runtime/names.py b/caikit/runtime/names.py
new file mode 100644
index 00000000..79c10dfc
--- /dev/null
+++ b/caikit/runtime/names.py
@@ -0,0 +1,321 @@
+# 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.
+"""
+! NOTE ! This file should not import any extra dependencies. It is intended for
+use by client libraries that do not necessarily use a specific runtime server
+type.
+"""
+
+# Standard
+from enum import Enum
+from typing import Optional, Type, Union
+import re
+
+# First Party
+import alog
+
+# Local
+from caikit.config import get_config
+from caikit.core.modules import ModuleBase
+from caikit.core.task import TaskBase
+from caikit.core.toolkit.name_tools import snake_to_upper_camel
+from caikit.interfaces.runtime.data_model import (
+    ModelInfoRequest,
+    ModelInfoResponse,
+    RuntimeInfoRequest,
+    RuntimeInfoResponse,
+    TrainingInfoRequest,
+    TrainingStatusResponse,
+)
+
+log = alog.use_channel("RNTM-NAMES")
+
+
+############# Serice Names ##############
+
+
+class ServiceType(Enum):
+    """Common class for describing service types"""
+
+    INFERENCE = 1  # Inference service for the GlobalPredictServicer
+    TRAINING = 2  # Training service for the GlobalTrainServicer
+    TRAINING_MANAGEMENT = 3
+    INFO = 4
+
+
+############# Serice Name Generation ##############
+
+
+##  Service Package Descriptors
+
+
+def get_ai_domain() -> str:
+    """Get the string name for the AI domain
+
+    Returns:
+        domain(str): The domain for this service
+    """
+    caikit_config = get_config()
+    lib = caikit_config.runtime.library
+    default_ai_domain_name = snake_to_upper_camel(lib.replace("caikit_", ""))
+    ai_domain_name = (
+        caikit_config.runtime.service_generation.domain or default_ai_domain_name
+    )
+    return ai_domain_name
+
+
+def get_service_package_name(service_type: Optional[ServiceType] = None) -> str:
+    """This helper will get the name of service package
+
+    Args:
+        service_type Optional[ServiceType]: The Service Type's package name to fetch defaults
+            to runtime
+
+    Returns:
+        str: The name of the service package
+    """
+
+    # If specific service_type was provided then return their packages
+    if service_type == ServiceType.INFO:
+        return INFO_SERVICE_PACKAGE
+    elif service_type == ServiceType.TRAINING_MANAGEMENT:
+        return TRAINING_MANAGEMENT_PACKAGE
+
+    caikit_config = get_config()
+    ai_domain_name = get_ai_domain()
+    default_package_name = f"caikit.runtime.{ai_domain_name}"
+    package_name = (
+        caikit_config.runtime.service_generation.package or default_package_name
+    )
+    return package_name
+
+
+def get_service_name(service_type: ServiceType) -> str:
+    """This helper will get the name of the service
+
+    Args:
+        service_type ServiceType: The Service Type whose name to fetch
+
+    Returns:
+        str: The name of the service
+    """
+    if service_type == ServiceType.INFERENCE:
+        return f"{get_ai_domain()}Service"
+    elif service_type == ServiceType.TRAINING:
+        return f"{get_ai_domain()}TrainingService"
+    elif service_type == ServiceType.TRAINING_MANAGEMENT:
+        return TRAINING_MANAGEMENT_SERVICE_NAME
+    elif service_type == ServiceType.INFO:
+        return INFO_SERVICE_NAME
+
+
+##  Service RPC Descriptors
+
+
+def get_train_rpc_name(module_class: Type[ModuleBase]) -> str:
+    """Helper function to convert from the name of a module to the name of the
+    request RPC function
+    """
+
+    # 🌶️🌶️🌶️ The naming scheme for training RPCs probably needs to change.
+    # This uses the first task from the `tasks` kwarg in the `@caikit.module` decorator.
+    # This is both:
+    # - Flaky, since re-ordering that list would be perfectly reasonable and valid to do except
+    #   for the side effect of breaking the training service api
+    # - Not very intuitive, since a module supporting multiple tasks will have a training
+    #   endpoint that lists only one of them
+    rpc_name = snake_to_upper_camel(
+        f"{next(iter(module_class.tasks)).__name__}_{module_class.__name__}_Train"
+    )
+
+    if len(module_class.tasks) > 1:
+        log.warning(
+            "<RUN35134050W>",
+            "Multiple tasks detected for training rpc. "
+            "Module: [%s], Tasks: [%s], RPC name: %s ",
+            module_class,
+            module_class.tasks,
+            rpc_name,
+        )
+
+    return rpc_name
+
+
+def get_task_predict_rpc_name(
+    task_or_module_class: Type[Union[ModuleBase, TaskBase]],
+    input_streaming: bool = False,
+    output_streaming: bool = False,
+) -> str:
+    """Helper function to get the name of a task's RPC"""
+    task_class = (
+        next(iter(task_or_module_class.tasks))
+        if issubclass(task_or_module_class, ModuleBase)
+        else task_or_module_class
+    )
+
+    if input_streaming and output_streaming:
+        return snake_to_upper_camel(f"BidiStreaming{task_class.__name__}_Predict")
+    if output_streaming:
+        return snake_to_upper_camel(f"ServerStreaming{task_class.__name__}_Predict")
+    if input_streaming:
+        return snake_to_upper_camel(f"ClientStreaming{task_class.__name__}_Predict")
+    return snake_to_upper_camel(f"{task_class.__name__}_Predict")
+
+
+##  Service DataModel Name Descriptors
+
+
+def get_train_request_name(module_class: Type[ModuleBase]) -> str:
+    """Helper function to get the request name of a Train Service"""
+    return f"{get_train_rpc_name(module_class)}Request"
+
+
+def get_train_parameter_name(module_class: Type[ModuleBase]) -> str:
+    """Helper function to get the inner request parameter  name of a Train Service"""
+    return f"{get_train_rpc_name(module_class)}Parameters"
+
+
+def get_task_predict_request_name(
+    task_or_module_class: Type[Union[ModuleBase, TaskBase]],
+    input_streaming: bool = False,
+    output_streaming: bool = False,
+) -> str:
+    """Helper function to get the name of an RPC's request data type"""
+
+    task_class = (
+        next(iter(task_or_module_class.tasks))
+        if issubclass(task_or_module_class, ModuleBase)
+        else task_or_module_class
+    )
+
+    if input_streaming and output_streaming:
+        return snake_to_upper_camel(f"BidiStreaming{task_class.__name__}_Request")
+    if output_streaming:
+        return snake_to_upper_camel(f"ServerStreaming{task_class.__name__}_Request")
+    if input_streaming:
+        return snake_to_upper_camel(f"ClientStreaming{task_class.__name__}_Request")
+    return snake_to_upper_camel(f"{task_class.__name__}_Request")
+
+
+##  Service Definitions
+
+TRAINING_MANAGEMENT_SERVICE_NAME = "TrainingManagement"
+TRAINING_MANAGEMENT_PACKAGE = "caikit.runtime.training"
+TRAINING_MANAGEMENT_SERVICE_SPEC = {
+    "service": {
+        "rpcs": [
+            {
+                "name": "GetTrainingStatus",
+                "input_type": TrainingInfoRequest.get_proto_class().DESCRIPTOR.full_name,
+                "output_type": TrainingStatusResponse.get_proto_class().DESCRIPTOR.full_name,
+            },
+            {
+                "name": "CancelTraining",
+                "input_type": TrainingInfoRequest.get_proto_class().DESCRIPTOR.full_name,
+                "output_type": TrainingStatusResponse.get_proto_class().DESCRIPTOR.full_name,
+            },
+        ]
+    }
+}
+
+INFO_SERVICE_NAME = "InfoService"
+INFO_SERVICE_PACKAGE = "caikit.runtime.info"
+INFO_SERVICE_SPEC = {
+    "service": {
+        "rpcs": [
+            {
+                "name": "GetRuntimeInfo",
+                "input_type": RuntimeInfoRequest.get_proto_class().DESCRIPTOR.full_name,
+                "output_type": RuntimeInfoResponse.get_proto_class().DESCRIPTOR.full_name,
+            },
+            {
+                "name": "GetModelsInfo",
+                "input_type": ModelInfoRequest.get_proto_class().DESCRIPTOR.full_name,
+                "output_type": ModelInfoResponse.get_proto_class().DESCRIPTOR.full_name,
+            },
+        ]
+    }
+}
+
+
+############### Server Names #############
+
+# Invocation metadata key for the model ID, provided by Model Mesh
+MODEL_MESH_MODEL_ID_KEY = "mm-model-id"
+
+
+## HTTP Server
+
+# Endpoint to use for health checks
+HEALTH_ENDPOINT = "/health"
+
+# Endpoint to use for server info
+RUNTIME_INFO_ENDPOINT = "/info/version"
+MODELS_INFO_ENDPOINT = "/info/models"
+
+# These keys are used to define the logical sections of the request and response
+# data structures.
+REQUIRED_INPUTS_KEY = "inputs"
+OPTIONAL_INPUTS_KEY = "parameters"
+MODEL_ID = "model_id"
+
+# Stream event types enum
+class StreamEventTypes(Enum):
+    MESSAGE = "message"
+    ERROR = "error"
+
+
+def get_http_route_name(rpc_name: str) -> str:
+    """Function to get the http route for a given rpc name
+
+    Args:
+        rpc_name (str): The name of the Caikit RPC
+
+    Raises:
+        NotImplementedError: If the RPC is not a Train or Predict RPC
+
+    Returns:
+        str: The name of the http route for RPC
+    """
+    if rpc_name.endswith("Predict"):
+        task_name = re.sub(
+            r"(?<!^)(?=[A-Z])",
+            "-",
+            re.sub("Task$", "", re.sub("Predict$", "", rpc_name)),
+        ).lower()
+        route = "/".join([get_config().runtime.http.route_prefix, "task", task_name])
+        if route[0] != "/":
+            route = "/" + route
+        return route
+    if rpc_name.endswith("Train"):
+        route = "/".join([get_config().runtime.http.route_prefix, rpc_name])
+        if route[0] != "/":
+            route = "/" + route
+        return route
+    raise NotImplementedError(f"Unknown RPC type for rpc name {rpc_name}")
+
+
+### GRPC Server
+
+
+def get_grpc_route_name(service_type: ServiceType, rpc_name: str) -> str:
+    """Function to get GRPC name for a given service type and rpc name
+
+    Args:
+        rpc_name (str): The name of the Caikit RPC
+
+    Returns:
+        str: The name of the GRPC route for RPC
+    """
+    return f"/{get_service_package_name(service_type)}.{get_service_name(service_type)}/{rpc_name}"
diff --git a/caikit/runtime/service_factory.py b/caikit/runtime/service_factory.py
index a43c2a0b..d533474c 100644
--- a/caikit/runtime/service_factory.py
+++ b/caikit/runtime/service_factory.py
@@ -13,7 +13,6 @@
 # limitations under the License.
 """This module is responsible for creating service objects for the runtime to consume"""
 # Standard
-from enum import Enum
 from types import ModuleType
 from typing import Callable, Dict, Set, Type, Union
 import dataclasses
@@ -46,6 +45,14 @@ from caikit.interfaces.runtime.data_model import (
     TrainingStatusResponse,
 )
 from caikit.runtime import service_generation
+from caikit.runtime.names import ServiceType as InterfaceServiceType
+from caikit.runtime.names import (
+    get_service_name,
+    get_service_package_name,
+    get_task_predict_request_name,
+    get_train_parameter_name,
+    get_train_request_name,
+)
 from caikit.runtime.service_generation.rpcs import CaikitRPCBase
 from caikit.runtime.utils import import_util
 
@@ -112,11 +119,7 @@ class ServicePackage:
 class ServicePackageFactory:
     """Factory responsible for yielding the correct concrete ServicePackage implementation"""
 
-    class ServiceType(Enum):
-        INFERENCE = 1  # Inference service for the GlobalPredictServicer
-        TRAINING = 2  # Training service for the GlobalTrainServicer
-        TRAINING_MANAGEMENT = 3
-        INFO = 4
+    ServiceType = InterfaceServiceType
 
     @classmethod
     def get_service_package(
@@ -174,8 +177,7 @@ class ServicePackageFactory:
         _ = import_util.get_data_model()
 
         # Get the names for the AI domain and the proto package
-        ai_domain_name = service_generation.get_ai_domain()
-        package_name = service_generation.get_runtime_service_package()
+        package_name = get_service_package_name(service_type)
 
         # Then do API introspection to come up with all the API definitions to support
         caikit_config = get_config()
@@ -183,6 +185,7 @@ class ServicePackageFactory:
             caikit_config, caikit_config.runtime.library, write_modules_file
         )
 
+        service_name = get_service_name(service_type)
         if service_type == cls.ServiceType.INFERENCE:
             # Assert for backwards compatibility, if enabled, when service type is INFERENCE
             ServicePackageFactory._check_backwards_compatibility(
@@ -192,10 +195,8 @@ class ServicePackageFactory:
             rpc_list = service_generation.create_inference_rpcs(
                 clean_modules, caikit_config
             )
-            service_name = f"{ai_domain_name}Service"
         else:  # service_type == cls.ServiceType.TRAINING
             rpc_list = service_generation.create_training_rpcs(clean_modules)
-            service_name = f"{ai_domain_name}TrainingService"
 
         rpc_list = [rpc for rpc in rpc_list if rpc.return_type is not None]
 
@@ -357,20 +358,12 @@ def get_inference_request(
         ModuleBase,
         TaskBase,
     )
-    task_class = (
-        next(iter(task_or_module_class.tasks))
-        if issubclass(task_or_module_class, ModuleBase)
-        else task_or_module_class
-    )
 
-    if input_streaming and output_streaming:
-        request_class_name = f"BidiStreaming{task_class.__name__}Request"
-    elif input_streaming:
-        request_class_name = f"ClientStreaming{task_class.__name__}Request"
-    elif output_streaming:
-        request_class_name = f"ServerStreaming{task_class.__name__}Request"
-    else:
-        request_class_name = f"{task_class.__name__}Request"
+    request_class_name = get_task_predict_request_name(
+        task_or_module_class,
+        input_streaming=input_streaming,
+        output_streaming=output_streaming,
+    )
     log.debug(
         "Request class name %s for class %s.", request_class_name, task_or_module_class
     )
@@ -384,10 +377,7 @@ def get_train_request(module_class: Type[ModuleBase]) -> Type[DataBase]:
         module_class,
         ModuleBase,
     )
-    # 🌶️🌶️🌶️ This is coupled to the naming scheme code in
-    # caikit.runtime.service_generation.rpcs, which is likely to change.
-    first_task = next(iter(module_class.tasks))
-    request_class_name = f"{first_task.__name__}{module_class.__name__}TrainRequest"
+    request_class_name = get_train_request_name(module_class)
     log.debug("Request class name %s for module %s.", request_class_name, module_class)
     return DataBase.get_class_for_name(request_class_name)
 
@@ -399,8 +389,7 @@ def get_train_params(module_class: Type[ModuleBase]) -> Type[DataBase]:
         module_class,
         ModuleBase,
     )
-    first_task = next(iter(module_class.tasks))
 
-    request_class_name = f"{first_task.__name__}{module_class.__name__}TrainParameters"
+    request_class_name = get_train_parameter_name(module_class)
     log.debug("Request class name %s for module %s.", request_class_name, module_class)
     return DataBase.get_class_for_name(request_class_name)
diff --git a/caikit/runtime/service_generation/__init__.py b/caikit/runtime/service_generation/__init__.py
index ebbe6280..18a79a7e 100644
--- a/caikit/runtime/service_generation/__init__.py
+++ b/caikit/runtime/service_generation/__init__.py
@@ -5,4 +5,3 @@ from .create_service import (
     create_inference_rpcs,
     create_training_rpcs,
 )
-from .proto_package import get_ai_domain, get_runtime_service_package
diff --git a/caikit/runtime/service_generation/data_stream_source.py b/caikit/runtime/service_generation/data_stream_source.py
index 1152a790..9b657668 100644
--- a/caikit/runtime/service_generation/data_stream_source.py
+++ b/caikit/runtime/service_generation/data_stream_source.py
@@ -41,7 +41,7 @@ from caikit.interfaces.common.data_model.stream_sources import (
     ListOfFileReferences,
     S3Files,
 )
-from caikit.runtime.service_generation.proto_package import get_runtime_service_package
+from caikit.runtime.names import get_service_package_name
 from caikit.runtime.types.caikit_runtime_exception import CaikitRuntimeException
 import caikit
 
@@ -321,7 +321,7 @@ class JsonDataStreamSourcePlugin(DataStreamSourcePlugin):
         if stream_message_type:
             return stream_message_type
 
-        package = get_runtime_service_package()
+        package = get_service_package_name()
         cls_name = _make_data_stream_source_type_name(element_type)
         JsonData = make_dataobject(
             package=package,
@@ -506,7 +506,7 @@ def make_data_stream_source(
     log.debug2("Looking for DataStreamSource[%s]", data_element_type)
     if data_element_type not in _DATA_STREAM_SOURCE_TYPES:
         cls_name = _make_data_stream_source_type_name(data_element_type)
-        package = get_runtime_service_package()
+        package = get_service_package_name()
 
         log.debug("Creating DataStreamSource[%s] -> %s", data_element_type, cls_name)
 
diff --git a/caikit/runtime/service_generation/proto_package.py b/caikit/runtime/service_generation/proto_package.py
deleted file mode 100644
index 99adb088..00000000
--- a/caikit/runtime/service_generation/proto_package.py
+++ /dev/null
@@ -1,47 +0,0 @@
-# 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 module holds the common helpers for managing the protobuf package name for
-the runtime
-"""
-
-# Local
-from ...config import get_config
-
-
-def snake_to_upper_camel(string: str) -> str:
-    """Simple snake -> upper camel conversion"""
-    return "".join([part[0].upper() + part[1:] for part in string.split("_")])
-
-
-def get_ai_domain() -> str:
-    """Get the string name for the AI domain"""
-    caikit_config = get_config()
-    lib = caikit_config.runtime.library
-    default_ai_domain_name = snake_to_upper_camel(lib.replace("caikit_", ""))
-    ai_domain_name = (
-        caikit_config.runtime.service_generation.domain or default_ai_domain_name
-    )
-    return ai_domain_name
-
-
-def get_runtime_service_package() -> str:
-    """This helper will get the common runtime package"""
-    caikit_config = get_config()
-    ai_domain_name = get_ai_domain()
-    default_package_name = f"caikit.runtime.{ai_domain_name}"
-    package_name = (
-        caikit_config.runtime.service_generation.package or default_package_name
-    )
-    return package_name
diff --git a/caikit/runtime/service_generation/rpcs.py b/caikit/runtime/service_generation/rpcs.py
index 39b74b49..ca44bab1 100644
--- a/caikit/runtime/service_generation/rpcs.py
+++ b/caikit/runtime/service_generation/rpcs.py
@@ -34,12 +34,18 @@ from ...interfaces.common.data_model.stream_sources import S3Path
 from . import protoable, type_helpers
 from .compatibility_checker import ApiFieldNames
 from .data_stream_source import make_data_stream_source
-from .proto_package import snake_to_upper_camel
 from caikit.core import ModuleBase, TaskBase
 from caikit.core.data_model.base import DataBase
 from caikit.core.data_model.dataobject import make_dataobject
 from caikit.core.signature_parsing import CaikitMethodSignature, CustomSignature
 from caikit.interfaces.runtime.data_model import ModelPointer, TrainingJob
+from caikit.runtime.names import (
+    get_task_predict_request_name,
+    get_task_predict_rpc_name,
+    get_train_parameter_name,
+    get_train_request_name,
+    get_train_rpc_name,
+)
 
 log = alog.use_channel("RPC-SERIALIZERS")
 
@@ -152,28 +158,7 @@ class ModuleClassTrainRPC(CaikitRPCBase):
         """Helper function to convert from the name of a module to the name of the
         request RPC function
         """
-        # 🌶️🌶️🌶️ The naming scheme for training RPCs probably needs to change.
-        # This uses the first task from the `tasks` kwarg in the `@caikit.module` decorator.
-        # This is both:
-        # - Flaky, since re-ordering that list would be perfectly reasonable and valid to do except
-        #   for the side effect of breaking the training service api
-        # - Not very intuitive, since a module supporting multiple tasks will have a training
-        #   endpoint that lists only one of them
-        rpc_name = snake_to_upper_camel(
-            f"{next(iter(module_class.tasks)).__name__}_{module_class.__name__}_Train"
-        )
-
-        if len(module_class.tasks) > 1:
-            log.warning(
-                "<RUN35134050W>",
-                "Multiple tasks detected for training rpc. "
-                "Module: [%s], Tasks: [%s], RPC name: %s ",
-                module_class,
-                module_class.tasks,
-                rpc_name,
-            )
-
-        return rpc_name
+        return get_train_rpc_name(module_class)
 
     @staticmethod
     def module_class_to_req_name(module_class: Type[ModuleBase]) -> str:
@@ -185,7 +170,7 @@ class ModuleClassTrainRPC(CaikitRPCBase):
         return: SampleTaskSampleModuleTrainRequest
 
         """
-        return f"{ModuleClassTrainRPC.module_class_to_rpc_name(module_class)}Request"
+        return get_train_request_name(module_class)
 
     @staticmethod
     def module_class_to_inner_request_name(module_class: Type[ModuleBase]) -> str:
@@ -196,7 +181,7 @@ class ModuleClassTrainRPC(CaikitRPCBase):
 
         return: SampleTaskSampleModuleTrainParameters
         """
-        return f"{ModuleClassTrainRPC.module_class_to_rpc_name(module_class)}Parameters"
+        return get_train_parameter_name(module_class)
 
     @staticmethod
     def _mutate_method_signature_for_training(
@@ -355,13 +340,9 @@ class TaskPredictRPC(CaikitRPCBase):
         """Helper function to convert the pair of library name and task name to
         a request message name
         """
-        if self._input_streaming and self._output_streaming:
-            return snake_to_upper_camel(f"BidiStreaming{self.task.__name__}_Request")
-        if self._output_streaming:
-            return snake_to_upper_camel(f"ServerStreaming{self.task.__name__}_Request")
-        if self._input_streaming:
-            return snake_to_upper_camel(f"ClientStreaming{self.task.__name__}_Request")
-        return snake_to_upper_camel(f"{self.task.__name__}_Request")
+        return get_task_predict_request_name(
+            self.task, self._input_streaming, self._output_streaming
+        )
 
     def _task_to_rpc_name(self) -> str:
         """Helper function to convert the pair of library name and task name
@@ -371,13 +352,9 @@ class TaskPredictRPC(CaikitRPCBase):
 
         return: SampleTaskPredict
         """
-        if self._input_streaming and self._output_streaming:
-            return snake_to_upper_camel(f"BidiStreaming{self.task.__name__}_Predict")
-        if self._output_streaming:
-            return snake_to_upper_camel(f"ServerStreaming{self.task.__name__}_Predict")
-        if self._input_streaming:
-            return snake_to_upper_camel(f"ClientStreaming{self.task.__name__}_Predict")
-        return snake_to_upper_camel(f"{self.task.__name__}_Predict")
+        return get_task_predict_rpc_name(
+            self.task, self._input_streaming, self._output_streaming
+        )
 
 
 class _RequestMessage:
diff --git a/caikit/runtime/servicers/global_predict_servicer.py b/caikit/runtime/servicers/global_predict_servicer.py
index 51c90a58..daf58168 100644
--- a/caikit/runtime/servicers/global_predict_servicer.py
+++ b/caikit/runtime/servicers/global_predict_servicer.py
@@ -34,6 +34,7 @@ from caikit.core.data_model import DataBase, DataStream
 from caikit.core.signature_parsing import CaikitMethodSignature
 from caikit.runtime.metrics.rpc_meter import RPCMeter
 from caikit.runtime.model_management.model_manager import ModelManager
+from caikit.runtime.names import MODEL_MESH_MODEL_ID_KEY
 from caikit.runtime.service_factory import ServicePackage
 from caikit.runtime.service_generation.rpcs import TaskPredictRPC
 from caikit.runtime.types.caikit_runtime_exception import CaikitRuntimeException
@@ -87,9 +88,6 @@ class GlobalPredictServicer:
     given request
     """
 
-    # Invocation metadata key for the model ID, provided by Model Mesh
-    MODEL_MESH_MODEL_ID_KEY = "mm-model-id"
-
     # Input size in code points, provided by orchestrator
     INPUT_SIZE_KEY = "input-length"
 
@@ -158,7 +156,7 @@ class GlobalPredictServicer:
                 A Caikit Library data model response object
         """
         # Make sure the request has a model before doing anything
-        model_id = get_metadata(context, self.MODEL_MESH_MODEL_ID_KEY)
+        model_id = get_metadata(context, MODEL_MESH_MODEL_ID_KEY)
         request_name = caikit_rpc.request.name
 
         with self._handle_predict_exceptions(model_id, request_name), alog.ContextLog(
diff --git a/pyproject.toml b/pyproject.toml
index 72e4fc70..9f7eae62 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -63,7 +63,7 @@ interfaces-ts = [
 
 interfaces-ts-pyspark = [
     "caikit[interfaces-ts]",
-    "pyspark>=3.3,<3.6",
+    "pyspark>=3.3,<3.5",
     "pyarrow>=8.0.0,<15"
 ]
 


Test patch:
diff --git a/tests/core/test_imports.py b/tests/core/test_imports.py
index 468cb2d3..a2911318 100644
--- a/tests/core/test_imports.py
+++ b/tests/core/test_imports.py
@@ -19,3 +19,7 @@ def test_caikit_core_has_DataValidationError():
 
 def test_caikit_core_has_error_handler():
     assert hasattr(caikit.core, "error_handler")
+    
+def test_caikit_core_has_runtime_anmes():
+    import caikit.runtime
+    assert hasattr(caikit.runtime, "names")
diff --git a/tests/examples/test_examples.py b/tests/examples/test_examples.py
index 8a18ae44..8050667f 100644
--- a/tests/examples/test_examples.py
+++ b/tests/examples/test_examples.py
@@ -13,7 +13,7 @@
 # limitations under the License.
 
 # Standard
-from os import path
+from os import getenv, path
 import asyncio
 import subprocess
 
diff --git a/tests/runtime/model_management/test_model_manager.py b/tests/runtime/model_management/test_model_manager.py
index 27825da8..be6fb557 100644
--- a/tests/runtime/model_management/test_model_manager.py
+++ b/tests/runtime/model_management/test_model_manager.py
@@ -1075,129 +1075,6 @@ def test_lazy_load_handles_temporary_errors():
                 assert model
 
 
-def test_lazy_load_true_local_models_dir_valid():
-    """When lazy_load_local_models is True and local_models_dir exists.
-    Check that the local_models_dir is pointing to the correct location
-    """
-
-    with TemporaryDirectory() as cache_dir:
-
-        ModelManager._ModelManager__instance = None
-        with temp_config(
-            {
-                "runtime": {
-                    "local_models_dir": cache_dir,
-                    "lazy_load_local_models": True,
-                }
-            },
-            merge_strategy="merge",
-        ):
-            MODEL_MANAGER = ModelManager()
-            assert len(MODEL_MANAGER.loaded_models) == 0
-            assert MODEL_MANAGER._local_models_dir == cache_dir
-
-
-def test_lazy_load_true_local_models_dir_invalid():
-    """When lazy_load_local_models is True and local_models_dir does not exist.
-    Raise ValueError with an appropriate message
-    """
-
-    with TemporaryDirectory() as cache_dir:
-
-        with pytest.raises(
-            ValueError,
-            match=(
-                "runtime.local_models_dir must be a valid path"
-                " if set with runtime.lazy_load_local_models. "
-                "Provided path: invalid"
-            ),
-        ):
-
-            ModelManager._ModelManager__instance = None
-            with temp_config(
-                {
-                    "runtime": {
-                        "local_models_dir": "invalid",
-                        "lazy_load_local_models": True,
-                    }
-                },
-                merge_strategy="merge",
-            ):
-                MODEL_MANAGER = ModelManager()
-
-
-def test_lazy_load_true_local_models_dir_none():
-    """When lazy_load_local_models is True and local_models_dir is not set in the config.
-    Raise ValueError with an appropriate message
-    """
-
-    with TemporaryDirectory() as cache_dir:
-
-        with pytest.raises(
-            ValueError,
-            match=(
-                "runtime.local_models_dir must be set"
-                " if using runtime.lazy_load_local_models. "
-            ),
-        ):
-
-            ModelManager._ModelManager__instance = None
-            with temp_config(
-                {
-                    "runtime": {
-                        "local_models_dir": None,
-                        "lazy_load_local_models": True,
-                    }
-                },
-                merge_strategy="merge",
-            ):
-                MODEL_MANAGER = ModelManager()
-
-
-def test_lazy_load_false_local_models_dir_valid():
-    """When lazy_load_local_models is False and local_models_dir exists.
-    Check that the local_models_dir is pointing to the correct location
-    """
-
-    with TemporaryDirectory() as cache_dir:
-
-        ModelManager._ModelManager__instance = None
-        with temp_config(
-            {
-                "runtime": {
-                    "local_models_dir": cache_dir,
-                    "lazy_load_local_models": False,
-                }
-            },
-            merge_strategy="merge",
-        ):
-            MODEL_MANAGER = ModelManager()
-            assert len(MODEL_MANAGER.loaded_models) == 0
-            assert MODEL_MANAGER._local_models_dir == cache_dir
-
-
-def test_lazy_load_false_local_models_dir_invalid():
-    """When lazy_load_local_models is False and local_models_dir does not exist.
-    Check that the local_models_dir is False / Empty
-    """
-
-    with TemporaryDirectory() as cache_dir:
-
-        ModelManager._ModelManager__instance = None
-        with temp_config(
-            {
-                "runtime": {
-                    "local_models_dir": "",
-                    "lazy_load_local_models": False,
-                }
-            },
-            merge_strategy="merge",
-        ):
-            MODEL_MANAGER = ModelManager()
-            assert len(MODEL_MANAGER.loaded_models) == 0
-            assert not MODEL_MANAGER._local_models_dir
-
-
 class NoModelFinder(ModelFinderBase):
     name = "NOMODEL"