fix: Add PipelineVariable support to ModelTrainer fields (fixes #5524) (#5608)

amitmodi · Amit Modi · web-flow · commit 0906526bc0f1 · 2026-03-09T14:02:23.000-07:00
* fix: Add PipelineVariable support to ModelTrainer fields (fixes #5524) Extend StrPipeVar type to ModelTrainer's direct fields: - training_image: Optional[str] -> Optional[StrPipeVar] - algorithm_name: Optional[str] -> Optional[StrPipeVar] - training_input_mode: Optional[str] -> Optional[StrPipeVar] - environment: Dict[str, str] -> Dict[str, StrPipeVar] This follows the existing V3 pattern already used by SourceCode, OutputDataConfig, and Compute (for instance_type). The StrPipeVar type alias and PipelineVariable.__get_pydantic_core_schema__() already exist in the codebase. This unblocks V2->V3 migration for SageMaker Pipelines users who need to pass ParameterString to ModelTrainer fields. Fixes #5524 * test: Add unit tests for PipelineVariable support + fix PipelineVariable-safe logging - Add test_model_trainer_pipeline_variable.py with 9 tests: - 4 PipelineVariable acceptance tests (training_image, algorithm_name, training_input_mode, environment) - 4 regression tests (real string values still work) - 1 invalid type rejection test - Fix PipelineVariable-safe logging in model_post_init (avoid __str__ on PipelineVariable which raises TypeError) All 57 tests pass (48 existing + 9 new, 0 regressions). --------- Co-authored-by: Amit Modi <modiamit@amazon.com>
diff --git a/sagemaker-train/src/sagemaker/train/model_trainer.py b/sagemaker-train/src/sagemaker/train/model_trainer.py
@@ -116,6 +116,7 @@
 from sagemaker.core.jumpstart.utils import get_eula_url
 from sagemaker.train.defaults import TrainDefaults, JumpStartTrainDefaults
 from sagemaker.core.workflow.pipeline_context import PipelineSession, runnable_by_pipeline
+from sagemaker.core.helper.pipeline_variable import StrPipeVar
 
 from sagemaker.train.local.local_container import _LocalContainer
 
@@ -235,14 +236,14 @@ class ModelTrainer(BaseModel):
     compute: Optional[Compute] = None
     networking: Optional[Networking] = None
     stopping_condition: Optional[StoppingCondition] = None
-    training_image: Optional[str] = None
+    training_image: Optional[StrPipeVar] = None
     training_image_config: Optional[TrainingImageConfig] = None
-    algorithm_name: Optional[str] = None
+    algorithm_name: Optional[StrPipeVar] = None
     output_data_config: Optional[shapes.OutputDataConfig] = None
     input_data_config: Optional[List[Union[Channel, InputData]]] = None
     checkpoint_config: Optional[shapes.CheckpointConfig] = None
-    training_input_mode: Optional[str] = "File"
-    environment: Optional[Dict[str, str]] = {}
+    training_input_mode: Optional[StrPipeVar] = "File"
+    environment: Optional[Dict[str, StrPipeVar]] = {}
     hyperparameters: Optional[Union[Dict[str, Any], str]] = {}
     tags: Optional[List[Tag]] = None
     local_container_root: Optional[str] = os.getcwd()
@@ -545,7 +546,11 @@ def model_post_init(self, __context: Any):
             )
 
         if self.training_image:
-            logger.info(f"Training image URI: {self.training_image}")
+            from sagemaker.core.helper.pipeline_variable import PipelineVariable
+            if isinstance(self.training_image, PipelineVariable):
+                logger.info("Training image URI: (PipelineVariable - resolved at pipeline execution)")
+            else:
+                logger.info(f"Training image URI: {self.training_image}")
     
 
     def _create_training_job_args(
diff --git a/sagemaker-train/tests/unit/train/test_model_trainer_pipeline_variable.py b/sagemaker-train/tests/unit/train/test_model_trainer_pipeline_variable.py
@@ -0,0 +1,178 @@
+# 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 PipelineVariable support in ModelTrainer (GH#5524).
+
+Verifies that ModelTrainer fields accept PipelineVariable objects
+(e.g., ParameterString) in addition to their concrete types, following
+the existing V3 pattern established by SourceCode and OutputDataConfig.
+
+See: https://github.com/aws/sagemaker-python-sdk/issues/5524
+"""
+from __future__ import absolute_import
+
+import pytest
+from pydantic import ValidationError
+from unittest.mock import patch, MagicMock
+
+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.train.model_trainer import ModelTrainer, Mode
+from sagemaker.train.configs import (
+    Compute,
+    StoppingCondition,
+    OutputDataConfig,
+)
+from sagemaker.train.defaults import DEFAULT_INSTANCE_TYPE
+
+
+DEFAULT_IMAGE = "000000000000.dkr.ecr.us-west-2.amazonaws.com/dummy-image:latest"
+DEFAULT_BUCKET = "sagemaker-us-west-2-000000000000"
+DEFAULT_ROLE = "arn:aws:iam::000000000000:role/test-role"
+DEFAULT_BUCKET_PREFIX = "sample-prefix"
+DEFAULT_REGION = "us-west-2"
+DEFAULT_COMPUTE = Compute(instance_type=DEFAULT_INSTANCE_TYPE, instance_count=1)
+DEFAULT_STOPPING = StoppingCondition(max_runtime_in_seconds=3600)
+DEFAULT_OUTPUT = OutputDataConfig(
+    s3_output_path=f"s3://{DEFAULT_BUCKET}/{DEFAULT_BUCKET_PREFIX}/test-job",
+)
+
+
+@pytest.fixture(scope="module", autouse=True)
+def modules_session():
+    with patch("sagemaker.train.Session", spec=Session) as session_mock:
+        session_instance = session_mock.return_value
+        session_instance.default_bucket.return_value = DEFAULT_BUCKET
+        session_instance.get_caller_identity_arn.return_value = DEFAULT_ROLE
+        session_instance.default_bucket_prefix = DEFAULT_BUCKET_PREFIX
+        session_instance.boto_session = MagicMock(spec="boto3.session.Session")
+        session_instance.boto_region_name = DEFAULT_REGION
+        yield session_instance
+
+
+class TestModelTrainerPipelineVariableAcceptance:
+    """Test that ModelTrainer fields accept PipelineVariable objects."""
+
+    def test_training_image_accepts_parameter_string(self):
+        """ModelTrainer.training_image should accept ParameterString (GH#5524)."""
+        param = ParameterString(name="TrainingImage", default_value=DEFAULT_IMAGE)
+        trainer = ModelTrainer(
+            training_image=param,
+            base_job_name="pipeline-test-job",  # Required: PipelineVariable can't generate job name
+            role=DEFAULT_ROLE,
+            compute=DEFAULT_COMPUTE,
+            stopping_condition=DEFAULT_STOPPING,
+            output_data_config=DEFAULT_OUTPUT,
+        )
+        assert trainer.training_image is param
+
+    def test_algorithm_name_accepts_parameter_string(self):
+        """ModelTrainer.algorithm_name should accept ParameterString."""
+        param = ParameterString(name="AlgorithmName", default_value="my-algo-arn")
+        trainer = ModelTrainer(
+            algorithm_name=param,
+            base_job_name="pipeline-test-job",  # Required: PipelineVariable can't generate job name
+            role=DEFAULT_ROLE,
+            compute=DEFAULT_COMPUTE,
+            stopping_condition=DEFAULT_STOPPING,
+            output_data_config=DEFAULT_OUTPUT,
+        )
+        assert trainer.algorithm_name is param
+
+    def test_training_input_mode_accepts_parameter_string(self):
+        """ModelTrainer.training_input_mode should accept ParameterString."""
+        param = ParameterString(name="InputMode", default_value="File")
+        trainer = ModelTrainer(
+            training_image=DEFAULT_IMAGE,
+            training_input_mode=param,
+            role=DEFAULT_ROLE,
+            compute=DEFAULT_COMPUTE,
+            stopping_condition=DEFAULT_STOPPING,
+            output_data_config=DEFAULT_OUTPUT,
+        )
+        assert trainer.training_input_mode is param
+
+    def test_environment_values_accept_parameter_string(self):
+        """ModelTrainer.environment dict values should accept ParameterString."""
+        param = ParameterString(name="DatasetVersion", default_value="v1")
+        trainer = ModelTrainer(
+            training_image=DEFAULT_IMAGE,
+            environment={"DATASET_VERSION": param, "STATIC_VAR": "hello"},
+            role=DEFAULT_ROLE,
+            compute=DEFAULT_COMPUTE,
+            stopping_condition=DEFAULT_STOPPING,
+            output_data_config=DEFAULT_OUTPUT,
+        )
+        assert trainer.environment["DATASET_VERSION"] is param
+        assert trainer.environment["STATIC_VAR"] == "hello"
+
+
+class TestModelTrainerRealValuesStillWork:
+    """Regression tests: verify that passing real values still works after the change."""
+
+    def test_training_image_accepts_real_string(self):
+        """ModelTrainer.training_image should still accept a plain string."""
+        trainer = ModelTrainer(
+            training_image=DEFAULT_IMAGE,
+            role=DEFAULT_ROLE,
+            compute=DEFAULT_COMPUTE,
+            stopping_condition=DEFAULT_STOPPING,
+            output_data_config=DEFAULT_OUTPUT,
+        )
+        assert trainer.training_image == DEFAULT_IMAGE
+
+    def test_algorithm_name_accepts_real_string(self):
+        """ModelTrainer.algorithm_name should still accept a plain string."""
+        trainer = ModelTrainer(
+            algorithm_name="arn:aws:sagemaker:us-west-2:000000000000:algorithm/my-algo",
+            role=DEFAULT_ROLE,
+            compute=DEFAULT_COMPUTE,
+            stopping_condition=DEFAULT_STOPPING,
+            output_data_config=DEFAULT_OUTPUT,
+        )
+        assert trainer.algorithm_name == "arn:aws:sagemaker:us-west-2:000000000000:algorithm/my-algo"
+
+    def test_training_input_mode_accepts_real_string(self):
+        """ModelTrainer.training_input_mode should still accept a plain string."""
+        trainer = ModelTrainer(
+            training_image=DEFAULT_IMAGE,
+            training_input_mode="Pipe",
+            role=DEFAULT_ROLE,
+            compute=DEFAULT_COMPUTE,
+            stopping_condition=DEFAULT_STOPPING,
+            output_data_config=DEFAULT_OUTPUT,
+        )
+        assert trainer.training_input_mode == "Pipe"
+
+    def test_environment_accepts_real_string_values(self):
+        """ModelTrainer.environment should still accept plain string values."""
+        trainer = ModelTrainer(
+            training_image=DEFAULT_IMAGE,
+            environment={"KEY1": "value1", "KEY2": "value2"},
+            role=DEFAULT_ROLE,
+            compute=DEFAULT_COMPUTE,
+            stopping_condition=DEFAULT_STOPPING,
+            output_data_config=DEFAULT_OUTPUT,
+        )
+        assert trainer.environment == {"KEY1": "value1", "KEY2": "value2"}
+
+    def test_training_image_rejects_invalid_type(self):
+        """ModelTrainer.training_image should still reject invalid types (e.g., int)."""
+        with pytest.raises(ValidationError):
+            ModelTrainer(
+                training_image=12345,
+                role=DEFAULT_ROLE,
+                compute=DEFAULT_COMPUTE,
+                stopping_condition=DEFAULT_STOPPING,
+                output_data_config=DEFAULT_OUTPUT,
+            )
