exp show: output executor info for successful exps when available

iterative · Jan 29, 2024 · 9c55949 · 9c55949
1 parent df5b582
commit 9c55949
Show file tree

Hide file tree

Showing 2 changed files with 36 additions and 2 deletions.
diff --git a/dvc/repo/experiments/collect.py b/dvc/repo/experiments/collect.py
@@ -261,6 +261,7 @@ def _collect_baseline(
         ref_it = (ref for ref in iter(refs) if ref.startswith(str(ref_info)))
     else:
         ref_it = repo.scm.iter_refs(base=str(ref_info))
+    executors = repo.experiments.celery_queue.collect_success_executors([baseline_rev])
     for ref in ref_it:
         try:
             ref_info = ExpRefInfo.from_ref(ref)
@@ -272,7 +273,11 @@ def _collect_baseline(
         exps = list(collect_branch(repo, exp_rev, baseline_rev, **kwargs))
         if exps:
             exps[0].name = ref_info.name
-            yield ExpRange(exps, name=ref_info.name)
+            yield ExpRange(
+                exps,
+                name=ref_info.name,
+                executor=executors.get(str(ref_info)),
+            )
 
 
 def collect(

diff --git a/dvc/repo/experiments/queue/celery.py b/dvc/repo/experiments/queue/celery.py
@@ -42,7 +42,7 @@
 
     from dvc.repo.experiments.executor.base import ExecutorResult
     from dvc.repo.experiments.refs import ExpRefInfo
-    from dvc.repo.experiments.serialize import ExpRange
+    from dvc.repo.experiments.serialize import ExpExecutor, ExpRange
     from dvc_task.app import FSApp
     from dvc_task.proc.manager import ProcessManager
     from dvc_task.worker import TemporaryWorker
@@ -632,3 +632,32 @@ def collect_failed_data(
                 )
             )
         return result
+
+    def collect_success_executors(
+        self,
+        baseline_revs: Optional[Collection[str]],
+        **kwargs,
+    ) -> Dict[str, "ExpExecutor"]:
+        """Map exp refs to any available successful executors."""
+        from dvc.repo.experiments.serialize import ExpExecutor, LocalExpExecutor
+
+        result: Dict[str, "ExpExecutor"] = {}
+        for entry, exec_result in self.iter_success():
+            if baseline_revs and entry.baseline_rev not in baseline_revs:
+                continue
+            if not (exec_result and exec_result.ref_info):
+                continue
+            proc_info = self.proc.get(entry.stash_rev)
+            if proc_info:
+                local_exec: Optional[LocalExpExecutor] = LocalExpExecutor(
+                    log=proc_info.stdout,
+                    pid=proc_info.pid,
+                    returncode=proc_info.returncode,
+                    task_id=entry.stash_rev,
+                )
+            else:
+                local_exec = None
+            result[str(exec_result.ref_info)] = ExpExecutor(
+                "success", name="dvc-task", local=local_exec
+            )
+        return result