handle estimator.create_model()

Author: laurenyu

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

@@ -218,6 +218,39 @@ def new_param_name(self):
         return "image_uri"
 
 
+class EstimatorCreateModelImageURIRenamer(MethodParamRenamer):
+    """A class to rename ``image`` to ``image_uri`` in estimator ``create_model()`` methods."""
+
+    @property
+    def calls_to_modify(self):
+        """A mapping of ``create_model`` to common variable names for estimators."""
+        return {
+            "create_model": (
+                "estimator",
+                "chainer",
+                "mxnet",
+                "mx",
+                "pytorch",
+                "rl",
+                "sklearn",
+                "tensorflow",
+                "tf",
+                "xgboost",
+                "xgb",
+            )
+        }
+
+    @property
+    def old_param_name(self):
+        """The previous name for the image URI argument."""
+        return "image"
+
+    @property
+    def new_param_name(self):
+        """The new name for the the image URI argument."""
+        return "image_uri"
+
+
 class SessionCreateModelImageURIRenamer(MethodParamRenamer):
     """A class to rename ``primary_container_image`` to ``image_uri``.
 

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

@@ -0,0 +1,62 @@
+# 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 renamed_params
+from tests.unit.sagemaker.cli.compatibility.v2.modifiers.ast_converter import ast_call
+
+ESTIMATORS = (
+    "estimator",
+    "chainer",
+    "mxnet",
+    "mx",
+    "pytorch",
+    "rl",
+    "sklearn",
+    "tensorflow",
+    "tf",
+    "xgboost",
+    "xgb",
+)
+
+
+def test_node_should_be_modified():
+    modifier = renamed_params.EstimatorCreateModelImageURIRenamer()
+
+    for estimator in ESTIMATORS:
+        call = "{}.create_model(image='my-image:latest')".format(estimator)
+        assert modifier.node_should_be_modified(ast_call(call))
+
+
+def test_node_should_be_modified_no_distribution():
+    modifier = renamed_params.EstimatorCreateModelImageURIRenamer()
+
+    for estimator in ESTIMATORS:
+        call = "{}.create_model()".format(estimator)
+        assert not modifier.node_should_be_modified(ast_call(call))
+
+
+def test_node_should_be_modified_random_function_call():
+    modifier = renamed_params.EstimatorCreateModelImageURIRenamer()
+    assert not modifier.node_should_be_modified(ast_call("create_model()"))
+
+
+def test_modify_node():
+    node = ast_call("estimator.create_model(image=my_image)")
+    modifier = renamed_params.EstimatorCreateModelImageURIRenamer()
+    modifier.modify_node(node)
+
+    expected = "estimator.create_model(image_uri=my_image)"
+    assert expected == pasta.dump(node)