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)