update docstring and remove unused algorithm name

Author: Chuyang Deng

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

@@ -10,9 +10,7 @@
 # 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 modify image_uris.retrieve() code to be compatible
-with version 2.0 and later of the SageMaker Python SDK.
-"""
+"""Classes to modify image uri retrieve methods for Python SDK v2.0 and later."""
 from __future__ import absolute_import
 
 import ast
@@ -33,8 +31,6 @@
 
 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",
@@ -49,12 +45,8 @@
     "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",
 }
@@ -97,13 +89,6 @@ def modify_node(self, node):
             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)
@@ -115,13 +100,6 @@ def modify_node(self, node):
         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)

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

@@ -22,10 +22,10 @@
 @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')",
+        "get_image_uri('us-west-2', 'sagemaker-xgboost')",
+        "sagemaker.get_image_uri(repo_region='us-west-2', repo_name='sagemaker-xgboost')",
+        "sagemaker.amazon_estimator.get_image_uri('us-west-2', repo_name='sagemaker-xgboost')",
+        "sagemaker.amazon.amazon_estimator.get_image_uri('us-west-2', 'sagemaker-xgboost', repo_version='1')",
     )
 
 
@@ -54,38 +54,30 @@ def test_methodnode_should_be_modified_random_call():
 def test_method_modify_node(methods, caplog):
     modifier = image_uris.ImageURIRetrieveRefactor()
 
-    method = "get_image_uri('us-west-2', 'sagemaker-scikit-learn')"
+    method = "get_image_uri('us-west-2', 'sagemaker-xgboost')"
     node = ast_call(method)
     modifier.modify_node(node)
-    assert "image_uris.retrieve('sklearn', 'us-west-2')" == pasta.dump(node)
+    assert "image_uris.retrieve('xgboost', 'us-west-2')" == pasta.dump(node)
 
-    method = "amazon_estimator.get_image_uri('us-west-2', 'sagemaker-scikit-learn')"
+    method = "amazon_estimator.get_image_uri('us-west-2', 'sagemaker-xgboost')"
     node = ast_call(method)
     modifier.modify_node(node)
-    assert "image_uris.retrieve('sklearn', 'us-west-2')" == pasta.dump(node)
+    assert "image_uris.retrieve('xgboost', 'us-west-2')" == pasta.dump(node)
 
-    method = "sagemaker.get_image_uri(repo_region='us-west-2', repo_name='sagemaker-scikit-learn')"
+    method = "sagemaker.get_image_uri(repo_region='us-west-2', repo_name='sagemaker-xgboost')"
     node = ast_call(method)
     modifier.modify_node(node)
-    assert "sagemaker.image_uris.retrieve('sklearn', 'us-west-2')" == pasta.dump(node)
+    assert "sagemaker.image_uris.retrieve('xgboost', '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')"
+    method = "sagemaker.amazon_estimator.get_image_uri('us-west-2', repo_name='sagemaker-xgboost')"
     node = ast_call(method)
     modifier.modify_node(node)
-    assert "sagemaker.image_uris.retrieve('sklearn', 'us-west-2', '1')" == pasta.dump(node)
+    assert "sagemaker.image_uris.retrieve('xgboost', 'us-west-2')" == pasta.dump(node)
 
-    method = "get_image_uri('us-west-2', 'sagemaker-rl-tensorflow')"
+    method = "sagemaker.amazon.amazon_estimator.get_image_uri('us-west-2', 'sagemaker-xgboost', repo_version='1')"
     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
+    assert "sagemaker.image_uris.retrieve('xgboost', 'us-west-2', '1')" == pasta.dump(node)
 
 
 def test_import_from_node_should_be_modified_image_uris_input(import_statements):