diff --git a/dvc/cli/parser.py b/dvc/cli/parser.py index 012b6cb737..dafa942f2b 100644 --- a/dvc/cli/parser.py +++ b/dvc/cli/parser.py @@ -33,6 +33,7 @@ move, params, plots, + queue, remote, remove, repro, @@ -50,6 +51,7 @@ COMMANDS = [ init, + queue, get, get_url, destroy, diff --git a/dvc/commands/queue/__init__.py b/dvc/commands/queue/__init__.py new file mode 100644 index 0000000000..4105aa8fb5 --- /dev/null +++ b/dvc/commands/queue/__init__.py @@ -0,0 +1,30 @@ +import argparse + +from dvc.cli.utils import append_doc_link, fix_subparsers +from dvc.commands.queue import kill, remove + +SUB_COMMANDS = [ + remove, + kill, +] + + +def add_parser(subparsers, parent_parser): + QUEUE_HELP = "Commands to manage dvc task queue." + + queue_parser = subparsers.add_parser( + "queue", + parents=[parent_parser], + description=append_doc_link(QUEUE_HELP, "queue"), + formatter_class=argparse.RawDescriptionHelpFormatter, + help=QUEUE_HELP, + ) + + queue_subparsers = queue_parser.add_subparsers( + dest="cmd", + help="Use `dvc queue CMD --help` to display " "command-specific help.", + ) + + fix_subparsers(queue_subparsers) + for cmd in SUB_COMMANDS: + cmd.add_parser(queue_subparsers, parent_parser) diff --git a/dvc/commands/queue/kill.py b/dvc/commands/queue/kill.py new file mode 100644 index 0000000000..703b984575 --- /dev/null +++ b/dvc/commands/queue/kill.py @@ -0,0 +1,34 @@ +import argparse +import logging + +from dvc.cli.command import CmdBase +from dvc.cli.utils import append_doc_link + +logger = logging.getLogger(__name__) + + +class CmdQueueKill(CmdBase): + """Kill exp task in queue.""" + + def run(self): + self.repo.experiments.celery_queue.kill(revs=self.args.experiment) + + return 0 + + +def add_parser(queue_subparsers, parent_parser): + QUEUE_KILL_HELP = "Kill experiments in queue" + queue_kill_parser = queue_subparsers.add_parser( + "kill", + parents=[parent_parser], + description=append_doc_link(QUEUE_KILL_HELP, "queue/kill"), + help=QUEUE_KILL_HELP, + formatter_class=argparse.RawDescriptionHelpFormatter, + ) + queue_kill_parser.add_argument( + "experiment", + nargs="*", + help="Experiments in queue to kill.", + metavar="", + ) + queue_kill_parser.set_defaults(func=CmdQueueKill) diff --git a/dvc/commands/queue/remove.py b/dvc/commands/queue/remove.py new file mode 100644 index 0000000000..6d9572e0c2 --- /dev/null +++ b/dvc/commands/queue/remove.py @@ -0,0 +1,47 @@ +import argparse +import logging + +from dvc.cli.command import CmdBase +from dvc.cli.utils import append_doc_link +from dvc.ui import ui + +logger = logging.getLogger(__name__) + + +class CmdQueueRemove(CmdBase): + """Remove exp in queue.""" + + def run(self): + if self.args.all: + removed_list = self.repo.experiments.celery_queue.clear() + else: + removed_list = self.repo.experiments.celery_queue.remove( + revs=self.args.experiment + ) + + removed = ", ".join(removed_list) + ui.write(f"Removed experiments in queue: {removed}") + + return 0 + + +def add_parser(queue_subparsers, parent_parser): + + QUEUE_REMOVE_HELP = "Remove experiments in queue" + queue_remove_parser = queue_subparsers.add_parser( + "remove", + parents=[parent_parser], + description=append_doc_link(QUEUE_REMOVE_HELP, "queue/remove"), + help=QUEUE_REMOVE_HELP, + formatter_class=argparse.RawDescriptionHelpFormatter, + ) + queue_remove_parser.add_argument( + "--all", action="store_true", help="Remove all experiments in queue." + ) + queue_remove_parser.add_argument( + "experiment", + nargs="*", + help="Experiments in queue to remove.", + metavar="", + ) + queue_remove_parser.set_defaults(func=CmdQueueRemove) diff --git a/tests/unit/command/test_queue.py b/tests/unit/command/test_queue.py new file mode 100644 index 0000000000..c475162511 --- /dev/null +++ b/tests/unit/command/test_queue.py @@ -0,0 +1,63 @@ +from dvc.cli import parse_args +from dvc.commands.queue.kill import CmdQueueKill +from dvc.commands.queue.remove import CmdQueueRemove + + +def test_experiments_remove(dvc, scm, mocker): + cli_args = parse_args( + [ + "queue", + "remove", + "--all", + ] + ) + assert cli_args.func == CmdQueueRemove + + cmd = cli_args.func(cli_args) + m = mocker.patch( + "dvc.repo.experiments.queue.local.LocalCeleryQueue.clear", + return_value={}, + ) + + assert cmd.run() == 0 + m.assert_called_once_with() + + cli_args = parse_args( + [ + "queue", + "remove", + "exp1", + "exp2", + ] + ) + assert cli_args.func == CmdQueueRemove + + cmd = cli_args.func(cli_args) + m = mocker.patch( + "dvc.repo.experiments.queue.local.LocalCeleryQueue.remove", + return_value={}, + ) + + assert cmd.run() == 0 + m.assert_called_once_with(revs=["exp1", "exp2"]) + + +def test_experiments_kill(dvc, scm, mocker): + cli_args = parse_args( + [ + "queue", + "kill", + "exp1", + "exp2", + ] + ) + assert cli_args.func == CmdQueueKill + + cmd = cli_args.func(cli_args) + m = mocker.patch( + "dvc.repo.experiments.queue.local.LocalCeleryQueue.kill", + return_value={}, + ) + + assert cmd.run() == 0 + m.assert_called_once_with(revs=["exp1", "exp2"])