Repository: caikit/caikit
Pull request #489 (2023-09-26 15:23:52)

fix: dump http services, add tests
<!--  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**:
closes: #433 
Fixes bug where dumping http services failed if the output directory didn't already exist.

**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/runtime/__init__.py b/caikit/runtime/__init__.py
index 4b57df6b..2068258b 100644
--- a/caikit/runtime/__init__.py
+++ b/caikit/runtime/__init__.py
@@ -11,6 +11,3 @@
 # 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.
-
-# Local
-from .service_factory import get_inference_request, get_train_params, get_train_request
diff --git a/caikit/runtime/dump_services.py b/caikit/runtime/dump_services.py
index 18d017ed..f2946073 100644
--- a/caikit/runtime/dump_services.py
+++ b/caikit/runtime/dump_services.py
@@ -75,6 +75,11 @@ def dump_http_services(output_dir: str):
     with TestClient(server.app) as client:
         response = client.get("/openapi.json")
         response.raise_for_status()
+
+        # create output dir if doesn't exist
+        if not os.path.isdir(output_dir):
+            os.makedirs(output_dir)
+
         with open(
             os.path.join(output_dir, "openapi.json"), "w", encoding="utf-8"
         ) as handle:
diff --git a/caikit/runtime/http_server/http_server.py b/caikit/runtime/http_server/http_server.py
index df2478d2..f3219de7 100644
--- a/caikit/runtime/http_server/http_server.py
+++ b/caikit/runtime/http_server/http_server.py
@@ -17,16 +17,12 @@ The server is responsible for binding caikit workloads to a consistent REST/SSE
 API based on the task definitions available at boot.
 """
 # Standard
-from contextlib import contextmanager
-from dataclasses import dataclass
 from functools import partial
 from typing import Any, Dict, Iterable, Optional, Type, get_args
 import asyncio
 import json
-import os
 import re
 import ssl
-import tempfile
 import threading
 import time
 
@@ -98,14 +94,6 @@ OPTIONAL_INPUTS_KEY = "parameters"
 # Endpoint to use for health checks
 HEALTH_ENDPOINT = "/health"
 
-# Small dataclass for consolidating TLS files
-@dataclass
-class _TlsFiles:
-    server_cert: Optional[str] = None
-    server_key: Optional[str] = None
-    client_cert: Optional[str] = None
-
-
 ## RuntimeHTTPServer ###########################################################
 
 
@@ -154,41 +142,34 @@ class RuntimeHTTPServer(RuntimeServerBase):
         )
 
         # Parse TLS configuration
-        # If any of the TLS values are not files, we assume that they're inline
-        # content. The python SslContext only takes files to load, so we use a
-        # temporary directory just long enough to load the config files.
-        with self._tls_files() as tls_files:
-            tls_kwargs = {}
-            if tls_files.server_key and tls_files.server_cert:
-                log.info("<RUN10001905I>", "Running with TLS")
-                tls_kwargs["ssl_keyfile"] = tls_files.server_key
-                tls_kwargs["ssl_certfile"] = tls_files.server_cert
-                if tls_files.client_cert:
-                    log.info("<RUN10001809I>", "Running with mutual TLS")
-                    tls_kwargs["ssl_ca_certs"] = tls_files.client_cert
-                    tls_kwargs["ssl_cert_reqs"] = ssl.CERT_REQUIRED
-            else:
-                log.info("<RUN10539515I>", "Running INSECURE")
-
-            # Start the server with a timeout_graceful_shutdown
-            # if not set in config, this is None and unvicorn accepts None or number of seconds
-            unvicorn_timeout_graceful_shutdown = (
-                get_config().runtime.http.server_shutdown_grace_period_seconds
-            )
-            config = uvicorn.Config(
-                self.app,
-                host="0.0.0.0",
-                port=self.port,
-                log_level=None,
-                log_config=None,
-                timeout_graceful_shutdown=unvicorn_timeout_graceful_shutdown,
-                **tls_kwargs,
-            )
-            # Make sure the config loads TLS files here so they can safely be
-            # deleted if they're ephemeral
-            config.load()
-
-        # Start the server with the loaded config
+        tls_kwargs = {}
+        if (
+            self.tls_config
+            and self.tls_config.server.key
+            and self.tls_config.server.cert
+        ):
+            log.info("<RUN10001905I>", "Running with TLS")
+            tls_kwargs["ssl_keyfile"] = self.tls_config.server.key
+            tls_kwargs["ssl_certfile"] = self.tls_config.server.cert
+            if self.tls_config.client.cert:
+                log.info("<RUN10001809I>", "Running with mutual TLS")
+                tls_kwargs["ssl_ca_certs"] = self.tls_config.client.cert
+                tls_kwargs["ssl_cert_reqs"] = ssl.CERT_REQUIRED
+
+        # Start the server with a timeout_graceful_shutdown
+        # if not set in config, this is None and unvicorn accepts None or number of seconds
+        unvicorn_timeout_graceful_shutdown = (
+            get_config().runtime.http.server_shutdown_grace_period_seconds
+        )
+        config = uvicorn.Config(
+            self.app,
+            host="0.0.0.0",
+            port=self.port,
+            log_level=None,
+            log_config=None,
+            timeout_graceful_shutdown=unvicorn_timeout_graceful_shutdown,
+            **tls_kwargs,
+        )
         self.server = uvicorn.Server(config=config)
 
         # Patch the exit handler to call this server's stop
@@ -551,50 +532,6 @@ class RuntimeHTTPServer(RuntimeServerBase):
         log.debug4("Server healthy")
         return "OK"
 
-    @contextmanager
-    def _tls_files(self) -> _TlsFiles:
-        """This contextmanager ensures that the tls config values are files on
-        disk since SslContext requires files
-
-        Returns:
-            tls_files (_TlsFiles): The set of configured TLS files
-        """
-        tls_cfg = self.tls_config or {}
-        tls_opts = {
-            "server_key": tls_cfg.get("server", {}).get("key", ""),
-            "server_cert": tls_cfg.get("server", {}).get("cert", ""),
-            "client_cert": tls_cfg.get("client", {}).get("cert", ""),
-        }
-        non_file_opts = {
-            key: val for key, val in tls_opts.items() if val and not os.path.isfile(val)
-        }
-        if not non_file_opts:
-            log.debug3("No TLS inline files configured")
-            yield _TlsFiles(**tls_opts)
-            return
-        # If any of the values are set and are not pointing to files, save them
-        # to a temporary directory
-        try:
-            with tempfile.TemporaryDirectory() as tls_dir:
-                updated_files = {
-                    key: val
-                    for key, val in tls_opts.items()
-                    if key not in non_file_opts and val
-                }
-                for fname, content in non_file_opts.items():
-                    temp_fname = os.path.join(tls_dir, fname)
-                    updated_files[fname] = temp_fname
-                    with open(temp_fname, "w", encoding="utf-8") as handle:
-                        handle.write(content)
-                yield _TlsFiles(**updated_files)
-        except OSError as err:
-            log.error(
-                "<RUN80977064E>",
-                "Cannot create temporary TLS files. Either pass config as file paths or run with write permissions.",
-                exc_info=True,
-            )
-            raise ValueError() from err
-
 
 def main(blocking: bool = True):
     server = RuntimeHTTPServer()
diff --git a/caikit/runtime/service_factory.py b/caikit/runtime/service_factory.py
index e6acab48..1970a82a 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, Set, Type, Union
+from typing import Callable, Dict, Set, Type
 import dataclasses
 
 # Third Party
@@ -31,10 +31,7 @@ import alog
 # Local
 from caikit import get_config
 from caikit.core import LocalBackend, ModuleBase, registries
-from caikit.core.data_model.base import DataBase
 from caikit.core.data_model.dataobject import _AUTO_GEN_PROTO_CLASSES
-from caikit.core.exceptions import error_handler
-from caikit.core.task import TaskBase
 from caikit.interfaces.runtime.data_model import (
     TrainingInfoRequest,
     TrainingStatusResponse,
@@ -44,7 +41,6 @@ from caikit.runtime.service_generation.rpcs import CaikitRPCBase
 from caikit.runtime.utils import import_util
 
 log = alog.use_channel("SVC-FACTORY")
-error = error_handler.get(log)
 
 TRAINING_MANAGEMENT_SERVICE_NAME = "TrainingManagement"
 TRAINING_MANAGEMENT_SERVICE_SPEC = {
@@ -270,63 +266,3 @@ class ServicePackageFactory:
             excluded_modules,
         )
         return clean_modules
-
-
-def get_inference_request(
-    task_or_module_class: Type[Union[ModuleBase, TaskBase]],
-    input_streaming: bool = False,
-    output_streaming: bool = False,
-) -> Type[DataBase]:
-    """Helper function to return the inference request DataModel for the Module or Task Class"""
-    error.subclass_check(
-        "<SVC98285724E>",
-        task_or_module_class,
-        ModuleBase,
-        TaskBase,
-    )
-    task_class = (
-        task_or_module_class.TASK_CLASS
-        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"
-    log.debug(
-        "Request class name %s for class %s.", request_class_name, task_or_module_class
-    )
-    return DataBase.get_class_for_name(request_class_name)
-
-
-def get_train_request(module_class: Type[ModuleBase]) -> Type[DataBase]:
-    """Helper function to return the train request DataModel for the Module Class"""
-    error.subclass_check(
-        "<SVC32285724E>",
-        module_class,
-        ModuleBase,
-    )
-    request_class_name = (
-        f"{module_class.TASK_CLASS.__name__}{module_class.__name__}TrainRequest"
-    )
-    log.debug("Request class name %s for module %s.", request_class_name, module_class)
-    return DataBase.get_class_for_name(request_class_name)
-
-
-def get_train_params(module_class: Type[ModuleBase]) -> Type[DataBase]:
-    """Helper function to return the train parameters DataModel for the Module Class"""
-    error.subclass_check(
-        "<SVC98435724E>",
-        module_class,
-        ModuleBase,
-    )
-    request_class_name = (
-        f"{module_class.TASK_CLASS.__name__}{module_class.__name__}TrainParameters"
-    )
-    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/data_stream_source.py b/caikit/runtime/service_generation/data_stream_source.py
index 19964c74..1a3b933a 100644
--- a/caikit/runtime/service_generation/data_stream_source.py
+++ b/caikit/runtime/service_generation/data_stream_source.py
@@ -167,11 +167,8 @@ class DataStreamSourceBase(DataStream):
         """Create a data stream object by deducing file extension
         and reading the file accordingly"""
 
-        _, extension = os.path.splitext(fname)
-        if not extension:
-            return cls._load_from_file_without_extension(fname)
-
         full_fname = cls._get_resolved_source_path(fname)
+        _, extension = os.path.splitext(fname)
         log.debug3("Pulling data stream from %s file [%s]", extension, full_fname)
 
         if not fname or not os.path.isfile(full_fname):
@@ -190,37 +187,6 @@ class DataStreamSourceBase(DataStream):
             f"Extension not supported! {extension}",
         )
 
-    @classmethod
-    def _load_from_file_without_extension(cls, fname) -> DataStream:
-        """Similar to _create_data_stream_from_file, but we don't have a file extension to work
-        with. Attempt to create a data stream using one of a few well-known formats"""
-        full_fname = cls._get_resolved_source_path(fname)
-        log.debug3("Attempting to guess file type for file: %s", full_fname)
-        for factory_method in (
-            DataStream.from_json_array,
-            DataStream.from_jsonl,
-            DataStream.from_header_csv,
-        ):
-            try:
-                stream = factory_method(full_fname).map(cls._to_element_type)
-                # Iterate once and assume we have the correct file type if this works
-                stream.peek()
-                return stream
-            except Exception as e:  # pylint: disable=broad-exception-caught
-                # Catch any exception: it's hard to know which all could be thrown by any of the
-                # formatters
-                log.debug3(
-                    "Failed to load file %s using data stream factory method %s: %s",
-                    full_fname,
-                    factory_method,
-                    e,
-                    exc_info=True,
-                )
-        raise CaikitRuntimeException(
-            grpc.StatusCode.INVALID_ARGUMENT,
-            f"Could not load input file with no extension: {full_fname}",
-        )
-
     @classmethod
     def _to_element_type(cls, raw_element: Any) -> "ElementType":
         """Stream adapter to adapt from the raw json object data representations
diff --git a/caikit/runtime/servicers/global_train_servicer.py b/caikit/runtime/servicers/global_train_servicer.py
index f672f7ac..b87b9d11 100644
--- a/caikit/runtime/servicers/global_train_servicer.py
+++ b/caikit/runtime/servicers/global_train_servicer.py
@@ -54,6 +54,9 @@ class GlobalTrainServicer:
     def __init__(self, training_service: ServicePackage):
         self._training_service = training_service
         self._model_manager = ModelManager.get_instance()
+        caikit_config = get_config()
+        self.training_output_dir = caikit_config.runtime.training.output_dir
+        self.save_with_id = caikit_config.runtime.training.save_with_id
 
         # TODO: think about if we really want to do this here:
         self.cdm = get_data_model()
@@ -66,7 +69,7 @@ class GlobalTrainServicer:
         # Or grab the `libraries` off of the `training_service` instead of config here?
         # Duplicate code in global_train_servicer
         # pylint: disable=duplicate-code
-        self.library = clean_lib_names(get_config().runtime.library)[0]
+        self.library = clean_lib_names(caikit_config.runtime.library)[0]
         try:
             lib_version = version(self.library)
         except Exception:  # pylint: disable=broad-exception-caught
@@ -80,14 +83,6 @@ class GlobalTrainServicer:
         )
         super()
 
-    @property
-    def training_output_dir(self) -> str:
-        return get_config().runtime.training.output_dir
-
-    @property
-    def save_with_id(self) -> str:
-        return get_config().runtime.training.save_with_id
-
     def Train(
         self,
         request: ProtoMessageType,
@@ -147,7 +142,7 @@ class GlobalTrainServicer:
         # pylint: disable=duplicate-code
         except (TypeError, ValueError) as e:
             log_dict = {
-                "log_code": "<RUN72924264W>",
+                "log_code": "<RUN490439039W>",
                 "message": repr(e),
                 "stack_trace": traceback.format_exc(),
             }
@@ -172,7 +167,7 @@ class GlobalTrainServicer:
         self,
         request: ProtoMessageType,
         module: Type[ModuleBase],
-        training_output_dir: Optional[str] = None,
+        training_output_dir: str,
         *,
         context: Optional[ServicerContext] = None,
         wait: bool = False,
@@ -184,9 +179,8 @@ class GlobalTrainServicer:
         Args:
             request (ProtoMessageType): The message that stimulated this request
             module (Type[ModuleBase]): The module class to train
-            training_output_dir (Optional[str]): The base directory where
-                the trained model should be saved, falling back to the
-                configured output dir if not given.
+            training_output_dir (str): The base directory where trained models
+                should be saved
 
         Kwargs:
             context (Optional[ServicerContext]): The grpc context for the
diff --git a/caikit/runtime/servicers/model_train_servicer.py b/caikit/runtime/servicers/model_train_servicer.py
index 697fc870..bf001122 100644
--- a/caikit/runtime/servicers/model_train_servicer.py
+++ b/caikit/runtime/servicers/model_train_servicer.py
@@ -25,7 +25,6 @@ import grpc
 import alog
 
 # Local
-from caikit import get_config
 from caikit.interfaces.common.data_model.stream_sources import (
     Directory,
     File,
@@ -99,25 +98,13 @@ class ModelTrainServicerImpl(process_pb2_grpc.ProcessServicer):
                 train_message_request.parameters, request.training_input_dir
             )
 
-            # Figure out the right training output dir to use for this request.
-            # Model Train does not allow this value in the request to be set by
-            # the inbound request to the frontend of MT, so the
-            # request.training_output_dir value will always point to a private
-            # emptyDir mount. In the case where the user wants to save to some
-            # other path (e.g. a mounted PVC), the user must set that in the
-            # global config as runtime.training.output_dir. This will cause all
-            # requests to be saved there and therefore S3 uploads will not be
-            # done.
-            output_dir = (
-                get_config().runtime.training.output_dir or request.training_output_dir
-            )
-
             # make the train call
-            log.debug("Training output dir: %s", output_dir)
+            log.debug("Training output dir: %s", request.training_output_dir)
+
             training_response = self._gts.run_training_job(
                 request=train_message_request,
                 module=train_module,
-                training_output_dir=output_dir,
+                training_output_dir=request.training_output_dir,
                 context=context,
                 wait=True,
                 # TODO: This usage of the server will sit behind Model Train and
@@ -125,12 +112,12 @@ class ModelTrainServicerImpl(process_pb2_grpc.ProcessServicer):
                 #   external_training_id override. If that is ever not the case
                 #   this will be brokenly passing that kwarg through to the
                 #   module's train function.
-                external_training_id=request.customTrainingID or request.trainingID,
+                external_training_id=request.trainingID,
             )
             log.debug("<RUN00837184D>", "training_response: %s", training_response)
             # return response
             process_response = process_pb2.ProcessResponse(
-                trainingID=request.trainingID,
+                trainingID=training_response.training_id,
                 customTrainingID=request.customTrainingID,
             )
             return process_response
diff --git a/examples/text-sentiment/client.py b/examples/text-sentiment/client.py
index 85f29134..7e13d7e0 100644
--- a/examples/text-sentiment/client.py
+++ b/examples/text-sentiment/client.py
@@ -21,9 +21,8 @@ import requests
 
 # Local
 from caikit.config.config import get_config
-from caikit.runtime import get_inference_request
+from caikit.core.data_model.base import DataBase
 from caikit.runtime.service_factory import ServicePackageFactory
-from text_sentiment.runtime_model import HuggingFaceSentimentModule
 import caikit
 
 if __name__ == "__main__":
@@ -52,9 +51,10 @@ if __name__ == "__main__":
 
         # Run inference for two sample prompts
         for text in ["I am not feeling well today!", "Today is a nice sunny day"]:
-            request = get_inference_request(HuggingFaceSentimentModule.TASK_CLASS)(
-                text_input=text
-            ).to_proto()
+            predict_class = DataBase.get_class_for_name(
+                "HuggingFaceSentimentTaskRequest"
+            )
+            request = predict_class(text_input=text).to_proto()
             response = client_stub.HuggingFaceSentimentTaskPredict(
                 request, metadata=[("mm-model-id", model_id)], timeout=1
             )


Test patch:
diff --git a/tests/fixtures/config/config.yml b/tests/fixtures/config/config.yml
index 1d4e55db..b8791d7f 100644
--- a/tests/fixtures/config/config.yml
+++ b/tests/fixtures/config/config.yml
@@ -1,5 +1,6 @@
 # Config for tests
 runtime:
+    library: sample_lib
     metrics:
         enabled: false
     metering:
diff --git a/tests/fixtures/sample_lib/__init__.py b/tests/fixtures/sample_lib/__init__.py
index a01314e4..f218fb0c 100644
--- a/tests/fixtures/sample_lib/__init__.py
+++ b/tests/fixtures/sample_lib/__init__.py
@@ -3,14 +3,7 @@ import os
 
 # Local
 from . import data_model, modules
-from .modules import (
-    CompositeModule,
-    InnerModule,
-    OtherModule,
-    SampleModule,
-    SamplePrimitiveModule,
-    StreamingModule,
-)
+from .modules import InnerModule, OtherModule, SampleModule, SamplePrimitiveModule
 from caikit.config import configure
 
 # Run configure for sample_lib configuration
diff --git a/tests/fixtures/sample_lib/modules/__init__.py b/tests/fixtures/sample_lib/modules/__init__.py
index 0b7d9b49..d3eaea6a 100644
--- a/tests/fixtures/sample_lib/modules/__init__.py
+++ b/tests/fixtures/sample_lib/modules/__init__.py
@@ -5,8 +5,6 @@ from .other_task import OtherModule
 from .sample_task import (
     CompositeModule,
     InnerModule,
-    ListModule,
     SampleModule,
     SamplePrimitiveModule,
-    StreamingModule,
 )
diff --git a/tests/runtime/http_server/test_http_server.py b/tests/runtime/http_server/test_http_server.py
index b69bc73b..5eff286e 100644
--- a/tests/runtime/http_server/test_http_server.py
+++ b/tests/runtime/http_server/test_http_server.py
@@ -31,7 +31,7 @@ import requests
 import tls_test_tools
 
 # Local
-from caikit.core import MODEL_MANAGER, DataObjectBase, dataobject
+from caikit.core import MODEL_MANAGER
 from caikit.runtime import http_server
 from tests.conftest import temp_config
 from tests.runtime.conftest import (
@@ -56,52 +56,40 @@ def save_key_cert_pair(prefix, workdir, key=None, cert=None):
     return crtfile, keyfile
 
 
-@dataobject
-class KeyPair(DataObjectBase):
+@dataclass
+class KeyPair:
     cert: str
     key: str
 
 
-@dataobject
-class TLSConfig(DataObjectBase):
+@dataclass
+class TLSConfig:
     server: KeyPair
     client: KeyPair
 
 
 @contextmanager
 def generate_tls_configs(
-    port: int,
-    tls: bool = False,
-    mtls: bool = False,
-    inline: bool = False,
-    **http_config_overrides,
+    port: int, tls: bool = False, mtls: bool = False, **http_config_overrides
 ) -> Dict[str, Dict]:
     """Helper to generate tls configs"""
     with tempfile.TemporaryDirectory() as workdir:
         config_overrides = {}
         client_keyfile, client_certfile, ca_certfile = None, None, None
-        ca_cert, server_cert, server_key = None, None, None
         if mtls or tls:
             ca_key = tls_test_tools.generate_key()[0]
             ca_cert = tls_test_tools.generate_ca_cert(ca_key)
             ca_certfile, _ = save_key_cert_pair("ca", workdir, cert=ca_cert)
-            server_key, server_cert = tls_test_tools.generate_derived_key_cert_pair(
-                ca_key=ca_key
-            )
             server_certfile, server_keyfile = save_key_cert_pair(
-                "server", workdir, server_key, server_cert
+                "server",
+                workdir,
+                *tls_test_tools.generate_derived_key_cert_pair(ca_key=ca_key),
             )
 
-            if inline:
-                tls_config = TLSConfig(
-                    server=KeyPair(cert=server_cert, key=server_key),
-                    client=KeyPair(cert=ca_cert if mtls else "", key=""),
-                )
-            else:
-                tls_config = TLSConfig(
-                    server=KeyPair(cert=server_certfile, key=server_keyfile),
-                    client=KeyPair(cert=ca_certfile if mtls else "", key=""),
-                )
+            tls_config = TLSConfig(
+                server=KeyPair(cert=server_certfile, key=server_keyfile),
+                client=KeyPair(cert="", key=""),
+            )
             # need to save this ca_certfile in config_overrides so the tls tests below can access it from client side
             config_overrides["use_in_test"] = {"ca_cert": ca_certfile}
 
@@ -120,11 +108,12 @@ def generate_tls_configs(
                     workdir,
                     *tls_test_tools.generate_derived_key_cert_pair(ca_key=ca_key),
                 )
+                tls_config.client = KeyPair(cert=ca_certfile, key="")
                 # need to save the client cert and key in config_overrides so the mtls test below can access it
                 config_overrides["use_in_test"]["client_cert"] = client_certfile
                 config_overrides["use_in_test"]["client_key"] = client_keyfile
 
-            config_overrides["runtime"] = {"tls": tls_config.to_dict()}
+            config_overrides["runtime"] = {"tls": tls_config}
         config_overrides.setdefault("runtime", {})["http"] = {
             "server_shutdown_grace_period_seconds": 0.01,  # this is so the server is killed after 0.1 if no test is running
             "port": port,
@@ -197,29 +186,6 @@ def test_mutual_tls_server(open_port):
             resp.raise_for_status()
 
 
-def test_mutual_tls_server_inline(open_port):
-    """Test that mutual TLS works when the TLS content is passed by value rather
-    than with files
-    """
-    with generate_tls_configs(
-        open_port, tls=True, mtls=True, inline=True, http_config_overrides={}
-    ) as config_overrides:
-        with runtime_http_test_server(
-            open_port,
-            tls_config_override=config_overrides,
-        ) as http_server_with_mtls:
-            # start a non-blocking http server with mutual tls
-            resp = requests.get(
-                f"https://localhost:{http_server_with_mtls.port}/docs",
-                verify=config_overrides["use_in_test"]["ca_cert"],
-                cert=(
-                    config_overrides["use_in_test"]["client_cert"],
-                    config_overrides["use_in_test"]["client_key"],
-                ),
-            )
-            resp.raise_for_status()
-
-
 def test_mutual_tls_server_with_wrong_cert(open_port):
     with generate_tls_configs(
         open_port, tls=True, mtls=True, http_config_overrides={}
diff --git a/tests/runtime/service_generation/test_data_stream_source.py b/tests/runtime/service_generation/test_data_stream_source.py
index e2168d10..182b03a8 100644
--- a/tests/runtime/service_generation/test_data_stream_source.py
+++ b/tests/runtime/service_generation/test_data_stream_source.py
@@ -4,7 +4,6 @@ from typing import List
 import json
 import os
 import pickle
-import shutil
 import tempfile
 
 # Third Party
@@ -337,23 +336,6 @@ def test_make_data_stream_source_jsonlfile(sample_train_service, sample_jsonl_fi
     validate_data_stream(data_stream, 2, SampleTrainingType)
 
 
-def test_make_data_stream_source_from_file_with_no_extension(
-    sample_train_service, sample_json_file, sample_jsonl_file, sample_csv_file, tmp_path
-):
-    stream_type = caikit.interfaces.common.data_model.DataStreamSourceSampleTrainingType
-    for file in [sample_json_file, sample_jsonl_file, sample_csv_file]:
-        no_extension_filename, _ = os.path.splitext(os.path.join(str(tmp_path), file))
-        shutil.copyfile(file, no_extension_filename)
-
-        ds = stream_type(file=stream_type.File(filename=no_extension_filename))
-        assert isinstance(ds, DataStreamSourceBase)
-
-        data_stream = ds.to_data_stream()
-        assert isinstance(data_stream, DataStream)
-
-        validate_data_stream(data_stream, 2, SampleTrainingType)
-
-
 #################
 # List of files
 #################
diff --git a/tests/runtime/servicers/test_global_predict_servicer_impl.py b/tests/runtime/servicers/test_global_predict_servicer_impl.py
index b1857606..644f3f51 100644
--- a/tests/runtime/servicers/test_global_predict_servicer_impl.py
+++ b/tests/runtime/servicers/test_global_predict_servicer_impl.py
@@ -16,7 +16,6 @@ from typing import Iterator
 from unittest.mock import MagicMock, patch
 
 # Local
-from caikit.runtime.service_factory import get_inference_request
 from sample_lib.modules.geospatial import GeoStreamingModule
 
 try:
@@ -35,6 +34,7 @@ import grpc
 import pytest
 
 # Local
+from caikit.core.data_model.base import DataBase
 from caikit.runtime.servicers.global_predict_servicer import GlobalPredictServicer
 from caikit.runtime.types.aborted_exception import AbortedException
 from caikit.runtime.types.caikit_runtime_exception import CaikitRuntimeException
@@ -57,7 +57,7 @@ def test_calling_predict_should_raise_if_module_raises(
 ):
     with pytest.raises(CaikitRuntimeException) as context:
         # SampleModules will raise a RuntimeError if the throw flag is set
-        predict_class = get_inference_request(SampleModule.TASK_CLASS)
+        predict_class = DataBase.get_class_for_name("SampleTaskRequest")
         request = predict_class(sample_input=HAPPY_PATH_INPUT_DM, throw=True).to_proto()
         sample_predict_servicer.Predict(
             request,
@@ -77,7 +77,7 @@ def test_invalid_input_to_a_valid_caikit_core_class_method_raises(
     """Test that a caikit.core module that gets an unexpected input value errors in an expected way"""
     with pytest.raises(CaikitRuntimeException) as context:
         # SampleModules will raise a ValueError if the poison pill name is given
-        predict_class = get_inference_request(SampleModule.TASK_CLASS)
+        predict_class = DataBase.get_class_for_name("SampleTaskRequest")
         request = predict_class(
             sample_input=SampleInputType(name=SampleModule.POISON_PILL_NAME)
         ).to_proto()
@@ -97,7 +97,7 @@ def test_global_predict_works_for_unary_rpcs(
     sample_task_unary_rpc,
 ):
     """Global predict of SampleTaskRequest returns a prediction"""
-    predict_class = get_inference_request(SampleModule.TASK_CLASS)
+    predict_class = DataBase.get_class_for_name("SampleTaskRequest")
     response = sample_predict_servicer.Predict(
         predict_class(sample_input=HAPPY_PATH_INPUT_DM).to_proto(),
         Fixtures.build_context(sample_task_model_id),
@@ -111,9 +111,7 @@ def test_global_predict_works_on_bidirectional_streaming_rpcs(
 ):
     """Simple test that our SampleModule's bidirectional stream inference fn is supported"""
 
-    predict_class = get_inference_request(
-        SampleModule.TASK_CLASS, input_streaming=True, output_streaming=True
-    )
+    predict_class = DataBase.get_class_for_name("BidiStreamingSampleTaskRequest")
 
     def req_iterator() -> Iterator[predict_class]:
         for i in range(100):
@@ -141,9 +139,7 @@ def test_global_predict_works_on_bidirectional_streaming_rpcs_with_multiple_stre
     mock_manager = MagicMock()
     mock_manager.retrieve_model.return_value = GeoStreamingModule()
 
-    predict_class = get_inference_request(
-        GeoStreamingModule.TASK_CLASS, input_streaming=True, output_streaming=True
-    )
+    predict_class = DataBase.get_class_for_name("BidiStreamingGeoSpatialTaskRequest")
 
     def req_iterator() -> Iterator[predict_class]:
         for i in range(100):
@@ -198,7 +194,7 @@ def test_global_predict_aborts_long_running_predicts(
     mock_manager.retrieve_model.return_value = dummy_model
 
     context = Fixtures.build_context("test-any-unresponsive-model")
-    predict_class = get_inference_request(SampleModule.TASK_CLASS)
+    predict_class = DataBase.get_class_for_name("SampleTaskRequest")
     predict_thread = threading.Thread(
         target=sample_predict_servicer.Predict,
         args=(
@@ -233,7 +229,7 @@ def test_metering_ignore_unsuccessful_calls(
         gps = GlobalPredictServicer(sample_inference_service)
         try:
             with patch.object(gps.rpc_meter, "update_metrics") as mock_update_func:
-                predict_class = get_inference_request(SampleModule.TASK_CLASS)
+                predict_class = DataBase.get_class_for_name("SampleTaskRequest")
                 request = predict_class(
                     sample_input=HAPPY_PATH_INPUT_DM, throw=True
                 ).to_proto()
@@ -257,7 +253,7 @@ def test_metering_predict_rpc_counter(
         sample_predict_servicer = GlobalPredictServicer(sample_inference_service)
         try:
             # Making 20 requests
-            predict_class = get_inference_request(SampleModule.TASK_CLASS)
+            predict_class = DataBase.get_class_for_name("SampleTaskRequest")
             for i in range(20):
                 sample_predict_servicer.Predict(
                     predict_class(sample_input=HAPPY_PATH_INPUT_DM).to_proto(),
@@ -298,7 +294,7 @@ def test_metering_write_to_metrics_file_twice(
         # need a new servicer to get a fresh new RPC meter
         sample_predict_servicer = GlobalPredictServicer(sample_inference_service)
         try:
-            predict_class = get_inference_request(SampleModule.TASK_CLASS)
+            predict_class = DataBase.get_class_for_name("SampleTaskRequest")
             sample_predict_servicer.Predict(
                 predict_class(sample_input=HAPPY_PATH_INPUT_DM).to_proto(),
                 Fixtures.build_context(sample_task_model_id),
diff --git a/tests/runtime/servicers/test_global_train_servicer_impl.py b/tests/runtime/servicers/test_global_train_servicer_impl.py
index f6826965..fbf5fa7c 100644
--- a/tests/runtime/servicers/test_global_train_servicer_impl.py
+++ b/tests/runtime/servicers/test_global_train_servicer_impl.py
@@ -27,13 +27,9 @@ import pytest
 # Local
 from caikit.config import get_config
 from caikit.core import MODEL_MANAGER
+from caikit.core.data_model.base import DataBase
 from caikit.core.data_model.producer import ProducerId
 from caikit.interfaces.common.data_model.stream_sources import S3Path
-from caikit.runtime.service_factory import (
-    get_inference_request,
-    get_train_params,
-    get_train_request,
-)
 from caikit.runtime.servicers.global_train_servicer import GlobalTrainServicer
 from caikit.runtime.types.caikit_runtime_exception import CaikitRuntimeException
 from sample_lib.data_model.sample import (
@@ -42,7 +38,7 @@ from sample_lib.data_model.sample import (
     SampleOutputType,
     SampleTrainingType,
 )
-from sample_lib.modules import CompositeModule, OtherModule, SampleModule
+from sample_lib.modules.sample_task.sample_implementation import SampleModule
 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
@@ -103,8 +99,10 @@ def test_global_train_sample_task(
         jsondata=stream_type.JsonData(data=[SampleTrainingType(1)])
     )
     model_name = random_test_id()
-    train_class = get_train_request(SampleModule)
-    train_request_params_class = get_train_params(SampleModule)
+    train_class = DataBase.get_class_for_name("SampleTaskSampleModuleTrainRequest")
+    train_request_params_class = DataBase.get_class_for_name(
+        "SampleTaskSampleModuleTrainParameters"
+    )
     train_request = train_class(
         model_name=model_name,
         parameters=train_request_params_class(
@@ -140,7 +138,7 @@ def test_global_train_sample_task(
         == "sample_lib.modules.sample_task.sample_implementation.SampleModule"
     )
 
-    predict_class = get_inference_request(SampleModule.TASK_CLASS)
+    predict_class = DataBase.get_class_for_name("SampleTaskRequest")
     inference_response = sample_predict_servicer.Predict(
         predict_class(sample_input=HAPPY_PATH_INPUT_DM).to_proto(),
         Fixtures.build_context(training_response.model_name),
@@ -167,8 +165,10 @@ def test_global_train_other_task(
     batch_size = 42
     stream_type = caikit.interfaces.common.data_model.DataStreamSourceInt
     training_data = stream_type(jsondata=stream_type.JsonData(data=[1]))
-    train_class = get_train_request(OtherModule)
-    train_request_params_class = get_train_params(OtherModule)
+    train_class = DataBase.get_class_for_name("OtherTaskOtherModuleTrainRequest")
+    train_request_params_class = DataBase.get_class_for_name(
+        "OtherTaskOtherModuleTrainParameters"
+    )
     train_request = train_class(
         model_name="Other module Training",
         parameters=train_request_params_class(
@@ -197,7 +197,7 @@ def test_global_train_other_task(
         == "sample_lib.modules.other_task.other_implementation.OtherModule"
     )
 
-    predict_class = get_inference_request(OtherModule.TASK_CLASS)
+    predict_class = DataBase.get_class_for_name("OtherTaskRequest")
     inference_response = sample_predict_servicer.Predict(
         predict_class(sample_input=HAPPY_PATH_INPUT_DM).to_proto(),
         Fixtures.build_context(training_response.model_name),
@@ -225,8 +225,10 @@ def test_global_train_Another_Widget_that_requires_SampleWidget_loaded_should_no
         model_id=sample_task_model_id
     )
 
-    train_class = get_train_request(CompositeModule)
-    train_request_params_class = get_train_params(CompositeModule)
+    train_class = DataBase.get_class_for_name("SampleTaskCompositeModuleTrainRequest")
+    train_request_params_class = DataBase.get_class_for_name(
+        "SampleTaskCompositeModuleTrainParameters"
+    )
     training_request = train_class(
         model_name="AnotherWidget_Training",
         parameters=train_request_params_class(sample_block=sample_model),
@@ -253,7 +255,7 @@ def test_global_train_Another_Widget_that_requires_SampleWidget_loaded_should_no
     )
 
     # make sure the trained model can run inference
-    predict_class = get_inference_request(SampleModule.TASK_CLASS)
+    predict_class = DataBase.get_class_for_name("SampleTaskRequest")
     inference_response = sample_predict_servicer.Predict(
         predict_class(sample_input=HAPPY_PATH_INPUT_DM).to_proto(),
         Fixtures.build_context(training_response.model_name),
@@ -278,8 +280,10 @@ def test_run_train_job_works_with_wait(
     training_data = stream_type(
         jsondata=stream_type.JsonData(data=[SampleTrainingType(1)])
     )
-    train_class = get_train_request(SampleModule)
-    train_request_params_class = get_train_params(SampleModule)
+    train_class = DataBase.get_class_for_name("SampleTaskSampleModuleTrainRequest")
+    train_request_params_class = DataBase.get_class_for_name(
+        "SampleTaskSampleModuleTrainParameters"
+    )
     train_request = train_class(
         model_name=random_test_id(),
         parameters=train_request_params_class(
@@ -302,7 +306,7 @@ def test_run_train_job_works_with_wait(
             training_response.training_id,
         )
 
-        predict_class = get_inference_request(SampleModule.TASK_CLASS)
+        predict_class = DataBase.get_class_for_name("SampleTaskRequest")
         inference_response = sample_predict_servicer.Predict(
             predict_class(sample_input=SampleInputType(name="Test")).to_proto(),
             Fixtures.build_context(training_response.model_name),
@@ -328,8 +332,10 @@ def test_global_train_Another_Widget_that_requires_SampleWidget_but_not_loaded_s
     model_id = random_test_id()
 
     sample_model = caikit.interfaces.runtime.data_model.ModelPointer(model_id=model_id)
-    train_class = get_train_request(CompositeModule)
-    train_request_params_class = get_train_params(CompositeModule)
+    train_class = DataBase.get_class_for_name("SampleTaskCompositeModuleTrainRequest")
+    train_request_params_class = DataBase.get_class_for_name(
+        "SampleTaskCompositeModuleTrainParameters"
+    )
     request = train_class(
         model_name="AnotherWidget_Training",
         parameters=train_request_params_class(sample_block=sample_model),
@@ -350,8 +356,10 @@ def test_global_train_Edge_Case_Widget_should_raise_when_error_surfaces_from_mod
         jsondata=stream_type.JsonData(data=[SampleTrainingType(1)])
     )
 
-    train_class = get_train_request(SampleModule)
-    train_request_params_class = get_train_params(SampleModule)
+    train_class = DataBase.get_class_for_name("SampleTaskSampleModuleTrainRequest")
+    train_request_params_class = DataBase.get_class_for_name(
+        "SampleTaskSampleModuleTrainParameters"
+    )
     train_request = train_class(
         model_name=random_test_id(),
         parameters=train_request_params_class(
@@ -379,8 +387,10 @@ def test_global_train_returns_exit_code_with_oom(
     training_data = stream_type(
         jsondata=stream_type.JsonData(data=[SampleTrainingType(1)])
     )
-    train_class = get_train_request(SampleModule)
-    train_request_params_class = get_train_params(SampleModule)
+    train_class = DataBase.get_class_for_name("SampleTaskSampleModuleTrainRequest")
+    train_request_params_class = DataBase.get_class_for_name(
+        "SampleTaskSampleModuleTrainParameters"
+    )
     train_request = train_class(
         model_name=random_test_id(),
         parameters=train_request_params_class(
@@ -411,8 +421,10 @@ def test_local_trainer_rejects_s3_output_paths(
     training_data = stream_type(
         jsondata=stream_type.JsonData(data=[SampleTrainingType(1)])
     )
-    train_class = get_train_request(SampleModule)
-    train_request_params_class = get_train_params(SampleModule)
+    train_class = DataBase.get_class_for_name("SampleTaskSampleModuleTrainRequest")
+    train_request_params_class = DataBase.get_class_for_name(
+        "SampleTaskSampleModuleTrainParameters"
+    )
     train_request = train_class(
         model_name=random_test_id(),
         output_path=S3Path(path="foo"),
@@ -445,8 +457,10 @@ def test_global_train_aborts_long_running_trains(
     )
     training_id = random_test_id()
 
-    train_class = get_train_request(SampleModule)
-    train_request_params_class = get_train_params(SampleModule)
+    train_class = DataBase.get_class_for_name("SampleTaskSampleModuleTrainRequest")
+    train_request_params_class = DataBase.get_class_for_name(
+        "SampleTaskSampleModuleTrainParameters"
+    )
     train_request = train_class(
         model_name=training_id,
         parameters=train_request_params_class(
diff --git a/tests/runtime/servicers/test_model_train_servicer_impl.py b/tests/runtime/servicers/test_model_train_servicer_impl.py
index 4077a31a..e4625692 100644
--- a/tests/runtime/servicers/test_model_train_servicer_impl.py
+++ b/tests/runtime/servicers/test_model_train_servicer_impl.py
@@ -185,95 +185,47 @@ def test_model_train_validation_error_raises(sample_model_train_servicer, output
 
 #####################################################################
 # Normal tests
-@pytest.mark.parametrize("output_in_config", (True, False))
-@pytest.mark.parametrize("save_with_id", (True, False))
-@pytest.mark.parametrize("use_custom_id", (True, False))
-def test_model_train_sample_widget(
-    sample_model_train_servicer,
-    output_dir,
-    output_in_config,
-    save_with_id,
-    use_custom_id,
-):
-    """This test tests end-to-end training. It includes verifying that the model
-    is saved in the right place. The place where the model is saved comes from:
-
-    1. get_config().runtime.training.output_dir
-    2. request.training_output_dir
-
-    IN THAT ORDER!
-    """
+def test_model_train_sample_widget(sample_model_train_servicer, output_dir):
     training_id = str(uuid.uuid4())
-    custom_id = f"custom-{uuid.uuid4()}" if use_custom_id else None
-    model_name = "abc"
-
-    # If getting the output path from config, use a bogus dir in the request to
-    # ensure that the config is preferred
-    if output_in_config:
-        config_output_dir = output_dir
-        req_output_dir = "/some/bad/path"
-    # Otherwise the config must be set to an empty string so the request is used
-    else:
-        config_output_dir = ""
-        req_output_dir = output_dir
-
-    with temp_config(
-        {
-            "runtime": {
-                "training": {
-                    "output_dir": config_output_dir,
-                    "save_with_id": save_with_id,
-                },
-            },
-        },
-        "merge",
-    ):
-        model_train_request = process_pb2.ProcessRequest(
-            trainingID=training_id,
-            customTrainingID=custom_id,
-            request_dict={
-                "train_module": "00110203-0405-0607-0809-0a0b02dd0e0f",
-                "training_params": json.dumps(
-                    {
-                        "model_name": model_name,
-                        "parameters": {
-                            "training_data": {
-                                "jsondata": {
-                                    "data": [
-                                        sample_lib.data_model.SampleTrainingType(
-                                            number=1
-                                        ).to_dict()
-                                    ]
-                                },
+
+    training_output_dir = os.path.join(output_dir, training_id)
+    model_train_request = process_pb2.ProcessRequest(
+        trainingID=training_id,
+        customTrainingID=str(uuid.uuid4()),
+        request_dict={
+            "train_module": "00110203-0405-0607-0809-0a0b02dd0e0f",
+            "training_params": json.dumps(
+                {
+                    "model_name": "abc",
+                    "parameters": {
+                        "training_data": {
+                            "jsondata": {
+                                "data": [
+                                    sample_lib.data_model.SampleTrainingType(
+                                        number=1
+                                    ).to_dict()
+                                ]
                             },
                         },
-                    }
-                ),
-            },
-            training_input_dir="training_input_dir",
-            training_output_dir=req_output_dir,
-        )
-        context = Fixtures.build_context("test-any-unresponsive-model")
-        training_response = sample_model_train_servicer.Run(
-            model_train_request, context
-        )
-        if save_with_id:
-            if use_custom_id:
-                model_save_path = os.path.join(output_dir, custom_id, model_name)
-            else:
-                model_save_path = os.path.join(output_dir, training_id, model_name)
-        else:
-            model_save_path = os.path.join(output_dir, model_name)
-        assert os.path.isdir(model_save_path)
+                    },
+                }
+            ),
+        },
+        training_input_dir="training_input_dir",
+        training_output_dir=training_output_dir,
+    )
+    context = Fixtures.build_context("test-any-unresponsive-model")
+    training_response = sample_model_train_servicer.Run(model_train_request, context)
+    assert os.path.isdir(training_output_dir)
 
     # Make sure that the request completed synchronously
-    model_future = MODEL_MANAGER.get_model_future(custom_id or training_id)
+    model_future = MODEL_MANAGER.get_model_future(training_response.trainingID)
     assert model_future.get_info().status == TrainingStatus.COMPLETED
 
     # Make sure that the return object looks right
     assert isinstance(training_response, process_pb2.ProcessResponse)
     assert training_response == process_pb2.ProcessResponse(
-        trainingID=training_id,
+        trainingID=model_train_request.trainingID,
         customTrainingID=model_train_request.customTrainingID,
     )
 
@@ -297,6 +249,7 @@ def test_files_from_training_input_dir_are_used(
 
     model_train_request = process_pb2.ProcessRequest(
         trainingID=training_id,
+        customTrainingID=str(uuid.uuid4()),
         request_dict={
             "train_module": "00110203-0405-0607-0809-0a0b02dd0e0f",
             "training_params": json.dumps(
diff --git a/tests/runtime/test_dump_services.py b/tests/runtime/test_dump_services.py
new file mode 100644
index 00000000..b4d0e412
--- /dev/null
+++ b/tests/runtime/test_dump_services.py
@@ -0,0 +1,69 @@
+# 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.
+# Standard
+import os
+import shutil
+import tempfile
+
+# First Party
+import alog
+
+# Local
+from caikit.runtime.dump_services import dump_grpc_services, dump_http_services
+
+## Helpers #####################################################################
+
+log = alog.use_channel("TEST-DUMP-I")
+
+
+def test_dump_grpc_services_dir_exists():
+    with tempfile.TemporaryDirectory() as workdir:
+        dump_grpc_services(workdir)
+        assert os.path.exists(workdir)
+
+        for file in os.listdir(workdir):
+            assert file.endswith(".proto")
+
+
+def test_dump_grpc_services_dir_does_not_exist():
+    fake_dir = "fake_dir"
+    dump_grpc_services(fake_dir)
+    assert os.path.exists(fake_dir)
+
+    for file in os.listdir(fake_dir):
+        assert file.endswith(".proto")
+
+    shutil.rmtree(fake_dir)
+
+
+def test_dump_http_services_dir_exists():
+    with tempfile.TemporaryDirectory() as workdir:
+        dump_http_services(workdir)
+        assert os.path.exists(workdir)
+
+        for file in os.listdir(workdir):
+            assert file == "openapi.json"
+            assert os.path.getsize(os.path.join(workdir, file)) > 0
+
+
+def test_dump_http_services_dir_does_not_exist():
+    fake_dir = "fake_dir"
+    dump_http_services(fake_dir)
+    assert os.path.exists(fake_dir)
+
+    for file in os.listdir(fake_dir):
+        assert file == "openapi.json"
+        assert os.path.getsize(os.path.join(fake_dir, file)) > 0
+
+    shutil.rmtree(fake_dir)
diff --git a/tests/runtime/test_grpc_server.py b/tests/runtime/test_grpc_server.py
index b71d4243..1c0a3de5 100644
--- a/tests/runtime/test_grpc_server.py
+++ b/tests/runtime/test_grpc_server.py
@@ -42,6 +42,7 @@ import alog
 # Local
 from caikit import get_config
 from caikit.core import MODEL_MANAGER
+from caikit.core.data_model.base import DataBase
 from caikit.core.data_model.producer import ProducerId
 from caikit.interfaces.runtime.data_model import (
     TrainingInfoRequest,
@@ -49,7 +50,6 @@ from caikit.interfaces.runtime.data_model import (
     TrainingStatus,
     TrainingStatusResponse,
 )
-from caikit.runtime import get_inference_request, get_train_params, get_train_request
 from caikit.runtime.grpc_server import RuntimeGRPCServer
 from caikit.runtime.model_management.model_manager import ModelManager
 from caikit.runtime.protobufs import (
@@ -60,14 +60,7 @@ from caikit.runtime.protobufs import (
 )
 from caikit.runtime.service_factory import ServicePackage, ServicePackageFactory
 from caikit.runtime.utils.servicer_util import build_caikit_library_request_dict
-from sample_lib import (
-    CompositeModule,
-    InnerModule,
-    OtherModule,
-    SampleModule,
-    SamplePrimitiveModule,
-    StreamingModule,
-)
+from sample_lib import InnerModule, SamplePrimitiveModule
 from sample_lib.data_model import (
     OtherOutputType,
     SampleInputType,
@@ -83,6 +76,7 @@ from tests.runtime.conftest import (
     runtime_grpc_test_server,
 )
 import caikit.interfaces.common
+import sample_lib
 
 ## Helpers #####################################################################
 
@@ -128,6 +122,7 @@ def test_model_train(runtime_grpc_server):
     model_name = "abc"
     model_train_request = process_pb2.ProcessRequest(
         trainingID=training_id,
+        customTrainingID=str(uuid.uuid4()),
         request_dict={
             "train_module": "00110203-0405-0607-0809-0a0b02dd0e0f",
             "training_params": json.dumps(
@@ -136,7 +131,11 @@ def test_model_train(runtime_grpc_server):
                     "parameters": {
                         "training_data": {
                             "jsondata": {
-                                "data": [SampleTrainingType(number=1).to_dict()]
+                                "data": [
+                                    sample_lib.data_model.SampleTrainingType(
+                                        number=1
+                                    ).to_dict()
+                                ]
                             },
                         },
                     },
@@ -213,9 +212,8 @@ def test_predict_sample_module_ok_response(
 ):
     """Test RPC CaikitRuntime.SampleTaskPredict successful response"""
     stub = sample_inference_service.stub_class(runtime_grpc_server.make_local_channel())
-    predict_request = get_inference_request(SampleModule.TASK_CLASS)(
-        sample_input=HAPPY_PATH_INPUT_DM
-    ).to_proto()
+    predict_class = DataBase.get_class_for_name("SampleTaskRequest")
+    predict_request = predict_class(sample_input=HAPPY_PATH_INPUT_DM).to_proto()
 
     actual_response = stub.SampleTaskPredict(
         predict_request, metadata=[("mm-model-id", sample_task_model_id)]
@@ -234,7 +232,7 @@ def test_global_predict_build_caikit_library_request_dict_creates_caikit_core_ru
     )
     proto_request_dict = build_caikit_library_request_dict(
         proto_request,
-        SampleModule.RUN_SIGNATURE,
+        sample_lib.modules.sample_task.SampleModule.RUN_SIGNATURE,
     )
 
     # unset fields not included
@@ -243,12 +241,12 @@ def test_global_predict_build_caikit_library_request_dict_creates_caikit_core_ru
     assert proto_expected_arguments == set(proto_request_dict.keys())
 
     # pythonic data model request
-    predict_class = get_inference_request(SampleModule.TASK_CLASS)
+    predict_class = DataBase.get_class_for_name("SampleTaskRequest")
     python_request = predict_class(sample_input=HAPPY_PATH_INPUT_DM).to_proto()
 
     python_sample_module_request_dict = build_caikit_library_request_dict(
         python_request,
-        SampleModule.RUN_SIGNATURE,
+        sample_lib.modules.sample_task.SampleModule.RUN_SIGNATURE,
     )
 
     # unset fields are included if they have defaults set
@@ -262,9 +260,7 @@ def test_predict_streaming_module(
 ):
     """Test RPC CaikitRuntime.StreamingTaskPredict successful response"""
     stub = sample_inference_service.stub_class(runtime_grpc_server.make_local_channel())
-    predict_class = get_inference_request(
-        StreamingModule.TASK_CLASS, input_streaming=False, output_streaming=True
-    )
+    predict_class = DataBase.get_class_for_name("ServerStreamingStreamingTaskRequest")
     predict_request = predict_class(sample_input=HAPPY_PATH_INPUT_DM).to_proto()
 
     stream = stub.ServerStreamingStreamingTaskPredict(
@@ -286,7 +282,7 @@ def test_predict_sample_module_error_response(
         stub = sample_inference_service.stub_class(
             runtime_grpc_server.make_local_channel()
         )
-        predict_class = get_inference_request(SampleModule.TASK_CLASS)
+        predict_class = DataBase.get_class_for_name("SampleTaskRequest")
         predict_request = predict_class(sample_input=HAPPY_PATH_INPUT_DM).to_proto()
 
         stub.SampleTaskPredict(
@@ -301,7 +297,7 @@ def test_rpc_validation_on_predict(
 ):
     """Check that the server catches models sent to the wrong task RPCs"""
     stub = sample_inference_service.stub_class(runtime_grpc_server.make_local_channel())
-    predict_class = get_inference_request(OtherModule.TASK_CLASS)
+    predict_class = DataBase.get_class_for_name("OtherTaskRequest")
     predict_request = predict_class(
         sample_input_sampleinputtype=HAPPY_PATH_INPUT_DM
     ).to_proto()
@@ -332,7 +328,7 @@ def test_rpc_validation_on_predict_for_unsupported_model(
         stub = sample_inference_service.stub_class(
             runtime_grpc_server.make_local_channel()
         )
-        predict_class = get_inference_request(SampleModule.TASK_CLASS)
+        predict_class = DataBase.get_class_for_name("SampleTaskRequest")
         predict_request = predict_class(sample_input=HAPPY_PATH_INPUT_DM).to_proto()
         with pytest.raises(grpc.RpcError) as context:
             stub.SampleTaskPredict(
@@ -364,10 +360,7 @@ def test_rpc_validation_on_predict_for_wrong_streaming_flavor(
         stub = sample_inference_service.stub_class(
             runtime_grpc_server.make_local_channel()
         )
-
-        predict_class = get_inference_request(
-            SampleModule.TASK_CLASS,
-        )
+        predict_class = DataBase.get_class_for_name("SampleTaskRequest")
         predict_request = predict_class(sample_input=HAPPY_PATH_INPUT_DM).to_proto()
         with pytest.raises(grpc.RpcError) as context:
             response = stub.ServerStreamingSampleTaskPredict(
@@ -404,9 +397,15 @@ def test_train_fake_module_ok_response_and_can_predict_with_trained_model(
         )
     )
     model_name = random_test_id()
-    train_request = get_train_request(SampleModule)(
+    train_request_class = DataBase.get_class_for_name(
+        "SampleTaskSampleModuleTrainRequest"
+    )
+    train_request_params_class = DataBase.get_class_for_name(
+        "SampleTaskSampleModuleTrainParameters"
+    )
+    train_request = train_request_class(
         model_name=model_name,
-        parameters=get_train_params(SampleModule)(
+        parameters=train_request_params_class(
             training_data=training_data,
             union_list=["str", "sequence"],
         ),
@@ -422,7 +421,7 @@ def test_train_fake_module_ok_response_and_can_predict_with_trained_model(
     )
 
     # make sure the trained model can run inference
-    predict_class = get_inference_request(SampleModule.TASK_CLASS)
+    predict_class = DataBase.get_class_for_name("SampleTaskRequest")
     predict_request = predict_class(sample_input=HAPPY_PATH_INPUT_DM).to_proto()
     inference_response = inference_stub.SampleTaskPredict(
         predict_request, metadata=[("mm-model-id", actual_response.model_name)]
@@ -444,9 +443,13 @@ def test_train_fake_module_ok_response_with_loaded_model_can_predict_with_traine
         model_id=sample_task_model_id
     )
     model_name = random_test_id()
-    train_request = get_train_request(CompositeModule)(
+    train_class = DataBase.get_class_for_name("SampleTaskCompositeModuleTrainRequest")
+    train_request_params_class = DataBase.get_class_for_name(
+        "SampleTaskCompositeModuleTrainParameters"
+    )
+    train_request = train_class(
         model_name=model_name,
-        parameters=get_train_params(CompositeModule)(sample_block=sample_model),
+        parameters=train_request_params_class(sample_block=sample_model),
     ).to_proto()
 
     actual_response = train_stub.SampleTaskCompositeModuleTrain(train_request)
@@ -458,7 +461,7 @@ def test_train_fake_module_ok_response_with_loaded_model_can_predict_with_traine
     )
 
     # make sure the trained model can run inference
-    predict_class = get_inference_request(CompositeModule.TASK_CLASS)
+    predict_class = DataBase.get_class_for_name("SampleTaskRequest")
     predict_request = predict_class(sample_input=HAPPY_PATH_INPUT_DM).to_proto()
     inference_response = inference_stub.SampleTaskPredict(
         predict_request, metadata=[("mm-model-id", actual_response.model_name)]
@@ -485,9 +488,13 @@ def test_train_fake_module_does_not_change_another_instance_model_of_block(
     stream_type = caikit.interfaces.common.data_model.DataStreamSourceInt
     training_data = stream_type(file=stream_type.File(filename=sample_int_file))
 
-    train_request = get_train_request(OtherModule)(
+    train_class = DataBase.get_class_for_name("OtherTaskOtherModuleTrainRequest")
+    train_request_params_class = DataBase.get_class_for_name(
+        "OtherTaskOtherModuleTrainParameters"
+    )
+    train_request = train_class(
         model_name="Bar Training",
-        parameters=get_train_params(OtherModule)(
+        parameters=train_request_params_class(
             sample_input_sampleinputtype=HAPPY_PATH_INPUT_DM,
             batch_size=100,
             training_data=training_data,
@@ -505,7 +512,7 @@ def test_train_fake_module_does_not_change_another_instance_model_of_block(
     )
 
     # make sure the trained model can run inference, and the batch size 100 was used
-    predict_class = get_inference_request(OtherModule.TASK_CLASS)
+    predict_class = DataBase.get_class_for_name("OtherTaskRequest")
     predict_request = predict_class(sample_input=HAPPY_PATH_INPUT_DM).to_proto()
     trained_inference_response = inference_stub.OtherTaskPredict(
         predict_request, metadata=[("mm-model-id", actual_response.model_name)]
@@ -536,10 +543,16 @@ def test_train_primitive_model(
     """Test that we can make a successful training and inference call to the primitive module using primitive inputs"""
 
     model_name = "primitive_trained_model"
+    train_request_class = DataBase.get_class_for_name(
+        "SampleTaskSamplePrimitiveModuleTrainRequest"
+    )
+    train_request_params_class = DataBase.get_class_for_name(
+        "SampleTaskSamplePrimitiveModuleTrainParameters"
+    )
 
-    train_request = get_train_request(SamplePrimitiveModule)(
+    train_request = train_request_class(
         model_name=model_name,
-        parameters=get_train_params(SamplePrimitiveModule)(
+        parameters=train_request_params_class(
             sample_input=HAPPY_PATH_INPUT_DM,
             simple_list=["hello", "world"],
             union_list=["str", "sequence"],
@@ -565,7 +578,7 @@ def test_train_primitive_model(
     )
 
     # make sure the trained model can run inference
-    predict_class = get_inference_request(SampleModule.TASK_CLASS)
+    predict_class = DataBase.get_class_for_name("SampleTaskRequest")
     predict_request = predict_class(sample_input=HAPPY_PATH_INPUT_DM).to_proto()
 
     inference_response = inference_stub.SampleTaskPredict(
@@ -594,10 +607,13 @@ def test_train_fake_module_ok_response_with_datastream_jsondata(
         )
     )
     model_name = random_test_id()
-
-    train_request = get_train_request(SampleModule)(
+    train_class = DataBase.get_class_for_name("SampleTaskSampleModuleTrainRequest")
+    train_request_params_class = DataBase.get_class_for_name(
+        "SampleTaskSampleModuleTrainParameters"
+    )
+    train_request = train_class(
         model_name=model_name,
-        parameters=get_train_params(SampleModule)(
+        parameters=train_request_params_class(
             batch_size=42,
             training_data=training_data,
         ),
@@ -612,7 +628,7 @@ def test_train_fake_module_ok_response_with_datastream_jsondata(
     )
 
     # make sure the trained model can run inference
-    predict_class = get_inference_request(SampleModule.TASK_CLASS)
+    predict_class = DataBase.get_class_for_name("SampleTaskRequest")
     predict_request = predict_class(sample_input=HAPPY_PATH_INPUT_DM).to_proto()
     inference_response = inference_stub.SampleTaskPredict(
         predict_request, metadata=[("mm-model-id", actual_response.model_name)]
@@ -633,10 +649,13 @@ def test_train_fake_module_ok_response_with_datastream_csv_file(
     stream_type = caikit.interfaces.common.data_model.DataStreamSourceSampleTrainingType
     training_data = stream_type(file=stream_type.File(filename=sample_csv_file))
     model_name = random_test_id()
-
-    train_request = get_train_request(SampleModule)(
+    train_class = DataBase.get_class_for_name("SampleTaskSampleModuleTrainRequest")
+    train_request_params_class = DataBase.get_class_for_name(
+        "SampleTaskSampleModuleTrainParameters"
+    )
+    train_request = train_class(
         model_name=model_name,
-        parameters=get_train_params(SampleModule)(
+        parameters=train_request_params_class(
             training_data=training_data,
         ),
     ).to_proto()
@@ -650,7 +669,7 @@ def test_train_fake_module_ok_response_with_datastream_csv_file(
     )
 
     # make sure the trained model can run inference
-    predict_class = get_inference_request(SampleModule.TASK_CLASS)
+    predict_class = DataBase.get_class_for_name("SampleTaskRequest")
     predict_request = predict_class(sample_input=HAPPY_PATH_INPUT_DM).to_proto()
     inference_response = inference_stub.SampleTaskPredict(
         predict_request, metadata=[("mm-model-id", actual_response.model_name)]
@@ -670,10 +689,13 @@ def test_train_and_successfully_cancel_training(
     )
     model_name = random_test_id()
     # start a training that sleeps for a long time, so I can cancel
-
-    train_request = get_train_request(SampleModule)(
+    train_class = DataBase.get_class_for_name("SampleTaskSampleModuleTrainRequest")
+    train_request_params_class = DataBase.get_class_for_name(
+        "SampleTaskSampleModuleTrainParameters"
+    )
+    train_request = train_class(
         model_name=model_name,
-        parameters=get_train_params(SampleModule)(
+        parameters=train_request_params_class(
             training_data=training_data, sleep_time=10
         ),
     ).to_proto()
@@ -708,9 +730,13 @@ def test_cancel_does_not_affect_other_models(
     )
     model_name = random_test_id()
     # start a training that sleeps for a long time, so I can cancel
-    train_request = get_train_request(SampleModule)(
+    train_class = DataBase.get_class_for_name("SampleTaskSampleModuleTrainRequest")
+    train_request_params_class = DataBase.get_class_for_name(
+        "SampleTaskSampleModuleTrainParameters"
+    )
+    train_request = train_class(
         model_name=model_name,
-        parameters=get_train_params(SampleModule)(
+        parameters=train_request_params_class(
             training_data=training_data, sleep_time=10
         ),
     ).to_proto()
@@ -731,9 +757,9 @@ def test_cancel_does_not_affect_other_models(
 
     # train another model
     model_name2 = random_test_id()
-    train_request2 = get_train_request(SampleModule)(
+    train_request2 = train_class(
         model_name=model_name2,
-        parameters=get_train_params(SampleModule)(training_data=training_data),
+        parameters=train_request_params_class(training_data=training_data),
     ).to_proto()
     train_response2 = train_stub.SampleTaskSampleModuleTrain(train_request2)
 
@@ -764,9 +790,15 @@ def test_train_fake_module_error_response_with_unloaded_model(
         sample_model = caikit.interfaces.runtime.data_model.ModelPointer(
             model_id=random_test_id()
         )
-        train_request = get_train_request(CompositeModule)(
+        train_class = DataBase.get_class_for_name(
+            "SampleTaskCompositeModuleTrainRequest"
+        )
+        train_request_params_class = DataBase.get_class_for_name(
+            "SampleTaskCompositeModuleTrainParameters"
+        )
+        train_request = train_class(
             model_name=random_test_id(),
-            parameters=get_train_params(CompositeModule)(sample_block=sample_model),
+            parameters=train_request_params_class(sample_block=sample_model),
         ).to_proto()
 
         train_stub.SampleTaskCompositeModuleTrain(train_request)
@@ -1115,7 +1147,7 @@ def test_metrics_stored_after_server_interrupt(
     with temp_config({"runtime": {"metering": {"enabled": True}}}, "merge"):
         with runtime_grpc_test_server(open_port) as server:
             stub = sample_inference_service.stub_class(server.make_local_channel())
-            predict_class = get_inference_request(SampleModule.TASK_CLASS)
+            predict_class = DataBase.get_class_for_name("SampleTaskRequest")
             predict_request = predict_class(sample_input=HAPPY_PATH_INPUT_DM).to_proto()
             _ = stub.SampleTaskPredict(
                 predict_request, metadata=[("mm-model-id", sample_task_model_id)]
diff --git a/tests/runtime/test_service_factory.py b/tests/runtime/test_service_factory.py
index 624ab202..36da1428 100644
--- a/tests/runtime/test_service_factory.py
+++ b/tests/runtime/test_service_factory.py
@@ -23,16 +23,11 @@ import pytest
 
 # Local
 from caikit.core.data_model import render_dataobject_protos
-from caikit.runtime.service_factory import (
-    ServicePackage,
-    ServicePackageFactory,
-    get_inference_request,
-    get_train_params,
-    get_train_request,
-)
+from caikit.core.data_model.base import DataBase
+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 import ListModule, OtherModule
+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
@@ -365,71 +360,8 @@ def test_backend_modules_included_in_service_generation(
     inference_service = ServicePackageFactory.get_service_package(
         ServicePackageFactory.ServiceType.INFERENCE
     )
-    predict_class = get_inference_request(SampleModule.TASK_CLASS)
+    predict_class = DataBase.get_class_for_name("SampleTaskRequest")
     sample_task_request = predict_class().to_proto()
 
     # 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()
-
-
-def test_get_inference_request_throws_wrong_type(runtime_grpc_server):
-    with pytest.raises(TypeError) as e:
-        get_inference_request(task_or_module_class="something random")
-    assert "subclass check failed" in e.value.args[0]
-
-
-def test_get_inference_request(runtime_grpc_server):
-    """Test that we are able to get inference request DM with either module or task class"""
-    assert get_inference_request(SampleModule).__name__ == "SampleTaskRequest"
-    assert (
-        get_inference_request(SampleModule.TASK_CLASS).__name__ == "SampleTaskRequest"
-    )
-    assert (
-        get_inference_request(SampleModule, output_streaming=True).__name__
-        == "ServerStreamingSampleTaskRequest"
-    )
-    assert (
-        get_inference_request(SampleModule.TASK_CLASS, output_streaming=True).__name__
-        == "ServerStreamingSampleTaskRequest"
-    )
-    assert (
-        get_inference_request(
-            SampleModule, input_streaming=True, output_streaming=True
-        ).__name__
-        == "BidiStreamingSampleTaskRequest"
-    )
-    assert (
-        get_inference_request(
-            SampleModule.TASK_CLASS, input_streaming=True, output_streaming=True
-        ).__name__
-        == "BidiStreamingSampleTaskRequest"
-    )
-
-
-def test_get_train_request_throws_wrong_type(runtime_grpc_server):
-    with pytest.raises(TypeError) as e:
-        get_train_request("not_a_module")
-    assert "subclass check failed" in e.value.args[0]
-
-
-def test_get_train_request(runtime_grpc_server):
-    assert (
-        get_train_request(SampleModule).__name__ == "SampleTaskSampleModuleTrainRequest"
-    )
-    assert get_train_request(OtherModule).__name__ == "OtherTaskOtherModuleTrainRequest"
-
-
-def test_get_train_params_throws_wrong_type(runtime_grpc_server):
-    with pytest.raises(TypeError) as e:
-        get_train_params("not_a_module")
-    assert "subclass check failed" in e.value.args[0]
-
-
-def test_get_train_params(runtime_grpc_server):
-    assert (
-        get_train_params(SampleModule).__name__
-        == "SampleTaskSampleModuleTrainParameters"
-    )
-    assert (
-        get_train_params(OtherModule).__name__ == "OtherTaskOtherModuleTrainParameters"
-    )