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..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: @@ -261,6 +262,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 +282,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 +636,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 +657,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 +881,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 @@ -1050,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", @@ -1124,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", @@ -1235,190 +1258,198 @@ 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, + }, + ) - 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", + 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={ - "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.exit_granted_at = time.monotonic() + 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..f14f4e4c 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: @@ -2000,16 +2005,23 @@ def test_spawn_children_defers_draining_child_at_min_concurrency() -> None: 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() -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 +2038,72 @@ 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")) + 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. + 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() + + +@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)