From bcc2f3b389d167602ad3d765d6716a28b080872d Mon Sep 17 00:00:00 2001
From: karajan1001 <mishanyo1001@gmail.com>
Date: Thu, 26 May 2022 17:55:38 +0800
Subject: [PATCH] Reorganized `get_queue_entry_by_names`

---
 dvc/repo/experiments/queue/base.py  | 34 +++++++++++++++--------------
 dvc/repo/experiments/queue/local.py | 14 +++---------
 dvc/repo/experiments/remove.py      |  7 ++++--
 3 files changed, 26 insertions(+), 29 deletions(-)

diff --git a/dvc/repo/experiments/queue/base.py b/dvc/repo/experiments/queue/base.py
index b30b8b062..5e58dd581 100644
--- a/dvc/repo/experiments/queue/base.py
+++ b/dvc/repo/experiments/queue/base.py
@@ -8,10 +8,12 @@ from typing import (
     Collection,
     Dict,
     Generator,
+    Iterable,
     List,
     Mapping,
     NamedTuple,
     Optional,
+    Tuple,
     Type,
 )
 
@@ -556,29 +558,29 @@ class BaseStashQueue(ABC):
 
         return results
 
-    def get_queue_entry_by_names(
+    def match_queue_entry_by_name(
         self,
         exp_names: Collection[str],
+        *entries: Iterable[QueueEntry],
     ) -> Dict[str, Optional[QueueEntry]]:
         from funcy import concat
-        from scmrepo.exceptions import RevError as InternalRevError
 
-        exp_name_set = set(exp_names)
+        entry_name_dict: Dict[str, QueueEntry] = {}
+        entry_rev_list: List[Tuple[str, QueueEntry]] = []
+        for queue_entry in concat(*entries):
+            entry_name_dict[queue_entry.name] = queue_entry
+            entry_rev_list.append((queue_entry.stash_rev, queue_entry))
+
         result: Dict[str, Optional[QueueEntry]] = {}
-        rev_entries = {}
-
-        for entry in concat(self.iter_queued(), self.iter_active()):
-            if entry.name in exp_name_set:
-                result[entry.name] = entry
+        for exp_name in exp_names:
+            if exp_name in entry_name_dict:
+                result[exp_name] = entry_name_dict[exp_name]
+                continue
+            for rev, entry in entry_rev_list:
+                if rev.startswith(exp_name):
+                    result[exp_name] = entry
+                    break
             else:
-                rev_entries[entry.stash_rev] = entry
-
-        for exp_name in exp_name_set.difference(result.keys()):
-            try:
-                rev = self.scm.resolve_rev(exp_name)
-                if rev in rev_entries:
-                    result[exp_name] = rev_entries[rev]
-            except InternalRevError:
                 result[exp_name] = None
 
         return result
diff --git a/dvc/repo/experiments/queue/local.py b/dvc/repo/experiments/queue/local.py
index b23614a9f..d2512b969 100644
--- a/dvc/repo/experiments/queue/local.py
+++ b/dvc/repo/experiments/queue/local.py
@@ -219,18 +219,14 @@ class LocalCeleryQueue(BaseStashQueue):
 
     def kill(self, revs: Collection[str]) -> None:
         to_kill: Set[QueueEntry] = set()
-        not_active: List[str] = []
         name_dict: Dict[
             str, Optional[QueueEntry]
-        ] = self.get_queue_entry_by_names(set(revs))
+        ] = self.match_queue_entry_by_name(set(revs), self.iter_active())
 
         missing_rev: List[str] = []
-        active_queue_entry = set(self.iter_active())
         for rev, queue_entry in name_dict.items():
             if queue_entry is None:
                 missing_rev.append(rev)
-            elif queue_entry not in active_queue_entry:
-                not_active.append(rev)
             else:
                 to_kill.add(queue_entry)
 
@@ -267,15 +263,11 @@ class LocalCeleryQueue(BaseStashQueue):
     ):
         from dvc.ui import ui
 
-        queue_entry: Optional[QueueEntry] = self.get_queue_entry_by_names(
-            {rev}
+        queue_entry: Optional[QueueEntry] = self.match_queue_entry_by_name(
+            {rev}, self.iter_active()
         ).get(rev)
         if queue_entry is None:
             raise UnresolvedExpNamesError([rev])
-        active_queue_entry = set(self.iter_active())
-        if queue_entry not in active_queue_entry:
-            ui.write("Cannot attach to an unstarted task")
-            return
         for line in self.proc.follow(queue_entry.stash_rev, encoding):
             ui.write(line)
 
diff --git a/dvc/repo/experiments/remove.py b/dvc/repo/experiments/remove.py
index 903bd9025..8e87bc86f 100644
--- a/dvc/repo/experiments/remove.py
+++ b/dvc/repo/experiments/remove.py
@@ -18,6 +18,7 @@ from .utils import (
 
 if TYPE_CHECKING:
     from dvc.repo import Repo
+    from dvc.repo.experiments.queue.local import LocalCeleryQueue
     from dvc.scm import Git
 
 
@@ -85,8 +86,10 @@ def _resolve_exp_by_name(
             commit_ref_dict[exp_ref] = exp_name
 
     if not git_remote:
-        _named_entries = (
-            repo.experiments.celery_queue.get_queue_entry_by_names(remained)
+        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: