exp: use full task chain when checking active tasks

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]]:

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

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:

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