change: make v2 migration script remove legacy run_tensorboard_locally parameter (#1537)

Author: laurenyu

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

@@ -20,6 +20,7 @@
 FUNCTION_CALL_MODIFIERS = [
     modifiers.framework_version.FrameworkVersionEnforcer(),
     modifiers.tf_legacy_mode.TensorFlowLegacyModeConstructorUpgrader(),
+    modifiers.tf_legacy_mode.TensorBoardParameterRemover(),
 ]
 
 

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

@@ -11,7 +11,6 @@
 # ANY KIND, either express or implied. See the License for the specific
 # language governing permissions and limitations under the License.
 """Classes to modify TensorFlow legacy mode code to be compatible with SageMaker Python SDK v2."""
-# TODO: handle fit(run_tensorboard_locally=True)
 from __future__ import absolute_import
 
 import ast
@@ -148,3 +147,38 @@ def _to_ast_keyword(self, hps):
             return ast.keyword(arg="hyperparameters", value=ast.Dict(keys=keys, values=values))
 
         return None
+
+
+class TensorBoardParameterRemover(Modifier):
+    """A class for removing the ``run_tensorboard_locally`` parameter from ``fit()``."""
+
+    def node_should_be_modified(self, node):
+        """Checks if the ``ast.Call`` node invokes a function named "fit" and
+        contains a keyword argument named "run_tensorboard_locally".
+
+        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 invoking a function named "fit" with
+                a parameter named "run_tensorboard_locally".
+        """
+        is_fit_call = isinstance(node.func, ast.Attribute) and node.func.attr == "fit"
+        if is_fit_call:
+            for kw in node.keywords:
+                if kw.arg == "run_tensorboard_locally":
+                    return True
+
+        return False
+
+    def modify_node(self, node):
+        """Removes ``run_tensorboard_locally`` from the ``ast.Call`` node's keywords.
+
+        Args:
+            node (ast.Call): a node that represents ``fit`` being called with
+                ``run_tensorboard_locally`` set.
+        """
+        for kw in node.keywords:
+            if kw.arg == "run_tensorboard_locally":
+                node.keywords.remove(kw)

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

@@ -0,0 +1,63 @@
+# 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
+
+from sagemaker.cli.compatibility.v2.modifiers import tf_legacy_mode
+
+
+def test_node_should_be_modified_fit_with_tensorboard():
+    fit_calls = (
+        "estimator.fit(run_tensorboard_locally=True)",
+        "tensorflow.fit(run_tensorboard_locally=False)",
+    )
+
+    modifier = tf_legacy_mode.TensorBoardParameterRemover()
+
+    for call in fit_calls:
+        node = _ast_call(call)
+        assert modifier.node_should_be_modified(node) is True
+
+
+def test_node_should_be_modified_fit_without_tensorboard():
+    fit_calls = ("estimator.fit()", "tensorflow.fit()")
+
+    modifier = tf_legacy_mode.TensorBoardParameterRemover()
+
+    for call in fit_calls:
+        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')")
+    modifier = tf_legacy_mode.TensorBoardParameterRemover()
+    assert modifier.node_should_be_modified(node) is False
+
+
+def test_modify_node():
+    fit_calls = (
+        "estimator.fit(run_tensorboard_locally=True)",
+        "estimator.fit(run_tensorboard_locally=False)",
+    )
+    modifier = tf_legacy_mode.TensorBoardParameterRemover()
+
+    for call in fit_calls:
+        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