fix: apply output KMS key to ModelTrainer source-code S3 uploads (#5956)

[mohamedzeidan2021]

Issue

Fixes #5956.

The v2 Estimator encrypted staged user code in S3 using output_kms_key (see _stage_user_code_in_s3). The v3 ModelTrainer uploads source code and driver files to S3 (via create_input_data_channel → Session.upload_data) with no KMS encryption and no way to configure one. In environments whose S3 bucket policies enforce SSE-KMS, this makes ModelTrainer unusable with custom training scripts — a regression from v2.

Fix

create_input_data_channel now encrypts local uploads with a KMS key when one is configured:

• New optional kms_key: Optional[str] = None argument (appended last; keyword-compatible).
• New _resolve_upload_extra_args: uses the explicit kms_key if given, else falls back to output_data_config.kms_key_id — mirroring v2, where the output KMS key also encrypted staged code. It returns S3 ExtraArgs {"ServerSideEncryption": "aws:kms", "SSEKMSKeyId": <key>}, or None when no usable string key is set (so a PipelineVariable/Unassigned/None key is ignored).
• Both upload_data calls now pass extra_args=.

Because the internal source_code, sm_drivers, and recipe channels all route through create_input_data_channel, they inherit the configured output_data_config.kms_key_id automatically. S3-URI / S3DataSource / FileSystemDataSource / local-container sources don't upload and are unaffected.

Validation

• Unit (test_model_trainer.py): new tests assert the KMS ExtraArgs are passed when output_data_config.kms_key_id is set, when an explicit kms_key overrides it, and that extra_args stays None by default. Fail on master, pass with the fix; full module green (72 passed).
• Real S3: an upload_data(..., extra_args={"ServerSideEncryption":"aws:kms","SSEKMSKeyId":...}) produced an object that head_object reports as ServerSideEncryption=aws:kms with a resolved KMS key ARN.

Backwards compatibility

Default behavior is byte-identical: with no KMS key configured, extra_args is None and uploads are unchanged. The new argument is optional and keyword-compatible; no public signature is broken.

[github-actions]

Claude finished @mohamedzeidan2021's task in 1m 28s —— View job

---

Code Review

• Read the PR diff
• Verify upload_data signature supports extra_args
• Verify output_data_config attribute & internal channel routing
• Check correctness, tests, backward compatibility
• Post review

(The inline-comment tool isn't available in this run, so findings are inline below with file:line references.)

---

Overall this is a clean, correct, well-scoped fix that restores v2 parity for encrypting staged user code. The design is sound: routing all local uploads through create_input_data_channel means the internal source_code, sm_drivers, and recipe channels inherit the KMS key automatically, and non-uploading sources (S3 URI / S3DataSource / FileSystemDataSource / local-container) are correctly untouched.

I verified the key claims:

• Session.upload_data(..., extra_args=None) accepts extra_args and forwards them to the S3 upload (sagemaker-core/.../session_helper.py:391), merging cleanly with the injected ExpectedBucketOwner — no conflict with the KMS keys.
• _populate_intelligent_defaults() (called at model_trainer.py:621, before the internal channels are built at 652/715/738) populates self.output_data_config at line 586, so the fallback resolves correctly for SDK-managed channels.
• The two upload_data call sites (model_trainer.py:1143, 1150) are the only upload paths in the module; both now pass extra_args.
• The new kms_key parameter is appended last and optional — no public signature break.

Minor observations (non-blocking)

1. Compatibility nuance in the PR description — model_trainer.py:1192-1194
   The body states "default behavior is byte-identical." That's true only when no output KMS key is set. For users who already set output_data_config.kms_key_id (common, since it governs output artifacts), source/driver uploads now become KMS-encrypted where they previously weren't. That is the intended v2 parity, but it's a genuine behavior change for those users: if the upload role lacks kms:GenerateDataKey/encrypt permission on that key, uploads that used to succeed will now fail. Worth calling out explicitly in the changelog/PR notes rather than framing it as fully transparent.

2. Untested branches — test_model_trainer.py
   The docstring promises a PipelineVariable/Unassigned key is ignored (_resolve_upload_extra_args returns None for non-str keys), but there's no test for that path, nor one asserting that an S3-URI/S3DataSource source performs no encrypted upload. The three added tests are good; a case for the non-string-key fallback would lock in the documented contract cheaply.

3. Minor: getattr defensiveness — model_trainer.py:1194
getattr(self.output_data_config, "kms_key_id", None) is slightly over-defensive since OutputDataConfig always defines kms_key_id; a direct attribute access would read fine. Not worth changing.

No correctness, security, or performance concerns. The KMS ExtraArgs shape (ServerSideEncryption: "aws:kms" + SSEKMSKeyId) is the correct boto3 S3 form.