fix: address review comments (iteration #2)

Author: aviruthen

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

@@ -1505,11 +1505,15 @@ def _build_training_job_definition(self, inputs):
                 )
 
         # Get environment variables from model_trainer.
-        # environment is a defined attribute on ModelTrainer (dict | None).
-        # We pass it through as-is; even an empty dict is valid for the API.
+        # environment is a defined attribute on ModelTrainer (typed as dict | None).
+        # We access it directly (consistent with how role, compute, etc. are accessed).
+        # We pass it through as-is when it's a dict — even an empty dict is valid for the API.
+        # When it's None or not a dict, we omit it from the constructor so the Pydantic
+        # model keeps its default (Unassigned), which is then excluded during serialization.
         env = model_trainer.environment
 
-        definition = HyperParameterTrainingJobDefinition(
+        # Build base kwargs for the definition
+        definition_kwargs = dict(
             algorithm_specification=algorithm_spec,
             role_arn=model_trainer.role,
             input_data_config=input_data_config if input_data_config else None,
@@ -1518,9 +1522,16 @@ def _build_training_job_definition(self, inputs):
             stopping_condition=stopping_condition,
             static_hyper_parameters=getattr(self, "static_hyperparameters", None) or {},
             enable_managed_spot_training=model_trainer.compute.enable_managed_spot_training,
-            environment=env,
         )
 
+        # Only include environment when it's a dict (including empty dict).
+        # This avoids Pydantic validation errors for non-dict values and keeps
+        # the field as Unassigned (excluded from serialization) when not set.
+        if isinstance(env, dict):
+            definition_kwargs["environment"] = env
+
+        definition = HyperParameterTrainingJobDefinition(**definition_kwargs)
+
         # Pass through VPC config from model_trainer
         networking = getattr(model_trainer, "networking", None)
         if networking and hasattr(networking, "_to_vpc_config"):

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

@@ -624,7 +624,13 @@ def test_build_training_job_definition_includes_environment_variables(self):
         }, "Environment variables should match those set on ModelTrainer"
 
     def test_build_training_job_definition_with_none_environment(self):
-        """Test that _build_training_job_definition handles None environment gracefully."""
+        """Test that _build_training_job_definition handles None environment gracefully.
+
+        When environment is None, it should not be passed to the Pydantic constructor,
+        so the field stays as Unassigned (excluded from serialization).
+        """
+        from sagemaker.core.utils.utils import Unassigned
+
         mock_trainer = _create_mock_model_trainer()
         mock_trainer.environment = None
 
@@ -636,7 +642,9 @@ def test_build_training_job_definition_with_none_environment(self):
 
         definition = tuner._build_training_job_definition(None)
 
-        assert definition.environment is None, "Environment should be None when not set"
+        assert isinstance(definition.environment, Unassigned), (
+            "Environment should be Unassigned when model_trainer.environment is None"
+        )
 
     def test_build_training_job_definition_with_empty_environment(self):
         """Test that _build_training_job_definition passes through empty environment.

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

@@ -405,8 +405,31 @@ def test_passes_environment_variables(self):
         definition = tuner._build_training_job_definition(inputs=None)
         assert definition.environment == {"MY_VAR": "value", "OTHER": "123"}
 
+    def test_passes_empty_environment(self):
+        """Should pass through empty dict environment as-is.
+
+        An empty dict is valid for the SageMaker API, so we pass it through
+        rather than silently converting it to None/Unassigned.
+        """
+        trainer = _mock_model_trainer(environment={})
+
+        tuner = HyperparameterTuner(
+            model_trainer=trainer,
+            objective_metric_name="accuracy",
+            hyperparameter_ranges=_hp_ranges(),
+        )
+
+        definition = tuner._build_training_job_definition(inputs=None)
+        assert definition.environment == {}, (
+            "Empty dict environment should be passed through as-is"
+        )
+
     def test_skips_environment_when_none(self):
-        """Should not set environment when model_trainer.environment is None."""
+        """Should not set environment when model_trainer.environment is None.
+
+        When environment is None, it is not passed to the Pydantic constructor,
+        so the field stays as Unassigned (excluded from serialization).
+        """
         trainer = _mock_model_trainer(environment=None)
 
         tuner = HyperparameterTuner(
@@ -416,10 +439,16 @@ def test_skips_environment_when_none(self):
         )
 
         definition = tuner._build_training_job_definition(inputs=None)
-        assert _is_unassigned(definition.environment)
+        assert _is_unassigned(definition.environment), (
+            "Environment should be Unassigned when model_trainer.environment is None"
+        )
 
     def test_skips_environment_when_not_dict(self):
-        """Should not set environment when it's not a dict (e.g. MagicMock)."""
+        """Should not set environment when it's not a dict (e.g. MagicMock).
+
+        Non-dict values are not passed to the Pydantic constructor to avoid
+        validation errors. The field stays as Unassigned.
+        """
         trainer = _mock_model_trainer(environment=MagicMock())
 
         tuner = HyperparameterTuner(
@@ -429,7 +458,9 @@ def test_skips_environment_when_not_dict(self):
         )
 
         definition = tuner._build_training_job_definition(inputs=None)
-        assert _is_unassigned(definition.environment)
+        assert _is_unassigned(definition.environment), (
+            "Environment should be Unassigned when model_trainer.environment is not a dict"
+        )
 
     def test_passes_vpc_config(self):
         """Should set definition.vpc_config from model_trainer.networking._to_vpc_config()."""