Repository: caikit/caikit
Pull request #434 (2023-08-24 00:17:12)

♻️  split training parameters
<!--  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**:

POC for #369 

**Special notes for your reviewer**:

This updates the training RPC to manage an inner and outer request message, and it seems to work pretty slick.
cc @gabe-l-hart, @gkumbhat and @prashantgupta24 I think this'll work if we wanted to rip the bandaid on doing this nesting so we can have the pydantic apis convert these data models directly on the http side as discussed.

Open problem now is getting the caikit user access to the inner data model type so they can construct one of these messages directly. But running via compiled protobuf or REST that won't be an issue.

> Update - the `client_module` has been updated to incorporate ALL DM objects, so that unit tests as well as users can fetch the inner Proto message from it.

## Additional context

Pending pydantic investigation: https://github.com/pydantic/pydantic/discussions/7239 - thinking what to do about `forward-compatibility` - right now adding `extra_forbid` on all `pydantic models` has the unindented consequence that we won't be able to pass in anything that's not in the pydantic model, which means if the clients update before the server does, they won't work.

One quick workaround was to only attach `extra_forbid` to DMs that are a `Union`, but then if we first come across a data model used in a non-union context, and then later it is used within a union, then we would only create the pydantic data model without the `extra_forbid` bit, which is incorrect.

Tracking forward-compatibility issue here: https://github.com/caikit/caikit/issues/450

**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/http_server.py b/caikit/runtime/http_server.py
index a848786d..eec8bc99 100644
--- a/caikit/runtime/http_server.py
+++ b/caikit/runtime/http_server.py
@@ -77,7 +77,6 @@ from caikit.runtime.service_generation.rpcs import (
 from caikit.runtime.servicers.global_predict_servicer import GlobalPredictServicer
 from caikit.runtime.servicers.global_train_servicer import GlobalTrainServicer
 from caikit.runtime.types.caikit_runtime_exception import CaikitRuntimeException
-import caikit
 
 ## Globals #####################################################################
 
@@ -129,6 +128,15 @@ HEALTH_ENDPOINT = "/health"
 ## RuntimeHTTPServer ###########################################################
 
 
+# Base class for pydantic models
+# We want to set the config to forbid extra attributes
+# while instantiating any pydantic models
+# This is done to make sure any oneofs can be
+# correctly infered by pydantic
+class ParentPydanticBaseModel(pydantic.BaseModel):
+    model_config = pydantic.ConfigDict(extra="forbid")
+
+
 class RuntimeHTTPServer(RuntimeServerBase):
     """An implementation of a FastAPI server that serves caikit runtimes"""
 
@@ -311,35 +319,13 @@ class RuntimeHTTPServer(RuntimeServerBase):
         # convert pydantic objects to our DM objects
         for param_name, param_value in request_params.items():
             if issubclass(type(param_value), pydantic.BaseModel):
-                request_params[param_name] = self.build_dm_object(param_value)
+                request_params[param_name] = self._build_dm_object(param_value)
         return request_params
 
-    def build_dm_object(self, pydantic_model: pydantic.BaseModel) -> DataBase:
-        """Convert pydantic objects to our DM objects"""
-        dm_class_to_build = PYDANTIC_TO_DM_MAPPING.get(type(pydantic_model))
-        dm_kwargs = {}
-
-        for field_name, field_value in pydantic_model:
-            # field could be a DM:
-            # pylint: disable=unidiomatic-typecheck
-            if type(field_value) in PYDANTIC_TO_DM_MAPPING:
-                dm_kwargs[field_name] = self.build_dm_object(field_value)
-            elif isinstance(field_value, list):
-                if all(type(val) in PYDANTIC_TO_DM_MAPPING for val in field_value):
-                    dm_kwargs[field_name] = [
-                        self.build_dm_object(val) for val in field_value
-                    ]
-                else:
-                    dm_kwargs[field_name] = field_value
-            else:
-                dm_kwargs[field_name] = field_value
-
-        return dm_class_to_build(**dm_kwargs)
-
     def _train_add_unary_input_unary_output_handler(self, rpc: CaikitRPCBase):
         """Add a unary:unary request handler for this RPC signature"""
         pydantic_request = self._dataobject_to_pydantic(
-            self._get_request_dataobject(rpc)
+            DataBase.get_class_for_name(rpc.request.name)
         )
         pydantic_response = self._dataobject_to_pydantic(
             self._get_response_dataobject(rpc)
@@ -354,26 +340,12 @@ class RuntimeHTTPServer(RuntimeServerBase):
             loop = asyncio.get_running_loop()
 
             # build request DM object
-            http_request_dm_object = self.build_dm_object(request)
+            http_request_dm_object = self._build_dm_object(request)
 
-            request_dm_class = caikit.core.data_model.DataBase.get_class_for_name(
-                rpc.request.name
-            )
-            request_dm_class_kwargs = {}
-            required_inputs = getattr(http_request_dm_object, REQUIRED_INPUTS_KEY, None)
-            optional_inputs = getattr(http_request_dm_object, OPTIONAL_INPUTS_KEY, None)
-            for input_type in (required_inputs, optional_inputs):
-                if input_type is not None:
-                    for field_name in input_type.fields:
-                        field_value = getattr(input_type, field_name, None)
-                        if field_value is not None:
-                            request_dm_class_kwargs[field_name] = field_value
-
-            request_dm_object = request_dm_class(**request_dm_class_kwargs)
             try:
                 call = partial(
                     self.global_train_servicer.run_training_job,
-                    request=request_dm_object.to_proto(),
+                    request=http_request_dm_object.to_proto(),
                     module=rpc.clz,
                     training_output_dir=None,  # pass None so that GTS picks up the config one # TODO: double-check?
                     # context=context,
@@ -604,6 +576,29 @@ class RuntimeHTTPServer(RuntimeServerBase):
         assert isinstance(dm_obj, type) and issubclass(dm_obj, DataBase)
         return dm_obj
 
+    @classmethod
+    def _build_dm_object(cls, pydantic_model: pydantic.BaseModel) -> DataBase:
+        """Convert pydantic objects to our DM objects"""
+        dm_class_to_build = PYDANTIC_TO_DM_MAPPING.get(type(pydantic_model))
+        dm_kwargs = {}
+
+        for field_name, field_value in pydantic_model:
+            # field could be a DM:
+            # pylint: disable=unidiomatic-typecheck
+            if type(field_value) in PYDANTIC_TO_DM_MAPPING:
+                dm_kwargs[field_name] = cls._build_dm_object(field_value)
+            elif isinstance(field_value, list):
+                if all(type(val) in PYDANTIC_TO_DM_MAPPING for val in field_value):
+                    dm_kwargs[field_name] = [
+                        cls._build_dm_object(val) for val in field_value
+                    ]
+                else:
+                    dm_kwargs[field_name] = field_value
+            else:
+                dm_kwargs[field_name] = field_value
+
+        return dm_class_to_build(**dm_kwargs)
+
     @classmethod
     # pylint: disable=too-many-return-statements
     def _get_pydantic_type(cls, field_type: type) -> type:
@@ -615,7 +610,7 @@ class RuntimeHTTPServer(RuntimeServerBase):
             return int
         if np.issubclass_(field_type, np.floating):
             return float
-        if field_type in (int, float, bool, str, bytes, type(None)):
+        if field_type in (int, float, bool, str, bytes, dict, type(None)):
             return field_type
         if isinstance(field_type, type) and issubclass(field_type, enum.Enum):
             return field_type
@@ -645,7 +640,10 @@ class RuntimeHTTPServer(RuntimeServerBase):
             return List[cls._get_pydantic_type(get_args(field_type)[0])]
 
         if get_origin(field_type) is dict:
-            return field_type
+            return Dict[
+                cls._get_pydantic_type(get_args(field_type)[0]),
+                cls._get_pydantic_type(get_args(field_type)[1]),
+            ]
 
         raise TypeError(f"Cannot get pydantic type for type [{field_type}]")
 
@@ -670,9 +668,9 @@ class RuntimeHTTPServer(RuntimeServerBase):
                 dm_class, localns=localns
             ).items()
         }
-        pydantic_model = type(pydantic.BaseModel)(
+        pydantic_model = type(ParentPydanticBaseModel)(
             dm_class.get_proto_class().DESCRIPTOR.full_name,
-            (pydantic.BaseModel,),
+            (ParentPydanticBaseModel,),
             {
                 "__annotations__": annotations,
                 **{
diff --git a/caikit/runtime/service_factory.py b/caikit/runtime/service_factory.py
index 8e0307de..1970a82a 100644
--- a/caikit/runtime/service_factory.py
+++ b/caikit/runtime/service_factory.py
@@ -31,6 +31,7 @@ import alog
 # Local
 from caikit import get_config
 from caikit.core import LocalBackend, ModuleBase, registries
+from caikit.core.data_model.dataobject import _AUTO_GEN_PROTO_CLASSES
 from caikit.interfaces.runtime.data_model import (
     TrainingInfoRequest,
     TrainingStatusResponse,
@@ -141,18 +142,17 @@ class ServicePackageFactory:
 
         rpc_list = [rpc for rpc in rpc_list if rpc.return_type is not None]
 
-        request_data_models = [
-            rpc.create_request_data_model(package_name) for rpc in rpc_list
-        ]
+        for rpc in rpc_list:
+            rpc.create_request_data_model(package_name)
 
         client_module = ModuleType(
             "ClientMessages",
             "Package with service message class implementations",
         )
 
-        for dm_class in request_data_models:
-            # We need the message class that data model serializes to
-            setattr(client_module, dm_class.__name__, type(dm_class().to_proto()))
+        for proto_class in _AUTO_GEN_PROTO_CLASSES:
+            # We need all the DM objects in the client_module for ease of use
+            setattr(client_module, proto_class.DESCRIPTOR.name, proto_class)
 
         rpc_jsons = [rpc.create_rpc_json(package_name) for rpc in rpc_list]
         service_json = {"service": {"rpcs": rpc_jsons}}
diff --git a/caikit/runtime/service_generation/rpcs.py b/caikit/runtime/service_generation/rpcs.py
index f1979a57..438ba3ea 100644
--- a/caikit/runtime/service_generation/rpcs.py
+++ b/caikit/runtime/service_generation/rpcs.py
@@ -63,33 +63,7 @@ class CaikitRPCBase(abc.ABC):
 
     def create_request_data_model(self, package_name: str) -> Type[DataBase]:
         """Dynamically create data model for this RPC's request message"""
-        properties = {
-            # triple e.g. ('caikit.interfaces.common.ProducerPriority', 'producer_id', 1)
-            # This does not take care of nested descriptors
-            triple[1]: Annotated[triple[0], FieldNumber(triple[2])]
-            for triple in self.request.triples
-            if triple[1] not in self.request.default_map
-        }
-        optional_properties = {
-            triple[1]: Annotated[Optional[triple[0]], FieldNumber(triple[2])]
-            for triple in self.request.triples
-            if triple[1] in self.request.default_map
-        }
-        attrs = copy.copy(self.request.default_map)
-
-        if not properties and not optional_properties:
-            log.warning(
-                "No arguments found for request %s. Cannot generate rpc",
-                self.request.name,
-            )
-            return None
-
-        return make_dataobject(
-            package=package_name,
-            name=self.request.name,
-            attrs=attrs,
-            annotations={**properties, **optional_properties},
-        )
+        return self.request.create_data_model(package_name)
 
     def create_rpc_json(self, package_name: str) -> Dict:
         """Return json snippet for the service definition of this RPC"""
@@ -128,12 +102,42 @@ class ModuleClassTrainRPC(CaikitRPCBase):
 
         # Store the input and output protobuf message types for this RPC
         self.return_type = self._method.return_type
-        self._req = _RequestMessage(
-            ModuleClassTrainRPC.module_class_to_req_name(self.clz),
+
+        # Inner / outer requests
+        # The actual training parameters from the method signature are nested under `parameters`
+        self._inner_request = _RequestMessage(
+            ModuleClassTrainRPC.module_class_to_inner_request_name(self.clz),
             self._method.parameters,
             self._method.default_parameters,
         )
 
+        params = {"model_name": str, "output_path": S3Path, "parameters": "PLACEHOLDER"}
+
+        self._req = _RequestMessage(
+            ModuleClassTrainRPC.module_class_to_req_name(self.clz),
+            params,
+            {},
+        )
+
+    def create_request_data_model(self, package_name: str):
+        """Partial override of CaikitRPCBase.create_request_data_model to take care of the inner
+        request data model"""
+        # Build the inner request data model
+        inner_request_data_model = self._inner_request.create_data_model(package_name)
+        # Insert the new type into the outer request
+        for triple_index, _ in enumerate(self._req.triples):
+            if self._req.triples[triple_index][1] == "parameters":
+                triple = self._req.triples[triple_index]
+                self._req.triples[triple_index] = (
+                    inner_request_data_model,
+                    triple[1],
+                    triple[2],
+                )
+        # Call the base class' method
+        return CaikitRPCBase.create_request_data_model(
+            self=self, package_name=package_name
+        )
+
     @property
     def module_list(self) -> List[Type[ModuleBase]]:
         """Returns a list containing the single caikit.core.module type that this RPC is for"""
@@ -164,6 +168,17 @@ class ModuleClassTrainRPC(CaikitRPCBase):
         """
         return f"{ModuleClassTrainRPC.module_class_to_rpc_name(module_class)}Request"
 
+    @staticmethod
+    def module_class_to_inner_request_name(module_class: Type[ModuleBase]) -> str:
+        """Helper function to convert from a module to the name of the inner message
+        containing all the training parameters
+
+        Example: self.clz._module__ = sample_lib.modules.sample_task.sample_implementation
+
+        return: SampleTaskSampleModuleTrainParameters
+        """
+        return f"{ModuleClassTrainRPC.module_class_to_rpc_name(module_class)}Parameters"
+
     @staticmethod
     def _mutate_method_signature_for_training(
         signature: CaikitMethodSignature,
@@ -171,10 +186,8 @@ class ModuleClassTrainRPC(CaikitRPCBase):
         # Change return type for async training interface
         return_type = TrainingJob
 
-        # Start with extra metaparameters
-        # - model_name: user-provided custom ID for the model to train
-        # - output_path: pointer to some storage where the model will be saved
-        new_params = {"model_name": str, "output_path": S3Path}
+        # Loop through params and edit as needed
+        new_params = {}
         for name, typ in signature.parameters.items():
             if type_helpers.has_data_stream(typ):
                 # Assume this is training data
@@ -387,3 +400,33 @@ class _RequestMessage:
             self.triples.append((typ, item_name, num))
 
         self.triples.sort(key=lambda x: x[2])
+
+    def create_data_model(self, package_name: str) -> Type[DataBase]:
+        """Dynamically create a data model for this request message"""
+        properties = {
+            # triple e.g. ('caikit.interfaces.common.ProducerPriority', 'producer_id', 1)
+            # This does not take care of nested descriptors
+            triple[1]: Annotated[triple[0], FieldNumber(triple[2])]
+            for triple in self.triples
+            if triple[1] not in self.default_map
+        }
+        optional_properties = {
+            triple[1]: Annotated[Optional[triple[0]], FieldNumber(triple[2])]
+            for triple in self.triples
+            if triple[1] in self.default_map
+        }
+        attrs = copy.copy(self.default_map)
+
+        if not properties and not optional_properties:
+            log.warning(
+                "No arguments found for request %s. Cannot generate rpc",
+                self.name,
+            )
+            return None
+
+        return make_dataobject(
+            package=package_name,
+            name=self.name,
+            attrs=attrs,
+            annotations={**properties, **optional_properties},
+        )
diff --git a/caikit/runtime/servicers/global_train_servicer.py b/caikit/runtime/servicers/global_train_servicer.py
index 9c87a74a..b87b9d11 100644
--- a/caikit/runtime/servicers/global_train_servicer.py
+++ b/caikit/runtime/servicers/global_train_servicer.py
@@ -213,7 +213,9 @@ class GlobalTrainServicer:
                 "save_path": model_path,
                 "save_with_id": self.save_with_id,
                 "model_name": request_data_model.model_name,
-                **build_caikit_library_request_dict(request, module.TRAIN_SIGNATURE),
+                **build_caikit_library_request_dict(
+                    request.parameters, module.TRAIN_SIGNATURE
+                ),
             }
         )
 
diff --git a/caikit/runtime/servicers/model_train_servicer.py b/caikit/runtime/servicers/model_train_servicer.py
index 79e6ec85..bf001122 100644
--- a/caikit/runtime/servicers/model_train_servicer.py
+++ b/caikit/runtime/servicers/model_train_servicer.py
@@ -95,7 +95,7 @@ class ModelTrainServicerImpl(process_pb2_grpc.ProcessServicer):
             )
             # update any file references relative to the training input dir
             self._update_file_references(
-                train_message_request, request.training_input_dir
+                train_message_request.parameters, request.training_input_dir
             )
 
             # make the train call


Test patch:
diff --git a/tests/runtime/service_generation/test_rpcs.py b/tests/runtime/service_generation/test_rpcs.py
index ae6ae90a..fc0ce676 100644
--- a/tests/runtime/service_generation/test_rpcs.py
+++ b/tests/runtime/service_generation/test_rpcs.py
@@ -257,3 +257,15 @@ def test_module_train_rpc():
     assert data_model is not None
 
     assert rpc.name == "TestTaskTestModuleTrain"
+
+    # Training RPCs nest the actual training params from the `.train` signature
+    assert hasattr(data_model, "model_name")
+    assert hasattr(data_model, "output_path")
+    assert hasattr(data_model, "parameters")
+
+    training_message = data_model.from_json(
+        {"model_name": "any_model_name", "parameters": {"int_val": 1, "str_val": "foo"}}
+    )
+
+    assert training_message.parameters.int_val == 1
+    assert training_message.parameters.str_val == "foo"
diff --git a/tests/runtime/servicers/test_global_train_servicer_impl.py b/tests/runtime/servicers/test_global_train_servicer_impl.py
index 4dc15f1e..c42f14c2 100644
--- a/tests/runtime/servicers/test_global_train_servicer_impl.py
+++ b/tests/runtime/servicers/test_global_train_servicer_impl.py
@@ -98,8 +98,10 @@ def test_global_train_sample_task(
     model_name = random_test_id()
     train_request = sample_train_service.messages.SampleTaskSampleModuleTrainRequest(
         model_name=model_name,
-        batch_size=42,
-        training_data=training_data,
+        parameters=sample_train_service.messages.SampleTaskSampleModuleTrainParameters(
+            batch_size=42,
+            training_data=training_data,
+        ),
     )
 
     training_response = sample_train_servicer.Train(
@@ -159,9 +161,11 @@ def test_global_train_other_task(
     training_data = stream_type(jsondata=stream_type.JsonData(data=[1])).to_proto()
     train_request = sample_train_service.messages.OtherTaskOtherModuleTrainRequest(
         model_name="Other module Training",
-        training_data=training_data,
-        sample_input_sampleinputtype=SampleInputType(name="Gabe").to_proto(),
-        batch_size=batch_size,
+        parameters=sample_train_service.messages.OtherTaskOtherModuleTrainParameters(
+            training_data=training_data,
+            sample_input_sampleinputtype=SampleInputType(name="Gabe").to_proto(),
+            batch_size=batch_size,
+        ),
     )
 
     training_response = sample_train_servicer.Train(
@@ -212,11 +216,11 @@ def test_global_train_Another_Widget_that_requires_SampleWidget_loaded_should_no
         model_id=sample_task_model_id
     ).to_proto()
 
-    training_request = (
-        sample_train_service.messages.SampleTaskCompositeModuleTrainRequest(
-            model_name="AnotherWidget_Training",
+    training_request = sample_train_service.messages.SampleTaskCompositeModuleTrainRequest(
+        model_name="AnotherWidget_Training",
+        parameters=sample_train_service.messages.SampleTaskCompositeModuleTrainParameters(
             sample_block=sample_model,
-        )
+        ),
     )
 
     training_response = sample_train_servicer.Train(
@@ -268,8 +272,10 @@ def test_run_train_job_works_with_wait(
     ).to_proto()
     train_request = sample_train_service.messages.SampleTaskSampleModuleTrainRequest(
         model_name=random_test_id(),
-        batch_size=42,
-        training_data=training_data,
+        parameters=sample_train_service.messages.SampleTaskSampleModuleTrainParameters(
+            batch_size=42,
+            training_data=training_data,
+        ),
     )
     servicer = GlobalTrainServicer(training_service=sample_train_service)
     with TemporaryDirectory() as tmp_dir:
@@ -317,7 +323,9 @@ def test_global_train_Another_Widget_that_requires_SampleWidget_but_not_loaded_s
     ).to_proto()
     request = sample_train_service.messages.SampleTaskCompositeModuleTrainRequest(
         model_name="AnotherWidget_Training",
-        sample_block=sample_model,
+        parameters=sample_train_service.messages.SampleTaskCompositeModuleTrainParameters(
+            sample_block=sample_model,
+        ),
     )
 
     with pytest.raises(CaikitRuntimeException) as context:
@@ -337,8 +345,10 @@ def test_global_train_Edge_Case_Widget_should_raise_when_error_surfaces_from_mod
 
     train_request = sample_train_service.messages.SampleTaskSampleModuleTrainRequest(
         model_name=random_test_id(),
-        batch_size=999,
-        training_data=training_data,
+        parameters=sample_train_service.messages.SampleTaskSampleModuleTrainParameters(
+            batch_size=999,
+            training_data=training_data,
+        ),
     )
 
     training_response = sample_train_servicer.Train(
@@ -362,9 +372,11 @@ def test_global_train_returns_exit_code_with_oom(
     ).to_proto()
     train_request = sample_train_service.messages.SampleTaskSampleModuleTrainRequest(
         model_name=random_test_id(),
-        batch_size=42,
-        training_data=training_data,
-        oom_exit=True,
+        parameters=sample_train_service.messages.SampleTaskSampleModuleTrainParameters(
+            batch_size=42,
+            training_data=training_data,
+            oom_exit=True,
+        ),
     )
 
     # Enable sub-processing for test
@@ -391,9 +403,11 @@ def test_local_trainer_rejects_s3_output_paths(
     train_request = sample_train_service.messages.SampleTaskSampleModuleTrainRequest(
         model_name=random_test_id(),
         output_path=S3Path(path="foo").to_proto(),
-        batch_size=42,
-        training_data=training_data,
-        oom_exit=True,
+        parameters=sample_train_service.messages.SampleTaskSampleModuleTrainParameters(
+            batch_size=42,
+            training_data=training_data,
+            oom_exit=True,
+        ),
     )
 
     with pytest.raises(
@@ -420,9 +434,11 @@ def test_global_train_aborts_long_running_trains(
 
     train_request = sample_train_service.messages.SampleTaskSampleModuleTrainRequest(
         model_name=training_id,
-        batch_size=42,
-        training_data=training_data,
-        oom_exit=False,
+        parameters=sample_train_service.messages.SampleTaskSampleModuleTrainParameters(
+            batch_size=42,
+            training_data=training_data,
+            oom_exit=False,
+        ),
     )
 
     if sample_train_servicer.use_subprocess:
@@ -468,7 +484,6 @@ def test_global_train_aborts_long_running_trains(
         f"{SampleModule.__module__}.{SampleModule.train.__qualname__}",
         never_respond,
     ):
-
         train_thread.start()
         # NB: assert is here to make sure we called the patched train
         assert test_event.wait(test_event_timeout)
diff --git a/tests/runtime/servicers/test_model_train_servicer_impl.py b/tests/runtime/servicers/test_model_train_servicer_impl.py
index 97c9581a..e4625692 100644
--- a/tests/runtime/servicers/test_model_train_servicer_impl.py
+++ b/tests/runtime/servicers/test_model_train_servicer_impl.py
@@ -91,7 +91,7 @@ def test_model_train_no_train_module_raises(sample_model_train_servicer, output_
         trainingID=training_id,
         customTrainingID=str(uuid.uuid4()),
         request_dict={
-            "training_params": '{"model_name": "abc", "training_data": [1]}',
+            "training_params": '{"model_name": "abc", "parameters": {"training_data": [1]}}',
         },
         training_input_dir="training_input_dir",
         training_output_dir=os.path.join(output_dir, training_id),
@@ -110,7 +110,7 @@ def test_model_train_incorrect_train_params_raises(
         customTrainingID=str(uuid.uuid4()),
         request_dict={
             "train_module": "00110203-0405-0607-0809-0a0b02dd0e0f",
-            "training_params": '{"model_name": "abc", "training_data": "blah"}',
+            "training_params": '{"model_name": "abc", "parameters": {"training_data": "blah"}}',
         },
         training_input_dir="training_input_dir",
         training_output_dir=os.path.join(output_dir, training_id),
@@ -150,7 +150,7 @@ def test_model_train_incorrect_train_request_raises(
         customTrainingID=str(uuid.uuid4()),
         request_dict={
             "train_module": "00110203-0405-0607-0809-0a0b02dd0e0f",
-            "training_params": '{"model_name": "abc", "training_data": [1]}',
+            "training_params": '{"model_name": "abc", "parameters": {"training_data": [1]}}',
         },
         training_input_dir="training_input_dir",
         training_output_dir=os.path.join(output_dir, training_id),
@@ -173,7 +173,7 @@ def test_model_train_validation_error_raises(sample_model_train_servicer, output
         customTrainingID=str(uuid.uuid4()),
         request_dict={
             "train_module": "00110203-0405-0607-0809-0a0b02dd0e0f",
-            "training_params": '{"model_name": "abc", "training_data": [1], "batch_size": 999}',
+            "training_params": '{"model_name": "abc", "parameters": {"training_data": [1], "batch_size": 999}}',
         },
         training_input_dir="training_input_dir",
         training_output_dir=os.path.join(output_dir, training_id),
@@ -197,13 +197,15 @@ def test_model_train_sample_widget(sample_model_train_servicer, output_dir):
             "training_params": json.dumps(
                 {
                     "model_name": "abc",
-                    "training_data": {
-                        "jsondata": {
-                            "data": [
-                                sample_lib.data_model.SampleTrainingType(
-                                    number=1
-                                ).to_dict()
-                            ]
+                    "parameters": {
+                        "training_data": {
+                            "jsondata": {
+                                "data": [
+                                    sample_lib.data_model.SampleTrainingType(
+                                        number=1
+                                    ).to_dict()
+                                ]
+                            },
                         },
                     },
                 }
@@ -253,9 +255,11 @@ def test_files_from_training_input_dir_are_used(
             "training_params": json.dumps(
                 {
                     "model_name": "abc",
-                    "training_data": {
-                        "file": {
-                            "filename": input_file_name  # This is relative to training_input_dir
+                    "parameters": {
+                        "training_data": {
+                            "file": {
+                                "filename": input_file_name  # This is relative to training_input_dir
+                            },
                         },
                     },
                 }
diff --git a/tests/runtime/test_grpc_server.py b/tests/runtime/test_grpc_server.py
index 8c1cbdf8..2c6303ef 100644
--- a/tests/runtime/test_grpc_server.py
+++ b/tests/runtime/test_grpc_server.py
@@ -126,13 +126,15 @@ def test_model_train(runtime_grpc_server):
             "training_params": json.dumps(
                 {
                     "model_name": model_name,
-                    "training_data": {
-                        "jsondata": {
-                            "data": [
-                                sample_lib.data_model.SampleTrainingType(
-                                    number=1
-                                ).to_dict()
-                            ]
+                    "parameters": {
+                        "training_data": {
+                            "jsondata": {
+                                "data": [
+                                    sample_lib.data_model.SampleTrainingType(
+                                        number=1
+                                    ).to_dict()
+                                ]
+                            },
                         },
                     },
                 }
@@ -364,10 +366,15 @@ def test_train_fake_module_ok_response_and_can_predict_with_trained_model(
     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,
-        training_data=training_data,
-        union_list=["str", "sequence"],
+        parameters=train_request_params_class(
+            training_data=training_data,
+            union_list=["str", "sequence"],
+        ),
     ).to_proto()
 
     actual_response = train_stub.SampleTaskSampleModuleTrain(train_request)
@@ -404,7 +411,10 @@ def test_train_fake_module_ok_response_with_loaded_model_can_predict_with_traine
     ).to_proto()
     model_name = random_test_id()
     train_request = sample_train_service.messages.SampleTaskCompositeModuleTrainRequest(
-        model_name=model_name, sample_block=sample_model
+        model_name=model_name,
+        parameters=sample_train_service.messages.SampleTaskCompositeModuleTrainParameters(
+            sample_block=sample_model
+        ),
     )
     actual_response = train_stub.SampleTaskCompositeModuleTrain(train_request)
     assert_training_successful(
@@ -447,9 +457,11 @@ def test_train_fake_module_does_not_change_another_instance_model_of_block(
 
     train_request = sample_train_service.messages.OtherTaskOtherModuleTrainRequest(
         model_name="Bar Training",
-        sample_input_sampleinputtype=SampleInputType(name="Gabe").to_proto(),
-        batch_size=100,
-        training_data=training_data,
+        parameters=sample_train_service.messages.OtherTaskOtherModuleTrainParameters(
+            sample_input_sampleinputtype=SampleInputType(name="Gabe").to_proto(),
+            batch_size=100,
+            training_data=training_data,
+        ),
     )
     actual_response = train_stub.OtherTaskOtherModuleTrain(train_request)
     assert_training_successful(
@@ -498,19 +510,24 @@ def test_train_primitive_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(
         model_name=model_name,
-        sample_input=SampleInputType(name="Gabe"),
-        simple_list=["hello", "world"],
-        union_list=["str", "sequence"],
-        union_list2=[1, 2],
-        union_list3=[True, False],
-        union_list4=123,
-        training_params_json_dict={"foo": {"bar": [1, 2, 3]}},
-        training_params_json_dict_list=[{"foo": {"bar": [1, 2, 3]}}],
-        training_params_dict={"layer_sizes": 100, "window_scaling": 200},
-        training_params_dict_int={1: 0.1, 2: 0.01},
+        parameters=train_request_params_class(
+            sample_input=SampleInputType(name="Gabe"),
+            simple_list=["hello", "world"],
+            union_list=["str", "sequence"],
+            union_list2=[1, 2],
+            union_list3=[True, False],
+            union_list4=123,
+            training_params_json_dict={"foo": {"bar": [1, 2, 3]}},
+            training_params_json_dict_list=[{"foo": {"bar": [1, 2, 3]}}],
+            training_params_dict={"layer_sizes": 100, "window_scaling": 200},
+            training_params_dict_int={1: 0.1, 2: 0.01},
+        ),
     ).to_proto()
 
     training_response = train_stub.SampleTaskSamplePrimitiveModuleTrain(train_request)
@@ -557,8 +574,10 @@ def test_train_fake_module_ok_response_with_datastream_jsondata(
     model_name = random_test_id()
     train_request = sample_train_service.messages.SampleTaskSampleModuleTrainRequest(
         model_name=model_name,
-        batch_size=42,
-        training_data=training_data,
+        parameters=sample_train_service.messages.SampleTaskSampleModuleTrainParameters(
+            batch_size=42,
+            training_data=training_data,
+        ),
     )
 
     actual_response = train_stub.SampleTaskSampleModuleTrain(train_request)
@@ -596,7 +615,9 @@ def test_train_fake_module_ok_response_with_datastream_csv_file(
     model_name = random_test_id()
     train_request = sample_train_service.messages.SampleTaskSampleModuleTrainRequest(
         model_name=model_name,
-        training_data=training_data,
+        parameters=sample_train_service.messages.SampleTaskSampleModuleTrainParameters(
+            training_data=training_data,
+        ),
     )
 
     actual_response = train_stub.SampleTaskSampleModuleTrain(train_request)
@@ -630,7 +651,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_request = sample_train_service.messages.SampleTaskSampleModuleTrainRequest(
-        model_name=model_name, training_data=training_data.to_proto(), sleep_time=10
+        model_name=model_name,
+        parameters=sample_train_service.messages.SampleTaskSampleModuleTrainParameters(
+            training_data=training_data.to_proto(), sleep_time=10
+        ),
     )
     train_response = train_stub.SampleTaskSampleModuleTrain(train_request)
 
@@ -664,7 +688,10 @@ 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 = sample_train_service.messages.SampleTaskSampleModuleTrainRequest(
-        model_name=model_name, training_data=training_data.to_proto(), sleep_time=10
+        model_name=model_name,
+        parameters=sample_train_service.messages.SampleTaskSampleModuleTrainParameters(
+            training_data=training_data.to_proto(), sleep_time=10
+        ),
     )
     train_response = train_stub.SampleTaskSampleModuleTrain(train_request)
 
@@ -684,7 +711,10 @@ def test_cancel_does_not_affect_other_models(
     # train another model
     model_name2 = random_test_id()
     train_request2 = sample_train_service.messages.SampleTaskSampleModuleTrainRequest(
-        model_name=model_name2, training_data=training_data.to_proto()
+        model_name=model_name2,
+        parameters=sample_train_service.messages.SampleTaskSampleModuleTrainParameters(
+            training_data=training_data.to_proto()
+        ),
     )
     train_response2 = train_stub.SampleTaskSampleModuleTrain(train_request2)
 
@@ -716,10 +746,11 @@ def test_train_fake_module_error_response_with_unloaded_model(
             model_id=random_test_id()
         ).to_proto()
 
-        train_request = (
-            sample_train_service.messages.SampleTaskCompositeModuleTrainRequest(
-                model_name=random_test_id(), sample_block=sample_model
-            )
+        train_request = sample_train_service.messages.SampleTaskCompositeModuleTrainRequest(
+            model_name=random_test_id(),
+            parameters=sample_train_service.messages.SampleTaskCompositeModuleTrainParameters(
+                sample_block=sample_model
+            ),
         )
         train_stub.SampleTaskCompositeModuleTrain(train_request)
     assert context.value.code() == grpc.StatusCode.NOT_FOUND
@@ -1171,7 +1202,6 @@ def test_grpc_sever_shutdown_with_model_poll(open_port):
             RUNTIME_METRICS_ENABLED="false",
         )
         with server_proc as proc:
-
             # Wait for the server to be up
             _assert_connection(
                 grpc.insecure_channel(f"localhost:{open_port}"), max_failures=500
diff --git a/tests/runtime/test_http_server.py b/tests/runtime/test_http_server.py
index 77d6ba02..f67326a1 100644
--- a/tests/runtime/test_http_server.py
+++ b/tests/runtime/test_http_server.py
@@ -17,7 +17,8 @@ Tests for the caikit HTTP server
 # Standard
 from contextlib import contextmanager
 from dataclasses import dataclass
-from typing import Dict, List
+from typing import Dict, List, Union, get_args
+import enum
 import json
 import os
 import signal
@@ -25,14 +26,24 @@ import tempfile
 
 # Third Party
 from fastapi.testclient import TestClient
+import numpy as np
+import pydantic
 import pytest
 import requests
 import tls_test_tools
 
+# First Party
+from py_to_proto.dataclass_to_proto import Annotated
+
 # Local
 from caikit.core import MODEL_MANAGER, DataObjectBase, dataobject
+from caikit.core.data_model.base import DataBase
+from caikit.interfaces.common.data_model.stream_sources import File
 from caikit.interfaces.nlp.data_model import GeneratedTextStreamResult, GeneratedToken
 from caikit.runtime import http_server
+from caikit.runtime.service_generation.data_stream_source import make_data_stream_source
+from sample_lib.data_model import SampleInputType, SampleOutputType
+from sample_lib.data_model.sample import SampleTrainingType
 from tests.conftest import temp_config
 from tests.runtime.conftest import (
     ModuleSubproc,
@@ -240,6 +251,235 @@ def test_services_disabled(open_port, enabled_services):
             # )
 
 
+## Functional Tests #######################################################################
+
+
+def test_build_dm_object_simple(runtime_http_server):
+    """Test building our simple DM objects through pydantic objects"""
+
+    # get our DM class
+    sample_input_dm_class = DataBase.get_class_for_name("SampleInputType")
+    # get pydantic model for our DM class
+    sample_input_pydantic_model = http_server.PYDANTIC_TO_DM_MAPPING.get(
+        sample_input_dm_class
+    )
+    # build our DM object using a pydantic object
+    sample_input_dm_obj = runtime_http_server._build_dm_object(
+        sample_input_pydantic_model(name="Hello world")
+    )
+
+    # assert it's our DM object, all fine and dandy
+    assert isinstance(sample_input_dm_obj, DataBase)
+    assert sample_input_dm_obj.to_json() == '{"name": "Hello world"}'
+
+
+def test_build_dm_object_datastream_jsondata(runtime_http_server):
+    """Test building our datastream DM objects through pydantic objects"""
+
+    # get our DM class
+    datastream_dm_class = DataBase.get_class_for_name(
+        "DataStreamSourceSampleTrainingType"
+    )
+    # get pydantic model for our DM class
+    datastream_pydantic_model = http_server.PYDANTIC_TO_DM_MAPPING.get(
+        datastream_dm_class
+    )
+    # build our DM Datastream JsonData object using a pydantic object
+    datastream_dm_obj = runtime_http_server._build_dm_object(
+        datastream_pydantic_model(data_stream={"data": [{"number": 1}, {"number": 2}]})
+    )
+
+    # assert it's our DM object, all fine and dandy
+    assert isinstance(datastream_dm_obj, DataBase)
+    assert (
+        datastream_dm_obj.to_json()
+        == '{"jsondata": {"data": [{"number": 1}, {"number": 2}]}}'
+    )
+
+
+def test_build_dm_object_datastream_file(runtime_http_server):
+    # get our DM class
+    datastream_dm_class = DataBase.get_class_for_name(
+        "DataStreamSourceSampleTrainingType"
+    )
+    # get pydantic model for our DM class
+    datastream_pydantic_model = http_server.PYDANTIC_TO_DM_MAPPING.get(
+        datastream_dm_class
+    )
+
+    # build our DM Datastream File object using a pydantic object
+    datastream_dm_obj = runtime_http_server._build_dm_object(
+        datastream_pydantic_model(data_stream={"filename": "hello"})
+    )
+
+    # assert it's our DM object, all fine and dandy
+    assert isinstance(datastream_dm_obj, DataBase)
+    assert isinstance(datastream_dm_obj.data_stream, File)
+    assert datastream_dm_obj.to_json() == '{"file": {"filename": "hello"}}'
+
+
+@pytest.mark.parametrize(
+    "input, output",
+    [
+        (np.integer, int),
+        (np.floating, float),
+        (int, int),
+        (float, float),
+        (bool, bool),
+        (str, str),
+        (bytes, bytes),
+        (type(None), type(None)),
+        (enum.Enum, enum.Enum),
+        (Annotated[str, "blah"], str),
+        (Union[str, int], Union[str, int]),
+        (List[str], List[str]),
+        (List[Annotated[str, "blah"]], List[str]),
+        (Dict[str, int], Dict[str, int]),
+        (Dict[Annotated[str, "blah"], int], Dict[str, int]),
+    ],
+)
+def test_get_pydantic_type(input, output):
+    assert http_server.RuntimeHTTPServer._get_pydantic_type(input) == output
+
+
+def test_get_pydantic_type_union():
+    union_type = Union[SampleInputType, SampleOutputType]
+    return_type = http_server.RuntimeHTTPServer._get_pydantic_type(union_type)
+    assert all(
+        issubclass(ret_type, pydantic.BaseModel) for ret_type in get_args(return_type)
+    )
+
+
+def test_get_pydantic_type_DM():
+    # DM case
+    sample_input_dm_class = DataBase.get_class_for_name("SampleInputType")
+    sample_input_pydantic_model = http_server.RuntimeHTTPServer._get_pydantic_type(
+        sample_input_dm_class
+    )
+
+    assert issubclass(sample_input_pydantic_model, pydantic.BaseModel)
+    assert sample_input_pydantic_model in http_server.PYDANTIC_TO_DM_MAPPING
+    assert sample_input_pydantic_model is http_server.PYDANTIC_TO_DM_MAPPING.get(
+        sample_input_dm_class
+    )
+
+
+def test_get_pydantic_type_throws_random_type():
+    # error case
+    with pytest.raises(TypeError):
+        http_server.RuntimeHTTPServer._get_pydantic_type("some_random_type")
+
+
+def test_pydantic_wrapping_with_enums():
+    """Check that the pydantic wrapping works on our data models when they have enums"""
+    # The NLP GeneratedTextStreamResult data model contains enums
+
+    # Check that our data model is fine and dandy
+    token = GeneratedToken(text="foo")
+    assert token.text == "foo"
+
+    # Wrap the containing data model in pydantic
+    http_server.RuntimeHTTPServer._dataobject_to_pydantic(GeneratedTextStreamResult)
+
+    # Check that our data model is _still_ fine and dandy
+    token = GeneratedToken(text="foo")
+    assert token.text == "foo"
+
+
+def test_pydantic_wrapping_with_lists(runtime_http_server):
+    """Check that pydantic wrapping works on data models with lists"""
+
+    @dataobject(package="http")
+    class BarTest(DataObjectBase):
+        baz: int
+
+    @dataobject(package="http")
+    class FooTest(DataObjectBase):
+        bars: List[BarTest]
+
+    foo = FooTest(bars=[BarTest(1)])
+    assert foo.bars[0].baz == 1
+
+    runtime_http_server._dataobject_to_pydantic(FooTest)
+
+    foo = FooTest(bars=[BarTest(1)])
+    assert foo.bars[0].baz == 1
+
+
+def test_dataobject_to_pydantic_simple_DM():
+    """Test that we can create a pydantic model from a simple DM"""
+    sample_input_dm_class = DataBase.get_class_for_name("SampleInputType")
+    sample_input_pydantic_model = http_server.RuntimeHTTPServer._dataobject_to_pydantic(
+        sample_input_dm_class
+    )
+    model_instance = sample_input_pydantic_model.model_validate_json(
+        '{"name": "world"}'
+    )
+    assert {"name": str} == sample_input_pydantic_model.__annotations__
+    assert issubclass(sample_input_pydantic_model, pydantic.BaseModel)
+    assert sample_input_pydantic_model in http_server.PYDANTIC_TO_DM_MAPPING
+    assert model_instance.name == "world"
+
+
+def test_dataobject_to_pydantic_simple_DM_extra_forbidden_throws():
+    """Test that if we forbid extra values, then we raise if we pass in extra values"""
+    sample_input_dm_class = DataBase.get_class_for_name("SampleInputType")
+
+    sample_input_pydantic_model = http_server.RuntimeHTTPServer._dataobject_to_pydantic(
+        sample_input_dm_class
+    )
+    # sample_input_pydantic_model_extra_forbidden doesn't allow anything extra
+    with pytest.raises(pydantic.ValidationError) as e1:
+        sample_input_pydantic_model.model_validate_json('{"blah": "world"}')
+    assert "Extra inputs are not permitted" in e1.value.errors()[0]["msg"]
+
+    with pytest.raises(pydantic.ValidationError) as e2:
+        sample_input_pydantic_model.model_validate_json(
+            '{"name": "world", "blah": "world"}'
+        )
+    assert "Extra inputs are not permitted" in e2.value.errors()[0]["msg"]
+
+
+def test_dataobject_to_pydantic_oneof():
+    """Test that we can create a pydantic model from a DM with a Union"""
+    make_data_stream_source(SampleTrainingType)
+    sample_input_dm_class = DataBase.get_class_for_name(
+        "DataStreamSourceSampleTrainingType"
+    )
+    data_stream_source_pydantic_model = (
+        http_server.RuntimeHTTPServer._dataobject_to_pydantic(sample_input_dm_class)
+    )
+
+    assert issubclass(data_stream_source_pydantic_model, pydantic.BaseModel)
+    assert {
+        "data_stream": Union[
+            http_server.PYDANTIC_TO_DM_MAPPING.get(sample_input_dm_class.JsonData),
+            http_server.PYDANTIC_TO_DM_MAPPING.get(sample_input_dm_class.File),
+            http_server.PYDANTIC_TO_DM_MAPPING.get(sample_input_dm_class.ListOfFiles),
+            http_server.PYDANTIC_TO_DM_MAPPING.get(sample_input_dm_class.Directory),
+            http_server.PYDANTIC_TO_DM_MAPPING.get(sample_input_dm_class.S3Files),
+        ]
+    } == data_stream_source_pydantic_model.__annotations__
+    assert data_stream_source_pydantic_model in http_server.PYDANTIC_TO_DM_MAPPING
+
+    assert issubclass(
+        http_server.PYDANTIC_TO_DM_MAPPING.get(sample_input_dm_class.JsonData),
+        type(
+            data_stream_source_pydantic_model.model_validate_json(
+                '{"data_stream": {"data": [{"number": 1}]}}'
+            ).data_stream
+        ),
+    )
+    assert issubclass(
+        http_server.PYDANTIC_TO_DM_MAPPING.get(sample_input_dm_class.File),
+        type(
+            data_stream_source_pydantic_model.model_validate_json(
+                '{"data_stream": {"filename": "file1"}}'
+            ).data_stream
+        ),
+    )
+
+
 ## Inference Tests #######################################################################
 
 
@@ -265,8 +505,8 @@ def test_inference_sample_task(sample_task_model_id, runtime_http_server):
             f"/api/v1/{sample_task_model_id}/task/sample",
             json=json_input,
         )
-        assert response.status_code == 200
         json_response = json.loads(response.content.decode(response.default_encoding))
+        assert response.status_code == 200, json_response
         assert json_response["greeting"] == "Hello world"
 
 
@@ -296,74 +536,11 @@ def test_inference_other_task(other_task_model_id, runtime_http_server):
             f"/api/v1/{other_task_model_id}/task/other",
             json=json_input,
         )
-        assert response.status_code == 200
         json_response = json.loads(response.content.decode(response.default_encoding))
+        assert response.status_code == 200, json_response
         assert json_response["farewell"] == "goodbye: world 42 times"
 
 
-def test_inference_other_task_request_throws_missing_field(
-    other_task_model_id, runtime_http_server
-):
-    """Test to check 400 in case of missing request fields"""
-    with TestClient(runtime_http_server.app) as client:
-        json_input = {}
-        response = client.post(
-            f"/api/v1/{other_task_model_id}/task/other",
-            json=json_input,
-        )
-        assert response.status_code == 400
-        json_response = json.loads(response.content.decode(response.default_encoding))
-        assert (
-            "missing 1 required positional argument: 'sample_input'"
-            in json_response["details"]
-        )
-
-
-def test_inference_other_task_response_throws_missing_producer_id(
-    other_task_model_id, runtime_http_server
-):
-    """Test to check we throw a 500 if response has fields that are missing"""
-    with TestClient(runtime_http_server.app) as client:
-        json_input = {
-            "inputs": {"name": "world"},
-            # this intentionally skips sending back a producer-id in our module
-            "parameters": {"include_producer_id": False},
-        }
-        response = client.post(
-            f"/api/v1/{other_task_model_id}/task/other",
-            json=json_input,
-        )
-        assert response.status_code == 500
-        json_response = json.loads(response.content.decode(response.default_encoding))
-        assert json_response["detail"][0]["loc"] == ["response", "producer_id"]
-
-
-def test_inference_other_task_request_throws_incorrect_field(
-    other_task_model_id, runtime_http_server
-):
-    """Test to check we throw a 422 if request has an incorrect field type"""
-    with TestClient(runtime_http_server.app) as client:
-        json_input = {
-            "inputs": {"name": "world"},
-            "parameters": {"include_producer_id": "wrong_type"},
-        }
-        response = client.post(
-            f"/api/v1/{other_task_model_id}/task/other",
-            json=json_input,
-        )
-        json_response = json.loads(response.content.decode(response.default_encoding))
-        assert json_response["detail"][0]["loc"] == [
-            "body",
-            "parameters",
-            "include_producer_id",
-        ]
-        assert (
-            json_response["detail"][0]["msg"]
-            == "Input should be a valid boolean, unable to interpret input"
-        )
-        assert response.status_code == 422
-
-
 def test_inference_streaming_sample_module(sample_task_model_id, runtime_http_server):
     """Simple check for testing a happy path unary-stream case"""
     with TestClient(runtime_http_server.app) as client:
@@ -397,17 +574,41 @@ def test_model_not_found(runtime_http_server):
         assert response.status_code == 404
 
 
-def test_inference_sample_task_throws_incorrect_input(
+def test_inference_sample_task_incorrect_input(
     sample_task_model_id, runtime_http_server
 ):
-    """error check for a request with incorrect input"""
+    """Test that with an incorrect input, the test doesn't throw but
+    instead returns None"""
     with TestClient(runtime_http_server.app) as client:
-        json_input = {"blah": {"sample_input": {"name": "world"}}}
+        json_input = {
+            "inputs": {"blah": "world"},
+        }
         response = client.post(
             f"/api/v1/{sample_task_model_id}/task/sample",
             json=json_input,
         )
-        assert response.status_code == 400
+        assert response.status_code == 422, response.content.decode(
+            response.default_encoding
+        )
+
+
+@pytest.mark.skip("Skipping since we're not tacking forward compatibility atm")
+def test_inference_sample_task_forward_compatibility(
+    sample_task_model_id, runtime_http_server
+):
+    """Test that clients can send in params that don't exist on server
+    without any error"""
+    with TestClient(runtime_http_server.app) as client:
+        json_input = {
+            "inputs": {"name": "world", "blah": "blah"},
+        }
+        response = client.post(
+            f"/api/v1/{sample_task_model_id}/task/sample",
+            json=json_input,
+        )
+        json_response = json.loads(response.content.decode(response.default_encoding))
+        assert response.status_code == 200, json_response
+        assert json_response["greeting"] == "Hello world"
 
 
 def test_health_check_ok(runtime_http_server):
@@ -418,69 +619,6 @@ def test_health_check_ok(runtime_http_server):
         assert response.text == "OK"
 
 
-def test_pydantic_wrapping_with_enums():
-    """Check that the pydantic wrapping works on our data models when they have enums"""
-    # The NLP GeneratedTextStreamResult data model contains enums
-
-    # Check that our data model is fine and dandy
-    token = GeneratedToken(text="foo")
-    assert token.text == "foo"
-
-    # Wrap the containing data model in pydantic
-    http_server.RuntimeHTTPServer._dataobject_to_pydantic(GeneratedTextStreamResult)
-
-    # Check that our data model is _still_ fine and dandy
-    token = GeneratedToken(text="foo")
-    assert token.text == "foo"
-
-
-def test_pydantic_wrapping_with_inheritance():
-    """Check that the pydantic wrapping works on our data models when they involve inheritance"""
-
-    @dataobject
-    class Base(DataObjectBase):
-        foo: int
-        bar: int
-
-    @dataobject
-    class Derived(Base):
-        bar: str
-        baz: str
-
-    # Check that our data model is fine and dandy
-    derived = Derived(bar="one", foo=2)
-    assert derived.bar == "one"
-    assert derived.foo == 2
-
-    # Wrap the containing data model in pydantic
-    pydantic_datamodel = http_server.RuntimeHTTPServer._dataobject_to_pydantic(Derived)
-
-    # Check that our data model is _still_ fine and dandy
-    new_derived = pydantic_datamodel(bar="one", foo=2)
-    assert new_derived.bar == "one"
-    assert new_derived.foo == 2
-
-
-def test_pydantic_wrapping_with_lists():
-    """Check that pydantic wrapping works on data models with lists"""
-
-    @dataobject(package="http")
-    class BarTest(DataObjectBase):
-        baz: int
-
-    @dataobject(package="http")
-    class FooTest(DataObjectBase):
-        bars: List[BarTest]
-
-    foo = FooTest(bars=[BarTest(1)])
-    assert foo.bars[0].baz == 1
-
-    http_server.RuntimeHTTPServer._dataobject_to_pydantic(FooTest)
-
-    foo = FooTest(bars=[BarTest(1)])
-    assert foo.bars[0].baz == 1
-
-
 def test_http_server_shutdown_with_model_poll(open_port):
     """Test that a SIGINT successfully shuts down the running server"""
     with tempfile.TemporaryDirectory() as workdir:
@@ -522,12 +660,11 @@ def test_train_sample_task(runtime_http_server):
     model_name = "sample_task_train"
     with TestClient(runtime_http_server.app) as client:
         json_input = {
-            "inputs": {
-                "model_name": model_name,
+            "model_name": model_name,
+            "parameters": {
                 "training_data": {"data_stream": {"data": [{"number": 1}]}},
-                # "training_data": {"data_stream": {"file": {"filename": "file1"}}},
+                "batch_size": 42,
             },
-            "parameters": {"batch_size": 42},
         }
         training_response = client.post(
             f"/api/v1/SampleTaskSampleModuleTrain",
@@ -535,12 +672,12 @@ def test_train_sample_task(runtime_http_server):
         )
 
         # assert training response
-        assert training_response.status_code == 200
         training_json_response = json.loads(
             training_response.content.decode(training_response.default_encoding)
         )
+        assert training_response.status_code == 200, training_json_response
         assert (training_id := training_json_response["training_id"])
-        assert (model_name := training_json_response["model_name"]) == model_name
+        assert training_json_response["model_name"] == model_name
 
         # assert trained model
         result = MODEL_MANAGER.get_model_future(training_id).load()
@@ -563,8 +700,8 @@ def test_train_sample_task(runtime_http_server):
             f"/api/v1/{model_name}/task/sample",
             json=json_input_inference,
         )
-        assert response.status_code == 200
         json_response = json.loads(response.content.decode(response.default_encoding))
+        assert response.status_code == 200, json_response
         assert json_response["greeting"] == "Hello world"
 
 
@@ -573,12 +710,12 @@ def test_train_sample_task_throws_s3_value_error(runtime_http_server):
     model_name = "sample_task_train"
     with TestClient(runtime_http_server.app) as client:
         json_input = {
-            "inputs": {
-                "model_name": model_name,
-                "training_data": {"data_stream": {"file": "hello"}},
-                "output_path": {"path": "non-existent path_to_s3"},
+            "model_name": model_name,
+            "output_path": {"path": "non-existent path_to_s3"},
+            "parameters": {
+                "training_data": {"data_stream": {"data": [{"number": 1}]}},
+                "batch_size": 42,
             },
-            "parameters": {"batch_size": 42},
         }
         training_response = client.post(
             f"/api/v1/SampleTaskSampleModuleTrain",
@@ -588,15 +725,17 @@ def test_train_sample_task_throws_s3_value_error(runtime_http_server):
             "S3 output path not supported by this runtime"
             in training_response.content.decode(training_response.default_encoding)
         )
-        assert training_response.status_code == 500
+        assert training_response.status_code == 500, training_response.content.decode(
+            training_response.default_encoding
+        )
 
 
 def test_train_primitive_task(runtime_http_server):
     model_name = "primitive_task_train"
     with TestClient(runtime_http_server.app) as client:
         json_input = {
-            "inputs": {
-                "model_name": model_name,
+            "model_name": model_name,
+            "parameters": {
                 "sample_input": {"name": "test"},
                 "simple_list": ["hello", "world"],
                 "union_list": ["hello", "world"],
@@ -604,8 +743,6 @@ def test_train_primitive_task(runtime_http_server):
                 "union_list3": ["hello", "world"],
                 "union_list4": 1,
                 "training_params_json_dict_list": [{"foo": {"bar": [1, 2, 3]}}],
-            },
-            "parameters": {
                 "training_params_json_dict": {"foo": {"bar": [1, 2, 3]}},
                 "training_params_dict": {"layer_sizes": 100, "window_scaling": 200},
                 "training_params_dict_int": {1: 0.1, 2: 0.01},
@@ -617,12 +754,12 @@ def test_train_primitive_task(runtime_http_server):
             json=json_input,
         )
         # assert training response
-        assert training_response.status_code == 200
         training_json_response = json.loads(
             training_response.content.decode(training_response.default_encoding)
         )
+        assert training_response.status_code == 200, training_json_response
         assert (training_id := training_json_response["training_id"])
-        assert (model_name := training_json_response["model_name"]) == model_name
+        assert training_json_response["model_name"] == model_name
 
         # assert trained model
         result = MODEL_MANAGER.get_model_future(training_id).load()
@@ -649,8 +786,8 @@ def test_train_primitive_task(runtime_http_server):
             f"/api/v1/{model_name}/task/sample",
             json=json_input_inference,
         )
-        assert response.status_code == 200
         json_response = json.loads(response.content.decode(response.default_encoding))
+        assert response.status_code == 200, json_response
         assert json_response["greeting"] == "hello: primitives! [1, 2, 3] 100"
 
 
@@ -658,11 +795,11 @@ def test_train_other_task(runtime_http_server):
     model_name = "other_task_train"
     with TestClient(runtime_http_server.app) as client:
         json_input = {
-            "inputs": {
-                "model_name": model_name,
+            "model_name": model_name,
+            "parameters": {
                 "training_data": {"data_stream": {"data": [1, 2]}},
                 "sample_input": {"name": "test"},
-            }
+            },
         }
 
         training_response = client.post(
@@ -670,12 +807,12 @@ def test_train_other_task(runtime_http_server):
             json=json_input,
         )
         # assert training response
-        assert training_response.status_code == 200
         training_json_response = json.loads(
             training_response.content.decode(training_response.default_encoding)
         )
+        assert training_response.status_code == 200, training_json_response
         assert (training_id := training_json_response["training_id"])
-        assert (model_name := training_json_response["model_name"]) == model_name
+        assert training_json_response["model_name"] == model_name
 
         # assert trained model
         result = MODEL_MANAGER.get_model_future(training_id).load()
@@ -698,6 +835,6 @@ def test_train_other_task(runtime_http_server):
             f"/api/v1/{model_name}/task/other",
             json=json_input_inference,
         )
-        assert response.status_code == 200
         json_response = json.loads(response.content.decode(response.default_encoding))
+        assert response.status_code == 200, json_response
         assert json_response["farewell"] == "goodbye: world 64 times"
diff --git a/tests/runtime/utils/test_servicer_util.py b/tests/runtime/utils/test_servicer_util.py
index fe7beb0c..c7a6adac 100644
--- a/tests/runtime/utils/test_servicer_util.py
+++ b/tests/runtime/utils/test_servicer_util.py
@@ -295,11 +295,14 @@ def test_global_train_build_caikit_library_request_dict_strips_empty_list_from_r
     stream_type = caikit.interfaces.common.data_model.DataStreamSourceSampleTrainingType
     training_data = stream_type(jsondata=stream_type.JsonData(data=[])).to_proto()
     train_request = sample_train_service.messages.SampleTaskSampleModuleTrainRequest(
-        model_name=random_test_id(), training_data=training_data
+        model_name=random_test_id(),
+        parameters=sample_train_service.messages.SampleTaskSampleModuleTrainParameters(
+            training_data=training_data
+        ),
     )
 
     caikit.core_request = build_caikit_library_request_dict(
-        train_request,
+        train_request.parameters,
         sample_lib.modules.sample_task.SampleModule.TRAIN_SIGNATURE,
     )
 
@@ -319,12 +322,15 @@ def test_global_train_build_caikit_library_request_dict_works_for_repeated_field
     training_data = stream_type(jsondata=stream_type.JsonData(data=[])).to_proto()
     train_request = sample_train_service.messages.SampleTaskListModuleTrainRequest(
         model_name=random_test_id(),
-        training_data=training_data,
-        poison_pills=["Bob Marley", "Bunny Livingston"],
+        parameters=sample_train_service.messages.SampleTaskListModuleTrainParameters(
+            training_data=training_data,
+            poison_pills=["Bob Marley", "Bunny Livingston"],
+        ),
     )
 
     caikit.core_request = build_caikit_library_request_dict(
-        train_request, sample_lib.modules.sample_task.ListModule.TRAIN_SIGNATURE
+        train_request.parameters,
+        sample_lib.modules.sample_task.ListModule.TRAIN_SIGNATURE,
     )
 
     # model_name is not expected to be passed through
@@ -346,12 +352,14 @@ def test_global_train_build_caikit_library_request_dict_ok_with_DataStreamSource
 
     train_request = sample_train_service.messages.OtherTaskOtherModuleTrainRequest(
         model_name="Bar Training",
-        sample_input_sampleinputtype=SampleInputType(name="Gabe").to_proto(),
-        batch_size=100,
-        training_data=training_data,
+        parameters=sample_train_service.messages.OtherTaskOtherModuleTrainParameters(
+            sample_input_sampleinputtype=SampleInputType(name="Gabe").to_proto(),
+            batch_size=100,
+            training_data=training_data,
+        ),
     )
     caikit.core_request = build_caikit_library_request_dict(
-        train_request,
+        train_request.parameters,
         sample_lib.modules.other_task.OtherModule.TRAIN_SIGNATURE,
     )
 
@@ -371,11 +379,13 @@ def test_global_train_build_caikit_library_request_dict_ok_with_data_stream_file
     ).to_proto()
     train_request = sample_train_service.messages.SampleTaskSampleModuleTrainRequest(
         model_name=random_test_id(),
-        training_data=training_data,
+        parameters=sample_train_service.messages.SampleTaskSampleModuleTrainParameters(
+            training_data=training_data,
+        ),
     )
 
     caikit.core_request = build_caikit_library_request_dict(
-        train_request,
+        train_request.parameters,
         sample_lib.modules.sample_task.SampleModule.TRAIN_SIGNATURE,
     )
 
@@ -395,12 +405,14 @@ def test_global_train_build_caikit_library_request_dict_ok_with_training_data_as
     ).to_proto()
     train_request = sample_train_service.messages.SampleTaskListModuleTrainRequest(
         model_name=random_test_id(),
-        training_data=training_data,
-        poison_pills=["Bob Marley", "Bunny Livingston"],
+        parameters=sample_train_service.messages.SampleTaskListModuleTrainParameters(
+            training_data=training_data,
+            poison_pills=["Bob Marley", "Bunny Livingston"],
+        ),
     )
 
     caikit.core_request = build_caikit_library_request_dict(
-        train_request,
+        train_request.parameters,
         sample_lib.modules.sample_task.ListModule.TRAIN_SIGNATURE,
     )
 
@@ -431,15 +443,15 @@ def test_build_caikit_library_request_dict_works_when_data_stream_directory_incl
         training_data = stream_type(
             directory=stream_type.Directory(dirname=tempdir, extension="json")
         ).to_proto()
-        train_request = (
-            sample_train_service.messages.SampleTaskSampleModuleTrainRequest(
-                model_name=random_test_id(),
+        train_request = sample_train_service.messages.SampleTaskSampleModuleTrainRequest(
+            model_name=random_test_id(),
+            parameters=sample_train_service.messages.SampleTaskSampleModuleTrainParameters(
                 training_data=training_data,
-            )
+            ),
         )
 
         # no error because at least 1 json file exists within the provided dir
         caikit.core_request = build_caikit_library_request_dict(
-            train_request,
+            train_request.parameters,
             sample_lib.modules.sample_task.SampleModule.TRAIN_SIGNATURE,
         )