From 4e82eb49781a050da9678a1581cca4a643d32444 Mon Sep 17 00:00:00 2001 From: karajan1001 Date: Fri, 3 Jun 2022 11:30:48 +0800 Subject: [PATCH] queue remove: add remove support for short revs(#7840) fix: #7840 Now queue remove can not accept short rev as name. 1. Refactor the queue entry match in `queue remove` 2. Better output if no queue task found for name. 3. short rev on kill can only be work for rev with lens longer than 4 --- dvc/commands/queue/remove.py | 9 ++++-- dvc/repo/experiments/queue/base.py | 46 ++++++++++++++++++----------- dvc/repo/experiments/queue/local.py | 18 +---------- dvc/repo/experiments/remove.py | 4 +-- 4 files changed, 38 insertions(+), 39 deletions(-) diff --git a/dvc/commands/queue/remove.py b/dvc/commands/queue/remove.py index 6d9572e0c..7b899ee44 100644 --- a/dvc/commands/queue/remove.py +++ b/dvc/commands/queue/remove.py @@ -19,8 +19,13 @@ class CmdQueueRemove(CmdBase): revs=self.args.experiment ) - removed = ", ".join(removed_list) - ui.write(f"Removed experiments in queue: {removed}") + if removed_list: + removed = ", ".join(removed_list) + ui.write(f"Removed experiments in queue: {removed}") + else: + ui.write( + f"No experiments found in queue named {self.args.experiment}" + ) return 0 diff --git a/dvc/repo/experiments/queue/base.py b/dvc/repo/experiments/queue/base.py index aef4902ea..ea1f27530 100644 --- a/dvc/repo/experiments/queue/base.py +++ b/dvc/repo/experiments/queue/base.py @@ -13,7 +13,6 @@ from typing import ( Mapping, NamedTuple, Optional, - Tuple, Type, Union, ) @@ -141,15 +140,18 @@ class BaseStashQueue(ABC): Returns: Revisions (or names) which were removed. """ - to_remove = {} removed: List[str] = [] - for stash_rev, stash_entry in self.stash.stash_revs.items(): - if stash_rev in revs: - to_remove[stash_rev] = stash_entry - removed.append(stash_rev) - elif stash_entry.name in revs: - to_remove[stash_rev] = stash_entry - removed.append(stash_entry.name) + to_remove: Dict[str, ExpStashEntry] = {} + queue_entries = self.match_queue_entry_by_name( + revs, self.iter_queued() + ) + for name, entry in queue_entries.items(): + if entry: + to_remove[entry.stash_rev] = self.stash.stash_revs[ + entry.stash_rev + ] + removed.append(name) + self._remove_revs(to_remove) return removed @@ -216,9 +218,17 @@ class BaseStashQueue(ABC): ) return result - @abstractmethod def _remove_revs(self, stash_revs: Mapping[str, ExpStashEntry]): """Remove the specified entries from the queue by stash revision.""" + for index in sorted( + ( + entry.stash_index + for entry in stash_revs.values() + if entry.stash_index is not None + ), + reverse=True, + ): + self.stash.drop(index) @abstractmethod def iter_queued(self) -> Generator[QueueEntry, None, None]: @@ -601,7 +611,7 @@ class BaseStashQueue(ABC): from funcy import concat entry_name_dict: Dict[str, QueueEntry] = {} - entry_rev_list: List[Tuple[str, QueueEntry]] = [] + entry_rev_dict: Dict[str, QueueEntry] = {} for entry in concat(*entries): if isinstance(entry, QueueDoneResult): queue_entry: QueueEntry = entry.entry @@ -617,19 +627,19 @@ class BaseStashQueue(ABC): name = queue_entry.name if name: entry_name_dict[name] = queue_entry - entry_rev_list.append((queue_entry.stash_rev, queue_entry)) + entry_rev_dict[queue_entry.stash_rev] = queue_entry result: Dict[str, Optional[QueueEntry]] = {} for exp_name in exp_names: + result[exp_name] = None 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: - result[exp_name] = None + if self.scm.is_sha(exp_name): + for rev in entry_rev_dict: + if rev.startswith(exp_name.lower()): + result[exp_name] = entry_rev_dict[rev] + break return result diff --git a/dvc/repo/experiments/queue/local.py b/dvc/repo/experiments/queue/local.py index 83d8a7152..aac45aa0d 100644 --- a/dvc/repo/experiments/queue/local.py +++ b/dvc/repo/experiments/queue/local.py @@ -149,17 +149,12 @@ class LocalCeleryQueue(BaseStashQueue): raise NotImplementedError def _remove_revs(self, stash_revs: Mapping[str, ExpStashEntry]): - to_drop: List[int] = [] try: for msg, queue_entry in self._iter_queued(): if queue_entry.stash_rev in stash_revs: self.celery.reject(msg.delivery_tag) - stash_entry = stash_revs[queue_entry.stash_rev] - assert stash_entry.stash_index is not None - to_drop.append(stash_entry.stash_index) finally: - for index in sorted(to_drop, reverse=True): - self.stash.drop(index) + super()._remove_revs(stash_revs) def iter_queued(self) -> Generator[QueueEntry, None, None]: for _, entry in self._iter_queued(): @@ -341,17 +336,6 @@ class WorkspaceQueue(BaseStashQueue): executor = self.setup_executor(self.repo.experiments, entry) return QueueGetResult(entry, executor) - def _remove_revs(self, stash_revs: Mapping[str, ExpStashEntry]): - for index in sorted( - ( - entry.stash_index - for entry in stash_revs.values() - if entry.stash_index is not None - ), - reverse=True, - ): - self.stash.drop(index) - def iter_queued(self) -> Generator[QueueEntry, None, None]: for rev, entry in self.stash.stash_revs: yield QueueEntry( diff --git a/dvc/repo/experiments/remove.py b/dvc/repo/experiments/remove.py index 8e87bc86f..de82ade93 100644 --- a/dvc/repo/experiments/remove.py +++ b/dvc/repo/experiments/remove.py @@ -155,6 +155,6 @@ def _remove_commited_exps( 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) + 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())