fix: Fix processing image uri param

src/sagemaker/utils.py

@@ -33,6 +33,7 @@
 
 from sagemaker import deprecations
 from sagemaker.session_settings import SessionSettings
+from sagemaker.workflow import is_pipeline_variable, is_pipeline_parameter_string
 
 
 ECR_URI_PATTERN = r"^(\d+)(\.)dkr(\.)ecr(\.)(.+)(\.)(.*)(/)(.*:.*)$"
@@ -99,9 +100,17 @@ def base_name_from_image(image):
     Returns:
         str: Algorithm name, as extracted from the image name.
     """
-    m = re.match("^(.+/)?([^:/]+)(:[^:]+)?$", image)
-    algo_name = m.group(2) if m else image
-    return algo_name
+    if is_pipeline_variable(image):
+        if is_pipeline_parameter_string(image) and image.default_value:
+            image_str = image.default_value
+        else:
+            return "base_name"
+    else:
+        image_str = image
+
+    m = re.match("^(.+/)?([^:/]+)(:[^:]+)?$", image_str)
+    base_name = m.group(2) if m else image_str
+    return base_name
 
 
 def base_from_name(name):

src/sagemaker/workflow/__init__.py

@@ -14,6 +14,7 @@
 from __future__ import absolute_import
 
 from sagemaker.workflow.entities import Expression
+from sagemaker.workflow.parameters import ParameterString
 
 
 def is_pipeline_variable(var: object) -> bool:
@@ -29,3 +30,14 @@ def is_pipeline_variable(var: object) -> bool:
     # as well as PipelineExperimentConfigProperty and PropertyFile
     # TODO: We should deprecate the Expression and replace it with PipelineVariable
     return isinstance(var, Expression)
+
+
+def is_pipeline_parameter_string(var: object) -> bool:
+    """Check if the variable is a pipeline parameter string
+
+    Args:
+        var (object): The variable to be verified.
+    Returns:
+         bool: True if it is, False otherwise.
+    """
+    return isinstance(var, ParameterString)

tests/unit/sagemaker/workflow/test_pipeline_session.py

@@ -17,8 +17,16 @@
 from mock import Mock, PropertyMock
 
 from sagemaker import Model
-from sagemaker.workflow.parameters import ParameterString
 from sagemaker.workflow.pipeline_context import PipelineSession
+from sagemaker.workflow import is_pipeline_variable, is_pipeline_parameter_string
+from sagemaker.workflow.parameters import (
+    ParameterString,
+    ParameterInteger,
+    ParameterBoolean,
+    ParameterFloat,
+)
+from sagemaker.workflow.functions import Join, JsonGet
+from tests.unit.sagemaker.workflow.helpers import CustomStep
 
 from botocore.config import Config
 
@@ -122,3 +130,43 @@ def test_pipeline_session_context_for_model_step(pipeline_session_mock):
     assert not register_step_args.create_model_request
     assert register_step_args.create_model_package_request
     assert len(register_step_args.need_runtime_repack) == 0
+
+
+@pytest.mark.parametrize(
+    "item",
+    [
+        (ParameterString(name="my-str"), True),
+        (ParameterBoolean(name="my-bool"), True),
+        (ParameterFloat(name="my-float"), True),
+        (ParameterInteger(name="my-int"), True),
+        (Join(on="/", values=["my", "value"]), True),
+        (JsonGet(step_name="my-step", property_file="pf", json_path="path"), True),
+        (CustomStep(name="my-step").properties.OutputDataConfig.S3OutputPath, True),
+        ("my-str", False),
+        (1, False),
+        (CustomStep(name="my-ste"), False),
+    ],
+)
+def test_is_pipeline_variable(item):
+    var, assertion = item
+    assert is_pipeline_variable(var) == assertion
+
+
+@pytest.mark.parametrize(
+    "item",
+    [
+        (ParameterString(name="my-str"), True),
+        (ParameterBoolean(name="my-bool"), False),
+        (ParameterFloat(name="my-float"), False),
+        (ParameterInteger(name="my-int"), False),
+        (Join(on="/", values=["my", "value"]), False),
+        (JsonGet(step_name="my-step", property_file="pf", json_path="path"), False),
+        (CustomStep(name="my-step").properties.OutputDataConfig.S3OutputPath, False),
+        ("my-str", False),
+        (1, False),
+        (CustomStep(name="my-ste"), False),
+    ],
+)
+def test_is_pipeline_parameter_string(item):
+    var, assertion = item
+    assert is_pipeline_parameter_string(var) == assertion

tests/unit/sagemaker/workflow/test_processing_step.py

@@ -328,17 +328,27 @@ def test_processing_step_with_processor(pipeline_session, processing_input):
     }
 
 
-def test_processing_step_with_processor_and_step_args(pipeline_session, processing_input):
+@pytest.mark.parametrize(
+    "image_uri",
+    [
+        IMAGE_URI,
+        ParameterString(name="MyImage"),
+        ParameterString(name="MyImage", default_value="my-image-uri"),
+        Join(on="/", values=["docker", "my-fake-image"]),
+    ],
+)
+def test_processing_step_with_processor_and_step_args(
+    pipeline_session, processing_input, image_uri
+):
     processor = Processor(
-        image_uri=IMAGE_URI,
+        image_uri=image_uri,
         role=ROLE,
         instance_count=1,
         instance_type=INSTANCE_TYPE,
         sagemaker_session=pipeline_session,
     )
 
     step_args = processor.run(inputs=processing_input)
-
     try:
         ProcessingStep(
             name="MyProcessingStep",