feature: add helper method to generate pipeline adjacency list

Author: Namrata Madan

src/sagemaker/workflow/_utils.py

@@ -365,7 +365,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

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/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/pipeline.py

@@ -322,6 +322,16 @@ def _interpolate_step_collection_name_in_depends_on(self, step_requests: dict):
                     depends_on.append(depend_step_name)
             step_request["DependsOn"] = depends_on
 
+    def _generate_adjacency_list(self):
+        """Generate an adjacency list representing the step dependency DAG in this pipeline."""
+        adjacency_list = {}
+        for step in self.steps:
+            if isinstance(step, Step):
+                adjacency_list[step.name] = step._find_step_dependencies()
+            elif isinstance(step, StepCollection):
+                adjacency_list.update(step._find_step_dependencies())
+        return adjacency_list
+
 
 def format_start_parameters(parameters: Dict[str, Any]) -> List[Dict[str, Any]]:
     """Formats start parameter overrides as a list of dicts.

src/sagemaker/workflow/properties.py

@@ -50,7 +50,8 @@ class Properties(PipelineVariable, metaclass=PropertiesMeta):
 
     def __init__(
         self,
-        path: str,
+        step_name: str,
+        path: str = None,
         shape_name: str = None,
         shape_names: List[str] = None,
         service_name: str = "sagemaker",
@@ -62,7 +63,9 @@ def __init__(
             shape_name (str): The botocore service model shape name.
             shape_names (str): A List of the botocore service model shape name.
         """
-        self._path = path
+        self.step_name = step_name
+        prefix = f"Steps.{step_name}"
+        self._path = prefix if path is None else f"{prefix}.{path}"
         shape_names = [] if shape_names is None else shape_names
         self._shape_names = shape_names if shape_name is None else [shape_name] + shape_names
 
@@ -78,15 +81,24 @@ def __init__(
                 for key, info in members.items():
                     if shapes.get(info["shape"], {}).get("type") == "list":
                         self.__dict__[key] = PropertiesList(
-                            f"{path}.{key}", info["shape"], service_name
+                            step_name=step_name,
+                            path=".".join(filter(None, (path, key))),
+                            shape_name=info["shape"],
+                            service_name=service_name,
                         )
                     elif shapes.get(info["shape"], {}).get("type") == "map":
                         self.__dict__[key] = PropertiesMap(
-                            f"{path}.{key}", info["shape"], service_name
+                            step_name=step_name,
+                            path=".".join(filter(None, (path, key))),
+                            shape_name=info["shape"],
+                            service_name=service_name,
                         )
                     else:
                         self.__dict__[key] = Properties(
-                            f"{path}.{key}", info["shape"], service_name=service_name
+                            step_name=step_name,
+                            path=".".join(filter(None, (path, key))),
+                            shape_name=info["shape"],
+                            service_name=service_name,
                         )
 
     @property
@@ -98,17 +110,21 @@ def expr(self):
 class PropertiesList(Properties):
     """PropertiesList for use in workflow expressions."""
 
-    def __init__(self, path: str, shape_name: str = None, service_name: str = "sagemaker"):
+    def __init__(
+        self, step_name: str, path: str, shape_name: str = None, service_name: str = "sagemaker"
+    ):
         """Create a PropertiesList instance representing the given shape.
 
         Args:
             path (str): The parent path of the PropertiesList instance.
             shape_name (str): The botocore service model shape name.
             service_name (str): The botocore service name.
         """
-        super(PropertiesList, self).__init__(path, shape_name)
+        super(PropertiesList, self).__init__(step_name, path, shape_name)
+        self.step_name = step_name
         self.shape_name = shape_name
         self.service_name = service_name
+        self.path = path
         self._items: Dict[Union[int, str], Properties] = dict()
 
     def __getitem__(self, item: Union[int, str]):
@@ -121,9 +137,9 @@ def __getitem__(self, item: Union[int, str]):
             shape = Properties._shapes_map.get(self.service_name, {}).get(self.shape_name)
             member = shape["member"]["shape"]
             if isinstance(item, str):
-                property_item = Properties(f"{self._path}['{item}']", member)
+                property_item = Properties(self.step_name, f"{self.path}['{item}']", member)
             else:
-                property_item = Properties(f"{self._path}[{item}]", member)
+                property_item = Properties(self.step_name, f"{self.path}[{item}]", member)
             self._items[item] = property_item
 
         return self._items.get(item)
@@ -132,17 +148,21 @@ def __getitem__(self, item: Union[int, str]):
 class PropertiesMap(Properties):
     """PropertiesMap for use in workflow expressions."""
 
-    def __init__(self, path: str, shape_name: str = None, service_name: str = "sagemaker"):
+    def __init__(
+        self, step_name: str, path: str, shape_name: str = None, service_name: str = "sagemaker"
+    ):
         """Create a PropertiesMap instance representing the given shape.
 
         Args:
             path (str): The parent path of the PropertiesMap instance.
             shape_name (str): The botocore sagemaker service model shape name.
             service_name (str): The botocore service name.
         """
-        super(PropertiesMap, self).__init__(path, shape_name)
+        super(PropertiesMap, self).__init__(step_name, path, shape_name)
+        self.step_name = step_name
         self.shape_name = shape_name
         self.service_name = service_name
+        self.path = path
         self._items: Dict[Union[int, str], Properties] = dict()
 
     def __getitem__(self, item: Union[int, str]):
@@ -155,9 +175,9 @@ def __getitem__(self, item: Union[int, str]):
             shape = Properties._shapes_map.get(self.service_name, {}).get(self.shape_name)
             member = shape["value"]["shape"]
             if isinstance(item, str):
-                property_item = Properties(f"{self._path}['{item}']", member)
+                property_item = Properties(self.step_name, f"{self.path}['{item}']", member)
             else:
-                property_item = Properties(f"{self._path}[{item}]", member)
+                property_item = Properties(self.step_name, f"{self.path}[{item}]", member)
             self._items[item] = property_item
 
         return self._items.get(item)

src/sagemaker/workflow/quality_check_step.py

@@ -205,19 +205,18 @@ def __init__(
             ],
         )
 
-        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__["CalculatedBaselineStatistics"] = Properties(
-            f"{root_path}.CalculatedBaselineStatistics"
+            step_name=name, path="CalculatedBaselineStatistics"
         )
         root_prop.__dict__["BaselineUsedForDriftCheckStatistics"] = Properties(
-            f"{root_path}.BaselineUsedForDriftCheckStatistics"
+            step_name=name, path="BaselineUsedForDriftCheckStatistics"
         )
         root_prop.__dict__["BaselineUsedForDriftCheckConstraints"] = Properties(
-            f"{root_path}.BaselineUsedForDriftCheckConstraints"
+            step_name=name, path="BaselineUsedForDriftCheckConstraints"
         )
         self._properties = root_prop
 

src/sagemaker/workflow/step_collections.py

@@ -54,6 +54,13 @@ def properties(self):
         size = len(self.steps)
         return self.steps[size - 1].properties
 
+    def _find_step_dependencies(self):
+        """Find the step names this step collection is dependent on."""
+        dependencies = {}
+        for step in self.steps:
+            dependencies[step.name] = step._find_step_dependencies()
+        return dependencies
+
 
 class RegisterModel(StepCollection):  # pragma: no cover
     """Register Model step collection for workflow."""

src/sagemaker/workflow/steps.py

@@ -148,6 +148,33 @@ def _resolve_depends_on(
                 raise ValueError(f"Invalid input step name: {step}")
         return depends_on
 
+    def _find_step_dependencies(self):
+        """Find the step names this step is dependent on."""
+
+        def _find_dependencies_in_step_arguments(obj):
+            dependencies = set()
+            if isinstance(obj, dict):
+                for value in obj.values():
+                    if isinstance(value, Properties):
+                        dependencies.add(value.step_name)
+                    dependencies.update(_find_dependencies_in_step_arguments(value))
+            if isinstance(obj, list):
+                for item in obj:
+                    if isinstance(item, Properties):
+                        dependencies.add(item.step_name)
+                    dependencies.update(_find_dependencies_in_step_arguments(item))
+            return dependencies
+
+        step_dependencies = set()
+        if self.depends_on:
+            for depends_on_step in self.depends_on:
+                if isinstance(depends_on_step, str):
+                    step_dependencies.add(depends_on_step)
+                elif isinstance(depends_on_step, Step):
+                    step_dependencies.add(depends_on_step.name)
+        step_dependencies.update(_find_dependencies_in_step_arguments(self.arguments))
+        return list(step_dependencies)
+
 
 @attr.s
 class CacheConfig:
@@ -290,9 +317,7 @@ def __init__(
         self.estimator = estimator
         self.inputs = inputs
 
-        self._properties = Properties(
-            path=f"Steps.{name}", shape_name="DescribeTrainingJobResponse"
-        )
+        self._properties = Properties(step_name=name, shape_name="DescribeTrainingJobResponse")
         self.cache_config = cache_config
 
         if self.cache_config:
@@ -430,7 +455,7 @@ def __init__(
         self.model = model
         self.inputs = inputs or CreateModelInput()
 
-        self._properties = Properties(path=f"Steps.{name}", shape_name="DescribeModelOutput")
+        self._properties = Properties(step_name=name, shape_name="DescribeModelOutput")
 
         # TODO: add public document link here once ready
         warnings.warn(
@@ -526,9 +551,7 @@ def __init__(
         self.transformer = transformer
         self.inputs = inputs
         self.cache_config = cache_config
-        self._properties = Properties(
-            path=f"Steps.{name}", shape_name="DescribeTransformJobResponse"
-        )
+        self._properties = Properties(step_name=name, shape_name="DescribeTransformJobResponse")
 
         if not self.step_args:
             if inputs is None:
@@ -652,9 +675,7 @@ def __init__(
         self.job_name = None
         self.kms_key = kms_key
         self.cache_config = cache_config
-        self._properties = Properties(
-            path=f"Steps.{name}", shape_name="DescribeProcessingJobResponse"
-        )
+        self._properties = Properties(step_name=name, shape_name="DescribeProcessingJobResponse")
 
         if not self.step_args:
             # Examine why run method in `sagemaker.processing.Processor`
@@ -806,7 +827,7 @@ def __init__(
         self.inputs = inputs
         self.job_arguments = job_arguments
         self._properties = Properties(
-            path=f"Steps.{name}",
+            step_name=name,
             shape_names=[
                 "DescribeHyperParameterTuningJobResponse",
                 "ListTrainingJobsForHyperParameterTuningJobResponse",

tests/unit/sagemaker/workflow/helpers.py

@@ -42,7 +42,7 @@ def __init__(self, name, display_name=None, description=None, depends_on=None):
         super(CustomStep, self).__init__(
             name, display_name, description, StepTypeEnum.TRAINING, depends_on
         )
-        self._properties = Properties(path=f"Steps.{name}")
+        self._properties = Properties(name)
 
     @property
     def arguments(self):

tests/unit/sagemaker/workflow/test_callback_step.py

@@ -21,7 +21,7 @@
 from sagemaker.workflow.parameters import ParameterInteger, ParameterString
 from sagemaker.workflow.pipeline import Pipeline
 from sagemaker.workflow.callback_step import CallbackStep, CallbackOutput, CallbackOutputTypeEnum
-from tests.unit.sagemaker.workflow.helpers import CustomStep
+from tests.unit.sagemaker.workflow.helpers import CustomStep, ordered
 
 
 @pytest.fixture
@@ -52,6 +52,7 @@ def test_callback_step():
         ],
         "Arguments": {"arg1": "foo", "arg2": 5, "arg3": param},
     }
+    assert ordered(cb_step._find_step_dependencies()) == ["SecondTestStep", "TestStep"]
 
 
 def test_callback_step_default_values():
@@ -75,6 +76,7 @@ def test_callback_step_default_values():
         ],
         "Arguments": {"arg1": "foo", "arg2": 5, "arg3": param},
     }
+    assert ordered(cb_step._find_step_dependencies()) == ["SecondTestStep", "TestStep"]
 
 
 def test_callback_step_output_expr():
@@ -95,6 +97,7 @@ def test_callback_step_output_expr():
     assert cb_step.properties.Outputs["output2"].expr == {
         "Get": "Steps.MyCallbackStep.OutputParameters['output2']"
     }
+    assert cb_step._find_step_dependencies() == ["TestStep"]
 
 
 def test_pipeline_interpolates_callback_outputs():
@@ -156,3 +159,6 @@ def test_pipeline_interpolates_callback_outputs():
             },
         ],
     }
+    assert ordered(pipeline._generate_adjacency_list()) == ordered(
+        {"MyCallbackStep1": ["TestStep"], "MyCallbackStep2": ["TestStep"], "TestStep": []}
+    )