diff --git a/sagemaker-train/src/sagemaker/train/tuner.py b/sagemaker-train/src/sagemaker/train/tuner.py index 119ec7b75a..f906c83317 100644 --- a/sagemaker-train/src/sagemaker/train/tuner.py +++ b/sagemaker-train/src/sagemaker/train/tuner.py @@ -590,7 +590,10 @@ def _build_driver_and_code_channels(cls, model_trainer): arcname = os.path.relpath(fpath, source_code.source_dir) tar.add(fpath, arcname=arcname) s3_client = session.boto_session.client("s3", region_name=session.boto_region_name) - s3_client.upload_file(tar_path, bucket, s3_key) + # Encrypt staged source with the trainer's output KMS key (parity with + # create_input_data_channel); None when no key is set. See issue #6555. + extra_args = model_trainer._resolve_upload_extra_args() + s3_client.upload_file(tar_path, bucket, s3_key, ExtraArgs=extra_args) model_trainer.hyperparameters["sagemaker_submit_directory"] = ( f"s3://{bucket}/{s3_key}" ) diff --git a/sagemaker-train/tests/unit/train/test_tuner_driver_channels.py b/sagemaker-train/tests/unit/train/test_tuner_driver_channels.py index 97b747dfd5..1b71aa8b06 100644 --- a/sagemaker-train/tests/unit/train/test_tuner_driver_channels.py +++ b/sagemaker-train/tests/unit/train/test_tuner_driver_channels.py @@ -277,6 +277,45 @@ def test_uploads_sourcedir_tar_gz(self, tmp_path): # sagemaker_submit_directory set assert trainer.hyperparameters["sagemaker_submit_directory"].startswith("s3://my-bucket/") + def _upload_trainer(self, tmp_path, extra_args): + """Mock trainer whose source upload returns ``extra_args`` from the resolver.""" + src_dir = str(tmp_path / "src") + os.makedirs(src_dir) + with open(os.path.join(src_dir, "train.py"), "w") as f: + f.write("print('hello')") + + mock_s3_client = MagicMock() + mock_session = MagicMock() + mock_session.default_bucket.return_value = "my-bucket" + mock_session.default_bucket_prefix = None + mock_session.boto_session.client.return_value = mock_s3_client + mock_session.boto_region_name = "us-west-2" + + trainer = _mock_model_trainer(source_code=_make_source_code(source_dir=src_dir)) + trainer.sagemaker_session = mock_session + trainer.create_input_data_channel.return_value = _create_channel("sm_drivers") + trainer.hyperparameters = {} + trainer._resolve_upload_extra_args.return_value = extra_args + return trainer, mock_s3_client + + def test_uploads_sourcedir_with_output_kms_key(self, tmp_path): + """Local source upload must forward the trainer's KMS ExtraArgs (issue #6555).""" + extra_args = {"ServerSideEncryption": "aws:kms", "SSEKMSKeyId": "my-kms-key"} + trainer, mock_s3_client = self._upload_trainer(tmp_path, extra_args) + + HyperparameterTuner._build_driver_and_code_channels(trainer) + + trainer._resolve_upload_extra_args.assert_called_once_with() + assert mock_s3_client.upload_file.call_args.kwargs["ExtraArgs"] == extra_args + + def test_uploads_sourcedir_without_kms_key(self, tmp_path): + """No KMS key resolved means ExtraArgs=None, preserving default behavior.""" + trainer, mock_s3_client = self._upload_trainer(tmp_path, None) + + HyperparameterTuner._build_driver_and_code_channels(trainer) + + assert mock_s3_client.upload_file.call_args.kwargs["ExtraArgs"] is None + def test_sets_submit_directory_for_s3_source_dir(self, tmp_path): """Should set sagemaker_submit_directory directly when source_dir is an S3 URI.""" source_code = _make_source_code(source_dir="s3://my-bucket/code/sourcedir.tar.gz")