checkpoints: allow missing outputs on checkout and stage.run (#4750)

dvc/output/base.py

@@ -282,6 +282,9 @@ class BaseOutput:
         self.hash_info = self.get_hash()
 
     def commit(self):
+        if not self.exists:
+            raise self.DoesNotExistError(self)
+
         assert self.hash_info
         if self.use_cache:
             self.cache.save(self.path_info, self.cache.tree, self.hash_info)
@@ -326,7 +329,7 @@ class BaseOutput:
         progress_callback=None,
         relink=False,
         filter_info=None,
-        allow_persist_missing=False,
+        allow_missing=False,
     ):
         if not self.use_cache:
             if progress_callback:
@@ -345,7 +348,7 @@ class BaseOutput:
                 filter_info=filter_info,
             )
         except CheckoutError:
-            if self.persist and allow_persist_missing:
+            if allow_missing:
                 return None
             raise
 

dvc/repo/checkout.py

@@ -46,7 +46,7 @@ def checkout(
     force=False,
     relink=False,
     recursive=False,
-    allow_persist_missing=False,
+    allow_missing=False,
 ):
     from dvc.stage.exceptions import (
         StageFileBadNameError,
@@ -97,7 +97,7 @@ def checkout(
                 progress_callback=pbar.update_msg,
                 relink=relink,
                 filter_info=filter_info,
-                allow_persist_missing=allow_persist_missing,
+                allow_missing=allow_missing,
             )
             for key, items in result.items():
                 stats[key].extend(_fspath_dir(path) for path in items)

dvc/repo/experiments/__init__.py

@@ -395,7 +395,7 @@ class Experiments:
         )
         exp_rev = first(results)
         if exp_rev is not None:
-            self.checkout_exp(exp_rev)
+            self.checkout_exp(exp_rev, allow_missing=checkpoint)
         return results
 
     def reproduce_queued(self, **kwargs):
@@ -727,7 +727,7 @@ class Experiments:
             raise UploadError(fails)
 
     @scm_locked
-    def checkout_exp(self, rev):
+    def checkout_exp(self, rev, allow_missing=False):
         """Checkout an experiment to the user's workspace."""
         from git.exc import GitCommandError
 
@@ -761,7 +761,7 @@ class Experiments:
                 self._unstash_workspace()
 
         if need_checkout:
-            dvc_checkout(self.repo)
+            dvc_checkout(self.repo, allow_missing=allow_missing)
 
     def _check_baseline(self, exp_rev):
         baseline_sha = self.repo.scm.get_rev()

dvc/repo/experiments/executor.py

@@ -168,8 +168,12 @@ class LocalExecutor(ExperimentExecutor):
             #   is not polluted with the (persistent) out from an unrelated
             #   experiment run
             checkpoint = kwargs.pop("checkpoint", False)
-            dvc.checkout(allow_persist_missing=checkpoint, force=checkpoint)
-            stages = dvc.reproduce(on_unchanged=filter_pipeline, **kwargs)
+            dvc.checkout(allow_missing=checkpoint, force=checkpoint)
+            stages = dvc.reproduce(
+                on_unchanged=filter_pipeline,
+                allow_missing=checkpoint,
+                **kwargs,
+            )
         finally:
             if old_cwd is not None:
                 os.chdir(old_cwd)

dvc/stage/__init__.py

@@ -8,6 +8,7 @@ from funcy import cached_property, project
 import dvc.dependency as dependency
 import dvc.prompt as prompt
 from dvc.exceptions import CheckoutError, DvcException, MergeError
+from dvc.output.base import OutputDoesNotExistError
 from dvc.utils import relpath
 
 from . import params
@@ -380,9 +381,9 @@ class Stage(params.StageParams):
         logger.debug(f"Computed {self} md5: '{m}'")
         return m
 
-    def save(self):
+    def save(self, allow_missing=False):
         self.save_deps()
-        self.save_outs()
+        self.save_outs(allow_missing=allow_missing)
         self.md5 = self.compute_md5()
 
         self.repo.stage_cache.save(self)
@@ -391,9 +392,13 @@ class Stage(params.StageParams):
         for dep in self.deps:
             dep.save()
 
-    def save_outs(self):
+    def save_outs(self, allow_missing=False):
         for out in self.outs:
-            out.save()
+            try:
+                out.save()
+            except OutputDoesNotExistError:
+                if not allow_missing:
+                    raise
 
     def ignore_outs(self):
         for out in self.outs:
@@ -416,12 +421,23 @@ class Stage(params.StageParams):
         )
 
     @rwlocked(write=["outs"])
-    def commit(self):
+    def commit(self, allow_missing=False):
         for out in self.outs:
-            out.commit()
+            try:
+                out.commit()
+            except OutputDoesNotExistError:
+                if not allow_missing:
+                    raise
 
     @rwlocked(read=["deps"], write=["outs"])
-    def run(self, dry=False, no_commit=False, force=False, **kwargs):
+    def run(
+        self,
+        dry=False,
+        no_commit=False,
+        force=False,
+        allow_missing=False,
+        **kwargs,
+    ):
         if (self.cmd or self.is_import) and not self.frozen and not dry:
             self.remove_outs(ignore_remove=False, force=False)
 
@@ -438,9 +454,9 @@ class Stage(params.StageParams):
                 check_missing_outputs(self)
 
         if not dry:
-            self.save()
+            self.save(allow_missing=allow_missing)
             if not no_commit:
-                self.commit()
+                self.commit(allow_missing=allow_missing)
 
     def filter_outs(self, path_info):
         def _func(o):

dvc/stage/run.py

@@ -169,7 +169,7 @@ def _checkpoint_run(stage, callback_func, done_cond, proc):
 
 @relock_repo
 def _run_callback(stage, callback_func):
-    stage.save()
+    stage.save(allow_missing=True)
     # TODO: do we need commit() (and check for --no-commit) here
     logger.debug("Running checkpoint callback for stage '%s'", stage)
     callback_func()