change: handle image_uri rename for Session methods

Author: laurenyu

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

@@ -68,6 +68,22 @@ def modify_node(self, node):
         keyword.arg = self.new_param_name
 
 
+class MethodParamRenamer(ParamRenamer):
+    """Abstract class to handle parameter renames for methods that belong to objects."""
+
+    def node_should_be_modified(self, node):
+        """Checks if the node matches any of the relevant functions and
+        contains the parameter to be renamed.
+
+        This looks for a call of the form ``<object>.<method>``, and
+        assumes the method cannot be called on its own.
+        """
+        if isinstance(node.func, ast.Name):
+            return False
+
+        return super(MethodParamRenamer, self).node_should_be_modified(node)
+
+
 class DistributionParameterRenamer(ParamRenamer):
     """A class to rename the ``distributions`` attribute to ``distrbution`` in
     MXNet and TensorFlow estimators.
@@ -100,7 +116,7 @@ def new_param_name(self):
         return "distribution"
 
 
-class S3SessionRenamer(ParamRenamer):
+class S3SessionRenamer(MethodParamRenamer):
     """A class to rename the ``session`` attribute to ``sagemaker_session`` in
     ``S3Uploader`` and ``S3Downloader``.
 
@@ -139,15 +155,6 @@ def new_param_name(self):
         """The new name for the SageMaker session argument."""
         return "sagemaker_session"
 
-    def node_should_be_modified(self, node):
-        """Checks if the node is one of the S3 utility functions and
-        contains the ``session`` parameter.
-        """
-        if isinstance(node.func, ast.Name):
-            return False
-
-        return super(S3SessionRenamer, self).node_should_be_modified(node)
-
 
 class EstimatorImageURIRenamer(ParamRenamer):
     """A class to rename the ``image_name`` attribute to ``image_uri`` in estimators."""
@@ -208,4 +215,60 @@ def old_param_name(self):
     @property
     def new_param_name(self):
         """The new name for the image URI argument."""
+
+
+class SessionCreateModelImageURIRenamer(MethodParamRenamer):
+    """A class to rename ``primary_container_image`` to ``image_uri``.
+
+    This looks for the following calls:
+
+    - ``sagemaker_session.create_model_from_job()``
+    - ``sess.create_model_from_job()``
+    """
+
+    @property
+    def calls_to_modify(self):
+        """A mapping of ``create_model_from_job`` to common variable names for Session."""
+        return {
+            "create_model_from_job": ("sagemaker_session", "sess"),
+        }
+
+    @property
+    def old_param_name(self):
+        """The previous name for the image URI argument."""
+        return "primary_container_image"
+
+    @property
+    def new_param_name(self):
+        """The new name for the the image URI argument."""
+        return "image_uri"
+
+
+class SessionCreateEndpointImageURIRenamer(MethodParamRenamer):
+    """A class to rename ``deployment_image`` to ``image_uri``.
+
+    This looks for the following calls:
+
+    - ``sagemaker_session.endpoint_from_job()``
+    - ``sess.endpoint_from_job()``
+    - ``sagemaker_session.endpoint_from_model_data()``
+    - ``sess.endpoint_from_model_data()``
+    """
+
+    @property
+    def calls_to_modify(self):
+        """A mapping of the ``endpoint_from_*`` functions to common variable names for Session."""
+        return {
+            "endpoint_from_job": ("sagemaker_session", "sess"),
+            "endpoint_from_model_data": ("sagemaker_session", "sess"),
+        }
+
+    @property
+    def old_param_name(self):
+        """The previous name for the image URI argument."""
+        return "deployment_image"
+
+    @property
+    def new_param_name(self):
+        """The new name for the the image URI argument."""
         return "image_uri"

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

@@ -0,0 +1,94 @@
+# 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
+
+CREATE_MODEL_TEMPLATES = (
+    "sagemaker_session.create_model_from_job({})",
+    "sess.create_model_from_job({})",
+)
+
+CREATE_ENDPOINT_TEMPLATES = (
+    "sagemaker_session.endpoint_from_job({})",
+    "sagemaker_session.endpoint_from_model_data({})",
+    "sess.endpoint_from_job({})",
+    "sess.endpoint_from_model_data({})",
+)
+
+
+def test_create_model_node_should_be_modified():
+    modifier = renamed_params.SessionCreateModelImageURIRenamer()
+
+    for template in CREATE_MODEL_TEMPLATES:
+        call = ast_call(template.format("primary_container_image=my_image"))
+        assert modifier.node_should_be_modified(call)
+
+
+def test_create_model_node_should_be_modified_no_image():
+    modifier = renamed_params.SessionCreateModelImageURIRenamer()
+
+    for template in CREATE_MODEL_TEMPLATES:
+        call = ast_call(template.format(""))
+        assert not modifier.node_should_be_modified(call)
+
+
+def test_create_model_node_should_be_modified_random_function_call():
+    modifier = renamed_params.SessionCreateModelImageURIRenamer()
+    assert not modifier.node_should_be_modified(ast_call("create_model()"))
+
+
+def test_create_model_modify_node():
+    modifier = renamed_params.SessionCreateModelImageURIRenamer()
+
+    for template in CREATE_MODEL_TEMPLATES:
+        call = ast_call(template.format("primary_container_image=my_image"))
+        modifier.modify_node(call)
+
+        expected = template.format("image_uri=my_image")
+        assert expected == pasta.dump(call)
+
+
+def test_create_endpoint_node_should_be_modified():
+    modifier = renamed_params.SessionCreateEndpointImageURIRenamer()
+
+    for template in CREATE_ENDPOINT_TEMPLATES:
+        call = ast_call(template.format("deployment_image=my_image"))
+        assert modifier.node_should_be_modified(call)
+
+
+def test_create_endpoint_node_should_be_modified_no_image():
+    modifier = renamed_params.SessionCreateEndpointImageURIRenamer()
+
+    for template in CREATE_ENDPOINT_TEMPLATES:
+        call = ast_call(template.format(""))
+        assert not modifier.node_should_be_modified(call)
+
+
+def test_create_endpoint_node_should_be_modified_random_function_call():
+    modifier = renamed_params.SessionCreateEndpointImageURIRenamer()
+    assert not modifier.node_should_be_modified(ast_call("create_endpoint()"))
+
+
+def test_create_endpoint_modify_node():
+    modifier = renamed_params.SessionCreateEndpointImageURIRenamer()
+
+    for template in CREATE_ENDPOINT_TEMPLATES:
+        call = ast_call(template.format("deployment_image=my_image"))
+        modifier.modify_node(call)
+
+        expected = template.format("image_uri=my_image")
+        assert expected == pasta.dump(call)