update output_data_config reference in HyperparameterTuner and add test for its inclusion

Author: MarvinPitt

sagemaker-train/src/sagemaker/train/tuner.py

@@ -253,7 +253,6 @@ def __init__(
         self.instance_configs_dict = None
         self.instance_configs = None
         self.autotune = autotune
-        self.output_data_config = model_trainer.output_data_config
 
     def override_resource_config(
         self,
@@ -1422,7 +1421,6 @@ def _build_training_job_definition(self, inputs):
                 if not any(c.channel_name == channel.channel_name for c in input_data_config):
                     input_data_config.append(channel)
 
-
         # Build resource config
         resource_config = ResourceConfig(
             instance_type=(
@@ -1448,7 +1446,7 @@ def _build_training_job_definition(self, inputs):
             algorithm_specification=algorithm_spec,
             role_arn=model_trainer.role,
             input_data_config=input_data_config if input_data_config else None,
-            output_data_config=self.output_data_config,
+            output_data_config=model_trainer.output_data_config,
             resource_config=resource_config,
             stopping_condition=stopping_condition,
             static_hyper_parameters=self.static_hyperparameters or {},

sagemaker-train/tests/unit/train/test_tuner.py

@@ -31,6 +31,7 @@
     Channel,
     DataSource,
     S3DataSource,
+    OutputDataConfig
 )
 
 
@@ -52,9 +53,7 @@ def _create_mock_model_trainer(with_internal_channels=False):
     trainer.training_image = "test-image:latest"
     trainer.training_input_mode = "File"
     trainer.role = "arn:aws:iam::123456789012:role/SageMakerRole"
-    trainer.output_data_config = MagicMock()
-    trainer.output_data_config.kms_key_id = None
-    trainer.output_data_config.s3_output_path = "s3://bucket/output"
+    trainer.output_data_config = OutputDataConfig(kms_key_id="arn:aws:kms:us-west-2:123456789012:key/abc123", s3_output_path="s3://bucket/output")
     trainer.compute = MagicMock()
     trainer.compute.instance_type = "ml.m5.xlarge"
     trainer.compute.instance_count = 1
@@ -575,3 +574,16 @@ def test_build_training_job_definition_includes_internal_channels(self):
         assert "train" in channel_names, "User 'train' channel should be included"
         assert "validation" in channel_names, "User 'validation' channel should be included"
         assert len(channel_names) == 4, "Should have exactly 4 channels"
+
+    def test_build_training_job_definition_includes_output_data_config(self):
+        """Test that _build_training_job_definition includes ModelTrainer's output data config."""
+        mock_trainer = _create_mock_model_trainer()
+        tuner = HyperparameterTuner(
+            model_trainer=mock_trainer,
+            objective_metric_name="accuracy",
+            hyperparameter_ranges=_create_single_hp_range(),
+        )
+
+        definition = tuner._build_training_job_definition([])
+
+        assert definition.output_data_config == mock_trainer.output_data_config