fix(train): enforce S3 ownership on ai_registry default bucket - #6275
Conversation
The ai_registry and finetune_utils modules derive a predictable default
bucket name (sagemaker-{region}-{account_id}) and read from / write to it
without verifying ownership, unlike sagemaker-core which enforces
ExpectedBucketOwner. Because S3 bucket names are globally unique, another
account could pre-create that name and the SDK would silently use it.
Add the ownership guard for the SDK-derived default bucket (explicitly
provided buckets are left untouched):
- air_hub.upload_to_s3 / download_from_s3 pass ExpectedBucketOwner via
ExtraArgs when the target is the derived default bucket.
- _create_mlflow_app (and _create_mlflow_app_as_upgrade) verify ownership
and pass ExpectedBucketOwner on the S3 probe/writes before registering
the MLflow ArtifactStoreUri.
- _validate_s3_path_exists verifies ownership of the derived default bucket
before it is created or used as OutputDataConfig.
Add unit tests asserting the guard is applied for the default bucket, left
off for explicit buckets, and fails safe on a foreign-owned bucket.
|
Claude finished @rsareddy0329's task in 2m 35s —— View job Code Review: enforce S3 ownership on ai_registry default bucket
The intent is sound and well-documented: guard the SDK-derived default bucket (
1.
|
Remove an internal ticket identifier from a test docstring; no behavior or test-logic change.
|
Claude finished @rsareddy0329's task in 1m 36s —— View job Code Review: enforce S3 ownership on ai_registry default bucket
The change is a genuine, well-scoped security improvement: it guards only the SDK-derived default bucket ( 1.
|
| error_code = str(e.response.get("Error", {}).get("Code", "")) | ||
| if error_code in ("404", "NoSuchBucket", "NotFound"): | ||
| return # Bucket does not exist yet; the caller creates it in-account. | ||
| if error_code in ("403", "AccessDenied"): |
There was a problem hiding this comment.
Suggestion (non-blocking): Consider making this error message neutral and omitting the account ID. A 403 from HeadBucket(..., ExpectedBucketOwner=...) can indicate either an ownership mismatch or missing s3:ListBucket permission, so it does not establish who owns the bucket. For example:
raise ValueError(
f"Refusing to use default bucket '{bucket_name}': ownership or access "
f"could not be verified. Confirm that the bucket belongs to your AWS "
f"account and that the caller has s3:ListBucket permission, or provide "
f"an explicit bucket you own."
)This avoids exposing account identifiers or directing users toward an ownership diagnosis that may be incorrect.
|
nit : Can we add a integ test in |
The ai_registry and finetune_utils modules derive a predictable default bucket name (sagemaker-{region}-{account_id}) and read from / write to it without verifying ownership, unlike sagemaker-core which enforces ExpectedBucketOwner. Because S3 bucket names are globally unique, another account could pre-create that name and the SDK would silently use it.
Add the ownership guard for the SDK-derived default bucket (explicitly provided buckets are left untouched):
Add unit tests asserting the guard is applied for the default bucket, left off for explicit buckets, and fails safe on a foreign-owned bucket.
Issue #, if available:
Description of changes:
By submitting this pull request, I confirm that you can use, modify, copy, and redistribute this contribution, under the terms of your choice.