Merge pull request #21482 from ankush/rq_worker_start

feat: RQ WorkerPool support

frappe/commands/scheduler.py

@@ -215,6 +215,27 @@ def start_worker(
 	)
 
 
+@click.command("worker-pool")
+@click.option(
+	"--queue",
+	type=str,
+	help="Queue to consume from. Multiple queues can be specified using comma-separated string. If not specified all queues are consumed.",
+)
+@click.option("--num-workers", type=int, default=2, help="Number of workers to spawn in pool.")
+@click.option("--quiet", is_flag=True, default=False, help="Hide Log Outputs")
+@click.option("--burst", is_flag=True, default=False, help="Run Worker in Burst mode.")
+def start_worker_pool(queue, quiet=False, num_workers=2, burst=False):
+	"""Start a backgrond worker"""
+	from frappe.utils.background_jobs import start_worker_pool
+
+	start_worker_pool(
+		queue=queue,
+		quiet=quiet,
+		burst=burst,
+		num_workers=num_workers,
+	)
+
+
 @click.command("ready-for-migration")
 @click.option("--site", help="site name")
 @pass_context
@@ -251,5 +272,6 @@ commands = [
 	show_pending_jobs,
 	start_scheduler,
 	start_worker,
+	start_worker_pool,
 	trigger_scheduler_event,
 ]

frappe/core/doctype/rq_job/test_rq_job.py

@@ -96,6 +96,17 @@ class TestRQJob(FrappeTestCase):
 		_, stderr = execute_in_shell("bench worker --queue short,default --burst", check_exit_code=True)
 		self.assertIn("quitting", cstr(stderr))
 
+	@timeout(20)
+	def test_multi_queue_burst_consumption_worker_pool(self):
+		for _ in range(3):
+			for q in ["default", "short"]:
+				frappe.enqueue(self.BG_JOB, sleep=1, queue=q)
+
+		_, stderr = execute_in_shell(
+			"bench worker-pool --queue short,default --burst --num-workers=4", check_exit_code=True
+		)
+		self.assertIn("quitting", cstr(stderr))
+
 	@timeout(20)
 	def test_job_id_dedup(self):
 		job_id = "test_dedup"

frappe/utils/background_jobs.py

@@ -4,16 +4,17 @@ import socket
 import time
 from collections import defaultdict
 from functools import lru_cache
-from typing import Any, Callable, Literal, NoReturn
+from typing import Any, Callable, NoReturn
 from uuid import uuid4
 
 import redis
 from redis.exceptions import BusyLoadingError, ConnectionError
-from rq import Connection, Queue, Worker
+from rq import Queue, Worker
 from rq.exceptions import NoSuchJobError
 from rq.job import Job, JobStatus
 from rq.logutils import setup_loghandlers
-from rq.worker import RandomWorker, RoundRobinWorker
+from rq.worker import DequeueStrategy
+from rq.worker_pool import WorkerPool
 from tenacity import retry, retry_if_exception_type, stop_after_attempt, wait_fixed
 
 import frappe
@@ -230,14 +231,14 @@ def start_worker(
 	rq_username: str | None = None,
 	rq_password: str | None = None,
 	burst: bool = False,
-	strategy: Literal["round_robin", "random"] | None = None,
+	strategy: DequeueStrategy | None = DequeueStrategy.DEFAULT,
 ) -> NoReturn | None:  # pragma: no cover
 	"""Wrapper to start rq worker. Connects to redis and monitors these queues."""
-	DEQUEUE_STRATEGIES = {"round_robin": RoundRobinWorker, "random": RandomWorker}
 
-	if frappe._tune_gc:
-		gc.collect()
-		gc.freeze()
+	if not strategy:
+		strategy = DequeueStrategy.DEFAULT
+
+	_freeze_gc()
 
 	with frappe.init_site():
 		# empty init is required to get redis_queue from common_site_config.json
@@ -251,19 +252,59 @@ def start_worker(
 	if os.environ.get("CI"):
 		setup_loghandlers("ERROR")
 
-	WorkerKlass = DEQUEUE_STRATEGIES.get(strategy, Worker)
+	logging_level = "INFO"
+	if quiet:
+		logging_level = "WARNING"
 
-	with Connection(redis_connection):
-		logging_level = "INFO"
-		if quiet:
-			logging_level = "WARNING"
-		worker = WorkerKlass(queues, name=get_worker_name(queue_name))
-		worker.work(
-			logging_level=logging_level,
-			burst=burst,
-			date_format="%Y-%m-%d %H:%M:%S",
-			log_format="%(asctime)s,%(msecs)03d %(message)s",
-		)
+	worker = Worker(queues, name=get_worker_name(queue_name), connection=redis_connection)
+	worker.work(
+		logging_level=logging_level,
+		burst=burst,
+		date_format="%Y-%m-%d %H:%M:%S",
+		log_format="%(asctime)s,%(msecs)03d %(message)s",
+		dequeue_strategy=strategy,
+	)
+
+
+def start_worker_pool(
+	queue: str | None = None,
+	num_workers: int = 1,
+	quiet: bool = False,
+	burst: bool = False,
+) -> NoReturn:
+	"""Start worker pool with specified number of workers.
+
+	WARNING: This feature is considered "EXPERIMENTAL".
+	"""
+
+	_freeze_gc()
+
+	with frappe.init_site():
+		redis_connection = get_redis_conn()
+
+		if queue:
+			queue = [q.strip() for q in queue.split(",")]
+		queues = get_queue_list(queue, build_queue_name=True)
+
+	if os.environ.get("CI"):
+		setup_loghandlers("ERROR")
+
+	logging_level = "INFO"
+	if quiet:
+		logging_level = "WARNING"
+
+	pool = WorkerPool(
+		queues=queues,
+		connection=redis_connection,
+		num_workers=num_workers,
+	)
+	pool.start(logging_level=logging_level, burst=burst)
+
+
+def _freeze_gc():
+	if frappe._tune_gc:
+		gc.collect()
+		gc.freeze()
 
 
 def get_worker_name(queue):