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

Author: aviruthen

sagemaker-train/src/sagemaker/train/model_trainer.py

@@ -410,14 +410,16 @@ def __del__(self):
                 self._temp_code_dir.cleanup()
 
     def _validate_training_image_and_algorithm_name(
-        self, training_image: Optional[str], algorithm_name: Optional[str]
+        self, training_image, algorithm_name
     ):
         """Validate that only one of 'training_image' or 'algorithm_name' is provided."""
-        if not training_image and not algorithm_name:
+        has_image = training_image is not None
+        has_algo = algorithm_name is not None
+        if not has_image and not has_algo:
             raise ValueError(
                 "Atleast one of 'training_image' or 'algorithm_name' must be provided.",
             )
-        if training_image and algorithm_name:
+        if has_image and has_algo:
             raise ValueError(
                 "Only one of 'training_image' or 'algorithm_name' must be provided.",
             )

sagemaker-train/src/sagemaker/train/utils.py

@@ -142,7 +142,7 @@ def _get_unique_name(base, max_length=63):
     return unique_name
 
 
-def _get_repo_name_from_image(image: str) -> str:
+def _get_repo_name_from_image(image) -> str:
     """Get the repository name from the image URI.
 
     Example:
@@ -152,11 +152,15 @@ def _get_repo_name_from_image(image: str) -> str:
     ```
 
     Args:
-        image (str): The image URI
+        image: The image URI (str or PipelineVariable)
 
     Returns:
-        str: The repository name
+        str: The repository name, or None if image is a PipelineVariable
     """
+    from sagemaker.core.helper.pipeline_variable import PipelineVariable
+
+    if isinstance(image, PipelineVariable):
+        return None
     return image.split("/")[-1].split(":")[0].split("@")[0]
 
 

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

@@ -33,7 +33,7 @@
     StoppingCondition,
     OutputDataConfig,
 )
-from sagemaker.train.defaults import DEFAULT_INSTANCE_TYPE
+from sagemaker.train.defaults import DEFAULT_INSTANCE_TYPE, TrainDefaults
 
 
 DEFAULT_IMAGE = "000000000000.dkr.ecr.us-west-2.amazonaws.com/dummy-image:latest"