From 3daeeaffdb43989283719d9cd8a7a4825dda19d5 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Peter=20Rowlands=20=28=EB=B3=80=EA=B8=B0=ED=98=B8=29?=
 <peter@pmrowla.com>
Date: Thu, 12 Nov 2020 19:28:28 +0900
Subject: [PATCH] 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                   |  6 +++++-
 dvc/repo/experiments/__init__.py | 18 +++++++++++++++++-
 2 files changed, 22 insertions(+), 2 deletions(-)

diff --git a/dvc/dvcfile.py b/dvc/dvcfile.py
index aacfb4222..2a46237d4 100644
--- a/dvc/dvcfile.py
+++ b/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(
diff --git a/dvc/repo/experiments/__init__.py b/dvc/repo/experiments/__init__.py
index 2aea83082..12ebfa02f 100644
--- a/dvc/repo/experiments/__init__.py
+++ b/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}"