From 39eee9b5a66a45e76a23cf3490502c3f76bffcdc Mon Sep 17 00:00:00 2001 From: Peter Rowlands Date: Mon, 20 Dec 2021 15:56:20 +0900 Subject: [PATCH] exp: implement git + DVC cache collection for ssh executors --- dvc/repo/experiments/__init__.py | 8 ++-- dvc/repo/experiments/executor/base.py | 35 ++++++++++++------ dvc/repo/experiments/executor/local.py | 15 +++++--- dvc/repo/experiments/executor/manager/base.py | 16 +++++--- .../experiments/executor/manager/local.py | 8 ++-- dvc/repo/experiments/executor/manager/ssh.py | 4 +- dvc/repo/experiments/executor/ssh.py | 37 ++++++++++++++++--- 7 files changed, 87 insertions(+), 36 deletions(-) diff --git a/dvc/repo/experiments/__init__.py b/dvc/repo/experiments/__init__.py index a2447df1d..811afb318 100644 --- a/dvc/repo/experiments/__init__.py +++ b/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( diff --git a/dvc/repo/experiments/executor/base.py b/dvc/repo/experiments/executor/base.py index 9943a9417..d1fa8b0e1 100644 --- a/dvc/repo/experiments/executor/base.py +++ b/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 diff --git a/dvc/repo/experiments/executor/local.py b/dvc/repo/experiments/executor/local.py index d8a9f6867..849110684 100644 --- a/dvc/repo/experiments/executor/local.py +++ b/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): diff --git a/dvc/repo/experiments/executor/manager/base.py b/dvc/repo/experiments/executor/manager/base.py index e5c621fbf..2f53802fe 100644 --- a/dvc/repo/experiments/executor/manager/base.py +++ b/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"): diff --git a/dvc/repo/experiments/executor/manager/local.py b/dvc/repo/experiments/executor/manager/local.py index 95daa2eaa..da139f759 100644 --- a/dvc/repo/experiments/executor/manager/local.py +++ b/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 diff --git a/dvc/repo/experiments/executor/manager/ssh.py b/dvc/repo/experiments/executor/manager/ssh.py index 5f1982651..a67c5c7e1 100644 --- a/dvc/repo/experiments/executor/manager/ssh.py +++ b/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 diff --git a/dvc/repo/experiments/executor/ssh.py b/dvc/repo/experiments/executor/ssh.py index 3b1df9503..9edae2c4e 100644 --- a/dvc/repo/experiments/executor/ssh.py +++ b/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)