From a0ae9d3f45f78d77fbd12441ffd903d623166bc3 Mon Sep 17 00:00:00 2001 From: Peter Rowlands Date: Mon, 7 Feb 2022 19:03:38 +0900 Subject: [PATCH] exp: migrate workspace repro to queues --- dvc/repo/experiments/__init__.py | 437 +++++------------------------ dvc/repo/experiments/gc.py | 4 +- dvc/repo/experiments/queue/base.py | 2 +- dvc/repo/experiments/refs.py | 2 + dvc/repo/experiments/remove.py | 69 +++-- dvc/repo/experiments/stash.py | 2 +- 6 files changed, 109 insertions(+), 407 deletions(-) diff --git a/dvc/repo/experiments/__init__.py b/dvc/repo/experiments/__init__.py index 256a73f7e..774c231f8 100644 --- a/dvc/repo/experiments/__init__.py +++ b/dvc/repo/experiments/__init__.py @@ -1,12 +1,10 @@ import logging import os import re -from typing import Dict, Iterable, List, Mapping, Optional, Type +from typing import Dict, Iterable, Mapping, Optional from funcy import cached_property, first -from dvc.dependency.param import MissingParamsError -from dvc.env import DVCLIVE_RESUME from dvc.exceptions import DvcException from dvc.utils import relpath @@ -22,22 +20,19 @@ from .executor.base import ( BaseExecutor, ExecutorInfo, ) -from .executor.manager.base import BaseExecutorManager -from .executor.manager.local import ( - TempDirExecutorManager, - WorkspaceExecutorManager, -) -from .executor.manager.ssh import SSHExecutorManager +from .queue.base import BaseStashQueue, QueueEntry +from .queue.local import LocalCeleryQueue, WorkspaceQueue from .refs import ( + CELERY_STASH, EXEC_APPLY, EXEC_BRANCH, EXEC_CHECKPOINT, EXEC_NAMESPACE, EXPS_NAMESPACE, - EXPS_STASH, + WORKSPACE_STASH, ExpRefInfo, ) -from .stash import ExpStash, ExpStashEntry +from .stash import ExpStashEntry from .utils import exp_refs_by_rev, scm_locked, unlocked_repo logger = logging.getLogger(__name__) @@ -82,255 +77,19 @@ class Experiments: return os.path.join(self.repo.tmp_dir, BaseExecutor.PACKED_ARGS_FILE) @cached_property - def stash(self) -> ExpStash: - return ExpStash(self.scm, EXPS_STASH) + def workspace_queue(self) -> WorkspaceQueue: + return WorkspaceQueue(self.repo, WORKSPACE_STASH) + + @cached_property + def celery_queue(self) -> LocalCeleryQueue: + return LocalCeleryQueue(self.repo, CELERY_STASH) @property def stash_revs(self) -> Dict[str, ExpStashEntry]: - return self.stash.stash_revs - - 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, - ): - """Stash changes from the workspace as an experiment. - - Args: - 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(untracked_files=False): - 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 - ): - msg = [ - ( - "Nothing to do for unchanged checkpoint " - f"'{resume_rev[:7]}'. " - ) - ] - if branch: - 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)" - ] - msg.append( - "To resume an experiment containing this " - "checkpoint, apply one of these heads:\n" - "\t{}".format(", ".join(names)) - ) - raise DvcException("".join(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 stash_rev - - 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, - ): - 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): - 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())) + revs = {} + for queue in (self.workspace_queue, self.celery_queue): + revs.update(queue.stash.stash_revs) + return revs def reproduce_one( self, @@ -372,26 +131,28 @@ class Experiments: else: checkpoint_resume = self._workspace_resume_rev() - stash_rev = self.new( - checkpoint_resume=checkpoint_resume, reset=reset, **kwargs + exp_queue = ( + self.celery_queue if (tmp_dir or queue) else self.workspace_queue + ) + entry = self.new( + exp_queue, + checkpoint_resume=checkpoint_resume, + reset=reset, + **kwargs, ) if queue: logger.info( - "Queued experiment '%s' for future execution.", stash_rev[:7] + "Queued experiment '%s' for future execution.", + entry.stash_rev[:7], ) - return [stash_rev] - if tmp_dir or queue: - manager_cls: Type = TempDirExecutorManager - elif machine: - manager_cls = SSHExecutorManager - else: - manager_cls = WorkspaceExecutorManager - results = self._reproduce_revs( - revs=[stash_rev], - keep_stash=False, - manager_cls=manager_cls, - machine=machine, - ) + return [entry.stash_rev] + if tmp_dir: + # TODO: start celery worker + raise NotImplementedError + if machine: + # TODO: decide how to handle queued remote execution + raise NotImplementedError + results = self._reproduce_queue(exp_queue) exp_rev = first(results) if exp_rev is not None: self._log_reproduced(results, tmp_dir=tmp_dir) @@ -405,10 +166,12 @@ class Experiments: return None def reproduce_queued(self, **kwargs): - results = self._reproduce_revs(**kwargs) - if results: - self._log_reproduced(results, tmp_dir=True) - return results + # TODO: start celery worker + raise NotImplementedError + # results = self._reproduce_revs(**kwargs) + # if results: + # self._log_reproduced(results, tmp_dir=True) + # return results def _log_reproduced(self, revs: Iterable[str], tmp_dir: bool = False): names = [] @@ -443,15 +206,20 @@ class Experiments: raise ExperimentExistsError(exp_ref.name) @scm_locked - def new(self, *args, checkpoint_resume: Optional[str] = None, **kwargs): - """Create a new experiment. + def new( + self, + queue: BaseStashQueue, + *args, + checkpoint_resume: Optional[str] = None, + **kwargs, + ) -> QueueEntry: + """Create and enqueue a new experiment. - Experiment will be reproduced and checked out into the user's - workspace. + Experiment will be derived from the current workspace. """ if checkpoint_resume is not None: return self._resume_checkpoint( - *args, resume_rev=checkpoint_resume, **kwargs + queue, *args, resume_rev=checkpoint_resume, **kwargs ) name = kwargs.get("name", None) @@ -464,16 +232,16 @@ class Experiments: if not (kwargs.get("force", False) or kwargs.get("reset", False)): raise err - return self._stash_exp(*args, **kwargs) + return queue.put(*args, **kwargs) def _resume_checkpoint( - self, *args, resume_rev: Optional[str] = None, **kwargs - ): - """Resume an existing (checkpoint) experiment. - - Experiment will be reproduced and checked out into the user's - workspace. - """ + self, + queue: BaseStashQueue, + *args, + resume_rev: Optional[str] = None, + **kwargs, + ) -> QueueEntry: + """Create and queue a resumed checkpoint experiment.""" assert resume_rev branch: Optional[str] = None @@ -504,7 +272,7 @@ class Experiments: resume_rev[:7], baseline_rev, ) - return self._stash_exp( + return queue.put( *args, resume_rev=resume_rev, baseline_rev=baseline_rev, @@ -540,90 +308,25 @@ class Experiments: self.scm.remove_ref(EXEC_CHECKPOINT) self.scm.remove_ref(EXEC_APPLY) - @scm_locked - def _reproduce_revs( - self, - revs: Optional[Iterable] = None, - keep_stash: Optional[bool] = True, - manager_cls: Type = TempDirExecutorManager, - machine: Optional[str] = None, - **kwargs, + @unlocked_repo + def _reproduce_queue( + self, queue: BaseStashQueue, **kwargs ) -> Mapping[str, str]: - """Reproduce the specified experiments. + """Reproduce queued experiments. - Args: - revs: If revs is not specified, all stashed experiments will be - reproduced. - keep_stash: If True, stashed experiments will be preserved if they - fail to reproduce successfully. + Arguments: + queue: Experiment queue. Returns: dict mapping successfully reproduced experiment revs to their - hashes. + results. """ - stash_revs = self.stash_revs + exec_results = queue.reproduce() - # to_run contains mapping of: - # input_rev: (stash_index, rev, baseline_rev) - # where input_rev contains the changes to execute (usually a stash - # commit), rev is the original SCM commit to be checked out, and - # baseline_rev is the experiment baseline. - if revs is None: - to_run = dict(stash_revs) - else: - to_run = { - rev: stash_revs[rev] - if rev in stash_revs - else ExpStashEntry(None, rev, rev, None, None) - for rev in revs - } - - logger.debug( - "Reproducing experiment revs '%s'", - ", ".join(rev[:7] for rev in to_run), - ) - - manager = manager_cls.from_stash_entries( - self.scm, - os.path.join(self.repo.tmp_dir, EXEC_TMP_DIR), - self.repo, - to_run, - machine_name=machine, - ) - try: - exec_results = {} - exec_results.update(self._executors_repro(manager, **kwargs)) - finally: - # only drop successfully run stashed experiments - to_drop: List[int] = [ - entry.stash_index - for rev, entry in to_run.items() - if ( - entry.stash_index is not None - and (not keep_stash or rev in exec_results) - ) - ] - for index in sorted(to_drop, reverse=True): - self.stash.drop(index) - - result: Dict[str, str] = {} + results: Dict[str, str] = {} for _, exp_result in exec_results.items(): - result.update(exp_result) - return result - - @unlocked_repo - def _executors_repro( - self, - manager: "BaseExecutorManager", - **kwargs, - ) -> Dict[str, Dict[str, str]]: - """Run dvc repro for the specified BaseExecutors in parallel. - - Returns: - dict mapping stash revs to the successfully executed experiments - for each stash rev. - """ - return manager.exec_queue(self.repo, **kwargs) + results.update(exp_result) + return results def check_baseline(self, exp_rev): baseline_sha = self.repo.scm.get_rev() diff --git a/dvc/repo/experiments/gc.py b/dvc/repo/experiments/gc.py index b9d360683..666d31f65 100644 --- a/dvc/repo/experiments/gc.py +++ b/dvc/repo/experiments/gc.py @@ -42,11 +42,11 @@ def gc( removed = len(to_remove) delete_stashes = [] - for _, entry in repo.experiments.stash_revs.items(): + for _, entry in repo.experiments.workspace_queue.stash.stash_revs.items(): if not queued or entry.baseline_rev not in keep_revs: delete_stashes.append(entry.stash_index) for index in sorted(delete_stashes, reverse=True): - repo.experiments.stash.drop(index) + repo.experiments.workspace_queue.stash.drop(index) removed += len(delete_stashes) return removed diff --git a/dvc/repo/experiments/queue/base.py b/dvc/repo/experiments/queue/base.py index 5bc545185..f54083730 100644 --- a/dvc/repo/experiments/queue/base.py +++ b/dvc/repo/experiments/queue/base.py @@ -165,7 +165,7 @@ class BaseStashQueue(ABC): branch_name = ExpRefInfo.from_ref(branch).name else: branch_name = f"{resume_rev[:7]}" - if self.scm.is_dirty(): + if self.scm.is_dirty(untracked_files=False): logger.info( "Modified checkpoint experiment based on " "'%s' will be created", diff --git a/dvc/repo/experiments/refs.py b/dvc/repo/experiments/refs.py index 7c1448496..c38c1bd9e 100644 --- a/dvc/repo/experiments/refs.py +++ b/dvc/repo/experiments/refs.py @@ -6,6 +6,8 @@ from .exceptions import InvalidExpRefError # refs/exps/01/234abcd.../ EXPS_NAMESPACE = "refs/exps" EXPS_STASH = f"{EXPS_NAMESPACE}/stash" +WORKSPACE_STASH = EXPS_STASH +CELERY_STASH = f"{EXPS_NAMESPACE}/celery/stash" EXEC_NAMESPACE = f"{EXPS_NAMESPACE}/exec" EXEC_APPLY = f"{EXEC_NAMESPACE}/EXEC_APPLY" EXEC_CHECKPOINT = f"{EXEC_NAMESPACE}/EXEC_CHECKPOINT" diff --git a/dvc/repo/experiments/remove.py b/dvc/repo/experiments/remove.py index eb2c9dfd6..181b81f5f 100644 --- a/dvc/repo/experiments/remove.py +++ b/dvc/repo/experiments/remove.py @@ -6,6 +6,7 @@ from dvc.repo.scm_context import scm_context from dvc.scm import iter_revs from .exceptions import UnresolvedExpNamesError +from .queue.base import QueueEntry from .refs import ExpRefInfo from .utils import ( exp_refs, @@ -39,16 +40,16 @@ def remove( return removed if queue: - removed.extend(_clear_stash(repo)) + removed.extend(_clear_queue(repo)) if all_commits: removed.extend(_clear_all_commits(repo, git_remote)) return removed commit_ref_dict: Dict[ExpRefInfo, str] = {} - queued_ref_dict: Dict[int, str] = {} + queue_entry_dict: Dict[str, QueueEntry] = {} if exp_names: _resolve_exp_by_name( - repo, exp_names, commit_ref_dict, queued_ref_dict, git_remote + repo, exp_names, commit_ref_dict, queue_entry_dict, git_remote ) if rev: @@ -59,8 +60,8 @@ def remove( _remove_commited_exps(repo.scm, commit_ref_dict, git_remote) ) - if queued_ref_dict: - removed.extend(_remove_queued_exps(repo, queued_ref_dict)) + if queue_entry_dict: + removed.extend(_remove_queued_exps(repo, queue_entry_dict)) return removed @@ -68,8 +69,8 @@ def remove( def _resolve_exp_by_name( repo: "Repo", exp_names: Union[str, List[str]], - commit_ref_dict: Dict[ExpRefInfo, str], - queued_ref_dict: Dict[int, str], + commit_ref_dict: Dict["ExpRefInfo", str], + queue_entry_dict: Dict[str, QueueEntry], git_remote: Optional[str], ): remained = set() @@ -84,10 +85,10 @@ def _resolve_exp_by_name( commit_ref_dict[exp_ref] = exp_name if not git_remote: - stash_index_dict = _get_queued_index_by_names(repo, remained) - for exp_name, stash_index in stash_index_dict.items(): - if stash_index is not None: - queued_ref_dict[stash_index] = exp_name + _named_entries = _get_queue_entry_by_names(repo, remained) + for exp_name, entry in _named_entries.items(): + if entry is not None: + queue_entry_dict[exp_name] = entry remained.remove(exp_name) if remained: @@ -111,15 +112,11 @@ def _resolve_exp_by_baseline( commit_ref_dict[ref_info] = ref_info.name -def _clear_stash(repo: "Repo") -> List[str]: +def _clear_queue(repo: "Repo") -> List[str]: removed_name_list = [] - for entry in repo.experiments.stash.list(): - new_sha = entry.new_sha.decode("utf8") - name = repo.experiments.get_exact_name(new_sha) - if not name: - name = new_sha[:7] - removed_name_list.append(name) - repo.experiments.stash.clear() + for entry in repo.experiments.celery_queue.iter_queued(): + removed_name_list.append(entry.name or entry.stash_rev[:7]) + repo.experiments.celery_queue.clear() return removed_name_list @@ -130,25 +127,25 @@ def _clear_all_commits(repo, git_remote) -> List: return _remove_commited_exps(repo.scm, ref_infos, git_remote) -def _get_queued_index_by_names( - repo, +def _get_queue_entry_by_names( + repo: "Repo", exp_name_set: Set[str], -) -> Mapping[str, Optional[int]]: +) -> Dict[str, Optional[QueueEntry]]: from scmrepo.exceptions import RevError as InternalRevError result = {} - stash_revs = repo.experiments.stash_revs - for _, entry in stash_revs.items(): + rev_entries = {} + for entry in repo.experiments.celery_queue.iter_queued(): if entry.name in exp_name_set: - result[entry.name] = entry.stash_index + result[entry.name] = entry + else: + rev_entries[entry.stash_rev] = entry - for exp_name in exp_name_set: - if exp_name in result: - continue + for exp_name in exp_name_set.difference(result.keys()): try: rev = repo.scm.resolve_rev(exp_name) - if rev in stash_revs: - result[exp_name] = stash_revs.get(rev).stash_index + if rev in rev_entries: + result[exp_name] = rev_entries[rev] except InternalRevError: result[exp_name] = None return result @@ -174,9 +171,9 @@ def _remove_commited_exps( return list(exp_ref_dict.values()) -def _remove_queued_exps(repo, indexes: Mapping[int, str]) -> List[str]: - index_list = list(indexes.keys()) - index_list.sort(reverse=True) - for index in index_list: - repo.experiments.stash.drop(index) - return list(indexes.values()) +def _remove_queued_exps( + repo: "Repo", named_entries: Mapping[str, QueueEntry] +) -> List[str]: + for entry in named_entries.values(): + repo.experiments.celery_queue.remove(entry.stash_rev) + return list(named_entries.keys()) diff --git a/dvc/repo/experiments/stash.py b/dvc/repo/experiments/stash.py index 71be33082..18f0c4abf 100644 --- a/dvc/repo/experiments/stash.py +++ b/dvc/repo/experiments/stash.py @@ -7,7 +7,7 @@ from scmrepo.git import Stash class ExpStashEntry(NamedTuple): """Experiment stash entry. - stash_index: EXPS_STASH index for this entry. Can be None if this commit + stash_index: Stash index for this entry. Can be None if this commit is not pushed onto the stash ref. head_rev: HEAD Git commit to be checked out for this experiment. baseline_rev: Experiment baseline commit.