Add region account map to marketplace test for canary (#688)

Author: yangaws

tests/integ/marketplace_utils.py

@@ -0,0 +1,29 @@
+# Copyright 2018 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
+
+REGION_ACCOUNT_MAP = {
+    'us-east-1': '865070037744',
+    'us-east-2': '057799348421',
+    'us-west-2': '594846645681',
+    'eu-west-1': '985815980388',
+    'eu-central-1': '446921602837',
+    'ap-northeast-1': '977537786026',
+    'ap-northeast-2': '745090734665',
+    'ap-southeast-2': '666831318237',
+    'ap-southeast-1': '192199979996',
+    'ap-south-1': '077584701553',
+    'ca-central-1': '470592106596',
+    'eu-west-2': '856760150666',
+    'us-west-1': '382657785993'
+}

tests/integ/test_marketplace.py

@@ -25,6 +25,7 @@
 from sagemaker.utils import sagemaker_timestamp
 from tests.integ import DATA_DIR
 from tests.integ.timeout import timeout, timeout_and_delete_endpoint_by_name
+from tests.integ.marketplace_utils import REGION_ACCOUNT_MAP
 
 
 # All these tests require a manual 1 time subscription to the following Marketplace items:
@@ -36,20 +37,23 @@
 #
 # Both are  written by Amazon and are free to subscribe.
 
-ALGORITHM_ARN = 'arn:aws:sagemaker:%s:594846645681:algorithm/scikit-decision-trees-' \
+ALGORITHM_ARN = 'arn:aws:sagemaker:%s:%s:algorithm/scikit-decision-trees-' \
                 '15423055-57b73412d2e93e9239e4e16f83298b8f'
 
-MODEL_PACKAGE_ARN = 'arn:aws:sagemaker:%s:594846645681:model-package/scikit-iris-detector-' \
+MODEL_PACKAGE_ARN = 'arn:aws:sagemaker:%s:%s:model-package/scikit-iris-detector-' \
                     '154230595-8f00905c1f927a512b73ea29dd09ae30'
 
 
 @pytest.mark.canary_quick
 def test_marketplace_estimator(sagemaker_session):
     with timeout(minutes=15):
         data_path = os.path.join(DATA_DIR, 'marketplace', 'training')
+        region = sagemaker_session.boto_region_name
+        account = REGION_ACCOUNT_MAP[region]
+        algorithm_arn = ALGORITHM_ARN % (region, account)
 
         algo = AlgorithmEstimator(
-            algorithm_arn=(ALGORITHM_ARN % sagemaker_session.boto_region_name),
+            algorithm_arn=algorithm_arn,
             role='SageMakerRole',
             train_instance_count=1,
             train_instance_type='ml.c4.xlarge',
@@ -78,9 +82,12 @@ def test_marketplace_estimator(sagemaker_session):
 def test_marketplace_attach(sagemaker_session):
     with timeout(minutes=15):
         data_path = os.path.join(DATA_DIR, 'marketplace', 'training')
+        region = sagemaker_session.boto_region_name
+        account = REGION_ACCOUNT_MAP[region]
+        algorithm_arn = ALGORITHM_ARN % (region, account)
 
         mktplace = AlgorithmEstimator(
-            algorithm_arn=(ALGORITHM_ARN % sagemaker_session.boto_region_name),
+            algorithm_arn=algorithm_arn,
             role='SageMakerRole',
             train_instance_count=1,
             train_instance_type='ml.c4.xlarge',
@@ -116,14 +123,17 @@ def test_marketplace_attach(sagemaker_session):
 
 @pytest.mark.canary_quick
 def test_marketplace_model(sagemaker_session):
+    region = sagemaker_session.boto_region_name
+    account = REGION_ACCOUNT_MAP[region]
+    model_package_arn = MODEL_PACKAGE_ARN % (region, account)
 
     def predict_wrapper(endpoint, session):
         return sagemaker.RealTimePredictor(
             endpoint, session, serializer=sagemaker.predictor.csv_serializer
         )
 
     model = ModelPackage(role='SageMakerRole',
-                         model_package_arn=(MODEL_PACKAGE_ARN % sagemaker_session.boto_region_name),
+                         model_package_arn=model_package_arn,
                          sagemaker_session=sagemaker_session,
                          predictor_cls=predict_wrapper)
 
@@ -144,9 +154,12 @@ def predict_wrapper(endpoint, session):
 
 def test_marketplace_tuning_job(sagemaker_session):
     data_path = os.path.join(DATA_DIR, 'marketplace', 'training')
+    region = sagemaker_session.boto_region_name
+    account = REGION_ACCOUNT_MAP[region]
+    algorithm_arn = ALGORITHM_ARN % (region, account)
 
     mktplace = AlgorithmEstimator(
-        algorithm_arn=(ALGORITHM_ARN % sagemaker_session.boto_region_name),
+        algorithm_arn=algorithm_arn,
         role='SageMakerRole',
         train_instance_count=1,
         train_instance_type='ml.c4.xlarge',
@@ -172,9 +185,12 @@ def test_marketplace_tuning_job(sagemaker_session):
 
 def test_marketplace_transform_job(sagemaker_session):
     data_path = os.path.join(DATA_DIR, 'marketplace', 'training')
+    region = sagemaker_session.boto_region_name
+    account = REGION_ACCOUNT_MAP[region]
+    algorithm_arn = ALGORITHM_ARN % (region, account)
 
     algo = AlgorithmEstimator(
-        algorithm_arn=(ALGORITHM_ARN % sagemaker_session.boto_region_name),
+        algorithm_arn=algorithm_arn,
         role='SageMakerRole',
         train_instance_count=1,
         train_instance_type='ml.c4.xlarge',
@@ -209,8 +225,12 @@ def test_marketplace_transform_job_from_model_package(sagemaker_session):
         TRANSFORM_WORKDIR,
         key_prefix='integ-test-data/marketplace/transform')
 
+    region = sagemaker_session.boto_region_name
+    account = REGION_ACCOUNT_MAP[region]
+    model_package_arn = MODEL_PACKAGE_ARN % (region, account)
+
     model = ModelPackage(role='SageMakerRole',
-                         model_package_arn=(MODEL_PACKAGE_ARN % sagemaker_session.boto_region_name),
+                         model_package_arn=model_package_arn,
                          sagemaker_session=sagemaker_session)
 
     transformer = model.transformer(1, 'ml.m4.xlarge')