diff --git a/dvc/repo/experiments/executor/base.py b/dvc/repo/experiments/executor/base.py index 822348ef2..4595507a9 100644 --- a/dvc/repo/experiments/executor/base.py +++ b/dvc/repo/experiments/executor/base.py @@ -35,6 +35,8 @@ from ..refs import ( EXEC_BASELINE, EXEC_BRANCH, EXEC_CHECKPOINT, + EXEC_HEAD, + EXEC_MERGE, EXEC_NAMESPACE, EXPS_NAMESPACE, EXPS_STASH, @@ -151,7 +153,13 @@ class BaseExecutor(ABC): self.result = result @abstractmethod - def init_git(self, scm: "Git", branch: Optional[str] = None): + def init_git( + self, + scm: "Git", + stash_rev: str, + entry: "ExpStashEntry", + branch: Optional[str] = None, + ): """Init git repo and populate it using exp refs from the specified SCM instance. """ @@ -222,7 +230,6 @@ class BaseExecutor(ABC): def from_stash_entry( cls: Type[_T], repo: "Repo", - stash_rev: str, entry: "ExpStashEntry", **kwargs, ) -> _T: @@ -232,7 +239,6 @@ class BaseExecutor(ABC): def _from_stash_entry( cls: Type[_T], repo: "Repo", - stash_rev: str, entry: "ExpStashEntry", root_dir: str, **kwargs, @@ -245,8 +251,6 @@ class BaseExecutor(ABC): wdir=relpath(os.getcwd(), repo.scm.root_dir), **kwargs, ) - executor.init_git(repo.scm, branch=entry.branch) - executor.init_cache(repo, stash_rev) return executor @staticmethod @@ -715,3 +719,23 @@ class BaseExecutor(ABC): disable_other_loggers() if level is not None: dvc_logger.setLevel(level) + + @contextmanager + def set_exec_refs( + self, scm: "Git", stash_rev: str, entry: "ExpStashEntry" + ): + 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) + yield + finally: + for ref in (EXEC_HEAD, EXEC_MERGE, EXEC_BASELINE): + 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 307b1a8ba..e977cd8f6 100644 --- a/dvc/repo/experiments/executor/local.py +++ b/dvc/repo/experiments/executor/local.py @@ -67,15 +67,23 @@ class TempDirExecutor(BaseLocalExecutor): QUIET = True DEFAULT_LOCATION = "tempdir" - def init_git(self, scm: "Git", branch: Optional[str] = None): + def init_git( + self, + scm: "Git", + stash_rev: str, + entry: "ExpStashEntry", + branch: Optional[str] = None, + ): from dulwich.repo import Repo as DulwichRepo from ..utils import push_refspec DulwichRepo.init(os.fspath(self.root_dir)) - refspec = f"{EXEC_NAMESPACE}/" - push_refspec(scm, self.git_url, refspec, refspec) + with self.set_exec_refs(scm, stash_rev, entry): + refspec = f"{EXEC_NAMESPACE}/" + push_refspec(scm, self.git_url, refspec, refspec) + if branch: push_refspec(scm, self.git_url, branch, branch) self.scm.set_ref(EXEC_BRANCH, branch, symbolic=True) @@ -84,17 +92,12 @@ class TempDirExecutor(BaseLocalExecutor): 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) - for ref in (EXEC_HEAD, EXEC_MERGE, EXEC_BASELINE): - if scm.get_ref(ref): - scm.remove_ref(ref) - try: self.scm.merge(merge_rev, squash=True, commit=False) except _SCMError as exc: @@ -123,7 +126,6 @@ class TempDirExecutor(BaseLocalExecutor): def from_stash_entry( cls, repo: "Repo", - stash_rev: str, entry: "ExpStashEntry", wdir: Optional[str] = None, **kwargs, @@ -132,9 +134,7 @@ class TempDirExecutor(BaseLocalExecutor): makedirs(parent_dir, exist_ok=True) tmp_dir = mkdtemp(dir=parent_dir) try: - executor = cls._from_stash_entry( - repo, stash_rev, entry, tmp_dir, **kwargs - ) + executor = cls._from_stash_entry(repo, entry, tmp_dir, **kwargs) logger.debug("Init temp dir executor in '%s'", tmp_dir) return executor except Exception: @@ -152,18 +152,26 @@ class WorkspaceExecutor(BaseLocalExecutor): def from_stash_entry( cls, repo: "Repo", - stash_rev: str, entry: "ExpStashEntry", **kwargs, ): root_dir = repo.scm.root_dir - executor = cls._from_stash_entry( - repo, stash_rev, entry, root_dir, **kwargs + executor: "WorkspaceExecutor" = cls._from_stash_entry( + repo, entry, root_dir, **kwargs ) logger.debug("Init workspace executor in '%s'", root_dir) return executor - def init_git(self, scm: "Git", branch: Optional[str] = None): + def init_git( + self, + scm: "Git", + stash_rev: str, + entry: "ExpStashEntry", + branch: Optional[str] = None, + ): + scm.set_ref(EXEC_HEAD, entry.head_rev) + scm.set_ref(EXEC_MERGE, stash_rev) + scm.set_ref(EXEC_BASELINE, entry.baseline_rev) self._detach_stack.enter_context( self.scm.detach_head( self.scm.get_ref(EXEC_HEAD), diff --git a/dvc/repo/experiments/executor/manager/base.py b/dvc/repo/experiments/executor/manager/base.py index 85e45dff6..b487a2da6 100644 --- a/dvc/repo/experiments/executor/manager/base.py +++ b/dvc/repo/experiments/executor/manager/base.py @@ -8,7 +8,7 @@ from typing import TYPE_CHECKING, Deque, Dict, Generator, Optional, Tuple, Type from dvc.proc.manager import ProcessManager from ...exceptions import CheckpointExistsError, ExperimentExistsError -from ...refs import EXEC_BASELINE, EXEC_HEAD, EXEC_MERGE, ExpRefInfo +from ...refs import ExpRefInfo from ...stash import ExpStashEntry from ..base import EXEC_PID_DIR, BaseExecutor from ..local import TempDirExecutor, WorkspaceExecutor @@ -120,28 +120,17 @@ class BaseExecutorManager(ABC, Mapping): # noqa: B024 to_run: Dict[str, ExpStashEntry], **kwargs, ): - try: - for stash_rev, entry in to_run.items(): - scm.set_ref(EXEC_HEAD, entry.head_rev) - scm.set_ref(EXEC_MERGE, stash_rev) - scm.set_ref(EXEC_BASELINE, entry.baseline_rev) + for stash_rev, entry in to_run.items(): + executor = self.EXECUTOR_CLS.from_stash_entry( + repo, + entry, + **kwargs, + ) - # 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 - executor = self.EXECUTOR_CLS.from_stash_entry( - repo, - stash_rev, - entry, - **kwargs, - ) - self.enqueue(stash_rev, executor) - finally: - for ref in (EXEC_HEAD, EXEC_MERGE, EXEC_BASELINE): - scm.remove_ref(ref) + executor.init_git(scm, stash_rev, entry, branch=entry.branch) + executor.init_cache(repo, stash_rev) + + self.enqueue(stash_rev, executor) def exec_queue( self, repo: "Repo", jobs: Optional[int] = 1, detach: bool = False diff --git a/dvc/repo/experiments/executor/manager/local.py b/dvc/repo/experiments/executor/manager/local.py index dd4c9b1e0..b4659a3b7 100644 --- a/dvc/repo/experiments/executor/manager/local.py +++ b/dvc/repo/experiments/executor/manager/local.py @@ -2,7 +2,7 @@ import logging from collections import defaultdict from typing import TYPE_CHECKING, Dict, Optional -from ...refs import EXEC_BASELINE, EXEC_BRANCH, EXEC_HEAD, EXEC_MERGE +from ...refs import EXEC_BRANCH from ...stash import ExpStashEntry from ..local import TempDirExecutor, WorkspaceExecutor from .base import BaseExecutorManager @@ -32,23 +32,17 @@ class WorkspaceExecutorManager(BaseExecutorManager): **kwargs, ): manager = cls(scm, wdir) - try: - assert len(to_run) == 1 - for stash_rev, entry in to_run.items(): - scm.set_ref(EXEC_HEAD, entry.head_rev) - scm.set_ref(EXEC_MERGE, stash_rev) - scm.set_ref(EXEC_BASELINE, entry.baseline_rev) + assert len(to_run) == 1 + for stash_rev, entry in to_run.items(): + executor = cls.EXECUTOR_CLS.from_stash_entry( + repo, + entry, + **kwargs, + ) + executor.init_git(repo.scm, stash_rev, entry, branch=entry.branch) + executor.init_cache(repo, stash_rev) - executor = cls.EXECUTOR_CLS.from_stash_entry( - repo, - stash_rev, - entry, - **kwargs, - ) - manager.enqueue(stash_rev, executor) - finally: - for ref in (EXEC_MERGE,): - scm.remove_ref(ref) + manager.enqueue(stash_rev, executor) return manager def _collect_executor(self, repo, executor, exec_result) -> Dict[str, str]: diff --git a/dvc/repo/experiments/executor/ssh.py b/dvc/repo/experiments/executor/ssh.py index 94b235b8c..0f706bdc3 100644 --- a/dvc/repo/experiments/executor/ssh.py +++ b/dvc/repo/experiments/executor/ssh.py @@ -23,8 +23,8 @@ if TYPE_CHECKING: from scmrepo.git import Git from dvc.repo import Repo - - from ..base import ExpRefInfo, ExpStashEntry + from dvc.repo.experiments.refs import ExpRefInfo + from dvc.repo.experiments.stash import ExpStashEntry logger = logging.getLogger(__name__) @@ -75,14 +75,12 @@ class SSHExecutor(BaseExecutor): def from_stash_entry( cls, repo: "Repo", - stash_rev: str, entry: "ExpStashEntry", **kwargs, ): machine_name: Optional[str] = kwargs.pop("machine_name", None) executor = cls._from_stash_entry( repo, - stash_rev, entry, cls.gen_dirname(entry.name), location=machine_name, @@ -119,7 +117,13 @@ class SSHExecutor(BaseExecutor): } return kwargs - def init_git(self, scm: "Git", branch: Optional[str] = None): + def init_git( + self, + scm: "Git", + stash_rev: str, + entry: "ExpStashEntry", + branch: Optional[str] = None, + ): from ..utils import push_refspec with self.sshfs() as fs: @@ -137,8 +141,11 @@ class SSHExecutor(BaseExecutor): # TODO: support multiple client key retries in git backends # (see https://github.com/iterative/dvc/issues/6508) kwargs = self._git_client_args(fs) - refspec = f"{EXEC_NAMESPACE}/" - push_refspec(scm, self.git_url, refspec, refspec, **kwargs) + + with self.set_exec_refs(scm, stash_rev, entry): + refspec = f"{EXEC_NAMESPACE}/" + push_refspec(scm, self.git_url, refspec, refspec, **kwargs) + if branch: push_refspec(scm, self.git_url, branch, branch, **kwargs) self._ssh_cmd(fs, f"git symbolic-ref {EXEC_BRANCH} {branch}") diff --git a/dvc/repo/experiments/queue/base.py b/dvc/repo/experiments/queue/base.py index c7995c798..d228e5929 100644 --- a/dvc/repo/experiments/queue/base.py +++ b/dvc/repo/experiments/queue/base.py @@ -33,7 +33,7 @@ from ..executor.base import ( ExecutorResult, ) from ..executor.local import WorkspaceExecutor -from ..refs import EXEC_BASELINE, EXEC_HEAD, EXEC_MERGE, ExpRefInfo +from ..refs import ExpRefInfo from ..stash import ExpStash, ExpStashEntry from ..utils import exp_refs_by_rev, scm_locked @@ -561,7 +561,7 @@ class BaseStashQueue(ABC): @staticmethod @scm_locked - def setup_executor( + def init_executor( exp: "Experiments", queue_entry: QueueEntry, executor_cls: Type[BaseExecutor] = WorkspaceExecutor, @@ -576,21 +576,17 @@ class BaseStashQueue(ABC): ) if stash_entry.stash_index is not None: stash.drop(stash_entry.stash_index) - - scm.set_ref(EXEC_HEAD, stash_entry.head_rev) - scm.set_ref(EXEC_MERGE, stash_rev) - scm.set_ref(EXEC_BASELINE, stash_entry.baseline_rev) - - # 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 - return executor_cls.from_stash_entry( - exp.repo, stash_rev, stash_entry, **kwargs + executor = executor_cls.from_stash_entry( + exp.repo, stash_entry, **kwargs ) + executor.init_git( + exp.repo.scm, stash_rev, stash_entry, branch=stash_entry.branch + ) + executor.init_cache(exp.repo, stash_rev) + + return executor + def get_infofile_path(self, name: str) -> str: return os.path.join( self.pid_dir, diff --git a/dvc/repo/experiments/queue/tasks.py b/dvc/repo/experiments/queue/tasks.py index 6936b798d..ac395371a 100644 --- a/dvc/repo/experiments/queue/tasks.py +++ b/dvc/repo/experiments/queue/tasks.py @@ -28,7 +28,7 @@ def setup_exp(entry_dict: Dict[str, Any]) -> str: repo = Repo(entry.dvc_root) # TODO: split executor.init_cache into separate subtask - we can release # exp.scm_lock before DVC push - executor = BaseStashQueue.setup_executor( + executor = BaseStashQueue.init_executor( repo.experiments, entry, TempDirExecutor, diff --git a/dvc/repo/experiments/queue/tempdir.py b/dvc/repo/experiments/queue/tempdir.py index 4133539cd..b2cbf3825 100644 --- a/dvc/repo/experiments/queue/tempdir.py +++ b/dvc/repo/experiments/queue/tempdir.py @@ -60,7 +60,7 @@ class TempDirQueue(WorkspaceQueue): stash_entry.name, stash_entry.head_rev, ) - executor = self.setup_executor( + executor = self.init_executor( self.repo.experiments, entry, TempDirExecutor, diff --git a/dvc/repo/experiments/queue/workspace.py b/dvc/repo/experiments/queue/workspace.py index ae0f6ca20..ff92dfcb1 100644 --- a/dvc/repo/experiments/queue/workspace.py +++ b/dvc/repo/experiments/queue/workspace.py @@ -41,7 +41,7 @@ class WorkspaceQueue(BaseStashQueue): stash_entry.name, stash_entry.head_rev, ) - executor = self.setup_executor(self.repo.experiments, entry) + executor = self.init_executor(self.repo.experiments, entry) return QueueGetResult(entry, executor) def iter_queued(self) -> Generator[QueueEntry, None, None]: diff --git a/tests/func/experiments/executor/test_ssh.py b/tests/func/experiments/executor/test_ssh.py index 49ef061e2..c5d930a3d 100644 --- a/tests/func/experiments/executor/test_ssh.py +++ b/tests/func/experiments/executor/test_ssh.py @@ -29,7 +29,6 @@ def test_init_from_stash(tmp_dir, scm, dvc, machine_instance, mocker): mock_entry.name = "" SSHExecutor.from_stash_entry( dvc, - "abc123", mock_entry, machine_name="foo", ) diff --git a/tests/unit/repo/experiments/queue/test_local.py b/tests/unit/repo/experiments/queue/test_local.py index e47c0e3bb..19e1dc1de 100644 --- a/tests/unit/repo/experiments/queue/test_local.py +++ b/tests/unit/repo/experiments/queue/test_local.py @@ -6,9 +6,7 @@ from flaky.flaky_decorator import flaky from dvc.exceptions import DvcException from dvc.repo.experiments.exceptions import UnresolvedExpNamesError -from dvc.repo.experiments.executor.local import TempDirExecutor from dvc.repo.experiments.queue.base import QueueDoneResult -from dvc.repo.experiments.refs import EXEC_BASELINE, EXEC_HEAD, EXEC_MERGE def test_shutdown_no_tasks(test_queue, mocker): @@ -110,20 +108,6 @@ def test_celery_queue_kill(test_queue, mocker): assert spy.called_once_with(mock_entry.stash_rev) -def test_queue_clean_workspace_refs(git_dir, tmp_dir): - - rev = git_dir.scm.get_rev() - exec_heads = (EXEC_MERGE, EXEC_BASELINE, EXEC_HEAD) - for ref in exec_heads: - git_dir.scm.set_ref(ref, rev) - - temp_dir_executor = TempDirExecutor(str(tmp_dir), str(tmp_dir), rev) - temp_dir_executor.init_git(git_dir.scm) - - for ref in exec_heads: - assert git_dir.scm.get_ref(ref) is None - - @pytest.mark.parametrize("status", ["FAILURE", "SUCCESS"]) def test_queue_iter_done_task(test_queue, mocker, status):