From c049d07588cccb1cc8aebe82e4849c0fe5032a11 Mon Sep 17 00:00:00 2001 From: Muhtasim-Munif-Fahim Date: Mon, 10 Aug 2026 01:31:58 +0600 Subject: [PATCH 1/2] fix(scheduler): preserve worker context when resizing --- .../interactive/blockallocation.py | 48 ++++++++++-------- .../interactive/test_blockallocation.py | 50 +++++++++++++++++-- 2 files changed, 74 insertions(+), 24 deletions(-) diff --git a/src/executorlib/task_scheduler/interactive/blockallocation.py b/src/executorlib/task_scheduler/interactive/blockallocation.py index 02501e786..c1f05c511 100644 --- a/src/executorlib/task_scheduler/interactive/blockallocation.py +++ b/src/executorlib/task_scheduler/interactive/blockallocation.py @@ -80,35 +80,36 @@ def __init__( executor_kwargs["restart_limit"] = restart_limit self._process_kwargs = executor_kwargs self._max_workers = max_workers - self_id = random.getrandbits(128) - self._self_id = self_id + self._self_id = random.getrandbits(128) _interrupt_bootup_dict[self._self_id] = False - alive_workers = [max_workers] - alive_workers_lock = Lock() - bootup_events = [Event() for _ in range(self._max_workers)] - bootup_events[0].set() + self._alive_workers = [max_workers] + self._alive_workers_lock = Lock() + self._bootup_events = [Event() for _ in range(self._max_workers)] + self._bootup_events[0].set() self._set_process( process=[ Thread( target=_execute_multiple_tasks, - kwargs=executor_kwargs - | { - "worker_id": worker_id, - "stop_function": lambda: _interrupt_bootup_dict[self_id], - "bootup_event": bootup_events[worker_id], - "next_bootup_event": ( - bootup_events[worker_id + 1] - if worker_id + 1 < self._max_workers - else None - ), - "alive_workers": alive_workers, - "alive_workers_lock": alive_workers_lock, - }, + kwargs=self._worker_kwargs(worker_id), ) for worker_id in range(self._max_workers) ], ) + def _worker_kwargs(self, worker_id: int) -> dict: + return self._process_kwargs | { + "worker_id": worker_id, + "stop_function": lambda: _interrupt_bootup_dict[self._self_id], + "bootup_event": self._bootup_events[worker_id], + "next_bootup_event": ( + self._bootup_events[worker_id + 1] + if worker_id + 1 < len(self._bootup_events) + else None + ), + "alive_workers": self._alive_workers, + "alive_workers_lock": self._alive_workers_lock, + } + @property def max_workers(self) -> int: return self._max_workers @@ -126,12 +127,17 @@ def max_workers(self, max_workers: int): process for process in self._process if process.is_alive() ] elif self._max_workers < max_workers: + old_max_workers = self._max_workers + self._bootup_events.extend( + Event() for _ in range(max_workers - old_max_workers) + ) + self._alive_workers[0] += max_workers - old_max_workers new_process_lst = [ Thread( target=_execute_multiple_tasks, - kwargs=self._process_kwargs, + kwargs=self._worker_kwargs(worker_id), ) - for _ in range(max_workers - self._max_workers) + for worker_id in range(old_max_workers, max_workers) ] for process_instance in new_process_lst: process_instance.start() diff --git a/tests/unit/task_scheduler/interactive/test_blockallocation.py b/tests/unit/task_scheduler/interactive/test_blockallocation.py index dbd784480..e30c3afcf 100644 --- a/tests/unit/task_scheduler/interactive/test_blockallocation.py +++ b/tests/unit/task_scheduler/interactive/test_blockallocation.py @@ -1,11 +1,55 @@ import queue import unittest -from threading import Lock from concurrent.futures import Future +from threading import Event, Lock +from unittest.mock import patch -from executorlib.task_scheduler.interactive.blockallocation import _drain_dead_worker -from executorlib.task_scheduler.interactive.shared import task_done from executorlib.standalone.interactive.communication import ExecutorlibSocketError +from executorlib.task_scheduler.interactive.blockallocation import ( + BlockAllocationTaskScheduler, + _drain_dead_worker, +) + + +class TestBlockAllocationResize(unittest.TestCase): + def test_increase_workers_passes_worker_context(self): + scheduler = object.__new__(BlockAllocationTaskScheduler) + scheduler._future_queue = queue.Queue() + scheduler._process = [] + scheduler._process_kwargs = {"future_queue": scheduler._future_queue} + scheduler._max_workers = 1 + scheduler._self_id = 1 + scheduler._alive_workers = [1] + scheduler._alive_workers_lock = Lock() + scheduler._bootup_events = [Event()] + + class FakeThread: + instances = [] + + def __init__(self, target, kwargs): + self.target = target + self.kwargs = kwargs + self.started = False + self.instances.append(self) + + def start(self): + self.started = True + + with patch( + "executorlib.task_scheduler.interactive.blockallocation.Thread", + FakeThread, + ): + scheduler._bootup_events.append(FakeThread) + scheduler.max_workers = 2 + + worker = FakeThread.instances[-1] + self.assertEqual(worker.kwargs["worker_id"], 1) + self.assertIn("stop_function", worker.kwargs) + self.assertIn("bootup_event", worker.kwargs) + self.assertIn("next_bootup_event", worker.kwargs) + self.assertIs(worker.kwargs["alive_workers"], scheduler._alive_workers) + self.assertTrue(worker.started) + self.assertEqual(scheduler._alive_workers[0], 2) class TestDrainDeadWorker(unittest.TestCase): From 2c5a622f979cf09e83745d28616a72d2a07a6441 Mon Sep 17 00:00:00 2001 From: Muhtasim-Munif-Fahim Date: Tue, 11 Aug 2026 23:33:36 +0600 Subject: [PATCH 2/2] fix(scheduler): prevent zombie threads on scheduler resize and shutdown - Fix lambda closure capturing self in _worker_kwargs causing reference cycles that prevent __del__ and block worker threads on future_queue.get() - Set bootup events for new workers when max_workers increases so threads don't block forever on bootup_event.wait() - Use .get('result') instead of ['result'] in communication.py for graceful shutdown when spawned process is already dead - Remove erroneous FakeThread class append in test blockallocation --- src/executorlib/standalone/interactive/communication.py | 2 +- .../task_scheduler/interactive/blockallocation.py | 5 ++++- .../unit/task_scheduler/interactive/test_blockallocation.py | 1 - 3 files changed, 5 insertions(+), 3 deletions(-) diff --git a/src/executorlib/standalone/interactive/communication.py b/src/executorlib/standalone/interactive/communication.py index a5623dd2c..379a757bb 100644 --- a/src/executorlib/standalone/interactive/communication.py +++ b/src/executorlib/standalone/interactive/communication.py @@ -152,7 +152,7 @@ def shutdown(self, wait: bool = True): if self._spawner.poll(): result = self.send_and_receive_dict( input_dict={"shutdown": True, "wait": wait} - )["result"] + ).get("result") self._spawner.shutdown(wait=wait) self._reset_socket() return result diff --git a/src/executorlib/task_scheduler/interactive/blockallocation.py b/src/executorlib/task_scheduler/interactive/blockallocation.py index c1f05c511..c0daf940b 100644 --- a/src/executorlib/task_scheduler/interactive/blockallocation.py +++ b/src/executorlib/task_scheduler/interactive/blockallocation.py @@ -97,9 +97,10 @@ def __init__( ) def _worker_kwargs(self, worker_id: int) -> dict: + self_id = self._self_id return self._process_kwargs | { "worker_id": worker_id, - "stop_function": lambda: _interrupt_bootup_dict[self._self_id], + "stop_function": lambda: _interrupt_bootup_dict[self_id], "bootup_event": self._bootup_events[worker_id], "next_bootup_event": ( self._bootup_events[worker_id + 1] @@ -131,6 +132,8 @@ def max_workers(self, max_workers: int): self._bootup_events.extend( Event() for _ in range(max_workers - old_max_workers) ) + for idx in range(old_max_workers, max_workers): + self._bootup_events[idx].set() self._alive_workers[0] += max_workers - old_max_workers new_process_lst = [ Thread( diff --git a/tests/unit/task_scheduler/interactive/test_blockallocation.py b/tests/unit/task_scheduler/interactive/test_blockallocation.py index e30c3afcf..4373edd25 100644 --- a/tests/unit/task_scheduler/interactive/test_blockallocation.py +++ b/tests/unit/task_scheduler/interactive/test_blockallocation.py @@ -39,7 +39,6 @@ def start(self): "executorlib.task_scheduler.interactive.blockallocation.Thread", FakeThread, ): - scheduler._bootup_events.append(FakeThread) scheduler.max_workers = 2 worker = FakeThread.instances[-1]