From ca4dd08e37ef8a4bda668259368c35ee26fd776b Mon Sep 17 00:00:00 2001 From: karajan1001 Date: Thu, 24 Nov 2022 21:42:37 +0800 Subject: [PATCH] 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 | 26 +++--------- dvc/repo/experiments/executor/local.py | 58 +++++++++++++++++++------- dvc/repo/experiments/executor/ssh.py | 15 +++++-- dvc/repo/experiments/refs.py | 1 + dvc/repo/experiments/remove.py | 3 +- dvc/repo/experiments/utils.py | 31 +++++++------- 6 files changed, 76 insertions(+), 58 deletions(-) diff --git a/dvc/repo/experiments/executor/base.py b/dvc/repo/experiments/executor/base.py index f85ab4e6a..bb2b2b0e3 100644 --- a/dvc/repo/experiments/executor/base.py +++ b/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) diff --git a/dvc/repo/experiments/executor/local.py b/dvc/repo/experiments/executor/local.py index 9ab2f18f9..e919a1f4f 100644 --- a/dvc/repo/experiments/executor/local.py +++ b/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( diff --git a/dvc/repo/experiments/executor/ssh.py b/dvc/repo/experiments/executor/ssh.py index 673a5bc1d..00e5e4226 100644 --- a/dvc/repo/experiments/executor/ssh.py +++ b/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( diff --git a/dvc/repo/experiments/refs.py b/dvc/repo/experiments/refs.py index d5e594e98..c7efc78b3 100644 --- a/dvc/repo/experiments/refs.py +++ b/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} diff --git a/dvc/repo/experiments/remove.py b/dvc/repo/experiments/remove.py index 2d6693a2d..9a5509723 100644 --- a/dvc/repo/experiments/remove.py +++ b/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: diff --git a/dvc/repo/experiments/utils.py b/dvc/repo/experiments/utils.py index 54d115f4c..f4307af2e 100644 --- a/dvc/repo/experiments/utils.py +++ b/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(