exp queue: remove queued experiments(#7591)

fix: #7591
1. Add a new command `dvc queue`
2. Add two sub-command `dvc queue remove` and `dvc queue kill`
3. Add a unit test to test them

Co-authored-by: Peter Rowlands (변기호) <[email protected]>

dvc/cli/parser.py

@@ -33,6 +33,7 @@ from dvc.commands import (
     move,
     params,
     plots,
+    queue,
     remote,
     remove,
     repro,
@@ -50,6 +51,7 @@ logger = logging.getLogger(__name__)
 
 COMMANDS = [
     init,
+    queue,
     get,
     get_url,
     destroy,

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)

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="<experiment>",
+    )
+    queue_kill_parser.set_defaults(func=CmdQueueKill)

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="<experiment>",
+    )
+    queue_remove_parser.set_defaults(func=CmdQueueRemove)

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"])