Repository: caikit/caikit
Pull request #643 (2024-01-19 01:44:34)

Fix union validation
<!--  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**:
Currently there's a bug in the way that we validate module signatures against task definitions when both contain union typed args. The main inconsistency is as follows:
- Say that a module has param named `inputs`, which is of type `Union[int, str, bytes]`
- Now say that this module implements task `Foo`. If `Foo` defines:
     - `inputs` as a an `int`, `str`, OR, `bytes`, everything is happy
     - `inputs` as a `Union` containing some subset of these types, everything explodes with a TypeError

This behavior should be allowed, because in some cases, the module may have args that aren't directly compatible with the base data model, e.g., types supported by the Image data model backend, namely PIL image. We should be able to pass these in at the library level, but ensure that any of the valid serializable arg types can be used in the task definition for support in runtime.

This PR fixes validation of union validation by ensuring that union typed parameters are further inspected, and are allowed if they are subsets of the corresponding union typed signature def.

**Special notes for your reviewer**:
Currently Caikit Computer Vision has `bytes` as the input type for its tasks, but we need this change to be able to support a wider variety of inputs that are supported by the image backend, e.g., paths to images, bytes, decoded images, etc.

**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/core/task.py b/caikit/core/task.py
index 4810d639..06a13a97 100644
--- a/caikit/core/task.py
+++ b/caikit/core/task.py
@@ -169,9 +169,17 @@ class TaskBase:
         ).items():
             signature_type = signature.parameters[parameter_name]
             if parameter_type != signature_type:
-                if typing.get_origin(
-                    signature_type
-                ) == typing.Union and parameter_type in typing.get_args(signature_type):
+                if typing.get_origin(signature_type) == typing.Union and (
+                    # Either our parameter type is not a union & is part of the union signature
+                    parameter_type in typing.get_args(signature_type)
+                    # Or our parameter type is a union that's a subset of the union signature
+                    or (
+                        typing.get_origin(parameter_type) == typing.Union
+                        and set(typing.get_args(parameter_type)).issubset(
+                            set(typing.get_args(signature_type))
+                        )
+                    )
+                ):
                     continue
                 if input_streaming and cls._is_iterable_type(parameter_type):
                     streaming_type = typing.get_args(parameter_type)[0]


Test patch:
diff --git a/tests/core/test_task.py b/tests/core/test_task.py
index 27c7e201..0e509bd8 100644
--- a/tests/core/test_task.py
+++ b/tests/core/test_task.py
@@ -563,6 +563,54 @@ def test_decorator_adds_default_run_method_to_modules():
     )
 
 
+def test_validation_allows_union_subsets():
+    """Validate that a task can take a union type that is a subset of implementing module types."""
+    # Task param types need to correctly map to proto types since they're used by runtime
+    @task(
+        unary_parameters={"sample_input": Union[str, int]},
+        unary_output_type=SampleOutputType,
+        streaming_output_type=Iterable[SampleOutputType],
+    )
+    class SomeTask(TaskBase):
+        pass
+
+    # But a module may consume types that are not backed by proto, e.g., PIL images
+    @caikit.core.module(
+        id=str(uuid.uuid4()),
+        name="SomeModule",
+        version="0.0.1",
+        task=SomeTask,
+    )
+    class SomeModule(caikit.core.ModuleBase):
+        def run(self, sample_input: Union[str, int, bytes]) -> SampleOutputType:
+            pass
+
+
+def test_validation_does_not_allow_union_supersets():
+    """Ensure that an implementing module cannot take a subset of param types of the task."""
+
+    @task(
+        unary_parameters={"sample_input": Union[str, int, bytes]},
+        unary_output_type=SampleOutputType,
+        streaming_output_type=Iterable[SampleOutputType],
+    )
+    class SomeTask(TaskBase):
+        pass
+
+    # If the task says bytes are okay, the module needs to be able to handle bytes also
+    with pytest.raises(TypeError):
+
+        @caikit.core.module(
+            id=str(uuid.uuid4()),
+            name="SomeModule",
+            version="0.0.1",
+            task=SomeTask,
+        )
+        class SomeModule(caikit.core.ModuleBase):
+            def run(self, sample_input: Union[str, int]) -> SampleOutputType:
+                pass
+
+
 # ----------- BACKWARDS COMPATIBILITY ------------------------------------------- ##