fix: support get sub steps of a specific step

Signed-off-by: zjgemi <[email protected]>

Author: zjgemi

src/dflow/argo_objects.py

@@ -6,7 +6,7 @@
 import time
 from collections import UserDict, UserList
 from copy import deepcopy
-from typing import Any, List, Union
+from typing import Any, List, Optional, Union
 
 from .common import jsonpickle
 from .config import config, s3_config
@@ -359,6 +359,7 @@ def get_step(
             phase: Union[str, List[str]] = None,
             id: Union[str, List[str]] = None,
             type: Union[str, List[str]] = None,
+            parent_id: Optional[str] = None,
             sort_by_generation: bool = False,
     ) -> List[ArgoStep]:
         if name is not None and not isinstance(name, list):
@@ -373,7 +374,11 @@ def get_step(
             type = [type]
         step_list = []
         if hasattr(self.status, "nodes"):
-            for step in self.status.nodes.values():
+            if parent_id is not None:
+                nodes = self.get_sub_nodes(parent_id)
+            else:
+                nodes = self.status.nodes.values()
+            for step in nodes:
                 if step["startedAt"] is None:
                     continue
                 if name is not None and not match(step["displayName"], name):
@@ -396,6 +401,8 @@ def get_step(
                     continue
                 step = ArgoStep(step, self.metadata.name)
                 step_list.append(step)
+        else:
+            return []
         if sort_by_generation:
             self.generation = {}
             self.record_generation(self.id, 0)
@@ -405,13 +412,35 @@ def get_step(
             step_list.sort(key=lambda x: x["startedAt"])
         return step_list
 
+    def get_sub_nodes(self, node_id):
+        assert node_id in self.status.nodes
+        node = self.status.nodes[node_id]
+        if node["type"] not in ["Steps", "DAG"]:
+            return [node]
+        if node.get("memoizationStatus", {}).get("hit", False):
+            return [node]
+        sub_nodes = []
+        outbound_nodes = node.get("outboundNodes", [])
+        children = node.get("children", [])
+        # order by generation (BFS)
+        current_generation = children
+        while len(current_generation) > 0:
+            for id in current_generation:
+                sub_nodes.append(self.status.nodes[id])
+            next_generation = []
+            for id in current_generation:
+                if id not in outbound_nodes:
+                    next_generation += self.status.nodes[id].get(
+                        "children", [])
+            current_generation = next_generation
+        return sub_nodes
+
     def record_generation(self, node_id, generation):
         self.generation[node_id] = generation
-        if "children" in self.status.nodes[node_id]:
-            for child in self.status.nodes[node_id]["children"]:
-                if child in self.generation:
-                    continue
-                self.record_generation(child, generation+1)
+        for child in self.status.nodes[node_id].get("children", []):
+            if child in self.generation:
+                continue
+            self.record_generation(child, generation+1)
 
     def get_duration(self) -> datetime.timedelta:
         return get_duration(self.status)

src/dflow/workflow.py

@@ -999,6 +999,7 @@ def query_step(
             phase: Union[str, List[str]] = None,
             id: Union[str, List[str]] = None,
             type: Union[str, List[str]] = None,
+            parent_id: Optional[str] = None,
             sort_by_generation: bool = False,
     ) -> List[ArgoStep]:
         """
@@ -1017,6 +1018,7 @@ def query_step(
             phase: filter by phase of step
             id: filter by id of step
             type: filter by type of step
+            parent_id: get sub steps of a specific step
             sort_by_generation: sort results by the number of generation from
                 the root node
         Returns:
@@ -1105,7 +1107,7 @@ def query_step(
 
         return self.query().get_step(
             name=name, key=key, phase=phase, id=id, type=type,
-            sort_by_generation=sort_by_generation)
+            parent_id=parent_id, sort_by_generation=sort_by_generation)
 
     def query_keys_of_steps(
             self,