Repository: caikit/caikit
Issue: Union within a Train function argument throws an exception

## Describe the bug
If we have a `train` function defined like so:
```
def train(
        cls,
        sample_input: Union[SampleInputType, str],
    )
```
I see an exception:
```
ValueError: Invalid input schema, cannot handle unions yet: typing.Union[sample_lib.data_model.sample.SampleInputType, str]
```

## Platform

Please provide details about the environment you are using, including the following:

- Interpreter version: 3.9
- Library version: 0.1.4

## Sample Code

Already provided

## Expected behavior

The introspection needs to pick the `primitive` or `primitive_data_model_types` from the Union.

## Observed behavior

What you see happening (error messages, stack traces, etc...)

## Additional context

Add any other context about the problem here.


Changed files:
tests/runtime/service_generation/test_primitives.py, caikit/runtime/service_generation/create_service.py, caikit/runtime/service_generation/primitives.py, tests/fixtures/sample_lib/blocks/other_task/other_implementation.py, caikit/runtime/service_generation/serializers.py, tests/runtime/servicers/test_global_train_servicer_impl.py

diff --git a/caikit/runtime/service_generation/create_service.py b/caikit/runtime/service_generation/create_service.py
index 6fbfa41..56b3037 100644
--- a/caikit/runtime/service_generation/create_service.py
+++ b/caikit/runtime/service_generation/create_service.py
@@ -74,6 +74,10 @@ def create_training_rpcs(modules: List[Type[ModuleBase]]) -> List[RPCSerializerB
 
     rpcs = []
 
+    primitive_data_model_types = (
+        ConfigParser.get_instance().service_generation.primitive_data_model_types
+    )
+
     for ck_module in modules:
         # If this train function has not been changed from the base, skip it as
         # a module that can't be trained
@@ -99,7 +103,7 @@ def create_training_rpcs(modules: List[Type[ModuleBase]]) -> List[RPCSerializerB
         )
         with alog.ContextLog(log.debug, "Generating train RPC for %s", ck_module):
             try:
-                rpcs.append(ModuleClassTrainRPC(signature))
+                rpcs.append(ModuleClassTrainRPC(signature, primitive_data_model_types))
                 log.debug("Successfully generated train RPC for %s", ck_module)
             except Exception as err:  # pylint: disable=broad-exception-caught
                 log.warning(
diff --git a/caikit/runtime/service_generation/primitives.py b/caikit/runtime/service_generation/primitives.py
index 00c9d05..2aad960 100644
--- a/caikit/runtime/service_generation/primitives.py
+++ b/caikit/runtime/service_generation/primitives.py
@@ -45,43 +45,52 @@ def to_primitive_signature(
     primitives = {}
     log.debug("Building primitive signature for %s", signature)
     for arg, arg_type in signature.items():
-        if _is_primitive_type(arg_type, primitive_data_model_types):
-            if typing.get_origin(arg_type) == Union:
-                union_primitives = [
-                    union_val
-                    for union_val in typing.get_args(arg_type)
-                    if _is_primitive_type(union_val, primitive_data_model_types)
-                ]
-                if len(union_primitives) == 1:
-                    primitives[arg] = union_primitives[0]
-                else:
-                    dm_types = [
-                        arg
-                        for arg in union_primitives
-                        if inspect.isclass(arg) and issubclass(arg, DataBase)
-                    ]
-                    if len(dm_types) > 0:
-                        log.debug2(
-                            "Picking first data model type %s in union primitives %s",
-                            dm_types,
-                            union_primitives,
-                        )
-                        primitives[arg] = dm_types[0]
-                    else:
-                        log.debug(
-                            "Just picking first primitive type %s in union",
-                            union_primitives[0],
-                        )
-                        primitives[arg] = union_primitives[0]
-
-            else:
-                primitives[arg] = arg_type
-        else:
-            log.debug("Skipping non-primitive argument [%s], type [%s]", arg, arg_type)
+        primitive_arg_type = extract_primitive_type_from_union(
+            primitive_data_model_types, arg_type
+        )
+        if primitive_arg_type:
+            primitives[arg] = primitive_arg_type
 
     return primitives
 
 
+def extract_primitive_type_from_union(
+    primitive_data_model_types: List[str], arg_type: Type
+) -> Type:
+    """Returns the primitive arg type from a Union if found"""
+    if _is_primitive_type(arg_type, primitive_data_model_types):
+        if typing.get_origin(arg_type) == Union:
+            union_primitives = [
+                union_val
+                for union_val in typing.get_args(arg_type)
+                if _is_primitive_type(union_val, primitive_data_model_types)
+            ]
+            # if there's only 1 primitive found, return that
+            if len(union_primitives) == 1:
+                return union_primitives[0]
+            # otherwise, try to get the primitive dm objects in the Union
+            dm_types = [
+                arg
+                for arg in union_primitives
+                if inspect.isclass(arg) and issubclass(arg, DataBase)
+            ]
+            # if there are multiple, pick the first one
+            if len(dm_types) > 0:
+                log.debug2(
+                    "Picking first data model type %s in union primitives %s",
+                    dm_types,
+                    union_primitives,
+                )
+                return dm_types[0]
+            log.debug(
+                "Just picking first primitive type %s in union",
+                union_primitives[0],
+            )
+            return union_primitives[0]
+        return arg_type
+    log.debug("Skipping non-primitive argument type [%s]", arg_type)
+
+
 def extract_data_model_type_from_union(arg_type: Type) -> Type:
     """Helper function that determines the right data model type to use from a Union"""
 
diff --git a/caikit/runtime/service_generation/serializers.py b/caikit/runtime/service_generation/serializers.py
index 09c37c2..9284387 100644
--- a/caikit/runtime/service_generation/serializers.py
+++ b/caikit/runtime/service_generation/serializers.py
@@ -56,16 +56,24 @@ class ModuleClassTrainRPC(RPCSerializerBase):
     for a given module class
     """
 
-    def __init__(self, method_signature: CaikitCoreModuleMethodSignature):
+    def __init__(
+        self,
+        method_signature: CaikitCoreModuleMethodSignature,
+        primitive_data_model_types: List[str],
+    ):
         """Initialize a .proto generator with a single module to convert
 
         Args:
             method_signature (CaikitCoreModuleMethodSignature): The module method signature to
             generate an RPC for
+
+            primitive_data_model_types: List[str]
+                List of primitive data model types for a caikit_* library, such as
+                caikit.interfaces.nlp.data_model.RawDocument for nlp domains
         """
         self.clz: Type[ModuleBase] = method_signature.module
         self._method = ModuleClassTrainRPC._mutate_method_signature_for_training(
-            method_signature
+            method_signature, primitive_data_model_types
         )
         self.name = self._module_class_to_rpc_name()
 
@@ -115,7 +123,7 @@ class ModuleClassTrainRPC(RPCSerializerBase):
 
     @staticmethod
     def _mutate_method_signature_for_training(
-        signature,
+        signature, primitive_data_model_types: List[str]
     ) -> Optional[CaikitCoreModuleMethodSignature]:
         # Change return type for async training interface
         return_type = TrainingJob
@@ -136,7 +144,10 @@ class ModuleClassTrainRPC(RPCSerializerBase):
                 # Found a model pointer
                 new_params[name] = ModelPointer
             else:
-                new_params[name] = typ
+                new_params[name] = primitives.extract_primitive_type_from_union(
+                    arg_type=typ,
+                    primitive_data_model_types=primitive_data_model_types,
+                )
 
         return CustomSignature(
             original_signature=signature, parameters=new_params, return_type=return_type
diff --git a/tests/fixtures/sample_lib/blocks/other_task/other_implementation.py b/tests/fixtures/sample_lib/blocks/other_task/other_implementation.py
index 254376e..1129e12 100644
--- a/tests/fixtures/sample_lib/blocks/other_task/other_implementation.py
+++ b/tests/fixtures/sample_lib/blocks/other_task/other_implementation.py
@@ -19,7 +19,9 @@ class OtherBlock(caikit.core.BlockBase):
         self.batch_size = batch_size
         self.learning_rate = learning_rate
 
-    def run(self, sample_input: SampleInputType) -> Union[OtherOutputType, str]:
+    def run(
+        self, sample_input: Union[SampleInputType, str]
+    ) -> Union[OtherOutputType, str]:
         return OtherOutputType(f"goodbye: {sample_input.name} {self.batch_size} times")
 
     @classmethod
@@ -47,7 +49,10 @@ class OtherBlock(caikit.core.BlockBase):
 
     @classmethod
     def train(
-        cls, training_data: DataStream[int], batch_size: int = 64
+        cls,
+        training_data: DataStream[int],
+        sample_input: Union[SampleInputType, str],
+        batch_size: int = 64,
     ) -> "OtherBlock":
         """Sample training method that produces a trained model"""
         # Barf if we were incorrectly passed data not in datastream format
diff --git a/tests/runtime/service_generation/test_primitives.py b/tests/runtime/service_generation/test_primitives.py
index 9d12883..eda443a 100644
--- a/tests/runtime/service_generation/test_primitives.py
+++ b/tests/runtime/service_generation/test_primitives.py
@@ -7,8 +7,48 @@ import pytest
 # Local
 from caikit.runtime.service_generation.primitives import (
     extract_data_model_type_from_union,
+    extract_primitive_type_from_union,
+    to_primitive_signature,
 )
-from sample_lib.data_model import SampleOutputType
+from sample_lib.data_model import SampleInputType, SampleOutputType
+
+
+def test_to_primitive_signature_raw():
+    assert to_primitive_signature(
+        signature={"name": str},
+        primitive_data_model_types=[],
+    ) == {"name": str}
+
+
+def test_to_primitive_signature_union_raw():
+    assert to_primitive_signature(
+        signature={"name": Union[str, int]},
+        primitive_data_model_types=[],
+    ) == {"name": str}
+
+
+def test_to_primitive_signature_dm():
+    assert to_primitive_signature(
+        signature={"name": SampleInputType},
+        primitive_data_model_types=["sample_lib.data_model.SampleInputType"],
+    ) == {"name": SampleInputType}
+
+
+def test_to_primitive_signature_union_dm():
+    assert to_primitive_signature(
+        signature={"name": Union[SampleInputType, str]},
+        primitive_data_model_types=["sample_lib.data_model.SampleInputType"],
+    ) == {"name": SampleInputType}
+
+
+def test_to_primitive_signature_no_primitive():
+    assert (
+        to_primitive_signature(
+            signature={"name": SampleInputType},
+            primitive_data_model_types=[],
+        )
+        == {}
+    )
 
 
 def test_to_output_dm_type_with_None():
@@ -35,3 +75,72 @@ def test_to_output_dm_type_with_union_optional_dm():
         extract_data_model_type_from_union(Union[Optional[SampleOutputType], str])
         == SampleOutputType
     )
+
+
+def test_extract_primitive_dm_type_from_union():
+    """If we provide a list of primitive_data_model_types,
+    then fetch that from the Union"""
+    assert (
+        extract_primitive_type_from_union(
+            primitive_data_model_types=["sample_lib.data_model.SampleInputType"],
+            arg_type=Union[SampleInputType, str],
+        )
+        == SampleInputType
+    )
+
+
+def test_extract_primitive_raw_type_from_union():
+    """If we don't provide a primitive_data_model_types, then
+    fetch the raw primitive from union"""
+    assert (
+        extract_primitive_type_from_union(
+            primitive_data_model_types=[],
+            arg_type=Union[SampleInputType, str],
+        )
+        == str
+    )
+
+
+def test_extract_primitive_multiple_raw_type_union():
+    """If multiple raw primitives exist, return the first one"""
+    assert (
+        extract_primitive_type_from_union(
+            primitive_data_model_types=[], arg_type=Union[str, int]
+        )
+        == str
+    )
+    assert (
+        extract_primitive_type_from_union(
+            primitive_data_model_types=[], arg_type=Union[int, str]
+        )
+        == int
+    )
+
+
+def test_extract_primitive_raw_type_no_union():
+    """If we provide a primitive without a union, we still get it back"""
+    assert (
+        extract_primitive_type_from_union(primitive_data_model_types=[], arg_type=str)
+        == str
+    )
+
+
+def test_extract_primitive_dm_type_no_union():
+    """If we provide a primitive dm without a union, we still get it back"""
+    assert (
+        extract_primitive_type_from_union(
+            primitive_data_model_types=["sample_lib.data_model.SampleInputType"],
+            arg_type=SampleInputType,
+        )
+        == SampleInputType
+    )
+
+
+def test_extract_primitive_type_no_primitive():
+    """If no primitives exist, nothing gets returned"""
+    assert (
+        extract_primitive_type_from_union(
+            primitive_data_model_types=[], arg_type=SampleInputType
+        )
+        == None
+    )
diff --git a/tests/runtime/servicers/test_global_train_servicer_impl.py b/tests/runtime/servicers/test_global_train_servicer_impl.py
index f51d2d7..4ce0be3 100644
--- a/tests/runtime/servicers/test_global_train_servicer_impl.py
+++ b/tests/runtime/servicers/test_global_train_servicer_impl.py
@@ -25,6 +25,7 @@ from caikit.runtime.servicers.global_train_servicer import GlobalTrainServicer
 from caikit.runtime.types.caikit_runtime_exception import CaikitRuntimeException
 from sample_lib.blocks.sample_task.sample_implementation import SampleBlock
 from sample_lib.data_model.sample import (
+    OtherOutputType,
     SampleInputType,
     SampleOutputType,
     SampleTrainingType,
@@ -87,7 +88,7 @@ def set_train_location(request):
 ##############
 
 
-def test_global_train_simple_Sample_Widget(
+def test_global_train_sample_task(
     sample_train_service,
     sample_train_servicer,
     sample_inference_service,
@@ -141,6 +142,58 @@ def test_global_train_simple_Sample_Widget(
     )
 
 
+def test_global_train_other_task(
+    sample_train_service,
+    sample_train_servicer,
+    sample_inference_service,
+    sample_predict_servicer,
+):
+    """Global train of TrainRequest returns a training job with the correct
+    model name, and some training id for a basic train function that doesn't
+    require any loaded model
+    """
+    batch_size = 42
+    stream_type = caikit.interfaces.common.data_model.DataStreamSourceInt
+    training_data = stream_type(jsondata=stream_type.JsonData(data=[1])).to_proto()
+    train_request = sample_train_service.messages.BlocksOtherTaskOtherBlockTrainRequest(
+        model_name="Other block Training",
+        training_data=training_data,
+        sample_input=SampleInputType(name="Gabe").to_proto(),
+        batch_size=batch_size,
+    )
+
+    training_response = sample_train_servicer.Train(train_request)
+    assert training_response.model_name == "Other block Training"
+    assert training_response.training_id is not None
+    assert isinstance(training_response.training_id, str)
+
+    result = sample_train_servicer.training_map.get(
+        training_response.training_id
+    ).result()
+    assert result.batch_size == batch_size
+    assert (
+        result.BLOCK_CLASS
+        == "sample_lib.blocks.other_task.other_implementation.OtherBlock"
+    )
+
+    # give the trained model time to load
+    # TODO: no sleeps in tests!
+    time.sleep(1)
+
+    inference_response = sample_predict_servicer.Predict(
+        sample_inference_service.messages.OtherTaskRequest(
+            sample_input=SampleInputType(name="Gabe").to_proto()
+        ),
+        Fixtures.build_context(training_response.model_name),
+    )
+    assert (
+        inference_response
+        == OtherOutputType(
+            farewell=f"goodbye: Gabe {batch_size} times",
+        ).to_proto()
+    )
+
+
 def test_global_train_Another_Widget_that_requires_SampleWidget_loaded_should_not_raise(
     loaded_model_id,
     sample_train_service,