exp: return exp results for celery runs

dvc/repo/experiments/__init__.py

@@ -144,13 +144,7 @@ class Experiments:
             )
             return [entry.stash_rev]
         if tmp_dir:
-            self.celery_queue.spawn_worker()
-            # wait for task execution to start
-            while not self.celery_queue.proc.get(entry.stash_rev):
-                time.sleep(1)
-            for line in self.celery_queue.proc.follow(entry.stash_rev):
-                ui.write(line, end="")
-            return
+            return self.reproduce_celery(entries=[entry])
         if machine:
             # TODO: decide how to handle queued remote execution
             raise NotImplementedError
@@ -167,14 +161,41 @@ class Experiments:
             return last_applied
         return None
 
-    def reproduce_queued(self, **kwargs):
+    def reproduce_celery(
+        self, entries: Optional[Iterable[QueueEntry]] = None, **kwargs
+    ) -> Dict[str, str]:
+        results: Dict[str, str] = {}
+        if entries is None:
+            entries = list(self.celery_queue.iter_queued())
+        if not entries:
+            return results
+
         self.celery_queue.spawn_worker()
-        # wait for task execution to start
-        while not first(self.celery_queue.iter_active()):
-            time.sleep(1)
-        for entry in self.celery_queue.iter_active():
+        failed = []
+        for entry in entries:
+            # wait for task execution to start
+            while not self.celery_queue.proc.get(entry.stash_rev):
+                time.sleep(1)
             for line in self.celery_queue.proc.follow(entry.stash_rev):
                 ui.write(line, end="")
+            # wait for task collection to complete
+            while True:
+                result = self.celery_queue.get_result(entry)
+                if result is not None:
+                    if result.exp_hash is None:
+                        name = entry.name or entry.stash_rev[:7]
+                        failed.append(name)
+                        break
+                    exp_rev = self.scm.get_ref(str(result.ref_info))
+                    if exp_rev is not None:
+                        results[exp_rev] = result.exp_hash
+                        break
+                time.sleep(1)
+        if failed:
+            names = ", ".join(name for name in failed)
+            ui.error(f"Failed to reproduce experiment(s) '{names}'")
+        self._log_reproduced((rev for rev in results), True)
+        return results
 
     def _log_reproduced(self, revs: Iterable[str], tmp_dir: bool = False):
         names = []

dvc/repo/experiments/queue/base.py

@@ -162,6 +162,10 @@ class BaseStashQueue(ABC):
     def reproduce(self) -> Mapping[str, Mapping[str, str]]:
         """Reproduce queued experiments sequentially."""
 
+    @abstractmethod
+    def get_result(self, entry: QueueEntry) -> Optional[ExecutorResult]:
+        """Return result of the specified item."""
+
     def _stash_exp(
         self,
         *args,

dvc/repo/experiments/queue/local.py

@@ -1,7 +1,16 @@
+import json
 import logging
 import os
 from collections import defaultdict
-from typing import TYPE_CHECKING, Dict, Generator, List, Mapping, NamedTuple
+from typing import (
+    TYPE_CHECKING,
+    Dict,
+    Generator,
+    List,
+    Mapping,
+    NamedTuple,
+    Optional,
+)
 
 from funcy import cached_property, first
 from kombu.message import Message
@@ -13,6 +22,7 @@ from ..executor.base import (
     EXEC_PID_DIR,
     EXEC_TMP_DIR,
     BaseExecutor,
+    ExecutorInfo,
     ExecutorResult,
 )
 from ..executor.local import WorkspaceExecutor
@@ -140,6 +150,15 @@ class LocalCeleryQueue(BaseStashQueue):
     def reproduce(self) -> Mapping[str, Mapping[str, str]]:
         raise NotImplementedError
 
+    def get_result(self, entry: QueueEntry) -> Optional[ExecutorResult]:
+        infofile = self.get_infofile_path(entry.stash_rev)
+        try:
+            with open(infofile, encoding="utf-8") as fobj:
+                executor_info = ExecutorInfo.from_dict(json.load(fobj))
+        except FileNotFoundError:
+            return None
+        return executor_info.result
+
 
 class WorkspaceQueue(BaseStashQueue):
     def put(self, *args, **kwargs) -> QueueEntry:
@@ -253,3 +272,6 @@ class WorkspaceQueue(BaseStashQueue):
             logger.debug("Collected experiment '%s'.", exp_rev[:7])
             results[exp_rev] = exec_result.exp_hash
         return results
+
+    def get_result(self, entry: QueueEntry) -> Optional[ExecutorResult]:
+        raise NotImplementedError