Skip to content

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

Merged
mufaddal-rohawala merged 2 commits into
aws:masterfrom
mufaddal-rohawala:fix/tuner-source-kms-6555
Nov 6, 2026
Merged

mufaddal-rohawala merged 2 commits into
aws:masterfrom
mufaddal-rohawala:fix/tuner-source-kms-6555

Conversation

@mufaddal-rohawala

Copy link
Copy Markdown
Member

…(#6555)

HyperparameterTuner._build_driver_and_code_channels() uploaded the local sourcedir.tar.gz with s3_client.upload_file() and no ExtraArgs, so the ModelTrainer's output_data_config.kms_key_id was never applied. The staged code could be encrypted with the wrong key, and uploads fail where a bucket policy requires explicit SSE-KMS headers.

Build ExtraArgs via ModelTrainer._resolve_upload_extra_args() and forward them to upload_file (parity with create_input_data_channel); the result is None when no key is configured, so default behavior is unchanged. Add regression tests covering the KMS and no-KMS cases.

Issue #, if available:

Description of changes:

By submitting this pull request, I confirm that you can use, modify, copy, and redistribute this contribution, under the terms of your choice.

…ws#6555)

HyperparameterTuner._build_driver_and_code_channels() uploaded the local
sourcedir.tar.gz with s3_client.upload_file() and no ExtraArgs, so the
ModelTrainer's output_data_config.kms_key_id was never applied. The staged
code could be encrypted with the wrong key, and uploads fail where a bucket
policy requires explicit SSE-KMS headers.

Build ExtraArgs via ModelTrainer._resolve_upload_extra_args() and forward them
to upload_file (parity with create_input_data_channel); the result is None when
no key is configured, so default behavior is unchanged. Add regression tests
covering the KMS and no-KMS cases.

@github-actions github-actions Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🤖 AI code review

This PR forwards the ModelTrainer's output KMS key to the tuner's sourcedir.tar.gz upload in HyperparameterTuner._build_driver_and_code_channels() by calling the existing _resolve_upload_extra_args() and passing the result as ExtraArgs to s3_client.upload_file(). This restores parity with create_input_data_channel (PR #6339) for staged code encryption. The resolver returns None by default, which matches boto3's upload_file default so unconfigured behavior is unchanged, and ExtraArgs=None is accepted by boto3. Regression tests cover both the KMS and no-KMS paths. The change is small, correct, and adequately tested; no issues found.

Reviewed commit fb43de7. Automated review; verify before acting.

@github-actions github-actions Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🤖 AI code review

This PR fixes HyperparameterTuner._build_driver_and_code_channels() so that the staged sourcedir.tar.gz upload applies the ModelTrainer's output KMS key via ModelTrainer._resolve_upload_extra_args(), matching the behavior of create_input_data_channel (precedent #6339). The change is minimal and correct: _resolve_upload_extra_args() returns None when no key is configured, and s3_client.upload_file(ExtraArgs=None) preserves default behavior. Regression tests cover both the KMS and no-KMS cases. The PR looks clean; no issues found.

Reviewed commit 6bf1805. Automated review; verify before acting.

@mufaddal-rohawala
mufaddal-rohawala merged commit 50fe39c into aws:master Nov 6, 2026
23 of 27 checks passed
jam-jee added a commit that referenced this pull request Nov 12, 2026
…6555) (#6578)

* fix(train): Stage tuner source in the trainer's bucket

HyperparameterTuner uploaded the local sourcedir.tar.gz straight to
session.default_bucket() under a key with no default_bucket_prefix.
ModelTrainer.create_input_data_channel, which stages the sm_drivers
channel for the same tuning job, resolves the bucket through
_resolve_staging_bucket() (falling back to the output-path bucket when
the training role cannot write the default bucket) and prefixes keys
with default_bucket_prefix. The source archive now follows the same
rules, so it lands next to sm_drivers and stays reachable by the role.

PR #6559 fixed the KMS half of issue #6555; this is the bucket half.

Closes #6555
X-AI-Prompt: Fix the open SDK issues from TT V2409806538 (#6555 bucket gap, #6554)
X-AI-Tool: kiro-cli

* fix(core): Qualify shape Tag annotations in resources

resources.py star-imports the shapes and later defines the resource
class Tag (ListTags). Annotations written above that class, such as
HyperParameterTuningJob.create(tags: Optional[List[Tag]]), meant the
shape only because Python <= 3.13 evaluates annotations eagerly. Python
3.14 evaluates them lazily (PEP 649), so they resolve to resources.Tag
and pydantic rejects the shapes.Tag objects HyperparameterTuner sends.
Annotations written below the class already resolved to resources.Tag
on every version (TrainingJob.create and eight others), which is why
ModelTrainer had to pass dicts.

Resource codegen now emits shapes listed in
RESOURCE_SHADOWED_SHAPE_NAMES (only Tag) as shapes.<Name> and imports
sagemaker.core.shapes; resources.py is regenerated, and the diff is
only the 75 tags annotations plus that import. Every generated tags
parameter now takes shapes.Tag or a dict on all Python versions. A
guard test fails if a bare annotation names a shape that a later
resource class shadows.

Closes #6554
X-AI-Prompt: Fix the open SDK issues from TT V2409806538 (#6555 bucket gap, #6554)
X-AI-Tool: kiro-cli

* style(core): Add blank line after test module docstring

black 26.3.1 (pinned in requirements/tox/black_requirements.txt)
requires an empty line after a module docstring; the
codestyle-doc-tests (sagemaker-core) check failed on it.

X-AI-Prompt: Fix the failing sagemaker-core doc test on PR 6578
X-AI-Tool: kiro-cli

This branch was previously deployed

1 inactive deployment
auto-approve — 6bf18058 Deployed Nov 5, 2026 by mufaddal-rohawala via wait-for-approval #2175
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants