exp queue: Add two command queue start and stop (#7589)

fix: #7589

1. Add two sub-command `dvc queue start` and `dvc queue stop`
2. Add a unit test to test them

dvc/commands/queue/__init__.py

@@ -1,11 +1,13 @@
 import argparse
 
 from dvc.cli.utils import append_doc_link, fix_subparsers
-from dvc.commands.queue import kill, remove
+from dvc.commands.queue import kill, remove, start, stop
 
 SUB_COMMANDS = [
     remove,
     kill,
+    start,
+    stop,
 ]
 
 

dvc/commands/queue/start.py

@@ -0,0 +1,54 @@
+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 CmdQueueStart(CmdBase):
+    """Start exp queue workers."""
+
+    def run(self):
+        for _ in range(self.args.jobs):
+            self.repo.experiments.celery_queue.spawn_worker()
+
+        suffix = "s" if self.args.jobs > 1 else ""
+        ui.write(
+            f"Start {self.args.jobs} queue worker{suffix} to process "
+            "the queue tasks"
+        )
+
+        return 0
+
+
+def job_type(job):
+    try:
+        job = int(job)
+        if job > 0:
+            return job
+    except ValueError:
+        pass
+    raise argparse.ArgumentTypeError("Worker number must be a natural number")
+
+
+def add_parser(queue_subparsers, parent_parser):
+
+    QUEUE_START_HELP = "Start experiments queue workers"
+    queue_start_parser = queue_subparsers.add_parser(
+        "start",
+        parents=[parent_parser],
+        description=append_doc_link(QUEUE_START_HELP, "queue/start"),
+        help=QUEUE_START_HELP,
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+    )
+    queue_start_parser.add_argument(
+        "-j",
+        "--jobs",
+        type=job_type,
+        default=1,
+        help="Number of queue workers to be started.",
+    )
+    queue_start_parser.set_defaults(func=CmdQueueStart)

dvc/commands/queue/stop.py

@@ -0,0 +1,46 @@
+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 CmdQueueStop(CmdBase):
+    """Stop exp queue workers."""
+
+    def run(self):
+        self.repo.experiments.celery_queue.shutdown(kill=self.args.kill)
+
+        if self.args.kill:
+            ui.write(
+                "All of the Queue tasked had already been killed, "
+                "Queue workers are stopping."
+            )
+        else:
+            ui.write(
+                "Queue workers will be stopped after current tasks finished."
+            )
+
+        return 0
+
+
+def add_parser(queue_subparsers, parent_parser):
+
+    QUEUE_STOP_HELP = "Stop experiment queue workers"
+    queue_stop_parser = queue_subparsers.add_parser(
+        "stop",
+        parents=[parent_parser],
+        description=append_doc_link(QUEUE_STOP_HELP, "queue/stop"),
+        help=QUEUE_STOP_HELP,
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+    )
+    queue_stop_parser.add_argument(
+        "--kill",
+        action="store_true",
+        help="Kill the current running experiments in queue "
+        "before shutting down.",
+    )
+    queue_stop_parser.set_defaults(func=CmdQueueStop)

tests/unit/command/test_queue.py

@@ -1,6 +1,8 @@
 from dvc.cli import parse_args
 from dvc.commands.queue.kill import CmdQueueKill
 from dvc.commands.queue.remove import CmdQueueRemove
+from dvc.commands.queue.start import CmdQueueStart
+from dvc.commands.queue.stop import CmdQueueStop
 
 
 def test_experiments_remove(dvc, scm, mocker):
@@ -61,3 +63,35 @@ def test_experiments_kill(dvc, scm, mocker):
 
     assert cmd.run() == 0
     m.assert_called_once_with(revs=["exp1", "exp2"])
+
+
+def test_experiments_start(dvc, scm, mocker):
+    cli_args = parse_args(["queue", "start", "-j", "3"])
+    assert cli_args.func == CmdQueueStart
+
+    cmd = cli_args.func(cli_args)
+    m = mocker.patch(
+        "dvc.repo.experiments.queue.local.LocalCeleryQueue.spawn_worker",
+    )
+
+    assert cmd.run() == 0
+    assert m.call_count == 3
+
+
+def test_experiments_stop(dvc, scm, mocker):
+    cli_args = parse_args(
+        [
+            "queue",
+            "stop",
+            "--kill",
+        ]
+    )
+    assert cli_args.func == CmdQueueStop
+
+    cmd = cli_args.func(cli_args)
+    m = mocker.patch(
+        "dvc.repo.experiments.queue.local.LocalCeleryQueue.shutdown",
+    )
+
+    assert cmd.run() == 0
+    m.assert_called_once_with(kill=True)