Merge branch 'zwei' into add-numpy-serializer

bveeramani · web-flow · commit 7c1762c0cf68 · 2020-07-09T18:28:37.000-05:00
diff --git a/src/sagemaker/amazon/common.py b/src/sagemaker/amazon/common.py
@@ -22,31 +22,37 @@
 
 from sagemaker.amazon.record_pb2 import Record
 from sagemaker.deserializers import BaseDeserializer
+from sagemaker.serializers import BaseSerializer
 from sagemaker.utils import DeferredError
 
 
-class numpy_to_record_serializer(object):
-    """Placeholder docstring"""
+class RecordSerializer(BaseSerializer):
+    """Serialize a NumPy array for an inference request."""
 
-    def __init__(self, content_type="application/x-recordio-protobuf"):
-        """
-        Args:
-            content_type:
-        """
-        self.content_type = content_type
+    CONTENT_TYPE = "application/x-recordio-protobuf"
+
+    def serialize(self, data):
+        """Serialize a NumPy array into a buffer containing RecordIO records.
 
-    def __call__(self, array):
-        """
         Args:
-            array:
+            data (numpy.ndarray): The data to serialize.
+
+        Returns:
+            io.BytesIO: A buffer containing the data serialized as records.
         """
-        if len(array.shape) == 1:
-            array = array.reshape(1, array.shape[0])
-        assert len(array.shape) == 2, "Expecting a 1 or 2 dimensional array"
-        buf = io.BytesIO()
-        write_numpy_to_dense_tensor(buf, array)
-        buf.seek(0)
-        return buf
+        if len(data.shape) == 1:
+            data = data.reshape(1, data.shape[0])
+
+        if len(data.shape) != 2:
+            raise ValueError(
+                "Expected a 1D or 2D array, but got a %dD array instead." % len(data.shape)
+            )
+
+        buffer = io.BytesIO()
+        write_numpy_to_dense_tensor(buffer, data)
+        buffer.seek(0)
+
+        return buffer
 
 
 class RecordDeserializer(BaseDeserializer):
diff --git a/src/sagemaker/amazon/factorization_machines.py b/src/sagemaker/amazon/factorization_machines.py
@@ -14,7 +14,7 @@
 from __future__ import absolute_import
 
 from sagemaker.amazon.amazon_estimator import AmazonAlgorithmEstimatorBase, registry
-from sagemaker.amazon.common import numpy_to_record_serializer, RecordDeserializer
+from sagemaker.amazon.common import RecordSerializer, RecordDeserializer
 from sagemaker.amazon.hyperparameter import Hyperparameter as hp  # noqa
 from sagemaker.amazon.validation import gt, isin, ge
 from sagemaker.predictor import Predictor
@@ -289,7 +289,7 @@ def __init__(self, endpoint_name, sagemaker_session=None):
         super(FactorizationMachinesPredictor, self).__init__(
             endpoint_name,
             sagemaker_session,
-            serializer=numpy_to_record_serializer(),
+            serializer=RecordSerializer(),
             deserializer=RecordDeserializer(),
         )
 
diff --git a/src/sagemaker/amazon/kmeans.py b/src/sagemaker/amazon/kmeans.py
@@ -14,7 +14,7 @@
 from __future__ import absolute_import
 
 from sagemaker.amazon.amazon_estimator import AmazonAlgorithmEstimatorBase, registry
-from sagemaker.amazon.common import numpy_to_record_serializer, RecordDeserializer
+from sagemaker.amazon.common import RecordSerializer, RecordDeserializer
 from sagemaker.amazon.hyperparameter import Hyperparameter as hp  # noqa
 from sagemaker.amazon.validation import gt, isin, ge, le
 from sagemaker.predictor import Predictor
@@ -222,7 +222,7 @@ def __init__(self, endpoint_name, sagemaker_session=None):
         super(KMeansPredictor, self).__init__(
             endpoint_name,
             sagemaker_session,
-            serializer=numpy_to_record_serializer(),
+            serializer=RecordSerializer(),
             deserializer=RecordDeserializer(),
         )
 
diff --git a/src/sagemaker/amazon/knn.py b/src/sagemaker/amazon/knn.py
@@ -14,7 +14,7 @@
 from __future__ import absolute_import
 
 from sagemaker.amazon.amazon_estimator import AmazonAlgorithmEstimatorBase, registry
-from sagemaker.amazon.common import numpy_to_record_serializer, RecordDeserializer
+from sagemaker.amazon.common import RecordSerializer, RecordDeserializer
 from sagemaker.amazon.hyperparameter import Hyperparameter as hp  # noqa
 from sagemaker.amazon.validation import ge, isin
 from sagemaker.predictor import Predictor
@@ -210,7 +210,7 @@ def __init__(self, endpoint_name, sagemaker_session=None):
         super(KNNPredictor, self).__init__(
             endpoint_name,
             sagemaker_session,
-            serializer=numpy_to_record_serializer(),
+            serializer=RecordSerializer(),
             deserializer=RecordDeserializer(),
         )
 
diff --git a/src/sagemaker/amazon/lda.py b/src/sagemaker/amazon/lda.py
@@ -14,7 +14,7 @@
 from __future__ import absolute_import
 
 from sagemaker.amazon.amazon_estimator import AmazonAlgorithmEstimatorBase, registry
-from sagemaker.amazon.common import numpy_to_record_serializer, RecordDeserializer
+from sagemaker.amazon.common import RecordSerializer, RecordDeserializer
 from sagemaker.amazon.hyperparameter import Hyperparameter as hp  # noqa
 from sagemaker.amazon.validation import gt
 from sagemaker.predictor import Predictor
@@ -194,7 +194,7 @@ def __init__(self, endpoint_name, sagemaker_session=None):
         super(LDAPredictor, self).__init__(
             endpoint_name,
             sagemaker_session,
-            serializer=numpy_to_record_serializer(),
+            serializer=RecordSerializer(),
             deserializer=RecordDeserializer(),
         )
 
diff --git a/src/sagemaker/amazon/linear_learner.py b/src/sagemaker/amazon/linear_learner.py
@@ -14,7 +14,7 @@
 from __future__ import absolute_import
 
 from sagemaker.amazon.amazon_estimator import AmazonAlgorithmEstimatorBase, registry
-from sagemaker.amazon.common import numpy_to_record_serializer, RecordDeserializer
+from sagemaker.amazon.common import RecordSerializer, RecordDeserializer
 from sagemaker.amazon.hyperparameter import Hyperparameter as hp  # noqa
 from sagemaker.amazon.validation import isin, gt, lt, ge, le
 from sagemaker.predictor import Predictor
@@ -453,7 +453,7 @@ def __init__(self, endpoint_name, sagemaker_session=None):
         super(LinearLearnerPredictor, self).__init__(
             endpoint_name,
             sagemaker_session,
-            serializer=numpy_to_record_serializer(),
+            serializer=RecordSerializer(),
             deserializer=RecordDeserializer(),
         )
 
diff --git a/src/sagemaker/amazon/ntm.py b/src/sagemaker/amazon/ntm.py
@@ -14,7 +14,7 @@
 from __future__ import absolute_import
 
 from sagemaker.amazon.amazon_estimator import AmazonAlgorithmEstimatorBase, registry
-from sagemaker.amazon.common import numpy_to_record_serializer, RecordDeserializer
+from sagemaker.amazon.common import RecordSerializer, RecordDeserializer
 from sagemaker.amazon.hyperparameter import Hyperparameter as hp  # noqa
 from sagemaker.amazon.validation import ge, le, isin
 from sagemaker.predictor import Predictor
@@ -224,7 +224,7 @@ def __init__(self, endpoint_name, sagemaker_session=None):
         super(NTMPredictor, self).__init__(
             endpoint_name,
             sagemaker_session,
-            serializer=numpy_to_record_serializer(),
+            serializer=RecordSerializer(),
             deserializer=RecordDeserializer(),
         )
 
diff --git a/src/sagemaker/amazon/pca.py b/src/sagemaker/amazon/pca.py
@@ -14,7 +14,7 @@
 from __future__ import absolute_import
 
 from sagemaker.amazon.amazon_estimator import AmazonAlgorithmEstimatorBase, registry
-from sagemaker.amazon.common import numpy_to_record_serializer, RecordDeserializer
+from sagemaker.amazon.common import RecordSerializer, RecordDeserializer
 from sagemaker.amazon.hyperparameter import Hyperparameter as hp  # noqa
 from sagemaker.amazon.validation import gt, isin
 from sagemaker.predictor import Predictor
@@ -206,7 +206,7 @@ def __init__(self, endpoint_name, sagemaker_session=None):
         super(PCAPredictor, self).__init__(
             endpoint_name,
             sagemaker_session,
-            serializer=numpy_to_record_serializer(),
+            serializer=RecordSerializer(),
             deserializer=RecordDeserializer(),
         )
 
diff --git a/src/sagemaker/amazon/randomcutforest.py b/src/sagemaker/amazon/randomcutforest.py
@@ -14,7 +14,7 @@
 from __future__ import absolute_import
 
 from sagemaker.amazon.amazon_estimator import AmazonAlgorithmEstimatorBase, registry
-from sagemaker.amazon.common import numpy_to_record_serializer, RecordDeserializer
+from sagemaker.amazon.common import RecordSerializer, RecordDeserializer
 from sagemaker.amazon.hyperparameter import Hyperparameter as hp  # noqa
 from sagemaker.amazon.validation import ge, le
 from sagemaker.predictor import Predictor
@@ -183,7 +183,7 @@ def __init__(self, endpoint_name, sagemaker_session=None):
         super(RandomCutForestPredictor, self).__init__(
             endpoint_name,
             sagemaker_session,
-            serializer=numpy_to_record_serializer(),
+            serializer=RecordSerializer(),
             deserializer=RecordDeserializer(),
         )
 
diff --git a/src/sagemaker/chainer/model.py b/src/sagemaker/chainer/model.py
@@ -24,7 +24,8 @@
 )
 from sagemaker.model import FrameworkModel, MODEL_SERVER_WORKERS_PARAM_NAME
 from sagemaker.chainer import defaults
-from sagemaker.predictor import Predictor, numpy_deserializer
+from sagemaker.deserializers import NumpyDeserializer
+from sagemaker.predictor import Predictor
 from sagemaker.serializers import NumpySerializer
 
 logger = logging.getLogger("sagemaker")
@@ -49,7 +50,7 @@ def __init__(self, endpoint_name, sagemaker_session=None):
                 using the default AWS configuration chain.
         """
         super(ChainerPredictor, self).__init__(
-            endpoint_name, sagemaker_session, NumpySerializer(), numpy_deserializer
+            endpoint_name, sagemaker_session, NumpySerializer(), NumpyDeserializer()
         )
 
 
diff --git a/src/sagemaker/deserializers.py b/src/sagemaker/deserializers.py
@@ -14,6 +14,11 @@
 from __future__ import absolute_import
 
 import abc
+import codecs
+import io
+import json
+
+import numpy as np
 
 
 class BaseDeserializer(abc.ABC):
@@ -111,3 +116,41 @@ def deserialize(self, data, content_type):
             tuple: A two-tuple containing the stream and content-type.
         """
         return data, content_type
+
+
+class NumpyDeserializer(BaseDeserializer):
+    """Deserialize a stream of data in the .npy format."""
+
+    ACCEPT = "application/x-npy"
+
+    def __init__(self, dtype=None):
+        """Initialize the dtype.
+
+        Args:
+            dtype (str): The dtype of the data.
+        """
+        self.dtype = dtype
+
+    def deserialize(self, data, content_type):
+        """Deserialize data from an inference endpoint into a NumPy array.
+
+        Args:
+            data (botocore.response.StreamingBody): Data to be deserialized.
+            content_type (str): The MIME type of the data.
+
+        Returns:
+            numpy.ndarray: The data deserialized into a NumPy array.
+        """
+        try:
+            if content_type == "text/csv":
+                return np.genfromtxt(
+                    codecs.getreader("utf-8")(data), delimiter=",", dtype=self.dtype
+                )
+            if content_type == "application/json":
+                return np.array(json.load(codecs.getreader("utf-8")(data)), dtype=self.dtype)
+            if content_type == "application/x-npy":
+                return np.load(io.BytesIO(data.read()))
+        finally:
+            data.close()
+
+        raise ValueError("%s cannot read content type %s." % (__class__.__name__, content_type))
diff --git a/src/sagemaker/predictor.py b/src/sagemaker/predictor.py
@@ -695,48 +695,4 @@ def __call__(self, stream, content_type):
             stream.close()
 
 
-json_deserializer = _JsonDeserializer()
-
-
-class _NumpyDeserializer(object):
-    """Placeholder docstring"""
-
-    def __init__(self, accept=CONTENT_TYPE_NPY, dtype=None):
-        """
-        Args:
-            accept:
-            dtype:
-        """
-        self.accept = accept
-        self.dtype = dtype
-
-    def __call__(self, stream, content_type=CONTENT_TYPE_NPY):
-        """Decode from serialized data into a Numpy array.
-
-        Args:
-            stream (stream): The response stream to be deserialized.
-            content_type (str): The content type of the response. Can accept
-                CSV, JSON, or NPY data.
-
-        Returns:
-            object: Body of the response deserialized into a Numpy array.
-        """
-        try:
-            if content_type == CONTENT_TYPE_CSV:
-                return np.genfromtxt(
-                    codecs.getreader("utf-8")(stream), delimiter=",", dtype=self.dtype
-                )
-            if content_type == CONTENT_TYPE_JSON:
-                return np.array(json.load(codecs.getreader("utf-8")(stream)), dtype=self.dtype)
-            if content_type == CONTENT_TYPE_NPY:
-                return np.load(BytesIO(stream.read()))
-        finally:
-            stream.close()
-        raise ValueError(
-            "content_type must be one of the following: CSV, JSON, NPY. content_type: {}".format(
-                content_type
-            )
-        )
-
-
-numpy_deserializer = _NumpyDeserializer()
+json_deserializer = _JsonDeserializer()
diff --git a/src/sagemaker/pytorch/model.py b/src/sagemaker/pytorch/model.py
@@ -17,6 +17,7 @@
 import packaging.version
 
 import sagemaker
+from sagemaker.deserializers import NumpyDeserializer
 from sagemaker.fw_utils import (
     create_image_uri,
     model_code_key_prefix,
@@ -25,7 +26,7 @@
 )
 from sagemaker.model import FrameworkModel, MODEL_SERVER_WORKERS_PARAM_NAME
 from sagemaker.pytorch import defaults
-from sagemaker.predictor import Predictor, numpy_deserializer
+from sagemaker.predictor import Predictor
 from sagemaker.serializers import NumpySerializer
 
 logger = logging.getLogger("sagemaker")
@@ -50,7 +51,7 @@ def __init__(self, endpoint_name, sagemaker_session=None):
                 using the default AWS configuration chain.
         """
         super(PyTorchPredictor, self).__init__(
-            endpoint_name, sagemaker_session, NumpySerializer(), numpy_deserializer
+            endpoint_name, sagemaker_session, NumpySerializer(), NumpyDeserializer()
         )
 
 
diff --git a/src/sagemaker/sklearn/model.py b/src/sagemaker/sklearn/model.py
@@ -16,10 +16,11 @@
 import logging
 
 import sagemaker
+from sagemaker.deserializers import NumpyDeserializer
 from sagemaker.fw_registry import default_framework_uri
 from sagemaker.fw_utils import model_code_key_prefix, validate_version_or_image_args
 from sagemaker.model import FrameworkModel, MODEL_SERVER_WORKERS_PARAM_NAME
-from sagemaker.predictor import Predictor, numpy_deserializer
+from sagemaker.predictor import Predictor
 from sagemaker.serializers import NumpySerializer
 from sagemaker.sklearn import defaults
 
@@ -45,7 +46,7 @@ def __init__(self, endpoint_name, sagemaker_session=None):
                 using the default AWS configuration chain.
         """
         super(SKLearnPredictor, self).__init__(
-            endpoint_name, sagemaker_session, NumpySerializer(), numpy_deserializer
+            endpoint_name, sagemaker_session, NumpySerializer(), NumpyDeserializer()
         )
 
 
diff --git a/tests/unit/sagemaker/test_deserializers.py b/tests/unit/sagemaker/test_deserializers.py
diff --git a/tests/unit/test_common.py b/tests/unit/test_common.py
diff --git a/tests/unit/test_predictor.py b/tests/unit/test_predictor.py
