exp: add ApplyStash implementation

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 = {}

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:

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<head_rev>[0-9a-f]+):(?P<rev>[0-9a-f]+)"
+        r":(?P<name>[^~^:\\?\[\]*]*)"
+    )
+
+    @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)

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,