debugging states and workflows

Author: tclose

pydra/conftest.py

@@ -20,7 +20,7 @@ def pytest_generate_tests(metafunc):
         if bool(shutil.which("sbatch")):
             Plugins = ["slurm"]
         else:
-            Plugins = ["cf"]
+            Plugins = ["debug"]  # ["debug", "cf"]
         try:
             if metafunc.config.getoption("dask"):
                 Plugins.append("dask")
@@ -50,7 +50,7 @@ def pytest_generate_tests(metafunc):
         elif bool(shutil.which("sbatch")):
             Plugins = ["slurm"]
         else:
-            Plugins = ["cf"]
+            Plugins = ["debug"]  # ["debug", "cf"]
         try:
             if metafunc.config.getoption("psij"):
                 Plugins.append("psij-" + metafunc.config.getoption("psij"))

pydra/design/tests/test_python.py

@@ -11,7 +11,7 @@
 sort_key = attrgetter("name")
 
 
-def test_interface_wrap_function():
+def test_interface_wrap_function(tmp_path):
     def func(a: int) -> float:
         """Sample function with inputs and outputs"""
         return a * 2
@@ -27,7 +27,7 @@ def func(a: int) -> float:
     ]
     assert outputs == [python.out(name="out", type=float)]
     definition = SampleDef(a=1)
-    outputs = definition()
+    outputs = definition(cache_dir=tmp_path)
     assert outputs.out == 2.0
     with pytest.raises(TypeError):
         SampleDef(a=1.5)

pydra/engine/audit.py

@@ -30,7 +30,7 @@ def __init__(self, audit_flags, messengers, messenger_args, develop=None):
             Base configuration of auditing.
         messengers : :obj:`pydra.util.messenger.Messenger`
             or list of :class:`pydra.util.messenger.Messenger`, optional
-            Defify types of messenger used by Audit to send a message.
+            Taskify types of messenger used by Audit to send a message.
             Could be `PrintMessenger`, `FileMessenger`, or `RemoteRESTMessenger`.
         messenger_args : :obj:`dict`, optional
             Optional arguments for the `Messenger.send` method.

pydra/engine/core.py

@@ -792,9 +792,8 @@ def execution_graph(self, submitter: "Submitter") -> DiGraph:
             node.graph = graph
         return graph
 
-    @property
-    def graph(self) -> DiGraph:
-        return self._create_graph(self.nodes, detailed=True)
+    def graph(self, detailed: bool = False) -> DiGraph:
+        return self._create_graph(self.nodes, detailed=detailed)
 
     def _create_graph(
         self, nodes: "list[Node | NodeExecution]", detailed: bool = False

pydra/engine/graph.py

@@ -81,8 +81,8 @@ def nodes(self) -> list[NodeType]:
     def nodes(self, nodes: ty.Iterable[NodeType]) -> None:
         if nodes:
             nodes = ensure_list(nodes)
-            if len(set(nodes)) != len(nodes):
-                raise Exception("nodes have repeated elements")
+            # if len(set(nodes)) != len(nodes):
+            #     raise Exception("nodes have repeated elements")
             self._nodes = nodes
 
     def node(self, name: str) -> NodeType:

pydra/engine/helpers.py

@@ -47,26 +47,17 @@ def plot_workflow(
 
     # Construct the workflow object
     wf = Workflow.construct(workflow_task)
-    graph = wf.graph
+
     if not name:
-        name = f"graph_{wf._node.name}"
+        name = f"graph_{type(workflow_task).__name__}"
     if type == "simple":
-        for task in graph.nodes:
-            wf.create_connections(task)
+        graph = wf.graph()
         dotfile = graph.create_dotfile_simple(outdir=out_dir, name=name)
     elif type == "nested":
-        for task in graph.nodes:
-            wf.create_connections(task)
+        graph = wf.graph()
         dotfile = graph.create_dotfile_nested(outdir=out_dir, name=name)
     elif type == "detailed":
-        # create connections with detailed=True
-        for task in graph.nodes:
-            wf.create_connections(task, detailed=True)
-        # adding wf outputs
-        for wf_out, lf in wf._connections:
-            graph.add_edges_description(
-                (wf._node.name, wf_out, lf._node.name, lf.field)
-            )
+        graph = wf.graph(detailed=True)
         dotfile = graph.create_dotfile_detailed(outdir=out_dir, name=name)
     else:
         raise Exception(

pydra/engine/lazy.py

@@ -151,7 +151,7 @@ def _get_value(
         value : Any
             the resolved value of the lazy-field
         """
-
+        state = self._node.state
         jobs = graph.node(self._node.name).get_jobs(state_index)
 
         def retrieve_from_job(job: "Task[DefType]") -> ty.Any:
@@ -184,25 +184,27 @@ def retrieve_from_job(job: "Task[DefType]") -> ty.Any:
             val = self._apply_cast(val)
             return val
 
-        if not isinstance(jobs, StateArray):
+        if not isinstance(jobs, StateArray):  # single job
             return retrieve_from_job(jobs)
-        elif not self._node.state or not self._node.state.depth(before_combine=True):
+        elif not state or not state.depth(before_combine=True):
             assert len(jobs) == 1
             return retrieve_from_job(jobs[0])
-        elif not self._node.state.keys_final:  # all states are combined over
-            return [retrieve_from_job(j) for j in jobs]
-        elif self._node.state.combiner:
-            sorted_values = {
-                frozenset(i.items()): [] for i in self._node.state.states_ind_final
-            }
-            assert len(jobs) == len(self._node.state.inputs_ind)
-            for ind, job in zip(self._node.state.inputs_ind, jobs):
-                sorted_values[
-                    frozenset((key, ind[key]) for key in self._node.state.keys_final)
-                ].append(retrieve_from_job(job))
-            return StateArray(sorted_values.values())
-        else:
-            return StateArray(retrieve_from_job(j) for j in jobs)
+        # elif state.combiner and state.keys_final:
+        #     # We initialise it here rather than using a defaultdict to ensure the order
+        #     # of the keys matches how it is defined in the state so we can return the
+        #     # values in the correct order
+        #     sorted_values = {frozenset(i.items()): [] for i in state.states_ind_final}
+        #     # Iterate through the jobs and append the values to the correct final state
+        #     # key
+        #     for job in jobs:
+        #         state_key = frozenset(
+        #             (key, state.states_ind[job.state_index][key])
+        #             for key in state.keys_final
+        #         )
+        #         sorted_values[state_key].append(retrieve_from_job(job))
+        #     return StateArray(sorted_values.values())
+        # else:
+        return [retrieve_from_job(j) for j in jobs]
 
     @property
     def _source(self):

pydra/engine/node.py

@@ -45,12 +45,6 @@ class Node(ty.Generic[OutputType]):
         init=False, default=None, eq=False, hash=False, repr=False
     )
     _state: State | None = attrs.field(init=False, default=NOT_SET)
-    # _cont_dim: dict[str, int] | None = attrs.field(
-    #     init=False, default=None
-    # )  # QUESTION: should this be included in the state?
-    # _inner_cont_dim: dict[str, int] = attrs.field(
-    #     init=False, factory=dict
-    # )  # QUESTION: should this be included in the state?
 
     def __attrs_post_init__(self):
         self._set_state()

pydra/engine/submitter.py

@@ -569,39 +569,39 @@ def tasks(self) -> ty.Generator["Task[DefType]", None, None]:
             self._tasks = {t.state_index: t for t in self._generate_tasks()}
         return self._tasks.values()
 
-    def get_jobs(
-        self, index: int | None = None, as_array: bool = False
-    ) -> "Task | StateArray[Task]":
+    def get_jobs(self, final_index: int | None = None) -> "Task | StateArray[Task]":
         """Get the jobs that match a given state index.
 
         Parameters
         ----------
-        index : int, optional
-            The index of the state of the task to get, by default None
-        as_array : bool, optional
-            Whether to return the tasks in a state-array object, by default if the index
-            matches
+        final_index : int, optional
+            The index of the output state array (i.e. after any combinations) of the
+            job to get, by default None
 
         Returns
         -------
         matching : Task | StateArray[Task]
             The task or tasks that match the given index
         """
-        matching = StateArray()
-        if self.tasks:
-            try:
-                task = self._tasks[index]
-            except KeyError:
-                if index is None:
-                    return StateArray(self._tasks.values())
-                # Select matching tasks and return them in nested state-array objects
-                for ind, task in self._tasks.items():
-                    matching.append(task)
-            else:
-                if not as_array:
-                    return task
-                matching.append(task)
-        return matching
+        if not self.tasks:  # No jobs, return empty state array
+            return StateArray()
+        if not self.node.state:  # Return the singular job
+            assert final_index is None
+            task = self._tasks[None]
+            return task
+        if final_index is None:  # return all jobs in a state array
+            return StateArray(self._tasks.values())
+        if not self.node.state.combiner:  # Select the job that matches the index
+            task = self._tasks[final_index]
+            return task
+        # Get a slice of the tasks that match the given index of the state array of the
+        # combined values
+        final_index = set(self.node.state.states_ind_final[final_index].items())
+        return StateArray(
+            self._tasks[i]
+            for i, ind in enumerate(self.node.state.states_ind)
+            if set(ind.items()).issuperset(final_index)
+        )
 
     @property
     def started(self) -> bool:
@@ -762,9 +762,23 @@ def get_runnable_tasks(self, graph: DiGraph) -> list["Task[DefType]"]:
         for index, task in list(self.blocked.items()):
             pred: NodeExecution
             is_runnable = True
+            states_ind = (
+                list(self.node.state.states_ind[index].items())
+                if self.node.state
+                else []
+            )
             for pred in graph.predecessors[self.node.name]:
-                pred_jobs: StateArray[Task] = pred.get_jobs(index, as_array=True)
-                pred_inds = [j.state_index for j in pred_jobs]
+                if pred.node.state:
+                    pred_states_ind = {
+                        (k, i) for k, i in states_ind if k.startswith(pred.name + ".")
+                    }
+                    pred_inds = [
+                        i
+                        for i, ind in enumerate(pred.node.state.states_ind)
+                        if set(ind.items()).issuperset(pred_states_ind)
+                    ]
+                else:
+                    pred_inds = [None]
                 if not all(i in pred.successful for i in pred_inds):
                     is_runnable = False
                     blocked = True