Repository navigation
fix(train): apply ModelTrainer output KMS key to tuner source upload … - #6559
Conversation
…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.
There was a problem hiding this comment.
🤖 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.
There was a problem hiding this comment.
🤖 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.
…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
…(#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.