fix(serve): honor request content types

sagemaker-serve/src/sagemaker/serve/builder/schema_builder.py

@@ -69,18 +69,17 @@ def serialize(self, data) -> bytes:
 class DeserializerWrapper(BaseDeserializer):
     """Wraps the deserializer to comply with the function signature."""
 
-    def __init__(self, deserializer, accept):
+    def __init__(self, deserializer, accept, honor_content_type=False):
         self._deserializer = deserializer
         self._accept = accept
+        self._honor_content_type = honor_content_type
 
     def deserialize(self, stream, content_type: str = None):
         """Deserialize stream into object"""
-        return self._deserializer.deserialize(
-            stream,
-            # We need to overwrite the accept type because model
-            # servers like XGBOOST always returns "text/html"
-            self._accept[0],
+        deserialization_content_type = (
+            content_type if self._honor_content_type and content_type else self._accept[0]
         )
+        return self._deserializer.deserialize(stream, deserialization_content_type)
 
     @property
     def ACCEPT(self):
@@ -133,7 +132,9 @@ def __init__(
             self.input_serializer = self._get_serializer(sample_input)
             self._input_deserializer = self._get_inverse(self.input_serializer)
             self.input_deserializer = DeserializerWrapper(
-                self._input_deserializer, self._input_deserializer.ACCEPT
+                self._input_deserializer,
+                self._input_deserializer.ACCEPT,
+                honor_content_type=True,
             )
 
         if output_translator:
@@ -148,7 +149,10 @@ def __init__(
             self.output_serializer = self._get_inverse(self._output_deserializer)
 
             self.output_deserializer = DeserializerWrapper(
-                self._output_deserializer, self._output_deserializer.ACCEPT
+                self._output_deserializer,
+                self._output_deserializer.ACCEPT,
+                # Some model servers return misleading response types such as text/html.
+                honor_content_type=False,
             )
 
     def _get_serializer(self, obj):

sagemaker-serve/tests/unit/builder/test_schema_builder.py

@@ -1,4 +1,7 @@
+import json
 import unittest
+from io import BytesIO
+
 import numpy as np
 from sagemaker.serve.builder.schema_builder import SchemaBuilder
 
@@ -54,6 +57,33 @@ def test_generate_marshalling_map(self):
         self.assertIn("input_serializer", mapping)
         self.assertIn("output_deserializer", mapping)
 
+    def test_numpy_input_deserializer_honors_json_content_type(self):
+        schema = SchemaBuilder(np.array([[0]]), np.array([[0]]))
+        stream = BytesIO(json.dumps([[1, 2], [3, 4]]).encode("utf-8"))
+
+        result = schema.input_deserializer.deserialize(stream, content_type="application/json")
+
+        np.testing.assert_array_equal(result, np.array([[1, 2], [3, 4]]))
+
+    def test_numpy_input_deserializer_honors_csv_content_type(self):
+        schema = SchemaBuilder(np.array([[0]]), np.array([[0]]))
+        stream = BytesIO(b"1,2\n3,4\n")
+
+        result = schema.input_deserializer.deserialize(stream, content_type="text/csv")
+
+        np.testing.assert_array_equal(result, np.array([[1, 2], [3, 4]]))
+
+    def test_numpy_output_deserializer_preserves_expected_accept_type(self):
+        expected = np.array([[0.1, 0.9]])
+        schema = SchemaBuilder(np.array([[0]]), expected)
+        stream = BytesIO()
+        np.save(stream, expected)
+        stream.seek(0)
+
+        result = schema.output_deserializer.deserialize(stream, content_type="text/html")
+
+        np.testing.assert_array_equal(result, expected)
+
 
 if __name__ == "__main__":
     unittest.main()