diff --git a/src/executorlib/standalone/queue.py b/src/executorlib/standalone/queue.py index 8dd789a0..09e4d1cb 100644 --- a/src/executorlib/standalone/queue.py +++ b/src/executorlib/standalone/queue.py @@ -1,4 +1,21 @@ import queue +from typing import Any + + +def put_front(que: queue.Queue, item: Any): + """ + Insert an item at the front of the queue, ahead of any items already waiting, and wake a consumer blocked in + get(). Equivalent to queue.Queue.put() except the item is placed at the front rather than the back; mutating + que.queue directly instead would silently skip the notification, leaving a waiting consumer asleep forever. + + Args: + que (queue.Queue): Queue with task objects which should be executed + item (Any): item to place at the front of the queue + """ + with que.not_full: + que.queue.insert(0, item) + que.unfinished_tasks += 1 + que.not_empty.notify() def cancel_items_in_queue(que: queue.Queue): diff --git a/src/executorlib/task_scheduler/interactive/blockallocation.py b/src/executorlib/task_scheduler/interactive/blockallocation.py index 227029a3..e9ee1fd2 100644 --- a/src/executorlib/task_scheduler/interactive/blockallocation.py +++ b/src/executorlib/task_scheduler/interactive/blockallocation.py @@ -15,7 +15,7 @@ interface_bootup, ) from executorlib.standalone.interactive.spawner import BaseSpawner, MpiExecSpawner -from executorlib.standalone.queue import cancel_items_in_queue +from executorlib.standalone.queue import cancel_items_in_queue, put_front from executorlib.task_scheduler.base import TaskSchedulerBase, validate_resource_dict from executorlib.task_scheduler.interactive.shared import ( execute_task_dict, @@ -122,7 +122,7 @@ def max_workers(self, max_workers: int): ): if self._max_workers > max_workers: for _ in range(self._max_workers - max_workers): - self._future_queue.queue.insert(0, {"shutdown": True, "wait": True}) + put_front(self._future_queue, {"shutdown": True, "wait": True}) while len(self._process) > max_workers: self._process = [ process for process in self._process if process.is_alive() diff --git a/tests/unit/standalone/test_queue.py b/tests/unit/standalone/test_queue.py index 7e6f40b3..6c3614f8 100644 --- a/tests/unit/standalone/test_queue.py +++ b/tests/unit/standalone/test_queue.py @@ -1,8 +1,10 @@ from concurrent.futures import Future, CancelledError from queue import Queue +from threading import Thread +import time import unittest -from executorlib.standalone.queue import cancel_items_in_queue +from executorlib.standalone.queue import cancel_items_in_queue, put_front class TestQueue(unittest.TestCase): @@ -21,3 +23,34 @@ def test_cancel_items_in_queue(self): with self.assertRaises(CancelledError): self.assertTrue(fs2.result()) q.join() + + def test_put_front_orders_ahead_of_existing_items(self): + q = Queue() + q.put("back") + put_front(q, "front") + self.assertEqual(q.get(), "front") + self.assertEqual(q.get(), "back") + q.task_done() + q.task_done() + q.join() + + def test_put_front_wakes_consumer_blocked_on_empty_queue(self): + # Mutating q.queue directly (q.queue.insert(0, item)) skips the not_empty + # notification, so a thread already parked in q.get() never wakes up and the + # item sits in the queue forever. put_front() must not have that problem. + q = Queue() + received = [] + + def consume(): + received.append(q.get()) + + consumer = Thread(target=consume) + consumer.start() + try: + time.sleep(0.2) # give the consumer time to block inside q.get() + put_front(q, "woken") + consumer.join(timeout=5) + self.assertFalse(consumer.is_alive(), "consumer stayed blocked in get()") + self.assertEqual(received, ["woken"]) + finally: + consumer.join(timeout=5) diff --git a/tests/unit/task_scheduler/interactive/test_blockallocation.py b/tests/unit/task_scheduler/interactive/test_blockallocation.py index 4373edd2..71a8397d 100644 --- a/tests/unit/task_scheduler/interactive/test_blockallocation.py +++ b/tests/unit/task_scheduler/interactive/test_blockallocation.py @@ -1,7 +1,8 @@ import queue +import time import unittest from concurrent.futures import Future -from threading import Event, Lock +from threading import Event, Lock, Thread from unittest.mock import patch from executorlib.standalone.interactive.communication import ExecutorlibSocketError @@ -50,6 +51,45 @@ def start(self): self.assertTrue(worker.started) self.assertEqual(scheduler._alive_workers[0], 2) + def test_shrink_wakes_worker_blocked_on_empty_queue(self): + """ + Regression test: shrinking max_workers while a worker is already blocked in + future_queue.get() used to hang forever, because the shutdown sentinel was + spliced directly into the queue's internal deque without waking the blocked + consumer (see executorlib.standalone.queue.put_front). + """ + future_queue = queue.Queue() + received = [] + + def worker_loop(): + received.append(future_queue.get()) + future_queue.task_done() + + worker = Thread(target=worker_loop, daemon=True) + worker.start() + time.sleep(0.2) # let the worker actually block inside future_queue.get() + + scheduler = object.__new__(BlockAllocationTaskScheduler) + scheduler._future_queue = future_queue + scheduler._process = [worker] + scheduler._max_workers = 1 + + shrink_done = Event() + + def shrink(): + scheduler.max_workers = 0 + shrink_done.set() + + shrink_thread = Thread(target=shrink, daemon=True) + shrink_thread.start() + shrink_thread.join(timeout=5) + + self.assertTrue(shrink_done.is_set(), "max_workers setter hung while shrinking") + worker.join(timeout=5) + self.assertFalse(worker.is_alive()) + self.assertEqual(received, [{"shutdown": True, "wait": True}]) + self.assertEqual(scheduler._process, []) + class TestDrainDeadWorker(unittest.TestCase): def test_fail_tasks_when_no_workers_remain(self):