debugging test_workflow

Author: tclose

pydra/engine/node.py

@@ -305,26 +305,6 @@ def _extract_input_el(self, inputs, inp_nm, ind):
         else:
             return getattr(inputs, inp_nm)[ind]
 
-    def _split_definition(self) -> dict[StateIndex, "TaskDef[OutputType]"]:
-        """Split the definition into the different states it will be run over"""
-        # TODO: doesn't work properly for more cmplicated wf (check if still an issue)
-        if not self.state:
-            return {None: self._definition}
-        split_defs = {}
-        for input_ind in self.state.inputs_ind:
-            inputs_dict = {}
-            for inp in set(self.input_names):
-                if f"{self.name}.{inp}" in input_ind:
-                    inputs_dict[inp] = self._extract_input_el(
-                        inputs=self._definition,
-                        inp_nm=inp,
-                        ind=input_ind[f"{self.name}.{inp}"],
-                    )
-            split_defs[StateIndex(input_ind)] = attrs.evolve(
-                self._definition, **inputs_dict
-            )
-        return split_defs
-
         # else:
         #     # todo it never gets here
         #     breakpoint()

pydra/engine/submitter.py

@@ -645,12 +645,9 @@ def _generate_tasks(self) -> ty.Iterable["Task[DefType]"]:
                 name=self.node.name,
             )
         else:
-            for index, split_defn in self.node._split_definition().items():
+            for index, split_defn in self._split_definition().items():
                 yield Task(
-                    definition=self._resolve_lazy_inputs(
-                        task_def=split_defn,
-                        state_index=index,
-                    ),
+                    definition=split_defn,
                     submitter=self.submitter,
                     environment=self.node._environment,
                     name=self.node.name,
@@ -686,6 +683,34 @@ def _resolve_lazy_inputs(
                 )
         return attrs.evolve(task_def, **resolved)
 
+    def _split_definition(self) -> dict[StateIndex, "TaskDef[OutputType]"]:
+        """Split the definition into the different states it will be run over"""
+        # TODO: doesn't work properly for more cmplicated wf (check if still an issue)
+        if not self.node.state:
+            return {None: self.node._definition}
+        split_defs = {}
+        for input_ind in self.node.state.inputs_ind:
+            inputs_dict = {}
+            for inp in set(self.node.input_names):
+                if f"{self.node.name}.{inp}" in input_ind:
+                    value = getattr(self.node._definition, inp)
+                    if isinstance(value, LazyField):
+                        inputs_dict[inp] = value._get_value(
+                            workflow=self.workflow,
+                            graph=self.graph,
+                            state_index=StateIndex(input_ind),
+                        )
+                    else:
+                        inputs_dict[inp] = self.node._extract_input_el(
+                            inputs=self.node._definition,
+                            inp_nm=inp,
+                            ind=input_ind[f"{self.node.name}.{inp}"],
+                        )
+            split_defs[StateIndex(input_ind)] = attrs.evolve(
+                self.node._definition, **inputs_dict
+            )
+        return split_defs
+
     def get_runnable_tasks(self, graph: DiGraph) -> list["Task[DefType]"]:
         """For a given node, check to see which tasks have been successfully run, are ready
         to run, can't be run due to upstream errors, or are blocked on other tasks to complete.

pydra/engine/tests/test_workflow.py

@@ -1163,7 +1163,7 @@ def Workflow(x, y):
         mult = workflow.add(Multiply(x=add2x.out, y=add2y.out))
         return mult.out
 
-    wf = Workflow.split(["x", "y"], x=[1, 2, 3], y=[11, 12]).combine("x")
+    wf = Workflow().split(["x", "y"], x=[1, 2, 3], y=[11, 12]).combine("x")
 
     with Submitter(worker=plugin, cache_dir=tmpdir) as sub:
         results = sub(wf)
@@ -1186,8 +1186,8 @@ def test_wf_3nd_ndst_2(plugin, tmpdir):
 
     @workflow.define
     def Workflow(x, y):
-        add2x = workflow.add(Add2().split("x", x=x))
-        add2y = workflow.add(Add2().split("x", x=y))
+        add2x = workflow.add(Add2().split("x", x=x), name="add2x")
+        add2y = workflow.add(Add2().split("x", x=y), name="add2y")
         mult = workflow.add(Multiply(x=add2x.out, y=add2y.out).combine("add2x.x"))
         return mult.out
 
@@ -1215,7 +1215,7 @@ def Workflow(x, y):
         mult = workflow.add(Multiply(x=add2x.out, y=add2y.out))
         return mult.out
 
-    wf = Workflow.split(["x", "y"], x=[1, 2, 3], y=[11, 12]).combine("y")
+    wf = Workflow().split(["x", "y"], x=[1, 2, 3], y=[11, 12]).combine("y")
 
     with Submitter(worker=plugin, cache_dir=tmpdir) as sub:
         results = sub(wf)
@@ -1245,7 +1245,7 @@ def Workflow(x, y):
 
     wf = Workflow(x=[1, 2, 3], y=[11, 12])
 
-    with Submitter(worker=plugin, cache_dir=tmpdir) as sub:
+    with Submitter(worker="debug", cache_dir=tmpdir) as sub:
         results = sub(wf)
 
     assert not results.errored, "\n".join(results.errors["error message"])
@@ -1268,7 +1268,7 @@ def Workflow(x, y):
         mult = workflow.add(Multiply(x=add2x.out, y=add2y.out))
         return mult.out
 
-    wf = Workflow.split(["x", "y"], x=[1, 2, 3], y=[11, 12]).combine(["x", "y"])
+    wf = Workflow().split(["x", "y"], x=[1, 2, 3], y=[11, 12]).combine(["x", "y"])
 
     with Submitter(worker=plugin, cache_dir=tmpdir) as sub:
         results = sub(wf)
@@ -1322,7 +1322,11 @@ def Workflow(x, y, z):
         addvar = workflow.add(FunAddVar3(a=add2x.out, b=add2y.out, c=z))
         return addvar.out
 
-    wf = Workflow.split(["x", "y", "z"], x=[2, 3], y=[11, 12], z=[10, 100]).combine("y")
+    wf = (
+        Workflow()
+        .split(["x", "y", "z"], x=[2, 3], y=[11, 12], z=[10, 100])
+        .combine("y")
+    )
 
     with Submitter(worker=plugin, cache_dir=tmpdir) as sub:
         results = sub(wf)
@@ -1828,7 +1832,7 @@ def test_wf_ndst_singl_1(plugin, tmpdir):
 
     @workflow.define
     def Workflow(x, y):
-        mult = workflow.add(Multiply(y=y).split("x", x=x))
+        mult = workflow.add(Multiply(y=y).split("x", x=x), name="mult")
         add2 = workflow.add(Add2(x=mult.out).combine("mult.x"))
         return add2.out
 
@@ -1855,7 +1859,7 @@ def Workflow(x, y):
         mult = workflow.add(Multiply(x=add2x.out, y=add2y.out))
         return mult.out
 
-    wf = Workflow.split("x", x=[1, 2, 3], y=11)
+    wf = Workflow().split("x", x=[1, 2, 3], y=11)
 
     with Submitter(worker=plugin, cache_dir=tmpdir) as sub:
         results = sub(wf)