Merge pull request #207 from djarecka/fix/wf_full_combiner

[fix] fixing workflows with first/intermediate tasks with a full combiner

Author: djarecka

pydra/engine/core.py

@@ -538,7 +538,11 @@ def _combined_output(self):
                 if result is None:
                     return None
                 combined_results[gr].append(result)
-        return combined_results
+        if len(combined_results) == 1 and self.state.splitter_rpn_final == []:
+            # in case it's full combiner, removing the nested structure
+            return combined_results[0]
+        else:
+            return combined_results
 
     def result(self, state_index=None):
         """
@@ -753,7 +757,10 @@ def create_connections(self, task):
                     self.graph.add_edges((getattr(self, val.name), task))
                     logger.debug("Connecting %s to %s", val.name, task.name)
 
-                    if getattr(self, val.name).state:
+                    if (
+                        getattr(self, val.name).state
+                        and getattr(self, val.name).state.splitter_rpn_final
+                    ):
                         # adding a state from the previous task to other_states
                         other_states[val.name] = (
                             getattr(self, val.name).state,

pydra/engine/tests/test_node_task.py

@@ -937,10 +937,10 @@ def test_task_state_comb_1(plugin):
     # checking the results
     results = nn.result()
 
-    combined_results = [[res.output.out for res in res_l] for res_l in results]
-    expected = [({}, [5, 7])]
-    for i, res in enumerate(expected):
-        assert combined_results[i] == res[1]
+    # fully combined (no nested list)
+    combined_results = [res.output.out for res in results]
+
+    assert combined_results == [5, 7]
     # checking the output_dir
     assert nn.output_dir
     for odir in nn.output_dir:
@@ -1072,11 +1072,11 @@ def test_task_state_comb_singl_1(plugin):
         sub(nn)
 
     # checking the results
-    expected = [({}, [13, 15])]
+    expected = ({}, [13, 15])
     results = nn.result()
-    combined_results = [[res.output.out for res in res_l] for res_l in results]
-    for i, res in enumerate(expected):
-        assert combined_results[i] == res[1]
+    # full combiner, no nested list
+    combined_results = [res.output.out for res in results]
+    assert combined_results == expected[1]
     # checking the output_dir
     assert nn.output_dir
     for odir in nn.output_dir:
@@ -1143,7 +1143,8 @@ def test_task_state_comb_order(plugin):
     assert nn_ab.state.combiner == ["NA.a", "NA.b"]
 
     results_ab = nn_ab()
-    combined_results_ab = [res.output.out for res in results_ab[0]]
+    # full combiner, no nested list
+    combined_results_ab = [res.output.out for res in results_ab]
     assert combined_results_ab == [13, 15, 23, 25]
 
     # combiner with both fields ["b", "a"] - will create the same list as nn_ab
@@ -1156,7 +1157,7 @@ def test_task_state_comb_order(plugin):
     assert nn_ba.state.combiner == ["NA.b", "NA.a"]
 
     results_ba = nn_ba()
-    combined_results_ba = [res.output.out for res in results_ba[0]]
+    combined_results_ba = [res.output.out for res in results_ba]
     assert combined_results_ba == [13, 15, 23, 25]
 
 

pydra/engine/tests/test_shelltask.py

@@ -174,8 +174,8 @@ def test_shell_cmd_5(plugin):
     assert shelly.cmdline == ["echo nipype", "echo pydra"]
     res = shelly(plugin=plugin)
 
-    assert res[0][0].output.stdout == "nipype\n"
-    assert res[0][1].output.stdout == "pydra\n"
+    assert res[0].output.stdout == "nipype\n"
+    assert res[1].output.stdout == "pydra\n"
 
 
 @pytest.mark.parametrize("plugin", Plugins)

pydra/engine/tests/test_workflow.py

@@ -532,9 +532,9 @@ def test_wf_st_2(plugin):
         sub(wf)
 
     results = wf.result()
-    # expected: [[({"test7.x": 1}, 3), ({"test7.x": 2}, 4)]]
-    assert results[0][0].output.out == 3
-    assert results[0][1].output.out == 4
+    # expected: [({"test7.x": 1}, 3), ({"test7.x": 2}, 4)]
+    assert results[0].output.out == 3
+    assert results[1].output.out == 4
     # checking all directories
     assert wf.output_dir
     for odir in wf.output_dir:
@@ -554,8 +554,8 @@ def test_wf_ndst_2(plugin):
         sub(wf)
 
     results = wf.result()
-    # expected: [[({"test7.x": 1}, 3), ({"test7.x": 2}, 4)]]
-    assert results.output.out[0] == [3, 4]
+    # expected: [({"test7.x": 1}, 3), ({"test7.x": 2}, 4)]
+    assert results.output.out == [3, 4]
     assert wf.output_dir.exists()
 
 
@@ -625,10 +625,10 @@ def test_wf_st_4(plugin):
 
     results = wf.result()
     # expected: [
-    #     [({"test7.x": 1, "test7.y": 11}, 13), ({"test7.x": 2, "test.y": 12}, 26)]
+    #     ({"test7.x": 1, "test7.y": 11}, 13), ({"test7.x": 2, "test.y": 12}, 26)
     # ]
-    assert results[0][0].output.out == 13
-    assert results[0][1].output.out == 26
+    assert results[0].output.out == 13
+    assert results[1].output.out == 26
     # checking all directories
     assert wf.output_dir
     for odir in wf.output_dir:
@@ -652,9 +652,9 @@ def test_wf_ndst_4(plugin):
 
     results = wf.result()
     # expected: [
-    #     [({"test7.x": 1, "test7.y": 11}, 13), ({"test7.x": 2, "test.y": 12}, 26)]
+    #     ({"test7.x": 1, "test7.y": 11}, 13), ({"test7.x": 2, "test.y": 12}, 26)
     # ]
-    assert results.output.out[0] == [13, 26]
+    assert results.output.out == [13, 26]
     # checking the output directory
     assert wf.output_dir.exists()
 
@@ -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 == [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 == [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
     """
@@ -950,21 +1020,21 @@ def test_wf_st_10(plugin):
         sub(wf)
 
     results = wf.result()
-    assert len(results) == 1
-    assert results[0][0].output.out == 39
-    assert results[0][1].output.out == 42
-    assert results[0][2].output.out == 52
-    assert results[0][3].output.out == 56
-    assert results[0][4].output.out == 65
-    assert results[0][5].output.out == 70
+    assert len(results) == 6
+    assert results[0].output.out == 39
+    assert results[1].output.out == 42
+    assert results[2].output.out == 52
+    assert results[3].output.out == 56
+    assert results[4].output.out == 65
+    assert results[5].output.out == 70
     # checking all directories
     assert wf.output_dir
     for odir in wf.output_dir:
         assert odir.exists()
 
 
 @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
     """
@@ -986,8 +1056,8 @@ def test_wf_ndst_10(plugin):
     # assert wf.output_dir.exists()
     results = wf.result()
 
-    assert len(results.output.out) == 1
-    assert results.output.out == [[39, 42, 52, 56, 65, 70]]
+    assert len(results.output.out) == 6
+    assert results.output.out == [39, 42, 52, 56, 65, 70]
     # checking the output directory
     assert wf.output_dir.exists()
 
@@ -1284,8 +1354,8 @@ def test_wf_st_singl_1(plugin):
         sub(wf)
 
     results = wf.result()
-    assert results[0][0].output.out == 13
-    assert results[0][1].output.out == 24
+    assert results[0].output.out == 13
+    assert results[1].output.out == 24
     # checking all directories
     assert wf.output_dir
     for odir in wf.output_dir:
@@ -1309,7 +1379,7 @@ def test_wf_ndst_singl_1(plugin):
         sub(wf)
 
     results = wf.result()
-    assert results.output.out[0] == [13, 24]
+    assert results.output.out == [13, 24]
     # checking the output directory
     assert wf.output_dir.exists()
 
@@ -2665,7 +2735,7 @@ def test_workflow_combine1(tmpdir):
 
     assert result.output.out_pow == [1, 1, 4, 8]
     assert result.output.out_iden1 == [[1, 4], [1, 8]]
-    assert result.output.out_iden2 == [[[1, 4], [1, 8]]]
+    assert result.output.out_iden2 == [[1, 4], [1, 8]]
 
 
 def test_workflow_combine2(tmpdir):
@@ -2679,7 +2749,7 @@ def test_workflow_combine2(tmpdir):
     result = wf1(plugin="cf")
 
     assert result.output.out_pow == [[1, 4], [1, 8]]
-    assert result.output.out_iden == [[[1, 4], [1, 8]]]
+    assert result.output.out_iden == [[1, 4], [1, 8]]
 
 
 # testing lzout.all to collect all of the results and let FunctionTask deal with it