feature: Add helper method to generate pipeline adjacency list (#3128)

Co-authored-by: Namrata Madan <[email protected]>

Author: nmadan

src/sagemaker/workflow/_utils.py

@@ -369,7 +369,7 @@ def __init__(
         self.kwargs = kwargs
         self.container_def_list = container_def_list
 
-        self._properties = Properties(path=f"Steps.{name}", shape_name="DescribeModelPackageOutput")
+        self._properties = Properties(step_name=name, shape_name="DescribeModelPackageOutput")
 
     @property
     def arguments(self) -> RequestType:

src/sagemaker/workflow/callback_step.py

@@ -112,13 +112,12 @@ def __init__(
         self.cache_config = cache_config
         self.inputs = inputs
 
-        root_path = f"Steps.{name}"
-        root_prop = Properties(path=root_path)
+        root_prop = Properties(step_name=name)
 
         property_dict = {}
         for output in outputs:
             property_dict[output.output_name] = Properties(
-                f"{root_path}.OutputParameters['{output.output_name}']"
+                step_name=name, path=f"OutputParameters['{output.output_name}']"
             )
 
         root_prop.__dict__["Outputs"] = property_dict

src/sagemaker/workflow/clarify_check_step.py

@@ -236,13 +236,12 @@ def __init__(
             self._generate_processing_job_analysis_config(), self._baselining_processor
         )
 
-        root_path = f"Steps.{name}"
-        root_prop = Properties(path=root_path)
+        root_prop = Properties(step_name=name)
         root_prop.__dict__["CalculatedBaselineConstraints"] = Properties(
-            f"{root_path}.CalculatedBaselineConstraints"
+            step_name=name, path="CalculatedBaselineConstraints"
         )
         root_prop.__dict__["BaselineUsedForDriftCheckConstraints"] = Properties(
-            f"{root_path}.BaselineUsedForDriftCheckConstraints"
+            step_name=name, path="BaselineUsedForDriftCheckConstraints"
         )
         self._properties = root_prop
 

src/sagemaker/workflow/condition_step.py

@@ -77,9 +77,8 @@ def __init__(
         self.if_steps = if_steps or []
         self.else_steps = else_steps or []
 
-        root_path = f"Steps.{name}"
-        root_prop = Properties(path=root_path)
-        root_prop.__dict__["Outcome"] = Properties(f"{root_path}.Outcome")
+        root_prop = Properties(step_name=name)
+        root_prop.__dict__["Outcome"] = Properties(step_name=name, path="Outcome")
         self._properties = root_prop
 
     @property
@@ -91,6 +90,11 @@ def arguments(self) -> RequestType:
             ElseSteps=list_to_request(self.else_steps),
         )
 
+    @property
+    def step_only_arguments(self):
+        """Argument dict pertaining to the step only, and not the `if_steps` or `else_steps`."""
+        return self.conditions
+
     @property
     def properties(self):
         """A simple Properties object with `Outcome` as the only property"""
@@ -126,5 +130,10 @@ def expr(self):
             }
         }
 
+    @property
+    def _referenced_steps(self) -> List[str]:
+        """List of step names that this function depends on."""
+        return [self.step.name]
+
 
 JsonGet = deprecated_class(JsonGet, "JsonGet")

src/sagemaker/workflow/conditions.py

@@ -17,6 +17,8 @@
 """
 from __future__ import absolute_import
 
+import abc
+
 from enum import Enum
 from typing import Dict, List, Union
 
@@ -33,6 +35,7 @@
 from sagemaker.workflow.execution_variables import ExecutionVariable
 from sagemaker.workflow.parameters import Parameter
 from sagemaker.workflow.properties import Properties
+from sagemaker.workflow.entities import PipelineVariable
 
 # TODO: consider base class for those with an expr method, rather than defining a type here
 ConditionValueType = Union[ExecutionVariable, Parameter, Properties]
@@ -61,6 +64,11 @@ class Condition(Entity):
 
     condition_type: ConditionTypeEnum = attr.ib(factory=ConditionTypeEnum.factory)
 
+    @property
+    @abc.abstractmethod
+    def _referenced_steps(self) -> List[str]:
+        """List of step names that this function depends on."""
+
 
 @attr.s
 class ConditionComparison(Condition):
@@ -84,6 +92,16 @@ def to_request(self) -> RequestType:
             "RightValue": primitive_or_expr(self.right),
         }
 
+    @property
+    def _referenced_steps(self) -> List[str]:
+        """List of step names that this function depends on."""
+        steps = []
+        if isinstance(self.left, PipelineVariable):
+            steps.extend(self.left._referenced_steps)
+        if isinstance(self.right, PipelineVariable):
+            steps.extend(self.right._referenced_steps)
+        return steps
+
 
 class ConditionEquals(ConditionComparison):
     """A condition for equality comparisons."""
@@ -213,6 +231,17 @@ def to_request(self) -> RequestType:
             "Values": [primitive_or_expr(in_value) for in_value in self.in_values],
         }
 
+    @property
+    def _referenced_steps(self) -> List[str]:
+        """List of step names that this function depends on."""
+        steps = []
+        if isinstance(self.value, PipelineVariable):
+            steps.extend(self.value._referenced_steps)
+        for in_value in self.in_values:
+            if isinstance(in_value, PipelineVariable):
+                steps.extend(in_value._referenced_steps)
+        return steps
+
 
 class ConditionNot(Condition):
     """A condition for negating another `Condition`."""
@@ -230,6 +259,11 @@ def to_request(self) -> RequestType:
         """Get the request structure for workflow service calls."""
         return {"Type": self.condition_type.value, "Expression": self.expression.to_request()}
 
+    @property
+    def _referenced_steps(self) -> List[str]:
+        """List of step names that this function depends on."""
+        return self.expression._referenced_steps
+
 
 class ConditionOr(Condition):
     """A condition for taking the logical OR of a list of `Condition` instances."""
@@ -250,6 +284,14 @@ def to_request(self) -> RequestType:
             "Conditions": [condition.to_request() for condition in self.conditions],
         }
 
+    @property
+    def _referenced_steps(self) -> List[str]:
+        """List of step names that this function depends on."""
+        steps = []
+        for condition in self.conditions:
+            steps.extend(condition._referenced_steps)
+        return steps
+
 
 def primitive_or_expr(
     value: Union[ExecutionVariable, Expression, PrimitiveType, Parameter, Properties]

src/sagemaker/workflow/emr_step.py

@@ -94,7 +94,7 @@ def __init__(
         self.args = emr_step_args
         self.cache_config = cache_config
 
-        root_property = Properties(path=f"Steps.{name}", shape_name="Step", service_name="emr")
+        root_property = Properties(step_name=name, shape_name="Step", service_name="emr")
         root_property.__dict__["ClusterId"] = cluster_id
         self._properties = root_property
 

src/sagemaker/workflow/entities.py

@@ -102,3 +102,8 @@ def to_string(self):
     @abc.abstractmethod
     def expr(self) -> RequestType:
         """Get the expression structure for workflow service calls."""
+
+    @property
+    @abc.abstractmethod
+    def _referenced_steps(self) -> List[str]:
+        """List of step names that this function depends on."""

src/sagemaker/workflow/execution_variables.py

@@ -13,6 +13,7 @@
 """Pipeline parameters and conditions for workflow."""
 from __future__ import absolute_import
 
+from typing import List
 from sagemaker.workflow.entities import (
     RequestType,
     PipelineVariable,
@@ -42,6 +43,11 @@ def expr(self) -> RequestType:
         """The 'Get' expression dict for an `ExecutionVariable`."""
         return {"Get": f"Execution.{self.name}"}
 
+    @property
+    def _referenced_steps(self) -> List[str]:
+        """List of step names that this function depends on."""
+        return []
+
 
 class ExecutionVariables:
     """Provide access to all available execution variables:

src/sagemaker/workflow/functions.py

@@ -64,6 +64,15 @@ def expr(self):
             },
         }
 
+    @property
+    def _referenced_steps(self) -> List[str]:
+        """List of step names that this function depends on."""
+        steps = []
+        for value in self.values:
+            if isinstance(value, PipelineVariable):
+                steps.extend(value._referenced_steps)
+        return steps
+
 
 @attr.s
 class JsonGet(PipelineVariable):
@@ -96,3 +105,8 @@ def expr(self):
                 "Path": self.json_path,
             }
         }
+
+    @property
+    def _referenced_steps(self) -> List[str]:
+        """List of step names that this function depends on."""
+        return [self.step_name]

src/sagemaker/workflow/lambda_step.py

@@ -115,13 +115,12 @@ def __init__(
         self.cache_config = cache_config
         self.inputs = inputs if inputs is not None else {}
 
-        root_path = f"Steps.{name}"
-        root_prop = Properties(path=root_path)
+        root_prop = Properties(step_name=name)
 
         property_dict = {}
         for output in self.outputs:
             property_dict[output.output_name] = Properties(
-                f"{root_path}.OutputParameters['{output.output_name}']"
+                step_name=name, path=f"OutputParameters['{output.output_name}']"
             )
 
         root_prop.__dict__["Outputs"] = property_dict