fix: address review comments (iteration #2)

Author: aviruthen

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

@@ -12,7 +12,7 @@
 # language governing permissions and limitations under the License.
 """Placeholder docstring"""
 
-from __future__ import absolute_import
+from __future__ import absolute_import, annotations
 
 import logging
 from enum import Enum
@@ -445,17 +445,18 @@ def _prepare_auto_parameters(self, static_hyperparameters, hyperparameters_to_ke
     @staticmethod
     def _get_model_trainer_environment(
         model_trainer: "ModelTrainer",
-    ) -> Optional[Dict[str, str]]:
+    ) -> dict[str, str] | None:
         """Extract environment variables from a ModelTrainer instance.
 
-        Returns the environment dict if it is non-empty, otherwise None.
+        Returns a copy of the environment dict if it is non-empty,
+        otherwise None.
 
         Args:
             model_trainer (ModelTrainer): ModelTrainer instance.
 
         Returns:
-            Optional[Dict[str, str]]: Environment variables dict,
-                or None if empty/not set.
+            dict[str, str] | None: A copy of the environment variables
+                dict, or None if empty/not set.
         """
         env = model_trainer.environment
         if env:

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

@@ -601,6 +601,11 @@ def test_build_training_job_definition_includes_environment_variables(self):
             f"Environment should be {env_vars}, "
             f"got {definition.environment}"
         )
+        # Verify defensive copy: the dict on the definition
+        # should not be the same object as the original
+        assert definition.environment is not env_vars, (
+            "Environment should be a copy, not the same object"
+        )
 
     def test_build_training_job_definition_with_empty_environment(self):
         """Test that empty env is not propagated to definition."""
@@ -656,6 +661,10 @@ def test_returns_environment_when_set(self):
             mock_trainer,
         )
         assert result == env_vars
+        # Verify it's a copy, not the same object
+        assert result is not env_vars, (
+            "Should return a defensive copy"
+        )
 
     def test_returns_none_when_empty(self):
         """Test that None is returned when environment is empty."""