fix: preserve volume KMS key in HyperparameterTuner training job definition (#6522)

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

Co-authored-by: Mufaddal Rohawala <mufi@amazon.com>

Author: mufaddal-rohawala

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:

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)."""
 

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