diff --git a/dvc/repo/experiments/__init__.py b/dvc/repo/experiments/__init__.py index 18f44202c..791b94b25 100644 --- a/dvc/repo/experiments/__init__.py +++ b/dvc/repo/experiments/__init__.py @@ -11,6 +11,7 @@ from typing import Iterable, Mapping, Optional from funcy import cached_property, first +from dvc.dvcfile import is_lock_file from dvc.exceptions import DvcException from dvc.path_info import PathInfo from dvc.stage.run import CheckpointKilledError @@ -147,11 +148,14 @@ class Experiments: the human-readable name in the experiment branch ref. Has no effect of branch is specified. """ - with self.scm.stash_workspace(include_untracked=True) as workspace: + with self.scm.stash_workspace( + include_untracked=detach_rev or branch + ) as workspace: # If we are not extending an existing branch, apply current # workspace changes to be made in new branch - if not branch and workspace: + if not (branch or detach_rev) and workspace: self.stash.apply(workspace) + self._prune_lockfiles() # checkout and detach at branch (or current HEAD) if detach_rev: @@ -186,12 +190,24 @@ class Experiments: baseline_rev[:7], ) - # Reset/clean any changes before prior workspace is unstashed - self.scm.gitpython.repo.git.reset(hard=True) - self.scm.gitpython.repo.git.clean(force=True) + # Reset any changes before prior workspace is unstashed + self.scm.reset(hard=True) return stash_rev + def _prune_lockfiles(self): + # NOTE: dirty DVC lock files must be restored to index state to + # avoid checking out incorrect persist or checkpoint outs + tree = self.scm.get_tree("HEAD") + lock_files = [ + str(fname) + for fname in tree.walk_files(self.scm.root_dir) + if is_lock_file(fname) + ] + if lock_files: + self.scm.reset(paths=lock_files) + self.scm.checkout_paths(lock_files, force=True) + def _stash_msg( self, rev: str, @@ -404,33 +420,29 @@ class Experiments: def _init_executors(self, to_run): executors = {} - with self.scm.stash_workspace(include_untracked=True): - with self.scm.detach_head(): - for stash_rev, item in to_run.items(): - self.scm.set_ref(EXEC_HEAD, item.rev) - self.scm.set_ref(EXEC_MERGE, stash_rev) - self.scm.set_ref(EXEC_BASELINE, item.baseline_rev) + for stash_rev, item in to_run.items(): + self.scm.set_ref(EXEC_HEAD, item.rev) + self.scm.set_ref(EXEC_MERGE, stash_rev) + self.scm.set_ref(EXEC_BASELINE, item.baseline_rev) - # Executor will be initialized with an empty git repo that - # we populate by pushing: - # EXEC_HEAD - the base commit for this experiment - # EXEC_MERGE - the unmerged changes (from our stash) - # to be reproduced - # EXEC_BASELINE - the baseline commit for this experiment - executor = LocalExecutor( - self.scm, - self.dvc_dir, - name=item.name, - branch=item.branch, - cache_dir=self.repo.cache.local.cache_dir, - ) - executors[item.rev] = executor + # Executor will be initialized with an empty git repo that + # we populate by pushing: + # EXEC_HEAD - the base commit for this experiment + # EXEC_MERGE - the unmerged changes (from our stash) + # to be reproduced + # EXEC_BASELINE - the baseline commit for this experiment + executor = LocalExecutor( + self.scm, + self.dvc_dir, + name=item.name, + branch=item.branch, + cache_dir=self.repo.cache.local.cache_dir, + ) + executors[item.rev] = executor - for ref in (EXEC_HEAD, EXEC_MERGE, EXEC_BASELINE): - self.scm.remove_ref(ref) + for ref in (EXEC_HEAD, EXEC_MERGE, EXEC_BASELINE): + self.scm.remove_ref(ref) - self.scm.gitpython.repo.git.reset(hard=True) - self.scm.gitpython.repo.git.clean(force=True) return executors def _reproduce( diff --git a/dvc/repo/experiments/executor.py b/dvc/repo/experiments/executor.py index 32b74aac6..70434fd5b 100644 --- a/dvc/repo/experiments/executor.py +++ b/dvc/repo/experiments/executor.py @@ -7,7 +7,6 @@ from typing import TYPE_CHECKING, Callable, Iterable, Optional, Tuple, Union from funcy import cached_property -from dvc.dvcfile import is_lock_file from dvc.path_info import PathInfo from dvc.repo import Repo from dvc.repo.experiments.base import ( @@ -23,7 +22,6 @@ from dvc.repo.experiments.base import ( UnchangedExperimentError, ) from dvc.scm import SCM -from dvc.scm.git import Git from dvc.stage import PipelineStage from dvc.stage.serialize import to_lockfile from dvc.utils import dict_sha256 @@ -32,6 +30,8 @@ from dvc.utils.fs import remove if TYPE_CHECKING: from multiprocessing import Queue + from dvc.scm.git import Git + logger = logging.getLogger(__name__) @@ -88,28 +88,9 @@ class BaseExecutor: self.scm.gitpython.repo.git.merge( EXEC_MERGE, squash=True, no_commit=True ) - self.scm.gitpython.repo.git.reset() - self._prune_lockfiles() finally: os.chdir(cwd) - def _prune_lockfiles(self): - # 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.gitpython.repo.index.diff(None) - ] - to_checkout = [fname for fname in dirty if is_lock_file(fname)] - self.scm.gitpython.repo.index.checkout(paths=to_checkout, force=True) - - untracked = self.scm.gitpython.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 - @cached_property def scm(self): return SCM(self.root_dir) @@ -260,10 +241,7 @@ class BaseExecutor: # experiment run dvc.checkout(force=True, quiet=True) - # We cannot use dvc.scm to make commits inside the executor since - # cached props are not picklable. - scm = Git() - checkpoint_func = partial(cls.checkpoint_callback, scm, name) + checkpoint_func = partial(cls.checkpoint_callback, dvc.scm, name) stages = dvc.reproduce( *args, on_unchanged=filter_pipeline, @@ -272,15 +250,14 @@ class BaseExecutor: ) exp_hash = cls.hash_exp(stages) - exp_rev = cls.commit(scm, exp_hash, exp_name=name) - if scm.get_ref(EXEC_CHECKPOINT): - scm.set_ref(EXEC_CHECKPOINT, exp_rev) + exp_rev = cls.commit(dvc.scm, exp_hash, exp_name=name) + if dvc.scm.get_ref(EXEC_CHECKPOINT): + dvc.scm.set_ref(EXEC_CHECKPOINT, exp_rev) except UnchangedExperimentError: pass finally: - if scm: - scm.close() - del scm + if dvc: + dvc.scm.close() if old_cwd: os.chdir(old_cwd) @@ -324,7 +301,7 @@ class BaseExecutor: old_ref = None logger.debug("Commit to new experiment branch '%s'", branch) - scm.gitpython.repo.git.add(A=True) + scm.gitpython.repo.git.add(update=True) scm.commit(f"dvc: commit experiment {exp_hash}") new_rev = scm.get_rev() scm.set_ref(branch, new_rev, old_ref=old_ref) @@ -354,7 +331,6 @@ class LocalExecutor(BaseExecutor): local_config = os.path.join(self.dvc_dir, "config.local") logger.debug("Writing experiments local config '%s'", local_config) with open(local_config, "w") as fobj: - fobj.write("[core]\n no_scm = true\n") fobj.write(f"[cache]\n dir = {cache_dir}") @property diff --git a/dvc/scm/git/__init__.py b/dvc/scm/git/__init__.py index ef12cb68b..864daafd6 100644 --- a/dvc/scm/git/__init__.py +++ b/dvc/scm/git/__init__.py @@ -326,6 +326,8 @@ class Git(Base): reflog_delete = partialmethod(_backend_func, "reflog_delete") describe = partialmethod(_backend_func, "describe") diff = partialmethod(_backend_func, "diff") + reset = partialmethod(_backend_func, "reset") + checkout_paths = partialmethod(_backend_func, "checkout_paths") def resolve_rev(self, rev: str) -> str: from dvc.repo.experiments.utils import exp_refs_by_name diff --git a/dvc/scm/git/backend/base.py b/dvc/scm/git/backend/base.py index dbf179654..49793823a 100644 --- a/dvc/scm/git/backend/base.py +++ b/dvc/scm/git/backend/base.py @@ -259,3 +259,11 @@ class BaseGitBackend(ABC): @abstractmethod def diff(self, rev_a: str, rev_b: str, binary=False) -> str: """Return the git diff for two commits.""" + + @abstractmethod + def reset(self, hard: bool = False, paths: Iterable[str] = None): + """Reset current git HEAD.""" + + @abstractmethod + def checkout_paths(self, paths: Iterable[str], force: bool = False): + """Checkout the specified paths from HEAD index.""" diff --git a/dvc/scm/git/backend/dulwich.py b/dvc/scm/git/backend/dulwich.py index 67ccef4a6..37888168b 100644 --- a/dvc/scm/git/backend/dulwich.py +++ b/dvc/scm/git/backend/dulwich.py @@ -340,3 +340,9 @@ class DulwichBackend(BaseGitBackend): # pylint:disable=abstract-method buf, self.repo.object_store, commit_a.tree, commit_b.tree ) return buf.getvalue().decode("utf-8") + + def reset(self, hard: bool = False, paths: Iterable[str] = None): + raise NotImplementedError + + def checkout_paths(self, paths: Iterable[str], force: bool = False): + raise NotImplementedError diff --git a/dvc/scm/git/backend/gitpython.py b/dvc/scm/git/backend/gitpython.py index 2c1afdb09..0a28739aa 100644 --- a/dvc/scm/git/backend/gitpython.py +++ b/dvc/scm/git/backend/gitpython.py @@ -394,3 +394,11 @@ class GitPythonBackend(BaseGitBackend): # pylint:disable=abstract-method def diff(self, rev_a: str, rev_b: str, binary=False) -> str: raise NotImplementedError + + def reset(self, hard: bool = False, paths: Iterable[str] = None): + self.repo.head.reset(index=True, working_tree=hard, paths=paths) + + def checkout_paths(self, paths: Iterable[str], force: bool = False): + """Checkout the specified paths from HEAD index.""" + if paths: + self.repo.index.checkout(paths=paths, force=force) diff --git a/dvc/stage/run.py b/dvc/stage/run.py index ad0d0b833..23a683c23 100644 --- a/dvc/stage/run.py +++ b/dvc/stage/run.py @@ -186,5 +186,9 @@ def _kill_nt(proc): def _run_callback(stage, callback_func): stage.save(allow_missing=True) stage.commit(allow_missing=True) + for out in stage.outs: + if not out.use_scm_ignore and out.is_in_repo: + stage.repo.scm.track_file(os.fspath(out.path_info)) + stage.repo.scm.track_changed_files() logger.debug("Running checkpoint callback for stage '%s'", stage) callback_func() diff --git a/tests/func/experiments/test_checkpoints.py b/tests/func/experiments/test_checkpoints.py index 1f04fc4fd..366bc92e1 100644 --- a/tests/func/experiments/test_checkpoints.py +++ b/tests/func/experiments/test_checkpoints.py @@ -13,11 +13,14 @@ def test_new_checkpoint(tmp_dir, scm, dvc, checkpoint_stage, mocker): exp = first(results) new_mock.assert_called_once() - tree = scm.get_tree(exp) - with tree.open(tmp_dir / "foo") as fobj: - assert fobj.read().strip() == "5" - with tree.open(tmp_dir / "metrics.yaml") as fobj: - assert fobj.read().strip() == "foo: 2" + for rev in dvc.brancher([exp]): + if rev == "workspace": + continue + tree = dvc.repo_tree + with tree.open(tmp_dir / "foo") as fobj: + assert fobj.read().strip() == "5" + with tree.open(tmp_dir / "metrics.yaml") as fobj: + assert fobj.read().strip() == "foo: 2" @pytest.mark.parametrize( @@ -49,11 +52,14 @@ def test_resume_checkpoint( ) exp = first(results) - tree = scm.get_tree(exp) - with tree.open(tmp_dir / "foo") as fobj: - assert fobj.read().strip() == "10" - with tree.open(tmp_dir / "metrics.yaml") as fobj: - assert fobj.read().strip() == "foo: 2" + for rev in dvc.brancher([exp]): + if rev == "workspace": + continue + tree = dvc.repo_tree + with tree.open(tmp_dir / "foo") as fobj: + assert fobj.read().strip() == "10" + with tree.open(tmp_dir / "metrics.yaml") as fobj: + assert fobj.read().strip() == "foo: 2" def test_reset_checkpoint(tmp_dir, scm, dvc, checkpoint_stage, caplog): @@ -73,11 +79,14 @@ def test_reset_checkpoint(tmp_dir, scm, dvc, checkpoint_stage, caplog): ) exp = first(results) - tree = scm.get_tree(exp) - with tree.open(tmp_dir / "foo") as fobj: - assert fobj.read().strip() == "5" - with tree.open(tmp_dir / "metrics.yaml") as fobj: - assert fobj.read().strip() == "foo: 2" + for rev in dvc.brancher([exp]): + if rev == "workspace": + continue + tree = dvc.repo_tree + with tree.open(tmp_dir / "foo") as fobj: + assert fobj.read().strip() == "5" + with tree.open(tmp_dir / "metrics.yaml") as fobj: + assert fobj.read().strip() == "foo: 2" def test_resume_branch(tmp_dir, scm, dvc, checkpoint_stage): @@ -98,17 +107,23 @@ def test_resume_branch(tmp_dir, scm, dvc, checkpoint_stage): ) checkpoint_b = first(results) - tree = scm.get_tree(checkpoint_a) - with tree.open(tmp_dir / "foo") as fobj: - assert fobj.read().strip() == "10" - with tree.open(tmp_dir / "metrics.yaml") as fobj: - assert fobj.read().strip() == "foo: 2" + for rev in dvc.brancher([checkpoint_a]): + if rev == "workspace": + continue + tree = dvc.repo_tree + with tree.open(tmp_dir / "foo") as fobj: + assert fobj.read().strip() == "10" + with tree.open(tmp_dir / "metrics.yaml") as fobj: + assert fobj.read().strip() == "foo: 2" - tree = scm.get_tree(checkpoint_b) - with tree.open(tmp_dir / "foo") as fobj: - assert fobj.read().strip() == "10" - with tree.open(tmp_dir / "metrics.yaml") as fobj: - assert fobj.read().strip() == "foo: 100" + for rev in dvc.brancher([checkpoint_b]): + if rev == "workspace": + continue + tree = dvc.repo_tree + with tree.open(tmp_dir / "foo") as fobj: + assert fobj.read().strip() == "10" + with tree.open(tmp_dir / "metrics.yaml") as fobj: + assert fobj.read().strip() == "foo: 100" with pytest.raises(MultipleBranchError): dvc.experiments.get_branch_by_rev(branch_rev) diff --git a/tests/func/experiments/test_experiments.py b/tests/func/experiments/test_experiments.py index 760869a2b..f24a60f64 100644 --- a/tests/func/experiments/test_experiments.py +++ b/tests/func/experiments/test_experiments.py @@ -337,3 +337,49 @@ def test_no_scm(tmp_dir): with pytest.raises(NoSCMError): dvc.experiments.gc() + + +def test_untracked(tmp_dir, scm, dvc, caplog): + tmp_dir.gen("copy.py", COPY_SCRIPT) + tmp_dir.gen("params.yaml", "foo: 1") + stage = dvc.run( + cmd="python copy.py params.yaml metrics.yaml", + metrics_no_cache=["metrics.yaml"], + params=["foo"], + name="copy-file", + ) + scm.add(["dvc.yaml", "dvc.lock", "params.yaml", "metrics.yaml"]) + scm.commit("init") + + # copy.py is untracked + with caplog.at_level(logging.ERROR): + results = dvc.experiments.run(stage.addressing, params=["foo=2"]) + assert "Failed to reproduce experiment" in caplog.text + assert not results + + # copy.py is staged as new file but not committed + scm.add(["copy.py"]) + results = dvc.experiments.run(stage.addressing, params=["foo=2"]) + exp = first(results) + tree = scm.get_tree(exp) + assert tree.exists("copy.py") + with tree.open(tmp_dir / "metrics.yaml") as fobj: + assert fobj.read().strip() == "foo: 2" + + +def test_dirty_lockfile(tmp_dir, scm, dvc, exp_stage): + from dvc.dvcfile import LockfileCorruptedError + + tmp_dir.gen("dvc.lock", "foo") + + with pytest.raises(LockfileCorruptedError): + dvc.reproduce(exp_stage.addressing) + + results = dvc.experiments.run(exp_stage.addressing, params=["foo=2"]) + exp = first(results) + + tree = scm.get_tree(exp) + with tree.open(tmp_dir / "metrics.yaml") as fobj: + assert fobj.read().strip() == "foo: 2" + + assert (tmp_dir / "dvc.lock").read_text() == "foo"