From 12fc8c4e7d6a8a581cdf95ac5ff4e0459ab2c391 Mon Sep 17 00:00:00 2001
From: karajan1001 <mishanyo1001@gmail.com>
Date: Mon, 16 May 2022 17:49:53 +0800
Subject: [PATCH] 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   |  4 ++-
 dvc/commands/queue/start.py      | 54 ++++++++++++++++++++++++++++++++
 dvc/commands/queue/stop.py       | 46 +++++++++++++++++++++++++++
 tests/unit/command/test_queue.py | 34 ++++++++++++++++++++
 4 files changed, 137 insertions(+), 1 deletion(-)
 create mode 100644 dvc/commands/queue/start.py
 create mode 100644 dvc/commands/queue/stop.py

diff --git a/dvc/commands/queue/__init__.py b/dvc/commands/queue/__init__.py
index 4105aa8fb..5c3998c89 100644
--- a/dvc/commands/queue/__init__.py
+++ b/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,
 ]
 
 
diff --git a/dvc/commands/queue/start.py b/dvc/commands/queue/start.py
new file mode 100644
index 000000000..a74864b9f
--- /dev/null
+++ b/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)
diff --git a/dvc/commands/queue/stop.py b/dvc/commands/queue/stop.py
new file mode 100644
index 000000000..caf40a2ea
--- /dev/null
+++ b/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)
diff --git a/tests/unit/command/test_queue.py b/tests/unit/command/test_queue.py
index c47516251..2180015c6 100644
--- a/tests/unit/command/test_queue.py
+++ b/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)