Use temp branch name to avoid locking the exec branches in temp dir executor

fix: #8448
In our current temp dir executor, we use will first set EXEC_* branches
to revisions retrieved from stash entry. And push them to remote, but
because the push operation is very slow, our locking operation will make
the whole repository unusable.

1. Now we set revisions from stash entry to some temp branches and then
   push them to remote.

dvc/repo/experiments/executor/base.py

@@ -35,14 +35,7 @@ from ..exceptions import (
     ExperimentExistsError,
     UnchangedExperimentError,
 )
-from ..refs import (
-    EXEC_BASELINE,
-    EXEC_BRANCH,
-    EXEC_CHECKPOINT,
-    EXEC_HEAD,
-    EXEC_MERGE,
-    ExpRefInfo,
-)
+from ..refs import EXEC_BASELINE, EXEC_BRANCH, EXEC_CHECKPOINT, ExpRefInfo
 
 if TYPE_CHECKING:
     from queue import Queue
@@ -744,21 +737,12 @@ class BaseExecutor(ABC):
             dvc_logger.setLevel(level)
 
     @contextmanager
-    def set_exec_refs(
-        self, scm: "Git", stash_rev: str, entry: "ExpStashEntry"
-    ):
+    def set_temp_refs(self, scm: "Git", temp_dict: Dict[str, str]):
         try:
-            # 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
-            scm.set_ref(EXEC_HEAD, entry.head_rev)
-            scm.set_ref(EXEC_MERGE, stash_rev)
-            scm.set_ref(EXEC_BASELINE, entry.baseline_rev)
+            for ref, rev in temp_dict.items():
+                scm.set_ref(ref, rev)
             yield
         finally:
-            for ref in (EXEC_HEAD, EXEC_MERGE, EXEC_BASELINE):
+            for ref in temp_dict:
                 if scm.get_ref(ref):
                     scm.remove_ref(ref)

dvc/repo/experiments/executor/local.py

@@ -6,6 +6,7 @@ from typing import TYPE_CHECKING, Optional
 
 from funcy import cached_property, retry
 from scmrepo.exceptions import SCMError as _SCMError
+from shortuuid import uuid
 
 from dvc.lock import LockError
 from dvc.scm import SCM, GitMergeError
@@ -19,6 +20,7 @@ from ..refs import (
     EXEC_HEAD,
     EXEC_MERGE,
     EXEC_NAMESPACE,
+    EXPS_TEMP,
 )
 from ..utils import EXEC_TMP_DIR, get_exp_rwlock
 from .base import BaseExecutor, TaskStatus
@@ -89,29 +91,53 @@ class TempDirExecutor(BaseLocalExecutor):
         if infofile:
             self.info.dump_json(infofile)
 
-        with get_exp_rwlock(repo, writes=[EXEC_NAMESPACE]):
-            with self.set_exec_refs(scm, stash_rev, entry):
-                refspec = f"{EXEC_NAMESPACE}/"
-                push_refspec(scm, self.git_url, refspec, refspec)
+        temp_head = f"{EXPS_TEMP}/head-{uuid()}"
+        temp_merge = f"{EXPS_TEMP}/merge-{uuid()}"
+        temp_baseline = f"{EXPS_TEMP}/baseline-{uuid()}"
+
+        temp_ref_dict = {
+            temp_head: entry.head_rev,
+            temp_merge: stash_rev,
+            temp_baseline: entry.baseline_rev,
+        }
+        with get_exp_rwlock(
+            repo, writes=[temp_head, temp_merge, temp_baseline]
+        ), self.set_temp_refs(scm, temp_ref_dict):
+            # 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
+            refspec = [
+                (temp_head, EXEC_HEAD),
+                (temp_merge, EXEC_MERGE),
+                (temp_baseline, EXEC_BASELINE),
+            ]
 
             if branch:
-                push_refspec(scm, self.git_url, branch, branch)
+                refspec.append((branch, branch))
+                with get_exp_rwlock(repo, reads=[branch]):
+                    push_refspec(scm, self.git_url, refspec)
                 self.scm.set_ref(EXEC_BRANCH, branch, symbolic=True)
-            elif self.scm.get_ref(EXEC_BRANCH):
-                self.scm.remove_ref(EXEC_BRANCH)
+            else:
+                push_refspec(scm, self.git_url, refspec)
+                if self.scm.get_ref(EXEC_BRANCH):
+                    self.scm.remove_ref(EXEC_BRANCH)
 
             if self.scm.get_ref(EXEC_CHECKPOINT):
                 self.scm.remove_ref(EXEC_CHECKPOINT)
-            # checkout EXEC_HEAD and apply EXEC_MERGE on top of it without
-            # committing
-            head = EXEC_BRANCH if branch else EXEC_HEAD
-            self.scm.checkout(head, detach=True)
-            merge_rev = self.scm.get_ref(EXEC_MERGE)
 
-            try:
-                self.scm.merge(merge_rev, squash=True, commit=False)
-            except _SCMError as exc:
-                raise GitMergeError(str(exc), scm=self.scm)
+        # checkout EXEC_HEAD and apply EXEC_MERGE on top of it without
+        # committing
+        head = EXEC_BRANCH if branch else EXEC_HEAD
+        self.scm.checkout(head, detach=True)
+        merge_rev = self.scm.get_ref(EXEC_MERGE)
+
+        try:
+            self.scm.merge(merge_rev, squash=True, commit=False)
+        except _SCMError as exc:
+            raise GitMergeError(str(exc), scm=self.scm)
 
     def _config(self, cache_dir):
         local_config = os.path.join(

dvc/repo/experiments/executor/ssh.py

@@ -9,6 +9,7 @@ from dvc_ssh import SSHFileSystem
 from funcy import first
 
 from ..refs import (
+    EXEC_BASELINE,
     EXEC_BRANCH,
     EXEC_CHECKPOINT,
     EXEC_HEAD,
@@ -148,12 +149,18 @@ class SSHExecutor(BaseExecutor):
             # (see https://github.com/iterative/dvc/issues/6508)
             kwargs = self._git_client_args(fs)
 
-            with self.set_exec_refs(scm, stash_rev, entry):
-                refspec = f"{EXEC_NAMESPACE}/"
-                push_refspec(scm, self.git_url, refspec, refspec, **kwargs)
+            ref_dict = {
+                EXEC_HEAD: entry.head_rev,
+                EXEC_MERGE: stash_rev,
+                EXEC_BASELINE: entry.baseline_rev,
+            }
+            with self.set_temp_refs(scm, ref_dict):
+                exec_namespace = f"{EXEC_NAMESPACE}/"
+                refspec = [(exec_namespace, exec_namespace)]
+                push_refspec(scm, self.git_url, refspec, **kwargs)
 
             if branch:
-                push_refspec(scm, self.git_url, branch, branch, **kwargs)
+                push_refspec(scm, self.git_url, [(branch, branch)], **kwargs)
                 self._ssh_cmd(fs, f"git symbolic-ref {EXEC_BRANCH} {branch}")
             else:
                 self._ssh_cmd(

dvc/repo/experiments/refs.py

@@ -16,6 +16,7 @@ EXEC_BRANCH = f"{EXEC_NAMESPACE}/EXEC_BRANCH"
 EXEC_BASELINE = f"{EXEC_NAMESPACE}/EXEC_BASELINE"
 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}
 
 

dvc/repo/experiments/remove.py

@@ -152,8 +152,7 @@ def _remove_commited_exps(
                 push_refspec(
                     scm,
                     remote,
-                    None,
-                    str(ref_info),
+                    [(None, str(ref_info))],
                     progress=pbar.update_git,
                 )
     else:

dvc/repo/experiments/utils.py

@@ -12,6 +12,7 @@ from typing import (
     Mapping,
     Optional,
     Set,
+    Tuple,
     Union,
 )
 
@@ -128,8 +129,7 @@ def iter_remote_refs(
 def push_refspec(
     scm: "Git",
     url: str,
-    src: Optional[str],
-    dest: str,
+    push_list=List[Tuple[Optional[str], str]],
     force: bool = False,
     on_diverged: Optional[Callable[[str, str], bool]] = None,
     **kwargs,
@@ -139,20 +139,21 @@ def push_refspec(
 
     from ...scm import GitAuthError, SCMError
 
-    if not src:
-        refspecs = [f":{dest}"]
-    elif src.endswith("/"):
-        refspecs = []
-        dest = dest.rstrip("/") + "/"
-        for ref in scm.iter_refs(base=src):
-            refname = ref.split("/")[-1]
-            refspecs.append(f"{ref}:{dest}{refname}")
-    else:
-        if dest.endswith("/"):
-            refname = src.split("/")[-1]
-            refspecs = [f"{src}:{dest}/{refname}"]
+    refspecs = []
+    for src, dest in push_list:
+        if not src:
+            refspecs.append(f":{dest}")
+        elif src.endswith("/"):
+            dest = dest.rstrip("/") + "/"
+            for ref in scm.iter_refs(base=src):
+                refname = ref.split("/")[-1]
+                refspecs.append(f"{ref}:{dest}{refname}")
         else:
-            refspecs = [f"{src}:{dest}"]
+            if dest.endswith("/"):
+                refname = src.split("/")[-1]
+                refspecs.append(f"{src}:{dest}/{refname}")
+            else:
+                refspecs.append(f"{src}:{dest}")
 
     try:
         results = scm.push_refspecs(