Merge pull request #1680 from basetenlabs/deepakn/BT-14764-support-partial-annotations

BT-14764: Truss should be permissive if either input or output doesn't match Pydantic expectations

Author: ndeepak-baseten

truss/templates/server/common/schema.py

@@ -24,28 +24,30 @@ class OutputType(BaseModel):
 
 
 class TrussSchema(BaseModel):
-    input_type: Type[BaseModel]
+    input_type: Optional[Type[BaseModel]]
     output_type: Optional[Type[BaseModel]]
-    supports_streaming: bool
+    supports_streaming: Optional[bool]
 
     @classmethod
     def from_signature(
-        cls, input_parameters: MappingProxyType, output_annotation: Any
+        cls, input_parameters: Optional[MappingProxyType], output_annotation: Any
     ) -> Optional["TrussSchema"]:
         """
         Create a TrussSchema from a function signature if annotated, else returns None
         """
 
-        input_type = _parse_input_type(input_parameters)
-        output_type = _parse_output_type(output_annotation)
+        input_type = _parse_input_type(input_parameters) if input_parameters else None
+        output_type = (
+            _parse_output_type(output_annotation) if output_annotation else None
+        )
 
-        if not input_type or not output_type:
+        if not input_type and not output_type:
             return None
 
         return cls(
             input_type=input_type,
-            output_type=output_type.type,
-            supports_streaming=output_type.supports_streaming,
+            output_type=output_type.type if output_type else None,
+            supports_streaming=output_type.supports_streaming if output_type else None,
         )
 
     def serialize(self) -> dict:
@@ -54,7 +56,7 @@ def serialize(self) -> dict:
         generating an OpenAPI spec for this Truss.
         """
         return {
-            "input_schema": self.input_type.schema(),
+            "input_schema": self.input_type.schema() if self.input_type else None,
             "output_schema": self.output_type.schema()
             if self.output_type is not None
             else None,

truss/templates/server/truss_server.py

@@ -145,7 +145,7 @@ async def _parse_body(
         if self.is_binary(request):
             with tracing.section_as_event(span, "binary-deserialize"):
                 inputs = serialization.truss_msgpack_deserialize(body_raw)
-            if truss_schema:
+            if truss_schema and truss_schema.input_type:
                 try:
                     with tracing.section_as_event(span, "parse-pydantic"):
                         inputs = truss_schema.input_type.parse_obj(inputs)
@@ -154,7 +154,7 @@ async def _parse_body(
                         errors.format_pydantic_validation_error(e)
                     ) from e
         else:
-            if truss_schema:
+            if truss_schema and truss_schema.input_type:
                 try:
                     with tracing.section_as_event(span, "parse-pydantic"):
                         inputs = truss_schema.input_type.parse_raw(body_raw)

truss/tests/templates/server/test_schema.py

@@ -47,6 +47,53 @@ def predict(self, request: ModelInput) -> ModelOutput:
     assert not schema.supports_streaming
 
 
+def test_truss_schema_pydantic_empty_input():
+    class Model:
+        def predict(self) -> ModelOutput:
+            return ModelOutput(output="hello")
+
+    model = Model()
+
+    input_signature = inspect.signature(model.predict).parameters
+    output_signature = inspect.signature(model.predict).return_annotation
+
+    schema = TrussSchema.from_signature(input_signature, output_signature)
+
+    assert schema.input_type is None
+    assert schema.output_type == ModelOutput
+
+
+def test_truss_schema_pydantic_empty_output():
+    class Model:
+        def predict(self, _: ModelInput) -> None:
+            return None
+
+    model = Model()
+
+    input_signature = inspect.signature(model.predict).parameters
+    output_signature = inspect.signature(model.predict).return_annotation
+
+    schema = TrussSchema.from_signature(input_signature, output_signature)
+
+    assert schema.input_type == ModelInput
+    assert schema.output_type is None
+
+
+def test_truss_schema_pydantic_empty_input_and_output():
+    class Model:
+        def predict(self) -> None:
+            return None
+
+    model = Model()
+
+    input_signature = inspect.signature(model.predict).parameters
+    output_signature = inspect.signature(model.predict).return_annotation
+
+    schema = TrussSchema.from_signature(input_signature, output_signature)
+
+    assert schema is None
+
+
 def test_truss_schema_non_pydantic_input():
     class Model:
         def predict(self, request: str) -> ModelOutput:
@@ -59,7 +106,8 @@ def predict(self, request: str) -> ModelOutput:
 
     schema = TrussSchema.from_signature(input_signature, output_signature)
 
-    assert schema is None
+    assert schema.input_type is None
+    assert schema.output_type == ModelOutput
 
 
 def test_truss_schema_non_pydantic_output():
@@ -74,7 +122,8 @@ def predict(self, request: ModelInput) -> str:
 
     schema = TrussSchema.from_signature(input_signature, output_signature)
 
-    assert schema is None
+    assert schema.input_type == ModelInput
+    assert schema.output_type is None
 
 
 def test_truss_schema_list_types():
@@ -218,7 +267,8 @@ async def predict(
     output_signature = inspect.signature(model.predict).return_annotation
 
     schema = TrussSchema.from_signature(input_signature, output_signature)
-    assert schema is None
+    assert schema.input_type == ModelInput
+    assert schema.output_type is None
 
 
 def test_truss_schema_union_non_pydantic():
@@ -233,7 +283,8 @@ def predict(self, request: ModelInput) -> Union[str, int]:
 
     schema = TrussSchema.from_signature(input_signature, output_signature)
 
-    assert schema is None
+    assert schema.input_type == ModelInput
+    assert schema.output_type is None
 
 
 def test_truss_schema_async_non_pydantic():
@@ -269,4 +320,5 @@ def predict(
 
     schema = TrussSchema.from_signature(input_signature, output_signature)
 
-    assert schema is None
+    assert schema.input_type == ModelInput
+    assert schema.output_type is None