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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
knakad merged 9 commits into from
Jan 29, 2020
Merged

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

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
9 commits
Select commit Hold shift + click to select a range
8dd4153
modify session and kms_utils to check for S3 bucket before creation
knakad Jan 29, 2020
c6f0aa7
add region name to s3 clients and resources
knakad Jan 29, 2020
7454f04
refactor s3 bucket creation
knakad Jan 29, 2020
93db859
fix method name typo
knakad Jan 29, 2020
7fffe78
address comments
knakad Jan 29, 2020
6d2786f
fix docstring
knakad Jan 29, 2020
4380bd6
update docstring
knakad Jan 29, 2020
200af80
Merge branch 'master' into 20200128-check-bucket-before-create
knakad Jan 29, 2020
ab9b497
fix docstring
knakad Jan 29, 2020
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
76 changes: 47 additions & 29 deletions src/sagemaker/session.py
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -342,40 +342,58 @@ def default_bucket(self):
).get_caller_identity()["Account"]
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}
)

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
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.

Comment thread
laurenyu marked this conversation as resolved.
Raises:
botocore.exceptions.ClientError: If S3 throws an unexpected exception during bucket
creation.
If the exception is due to the bucket already existing or
already being created, no exception is raised.

"""
bucket = self.boto_session.resource("s3", region_name=region).Bucket(name=bucket_name)
if bucket.creation_date is None:
try:
s3 = self.boto_session.resource("s3", region_name=region)
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}
)

LOGGER.info("Created S3 bucket: %s", bucket_name)
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
else:
raise

def train( # noqa: C901
self,
input_mode,
Expand Down
29 changes: 9 additions & 20 deletions tests/integ/kms_utils.py
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -15,8 +15,6 @@
import contextlib
import json

from botocore import exceptions

from sagemaker import utils

PRINCIPAL_TEMPLATE = (
Expand DownExpand Up@@ -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)
Expand All@@ -173,22 +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")
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

s3.put_bucket_encryption(
sagemaker_session._create_s3_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": [
Expand All@@ -202,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

Expand Down
3 changes: 1 addition & 2 deletions tests/integ/test_tf_script_mode.py
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -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")
)
Expand Down
5 changes: 3 additions & 2 deletions tests/unit/test_default_bucket.py
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -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
Comment thread
knakad marked this conversation as resolved.
return sagemaker_session


def test_default_bucket_s3_create_call(sagemaker_session):
Expand Down
, 'i'); if (__m === '*' || __re.test(location.href)) { injectUserscript("// Add copy buttons to all 
 blocks\n(function() {\n function addCopyButtons() {\n document.querySelectorAll('pre code').forEach(function(codeBlock) {\n if (codeBlock.parentElement.hasAttribute('data-copy-added')) return;\n codeBlock.parentElement.setAttribute('data-copy-added', 'true');\n \n var btn = document.createElement('button');\n btn.textContent = 'Copy';\n btn.style.cssText = 'position:absolute;top:4px;right:4px;padding:2px 8px;font-size:11px;background:#4ecdc4;border:none;border-radius:4px;color:#1a1a2e;cursor:pointer;opacity:0.7;transition:opacity 0.2s;';\n btn.onmouseover = function() { this.style.opacity = '1'; };\n btn.onmouseout = function() { this.style.opacity = '0.7'; };\n btn.onclick = function() {\n navigator.clipboard.writeText(codeBlock.textContent).then(function() {\n btn.textContent = 'Copied!';\n setTimeout(function() { btn.textContent = 'Copy'; }, 1500);\n });\n };\n codeBlock.parentElement.style.position = 'relative';\n codeBlock.parentElement.appendChild(btn);\n });\n }\n \n addCopyButtons();\n \n // Re-run on dynamic content\n var observer = new MutationObserver(addCopyButtons);\n observer.observe(document.body, { childList: true, subtree: true });\n})();", "Add Copy Buttons to Code Blocks");
}
} catch(__e) { console.warn('[Userscript:Add Copy Buttons to Code Blocks]', __e); }
})();
(function(){
try {
var __m = "github.com";
var __re = new RegExp('^' + "github\\.com" + '

Skip to content
fix: modify session and kms_utils to check for S3 bucket before creation#1271

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

 Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and
privacy statement. We’ll occasionally send you account related emails.

Already on GitHub?
Sign in
to your account

 Merged

knakad
merged 9 commits into

from

Jan 29, 2020
 Merged


fix: modify session and kms_utils to check for S3 bucket before creation
#1271
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


9 commits


Select commit
Hold shift + click to select a range
8dd4153

modify session and kms_utils to check for S3 bucket before creation


knakad Jan 29, 2020
c6f0aa7

add region name to s3 clients and resources


knakad Jan 29, 2020
7454f04

refactor s3 bucket creation


knakad Jan 29, 2020
93db859

fix method name typo


knakad Jan 29, 2020
7fffe78

address comments


knakad Jan 29, 2020
6d2786f

fix docstring


knakad Jan 29, 2020
4380bd6

update docstring


knakad Jan 29, 2020
200af80

Merge branch 'master' into 20200128-check-bucket-before-create


knakad Jan 29, 2020
ab9b497

fix docstring


knakad Jan 29, 2020

File filter

Filter by extension
Filter by extension
Conversations

Failed to load comments. 
Loading
Jump to

Jump to file


Failed to load files. 
Loading
 Diff view

 Diff view


76 changes: 47 additions & 29 deletions
src/sagemaker/session.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -342,40 +342,58 @@ def default_bucket(self):
 ).get_caller_identity()["Account"]
 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}
 )

 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
 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.

Comment thread
laurenyu marked this conversation as resolved.

 Raises:
 botocore.exceptions.ClientError: If S3 throws an unexpected exception during bucket
 creation.
 If the exception is due to the bucket already existing or
 already being created, no exception is raised.

 """
 bucket = self.boto_session.resource("s3", region_name=region).Bucket(name=bucket_name)
 if bucket.creation_date is None:
 try:
 s3 = self.boto_session.resource("s3", region_name=region)
 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}
 )

 LOGGER.info("Created S3 bucket: %s", bucket_name)
 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
 else:
 raise

 def train( # noqa: C901
 self,
 input_mode,
Expand Down

29 changes: 9 additions & 20 deletions
tests/integ/kms_utils.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -15,8 +15,6 @@
import contextlib
import json

from botocore import exceptions

from sagemaker import utils

PRINCIPAL_TEMPLATE = (
Expand DownExpand Up@@ -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)
Expand All@@ -173,22 +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")
 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

 s3.put_bucket_encryption(
 sagemaker_session._create_s3_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": [
Expand All@@ -202,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

Expand Down

3 changes: 1 addition & 2 deletions
tests/integ/test_tf_script_mode.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -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")
 )
Expand Down

5 changes: 3 additions & 2 deletions
tests/unit/test_default_bucket.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -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
Comment thread
knakad marked this conversation as resolved.

 return sagemaker_session


def test_default_bucket_s3_create_call(sagemaker_session):
Expand Down

, 'i');
if (__m === '*' || __re.test(location.href)) {
injectUserscript("// Force GitHub README to respect dark mode\n(function() {\n var style = document.createElement('style');\n style.textContent = '\n .markdown-body {\n color-scheme: dark light;\n }\n .markdown-body pre { background: #161b22 !important; }\n .markdown-body code { background: rgba(110, 118, 129, 0.4) !important; }\n .markdown-body table th, .markdown-body table td { border-color: #30363d !important; }\n .markdown-body img { background: #0d1117; }\n .markdown-body blockquote { border-left-color: #8b949e; }\n .markdown-body hr { border-color: #30363d; }\n ';\n document.head.appendChild(style);\n})();", "GitHub Dark Mode README Fix");
}
} catch(__e) { console.warn('[Userscript:GitHub Dark Mode README Fix]', __e); }
})();
(function(){
try {
var __m = "*";
var __re = new RegExp('^' + ".*" + '

Skip to content
fix: modify session and kms_utils to check for S3 bucket before creation#1271

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

 Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and
privacy statement. We’ll occasionally send you account related emails.

Already on GitHub?
Sign in
to your account

 Merged

knakad
merged 9 commits into

from

Jan 29, 2020
 Merged


fix: modify session and kms_utils to check for S3 bucket before creation
#1271
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


9 commits


Select commit
Hold shift + click to select a range
8dd4153

modify session and kms_utils to check for S3 bucket before creation


knakad Jan 29, 2020
c6f0aa7

add region name to s3 clients and resources


knakad Jan 29, 2020
7454f04

refactor s3 bucket creation


knakad Jan 29, 2020
93db859

fix method name typo


knakad Jan 29, 2020
7fffe78

address comments


knakad Jan 29, 2020
6d2786f

fix docstring


knakad Jan 29, 2020
4380bd6

update docstring


knakad Jan 29, 2020
200af80

Merge branch 'master' into 20200128-check-bucket-before-create


knakad Jan 29, 2020
ab9b497

fix docstring


knakad Jan 29, 2020

File filter

Filter by extension
Filter by extension
Conversations

Failed to load comments. 
Loading
Jump to

Jump to file


Failed to load files. 
Loading
 Diff view

 Diff view


76 changes: 47 additions & 29 deletions
src/sagemaker/session.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -342,40 +342,58 @@ def default_bucket(self):
 ).get_caller_identity()["Account"]
 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}
 )

 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
 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.

Comment thread
laurenyu marked this conversation as resolved.

 Raises:
 botocore.exceptions.ClientError: If S3 throws an unexpected exception during bucket
 creation.
 If the exception is due to the bucket already existing or
 already being created, no exception is raised.

 """
 bucket = self.boto_session.resource("s3", region_name=region).Bucket(name=bucket_name)
 if bucket.creation_date is None:
 try:
 s3 = self.boto_session.resource("s3", region_name=region)
 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}
 )

 LOGGER.info("Created S3 bucket: %s", bucket_name)
 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
 else:
 raise

 def train( # noqa: C901
 self,
 input_mode,
Expand Down

29 changes: 9 additions & 20 deletions
tests/integ/kms_utils.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -15,8 +15,6 @@
import contextlib
import json

from botocore import exceptions

from sagemaker import utils

PRINCIPAL_TEMPLATE = (
Expand DownExpand Up@@ -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)
Expand All@@ -173,22 +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")
 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

 s3.put_bucket_encryption(
 sagemaker_session._create_s3_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": [
Expand All@@ -202,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

Expand Down

3 changes: 1 addition & 2 deletions
tests/integ/test_tf_script_mode.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -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")
 )
Expand Down

5 changes: 3 additions & 2 deletions
tests/unit/test_default_bucket.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -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
Comment thread
knakad marked this conversation as resolved.

 return sagemaker_session


def test_default_bucket_s3_create_call(sagemaker_session):
Expand Down

, 'i');
if (__m === '*' || __re.test(location.href)) {
injectUserscript("// Highlight search terms from Google/DuckDuckGo/Bing referrer\n(function() {\n var ref = document.referrer;\n var terms = [];\n \n if (ref.includes('google.com') || ref.includes('duckduckgo.com') || ref.includes('bing.com')) {\n var url = new URL(ref);\n var q = url.searchParams.get('q') || url.searchParams.get('p');\n if (q) {\n terms = q.split(/\\s+/).filter(function(t) { return t.length > 2; });\n }\n }\n \n if (terms.length === 0) return;\n \n var style = document.createElement('style');\n style.textContent = '.userscript-highlight { background: #fbbf24; color: #1a1a2e; padding: 1px 3px; border-radius: 2px; }';\n document.head.appendChild(style);\n \n function highlight(node) {\n if (node.nodeType === 3) { // text node\n var text = node.textContent;\n var found = false;\n terms.forEach(function(term) {\n var regex = new RegExp('(' + term.replace(/[.*+?^${}()|[\\]\\\\]/g, '\\\\') + ')', 'gi');\n if (regex.test(text)) {\n found = true;\n var frag = document.createDocumentFragment();\n var parts = text.split(regex);\n parts.forEach(function(part, i) {\n if (i % 2 === 0) {\n frag.appendChild(document.createTextNode(part));\n } else {\n var span = document.createElement('span');\n span.className = 'userscript-highlight';\n span.textContent = part;\n frag.appendChild(span);\n }\n });\n node.parentNode.replaceChild(frag, node);\n }\n });\n } else if (node.nodeType === 1 && node.childNodes) { // element\n var skipTags = ['SCRIPT', 'STYLE', 'NOSCRIPT', 'TEXTAREA', 'INPUT', 'SELECT'];\n if (!skipTags.includes(node.tagName)) {\n Array.from(node.childNodes).forEach(highlight);\n }\n }\n }\n \n highlight(document.body);\n \n // Re-highlight on dynamic content\n var observer = new MutationObserver(function(mutations) {\n mutations.forEach(function(m) {\n m.addedNodes.forEach(function(node) {\n if (node.nodeType === 1 || node.nodeType === 3) highlight(node);\n });\n });\n });\n observer.observe(document.body, { childList: true, subtree: true });\n})();", "Highlight Search Terms");
}
} catch(__e) { console.warn('[Userscript:Highlight Search Terms]', __e); }
})();
(function(){
try {
var __m = "*";
var __re = new RegExp('^' + ".*" + '

Skip to content
fix: modify session and kms_utils to check for S3 bucket before creation#1271

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

 Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and
privacy statement. We’ll occasionally send you account related emails.

Already on GitHub?
Sign in
to your account

 Merged

knakad
merged 9 commits into

from

Jan 29, 2020
 Merged


fix: modify session and kms_utils to check for S3 bucket before creation
#1271
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


9 commits


Select commit
Hold shift + click to select a range
8dd4153

modify session and kms_utils to check for S3 bucket before creation


knakad Jan 29, 2020
c6f0aa7

add region name to s3 clients and resources


knakad Jan 29, 2020
7454f04

refactor s3 bucket creation


knakad Jan 29, 2020
93db859

fix method name typo


knakad Jan 29, 2020
7fffe78

address comments


knakad Jan 29, 2020
6d2786f

fix docstring


knakad Jan 29, 2020
4380bd6

update docstring


knakad Jan 29, 2020
200af80

Merge branch 'master' into 20200128-check-bucket-before-create


knakad Jan 29, 2020
ab9b497

fix docstring


knakad Jan 29, 2020

File filter

Filter by extension
Filter by extension
Conversations

Failed to load comments. 
Loading
Jump to

Jump to file


Failed to load files. 
Loading
 Diff view

 Diff view


76 changes: 47 additions & 29 deletions
src/sagemaker/session.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -342,40 +342,58 @@ def default_bucket(self):
 ).get_caller_identity()["Account"]
 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}
 )

 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
 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.

Comment thread
laurenyu marked this conversation as resolved.

 Raises:
 botocore.exceptions.ClientError: If S3 throws an unexpected exception during bucket
 creation.
 If the exception is due to the bucket already existing or
 already being created, no exception is raised.

 """
 bucket = self.boto_session.resource("s3", region_name=region).Bucket(name=bucket_name)
 if bucket.creation_date is None:
 try:
 s3 = self.boto_session.resource("s3", region_name=region)
 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}
 )

 LOGGER.info("Created S3 bucket: %s", bucket_name)
 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
 else:
 raise

 def train( # noqa: C901
 self,
 input_mode,
Expand Down

29 changes: 9 additions & 20 deletions
tests/integ/kms_utils.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -15,8 +15,6 @@
import contextlib
import json

from botocore import exceptions

from sagemaker import utils

PRINCIPAL_TEMPLATE = (
Expand DownExpand Up@@ -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)
Expand All@@ -173,22 +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")
 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

 s3.put_bucket_encryption(
 sagemaker_session._create_s3_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": [
Expand All@@ -202,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

Expand Down

3 changes: 1 addition & 2 deletions
tests/integ/test_tf_script_mode.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -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")
 )
Expand Down

5 changes: 3 additions & 2 deletions
tests/unit/test_default_bucket.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -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
Comment thread
knakad marked this conversation as resolved.

 return sagemaker_session


def test_default_bucket_s3_create_call(sagemaker_session):
Expand Down

, 'i');
if (__m === '*' || __re.test(location.href)) {
injectUserscript("// Strip utm_, fbclid, gclid, etc. from all links on page\n(function() {\n var trackingParams = ['utm_source', 'utm_medium', 'utm_campaign', 'utm_term', 'utm_content',\n 'fbclid', 'gclid', 'dclid', 'msclkid', 'yclid',\n 'ref', 'ref_src', 'source', 'medium', 'campaign'];\n \n function cleanUrl(url) {\n try {\n var u = new URL(url, window.location.origin);\n var changed = false;\n trackingParams.forEach(function(p) {\n if (u.searchParams.has(p)) {\n u.searchParams.delete(p);\n changed = true;\n }\n });\n return changed ? u.toString() : url;\n } catch (e) {\n return url;\n }\n }\n \n function cleanLinks() {\n document.querySelectorAll('a[href]').forEach(function(a) {\n var clean = cleanUrl(a.href);\n if (clean !== a.href) a.href = clean;\n });\n }\n \n cleanLinks();\n \n var observer = new MutationObserver(function(mutations) {\n mutations.forEach(function(m) {\n m.addedNodes.forEach(function(node) {\n if (node.nodeType === 1) {\n if (node.tagName === 'A') cleanLinks();\n node.querySelectorAll('a[href]').forEach(function(a) {\n var clean = cleanUrl(a.href);\n if (clean !== a.href) a.href = clean;\n });\n }\n });\n });\n });\n observer.observe(document.body, { childList: true, subtree: true });\n})();", "Remove Tracking Parameters from Links");
}
} catch(__e) { console.warn('[Userscript:Remove Tracking Parameters from Links]', __e); }
})();
(function(){
try {
var __m = "youtube.com";
var __re = new RegExp('^' + "youtube\\.com" + '

Skip to content
fix: modify session and kms_utils to check for S3 bucket before creation#1271

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

 Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and
privacy statement. We’ll occasionally send you account related emails.

Already on GitHub?
Sign in
to your account

 Merged

knakad
merged 9 commits into

from

Jan 29, 2020
 Merged


fix: modify session and kms_utils to check for S3 bucket before creation
#1271
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


9 commits


Select commit
Hold shift + click to select a range
8dd4153

modify session and kms_utils to check for S3 bucket before creation


knakad Jan 29, 2020
c6f0aa7

add region name to s3 clients and resources


knakad Jan 29, 2020
7454f04

refactor s3 bucket creation


knakad Jan 29, 2020
93db859

fix method name typo


knakad Jan 29, 2020
7fffe78

address comments


knakad Jan 29, 2020
6d2786f

fix docstring


knakad Jan 29, 2020
4380bd6

update docstring


knakad Jan 29, 2020
200af80

Merge branch 'master' into 20200128-check-bucket-before-create


knakad Jan 29, 2020
ab9b497

fix docstring


knakad Jan 29, 2020

File filter

Filter by extension
Filter by extension
Conversations

Failed to load comments. 
Loading
Jump to

Jump to file


Failed to load files. 
Loading
 Diff view

 Diff view


76 changes: 47 additions & 29 deletions
src/sagemaker/session.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -342,40 +342,58 @@ def default_bucket(self):
 ).get_caller_identity()["Account"]
 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}
 )

 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
 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.

Comment thread
laurenyu marked this conversation as resolved.

 Raises:
 botocore.exceptions.ClientError: If S3 throws an unexpected exception during bucket
 creation.
 If the exception is due to the bucket already existing or
 already being created, no exception is raised.

 """
 bucket = self.boto_session.resource("s3", region_name=region).Bucket(name=bucket_name)
 if bucket.creation_date is None:
 try:
 s3 = self.boto_session.resource("s3", region_name=region)
 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}
 )

 LOGGER.info("Created S3 bucket: %s", bucket_name)
 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
 else:
 raise

 def train( # noqa: C901
 self,
 input_mode,
Expand Down

29 changes: 9 additions & 20 deletions
tests/integ/kms_utils.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -15,8 +15,6 @@
import contextlib
import json

from botocore import exceptions

from sagemaker import utils

PRINCIPAL_TEMPLATE = (
Expand DownExpand Up@@ -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)
Expand All@@ -173,22 +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")
 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

 s3.put_bucket_encryption(
 sagemaker_session._create_s3_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": [
Expand All@@ -202,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

Expand Down

3 changes: 1 addition & 2 deletions
tests/integ/test_tf_script_mode.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -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")
 )
Expand Down

5 changes: 3 additions & 2 deletions
tests/unit/test_default_bucket.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -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
Comment thread
knakad marked this conversation as resolved.

 return sagemaker_session


def test_default_bucket_s3_create_call(sagemaker_session):
Expand Down

, 'i');
if (__m === '*' || __re.test(location.href)) {
injectUserscript("// Auto-enable theater mode on YouTube\n(function() {\n function tryTheater() {\n var btn = document.querySelector('button[aria-label=\"Theater mode\"], ytd-player #player button[title=\"Theater mode\"]');\n if (btn && !btn.classList.contains('activated')) {\n btn.click();\n }\n }\n \n // Try immediately\n tryTheater();\n \n // Try after navigation (SPA)\n var lastUrl = location.href;\n setInterval(function() {\n if (location.href !== lastUrl) {\n lastUrl = location.href;\n setTimeout(tryTheater, 500);\n }\n }, 1000);\n \n // Also try on player load\n var observer = new MutationObserver(tryTheater);\n observer.observe(document.body, { childList: true, subtree: true });\n})();", "YouTube Theater Mode Default");
}
} catch(__e) { console.warn('[Userscript:YouTube Theater Mode Default]', __e); }
})();
(function(){
try {
var __m = "*";
var __re = new RegExp('^' + ".*" + '

Skip to content
fix: modify session and kms_utils to check for S3 bucket before creation#1271

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

 Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and
privacy statement. We’ll occasionally send you account related emails.

Already on GitHub?
Sign in
to your account

 Merged

knakad
merged 9 commits into

from

Jan 29, 2020
 Merged


fix: modify session and kms_utils to check for S3 bucket before creation
#1271
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


9 commits


Select commit
Hold shift + click to select a range
8dd4153

modify session and kms_utils to check for S3 bucket before creation


knakad Jan 29, 2020
c6f0aa7

add region name to s3 clients and resources


knakad Jan 29, 2020
7454f04

refactor s3 bucket creation


knakad Jan 29, 2020
93db859

fix method name typo


knakad Jan 29, 2020
7fffe78

address comments


knakad Jan 29, 2020
6d2786f

fix docstring


knakad Jan 29, 2020
4380bd6

update docstring


knakad Jan 29, 2020
200af80

Merge branch 'master' into 20200128-check-bucket-before-create


knakad Jan 29, 2020
ab9b497

fix docstring


knakad Jan 29, 2020

File filter

Filter by extension
Filter by extension
Conversations

Failed to load comments. 
Loading
Jump to

Jump to file


Failed to load files. 
Loading
 Diff view

 Diff view


76 changes: 47 additions & 29 deletions
src/sagemaker/session.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -342,40 +342,58 @@ def default_bucket(self):
 ).get_caller_identity()["Account"]
 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}
 )

 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
 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.

Comment thread
laurenyu marked this conversation as resolved.

 Raises:
 botocore.exceptions.ClientError: If S3 throws an unexpected exception during bucket
 creation.
 If the exception is due to the bucket already existing or
 already being created, no exception is raised.

 """
 bucket = self.boto_session.resource("s3", region_name=region).Bucket(name=bucket_name)
 if bucket.creation_date is None:
 try:
 s3 = self.boto_session.resource("s3", region_name=region)
 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}
 )

 LOGGER.info("Created S3 bucket: %s", bucket_name)
 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
 else:
 raise

 def train( # noqa: C901
 self,
 input_mode,
Expand Down

29 changes: 9 additions & 20 deletions
tests/integ/kms_utils.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -15,8 +15,6 @@
import contextlib
import json

from botocore import exceptions

from sagemaker import utils

PRINCIPAL_TEMPLATE = (
Expand DownExpand Up@@ -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)
Expand All@@ -173,22 +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")
 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

 s3.put_bucket_encryption(
 sagemaker_session._create_s3_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": [
Expand All@@ -202,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

Expand Down

3 changes: 1 addition & 2 deletions
tests/integ/test_tf_script_mode.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -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")
 )
Expand Down

5 changes: 3 additions & 2 deletions
tests/unit/test_default_bucket.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -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
Comment thread
knakad marked this conversation as resolved.

 return sagemaker_session


def test_default_bucket_s3_create_call(sagemaker_session):
Expand Down

, 'i');
if (__m === '*' || __re.test(location.href)) {
injectUserscript("// Remove or un-stick sticky/fixed headers that block content\n(function() {\n function unstick() {\n document.querySelectorAll('header, nav, [role=\"banner\"], .header, .navbar, .sticky, .fixed-top, [style*=\"position: fixed\"], [style*=\"position:sticky\"]').forEach(function(el) {\n if (el.style.position === 'fixed' || el.style.position === 'sticky' || \n getComputedStyle(el).position === 'fixed' || getComputedStyle(el).position === 'sticky') {\n el.style.position = 'static';\n el.style.top = 'auto';\n el.style.zIndex = 'auto';\n }\n });\n }\n \n unstick();\n \n var observer = new MutationObserver(unstick);\n observer.observe(document.body, { childList: true, subtree: true, attributes: true, attributeFilter: ['style', 'class'] });\n})();", "Kill Sticky Headers");
}
} catch(__e) { console.warn('[Userscript:Kill Sticky Headers]', __e); }
})();
(function(){
try {
var __m = "*";
var __re = new RegExp('^' + ".*" + '

Skip to content
fix: modify session and kms_utils to check for S3 bucket before creation#1271

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

 Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and
privacy statement. We’ll occasionally send you account related emails.

Already on GitHub?
Sign in
to your account

 Merged

knakad
merged 9 commits into

from

Jan 29, 2020
 Merged


fix: modify session and kms_utils to check for S3 bucket before creation
#1271
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


9 commits


Select commit
Hold shift + click to select a range
8dd4153

modify session and kms_utils to check for S3 bucket before creation


knakad Jan 29, 2020
c6f0aa7

add region name to s3 clients and resources


knakad Jan 29, 2020
7454f04

refactor s3 bucket creation


knakad Jan 29, 2020
93db859

fix method name typo


knakad Jan 29, 2020
7fffe78

address comments


knakad Jan 29, 2020
6d2786f

fix docstring


knakad Jan 29, 2020
4380bd6

update docstring


knakad Jan 29, 2020
200af80

Merge branch 'master' into 20200128-check-bucket-before-create


knakad Jan 29, 2020
ab9b497

fix docstring


knakad Jan 29, 2020

File filter

Filter by extension
Filter by extension
Conversations

Failed to load comments. 
Loading
Jump to

Jump to file


Failed to load files. 
Loading
 Diff view

 Diff view


76 changes: 47 additions & 29 deletions
src/sagemaker/session.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -342,40 +342,58 @@ def default_bucket(self):
 ).get_caller_identity()["Account"]
 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}
 )

 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
 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.

Comment thread
laurenyu marked this conversation as resolved.

 Raises:
 botocore.exceptions.ClientError: If S3 throws an unexpected exception during bucket
 creation.
 If the exception is due to the bucket already existing or
 already being created, no exception is raised.

 """
 bucket = self.boto_session.resource("s3", region_name=region).Bucket(name=bucket_name)
 if bucket.creation_date is None:
 try:
 s3 = self.boto_session.resource("s3", region_name=region)
 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}
 )

 LOGGER.info("Created S3 bucket: %s", bucket_name)
 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
 else:
 raise

 def train( # noqa: C901
 self,
 input_mode,
Expand Down

29 changes: 9 additions & 20 deletions
tests/integ/kms_utils.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -15,8 +15,6 @@
import contextlib
import json

from botocore import exceptions

from sagemaker import utils

PRINCIPAL_TEMPLATE = (
Expand DownExpand Up@@ -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)
Expand All@@ -173,22 +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")
 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

 s3.put_bucket_encryption(
 sagemaker_session._create_s3_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": [
Expand All@@ -202,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

Expand Down

3 changes: 1 addition & 2 deletions
tests/integ/test_tf_script_mode.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -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")
 )
Expand Down

5 changes: 3 additions & 2 deletions
tests/unit/test_default_bucket.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -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
Comment thread
knakad marked this conversation as resolved.

 return sagemaker_session


def test_default_bucket_s3_create_call(sagemaker_session):
Expand Down

, 'i');
if (__m === '*' || __re.test(location.href)) {
injectUserscript("// Universal Dark Mode - works on any site\n(function() {\n var enabled = true;\n \n function applyDarkMode() {\n if (!enabled) return;\n \n // Create style element if it doesn't exist\n var style = document.getElementById('universal-dark-mode-style');\n if (!style) {\n style = document.createElement('style');\n style.id = 'universal-dark-mode-style';\n document.head.appendChild(style);\n }\n \n // Dark mode CSS - inverts colors but preserves images/video\n style.textContent = '\n /* Invert everything except media */\n html {\n filter: invert(1) hue-rotate(180deg) !important;\n background: #1a1a2e !important;\n }\n \n /* Restore images, videos, iframes, canvas */\n img, video, iframe, canvas, svg, picture, [style*=\"background-image\"] {\n filter: invert(1) hue-rotate(180deg) !important;\n }\n \n /* Preserve specific elements that should not be inverted */\n .no-dark-mode, .no-dark-mode *,\n [data-theme=\"light\"], [data-theme=\"light\"],\n .ace_editor, .ace_editor *,\n .CodeMirror, .CodeMirror *,\n .monaco-editor, .monaco-editor *,\n .markdown-body pre, .markdown-body pre *,\n .highlight, .highlight *,\n pre code, pre code * {\n filter: none !important;\n }\n \n /* Fix common UI elements */\n .modal, .popup, .dropdown-menu, .tooltip, .popover {\n filter: invert(1) hue-rotate(180deg) !important;\n background: #2d2d44 !important;\n border-color: #444 !important;\n }\n \n /* Scrollbars */\n ::-webkit-scrollbar { background: #1a1a2e !important; }\n ::-webkit-scrollbar-thumb { background: #444 !important; }\n ::-webkit-scrollbar-thumb:hover { background: #555 !important; }\n \n /* Selection */\n ::selection { background: #4ecdc4 !important; color: #1a1a2e !important; }\n ::-moz-selection { background: #4ecdc4 !important; color: #1a1a2e !important; }\n ';\n }\n \n function removeDarkMode() {\n var style = document.getElementById('universal-dark-mode-style');\n if (style) style.remove();\n }\n \n // Toggle with Alt+Shift+D\n document.addEventListener('keydown', function(e) {\n if (e.altKey && e.shiftKey && e.key === 'D') {\n e.preventDefault();\n enabled = !enabled;\n if (enabled) {\n applyDarkMode();\n console.log('[Universal Dark Mode] Enabled');\n } else {\n removeDarkMode();\n console.log('[Universal Dark Mode] Disabled');\n }\n }\n });\n \n // Apply on load\n applyDarkMode();\n \n // Re-apply on dynamic content\n var observer = new MutationObserver(function(mutations) {\n if (enabled && !document.getElementById('universal-dark-mode-style')) {\n applyDarkMode();\n }\n });\n observer.observe(document.head, { childList: true });\n \n console.log('[Universal Dark Mode] Loaded - Press Alt+Shift+D to toggle');\n})();", "Universal Dark Mode");
}
} catch(__e) { console.warn('[Userscript:Universal Dark Mode]', __e); }
})();
})();

Skip to content
fix: modify session and kms_utils to check for S3 bucket before creation#1271

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

 Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and
privacy statement. We’ll occasionally send you account related emails.

Already on GitHub?
Sign in
to your account

 Merged

knakad
merged 9 commits into

from

Jan 29, 2020
 Merged


fix: modify session and kms_utils to check for S3 bucket before creation
#1271
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


9 commits


Select commit
Hold shift + click to select a range
8dd4153

modify session and kms_utils to check for S3 bucket before creation


knakad Jan 29, 2020
c6f0aa7

add region name to s3 clients and resources


knakad Jan 29, 2020
7454f04

refactor s3 bucket creation


knakad Jan 29, 2020
93db859

fix method name typo


knakad Jan 29, 2020
7fffe78

address comments


knakad Jan 29, 2020
6d2786f

fix docstring


knakad Jan 29, 2020
4380bd6

update docstring


knakad Jan 29, 2020
200af80

Merge branch 'master' into 20200128-check-bucket-before-create


knakad Jan 29, 2020
ab9b497

fix docstring


knakad Jan 29, 2020

File filter

Filter by extension
Filter by extension
Conversations

Failed to load comments. 
Loading
Jump to

Jump to file


Failed to load files. 
Loading
 Diff view

 Diff view


76 changes: 47 additions & 29 deletions
src/sagemaker/session.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -342,40 +342,58 @@ def default_bucket(self):
 ).get_caller_identity()["Account"]
 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}
 )

 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
 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.

Comment thread
laurenyu marked this conversation as resolved.

 Raises:
 botocore.exceptions.ClientError: If S3 throws an unexpected exception during bucket
 creation.
 If the exception is due to the bucket already existing or
 already being created, no exception is raised.

 """
 bucket = self.boto_session.resource("s3", region_name=region).Bucket(name=bucket_name)
 if bucket.creation_date is None:
 try:
 s3 = self.boto_session.resource("s3", region_name=region)
 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}
 )

 LOGGER.info("Created S3 bucket: %s", bucket_name)
 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
 else:
 raise

 def train( # noqa: C901
 self,
 input_mode,
Expand Down

29 changes: 9 additions & 20 deletions
tests/integ/kms_utils.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -15,8 +15,6 @@
import contextlib
import json

from botocore import exceptions

from sagemaker import utils

PRINCIPAL_TEMPLATE = (
Expand DownExpand Up@@ -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)
Expand All@@ -173,22 +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")
 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

 s3.put_bucket_encryption(
 sagemaker_session._create_s3_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": [
Expand All@@ -202,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

Expand Down

3 changes: 1 addition & 2 deletions
tests/integ/test_tf_script_mode.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -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")
 )
Expand Down

5 changes: 3 additions & 2 deletions
tests/unit/test_default_bucket.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -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
Comment thread
knakad marked this conversation as resolved.

 return sagemaker_session


def test_default_bucket_s3_create_call(sagemaker_session):
Expand Down