change: Implement overide solution for pipeline varialbes

Author: Dewen Qi

src/sagemaker/parameter.py

@@ -14,9 +14,8 @@
 from __future__ import absolute_import
 
 import json
-from sagemaker.workflow.parameters import Parameter as PipelineParameter
-from sagemaker.workflow.functions import JsonGet as PipelineJsonGet
-from sagemaker.workflow.functions import Join as PipelineJoin
+
+from sagemaker.workflow.entities import PipelineVariable
 
 
 class ParameterRange(object):
@@ -73,11 +72,11 @@ def as_tuning_range(self, name):
         return {
             "Name": name,
             "MinValue": str(self.min_value)
-            if not isinstance(self.min_value, (PipelineParameter, PipelineJsonGet, PipelineJoin))
-            else self.min_value,
+            if not isinstance(self.min_value, PipelineVariable)
+            else self.min_value.to_string(),
             "MaxValue": str(self.max_value)
-            if not isinstance(self.max_value, (PipelineParameter, PipelineJsonGet, PipelineJoin))
-            else self.max_value,
+            if not isinstance(self.max_value, PipelineVariable)
+            else self.max_value.to_string(),
             "ScalingType": self.scaling_type,
         }
 
@@ -112,8 +111,7 @@ def __init__(self, values):  # pylint: disable=super-init-not-called
         """
         values = values if isinstance(values, list) else [values]
         self.values = [
-            str(v) if not isinstance(v, (PipelineParameter, PipelineJsonGet, PipelineJoin)) else v
-            for v in values
+            str(v) if not isinstance(v, PipelineVariable) else v.to_string() for v in values
         ]
 
     def as_tuning_range(self, name):

src/sagemaker/tuner.py

@@ -38,6 +38,7 @@
     IntegerParameter,
     ParameterRange,
 )
+from sagemaker.workflow.entities import PipelineVariable
 from sagemaker.workflow.parameters import Parameter as PipelineParameter
 from sagemaker.workflow.functions import JsonGet as PipelineJsonGet
 from sagemaker.workflow.functions import Join as PipelineJoin
@@ -376,9 +377,7 @@ def _prepare_static_hyperparameters(
         """Prepare static hyperparameters for one estimator before tuning."""
         # Remove any hyperparameter that will be tuned
         static_hyperparameters = {
-            str(k): str(v)
-            if not isinstance(v, (PipelineParameter, PipelineJsonGet, PipelineJoin))
-            else v
+            str(k): str(v) if not isinstance(v, PipelineVariable) else v.to_string()
             for (k, v) in estimator.hyperparameters().items()
         }
         for hyperparameter_name in hyperparameter_ranges.keys():

src/sagemaker/workflow/entities.py

@@ -16,7 +16,7 @@
 import abc
 
 from enum import EnumMeta
-from typing import Any, Dict, List, Union
+from typing import Any, Dict, List, Union, Optional
 
 PrimitiveType = Union[str, int, bool, float, None]
 RequestType = Union[Dict[str, Any], List[Dict[str, Any]]]
@@ -57,3 +57,62 @@ class Expression(abc.ABC):
     @abc.abstractmethod
     def expr(self) -> RequestType:
         """Get the expression structure for workflow service calls."""
+
+
+class PipelineVariable(Expression):
+    """Base object for pipeline variables
+
+    PipelineVariables must implement the expr property.
+    """
+
+    def __add__(self, other: Union[Expression, PrimitiveType]):
+        """Add function for PipelineVariable
+
+        Args:
+            other (Union[Expression, PrimitiveType]): The other object to be concatenated.
+
+        Always raise an error since pipeline variables do not support concatenation
+        """
+
+        raise TypeError("Pipeline variables do not support concatenation.")
+
+    def __str__(self):
+        """Override built-in String function for PipelineVariable"""
+        raise TypeError("Pipeline variables do not support __str__ operation.")
+
+    def __int__(self):
+        """Override built-in Integer function for PipelineVariable"""
+        raise TypeError("Pipeline variables do not support __int__ operation.")
+
+    def __float__(self):
+        """Override built-in Float function for PipelineVariable"""
+        raise TypeError("Pipeline variables do not support __float__ operation.")
+
+    def to_string(self):
+        """Prompt the pipeline to convert the pipeline variable to String in runtime"""
+        from sagemaker.workflow.functions import Join
+
+        return Join(on="", values=[self])
+
+    @property
+    @abc.abstractmethod
+    def expr(self) -> RequestType:
+        """Get the expression structure for workflow service calls."""
+
+    def startswith(
+        self,
+        prefix: Union[str, tuple],  # pylint: disable=unused-argument
+        start: Optional[int] = None,  # pylint: disable=unused-argument
+        end: Optional[int] = None,  # pylint: disable=unused-argument
+    ) -> bool:
+        """Simulate the Python string's built-in method: startswith
+
+        Args:
+            prefix (str, tuple): The (tuple of) string to be checked.
+            start (int): To set the start index of the matching boundary (default: None).
+            end (int): To set the end index of the matching boundary (default: None).
+
+        Return:
+            bool: Always return False as Pipeline variables are parsed during execution runtime
+        """
+        return False

src/sagemaker/workflow/execution_variables.py

@@ -14,12 +14,12 @@
 from __future__ import absolute_import
 
 from sagemaker.workflow.entities import (
-    Expression,
     RequestType,
+    PipelineVariable,
 )
 
 
-class ExecutionVariable(Expression):
+class ExecutionVariable(PipelineVariable):
     """Pipeline execution variables for workflow."""
 
     def __init__(self, name: str):
@@ -30,6 +30,13 @@ def __init__(self, name: str):
         """
         self.name = name
 
+    def to_string(self) -> PipelineVariable:
+        """Prompt the pipeline to convert the pipeline variable to String in runtime
+
+        As ExecutionVariable is treated as String in runtime, no extra actions are needed.
+        """
+        return self
+
     @property
     def expr(self) -> RequestType:
         """The 'Get' expression dict for an `ExecutionVariable`."""

src/sagemaker/workflow/functions.py

@@ -17,12 +17,12 @@
 
 import attr
 
-from sagemaker.workflow.entities import Expression
+from sagemaker.workflow.entities import PipelineVariable
 from sagemaker.workflow.properties import PropertyFile
 
 
 @attr.s
-class Join(Expression):
+class Join(PipelineVariable):
     """Join together properties.
 
     Examples:
@@ -38,15 +38,23 @@ class Join(Expression):
     Attributes:
         values (List[Union[PrimitiveType, Parameter, Expression]]):
             The primitive type values, parameters, step properties, expressions to join.
-        on_str (str): The string to join the values on (Defaults to "").
+        on (str): The string to join the values on (Defaults to "").
     """
 
     on: str = attr.ib(factory=str)
     values: List = attr.ib(factory=list)
 
+    def to_string(self) -> PipelineVariable:
+        """Prompt the pipeline to convert the pipeline variable to String in runtime
+
+        As Join is treated as String in runtime, no extra actions are needed.
+        """
+        return self
+
     @property
     def expr(self):
         """The expression dict for a `Join` function."""
+
         return {
             "Std:Join": {
                 "On": self.on,
@@ -58,7 +66,7 @@ def expr(self):
 
 
 @attr.s
-class JsonGet(Expression):
+class JsonGet(PipelineVariable):
     """Get JSON properties from PropertyFiles.
 
     Attributes:

src/sagemaker/workflow/parameters.py

@@ -24,6 +24,7 @@
     Entity,
     PrimitiveType,
     RequestType,
+    PipelineVariable,
 )
 
 
@@ -48,7 +49,7 @@ def python_type(self) -> Type:
 
 
 @attr.s
-class Parameter(Entity):
+class Parameter(PipelineVariable, Entity):
     """Pipeline parameter for workflow.
 
     Attributes:
@@ -170,6 +171,13 @@ def __hash__(self):
         """Hash function for parameter types"""
         return hash(tuple(self.to_request()))
 
+    def to_string(self) -> PipelineVariable:
+        """Prompt the pipeline to convert the pipeline variable to String in runtime
+
+        As ParameterString is treated as String in runtime, no extra actions are needed.
+        """
+        return self
+
     def to_request(self) -> RequestType:
         """Get the request structure for workflow service calls."""
         request_dict = super(ParameterString, self).to_request()

src/sagemaker/workflow/properties.py

@@ -13,16 +13,17 @@
 """The properties definitions for workflow."""
 from __future__ import absolute_import
 
+from abc import ABCMeta
 from typing import Dict, Union, List
 
 import attr
 
 import botocore.loaders
 
-from sagemaker.workflow.entities import Expression
+from sagemaker.workflow.entities import Expression, PipelineVariable
 
 
-class PropertiesMeta(type):
+class PropertiesMeta(ABCMeta):
     """Load an internal shapes attribute from the botocore service model
 
     for sagemaker and emr service.
@@ -44,7 +45,7 @@ def __new__(mcs, *args, **kwargs):
         return super().__new__(mcs, *args, **kwargs)
 
 
-class Properties(metaclass=PropertiesMeta):
+class Properties(PipelineVariable, metaclass=PropertiesMeta):
     """Properties for use in workflow expressions."""
 
     def __init__(

tests/integ/sagemaker/workflow/test_pipeline_var_behaviors.py

@@ -0,0 +1,119 @@
+# 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.
+from __future__ import absolute_import
+
+import pytest
+from botocore.exceptions import WaiterError
+
+from sagemaker import get_execution_role, utils
+from sagemaker.workflow.condition_step import ConditionStep
+from sagemaker.workflow.conditions import ConditionGreaterThan
+from sagemaker.workflow.fail_step import FailStep
+from sagemaker.workflow.functions import Join
+from sagemaker.workflow.parameters import ParameterString, ParameterInteger
+from sagemaker.workflow.pipeline import Pipeline
+
+
+@pytest.fixture
+def role(sagemaker_session):
+    return get_execution_role(sagemaker_session)
+
+
+@pytest.fixture
+def pipeline_name():
+    return utils.unique_name_from_base("my-pipeline-vars")
+
+
+def test_ppl_var_to_string_and_add(sagemaker_session, role, pipeline_name):
+    param_str = ParameterString(name="MyString", default_value="1")
+    param_int = ParameterInteger(name="MyInteger", default_value=3)
+
+    cond = ConditionGreaterThan(left=param_str, right=param_int.to_string())
+    step_cond = ConditionStep(
+        name="CondStep",
+        conditions=[cond],
+        if_steps=[],
+        else_steps=[],
+    )
+    join_fn1 = Join(
+        on=" ",
+        values=[
+            "condition greater than check return:",
+            step_cond.properties.Outcome.to_string(),
+            "and left side param str is",
+            param_str,
+            "and right side param int is",
+            param_int,
+        ],
+    )
+
+    step_fail = FailStep(
+        name="FailStep",
+        error_message=join_fn1,
+    )
+    pipeline = Pipeline(
+        name=pipeline_name,
+        parameters=[param_str, param_int],
+        steps=[step_cond, step_fail],
+        sagemaker_session=sagemaker_session,
+    )
+
+    try:
+        response = pipeline.create(role)
+        pipeline_arn = response["PipelineArn"]
+        execution = pipeline.start()
+        response = execution.describe()
+        assert response["PipelineArn"] == pipeline_arn
+
+        try:
+            execution.wait(delay=30, max_attempts=60)
+        except WaiterError:
+            pass
+        execution_steps = execution.list_steps()
+
+        assert len(execution_steps) == 2
+        for execution_step in execution_steps:
+            if execution_step["StepName"] == "CondStep":
+                assert execution_step["StepStatus"] == "Succeeded"
+                continue
+            assert execution_step["StepName"] == "FailStep"
+            assert execution_step["StepStatus"] == "Failed"
+            assert (
+                execution_step["FailureReason"] == "condition greater than check return: false "
+                "and left side param str is 1 and right side param int is 3"
+            )
+
+        # Update int param to update cond step outcome
+        execution = pipeline.start(parameters={"MyInteger": 0})
+        try:
+            execution.wait(delay=30, max_attempts=60)
+        except WaiterError:
+            pass
+        execution_steps = execution.list_steps()
+
+        assert len(execution_steps) == 2
+        for execution_step in execution_steps:
+            if execution_step["StepName"] == "CondStep":
+                assert execution_step["StepStatus"] == "Succeeded"
+                continue
+            assert execution_step["StepName"] == "FailStep"
+            assert execution_step["StepStatus"] == "Failed"
+            assert (
+                execution_step["FailureReason"] == "condition greater than check return: true "
+                "and left side param str is 1 and right side param int is 0"
+            )
+    finally:
+        try:
+            pipeline.delete()
+        except Exception:
+            pass