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

Author: aviruthen

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

@@ -410,14 +410,18 @@ 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:
+        from sagemaker.core.helper.pipeline_variable import PipelineVariable as _PV
+        # PipelineVariables are truthy for validation purposes
+        has_image = isinstance(training_image, _PV) or bool(training_image)
+        has_algo = isinstance(algorithm_name, _PV) or bool(algorithm_name)
+        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,13 @@ 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
     """
+    if isinstance(image, PipelineVariable):
+        return "pipeline-variable-image"
     return image.split("/")[-1].split(":")[0].split("@")[0]