fix: modify session and kms_utils to check for S3 bucket before creation

src/sagemaker/session.py

@@ -343,34 +343,35 @@ def default_bucket(self):
             default_bucket = "sagemaker-{}-{}".format(region, account)
 
         s3 = self.boto_session.resource("s3")
-        try:
-            # 'us-east-1' cannot be specified because it is the default region:
-            # https://github.com/boto/boto3/issues/125
-            if region == "us-east-1":
-                s3.create_bucket(Bucket=default_bucket)
-            else:
-                s3.create_bucket(
-                    Bucket=default_bucket, CreateBucketConfiguration={"LocationConstraint": region}
-                )
+        bucket = self.boto_session.resource("s3").Bucket(name=default_bucket)
+        if bucket.creation_date is None:
+            try:
+                if region == "us-east-1":
+                    # 'us-east-1' cannot be specified because it is the default region:
+                    # https://github.com/boto/boto3/issues/125
+                    s3.create_bucket(Bucket=default_bucket)
+                else:
+                    s3.create_bucket(
+                        Bucket=default_bucket,
+                        CreateBucketConfiguration={"LocationConstraint": region},
+                    )
 
-            LOGGER.info("Created S3 bucket: %s", default_bucket)
-        except ClientError as e:
-            error_code = e.response["Error"]["Code"]
-            message = e.response["Error"]["Message"]
+                LOGGER.info("Created S3 bucket: %s", default_bucket)
+            except ClientError as e:
+                error_code = e.response["Error"]["Code"]
+                message = e.response["Error"]["Message"]
 
-            if error_code == "BucketAlreadyOwnedByYou":
-                pass
-            elif (
-                error_code == "OperationAborted" and "conflicting conditional operation" in message
-            ):
-                # If this bucket is already being concurrently created, we don't need to create it
-                # again.
-                pass
-            elif error_code == "TooManyBuckets":
-                # Succeed if the default bucket exists
-                s3.meta.client.head_bucket(Bucket=default_bucket)
-            else:
-                raise
+                if error_code == "BucketAlreadyOwnedByYou":
+                    pass
+                elif (
+                    error_code == "OperationAborted"
+                    and "conflicting conditional operation" in message
+                ):
+                    # If this bucket is already being concurrently created, we don't need to create
+                    # it again.
+                    pass
+                else:
+                    raise
 
         self._default_bucket = default_bucket
 

tests/integ/kms_utils.py

@@ -174,19 +174,31 @@ def bucket_with_encryption(boto_session, sagemaker_role):
     bucket_name = "sagemaker-{}-{}-with-kms".format(region, account)
 
     s3 = boto_session.client("s3")
-    try:
-        # 'us-east-1' cannot be specified because it is the default region:
-        # https://github.com/boto/boto3/issues/125
-        if region == "us-east-1":
-            s3.create_bucket(Bucket=bucket_name)
-        else:
-            s3.create_bucket(
-                Bucket=bucket_name, CreateBucketConfiguration={"LocationConstraint": region}
-            )
-
-    except exceptions.ClientError as e:
-        if e.response["Error"]["Code"] != "BucketAlreadyOwnedByYou":
-            raise
+    bucket = boto_session.resource("s3").Bucket(name=bucket_name)
+    if bucket.creation_date is None:
+        try:
+            if region == "us-east-1":
+                # 'us-east-1' cannot be specified because it is the default region:
+                # https://github.com/boto/boto3/issues/125
+                s3.create_bucket(Bucket=bucket_name)
+            else:
+                s3.create_bucket(
+                    Bucket=bucket_name, CreateBucketConfiguration={"LocationConstraint": region}
+                )
+        except exceptions.ClientError as e:
+            error_code = e.response["Error"]["Code"]
+            message = e.response["Error"]["Message"]
+
+            if error_code == "BucketAlreadyOwnedByYou":
+                pass
+            elif (
+                error_code == "OperationAborted" and "conflicting conditional operation" in message
+            ):
+                # If this bucket is already being concurrently created, we don't need to create it
+                # again.
+                pass
+            else:
+                raise
 
     s3.put_bucket_encryption(
         Bucket=bucket_name,

tests/unit/test_default_bucket.py

@@ -26,8 +26,9 @@
 def sagemaker_session():
     boto_mock = Mock(name="boto_session", region_name=REGION)
     boto_mock.client("sts").get_caller_identity.return_value = {"Account": ACCOUNT_ID}
-    ims = sagemaker.Session(boto_session=boto_mock)
-    return ims
+    sagemaker_session = sagemaker.Session(boto_session=boto_mock)
+    sagemaker_session.boto_session.resource("s3").Bucket().creation_date = None
+    return sagemaker_session
 
 
 def test_default_bucket_s3_create_call(sagemaker_session):