infra: programmatically determine partition based on region (#1316)

Author: knakad

src/sagemaker/utils.py

@@ -15,6 +15,7 @@
 
 import contextlib
 import errno
+import logging
 import os
 import random
 import re
@@ -38,6 +39,8 @@
 HTTPS_PREFIX = "https://"
 DEFAULT_SLEEP_TIME_SECONDS = 10
 
+logger = logging.getLogger(__name__)
+
 
 # Use the base name of the image as the job name if the user doesn't give us one
 def name_from_image(image):
@@ -667,6 +670,21 @@ def _botocore_resolver():
     return botocore.regions.EndpointResolver(loader.load_data("endpoints"))
 
 
+def _aws_partition(region):
+    """
+    Given a region name (ex: "cn-north-1"), return the corresponding aws partition ("aws-cn").
+
+    Args:
+        region (str): The region name for which to return the corresponding partition.
+        Ex: "cn-north-1"
+
+    Returns:
+        str: partition corresponding to the region name passed in. Ex: "aws-cn"
+    """
+    endpoint_data = _botocore_resolver().construct_endpoint("sts", region)
+    return endpoint_data["partition"]
+
+
 class DeferredError(object):
     """Stores an exception and raises it at a later time if this object is
     accessed in any way. Useful to allow soft-dependencies on imports, so that

tests/integ/kms_utils.py

@@ -18,7 +18,7 @@
 from sagemaker import utils
 
 PRINCIPAL_TEMPLATE = (
-    '["{account_id}", "{role_arn}", ' '"arn:aws:iam::{account_id}:role/{sagemaker_role}"] '
+    '["{account_id}", "{role_arn}", ' '"arn:{partition}:iam::{account_id}:role/{sagemaker_role}"] '
 )
 
 KEY_ALIAS = "SageMakerTestKMSKey"
@@ -60,11 +60,14 @@ def _get_kms_key_id(kms_client, alias):
 
 
 def _create_kms_key(
-    kms_client, account_id, role_arn=None, sagemaker_role="SageMakerRole", alias=KEY_ALIAS
+    kms_client, account_id, region, role_arn=None, sagemaker_role="SageMakerRole", alias=KEY_ALIAS
 ):
     if role_arn:
         principal = PRINCIPAL_TEMPLATE.format(
-            account_id=account_id, role_arn=role_arn, sagemaker_role=sagemaker_role
+            partition=utils._aws_partition(region),
+            account_id=account_id,
+            role_arn=role_arn,
+            sagemaker_role=sagemaker_role,
         )
     else:
         principal = '"{account_id}"'.format(account_id=account_id)
@@ -83,7 +86,7 @@ def _create_kms_key(
 
 
 def _add_role_to_policy(
-    kms_client, account_id, role_arn, alias=KEY_ALIAS, sagemaker_role="SageMakerRole"
+    kms_client, account_id, role_arn, region, alias=KEY_ALIAS, sagemaker_role="SageMakerRole"
 ):
     key_id = _get_kms_key_id(kms_client, alias)
     policy = kms_client.get_key_policy(KeyId=key_id, PolicyName=POLICY_NAME)
@@ -92,7 +95,10 @@ def _add_role_to_policy(
 
     if role_arn not in principal or sagemaker_role not in principal:
         principal = PRINCIPAL_TEMPLATE.format(
-            account_id=account_id, role_arn=role_arn, sagemaker_role=sagemaker_role
+            partition=utils._aws_partition(region),
+            account_id=account_id,
+            role_arn=role_arn,
+            sagemaker_role=sagemaker_role,
         )
 
         kms_client.put_key_policy(
@@ -115,44 +121,44 @@ def get_or_create_kms_key(
     account_id = sts_client.get_caller_identity()["Account"]
 
     if kms_key_arn is None:
-        return _create_kms_key(kms_client, account_id, role_arn, sagemaker_role, alias)
+        return _create_kms_key(kms_client, account_id, region, role_arn, sagemaker_role, alias)
 
     if role_arn:
-        _add_role_to_policy(kms_client, account_id, role_arn, alias, sagemaker_role)
+        _add_role_to_policy(kms_client, account_id, role_arn, region, alias, sagemaker_role)
 
     return kms_key_arn
 
 
-KMS_BUCKET_POLICY = """{
+KMS_BUCKET_POLICY = """{{
   "Version": "2012-10-17",
   "Id": "PutObjPolicy",
   "Statement": [
-    {
+    {{
       "Sid": "DenyIncorrectEncryptionHeader",
       "Effect": "Deny",
       "Principal": "*",
       "Action": "s3:PutObject",
-      "Resource": "arn:aws:s3:::%s/*",
-      "Condition": {
-        "StringNotEquals": {
-          "s3:x-amz-server-side-encryption": "aws:kms"
-        }
-      }
-    },
-    {
+      "Resource": "arn:{partition}:s3:::{bucket_name}/*",
+      "Condition": {{
+        "StringNotEquals": {{
+          "s3:x-amz-server-side-encryption": "{partition}:kms"
+        }}
+      }}
+    }},
+    {{
       "Sid": "DenyUnEncryptedObjectUploads",
       "Effect": "Deny",
       "Principal": "*",
       "Action": "s3:PutObject",
-      "Resource": "arn:aws:s3:::%s/*",
-      "Condition": {
-        "Null": {
+      "Resource": "arn:{partition}:s3:::{bucket_name}/*",
+      "Condition": {{
+        "Null": {{
           "s3:x-amz-server-side-encryption": "true"
-        }
-      }
-    }
+        }}
+      }}
+    }}
   ]
-}"""
+}}"""
 
 
 @contextlib.contextmanager
@@ -167,7 +173,7 @@ def bucket_with_encryption(sagemaker_session, sagemaker_role):
     role_arn = sts_client.get_caller_identity()["Arn"]
 
     kms_client = boto_session.client("kms")
-    kms_key_arn = _create_kms_key(kms_client, account, role_arn, sagemaker_role, None)
+    kms_key_arn = _create_kms_key(kms_client, account, region, role_arn, sagemaker_role, None)
 
     region = boto_session.region_name
     bucket_name = "sagemaker-{}-{}-with-kms".format(region, account)
@@ -181,7 +187,9 @@ def bucket_with_encryption(sagemaker_session, sagemaker_role):
             "Rules": [
                 {
                     "ApplyServerSideEncryptionByDefault": {
-                        "SSEAlgorithm": "aws:kms",
+                        "SSEAlgorithm": "{partition}:kms".format(
+                            partition=utils._aws_partition(region)
+                        ),
                         "KMSMasterKeyID": kms_key_arn,
                     }
                 }
@@ -190,7 +198,10 @@ def bucket_with_encryption(sagemaker_session, sagemaker_role):
     )
 
     s3_client.put_bucket_policy(
-        Bucket=bucket_name, Policy=KMS_BUCKET_POLICY % (bucket_name, bucket_name)
+        Bucket=bucket_name,
+        Policy=KMS_BUCKET_POLICY.format(
+            partition=utils._aws_partition(region), bucket_name=bucket_name
+        ),
     )
 
     yield "s3://" + bucket_name, kms_key_arn

tests/integ/test_marketplace.py

@@ -24,6 +24,7 @@
 from sagemaker import AlgorithmEstimator, ModelPackage
 from sagemaker.tuner import IntegerParameter, HyperparameterTuner
 from sagemaker.utils import sagemaker_timestamp
+from sagemaker.utils import _aws_partition
 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
@@ -39,12 +40,12 @@
 # Both are written by Amazon and are free to subscribe.
 
 ALGORITHM_ARN = (
-    "arn:aws:sagemaker:%s:%s:algorithm/scikit-decision-trees-"
+    "arn:{partition}:sagemaker:{region}:{account}:algorithm/scikit-decision-trees-"
     "15423055-57b73412d2e93e9239e4e16f83298b8f"
 )
 
 MODEL_PACKAGE_ARN = (
-    "arn:aws:sagemaker:%s:%s:model-package/scikit-iris-detector-"
+    "arn:{partition}:sagemaker:{region}:{account}:model-package/scikit-iris-detector-"
     "154230595-8f00905c1f927a512b73ea29dd09ae30"
 )
 
@@ -63,7 +64,9 @@ def test_marketplace_estimator(sagemaker_session, cpu_instance_type):
         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)
+        algorithm_arn = ALGORITHM_ARN.format(
+            partition=_aws_partition(region), region=region, account=account
+        )
 
         algo = AlgorithmEstimator(
             algorithm_arn=algorithm_arn,
@@ -103,7 +106,9 @@ def test_marketplace_attach(sagemaker_session, cpu_instance_type):
         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)
+        algorithm_arn = ALGORITHM_ARN.format(
+            partition=_aws_partition(region), region=region, account=account
+        )
 
         mktplace = AlgorithmEstimator(
             algorithm_arn=algorithm_arn,
@@ -155,7 +160,9 @@ def test_marketplace_attach(sagemaker_session, cpu_instance_type):
 def test_marketplace_model(sagemaker_session, cpu_instance_type):
     region = sagemaker_session.boto_region_name
     account = REGION_ACCOUNT_MAP[region]
-    model_package_arn = MODEL_PACKAGE_ARN % (region, account)
+    model_package_arn = MODEL_PACKAGE_ARN.format(
+        partition=_aws_partition(region), region=region, account=account
+    )
 
     def predict_wrapper(endpoint, session):
         return sagemaker.RealTimePredictor(
@@ -192,7 +199,9 @@ def test_marketplace_tuning_job(sagemaker_session, cpu_instance_type):
     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)
+    algorithm_arn = ALGORITHM_ARN.format(
+        partition=_aws_partition(region), region=region, account=account
+    )
 
     mktplace = AlgorithmEstimator(
         algorithm_arn=algorithm_arn,
@@ -233,7 +242,9 @@ def test_marketplace_transform_job(sagemaker_session, cpu_instance_type):
     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)
+    algorithm_arn = ALGORITHM_ARN.format(
+        partition=_aws_partition(region), region=region, account=account
+    )
 
     algo = AlgorithmEstimator(
         algorithm_arn=algorithm_arn,
@@ -279,7 +290,9 @@ def test_marketplace_transform_job_from_model_package(sagemaker_session, cpu_ins
 
     region = sagemaker_session.boto_region_name
     account = REGION_ACCOUNT_MAP[region]
-    model_package_arn = MODEL_PACKAGE_ARN % (region, account)
+    model_package_arn = MODEL_PACKAGE_ARN.format(
+        partition=_aws_partition(region), region=region, account=account
+    )
 
     model = ModelPackage(
         role="SageMakerRole",

tests/unit/test_utils.py

@@ -699,3 +699,11 @@ def test_sts_regional_endpoint():
     endpoint = sagemaker.utils.sts_regional_endpoint("us-iso-east-1")
     assert endpoint == "https://sts.us-iso-east-1.c2s.ic.gov"
     assert botocore.utils.is_valid_endpoint_url(endpoint)
+
+
+def test_partition_by_region():
+    assert sagemaker.utils._aws_partition("us-west-2") == "aws"
+    assert sagemaker.utils._aws_partition("cn-north-1") == "aws-cn"
+    assert sagemaker.utils._aws_partition("us-gov-east-1") == "aws-us-gov"
+    assert sagemaker.utils._aws_partition("us-iso-east-1") == "aws-iso"
+    assert sagemaker.utils._aws_partition("us-isob-east-1") == "aws-iso-b"