change: create ASTTransformer class to handle migrating Python SDK code for v2

tools/__init__.py

@@ -0,0 +1,13 @@
+# 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.
+from __future__ import absolute_import

tools/compatibility/__init__.py

@@ -0,0 +1,13 @@
+# 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.
+from __future__ import absolute_import

tools/compatibility/v2/__init__.py

@@ -0,0 +1,13 @@
+# 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.
+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 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/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

tools/compatibility/v2/modifiers/framework_version.py

@@ -0,0 +1,123 @@
+# 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.
+"""A class to ensure that ``framework_version`` is defined when constructing framework classes."""
+from __future__ import absolute_import
+
+import ast
+
+from modifiers.modifier import Modifier
+
+FRAMEWORK_DEFAULTS = {
+    "Chainer": "4.1.0",
+    "MXNet": "1.2.0",
+    "PyTorch": "0.4.0",
+    "SKLearn": "0.20.0",
+    "TensorFlow": "1.11.0",
+}
+
+FRAMEWORKS = list(FRAMEWORK_DEFAULTS.keys())
+# TODO: check for sagemaker.tensorflow.serving.Model
+FRAMEWORK_CLASSES = FRAMEWORKS + ["{}Model".format(fw) for fw in FRAMEWORKS]
+FRAMEWORK_MODULES = [fw.lower() for fw in FRAMEWORKS]
+
+
+class FrameworkVersionEnforcer(Modifier):
+    def node_should_be_modified(self, node):
+        """Check if the ast.Call node instantiates a framework estimator or model,
+        but doesn't specify the framework_version parameter.
+
+        This looks for the following formats:
+
+        - ``TensorFlow``
+        - ``sagemaker.tensorflow.TensorFlow``
+
+        where "TensorFlow" can be Chainer, MXNet, PyTorch, SKLearn, or TensorFlow.
+
+        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`` is instantiating a framework class that
+                should specify ``framework_version``, but doesn't.
+        """
+        if self._is_framework_constructor(node):
+            return not self._fw_version_in_keywords(node)
+
+        return False
+
+    def _is_framework_constructor(self, node):
+        """Check if the ``ast.Call`` node represents a call of the form
+        <Framework> or sagemaker.<framework>.<Framework>.
+        """
+        if isinstance(node.func, ast.Name):
+            if node.func.id in FRAMEWORK_CLASSES:
+                return True
+
+        if (
+            isinstance(node.func, ast.Attribute)
+            and node.func.attr in FRAMEWORK_CLASSES
+            and isinstance(node.func.value, ast.Attribute)
+            and node.func.value.attr in FRAMEWORK_MODULES
+            and isinstance(node.func.value.value, ast.Name)
+            and node.func.value.value.id == "sagemaker"
+        ):
+            return True
+
+        return False
+
+    def _fw_version_in_keywords(self, node):
+        """Check if the ``ast.Call`` node's keywords contain ``framework_version``."""
+        for kw in node.keywords:
+            if kw.arg == "framework_version" and kw.value:
+                return True
+        return False
+
+    def modify_node(self, node):
+        """Modify the ``ast.Call`` node's keywords to include ``framework_version``.
+
+        The ``framework_version`` value is determined by the framework:
+
+        - Chainer: "4.1.0"
+        - MXNet: "1.2.0"
+        - PyTorch: "0.4.0"
+        - SKLearn: "0.20.0"
+        - TensorFlow: "1.11.0"
+
+        Args:
+            node (ast.Call): a node that represents the constructor of a framework class.
+        """
+        framework = self._framework_name_from_node(node)
+        node.keywords.append(
+            ast.keyword(arg="framework_version", value=ast.Str(s=FRAMEWORK_DEFAULTS[framework]))
+        )
+
+    def _framework_name_from_node(self, node):
+        """Retrieve the framework name based on the function call.
+
+        Args:
+            node (ast.Call): a node that represents the constructor of a framework class.
+                This can represent either <Framework> or sagemaker.<framework>.<Framework>.
+
+        Returns:
+            str: the (capitalized) framework name.
+        """
+        if isinstance(node.func, ast.Name):
+            framework = node.func.id
+        elif isinstance(node.func, ast.Attribute):
+            framework = node.func.attr
+
+        if framework.endswith("Model"):
+            framework = framework[: framework.find("Model")]
+
+        return framework

tools/compatibility/v2/modifiers/modifier.py

@@ -0,0 +1,35 @@
+# 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.
+"""Abstract class for modifying AST nodes."""
+from __future__ import absolute_import
+
+from abc import abstractmethod
+
+
+class Modifier(object):
+    """Abstract class to take in an AST node, check if it needs modification,
+    and potentially modify the node.
+    """
+
+    def check_and_modify_node(self, node):
+        """Check an AST node, and modify it if applicable."""
+        if self.node_should_be_modified(node):
+            self.modify_node(node)
+
+    @abstractmethod
+    def node_should_be_modified(self, node):
+        """Check if an AST node should be modified."""
+
+    @abstractmethod
+    def modify_node(self, node):
+        """Modify an AST node."""