diff --git a/dvc/commands/experiments/remove.py b/dvc/commands/experiments/remove.py index bd7bf96e0..9814f1400 100644 --- a/dvc/commands/experiments/remove.py +++ b/dvc/commands/experiments/remove.py @@ -10,7 +10,7 @@ logger = logging.getLogger(__name__) class CmdExperimentsRemove(CmdBase): - def argument_checking(self): + def check_arguments(self): if not any( [ self.args.all_commits, @@ -25,7 +25,7 @@ class CmdExperimentsRemove(CmdBase): def run(self): - self.argument_checking() + self.check_arguments() removed_list = self.repo.experiments.remove( exp_names=self.args.experiment, diff --git a/dvc/commands/queue/remove.py b/dvc/commands/queue/remove.py index 760ab76ce..6699b5fd2 100644 --- a/dvc/commands/queue/remove.py +++ b/dvc/commands/queue/remove.py @@ -12,7 +12,7 @@ logger = logging.getLogger(__name__) class CmdQueueRemove(CmdBase): """Remove exp in queue.""" - def argument_checking(self): + def check_arguments(self): clear_flag = any( [ self.args.all, @@ -29,7 +29,7 @@ class CmdQueueRemove(CmdBase): def run(self): - self.argument_checking() + self.check_arguments() if self.args.all: self.args.queued = True diff --git a/dvc/repo/experiments/queue/remove.py b/dvc/repo/experiments/queue/remove.py index ca757dd18..605e5effa 100644 --- a/dvc/repo/experiments/queue/remove.py +++ b/dvc/repo/experiments/queue/remove.py @@ -23,17 +23,15 @@ def remove_tasks( stash_revs: Dict[str, "ExpStashEntry"] = {} failed_stash_revs: List["ExpStashEntry"] = [] done_entry_set: Set["QueueEntry"] = set() + stash_rev_all = celery_queue.stash.stash_revs + failed_rev_all = celery_queue.failed_stash.stash_revs for entry in queue_entries: - if entry.stash_rev in celery_queue.stash.stash_revs: - stash_revs[entry.stash_rev] = celery_queue.stash.stash_revs[ - entry.stash_rev - ] + if entry.stash_rev in stash_rev_all: + stash_revs[entry.stash_rev] = stash_rev_all[entry.stash_rev] else: done_entry_set.add(entry) - if entry.stash_rev in celery_queue.failed_stash.stash_revs: - failed_stash_revs.append( - celery_queue.failed_stash.stash_revs[entry.stash_rev] - ) + if entry.stash_rev in failed_rev_all: + failed_stash_revs.append(failed_rev_all[entry.stash_rev]) try: for ( diff --git a/dvc/repo/experiments/remove.py b/dvc/repo/experiments/remove.py index 8b353ce1c..a64503e13 100644 --- a/dvc/repo/experiments/remove.py +++ b/dvc/repo/experiments/remove.py @@ -1,12 +1,20 @@ import logging -from typing import TYPE_CHECKING, Dict, List, Mapping, Optional, Union +from typing import ( + TYPE_CHECKING, + Dict, + Iterable, + List, + NamedTuple, + Optional, + Union, +) from dvc.repo import locked from dvc.repo.scm_context import scm_context from dvc.scm import iter_revs from .exceptions import UnresolvedExpNamesError -from .queue.base import QueueEntry +from .queue.remove import remove_tasks from .refs import ExpRefInfo from .utils import ( exp_refs, @@ -21,10 +29,56 @@ if TYPE_CHECKING: from dvc.repo.experiments.queue.celery import LocalCeleryQueue from dvc.scm import Git + from .queue.base import QueueEntry + logger = logging.getLogger(__name__) +class ExpRefAndQueueEntry(NamedTuple): + exp_ref_list: List["ExpRefInfo"] + queue_entry_list: List["QueueEntry"] + removed: List["str"] + + +def _get_ref_and_entry_by_names( + exp_names: Union[str, List[str]], + scm: "Git", + celery_queue: "LocalCeleryQueue", + git_remote: Optional[str], +) -> ExpRefAndQueueEntry: + exp_ref_list: List["ExpRefInfo"] = [] + queue_entry_list: List["QueueEntry"] = [] + removed: List[str] = [] + if isinstance(exp_names, str): + exp_names = [exp_names] + exp_ref_match: Dict[str, Optional["ExpRefInfo"]] = resolve_name( + scm, exp_names, git_remote + ) + if not git_remote: + queue_entry_match: Dict[ + str, Optional["QueueEntry"] + ] = celery_queue.match_queue_entry_by_name( + exp_names, celery_queue.iter_queued() + ) + + remained = [] + for exp_name in exp_names: + exp_ref = exp_ref_match[exp_name] + queue_entry = None if git_remote else queue_entry_match[exp_name] + if exp_ref or queue_entry: + if exp_ref: + exp_ref_list.append(exp_ref) + if queue_entry: + queue_entry_list.append(queue_entry) + removed.append(exp_name) + else: + remained.append(exp_name) + if remained: + raise UnresolvedExpNamesError(remained) + return ExpRefAndQueueEntry(exp_ref_list, queue_entry_list, removed) + + @locked @scm_context def remove( @@ -39,98 +93,64 @@ def remove( removed: List[str] = [] if not any([exp_names, queue, all_commits, rev]): return removed + celery_queue: "LocalCeleryQueue" = repo.experiments.celery_queue if queue: - removed.extend(repo.experiments.celery_queue.clear(queued=True)) + removed.extend(celery_queue.clear(queued=True)) if all_commits: - removed.extend(_clear_all_commits(repo, git_remote)) + removed.extend( + _remove_commited_exps( + repo.scm, list(exp_refs(repo.scm, git_remote)), git_remote + ) + ) return removed - commit_ref_dict: Dict[ExpRefInfo, str] = {} - queue_entry_dict: Dict[str, QueueEntry] = {} + exp_ref_list: List["ExpRefInfo"] = [] + queue_entry_list: List["QueueEntry"] = [] if exp_names: - _resolve_exp_by_name( - repo, exp_names, commit_ref_dict, queue_entry_dict, git_remote + result = _get_ref_and_entry_by_names( + exp_names, repo.scm, celery_queue, git_remote ) + removed.extend(exp_names) + exp_ref_list.extend(result.exp_ref_list) + queue_entry_list.extend(result.queue_entry_list) + elif rev: + exp_ref_dict = _resolve_exp_by_baseline(repo, rev, num, git_remote) + removed.extend(exp_ref_dict.keys()) + exp_ref_list.extend(exp_ref_dict.values()) - if rev: - _resolve_exp_by_baseline(repo, rev, num, commit_ref_dict, git_remote) + if exp_ref_list: + _remove_commited_exps(repo.scm, exp_ref_list, git_remote) - if commit_ref_dict: - removed.extend( - _remove_commited_exps(repo.scm, commit_ref_dict, git_remote) - ) - - if queue_entry_dict: - removed.extend(_remove_queued_exps(repo, queue_entry_dict)) + if queue_entry_list: + remove_tasks(celery_queue, queue_entry_list) return removed -def _resolve_exp_by_name( - repo: "Repo", - exp_names: Union[str, List[str]], - commit_ref_dict: Dict["ExpRefInfo", str], - queue_entry_dict: Dict[str, QueueEntry], - git_remote: Optional[str], -): - remained = set() - if isinstance(exp_names, str): - exp_names = [exp_names] - - exp_ref_dict = resolve_name(repo.scm, exp_names, git_remote) - for exp_name, exp_ref in exp_ref_dict.items(): - if exp_ref is None: - remained.add(exp_name) - else: - commit_ref_dict[exp_ref] = exp_name - - if not git_remote: - celery_queue: "LocalCeleryQueue" = repo.experiments.celery_queue - - _named_entries = celery_queue.match_queue_entry_by_name( - remained, celery_queue.iter_queued(), celery_queue.iter_active() - ) - 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: - raise UnresolvedExpNamesError(remained) - - def _resolve_exp_by_baseline( repo, rev: str, num: int, - commit_ref_dict: Dict[ExpRefInfo, str], git_remote: Optional[str] = None, -): +) -> Dict[str, "ExpRefInfo"]: + commit_ref_dict: Dict[str, "ExpRefInfo"] = {} rev_dict = iter_revs(repo.scm, [rev], num) rev_set = set(rev_dict.keys()) ref_info_dict = exp_refs_by_baseline(repo.scm, rev_set, git_remote) - for _, ref_info_list in ref_info_dict.items(): for ref_info in ref_info_list: - if ref_info not in commit_ref_dict: - commit_ref_dict[ref_info] = ref_info.name - - -def _clear_all_commits(repo, git_remote) -> List: - ref_infos = { - ref_info: ref_info.name for ref_info in exp_refs(repo.scm, git_remote) - } - return _remove_commited_exps(repo.scm, ref_infos, git_remote) + commit_ref_dict[ref_info.name] = ref_info + return commit_ref_dict def _remove_commited_exps( - scm: "Git", exp_ref_dict: Mapping[ExpRefInfo, str], remote: Optional[str] + scm: "Git", exp_refs_list: Iterable[ExpRefInfo], remote: Optional[str] ) -> List[str]: if remote: from dvc.scm import TqdmGit - for ref_info in exp_ref_dict: + for ref_info in exp_refs_list: with TqdmGit(desc="Pushing git refs") as pbar: push_refspec( scm, @@ -140,13 +160,5 @@ def _remove_commited_exps( progress=pbar.update_git, ) else: - remove_exp_refs(scm, exp_ref_dict) - return list(exp_ref_dict.values()) - - -def _remove_queued_exps( - repo: "Repo", named_entries: Mapping[str, QueueEntry] -) -> List[str]: - stash_rev_list = [entry.stash_rev for entry in named_entries.values()] - repo.experiments.celery_queue.remove(stash_rev_list) - return list(named_entries.keys()) + remove_exp_refs(scm, exp_refs_list) + return [exp_ref.name for exp_ref in exp_refs_list]