debugged workflow graph plotting

Author: tclose

pydra/engine/core.py

@@ -18,7 +18,7 @@
 import attrs
 from filelock import SoftFileLock
 from pydra.engine.specs import TaskDef, WorkflowDef, TaskOutputs, WorkflowOutputs
-from pydra.engine.graph import DiGraph
+from pydra.engine.graph import DiGraph, INPUTS_NODE_NAME, OUTPUTS_NODE_NAME
 from pydra.engine import state
 from .lazy import LazyInField, LazyOutField
 from pydra.utils.hash import hash_function, Cache
@@ -598,16 +598,32 @@ def clear_cache(
 
     @classmethod
     def construct(
-        cls, definition: WorkflowDef[WorkflowOutputsType], dont_cache: bool = False
+        cls,
+        definition: WorkflowDef[WorkflowOutputsType],
+        dont_cache: bool = False,
+        lazy: ty.Sequence[str] = (),
     ) -> Self:
-        """Construct a workflow from a definition, caching the constructed worklow"""
+        """Construct a workflow from a definition, caching the constructed worklow
+
+        Parameters
+        ----------
+        definition : WorkflowDef
+            The definition of the workflow to construct
+        dont_cache : bool, optional
+            Whether to cache the constructed workflow, by default False
+        lazy : Sequence[str], optional
+            The names of the inputs to the workflow to be considered lazy even if they
+            have values in the given definition, by default ()
+        """
 
         # Check the previously constructed workflows to see if a workflow has been
         # constructed for the given set of inputs, or a less-specific set (i.e. with a
         # super-set of lazy inputs), and use that if it exists
 
         non_lazy_vals = {
-            n: v for n, v in attrs_values(definition).items() if not is_lazy(v)
+            n: v
+            for n, v in attrs_values(definition).items()
+            if not is_lazy(v) and n not in lazy
         }
         non_lazy_keys = frozenset(non_lazy_vals)
         hash_cache = Cache()  # share the hash cache to avoid recalculations
@@ -821,56 +837,55 @@ def _create_graph(
         DiGraph
             The graph of the workflow
         """
-        graph: DiGraph = DiGraph()
+        graph: DiGraph = DiGraph(name=self.name)
         for node in nodes:
             graph.add_nodes(node)
         # TODO: create connection is run twice
         for node in nodes:
             other_states = {}
-            for field in attrs_fields(node.inputs):
+            for field in list_fields(node._definition):
                 lf = node._definition[field.name]
                 if isinstance(lf, LazyOutField):
                     # adding an edge to the graph if task id expecting output from a different task
-                    if lf._node.name != self.name:
-                        # checking if the connection is already in the graph
-                        if (graph.node(lf._node.name), node) not in graph.edges:
-                            graph.add_edges((graph.node(lf._node.name), node))
-                        if detailed:
-                            graph.add_edges_description(
-                                (node.name, field.name, lf._node.name, lf._field)
-                            )
-                        logger.debug("Connecting %s to %s", lf._node.name, node.name)
-                        # adding a state from the previous task to other_states
+
+                    # checking if the connection is already in the graph
+                    if (graph.node(lf._node.name), node) not in graph.edges:
+                        graph.add_edges((graph.node(lf._node.name), node))
+                    if detailed:
+                        graph.add_edges_description(
+                            (node.name, field.name, lf._node.name, lf._field)
+                        )
+                    logger.debug("Connecting %s to %s", lf._node.name, node.name)
+                    # adding a state from the previous task to other_states
+                    if (
+                        graph.node(lf._node.name).state
+                        and graph.node(lf._node.name).state.splitter_rpn_final
+                    ):
+                        # variables that are part of inner splitters should be
+                        # treated as a containers
                         if (
-                            graph.node(lf._node.name).state
-                            and graph.node(lf._node.name).state.splitter_rpn_final
+                            node.state
+                            and f"{node.name}.{field.name}"
+                            in node.state._current_splitter_rpn
                         ):
-                            # variables that are part of inner splitters should be
-                            # treated as a containers
-                            if (
-                                node.state
-                                and f"{node.name}.{field.name}"
-                                in node.state._current_splitter_rpn
-                            ):
-                                node.state._inner_cont_dim[
-                                    f"{node.name}.{field.name}"
-                                ] = 1
-                            # adding task_name: (task.state, [a field from the connection]
-                            if lf._node.name not in other_states:
-                                other_states[lf._node.name] = (
-                                    graph.node(lf._node.name).state,
-                                    [field.name],
-                                )
-                            else:
-                                # if the task already exist in other_state,
-                                # additional field name should be added to the list of fields
-                                other_states[lf._node.name][1].append(field.name)
-                    else:  # LazyField with the wf input
-                        # connections with wf input should be added to the detailed graph description
-                        if detailed:
-                            graph.add_edges_description(
-                                (node.name, field.name, lf._node.name, lf._field)
+                            node.state._inner_cont_dim[f"{node.name}.{field.name}"] = 1
+                        # adding task_name: (task.state, [a field from the connection]
+                        if lf._node.name not in other_states:
+                            other_states[lf._node.name] = (
+                                graph.node(lf._node.name).state,
+                                [field.name],
                             )
+                        else:
+                            # if the task already exist in other_state,
+                            # additional field name should be added to the list of fields
+                            other_states[lf._node.name][1].append(field.name)
+                elif (
+                    isinstance(lf, LazyInField) and detailed
+                ):  # LazyField with the wf input
+                    # connections with wf input should be added to the detailed graph description
+                    graph.add_edges_description(
+                        (node.name, field.name, INPUTS_NODE_NAME, lf._field)
+                    )
 
             # if task has connections state has to be recalculated
             if other_states:
@@ -890,6 +905,12 @@ def _create_graph(
                         other_states=other_states,
                         combiner=combiner,
                     )
+        if detailed:
+            lf: LazyOutField
+            for outpt_name, lf in attrs_values(self.outputs).items():
+                graph.add_edges_description(
+                    (OUTPUTS_NODE_NAME, outpt_name, lf._node.name, lf._field)
+                )
         return graph
 
 

pydra/engine/graph.py

@@ -10,6 +10,9 @@
 
 NodeType = ty.TypeVar("NodeType")
 
+INPUTS_NODE_NAME = "__INPUTS__"
+OUTPUTS_NODE_NAME = "__OUTPUTS__"
+
 
 class DiGraph(ty.Generic[NodeType]):
     """A simple Directed Graph object."""
@@ -395,7 +398,7 @@ def create_dotfile_simple(self, outdir, name="graph"):
 
         dotstr = "digraph G {\n"
         for nd in self.nodes:
-            if is_workflow(nd):
+            if is_workflow(nd._definition):
                 if nd.state:
                     # adding color for wf with a state
                     dotstr += f"{nd.name} [shape=box, color=blue]\n"
@@ -430,27 +433,29 @@ def create_dotfile_detailed(self, outdir, name="graph_det"):
         if not self._nodes_details:
             raise Exception("node_details is empty, detailed dotfile can't be created")
         for nd_nm, nd_det in self.nodes_details.items():
-            if nd_nm == self.name:  # the main workflow itself
+            if nd_nm == INPUTS_NODE_NAME:  # the main workflow itself
                 # wf inputs
                 wf_inputs_str = f'{{<{nd_det["outputs"][0]}> {nd_det["outputs"][0]}'
                 for el in nd_det["outputs"][1:]:
                     wf_inputs_str += f" | <{el}> {el}"
                 wf_inputs_str += "}"
-                dotstr += f'struct_{nd_nm} [color=red, label="{{WORKFLOW INPUT: | {wf_inputs_str}}}"];\n'
+                dotstr += (
+                    f"struct_{self.name} [color=red, "
+                    f'label="{{WORKFLOW INPUT: | {wf_inputs_str}}}"];\n'
+                )
+            elif nd_nm == OUTPUTS_NODE_NAME:
                 # wf outputs
                 wf_outputs_str = f'{{<{nd_det["inputs"][0]}> {nd_det["inputs"][0]}'
                 for el in nd_det["inputs"][1:]:
                     wf_outputs_str += f" | <{el}> {el}"
                 wf_outputs_str += "}"
                 dotstr += (
-                    f"struct_{nd_nm}_out "
+                    f"struct_{self.name}_out "
                     f'[color=red, label="{{WORKFLOW OUTPUT: | {wf_outputs_str}}}"];\n'
                 )
                 # connections to the wf outputs
                 for con in nd_det["connections"]:
-                    dotstr += (
-                        f"struct_{con[1]}:{con[2]} -> struct_{nd_nm}_out:{con[0]};\n"
-                    )
+                    dotstr += f"struct_{con[1]}:{con[2]} -> struct_{self.name}_out:{con[0]};\n"
             else:  # elements of the main workflow
                 inputs_str = "{INPUT:"
                 for inp in nd_det["inputs"]:
@@ -466,7 +471,11 @@ def create_dotfile_detailed(self, outdir, name="graph_det"):
                 )
                 # connections between elements
                 for con in nd_det["connections"]:
-                    dotstr += f"struct_{con[1]}:{con[2]} -> struct_{nd_nm}:{con[0]};\n"
+                    in_conn = self.name if con[1] == INPUTS_NODE_NAME else con[1]
+                    out_conn = self.name if con[0] == OUTPUTS_NODE_NAME else con[0]
+                    dotstr += (
+                        f"struct_{in_conn}:{con[2]} -> struct_{nd_nm}:{out_conn};\n"
+                    )
         dotstr += "}"
         Path(outdir).mkdir(parents=True, exist_ok=True)
         dotfile = Path(outdir) / f"{name}.dot"
@@ -486,16 +495,17 @@ def create_dotfile_nested(self, outdir, name="graph"):
     def _create_dotfile_single_graph(self, nodes, edges):
         from .core import is_workflow
 
-        wf_asnd = []
+        wf_asnd = {}
         dotstr = ""
         for nd in nodes:
-            if is_workflow(nd):
-                wf_asnd.append(nd.name)
-                for task in nd.graph.nodes:
-                    nd.create_connections(task)
+            if is_workflow(nd._definition):
+                nd_graph = nd._definition.construct().graph()
+                wf_asnd[nd.name] = nd_graph
+                # for task in nd_graph.nodes:
+                #     nd.create_connections(task)
                 dotstr += f"subgraph cluster_{nd.name} {{\n" f"label = {nd.name} \n"
                 dotstr += self._create_dotfile_single_graph(
-                    nodes=nd.graph.nodes, edges=nd.graph.edges
+                    nodes=nd_graph.nodes, edges=nd_graph.edges
                 )
                 if nd.state:
                     dotstr += "color=blue\n"
@@ -517,12 +527,14 @@ def _create_dotfile_single_graph(self, nodes, edges):
                     f"lhead=cluster_{ed[1].name}]\n"
                 )
             elif ed[0].name in wf_asnd:
-                tail_nd = list(ed[0].nodes)[-1].name
+                nd_nodes = wf_asnd[ed[0].name].nodes
+                tail_nd = list(nd_nodes)[-1].name
                 dotstr_edg += (
                     f"{tail_nd} -> {ed[1].name} [ltail=cluster_{ed[0].name}]\n"
                 )
             elif ed[1].name in wf_asnd:
-                head_nd = list(ed[1].nodes)[0].name
+                nd_nodes = wf_asnd[ed[1].name].nodes
+                head_nd = list(nd_nodes)[0].name
                 dotstr_edg += (
                     f"{ed[0].name} -> {head_nd} [lhead=cluster_{ed[1].name}]\n"
                 )

pydra/engine/helpers.py

@@ -34,35 +34,36 @@
 def plot_workflow(
     workflow_task: "WorkflowDef",
     out_dir: Path,
-    type="simple",
-    export=None,
-    name=None,
-    output_dir=None,
+    plot_type: str = "simple",
+    export: ty.Sequence[str] | None = None,
+    name: str | None = None,
+    output_dir: Path | None = None,
+    lazy: ty.Sequence[str] | ty.Set[str] = (),
 ):
     """creating a graph - dotfile and optionally exporting to other formats"""
     from .core import Workflow
 
     # Create output directory
     out_dir.mkdir(parents=True, exist_ok=True)
 
-    # Construct the workflow object
-    wf = Workflow.construct(workflow_task)
+    # Construct the workflow object with all of the fields lazy
+    wf = Workflow.construct(workflow_task, lazy=lazy)
 
     if not name:
         name = f"graph_{type(workflow_task).__name__}"
-    if type == "simple":
+    if plot_type == "simple":
         graph = wf.graph()
         dotfile = graph.create_dotfile_simple(outdir=out_dir, name=name)
-    elif type == "nested":
+    elif plot_type == "nested":
         graph = wf.graph()
         dotfile = graph.create_dotfile_nested(outdir=out_dir, name=name)
-    elif type == "detailed":
+    elif plot_type == "detailed":
         graph = wf.graph(detailed=True)
         dotfile = graph.create_dotfile_detailed(outdir=out_dir, name=name)
     else:
         raise Exception(
             f"type of the graph can be simple, detailed or nested, "
-            f"but {type} provided"
+            f"but {plot_type} provided"
         )
     if not export:
         return dotfile

pydra/engine/node.py

@@ -205,11 +205,11 @@ def _set_state(self) -> None:
                 if not_split := [
                     c
                     for c in combiner
-                    if not any(c in s for s in self.state.splitter_rpn)
+                    if not any(c in s for s in self.state.splitter_rpn) and "." not in c
                 ]:
                     raise ValueError(
                         f"Combiner fields {not_split} for Node {self.name!r} are not in the "
-                        f"splitter fields {self.state.splitter_rpn}"
+                        f"splitter  {self.state.splitter_rpn}"
                     )
         else:
             self._state = None

pydra/engine/state.py

@@ -855,13 +855,17 @@ def splitter_validation(self):
 
     def combiner_validation(self):
         """validating if the combiner is correct (after all states are connected)"""
-        if self.combiner:
+        if local_names := set(c for c in self.combiner if "." not in c):
             if not self.splitter:
                 raise hlpst.PydraStateError(
-                    "splitter has to be set before setting combiner"
+                    "splitter has to be set before setting combiner with field names "
+                    f"in the current node {list(local_names)}"
+                )
+            if missing := local_names - set(self.splitter_rpn):
+                raise hlpst.PydraStateError(
+                    "all field names from the current node referenced in the combiner, "
+                    f"{list(missing)} are missing in the splitter"
                 )
-            if set(self._combiner) - set(self.splitter_rpn):
-                raise hlpst.PydraStateError("all combiners have to be in the splitter")
 
     def prepare_states(
         self,