Refactor: move get_running_exps to queue

Move logic of `get_running_exps` from `repo.experiments` to each queue.

dvc/repo/experiments/__init__.py

@@ -2,7 +2,7 @@ import logging
 import os
 import re
 import time
-from typing import Any, Dict, Iterable, Optional
+from typing import Dict, Iterable, Optional
 
 from funcy import cached_property, first
 
@@ -16,7 +16,7 @@ from .exceptions import (
     InvalidExpRefError,
     MultipleBranchError,
 )
-from .executor.base import BaseExecutor, ExecutorInfo, TaskStatus
+from .executor.base import BaseExecutor
 from .queue.base import BaseStashQueue, QueueEntry
 from .queue.celery import LocalCeleryQueue
 from .queue.tempdir import TempDirQueue
@@ -25,7 +25,6 @@ from .refs import (
     CELERY_FAILED_STASH,
     CELERY_STASH,
     EXEC_APPLY,
-    EXEC_BRANCH,
     EXEC_CHECKPOINT,
     EXEC_NAMESPACE,
     EXPS_NAMESPACE,
@@ -438,74 +437,15 @@ class Experiments:
             result[rev] = name
         return result
 
-    def get_running_exps(self, fetch_refs: bool = True) -> Dict[str, Any]:
+    def get_running_exps(self, fetch_refs: bool = True) -> Dict[str, Dict]:
         """Return info for running experiments."""
         result = {}
-        infofile = self.workspace_queue.get_infofile_path("workspace")
-        result.update(
-            self._fetch_running_exp("workspace", infofile, fetch_refs)
-        )
-        for queue in (self.tempdir_queue, self.celery_queue):
-            for entry in queue.iter_active():
-                infofile = queue.get_infofile_path(entry.stash_rev)
-                result.update(
-                    self._fetch_running_exp(
-                        entry.stash_rev, infofile, fetch_refs
-                    )
-                )
-        return result
-
-    def _fetch_running_exp(
-        self, rev: str, infofile: str, fetch_refs: bool
-    ) -> Dict[str, Any]:
-        from dvc.scm import InvalidRemoteSCMRepo
-        from dvc.utils.serialize import load_json
-
-        from .executor.local import TempDirExecutor
-
-        result: Dict[str, Any] = {}
-        try:
-            info = ExecutorInfo.from_dict(load_json(infofile))
-        except OSError:
-            return result
-        if info.status < TaskStatus.FAILED:
-            if rev == "workspace":
-                # If we are appending to a checkpoint branch in a workspace
-                # run, show the latest checkpoint as running.
-                if info.status == TaskStatus.SUCCESS:
-                    return result
-                last_rev = self.scm.get_ref(EXEC_BRANCH)
-                if last_rev:
-                    result[last_rev] = info.asdict()
-                else:
-                    result[rev] = info.asdict()
-            else:
-                result[rev] = info.asdict()
-                if (
-                    info.git_url
-                    and fetch_refs
-                    and info.status > TaskStatus.PREPARING
-                ):
-
-                    def on_diverged(_ref: str, _checkpoint: bool):
-                        return False
-
-                    executor = TempDirExecutor.from_info(info)
-                    try:
-                        for ref in executor.fetch_exps(
-                            self.scm,
-                            on_diverged=on_diverged,
-                        ):
-                            logger.debug(
-                                "Updated running experiment '%s'.", ref
-                            )
-                            last_rev = self.scm.get_ref(ref)
-                            result[rev]["last"] = last_rev
-                            if last_rev:
-                                result[last_rev] = info.asdict()
-                    except InvalidRemoteSCMRepo:
-                        # ignore stale info files
-                        del result[rev]
+        for queue in (
+            self.workspace_queue,
+            self.tempdir_queue,
+            self.celery_queue,
+        ):
+            result.update(queue.get_running_exps(fetch_refs))
         return result
 
     def apply(self, *args, **kwargs):

dvc/repo/experiments/queue/base.py

@@ -31,7 +31,9 @@ from ..executor.base import (
     EXEC_PID_DIR,
     EXEC_TMP_DIR,
     BaseExecutor,
+    ExecutorInfo,
     ExecutorResult,
+    TaskStatus,
 )
 from ..executor.local import WorkspaceExecutor
 from ..refs import ExpRefInfo
@@ -706,3 +708,57 @@ class BaseStashQueue(ABC):
                 entry.stash_rev,
                 message=f"commit: {msg}",
             )
+
+    def _fetch_running_exp(
+        self, rev: str, infofile: str, fetch_refs: bool
+    ) -> Dict[str, Dict]:
+        from dvc.scm import InvalidRemoteSCMRepo
+        from dvc.utils.serialize import load_json
+
+        from ..executor.local import TempDirExecutor
+
+        result: Dict[str, Dict] = {}
+        try:
+            info = ExecutorInfo.from_dict(load_json(infofile))
+        except OSError:
+            return result
+        if info.status < TaskStatus.FAILED:
+            result[rev] = info.asdict()
+            if (
+                info.git_url
+                and fetch_refs
+                and info.status > TaskStatus.PREPARING
+            ):
+
+                def on_diverged(_ref: str, _checkpoint: bool):
+                    return False
+
+                executor = TempDirExecutor.from_info(info)
+                try:
+                    for ref in executor.fetch_exps(
+                        self.scm,
+                        on_diverged=on_diverged,
+                    ):
+                        logger.debug("Updated running experiment '%s'.", ref)
+                        last_rev = self.scm.get_ref(ref)
+                        result[rev]["last"] = last_rev
+                        if last_rev:
+                            result[last_rev] = info.asdict()
+                except InvalidRemoteSCMRepo:
+                    # ignore stale info files
+                    del result[rev]
+        return result
+
+    def get_running_exps(self, fetch_refs: bool = True) -> Dict[str, Dict]:
+        """Get the execution info of the currently running experiments
+
+        Args:
+            fetch_ref (bool): fetch completed checkpoints or not.
+        """
+        result: Dict[str, Dict] = {}
+        for entry in self.iter_active():
+            infofile = self.get_infofile_path(entry.stash_rev)
+            result.update(
+                self._fetch_running_exp(entry.stash_rev, infofile, fetch_refs)
+            )
+        return result

dvc/repo/experiments/queue/tempdir.py

@@ -98,3 +98,6 @@ class TempDirQueue(WorkspaceQueue):
         exec_result: ExecutorResult,
     ) -> Dict[str, str]:
         return BaseStashQueue.collect_executor(exp, executor, exec_result)
+
+    def get_running_exps(self, fetch_refs: bool = True) -> Dict[str, Dict]:
+        return super(WorkspaceQueue, self).get_running_exps(fetch_refs)

dvc/repo/experiments/queue/workspace.py

@@ -7,7 +7,12 @@ from funcy import first
 from dvc.exceptions import DvcException
 
 from ..exceptions import ExpQueueEmptyError
-from ..executor.base import BaseExecutor, ExecutorResult
+from ..executor.base import (
+    BaseExecutor,
+    ExecutorInfo,
+    ExecutorResult,
+    TaskStatus,
+)
 from ..executor.local import WorkspaceExecutor
 from ..refs import EXEC_BRANCH
 from .base import BaseStashQueue, QueueDoneResult, QueueEntry, QueueGetResult
@@ -150,3 +155,28 @@ class WorkspaceQueue(BaseStashQueue):
         follow: bool = False,
     ):
         raise NotImplementedError
+
+    def get_running_exps(self, fetch_refs: bool = True) -> Dict[str, Dict]:
+        from dvc.utils.serialize import load_json
+
+        assert self._EXEC_NAME
+
+        result: Dict[str, Dict] = {}
+        infofile = self.get_infofile_path(self._EXEC_NAME)
+
+        try:
+            info = ExecutorInfo.from_dict(load_json(infofile))
+        except OSError:
+            return result
+
+        if info.status < TaskStatus.FAILED:
+            # If we are appending to a checkpoint branch in a workspace
+            # run, show the latest checkpoint as running.
+            if info.status == TaskStatus.SUCCESS:
+                return result
+            last_rev = self.scm.get_ref(EXEC_BRANCH)
+            if last_rev:
+                result[last_rev] = info.asdict()
+            else:
+                result[self._EXEC_NAME] = info.asdict()
+        return result

dvc/repo/experiments/show.py

@@ -287,7 +287,9 @@ def show(
         iter_revs(repo.scm, revs, num, all_branches, all_tags, all_commits)
     )
 
-    running = repo.experiments.get_running_exps(fetch_refs=fetch_running)
+    running: Dict[str, Dict] = repo.experiments.get_running_exps(
+        fetch_refs=fetch_running
+    )
 
     queued_experiment = (
         _collect_queued_experiment(

tests/func/experiments/test_show.py

@@ -487,7 +487,7 @@ def test_show_running_tempdir(tmp_dir, scm, dvc, exp_stage, mocker):
     makedirs(os.path.dirname(pidfile), True)
     (tmp_dir / pidfile).dump_json(info.asdict())
     mock_fetch = mocker.patch.object(
-        dvc.experiments,
+        dvc.experiments.tempdir_queue,
         "_fetch_running_exp",
         return_value={exp_rev: info.asdict()},
     )