experiments: ignore untracked files (#5029)

* git: add reset()

* experiments: exclude untracked files

* experiments: add test for untracked & staged files

* git: support checkout/reset for specific paths to HEAD state

* stage: track changed files when calling checkpoint callback

* experiments: ignore untracked files when generating exp commits

* add tests for dirty lock files

* stage: track uncached outs when generating checkpoints

* experiments: update checkpoints scm handling

* update checkpoints tests

dvc/repo/experiments/__init__.py

@@ -11,6 +11,7 @@ from typing import Iterable, Mapping, Optional
 
 from funcy import cached_property, first
 
+from dvc.dvcfile import is_lock_file
 from dvc.exceptions import DvcException
 from dvc.path_info import PathInfo
 from dvc.stage.run import CheckpointKilledError
@@ -147,11 +148,14 @@ class Experiments:
                 the human-readable name in the experiment branch ref. Has no
                 effect of branch is specified.
         """
-        with self.scm.stash_workspace(include_untracked=True) as workspace:
+        with self.scm.stash_workspace(
+            include_untracked=detach_rev or branch
+        ) as workspace:
             # If we are not extending an existing branch, apply current
             # workspace changes to be made in new branch
-            if not branch and workspace:
+            if not (branch or detach_rev) and workspace:
                 self.stash.apply(workspace)
+                self._prune_lockfiles()
 
             # checkout and detach at branch (or current HEAD)
             if detach_rev:
@@ -186,12 +190,24 @@ class Experiments:
                 baseline_rev[:7],
             )
 
-            # Reset/clean any changes before prior workspace is unstashed
-            self.scm.gitpython.repo.git.reset(hard=True)
-            self.scm.gitpython.repo.git.clean(force=True)
+            # Reset any changes before prior workspace is unstashed
+            self.scm.reset(hard=True)
 
         return stash_rev
 
+    def _prune_lockfiles(self):
+        # NOTE: dirty DVC lock files must be restored to index state to
+        # avoid checking out incorrect persist or checkpoint outs
+        tree = self.scm.get_tree("HEAD")
+        lock_files = [
+            str(fname)
+            for fname in tree.walk_files(self.scm.root_dir)
+            if is_lock_file(fname)
+        ]
+        if lock_files:
+            self.scm.reset(paths=lock_files)
+            self.scm.checkout_paths(lock_files, force=True)
+
     def _stash_msg(
         self,
         rev: str,
@@ -404,33 +420,29 @@ class Experiments:
 
     def _init_executors(self, to_run):
         executors = {}
-        with self.scm.stash_workspace(include_untracked=True):
-            with self.scm.detach_head():
-                for stash_rev, item in to_run.items():
-                    self.scm.set_ref(EXEC_HEAD, item.rev)
-                    self.scm.set_ref(EXEC_MERGE, stash_rev)
-                    self.scm.set_ref(EXEC_BASELINE, item.baseline_rev)
+        for stash_rev, item in to_run.items():
+            self.scm.set_ref(EXEC_HEAD, item.rev)
+            self.scm.set_ref(EXEC_MERGE, stash_rev)
+            self.scm.set_ref(EXEC_BASELINE, item.baseline_rev)
 
-                    # Executor will be initialized with an empty git repo that
-                    # we populate by pushing:
-                    #   EXEC_HEAD - the base commit for this experiment
-                    #   EXEC_MERGE - the unmerged changes (from our stash)
-                    #       to be reproduced
-                    #   EXEC_BASELINE - the baseline commit for this experiment
-                    executor = LocalExecutor(
-                        self.scm,
-                        self.dvc_dir,
-                        name=item.name,
-                        branch=item.branch,
-                        cache_dir=self.repo.cache.local.cache_dir,
-                    )
-                    executors[item.rev] = executor
+            # Executor will be initialized with an empty git repo that
+            # we populate by pushing:
+            #   EXEC_HEAD - the base commit for this experiment
+            #   EXEC_MERGE - the unmerged changes (from our stash)
+            #       to be reproduced
+            #   EXEC_BASELINE - the baseline commit for this experiment
+            executor = LocalExecutor(
+                self.scm,
+                self.dvc_dir,
+                name=item.name,
+                branch=item.branch,
+                cache_dir=self.repo.cache.local.cache_dir,
+            )
+            executors[item.rev] = executor
 
-                for ref in (EXEC_HEAD, EXEC_MERGE, EXEC_BASELINE):
-                    self.scm.remove_ref(ref)
+        for ref in (EXEC_HEAD, EXEC_MERGE, EXEC_BASELINE):
+            self.scm.remove_ref(ref)
 
-            self.scm.gitpython.repo.git.reset(hard=True)
-            self.scm.gitpython.repo.git.clean(force=True)
         return executors
 
     def _reproduce(

dvc/repo/experiments/executor.py

@@ -7,7 +7,6 @@ from typing import TYPE_CHECKING, Callable, Iterable, Optional, Tuple, Union
 
 from funcy import cached_property
 
-from dvc.dvcfile import is_lock_file
 from dvc.path_info import PathInfo
 from dvc.repo import Repo
 from dvc.repo.experiments.base import (
@@ -23,7 +22,6 @@ from dvc.repo.experiments.base import (
     UnchangedExperimentError,
 )
 from dvc.scm import SCM
-from dvc.scm.git import Git
 from dvc.stage import PipelineStage
 from dvc.stage.serialize import to_lockfile
 from dvc.utils import dict_sha256
@@ -32,6 +30,8 @@ from dvc.utils.fs import remove
 if TYPE_CHECKING:
     from multiprocessing import Queue
 
+    from dvc.scm.git import Git
+
 logger = logging.getLogger(__name__)
 
 
@@ -88,28 +88,9 @@ class BaseExecutor:
             self.scm.gitpython.repo.git.merge(
                 EXEC_MERGE, squash=True, no_commit=True
             )
-            self.scm.gitpython.repo.git.reset()
-            self._prune_lockfiles()
         finally:
             os.chdir(cwd)
 
-    def _prune_lockfiles(self):
-        # NOTE: dirty DVC lock files must be restored to index state to
-        # avoid checking out incorrect persist or checkpoint outs
-        dirty = [
-            diff.a_path for diff in self.scm.gitpython.repo.index.diff(None)
-        ]
-        to_checkout = [fname for fname in dirty if is_lock_file(fname)]
-        self.scm.gitpython.repo.index.checkout(paths=to_checkout, force=True)
-
-        untracked = self.scm.gitpython.repo.untracked_files
-        to_remove = [fname for fname in untracked if is_lock_file(fname)]
-        for fname in to_remove:
-            remove(fname)
-        return (
-            len(dirty) - len(to_checkout) + len(untracked) - len(to_remove)
-        ) != 0
-
     @cached_property
     def scm(self):
         return SCM(self.root_dir)
@@ -260,10 +241,7 @@ class BaseExecutor:
             #   experiment run
             dvc.checkout(force=True, quiet=True)
 
-            # We cannot use dvc.scm to make commits inside the executor since
-            # cached props are not picklable.
-            scm = Git()
-            checkpoint_func = partial(cls.checkpoint_callback, scm, name)
+            checkpoint_func = partial(cls.checkpoint_callback, dvc.scm, name)
             stages = dvc.reproduce(
                 *args,
                 on_unchanged=filter_pipeline,
@@ -272,15 +250,14 @@ class BaseExecutor:
             )
 
             exp_hash = cls.hash_exp(stages)
-            exp_rev = cls.commit(scm, exp_hash, exp_name=name)
-            if scm.get_ref(EXEC_CHECKPOINT):
-                scm.set_ref(EXEC_CHECKPOINT, exp_rev)
+            exp_rev = cls.commit(dvc.scm, exp_hash, exp_name=name)
+            if dvc.scm.get_ref(EXEC_CHECKPOINT):
+                dvc.scm.set_ref(EXEC_CHECKPOINT, exp_rev)
         except UnchangedExperimentError:
             pass
         finally:
-            if scm:
-                scm.close()
-                del scm
+            if dvc:
+                dvc.scm.close()
             if old_cwd:
                 os.chdir(old_cwd)
 
@@ -324,7 +301,7 @@ class BaseExecutor:
             old_ref = None
             logger.debug("Commit to new experiment branch '%s'", branch)
 
-        scm.gitpython.repo.git.add(A=True)
+        scm.gitpython.repo.git.add(update=True)
         scm.commit(f"dvc: commit experiment {exp_hash}")
         new_rev = scm.get_rev()
         scm.set_ref(branch, new_rev, old_ref=old_ref)
@@ -354,7 +331,6 @@ class LocalExecutor(BaseExecutor):
         local_config = os.path.join(self.dvc_dir, "config.local")
         logger.debug("Writing experiments local config '%s'", local_config)
         with open(local_config, "w") as fobj:
-            fobj.write("[core]\n    no_scm = true\n")
             fobj.write(f"[cache]\n    dir = {cache_dir}")
 
     @property

dvc/scm/git/__init__.py

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

dvc/scm/git/backend/base.py

@@ -259,3 +259,11 @@ class BaseGitBackend(ABC):
     @abstractmethod
     def diff(self, rev_a: str, rev_b: str, binary=False) -> str:
         """Return the git diff for two commits."""
+
+    @abstractmethod
+    def reset(self, hard: bool = False, paths: Iterable[str] = None):
+        """Reset current git HEAD."""
+
+    @abstractmethod
+    def checkout_paths(self, paths: Iterable[str], force: bool = False):
+        """Checkout the specified paths from HEAD index."""

dvc/scm/git/backend/dulwich.py

@@ -340,3 +340,9 @@ class DulwichBackend(BaseGitBackend):  # pylint:disable=abstract-method
             buf, self.repo.object_store, commit_a.tree, commit_b.tree
         )
         return buf.getvalue().decode("utf-8")
+
+    def reset(self, hard: bool = False, paths: Iterable[str] = None):
+        raise NotImplementedError
+
+    def checkout_paths(self, paths: Iterable[str], force: bool = False):
+        raise NotImplementedError

dvc/scm/git/backend/gitpython.py

@@ -394,3 +394,11 @@ class GitPythonBackend(BaseGitBackend):  # pylint:disable=abstract-method
 
     def diff(self, rev_a: str, rev_b: str, binary=False) -> str:
         raise NotImplementedError
+
+    def reset(self, hard: bool = False, paths: Iterable[str] = None):
+        self.repo.head.reset(index=True, working_tree=hard, paths=paths)
+
+    def checkout_paths(self, paths: Iterable[str], force: bool = False):
+        """Checkout the specified paths from HEAD index."""
+        if paths:
+            self.repo.index.checkout(paths=paths, force=force)

dvc/stage/run.py

@@ -186,5 +186,9 @@ def _kill_nt(proc):
 def _run_callback(stage, callback_func):
     stage.save(allow_missing=True)
     stage.commit(allow_missing=True)
+    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))
+    stage.repo.scm.track_changed_files()
     logger.debug("Running checkpoint callback for stage '%s'", stage)
     callback_func()

tests/func/experiments/test_checkpoints.py

@@ -13,11 +13,14 @@ def test_new_checkpoint(tmp_dir, scm, dvc, checkpoint_stage, mocker):
     exp = first(results)
 
     new_mock.assert_called_once()
-    tree = scm.get_tree(exp)
-    with tree.open(tmp_dir / "foo") as fobj:
-        assert fobj.read().strip() == "5"
-    with tree.open(tmp_dir / "metrics.yaml") as fobj:
-        assert fobj.read().strip() == "foo: 2"
+    for rev in dvc.brancher([exp]):
+        if rev == "workspace":
+            continue
+        tree = dvc.repo_tree
+        with tree.open(tmp_dir / "foo") as fobj:
+            assert fobj.read().strip() == "5"
+        with tree.open(tmp_dir / "metrics.yaml") as fobj:
+            assert fobj.read().strip() == "foo: 2"
 
 
 @pytest.mark.parametrize(
@@ -49,11 +52,14 @@ def test_resume_checkpoint(
     )
     exp = first(results)
 
-    tree = scm.get_tree(exp)
-    with tree.open(tmp_dir / "foo") as fobj:
-        assert fobj.read().strip() == "10"
-    with tree.open(tmp_dir / "metrics.yaml") as fobj:
-        assert fobj.read().strip() == "foo: 2"
+    for rev in dvc.brancher([exp]):
+        if rev == "workspace":
+            continue
+        tree = dvc.repo_tree
+        with tree.open(tmp_dir / "foo") as fobj:
+            assert fobj.read().strip() == "10"
+        with tree.open(tmp_dir / "metrics.yaml") as fobj:
+            assert fobj.read().strip() == "foo: 2"
 
 
 def test_reset_checkpoint(tmp_dir, scm, dvc, checkpoint_stage, caplog):
@@ -73,11 +79,14 @@ def test_reset_checkpoint(tmp_dir, scm, dvc, checkpoint_stage, caplog):
     )
     exp = first(results)
 
-    tree = scm.get_tree(exp)
-    with tree.open(tmp_dir / "foo") as fobj:
-        assert fobj.read().strip() == "5"
-    with tree.open(tmp_dir / "metrics.yaml") as fobj:
-        assert fobj.read().strip() == "foo: 2"
+    for rev in dvc.brancher([exp]):
+        if rev == "workspace":
+            continue
+        tree = dvc.repo_tree
+        with tree.open(tmp_dir / "foo") as fobj:
+            assert fobj.read().strip() == "5"
+        with tree.open(tmp_dir / "metrics.yaml") as fobj:
+            assert fobj.read().strip() == "foo: 2"
 
 
 def test_resume_branch(tmp_dir, scm, dvc, checkpoint_stage):
@@ -98,17 +107,23 @@ def test_resume_branch(tmp_dir, scm, dvc, checkpoint_stage):
     )
     checkpoint_b = first(results)
 
-    tree = scm.get_tree(checkpoint_a)
-    with tree.open(tmp_dir / "foo") as fobj:
-        assert fobj.read().strip() == "10"
-    with tree.open(tmp_dir / "metrics.yaml") as fobj:
-        assert fobj.read().strip() == "foo: 2"
+    for rev in dvc.brancher([checkpoint_a]):
+        if rev == "workspace":
+            continue
+        tree = dvc.repo_tree
+        with tree.open(tmp_dir / "foo") as fobj:
+            assert fobj.read().strip() == "10"
+        with tree.open(tmp_dir / "metrics.yaml") as fobj:
+            assert fobj.read().strip() == "foo: 2"
 
-    tree = scm.get_tree(checkpoint_b)
-    with tree.open(tmp_dir / "foo") as fobj:
-        assert fobj.read().strip() == "10"
-    with tree.open(tmp_dir / "metrics.yaml") as fobj:
-        assert fobj.read().strip() == "foo: 100"
+    for rev in dvc.brancher([checkpoint_b]):
+        if rev == "workspace":
+            continue
+        tree = dvc.repo_tree
+        with tree.open(tmp_dir / "foo") as fobj:
+            assert fobj.read().strip() == "10"
+        with tree.open(tmp_dir / "metrics.yaml") as fobj:
+            assert fobj.read().strip() == "foo: 100"
 
     with pytest.raises(MultipleBranchError):
         dvc.experiments.get_branch_by_rev(branch_rev)

tests/func/experiments/test_experiments.py

@@ -337,3 +337,49 @@ def test_no_scm(tmp_dir):
 
     with pytest.raises(NoSCMError):
         dvc.experiments.gc()
+
+
+def test_untracked(tmp_dir, scm, dvc, caplog):
+    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"],
+        name="copy-file",
+    )
+    scm.add(["dvc.yaml", "dvc.lock", "params.yaml", "metrics.yaml"])
+    scm.commit("init")
+
+    # copy.py is untracked
+    with caplog.at_level(logging.ERROR):
+        results = dvc.experiments.run(stage.addressing, params=["foo=2"])
+        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"])
+    results = dvc.experiments.run(stage.addressing, params=["foo=2"])
+    exp = first(results)
+    tree = scm.get_tree(exp)
+    assert tree.exists("copy.py")
+    with tree.open(tmp_dir / "metrics.yaml") as fobj:
+        assert fobj.read().strip() == "foo: 2"
+
+
+def test_dirty_lockfile(tmp_dir, scm, dvc, exp_stage):
+    from dvc.dvcfile import LockfileCorruptedError
+
+    tmp_dir.gen("dvc.lock", "foo")
+
+    with pytest.raises(LockfileCorruptedError):
+        dvc.reproduce(exp_stage.addressing)
+
+    results = dvc.experiments.run(exp_stage.addressing, params=["foo=2"])
+    exp = first(results)
+
+    tree = scm.get_tree(exp)
+    with tree.open(tmp_dir / "metrics.yaml") as fobj:
+        assert fobj.read().strip() == "foo: 2"
+
+    assert (tmp_dir / "dvc.lock").read_text() == "foo"