fix: ModelTrainer and HyperparameterTuner missing environment variables (5613)

Author: aviruthen

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

@@ -1504,6 +1504,11 @@ def _build_training_job_definition(self, inputs):
                     model_trainer.stopping_condition.max_wait_time_in_seconds
                 )
 
+        # Get environment variables from model_trainer
+        env = getattr(model_trainer, "environment", None)
+        if not env or not isinstance(env, dict):
+            env = None
+
         definition = HyperParameterTrainingJobDefinition(
             algorithm_specification=algorithm_spec,
             role_arn=model_trainer.role,
@@ -1513,13 +1518,9 @@ 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,
         )
 
-        # Pass through environment variables from model_trainer
-        env = getattr(model_trainer, "environment", None)
-        if env and isinstance(env, dict):
-            definition.environment = env
-
         # 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

@@ -596,3 +596,61 @@ def test_build_training_job_definition_includes_spot_params(self):
         assert isinstance(
             definition.stopping_condition.max_wait_time_in_seconds, int
         ), "Max wait time should be set"
+
+    def test_build_training_job_definition_includes_environment_variables(self):
+        """Test that _build_training_job_definition includes environment variables.
+
+        This test verifies the fix for GitHub issue #5613 where tuning jobs were
+        missing environment variables that were set on the ModelTrainer.
+        """
+        mock_trainer = _create_mock_model_trainer()
+        mock_trainer.environment = {
+            "FOO": "bar",
+            "RANDOM_STATE": "42",
+        }
+
+        tuner = HyperparameterTuner(
+            model_trainer=mock_trainer,
+            objective_metric_name="accuracy",
+            hyperparameter_ranges=_create_single_hp_range(),
+        )
+
+        definition = tuner._build_training_job_definition(None)
+
+        assert definition.environment is not None, "Environment should not be None"
+        assert definition.environment == {
+            "FOO": "bar",
+            "RANDOM_STATE": "42",
+        }, "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."""
+        mock_trainer = _create_mock_model_trainer()
+        mock_trainer.environment = None
+
+        tuner = HyperparameterTuner(
+            model_trainer=mock_trainer,
+            objective_metric_name="accuracy",
+            hyperparameter_ranges=_create_single_hp_range(),
+        )
+
+        definition = tuner._build_training_job_definition(None)
+
+        assert definition.environment is None, "Environment should be None when not set"
+
+    def test_build_training_job_definition_with_empty_environment(self):
+        """Test that _build_training_job_definition handles empty environment gracefully."""
+        mock_trainer = _create_mock_model_trainer()
+        mock_trainer.environment = {}
+
+        tuner = HyperparameterTuner(
+            model_trainer=mock_trainer,
+            objective_metric_name="accuracy",
+            hyperparameter_ranges=_create_single_hp_range(),
+        )
+
+        definition = tuner._build_training_job_definition(None)
+
+        assert definition.environment is None, (
+            "Environment should be None when empty dict is provided"
+        )