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(