diff --git a/clients/python/src/taskbroker_client/worker/worker.py b/clients/python/src/taskbroker_client/worker/worker.py index 794c7d8c..7f89d5d1 100644 --- a/clients/python/src/taskbroker_client/worker/worker.py +++ b/clients/python/src/taskbroker_client/worker/worker.py @@ -7,11 +7,15 @@ import signal import threading import time +from collections import deque 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, Deque, Dict, List, Literal +from uuid import UUID, uuid4 import grpc import prometheus_client @@ -43,7 +47,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] @@ -144,6 +148,17 @@ class RequeueException(Exception): pass +ChildState = Literal["pending", "running", "exiting"] + + +@dataclass +class TrackedChild: + process: BaseProcess + state: ChildState + release: Event + busy: bool = False + + class PushTaskWorker: _mp_context: ForkContext | SpawnContext | ForkServerContext @@ -154,6 +169,7 @@ def __init__( max_child_task_count: int | None = None, namespace: str | None = None, concurrency: int = 1, + min_concurrency: int = 0, child_tasks_queue_maxsize: int = DEFAULT_WORKER_QUEUE_SIZE, result_queue_maxsize: int = DEFAULT_WORKER_QUEUE_SIZE, rebalance_after: int = DEFAULT_REBALANCE_AFTER, @@ -187,6 +203,7 @@ def __init__( send_result_fn=self._send_results, max_child_task_count=max_child_task_count, concurrency=concurrency, + min_concurrency=min_concurrency, child_tasks_queue_maxsize=child_tasks_queue_maxsize, result_queue_maxsize=result_queue_maxsize, processing_pool_name=processing_pool_name, @@ -587,6 +604,7 @@ def __init__( max_child_task_count: int | None = None, namespace: str | None = None, concurrency: int = 1, + min_concurrency: int = 0, child_tasks_queue_maxsize: int = DEFAULT_WORKER_QUEUE_SIZE, result_queue_maxsize: int = DEFAULT_WORKER_QUEUE_SIZE, rebalance_after: int = DEFAULT_REBALANCE_AFTER, @@ -615,6 +633,7 @@ def __init__( send_result_fn=self._send_results, max_child_task_count=max_child_task_count, concurrency=concurrency, + min_concurrency=min_concurrency, child_tasks_queue_maxsize=child_tasks_queue_maxsize, result_queue_maxsize=result_queue_maxsize, processing_pool_name=processing_pool_name, @@ -798,6 +817,7 @@ def __init__( mp_context: ForkContext | SpawnContext | ForkServerContext, max_child_task_count: int | None = None, concurrency: int = 1, + min_concurrency: int = 0, child_tasks_queue_maxsize: int = DEFAULT_WORKER_QUEUE_SIZE, result_queue_maxsize: int = DEFAULT_WORKER_QUEUE_SIZE, processing_pool_name: str | None = None, @@ -809,6 +829,12 @@ def __init__( future_checking_frequency: float = 0.1, ) -> None: self._concurrency = concurrency + + if min_concurrency < concurrency: + self._min_concurrency = min_concurrency + else: + raise ValueError("Minimum concurrency must be strictly below concurrency") + self._processing_pool_name = processing_pool_name or "unknown" self._pod_name = pod_name or "unknown" self._update_in_batches = update_in_batches @@ -831,10 +857,10 @@ 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._exiting_children: Deque[UUID] = deque() + self._children_lock = threading.Lock() self._shutdown_event = self._mp_context.Event() - self._ready_counter = self._mp_context.Value("i", 0) - self._busy_counter = self._mp_context.Value("i", 0) self._prometheus_port = prometheus_port self._prom: WorkerPrometheusMetrics | None = None self._result_thread: threading.Thread | None = None @@ -844,7 +870,73 @@ 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 == "running") + + def _emit_periodic_metrics(self) -> None: + tags = { + "processing_pool": self._processing_pool_name, + "pod_name": self._pod_name, + } + + # Emit queue size metrics + try: + # Method 'qsize' not implemented on all platforms, such as macOS + self._metrics.gauge( + "taskworker.child_tasks.size", + float(self._child_tasks.qsize()), + tags=tags, + ) + + self._metrics.gauge( + "taskworker.processed_tasks.size", + float(self._processed_tasks.qsize()), + tags=tags, + ) + except Exception as e: + logger.debug( + "taskworker.worker.queue_gauges.error", + extra={"error": e, "processing_pool": self._processing_pool_name}, + ) + + # Count the number of children in each state and waiting for exit + with self._children_lock: + state_counts: dict[ChildState, int] = { + "pending": 0, + "running": 0, + "exiting": 0, + } + + for child in self._children.values(): + state_counts[child.state] += 1 + + busy = sum(1 for child in self._children.values() if child.busy) + exiting_children = len(self._exiting_children) + + bounded_busy = max(0, min(busy, self._concurrency)) + occupancy = bounded_busy / self._concurrency if self._concurrency else 0.0 + self._metrics.gauge( + "taskworker.worker.occupancy", + occupancy, + tags=tags, + ) + if self._prom is not None: + self._prom.occupancy.labels(processing_pool=self._processing_pool_name).set(occupancy) + + # Emit number of children in each state + for state, count in state_counts.items(): + self._metrics.gauge( + "taskworker.worker.children", + float(count), + tags={**tags, "state": state}, + ) + + # Emit number of children waiting to exit + self._metrics.gauge( + "taskworker.worker.exiting_children.size", + float(exiting_children), + tags=tags, + ) def send_results(self, results: list[ProcessingResult], is_draining: bool = False) -> None: """ @@ -874,37 +966,9 @@ def start_metrics_thread(self) -> None: self._prom = WorkerPrometheusMetrics(self._prometheus_port) def metrics_thread() -> None: - tags = { - "processing_pool": self._processing_pool_name, - "pod_name": self._pod_name, - } - while True: try: - busy = max(0, min(self._busy_counter.value, self._concurrency)) - occupancy = busy / self._concurrency if self._concurrency else 0.0 - self._metrics.gauge( - "taskworker.worker.occupancy", - occupancy, - tags=tags, - ) - if self._prom is not None: - self._prom.occupancy.labels(processing_pool=self._processing_pool_name).set( - occupancy - ) - - # 'qsize' is not implemented on all platforms, such as macOS - self._metrics.gauge( - "taskworker.child_tasks.size", - float(self._child_tasks.qsize()), - tags=tags, - ) - - self._metrics.gauge( - "taskworker.processed_tasks.size", - float(self._processed_tasks.qsize()), - tags=tags, - ) + self._emit_periodic_metrics() except Exception as e: logger.debug( "taskworker.worker.queue_gauges.error", @@ -965,16 +1029,110 @@ 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 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 + + # This child is now running + if message.event == "running": + child.state = "running" + + # 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) + + # This child is executing a task + elif message.event == "busy": + child.busy = True + + # This child isn't doing anything right now + elif message.event == "idle": + child.busy = False + + 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" + 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() + 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, @@ -984,21 +1142,61 @@ def spawn_children_thread() -> None: self._process_type, self._skip_awaiting_futures, self._future_checking_frequency, - self._ready_counter, - self._busy_counter, + messages, + release, ), ) - process.start() - self._children.append(process) + + try: + process.start() + + with self._children_lock: + child = TrackedChild( + process=process, + state="pending", + release=release, + ) + + self._children[child_id] = child + 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 + 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 ) @@ -1055,9 +1253,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 1483b492..86eb6905 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,21 +165,14 @@ def _log_task_retry_exhausted( ) -def _adjust_busy(counter: "Synchronized[int] | None", delta: int) -> None: - """ - Adjust the shared count of children currently executing a task. - - The parent worker pool divides this by concurrency to emit occupancy, the - autoscaling signal. A child that is hard-killed (e.g. OOM) mid-task leaks - its increment; the parent clamps occupancy to [0, 1] to bound the drift. - """ - if counter is None: - return - with counter.get_lock(): - counter.value += delta +@dataclass(frozen=True) +class ChildMessage: + child_id: UUID + event: Literal["running", "exiting", "busy", "idle"] def child_process( + child_id: UUID, app_module: str, child_tasks: queue.Queue[InflightTaskActivation], processed_tasks: queue.Queue[ProcessingResult], @@ -190,8 +182,8 @@ def child_process( process_type: str, skip_awaiting_futures: bool, future_checking_frequency: float, - ready_counter: "Synchronized[int] | None" = None, - busy_counter: "Synchronized[int] | None" = None, + messages: multiprocessing.Queue[ChildMessage], + parent_release: Event, ) -> None: """ The entrypoint for spawned worker children. @@ -375,18 +367,37 @@ def check_task_future_completion( ) _future_completion_thread.start() + # Track when exit was initiated for metrics + exit_initiated: float | None = None + 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 exit_initiated is None: + 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")) + exit_initiated = time.monotonic() + + if parent_release.is_set(): + metrics.distribution( + "taskworker.worker.child.parent_release_wait_duration", + time.monotonic() - exit_initiated, + tags={"processing_pool": processing_pool_name}, + ) + + local_shutdown.set() + break try: inflight = child_tasks.get(timeout=1.0) @@ -453,7 +464,7 @@ def check_task_future_completion( next_state = TASK_ACTIVATION_STATUS_FAILURE # Use time.time() so we can measure against activation.received_at execution_start_time = time.time() - _adjust_busy(busy_counter, 1) + messages.put_nowait(ChildMessage(child_id, "busy")) try: with timeout_alarm(inflight.activation.processing_deadline_duration, handle_alarm): _execute_activation(task_func, inflight.activation, app.context_hooks) @@ -526,7 +537,7 @@ def check_task_future_completion( ): _log_task_failed(inflight.activation, err, processing_pool_name) finally: - _adjust_busy(busy_counter, -1) + messages.put_nowait(ChildMessage(child_id, "idle")) clear_current_task() processed_task_count += 1 @@ -833,11 +844,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, "running")) # 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 d06f1d02..e3653905 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,38 @@ 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, + future_checking_frequency: float, +) -> 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, + future_checking_frequency, + messages, + parent_release, + ) + + class _SendResultCapture: def __init__(self) -> None: self.send_calls: list[tuple[list[ProcessingResult], bool]] = [] @@ -302,6 +337,89 @@ 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, + min_concurrency: int = 0, +) -> 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, + min_concurrency=min_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 +477,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 +517,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 +779,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 +806,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 +820,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 +857,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 +891,107 @@ 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_above_min_concurrency() -> None: + fake_context = _FakeContext() + pool = _make_fake_context_pool(fake_context, concurrency=2, min_concurrency=1) + + pool.start_spawn_children_thread() + try: + _wait_for(lambda: len(fake_context.processes) == 2) + 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, "running")) + second_process = fake_context.processes[1] + second_child_id = second_process.args[0] + messages.put(ChildMessage(second_child_id, "running")) + _wait_for(lambda: pool.ready_count == 2) + + messages.put(ChildMessage(first_child_id, "exiting")) + _wait_for(first_release.is_set) + + _wait_for(lambda: len(fake_context.processes) == 3) + finally: + pool.shutdown() + + +def test_spawn_children_defers_draining_child_at_min_concurrency() -> None: + fake_context = _FakeContext() + pool = _make_fake_context_pool(fake_context, concurrency=2, min_concurrency=1) + + pool.start_spawn_children_thread() + try: + _wait_for(lambda: len(fake_context.processes) == 2) + messages = fake_context.queues[-1] + first_process = fake_context.processes[0] + first_child_id = first_process.args[0] + first_release = first_process.args[-1] + + second_process = fake_context.processes[1] + second_child_id = second_process.args[0] + + messages.put(ChildMessage(first_child_id, "running")) + _wait_for(lambda: pool.ready_count == 1) + + messages.put(ChildMessage(first_child_id, "exiting")) + time.sleep(0.25) + assert not first_release.is_set() + + messages.put(ChildMessage(second_child_id, "running")) + _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( @@ -894,15 +1118,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_running_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, @@ -912,22 +1140,91 @@ def test_child_process_increments_ready_counter() -> None: process_type="fork", skip_awaiting_futures=False, future_checking_frequency=0.1, - 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, "running") + + +@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, + 0.1, + messages, + parent_release, + ), ) + process.start() + try: + running_message = messages.get(timeout=5) + busy_message = messages.get(timeout=5) + idle_message = messages.get(timeout=5) + exiting_message = messages.get(timeout=5) + + assert running_message == ChildMessage(child_id, "running") + assert busy_message == ChildMessage(child_id, "busy") + assert idle_message == ChildMessage(child_id, "idle") + assert exiting_message == ChildMessage(child_id, "exiting") + assert processed.get(timeout=5).task_id == SIMPLE_TASK.activation.id + + todo.put(SIMPLE_TASK) + assert processed.get(timeout=5).task_id == SIMPLE_TASK.activation.id + assert messages.get(timeout=5) == ChildMessage(child_id, "busy") + assert messages.get(timeout=5) == ChildMessage(child_id, "idle") - # The child increments the counter once warmup is done, before consuming. - assert ready_counter.value == 1 + 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_busy_counter_returns_to_zero() -> None: +def test_child_process_emits_busy_and_idle_messages() -> None: todo: queue.Queue[InflightTaskActivation] = queue.Queue() processed: queue.Queue[ProcessingResult] = queue.Queue() shutdown = Event() ctx = get_context("fork") - busy_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, @@ -937,12 +1234,14 @@ def test_child_process_busy_counter_returns_to_zero() -> None: process_type="fork", skip_awaiting_futures=False, future_checking_frequency=0.1, - busy_counter=busy_counter, + messages=messages, + parent_release=parent_release, ) - # Incremented while executing, decremented in the finally afterwards, so the - # slot is released back to idle once the task is done. - assert busy_counter.value == 0 + assert messages.get(timeout=1) == ChildMessage(child_id, "running") + assert messages.get(timeout=1) == ChildMessage(child_id, "busy") + assert messages.get(timeout=1) == ChildMessage(child_id, "idle") + assert processed.get(timeout=1).task_id == SIMPLE_TASK.activation.id def test_child_process_remove_start_time_kwargs() -> None: