From b1bc0dea3812de3efbcc5230d88809b8a749e1d8 Mon Sep 17 00:00:00 2001 From: Akshay Shah Date: Fri, 9 Oct 2026 12:56:25 -0500 Subject: [PATCH] fix(core): don't send kms_key as an SSE-C key in S3Downloader.download S3Downloader.download passed kms_key as extra_args={'SSECustomerKey': kms_key}. SSECustomerKey is the SSE-C header, which expects a raw 256-bit customer key plus SSECustomerAlgorithm, not a KMS key ARN, so every download with kms_key failed with 400 Bad Request. S3 decrypts SSE-KMS objects transparently, so no key is needed on GET. This breaks DatasetBuilder.to_dataframe() for KMS-encrypted feature store output, as reported in #4034 (closed, but the v3 code is unchanged). Keep the kms_key parameter for backwards compatibility but no longer send it. --- sagemaker-core/src/sagemaker/core/s3/client.py | 13 +++++++------ sagemaker-core/tests/unit/s3/test_client.py | 4 +++- 2 files changed, 10 insertions(+), 7 deletions(-) diff --git a/sagemaker-core/src/sagemaker/core/s3/client.py b/sagemaker-core/src/sagemaker/core/s3/client.py index 15a9a8bb5f..8122b1efb3 100644 --- a/sagemaker-core/src/sagemaker/core/s3/client.py +++ b/sagemaker-core/src/sagemaker/core/s3/client.py @@ -143,7 +143,9 @@ def download(s3_uri, local_path, kms_key=None, sagemaker_session=None): Args: s3_uri (str): An S3 uri to download from. local_path (str): A local path to download the file(s) to. - kms_key (str): The KMS key to use to decrypt the files. + kms_key (str): Not used. S3 decrypts objects encrypted with SSE-KMS + transparently, so no key needs to be sent with the download request. + Kept for backwards compatibility. sagemaker_session (sagemaker.core.helper.session_helper.Session): Session object which manages interactions with Amazon SageMaker APIs and any other AWS services needed. If not specified, one is created @@ -154,13 +156,12 @@ def download(s3_uri, local_path, kms_key=None, sagemaker_session=None): """ sagemaker_session = sagemaker_session or Session() bucket, key_prefix = parse_s3_url(url=s3_uri) - if kms_key is not None: - extra_args = {"SSECustomerKey": kms_key} - else: - extra_args = None + # Do not pass kms_key as SSECustomerKey: that header is for SSE-C (a raw + # customer-provided key), and sending a KMS key ARN in it makes S3 reject the + # request with 400 Bad Request. return sagemaker_session.download_data( - path=local_path, bucket=bucket, key_prefix=key_prefix, extra_args=extra_args + path=local_path, bucket=bucket, key_prefix=key_prefix, extra_args=None ) @staticmethod diff --git a/sagemaker-core/tests/unit/s3/test_client.py b/sagemaker-core/tests/unit/s3/test_client.py index 0fb5f72bfd..aff3b4e631 100644 --- a/sagemaker-core/tests/unit/s3/test_client.py +++ b/sagemaker-core/tests/unit/s3/test_client.py @@ -107,7 +107,9 @@ def test_download_with_kms_key(sagemaker_session): path="/path/for/download/", bucket=BUCKET_NAME, key_prefix=os.path.join(CURRENT_JOB_NAME, SOURCE_NAME), - extra_args={"SSECustomerKey": KMS_KEY}, + # SSE-KMS objects are decrypted by S3 transparently; the KMS key must not be + # sent as an SSE-C customer key. + extra_args=None, )