From dc56c4c9b031fbe4f56ff78db6c503eeae636959 Mon Sep 17 00:00:00 2001
From: Peter Rowlands <peter@pmrowla.com>
Date: Tue, 24 May 2022 16:58:12 +0900
Subject: [PATCH] exp: use full task chain when checking active tasks

---
 dvc/repo/experiments/queue/base.py  | 11 +++++++++--
 dvc/repo/experiments/queue/local.py | 21 +++++++++++++++++++--
 dvc/repo/experiments/queue/tasks.py |  9 +++++----
 dvc/repo/experiments/show.py        |  2 +-
 4 files changed, 34 insertions(+), 9 deletions(-)

diff --git a/dvc/repo/experiments/queue/base.py b/dvc/repo/experiments/queue/base.py
index 1f895a64b..48d9b5a84 100644
--- a/dvc/repo/experiments/queue/base.py
+++ b/dvc/repo/experiments/queue/base.py
@@ -156,8 +156,15 @@ class BaseStashQueue(ABC):
         """Iterate over items in the queue."""
 
     @abstractmethod
-    def iter_active(self) -> Generator[QueueEntry, None, None]:
-        """Iterate over items which are being actively processed."""
+    def iter_active(
+        self, ignore_collected: bool = False
+    ) -> Generator[QueueEntry, None, None]:
+        """Iterate over items which are being actively processed.
+
+        Arguments:
+            ignore_collected: Ignore experiments which have been collected (but
+                not cleaned up).
+        """
 
     @abstractmethod
     def reproduce(self) -> Mapping[str, Mapping[str, str]]:
diff --git a/dvc/repo/experiments/queue/local.py b/dvc/repo/experiments/queue/local.py
index 944249cb0..99d82016f 100644
--- a/dvc/repo/experiments/queue/local.py
+++ b/dvc/repo/experiments/queue/local.py
@@ -181,9 +181,24 @@ class LocalCeleryQueue(BaseStashQueue):
             result = AsyncResult(task_id)
             if not result.ready():
                 yield _TaskEntry(task_id, entry)
+                continue
+            cleanup_id = result.get()
+            cleanup_result = AsyncResult(cleanup_id)
+            if not cleanup_result.ready():
+                yield _TaskEntry(task_id, entry)
 
-    def iter_active(self) -> Generator[QueueEntry, None, None]:
+    def iter_active(
+        self, ignore_collected: bool = False
+    ) -> Generator[QueueEntry, None, None]:
         for _, entry in self._iter_active_tasks():
+            if ignore_collected:
+                infofile = self.get_infofile_path(entry.stash_rev)
+                try:
+                    executor_info = ExecutorInfo.load_json(infofile)
+                    if executor_info.collected:
+                        continue
+                except FileNotFoundError:
+                    pass
             yield entry
 
     def reproduce(self) -> Mapping[str, Mapping[str, str]]:
@@ -288,7 +303,9 @@ class WorkspaceQueue(BaseStashQueue):
                 entry.name,
             )
 
-    def iter_active(self) -> Generator[QueueEntry, None, None]:
+    def iter_active(
+        self, ignore_collected: bool = False
+    ) -> 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
diff --git a/dvc/repo/experiments/queue/tasks.py b/dvc/repo/experiments/queue/tasks.py
index e3fc60dae..1c69a8374 100644
--- a/dvc/repo/experiments/queue/tasks.py
+++ b/dvc/repo/experiments/queue/tasks.py
@@ -20,7 +20,8 @@ def setup_exp(entry_dict: Dict[str, Any]) -> str:
         entry_dict: Serialized QueueEntry for this experiment.
 
     Returns:
-        Celery task ID for the queued experiment task chain.
+        Celery task ID for the final task in the queued task chain (experiment
+        cleanup task ID).
     """
     from dvc.repo import Repo
     from dvc_task.proc.tasks import run
@@ -43,9 +44,9 @@ def setup_exp(entry_dict: Dict[str, Any]) -> str:
         collect_exp.s(entry.asdict(), infofile),
         cleanup_exp.si(entry.asdict(), executor.root_dir),
     )
-    exp_chain.freeze()
+    cleanup_task = exp_chain.freeze()
     exp_chain.delay()
-    return exp_chain.id
+    return cleanup_task.id
 
 
 @shared_task
@@ -82,7 +83,7 @@ def collect_exp(
         executor_info.dump_json(infofile)
 
 
-@shared_task(ignore_result=True)
+@shared_task
 def cleanup_exp(  # pylint: disable=unused-argument
     entry_dict: Dict[str, Any], tmp_dir: str
 ) -> None:
diff --git a/dvc/repo/experiments/show.py b/dvc/repo/experiments/show.py
index 1ed7d76c8..3e2fbf445 100644
--- a/dvc/repo/experiments/show.py
+++ b/dvc/repo/experiments/show.py
@@ -195,7 +195,7 @@ def show(
             )
         # collect celery experiments
         for entry in chain(
-            repo.experiments.celery_queue.iter_active(),
+            repo.experiments.celery_queue.iter_active(ignore_collected=True),
             repo.experiments.celery_queue.iter_queued(),
         ):
             stash_rev = entry.stash_rev