change: Support multiple Accept types

src/sagemaker/deserializers.py

@@ -52,13 +52,13 @@ def deserialize(self, stream, content_type):
     @property
     @abc.abstractmethod
     def ACCEPT(self):
-        """The content type that is expected from the inference endpoint."""
+        """The content types that are expected from the inference endpoint."""
 
 
 class StringDeserializer(BaseDeserializer):
     """Deserialize data from an inference endpoint into a decoded string."""
 
-    ACCEPT = "application/json"
+    ACCEPT = ["application/json", "text/csv"]
 
     def __init__(self, encoding="UTF-8"):
         """Initialize the string encoding.
@@ -87,7 +87,7 @@ def deserialize(self, stream, content_type):
 class BytesDeserializer(BaseDeserializer):
     """Deserialize a stream of bytes into a bytes object."""
 
-    ACCEPT = "*/*"
+    ACCEPT = ["*/*"]
 
     def deserialize(self, stream, content_type):
         """Read a stream of bytes returned from an inference endpoint.
@@ -108,7 +108,7 @@ def deserialize(self, stream, content_type):
 class CSVDeserializer(BaseDeserializer):
     """Deserialize a stream of bytes into a list of lists."""
 
-    ACCEPT = "text/csv"
+    ACCEPT = ["text/csv"]
 
     def __init__(self, encoding="utf-8"):
         """Initialize the string encoding.
@@ -143,7 +143,7 @@ class StreamDeserializer(BaseDeserializer):
     reading it.
     """
 
-    ACCEPT = "*/*"
+    ACCEPT = ["*/*"]
 
     def deserialize(self, stream, content_type):
         """Returns a stream of the response body and the MIME type of the data.
@@ -161,7 +161,7 @@ def deserialize(self, stream, content_type):
 class NumpyDeserializer(BaseDeserializer):
     """Deserialize a stream of data in the .npy format."""
 
-    ACCEPT = "application/x-npy"
+    ACCEPT = ["application/x-npy", "text/csv", "application/json"]
 
     def __init__(self, dtype=None, allow_pickle=True):
         """Initialize the dtype and allow_pickle arguments.
@@ -201,7 +201,7 @@ def deserialize(self, stream, content_type):
 class JSONDeserializer(BaseDeserializer):
     """Deserialize JSON data from an inference endpoint into a Python object."""
 
-    ACCEPT = "application/json"
+    ACCEPT = ["application/json"]
 
     def deserialize(self, stream, content_type):
         """Deserialize JSON data from an inference endpoint into a Python object.
@@ -222,7 +222,7 @@ def deserialize(self, stream, content_type):
 class PandasDeserializer(BaseDeserializer):
     """Deserialize CSV or JSON data from an inference endpoint into a pandas dataframe."""
 
-    ACCEPT = "text/csv"
+    ACCEPT = ["text/csv", "application/json"]
 
     def deserialize(self, stream, content_type):
         """Deserialize CSV or JSON data from an inference endpoint into a pandas
@@ -250,7 +250,7 @@ def deserialize(self, stream, content_type):
 class JSONLinesDeserializer(BaseDeserializer):
     """Deserialize JSON lines data from an inference endpoint."""
 
-    ACCEPT = "application/jsonlines"
+    ACCEPT = ["application/jsonlines"]
 
     def deserialize(self, stream, content_type):
         """Deserialize JSON lines data from an inference endpoint.

src/sagemaker/predictor.py

@@ -121,7 +121,7 @@ def _create_request_args(self, data, initial_args=None, target_model=None, targe
             args["ContentType"] = self.content_type
 
         if "Accept" not in args:
-            args["Accept"] = self.accept
+            args["Accept"] = ", ".join(self.accept)
 
         if target_model:
             args["TargetModel"] = target_model

tests/unit/test_predictor.py

@@ -17,6 +17,7 @@
 import pytest
 from mock import Mock, call, patch
 
+from sagemaker.deserializers import CSVDeserializer, StringDeserializer
 from sagemaker.predictor import Predictor
 from sagemaker.serializers import JSONSerializer, CSVSerializer
 
@@ -132,7 +133,7 @@ def json_sagemaker_session():
     response_body.close = Mock("close", return_value=None)
     ims.sagemaker_runtime_client.invoke_endpoint = Mock(
         name="invoke_endpoint",
-        return_value={"Body": response_body, "ContentType": DEFAULT_CONTENT_TYPE},
+        return_value={"Body": response_body, "ContentType": "application/json"},
     )
     return ims
 
@@ -169,7 +170,7 @@ def ret_csv_sagemaker_session():
     ims.sagemaker_client.describe_endpoint_config = Mock(return_value=ENDPOINT_CONFIG_DESC)
 
     response_body = Mock("body")
-    response_body.read = Mock("read", return_value=CSV_RETURN_VALUE)
+    response_body.read = Mock("read", return_value=bytes(CSV_RETURN_VALUE, "utf-8"))
     response_body.close = Mock("close", return_value=None)
     ims.sagemaker_runtime_client.invoke_endpoint = Mock(
         name="invoke_endpoint",
@@ -180,23 +181,48 @@ def ret_csv_sagemaker_session():
 
 def test_predict_call_with_csv():
     sagemaker_session = ret_csv_sagemaker_session()
-    predictor = Predictor(ENDPOINT, sagemaker_session, serializer=CSVSerializer())
+    predictor = Predictor(
+        ENDPOINT, sagemaker_session, serializer=CSVSerializer(), deserializer=CSVDeserializer()
+    )
 
     data = [1, 2]
     result = predictor.predict(data)
 
     assert sagemaker_session.sagemaker_runtime_client.invoke_endpoint.called
 
     expected_request_args = {
-        "Accept": DEFAULT_ACCEPT,
+        "Accept": CSV_CONTENT_TYPE,
+        "Body": "1,2",
+        "ContentType": CSV_CONTENT_TYPE,
+        "EndpointName": ENDPOINT,
+    }
+    call_args, kwargs = sagemaker_session.sagemaker_runtime_client.invoke_endpoint.call_args
+    assert kwargs == expected_request_args
+
+    assert result == [["1", "2", "3"]]
+
+
+def test_predict_call_with_multiple_accept_types():
+    sagemaker_session = ret_csv_sagemaker_session()
+    predictor = Predictor(
+        ENDPOINT, sagemaker_session, serializer=CSVSerializer(), deserializer=StringDeserializer()
+    )
+
+    data = [1, 2]
+    result = predictor.predict(data)
+
+    assert sagemaker_session.sagemaker_runtime_client.invoke_endpoint.called
+
+    expected_request_args = {
+        "Accept": "application/json, text/csv",
         "Body": "1,2",
         "ContentType": CSV_CONTENT_TYPE,
         "EndpointName": ENDPOINT,
     }
     call_args, kwargs = sagemaker_session.sagemaker_runtime_client.invoke_endpoint.call_args
     assert kwargs == expected_request_args
 
-    assert result == CSV_RETURN_VALUE
+    assert result == "1,2,3\r\n"
 
 
 @patch("sagemaker.predictor.name_from_base")