handle partially written result files and lockfile getting deleted before they are detected

Author: tclose

pydra/engine/job.py

@@ -85,6 +85,7 @@ class Job(ty.Generic[TaskType]):
     bindings: dict[str, ty.Any] | None = None  # Bindings for the job environment
 
     _inputs: dict[str, ty.Any] | None = None
+    _run_start_time: datetime | None
 
     def __init__(
         self,
@@ -145,6 +146,7 @@ def __init__(
         self.audit = submitter.audit
         self.cache_dir = submitter.cache_dir
         self.cache_locations = submitter.cache_locations
+        self._run_start_time = None
 
     @property
     def cache_dir(self):
@@ -430,9 +432,12 @@ def done(self):
     @property
     def run_start_time(self) -> datetime | None:
         """Check whether the job is currently running."""
+        if self._run_start_time is not None:
+            return self._run_start_time
         if not self.lockfile.exists():
             return None
-        return datetime.fromtimestamp(self.lockfile.stat().st_ctime)
+        self._run_start_time = datetime.fromtimestamp(self.lockfile.stat().st_ctime)
+        return self._run_start_time
 
     def _combined_output(self, return_inputs=False):
         combined_results = []

pydra/engine/result.py

@@ -3,6 +3,8 @@
 from pathlib import Path
 import typing as ty
 import attrs
+import pickle
+import time
 import os
 import cloudpickle as cp
 import getpass
@@ -127,7 +129,12 @@ class RuntimeSpec:
     network: bool = False
 
 
-def load_result(checksum, cache_locations):
+def load_result(
+    checksum: str,
+    cache_locations: list[Path],
+    retries: int = 10,
+    polling_interval: float = 0.1,
+) -> Result | None:
     """
     Restore a result from the cache.
 
@@ -139,6 +146,11 @@ def load_result(checksum, cache_locations):
         List of cache directories, in order of priority, where
         the checksum will be looked for.
 
+    Returns
+    -------
+    result : :obj:`Result` | None
+        The result object if found, otherwise None.
+
     """
     if not cache_locations:
         return None
@@ -148,8 +160,16 @@ def load_result(checksum, cache_locations):
         if (location / checksum).exists():
             result_file = location / checksum / "_result.pklz"
             if result_file.exists() and result_file.stat().st_size > 0:
-                with open(result_file, "rb") as fp:
-                    return cp.load(fp)
+                # Load the result file, retrying if necessary while waiting for the file
+                # to be written completely.
+                for _ in range(retries):
+                    try:
+                        with open(result_file, "rb") as fp:
+                            return cp.load(fp)
+                    except (pickle.UnpicklingError, EOFError):
+                        # if the file is not finished writing
+                        # wait and retry
+                        time.sleep(polling_interval)
             return None
     return None
 

pydra/engine/submitter.py

@@ -738,7 +738,7 @@ def update_status(self) -> None:
                     job.run_start_time,
                 )
         # Check to see if any previously running tasks have completed
-        for index, (job, start_time) in list(self.running.items()):
+        for index, (job, _) in list(self.running.items()):
             if job.done:
                 self.successful[job.state_index] = self.running.pop(index)[0]
             elif job.errored: