exp: add queue.iter_active for iterating over active exps

dvc/repo/experiments/queue/base.py

@@ -116,6 +116,10 @@ class BaseStashQueue(ABC):
     def iter_queued(self) -> Generator[QueueEntry, None, None]:
         """Iterate over items in the queue."""
 
+    @abstractmethod
+    def iter_active(self) -> Generator[QueueEntry, None, None]:
+        """Iterate over items which are being actively processed."""
+
     @abstractmethod
     def reproduce(self) -> Mapping[str, Mapping[str, str]]:
         """Reproduce queued experiments sequentially."""

dvc/repo/experiments/queue/local.py

@@ -8,7 +8,12 @@ from funcy import cached_property, first
 from dvc.exceptions import DvcException
 
 from ..exceptions import ExpQueueEmptyError
-from ..executor.base import EXEC_TMP_DIR, BaseExecutor, ExecutorResult
+from ..executor.base import (
+    EXEC_PID_DIR,
+    EXEC_TMP_DIR,
+    BaseExecutor,
+    ExecutorResult,
+)
 from ..executor.local import WorkspaceExecutor
 from ..refs import EXEC_BRANCH
 from .base import BaseStashQueue, QueueEntry, QueueGetResult
@@ -17,6 +22,7 @@ from .tasks import setup_exp
 if TYPE_CHECKING:
     from dvc.repo.experiments import Experiments
     from dvc_task.app import FSApp
+    from dvc_task.proc.manager import ProcessManager
     from dvc_task.worker import TemporaryWorker
 
 logger = logging.getLogger(__name__)
@@ -49,6 +55,13 @@ class LocalCeleryQueue(BaseStashQueue):
         )
         return app
 
+    @cached_property
+    def proc(self) -> "ProcessManager":
+        from dvc_task.proc.manager import ProcessManager
+
+        pid_dir = os.path.join(self.repo.tmp_dir, EXEC_TMP_DIR, EXEC_PID_DIR)
+        return ProcessManager(pid_dir)
+
     @cached_property
     def worker(self) -> "TemporaryWorker":
         from dvc_task.worker import TemporaryWorker
@@ -86,6 +99,20 @@ class LocalCeleryQueue(BaseStashQueue):
             entry_dict = kwargs.get("entry_dict", args[0])
             yield QueueEntry.from_dict(entry_dict)
 
+    def _iter_processed(self) -> Generator[QueueEntry, None, None]:
+        for msg in self.celery.iter_processed():
+            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 iter_active(self) -> Generator[QueueEntry, None, None]:
+        for entry in self._iter_processed():
+            proc_info = self.proc.get(entry.stash_rev)
+            if proc_info is not None and proc_info.returncode is None:
+                yield entry
+
     def reproduce(self) -> Mapping[str, Mapping[str, str]]:
         raise NotImplementedError
 
@@ -123,6 +150,11 @@ class WorkspaceQueue(BaseStashQueue):
                 entry.name,
             )
 
+    def iter_active(self) -> Generator[QueueEntry, None, None]:
+        # Workspace run state is reflected in the workspace itself and does not
+        # need to be handled via the queue
+        raise NotImplementedError
+
     def reproduce(self) -> Dict[str, Dict[str, str]]:
         results: Dict[str, Dict[str, str]] = defaultdict(dict)
         try:

dvc/repo/experiments/queue/tasks.py

@@ -7,7 +7,7 @@ from celery.utils.log import get_task_logger
 
 from dvc.utils.fs import makedirs, remove
 
-from ..executor.base import EXEC_PID_DIR, EXEC_TMP_DIR, ExecutorInfo
+from ..executor.base import ExecutorInfo
 from ..executor.local import TempDirExecutor
 from .base import BaseStashQueue, QueueEntry
 
@@ -26,15 +26,8 @@ def setup_exp(entry_dict: Dict[str, Any]) -> None:
     executor = BaseStashQueue.setup_executor(
         repo.experiments, entry, TempDirExecutor
     )
-    pid_dir = os.path.join(
-        repo.tmp_dir,
-        EXEC_TMP_DIR,
-        EXEC_PID_DIR,
-    )
-    infofile = os.path.join(
-        pid_dir,
-        f"{entry.stash_rev}{TempDirExecutor.INFOFILE_EXT}",
-    )
+    proc = repo.experiments.celery_queue.proc
+    infofile = repo.experiments.celery_queue.get_infofile_path(entry.stash_rev)
     makedirs(os.path.dirname(infofile), exist_ok=True)
     with open(infofile, "w", encoding="utf-8") as fobj:
         json.dump(executor.info.asdict(), fobj)
@@ -42,7 +35,7 @@ def setup_exp(entry_dict: Dict[str, Any]) -> None:
 
     # schedule execution + cleanup
     chain(
-        run.si(cmd, wdir=pid_dir, name=entry.stash_rev),
+        proc.run(cmd, name=entry.stash_rev),
         collect_exp.s(entry.asdict(), infofile),
         cleanup_exp.si(entry.asdict(), executor.root_dir),
     ).delay()