From 077036c9d777838a7502039720a5590381d676ab Mon Sep 17 00:00:00 2001 From: Jamie <2119834+jamieQ@users.noreply.github.com> Date: Mon, 9 Nov 2026 10:10:30 -0600 Subject: [PATCH 1/2] feat(taskworker): Add opt-in wait for child exit during recycling Count retiring children toward concurrency until they exit and are reaped when wait_for_child_exit is enabled. Expose the option in both push and pull workers, preserving the existing default behavior. --- clients/python/src/taskbroker_client/cli.py | 14 + .../src/taskbroker_client/worker/worker.py | 353 +++++++++--------- clients/python/tests/worker/test_worker.py | 61 ++- 3 files changed, 254 insertions(+), 174 deletions(-) diff --git a/clients/python/src/taskbroker_client/cli.py b/clients/python/src/taskbroker_client/cli.py index b1ddb465..e207f4b6 100644 --- a/clients/python/src/taskbroker_client/cli.py +++ b/clients/python/src/taskbroker_client/cli.py @@ -111,6 +111,11 @@ def scheduler() -> None: default=0, type=int, ) +@click.option( + "--wait-for-child-exit/--no-wait-for-child-exit", + help="Count retiring children toward concurrency until they exit. May reduce recycling throughput.", + default=False, +) @click.option( "--child-tasks-queue-maxsize", help="Maximum number of pending tasks buffered for child processes.", @@ -174,6 +179,7 @@ def pull( namespace: str | None, concurrency: int, min_concurrency: int, + wait_for_child_exit: bool, child_tasks_queue_maxsize: int, result_queue_maxsize: int, rebalance_after: int, @@ -197,6 +203,7 @@ def pull( namespace=namespace, concurrency=concurrency, min_concurrency=min_concurrency, + wait_for_child_exit=wait_for_child_exit, child_tasks_queue_maxsize=child_tasks_queue_maxsize, result_queue_maxsize=result_queue_maxsize, rebalance_after=rebalance_after, @@ -245,6 +252,11 @@ def pull( default=0, type=int, ) +@click.option( + "--wait-for-child-exit/--no-wait-for-child-exit", + help="Count retiring children toward concurrency until they exit. May reduce recycling throughput.", + default=False, +) @click.option( "--child-tasks-queue-maxsize", help="Maximum number of pending tasks buffered for child processes.", @@ -343,6 +355,7 @@ def push( namespace: str | None, concurrency: int, min_concurrency: int, + wait_for_child_exit: bool, child_tasks_queue_maxsize: int, result_queue_maxsize: int, rebalance_after: int, @@ -372,6 +385,7 @@ def push( namespace=namespace, concurrency=concurrency, min_concurrency=min_concurrency, + wait_for_child_exit=wait_for_child_exit, child_tasks_queue_maxsize=child_tasks_queue_maxsize, result_queue_maxsize=result_queue_maxsize, rebalance_after=rebalance_after, diff --git a/clients/python/src/taskbroker_client/worker/worker.py b/clients/python/src/taskbroker_client/worker/worker.py index cdd71f71..8938a5d8 100644 --- a/clients/python/src/taskbroker_client/worker/worker.py +++ b/clients/python/src/taskbroker_client/worker/worker.py @@ -261,6 +261,7 @@ def __init__( prometheus_port: int | None = None, future_checking_frequency: float = 0.1, drain_file_path: str | None = DEFAULT_WORKER_DRAIN_FILE_PATH, + wait_for_child_exit: bool = False, ) -> None: app = import_app(app_module) @@ -280,6 +281,7 @@ def __init__( max_child_task_count=max_child_task_count, concurrency=concurrency, min_concurrency=min_concurrency, + wait_for_child_exit=wait_for_child_exit, child_tasks_queue_maxsize=child_tasks_queue_maxsize, result_queue_maxsize=result_queue_maxsize, processing_pool_name=processing_pool_name, @@ -633,6 +635,7 @@ def __init__( health_check_sec_per_touch: float = DEFAULT_WORKER_HEALTH_CHECK_SEC_PER_TOUCH, skip_awaiting_futures: bool = True, future_checking_frequency: float = 0.1, + wait_for_child_exit: bool = False, ) -> None: self._namespace = namespace app = import_app(app_module) @@ -653,6 +656,7 @@ def __init__( max_child_task_count=max_child_task_count, concurrency=concurrency, min_concurrency=min_concurrency, + wait_for_child_exit=wait_for_child_exit, child_tasks_queue_maxsize=child_tasks_queue_maxsize, result_queue_maxsize=result_queue_maxsize, processing_pool_name=processing_pool_name, @@ -876,8 +880,10 @@ def __init__( prometheus_port: int | None = None, future_checking_frequency: float = 0.1, drain_file_path: str | None = None, + wait_for_child_exit: bool = False, ) -> None: self._concurrency = concurrency + self._wait_for_child_exit = wait_for_child_exit if min_concurrency < concurrency: self._min_concurrency = min_concurrency @@ -1235,190 +1241,197 @@ def result_thread() -> None: ) self._result_thread.start() - def start_spawn_children_thread(self) -> None: - def spawn_children_thread() -> None: - logger.debug("taskworker.worker.spawn_children_thread.started") + def _manage_children_once(self, messages: multiprocessing.Queue[ChildMessage]) -> None: + # Read any events that may have come in since the last loop iteration + received: List[ChildMessage] = [] - # Queue of incoming message from children - messages: multiprocessing.Queue[ChildMessage] = self._mp_context.Queue() + while True: + try: + message = messages.get(block=False) + received.append(message) + except queue.Empty: + break - while not self._shutdown_event.is_set(): - # Read any events that may have come in since the last loop iteration - received: List[ChildMessage] = [] - - while True: - try: - message = messages.get(block=False) - received.append(message) - except queue.Empty: - break - - # Lifecycle-queue lag. A rising line means this thread is behind. - if received: - drain_at = time.monotonic() - self._metrics.distribution( - "taskworker.worker.child_message.age", - drain_at - min(m.timestamp for m in received), - tags={ - "processing_pool": self._processing_pool_name, - "pod_name": self._pod_name, - }, - ) + # Lifecycle-queue lag. A rising line means this thread is behind. + if received: + drain_at = time.monotonic() + self._metrics.distribution( + "taskworker.worker.child_message.age", + drain_at - min(m.timestamp for m in received), + tags={ + "processing_pool": self._processing_pool_name, + "pod_name": self._pod_name, + }, + ) - with self._children_lock: - children = list(self._children.items()) - - for cid, c in children: - if c.process.is_alive(): - continue - - c.process.join(timeout=0) - self._children.pop(cid) - - logger.info( - "taskworker.child.exited", - extra={ - "pid": c.process.pid, - "cid": str(cid), - "exitcode": c.process.exitcode, - "state": c.state, - "processing_pool": self._processing_pool_name, - }, - ) - - for message in received: - child = self._children.get(message.child_id) - - # If we received a message from a child, we MUST be tracking that child - if child is None: - logger.warning( - "taskworker.child_message.unknown_child", - extra={ - "cid": str(message.child_id), - "event": message.event, - "processing_pool": self._processing_pool_name, - }, - ) - - continue - - # Baseline here, where it also enters `running_count`. - if message.event == "running": - child.state = "running" - child.timing.mark_running(time.monotonic()) - - # This child wants to exit, but we may not have enough running children to shut down right away - elif message.event == "exiting": - self._exiting_children.append(message.child_id) + with self._children_lock: + children = list(self._children.items()) - while True: - # Compute how many children are still running - running = sum(1 for c in self._children.values() if c.state == "running") + for cid, c in children: + if c.process.is_alive(): + continue - if running <= self._min_concurrency: - # Cannot shut down any more children without falling below minimum concurrency (guaranteed < concurrency) - break + c.process.join(timeout=0) + self._children.pop(cid) - if not self._exiting_children: - # No children are waiting to exit - break + logger.info( + "taskworker.child.exited", + extra={ + "pid": c.process.pid, + "cid": str(cid), + "exitcode": c.process.exitcode, + "state": c.state, + "processing_pool": self._processing_pool_name, + }, + ) + + for message in received: + child = self._children.get(message.child_id) - child_id = self._exiting_children.popleft() - child = self._children.get(child_id) - - # Child may have died already - if child is None: - continue - - child.state = "exiting" - # Stop counting it, but keep the slot: the child writes - # once more on its way out, and its replacement can be - # handed the same memory the moment this one is freed. - child.timing.mark_stopped() - child.release.set() - - spawned = sum(1 for c in self._children.values() if c.state != "exiting") - - # How many children do we need to spawn? - needed = max(self._concurrency - spawned, 0) - - for _ in range(needed): - child_id = uuid4() - release = self._mp_context.Event() - - try: - # Inside the try so an allocation failure retries next - # pass instead of killing the spawn thread. - timing_shm = new_slot(self._mp_context) - - process = self._mp_context.Process( - name=f"taskworker-child-{child_id}", - target=child_process, - args=( - child_id, - self._app_module, - self._child_tasks, - self._processed_tasks, - self._shutdown_event, - self._max_child_task_count, - self._processing_pool_name, - self._process_type, - self._skip_awaiting_futures, - self._future_checking_frequency, - messages, - release, - timing_shm, - ), - ) - process.start() - - with self._children_lock: - child = TrackedChild( - process=process, - state="pending", - release=release, - timing=ChildTimeAccounting(shm=timing_shm), - ) - - self._children[child_id] = child - except Exception as e: - # `timing_shm` dies with this frame, so nothing leaks. - logger.exception( - "taskworker.child.spawn.failed", - extra={ - "cid": str(child_id), - "error": e, - "processing_pool": self._processing_pool_name, - }, - ) - - self._metrics.incr( - "taskworker.worker.child.spawn", - tags={ - "processing_pool": self._processing_pool_name, - "result": "failure", - }, - ) - - continue - - logger.info( - "taskworker.spawn_child", + # If we received a message from a child, we MUST be tracking that child + if child is None: + logger.warning( + "taskworker.child_message.unknown_child", extra={ - "pid": process.pid, - "cid": str(child_id), + "cid": str(message.child_id), + "event": message.event, "processing_pool": self._processing_pool_name, }, ) - self._metrics.incr( - "taskworker.worker.spawn_child", - tags={ - "processing_pool": self._processing_pool_name, - "result": "success", - }, + continue + + # Baseline here, where it also enters `running_count`. + if message.event == "running": + child.state = "running" + child.timing.mark_running(time.monotonic()) + + # This child wants to exit, but we may not have enough running children to shut down right away + elif message.event == "exiting": + self._exiting_children.append(message.child_id) + + while True: + # Compute how many children are still running + running = sum(1 for c in self._children.values() if c.state == "running") + + if running <= self._min_concurrency: + # Cannot shut down any more children without falling below minimum concurrency (guaranteed < concurrency) + break + + if not self._exiting_children: + # No children are waiting to exit + break + + child_id = self._exiting_children.popleft() + child = self._children.get(child_id) + + # Child may have died already + if child is None: + continue + + child.state = "exiting" + # Stop counting it, but keep the slot: the child writes + # once more on its way out, and its replacement can be + # handed the same memory the moment this one is freed. + child.timing.mark_stopped() + child.release.set() + + if self._wait_for_child_exit: + # Granting exit permission does not mean the child has exited. + # Keep counting it until we observe its exit and remove it above. + spawned = len(self._children) + else: + spawned = sum(1 for c in self._children.values() if c.state != "exiting") + + # How many children do we need to spawn? + needed = max(self._concurrency - spawned, 0) + + for _ in range(needed): + child_id = uuid4() + release = self._mp_context.Event() + + try: + # Inside the try so an allocation failure retries next + # pass instead of killing the spawn thread. + timing_shm = new_slot(self._mp_context) + + process = self._mp_context.Process( + name=f"taskworker-child-{child_id}", + target=child_process, + args=( + child_id, + self._app_module, + self._child_tasks, + self._processed_tasks, + self._shutdown_event, + self._max_child_task_count, + self._processing_pool_name, + self._process_type, + self._skip_awaiting_futures, + self._future_checking_frequency, + messages, + release, + timing_shm, + ), + ) + process.start() + + with self._children_lock: + child = TrackedChild( + process=process, + state="pending", + release=release, + timing=ChildTimeAccounting(shm=timing_shm), ) + self._children[child_id] = child + except Exception as e: + # `timing_shm` dies with this frame, so nothing leaks. + logger.exception( + "taskworker.child.spawn.failed", + extra={ + "cid": str(child_id), + "error": e, + "processing_pool": self._processing_pool_name, + }, + ) + + self._metrics.incr( + "taskworker.worker.child.spawn", + tags={ + "processing_pool": self._processing_pool_name, + "result": "failure", + }, + ) + + continue + + logger.info( + "taskworker.spawn_child", + extra={ + "pid": process.pid, + "cid": str(child_id), + "processing_pool": self._processing_pool_name, + }, + ) + + self._metrics.incr( + "taskworker.worker.spawn_child", + tags={ + "processing_pool": self._processing_pool_name, + "result": "success", + }, + ) + + def start_spawn_children_thread(self) -> None: + def spawn_children_thread() -> None: + logger.debug("taskworker.worker.spawn_children_thread.started") + + # Queue of incoming message from children + messages: multiprocessing.Queue[ChildMessage] = self._mp_context.Queue() + + while not self._shutdown_event.is_set(): + self._manage_children_once(messages) time.sleep(0.1) self._spawn_children_thread = threading.Thread( diff --git a/clients/python/tests/worker/test_worker.py b/clients/python/tests/worker/test_worker.py index 56079841..89776846 100644 --- a/clients/python/tests/worker/test_worker.py +++ b/clients/python/tests/worker/test_worker.py @@ -447,6 +447,7 @@ def _make_fake_context_pool( *, concurrency: int = 1, min_concurrency: int = 0, + wait_for_child_exit: bool = False, ) -> TaskWorkerProcessingPool: return TaskWorkerProcessingPool( app_module="examples.app:app", @@ -455,6 +456,7 @@ def _make_fake_context_pool( max_child_task_count=1, concurrency=concurrency, min_concurrency=min_concurrency, + wait_for_child_exit=wait_for_child_exit, processing_pool_name="test", process_type="fork", ) @@ -1979,9 +1981,12 @@ def test_spawn_children_releases_draining_child_above_min_concurrency() -> None: pool.shutdown() -def test_spawn_children_defers_draining_child_at_min_concurrency() -> None: +@pytest.mark.parametrize("wait_for_child_exit", [False, True]) +def test_spawn_children_defers_draining_child_at_min_concurrency(wait_for_child_exit: bool) -> None: fake_context = _FakeContext() - pool = _make_fake_context_pool(fake_context, concurrency=2, min_concurrency=1) + pool = _make_fake_context_pool( + fake_context, concurrency=2, min_concurrency=1, wait_for_child_exit=wait_for_child_exit + ) pool.start_spawn_children_thread() try: @@ -2007,9 +2012,14 @@ def test_spawn_children_defers_draining_child_at_min_concurrency() -> None: pool.shutdown() -def test_spawn_children_replaces_pending_child_that_dies_before_ready() -> None: +@pytest.mark.parametrize("wait_for_child_exit", [False, True]) +def test_spawn_children_replaces_pending_child_that_dies_before_ready( + wait_for_child_exit: bool, +) -> None: fake_context = _FakeContext() - pool = _make_fake_context_pool(fake_context, concurrency=1) + pool = _make_fake_context_pool( + fake_context, concurrency=1, wait_for_child_exit=wait_for_child_exit + ) pool.start_spawn_children_thread() try: @@ -2026,6 +2036,49 @@ def test_spawn_children_replaces_pending_child_that_dies_before_ready() -> None: pool.shutdown() +@pytest.mark.parametrize("wait_for_child_exit", [False, True]) +def test_spawn_children_replacement_timing(wait_for_child_exit: bool) -> None: + fake_context = _FakeContext() + pool = _make_fake_context_pool( + fake_context, concurrency=1, wait_for_child_exit=wait_for_child_exit + ) + messages = pool._mp_context.Queue() + + try: + pool._manage_children_once(messages) + assert len(fake_context.processes) == 1 + process = fake_context.processes[0] + child_id = process.args[0] + child = pool._children[child_id] + + messages.put(ChildMessage(child_id, "running")) + pool._manage_children_once(messages) + assert pool.ready_count == 1 + + messages.put(ChildMessage(child_id, "exiting")) + pool._manage_children_once(messages) + assert child.release.is_set() + assert process.is_alive() + + # By default the replacement overlaps with the retiring child. + expected_count = 1 if wait_for_child_exit else 2 + assert len(fake_context.processes) == expected_count + assert all(p.is_alive() for p in fake_context.processes) + + pool._manage_children_once(messages) + assert len(fake_context.processes) == expected_count + + process.exitcode = 0 + process.alive = False + pool._manage_children_once(messages) + assert process.join_calls == [0] + assert child_id not in pool._children + assert len(fake_context.processes) == 2 + assert fake_context.processes[1].is_alive() + finally: + pool.shutdown() + + def test_shutdown_terminates_all_tracked_children() -> None: fake_context = _FakeContext() pool = _make_fake_context_pool(fake_context, concurrency=2) From 9ed5292570d4f8a436e4cead8cd4d67f0babab53 Mon Sep 17 00:00:00 2001 From: Jamie <2119834+jamieQ@users.noreply.github.com> Date: Wed, 11 Nov 2026 09:29:05 -0600 Subject: [PATCH 2/2] feat(taskworker): Report oldest exiting child age Record when child exit permission is granted and emit the oldest exiting child age in seconds per pool and pod. Report in both recycling modes, with zero when no children are exiting, so prolonged exits can be monitored without a hardcoded threshold or termination policy. --- .../src/taskbroker_client/worker/worker.py | 18 +++++++++++++ clients/python/tests/worker/test_worker.py | 25 +++++++++++++++++++ 2 files changed, 43 insertions(+) diff --git a/clients/python/src/taskbroker_client/worker/worker.py b/clients/python/src/taskbroker_client/worker/worker.py index 8938a5d8..9ebc02e4 100644 --- a/clients/python/src/taskbroker_client/worker/worker.py +++ b/clients/python/src/taskbroker_client/worker/worker.py @@ -232,6 +232,7 @@ class TrackedChild: release: Event # Holds the child's slot. Dropping this entry is what frees the memory. timing: ChildTimeAccounting + exit_granted_at: float | None = None class PushTaskWorker: @@ -1056,6 +1057,16 @@ def _emit_periodic_metrics(self) -> None: exiting_children = len(self._exiting_children) + now = time.monotonic() + oldest_exiting_child_age = max( + ( + now - child.exit_granted_at + for child in self._children.values() + if child.exit_granted_at is not None + ), + default=0.0, + ) + # Emitted during warmup too: zero is correct for a counter. self._metrics.distribution( "taskworker.worker.child_busy_seconds", @@ -1130,6 +1141,12 @@ def _emit_periodic_metrics(self) -> None: tags={**tags, "state": state}, ) + self._metrics.gauge( + "taskworker.worker.oldest_exiting_child_age", + oldest_exiting_child_age, + tags=tags, + ) + # Emit number of children waiting to exit self._metrics.gauge( "taskworker.worker.exiting_children.size", @@ -1334,6 +1351,7 @@ def _manage_children_once(self, messages: multiprocessing.Queue[ChildMessage]) - # once more on its way out, and its replacement can be # handed the same memory the moment this one is freed. child.timing.mark_stopped() + child.exit_granted_at = time.monotonic() child.release.set() if self._wait_for_child_exit: diff --git a/clients/python/tests/worker/test_worker.py b/clients/python/tests/worker/test_worker.py index 89776846..f14f4e4c 100644 --- a/clients/python/tests/worker/test_worker.py +++ b/clients/python/tests/worker/test_worker.py @@ -2005,9 +2005,11 @@ def test_spawn_children_defers_draining_child_at_min_concurrency(wait_for_child_ messages.put(ChildMessage(first_child_id, "exiting")) time.sleep(0.25) assert not first_release.is_set() + assert pool._children[first_child_id].exit_granted_at is None messages.put(ChildMessage(second_child_id, "running")) _wait_for(first_release.is_set) + assert pool._children[first_child_id].exit_granted_at is not None finally: pool.shutdown() @@ -2056,8 +2058,11 @@ def test_spawn_children_replacement_timing(wait_for_child_exit: bool) -> None: assert pool.ready_count == 1 messages.put(ChildMessage(child_id, "exiting")) + before_release = time.monotonic() pool._manage_children_once(messages) assert child.release.is_set() + assert child.exit_granted_at is not None + assert before_release <= child.exit_granted_at <= time.monotonic() assert process.is_alive() # By default the replacement overlaps with the retiring child. @@ -2079,6 +2084,26 @@ def test_spawn_children_replacement_timing(wait_for_child_exit: bool) -> None: pool.shutdown() +@pytest.mark.parametrize( + "exit_granted_times,expected", + [([], 0), ([None], 0), ([150, None, 0, 100], 200)], +) +def test_oldest_exiting_child_age(exit_granted_times: list[float | None], expected: float) -> None: + pool = _make_fake_context_pool(_FakeContext()) + pool._metrics = mock.Mock() + for granted_at in exit_granted_times: + child = _make_tracked_child("pending" if granted_at is None else "exiting") + child.exit_granted_at = granted_at + pool._children[uuid4()] = child + + with mock.patch("taskbroker_client.worker.worker.time.monotonic", return_value=200): + pool._emit_periodic_metrics() + + calls = _gauge_calls(pool._metrics, "taskworker.worker.oldest_exiting_child_age") + assert len(calls) == 1 + assert calls[0].args[1] == expected + + def test_shutdown_terminates_all_tracked_children() -> None: fake_context = _FakeContext() pool = _make_fake_context_pool(fake_context, concurrency=2)