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

Author: laurenyu

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 TensorFlowScriptModeParamRemover(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 = (
@@ -89,7 +92,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 +108,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 +121,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 +159,20 @@ 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
+
+        region = boto3.Session().region_name
+        return fw_utils.create_image_uri(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/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

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

@@ -16,8 +16,13 @@
 
 import pasta
 import pytest
+from mock import patch
 
 from sagemaker.cli.compatibility.v2.modifiers import tf_legacy_mode
+from tests.unit.sagemaker.cli.compatibility.v2.modifiers.ast_converter import ast_call
+
+IMAGE_URI = "sagemaker-tensorflow:latest"
+REGION_NAME = "us-west-2"
 
 
 @pytest.fixture(autouse=True)
@@ -42,7 +47,7 @@ def test_node_should_be_modified_tf_constructor_legacy_mode():
     modifier = tf_legacy_mode.TensorFlowLegacyModeConstructorUpgrader()
 
     for constructor in tf_legacy_mode_constructors:
-        node = _ast_call(constructor)
+        node = ast_call(constructor)
         assert modifier.node_should_be_modified(node) is True
 
 
@@ -61,28 +66,56 @@ def test_node_should_be_modified_tf_constructor_script_mode():
     modifier = tf_legacy_mode.TensorFlowLegacyModeConstructorUpgrader()
 
     for constructor in tf_script_mode_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("MXNet(py_version='py3')")
+    node = ast_call("MXNet(py_version='py3')")
     modifier = tf_legacy_mode.TensorFlowLegacyModeConstructorUpgrader()
     assert modifier.node_should_be_modified(node) is False
 
 
-def test_modify_node_set_script_mode_false():
+@patch("boto3.Session")
+@patch("sagemaker.fw_utils.create_image_uri", return_value=IMAGE_URI)
+def test_modify_node_set_model_dir_and_image_name(create_image_uri, boto_session):
+    boto_session.return_value.region_name = REGION_NAME
+
     tf_constructors = (
         "TensorFlow()",
         "TensorFlow(script_mode=False)",
-        "TensorFlow(script_mode=None)",
+        "TensorFlow(model_dir='s3//bucket/model')",
     )
     modifier = tf_legacy_mode.TensorFlowLegacyModeConstructorUpgrader()
 
     for constructor in tf_constructors:
-        node = _ast_call(constructor)
+        node = ast_call(constructor)
         modifier.modify_node(node)
-        assert "TensorFlow(script_mode=False)" == pasta.dump(node)
+
+        assert "TensorFlow(image_name='{}', model_dir=False)".format(IMAGE_URI) == pasta.dump(node)
+        create_image_uri.assert_called_with(
+            REGION_NAME, "tensorflow", "ml.m4.xlarge", "1.11.0", "py2"
+        )
+
+
+@patch("boto3.Session")
+@patch("sagemaker.fw_utils.create_image_uri", return_value=IMAGE_URI)
+def test_modify_node_set_image_name_from_args(create_image_uri, boto_session):
+    boto_session.return_value.region_name = REGION_NAME
+
+    tf_constructor = "TensorFlow(train_instance_type='ml.p2.xlarge', framework_version='1.4.0')"
+
+    node = ast_call(tf_constructor)
+    modifier = tf_legacy_mode.TensorFlowLegacyModeConstructorUpgrader()
+    modifier.modify_node(node)
+
+    create_image_uri.assert_called_with(REGION_NAME, "tensorflow", "ml.p2.xlarge", "1.4.0", "py2")
+
+    expected_string = (
+        "TensorFlow(train_instance_type='ml.p2.xlarge', framework_version='1.4.0', "
+        "image_name='{}', model_dir=False)".format(IMAGE_URI)
+    )
+    assert expected_string == pasta.dump(node)
 
 
 def test_modify_node_set_hyperparameters():
@@ -93,7 +126,7 @@ def test_modify_node_set_hyperparameters():
         requirements_file='source/requirements.txt',
     )"""
 
-    node = _ast_call(tf_constructor)
+    node = ast_call(tf_constructor)
     modifier = tf_legacy_mode.TensorFlowLegacyModeConstructorUpgrader()
     modifier.modify_node(node)
 
@@ -115,7 +148,7 @@ def test_modify_node_preserve_other_hyperparameters():
         hyperparameters={'optimizer': 'sgd', 'lr': 0.1, 'checkpoint_path': 's3://foo/bar'},
     )"""
 
-    node = _ast_call(tf_constructor)
+    node = ast_call(tf_constructor)
     modifier = tf_legacy_mode.TensorFlowLegacyModeConstructorUpgrader()
     modifier.modify_node(node)
 
@@ -138,7 +171,7 @@ def test_modify_node_prefer_param_over_hyperparameter():
         hyperparameters={'training_steps': 10, 'sagemaker_requirements': 'foo.txt'},
     )"""
 
-    node = _ast_call(tf_constructor)
+    node = ast_call(tf_constructor)
     modifier = tf_legacy_mode.TensorFlowLegacyModeConstructorUpgrader()
     modifier.modify_node(node)
 
@@ -156,7 +189,3 @@ def _hyperparameters_from_node(node):
             keys = [k.s for k in kw.value.keys]
             values = [getattr(v, v._fields[0]) for v in kw.value.values]
             return dict(zip(keys, values))
-
-
-def _ast_call(code):
-    return pasta.parse(code).body[0].value