exp: implement git + DVC cache collection for ssh executors

dvc/repo/experiments/__init__.py

@@ -672,7 +672,7 @@ class Experiments:
             dict mapping stash revs to the successfully executed experiments
             for each stash rev.
         """
-        return manager.exec_queue(**kwargs)
+        return manager.exec_queue(self.repo, **kwargs)
 
     def check_baseline(self, exp_rev):
         baseline_sha = self.repo.scm.get_rev()
@@ -747,6 +747,8 @@ class Experiments:
         from dvc.scm import InvalidRemoteSCMRepo
         from dvc.utils.serialize import load_json
 
+        from .executor.local import TempDirExecutor
+
         result = {}
         pid_dir = os.path.join(
             self.repo.tmp_dir,
@@ -777,10 +779,10 @@ class Experiments:
                         def on_diverged(_ref: str, _checkpoint: bool):
                             return False
 
+                        executor = TempDirExecutor.from_info(info)
                         try:
-                            for ref in BaseExecutor.fetch_exps(
+                            for ref in executor.fetch_exps(
                                 self.scm,
-                                info.git_url,
                                 on_diverged=on_diverged,
                             ):
                                 logger.debug(

dvc/repo/experiments/executor/base.py

@@ -143,8 +143,14 @@ class BaseExecutor(ABC):
         pass
 
     @abstractmethod
-    def init_cache(self, dvc: "Repo", rev: str, run_cache: bool = True):
-        """Initialize DVC (cache)."""
+    def init_cache(self, repo: "Repo", rev: str, run_cache: bool = True):
+        """Initialize DVC cache."""
+
+    @abstractmethod
+    def collect_cache(
+        self, repo: "Repo", exp_ref: "ExpRefInfo", run_cache: bool = True
+    ):
+        """Collect DVC cache."""
 
     @property
     def info(self) -> "ExecutorInfo":
@@ -266,28 +272,33 @@ class BaseExecutor(ABC):
             data = pickle.load(fobj)
         return data["args"], data["kwargs"]
 
-    @classmethod
     def fetch_exps(
-        cls,
+        self,
         dest_scm: "Git",
-        url: str,
         force: bool = False,
         on_diverged: Callable[[str, bool], None] = None,
+        **kwargs,
     ) -> Iterable[str]:
-        """Fetch reproduced experiments into the specified SCM.
+        """Fetch reproduced experiment refs into the specified SCM.
 
         Args:
             dest_scm: Destination Git instance.
-            url: Git remote URL to fetch from.
             force: If True, diverged refs will be overwritten
             on_diverged: Callback in the form on_diverged(ref, is_checkpoint)
                 to be called when an experiment ref has diverged.
+
+        Extra kwargs will be passed into the remote git client.
         """
         from ..utils import iter_remote_refs
 
         refs = []
         has_checkpoint = False
-        for ref in iter_remote_refs(dest_scm, url, base=EXPS_NAMESPACE):
+        for ref in iter_remote_refs(
+            dest_scm,
+            self.git_url,
+            base=EXPS_NAMESPACE,
+            **kwargs,
+        ):
             if ref == EXEC_CHECKPOINT:
                 has_checkpoint = True
             elif not ref.startswith(EXEC_NAMESPACE) and ref != EXPS_STASH:
@@ -298,7 +309,7 @@ class BaseExecutor(ABC):
                 logger.debug("Replacing existing experiment '%s'", orig_ref)
                 return True
 
-            cls._raise_ref_conflict(
+            self._raise_ref_conflict(
                 dest_scm, orig_ref, new_rev, has_checkpoint
             )
             if on_diverged:
@@ -308,17 +319,19 @@ class BaseExecutor(ABC):
 
         # fetch experiments
         dest_scm.fetch_refspecs(
-            url,
+            self.git_url,
             [f"{ref}:{ref}" for ref in refs],
             on_diverged=on_diverged_ref,
             force=force,
+            **kwargs,
         )
         # update last run checkpoint (if it exists)
         if has_checkpoint:
             dest_scm.fetch_refspecs(
-                url,
+                self.git_url,
                 [f"{EXEC_CHECKPOINT}:{EXEC_CHECKPOINT}"],
                 force=True,
+                **kwargs,
             )
         return refs
 

dvc/repo/experiments/executor/local.py

@@ -25,7 +25,7 @@ if TYPE_CHECKING:
 
     from dvc.repo import Repo
 
-    from ..base import ExpStashEntry
+    from ..base import ExpRefInfo, ExpStashEntry
 
 logger = logging.getLogger(__name__)
 
@@ -49,6 +49,11 @@ class BaseLocalExecutor(BaseExecutor):
         self.scm.close()
         del self.scm
 
+    def collect_cache(
+        self, repo: "Repo", exp_ref: "ExpRefInfo", run_cache: bool = True
+    ):
+        """Collect DVC cache."""
+
 
 class TempDirExecutor(BaseLocalExecutor):
     """Temp directory experiment executor."""
@@ -91,9 +96,9 @@ class TempDirExecutor(BaseLocalExecutor):
         with open(local_config, "w", encoding="utf-8") as fobj:
             fobj.write(f"[cache]\n    dir = {cache_dir}")
 
-    def init_cache(self, dvc: "Repo", rev: str, run_cache: bool = True):
-        """Initialize DVC (cache)."""
-        self._config(dvc.odb.local.cache_dir)
+    def init_cache(self, repo: "Repo", rev: str, run_cache: bool = True):
+        """Initialize DVC cache."""
+        self._config(repo.odb.local.cache_dir)
 
     def cleanup(self):
         super().cleanup()
@@ -152,7 +157,7 @@ class WorkspaceExecutor(BaseLocalExecutor):
         elif scm.get_ref(EXEC_BRANCH):
             self.scm.remove_ref(EXEC_BRANCH)
 
-    def init_cache(self, dvc: "Repo", rev: str, run_cache: bool = True):
+    def init_cache(self, repo: "Repo", rev: str, run_cache: bool = True):
         pass
 
     def cleanup(self):

dvc/repo/experiments/executor/manager/base.py

@@ -149,14 +149,16 @@ class BaseExecutorManager(ABC, Mapping):
             for ref in (EXEC_HEAD, EXEC_MERGE, EXEC_BASELINE):
                 scm.remove_ref(ref)
 
-    def exec_queue(self, jobs: Optional[int] = 1, detach: bool = False):
+    def exec_queue(
+        self, repo: "Repo", jobs: Optional[int] = 1, detach: bool = False
+    ):
         """Run dvc repro for queued executors in parallel."""
         if detach:
             raise NotImplementedError
             # TODO use ProcessManager.spawn() to support detached runs
-        return self._exec_attached(jobs=jobs)
+        return self._exec_attached(repo, jobs=jobs)
 
-    def _exec_attached(self, jobs: Optional[int] = 1):
+    def _exec_attached(self, repo: "Repo", jobs: Optional[int] = 1):
         import signal
         from concurrent.futures import (
             CancelledError,
@@ -214,7 +216,7 @@ class BaseExecutorManager(ABC, Mapping):
                     if exc is None:
                         exec_result = future.result()
                         result[rev].update(
-                            self._collect_executor(executor, exec_result)
+                            self._collect_executor(repo, executor, exec_result)
                         )
                     elif not isinstance(exc, CheckpointKilledError):
                         logger.error(
@@ -231,7 +233,7 @@ class BaseExecutorManager(ABC, Mapping):
 
         return result
 
-    def _collect_executor(self, executor, exec_result) -> Dict[str, str]:
+    def _collect_executor(self, repo, executor, exec_result) -> Dict[str, str]:
         # NOTE: GitPython Repo instances cannot be re-used
         # after process has received SIGINT or SIGTERM, so we
         # need this hack to re-instantiate git instances after
@@ -249,7 +251,6 @@ class BaseExecutorManager(ABC, Mapping):
 
         for ref in executor.fetch_exps(
             self.scm,
-            executor.git_url,
             force=exec_result.force,
             on_diverged=on_diverged,
         ):
@@ -258,6 +259,9 @@ class BaseExecutorManager(ABC, Mapping):
                 logger.debug("Collected experiment '%s'.", exp_rev[:7])
                 results[exp_rev] = exec_result.exp_hash
 
+        if exec_result.ref_info is not None:
+            executor.collect_cache(repo, exec_result.ref_info)
+
         return results
 
     def cleanup_executor(self, rev: str, executor: "BaseExecutor"):

dvc/repo/experiments/executor/manager/local.py

@@ -56,7 +56,7 @@ class WorkspaceExecutorManager(BaseExecutorManager):
                 scm.remove_ref(ref)
         return manager
 
-    def _collect_executor(self, executor, exec_result) -> Dict[str, str]:
+    def _collect_executor(self, repo, executor, exec_result) -> Dict[str, str]:
         results = {}
         exp_rev = self.scm.get_ref(EXEC_BRANCH)
         if exp_rev:
@@ -64,7 +64,9 @@ class WorkspaceExecutorManager(BaseExecutorManager):
             results[exp_rev] = exec_result.exp_hash
         return results
 
-    def exec_queue(self, jobs: Optional[int] = 1, detach: bool = False):
+    def exec_queue(
+        self, repo: "Repo", jobs: Optional[int] = 1, detach: bool = False
+    ):
         """Run a single WorkspaceExecutor.
 
         Workspace execution is done within the main DVC process
@@ -93,7 +95,7 @@ class WorkspaceExecutorManager(BaseExecutorManager):
                 )
             if exec_result.ref_info:
                 result[rev].update(
-                    self._collect_executor(executor, exec_result)
+                    self._collect_executor(repo, executor, exec_result)
                 )
         except CheckpointKilledError:
             # Checkpoint errors have already been logged

dvc/repo/experiments/executor/manager/ssh.py

@@ -62,7 +62,7 @@ class SSHExecutorManager(BaseExecutorManager):
     def get_infofile_path(self, name: str) -> str:
         return f"{name}{BaseExecutor.INFOFILE_EXT}"
 
-    def _exec_attached(self, jobs: Optional[int] = 1):
+    def _exec_attached(self, repo: "Repo", jobs: Optional[int] = 1):
         from dvc.exceptions import DvcException
         from dvc.stage.monitor import CheckpointKilledError
 
@@ -90,7 +90,7 @@ class SSHExecutorManager(BaseExecutorManager):
                 )
             if exec_result.ref_info:
                 result[rev].update(
-                    self._collect_executor(executor, exec_result)
+                    self._collect_executor(repo, executor, exec_result)
                 )
         except CheckpointKilledError:
             # Checkpoint errors have already been logged

dvc/repo/experiments/executor/ssh.py

@@ -1,7 +1,7 @@
 import logging
 import posixpath
 from contextlib import contextmanager
-from typing import TYPE_CHECKING, Callable, Optional
+from typing import TYPE_CHECKING, Callable, Iterable, Optional
 
 from funcy import first
 
@@ -22,7 +22,7 @@ if TYPE_CHECKING:
     from dvc.machine import MachineManager
     from dvc.repo import Repo
 
-    from ..base import ExpStashEntry
+    from ..base import ExpRefInfo, ExpStashEntry
 
 logger = logging.getLogger(__name__)
 
@@ -155,10 +155,31 @@ class SSHExecutor(BaseExecutor):
         working_dir = chdir or self.root_dir
         return sshfs.fs.execute(f"cd {working_dir};{cmd}", **kwargs)
 
-    def init_cache(self, dvc: "Repo", rev: str, run_cache: bool = True):
+    def init_cache(self, repo: "Repo", rev: str, run_cache: bool = True):
+        from dvc.repo.push import push
+
+        with self.get_odb() as odb:
+            push(
+                repo,
+                revs=[rev],
+                run_cache=run_cache,
+                odb=odb,
+                include_imports=True,
+            )
+
+    def collect_cache(
+        self, repo: "Repo", exp_ref: "ExpRefInfo", run_cache: bool = True
+    ):
+        """Collect DVC cache."""
+        from dvc.repo.experiments.pull import _pull_cache
+
+        with self.get_odb() as odb:
+            _pull_cache(repo, exp_ref, run_cache=run_cache, odb=odb)
+
+    @contextmanager
+    def get_odb(self):
         from dvc.data.db import ODBManager, get_odb
         from dvc.repo import Repo
-        from dvc.repo.push import push
 
         cache_path = posixpath.join(
             self._repo_abspath,
@@ -167,5 +188,9 @@ class SSHExecutor(BaseExecutor):
         )
 
         with self.sshfs() as fs:
-            odb = get_odb(fs, cache_path, **fs.config)
-            push(dvc, revs=[rev], run_cache=run_cache, odb=odb)
+            yield get_odb(fs, cache_path, **fs.config)
+
+    def fetch_exps(self, *args, **kwargs) -> Iterable[str]:
+        with self.sshfs() as fs:
+            kwargs.update(self._git_client_args(fs))
+            return super().fetch_exps(*args, **kwargs)