fix: honor request content types

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):

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

@@ -12,6 +12,8 @@
 # language governing permissions and limitations under the License.
 from __future__ import absolute_import
 
+import json
+
 import pytest
 
 from io import BytesIO
@@ -178,3 +180,33 @@ def test_schema_builder_with_payload_converter_invalid_payload(custom_translator
             input_translator=custom_translator,
             output_translator=custom_translator,
         )
+
+
+def test_numpy_input_deserializer_honors_json_content_type():
+    schema_builder = SchemaBuilder(np.array([[0]]), np.array([[0]]))
+    stream = BytesIO(json.dumps([[1, 2], [3, 4]]).encode("utf-8"))
+
+    result = schema_builder.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():
+    schema_builder = SchemaBuilder(np.array([[0]]), np.array([[0]]))
+    stream = BytesIO(b"1,2\n3,4\n")
+
+    result = schema_builder.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():
+    expected = np.array([[0.1, 0.9]])
+    schema_builder = SchemaBuilder(np.array([[0]]), expected)
+    stream = BytesIO()
+    np.save(stream, expected)
+    stream.seek(0)
+
+    result = schema_builder.output_deserializer.deserialize(stream, content_type="text/html")
+
+    np.testing.assert_array_equal(result, expected)