feature: add helper method to generate pipeline adjacency list

Author: Namrata Madan

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.arguments["Conditions"]
+
     @property
     def properties(self):
         """A simple Properties object with `Outcome` as the only property"""
@@ -126,5 +130,10 @@ def expr(self):
             }
         }
 
+    @property
+    def depends_on(self) -> List[str]:
+        """List of step names that this function depends on."""
+        return [self.step.name]
+
 
 JsonGet = deprecated_class(JsonGet, "JsonGet")

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 depends_on(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 depends_on(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 depends_on(self) -> List[str]:
+        """List of step names that this function depends on."""
+        steps = []
+        for value in self.values:
+            if hasattr(value, "depends_on"):
+                steps.extend(value.depends_on)
+        return steps
+
 
 @attr.s
 class JsonGet(PipelineVariable):
@@ -96,3 +105,8 @@ def expr(self):
                 "Path": self.json_path,
             }
         }
+
+    @property
+    def depends_on(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

src/sagemaker/workflow/parameters.py

@@ -90,6 +90,11 @@ def expr(self) -> Dict[str, str]:
         """The 'Get' expression dict for a `Parameter`."""
         return Parameter._expr(self.name)
 
+    @property
+    def depends_on(self) -> List[str]:
+        """List of step names that this function depends on."""
+        return []
+
     @classmethod
     def _expr(cls, name):
         """An internal classmethod for the 'Get' expression dict for a `Parameter`.

src/sagemaker/workflow/pipeline.py

@@ -16,7 +16,7 @@
 import json
 
 from copy import deepcopy
-from typing import Any, Dict, List, Sequence, Union, Optional
+from typing import Any, Dict, List, Sequence, Union, Optional, Set
 
 import attr
 import botocore
@@ -39,6 +39,7 @@
 from sagemaker.workflow.properties import Properties
 from sagemaker.workflow.steps import Step
 from sagemaker.workflow.step_collections import StepCollection
+from sagemaker.workflow.condition_step import ConditionStep
 from sagemaker.workflow.utilities import list_to_request
 
 
@@ -534,3 +535,138 @@ def wait(self, delay=30, max_attempts=60):
             waiter_id, model, self.sagemaker_session.sagemaker_client
         )
         waiter.wait(PipelineExecutionArn=self.arn)
+
+
+class PipelineGraph:
+    """Helper class representing the Pipeline Directed Acyclic Graph (DAG)
+
+    Attributes:
+        steps (Sequence[Union[Step, StepCollection]]): Sequence of `Step`s and/or `StepCollection`s
+            that represent each node in the pipeline DAG
+    """
+
+    def __init__(self, steps: Sequence[Union[Step, StepCollection]]):
+        if self.is_duplicate_step_name(steps):
+            raise ValueError("Pipeline steps cannot have duplicate names.")
+        self.step_map = PipelineGraph._generate_step_map(steps)
+        self.adjacency_list = self._initialize_adjacency_list()
+        if self.is_cyclic():
+            raise ValueError("Cycle detected in pipeline step graph.")
+
+    @staticmethod
+    def _generate_step_map(
+        steps: Sequence[Union[Step, StepCollection]]
+    ) -> Dict[str, Union[Step, StepCollection]]:
+        """Helper method to create a mapping from Step/Step Collection name to itself."""
+        step_map = {}
+        for step in steps:
+            if isinstance(step, Step):
+                step_map[step.name] = step
+                if isinstance(step, ConditionStep):
+                    step_map.update(
+                        PipelineGraph._generate_step_map(step.if_steps + step.else_steps)
+                    )
+            elif isinstance(step, StepCollection):
+                step_map[step.name] = step
+                for inner_step in step.steps:
+                    step_map[inner_step.name] = inner_step
+        return step_map
+
+    @classmethod
+    def from_pipeline(cls, pipeline: Pipeline):
+        """Create a PipelineGraph object from the Pipeline object."""
+        return cls(pipeline.steps)
+
+    def _initialize_adjacency_list(self) -> Dict[str, Set[str]]:
+        """Generate an adjacency list representing the step dependency DAG in this pipeline."""
+        from collections import defaultdict
+
+        dependency_list = defaultdict(set)
+        for step in self.step_map.values():
+            if isinstance(step, Step):
+                dependency_list[step.name].update(step._find_step_dependencies(self.step_map))
+            elif isinstance(step, StepCollection):
+                step_collection_dependencies = step._find_step_dependencies(self.step_map)
+                for k, v in step_collection_dependencies.items():
+                    dependency_list[k].update(v)
+
+            if isinstance(step, ConditionStep):
+                for child_step in step.if_steps + step.else_steps:
+                    if isinstance(child_step, Step):
+                        dependency_list[child_step.name].add(step.name)
+                    elif isinstance(child_step, StepCollection):
+                        child_first_step = self.step_map[child_step.name].steps[0].name
+                        dependency_list[child_first_step].add(step.name)
+
+        adjacency_list = {}
+        for step in dependency_list:
+            for step_dependency in dependency_list[step]:
+                adjacency_list[step_dependency] = list(
+                    set(adjacency_list.get(step_dependency, []) + [step])
+                )
+        for step in dependency_list:
+            if step not in adjacency_list:
+                adjacency_list[step] = []
+        return adjacency_list
+
+    def is_duplicate_step_name(self, steps: Sequence[Union[Step, StepCollection]]) -> bool:
+        """Check if the provided step names have any duplicates.
+
+        Returns true if there are one or more duplicates, false otherwise.
+        """
+        step_names = set()
+        for step in steps:
+            if step.name in step_names:
+                return True
+            step_names.add(step.name)
+        return False
+
+    def is_cyclic(self) -> bool:
+        """Check if this pipeline graph is cyclic.
+
+        Returns true if it is cyclic, false otherwise.
+        """
+
+        def is_cyclic_helper(current_step):
+            visited_steps.add(current_step)
+            recurse_steps.add(current_step)
+            for child_step in self.adjacency_list[current_step]:
+                if child_step in recurse_steps:
+                    return True
+                if child_step not in visited_steps:
+                    if is_cyclic_helper(child_step):
+                        return True
+            recurse_steps.remove(current_step)
+            return False
+
+        visited_steps = set()
+        recurse_steps = set()
+        for step in self.adjacency_list:
+            if step not in visited_steps:
+                if is_cyclic_helper(step):
+                    return True
+        return False
+
+    def __iter__(self):
+        """Perform topological sort traversal of the Pipeline Graph."""
+
+        def topological_sort(current_step):
+            visited_steps.add(current_step)
+            for child_step in self.adjacency_list[current_step]:
+                if child_step not in visited_steps:
+                    topological_sort(child_step)
+            self.stack.append(current_step)
+
+        visited_steps = set()
+        self.stack = []  # pylint: disable=W0201
+        for step in self.adjacency_list:
+            if step not in visited_steps:
+                topological_sort(step)
+        return self
+
+    def __next__(self) -> Step:
+        """Return the next Step node from the Topological sort order."""
+
+        while self.stack:
+            return self.step_map.get(self.stack.pop())
+        raise StopIteration