backed out of trying to allow execution of downstream nodes before/when upstream nodes are partially incompleted/have errored

Author: tclose

pydra/engine/core.py

@@ -844,7 +844,8 @@ def _create_graph(
                             # treated as a containers
                             if (
                                 node.state
-                                and f"{node.name}.{field.name}" in node.state.splitter
+                                and f"{node.name}.{field.name}"
+                                in node.state._current_splitter_rpn
                             ):
                                 node.state._inner_cont_dim[
                                     f"{node.name}.{field.name}"

pydra/engine/node.py

@@ -225,15 +225,15 @@ def _get_upstream_states(self) -> dict[str, tuple["State", list[str]]]:
             ):
                 node: Node = val._node
                 # variables that are part of inner splitters should be treated as a containers
-                if node.state and f"{node.name}.{inpt_name}" in node.state.splitter:
-                    node.state._inner_cont_dim[f"{node.name}.{inpt_name}"] = 1
+                if node.state and f"{node.name}.{val._field}" in node.state.splitter:
+                    node.state._inner_cont_dim[f"{node.name}.{val._field}"] = 1
                 # adding task_name: (task.state, [a field from the connection]
                 if node.name not in upstream_states:
-                    upstream_states[node.name] = (node.state, [inpt_name])
+                    upstream_states[node.name] = (node.state, [val._field])
                 else:
                     # if the task already exist in other_state,
                     # additional field name should be added to the list of fields
-                    upstream_states[node.name][1].append(inpt_name)
+                    upstream_states[node.name][1].append(val._field)
         return upstream_states
 
         # else:

pydra/engine/state.py

@@ -885,7 +885,7 @@ def prepare_states(
             st: State
             for nm, (st, _) in self.other_states.items():
                 self.inputs.update(st.inputs)
-                self.cont_dim.update(st.cont_dim)
+                self.cont_dim.update(st.cont_dim_all)
 
         self.prepare_states_ind()
         self.prepare_states_val()
@@ -1225,11 +1225,11 @@ def _get_element(self, value: ty.Any, field_name: str, ind: int) -> ty.Any:
         Any
             specific element of the input field
         """
-        if f"{self.name}.{field_name}" in self.cont_dim:
+        if f"{self.name}.{field_name}" in self.cont_dim_all:
             return list(
                 hlpst.flatten(
                     ensure_list(value),
-                    max_depth=self.cont_dim[f"{self.name}.{field_name}"],
+                    max_depth=self.cont_dim_all[f"{self.name}.{field_name}"],
                 )
             )[ind]
         else:

pydra/engine/submitter.py

@@ -570,7 +570,9 @@ def start(self) -> None:
         if self.state:
             values = {}
             for name, value in self.node.state_values.items():
-                if name in self.node.state.names:
+                if name in self.node.state.current_splitter_rpn:
+                    if name in ("*", "."):
+                        continue
                     if isinstance(value, LazyField):
                         values[name] = value._get_value(
                             workflow=self.workflow, graph=self.graph
@@ -692,23 +694,20 @@ def _split_definition(self) -> dict[int, "TaskDef[OutputType]"]:
         if not self.node.state:
             return {None: self.node._definition}
         split_defs = []
-        for input_ind in self.node.state.inputs_ind:
+        for index, vals in zip(self.node.state.inputs_ind, self.node.state.states_val):
             resolved = {}
             for inpt_name in set(self.node.input_names):
                 value = getattr(self._definition, inpt_name)
                 state_key = f"{self.node.name}.{inpt_name}"
-                if isinstance(value, LazyField):
-                    resolved[inpt_name] = value._get_value(
-                        workflow=self.workflow,
-                        graph=self.graph,
-                        state_index=input_ind.get(state_key),
-                    )
-                elif state_key in input_ind:
-                    resolved[inpt_name] = self.node.state._get_element(
-                        value=value,
-                        field_name=inpt_name,
-                        ind=input_ind[state_key],
-                    )
+                try:
+                    resolved[inpt_name] = vals[state_key]
+                except KeyError:
+                    if isinstance(value, LazyField):
+                        resolved[inpt_name] = value._get_value(
+                            workflow=self.workflow,
+                            graph=self.graph,
+                            state_index=index.get(state_key),
+                        )
             split_defs.append(attrs.evolve(self.node._definition, **resolved))
         return split_defs
 
@@ -736,21 +735,28 @@ 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 []
-            )
+            # This is required for the commented-out code below
+            # states_ind = (
+            #     list(self.node.state.states_ind[index].items())
+            #     if self.node.state
+            #     else []
+            # )
             for pred in graph.predecessors[self.node.name]:
                 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)
-                    ]
+                    # FIXME: These should be the only predecessor jobs that are required to have
+                    # completed before the job can be run, however, due to how the state
+                    # is currently built, all predecessors are required to have completed.
+                    # If/when this is relaxed, then the following code should be used instead.
+                    #
+                    # 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)
+                    # ]
+                    pred_inds = list(range(len(pred.node.state.states_ind)))
                 else:
                     pred_inds = [None]
                 if not all(i in pred.successful for i in pred_inds):