You cannot select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
rq/tests/test_worker_pool.py

142 lines
5.7 KiB
Python

Worker pool (#1874) * First stab at implementating worker pool * Use process.is_alive() to check whether a process is still live * Handle shutdown signal * Check worker loop done * First working version of `WorkerPool`. * Added test for check_workers() * Added test for pool.start() * Better shutdown process * Comment out test_start() to see if it fixes CI * Make tests pass * Make CI pass * Comment out some tests * Comment out more tests * Re-enable a test * Re-enable another test * Uncomment check_workers test * Added run_worker test * Minor modification to dead worker detection * More test cases * Better process name for workers * Added back pool.stop_workers() when signal is received * Cleaned up cli.py * WIP on worker-pool command * Fix test * Test that worker pool ignores consecutive shutdown signals * Added test for worker-pool CLI command. * Added timeout to CI jobs * Fix worker pool test * Comment out test_scheduler.py * Fixed worker-pool in burst mode * Increase test coverage * Exclude tests directory from coverage.py * Improve test coverage * Renamed `Pool(num_workers=2) to `Pool(size=2)` * Revert "Renamed `Pool(num_workers=2) to `Pool(size=2)`" This reverts commit a1306f89ad0d8686c6bde447bff75e2f71f0733b. * Renamed Pool to WorkerPool * Added a new TestCase that doesn't use LocalStack * Added job_class, worker_class and serializer arguments to WorkerPool * Use parse_connection() in WorkerPool.__init__ * Added CLI arguments for worker-pool * Minor WorkerPool and test fixes * Fixed failing CLI test * Document WorkerPool
2 years ago
import os
import signal
from multiprocessing import Process
from time import sleep
from rq.job import JobStatus
from tests import TestCase
from tests.fixtures import CustomJob, _send_shutdown_command, long_running_job, say_hello
from rq.connections import parse_connection
Worker pool (#1874) * First stab at implementating worker pool * Use process.is_alive() to check whether a process is still live * Handle shutdown signal * Check worker loop done * First working version of `WorkerPool`. * Added test for check_workers() * Added test for pool.start() * Better shutdown process * Comment out test_start() to see if it fixes CI * Make tests pass * Make CI pass * Comment out some tests * Comment out more tests * Re-enable a test * Re-enable another test * Uncomment check_workers test * Added run_worker test * Minor modification to dead worker detection * More test cases * Better process name for workers * Added back pool.stop_workers() when signal is received * Cleaned up cli.py * WIP on worker-pool command * Fix test * Test that worker pool ignores consecutive shutdown signals * Added test for worker-pool CLI command. * Added timeout to CI jobs * Fix worker pool test * Comment out test_scheduler.py * Fixed worker-pool in burst mode * Increase test coverage * Exclude tests directory from coverage.py * Improve test coverage * Renamed `Pool(num_workers=2) to `Pool(size=2)` * Revert "Renamed `Pool(num_workers=2) to `Pool(size=2)`" This reverts commit a1306f89ad0d8686c6bde447bff75e2f71f0733b. * Renamed Pool to WorkerPool * Added a new TestCase that doesn't use LocalStack * Added job_class, worker_class and serializer arguments to WorkerPool * Use parse_connection() in WorkerPool.__init__ * Added CLI arguments for worker-pool * Minor WorkerPool and test fixes * Fixed failing CLI test * Document WorkerPool
2 years ago
from rq.queue import Queue
from rq.serializers import JSONSerializer
from rq.worker import SimpleWorker
from rq.worker_pool import run_worker, WorkerPool
def wait_and_send_shutdown_signal(pid, time_to_wait=0.0):
sleep(time_to_wait)
os.kill(pid, signal.SIGTERM)
class TestWorkerPool(TestCase):
def test_queues(self):
"""Test queue parsing"""
pool = WorkerPool(['default', 'foo'], connection=self.connection)
self.assertEqual(
set(pool.queues), {Queue('default', connection=self.connection), Queue('foo', connection=self.connection)}
)
# def test_spawn_workers(self):
# """Test spawning workers"""
# pool = WorkerPool(['default', 'foo'], connection=self.connection, num_workers=2)
# pool.start_workers(burst=False)
# self.assertEqual(len(pool.worker_dict.keys()), 2)
# pool.stop_workers()
def test_check_workers(self):
"""Test check_workers()"""
pool = WorkerPool(['default'], connection=self.connection, num_workers=2)
pool.start_workers(burst=False)
# There should be two workers
pool.check_workers()
self.assertEqual(len(pool.worker_dict.keys()), 2)
worker_data = list(pool.worker_dict.values())[0]
_send_shutdown_command(worker_data.name, self.connection.connection_pool.connection_kwargs.copy(), delay=0)
# 1 worker should be dead since we sent a shutdown command
sleep(0.2)
pool.check_workers(respawn=False)
self.assertEqual(len(pool.worker_dict.keys()), 1)
# If we call `check_workers` with `respawn=True`, the worker should be respawned
pool.check_workers(respawn=True)
self.assertEqual(len(pool.worker_dict.keys()), 2)
pool.stop_workers()
def test_reap_workers(self):
"""Dead workers are removed from worker_dict"""
pool = WorkerPool(['default'], connection=self.connection, num_workers=2)
pool.start_workers(burst=False)
# There should be two workers
pool.reap_workers()
self.assertEqual(len(pool.worker_dict.keys()), 2)
worker_data = list(pool.worker_dict.values())[0]
_send_shutdown_command(worker_data.name, self.connection.connection_pool.connection_kwargs.copy(), delay=0)
# 1 worker should be dead since we sent a shutdown command
sleep(0.2)
pool.reap_workers()
self.assertEqual(len(pool.worker_dict.keys()), 1)
pool.stop_workers()
def test_start(self):
"""Test start()"""
pool = WorkerPool(['default'], connection=self.connection, num_workers=2)
p = Process(target=wait_and_send_shutdown_signal, args=(os.getpid(), 0.5))
p.start()
pool.start()
self.assertEqual(pool.status, pool.Status.STOPPED)
self.assertTrue(pool.all_workers_have_stopped())
# We need this line so the test doesn't hang
pool.stop_workers()
def test_pool_ignores_consecutive_shutdown_signals(self):
"""If two shutdown signals are sent within one second, only the first one is processed"""
# Send two shutdown signals within one second while the worker is
# working on a long running job. The job should still complete (not killed)
pool = WorkerPool(['foo'], connection=self.connection, num_workers=2)
process_1 = Process(target=wait_and_send_shutdown_signal, args=(os.getpid(), 0.5))
process_1.start()
process_2 = Process(target=wait_and_send_shutdown_signal, args=(os.getpid(), 0.5))
process_2.start()
queue = Queue('foo', connection=self.connection)
job = queue.enqueue(long_running_job, 1)
pool.start(burst=True)
self.assertEqual(job.get_status(refresh=True), JobStatus.FINISHED)
# We need this line so the test doesn't hang
pool.stop_workers()
def test_run_worker(self):
"""Ensure run_worker() properly spawns a Worker"""
queue = Queue('foo', connection=self.connection)
queue.enqueue(say_hello)
connection_class, pool_class, pool_kwargs = parse_connection(self.connection)
Worker pool (#1874) * First stab at implementating worker pool * Use process.is_alive() to check whether a process is still live * Handle shutdown signal * Check worker loop done * First working version of `WorkerPool`. * Added test for check_workers() * Added test for pool.start() * Better shutdown process * Comment out test_start() to see if it fixes CI * Make tests pass * Make CI pass * Comment out some tests * Comment out more tests * Re-enable a test * Re-enable another test * Uncomment check_workers test * Added run_worker test * Minor modification to dead worker detection * More test cases * Better process name for workers * Added back pool.stop_workers() when signal is received * Cleaned up cli.py * WIP on worker-pool command * Fix test * Test that worker pool ignores consecutive shutdown signals * Added test for worker-pool CLI command. * Added timeout to CI jobs * Fix worker pool test * Comment out test_scheduler.py * Fixed worker-pool in burst mode * Increase test coverage * Exclude tests directory from coverage.py * Improve test coverage * Renamed `Pool(num_workers=2) to `Pool(size=2)` * Revert "Renamed `Pool(num_workers=2) to `Pool(size=2)`" This reverts commit a1306f89ad0d8686c6bde447bff75e2f71f0733b. * Renamed Pool to WorkerPool * Added a new TestCase that doesn't use LocalStack * Added job_class, worker_class and serializer arguments to WorkerPool * Use parse_connection() in WorkerPool.__init__ * Added CLI arguments for worker-pool * Minor WorkerPool and test fixes * Fixed failing CLI test * Document WorkerPool
2 years ago
run_worker(
'test-worker', ['foo'], connection_class, pool_class, pool_kwargs
Worker pool (#1874) * First stab at implementating worker pool * Use process.is_alive() to check whether a process is still live * Handle shutdown signal * Check worker loop done * First working version of `WorkerPool`. * Added test for check_workers() * Added test for pool.start() * Better shutdown process * Comment out test_start() to see if it fixes CI * Make tests pass * Make CI pass * Comment out some tests * Comment out more tests * Re-enable a test * Re-enable another test * Uncomment check_workers test * Added run_worker test * Minor modification to dead worker detection * More test cases * Better process name for workers * Added back pool.stop_workers() when signal is received * Cleaned up cli.py * WIP on worker-pool command * Fix test * Test that worker pool ignores consecutive shutdown signals * Added test for worker-pool CLI command. * Added timeout to CI jobs * Fix worker pool test * Comment out test_scheduler.py * Fixed worker-pool in burst mode * Increase test coverage * Exclude tests directory from coverage.py * Improve test coverage * Renamed `Pool(num_workers=2) to `Pool(size=2)` * Revert "Renamed `Pool(num_workers=2) to `Pool(size=2)`" This reverts commit a1306f89ad0d8686c6bde447bff75e2f71f0733b. * Renamed Pool to WorkerPool * Added a new TestCase that doesn't use LocalStack * Added job_class, worker_class and serializer arguments to WorkerPool * Use parse_connection() in WorkerPool.__init__ * Added CLI arguments for worker-pool * Minor WorkerPool and test fixes * Fixed failing CLI test * Document WorkerPool
2 years ago
)
# Worker should have processed the job
self.assertEqual(len(queue), 0)
def test_worker_pool_arguments(self):
"""Ensure arguments are properly used to create the right workers"""
queue = Queue('foo', connection=self.connection)
job = queue.enqueue(say_hello)
pool = WorkerPool([queue], connection=self.connection, num_workers=2, worker_class=SimpleWorker)
pool.start(burst=True)
# Worker should have processed the job
self.assertEqual(job.get_status(refresh=True), JobStatus.FINISHED)
queue = Queue('json', connection=self.connection, serializer=JSONSerializer)
job = queue.enqueue(say_hello, 'Hello')
pool = WorkerPool(
[queue], connection=self.connection, num_workers=2, worker_class=SimpleWorker, serializer=JSONSerializer
)
pool.start(burst=True)
# Worker should have processed the job
self.assertEqual(job.get_status(refresh=True), JobStatus.FINISHED)
pool = WorkerPool([queue], connection=self.connection, num_workers=2, job_class=CustomJob)
pool.start(burst=True)
# Worker should have processed the job
self.assertEqual(job.get_status(refresh=True), JobStatus.FINISHED)