fix: address review comments (iteration #3)

Author: aviruthen

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

@@ -24,11 +24,7 @@
 
 from sagemaker.core.shapes import Unassigned
 from sagemaker.core.modules import logger
-
-try:
-    from sagemaker.core.helper.pipeline_variable import PipelineVariable
-except ImportError:
-    PipelineVariable = None
+from sagemaker.core.helper.pipeline_variable import PipelineVariable
 
 
 def _is_valid_s3_uri(path: str, path_type: Literal["File", "Directory", "Any"] = "Any") -> bool:
@@ -145,11 +141,12 @@ def safe_serialize(data):
         data (Any): The data to serialize.
 
     Returns:
-        str: The serialized JSON-compatible string or the string representation of the input.
+        str | PipelineVariable: The serialized JSON-compatible string, the string
+            representation of the input, or the PipelineVariable object as-is.
     """
     if isinstance(data, str):
         return data
-    elif PipelineVariable is not None and isinstance(data, PipelineVariable):
+    elif isinstance(data, PipelineVariable):
         return data
     try:
         return json.dumps(data)

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

@@ -30,27 +30,16 @@
 class TestSafeSerializeWithPipelineVariables:
     """Test safe_serialize handles PipelineVariable objects correctly."""
 
-    def test_safe_serialize_with_parameter_integer(self):
-        """ParameterInteger should be returned as-is (identity preserved)."""
-        param = ParameterInteger(name="MaxDepth", default_value=5)
+    @pytest.mark.parametrize("param", [
+        ParameterInteger(name="MaxDepth", default_value=5),
+        ParameterString(name="Algorithm", default_value="xgboost"),
+    ])
+    def test_safe_serialize_returns_pipeline_variable_as_is(self, param):
+        """PipelineVariable objects should be returned as-is (identity preserved)."""
         result = safe_serialize(param)
         assert result is param
         assert isinstance(result, PipelineVariable)
 
-    def test_safe_serialize_with_parameter_string(self):
-        """ParameterString should be returned as-is (identity preserved)."""
-        param = ParameterString(name="Algorithm", default_value="xgboost")
-        result = safe_serialize(param)
-        assert result is param
-        assert isinstance(result, PipelineVariable)
-
-    def test_safe_serialize_does_not_call_str_on_pipeline_variable(self):
-        """Verify that PipelineVariable.__str__ is never invoked (would raise TypeError)."""
-        param = ParameterInteger(name="TestParam", default_value=10)
-        # This should NOT raise TypeError
-        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)

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

@@ -204,7 +204,7 @@ def test_hyperparameters_with_parameter_integer(self):
         assert args["hyper_parameters"]["max_depth"] is max_depth
 
     def test_hyperparameters_with_parameter_string(self):
-        """ParameterString in hyperparameters should be preserved through _create_training_job_args."""
+        """ParameterString in hyperparameters should be preserved."""
         algo = ParameterString(name="Algorithm", default_value="xgboost")
         trainer = ModelTrainer(
             training_image=DEFAULT_IMAGE,