fix: backup output files marked for update and restore them in case of job failures

johanneskoester · johanneskoester · commit 03973c328496 · 2025-11-07T12:44:24.000+01:00
diff --git a/docs/snakefiles/rules.rst b/docs/snakefiles/rules.rst
@@ -3371,7 +3371,8 @@ Consider the following example:
 
 Here, the statement ``test`` is appended to the output file ``test.txt``.
 Hence, we declare it as being updated via the ``update`` flag.
-This way, Snakemake will not delete the file before the job is executed.
+This way, Snakemake will not delete the file or directory before the job is executed.
+Furthermore, Snakemake will restore the previous version of the file/directory if the job fails.
 
 If such a file/directory has to be considered as input **before the update** for another rule
 it can be marked as ``before_update``.
diff --git a/src/snakemake/jobs.py b/src/snakemake/jobs.py
@@ -20,6 +20,7 @@
 from collections.abc import AsyncGenerator
 from abc import abstractmethod
 from snakemake import wrapper
+from snakemake.rules import Rule
 from snakemake.settings.types import DeploymentMethod
 
 from snakemake.template_rendering import check_template_output
@@ -229,7 +230,7 @@ def __init__(
         targetfile=None,
         groupid=None,
     ):
-        self.rule = rule
+        self.rule: Rule = rule
         self.dag = dag
 
         # the targetfile that led to the job
@@ -837,7 +838,9 @@ async def prepare(self):
             if self.resources.get("tmpdir"):
                 os.makedirs(self.resources.tmpdir, exist_ok=True)
 
-            for f, f_ in zip(self.output, self.rule.output):
+            for f in self.output:
+                if is_flagged(f, "update"):
+                    self.rule.workflow.persistence.backup_output(Path(f))
                 f.prepare()
 
             for f in self.log:
@@ -950,7 +953,7 @@ def is_pipe_or_service_input(self, path) -> bool:
 
     async def cleanup(self):
         """Cleanup output files."""
-        to_remove = [f for f in self.output if await f.exists()]
+        to_remove = [f for f in self.output if await f.exists() and not is_flagged(f, "update")]
         to_remove.extend(
             [
                 f
@@ -1229,6 +1232,13 @@ async def postprocess(
                 or (self.dag.workflow.remote_exec and not shared_input_output)
                 or self.is_local
             ):
+                for f in self.output:
+                    if is_flagged(f, "update"):
+                        if error:
+                            logger.warning(f"Restoring previous version of {fmt_iofile(f)} (updating job failed).")
+                            self.dag.workflow.persistence.restore_output(Path(f))
+                        else:
+                            self.dag.workflow.persistence.cleanup_backup(Path(f))
                 if not error and handle_touch:
                     self.dag.handle_touch(self)
                 if handle_log:
diff --git a/src/snakemake/persistence.py b/src/snakemake/persistence.py
@@ -70,6 +70,8 @@ def __init__(
         self.dag = dag
         self._lockfile = dict()
 
+        self._backup_path = self.path / "backups"
+
         self._metadata_path = os.path.join(self.path, "metadata")
         self._incomplete_path = os.path.join(self.path, "incomplete")
 
@@ -298,6 +300,41 @@ def conda_cleanup_envs(self):
             if d not in in_use:
                 shutil.rmtree(os.path.join(self.conda_env_archive_path, d))
 
+    def backup_output(self, path: Path) -> None:
+        backup_path = self._get_backup_path(path)
+        backup_path.parent.mkdir(parents=True, exist_ok=True)
+        if path.is_dir():
+            shutil.copytree(path, backup_path)
+        else:
+            shutil.copy(path, backup_path)
+
+    def restore_output(self, path: Path) -> None:
+        backup_path = self._get_backup_path(path)
+        if not backup_path.exists():
+            raise WorkflowError(f"Cannot restore {path}: no backup found.")
+        if backup_path.is_dir():
+            if path.exists():
+                shutil.rmtree(path)
+            shutil.copytree(backup_path, path)
+            shutil.rmtree(backup_path)
+        else:
+            shutil.copy(backup_path, path)
+            backup_path.unlink()
+
+    def cleanup_backup(self, path: Path) -> None:
+        backup_path = self._get_backup_path(path)
+        if backup_path.exists():
+            if backup_path.is_dir():
+                shutil.rmtree(backup_path)
+            else:
+                backup_path.unlink()
+
+    def _get_backup_path(self, path: Path) -> Path:
+        if path.is_absolute():
+            path = path.relative_to(path.parents[-1])
+        backup_path = self._backup_path / path
+        return backup_path
+
     def started(self, job, external_jobid: Optional[str] = None):
         for f in job.output:
             self._record(self._incomplete_path, {"external_jobid": external_jobid}, f)
diff --git a/tests/tests.py b/tests/tests.py
@@ -2301,6 +2301,10 @@ def test_default_flags():
     run(dpath("test_default_flags"), executor="dryrun", check_results=False)
 
 
+def test_update_flag_fail():
+    run(dpath("test_update_flag_fail"), shouldfail=True, check_results=True)
+
+
 @skip_on_windows
 @apptainer
 def test_shell_exec_singularity():
