diff --git a/clients/python/src/taskbroker_client/worker/worker.py b/clients/python/src/taskbroker_client/worker/worker.py index 690b6514..1c06f0ee 100644 --- a/clients/python/src/taskbroker_client/worker/worker.py +++ b/clients/python/src/taskbroker_client/worker/worker.py @@ -8,10 +8,13 @@ import threading import time from concurrent.futures import ThreadPoolExecutor +from dataclasses import dataclass from multiprocessing.context import ForkContext, ForkServerContext, SpawnContext from multiprocessing.process import BaseProcess +from multiprocessing.synchronize import Event from pathlib import Path -from typing import TYPE_CHECKING, Any, Callable, List +from typing import TYPE_CHECKING, Any, Callable, Dict, List, Literal +from uuid import UUID, uuid4 import grpc from grpc_health.v1 import health, health_pb2, health_pb2_grpc @@ -42,7 +45,7 @@ parse_rpc_secret_list, ) from taskbroker_client.worker.push_clients import BatchPushTaskbrokerClient, PushTaskbrokerClient -from taskbroker_client.worker.workerchild import child_process +from taskbroker_client.worker.workerchild import ChildMessage, child_process if TYPE_CHECKING: ServerInterceptor = grpc.ServerInterceptor[Any, Any] @@ -122,6 +125,16 @@ class RequeueException(Exception): pass +ChildState = Literal["pending", "ready", "exiting"] + + +@dataclass +class TrackedChild: + process: BaseProcess + state: ChildState + release: Event + + class PushTaskWorker: _mp_context: ForkContext | SpawnContext | ForkServerContext @@ -800,9 +813,9 @@ def __init__( self._processed_tasks: multiprocessing.Queue[ProcessingResult] = self._mp_context.Queue( maxsize=result_queue_maxsize ) - self._children: list[BaseProcess] = [] + self._children: Dict[UUID, TrackedChild] = {} + self._children_lock = threading.Lock() self._shutdown_event = self._mp_context.Event() - self._ready_counter = self._mp_context.Value("i", 0) self._result_thread: threading.Thread | None = None self._metrics_thread: threading.Thread | None = None self._spawn_children_thread: threading.Thread | None = None @@ -810,7 +823,41 @@ def __init__( @property def ready_count(self) -> int: """Number of children that have finished warming up and are consuming.""" - return self._ready_counter.value + with self._children_lock: + return sum(1 for c in self._children.values() if c.state == "ready") + + def _plan_child_pool_changes(self) -> tuple[int, list[TrackedChild]]: + """ + Determine how to move the child pool toward the desired concurrency. + + The caller must hold `_children_lock`. + + Exiting children are still consuming tasks until released, so they count + as serving capacity. Pending children are not serving yet, but they count + as replacement capacity that has already been requested. + """ + pending_children = [c for c in self._children.values() if c.state == "pending"] + ready_children = [c for c in self._children.values() if c.state == "ready"] + exiting_children = [c for c in self._children.values() if c.state == "exiting"] + + desired_serving_capacity = self._concurrency + current_serving_capacity = len(ready_children) + len(exiting_children) + + excess_serving_capacity = max(current_serving_capacity - desired_serving_capacity, 0) + children_to_release = exiting_children[:excess_serving_capacity] + + kept_exiting_count = len(exiting_children) - len(children_to_release) + + # Keep enough pending children to both reach desired concurrency and replace + # exiting children that are still serving during their grace period + target_pending_count = max( + desired_serving_capacity - len(ready_children), + kept_exiting_count, + ) + + children_to_spawn = max(target_pending_count - len(pending_children), 0) + + return children_to_spawn, children_to_release def send_results(self, results: list[ProcessingResult], is_draining: bool = False) -> None: """ @@ -917,16 +964,97 @@ def result_thread() -> None: 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._children = [child for child in self._children if child.is_alive()] - if len(self._children) >= self._concurrency: - time.sleep(0.1) - continue - for i in range(self._concurrency - len(self._children)): + # 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 + + 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 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 + + if message.event == "ready": + if child.state != "pending": + logger.warning( + "taskworker.child.ready_not_pending", + extra={ + "cid": str(message.child_id), + "state": child.state, + "processing_pool": self._processing_pool_name, + }, + ) + + child.state = "ready" + + elif message.event == "exiting": + if child.state != "ready": + logger.warning( + "taskworker.child.exiting_not_ready", + extra={ + "cid": str(message.child_id), + "state": child.state, + "processing_pool": self._processing_pool_name, + }, + ) + + child.state = "exiting" + + children_to_spawn, children_to_release = self._plan_child_pool_changes() + + for child in children_to_release: + child.release.set() + + for _ in range(children_to_spawn): + child_id = uuid4() + release = self._mp_context.Event() + process = self._mp_context.Process( - name=f"taskworker-child-{i}", + name=f"taskworker-child-{child_id}", target=child_process, args=( + child_id, self._app_module, self._child_tasks, self._processed_tasks, @@ -935,20 +1063,59 @@ def spawn_children_thread() -> None: self._processing_pool_name, self._process_type, self._skip_awaiting_futures, - self._ready_counter, + messages, + release, ), ) - process.start() - self._children.append(process) + + try: + process.start() + except Exception as e: + 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 + + with self._children_lock: + self._children[child_id] = TrackedChild( + process=process, + state="pending", + release=release, + ) + logger.info( "taskworker.spawn_child", - extra={"pid": process.pid, "processing_pool": self._processing_pool_name}, + 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}, + tags={ + "processing_pool": self._processing_pool_name, + "result": "success", + }, ) + time.sleep(0.1) + self._spawn_children_thread = threading.Thread( name="spawn-children", target=spawn_children_thread, daemon=True ) @@ -1005,9 +1172,12 @@ def shutdown(self) -> None: self._spawn_children_thread.join() logger.info("taskworker.worker.shutdown.children") - for child in self._children: + with self._children_lock: + children = [tracked_child.process for tracked_child in self._children.values()] + + for child in children: child.terminate() - for child in self._children: + for child in children: child.join(WORKER_CHILD_JOIN_TIMEOUT_SEC) if child.is_alive(): child.kill() diff --git a/clients/python/src/taskbroker_client/worker/workerchild.py b/clients/python/src/taskbroker_client/worker/workerchild.py index 5395920c..cc4fffb1 100644 --- a/clients/python/src/taskbroker_client/worker/workerchild.py +++ b/clients/python/src/taskbroker_client/worker/workerchild.py @@ -2,6 +2,7 @@ import contextlib import logging +import multiprocessing import queue import signal import threading @@ -11,10 +12,8 @@ from functools import partial from multiprocessing.synchronize import Event from types import FrameType -from typing import TYPE_CHECKING, Any - -if TYPE_CHECKING: - from multiprocessing.sharedctypes import Synchronized +from typing import Any, Literal +from uuid import UUID # XXX: Don't import any modules that will import django here, do those within child_process import msgpack @@ -166,7 +165,14 @@ def _log_task_retry_exhausted( ) +@dataclass(frozen=True) +class ChildMessage: + child_id: UUID + event: Literal["ready"] | Literal["exiting"] + + def child_process( + child_id: UUID, app_module: str, child_tasks: queue.Queue[InflightTaskActivation], processed_tasks: queue.Queue[ProcessingResult], @@ -175,7 +181,8 @@ def child_process( processing_pool_name: str, process_type: str, skip_awaiting_futures: bool, - ready_counter: "Synchronized[int] | None" = None, + messages: multiprocessing.Queue[ChildMessage], + parent_release: Event, ) -> None: """ The entrypoint for spawned worker children. @@ -351,18 +358,30 @@ def check_task_future_completion( ) _future_completion_thread.start() + waiting_for_parent_release = False + while not shutdown_event.is_set() and not local_shutdown.is_set(): if max_task_count and processed_task_count >= max_task_count: - metrics.incr( - "taskworker.worker.max_task_count_reached", - tags={"count": processed_task_count, "processing_pool": processing_pool_name}, - ) - logger.info( - "taskworker.max_task_count_reached", extra={"count": processed_task_count} - ) - # Still set the shutdown signal to trigger shutdown of the future checker thread - local_shutdown.set() - break + if not waiting_for_parent_release: + metrics.incr( + "taskworker.worker.max_task_count_reached", + tags={ + "count": processed_task_count, + "processing_pool": processing_pool_name, + }, + ) + + logger.info( + "taskworker.max_task_count_reached", extra={"count": processed_task_count} + ) + + # Tell the parent this child can be replaced. + messages.put_nowait(ChildMessage(child_id, "exiting")) + waiting_for_parent_release = True + + if parent_release.is_set(): + local_shutdown.set() + break try: inflight = child_tasks.get(timeout=1.0) @@ -806,11 +825,8 @@ def _task_execution_complete( futures_start_time, ) - # Signal that this child has finished warmup and ready to consume tasks. The parent uses this - # to gate the gRPC SERVING health signal. Monotonic by design - if ready_counter is not None: - with ready_counter.get_lock(): - ready_counter.value += 1 + # Tell the parent that this child has warmed up and is ready to consume tasks + messages.put_nowait(ChildMessage(child_id, "ready")) # Run the worker loop run_worker( diff --git a/clients/python/tests/worker/test_worker.py b/clients/python/tests/worker/test_worker.py index bcf6186a..0f68a397 100644 --- a/clients/python/tests/worker/test_worker.py +++ b/clients/python/tests/worker/test_worker.py @@ -8,9 +8,11 @@ from concurrent.futures import Future from datetime import datetime from multiprocessing import Event, get_context +from multiprocessing.synchronize import Event as MultiprocessingEvent from pathlib import Path -from typing import Any +from typing import Any, Callable from unittest import TestCase, mock +from uuid import uuid4 import grpc import msgpack @@ -45,7 +47,8 @@ TaskWorkerProcessingPool, WorkerServicer, ) -from taskbroker_client.worker.workerchild import child_process +from taskbroker_client.worker.workerchild import ChildMessage +from taskbroker_client.worker.workerchild import child_process as _child_process SIMPLE_TASK = InflightTaskActivation( host="localhost:50051", @@ -266,6 +269,36 @@ def _make_processing_result(task_id: str) -> ProcessingResult: ) +def child_process( + app_module: str, + child_tasks: queue.Queue[InflightTaskActivation], + processed_tasks: queue.Queue[ProcessingResult], + shutdown_event: MultiprocessingEvent, + max_task_count: int | None, + processing_pool_name: str, + process_type: str, + skip_awaiting_futures: bool, +) -> None: + ctx = get_context("fork") + messages = ctx.Queue() + parent_release = ctx.Event() + parent_release.set() + + _child_process( + uuid4(), + app_module, + child_tasks, + processed_tasks, + shutdown_event, + max_task_count, + processing_pool_name, + process_type, + skip_awaiting_futures, + messages, + parent_release, + ) + + class _SendResultCapture: def __init__(self) -> None: self.send_calls: list[tuple[list[ProcessingResult], bool]] = [] @@ -302,6 +335,87 @@ def _make_result_thread_pool( ) +class _FakeProcess: + _next_pid = 1 + + def __init__(self, *args: Any, **kwargs: Any) -> None: + self.name = kwargs["name"] + self.target = kwargs["target"] + self.args = kwargs["args"] + self.pid = _FakeProcess._next_pid + _FakeProcess._next_pid += 1 + self.exitcode: int | None = None + self.started = False + self.alive = False + self.terminated = False + self.killed = False + self.join_calls: list[float | None] = [] + + def start(self) -> None: + self.started = True + self.alive = True + + def is_alive(self) -> bool: + return self.alive + + def join(self, timeout: float | None = None) -> None: + self.join_calls.append(timeout) + + def terminate(self) -> None: + self.terminated = True + self.alive = False + self.exitcode = -signal.SIGTERM + + def kill(self) -> None: + self.killed = True + self.alive = False + self.exitcode = -signal.SIGKILL + + +class _FakeContext: + def __init__(self) -> None: + self.processes: list[_FakeProcess] = [] + self.queues: list[queue.Queue[Any]] = [] + + def Queue(self, maxsize: int = 0) -> queue.Queue[Any]: + created: queue.Queue[Any] = queue.Queue(maxsize=maxsize) + self.queues.append(created) + return created + + def Event(self) -> threading.Event: + return threading.Event() + + def Process(self, *args: Any, **kwargs: Any) -> _FakeProcess: + process = _FakeProcess(*args, **kwargs) + self.processes.append(process) + return process + + +def _make_fake_context_pool( + fake_context: _FakeContext, + *, + concurrency: int = 1, +) -> TaskWorkerProcessingPool: + return TaskWorkerProcessingPool( + app_module="examples.app:app", + send_result_fn=lambda x, y: None, + mp_context=fake_context, # type: ignore[arg-type] + max_child_task_count=1, + concurrency=concurrency, + processing_pool_name="test", + process_type="fork", + ) + + +def _wait_for(condition: Callable[[], bool], timeout: float = 5) -> None: + start = time.time() + while time.time() - start < timeout: + if condition(): + return + time.sleep(0.01) + raise AssertionError("Timed out waiting for condition") + + class TestTaskWorker(TestCase): def test_fetch_task(self) -> None: taskworker = TaskWorker( @@ -359,13 +473,11 @@ def test_run_once_no_next_task(self) -> None: taskworker.shutdown() assert mock_client.get_task.called - assert mock_client.update_task.call_count == 1 - assert mock_client.update_task.call_args.args[0].host == "localhost:50051" - assert mock_client.update_task.call_args.args[0].task_id == SIMPLE_TASK.activation.id - assert ( - mock_client.update_task.call_args.args[0].status == TASK_ACTIVATION_STATUS_COMPLETE - ) - assert mock_client.update_task.call_args.args[1] is None + assert mock_client.update_task.call_count >= 1 + first_update = mock_client.update_task.call_args_list[0] + assert first_update.args[0].host == "localhost:50051" + assert first_update.args[0].task_id == SIMPLE_TASK.activation.id + assert first_update.args[0].status == TASK_ACTIVATION_STATUS_COMPLETE def test_run_once_with_next_task(self) -> None: # Cover the scenario where update_task returns the next task which should @@ -401,13 +513,11 @@ def update_task_response(*args: Any, **kwargs: Any) -> InflightTaskActivation | taskworker.shutdown() assert mock_client.get_task.called - assert mock_client.update_task.call_count == 2 - assert mock_client.update_task.call_args.args[0].host == "localhost:50051" - assert mock_client.update_task.call_args.args[0].task_id == SIMPLE_TASK.activation.id - assert ( - mock_client.update_task.call_args.args[0].status == TASK_ACTIVATION_STATUS_COMPLETE - ) - assert mock_client.update_task.call_args.args[1] is None + assert mock_client.update_task.call_count >= 2 + for update_call in mock_client.update_task.call_args_list[:2]: + assert update_call.args[0].host == "localhost:50051" + assert update_call.args[0].task_id == SIMPLE_TASK.activation.id + assert update_call.args[0].status == TASK_ACTIVATION_STATUS_COMPLETE def test_run_once_with_update_failure(self) -> None: # Cover the scenario where update_task fails a few times in a row @@ -665,9 +775,14 @@ def _make_push_worker(**kwargs: Any) -> PushTaskWorker: def test_await_children_warm_returns_when_ready() -> None: taskworker = _make_push_worker(concurrency=4, warmup_timeout=5) - taskworker.worker_pool._ready_counter.value = 4 - with mock.patch.object(taskworker, "_metrics") as mock_metrics: + with ( + mock.patch.object(taskworker, "_metrics") as mock_metrics, + mock.patch.object( + TaskWorkerProcessingPool, "ready_count", new_callable=mock.PropertyMock + ) as ready_count, + ): + ready_count.return_value = 4 start = time.time() taskworker._await_children_warm() elapsed = time.time() - start @@ -687,8 +802,6 @@ def test_await_children_warm_returns_when_ready() -> None: def test_await_children_warm_times_out() -> None: taskworker = _make_push_worker(concurrency=4, warmup_timeout=0.1) - # Never becomes ready. - taskworker.worker_pool._ready_counter.value = 0 with mock.patch.object(taskworker, "_metrics") as mock_metrics: start = time.time() @@ -703,16 +816,23 @@ def test_await_children_warm_times_out() -> None: def test_await_children_warm_unblocks_when_children_warm() -> None: taskworker = _make_push_worker(concurrency=2, warmup_timeout=5) - taskworker.worker_pool._ready_counter.value = 0 + ready_child_count = 0 def warm_up() -> None: + nonlocal ready_child_count time.sleep(0.2) - taskworker.worker_pool._ready_counter.value = 2 + ready_child_count = 2 warmer = threading.Thread(target=warm_up) warmer.start() try: - with mock.patch.object(taskworker, "_metrics") as mock_metrics: + with ( + mock.patch.object(taskworker, "_metrics") as mock_metrics, + mock.patch.object( + TaskWorkerProcessingPool, "ready_count", new_callable=mock.PropertyMock + ) as ready_count, + ): + ready_count.side_effect = lambda: ready_child_count start = time.time() taskworker._await_children_warm() elapsed = time.time() - start @@ -733,7 +853,6 @@ def test_start_does_not_serve_when_shutdown_during_warmup() -> None: taskworker = _make_push_worker(concurrency=2, warmup_timeout=5) # Children never warm, and shutdown is requested before start() runs. - taskworker.worker_pool._ready_counter.value = 0 taskworker._grpc_sync_event.set() fake_health = mock.MagicMock() @@ -768,6 +887,82 @@ def test_start_does_not_serve_when_shutdown_during_warmup() -> None: fake_server.wait_for_termination.assert_not_called() +def test_spawn_children_counts_pending_children_toward_concurrency() -> None: + fake_context = _FakeContext() + pool = _make_fake_context_pool(fake_context, concurrency=2) + + pool.start_spawn_children_thread() + try: + _wait_for(lambda: len(fake_context.processes) == 2) + time.sleep(0.25) + + assert len(fake_context.processes) == 2 + assert pool.ready_count == 0 + finally: + pool.shutdown() + + +def test_spawn_children_releases_draining_child_after_replacement_ready() -> None: + fake_context = _FakeContext() + pool = _make_fake_context_pool(fake_context, concurrency=1) + + pool.start_spawn_children_thread() + try: + _wait_for(lambda: len(fake_context.processes) == 1) + messages = fake_context.queues[-1] + first_process = fake_context.processes[0] + first_child_id = first_process.args[0] + first_release = first_process.args[-1] + + messages.put(ChildMessage(first_child_id, "ready")) + _wait_for(lambda: pool.ready_count == 1) + + messages.put(ChildMessage(first_child_id, "exiting")) + _wait_for(lambda: len(fake_context.processes) == 2) + + assert not first_release.is_set() + + second_process = fake_context.processes[1] + second_child_id = second_process.args[0] + messages.put(ChildMessage(second_child_id, "ready")) + + _wait_for(first_release.is_set) + finally: + pool.shutdown() + + +def test_spawn_children_replaces_pending_child_that_dies_before_ready() -> None: + fake_context = _FakeContext() + pool = _make_fake_context_pool(fake_context, concurrency=1) + + pool.start_spawn_children_thread() + try: + _wait_for(lambda: len(fake_context.processes) == 1) + first_process = fake_context.processes[0] + first_process.alive = False + first_process.exitcode = 1 + + _wait_for(lambda: len(fake_context.processes) == 2) + + assert first_process.join_calls + assert fake_context.processes[1].started + finally: + pool.shutdown() + + +def test_shutdown_terminates_all_tracked_children() -> None: + fake_context = _FakeContext() + pool = _make_fake_context_pool(fake_context, concurrency=2) + + pool.start_spawn_children_thread() + _wait_for(lambda: len(fake_context.processes) == 2) + + pool.shutdown() + + assert all(process.terminated for process in fake_context.processes) + assert all(process.join_calls for process in fake_context.processes) + + class TestWorkerServicer(TestCase): def test_push_task_success(self) -> None: taskworker = PushTaskWorker( @@ -893,15 +1088,19 @@ def test_child_process_complete(mock_capture_checkin: mock.MagicMock) -> None: assert mock_capture_checkin.call_count == 0 -def test_child_process_increments_ready_counter() -> None: +def test_child_process_emits_ready_message() -> None: todo: queue.Queue[InflightTaskActivation] = queue.Queue() processed: queue.Queue[ProcessingResult] = queue.Queue() shutdown = Event() ctx = get_context("fork") - ready_counter = ctx.Value("i", 0) + child_id = uuid4() + messages = ctx.Queue() + parent_release = ctx.Event() + parent_release.set() todo.put(SIMPLE_TASK) - child_process( + _child_process( + child_id, "examples.app:app", todo, processed, @@ -910,11 +1109,69 @@ def test_child_process_increments_ready_counter() -> None: processing_pool_name="test", process_type="fork", skip_awaiting_futures=False, - ready_counter=ready_counter, + messages=messages, + parent_release=parent_release, + ) + + # The child signals readiness once warmup is done, before consuming. + message = messages.get(timeout=1) + assert message == ChildMessage(child_id, "ready") + + +@mock.patch("taskbroker_client.worker.workerchild.capture_checkin") +def test_child_process_emits_exiting_once_and_continues_until_release( + mock_capture_checkin: mock.MagicMock, +) -> None: + shutdown = Event() + ctx = get_context("fork") + child_id = uuid4() + todo = ctx.Queue() + processed = ctx.Queue() + messages = ctx.Queue() + parent_release = ctx.Event() + + todo.put(SIMPLE_TASK) + process = ctx.Process( + target=_child_process, + args=( + child_id, + "examples.app:app", + todo, + processed, + shutdown, + 1, + "test", + "fork", + False, + messages, + parent_release, + ), ) + process.start() + try: + ready_message = messages.get(timeout=5) + exiting_message = messages.get(timeout=5) + + assert ready_message == ChildMessage(child_id, "ready") + assert exiting_message == ChildMessage(child_id, "exiting") + assert processed.get(timeout=5).task_id == SIMPLE_TASK.activation.id - # The child increments the counter once warmup is done, before consuming. - assert ready_counter.value == 1 + todo.put(SIMPLE_TASK) + assert processed.get(timeout=5).task_id == SIMPLE_TASK.activation.id + + time.sleep(0.2) + assert process.is_alive() + assert messages.empty() + + parent_release.set() + process.join(timeout=5) + assert not process.is_alive() + finally: + if process.is_alive(): + process.terminate() + process.join(timeout=5) + + assert mock_capture_checkin.call_count == 0 def test_child_process_remove_start_time_kwargs() -> None: