fix(train): apply ModelTrainer output KMS key to tuner source upload …

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}"
                 )

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")