fix: use sts regional endpoint when creating default bucket (#1086)

aws · Oct 12, 2019 · 8a6b784 · 8a6b784
1 parent a98f488
commit 8a6b784
Showing 1 changed file with 6 additions and 2 deletions.
diff --git a/src/sagemaker/session.py b/src/sagemaker/session.py
@@ -208,8 +208,10 @@ def default_bucket(self):
         if self._default_bucket:
             return self._default_bucket
 
-        account = self.boto_session.client("sts").get_caller_identity()["Account"]
         region = self.boto_session.region_name
+        account = self.boto_session.client(
+            "sts", region_name=region, endpoint_url=sts_regional_endpoint(region)
+        ).get_caller_identity()["Account"]
         default_bucket = "sagemaker-{}-{}".format(region, account)
 
         s3 = self.boto_session.resource("s3")
@@ -1400,7 +1402,9 @@ def get_caller_identity_arn(self):
                 )
 
         assumed_role = self.boto_session.client(
-            "sts", endpoint_url=sts_regional_endpoint(self.boto_region_name)
+            "sts",
+            region_name=self.boto_region_name,
+            endpoint_url=sts_regional_endpoint(self.boto_region_name),
         ).get_caller_identity()["Arn"]
 
         if "AmazonSageMaker-ExecutionRole" in assumed_role: