adding task_rerun to the task attributes; for wf it doesn't propagate to all the nodes, but it starts rerunning the workflow

Author: djarecka

pydra/engine/core.py

@@ -84,6 +84,7 @@ def __init__(
         inputs: ty.Union[ty.Text, File, ty.Dict, None] = None,
         messenger_args=None,
         messengers=None,
+        rerun=False,
     ):
         """
         Initialize a task.
@@ -173,6 +174,8 @@ def __init__(
         self.cache_locations = cache_locations
         self.allow_cache_override = True
         self._checksum = None
+        # if True the results are not checked (does not propagate to nodes)
+        self.task_rerun = rerun
 
         self.plugin = None
         self.hooks = TaskHook()
@@ -365,7 +368,7 @@ def _run(self, rerun=False, **kwargs):
         self.hooks.pre_run(self)
         # TODO add signal handler for processes killed after lock acquisition
         with SoftFileLock(lockfile):
-            if not rerun:
+            if not (rerun or self.task_rerun):
                 result = self.result()
                 if result is not None:
                     return result
@@ -610,6 +613,7 @@ def __init__(
         messenger_args=None,
         messengers=None,
         output_spec: ty.Optional[BaseSpec] = None,
+        rerun=False,
         **kwargs,
     ):
         """
@@ -671,6 +675,7 @@ def __init__(
             audit_flags=audit_flags,
             messengers=messengers,
             messenger_args=messenger_args,
+            rerun=rerun,
         )
 
         self.graph = DiGraph()
@@ -789,7 +794,7 @@ async def _run(self, submitter=None, rerun=False, **kwargs):
         checksum = self.checksum
         lockfile = self.cache_dir / (checksum + ".lock")
         # Eagerly retrieve cached
-        if not rerun:
+        if not (rerun or self.task_rerun):
             result = self.result()
             if result is not None:
                 return result

pydra/engine/task.py

@@ -75,6 +75,7 @@ def __init__(
         messengers=None,
         name=None,
         output_spec: ty.Optional[BaseSpec] = None,
+        rerun=False,
         **kwargs,
     ):
         """
@@ -143,6 +144,7 @@ def __init__(
             messenger_args=messenger_args,
             cache_dir=cache_dir,
             cache_locations=cache_locations,
+            rerun=rerun,
         )
         if output_spec is None:
             if "return" not in func.__annotations__:
@@ -242,6 +244,7 @@ def __init__(
         messengers=None,
         name=None,
         output_spec: ty.Optional[SpecInfo] = None,
+        rerun=False,
         strip=False,
         **kwargs,
     ):
@@ -283,6 +286,7 @@ def __init__(
             messengers=messengers,
             messenger_args=messenger_args,
             cache_dir=cache_dir,
+            rerun=rerun,
         )
         self.strip = strip
 
@@ -412,6 +416,7 @@ def __init__(
         messengers=None,
         output_cpath="/output_pydra",
         output_spec: ty.Optional[SpecInfo] = None,
+        rerun=False,
         strip=False,
         **kwargs,
     ):
@@ -452,6 +457,7 @@ def __init__(
             messenger_args=messenger_args,
             cache_dir=cache_dir,
             strip=strip,
+            rerun=rerun,
             **kwargs,
         )
 
@@ -523,6 +529,7 @@ def __init__(
         messengers=None,
         output_cpath="/output_pydra",
         output_spec: ty.Optional[SpecInfo] = None,
+        rerun=False,
         strip=False,
         **kwargs,
     ):
@@ -565,6 +572,7 @@ def __init__(
                 cache_dir=cache_dir,
                 strip=strip,
                 output_cpath=output_cpath,
+                rerun=rerun,
                 **kwargs,
             )
             self.inputs.container_xargs = ["--rm"]
@@ -619,6 +627,7 @@ def __init__(
         messenger_args=None,
         messengers=None,
         output_spec: ty.Optional[SpecInfo] = None,
+        rerun=False,
         strip=False,
         **kwargs,
     ):
@@ -658,6 +667,7 @@ def __init__(
                 messenger_args=messenger_args,
                 cache_dir=cache_dir,
                 strip=strip,
+                rerun=rerun,
                 **kwargs,
             )
             self.init = True

pydra/engine/tests/test_workflow.py

@@ -1988,6 +1988,126 @@ def test_wf_nostate_cachelocations_forcererun(plugin, tmpdir):
     assert wf2.output_dir.exists()
 
 
+@pytest.mark.parametrize("plugin", Plugins)
+def test_wf_nostate_cachelocations_wftaskrerun(plugin, tmpdir):
+    """
+    Two identical wfs with provided cache_dir;
+    the second wf has cache_locations and rerun=True,
+    so the workflow is rerun, but it doesn't propagate to nodes,
+    so none of the node has to be recalculated
+    """
+    cache_dir1 = tmpdir.mkdir("test_wf_cache3")
+    cache_dir2 = tmpdir.mkdir("test_wf_cache4")
+
+    wf1 = Workflow(name="wf", input_spec=["x", "y"], cache_dir=cache_dir1)
+    wf1.add(multiply(name="mult", x=wf1.lzin.x, y=wf1.lzin.y))
+    wf1.add(add2_wait(name="add2", x=wf1.mult.lzout.out))
+    wf1.set_output([("out", wf1.add2.lzout.out)])
+    wf1.inputs.x = 2
+    wf1.inputs.y = 3
+    wf1.plugin = plugin
+
+    t0 = time.time()
+    with Submitter(plugin=plugin) as sub:
+        sub(wf1)
+    t1 = time.time() - t0
+
+    results1 = wf1.result()
+    assert 8 == results1.output.out
+
+    wf2 = Workflow(
+        name="wf",
+        input_spec=["x", "y"],
+        cache_dir=cache_dir2,
+        cache_locations=cache_dir1,
+        rerun=True,
+    )
+    wf2.add(multiply(name="mult", x=wf2.lzin.x, y=wf2.lzin.y))
+    wf2.add(add2_wait(name="add2", x=wf2.mult.lzout.out))
+    wf2.set_output([("out", wf2.add2.lzout.out)])
+    wf2.inputs.x = 2
+    wf2.inputs.y = 3
+    wf2.plugin = plugin
+
+    t0 = time.time()
+    with Submitter(plugin=plugin) as sub:
+        sub(wf2)
+    t2 = time.time() - t0
+
+    results2 = wf2.result()
+    assert 8 == results2.output.out
+
+    # checking if the second wf runs again
+    assert wf1.output_dir.exists()
+    assert wf2.output_dir.exists()
+
+    # even if the second wf is recomputed the nodes are not, so it's fast
+    assert len(list(Path(cache_dir1).glob("F*"))) == 2
+    assert len(list(Path(cache_dir2).glob("F*"))) == 0
+    assert t1 > 3
+    assert t2 < 1
+
+
+@pytest.mark.parametrize("plugin", Plugins)
+def test_wf_nostate_cachelocations_taskrerun(plugin, tmpdir):
+    """
+    Two identical wfs with provided cache_dir;
+    the second wf has cache_locations and rerun=True
+    and one of the task also has rerun=True;
+    so the workflow is rerun, and one of the node also has to be rerun,
+    """
+    cache_dir1 = tmpdir.mkdir("test_wf_cache3")
+    cache_dir2 = tmpdir.mkdir("test_wf_cache4")
+
+    wf1 = Workflow(name="wf", input_spec=["x", "y"], cache_dir=cache_dir1)
+    wf1.add(multiply(name="mult", x=wf1.lzin.x, y=wf1.lzin.y))
+    wf1.add(add2_wait(name="add2", x=wf1.mult.lzout.out))
+    wf1.set_output([("out", wf1.add2.lzout.out)])
+    wf1.inputs.x = 2
+    wf1.inputs.y = 3
+    wf1.plugin = plugin
+
+    t0 = time.time()
+    with Submitter(plugin=plugin) as sub:
+        sub(wf1)
+    t1 = time.time() - t0
+
+    results1 = wf1.result()
+    assert 8 == results1.output.out
+
+    wf2 = Workflow(
+        name="wf",
+        input_spec=["x", "y"],
+        cache_dir=cache_dir2,
+        cache_locations=cache_dir1,
+        rerun=True,
+    )
+    wf2.add(multiply(name="mult", x=wf2.lzin.x, y=wf2.lzin.y))
+    wf2.add(add2_wait(name="add2", x=wf2.mult.lzout.out, rerun=True))
+    wf2.set_output([("out", wf2.add2.lzout.out)])
+    wf2.inputs.x = 2
+    wf2.inputs.y = 3
+    wf2.plugin = plugin
+
+    t0 = time.time()
+    with Submitter(plugin=plugin) as sub:
+        sub(wf2)
+    t2 = time.time() - t0
+
+    results2 = wf2.result()
+    assert 8 == results2.output.out
+
+    # checking if the second wf runs again
+    assert wf1.output_dir.exists()
+    assert wf2.output_dir.exists()
+
+    # the second task also has to be recomputed this time
+    assert len(list(Path(cache_dir1).glob("F*"))) == 2
+    assert len(list(Path(cache_dir2).glob("F*"))) == 1
+    assert t1 > 3
+    assert t2 > 3
+
+
 @pytest.mark.parametrize("plugin", Plugins)
 def test_wf_nostate_nodecachelocations(plugin, tmpdir):
     """
@@ -2408,7 +2528,7 @@ def test_wf_nostate_cachelocations_recompute(plugin, tmpdir):
 
     # checking execution time (second task shouldn't be recompute, t2 should be small)
     assert t1 > 3
-    assert t2 < 0.5
+    assert t2 < 1
 
 
 @pytest.mark.parametrize("plugin", Plugins)