diff --git a/sagemaker-serve/src/sagemaker/serve/builder/schema_builder.py b/sagemaker-serve/src/sagemaker/serve/builder/schema_builder.py index faa8066d52..5505e4b90f 100644 --- a/sagemaker-serve/src/sagemaker/serve/builder/schema_builder.py +++ b/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): diff --git a/sagemaker-serve/tests/unit/builder/test_schema_builder.py b/sagemaker-serve/tests/unit/builder/test_schema_builder.py index ce812bb6d4..b2d557acc6 100644 --- a/sagemaker-serve/tests/unit/builder/test_schema_builder.py +++ b/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()