change: make v2 migration script remove script_mode param and set model_dir=False if needed (#1540)

Author: laurenyu

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

@@ -21,6 +21,7 @@
     modifiers.framework_version.FrameworkVersionEnforcer(),
     modifiers.tf_legacy_mode.TensorFlowLegacyModeConstructorUpgrader(),
     modifiers.tf_legacy_mode.TensorBoardParameterRemover(),
+    modifiers.deprecated_params.TensorFlowScriptModeParameterRemover(),
 ]
 
 

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

@@ -14,6 +14,7 @@
 from __future__ import absolute_import
 
 from sagemaker.cli.compatibility.v2.modifiers import (  # noqa: F401 (imported but unused)
+    deprecated_params,
     framework_version,
     tf_legacy_mode,
 )

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

@@ -0,0 +1,80 @@
+# 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 remove deprecated parameters."""
+from __future__ import absolute_import
+
+import ast
+
+from sagemaker.cli.compatibility.v2.modifiers.modifier import Modifier
+
+
+class TensorFlowScriptModeParameterRemover(Modifier):
+    """A class to remove ``script_mode`` from TensorFlow estimators (because it's the only mode)."""
+
+    def node_should_be_modified(self, node):
+        """Checks if the ``ast.Call`` node instantiates a TensorFlow estimator with
+        ``script_mode`` set.
+
+        This looks for the following formats:
+
+        - ``TensorFlow``
+        - ``sagemaker.tensorflow.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 TensorFlow estimator with ``script_mode``.
+        """
+        return self._is_tf_constructor(node) and self._has_script_mode_param(node)
+
+    def _is_tf_constructor(self, node):
+        """Checks if the ``ast.Call`` node represents a call of the form
+        ``TensorFlow`` or ``sagemaker.tensorflow.TensorFlow``.
+        """
+        # Check for TensorFlow()
+        if isinstance(node.func, ast.Name):
+            return node.func.id == "TensorFlow"
+
+        # Check for sagemaker.tensorflow.TensorFlow()
+        ends_with_tensorflow_constructor = (
+            isinstance(node.func, ast.Attribute) and node.func.attr == "TensorFlow"
+        )
+
+        is_in_tensorflow_module = (
+            isinstance(node.func.value, ast.Attribute)
+            and node.func.value.attr == "tensorflow"
+            and isinstance(node.func.value.value, ast.Name)
+            and node.func.value.value.id == "sagemaker"
+        )
+
+        return ends_with_tensorflow_constructor and is_in_tensorflow_module
+
+    def _has_script_mode_param(self, node):
+        """Checks if the ``ast.Call`` node's keywords include ``script_mode``."""
+        for kw in node.keywords:
+            if kw.arg == "script_mode":
+                return True
+
+        return False
+
+    def modify_node(self, node):
+        """Modifies the ``ast.Call`` node's keywords to remove ``script_mode``.
+
+        Args:
+            node (ast.Call): a node that represents a TensorFlow constructor.
+        """
+        for kw in node.keywords:
+            if kw.arg == "script_mode":
+                node.keywords.remove(kw)

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

@@ -15,14 +15,17 @@
 
 import ast
 
+import boto3
 import six
 
+from sagemaker.cli.compatibility.v2.modifiers import framework_version
 from sagemaker.cli.compatibility.v2.modifiers.modifier import Modifier
+from sagemaker import fw_utils
 
 
 class TensorFlowLegacyModeConstructorUpgrader(Modifier):
-    """A class to turn legacy mode parameters into hyperparameters when
-    instantiating a TensorFlow estimator.
+    """A class to turn legacy mode parameters into hyperparameters, disable the ``model_dir``
+    hyperparameter, and set the image URI when instantiating a TensorFlow estimator.
     """
 
     LEGACY_MODE_PARAMETERS = (
@@ -32,6 +35,18 @@ class TensorFlowLegacyModeConstructorUpgrader(Modifier):
         "training_steps",
     )
 
+    def __init__(self):
+        """Initializes a ``TensorFlowLegacyModeConstructorUpgrader``."""
+        self._region = None
+
+    @property
+    def region(self):
+        """Returns the AWS region for constructing an ECR image URI."""
+        if self._region is None:
+            self._region = boto3.Session().region_name
+
+        return self._region
+
     def node_should_be_modified(self, node):
         """Checks if the ``ast.Call`` node instantiates a TensorFlow estimator with legacy mode.
 
@@ -89,7 +104,7 @@ def _is_legacy_mode(self, node):
 
     def modify_node(self, node):
         """Modifies the ``ast.Call`` node's keywords to turn TensorFlow legacy mode parameters
-        into hyperparameters and set ``script_mode=False``.
+        into hyperparameters and sets ``model_dir=False``.
 
         The parameters that are converted into hyperparameters:
 
@@ -105,9 +120,11 @@ def modify_node(self, node):
         additional_hps = {}
         kw_to_remove = []  # remove keyword args after so that none are skipped during iteration
 
+        add_image_uri = True
+
         for kw in node.keywords:
-            if kw.arg == "script_mode":
-                # remove here because is set to False later regardless of current value
+            if kw.arg in ("script_mode", "model_dir"):
+                # model_dir is removed so that it can be set to False later
                 kw_to_remove.append(kw)
             if kw.arg == "hyperparameters" and kw.value:
                 base_hps = dict(zip(kw.value.keys, kw.value.values))
@@ -116,11 +133,17 @@ def modify_node(self, node):
                 hp_key = self._hyperparameter_key_for_param(kw.arg)
                 additional_hps[hp_key] = kw.value
                 kw_to_remove.append(kw)
+            if kw.arg == "image_name":
+                add_image_uri = False
 
         self._remove_keywords(node, kw_to_remove)
         self._add_updated_hyperparameters(node, base_hps, additional_hps)
 
-        node.keywords.append(ast.keyword(arg="script_mode", value=ast.NameConstant(value=False)))
+        if add_image_uri:
+            image_uri = self._image_uri_from_args(node.keywords)
+            node.keywords.append(ast.keyword(arg="image_name", value=ast.Str(s=image_uri)))
+
+        node.keywords.append(ast.keyword(arg="model_dir", value=ast.NameConstant(value=False)))
 
     def _hyperparameter_key_for_param(self, arg):
         """Returns an ``ast.Str`` for a hyperparameter key replacing a legacy mode parameter."""
@@ -148,6 +171,21 @@ 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."""
+        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
+            if kw.arg == "train_instance_type":
+                instance_type = kw.value.s
+
+        return fw_utils.create_image_uri(
+            self.region, "tensorflow", instance_type, tf_version, "py2"
+        )
+
 
 class TensorBoardParameterRemover(Modifier):
     """A class for removing the ``run_tensorboard_locally`` parameter from ``fit()``."""

tests/unit/sagemaker/cli/__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

tests/unit/sagemaker/cli/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

tests/unit/sagemaker/cli/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

tests/unit/sagemaker/cli/compatibility/v2/modifiers/__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

tests/unit/sagemaker/cli/compatibility/v2/modifiers/ast_converter.py

@@ -0,0 +1,19 @@
+# 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
+
+import pasta
+
+
+def ast_call(code):
+    return pasta.parse(code).body[0].value

tests/unit/sagemaker/cli/compatibility/v2/modifiers/test_framework_version.py

@@ -18,6 +18,7 @@
 import pytest
 
 from sagemaker.cli.compatibility.v2.modifiers import framework_version
+from tests.unit.sagemaker.cli.compatibility.v2.modifiers.ast_converter import ast_call
 
 
 @pytest.fixture(autouse=True)
@@ -54,7 +55,7 @@ def test_node_should_be_modified_fw_constructor_no_fw_version():
     modifier = framework_version.FrameworkVersionEnforcer()
 
     for constructor in fw_constructors:
-        node = _ast_call(constructor)
+        node = ast_call(constructor)
         assert modifier.node_should_be_modified(node) is True
 
 
@@ -85,12 +86,12 @@ def test_node_should_be_modified_fw_constructor_with_fw_version():
     modifier = framework_version.FrameworkVersionEnforcer()
 
     for constructor in fw_constructors:
-        node = _ast_call(constructor)
+        node = ast_call(constructor)
         assert modifier.node_should_be_modified(node) is False
 
 
 def test_node_should_be_modified_random_function_call():
-    node = _ast_call("sagemaker.session.Session()")
+    node = ast_call("sagemaker.session.Session()")
     modifier = framework_version.FrameworkVersionEnforcer()
     assert modifier.node_should_be_modified(node) is False
 
@@ -139,14 +140,10 @@ def test_modify_node_sklearn():
     _test_modify_node(classes, "0.20.0")
 
 
-def _ast_call(code):
-    return pasta.parse(code).body[0].value
-
-
 def _test_modify_node(classes, default_version):
     modifier = framework_version.FrameworkVersionEnforcer()
     for cls in classes:
-        node = _ast_call("{}()".format(cls))
+        node = ast_call("{}()".format(cls))
         modifier.modify_node(node)
 
         expected_result = "{}(framework_version='{}')".format(cls, default_version)

tests/unit/sagemaker/cli/compatibility/v2/modifiers/test_tensorboard_remove_param.py

@@ -15,6 +15,7 @@
 import pasta
 
 from sagemaker.cli.compatibility.v2.modifiers import tf_legacy_mode
+from tests.unit.sagemaker.cli.compatibility.v2.modifiers.ast_converter import ast_call
 
 
 def test_node_should_be_modified_fit_with_tensorboard():
@@ -26,7 +27,7 @@ def test_node_should_be_modified_fit_with_tensorboard():
     modifier = tf_legacy_mode.TensorBoardParameterRemover()
 
     for call in fit_calls:
-        node = _ast_call(call)
+        node = ast_call(call)
         assert modifier.node_should_be_modified(node) is True
 
 
@@ -36,12 +37,12 @@ def test_node_should_be_modified_fit_without_tensorboard():
     modifier = tf_legacy_mode.TensorBoardParameterRemover()
 
     for call in fit_calls:
-        node = _ast_call(call)
+        node = ast_call(call)
         assert modifier.node_should_be_modified(node) is False
 
 
 def test_node_should_be_modified_random_function_call():
-    node = _ast_call("estimator.deploy(1, 'local')")
+    node = ast_call("estimator.deploy(1, 'local')")
     modifier = tf_legacy_mode.TensorBoardParameterRemover()
     assert modifier.node_should_be_modified(node) is False
 
@@ -54,10 +55,6 @@ def test_modify_node():
     modifier = tf_legacy_mode.TensorBoardParameterRemover()
 
     for call in fit_calls:
-        node = _ast_call(call)
+        node = ast_call(call)
         modifier.modify_node(node)
         assert "estimator.fit()" == pasta.dump(node)
-
-
-def _ast_call(code):
-    return pasta.parse(code).body[0].value