refactor s3 bucket creation

Author: knakad

src/sagemaker/session.py

@@ -342,21 +342,36 @@ def default_bucket(self):
             ).get_caller_identity()["Account"]
             default_bucket = "sagemaker-{}-{}".format(region, account)
 
+        self.create_s3_bucket_if_it_does_not_exist(bucket_name=default_bucket, region=region)
+
+        self._default_bucket = default_bucket
+
+        return self._default_bucket
+
+    def create_s3_bucket_if_it_does_not_exist(self, bucket_name, region):
+        """Creates an S3 Bucket if it does not exist.
+        Also swallows a few common exceptions that indicate that the bucket already exists, or
+        that it is being created.
+
+        Args:
+            bucket_name (str): Name of the S3 bucket to be created.
+            region (str): The region in which to create the bucket.
+
+        """
         s3 = self.boto_session.resource("s3", region_name=region)
-        bucket = self.boto_session.resource("s3", region_name=region).Bucket(name=default_bucket)
+        bucket = self.boto_session.resource("s3", region_name=region).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=default_bucket)
+                    s3.create_bucket(Bucket=bucket_name)
                 else:
                     s3.create_bucket(
-                        Bucket=default_bucket,
-                        CreateBucketConfiguration={"LocationConstraint": region},
+                        Bucket=bucket_name, CreateBucketConfiguration={"LocationConstraint": region}
                     )
 
-                LOGGER.info("Created S3 bucket: %s", default_bucket)
+                LOGGER.info("Created S3 bucket: %s", bucket_name)
             except ClientError as e:
                 error_code = e.response["Error"]["Code"]
                 message = e.response["Error"]["Message"]
@@ -373,10 +388,6 @@ def default_bucket(self):
                 else:
                     raise
 
-        self._default_bucket = default_bucket
-
-        return self._default_bucket
-
     def train(  # noqa: C901
         self,
         input_mode,

tests/integ/kms_utils.py

@@ -15,8 +15,6 @@
 import contextlib
 import json
 
-from botocore import exceptions
-
 from sagemaker import utils
 
 PRINCIPAL_TEMPLATE = (
@@ -158,7 +156,8 @@ def get_or_create_kms_key(
 
 
 @contextlib.contextmanager
-def bucket_with_encryption(boto_session, sagemaker_role):
+def bucket_with_encryption(sagemaker_session, sagemaker_role):
+    boto_session = sagemaker_session.boto_session
     region = boto_session.region_name
     sts_client = boto_session.client(
         "sts", region_name=region, endpoint_url=utils.sts_regional_endpoint(region)
@@ -173,34 +172,10 @@ def bucket_with_encryption(boto_session, sagemaker_role):
     region = boto_session.region_name
     bucket_name = "sagemaker-{}-{}-with-kms".format(region, account)
 
-    s3 = boto_session.client("s3", region_name=region)
-    bucket = boto_session.resource("s3", region_name=region).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(
+    sagemaker_session.create_bucket_if_it_does_not_exist(bucket_name=bucket_name, region=region)
+
+    s3_client = boto_session.client("s3", region_name=region)
+    s3_client.put_bucket_encryption(
         Bucket=bucket_name,
         ServerSideEncryptionConfiguration={
             "Rules": [
@@ -214,7 +189,9 @@ def bucket_with_encryption(boto_session, sagemaker_role):
         },
     )
 
-    s3.put_bucket_policy(Bucket=bucket_name, Policy=KMS_BUCKET_POLICY % (bucket_name, bucket_name))
+    s3_client.put_bucket_policy(
+        Bucket=bucket_name, Policy=KMS_BUCKET_POLICY % (bucket_name, bucket_name)
+    )
 
     yield "s3://" + bucket_name, kms_key_arn
 

tests/integ/test_tf_script_mode.py

@@ -83,8 +83,7 @@ def test_mnist_with_checkpoint_config(sagemaker_session, instance_type, tf_full_
 
 
 def test_server_side_encryption(sagemaker_session, tf_full_version):
-    boto_session = sagemaker_session.boto_session
-    with kms_utils.bucket_with_encryption(boto_session, ROLE) as (bucket_with_kms, kms_key):
+    with kms_utils.bucket_with_encryption(sagemaker_session, ROLE) as (bucket_with_kms, kms_key):
         output_path = os.path.join(
             bucket_with_kms, "test-server-side-encryption", time.strftime("%y%m%d-%H%M")
         )