Merge pull request #208 from djarecka/fix/usecache_nodelevel

[fix] passing wf.cache_locations and submitter.rerun

Author: satra

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,8 @@ def __init__(
         messenger_args=None,
         messengers=None,
         output_spec: ty.Optional[BaseSpec] = None,
+        rerun=False,
+        propagate_rerun=True,
         **kwargs,
     ):
         """
@@ -671,13 +676,16 @@ def __init__(
             audit_flags=audit_flags,
             messengers=messengers,
             messenger_args=messenger_args,
+            rerun=rerun,
         )
 
         self.graph = DiGraph()
         self.name2obj = {}
 
         # store output connections
         self._connections = None
+        # propagating rerun if task_rerun=True
+        self.propagate_rerun = propagate_rerun
 
     def __getattr__(self, name):
         if name == "lzin":
@@ -789,12 +797,20 @@ 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
         # creating connections that were defined after adding tasks to the wf
         for task in self.graph.nodes:
+            # if workflow has task_rerun=True and propagate_rerun=True,
+            # it should be passed to the tasks
+            if self.task_rerun and self.propagate_rerun:
+                task.task_rerun = self.task_rerun
+                # if the task is a wf, than the propagate_rerun should be also set
+                if is_workflow(task):
+                    task.propagate_rerun = self.propagate_rerun
+            task.cache_locations = task._cache_locations + self.cache_locations
             self.create_connections(task)
         # TODO add signal handler for processes killed after lock acquisition
         self.hooks.pre_run(self)
@@ -810,7 +826,7 @@ async def _run(self, submitter=None, rerun=False, **kwargs):
             self.hooks.pre_run_task(self)
             try:
                 self.audit.monitor()
-                await self._run_task(submitter)
+                await self._run_task(submitter, rerun=rerun)
                 result.output = self._collect_outputs()
             except Exception as e:
                 record_error(self.output_dir, e)
@@ -824,11 +840,11 @@ async def _run(self, submitter=None, rerun=False, **kwargs):
         self.hooks.post_run(self, result)
         return result
 
-    async def _run_task(self, submitter):
+    async def _run_task(self, submitter, rerun=False):
         if not submitter:
             raise Exception("Submitter should already be set.")
         # at this point Workflow is stateless so this should be fine
-        await submitter._run_workflow(self)
+        await submitter._run_workflow(self, rerun=rerun)
 
     def set_output(self, connections):
         """

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