fix: address review comments (iteration #1)

Author: aviruthen

sagemaker-core/src/sagemaker/core/modules/utils.py

@@ -24,7 +24,11 @@
 
 from sagemaker.core.shapes import Unassigned
 from sagemaker.core.modules import logger
-from sagemaker.core.helper.pipeline_variable import PipelineVariable
+
+try:
+    from sagemaker.core.helper.pipeline_variable import PipelineVariable
+except ImportError:
+    PipelineVariable = None
 
 
 def _is_valid_s3_uri(path: str, path_type: Literal["File", "Directory", "Any"] = "Any") -> bool:
@@ -145,7 +149,7 @@ def safe_serialize(data):
     """
     if isinstance(data, str):
         return data
-    elif isinstance(data, PipelineVariable):
+    elif PipelineVariable is not None and isinstance(data, PipelineVariable):
         return data
     try:
         return json.dumps(data)

sagemaker-core/tests/unit/core/modules/test_utils_safe_serialize.py

@@ -18,7 +18,7 @@
 
 See: https://github.com/aws/sagemaker-python-sdk/issues/5504
 """
-from __future__ import absolute_import
+from __future__ import annotations
 
 import pytest
 
@@ -51,6 +51,12 @@ def test_safe_serialize_does_not_call_str_on_pipeline_variable(self):
         result = safe_serialize(param)
         assert result is param
 
+    def test_pipeline_variable_str_raises_type_error(self):
+        """Confirm PipelineVariable.__str__ raises TypeError (the root cause of the bug)."""
+        param = ParameterInteger(name="TestParam", default_value=10)
+        with pytest.raises(TypeError):
+            str(param)
+
 
 class TestSafeSerializeBasicTypes:
     """Regression tests: verify basic types still work after PipelineVariable support."""

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

@@ -210,6 +210,32 @@ def test_hyperparameters_with_parameter_string(self):
         args = trainer._create_training_job_args()
         assert args["hyper_parameters"]["algorithm"] is algo
 
+    def test_hyperparameters_with_parameter_integer_does_not_raise(self):
+        """Verify ParameterInteger in hyperparameters does NOT raise TypeError.
+
+        This test documents the exact bug scenario from GH#5504: safe_serialize
+        would fall back to str(data) for PipelineVariable objects, but
+        PipelineVariable.__str__ intentionally raises TypeError.
+        """
+        max_depth = ParameterInteger(name="MaxDepth", default_value=5)
+        trainer = ModelTrainer(
+            training_image=DEFAULT_IMAGE,
+            role=DEFAULT_ROLE,
+            compute=DEFAULT_COMPUTE,
+            stopping_condition=DEFAULT_STOPPING,
+            output_data_config=DEFAULT_OUTPUT,
+            hyperparameters={"max_depth": max_depth},
+        )
+        # This call would have raised TypeError before the fix
+        try:
+            args = trainer._create_training_job_args()
+        except TypeError:
+            pytest.fail(
+                "safe_serialize raised TypeError on PipelineVariable - "
+                "this is the bug described in GH#5504"
+            )
+        assert args["hyper_parameters"]["max_depth"] is max_depth
+
     def test_hyperparameters_with_mixed_pipeline_and_static_values(self):
         """Mixed PipelineVariable and static values should both be handled correctly."""
         max_depth = ParameterInteger(name="MaxDepth", default_value=5)

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

@@ -19,7 +19,7 @@
 
 See: https://github.com/aws/sagemaker-python-sdk/issues/5504
 """
-from __future__ import absolute_import
+from __future__ import annotations
 
 import pytest
 
@@ -52,6 +52,12 @@ def test_safe_serialize_does_not_call_str_on_pipeline_variable(self):
         result = safe_serialize(param)
         assert result is param
 
+    def test_pipeline_variable_str_raises_type_error(self):
+        """Confirm PipelineVariable.__str__ raises TypeError (the root cause of the bug)."""
+        param = ParameterInteger(name="TestParam", default_value=10)
+        with pytest.raises(TypeError):
+            str(param)
+
 
 class TestSafeSerializeBasicTypes:
     """Regression tests: verify basic types still work after PipelineVariable support."""