From 7ad6acfeed154d52df61df5433dbcbeced7cba9b Mon Sep 17 00:00:00 2001 From: Kartik Shenoy Date: Fri, 5 Jun 2026 07:47:34 -0700 Subject: [PATCH 1/3] fix: use SSE-KMS in _validate_s3_path_exists when kms_key_id is set When a target S3 bucket enforces SSE-KMS, _validate_s3_path_exists previously failed with access denied because its put_object call to create a missing prefix did not specify encryption. Thread the SFTTrainer kms_key_id (already available in _create_output_config) into _validate_s3_path_exists and include ServerSideEncryption=aws:kms and SSEKMSKeyId when it is provided. Fixes #5930 --- .../train/common_utils/finetune_utils.py | 20 ++++++++-- .../train/common_utils/test_finetune_utils.py | 40 ++++++++++++++++++- 2 files changed, 55 insertions(+), 5 deletions(-) diff --git a/sagemaker-train/src/sagemaker/train/common_utils/finetune_utils.py b/sagemaker-train/src/sagemaker/train/common_utils/finetune_utils.py index f8d1d2e5b9..ab3d320dc3 100644 --- a/sagemaker-train/src/sagemaker/train/common_utils/finetune_utils.py +++ b/sagemaker-train/src/sagemaker/train/common_utils/finetune_utils.py @@ -1385,7 +1385,7 @@ def _create_output_config( s3_output_path = _get_default_s3_output_path(sagemaker_session) # Validate S3 path exists - _validate_s3_path_exists(s3_output_path, sagemaker_session) + _validate_s3_path_exists(s3_output_path, sagemaker_session, kms_key_id=kms_key_id) config_kwargs = { "s3_output_path": s3_output_path, @@ -1486,8 +1486,16 @@ def _validate_eula_for_gated_model(model, accept_eula, is_gated_model): return accept_eula -def _validate_s3_path_exists(s3_path: str, sagemaker_session): - """Validate S3 path and create bucket/prefix if they don't exist.""" +def _validate_s3_path_exists(s3_path: str, sagemaker_session, kms_key_id: Optional[str] = None): + """Validate S3 path and create bucket/prefix if they don't exist. + + Args: + s3_path: S3 path to validate. + sagemaker_session: SageMaker session used to create the S3 client. + kms_key_id: Optional KMS key ID. When provided, objects created to + initialize a missing prefix are written with SSE-KMS encryption so + that validation does not fail on buckets enforcing KMS encryption. + """ if not s3_path.startswith("s3://"): raise ValueError(f"Invalid S3 path format: {s3_path}") @@ -1535,7 +1543,11 @@ def _validate_s3_path_exists(s3_path: str, sagemaker_session): # Create the prefix by putting an empty object if not prefix.endswith("/"): prefix += "/" - s3_client.put_object(Bucket=bucket_name, Key=prefix, Body=b"") + put_object_kwargs = {"Bucket": bucket_name, "Key": prefix, "Body": b""} + if kms_key_id: + put_object_kwargs["ServerSideEncryption"] = "aws:kms" + put_object_kwargs["SSEKMSKeyId"] = kms_key_id + s3_client.put_object(**put_object_kwargs) except Exception as e: raise ValueError(f"Failed to validate/create S3 path '{s3_path}': {str(e)}") diff --git a/sagemaker-train/tests/unit/train/common_utils/test_finetune_utils.py b/sagemaker-train/tests/unit/train/common_utils/test_finetune_utils.py index 2f948d3b5f..1dec6e4e9b 100644 --- a/sagemaker-train/tests/unit/train/common_utils/test_finetune_utils.py +++ b/sagemaker-train/tests/unit/train/common_utils/test_finetune_utils.py @@ -674,7 +674,9 @@ def test__create_output_config(self, mock_validate_s3): assert config.s3_output_path == "s3://bucket/output" assert config.kms_key_id == "kms-key" - mock_validate_s3.assert_called_once_with("s3://bucket/output", mock_session) + mock_validate_s3.assert_called_once_with( + "s3://bucket/output", mock_session, kms_key_id="kms-key" + ) def test__convert_input_data_to_channels(self): @@ -879,6 +881,42 @@ def test__validate_s3_path_exists_with_prefix_not_exists(self, mock_boto_client) Bucket="test-bucket", Key="prefix/", Body=b"" ) + @patch('boto3.client') + def test__validate_s3_path_exists_with_prefix_not_exists_kms(self, mock_boto_client): + """Test S3 path validation uses SSE-KMS when kms_key_id is provided""" + mock_session = Mock() + mock_s3_client = Mock() + mock_session.boto_session.client.return_value = mock_s3_client + mock_s3_client.list_objects_v2.return_value = {} # No contents + + _validate_s3_path_exists("s3://test-bucket/prefix", mock_session, kms_key_id="kms-key") + + mock_s3_client.put_object.assert_called_once_with( + Bucket="test-bucket", + Key="prefix/", + Body=b'', + ServerSideEncryption="aws:kms", + SSEKMSKeyId="kms-key", + ) + + @patch('boto3.client') + def test__validate_s3_path_exists_with_prefix_not_exists_kms(self, mock_boto_client): + """Test S3 path validation uses SSE-KMS when kms_key_id is provided""" + mock_session = Mock() + mock_s3_client = Mock() + mock_session.boto_session.client.return_value = mock_s3_client + mock_s3_client.list_objects_v2.return_value = {} # No contents + + _validate_s3_path_exists("s3://test-bucket/prefix", mock_session, kms_key_id="kms-key") + + mock_s3_client.put_object.assert_called_once_with( + Bucket="test-bucket", + Key="prefix/", + Body=b'', + ServerSideEncryption="aws:kms", + SSEKMSKeyId="kms-key", + ) + class TestMlflowVersionMeetsMinimum: def test_meets_minimum(self): From 26901707f11edb30dccf1c348064b3c585787d9f Mon Sep 17 00:00:00 2001 From: Kartik Shenoy Date: Fri, 5 Jun 2026 08:11:43 -0700 Subject: [PATCH 2/3] Removed duplicate --- .../train/common_utils/test_finetune_utils.py | 18 ------------------ 1 file changed, 18 deletions(-) diff --git a/sagemaker-train/tests/unit/train/common_utils/test_finetune_utils.py b/sagemaker-train/tests/unit/train/common_utils/test_finetune_utils.py index 1dec6e4e9b..ad2001f076 100644 --- a/sagemaker-train/tests/unit/train/common_utils/test_finetune_utils.py +++ b/sagemaker-train/tests/unit/train/common_utils/test_finetune_utils.py @@ -899,24 +899,6 @@ def test__validate_s3_path_exists_with_prefix_not_exists_kms(self, mock_boto_cli SSEKMSKeyId="kms-key", ) - @patch('boto3.client') - def test__validate_s3_path_exists_with_prefix_not_exists_kms(self, mock_boto_client): - """Test S3 path validation uses SSE-KMS when kms_key_id is provided""" - mock_session = Mock() - mock_s3_client = Mock() - mock_session.boto_session.client.return_value = mock_s3_client - mock_s3_client.list_objects_v2.return_value = {} # No contents - - _validate_s3_path_exists("s3://test-bucket/prefix", mock_session, kms_key_id="kms-key") - - mock_s3_client.put_object.assert_called_once_with( - Bucket="test-bucket", - Key="prefix/", - Body=b'', - ServerSideEncryption="aws:kms", - SSEKMSKeyId="kms-key", - ) - class TestMlflowVersionMeetsMinimum: def test_meets_minimum(self): From 5c77980c2a76b7d36fd7256daeeec2e7d80e6494 Mon Sep 17 00:00:00 2001 From: Kartik Shenoy Date: Thu, 8 Oct 2026 22:27:49 -0700 Subject: [PATCH 3/3] style: black-format SSE-KMS test in test_finetune_utils Co-Authored-By: Claude Opus 5.5 --- .../tests/unit/train/common_utils/test_finetune_utils.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/sagemaker-train/tests/unit/train/common_utils/test_finetune_utils.py b/sagemaker-train/tests/unit/train/common_utils/test_finetune_utils.py index ad2001f076..a2f88ac9de 100644 --- a/sagemaker-train/tests/unit/train/common_utils/test_finetune_utils.py +++ b/sagemaker-train/tests/unit/train/common_utils/test_finetune_utils.py @@ -881,7 +881,7 @@ def test__validate_s3_path_exists_with_prefix_not_exists(self, mock_boto_client) Bucket="test-bucket", Key="prefix/", Body=b"" ) - @patch('boto3.client') + @patch("boto3.client") def test__validate_s3_path_exists_with_prefix_not_exists_kms(self, mock_boto_client): """Test S3 path validation uses SSE-KMS when kms_key_id is provided""" mock_session = Mock() @@ -894,7 +894,7 @@ def test__validate_s3_path_exists_with_prefix_not_exists_kms(self, mock_boto_cli mock_s3_client.put_object.assert_called_once_with( Bucket="test-bucket", Key="prefix/", - Body=b'', + Body=b"", ServerSideEncryption="aws:kms", SSEKMSKeyId="kms-key", )