Merge branch 'zwei' into add-json-deserializer

Author: Balaji Veeramani

MANIFEST.in

@@ -1,4 +1,6 @@
-recursive-include src/sagemaker *
+recursive-include src/sagemaker *.py
+
+include src/sagemaker/image_uri_config/*.json
 
 include VERSION
 include LICENSE.txt

doc/frameworks/tensorflow/upgrade_from_legacy.rst

@@ -245,11 +245,10 @@ For example, if you want to use JSON serialization and deserialization:
 
 .. code:: python
 
-    from sagemaker.predictor import json_serializer
     from sagemaker.deserializers import JSONDeserializer
+    from sagemaker.serializers import JSONSerializer
 
-    predictor.content_type = "application/json"
-    predictor.serializer = json_serializer
+    predictor.serializer = JSONSerializer()
     predictor.deserializer = JSONDeserializer()
 
     predictor.predict(data)

setup.py

@@ -83,6 +83,7 @@ def read_version():
     packages=find_packages("src"),
     package_dir={"": "src"},
     py_modules=[os.path.splitext(os.path.basename(path))[0] for path in glob("src/*.py")],
+    include_package_data=True,
     long_description=read("README.rst"),
     author="Amazon Web Services",
     url="https://github.com/aws/sagemaker-python-sdk/",

src/sagemaker/chainer/model.py

@@ -25,7 +25,8 @@
 from sagemaker.model import FrameworkModel, MODEL_SERVER_WORKERS_PARAM_NAME
 from sagemaker.chainer import defaults
 from sagemaker.deserializers import NumpyDeserializer
-from sagemaker.predictor import Predictor, npy_serializer
+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, npy_serializer, NumpyDeserializer()
+            endpoint_name, sagemaker_session, NumpySerializer(), NumpyDeserializer()
         )
 
 

src/sagemaker/cli/compatibility/v2/ast_transformer.py

@@ -34,13 +34,15 @@
     modifiers.renamed_params.SessionCreateModelImageURIRenamer(),
     modifiers.renamed_params.SessionCreateEndpointImageURIRenamer(),
     modifiers.training_params.TrainPrefixRemover(),
+    modifiers.training_input.TrainingInputConstructorRefactor(),
 ]
 
 IMPORT_MODIFIERS = [modifiers.tfs.TensorFlowServingImportRenamer()]
 
 IMPORT_FROM_MODIFIERS = [
     modifiers.predictors.PredictorImportFromRenamer(),
     modifiers.tfs.TensorFlowServingImportFromRenamer(),
+    modifiers.training_input.TrainingInputImportFromRenamer(),
 ]
 
 

src/sagemaker/cli/compatibility/v2/modifiers/__init__.py

@@ -22,4 +22,5 @@
     tf_legacy_mode,
     tfs,
     training_params,
+    training_input,
 )

src/sagemaker/cli/compatibility/v2/modifiers/framework_version.py

@@ -15,6 +15,8 @@
 
 import ast
 
+from packaging.version import InvalidVersion, Version
+
 from sagemaker.cli.compatibility.v2.modifiers import matching, parsing
 from sagemaker.cli.compatibility.v2.modifiers.modifier import Modifier
 
@@ -135,10 +137,15 @@ def _tf_py_version_default(framework_version):
     """Gets the py_version default based on framework_version for TensorFlow."""
     if not framework_version:
         return "py2"
-    version = [int(s) for s in framework_version.split(".")]
-    if version < [1, 12]:
+
+    try:
+        version = Version(framework_version)
+    except InvalidVersion:
+        return "py2"
+
+    if version < Version("1.12"):
         return "py2"
-    if version < [2, 2]:
+    if version < Version("2.2"):
         return "py3"
     return "py37"
 
@@ -186,7 +193,6 @@ def _version_args_needed(node):
     framework, is_model = _framework_from_node(node)
     expecting_py_version = _py_version_defaults(framework, framework_version, is_model)
     if expecting_py_version:
-        py_version = parsing.arg_value(node, PY_ARG)
-        return py_version is None
+        return not matching.has_arg(node, PY_ARG)
 
     return False

src/sagemaker/cli/compatibility/v2/modifiers/tf_legacy_mode.py

@@ -79,9 +79,11 @@ def _is_legacy_mode(self, node):
 
         for kw in node.keywords:
             if kw.arg == "script_mode":
-                script_mode = bool(kw.value.value)
+                script_mode = (
+                    bool(kw.value.value) if isinstance(kw.value, ast.NameConstant) else True
+                )
             if kw.arg == "py_version":
-                py_version = kw.value.s
+                py_version = kw.value.s if isinstance(kw.value, ast.Str) else "py3"
 
         return not (py_version.startswith("py3") or script_mode)
 
@@ -124,7 +126,8 @@ def modify_node(self, node):
 
         if add_image_uri:
             image_uri = self._image_uri_from_args(node.keywords)
-            node.keywords.append(ast.keyword(arg="image_uri", value=ast.Str(s=image_uri)))
+            if image_uri:
+                node.keywords.append(ast.keyword(arg="image_uri", value=ast.Str(s=image_uri)))
 
         node.keywords.append(ast.keyword(arg="model_dir", value=ast.NameConstant(value=False)))
 
@@ -155,19 +158,22 @@ def _to_ast_keyword(self, hps):
         return None
 
     def _image_uri_from_args(self, keywords):
-        """Returns a legacy TensorFlow image URI based on the estimator arguments."""
+        """Returns a legacy TensorFlow image URI based on the estimator arguments if possible."""
         tf_version = framework_version.FRAMEWORK_DEFAULTS["TensorFlow"]
         instance_type = "ml.m4.xlarge"  # CPU default (exact type doesn't matter)
 
         for kw in keywords:
             if kw.arg == "framework_version":
-                tf_version = kw.value.s
+                tf_version = kw.value.s if isinstance(kw.value, ast.Str) else None
             if kw.arg == "train_instance_type":
-                instance_type = kw.value.s
+                instance_type = kw.value.s if isinstance(kw.value, ast.Str) else None
 
-        return fw_utils.create_image_uri(
-            self.region, "tensorflow", instance_type, tf_version, "py2"
-        )
+        if tf_version and instance_type:
+            return fw_utils.create_image_uri(
+                self.region, "tensorflow", instance_type, tf_version, "py2"
+            )
+
+        return None
 
 
 class TensorBoardParameterRemover(Modifier):

src/sagemaker/cli/compatibility/v2/modifiers/training_input.py

@@ -0,0 +1,97 @@
+# Copyright 2020 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"). You
+# may not use this file except in compliance with the License. A copy of
+# the License is located at
+#
+#     http://aws.amazon.com/apache2.0/
+#
+# or in the "license" file accompanying this file. This file is
+# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
+# ANY KIND, either express or implied. See the License for the specific
+# language governing permissions and limitations under the License.
+"""Classes to modify TrainingInput code to be compatible
+with version 2.0 and later of the SageMaker Python SDK.
+"""
+from __future__ import absolute_import
+
+import ast
+
+from sagemaker.cli.compatibility.v2.modifiers import matching
+from sagemaker.cli.compatibility.v2.modifiers.modifier import Modifier
+
+S3_INPUT_NAME = "s3_input"
+S3_INPUT_NAMESPACES = ("sagemaker", "sagemaker.inputs", "sagemaker.session")
+
+
+class TrainingInputConstructorRefactor(Modifier):
+    """A class to refactor *s3_input class."""
+
+    def node_should_be_modified(self, node):
+        """Checks if the ``ast.Call`` node instantiates a class of interest.
+
+        This looks for the following calls:
+
+        - ``sagemaker.s3_input``
+        - ``sagemaker.session.s3_input``
+        - ``s3_input``
+
+        Args:
+            node (ast.Call): a node that represents a function call. For more,
+                see https://docs.python.org/3/library/ast.html#abstract-grammar.
+
+        Returns:
+            bool: If the ``ast.Call`` instantiates a class of interest.
+        """
+        return matching.matches_name_or_namespaces(node, S3_INPUT_NAME, S3_INPUT_NAMESPACES)
+
+    def modify_node(self, node):
+        """Modifies the ``ast.Call`` node to call ``TrainingInput`` instead.
+
+        Args:
+            node (ast.Call): a node that represents a *TrainingInput constructor.
+        """
+        if matching.matches_name(node, S3_INPUT_NAME):
+            node.func.id = "TrainingInput"
+        elif matching.matches_attr(node, S3_INPUT_NAME):
+            node.func.attr = "TrainingInput"
+            _rename_namespace(node, "session")
+
+
+def _rename_namespace(node, name):
+    """Rename namespace ``session`` to ``inputs`` """
+    if isinstance(node.func.value, ast.Attribute) and node.func.value.attr == name:
+        node.func.value.attr = "inputs"
+    elif isinstance(node.func.value, ast.Name) and node.func.value.id == name:
+        node.func.value.id = "inputs"
+
+
+class TrainingInputImportFromRenamer(Modifier):
+    """A class to update import statements of ``s3_input``."""
+
+    def node_should_be_modified(self, node):
+        """Checks if the import statement imports ``s3_input`` from the correct module.
+
+        Args:
+            node (ast.ImportFrom): a node that represents a ``from ... import ... `` statement.
+                For more, see https://docs.python.org/3/library/ast.html#abstract-grammar.
+
+        Returns:
+            bool: If the import statement imports ``s3_input`` from the correct module.
+        """
+        return node.module in S3_INPUT_NAMESPACES and any(
+            name.name == S3_INPUT_NAME for name in node.names
+        )
+
+    def modify_node(self, node):
+        """Changes the ``ast.ImportFrom`` node's name from ``s3_input`` to ``TrainingInput``.
+
+        Args:
+            node (ast.ImportFrom): a node that represents a ``from ... import ... `` statement.
+                For more, see https://docs.python.org/3/library/ast.html#abstract-grammar.
+        """
+        for name in node.names:
+            if name.name == S3_INPUT_NAME:
+                name.name = "TrainingInput"
+            if node.module == "sagemaker.session":
+                node.module = "sagemaker.inputs"

src/sagemaker/deserializers.py

@@ -13,6 +13,8 @@
 """Implements methods for deserializing data returned from an inference endpoint."""
 from __future__ import absolute_import
 
+import csv
+
 import abc
 import codecs
 import io
@@ -96,6 +98,37 @@ def deserialize(self, data, content_type):
             data.close()
 
 
+class CSVDeserializer(BaseDeserializer):
+    """Deserialize a stream of bytes into a list of lists."""
+
+    ACCEPT = "text/csv"
+
+    def __init__(self, encoding="utf-8"):
+        """Initialize the string encoding.
+
+        Args:
+            encoding (str): The string encoding to use (default: "utf-8").
+        """
+        self.encoding = encoding
+
+    def deserialize(self, data, content_type):
+        """Deserialize data from an inference endpoint into a list of lists.
+
+        Args:
+            data (botocore.response.StreamingBody): Data to be deserialized.
+            content_type (str): The MIME type of the data.
+
+        Returns:
+            list: The data deserialized into a list of lists representing the
+                contents of a CSV file.
+        """
+        try:
+            decoded_string = data.read().decode(self.encoding)
+            return list(csv.reader(decoded_string.splitlines()))
+        finally:
+            data.close()
+
+
 class StreamDeserializer(BaseDeserializer):
     """Returns the data and content-type received from an inference endpoint.