From 00a2ae9fa1963a0fbf00955991928a6d87ae3ab0 Mon Sep 17 00:00:00 2001 From: Mufaddal Rohawala Date: Thu, 22 Oct 2026 11:46:01 -0700 Subject: [PATCH] fix: preserve volume KMS key in HyperparameterTuner training job definition HyperparameterTuner._build_training_job_definition() built the tuning job's ResourceConfig from ModelTrainer.compute but never copied volume_kms_key_id. HPO child training jobs therefore dropped the configured customer-managed KMS key for their compute volumes, and the request succeeded without error, so the loss was silent. Copy volume_kms_key_id from the trainer's compute when it is assigned, and leave the field Unassigned otherwise so it stays out of the serialized CreateHyperParameterTuningJob request. Explicit keys, plain strings, and PipelineVariable values are all passed through unchanged. Add unit tests covering explicit keys, plain-string keys, pipeline variables, and omitted keys. Closes #6521 --- sagemaker-train/src/sagemaker/train/tuner.py | 11 ++++ .../tests/unit/train/test_tuner.py | 61 +++++++++++++++++++ .../unit/train/test_tuner_driver_channels.py | 1 + 3 files changed, 73 insertions(+) diff --git a/sagemaker-train/src/sagemaker/train/tuner.py b/sagemaker-train/src/sagemaker/train/tuner.py index 0931ae1816..54eb65feea 100644 --- a/sagemaker-train/src/sagemaker/train/tuner.py +++ b/sagemaker-train/src/sagemaker/train/tuner.py @@ -1413,6 +1413,7 @@ def _build_training_job_definition(self, inputs): DataSource, S3DataSource, ) + from sagemaker.core.utils.utils import Unassigned model_trainer = self.model_trainer @@ -1520,6 +1521,16 @@ def _build_training_job_definition(self, inputs): ), ) + # Preserve the volume KMS key configured on the ModelTrainer's compute so that + # HPO child training jobs encrypt their storage volumes with the customer-managed + # key. Omit the field entirely when no key is configured; leaving it Unassigned + # keeps it out of the serialized request (issue #6521). A PipelineVariable or a + # plain string is passed through unchanged. + if model_trainer.compute: + volume_kms_key_id = model_trainer.compute.volume_kms_key_id + if volume_kms_key_id is not None and not isinstance(volume_kms_key_id, Unassigned): + resource_config.volume_kms_key_id = volume_kms_key_id + # Build stopping condition stopping_condition = StoppingCondition() if model_trainer.stopping_condition: diff --git a/sagemaker-train/tests/unit/train/test_tuner.py b/sagemaker-train/tests/unit/train/test_tuner.py index ff4800be51..f502fbbb26 100644 --- a/sagemaker-train/tests/unit/train/test_tuner.py +++ b/sagemaker-train/tests/unit/train/test_tuner.py @@ -70,6 +70,7 @@ def _create_mock_model_trainer(with_internal_channels=False, with_spot_training= trainer.compute.instance_type = "ml.m5.xlarge" trainer.compute.instance_count = 1 trainer.compute.volume_size_in_gb = 30 + trainer.compute.volume_kms_key_id = None trainer.stopping_condition = MagicMock() trainer.stopping_condition.max_runtime_in_seconds = 3600 trainer.input_data_config = None @@ -865,6 +866,66 @@ def test_omitted_values_stay_unassigned(self): assert not isinstance(definition.retry_strategy, RetryStrategy) +class TestBuildTrainingJobDefinitionVolumeKmsKey: + """Regression for #6521. + + ``_build_training_job_definition`` dropped ``volume_kms_key_id`` from + ``ModelTrainer.compute``, so HPO child training jobs silently failed to + encrypt their storage volumes with the configured customer-managed KMS key. + """ + + @staticmethod + def _tuner(trainer): + return HyperparameterTuner( + model_trainer=trainer, + objective_metric_name="accuracy", + hyperparameter_ranges=_create_single_hp_range(), + ) + + def test_explicit_key_is_propagated(self): + trainer = _create_mock_model_trainer() + trainer.compute.volume_kms_key_id = ( + "arn:aws:kms:us-east-2:123456789012:key/00000000-0000-0000-0000-000000000000" + ) + + definition = self._tuner(trainer)._build_training_job_definition(None) + + assert definition.resource_config.volume_kms_key_id == ( + "arn:aws:kms:us-east-2:123456789012:key/00000000-0000-0000-0000-000000000000" + ) + + def test_plain_string_key_is_propagated(self): + trainer = _create_mock_model_trainer() + trainer.compute.volume_kms_key_id = "1234abcd-12ab-34cd-56ef-1234567890ab" + + definition = self._tuner(trainer)._build_training_job_definition(None) + + assert ( + definition.resource_config.volume_kms_key_id == "1234abcd-12ab-34cd-56ef-1234567890ab" + ) + + def test_pipeline_variable_key_is_propagated(self): + trainer = _create_mock_model_trainer() + kms_param = ParameterString(name="VolumeKmsKeyId") + trainer.compute.volume_kms_key_id = kms_param + + definition = self._tuner(trainer)._build_training_job_definition(None) + + assert definition.resource_config.volume_kms_key_id is kms_param + + def test_omitted_key_stays_unassigned(self): + from sagemaker.core.utils.utils import Unassigned + + trainer = _create_mock_model_trainer() + trainer.compute.volume_kms_key_id = None + + definition = self._tuner(trainer)._build_training_job_definition(None) + + # Leaving the field Unassigned keeps VolumeKmsKeyId out of the + # serialized CreateHyperParameterTuningJob request. + assert isinstance(definition.resource_config.volume_kms_key_id, Unassigned) + + class TestHyperparameterTunerTagInheritance: """Test tag inheritance from ModelTrainer (issue #6472).""" diff --git a/sagemaker-train/tests/unit/train/test_tuner_driver_channels.py b/sagemaker-train/tests/unit/train/test_tuner_driver_channels.py index d939611bc9..97b747dfd5 100644 --- a/sagemaker-train/tests/unit/train/test_tuner_driver_channels.py +++ b/sagemaker-train/tests/unit/train/test_tuner_driver_channels.py @@ -79,6 +79,7 @@ def _mock_model_trainer(**overrides): trainer.compute.instance_type = "ml.m5.xlarge" trainer.compute.instance_count = 1 trainer.compute.volume_size_in_gb = 30 + trainer.compute.volume_kms_key_id = None trainer.stopping_condition = MagicMock() trainer.stopping_condition.max_runtime_in_seconds = 3600 trainer.input_data_config = None