V2 compability tests

tests/unit/v2/test_transformer.py

@@ -0,0 +1,38 @@
+from __future__ import absolute_import
+
+import ast
+from sagemaker.tools.compatibility.v2.ast_transformer import ASTTransformer
+import pasta
+
+
+def test_code_needs_transform():
+    simple = """
+TensorFlow(entry_point="foo.py")
+sagemaker.tensorflow.TensorFlow()
+m = MXNet()
+sagemaker.mxnet.MXNet()
+"""
+
+    transformer_class = ASTTransformer()
+    rewrite = transformer_class.visit(ast.parse(simple))
+    expected = """TensorFlow(entry_point='foo.py', framework_version='1.11.0')
+sagemaker.tensorflow.TensorFlow(framework_version='1.11.0')
+m = MXNet(framework_version='1.2.0')
+sagemaker.mxnet.MXNet(framework_version='1.2.0')\n"""
+
+    assert pasta.dump(rewrite) == expected
+
+
+def test_code_does_not_need_transform():
+    simple = """TensorFlow(entry_point='foo.py', framework_version='1.11.0')
+sagemaker.tensorflow.TensorFlow(framework_version='1.11.0')
+m = MXNet(framework_version='1.2.0')
+sagemaker.mxnet.MXNet(framework_version='1.2.0')\n"""
+    transformer_class = ASTTransformer()
+    rewrite = transformer_class.visit(ast.parse(simple))
+    expected = """TensorFlow(entry_point='foo.py', framework_version='1.11.0')
+sagemaker.tensorflow.TensorFlow(framework_version='1.11.0')
+m = MXNet(framework_version='1.2.0')
+sagemaker.mxnet.MXNet(framework_version='1.2.0')\n"""
+
+    assert pasta.dump(rewrite) == expected

tools/__init__.py

@@ -0,0 +1,14 @@
+# 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.
+"""Tools to assist with using the SageMake Python SDK."""
+from __future__ import absolute_import

tools/compatibility/__init__.py

@@ -0,0 +1,14 @@
+# 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.
+"""Tools to assist with compatibility between SageMaker Python SDK versions."""
+from __future__ import absolute_import

tools/compatibility/v2/__init__.py

@@ -0,0 +1,14 @@
+# 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.
+"""Tools to assist with upgrading to v2 of the SageMaker Python SDK."""
+from __future__ import absolute_import

tools/compatibility/v2/ast_transformer.py

@@ -0,0 +1,41 @@
+# 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.
+"""An ast.NodeTransformer subclass for updating SageMaker Python SDK code."""
+from __future__ import absolute_import
+
+import ast
+
+from tools.compatibility.v2.modifiers import framework_version
+
+FUNCTION_CALL_MODIFIERS = [framework_version.FrameworkVersionEnforcer()]
+
+
+class ASTTransformer(ast.NodeTransformer):
+    """An ``ast.NodeTransformer`` subclass that walks the abstract syntax tree and
+    modifies nodes to upgrade the given SageMaker Python SDK code.
+    """
+
+    def visit_Call(self, node):
+        """Visits an ``ast.Call`` node and returns a modified node, if needed.
+        See https://docs.python.org/3/library/ast.html#ast.NodeTransformer.
+
+        Args:
+            node (ast.Call): a node that represents a function call.
+
+        Returns:
+            ast.Call: a node that represents a function call, which has
+                potentially been modified from the original input.
+        """
+        for function_checker in FUNCTION_CALL_MODIFIERS:
+            function_checker.check_and_modify_node(node)
+        return node

tools/compatibility/v2/files.py

@@ -0,0 +1,180 @@
+# 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 for updating code in files."""
+from __future__ import absolute_import
+
+from abc import abstractmethod
+import json
+import logging
+import os
+
+import pasta
+
+from ast_transformer import ASTTransformer
+
+LOGGER = logging.getLogger(__name__)
+
+
+class FileUpdater(object):
+    """An abstract class for updating files."""
+
+    def __init__(self, input_path, output_path):
+        """Creates a ``FileUpdater`` for updating a file so that
+        it is compatible with v2 of the SageMaker Python SDK.
+
+        Args:
+            input_path (str): Location of the input file.
+            output_path (str): Desired location for the output file.
+                If the directories don't already exist, then they are created.
+                If a file exists at ``output_path``, then it is overwritten.
+        """
+        self.input_path = input_path
+        self.output_path = output_path
+
+    @abstractmethod
+    def update(self):
+        """Reads the input file, updates the code so that it is
+        compatible with v2 of the SageMaker Python SDK, and writes the
+        updated code to an output file.
+        """
+
+
+class PyFileUpdater(FileUpdater):
+    """A class for updating Python (``*.py``) files."""
+
+    def update(self):
+        """Reads the input Python file, updates the code so that it is
+        compatible with v2 of the SageMaker Python SDK, and writes the
+        updated code to an output file.
+        """
+        output = self._update_ast(self._read_input_file())
+        self._write_output_file(output)
+
+    def _update_ast(self, input_ast):
+        """Updates an abstract syntax tree (AST) so that it is compatible
+        with v2 of the SageMaker Python SDK.
+
+        Args:
+            input_ast (ast.Module): AST to be updated for use with Python SDK v2.
+
+        Returns:
+            ast.Module: Updated AST that is compatible with Python SDK v2.
+        """
+        return ASTTransformer().visit(input_ast)
+
+    def _read_input_file(self):
+        """Reads input file and parses it as an abstract syntax tree (AST).
+
+        Returns:
+            ast.Module: AST representation of the input file.
+        """
+        with open(self.input_path) as input_file:
+            return pasta.parse(input_file.read())
+
+    def _write_output_file(self, output):
+        """Writes abstract syntax tree (AST) to output file.
+        Creates the directories for the output path, if needed.
+
+        Args:
+            output (ast.Module): AST to save as the output file.
+        """
+        output_dir = os.path.dirname(self.output_path)
+        if output_dir and not os.path.exists(output_dir):
+            os.makedirs(output_dir)
+
+        if os.path.exists(self.output_path):
+            LOGGER.warning("Overwriting file %s", self.output_path)
+
+        with open(self.output_path, "w") as output_file:
+            output_file.write(pasta.dump(output))
+
+
+class JupyterNotebookFileUpdater(FileUpdater):
+    """A class for updating Jupyter notebook (``*.ipynb``) files.
+
+    For more on this file format, see
+    https://ipython.org/ipython-doc/dev/notebook/nbformat.html#nbformat.
+    """
+
+    def update(self):
+        """Reads the input Jupyter notebook file, updates the code so that it is
+        compatible with v2 of the SageMaker Python SDK, and writes the
+        updated code to an output file.
+        """
+        nb_json = self._read_input_file()
+        for cell in nb_json["cells"]:
+            if cell["cell_type"] == "code":
+                updated_source = self._update_code_from_cell(cell)
+                cell["source"] = updated_source
+
+        self._write_output_file(nb_json)
+
+    def _update_code_from_cell(self, cell):
+        """Updates the code from a code cell so that it is
+        compatible with v2 of the SageMaker Python SDK.
+
+        Args:
+            cell (dict): A dictionary representation of a code cell from
+                a Jupyter notebook. For more info, see
+                https://ipython.org/ipython-doc/dev/notebook/nbformat.html#code-cells.
+
+        Returns:
+            list[str]: A list of strings containing the lines of updated code that
+                can be used for the "source" attribute of a Jupyter notebook code cell.
+        """
+        code = "".join(cell["source"])
+        updated_ast = ASTTransformer().visit(pasta.parse(code))
+        updated_code = pasta.dump(updated_ast)
+        return self._code_str_to_source_list(updated_code)
+
+    def _code_str_to_source_list(self, code):
+        """Converts a string of code into a list for a Jupyter notebook code cell.
+
+        Args:
+            code (str): Code to be converted.
+
+        Returns:
+            list[str]: A list of strings containing the lines of code that
+                can be used for the "source" attribute of a Jupyter notebook code cell.
+                Each element of the list (i.e. line of code) contains a
+                trailing newline character ("\n") except for the last element.
+        """
+        source_list = ["{}\n".format(s) for s in code.split("\n")]
+        source_list[-1] = source_list[-1].rstrip("\n")
+        return source_list
+
+    def _read_input_file(self):
+        """Reads input file and parses it as JSON.
+
+        Returns:
+            dict: JSON representation of the input file.
+        """
+        with open(self.input_path) as input_file:
+            return json.load(input_file)
+
+    def _write_output_file(self, output):
+        """Writes JSON to output file. Creates the directories for the output path, if needed.
+
+        Args:
+            output (dict): JSON to save as the output file.
+        """
+        output_dir = os.path.dirname(self.output_path)
+        if output_dir and not os.path.exists(output_dir):
+            os.makedirs(output_dir)
+
+        if os.path.exists(self.output_path):
+            LOGGER.warning("Overwriting file %s", self.output_path)
+
+        with open(self.output_path, "w") as output_file:
+            json.dump(output, output_file, indent=1)
+            output_file.write("\n")  # json.dump does not write trailing newline

tools/compatibility/v2/modifiers/__init__.py

@@ -0,0 +1,14 @@
+# 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 for modifying AST nodes"""
+from __future__ import absolute_import