From e45ca1e83048d6ceaad53dad63fbfbe0812f8505 Mon Sep 17 00:00:00 2001 From: Peter Rowlands Date: Wed, 2 Feb 2022 17:26:22 +0900 Subject: [PATCH] exp: start moving queueing/execution from __init__ --- dvc/repo/experiments/queue/base.py | 336 ++++++++++++++++++++++++++++ dvc/repo/experiments/queue/local.py | 61 +++++ 2 files changed, 397 insertions(+) create mode 100644 dvc/repo/experiments/queue/base.py create mode 100644 dvc/repo/experiments/queue/local.py diff --git a/dvc/repo/experiments/queue/base.py b/dvc/repo/experiments/queue/base.py new file mode 100644 index 000000000..f7d84e27b --- /dev/null +++ b/dvc/repo/experiments/queue/base.py @@ -0,0 +1,336 @@ +import logging +import os +from dataclasses import asdict, dataclass +from typing import TYPE_CHECKING, Any, Dict, Optional + +from funcy import cached_property + +from dvc.dependency.param import MissingParamsError +from dvc.env import DVCLIVE_RESUME +from dvc.exceptions import DvcException + +from ..base import ExpRefInfo +from ..executor.base import BaseExecutor +from ..stash import ExpStash +from ..utils import exp_refs_by_rev + +if TYPE_CHECKING: + from scmrepo.git import Git + + from dvc.repo import Repo + +logger = logging.getLogger(__name__) + + +@dataclass +class QueueEntry: + dvc_root: str + scm_root: str + stash_ref: str + stash_rev: str + baseline_rev: str + branch: Optional[str] + name: Optional[str] + + def __eq__(self, other: object): + return ( + isinstance(other, QueueEntry) + and self.dvc_root == other.dvc_root + and self.scm_root == other.scm_root + and self.stash_ref == other.stash_ref + and self.stash_rev == other.stash_rev + ) + + def asdict(self) -> Dict[str, Any]: + return asdict(self) + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> "QueueEntry": + return cls(**d) + + +class StashQueue: + """Naive Git-stash based experiment queue. + + Maps queued experiments to (Git) stash reflog entries. + """ + + def __init__(self, repo: "Repo", ref: str): + """Construct a queue. + + Arguments: + scm: Git SCM instance for this queue. + ref: Git stash ref for this queue. + """ + self.repo = repo + self.ref = ref + + @property + def scm(self) -> "Git": + return self.repo.scm + + @cached_property + def stash(self) -> ExpStash: + return ExpStash(self.scm, self.ref) + + @cached_property + def args_file(self): + return os.path.join(self.repo.tmp_dir, BaseExecutor.PACKED_ARGS_FILE) + + def put(self, *args, **kwargs) -> QueueEntry: + """Stash an experiment and add it to the queue.""" + return self._stash_exp(*args, **kwargs) + + def get(self) -> QueueEntry: + """Pop and return the first item in the queue.""" + + def _stash_exp( + self, + *args, + params: Optional[dict] = None, + resume_rev: Optional[str] = None, + baseline_rev: Optional[str] = None, + branch: Optional[str] = None, + name: Optional[str] = None, + **kwargs, + ) -> QueueEntry: + """Stash changes from the workspace as an experiment. + + Arguments: + params: Optional dictionary of parameter values to be used. + Values take priority over any parameters specified in the + user's workspace. + resume_rev: Optional checkpoint resume rev. + baseline_rev: Optional baseline rev for this experiment, defaults + to the current SCM rev. + branch: Optional experiment branch name. If specified, the + experiment will be added to `branch` instead of creating + a new branch. + name: Optional experiment name. If specified this will be used as + the human-readable name in the experiment branch ref. Has no + effect of branch is specified. + """ + with self.scm.detach_head(client="dvc") as orig_head: + stash_head = orig_head + if baseline_rev is None: + baseline_rev = orig_head + + with self.scm.stash_workspace() as workspace: + try: + if workspace: + self.stash.apply(workspace) + + if resume_rev: + # move HEAD to the resume rev so that the stashed diff + # only contains changes relative to resume rev + stash_head = resume_rev + self.scm.set_ref( + "HEAD", + resume_rev, + message=f"dvc: resume from HEAD {resume_rev[:7]}", + ) + self.scm.reset() + + # update experiment params from command line + if params: + self._update_params(params) + + # DVC commit data deps to preserve state across workspace + # & tempdir runs + self._stash_commit_deps(*args, **kwargs) + + if resume_rev: + if branch: + branch_name = ExpRefInfo.from_ref(branch).name + else: + branch_name = f"{resume_rev[:7]}" + if self.scm.is_dirty(): + logger.info( + "Modified checkpoint experiment based on " + "'%s' will be created", + branch_name, + ) + branch = None + elif ( + not branch + or self.scm.get_ref(branch) != resume_rev + ): + err_msg = [ + ( + "Nothing to do for unchanged checkpoint " + f"'{resume_rev[:7]}'. " + ) + ] + if branch: + err_msg.append( + "To resume from the head of this " + "experiment, use " + f"'dvc exp apply {branch_name}'." + ) + else: + names = [ + ref_info.name + for ref_info in exp_refs_by_rev( + self.scm, resume_rev + ) + ] + if len(names) > 3: + names[3:] = [ + f"... ({len(names) - 3} more)" + ] + err_msg.append( + "To resume an experiment containing this " + "checkpoint, apply one of these heads:\n" + "\t{}".format(", ".join(names)) + ) + raise DvcException("".join(err_msg)) + else: + logger.info( + "Existing checkpoint experiment '%s' will be " + "resumed", + branch_name, + ) + if name: + logger.warning( + "Ignoring option '--name %s' for resumed " + "experiment. Existing experiment name will" + "be preserved instead.", + name, + ) + + # save additional repro command line arguments + run_env = {DVCLIVE_RESUME: "1"} if resume_rev else {} + self._pack_args(*args, run_env=run_env, **kwargs) + + # save experiment as a stash commit + msg = self._stash_msg( + stash_head, + baseline_rev=baseline_rev, + branch=branch, + name=name, + ) + stash_rev = self.stash.push(message=msg) + logger.debug( + ( + "Stashed experiment '%s' with baseline '%s' " + "for future execution." + ), + stash_rev[:7], + baseline_rev[:7], + ) + finally: + if resume_rev: + # NOTE: this set_ref + reset() is equivalent to + # `git reset orig_head` (our SCM reset() only operates + # on HEAD rather than any arbitrary commit) + self.scm.set_ref( + "HEAD", orig_head, message="dvc: restore HEAD" + ) + self.scm.reset() + # Revert any of our changes before prior unstashing + self.scm.reset(hard=True) + + return QueueEntry( + self.repo.root_dir, + self.scm.root_dir, + self.ref, + stash_rev, + baseline_rev, + branch, + name, + ) + + def _stash_commit_deps(self, *args, **kwargs): + if len(args): + targets = args[0] + else: + targets = kwargs.get("targets") + if isinstance(targets, str): + targets = [targets] + elif not targets: + targets = [None] + for target in targets: + self.repo.commit( + target, + with_deps=True, + recursive=kwargs.get("recursive", False), + force=True, + allow_missing=True, + data_only=True, + ) + + def _stash_msg( + self, + rev: str, + baseline_rev: str, + branch: Optional[str] = None, + name: Optional[str] = None, + ) -> str: + if not baseline_rev: + baseline_rev = rev + msg = ExpStash.format_message(rev, baseline_rev, name) + if branch: + return f"{msg}:{branch}" + return msg + + def _pack_args(self, *args, **kwargs) -> None: + import pickle + + if os.path.exists(self.args_file) and self.scm.is_tracked( + self.args_file + ): + logger.warning( + ( + "Temporary DVC file '.dvc/tmp/%s' exists and was " + "likely committed to Git by mistake. It should be removed " + "with:\n" + "\tgit rm .dvc/tmp/%s" + ), + BaseExecutor.PACKED_ARGS_FILE, + BaseExecutor.PACKED_ARGS_FILE, + ) + with open(self.args_file, "rb") as fobj: + try: + data = pickle.load(fobj) + except Exception: # pylint: disable=broad-except + data = {} + extra = int(data.get("extra", 0)) + 1 + else: + extra = None + BaseExecutor.pack_repro_args( + self.args_file, *args, extra=extra, **kwargs + ) + self.scm.add(self.args_file) + + def _format_new_params_msg(self, new_params, config_path): + """Format an error message for when new parameters are identified""" + new_param_count = len(new_params) + pluralise = "s are" if new_param_count > 1 else " is" + param_list = ", ".join(new_params) + return ( + f"{new_param_count} parameter{pluralise} missing " + f"from '{config_path}': {param_list}" + ) + + def _update_params(self, params: dict): + """Update experiment params files with the specified values.""" + from dvc.utils.collections import NewParamsFound, merge_params + from dvc.utils.serialize import MODIFIERS + + logger.debug("Using experiment params '%s'", params) + + for path in params: + suffix = self.repo.fs.path.suffix(path).lower() + modify_data = MODIFIERS[suffix] + with modify_data(path, fs=self.repo.fs) as data: + try: + merge_params(data, params[path], allow_new=False) + except NewParamsFound as e: + msg = self._format_new_params_msg(e.new_params, path) + raise MissingParamsError(msg) + + # Force params file changes to be staged in git + # Otherwise in certain situations the changes to params file may be + # ignored when we `git stash` them since mtime is used to determine + # whether the file is dirty + self.scm.add(list(params.keys())) diff --git a/dvc/repo/experiments/queue/local.py b/dvc/repo/experiments/queue/local.py new file mode 100644 index 000000000..c94228edf --- /dev/null +++ b/dvc/repo/experiments/queue/local.py @@ -0,0 +1,61 @@ +import os +from typing import TYPE_CHECKING + +from funcy import cached_property + +from ..executor.base import EXEC_TMP_DIR +from .base import QueueEntry, StashQueue +from .tasks import setup_exp + +if TYPE_CHECKING: + from celery import Celery, Signature + from dvc_task.worker import TemporaryWorker + + +class LocalCeleryQueue(StashQueue): + """DVC experiment queue. + + Maps queued experiments to (Git) stash reflog entries. + """ + + CELERY_DIR = "celery" + + @cached_property + def wdir(self) -> str: + return os.path.join(self.repo.tmp_dir, EXEC_TMP_DIR, self.CELERY_DIR) + + @cached_property + def celery(self) -> "Celery": + from celery import Celery + + from .celery import get_config + + app = Celery() + app.conf.update(get_config(self.wdir)) + return app + + @cached_property + def worker(self) -> "TemporaryWorker": + from dvc_task.worker import TemporaryWorker + + return TemporaryWorker(self.celery, concurrency=1) + + def put(self, *args, **kwargs) -> QueueEntry: + """Stash an experiment and add it to the queue.""" + entry = super()._stash_exp(*args, **kwargs) + # self.celery.tasks[exp_run.name].delay(entry.asdict()) + self.exp_chain(entry).delay() + return entry + + def get(self) -> QueueEntry: + """Pop and return the first item in the queue.""" + # Queue consumption should not be done directly. Celery worker(s) will + # automatically consume available experiments. + raise NotImplementedError + + def exp_chain(self, entry: QueueEntry) -> "Signature": + from celery import chain + + chain( + self.celery.tasks[setup_exp.name](entry.asdict()), + )