experiments: fix checkout error for unexisting outputs (#5175)

* stage: support explicit allow_missing in checkout

* exp: explicitly clear EXEC_BRANCH and checkout with allow_missing on
starting experiment run

* git: wrap stash apply errors with SCMError

* git: add `git status` to backends

* exp: reset staged git changes before workspace exp runs to avoid merge
conflicts

* test untracked (and staged) pipeline file and deps for both workspace
and tmpdir runs

* experiments: explicitly track pipeline files and uncached deps for exp
commmits

* fix exp typing errors

dvc/repo/experiments/__init__.py

@@ -328,6 +328,15 @@ class Experiments:
 
     def reproduce_one(self, queue=False, tmp_dir=False, **kwargs):
         """Reproduce and checkout a single experiment."""
+        if not (queue or tmp_dir):
+            staged, _, _ = self.scm.status()
+            if staged:
+                logger.warning(
+                    "Your workspace contains staged Git changes which will be "
+                    "unstaged before running this experiment."
+                )
+                self.scm.reset()
+
         stash_rev = self.new(**kwargs)
         if queue:
             logger.info(

dvc/repo/experiments/base.py

@@ -102,7 +102,7 @@ class ExpRefInfo:
         self, baseline_sha: Optional[str] = None, name: Optional[str] = None,
     ):
         self.baseline_sha = baseline_sha
-        self.name = name
+        self.name: str = name if name else ""
 
     def __str__(self):
         return "/".join(self.parts)

dvc/repo/experiments/executor/base.py

@@ -92,6 +92,8 @@ class BaseExecutor(ABC):
             if branch:
                 scm.push_refspec(self.git_url, branch, branch)
                 self.scm.set_ref(EXEC_BRANCH, branch, symbolic=True)
+            elif self.scm.get_ref(EXEC_BRANCH):
+                self.scm.remove_ref(EXEC_BRANCH)
 
             if self.scm.get_ref(EXEC_CHECKPOINT):
                 self.scm.remove_ref(EXEC_CHECKPOINT)
@@ -120,7 +122,7 @@ class BaseExecutor(ABC):
         return os.path.join(self.root_dir, self._dvc_dir)
 
     @staticmethod
-    def hash_exp(stages: Iterable["PipelineStage"]):
+    def hash_exp(stages: Iterable["PipelineStage"]) -> str:
         exp_data = {}
         for stage in stages:
             if isinstance(stage, PipelineStage):
@@ -245,7 +247,7 @@ class BaseExecutor(ABC):
         try:
             dvc = Repo(dvc_dir)
             if dvc_dir is not None:
-                old_cwd = os.getcwd()
+                old_cwd: Optional[str] = os.getcwd()
                 if rel_cwd:
                     os.chdir(os.path.join(dvc.root_dir, rel_cwd))
                 else:
@@ -278,7 +280,7 @@ class BaseExecutor(ABC):
             #   be removed/does not yet exist) so that our executor workspace
             #   is not polluted with the (persistent) out from an unrelated
             #   experiment run
-            dvc_checkout(dvc, force=True, quiet=True)
+            dvc_checkout(dvc, force=True, quiet=True, allow_missing=True)
 
             checkpoint_func = partial(
                 cls.checkpoint_callback, dvc.scm, name, repro_force

dvc/repo/reproduce.py

@@ -60,6 +60,10 @@ def _dump_stage(stage):
 
 
 def _track_stage(stage):
+    stage.repo.scm.track_file(stage.dvcfile.path)
+    for dep in stage.deps:
+        if not dep.use_scm_ignore and dep.is_in_repo:
+            stage.repo.scm.track_file(os.fspath(dep.path_info))
     for out in stage.outs:
         if not out.use_scm_ignore and out.is_in_repo:
             stage.repo.scm.track_file(os.fspath(out.path_info))

dvc/scm/git/__init__.py

@@ -355,6 +355,7 @@ class Git(Base):
     diff = partialmethod(_backend_func, "diff")
     reset = partialmethod(_backend_func, "reset")
     checkout_paths = partialmethod(_backend_func, "checkout_paths")
+    status = partialmethod(_backend_func, "status")
 
     def resolve_rev(self, rev: str) -> str:
         from dvc.repo.experiments.utils import exp_refs_by_name

dvc/scm/git/backend/base.py

@@ -1,6 +1,6 @@
 import os
 from abc import ABC, abstractmethod
-from typing import Callable, Iterable, Optional, Tuple
+from typing import Callable, Iterable, Mapping, Optional, Tuple
 
 from dvc.scm.base import SCMError
 
@@ -292,3 +292,16 @@ class BaseGitBackend(ABC):
     @abstractmethod
     def checkout_paths(self, paths: Iterable[str], force: bool = False):
         """Checkout the specified paths from HEAD index."""
+
+    @abstractmethod
+    def status(
+        self, ignored: bool = False
+    ) -> Tuple[Mapping[str, Iterable[str]], Iterable[str], Iterable[str]]:
+        """Return tuple of (staged_files, unstaged_files, untracked_files).
+
+        staged_files will be a dict mapping status (add, delete, modify) to a
+        list of paths.
+
+        If ignored is True, gitignored files will be included in
+        untracked_paths.
+        """

dvc/scm/git/backend/dulwich.py

@@ -4,7 +4,7 @@ import logging
 import os
 import stat
 from io import BytesIO, StringIO
-from typing import Callable, Dict, Iterable, Optional, Tuple
+from typing import Callable, Dict, Iterable, Mapping, Optional, Tuple
 
 from dvc.path_info import PathInfo
 from dvc.progress import Tqdm
@@ -172,7 +172,8 @@ class DulwichBackend(BaseGitBackend):  # pylint:disable=abstract-method
         raise NotImplementedError
 
     def untracked_files(self) -> Iterable[str]:
-        raise NotImplementedError
+        _staged, _unstaged, untracked = self.status()
+        return untracked
 
     def is_tracked(self, path: str) -> bool:
         rel = PathInfo(path).relative_to(self.root_dir).as_posix().encode()
@@ -531,3 +532,21 @@ class DulwichBackend(BaseGitBackend):  # pylint:disable=abstract-method
 
     def checkout_paths(self, paths: Iterable[str], force: bool = False):
         raise NotImplementedError
+
+    def status(
+        self, ignored: bool = False
+    ) -> Tuple[Mapping[str, Iterable[str]], Iterable[str], Iterable[str]]:
+        from dulwich.porcelain import status as git_status
+
+        staged, unstaged, untracked = git_status(
+            self.root_dir, ignored=ignored
+        )
+        return (
+            {
+                status: [os.fsdecode(name)]
+                for status, paths in staged.items()
+                for name in paths
+            },
+            [os.fsdecode(name) for name in unstaged],
+            [os.fsdecode(name) for name in untracked],
+        )

dvc/scm/git/backend/gitpython.py

@@ -3,7 +3,7 @@ import locale
 import logging
 import os
 from functools import partial
-from typing import Callable, Iterable, Optional, Tuple
+from typing import Callable, Iterable, Mapping, Optional, Tuple
 
 from funcy import first
 
@@ -516,3 +516,8 @@ class GitPythonBackend(BaseGitBackend):  # pylint:disable=abstract-method
         """Checkout the specified paths from HEAD index."""
         if paths:
             self.repo.index.checkout(paths=paths, force=force)
+
+    def status(
+        self, ignored: bool = False
+    ) -> Tuple[Mapping[str, Iterable[str]], Iterable[str], Iterable[str]]:
+        raise NotImplementedError

dvc/scm/git/backend/pygit2.py

@@ -2,7 +2,7 @@ import locale
 import logging
 import os
 from io import BytesIO, StringIO
-from typing import Callable, Iterable, Optional, Tuple
+from typing import Callable, Iterable, Mapping, Optional, Tuple
 
 from dvc.scm.base import SCMError
 
@@ -229,3 +229,8 @@ class Pygit2Backend(BaseGitBackend):  # pylint:disable=abstract-method
 
     def iter_remote_refs(self, url: str, base: Optional[str] = None):
         raise NotImplementedError
+
+    def status(
+        self, ignored: bool = False
+    ) -> Tuple[Mapping[str, Iterable[str]], Iterable[str], Iterable[str]]:
+        raise NotImplementedError

dvc/scm/git/stash.py

@@ -50,7 +50,10 @@ class Stash:
         logger.debug("Popping from stash '%s'", self.ref)
         ref = "{0}@{{0}}".format(self.ref)
         rev = self.scm.resolve_rev(ref)
-        self.apply(rev)
+        try:
+            self.apply(rev)
+        except Exception as exc:
+            raise SCMError("Could not apply stash commit") from exc
         self.drop()
         return rev
 

dvc/stage/__init__.py

@@ -520,11 +520,14 @@ class Stage(params.StageParams):
         return filter(_func, self.outs) if path_info else self.outs
 
     @rwlocked(write=["outs"])
-    def checkout(self, **kwargs):
+    def checkout(self, allow_missing=False, **kwargs):
         stats = defaultdict(list)
-        kwargs["allow_missing"] = self.is_checkpoint
         for out in self.filter_outs(kwargs.get("filter_info")):
-            key, outs = self._checkout(out, **kwargs)
+            key, outs = self._checkout(
+                out,
+                allow_missing=allow_missing or self.is_checkpoint,
+                **kwargs,
+            )
             if key:
                 stats[key].extend(outs)
         return stats

tests/func/experiments/conftest.py

@@ -50,6 +50,7 @@ def exp_stage(tmp_dir, scm, dvc):
         metrics_no_cache=["metrics.yaml"],
         params=["foo"],
         name="copy-file",
+        deps=["copy.py"],
     )
     scm.add(
         [
@@ -74,6 +75,7 @@ def checkpoint_stage(tmp_dir, scm, dvc):
         metrics_no_cache=["metrics.yaml"],
         params=["foo"],
         checkpoints=["foo"],
+        deps=["checkpoint.py"],
         no_exec=True,
         name="checkpoint-file",
     )

tests/func/experiments/test_experiments.py

@@ -379,17 +379,18 @@ def test_no_scm(tmp_dir):
             getattr(dvc.experiments, cmd)()
 
 
-def test_untracked(tmp_dir, scm, dvc, caplog):
+@pytest.mark.parametrize("workspace", [True, False])
+def test_untracked(tmp_dir, scm, dvc, caplog, workspace):
     tmp_dir.gen("copy.py", COPY_SCRIPT)
     tmp_dir.gen("params.yaml", "foo: 1")
     stage = dvc.run(
         cmd="python copy.py params.yaml metrics.yaml",
         metrics_no_cache=["metrics.yaml"],
         params=["foo"],
+        deps=["copy.py"],
         name="copy-file",
+        no_exec=True,
     )
-    scm.add(["dvc.yaml", "dvc.lock", "params.yaml", "metrics.yaml"])
-    scm.commit("init")
 
     # copy.py is untracked
     with caplog.at_level(logging.ERROR):
@@ -399,13 +400,15 @@ def test_untracked(tmp_dir, scm, dvc, caplog):
         assert "Failed to reproduce experiment" in caplog.text
         assert not results
 
-    # copy.py is staged as new file but not committed
-    scm.add(["copy.py"])
+    # dvc.yaml, copy.py are staged as new file but not committed
+    scm.add(["dvc.yaml", "copy.py"])
     results = dvc.experiments.run(
-        stage.addressing, params=["foo=2"], tmp_dir=True
+        stage.addressing, params=["foo=2"], tmp_dir=not workspace
     )
     exp = first(results)
     tree = scm.get_tree(exp)
+    assert tree.exists("dvc.yaml")
+    assert tree.exists("dvc.lock")
     assert tree.exists("copy.py")
     with tree.open(tmp_dir / "metrics.yaml") as fobj:
         assert fobj.read().strip() == "foo: 2"