Merge remote-tracking branch 'aws/zwei' into endpoint-rename

Author: metrizable

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

@@ -23,6 +23,7 @@
     modifiers.tf_legacy_mode.TensorBoardParameterRemover(),
     modifiers.deprecated_params.TensorFlowScriptModeParameterRemover(),
     modifiers.tfs.TensorFlowServingConstructorRenamer(),
+    modifiers.airflow.ModelConfigArgModifier(),
 ]
 
 IMPORT_MODIFIERS = [modifiers.tfs.TensorFlowServingImportRenamer()]

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

@@ -14,6 +14,7 @@
 from __future__ import absolute_import
 
 from sagemaker.cli.compatibility.v2.modifiers import (  # noqa: F401 (imported but unused)
+    airflow,
     deprecated_params,
     framework_version,
     tf_legacy_mode,

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

@@ -0,0 +1,80 @@
+# 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.
+"""A class to handle argument changes for Airflow functions."""
+from __future__ import absolute_import
+
+import ast
+
+from sagemaker.cli.compatibility.v2.modifiers.modifier import Modifier
+
+
+class ModelConfigArgModifier(Modifier):
+    """A class to handle argument changes for Airflow model config functions."""
+
+    FUNCTION_NAMES = ("model_config", "model_config_from_estimator")
+
+    def node_should_be_modified(self, node):
+        """Checks if the ``ast.Call`` node creates an Airflow model config and
+        contains positional arguments.
+
+        This looks for the following formats:
+
+        - ``model_config``
+        - ``airflow.model_config``
+        - ``workflow.airflow.model_config``
+        - ``sagemaker.workflow.airflow.model_config``
+
+        where ``model_config`` is either ``model_config`` or ``model_config_from_estimator``.
+
+        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 either a ``model_config`` call or
+                a ``model_config_from_estimator`` call and has positional arguments.
+        """
+        return self._is_model_config_call(node) and len(node.args) > 0
+
+    def _is_model_config_call(self, node):
+        """Checks if the node is a ``model_config`` or  ``model_config_from_estimator`` call."""
+        if isinstance(node.func, ast.Name):
+            return node.func.id in self.FUNCTION_NAMES
+
+        if not (isinstance(node.func, ast.Attribute) and node.func.attr in self.FUNCTION_NAMES):
+            return False
+
+        return self._is_in_module(node.func, "sagemaker.workflow.airflow".split("."))
+
+    def _is_in_module(self, node, module):
+        """Checks if the node is in the module, including partial matches to the module path."""
+        if isinstance(node.value, ast.Name):
+            return node.value.id == module[-1]
+
+        if isinstance(node.value, ast.Attribute) and node.value.attr == module[-1]:
+            return self._is_in_module(node.value, module[:-1])
+
+        return False
+
+    def modify_node(self, node):
+        """Modifies the ``ast.Call`` node's arguments.
+
+        The first argument, the instance type, is turned into a keyword arg,
+        leaving the second argument, the model, to be the first argument.
+
+        Args:
+            node (ast.Call): a node that represents either a ``model_config`` call or
+                a ``model_config_from_estimator`` call.
+        """
+        instance_type = node.args.pop(0)
+        node.keywords.append(ast.keyword(arg="instance_type", value=instance_type))

src/sagemaker/workflow/airflow.py

@@ -557,29 +557,28 @@ def prepare_framework_container_def(model, instance_type, s3_operations):
     return sagemaker.container_def(deploy_image, model.model_data, deploy_env)
 
 
-def model_config(instance_type, model, role=None, image=None):
+def model_config(model, instance_type=None, role=None, image=None):
     """Export Airflow model config from a SageMaker model
 
     Args:
+        model (sagemaker.model.Model): The Model object from which to export the Airflow config
         instance_type (str): The EC2 instance type to deploy this Model to. For
             example, 'ml.p2.xlarge'
-        model (sagemaker.model.FrameworkModel): The SageMaker model to export
-            Airflow config from
         role (str): The ``ExecutionRoleArn`` IAM Role ARN for the model
         image (str): An container image to use for deploying the model
 
     Returns:
         dict: Model config that can be directly used by SageMakerModelOperator
-        in Airflow. It can also be part of the config used by
-        SageMakerEndpointOperator and SageMakerTransformOperator in Airflow.
+            in Airflow. It can also be part of the config used by
+            SageMakerEndpointOperator and SageMakerTransformOperator in Airflow.
     """
     s3_operations = {}
     model.image = image or model.image
 
     if isinstance(model, sagemaker.model.FrameworkModel):
         container_def = prepare_framework_container_def(model, instance_type, s3_operations)
     else:
-        container_def = model.prepare_container_def(instance_type)
+        container_def = model.prepare_container_def()
         base_name = utils.base_name_from_image(container_def["Image"])
         model.name = model.name or utils.name_from_base(base_name)
 
@@ -601,10 +600,10 @@ def model_config(instance_type, model, role=None, image=None):
 
 
 def model_config_from_estimator(
-    instance_type,
     estimator,
     task_id,
     task_type,
+    instance_type=None,
     role=None,
     image=None,
     name=None,
@@ -614,8 +613,6 @@ def model_config_from_estimator(
     """Export Airflow model config from a SageMaker estimator
 
     Args:
-        instance_type (str): The EC2 instance type to deploy this Model to. For
-            example, 'ml.p2.xlarge'
         estimator (sagemaker.model.EstimatorBase): The SageMaker estimator to
             export Airflow config from. It has to be an estimator associated
             with a training job.
@@ -627,6 +624,8 @@ def model_config_from_estimator(
         task_type (str): Whether the task is from SageMakerTrainingOperator or
             SageMakerTuningOperator. Values can be 'training', 'tuning' or None
             (which means training job is not from any task).
+        instance_type (str): The EC2 instance type to deploy this Model to. For
+            example, 'ml.p2.xlarge'
         role (str): The ``ExecutionRoleArn`` IAM Role ARN for the model
         image (str): An container image to use for deploying the model
         name (str): Name of the model
@@ -667,7 +666,7 @@ def model_config_from_estimator(
         )
     model.name = name
 
-    return model_config(instance_type, model, role, image)
+    return model_config(model, instance_type, role, image)
 
 
 def transform_config(
@@ -914,10 +913,10 @@ def transform_config_from_estimator(
         SageMakerTransformOperator in Airflow.
     """
     model_base_config = model_config_from_estimator(
-        instance_type=instance_type,
         estimator=estimator,
         task_id=task_id,
         task_type=task_type,
+        instance_type=instance_type,
         role=role,
         image=image,
         name=model_name,
@@ -997,7 +996,7 @@ def deploy_config(model, initial_instance_count, instance_type, endpoint_name=No
         dict: Deploy config that can be directly used by
         SageMakerEndpointOperator in Airflow.
     """
-    model_base_config = model_config(instance_type, model)
+    model_base_config = model_config(model, instance_type)
 
     production_variant = sagemaker.production_variant(
         model.name, instance_type, initial_instance_count

tests/conftest.py

@@ -44,22 +44,6 @@ def pytest_addoption(parser):
     parser.addoption("--sagemaker-client-config", action="store", default=None)
     parser.addoption("--sagemaker-runtime-config", action="store", default=None)
     parser.addoption("--boto-config", action="store", default=None)
-    parser.addoption("--chainer-full-version", action="store", default="5.0.0")
-    parser.addoption("--ei-mxnet-full-version", action="store", default="1.5.1")
-    parser.addoption(
-        "--rl-coach-mxnet-full-version",
-        action="store",
-        default=RLEstimator.COACH_LATEST_VERSION_MXNET,
-    )
-    parser.addoption(
-        "--rl-coach-tf-full-version", action="store", default=RLEstimator.COACH_LATEST_VERSION_TF
-    )
-    parser.addoption(
-        "--rl-ray-full-version", action="store", default=RLEstimator.RAY_LATEST_VERSION
-    )
-    parser.addoption("--sklearn-full-version", action="store", default="0.20.0")
-    parser.addoption("--ei-tf-full-version", action="store")
-    parser.addoption("--xgboost-full-version", action="store", default="1.0-1")
 
 
 def pytest_configure(config):
@@ -249,8 +233,13 @@ def rl_ray_version(request):
 
 
 @pytest.fixture(scope="module")
-def chainer_full_version(request):
-    return request.config.getoption("--chainer-full-version")
+def chainer_full_version():
+    return "5.0.0"
+
+
+@pytest.fixture(scope="module")
+def chainer_full_py_version():
+    return "py3"
 
 
 @pytest.fixture(scope="module")
@@ -264,8 +253,8 @@ def mxnet_full_py_version():
 
 
 @pytest.fixture(scope="module")
-def ei_mxnet_full_version(request):
-    return request.config.getoption("--ei-mxnet-full-version")
+def ei_mxnet_full_version():
+    return "1.5.1"
 
 
 @pytest.fixture(scope="module")
@@ -284,23 +273,28 @@ def pytorch_full_ei_version():
 
 
 @pytest.fixture(scope="module")
-def rl_coach_mxnet_full_version(request):
-    return request.config.getoption("--rl-coach-mxnet-full-version")
+def rl_coach_mxnet_full_version():
+    return RLEstimator.COACH_LATEST_VERSION_MXNET
+
+
+@pytest.fixture(scope="module")
+def rl_coach_tf_full_version():
+    return RLEstimator.COACH_LATEST_VERSION_TF
 
 
 @pytest.fixture(scope="module")
-def rl_coach_tf_full_version(request):
-    return request.config.getoption("--rl-coach-tf-full-version")
+def rl_ray_full_version():
+    return RLEstimator.RAY_LATEST_VERSION
 
 
 @pytest.fixture(scope="module")
-def rl_ray_full_version(request):
-    return request.config.getoption("--rl-ray-full-version")
+def sklearn_full_version():
+    return "0.20.0"
 
 
 @pytest.fixture(scope="module")
-def sklearn_full_version(request):
-    return request.config.getoption("--sklearn-full-version")
+def sklearn_full_py_version():
+    return "py3"
 
 
 @pytest.fixture(scope="module")
@@ -343,13 +337,19 @@ def tf_full_py_version(tf_full_version):
     return "py37"
 
 
-@pytest.fixture(scope="module", params=["1.15.0", "2.0.0"])
-def ei_tf_full_version(request):
-    tf_ei_version = request.config.getoption("--ei-tf-full-version")
-    if tf_ei_version is None:
-        return request.param
-    else:
-        tf_ei_version
+@pytest.fixture(scope="module")
+def ei_tf_full_version():
+    return "2.0.0"
+
+
+@pytest.fixture(scope="module")
+def xgboost_full_version():
+    return "1.0-1"
+
+
+@pytest.fixture(scope="module")
+def xgboost_full_py_version():
+    return "py3"
 
 
 @pytest.fixture(scope="session")
@@ -405,8 +405,3 @@ def pytest_generate_tests(metafunc):
         ):
             params.append("ml.p2.xlarge")
         metafunc.parametrize("instance_type", params, scope="session")
-
-
-@pytest.fixture(scope="module")
-def xgboost_full_version(request):
-    return request.config.getoption("--xgboost-full-version")

tests/integ/__init__.py

@@ -14,7 +14,6 @@
 
 import logging
 import os
-import sys
 
 import boto3
 
@@ -23,7 +22,6 @@
 TUNING_DEFAULT_TIMEOUT_MINUTES = 20
 TRANSFORM_DEFAULT_TIMEOUT_MINUTES = 20
 AUTO_ML_DEFAULT_TIMEMOUT_MINUTES = 60
-PYTHON_VERSION = "py{}".format(sys.version_info.major)
 
 # these regions have some p2 and p3 instances, but not enough for continuous testing
 HOSTING_NO_P2_REGIONS = [