exp queue: implement queue.logs() support

dvc/repo/experiments/__init__.py

@@ -176,8 +176,7 @@ class Experiments:
             # 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="")
+            self.celery_queue.follow(entry)
             # wait for task collection to complete
             result = self.celery_queue.get_result(entry)
             if result is None or result.exp_hash is None:

dvc/repo/experiments/exceptions.py

@@ -105,18 +105,20 @@ class AmbiguousExpRefInfo(InvalidArgumentError):
 
 
 class UnresolvedExpNamesError(InvalidArgumentError):
+    NAME = "experiment name"
+
     def __init__(
         self, unresolved_list: Collection[str], *args, git_remote: str = None
     ):
-        unresolved_names = ";".join(unresolved_list)
+        unresolved_names = "; ".join(unresolved_list)
         if not git_remote:
-            if len(unresolved_names) > 1:
+            if len(unresolved_list) > 1:
                 super().__init__(
-                    f"'{unresolved_names}' are not valid experiment names"
+                    f"'{unresolved_names}' are not valid {self.NAME}s"
                 )
             else:
                 super().__init__(
-                    f"'{unresolved_names}' is not a valid experiment name"
+                    f"'{unresolved_names}' is not a valid {self.NAME}"
                 )
         else:
             super().__init__(
@@ -125,5 +127,9 @@ class UnresolvedExpNamesError(InvalidArgumentError):
             )
 
 
+class UnresolvedQueueExpNamesError(UnresolvedExpNamesError):
+    NAME = "queued experiment name"
+
+
 class ExpQueueEmptyError(DvcException):
     pass

dvc/repo/experiments/queue/base.py

@@ -15,6 +15,7 @@ from typing import (
     Optional,
     Tuple,
     Type,
+    Union,
 )
 
 from funcy import cached_property
@@ -264,7 +265,7 @@ class BaseStashQueue(ABC):
         rev: str,
         encoding: Optional[str] = None,
     ):
-        """Iterate over lines in redirected output for a process.
+        """Attach to a running exp process and print redirected output.
 
         Args:
             rev: Stash rev or running exp name to be attached.
@@ -272,6 +273,22 @@ class BaseStashQueue(ABC):
                 `locale.getpreferredencoding()`.
         """
 
+    def logs(
+        self,
+        rev: str,
+        encoding: Optional[str] = None,
+        follow: bool = False,
+    ):
+        """Print redirected output logs for an exp process.
+
+        Args:
+            rev: Stash rev or exp name.
+            encoding: Text encoding for redirected output. Defaults to
+                `locale.getpreferredencoding()`.
+            follow: Attach to running exp process and follow additional
+                output.
+        """
+
     def _stash_exp(
         self,
         *args,
@@ -592,14 +609,27 @@ class BaseStashQueue(ABC):
     def match_queue_entry_by_name(
         self,
         exp_names: Collection[str],
-        *entries: Iterable[QueueEntry],
+        *entries: Iterable[Union[QueueEntry, QueueDoneResult]],
     ) -> Dict[str, Optional[QueueEntry]]:
         from funcy import concat
 
         entry_name_dict: Dict[str, QueueEntry] = {}
         entry_rev_list: List[Tuple[str, QueueEntry]] = []
-        for queue_entry in concat(*entries):
-            entry_name_dict[queue_entry.name] = queue_entry
+        for entry in concat(*entries):
+            if isinstance(entry, QueueDoneResult):
+                queue_entry: QueueEntry = entry.entry
+                if (
+                    entry.result is not None
+                    and entry.result.ref_info is not None
+                ):
+                    name: Optional[str] = entry.result.ref_info.name
+                else:
+                    name = queue_entry.name
+            else:
+                queue_entry = entry
+                name = queue_entry.name
+            if name:
+                entry_name_dict[name] = queue_entry
             entry_rev_list.append((queue_entry.stash_rev, queue_entry))
 
         result: Dict[str, Optional[QueueEntry]] = {}

dvc/repo/experiments/queue/local.py

@@ -1,4 +1,5 @@
 import hashlib
+import locale
 import logging
 import os
 from collections import defaultdict
@@ -19,8 +20,9 @@ from kombu.message import Message
 
 from dvc.daemon import daemonize
 from dvc.exceptions import DvcException
+from dvc.ui import ui
 
-from ..exceptions import ExpQueueEmptyError, UnresolvedExpNamesError
+from ..exceptions import ExpQueueEmptyError, UnresolvedQueueExpNamesError
 from ..executor.base import (
     EXEC_PID_DIR,
     EXEC_TMP_DIR,
@@ -252,7 +254,7 @@ class LocalCeleryQueue(BaseStashQueue):
                 to_kill.add(queue_entry)
 
         if missing_rev:
-            raise UnresolvedExpNamesError(missing_rev)
+            raise UnresolvedQueueExpNamesError(missing_rev)
 
         for queue_entry in to_kill:
             self.proc.kill(queue_entry.stash_rev)
@@ -271,15 +273,57 @@ class LocalCeleryQueue(BaseStashQueue):
         rev: str,
         encoding: Optional[str] = None,
     ):
-        from dvc.ui import ui
-
         queue_entry: Optional[QueueEntry] = self.match_queue_entry_by_name(
             {rev}, self.iter_active()
         ).get(rev)
         if queue_entry is None:
-            raise UnresolvedExpNamesError([rev])
-        for line in self.proc.follow(queue_entry.stash_rev, encoding):
-            ui.write(line)
+            raise UnresolvedQueueExpNamesError([rev])
+        return self.follow(queue_entry, encoding=encoding)
+
+    def follow(
+        self,
+        entry: QueueEntry,
+        encoding: Optional[str] = None,
+    ):
+        for line in self.proc.follow(entry.stash_rev, encoding):
+            ui.write(line, end="")
+
+    def logs(
+        self,
+        rev: str,
+        encoding: Optional[str] = None,
+        follow: bool = False,
+    ):
+        queue_entry: Optional[QueueEntry] = self.match_queue_entry_by_name(
+            {rev}, self.iter_active(), self.iter_done()
+        ).get(rev)
+        if queue_entry is None:
+            if rev in self.match_queue_entry_by_name(
+                {rev}, self.iter_queued()
+            ):
+                raise DvcException(
+                    f"Experiment '{rev}' is in queue but has not been started"
+                )
+            raise UnresolvedQueueExpNamesError([rev])
+        if follow:
+            ui.write(
+                f"Following logs for experiment '{rev}'. Use Ctrl+C to stop "
+                "following logs (experiment execution will continue).\n"
+            )
+            try:
+                self.follow(queue_entry)
+            except KeyboardInterrupt:
+                pass
+            return
+        try:
+            proc_info = self.proc[queue_entry.stash_rev]
+        except KeyError:
+            raise DvcException(f"No output logs found for experiment '{rev}'")
+        with open(
+            proc_info.stdout,
+            encoding=encoding or locale.getpreferredencoding(),
+        ) as fobj:
+            ui.write(fobj.read())
 
 
 class WorkspaceQueue(BaseStashQueue):
@@ -415,3 +459,11 @@ class WorkspaceQueue(BaseStashQueue):
         encoding: Optional[str] = None,
     ):
         raise NotImplementedError
+
+    def logs(
+        self,
+        rev: str,
+        encoding: Optional[str] = None,
+        follow: bool = False,
+    ):
+        raise NotImplementedError

tests/func/experiments/conftest.py

@@ -97,9 +97,9 @@ def failed_exp_stage(tmp_dir, scm, dvc):
     tmp_dir.gen("params.yaml", "foo: 1")
     stage = dvc.stage.add(
         cmd="python -c 'import sys; sys.exit(1)'",
-        metrics_no_cache=["metrics.yaml"],
+        metrics_no_cache=["failed-metrics.yaml"],
         params=["foo"],
-        name="copy-file",
+        name="failed-copy-file",
         deps=["copy.py"],
     )
     scm.add(
@@ -108,7 +108,7 @@ def failed_exp_stage(tmp_dir, scm, dvc):
             "dvc.lock",
             "copy.py",
             "params.yaml",
-            "metrics.yaml",
+            "failed-metrics.yaml",
             ".gitignore",
         ]
     )

tests/func/experiments/test_experiments.py

@@ -122,6 +122,29 @@ def test_failed_exp_celery(
     )
 
 
+@pytest.mark.parametrize("follow", [True, False])
+def test_celery_logs(
+    tmp_dir,
+    scm,
+    dvc,
+    failed_exp_stage,
+    test_queue,
+    follow,
+    mocker,
+    capsys,
+):
+    dvc.experiments.run(failed_exp_stage.addressing, queue=True)
+    dvc.experiments.run(run_all=True)
+
+    queue = dvc.experiments.celery_queue
+    done_result = first(queue.iter_done())
+    name = done_result.entry.stash_rev
+    captured = capsys.readouterr()
+    queue.logs(name, follow=follow)
+    captured = capsys.readouterr()
+    assert "failed to reproduce 'failed-copy-file'" in captured.out
+
+
 @pytest.mark.parametrize(
     "changes, expected",
     [