fixing workflows with first/intermediate task that has a full combiner; adding tests

Author: djarecka

pydra/engine/core.py

@@ -222,7 +222,7 @@ def checksum(self):
             self.inputs._graph_checksums = [nd.checksum for nd in self.graph_sorted]
 
         input_hash = self.inputs.hash
-        if self.state is None:
+        if self.state is None or self.state.splitter_rpn == []:
             self._checksum = create_checksum(self.__class__.__name__, input_hash)
         else:
             # including splitter in the hash
@@ -329,7 +329,7 @@ def cache_locations(self, locations):
     @property
     def output_dir(self):
         """Get the filesystem path where outputs will be written."""
-        if self.state:
+        if self.state and self.state.splitter_rpn:
             return [self._cache_dir / checksum for checksum in self.checksum_states()]
         return self._cache_dir / self.checksum
 
@@ -342,7 +342,7 @@ def __call__(self, submitter=None, plugin=None, rerun=False, **kwargs):
         plugin = plugin or self.plugin
         if plugin:
             submitter = Submitter(plugin=plugin)
-        elif self.state:
+        elif self.state and self.state.splitter_rpn:
             submitter = Submitter()
 
         if submitter:
@@ -512,7 +512,7 @@ def done(self):
         # if any of the field is lazy, there is no need to check results
         if is_lazy(self.inputs):
             return False
-        if self.state:
+        if self.state and self.state.splitter_rpn:
             # TODO: only check for needed state result
             if self.result() and all(self.result()):
                 return True
@@ -556,7 +556,7 @@ def result(self, state_index=None):
         """
         # TODO: check if result is available in load_result and
         # return a future if not
-        if self.state:
+        if self.state and self.state.splitter_rpn:
             if state_index is None:
                 # if state_index=None, collecting all results
                 if self.state.combiner:

pydra/engine/submitter.py

@@ -50,7 +50,9 @@ def __call__(self, runnable, cache_locations=None, rerun=False):
             runnable.inputs._graph_checksums = [
                 nd.checksum for nd in runnable.graph_sorted
             ]
-        if is_workflow(runnable) and runnable.state is None:
+        if is_workflow(runnable) and (
+            runnable.state is None or runnable.state.splitter_rpn == []
+        ):
             self.loop.run_until_complete(self.submit_workflow(runnable, rerun=rerun))
         else:
             self.loop.run_until_complete(self.submit(runnable, wait=True, rerun=rerun))
@@ -91,7 +93,7 @@ async def submit(self, runnable, wait=False, rerun=False):
 
         """
         futures = set()
-        if runnable.state:
+        if runnable.state and runnable.state.splitter_rpn:
             runnable.state.prepare_states(runnable.inputs)
             runnable.state.prepare_inputs()
             logger.debug(
@@ -154,7 +156,9 @@ async def _run_workflow(self, wf, rerun=False):
                 task.inputs.retrieve_values(wf)
                 # checksum has to be updated, so resetting
                 task._checksum = None
-                if is_workflow(task) and not task.state:
+                if is_workflow(task) and (
+                    not task.state or task.state.splitter_rpn == []
+                ):
                     await self.submit_workflow(task, rerun=rerun)
                 else:
                     for fut in await self.submit(task, rerun=rerun):

pydra/engine/tests/test_workflow.py

@@ -757,11 +757,81 @@ def test_wf_ndst_6(plugin):
     assert wf.output_dir.exists()
 
 
+@pytest.mark.parametrize("plugin", Plugins)
+def test_wf_ndst_7(plugin):
+    """ workflow with two tasks, outer splitter and (full) combiner for first node only"""
+    wf = Workflow(name="wf_ndst_6", input_spec=["x", "y"])
+    wf.add(multiply(name="mult", x=wf.lzin.x, y=wf.lzin.y).split("x").combine("x"))
+    wf.add(identity(name="iden", x=wf.mult.lzout.out))
+    wf.inputs.x = [1, 2, 3]
+    wf.inputs.y = 11
+    wf.set_output([("out", wf.iden.lzout.out)])
+    wf.plugin = plugin
+
+    with Submitter(plugin=plugin) as sub:
+        sub(wf)
+
+    results = wf.result()
+    assert results.output.out[0] == [11, 22, 33]
+
+    # checking the output directory
+    assert wf.output_dir.exists()
+
+
+@pytest.mark.parametrize("plugin", Plugins)
+def test_wf_ndst_8(plugin):
+    """ workflow with two tasks, outer splitter and (partial) combiner for first task only"""
+    wf = Workflow(name="wf_ndst_6", input_spec=["x", "y"])
+    wf.add(
+        multiply(name="mult", x=wf.lzin.x, y=wf.lzin.y).split(["x", "y"]).combine("x")
+    )
+    wf.add(identity(name="iden", x=wf.mult.lzout.out))
+    wf.inputs.x = [1, 2, 3]
+    wf.inputs.y = [11, 12]
+    wf.set_output([("out", wf.iden.lzout.out)])
+    wf.plugin = plugin
+
+    with Submitter(plugin=plugin) as sub:
+        sub(wf)
+
+    results = wf.result()
+    assert results.output.out[0] == [11, 22, 33]
+    assert results.output.out[1] == [12, 24, 36]
+
+    # checking the output directory
+    assert wf.output_dir.exists()
+
+
+@pytest.mark.parametrize("plugin", Plugins)
+def test_wf_ndst_9(plugin):
+    """ workflow with two tasks, outer splitter and (full) combiner for first task only"""
+    wf = Workflow(name="wf_ndst_6", input_spec=["x", "y"])
+    wf.add(
+        multiply(name="mult", x=wf.lzin.x, y=wf.lzin.y)
+        .split(["x", "y"])
+        .combine(["x", "y"])
+    )
+    wf.add(identity(name="iden", x=wf.mult.lzout.out))
+    wf.inputs.x = [1, 2, 3]
+    wf.inputs.y = [11, 12]
+    wf.set_output([("out", wf.iden.lzout.out)])
+    wf.plugin = plugin
+
+    with Submitter(plugin=plugin) as sub:
+        sub(wf)
+
+    results = wf.result()
+    assert results.output.out[0] == [11, 12, 22, 24, 33, 36]
+
+    # checking the output directory
+    assert wf.output_dir.exists()
+
+
 # workflows with structures A -> C, B -> C
 
 
 @pytest.mark.parametrize("plugin", Plugins)
-def test_wf_st_7(plugin):
+def test_wf_3nd_st_1(plugin):
     """ workflow with three tasks, third one connected to two previous tasks,
         splitter on the workflow level
     """
@@ -789,7 +859,7 @@ def test_wf_st_7(plugin):
 
 
 @pytest.mark.parametrize("plugin", Plugins)
-def test_wf_ndst_7(plugin):
+def test_wf_3nd_ndst_1(plugin):
     """ workflow with three tasks, third one connected to two previous tasks,
         splitter on the tasks levels
     """
@@ -813,7 +883,7 @@ def test_wf_ndst_7(plugin):
 
 
 @pytest.mark.parametrize("plugin", Plugins)
-def test_wf_st_8(plugin):
+def test_wf_3nd_st_2(plugin):
     """ workflow with three tasks, third one connected to two previous tasks,
         splitter and partial combiner on the workflow level
     """
@@ -844,7 +914,7 @@ def test_wf_st_8(plugin):
 
 
 @pytest.mark.parametrize("plugin", Plugins)
-def test_wf_ndst_8(plugin):
+def test_wf_3nd_ndst_2(plugin):
     """ workflow with three tasks, third one connected to two previous tasks,
         splitter and partial combiner on the tasks levels
     """
@@ -873,7 +943,7 @@ def test_wf_ndst_8(plugin):
 
 
 @pytest.mark.parametrize("plugin", Plugins)
-def test_wf_st_9(plugin):
+def test_wf_3nd_st_3(plugin):
     """ workflow with three tasks, third one connected to two previous tasks,
         splitter and partial combiner (from the second task) on the workflow level
     """
@@ -904,7 +974,7 @@ def test_wf_st_9(plugin):
 
 
 @pytest.mark.parametrize("plugin", Plugins)
-def test_wf_ndst_9(plugin):
+def test_wf_3nd_ndst_3(plugin):
     """ workflow with three tasks, third one connected to two previous tasks,
         splitter and partial combiner (from the second task) on the tasks levels
     """
@@ -934,7 +1004,7 @@ def test_wf_ndst_9(plugin):
 
 
 @pytest.mark.parametrize("plugin", Plugins)
-def test_wf_st_10(plugin):
+def test_wf_3nd_st_4(plugin):
     """ workflow with three tasks, third one connected to two previous tasks,
         splitter and full combiner on the workflow level
     """
@@ -964,7 +1034,7 @@ def test_wf_st_10(plugin):
 
 
 @pytest.mark.parametrize("plugin", Plugins)
-def test_wf_ndst_10(plugin):
+def test_wf_3nd_ndst_4(plugin):
     """ workflow with three tasks, third one connected to two previous tasks,
         splitter and full combiner on the tasks levels
     """