feature: add 1p algorithm image_uris migration tool

Author: Chuyang Deng

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

@@ -36,6 +36,7 @@
     modifiers.training_params.TrainPrefixRemover(),
     modifiers.training_input.TrainingInputConstructorRefactor(),
     modifiers.serde.SerdeConstructorRenamer(),
+    modifiers.image_uris.ImageURIRetrieveRefactor(),
 ]
 
 IMPORT_MODIFIERS = [modifiers.tfs.TensorFlowServingImportRenamer()]
@@ -53,6 +54,7 @@
     modifiers.training_input.TrainingInputImportFromRenamer(),
     modifiers.serde.SerdeImportFromAmazonCommonRenamer(),
     modifiers.serde.SerdeImportFromPredictorRenamer(),
+    modifiers.image_uris.ImageURIRetrieveImportFromRenamer(),
 ]
 
 

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

@@ -24,4 +24,5 @@
     tfs,
     training_params,
     training_input,
+    image_uris,
 )

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

@@ -16,12 +16,50 @@
 from __future__ import absolute_import
 
 import ast
+import logging
+import pasta
 
 from sagemaker.cli.compatibility.v2.modifiers import matching
 from sagemaker.cli.compatibility.v2.modifiers.modifier import Modifier
 
 GET_IMAGE_URI_NAME = "get_image_uri"
-GET_IMAGE_URI_NAMESPACES = ("sagemaker", "sagemaker.amazon_estimator")
+GET_IMAGE_URI_NAMESPACES = (
+    "sagemaker",
+    "sagemaker.amazon_estimator",
+    "sagemaker.amazon.amazon_estimator",
+    "amazon_estimator",
+    "amazon.amazon_estimator",
+)
+
+ALGORITHM_NAME_FROM_REPO = {
+    "blazingtext": "blazingtext",
+    "sagemaker-rl-mxnet": "coach-mxnet",
+    "sagemaker-rl-tensorflow": ["coach-tensorflow", "ray-tensorflow"],
+    "factorization-machine": "factorization-machines",
+    "forecasting-deepar": "forecasting-deepar",
+    "image-classification": "image-classification",
+    "image-classification-neo": "image-classification-neo",
+    "ipinsights": "ipinsights",
+    "kmeans": "kmeans",
+    "knn": "knn",
+    "lda": "lda",
+    "linear-learner": "linear-learner",
+    "ntm": "ntm",
+    "object2vec": "object2vec",
+    "object-detection": "object-detection",
+    "pca": "pca",
+    "randomcutforest": "randomcutforest",
+    "sagemaker-rl-ray-container": "ray-pytorch",
+    "semantic-segmentation": "semantic-segmentation",
+    "seq2seq": "seq2seq",
+    "sagemaker-scikit-learn": "sklearn",
+    "sagemaker-sparkml-serving": "sparkml-serving",
+    "sagemaker-rl-vw-container": "vw",
+    "sagemaker-xgboost": "xgboost",
+    "xgboost-neo": "xgboost-neo",
+}
+
+logger = logging.getLogger("sagemaker")
 
 
 class ImageURIRetrieveRefactor(Modifier):
@@ -43,7 +81,9 @@ def node_should_be_modified(self, node):
         Returns:
             bool: If the ``ast.Call`` instantiates a class of interest.
         """
-        return matching.matches_name_or_namespaces(node, GET_IMAGE_URI_NAME, GET_IMAGE_URI_NAMESPACES)
+        return matching.matches_name_or_namespaces(
+            node, GET_IMAGE_URI_NAME, GET_IMAGE_URI_NAMESPACES
+        )
 
     def modify_node(self, node):
         """Modifies the ``ast.Call`` node to call ``image_uris.retrieve`` instead.
@@ -52,12 +92,59 @@ def modify_node(self, node):
         Args:
             node (ast.Call): a node that represents a *image_uris.retrieve call.
         """
-        if matching.matches_name(node, GET_IMAGE_URI_NAME):
-            node.func.id = "image_uris.retrieve"
-            node.func.params.argOne, node.func.params.argTwo = node.func.params.argTwo, node.func.params.argOne
-        elif matching.matches_attr(node, GET_IMAGE_URI_NAME):
-            node.func.attr = "image_uris.retrieve"
-            node.func.params.argOne, node.func.params.argTwo = node.func.params.argTwo, node.func.params.argOne
+        original_args = [None] * 3
+        for kw in node.keywords:
+            if kw.arg == "repo_name":
+                arg = kw.value.s
+                modified_arg = ALGORITHM_NAME_FROM_REPO[arg]
+                if isinstance(modified_arg, list):
+                    logger.warning(
+                        "There are more than one value mapping to {}, {} will be used".format(
+                            arg, modified_arg[0]
+                        )
+                    )
+                    modified_arg = modified_arg[0]
+                original_args[0] = ast.Str(modified_arg)
+            elif kw.arg == "repo_region":
+                original_args[1] = ast.Str(kw.value.s)
+            elif kw.arg == "repo_version":
+                original_args[2] = ast.Str(kw.value.s)
+
+        if len(node.args) > 0:
+            original_args[1] = ast.Str(node.args[0].s)
+        if len(node.args) > 1:
+            arg = node.args[1].s
+            modified_arg = ALGORITHM_NAME_FROM_REPO[arg]
+            if isinstance(modified_arg, list):
+                logger.warning(
+                    "There are more than one value mapping to {}, {} will be used".format(
+                        arg, modified_arg[0]
+                    )
+                )
+                modified_arg = modified_arg[0]
+            original_args[0] = ast.Str(modified_arg)
+        if len(node.args) > 2:
+            original_args[2] = ast.Str(node.args[2].s)
+
+        args = []
+        for arg in original_args:
+            if arg:
+                args.append(arg)
+
+        if matching.matches_name(node, GET_IMAGE_URI_NAME) or matching.matches_attr(
+            node, GET_IMAGE_URI_NAME
+        ):
+            node_components = list(pasta.dump(node).split("."))
+            node_modules = node_components[: len(node_components) - 1]
+            if "sagemaker" in node_modules:
+                node.func = ast.Attribute(
+                    value=ast.Attribute(attr="image_uris", value=ast.Name(id="sagemaker")),
+                    attr="retrieve",
+                )
+            else:
+                node.func = ast.Attribute(value=ast.Name(id="image_uris"), attr="retrieve")
+            node.args = args
+            node.keywords = []
         return node
 
 
@@ -75,7 +162,7 @@ def node_should_be_modified(self, node):
             bool: If the import statement imports ``get_image_uri`` from the correct module.
         """
         return node.module in GET_IMAGE_URI_NAMESPACES and any(
-            name.name == GET_IMAGE_URI_NAMESPACES for name in node.names
+            name.name == GET_IMAGE_URI_NAME for name in node.names
         )
 
     def modify_node(self, node):
@@ -91,6 +178,6 @@ def modify_node(self, node):
         for name in node.names:
             if name.name == GET_IMAGE_URI_NAME:
                 name.name = "image_uris"
-            if node.module == "sagemaker.amazon_estimator":
+            if node.module in GET_IMAGE_URI_NAMESPACES:
                 node.module = "sagemaker"
         return node

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

@@ -0,0 +1,120 @@
+# 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
+import pytest
+
+from sagemaker.cli.compatibility.v2.modifiers import image_uris
+from tests.unit.sagemaker.cli.compatibility.v2.modifiers.ast_converter import ast_call, ast_import
+
+
+@pytest.fixture
+def methods():
+    return (
+        "get_image_uri('us-west-2', 'kmeans')",
+        "sagemaker.get_image_uri(repo_region='us-west-2', repo_name='sagemaker-scikil-learn')",
+        "sagemaker.amazon_estimator.get_image_uri('us-west-2', repo_name='sagemaker-scikil-learn')",
+        "sagemaker.amazon.amazon_estimator.get_image_uri('us-west-2', 'sagemaker-scikil-learn', repo_version='1')",
+    )
+
+
+@pytest.fixture
+def import_statements():
+    return (
+        "from sagemaker import get_image_uri",
+        "from sagemaker.amazon_estimator import get_image_uri",
+        "from sagemaker.amazon.amazon_estimator import get_image_uri",
+    )
+
+
+def test_method_node_should_be_modified(methods):
+    modifier = image_uris.ImageURIRetrieveRefactor()
+    for method in methods:
+        node = ast_call(method)
+        assert modifier.node_should_be_modified(node)
+
+
+def test_methodnode_should_be_modified_random_call():
+    modifier = image_uris.ImageURIRetrieveRefactor()
+    node = ast_call("create_image_uri()")
+    assert not modifier.node_should_be_modified(node)
+
+
+def test_method_modify_node(methods, caplog):
+    modifier = image_uris.ImageURIRetrieveRefactor()
+
+    method = "get_image_uri('us-west-2', 'sagemaker-scikit-learn')"
+    node = ast_call(method)
+    modifier.modify_node(node)
+    assert "image_uris.retrieve('sklearn', 'us-west-2')" == pasta.dump(node)
+
+    method = "amazon_estimator.get_image_uri('us-west-2', 'sagemaker-scikit-learn')"
+    node = ast_call(method)
+    modifier.modify_node(node)
+    assert "image_uris.retrieve('sklearn', 'us-west-2')" == pasta.dump(node)
+
+    method = "sagemaker.get_image_uri(repo_region='us-west-2', repo_name='sagemaker-scikit-learn')"
+    node = ast_call(method)
+    modifier.modify_node(node)
+    assert "sagemaker.image_uris.retrieve('sklearn', 'us-west-2')" == pasta.dump(node)
+
+    method = (
+        "sagemaker.amazon_estimator.get_image_uri('us-west-2', repo_name='sagemaker-scikit-learn')"
+    )
+    node = ast_call(method)
+    modifier.modify_node(node)
+    assert "sagemaker.image_uris.retrieve('sklearn', 'us-west-2')" == pasta.dump(node)
+
+    method = "sagemaker.amazon.amazon_estimator.get_image_uri('us-west-2', 'sagemaker-scikit-learn', repo_version='1')"
+    node = ast_call(method)
+    modifier.modify_node(node)
+    assert "sagemaker.image_uris.retrieve('sklearn', 'us-west-2', '1')" == pasta.dump(node)
+
+    method = "get_image_uri('us-west-2', 'sagemaker-rl-tensorflow')"
+    node = ast_call(method)
+    modifier.modify_node(node)
+    assert "image_uris.retrieve('coach-tensorflow', 'us-west-2')"
+    assert "There are more than one value mapping to" in caplog.text
+
+
+def test_import_from_node_should_be_modified_image_uris_input(import_statements):
+    modifier = image_uris.ImageURIRetrieveImportFromRenamer()
+
+    statement = "from sagemaker import get_image_uri"
+    node = ast_import(statement)
+    assert modifier.node_should_be_modified(node)
+
+    statement = "from sagemaker.amazon_estimator import get_image_uri"
+    node = ast_import(statement)
+    assert modifier.node_should_be_modified(node)
+
+    statement = "from sagemaker.amazon.amazon_estimator import get_image_uri"
+    node = ast_import(statement)
+    assert modifier.node_should_be_modified(node)
+
+
+def test_import_from_node_should_be_modified_random_import():
+    modifier = image_uris.ImageURIRetrieveImportFromRenamer()
+    node = ast_import("from sagemaker.amazon_estimator import registry")
+    assert not modifier.node_should_be_modified(node)
+
+
+def test_import_from_modify_node(import_statements):
+    modifier = image_uris.ImageURIRetrieveImportFromRenamer()
+    expected_result = "from sagemaker import image_uris"
+
+    for import_statement in import_statements:
+        node = ast_import(import_statement)
+        modifier.modify_node(node)
+        assert expected_result == pasta.dump(node)