[mohamedzeidan2021]

Issue #, if available

Fixes #6536

Description

ModelBuilder.build() defers source-code repacking to an internal _RepackModelStep when the model data is a pipeline variable. That step constructed its Compute with only an instance type:

compute = Compute(
    instance_type=kwargs.pop("instance_type", None) or INSTANCE_TYPE,
)

so ModelStep.repack_model_step_settings could not set the compute-volume KMS key — any volume_kms_key_id was forwarded to ModelTrainer as a stray kwarg instead of being applied to the job's Compute. As a result the runtime repack training job could not encrypt its storage volume with a customer-managed KMS key, and the generated training request contained no ResourceConfig.VolumeKmsKeyId.

Fix

• _utils.py — _RepackModelStep now pops volume_kms_key_id from kwargs into Compute (mirroring how instance_type is already handled), so it lands on ResourceConfig.VolumeKmsKeyId. When omitted it stays None and the field is not emitted (legacy behavior preserved).
• model_step.py — volume_kms_key_id is defaulted from the model's KMS key using setdefault, mirroring the existing output_data_config logic, so an explicit repack_model_step_settings override still wins.

Behavior

model_kms_key	repack_model_step_settings	Resulting ResourceConfig.VolumeKmsKeyId
set	—	model_kms_key
unset	—	absent (unchanged)
set	{"volume_kms_key_id": K}	K (override wins)
—	{"volume_kms_key_id": <PipelineVariable>}	forwarded unchanged

Testing done

• Reproduced the issue's offline script before the change (VolumeKmsKeyId absent → KeyError) and confirmed it passes after.
• Added unit tests covering an explicit key, an omitted key, a user override, and a pipeline-variable key, at both the ModelStep and _RepackModelStep layers.
• pytest sagemaker-mlops/tests/unit/workflow/ → 418 passed.
• black --check and flake8 (repo config) clean on all changed files.

Merge Checklist

Put an x in the boxes that apply.

• I have read the CONTRIBUTING doc
• I used the commit message format described in CONTRIBUTING
• I have passed the region in to all S3 and STS clients that I've initialized as part of this change (N/A)
• I have updated any necessary documentation, including READMEs and API docs (if appropriate)

Tests

• I have added tests that prove my fix is effective or that my feature works (if appropriate)
• I have added unit and/or integration tests as appropriate to ensure backward compatibility of the changes

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

[github-actions]

🤖 AI code review

This PR fixes GH #6536 by threading volume_kms_key_id through _RepackModelStep into the repack job's Compute (so it lands on ResourceConfig.VolumeKmsKeyId), and defaulting it from the model's KMS key in ModelStep._append_repack_model_step via setdefault so explicit overrides still win. The change mirrors the existing output_data_config wiring pattern (PR #6338), Compute from sagemaker.core.training.configs does accept volume_kms_key_id, and _repack_model_step_settings is always a dict so setdefault is safe. Tests cover the explicit-key, omitted-key, user-override, and pipeline-variable cases at both layers. The change looks clean and I found no issues.

_{Reviewed commit d604bd0. Automated review; verify before acting.}