diff --git a/dvc/repo/experiments/__init__.py b/dvc/repo/experiments/__init__.py index e4f3eec23..39c4aaf4b 100644 --- a/dvc/repo/experiments/__init__.py +++ b/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: diff --git a/dvc/repo/experiments/exceptions.py b/dvc/repo/experiments/exceptions.py index 20c0b64fb..616b692b3 100644 --- a/dvc/repo/experiments/exceptions.py +++ b/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 diff --git a/dvc/repo/experiments/queue/base.py b/dvc/repo/experiments/queue/base.py index a705031e3..69d16e337 100644 --- a/dvc/repo/experiments/queue/base.py +++ b/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]] = {} diff --git a/dvc/repo/experiments/queue/local.py b/dvc/repo/experiments/queue/local.py index 423a55023..a0a5ae789 100644 --- a/dvc/repo/experiments/queue/local.py +++ b/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 diff --git a/tests/func/experiments/conftest.py b/tests/func/experiments/conftest.py index d69cb0081..cb295c816 100644 --- a/tests/func/experiments/conftest.py +++ b/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", ] ) diff --git a/tests/func/experiments/test_experiments.py b/tests/func/experiments/test_experiments.py index 6d5574c77..2cde0feaf 100644 --- a/tests/func/experiments/test_experiments.py +++ b/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", [