diff --git a/dvc/repo/experiments/__init__.py b/dvc/repo/experiments/__init__.py index a006c25b8..361757b42 100644 --- a/dvc/repo/experiments/__init__.py +++ b/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 diff --git a/dvc/scm/git.py b/dvc/scm/git.py index 09f56b5ef..22318d781 100644 --- a/dvc/scm/git.py +++ b/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 diff --git a/dvc/stage/run.py b/dvc/stage/run.py index 01c064ccb..d839df2a3 100644 --- a/dvc/stage/run.py +++ b/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 diff --git a/tests/func/test_repro.py b/tests/func/test_repro.py index 7b2213e3d..26d353082 100644 --- a/tests/func/test_repro.py +++ b/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) diff --git a/tests/func/test_repro_multistage.py b/tests/func/test_repro_multistage.py index d1bb897fb..e658c7608 100644 --- a/tests/func/test_repro_multistage.py +++ b/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): diff --git a/tests/func/test_stage.py b/tests/func/test_stage.py index 413fcf24c..362c5aa7a 100644 --- a/tests/func/test_stage.py +++ b/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)