diff --git a/dvc/commands/queue/__init__.py b/dvc/commands/queue/__init__.py index bc1c0d16f..efd60a5ed 100644 --- a/dvc/commands/queue/__init__.py +++ b/dvc/commands/queue/__init__.py @@ -1,7 +1,7 @@ import argparse from dvc.cli.utils import append_doc_link, fix_subparsers -from dvc.commands.queue import attach, kill, remove, start, status, stop +from dvc.commands.queue import attach, kill, logs, remove, start, status, stop SUB_COMMANDS = [ remove, @@ -10,6 +10,7 @@ SUB_COMMANDS = [ stop, attach, status, + logs, ] diff --git a/dvc/commands/queue/logs.py b/dvc/commands/queue/logs.py new file mode 100644 index 000000000..40d401d70 --- /dev/null +++ b/dvc/commands/queue/logs.py @@ -0,0 +1,53 @@ +import argparse +import logging + +from dvc.cli.command import CmdBase +from dvc.cli.utils import append_doc_link + +logger = logging.getLogger(__name__) + + +class CmdQueueLogs(CmdBase): + """Show output logs for a queued experiment.""" + + def run(self): + self.repo.experiments.celery_queue.logs( + rev=self.args.experiment, + encoding=self.args.encoding, + follow=self.args.follow, + ) + + return 0 + + +def add_parser(queue_subparsers, parent_parser): + QUEUE_LOGS_HELP = "Show output logs for a queued experiment." + queue_logs_parser = queue_subparsers.add_parser( + "logs", + parents=[parent_parser], + description=append_doc_link(QUEUE_LOGS_HELP, "queue/logs"), + help=QUEUE_LOGS_HELP, + formatter_class=argparse.RawDescriptionHelpFormatter, + ) + queue_logs_parser.add_argument( + "-e", + "--encoding", + help=( + "Text encoding for log output. Defaults to system locale encoding." + ), + ) + queue_logs_parser.add_argument( + "-f", + "--follow", + help=( + "Attach to experiment and follow additional live output. Only " + "applicable if the experiment is still running." + ), + action="store_true", + ) + queue_logs_parser.add_argument( + "experiment", + help="Experiment to show.", + metavar="", + ) + queue_logs_parser.set_defaults(func=CmdQueueLogs) diff --git a/tests/unit/command/test_queue.py b/tests/unit/command/test_queue.py index bc36d75e1..e46a49a87 100644 --- a/tests/unit/command/test_queue.py +++ b/tests/unit/command/test_queue.py @@ -1,6 +1,7 @@ from dvc.cli import parse_args from dvc.commands.queue.attach import CmdQueueAttach from dvc.commands.queue.kill import CmdQueueKill +from dvc.commands.queue.logs import CmdQueueLogs from dvc.commands.queue.remove import CmdQueueRemove from dvc.commands.queue.start import CmdQueueStart from dvc.commands.queue.status import CmdQueueStatus @@ -137,3 +138,17 @@ def test_experiments_attach(dvc, scm, mocker): assert cmd.run() == 0 m.assert_called_once_with(rev="exp1", encoding="utf8") + + +def test_queue_logs(dvc, scm, mocker): + cli_args = parse_args(["queue", "logs", "exp1", "-e", "utf8", "-f"]) + assert cli_args.func == CmdQueueLogs + + cmd = cli_args.func(cli_args) + m = mocker.patch( + "dvc.repo.experiments.queue.local.LocalCeleryQueue.logs", + return_value={}, + ) + + assert cmd.run() == 0 + m.assert_called_once_with(rev="exp1", encoding="utf8", follow=True)