exp: wait for task collection when getting exp result

dvc/repo/experiments/__init__.py

@@ -177,22 +177,18 @@ class Experiments:
             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)
+            result = self.celery_queue.get_result(entry)
+            if result is None or result.exp_hash is None:
+                name = entry.name or entry.stash_rev[:7]
+                failed.append(name)
+            elif result.ref_info:
+                exp_rev = self.scm.get_ref(str(result.ref_info))
+                results[exp_rev] = result.exp_hash
         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)
+        if results:
+            self._log_reproduced((rev for rev in results), True)
         return results
 
     def _log_reproduced(self, revs: Iterable[str], tmp_dir: bool = False):

dvc/repo/experiments/executor/base.py

@@ -73,6 +73,7 @@ class ExecutorInfo:
     dvc_dir: str
     name: Optional[str] = None
     wdir: Optional[str] = None
+    collected: bool = False
     result_hash: Optional[str] = None
     result_ref: Optional[str] = None
     result_force: bool = False
@@ -102,6 +103,12 @@ class ExecutorInfo:
         with modify_json(filename) as d:
             d.update(self.asdict())
 
+    @classmethod
+    def load_json(cls, filename: str) -> "ExecutorInfo":
+        from dvc.utils.serialize import load_json
+
+        return cls.from_dict(load_json(filename))
+
 
 _T = TypeVar("_T", bound="BaseExecutor")
 

dvc/repo/experiments/queue/base.py

@@ -165,7 +165,10 @@ class BaseStashQueue(ABC):
 
     @abstractmethod
     def get_result(self, entry: QueueEntry) -> Optional[ExecutorResult]:
-        """Return result of the specified item."""
+        """Return result of the specified item.
+
+        This method blocks until the specified item has been collected.
+        """
 
     def _stash_exp(
         self,

dvc/repo/experiments/queue/local.py

@@ -1,7 +1,7 @@
 import hashlib
-import json
 import logging
 import os
+import time
 from collections import defaultdict
 from typing import (
     TYPE_CHECKING,
@@ -155,12 +155,15 @@ class LocalCeleryQueue(BaseStashQueue):
 
     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
+        while True:
+            try:
+                executor_info = ExecutorInfo.load_json(infofile)
+                if executor_info.collected:
+                    return executor_info.result
+            except FileNotFoundError:
+                # Infofile will not be created until execution begins
+                pass
+            time.sleep(1)
 
 
 class WorkspaceQueue(BaseStashQueue):

dvc/repo/experiments/queue/tasks.py

@@ -1,11 +1,9 @@
-import json
-import os
 from typing import Any, Dict
 
 from celery import chain, shared_task
 from celery.utils.log import get_task_logger
 
-from dvc.utils.fs import makedirs, remove
+from dvc.utils.fs import remove
 
 from ..executor.base import ExecutorInfo
 from ..executor.local import TempDirExecutor
@@ -28,9 +26,7 @@ def setup_exp(entry_dict: Dict[str, Any]) -> None:
     )
     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)
+    executor.info.dump_json(infofile)
     cmd = ["dvc", "exp", "exec-run", "--infofile", infofile]
 
     # schedule execution + cleanup
@@ -57,19 +53,22 @@ def collect_exp(
 
     entry = QueueEntry.from_dict(entry_dict)
     repo = Repo(entry.dvc_root)
-    with open(infofile, encoding="utf-8") as fobj:
-        executor_info = ExecutorInfo.from_dict(json.load(fobj))
+    executor_info = ExecutorInfo.load_json(infofile)
     logger.debug("Collecting experiment info '%s'", str(executor_info))
     executor = TempDirExecutor.from_info(executor_info)
     exec_result = executor_info.result
-    if exec_result is not None:
-        results = BaseStashQueue.collect_executor(
-            repo.experiments, executor, exec_result
-        )
-        for rev in results:
-            logger.debug("Collected experiment '%s'", rev[:7])
-    else:
-        logger.debug("Exec result was None")
+    try:
+        if exec_result is not None:
+            results = BaseStashQueue.collect_executor(
+                repo.experiments, executor, exec_result
+            )
+            for rev in results:
+                logger.debug("Collected experiment '%s'", rev[:7])
+        else:
+            logger.debug("Exec result was None")
+    finally:
+        executor_info.collected = True
+        executor_info.dump_json(infofile)
 
 
 @shared_task(ignore_result=True)