From 2407cf08e510d0f2d950972eed535d7ecd6675c8 Mon Sep 17 00:00:00 2001 From: Lucas Jia Date: Tue, 13 Oct 2026 13:02:17 -0700 Subject: [PATCH] fix: honor request content types Pass caller-provided content types to request deserializers so JSON and CSV NumPy payloads are decoded correctly. Keep response deserializers pinned to configured Accept types for model servers that return misleading headers. Fixes #5006 --- src/sagemaker/serve/builder/schema_builder.py | 20 +++++++----- .../serve/builder/test_schema_builder.py | 32 +++++++++++++++++++ 2 files changed, 44 insertions(+), 8 deletions(-) diff --git a/src/sagemaker/serve/builder/schema_builder.py b/src/sagemaker/serve/builder/schema_builder.py index 7f70e98747..bb5a228182 100644 --- a/src/sagemaker/serve/builder/schema_builder.py +++ b/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): diff --git a/tests/unit/sagemaker/serve/builder/test_schema_builder.py b/tests/unit/sagemaker/serve/builder/test_schema_builder.py index cff2f9e6de..bf35aa05c1 100644 --- a/tests/unit/sagemaker/serve/builder/test_schema_builder.py +++ b/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)