feature: support cn-north-1 and cn-northwest-1 (#1380)

knakad · web-flow · commit c5890d25116d · 2020-03-31T12:01:20.000-07:00
diff --git a/.github/PULL_REQUEST_TEMPLATE.md b/.github/PULL_REQUEST_TEMPLATE.md
@@ -12,7 +12,7 @@ _Put an `x` in the boxes that apply. You can also fill these out after creating
 
 - [ ] I have read the [CONTRIBUTING](https://github.com/aws/sagemaker-python-sdk/blob/master/CONTRIBUTING.md) doc
 - [ ] I used the commit message format described in [CONTRIBUTING](https://github.com/aws/sagemaker-python-sdk/blob/master/CONTRIBUTING.md#committing-your-change)
-- [ ] I have used the regional endpoint when creating S3 and/or STS clients (if appropriate)
+- [ ] I have passed the region in to any/all clients that I've initialized as part of this change.
 - [ ] I have updated any necessary documentation, including [READMEs](https://github.com/aws/sagemaker-python-sdk/blob/master/README.rst) and [API docs](https://github.com/aws/sagemaker-python-sdk/tree/master/doc) (if appropriate)
 
 #### Tests
diff --git a/src/sagemaker/amazon/amazon_estimator.py b/src/sagemaker/amazon/amazon_estimator.py
@@ -281,7 +281,9 @@ def record_set(self, train, labels=None, channel="train", encrypt=False):
             RecordSet: A RecordSet referencing the encoded, uploading training
             and label data.
         """
-        s3 = self.sagemaker_session.boto_session.resource("s3")
+        s3 = self.sagemaker_session.boto_session.resource(
+            "s3", region_name=self.sagemaker_session.boto_region_name
+        )
         parsed_s3_url = urlparse(self.data_location)
         bucket, key_prefix = parsed_s3_url.netloc, parsed_s3_url.path
         key_prefix = key_prefix + "{}-{}/".format(type(self).__name__, sagemaker_timestamp())
@@ -467,9 +469,14 @@ def registry(region_name, algorithm=None):
     https://github.com/aws/sagemaker-python-sdk/tree/master/src/sagemaker/amazon
 
     Args:
-        region_name:
-        algorithm:
+        region_name (str): The region name for the account.
+        algorithm (str): The algorithm for the account.
+
+    Raises:
+        ValueError: If invalid algorithm passed in or if mapping does not exist for given algorithm
+            and region.
     """
+    region_to_accounts = {}
     if algorithm in [
         None,
         "pca",
@@ -482,7 +489,7 @@ def registry(region_name, algorithm=None):
         "object2vec",
         "ipinsights",
     ]:
-        account_id = {
+        region_to_accounts = {
             "us-east-1": "382416733822",
             "us-east-2": "404615174143",
             "us-west-2": "174872318107",
@@ -503,9 +510,11 @@ def registry(region_name, algorithm=None):
             "eu-west-3": "749696950732",
             "sa-east-1": "855470959533",
             "me-south-1": "249704162688",
-        }[region_name]
+            "cn-north-1": "390948362332",
+            "cn-northwest-1": "387376663083",
+        }
     elif algorithm in ["lda"]:
-        account_id = {
+        region_to_accounts = {
             "us-east-1": "766337827248",
             "us-east-2": "999911452149",
             "us-west-2": "266724342769",
@@ -521,9 +530,9 @@ def registry(region_name, algorithm=None):
             "eu-west-2": "644912444149",
             "us-west-1": "632365934929",
             "us-iso-east-1": "490574956308",
-        }[region_name]
+        }
     elif algorithm in ["forecasting-deepar"]:
-        account_id = {
+        region_to_accounts = {
             "us-east-1": "522234722520",
             "us-east-2": "566113047672",
             "us-west-2": "156387875391",
@@ -544,7 +553,9 @@ def registry(region_name, algorithm=None):
             "eu-west-3": "749696950732",
             "sa-east-1": "855470959533",
             "me-south-1": "249704162688",
-        }[region_name]
+            "cn-north-1": "390948362332",
+            "cn-northwest-1": "387376663083",
+        }
     elif algorithm in [
         "xgboost",
         "seq2seq",
@@ -553,7 +564,7 @@ def registry(region_name, algorithm=None):
         "object-detection",
         "semantic-segmentation",
     ]:
-        account_id = {
+        region_to_accounts = {
             "us-east-1": "811284229777",
             "us-east-2": "825641698319",
             "us-west-2": "433757028032",
@@ -574,15 +585,25 @@ def registry(region_name, algorithm=None):
             "eu-west-3": "749696950732",
             "sa-east-1": "855470959533",
             "me-south-1": "249704162688",
-        }[region_name]
+            "cn-north-1": "390948362332",
+            "cn-northwest-1": "387376663083",
+        }
     elif algorithm in ["image-classification-neo", "xgboost-neo"]:
-        account_id = NEO_IMAGE_ACCOUNT[region_name]
+        region_to_accounts = NEO_IMAGE_ACCOUNT
     else:
         raise ValueError(
             "Algorithm class:{} does not have mapping to account_id with images".format(algorithm)
         )
 
-    return get_ecr_image_uri_prefix(account_id, region_name)
+    if region_name in region_to_accounts:
+        account_id = region_to_accounts[region_name]
+        return get_ecr_image_uri_prefix(account_id, region_name)
+
+    raise ValueError(
+        "Algorithm ({algorithm}) is unsupported for region ({region_name}).".format(
+            algorithm=algorithm, region_name=region_name
+        )
+    )
 
 
 def get_image_uri(region_name, repo_name, repo_version=1):
diff --git a/src/sagemaker/debugger.py b/src/sagemaker/debugger.py
@@ -23,6 +23,7 @@
 
 import smdebug_rulesconfig as rule_configs  # noqa: F401 # pylint: disable=unused-import
 
+from sagemaker.utils import get_ecr_image_uri_prefix
 
 RULES_ECR_REPO_NAME = "sagemaker-debugger-rules"
 
@@ -45,6 +46,8 @@
     "ap-southeast-1": {RULES_ECR_REPO_NAME: "972752614525"},
     "ap-southeast-2": {RULES_ECR_REPO_NAME: "184798709955"},
     "ca-central-1": {RULES_ECR_REPO_NAME: "519511493484"},
+    "cn-north-1": {RULES_ECR_REPO_NAME: "618459771430"},
+    "cn-northwest-1": {RULES_ECR_REPO_NAME: "658757709296"},
 }
 
 
@@ -59,7 +62,8 @@ def get_rule_container_image_uri(region):
         str: Formatted image uri for the given region and the rule container type
     """
     registry_id = SAGEMAKER_RULE_CONTAINERS_ACCOUNTS_MAP.get(region).get(RULES_ECR_REPO_NAME)
-    return "{}.dkr.ecr.{}.amazonaws.com/{}:latest".format(registry_id, region, RULES_ECR_REPO_NAME)
+    image_uri_prefix = get_ecr_image_uri_prefix(registry_id, region)
+    return "{}/{}:latest".format(image_uri_prefix, RULES_ECR_REPO_NAME)
 
 
 class Rule(object):
diff --git a/src/sagemaker/fw_registry.py b/src/sagemaker/fw_registry.py
@@ -117,6 +117,16 @@
         "scikit-learn": "801668240914",
         "xgboost": "801668240914",
     },
+    "cn-north-1": {
+        "sparkml-serving": "450853457545",
+        "scikit-learn": "450853457545",
+        "xgboost": "450853457545",
+    },
+    "cn-northwest-1": {
+        "sparkml-serving": "451049120500",
+        "scikit-learn": "451049120500",
+        "xgboost": "451049120500",
+    },
 }
 
 
diff --git a/src/sagemaker/fw_utils.py b/src/sagemaker/fw_utils.py
@@ -71,8 +71,18 @@
     "pytorch-serving",
 ]
 PY2_RESTRICTED_EIA_FRAMEWORKS = ["pytorch-serving"]
-VALID_ACCOUNTS_BY_REGION = {"us-gov-west-1": "246785580436", "us-iso-east-1": "744548109606"}
-ASIMOV_VALID_ACCOUNTS_BY_REGION = {"us-gov-west-1": "442386744353", "us-iso-east-1": "886529160074"}
+VALID_ACCOUNTS_BY_REGION = {
+    "us-gov-west-1": "246785580436",
+    "us-iso-east-1": "744548109606",
+    "cn-north-1": "422961961927",
+    "cn-northwest-1": "423003514399",
+}
+ASIMOV_VALID_ACCOUNTS_BY_REGION = {
+    "us-gov-west-1": "442386744353",
+    "us-iso-east-1": "886529160074",
+    "cn-north-1": "727897471807",
+    "cn-northwest-1": "727897471807",
+}
 OPT_IN_ACCOUNTS_BY_REGION = {"ap-east-1": "057415533634", "me-south-1": "724002660598"}
 ASIMOV_OPT_IN_ACCOUNTS_BY_REGION = {"ap-east-1": "871362719292", "me-south-1": "217643126080"}
 DEFAULT_ACCOUNT = "520713654638"
diff --git a/src/sagemaker/model.py b/src/sagemaker/model.py
@@ -47,6 +47,8 @@
     "sa-east-1": "756306329178",
     "ca-central-1": "464438896020",
     "me-south-1": "836785723513",
+    "cn-north-1": "472730292857",
+    "cn-northwest-1": "474822919863",
     "us-gov-west-1": "263933020539",
 }
 
diff --git a/src/sagemaker/model_monitor/model_monitoring.py b/src/sagemaker/model_monitor/model_monitoring.py
@@ -32,11 +32,9 @@
 from sagemaker.processing import Processor, ProcessingInput, ProcessingJob, ProcessingOutput
 from sagemaker.s3 import S3Uploader
 from sagemaker.session import Session
-from sagemaker.utils import name_from_base, retries
+from sagemaker.utils import name_from_base, retries, get_ecr_image_uri_prefix
 
-_DEFAULT_MONITOR_IMAGE_URI_WITH_PLACEHOLDERS = (
-    "{}.dkr.ecr.{}.amazonaws.com/sagemaker-model-monitor-analyzer"
-)
+_DEFAULT_MONITOR_IMAGE_URI_WITH_PLACEHOLDERS = "{}/sagemaker-model-monitor-analyzer"
 
 _DEFAULT_MONITOR_IMAGE_REGION_ACCOUNT_MAPPING = {
     "eu-north-1": "895015795356",
@@ -57,6 +55,8 @@
     "ap-southeast-1": "245545462676",
     "ap-southeast-2": "563025443158",
     "ca-central-1": "536280801234",
+    "cn-north-1": "453000072557",
+    "cn-northwest-1": "453252182341",
 }
 
 STATISTICS_JSON_DEFAULT_FILE_NAME = "statistics.json"
@@ -1761,7 +1761,7 @@ def _get_default_image_uri(region):
             str: The Default Model Monitoring image uri based on the region.
         """
         return _DEFAULT_MONITOR_IMAGE_URI_WITH_PLACEHOLDERS.format(
-            _DEFAULT_MONITOR_IMAGE_REGION_ACCOUNT_MAPPING[region], region
+            get_ecr_image_uri_prefix(_DEFAULT_MONITOR_IMAGE_REGION_ACCOUNT_MAPPING[region], region)
         )
 
 
diff --git a/src/sagemaker/utils.py b/src/sagemaker/utils.py
@@ -509,9 +509,9 @@ def _save_model(repacked_model_uri, tmp_model_path, sagemaker_session, kms_key):
             extra_args = {"ServerSideEncryption": "aws:kms", "SSEKMSKeyId": kms_key}
         else:
             extra_args = None
-        sagemaker_session.boto_session.resource("s3").Object(bucket, new_key).upload_file(
-            tmp_model_path, ExtraArgs=extra_args
-        )
+        sagemaker_session.boto_session.resource(
+            "s3", region_name=sagemaker_session.boto_region_name
+        ).Object(bucket, new_key).upload_file(tmp_model_path, ExtraArgs=extra_args)
     else:
         shutil.move(tmp_model_path, repacked_model_uri.replace("file://", ""))
 
@@ -604,7 +604,7 @@ def download_file(bucket_name, path, target, sagemaker_session):
     path = path.lstrip("/")
     boto_session = sagemaker_session.boto_session
 
-    s3 = boto_session.resource("s3")
+    s3 = boto_session.resource("s3", region_name=sagemaker_session.boto_region_name)
     bucket = s3.Bucket(bucket_name)
     bucket.download_file(path, target)
 
diff --git a/tests/integ/kms_utils.py b/tests/integ/kms_utils.py
@@ -141,7 +141,7 @@ def get_or_create_kms_key(
       "Resource": "arn:{partition}:s3:::{bucket_name}/*",
       "Condition": {{
         "StringNotEquals": {{
-          "s3:x-amz-server-side-encryption": "{partition}:kms"
+          "s3:x-amz-server-side-encryption": "aws:kms"
         }}
       }}
     }},
@@ -172,7 +172,7 @@ def bucket_with_encryption(sagemaker_session, sagemaker_role):
     account = sts_client.get_caller_identity()["Account"]
     role_arn = sts_client.get_caller_identity()["Arn"]
 
-    kms_client = boto_session.client("kms")
+    kms_client = boto_session.client("kms", region_name=region)
     kms_key_arn = _create_kms_key(kms_client, account, region, role_arn, sagemaker_role, None)
 
     region = boto_session.region_name
@@ -187,9 +187,7 @@ def bucket_with_encryption(sagemaker_session, sagemaker_role):
             "Rules": [
                 {
                     "ApplyServerSideEncryptionByDefault": {
-                        "SSEAlgorithm": "{partition}:kms".format(
-                            partition=utils._aws_partition(region)
-                        ),
+                        "SSEAlgorithm": "aws:kms",
                         "KMSMasterKeyID": kms_key_arn,
                     }
                 }
diff --git a/tests/integ/marketplace_utils.py b/tests/integ/marketplace_utils.py
@@ -30,4 +30,7 @@
     "eu-north-1": "136758871317",
     "sa-east-1": "270155090741",
     "ap-east-1": "822005858737",
+    "me-south-1": "335155493544",
+    "cn-north-1": "295401494951",
+    "cn-northwest-1": "304690803264",
 }
diff --git a/tests/integ/test_debugger.py b/tests/integ/test_debugger.py
@@ -53,6 +53,8 @@
     "us-east-2": "840043622174",
     "us-west-1": "952348334681",
     "us-west-2": "759209512951",
+    "cn-north-1": "617202126805",
+    "cn-northwest-1": "658559488188",
 }
 
 # TODO-reinvent-2019: test get_debugger_artifacts_path and get_tensorboard_artifacts_path
diff --git a/tests/integ/test_horovod.py b/tests/integ/test_horovod.py
@@ -86,9 +86,9 @@ def read_json(file, tmp):
         return json.load(f)
 
 
-def extract_files_from_s3(s3_url, tmpdir):
+def extract_files_from_s3(s3_url, tmpdir, sagemaker_session):
     parsed_url = urlparse(s3_url)
-    s3 = boto3.resource("s3")
+    s3 = boto3.resource("s3", region_name=sagemaker_session.boto_region_name)
 
     model = os.path.join(tmpdir, "model")
     s3.Bucket(parsed_url.netloc).download_file(parsed_url.path.lstrip("/"), model)
@@ -115,7 +115,7 @@ def _create_and_fit_estimator(sagemaker_session, instance_type, tmpdir):
         estimator.fit(job_name=job_name)
 
         tmp = str(tmpdir)
-        extract_files_from_s3(estimator.model_data, tmp)
+        extract_files_from_s3(estimator.model_data, tmp, sagemaker_session)
 
         for rank in range(2):
             assert read_json("rank-%s" % rank, tmp)["rank"] == rank
diff --git a/tests/integ/test_model_monitor.py b/tests/integ/test_model_monitor.py
@@ -61,7 +61,7 @@
 DEFAULT_VOLUME_SIZE_IN_GB = 30
 DEFAULT_BASELINING_MAX_RUNTIME_IN_SECONDS = 86400
 DEFAULT_EXECUTION_MAX_RUNTIME_IN_SECONDS = 3600
-DEFAULT_IMAGE_SUFFIX = ".com/sagemaker-model-monitor-analyzer"
+DEFAULT_IMAGE_SUFFIX = "/sagemaker-model-monitor-analyzer"
 
 UPDATED_ROLE = "SageMakerRole"
 UPDATED_INSTANCE_COUNT = 2
diff --git a/tests/integ/test_multidatamodel.py b/tests/integ/test_multidatamodel.py
@@ -27,7 +27,7 @@
 from sagemaker.multidatamodel import MultiDataModel
 from sagemaker.mxnet import MXNet
 from sagemaker.predictor import RealTimePredictor, StringDeserializer, npy_serializer
-from sagemaker.utils import sagemaker_timestamp, unique_name_from_base
+from sagemaker.utils import sagemaker_timestamp, unique_name_from_base, get_ecr_image_uri_prefix
 from tests.integ import DATA_DIR, PYTHON_VERSION, TRAINING_DEFAULT_TIMEOUT_MINUTES
 from tests.integ.retry import retries
 from tests.integ.timeout import timeout, timeout_and_delete_endpoint_by_name
@@ -49,8 +49,9 @@ def container_image(sagemaker_session):
     )
     account_id = sts_client.get_caller_identity()["Account"]
     algorithm_name = "sagemaker-multimodel-integ-test-{}".format(sagemaker_timestamp())
-    ecr_image = "{account}.dkr.ecr.{region}.amazonaws.com/{algorithm_name}:latest".format(
-        account=account_id, region=region, algorithm_name=algorithm_name
+    ecr_image_uri_prefix = get_ecr_image_uri_prefix(account=account_id, region=region)
+    ecr_image = "{prefix}/{algorithm_name}:latest".format(
+        prefix=ecr_image_uri_prefix, algorithm_name=algorithm_name
     )
 
     # Build and tag docker image locally
diff --git a/tests/integ/test_session.py b/tests/integ/test_session.py
@@ -46,5 +46,5 @@ def test_sagemaker_session_does_not_create_bucket_on_init(
         default_bucket=CUSTOM_BUCKET_NAME,
     )
 
-    s3 = boto3.resource("s3")
+    s3 = boto3.resource("s3", region_name=DEFAULT_REGION)
     assert s3.Bucket(CUSTOM_BUCKET_NAME).creation_date is None
diff --git a/tests/unit/test_amazon_estimator.py b/tests/unit/test_amazon_estimator.py
@@ -462,3 +462,15 @@ def test_get_xgboost_image_uri():
         updated_xgb_image_uri_v2
         == "246618743249.dkr.ecr.us-west-2.amazonaws.com/sagemaker-xgboost:0.90-2-cpu-py3"
     )
+
+
+def test_regitry_throws_error_if_mapping_does_not_exist_for_lda():
+    with pytest.raises(ValueError) as error:
+        registry("cn-north-1", "lda")
+    assert "Algorithm (lda) is unsupported for region (cn-north-1)." in str(error)
+
+
+def test_regitry_throws_error_if_mapping_does_not_exist_for_default_algorithm():
+    with pytest.raises(ValueError) as error:
+        registry("broken_region_name")
+    assert "Algorithm (None) is unsupported for region (broken_region_name)." in str(error)
diff --git a/tests/unit/test_fw_utils.py b/tests/unit/test_fw_utils.py
@@ -279,6 +279,24 @@ def test_create_image_uri_bah():
     }
 
 
+def test_create_image_uri_cn_north_1():
+    image_uri = fw_utils.create_image_uri(
+        "cn-north-1", MOCK_FRAMEWORK, "ml.p3.2xlarge", "1.0rc", "py3"
+    )
+    assert {
+        image_uri == "727897471807.dkr.ecr.me-south-1.amazonaws.com/sagemaker-mlfw:1.0rc-gpu-py3"
+    }
+
+
+def test_create_image_uri_cn_northwest_1():
+    image_uri = fw_utils.create_image_uri(
+        "cn-northwest-1", MOCK_FRAMEWORK, "ml.p3.2xlarge", "1.0rc", "py3"
+    )
+    assert {
+        image_uri == "727897471807.dkr.ecr.me-south-1.amazonaws.com/sagemaker-mlfw:1.0rc-gpu-py3"
+    }
+
+
 def test_tf_eia_images():
     image_uri = fw_utils.create_image_uri(
         "us-west-2",
diff --git a/tests/unit/test_model.py b/tests/unit/test_model.py

Original file line number	Diff line number	Diff line change
`@@ -47,6 +47,8 @@`
`47`	`47`	`"sa-east-1": "756306329178",`
`48`	`48`	`"ca-central-1": "464438896020",`
`49`	`49`	`"me-south-1": "836785723513",`
	`50`	`+ "cn-north-1": "472730292857",`
	`51`	`+ "cn-northwest-1": "474822919863",`
`50`	`52`	`"us-gov-west-1": "263933020539",`
`51`	`53`	`}`
`52`	`54`
Original file line number	Diff line number	Diff line change
`@@ -30,4 +30,7 @@`
`30`	`30`	`"eu-north-1": "136758871317",`
`31`	`31`	`"sa-east-1": "270155090741",`
`32`	`32`	`"ap-east-1": "822005858737",`
	`33`	`+ "me-south-1": "335155493544",`
	`34`	`+ "cn-north-1": "295401494951",`
	`35`	`+ "cn-northwest-1": "304690803264",`
`33`	`36`	`}`