diff --git a/dvc/repo/experiments/queue/base.py b/dvc/repo/experiments/queue/base.py index bfd836150..6047f27a5 100644 --- a/dvc/repo/experiments/queue/base.py +++ b/dvc/repo/experiments/queue/base.py @@ -6,6 +6,7 @@ from typing import ( TYPE_CHECKING, Any, Dict, + Generator, Mapping, NamedTuple, Optional, @@ -111,6 +112,10 @@ class BaseStashQueue(ABC): def get(self) -> QueueGetResult: """Pop and return the first item in the queue.""" + @abstractmethod + def iter_queued(self) -> Generator[QueueEntry, None, None]: + """Iterate over items in the queue.""" + @abstractmethod def reproduce(self) -> Mapping[str, Mapping[str, str]]: """Reproduce queued experiments sequentially.""" diff --git a/dvc/repo/experiments/queue/local.py b/dvc/repo/experiments/queue/local.py index d172c3ef7..d0d3b65a0 100644 --- a/dvc/repo/experiments/queue/local.py +++ b/dvc/repo/experiments/queue/local.py @@ -1,7 +1,7 @@ import logging import os from collections import defaultdict -from typing import TYPE_CHECKING, Dict, Mapping +from typing import TYPE_CHECKING, Dict, Generator, Mapping from funcy import cached_property, first @@ -78,6 +78,14 @@ class LocalCeleryQueue(BaseStashQueue): def get(self) -> QueueGetResult: raise NotImplementedError + def iter_queued(self) -> Generator[QueueEntry, None, None]: + for msg in self.celery.iter_queued(): + if msg.headers.get("task") != setup_exp.name: + continue + args, kwargs, _embed = msg.decode() + entry_dict = kwargs.get("entry_dict", args[0]) + yield QueueEntry.from_dict(entry_dict) + def reproduce(self) -> Mapping[str, Mapping[str, str]]: raise NotImplementedError @@ -103,6 +111,18 @@ class WorkspaceQueue(BaseStashQueue): executor = self.setup_executor(self.repo.experiments, entry) return QueueGetResult(entry, executor) + def iter_queued(self) -> Generator[QueueEntry, None, None]: + for rev, entry in self.stash.stash_revs: + yield QueueEntry( + self.repo.root_dir, + self.scm.root_dir, + self.ref, + rev, + entry.baseline_rev, + entry.branch, + entry.name, + ) + def reproduce(self) -> Dict[str, Dict[str, str]]: results: Dict[str, Dict[str, str]] = defaultdict(dict) try: diff --git a/dvc/repo/experiments/show.py b/dvc/repo/experiments/show.py index 609c0de0c..7558ae025 100644 --- a/dvc/repo/experiments/show.py +++ b/dvc/repo/experiments/show.py @@ -193,7 +193,8 @@ def show( onerror=onerror, ) # collect queued (not yet reproduced) experiments - for stash_rev, entry in repo.experiments.stash_revs.items(): + for entry in repo.experiments.celery_queue.iter_queued(): + stash_rev = entry.stash_rev if entry.baseline_rev in found_revs: if stash_rev not in running or not running[stash_rev].get( "last"