fix: set VolumeKmsKeyId on ModelStep runtime repack job (#6537)

ModelBuilder.build() defers source-code repacking to a _RepackModelStep
when model data is a pipeline variable. That step built its Compute with
only an instance type, so repack_model_step_settings could not set the
compute-volume KMS key: unknown settings were forwarded to ModelTrainer
rather than the job's Compute. As a result the runtime repack training
job could not encrypt its storage volume with a customer-managed key,
and the generated request had no ResourceConfig.VolumeKmsKeyId.

Thread volume_kms_key_id from repack_model_step_settings into the repack
job's Compute (mirroring how instance_type is handled), and default it
from the model's KMS key using setdefault so an explicit override wins
(mirroring the existing output_data_config logic).

Add unit tests covering an explicit key, an omitted key, a user
override, and a pipeline-variable key, at both the ModelStep and
_RepackModelStep layers.

Fixes #6536

Co-authored-by: Mohamed Zeidan <zeidmo@amazon.com>

Author: mohamedzeidan2021

sagemaker-mlops/src/sagemaker/mlops/workflow/_utils.py

@@ -173,7 +173,10 @@ def __init__(
                     job will be created without VPC config (default: None).
             security_group_ids (list[str]): List of security group ids. If not
                 specified, the re-packing job will be created without VPC config (default: None).
-            **kwargs: additional arguments for the repacking job.
+            **kwargs: additional arguments for the repacking job. ``volume_kms_key_id``
+                is consumed here and applied to the repack job's ``Compute`` so the
+                training volume is encrypted with the given KMS key; any remaining
+                keys are forwarded to the underlying ``ModelTrainer``.
         """
         self._model_data = model_data
         self.sagemaker_session = sagemaker_session
@@ -207,6 +210,7 @@ def __init__(
 
         compute = Compute(
             instance_type=kwargs.pop("instance_type", None) or INSTANCE_TYPE,
+            volume_kms_key_id=kwargs.pop("volume_kms_key_id", None),
         )
 
         networking = None

sagemaker-mlops/src/sagemaker/mlops/workflow/model_step.py

@@ -291,6 +291,15 @@ def _append_repack_model_step(self):
                             kms_key_id=fields["model_kms_key"],
                         ),
                     )
+                # Encrypt the repack job's compute volume with the same customer-managed
+                # key as the model. This reaches _RepackModelStep via **kwargs, which
+                # applies it to the job's Compute (ResourceConfig.VolumeKmsKeyId). Without
+                # it the runtime repack cannot use a configured CMK for its volume (GH
+                # #6536). setdefault so an explicit repack_model_step_settings override wins.
+                if fields["model_kms_key"]:
+                    self._repack_model_step_settings.setdefault(
+                        "volume_kms_key_id", fields["model_kms_key"]
+                    )
                 repack_model_step = _RepackModelStep(
                     name="{}-{}-{}".format(self.name, _REPACK_MODEL_NAME_BASE, name_base),
                     sagemaker_session=(

sagemaker-mlops/tests/unit/workflow/test_model_step.py

@@ -161,6 +161,137 @@ def test_repack_step_output_config_respects_user_override():
     assert kwargs["output_data_config"] is mine
 
 
+def test_repack_step_gets_volume_kms_key_from_model():
+    """GH #6536: the repack job's compute volume must be encrypted with the model's KMS
+    key. This reaches _RepackModelStep via volume_kms_key_id, which applies it to the
+    job's Compute (ResourceConfig.VolumeKmsKeyId). Without it the runtime repack cannot
+    use a configured CMK for its volume."""
+    from sagemaker.serve.model_builder import ModelBuilder
+    from sagemaker.mlops.workflow import model_step as ms
+
+    ps = _pipeline_session()
+    builder = Mock(spec=ModelBuilder)
+    builder.sagemaker_session = ps
+    builder.model_name = "my-model"
+    builder.role_arn = "arn:aws:iam::111122223333:role/R"
+    builder.s3_model_data_url = "s3://orig/model.tar.gz"
+    builder.entry_point = "inference.py"
+    builder.source_dir = "/code"
+    builder.source_code = Mock(requirements="requirements.txt")
+    builder.vpc_config = None
+    builder.model_kms_key = "arn:aws:kms:us-west-2:111122223333:key/abc"
+
+    step_args = _FakeModelStepArgs(builder, {id(builder)})
+    fake_repack = Mock()
+    fake_repack.properties.ModelArtifacts.S3ModelArtifacts = "s3://repacked/model.tar.gz"
+
+    with patch("sagemaker.core.workflow.utilities.validate_step_args_input"):
+        with patch.object(ms, "_RepackModelStep", return_value=fake_repack) as mock_repack:
+            ms.ModelStep(name="step", step_args=step_args)
+
+    _, kwargs = mock_repack.call_args
+    assert kwargs["volume_kms_key_id"] == "arn:aws:kms:us-west-2:111122223333:key/abc"
+
+
+def test_repack_step_omits_volume_kms_key_when_model_has_none():
+    """With no model KMS key, volume_kms_key_id must not be sent, leaving the repack job's
+    ResourceConfig.VolumeKmsKeyId unset (legacy behavior preserved)."""
+    from sagemaker.serve.model_builder import ModelBuilder
+    from sagemaker.mlops.workflow import model_step as ms
+
+    ps = _pipeline_session()
+    builder = Mock(spec=ModelBuilder)
+    builder.sagemaker_session = ps
+    builder.model_name = "my-model"
+    builder.role_arn = "arn:aws:iam::111122223333:role/R"
+    builder.s3_model_data_url = "s3://orig/model.tar.gz"
+    builder.entry_point = "inference.py"
+    builder.source_dir = "/code"
+    builder.source_code = Mock(requirements="requirements.txt")
+    builder.vpc_config = None
+    builder.model_kms_key = None
+
+    step_args = _FakeModelStepArgs(builder, {id(builder)})
+    fake_repack = Mock()
+    fake_repack.properties.ModelArtifacts.S3ModelArtifacts = "s3://repacked/model.tar.gz"
+
+    with patch("sagemaker.core.workflow.utilities.validate_step_args_input"):
+        with patch.object(ms, "_RepackModelStep", return_value=fake_repack) as mock_repack:
+            ms.ModelStep(name="step", step_args=step_args)
+
+    _, kwargs = mock_repack.call_args
+    assert "volume_kms_key_id" not in kwargs
+
+
+def test_repack_step_volume_kms_key_respects_user_override():
+    """An explicit volume_kms_key_id in repack_model_step_settings must win over the model
+    KMS key default."""
+    from sagemaker.serve.model_builder import ModelBuilder
+    from sagemaker.mlops.workflow import model_step as ms
+
+    ps = _pipeline_session()
+    builder = Mock(spec=ModelBuilder)
+    builder.sagemaker_session = ps
+    builder.model_name = "my-model"
+    builder.role_arn = "arn:aws:iam::111122223333:role/R"
+    builder.s3_model_data_url = "s3://orig/model.tar.gz"
+    builder.entry_point = "inference.py"
+    builder.source_dir = "/code"
+    builder.source_code = Mock(requirements="requirements.txt")
+    builder.vpc_config = None
+    builder.model_kms_key = "arn:aws:kms:us-west-2:111122223333:key/abc"
+
+    step_args = _FakeModelStepArgs(builder, {id(builder)})
+    fake_repack = Mock()
+    fake_repack.properties.ModelArtifacts.S3ModelArtifacts = "s3://repacked/model.tar.gz"
+
+    with patch("sagemaker.core.workflow.utilities.validate_step_args_input"):
+        with patch.object(ms, "_RepackModelStep", return_value=fake_repack) as mock_repack:
+            ms.ModelStep(
+                name="step",
+                step_args=step_args,
+                repack_model_step_settings={"volume_kms_key_id": "my-volume-key"},
+            )
+
+    _, kwargs = mock_repack.call_args
+    assert kwargs["volume_kms_key_id"] == "my-volume-key"
+
+
+def test_repack_step_volume_kms_key_accepts_pipeline_variable():
+    """A pipeline variable supplied as volume_kms_key_id must be forwarded unchanged."""
+    from sagemaker.serve.model_builder import ModelBuilder
+    from sagemaker.core.workflow import ParameterString
+    from sagemaker.mlops.workflow import model_step as ms
+
+    ps = _pipeline_session()
+    builder = Mock(spec=ModelBuilder)
+    builder.sagemaker_session = ps
+    builder.model_name = "my-model"
+    builder.role_arn = "arn:aws:iam::111122223333:role/R"
+    builder.s3_model_data_url = "s3://orig/model.tar.gz"
+    builder.entry_point = "inference.py"
+    builder.source_dir = "/code"
+    builder.source_code = Mock(requirements="requirements.txt")
+    builder.vpc_config = None
+    builder.model_kms_key = None
+
+    volume_key = ParameterString(name="volume-kms")
+    step_args = _FakeModelStepArgs(builder, {id(builder)})
+    fake_repack = Mock()
+    fake_repack.properties.ModelArtifacts.S3ModelArtifacts = "s3://repacked/model.tar.gz"
+
+    with patch("sagemaker.core.workflow.utilities.validate_step_args_input"):
+        with patch.object(ms, "_RepackModelStep", return_value=fake_repack) as mock_repack:
+            ms.ModelStep(
+                name="step",
+                step_args=step_args,
+                repack_model_step_settings={"volume_kms_key_id": volume_key},
+            )
+
+    _, kwargs = mock_repack.call_args
+    assert kwargs["volume_kms_key_id"] is volume_key
+
+
 def test_no_repack_step_for_unrecognized_model_type():
     """An object that is neither a core Model nor a ModelBuilder yields no repack step."""
     from sagemaker.mlops.workflow import model_step as ms

sagemaker-mlops/tests/unit/workflow/test_utils.py

@@ -207,6 +207,63 @@ def test_init_with_custom_instance_type(self, mock_session, temp_entry_point):
                     call_kwargs = mock_retrieve.call_args[1]
                     assert call_kwargs["instance_type"] == "ml.p3.2xlarge"
 
+    def test_init_threads_volume_kms_key_id_to_compute(self, mock_session, temp_entry_point):
+        """GH #6536: volume_kms_key_id must be applied to the repack job's Compute
+        (ResourceConfig.VolumeKmsKeyId) and consumed here rather than forwarded to
+        ModelTrainer."""
+        with patch("sagemaker.mlops.workflow._utils.image_uris.retrieve") as mock_retrieve:
+            with patch("sagemaker.train.ModelTrainer") as mock_trainer:
+                with patch("sagemaker.mlops.workflow._utils.TrainingStep.__init__") as mock_super:
+                    mock_retrieve.return_value = "sklearn-image:latest"
+                    mock_trainer_instance = Mock()
+                    mock_trainer_instance.train = Mock(return_value=Mock())
+                    mock_trainer.return_value = mock_trainer_instance
+                    mock_super.return_value = None
+
+                    _RepackModelStep(
+                        name="repack-step",
+                        sagemaker_session=mock_session,
+                        role="arn:aws:iam::123456789012:role/SageMakerRole",
+                        model_data="s3://bucket/model.tar.gz",
+                        entry_point=temp_entry_point,
+                        volume_kms_key_id="arn:aws:kms:us-west-2:111122223333:key/abc",
+                    )
+
+                    mock_trainer.assert_called_once()
+                    call_kwargs = mock_trainer.call_args[1]
+                    assert (
+                        call_kwargs["compute"].volume_kms_key_id
+                        == "arn:aws:kms:us-west-2:111122223333:key/abc"
+                    )
+                    # Consumed by Compute, not forwarded to ModelTrainer as a stray kwarg.
+                    assert "volume_kms_key_id" not in call_kwargs
+
+    def test_init_without_volume_kms_key_id_leaves_compute_unset(
+        self, mock_session, temp_entry_point
+    ):
+        """Omitting volume_kms_key_id must leave Compute.volume_kms_key_id unset so the
+        repack job's ResourceConfig carries no VolumeKmsKeyId (legacy behavior)."""
+        with patch("sagemaker.mlops.workflow._utils.image_uris.retrieve") as mock_retrieve:
+            with patch("sagemaker.train.ModelTrainer") as mock_trainer:
+                with patch("sagemaker.mlops.workflow._utils.TrainingStep.__init__") as mock_super:
+                    mock_retrieve.return_value = "sklearn-image:latest"
+                    mock_trainer_instance = Mock()
+                    mock_trainer_instance.train = Mock(return_value=Mock())
+                    mock_trainer.return_value = mock_trainer_instance
+                    mock_super.return_value = None
+
+                    _RepackModelStep(
+                        name="repack-step",
+                        sagemaker_session=mock_session,
+                        role="arn:aws:iam::123456789012:role/SageMakerRole",
+                        model_data="s3://bucket/model.tar.gz",
+                        entry_point=temp_entry_point,
+                    )
+
+                    mock_trainer.assert_called_once()
+                    compute = mock_trainer.call_args[1]["compute"]
+                    assert compute.volume_kms_key_id is None
+
     def test_init_with_depends_on(self, mock_session, temp_entry_point):
         """Test _RepackModelStep initialization with depends_on."""
         with patch("sagemaker.mlops.workflow._utils.image_uris.retrieve") as mock_retrieve: