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