From d604bd0ea120df936aa37c6be088755013a1535e Mon Sep 17 00:00:00 2001 From: Mohamed Zeidan Date: Mon, 26 Oct 2026 22:54:55 -0700 Subject: [PATCH] fix: set VolumeKmsKeyId on ModelStep runtime repack job 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 --- .../src/sagemaker/mlops/workflow/_utils.py | 6 +- .../sagemaker/mlops/workflow/model_step.py | 9 ++ .../tests/unit/workflow/test_model_step.py | 131 ++++++++++++++++++ .../tests/unit/workflow/test_utils.py | 57 ++++++++ 4 files changed, 202 insertions(+), 1 deletion(-) diff --git a/sagemaker-mlops/src/sagemaker/mlops/workflow/_utils.py b/sagemaker-mlops/src/sagemaker/mlops/workflow/_utils.py index 7a46ecd0ff..eda891813f 100644 --- a/sagemaker-mlops/src/sagemaker/mlops/workflow/_utils.py +++ b/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 diff --git a/sagemaker-mlops/src/sagemaker/mlops/workflow/model_step.py b/sagemaker-mlops/src/sagemaker/mlops/workflow/model_step.py index aff3662b93..844e439b7a 100644 --- a/sagemaker-mlops/src/sagemaker/mlops/workflow/model_step.py +++ b/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=( diff --git a/sagemaker-mlops/tests/unit/workflow/test_model_step.py b/sagemaker-mlops/tests/unit/workflow/test_model_step.py index 7e4c1ce4e8..47d878babe 100644 --- a/sagemaker-mlops/tests/unit/workflow/test_model_step.py +++ b/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 diff --git a/sagemaker-mlops/tests/unit/workflow/test_utils.py b/sagemaker-mlops/tests/unit/workflow/test_utils.py index c5927bd2a3..e4172bfd89 100644 --- a/sagemaker-mlops/tests/unit/workflow/test_utils.py +++ b/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: