From 303da84b18872aa0d1b7cc3dbf048ef5a67d36ec Mon Sep 17 00:00:00 2001 From: Peter Rowlands Date: Thu, 9 Feb 2023 16:46:45 +0900 Subject: [PATCH] exp: add ApplyStash implementation --- dvc/repo/experiments/__init__.py | 6 ++ dvc/repo/experiments/refs.py | 4 + dvc/repo/experiments/stash.py | 131 ++++++++++++++++++++++++++++++- dvc/repo/experiments/utils.py | 19 +++-- 4 files changed, 153 insertions(+), 7 deletions(-) diff --git a/dvc/repo/experiments/__init__.py b/dvc/repo/experiments/__init__.py index b995febb7..20c26f784 100644 --- a/dvc/repo/experiments/__init__.py +++ b/dvc/repo/experiments/__init__.py @@ -22,6 +22,7 @@ from .queue.celery import LocalCeleryQueue from .queue.tempdir import TempDirQueue from .queue.workspace import WorkspaceQueue from .refs import ( + APPLY_STASH, CELERY_FAILED_STASH, CELERY_STASH, EXEC_APPLY, @@ -31,6 +32,7 @@ from .refs import ( WORKSPACE_STASH, ExpRefInfo, ) +from .stash import ApplyStash from .utils import check_ref_format, exp_refs_by_rev, unlocked_repo if TYPE_CHECKING: @@ -86,6 +88,10 @@ class Experiments: def celery_queue(self) -> LocalCeleryQueue: return LocalCeleryQueue(self.repo, CELERY_STASH, CELERY_FAILED_STASH) + @cached_property + def apply_stash(self) -> ApplyStash: + return ApplyStash(self.scm, APPLY_STASH) + @property def stash_revs(self) -> Dict[str, "ExpStashEntry"]: revs = {} diff --git a/dvc/repo/experiments/refs.py b/dvc/repo/experiments/refs.py index 87ef230bc..dd9c4b215 100644 --- a/dvc/repo/experiments/refs.py +++ b/dvc/repo/experiments/refs.py @@ -7,6 +7,9 @@ from .exceptions import InvalidExpRefError EXPS_NAMESPACE = "refs/exps" EXPS_STASH = f"{EXPS_NAMESPACE}/stash" WORKSPACE_STASH = EXPS_STASH +APPLY_NAMESPACE = f"{EXPS_NAMESPACE}/apply" +APPLY_HEAD = f"{APPLY_NAMESPACE}/ORIG_HEAD" +APPLY_STASH = f"{APPLY_NAMESPACE}/stash" CELERY_STASH = f"{EXPS_NAMESPACE}/celery/stash" CELERY_FAILED_STASH = f"{EXPS_NAMESPACE}/celery/failed" EXEC_NAMESPACE = f"{EXPS_NAMESPACE}/exec" @@ -18,6 +21,7 @@ EXEC_HEAD = f"{EXEC_NAMESPACE}/EXEC_HEAD" EXEC_MERGE = f"{EXEC_NAMESPACE}/EXEC_MERGE" EXPS_TEMP = f"{EXPS_NAMESPACE}/temp" STASHES = {WORKSPACE_STASH, CELERY_STASH} +ITER_SKIP_NAMESPACES = {APPLY_NAMESPACE, EXEC_NAMESPACE} class ExpRefInfo: diff --git a/dvc/repo/experiments/stash.py b/dvc/repo/experiments/stash.py index 7678c8d28..332701bdc 100644 --- a/dvc/repo/experiments/stash.py +++ b/dvc/repo/experiments/stash.py @@ -1,8 +1,18 @@ +import logging import re -from typing import Dict, Iterable, NamedTuple, Optional +from contextlib import contextmanager +from typing import Dict, Iterable, Iterator, NamedTuple, Optional from scmrepo.git import Stash +from dvc.exceptions import DvcException +from dvc_objects.fs.local import localfs +from dvc_objects.fs.utils import as_atomic + +from .refs import APPLY_HEAD, APPLY_STASH + +logger = logging.getLogger(__name__) + class ExpStashEntry(NamedTuple): """Experiment stash entry. @@ -72,3 +82,122 @@ class ExpStash(Stash): reverse=True, ): self.drop(index) + + +class ApplyStashEntry(NamedTuple): + """Apply stash entry. + + stash_index: Stash index for this entry. Can be None if this commit + is not pushed onto the stash ref. + head_rev: HEAD Git commit prior to exp apply. + rev: Applied experiment commit. + name: Optional applied exp name. + """ + + stash_index: Optional[int] + head_rev: str + rev: str + name: Optional[str] + + +class ApplyStash(Stash): + DEFAULT_STASH = APPLY_STASH + MESSAGE_FORMAT = "dvc-exp-apply:{head_rev}:{rev}:{name}" + MESSAGE_RE = re.compile( + r"(?:commit: )" + r"dvc-exp-apply:(?P[0-9a-f]+):(?P[0-9a-f]+)" + r":(?P[^~^:\\?\[\]*]*)" + ) + + @property + def stash_revs(self) -> Dict[str, ApplyStashEntry]: + revs = {} + for i, entry in enumerate(self): + msg = entry.message.decode("utf-8").strip() + m = self.MESSAGE_RE.match(msg) + if m: + revs[entry.new_sha.decode("utf-8")] = ApplyStashEntry( + i, + m.group("head_rev"), + m.group("rev"), + m.group("name"), + ) + return revs + + @classmethod + def format_message( + cls, + head_rev: str, + rev: str, + name: Optional[str] = None, + ) -> str: + return cls.MESSAGE_FORMAT.format( + head_rev=head_rev, rev=rev, name=name if name else "" + ) + + @contextmanager + def preserve_workspace( + self, rev: str, name: Optional[str] = None + ) -> Iterator[Optional[str]]: + if len(self): + logger.debug("Clearing existing exp-apply stash") + self.clear() + head = self.scm.get_rev() + self.scm.set_ref(APPLY_HEAD, head) + message = self.format_message(head, rev, name=name) + stash_rev = self.push(message=message, include_untracked=True) + try: + yield stash_rev + if stash_rev: + self._apply_difference(stash_rev, rev) + except Exception: # pylint: disable=broad-except + self.revert_workspace() + raise + + def _apply_difference(self, stash_rev: str, rev: str): + """Selectively apply changes from stash_rev. + + Only changes to files from left which do not exist in right will be applied. + """ + self._copy_difference(stash_rev, rev) + commit = self.scm.resolve_commit(stash_rev) + for parent_rev in commit.parents: + parent_commit = self.scm.resolve_commit(parent_rev) + if parent_commit.message.startswith("untracked files on "): + self._copy_difference(parent_rev, rev) + + def _copy_difference(self, left_rev: str, right_rev: str): + left_fs = self.scm.get_fs(left_rev) + right_fs = self.scm.get_fs(right_rev) + paths = [path for path in left_fs.find("/") if not right_fs.exists(path)] + dest_paths = [ + localfs.path.join(self.scm.root_dir, left_fs.path.relpath(path, "/")) + for path in paths + ] + for src, dest in zip(paths, dest_paths): + with as_atomic(localfs, dest, create_parents=True) as tmp_file: + left_fs.get_file(src, tmp_file) + + def revert_workspace(self): + apply_head = self.scm.get_ref(self.ref) + head = self.scm.get_rev() + if apply_head != head: + raise DvcException( + f"Cannot revert workspace, current HEAD '{head[:7]}' does not match the" + f" pre-apply HEAD '{apply_head[:7]}'" + ) + self.scm.reset(hard=True) + if len(self): + # In the event that the apply-stash and current workspace contain + # conflicting untracked files, we do: + # 1. stash the current untracked files + # 2. restore/pop the apply-stash (with untracked files) + # 3. restore/pop the untracked files from (1) and ignore any conflicts + # (forcefully reverting to the apply-stash version) + workspace_rev = self.scm.stash.push(include_untracked=True) + try: + self.pop() + finally: + if workspace_rev: + self.scm.stash.pop(skip_conflicts=True) + self.scm.remove_ref(self.ref) diff --git a/dvc/repo/experiments/utils.py b/dvc/repo/experiments/utils.py index 4d5a14480..9f3902190 100644 --- a/dvc/repo/experiments/utils.py +++ b/dvc/repo/experiments/utils.py @@ -27,8 +27,8 @@ from .refs import ( EXEC_BASELINE, EXEC_BRANCH, EXEC_CHECKPOINT, - EXEC_NAMESPACE, EXPS_NAMESPACE, + ITER_SKIP_NAMESPACES, STASHES, ExpRefInfo, ) @@ -78,6 +78,13 @@ def unlocked_repo(f): return wrapper +def _ignore_ref(ref: str) -> bool: + return ( + any(ref.startswith(namespace) for namespace in ITER_SKIP_NAMESPACES) + or ref in STASHES + ) + + def exp_refs( scm: "Git", url: Optional[str] = None ) -> Generator["ExpRefInfo", None, None]: @@ -88,7 +95,7 @@ def exp_refs( else scm.iter_refs(base=EXPS_NAMESPACE) ) for ref in ref_gen: - if ref.startswith(EXEC_NAMESPACE) or ref in STASHES: + if _ignore_ref(ref): continue yield ExpRefInfo.from_ref(ref) @@ -96,7 +103,7 @@ def exp_refs( def exp_refs_by_rev(scm: "Git", rev: str) -> Generator[ExpRefInfo, None, None]: """Iterate over all experiment refs pointing to the specified revision.""" for ref in scm.get_refs_containing(rev, EXPS_NAMESPACE): - if not (ref.startswith(EXEC_NAMESPACE) or ref in STASHES): + if not _ignore_ref(ref): yield ExpRefInfo.from_ref(ref) @@ -169,7 +176,7 @@ def push_refspec( def remote_exp_refs(scm: "Git", url: str) -> Generator[ExpRefInfo, None, None]: """Iterate over all remote experiment refs.""" for ref in iter_remote_refs(scm, url, base=EXPS_NAMESPACE): - if ref.startswith(EXEC_NAMESPACE) or ref in STASHES: + if _ignore_ref(ref): continue yield ExpRefInfo.from_ref(ref) @@ -193,7 +200,7 @@ def remote_exp_refs_by_baseline( """Iterate over all remote experiment refs with the specified baseline.""" ref_info = ExpRefInfo(baseline_sha=rev) for ref in iter_remote_refs(scm, url, base=str(ref_info)): - if ref.startswith(EXEC_NAMESPACE) or ref in STASHES: + if _ignore_ref(ref): continue yield ExpRefInfo.from_ref(ref) @@ -283,7 +290,7 @@ def fetch_all_exps(scm: "Git", url: str, progress: Optional[Callable] = None): refspecs = [ f"{ref}:{ref}" for ref in iter_remote_refs(scm, url, base=EXPS_NAMESPACE) - if not (ref.startswith(EXEC_NAMESPACE) or ref in STASHES) + if not _ignore_ref(ref) ] scm.fetch_refspecs( url,