change: use regional endpoint for STS in builds and tests (#1102)

Author: laurenyu

buildspec.yml

@@ -25,7 +25,7 @@ phases:
       - |
         if [ -d "ci-lock" ]; then
           FILENAME=$(ls ci-lock/ || true)
-          ACCOUNT=$(aws sts get-caller-identity --region us-west-2 --output text | awk '{print $1}')
+          ACCOUNT=$(aws --endpoint-url https://sts.us-west-2.amazonaws.com sts get-caller-identity --region us-west-2 --output text | awk '{print $1}')
           S3_BUCKET_DIR=s3://sagemaker-us-west-2-${ACCOUNT}/ci-lock/
           aws s3 rm ${S3_BUCKET_DIR}${FILENAME}
         fi

ci-scripts/queue_build.py

@@ -16,7 +16,9 @@
 import time
 import boto3
 
-account = boto3.client("sts").get_caller_identity()["Account"]
+account = boto3.client(
+    "sts", region_name="us-west-2", endpoint_url="https://sts.us-west-2.amazonaws.com"
+).get_caller_identity()["Account"]
 bucket_name = "sagemaker-us-west-2-%s" % account
 
 

tests/integ/kms_utils.py

@@ -17,6 +17,8 @@
 
 from botocore import exceptions
 
+from sagemaker import utils
+
 PRINCIPAL_TEMPLATE = (
     '["{account_id}", "{role_arn}", ' '"arn:aws:iam::{account_id}:role/{sagemaker_role}"] '
 )
@@ -108,7 +110,10 @@ def get_or_create_kms_key(
     kms_client = sagemaker_session.boto_session.client("kms")
     kms_key_arn = _get_kms_key_arn(kms_client, alias)
 
-    sts_client = sagemaker_session.boto_session.client("sts")
+    region = sagemaker_session.boto_region_name
+    sts_client = sagemaker_session.boto_session.client(
+        "sts", region_name=region, endpoint_url=utils.sts_regional_endpoint(region)
+    )
     account_id = sts_client.get_caller_identity()["Account"]
 
     if kms_key_arn is None:
@@ -154,8 +159,13 @@ def get_or_create_kms_key(
 
 @contextlib.contextmanager
 def bucket_with_encryption(boto_session, sagemaker_role):
-    account = boto_session.client("sts").get_caller_identity()["Account"]
-    role_arn = boto_session.client("sts").get_caller_identity()["Arn"]
+    region = boto_session.region_name
+    sts_client = boto_session.client(
+        "sts", region_name=region, endpoint_url=utils.sts_regional_endpoint(region)
+    )
+
+    account = sts_client.get_caller_identity()["Account"]
+    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)