experiments: support exp run --checkpoint for baseline commits (#4701)

* handle untracked files when `git reset --hard` in experiments workspace

* error on checkpoint branch already exists

* terminate stage process on checkpoint errors

* update tests

* cleanup error message

dvc/repo/experiments/__init__.py

@@ -173,6 +173,7 @@ class Experiments:
 
     def _scm_checkout(self, rev):
         self.scm.repo.git.reset(hard=True)
+        self.scm.repo.git.clean(force=True)
         if self.scm.repo.head.is_detached:
             self._checkout_default_branch()
         if not Git.is_sha(rev) or not self.scm.has_rev(rev):
@@ -185,6 +186,8 @@ class Experiments:
 
         # switch to default branch
         git_repo = self.scm.repo
+        git_repo.git.reset(hard=True)
+        git_repo.git.clean(force=True)
         origin_refs = git_repo.remotes["origin"].refs
 
         # origin/HEAD will point to tip of the default branch unless we
@@ -309,14 +312,23 @@ class Experiments:
         self, exp_hash, check_exists=True, create_branch=True, checkpoint=False
     ):
         """Commit stages as an experiment and return the commit SHA."""
-        if not self.scm.is_dirty():
+        if not self.scm.is_dirty(untracked_files=True):
             raise UnchangedExperimentError(self.scm.get_rev())
 
         rev = self.scm.get_rev()
         checkpoint = "-checkpoint" if checkpoint else ""
         exp_name = f"{rev[:7]}-{exp_hash}{checkpoint}"
         if create_branch:
-            if check_exists and exp_name in self.scm.list_branches():
+            if (
+                check_exists or checkpoint
+            ) and exp_name in self.scm.list_branches():
+                if checkpoint:
+                    raise DvcException(
+                        f"Checkpoint experiment containing '{rev[:7]}' "
+                        "already exists. To resume the experiment, "
+                        "run:\n\n"
+                        f"\tdvc exp run --continue {rev[:7]}"
+                    )
                 logger.debug("Using existing experiment branch '%s'", exp_name)
                 return self.scm.resolve_rev(exp_name)
             self.scm.checkout(exp_name, create_new=True)
@@ -696,6 +708,9 @@ class Experiments:
 
     def _check_baseline(self, exp_rev):
         baseline_sha = self.repo.scm.get_rev()
+        if exp_rev == baseline_sha:
+            return exp_rev
+
         exp_baseline = self._get_baseline(exp_rev)
         if exp_baseline is None:
             # if we can't tell from branch name, fall back to parent commit

dvc/scm/git.py

@@ -288,8 +288,8 @@ class Git(Base):
     def is_tracked(self, path):
         return bool(self.repo.git.ls_files(path))
 
-    def is_dirty(self):
-        return self.repo.is_dirty()
+    def is_dirty(self, **kwargs):
+        return self.repo.is_dirty(**kwargs)
 
     def active_branch(self):
         return self.repo.active_branch.name

dvc/stage/run.py

@@ -39,9 +39,9 @@ def warn_if_fish(executable):
 
 
 @unlocked_repo
-def cmd_run(stage, *args, checkpoint=False, **kwargs):
+def cmd_run(stage, *args, checkpoint_func=None, **kwargs):
     kwargs = {"cwd": stage.wdir, "env": fix_env(None), "close_fds": True}
-    if checkpoint:
+    if checkpoint_func:
         # indicate that checkpoint cmd is being run inside DVC
         kwargs["env"].update({"DVC_CHECKPOINT": "1"})
 
@@ -78,7 +78,9 @@ def cmd_run(stage, *args, checkpoint=False, **kwargs):
         p = subprocess.Popen(cmd, **kwargs)
         if main_thread:
             old_handler = signal.signal(signal.SIGINT, signal.SIG_IGN)
-        p.communicate()
+
+        with checkpoint_monitor(stage, checkpoint_func, p):
+            p.communicate()
     finally:
         if old_handler:
             signal.signal(signal.SIGINT, old_handler)
@@ -106,8 +108,7 @@ def run_stage(stage, dry=False, force=False, checkpoint_func=None, **kwargs):
     )
     logger.info("\t%s", stage.cmd)
     if not dry:
-        with checkpoint_monitor(stage, checkpoint_func) as monitor:
-            cmd_run(stage, checkpoint=monitor is not None)
+        cmd_run(stage, checkpoint_func=checkpoint_func)
 
 
 class CheckpointCond:
@@ -126,14 +127,17 @@ class CheckpointCond:
 
 
 @contextmanager
-def checkpoint_monitor(stage, callback_func):
+def checkpoint_monitor(stage, callback_func, proc):
     if not callback_func:
         yield None
         return
 
+    logger.debug(
+        "Monitoring checkpoint stage '%s' with cmd process '%s'", stage, proc
+    )
     done_cond = CheckpointCond()
     monitor_thread = threading.Thread(
-        target=_checkpoint_run, args=(stage, callback_func, done_cond),
+        target=_checkpoint_run, args=(stage, callback_func, done_cond, proc),
     )
 
     try:
@@ -144,14 +148,21 @@ def checkpoint_monitor(stage, callback_func):
         monitor_thread.join()
 
 
-def _checkpoint_run(stage, callback_func, done_cond):
+def _checkpoint_run(stage, callback_func, done_cond, proc):
     """Run callback_func whenever checkpoint signal file is present."""
     signal_path = os.path.join(stage.repo.tmp_dir, CHECKPOINT_SIGNAL_FILE)
     while True:
         if os.path.exists(signal_path):
-            _run_callback(stage, callback_func)
-            logger.debug("Remove checkpoint signal file")
-            os.remove(signal_path)
+            try:
+                _run_callback(stage, callback_func)
+            except Exception:  # pylint: disable=broad-except
+                logger.exception(
+                    "Error generating checkpoint, %s will be aborted", stage
+                )
+                proc.terminate()
+            finally:
+                logger.debug("Remove checkpoint signal file")
+                os.remove(signal_path)
         if done_cond.wait(1):
             return
 

tests/func/test_repro.py

@@ -1302,4 +1302,4 @@ def test_repro_when_cmd_changes(tmp_dir, dvc, run_copy, mocker):
         stage.addressing: ["changed checksum"]
     }
     assert dvc.reproduce(stage.addressing)[0] == stage
-    m.assert_called_once_with(stage, checkpoint=False)
+    m.assert_called_once_with(stage, checkpoint_func=None)

tests/func/test_repro_multistage.py

@@ -286,7 +286,7 @@ def test_repro_when_cmd_changes(tmp_dir, dvc, run_copy, mocker):
 
     assert dvc.status([target]) == {target: ["changed command"]}
     assert dvc.reproduce(target)[0] == stage
-    m.assert_called_once_with(stage, checkpoint=False)
+    m.assert_called_once_with(stage, checkpoint_func=None)
 
 
 def test_repro_when_new_deps_is_added_in_dvcfile(tmp_dir, dvc, run_copy):

tests/func/test_stage.py

@@ -288,4 +288,4 @@ def test_stage_run_checkpoint(tmp_dir, dvc, mocker, checkpoint):
     else:
         callback = None
     run_stage(stage, checkpoint_func=callback)
-    mock_cmd_run.assert_called_with(stage, checkpoint=checkpoint)
+    mock_cmd_run.assert_called_with(stage, checkpoint_func=callback)