fix: [Bug] Pipeline parameters (ParameterInteger, ParameterString) fail in ModelTrain (5504)

Author: aviruthen

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

@@ -24,6 +24,7 @@
 
 from sagemaker.core.shapes import Unassigned
 from sagemaker.core.modules import logger
+from sagemaker.core.helper.pipeline_variable import PipelineVariable
 
 
 def _is_valid_s3_uri(path: str, path_type: Literal["File", "Directory", "Any"] = "Any") -> bool:
@@ -129,9 +130,11 @@ def safe_serialize(data):
 
     This function handles the following cases:
     1. If `data` is a string, it returns the string as-is without wrapping in quotes.
-    2. If `data` is serializable (e.g., a dictionary, list, int, float), it returns
+    2. If `data` is of type `PipelineVariable`, it returns the PipelineVariable object
+       as-is for pipeline serialization.
+    3. If `data` is serializable (e.g., a dictionary, list, int, float), it returns
        the JSON-encoded string using `json.dumps()`.
-    3. If `data` cannot be serialized (e.g., a custom object), it returns the string
+    4. If `data` cannot be serialized (e.g., a custom object), it returns the string
        representation of the data using `str(data)`.
 
     Args:
@@ -142,6 +145,8 @@ def safe_serialize(data):
     """
     if isinstance(data, str):
         return data
+    elif isinstance(data, PipelineVariable):
+        return data
     try:
         return json.dumps(data)
     except TypeError:

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

@@ -0,0 +1,87 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"). You
+# may not use this file except in compliance with the License. A copy of
+# the License is located at
+#
+#     http://aws.amazon.com/apache2.0/
+#
+# or in the "license" file accompanying this file. This file is
+# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
+# ANY KIND, either express or implied. See the License for the specific
+# language governing permissions and limitations under the License.
+"""Tests for safe_serialize in sagemaker.core.modules.utils with PipelineVariable support.
+
+Verifies that safe_serialize correctly handles PipelineVariable objects
+(e.g., ParameterInteger, ParameterString) by returning them as-is rather
+than attempting str() conversion which would raise TypeError.
+
+See: https://github.com/aws/sagemaker-python-sdk/issues/5504
+"""
+from __future__ import absolute_import
+
+import pytest
+
+from sagemaker.core.modules.utils import safe_serialize
+from sagemaker.core.helper.pipeline_variable import PipelineVariable
+from sagemaker.core.workflow.parameters import ParameterInteger, ParameterString
+
+
+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)
+        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
+
+
+class TestSafeSerializeBasicTypes:
+    """Regression tests: verify basic types still work after PipelineVariable support."""
+
+    def test_safe_serialize_with_string(self):
+        """Strings should be returned as-is without JSON wrapping."""
+        assert safe_serialize("hello") == "hello"
+
+    def test_safe_serialize_with_int(self):
+        """Integers should be JSON-serialized to string."""
+        assert safe_serialize(42) == "42"
+
+    def test_safe_serialize_with_dict(self):
+        """Dicts should be JSON-serialized."""
+        result = safe_serialize({"key": "val"})
+        assert result == '{"key": "val"}'
+
+    def test_safe_serialize_with_bool(self):
+        """Booleans should be JSON-serialized."""
+        assert safe_serialize(True) == "true"
+        assert safe_serialize(False) == "false"
+
+    def test_safe_serialize_with_none(self):
+        """None should be JSON-serialized to 'null'."""
+        assert safe_serialize(None) == "null"
+
+    def test_safe_serialize_with_custom_object(self):
+        """Custom objects should fall back to str()."""
+
+        class CustomObj:
+            def __str__(self):
+                return "custom"
+
+        assert safe_serialize(CustomObj()) == "custom"

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

@@ -26,7 +26,7 @@
 
 from sagemaker.core.helper.session_helper import Session
 from sagemaker.core.helper.pipeline_variable import PipelineVariable, StrPipeVar
-from sagemaker.core.workflow.parameters import ParameterString
+from sagemaker.core.workflow.parameters import ParameterString, ParameterInteger
 from sagemaker.train.model_trainer import ModelTrainer, Mode
 from sagemaker.train.configs import (
     Compute,
@@ -176,3 +176,61 @@ def test_training_image_rejects_invalid_type(self):
                 stopping_condition=DEFAULT_STOPPING,
                 output_data_config=DEFAULT_OUTPUT,
             )
+
+
+class TestModelTrainerPipelineVariableHyperparameters:
+    """Test that PipelineVariable objects work correctly in ModelTrainer hyperparameters."""
+
+    def test_hyperparameters_with_parameter_integer(self):
+        """ParameterInteger in hyperparameters should be preserved through _create_training_job_args."""
+        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},
+        )
+        args = trainer._create_training_job_args()
+        # PipelineVariable should be preserved as-is, not stringified
+        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."""
+        algo = ParameterString(name="Algorithm", default_value="xgboost")
+        trainer = ModelTrainer(
+            training_image=DEFAULT_IMAGE,
+            role=DEFAULT_ROLE,
+            compute=DEFAULT_COMPUTE,
+            stopping_condition=DEFAULT_STOPPING,
+            output_data_config=DEFAULT_OUTPUT,
+            hyperparameters={"algorithm": algo},
+        )
+        args = trainer._create_training_job_args()
+        assert args["hyper_parameters"]["algorithm"] is algo
+
+    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)
+        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,
+                "eta": 0.1,
+                "objective": "binary:logistic",
+                "num_round": 100,
+            },
+        )
+        args = trainer._create_training_job_args()
+        hp = args["hyper_parameters"]
+        # PipelineVariable preserved as-is
+        assert hp["max_depth"] is max_depth
+        # Static values serialized to strings
+        assert hp["eta"] == "0.1"
+        assert hp["objective"] == "binary:logistic"
+        assert hp["num_round"] == "100"

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

@@ -0,0 +1,97 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"). You
+# may not use this file except in compliance with the License. A copy of
+# the License is located at
+#
+#     http://aws.amazon.com/apache2.0/
+#
+# or in the "license" file accompanying this file. This file is
+# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
+# ANY KIND, either express or implied. See the License for the specific
+# language governing permissions and limitations under the License.
+"""Tests for safe_serialize with PipelineVariable support.
+
+Verifies that safe_serialize in sagemaker.train.utils correctly handles
+PipelineVariable objects (e.g., ParameterInteger, ParameterString) by
+returning them as-is rather than attempting str() conversion which would
+raise TypeError.
+
+See: https://github.com/aws/sagemaker-python-sdk/issues/5504
+"""
+from __future__ import absolute_import
+
+import pytest
+
+from sagemaker.train.utils import safe_serialize
+from sagemaker.core.helper.pipeline_variable import PipelineVariable
+from sagemaker.core.workflow.parameters import ParameterInteger, ParameterString
+
+
+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)
+        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
+
+
+class TestSafeSerializeBasicTypes:
+    """Regression tests: verify basic types still work after PipelineVariable support."""
+
+    def test_safe_serialize_with_string(self):
+        """Strings should be returned as-is without JSON wrapping."""
+        assert safe_serialize("hello") == "hello"
+        assert safe_serialize("12345") == "12345"
+
+    def test_safe_serialize_with_int(self):
+        """Integers should be JSON-serialized to string."""
+        assert safe_serialize(42) == "42"
+
+    def test_safe_serialize_with_float(self):
+        """Floats should be JSON-serialized to string."""
+        assert safe_serialize(3.14) == "3.14"
+
+    def test_safe_serialize_with_dict(self):
+        """Dicts should be JSON-serialized."""
+        result = safe_serialize({"key": "val"})
+        assert result == '{"key": "val"}'
+
+    def test_safe_serialize_with_bool(self):
+        """Booleans should be JSON-serialized."""
+        assert safe_serialize(True) == "true"
+        assert safe_serialize(False) == "false"
+
+    def test_safe_serialize_with_none(self):
+        """None should be JSON-serialized to 'null'."""
+        assert safe_serialize(None) == "null"
+
+    def test_safe_serialize_with_list(self):
+        """Lists should be JSON-serialized."""
+        assert safe_serialize([1, 2, 3]) == "[1, 2, 3]"
+
+    def test_safe_serialize_with_custom_object(self):
+        """Custom objects should fall back to str()."""
+
+        class CustomObj:
+            def __str__(self):
+                return "custom"
+
+        assert safe_serialize(CustomObj()) == "custom"