experiments: fix checkpoint/persist related checkout issues (#4885)

* dvcfile: add is_lock_file to identify pipeline lock files

* experiments: clean dirty lock files before stashing experiments

dvc/dvcfile.py

@@ -48,10 +48,14 @@ def is_valid_filename(path):
 
 def is_dvc_file(path):
     return os.path.isfile(path) and (
-        is_valid_filename(path) or os.path.basename(path) == PIPELINE_LOCK
+        is_valid_filename(path) or is_lock_file(path)
     )
 
 
+def is_lock_file(path):
+    return os.path.basename(path) == PIPELINE_LOCK
+
+
 def check_dvc_filename(path):
     if not is_valid_filename(path):
         raise StageFileBadNameError(

dvc/repo/experiments/__init__.py

@@ -12,6 +12,7 @@ from typing import Iterable, Optional
 
 from funcy import cached_property, first
 
+from dvc.dvcfile import is_lock_file
 from dvc.exceptions import DownloadError, DvcException, UploadError
 from dvc.path_info import PathInfo
 from dvc.progress import Tqdm
@@ -284,7 +285,7 @@ class Experiments:
         if params:
             self._update_params(params)
 
-        if not self.scm.is_dirty(untracked_files=True) and not allow_unchanged:
+        if not self._check_dirty() and not allow_unchanged:
             # experiment matches original baseline
             raise UnchangedExperimentError(rev)
 
@@ -298,6 +299,21 @@ class Experiments:
         self.scm.repo.git.stash("push", "-m", msg)
         return self.scm.resolve_rev("stash@{0}")
 
+    def _check_dirty(self) -> bool:
+        # 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.repo.index.diff(None)]
+        to_checkout = [fname for fname in dirty if is_lock_file(fname)]
+        self.scm.repo.index.checkout(paths=to_checkout, force=True)
+
+        untracked = self.scm.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
+
     def _stash_msg(self, rev, branch=None):
         if branch:
             return f"{self.STASH_MSG_PREFIX}{rev}:{branch}"