Refactor seperate the initialization of executor and setup environment

Reason:
1. Need some more granular control over initialization progress of
   executor.

What done:
1. Seperate `init_git` and `init_cache` progress out from setup_executor
2. Move `set_ref` into `init_git`

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)

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),

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

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]:

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}")

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,

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,

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,

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]:

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",
     )

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):