Repository: caikit/caikit
Pull request #475 (2023-09-15 19:57:27)

✨ add a better way to get request DMs
<!--  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**:

for https://github.com/caikit/caikit/issues/476.

Essentially, we should now be able to do this:
```
caikit.runtime.get_train_request(sample_module.modules.SampleModule)(
    model_name=model_id,
    parameters=get_train_params(sample_lib.modules.SampleModule)(
        training_data=training_data
    ),
)
```

**Special notes for your reviewer**:

There are still some tests left that will need to be changed from `DataBase.get_class...` to this

**If applicable**:
- [ ] this PR contains documentation
- [ ] 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 2068258b..4b57df6b 100644
--- a/caikit/runtime/__init__.py
+++ b/caikit/runtime/__init__.py
@@ -11,3 +11,6 @@
 # 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/service_factory.py b/caikit/runtime/service_factory.py
index 1970a82a..e6acab48 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
+from typing import Callable, Dict, Set, Type, Union
 import dataclasses
 
 # Third Party
@@ -31,7 +31,10 @@ 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,
@@ -41,6 +44,7 @@ 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 = {
@@ -266,3 +270,63 @@ 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/examples/text-sentiment/client.py b/examples/text-sentiment/client.py
index 7e13d7e0..85f29134 100644
--- a/examples/text-sentiment/client.py
+++ b/examples/text-sentiment/client.py
@@ -21,8 +21,9 @@ import requests
 
 # Local
 from caikit.config.config import get_config
-from caikit.core.data_model.base import DataBase
+from caikit.runtime import get_inference_request
 from caikit.runtime.service_factory import ServicePackageFactory
+from text_sentiment.runtime_model import HuggingFaceSentimentModule
 import caikit
 
 if __name__ == "__main__":
@@ -51,10 +52,9 @@ if __name__ == "__main__":
 
         # Run inference for two sample prompts
         for text in ["I am not feeling well today!", "Today is a nice sunny day"]:
-            predict_class = DataBase.get_class_for_name(
-                "HuggingFaceSentimentTaskRequest"
-            )
-            request = predict_class(text_input=text).to_proto()
+            request = get_inference_request(HuggingFaceSentimentModule.TASK_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/sample_lib/__init__.py b/tests/fixtures/sample_lib/__init__.py
index f218fb0c..a01314e4 100644
--- a/tests/fixtures/sample_lib/__init__.py
+++ b/tests/fixtures/sample_lib/__init__.py
@@ -3,7 +3,14 @@ import os
 
 # Local
 from . import data_model, modules
-from .modules import InnerModule, OtherModule, SampleModule, SamplePrimitiveModule
+from .modules import (
+    CompositeModule,
+    InnerModule,
+    OtherModule,
+    SampleModule,
+    SamplePrimitiveModule,
+    StreamingModule,
+)
 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 d3eaea6a..0b7d9b49 100644
--- a/tests/fixtures/sample_lib/modules/__init__.py
+++ b/tests/fixtures/sample_lib/modules/__init__.py
@@ -5,6 +5,8 @@ from .other_task import OtherModule
 from .sample_task import (
     CompositeModule,
     InnerModule,
+    ListModule,
     SampleModule,
     SamplePrimitiveModule,
+    StreamingModule,
 )
diff --git a/tests/runtime/servicers/test_global_predict_servicer_impl.py b/tests/runtime/servicers/test_global_predict_servicer_impl.py
index 644f3f51..b1857606 100644
--- a/tests/runtime/servicers/test_global_predict_servicer_impl.py
+++ b/tests/runtime/servicers/test_global_predict_servicer_impl.py
@@ -16,6 +16,7 @@ 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:
@@ -34,7 +35,6 @@ 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 = DataBase.get_class_for_name("SampleTaskRequest")
+        predict_class = get_inference_request(SampleModule.TASK_CLASS)
         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 = DataBase.get_class_for_name("SampleTaskRequest")
+        predict_class = get_inference_request(SampleModule.TASK_CLASS)
         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 = DataBase.get_class_for_name("SampleTaskRequest")
+    predict_class = get_inference_request(SampleModule.TASK_CLASS)
     response = sample_predict_servicer.Predict(
         predict_class(sample_input=HAPPY_PATH_INPUT_DM).to_proto(),
         Fixtures.build_context(sample_task_model_id),
@@ -111,7 +111,9 @@ def test_global_predict_works_on_bidirectional_streaming_rpcs(
 ):
     """Simple test that our SampleModule's bidirectional stream inference fn is supported"""
 
-    predict_class = DataBase.get_class_for_name("BidiStreamingSampleTaskRequest")
+    predict_class = get_inference_request(
+        SampleModule.TASK_CLASS, input_streaming=True, output_streaming=True
+    )
 
     def req_iterator() -> Iterator[predict_class]:
         for i in range(100):
@@ -139,7 +141,9 @@ def test_global_predict_works_on_bidirectional_streaming_rpcs_with_multiple_stre
     mock_manager = MagicMock()
     mock_manager.retrieve_model.return_value = GeoStreamingModule()
 
-    predict_class = DataBase.get_class_for_name("BidiStreamingGeoSpatialTaskRequest")
+    predict_class = get_inference_request(
+        GeoStreamingModule.TASK_CLASS, input_streaming=True, output_streaming=True
+    )
 
     def req_iterator() -> Iterator[predict_class]:
         for i in range(100):
@@ -194,7 +198,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 = DataBase.get_class_for_name("SampleTaskRequest")
+    predict_class = get_inference_request(SampleModule.TASK_CLASS)
     predict_thread = threading.Thread(
         target=sample_predict_servicer.Predict,
         args=(
@@ -229,7 +233,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 = DataBase.get_class_for_name("SampleTaskRequest")
+                predict_class = get_inference_request(SampleModule.TASK_CLASS)
                 request = predict_class(
                     sample_input=HAPPY_PATH_INPUT_DM, throw=True
                 ).to_proto()
@@ -253,7 +257,7 @@ def test_metering_predict_rpc_counter(
         sample_predict_servicer = GlobalPredictServicer(sample_inference_service)
         try:
             # Making 20 requests
-            predict_class = DataBase.get_class_for_name("SampleTaskRequest")
+            predict_class = get_inference_request(SampleModule.TASK_CLASS)
             for i in range(20):
                 sample_predict_servicer.Predict(
                     predict_class(sample_input=HAPPY_PATH_INPUT_DM).to_proto(),
@@ -294,7 +298,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 = DataBase.get_class_for_name("SampleTaskRequest")
+            predict_class = get_inference_request(SampleModule.TASK_CLASS)
             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 fbf5fa7c..f6826965 100644
--- a/tests/runtime/servicers/test_global_train_servicer_impl.py
+++ b/tests/runtime/servicers/test_global_train_servicer_impl.py
@@ -27,9 +27,13 @@ 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 (
@@ -38,7 +42,7 @@ from sample_lib.data_model.sample import (
     SampleOutputType,
     SampleTrainingType,
 )
-from sample_lib.modules.sample_task.sample_implementation import SampleModule
+from sample_lib.modules import CompositeModule, OtherModule, 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
@@ -99,10 +103,8 @@ def test_global_train_sample_task(
         jsondata=stream_type.JsonData(data=[SampleTrainingType(1)])
     )
     model_name = random_test_id()
-    train_class = DataBase.get_class_for_name("SampleTaskSampleModuleTrainRequest")
-    train_request_params_class = DataBase.get_class_for_name(
-        "SampleTaskSampleModuleTrainParameters"
-    )
+    train_class = get_train_request(SampleModule)
+    train_request_params_class = get_train_params(SampleModule)
     train_request = train_class(
         model_name=model_name,
         parameters=train_request_params_class(
@@ -138,7 +140,7 @@ def test_global_train_sample_task(
         == "sample_lib.modules.sample_task.sample_implementation.SampleModule"
     )
 
-    predict_class = DataBase.get_class_for_name("SampleTaskRequest")
+    predict_class = get_inference_request(SampleModule.TASK_CLASS)
     inference_response = sample_predict_servicer.Predict(
         predict_class(sample_input=HAPPY_PATH_INPUT_DM).to_proto(),
         Fixtures.build_context(training_response.model_name),
@@ -165,10 +167,8 @@ 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 = DataBase.get_class_for_name("OtherTaskOtherModuleTrainRequest")
-    train_request_params_class = DataBase.get_class_for_name(
-        "OtherTaskOtherModuleTrainParameters"
-    )
+    train_class = get_train_request(OtherModule)
+    train_request_params_class = get_train_params(OtherModule)
     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 = DataBase.get_class_for_name("OtherTaskRequest")
+    predict_class = get_inference_request(OtherModule.TASK_CLASS)
     inference_response = sample_predict_servicer.Predict(
         predict_class(sample_input=HAPPY_PATH_INPUT_DM).to_proto(),
         Fixtures.build_context(training_response.model_name),
@@ -225,10 +225,8 @@ def test_global_train_Another_Widget_that_requires_SampleWidget_loaded_should_no
         model_id=sample_task_model_id
     )
 
-    train_class = DataBase.get_class_for_name("SampleTaskCompositeModuleTrainRequest")
-    train_request_params_class = DataBase.get_class_for_name(
-        "SampleTaskCompositeModuleTrainParameters"
-    )
+    train_class = get_train_request(CompositeModule)
+    train_request_params_class = get_train_params(CompositeModule)
     training_request = train_class(
         model_name="AnotherWidget_Training",
         parameters=train_request_params_class(sample_block=sample_model),
@@ -255,7 +253,7 @@ def test_global_train_Another_Widget_that_requires_SampleWidget_loaded_should_no
     )
 
     # make sure the trained model can run inference
-    predict_class = DataBase.get_class_for_name("SampleTaskRequest")
+    predict_class = get_inference_request(SampleModule.TASK_CLASS)
     inference_response = sample_predict_servicer.Predict(
         predict_class(sample_input=HAPPY_PATH_INPUT_DM).to_proto(),
         Fixtures.build_context(training_response.model_name),
@@ -280,10 +278,8 @@ def test_run_train_job_works_with_wait(
     training_data = stream_type(
         jsondata=stream_type.JsonData(data=[SampleTrainingType(1)])
     )
-    train_class = DataBase.get_class_for_name("SampleTaskSampleModuleTrainRequest")
-    train_request_params_class = DataBase.get_class_for_name(
-        "SampleTaskSampleModuleTrainParameters"
-    )
+    train_class = get_train_request(SampleModule)
+    train_request_params_class = get_train_params(SampleModule)
     train_request = train_class(
         model_name=random_test_id(),
         parameters=train_request_params_class(
@@ -306,7 +302,7 @@ def test_run_train_job_works_with_wait(
             training_response.training_id,
         )
 
-        predict_class = DataBase.get_class_for_name("SampleTaskRequest")
+        predict_class = get_inference_request(SampleModule.TASK_CLASS)
         inference_response = sample_predict_servicer.Predict(
             predict_class(sample_input=SampleInputType(name="Test")).to_proto(),
             Fixtures.build_context(training_response.model_name),
@@ -332,10 +328,8 @@ 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 = DataBase.get_class_for_name("SampleTaskCompositeModuleTrainRequest")
-    train_request_params_class = DataBase.get_class_for_name(
-        "SampleTaskCompositeModuleTrainParameters"
-    )
+    train_class = get_train_request(CompositeModule)
+    train_request_params_class = get_train_params(CompositeModule)
     request = train_class(
         model_name="AnotherWidget_Training",
         parameters=train_request_params_class(sample_block=sample_model),
@@ -356,10 +350,8 @@ def test_global_train_Edge_Case_Widget_should_raise_when_error_surfaces_from_mod
         jsondata=stream_type.JsonData(data=[SampleTrainingType(1)])
     )
 
-    train_class = DataBase.get_class_for_name("SampleTaskSampleModuleTrainRequest")
-    train_request_params_class = DataBase.get_class_for_name(
-        "SampleTaskSampleModuleTrainParameters"
-    )
+    train_class = get_train_request(SampleModule)
+    train_request_params_class = get_train_params(SampleModule)
     train_request = train_class(
         model_name=random_test_id(),
         parameters=train_request_params_class(
@@ -387,10 +379,8 @@ def test_global_train_returns_exit_code_with_oom(
     training_data = stream_type(
         jsondata=stream_type.JsonData(data=[SampleTrainingType(1)])
     )
-    train_class = DataBase.get_class_for_name("SampleTaskSampleModuleTrainRequest")
-    train_request_params_class = DataBase.get_class_for_name(
-        "SampleTaskSampleModuleTrainParameters"
-    )
+    train_class = get_train_request(SampleModule)
+    train_request_params_class = get_train_params(SampleModule)
     train_request = train_class(
         model_name=random_test_id(),
         parameters=train_request_params_class(
@@ -421,10 +411,8 @@ def test_local_trainer_rejects_s3_output_paths(
     training_data = stream_type(
         jsondata=stream_type.JsonData(data=[SampleTrainingType(1)])
     )
-    train_class = DataBase.get_class_for_name("SampleTaskSampleModuleTrainRequest")
-    train_request_params_class = DataBase.get_class_for_name(
-        "SampleTaskSampleModuleTrainParameters"
-    )
+    train_class = get_train_request(SampleModule)
+    train_request_params_class = get_train_params(SampleModule)
     train_request = train_class(
         model_name=random_test_id(),
         output_path=S3Path(path="foo"),
@@ -457,10 +445,8 @@ def test_global_train_aborts_long_running_trains(
     )
     training_id = random_test_id()
 
-    train_class = DataBase.get_class_for_name("SampleTaskSampleModuleTrainRequest")
-    train_request_params_class = DataBase.get_class_for_name(
-        "SampleTaskSampleModuleTrainParameters"
-    )
+    train_class = get_train_request(SampleModule)
+    train_request_params_class = get_train_params(SampleModule)
     train_request = train_class(
         model_name=training_id,
         parameters=train_request_params_class(
diff --git a/tests/runtime/test_grpc_server.py b/tests/runtime/test_grpc_server.py
index 1c0a3de5..6910f3fb 100644
--- a/tests/runtime/test_grpc_server.py
+++ b/tests/runtime/test_grpc_server.py
@@ -42,7 +42,6 @@ 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,
@@ -50,6 +49,7 @@ 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,7 +60,14 @@ 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 InnerModule, SamplePrimitiveModule
+from sample_lib import (
+    CompositeModule,
+    InnerModule,
+    OtherModule,
+    SampleModule,
+    SamplePrimitiveModule,
+    StreamingModule,
+)
 from sample_lib.data_model import (
     OtherOutputType,
     SampleInputType,
@@ -76,7 +83,6 @@ from tests.runtime.conftest import (
     runtime_grpc_test_server,
 )
 import caikit.interfaces.common
-import sample_lib
 
 ## Helpers #####################################################################
 
@@ -131,11 +137,7 @@ def test_model_train(runtime_grpc_server):
                     "parameters": {
                         "training_data": {
                             "jsondata": {
-                                "data": [
-                                    sample_lib.data_model.SampleTrainingType(
-                                        number=1
-                                    ).to_dict()
-                                ]
+                                "data": [SampleTrainingType(number=1).to_dict()]
                             },
                         },
                     },
@@ -212,8 +214,9 @@ 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_class = DataBase.get_class_for_name("SampleTaskRequest")
-    predict_request = predict_class(sample_input=HAPPY_PATH_INPUT_DM).to_proto()
+    predict_request = get_inference_request(SampleModule.TASK_CLASS)(
+        sample_input=HAPPY_PATH_INPUT_DM
+    ).to_proto()
 
     actual_response = stub.SampleTaskPredict(
         predict_request, metadata=[("mm-model-id", sample_task_model_id)]
@@ -232,7 +235,7 @@ def test_global_predict_build_caikit_library_request_dict_creates_caikit_core_ru
     )
     proto_request_dict = build_caikit_library_request_dict(
         proto_request,
-        sample_lib.modules.sample_task.SampleModule.RUN_SIGNATURE,
+        SampleModule.RUN_SIGNATURE,
     )
 
     # unset fields not included
@@ -241,12 +244,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 = DataBase.get_class_for_name("SampleTaskRequest")
+    predict_class = get_inference_request(SampleModule.TASK_CLASS)
     python_request = predict_class(sample_input=HAPPY_PATH_INPUT_DM).to_proto()
 
     python_sample_module_request_dict = build_caikit_library_request_dict(
         python_request,
-        sample_lib.modules.sample_task.SampleModule.RUN_SIGNATURE,
+        SampleModule.RUN_SIGNATURE,
     )
 
     # unset fields are included if they have defaults set
@@ -260,7 +263,9 @@ 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 = DataBase.get_class_for_name("ServerStreamingStreamingTaskRequest")
+    predict_class = get_inference_request(
+        StreamingModule.TASK_CLASS, input_streaming=False, output_streaming=True
+    )
     predict_request = predict_class(sample_input=HAPPY_PATH_INPUT_DM).to_proto()
 
     stream = stub.ServerStreamingStreamingTaskPredict(
@@ -282,7 +287,7 @@ def test_predict_sample_module_error_response(
         stub = sample_inference_service.stub_class(
             runtime_grpc_server.make_local_channel()
         )
-        predict_class = DataBase.get_class_for_name("SampleTaskRequest")
+        predict_class = get_inference_request(SampleModule.TASK_CLASS)
         predict_request = predict_class(sample_input=HAPPY_PATH_INPUT_DM).to_proto()
 
         stub.SampleTaskPredict(
@@ -297,7 +302,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 = DataBase.get_class_for_name("OtherTaskRequest")
+    predict_class = get_inference_request(OtherModule.TASK_CLASS)
     predict_request = predict_class(
         sample_input_sampleinputtype=HAPPY_PATH_INPUT_DM
     ).to_proto()
@@ -328,7 +333,7 @@ def test_rpc_validation_on_predict_for_unsupported_model(
         stub = sample_inference_service.stub_class(
             runtime_grpc_server.make_local_channel()
         )
-        predict_class = DataBase.get_class_for_name("SampleTaskRequest")
+        predict_class = get_inference_request(SampleModule.TASK_CLASS)
         predict_request = predict_class(sample_input=HAPPY_PATH_INPUT_DM).to_proto()
         with pytest.raises(grpc.RpcError) as context:
             stub.SampleTaskPredict(
@@ -360,7 +365,10 @@ def test_rpc_validation_on_predict_for_wrong_streaming_flavor(
         stub = sample_inference_service.stub_class(
             runtime_grpc_server.make_local_channel()
         )
-        predict_class = DataBase.get_class_for_name("SampleTaskRequest")
+
+        predict_class = get_inference_request(
+            SampleModule.TASK_CLASS,
+        )
         predict_request = predict_class(sample_input=HAPPY_PATH_INPUT_DM).to_proto()
         with pytest.raises(grpc.RpcError) as context:
             response = stub.ServerStreamingSampleTaskPredict(
@@ -397,15 +405,9 @@ def test_train_fake_module_ok_response_and_can_predict_with_trained_model(
         )
     )
     model_name = random_test_id()
-    train_request_class = DataBase.get_class_for_name(
-        "SampleTaskSampleModuleTrainRequest"
-    )
-    train_request_params_class = DataBase.get_class_for_name(
-        "SampleTaskSampleModuleTrainParameters"
-    )
-    train_request = train_request_class(
+    train_request = get_train_request(SampleModule)(
         model_name=model_name,
-        parameters=train_request_params_class(
+        parameters=get_train_params(SampleModule)(
             training_data=training_data,
             union_list=["str", "sequence"],
         ),
@@ -421,7 +423,7 @@ def test_train_fake_module_ok_response_and_can_predict_with_trained_model(
     )
 
     # make sure the trained model can run inference
-    predict_class = DataBase.get_class_for_name("SampleTaskRequest")
+    predict_class = get_inference_request(SampleModule.TASK_CLASS)
     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)]
@@ -443,13 +445,9 @@ 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_class = DataBase.get_class_for_name("SampleTaskCompositeModuleTrainRequest")
-    train_request_params_class = DataBase.get_class_for_name(
-        "SampleTaskCompositeModuleTrainParameters"
-    )
-    train_request = train_class(
+    train_request = get_train_request(CompositeModule)(
         model_name=model_name,
-        parameters=train_request_params_class(sample_block=sample_model),
+        parameters=get_train_params(CompositeModule)(sample_block=sample_model),
     ).to_proto()
 
     actual_response = train_stub.SampleTaskCompositeModuleTrain(train_request)
@@ -461,7 +459,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 = DataBase.get_class_for_name("SampleTaskRequest")
+    predict_class = get_inference_request(CompositeModule.TASK_CLASS)
     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)]
@@ -488,13 +486,9 @@ 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_class = DataBase.get_class_for_name("OtherTaskOtherModuleTrainRequest")
-    train_request_params_class = DataBase.get_class_for_name(
-        "OtherTaskOtherModuleTrainParameters"
-    )
-    train_request = train_class(
+    train_request = get_train_request(OtherModule)(
         model_name="Bar Training",
-        parameters=train_request_params_class(
+        parameters=get_train_params(OtherModule)(
             sample_input_sampleinputtype=HAPPY_PATH_INPUT_DM,
             batch_size=100,
             training_data=training_data,
@@ -512,7 +506,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 = DataBase.get_class_for_name("OtherTaskRequest")
+    predict_class = get_inference_request(OtherModule.TASK_CLASS)
     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)]
@@ -543,16 +537,10 @@ 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 = train_request_class(
+    train_request = get_train_request(SamplePrimitiveModule)(
         model_name=model_name,
-        parameters=train_request_params_class(
+        parameters=get_train_params(SamplePrimitiveModule)(
             sample_input=HAPPY_PATH_INPUT_DM,
             simple_list=["hello", "world"],
             union_list=["str", "sequence"],
@@ -578,7 +566,7 @@ def test_train_primitive_model(
     )
 
     # make sure the trained model can run inference
-    predict_class = DataBase.get_class_for_name("SampleTaskRequest")
+    predict_class = get_inference_request(SampleModule.TASK_CLASS)
     predict_request = predict_class(sample_input=HAPPY_PATH_INPUT_DM).to_proto()
 
     inference_response = inference_stub.SampleTaskPredict(
@@ -607,13 +595,10 @@ def test_train_fake_module_ok_response_with_datastream_jsondata(
         )
     )
     model_name = random_test_id()
-    train_class = DataBase.get_class_for_name("SampleTaskSampleModuleTrainRequest")
-    train_request_params_class = DataBase.get_class_for_name(
-        "SampleTaskSampleModuleTrainParameters"
-    )
-    train_request = train_class(
+
+    train_request = get_train_request(SampleModule)(
         model_name=model_name,
-        parameters=train_request_params_class(
+        parameters=get_train_params(SampleModule)(
             batch_size=42,
             training_data=training_data,
         ),
@@ -628,7 +613,7 @@ def test_train_fake_module_ok_response_with_datastream_jsondata(
     )
 
     # make sure the trained model can run inference
-    predict_class = DataBase.get_class_for_name("SampleTaskRequest")
+    predict_class = get_inference_request(SampleModule.TASK_CLASS)
     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)]
@@ -649,13 +634,10 @@ 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_class = DataBase.get_class_for_name("SampleTaskSampleModuleTrainRequest")
-    train_request_params_class = DataBase.get_class_for_name(
-        "SampleTaskSampleModuleTrainParameters"
-    )
-    train_request = train_class(
+
+    train_request = get_train_request(SampleModule)(
         model_name=model_name,
-        parameters=train_request_params_class(
+        parameters=get_train_params(SampleModule)(
             training_data=training_data,
         ),
     ).to_proto()
@@ -669,7 +651,7 @@ def test_train_fake_module_ok_response_with_datastream_csv_file(
     )
 
     # make sure the trained model can run inference
-    predict_class = DataBase.get_class_for_name("SampleTaskRequest")
+    predict_class = get_inference_request(SampleModule.TASK_CLASS)
     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)]
@@ -689,13 +671,10 @@ 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_class = DataBase.get_class_for_name("SampleTaskSampleModuleTrainRequest")
-    train_request_params_class = DataBase.get_class_for_name(
-        "SampleTaskSampleModuleTrainParameters"
-    )
-    train_request = train_class(
+
+    train_request = get_train_request(SampleModule)(
         model_name=model_name,
-        parameters=train_request_params_class(
+        parameters=get_train_params(SampleModule)(
             training_data=training_data, sleep_time=10
         ),
     ).to_proto()
@@ -730,13 +709,9 @@ 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_class = DataBase.get_class_for_name("SampleTaskSampleModuleTrainRequest")
-    train_request_params_class = DataBase.get_class_for_name(
-        "SampleTaskSampleModuleTrainParameters"
-    )
-    train_request = train_class(
+    train_request = get_train_request(SampleModule)(
         model_name=model_name,
-        parameters=train_request_params_class(
+        parameters=get_train_params(SampleModule)(
             training_data=training_data, sleep_time=10
         ),
     ).to_proto()
@@ -757,9 +732,9 @@ def test_cancel_does_not_affect_other_models(
 
     # train another model
     model_name2 = random_test_id()
-    train_request2 = train_class(
+    train_request2 = get_train_request(SampleModule)(
         model_name=model_name2,
-        parameters=train_request_params_class(training_data=training_data),
+        parameters=get_train_params(SampleModule)(training_data=training_data),
     ).to_proto()
     train_response2 = train_stub.SampleTaskSampleModuleTrain(train_request2)
 
@@ -790,15 +765,9 @@ def test_train_fake_module_error_response_with_unloaded_model(
         sample_model = caikit.interfaces.runtime.data_model.ModelPointer(
             model_id=random_test_id()
         )
-        train_class = DataBase.get_class_for_name(
-            "SampleTaskCompositeModuleTrainRequest"
-        )
-        train_request_params_class = DataBase.get_class_for_name(
-            "SampleTaskCompositeModuleTrainParameters"
-        )
-        train_request = train_class(
+        train_request = get_train_request(CompositeModule)(
             model_name=random_test_id(),
-            parameters=train_request_params_class(sample_block=sample_model),
+            parameters=get_train_params(CompositeModule)(sample_block=sample_model),
         ).to_proto()
 
         train_stub.SampleTaskCompositeModuleTrain(train_request)
@@ -1147,7 +1116,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 = DataBase.get_class_for_name("SampleTaskRequest")
+            predict_class = get_inference_request(SampleModule.TASK_CLASS)
             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 36da1428..624ab202 100644
--- a/tests/runtime/test_service_factory.py
+++ b/tests/runtime/test_service_factory.py
@@ -23,11 +23,16 @@ import pytest
 
 # Local
 from caikit.core.data_model import render_dataobject_protos
-from caikit.core.data_model.base import DataBase
-from caikit.runtime.service_factory import ServicePackage, ServicePackageFactory
+from caikit.runtime.service_factory import (
+    ServicePackage,
+    ServicePackageFactory,
+    get_inference_request,
+    get_train_params,
+    get_train_request,
+)
 from sample_lib import SampleModule
 from sample_lib.data_model import SampleInputType, SampleOutputType
-from sample_lib.modules.sample_task import ListModule
+from sample_lib.modules import ListModule, OtherModule
 from tests.conftest import temp_config
 from tests.core.helpers import MockBackend
 from tests.data_model_helpers import reset_global_protobuf_registry, temp_dpool
@@ -360,8 +365,71 @@ def test_backend_modules_included_in_service_generation(
     inference_service = ServicePackageFactory.get_service_package(
         ServicePackageFactory.ServiceType.INFERENCE
     )
-    predict_class = DataBase.get_class_for_name("SampleTaskRequest")
+    predict_class = get_inference_request(SampleModule.TASK_CLASS)
     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"
+    )