Refactor output data configuration handling in tuner

Author: MarvinBeGood

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

@@ -253,6 +253,7 @@ 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,
@@ -1331,7 +1332,6 @@ def _build_training_job_definition(self, inputs):
         from sagemaker.core.shapes import (
             HyperParameterTrainingJobDefinition,
             HyperParameterAlgorithmSpecification,
-            OutputDataConfig,
             ResourceConfig,
             StoppingCondition,
             Channel,
@@ -1422,14 +1422,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 output data config
-        output_config = OutputDataConfig(
-            s3_output_path=(
-                model_trainer.output_data_config.s3_output_path
-                if model_trainer.output_data_config
-                else None
-            )
-        )
 
         # Build resource config
         resource_config = ResourceConfig(
@@ -1456,7 +1448,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=output_config,
+            output_data_config=self.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

@@ -53,6 +53,7 @@ def _create_mock_model_trainer(with_internal_channels=False):
     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.compute = MagicMock()
     trainer.compute.instance_type = "ml.m5.xlarge"