Bug fix: #123

Author: mvsusp

setup.py

@@ -49,7 +49,7 @@ def read(fname):
 
       extras_require={
           'test': ['tox', 'flake8', 'pytest', 'pytest-cov', 'pytest-xdist',
-                   'mock', 'tensorflow>=1.3.0', 'contextlib2', 'awslogs', 'pandas']},
+                   'mock', 'tensorflow>=1.3.0', 'protobuf3-to-dict', 'contextlib2', 'awslogs', 'pandas']},
 
       entry_points={
           'console_scripts': ['sagemaker=sagemaker.cli.main:main'],

src/sagemaker/tensorflow/predictor.py

@@ -16,13 +16,18 @@
 
 import google.protobuf.json_format as json_format
 from google.protobuf.message import DecodeError
+from protobuf_to_dict import protobuf_to_dict
 from tensorflow.core.framework import tensor_pb2
 from tensorflow.python.framework import tensor_util
 
 from sagemaker.content_types import CONTENT_TYPE_JSON, CONTENT_TYPE_OCTET_STREAM, CONTENT_TYPE_CSV
 from sagemaker.predictor import json_serializer, csv_serializer
 from tensorflow_serving.apis import predict_pb2, classification_pb2, inference_pb2, regression_pb2
 
+_POSSIBLE_RESPONSES = [predict_pb2.PredictResponse, classification_pb2.ClassificationResponse,
+                       inference_pb2.MultiInferenceResponse, regression_pb2.RegressionResponse,
+                       tensor_pb2.TensorProto]
+
 REGRESSION_REQUEST = 'RegressionRequest'
 MULTI_INFERENCE_REQUEST = 'MultiInferenceRequest'
 CLASSIFICATION_REQUEST = 'ClassificationRequest'
@@ -53,17 +58,12 @@ def __init__(self):
         self.accept = CONTENT_TYPE_OCTET_STREAM
 
     def __call__(self, stream, content_type):
-        possible_responses = [predict_pb2.PredictResponse,
-                              classification_pb2.ClassificationResponse,
-                              inference_pb2.MultiInferenceResponse,
-                              regression_pb2.RegressionResponse]
-
         try:
             data = stream.read()
         finally:
             stream.close()
 
-        for possible_response in possible_responses:
+        for possible_response in _POSSIBLE_RESPONSES:
             try:
                 response = possible_response()
                 response.ParseFromString(data)
@@ -101,10 +101,15 @@ def __call__(self, stream, content_type):
             data = stream.read()
         finally:
             stream.close()
-        try:
-            return json_format.Parse(data, tensor_pb2.TensorProto())
-        except json_format.ParseError:
-            return json.loads(data.decode())
+
+        for possible_response in _POSSIBLE_RESPONSES:
+            try:
+                return protobuf_to_dict(json_format.Parse(data, possible_response()))
+            except (UnicodeDecodeError, DecodeError, json_format.ParseError):
+                # given that the payload does not have the response type, there no way to infer
+                # the response without keeping state, so I'm iterating all the options.
+                pass
+        return json.loads(data.decode())
 
 
 tf_json_deserializer = _TFJsonDeserializer()

tests/unit/test_tf_predictor.py

@@ -21,6 +21,7 @@
 from google.protobuf import json_format
 import tensorflow as tf
 from mock import Mock
+from six import StringIO
 from tensorflow.python.saved_model.signature_constants import DEFAULT_SERVING_SIGNATURE_DEF_KEY, PREDICT_INPUTS
 
 from sagemaker.predictor import RealTimePredictor
@@ -139,6 +140,72 @@ def test_classification_request_csv(sagemaker_session):
 """
 
 
+def test_json_deserializer_should_work_with_predict_response():
+    data = b"""{
+"outputs": {
+    "example_strings": {
+      "dtype": "DT_STRING", 
+      "tensorShape": {
+        "dim": [
+          {
+            "size": "3"
+          }
+        ]
+      }, 
+      "stringVal": [
+        "YXBwbGU=", 
+        "YmFuYW5h", 
+        "b3Jhbmdl"
+      ]
+    }, 
+    "ages": {
+      "dtype": "DT_FLOAT", 
+      "floatVal": [
+        4.954165935516357
+      ], 
+      "tensorShape": {
+        "dim": [
+          {
+            "size": "1"
+          }
+        ]
+      }
+    }
+  }, 
+  "modelSpec": {
+    "version": "1531758457", 
+    "name": "generic_model", 
+    "signatureName": "serving_default"
+  }
+}"""
+
+    stream = StringIO(data)
+
+    response = tf_json_deserializer(stream, 'application/json')
+
+    assert response == {
+        'model_spec': {
+            'name':           u'generic_model',
+            'signature_name': u'serving_default',
+            'version':        {'value': 1531758457L}
+        },
+        'outputs':    {
+            u'ages':            {
+                'dtype':        1,
+                'float_val':    [4.954165935516357],
+                'tensor_shape': {'dim': [{'size': 1L}]}
+            },
+            u'example_strings': {
+                'dtype':        7,
+                'string_val':   ['apple',
+                                 'banana',
+                                 'orange'],
+                'tensor_shape': {'dim': [{'size': 3L}]}
+            }
+        }
+    }
+
+
 def test_classification_request_pb(sagemaker_session):
     request = classification_pb2.ClassificationRequest()
     request.model_spec.name = "generic_model"