diff --git a/dvc/repo/experiments/__init__.py b/dvc/repo/experiments/__init__.py index 791b94b25..2e5030101 100644 --- a/dvc/repo/experiments/__init__.py +++ b/dvc/repo/experiments/__init__.py @@ -164,34 +164,38 @@ class Experiments: head = branch else: head = None - with self.scm.detach_head(head) as rev: - if baseline_rev is None: - baseline_rev = rev - # update experiment params from command line - if params: - self._update_params(params) + try: + with self.scm.detach_head(head) as rev: + if baseline_rev is None: + baseline_rev = rev - # save additional repro command line arguments - self._pack_args(*args, **kwargs) + # update experiment params from command line + if params: + self._update_params(params) - # save experiment as a stash commit - msg = self._stash_msg( - rev, baseline_rev=baseline_rev, branch=branch, name=name - ) - stash_rev = self.stash.push(message=msg) + # save additional repro command line arguments + self._pack_args(*args, **kwargs) - logger.debug( - ( - "Stashed experiment '%s' with baseline '%s' " - "for future execution." - ), - stash_rev[:7], - baseline_rev[:7], - ) - - # Reset any changes before prior workspace is unstashed - self.scm.reset(hard=True) + # save experiment as a stash commit + msg = self._stash_msg( + rev, + baseline_rev=baseline_rev, + branch=branch, + name=name, + ) + stash_rev = self.stash.push(message=msg) + logger.debug( + ( + "Stashed experiment '%s' with baseline '%s' " + "for future execution." + ), + stash_rev[:7], + baseline_rev[:7], + ) + finally: + # Reset any changes before prior workspace is unstashed + self.scm.reset(hard=True) return stash_rev @@ -225,7 +229,32 @@ class Experiments: return msg def _pack_args(self, *args, **kwargs): - BaseExecutor.pack_repro_args(self.args_file, *args, **kwargs) + import pickle + + if os.path.exists(self.args_file) and self.scm.is_tracked( + self.args_file + ): + logger.warning( + ( + "Temporary DVC file '.dvc/tmp/%s' exists and was " + "likely committed to Git by mistake. It should be removed " + "with:\n" + "\tgit rm .dvc/tmp/%s" + ), + BaseExecutor.PACKED_ARGS_FILE, + BaseExecutor.PACKED_ARGS_FILE, + ) + with open(self.args_file, "rb") as fobj: + try: + data = pickle.load(fobj) + except Exception: # pylint: disable=broad-except + data = {} + extra = int(data.get("extra", 0)) + 1 + else: + extra = None + BaseExecutor.pack_repro_args( + self.args_file, *args, extra=extra, **kwargs + ) self.scm.add(self.args_file) def _update_params(self, params: dict): diff --git a/dvc/repo/experiments/executor.py b/dvc/repo/experiments/executor.py index 70434fd5b..713c5e417 100644 --- a/dvc/repo/experiments/executor.py +++ b/dvc/repo/experiments/executor.py @@ -117,7 +117,7 @@ class BaseExecutor: # TODO: come up with better way to stash repro arguments @staticmethod - def pack_repro_args(path, *args, tree=None, **kwargs): + def pack_repro_args(path, *args, tree=None, extra=None, **kwargs): dpath = os.path.dirname(path) if tree: open_func = tree.open @@ -127,7 +127,10 @@ class BaseExecutor: open_func = open makedirs(dpath, exist_ok=True) + data = {"args": args, "kwargs": kwargs} + if extra is not None: + data["extra"] = extra with open_func(path, "wb") as fobj: pickle.dump(data, fobj) diff --git a/dvc/scm/git/backend/dulwich.py b/dvc/scm/git/backend/dulwich.py index 37888168b..7228fa4b3 100644 --- a/dvc/scm/git/backend/dulwich.py +++ b/dvc/scm/git/backend/dulwich.py @@ -302,8 +302,8 @@ class DulwichBackend(BaseGitBackend): # pylint:disable=abstract-method stash = self._get_stash(ref) message_b = message.encode("utf-8") if message else None - stash.push(message=message_b) - return os.fsdecode(stash[0].new_sha), True + rev = stash.push(message=message_b) + return os.fsdecode(rev), True def _stash_apply(self, rev: str): raise NotImplementedError diff --git a/dvc/scm/git/stash.py b/dvc/scm/git/stash.py index 476d9b7e3..00de7cb80 100644 --- a/dvc/scm/git/stash.py +++ b/dvc/scm/git/stash.py @@ -32,9 +32,7 @@ class Stash: return list(iter(self)) def push( - self, - message: Optional[str] = None, - include_untracked: Optional[bool] = False, + self, message: Optional[str] = None, include_untracked: bool = False, ) -> Optional[str]: if not self.scm.is_dirty(untracked_files=include_untracked): logger.debug("No changes to stash") diff --git a/tests/func/experiments/test_experiments.py b/tests/func/experiments/test_experiments.py index f24a60f64..3036cddc5 100644 --- a/tests/func/experiments/test_experiments.py +++ b/tests/func/experiments/test_experiments.py @@ -383,3 +383,17 @@ def test_dirty_lockfile(tmp_dir, scm, dvc, exp_stage): assert fobj.read().strip() == "foo: 2" assert (tmp_dir / "dvc.lock").read_text() == "foo" + + +def test_packed_args_exists(tmp_dir, scm, dvc, exp_stage, caplog): + from dvc.repo.experiments.executor import BaseExecutor + + tmp_dir.scm_gen( + tmp_dir / ".dvc" / "tmp" / BaseExecutor.PACKED_ARGS_FILE, + "", + commit="commit args file", + ) + + with caplog.at_level(logging.WARNING): + dvc.experiments.run(exp_stage.addressing) + assert "Temporary DVC file" in caplog.text