attempting to debug state splitting/combining logic

Author: tclose

pydra/engine/lazy.py

@@ -1,6 +1,7 @@
 import typing as ty
 import abc
 import attrs
+from typing import Self
 from pydra.utils.typing import StateArray
 from pydra.utils.hash import hash_single
 from . import node
@@ -68,11 +69,29 @@ def _get_value(
         """
         raise NotImplementedError("LazyField is an abstract class")
 
-    def split(self) -> "LazyField":
+    def split(self) -> Self:
         """ "Splits" the lazy field over an array of nodes by replacing the sequence type
         of the lazy field with StateArray to signify that it will be "split" across
         """
-        raise NotImplementedError("LazyField is an abstract class")
+        from ..utils.typing import TypeParser  # pylint: disable=import-outside-toplevel
+
+        # Modify the type of the lazy field to include the split across a state-array
+        inner_type, prev_split_depth = TypeParser.strip_splits(self._type)
+        assert prev_split_depth <= 1
+        if inner_type is ty.Any:
+            type_ = StateArray[ty.Any]
+        elif TypeParser.matches_type(inner_type, list):
+            item_type = TypeParser.get_item_type(inner_type)
+            type_ = StateArray[item_type]
+        else:
+            raise TypeError(
+                f"Cannot split non-sequence field {self}  of type {inner_type}"
+            )
+        if prev_split_depth:
+            type_ = StateArray[
+                type_
+            ]  # FIXME: This nesting of StateArray is probably unnecessary
+        return attrs.evolve(self, type=type_)
 
 
 @attrs.define(kw_only=True)
@@ -122,25 +141,6 @@ def _get_value(
         value = self._apply_cast(value)
         return value
 
-    def split(self) -> "LazyField":
-        """ "Splits" the lazy field over an array of nodes by replacing the sequence type
-        of the lazy field with StateArray to signify that it will be "split" across
-        """
-        from ..utils.typing import TypeParser  # pylint: disable=import-outside-toplevel
-
-        assert not isinstance(self, LazyInField)
-
-        if not TypeParser.matches_type(self.type, list):
-            raise TypeError(
-                f"Cannot split non-sequence field {self}  of type {self.type}"
-            )
-
-        return type(self)(
-            name=self.name,
-            field=self.field,
-            type=StateArray[TypeParser.get_item_type(self.type)],
-        )
-
 
 @attrs.define(kw_only=True)
 class LazyOutField(LazyField[T]):
@@ -214,51 +214,8 @@ def retrieve_from_job(job: "Task[DefType]") -> ty.Any:
         elif not state or not state.depth(before_combine=True):
             assert len(jobs) == 1
             return retrieve_from_job(jobs[0])
-        # 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]
 
-    def split(self) -> "LazyField":
-        """ "Splits" the lazy field over an array of nodes by replacing the sequence type
-        of the lazy field with StateArray to signify that it will be "split" across
-        """
-        from ..utils.typing import TypeParser  # pylint: disable=import-outside-toplevel
-
-        # Modify the type of the lazy field to include the split across a state-array
-        inner_type, prev_split_depth = TypeParser.strip_splits(self.type)
-        assert prev_split_depth <= 1
-        if inner_type is ty.Any:
-            type_ = StateArray[ty.Any]
-        elif TypeParser.matches_type(inner_type, list):
-            item_type = TypeParser.get_item_type(inner_type)
-            type_ = StateArray[item_type]
-        else:
-            raise TypeError(
-                f"Cannot split non-sequence field {self}  of type {inner_type}"
-            )
-        if prev_split_depth:
-            type_ = StateArray[
-                type_
-            ]  # FIXME: This nesting of StateArray is probably unnecessary
-        return type(self)[type_](
-            name=self.name,
-            field=self.field,
-            type=type_,
-        )
-
     @property
     def _source(self):
         return self._node

pydra/engine/specs.py

@@ -345,7 +345,7 @@ def split(
         split_inputs = {}
         for name, value in inputs.items():
             if isinstance(value, lazy.LazyField):
-                split_val = value.split(splitter)
+                split_val = value.split()
             elif isinstance(value, ty.Iterable) and not isinstance(
                 value, (ty.Mapping, str)
             ):

pydra/engine/submitter.py

@@ -536,17 +536,14 @@ def __init__(
         self.queued = {}
         self.running = {}  # Not used in logic, but may be useful for progress tracking
         self.unrunnable = defaultdict(list)
-        # Prepare the state to be run
-        if node.state:
-            self.state = node.state
-            self.state.prepare_states(self.node.state_values)
-            self.state.prepare_inputs()
-        else:
-            self.state = None
         self.state_names = self.node.state.names if self.node.state else []
         self.workflow = workflow
         self.graph = None
 
+    @property
+    def state(self):
+        return self.node.state
+
     def __repr__(self):
         return (
             f"NodeExecution(name={self.name!r}, blocked={list(self.blocked)}, "
@@ -566,7 +563,7 @@ def _definition(self) -> "Node":
     @property
     def tasks(self) -> ty.Generator["Task[DefType]", None, None]:
         if self._tasks is None:
-            self._tasks = {t.state_index: t for t in self._generate_tasks()}
+            raise RuntimeError("Tasks have not been generated")
         return self._tasks.values()
 
     def get_jobs(self, final_index: int | None = None) -> "Task | StateArray[Task]":
@@ -586,14 +583,11 @@ def get_jobs(self, final_index: int | None = None) -> "Task | StateArray[Task]":
         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
+            return self._tasks[None]
         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
+            return self._tasks[final_index]
         # 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())
@@ -603,6 +597,38 @@ def get_jobs(self, final_index: int | None = None) -> "Task | StateArray[Task]":
             if set(ind.items()).issuperset(final_index)
         )
 
+    def start(self) -> None:
+        """Prepare the execution node so that it can be processed"""
+        self._tasks = {}
+        if self.state:
+            values = {}
+            for name, value in self.node.state_values.items():
+                if name in self.node.state.names and isinstance(value, LazyField):
+                    values[name] = value._get_value(
+                        workflow=self.workflow, graph=self.graph
+                    )
+            self.state.prepare_states(values)
+            self.state.prepare_inputs()
+            # Generate the tasks
+            for index, split_defn in enumerate(self._split_definition()):
+                self._tasks[index] = Task(
+                    definition=split_defn,
+                    submitter=self.submitter,
+                    environment=self.node._environment,
+                    name=self.node.name,
+                    hooks=self.node._hooks,
+                    state_index=index,
+                )
+        else:
+            self._tasks[None] = Task(
+                definition=self._resolve_lazy_inputs(task_def=self.node._definition),
+                submitter=self.submitter,
+                environment=self.node._environment,
+                hooks=self.node._hooks,
+                name=self.node.name,
+            )
+        self.blocked = copy(self._tasks)
+
     @property
     def started(self) -> bool:
         return (
@@ -656,26 +682,6 @@ def all_failed(self) -> bool:
             self.successful or self.blocked or self.queued
         )
 
-    def _generate_tasks(self) -> ty.Iterable["Task[DefType]"]:
-        if not self.node.state:
-            yield Task(
-                definition=self._resolve_lazy_inputs(task_def=self.node._definition),
-                submitter=self.submitter,
-                environment=self.node._environment,
-                hooks=self.node._hooks,
-                name=self.node.name,
-            )
-        else:
-            for index, split_defn in enumerate(self._split_definition()):
-                yield Task(
-                    definition=split_defn,
-                    submitter=self.submitter,
-                    environment=self.node._environment,
-                    name=self.node.name,
-                    hooks=self.node._hooks,
-                    state_index=index,
-                )
-
     def _resolve_lazy_inputs(
         self,
         task_def: "TaskDef",
@@ -754,10 +760,8 @@ def get_runnable_tasks(self, graph: DiGraph) -> list["Task[DefType]"]:
             List of tasks that are ready to run
         """
         runnable: list["Task[DefType]"] = []
-        self.tasks  # Ensure tasks are loaded
         if not self.started:
-            assert self._tasks is not None
-            self.blocked = copy(self._tasks)
+            self.start()
         # Check to see if any blocked tasks are now runnable/unrunnable
         for index, task in list(self.blocked.items()):
             pred: NodeExecution