Refactor dvc.repo.experiment.remove to make matching exp by name also

return its queue entry info.

Some review changes.

Co-authored-by: Peter Rowlands (변기호) <[email protected]>

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,

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

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 (

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]