add 'remove_on' argument (#464)

* add 'remove_on' argument

* update docstring

tests/integration_tests/test_operating_directory.py

@@ -30,12 +30,37 @@ def run(self):
                 self.file.write_text(f"{text} there")
 
 
+class RemoveOnError(Node):
+    data: list = zn.outs()
+
+    def run(self):
+        with self.operating_directory(remove_on=(TypeError, ValueError)):
+            raise ValueError("Execution was interrupted")
+
+
+def test_remove_on_error(proj_path):
+    node = RemoveOnError()
+    node.write_graph()
+
+    node = RemoveOnError.load()
+    with pytest.raises(ValueError):
+        node.run_and_save()
+
+    nwd_new = node.nwd.with_name(f"ckpt_{node.nwd.name}")
+    assert not nwd_new.exists()
+
+
 def test_ListOfDataNode(proj_path):
     ListOfDataNode().write_graph()
 
+    node = ListOfDataNode.load()
+
     with pytest.raises(exceptions.DVCProcessError):
         utils.run_dvc_cmd(["repro"])
+    nwd_new = node.nwd.with_name(f"ckpt_{node.nwd.name}")
+    assert nwd_new.exists()
     utils.run_dvc_cmd(["repro"])
+    assert not nwd_new.exists()
 
     assert ListOfDataNode.load().data == list(range(10))
 

zntrack/core/base.py

@@ -30,6 +30,10 @@
 
 log = logging.getLogger(__name__)
 
+EXCEPTION_OR_LST_EXCEPTIONS = typing.Union[
+    typing.Type[Exception], typing.Collection[typing.Type[Exception]]
+]
+
 
 def update_dependency_options(value):
     """Handle Node dependencies.
@@ -666,7 +670,9 @@ def write_graph(  # noqa: C901
             utils.run_dvc_cmd(["repro", self.node_name])
 
     @contextlib.contextmanager
-    def operating_directory(self, prefix="ckpt") -> bool:
+    def operating_directory(
+        self, prefix="ckpt", remove_on: EXCEPTION_OR_LST_EXCEPTIONS = None
+    ) -> bool:
         """Work in an operating directory until successfully finished.
 
         This context manager will replace $nwd$ with 'prefix_$nwd$' and move the files
@@ -680,6 +686,14 @@ def operating_directory(self, prefix="ckpt") -> bool:
         - checkpoints are not versioned. If you want to checkpoint e.g., model training,
             use 'dvc checkpoints'.
 
+        Parameters
+        ----------
+        prefix: str, default = 'ckpt'
+            Prefix for the temporary directory
+        remove_on: Exception or list of Exceptions, default = None
+            If one of the exceptions in 'remove_on' is raised, the operating directory
+             will be removed. Otherwise, it will remain and reused upton the next run.
+
 
         Yields
         ------
@@ -690,6 +704,10 @@ def operating_directory(self, prefix="ckpt") -> bool:
         nwd_new = self.nwd.with_name(f"{prefix}_{self.nwd.name}")
         nwd_is_new = not nwd_new.exists()
 
+        remove = False
+        if not isinstance(remove_on, list):
+            remove_on = [remove_on] if remove_on else []
+
         if self._run_and_save:
             utils.update_gitignore(prefix=prefix)
 
@@ -710,11 +728,16 @@ def operating_directory(self, prefix="ckpt") -> bool:
                 yield nwd_is_new
             except Exception as err:
                 log.warning("Node execution was interrupted.")
+                if any(isinstance(err, e) for e in remove_on):
+                    remove = True
                 raise err
             finally:
                 # Save e.g. `zn.outs` before stopping.
                 self.save(results=True)
                 self.nwd = nwd
+                if remove:
+                    log.info(f"Removing operating directory: {nwd_new}")
+                    shutil.rmtree(nwd_new)
 
             log.info(f"Finished successfully. Moving files from {nwd_new} to {nwd}")
             shutil.rmtree(nwd)