From 8c73b6f6f37d0f029edf8b509abc83a95d4dca9b Mon Sep 17 00:00:00 2001 From: Enoch Tang Date: Thu, 25 Jun 2026 12:59:08 -0400 Subject: [PATCH 1/3] Defer SERVING until child processes are warm --- .../python/src/taskbroker_client/constants.py | 6 + .../src/taskbroker_client/worker/worker.py | 68 +++++++++ .../taskbroker_client/worker/workerchild.py | 12 +- clients/python/tests/worker/test_worker.py | 135 ++++++++++++++++++ 4 files changed, 220 insertions(+), 1 deletion(-) diff --git a/clients/python/src/taskbroker_client/constants.py b/clients/python/src/taskbroker_client/constants.py index 817a70fb..7d5cd08c 100644 --- a/clients/python/src/taskbroker_client/constants.py +++ b/clients/python/src/taskbroker_client/constants.py @@ -57,6 +57,12 @@ The number of gRPC requests before touching the health check file """ +DEFAULT_WORKER_WARMUP_TIMEOUT_SEC = 90.0 +""" +Max seconds PushTaskWorker waits for >= min_ready children to warm up +before flipping gRPC health to SERVING anyway. +""" + ALWAYS_EAGER = False """ diff --git a/clients/python/src/taskbroker_client/worker/worker.py b/clients/python/src/taskbroker_client/worker/worker.py index 1ae6aa3d..6ce61222 100644 --- a/clients/python/src/taskbroker_client/worker/worker.py +++ b/clients/python/src/taskbroker_client/worker/worker.py @@ -28,6 +28,7 @@ DEFAULT_REBALANCE_AFTER, DEFAULT_WORKER_HEALTH_CHECK_SEC_PER_TOUCH, DEFAULT_WORKER_QUEUE_SIZE, + DEFAULT_WORKER_WARMUP_TIMEOUT_SEC, MAX_BACKOFF_SECONDS_WHEN_HOST_UNAVAILABLE, WORKER_CHILD_JOIN_TIMEOUT_SEC, ) @@ -143,6 +144,8 @@ def __init__( push_task_timeout: float = 5, update_in_batches: bool = False, skip_awaiting_futures: bool = True, + min_ready: int | None = None, + warmup_timeout: float = DEFAULT_WORKER_WARMUP_TIMEOUT_SEC, ) -> None: app = import_app(app_module) @@ -202,6 +205,9 @@ def __init__( self._grpc_secrets = parse_rpc_secret_list(app.config["rpc_secret"]) self._push_task_timeout = push_task_timeout + self._min_ready = concurrency if min_ready is None else min_ready + self._warmup_timeout = warmup_timeout + def _create_client( self, service: str, @@ -310,6 +316,57 @@ def _stop_health_check_thread(self) -> None: self._health_check_thread.join(timeout=5) self._health_check_thread = None + def _await_children_warm(self) -> None: + """ + Block until at least min_ready children have warmed up or warmup_timeout elapses. + + On timeout we fall through and serve anyway, a degraded-but-routable pod + beats one that never becomes ready. The min_ready is clamped to concurrency and a value + <= 0 disables the gate entirely. + """ + min_ready = min(self._min_ready, self._concurrency) + if min_ready <= 0: + return + + warmup_start = time.monotonic() + deadline = warmup_start + self._warmup_timeout + timed_out = False + while self.worker_pool.ready_count < min_ready: + if time.monotonic() >= deadline: + timed_out = True + self._metrics.incr( + "taskworker.worker.warmup_timeout", + tags={"processing_pool": self._processing_pool_name}, + ) + logger.warning( + "taskworker.worker.warmup_timeout", + extra={ + "processing_pool": self._processing_pool_name, + "ready_count": self.worker_pool.ready_count, + "min_ready": min_ready, + "warmup_timeout": self._warmup_timeout, + }, + ) + break + # Sleep and break early if shutdown was requested via shutdown(). + if self._grpc_sync_event.wait(0.25): + break + + self._metrics.distribution( + "taskworker.worker.warmup_duration", + time.monotonic() - warmup_start, + tags={"processing_pool": self._processing_pool_name}, + ) + logger.info( + "taskworker.worker.warmup_complete", + extra={ + "processing_pool": self._processing_pool_name, + "ready_count": self.worker_pool.ready_count, + "min_ready": min_ready, + "timed_out": timed_out, + }, + ) + def start(self) -> int: """ This starts the worker gRPC server. @@ -364,6 +421,10 @@ def signal_handler(*args: Any) -> None: server.add_insecure_port(f"[::]:{self._grpc_port}") server.start() + # Hold NOT_SERVING until children are warm so the pod stays out of + # the NEG/readiness set while its child processes are still loading. + self._await_children_warm() + # Indicate that the server is ready health_servicer.set("", health_pb2.HealthCheckResponse.SERVING) health_servicer.set(WORKER_SERVICE_NAME, health_pb2.HealthCheckResponse.SERVING) @@ -739,10 +800,16 @@ def __init__( ) self._children: list[BaseProcess] = [] 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 + @property + def ready_count(self) -> int: + """Number of children that have finished warming up and are consuming.""" + return self._ready_counter.value + def send_results(self, results: list[ProcessingResult], is_draining: bool = False) -> None: """ Call the passed in function. If is_draining is True, the function should not fetch a new task. @@ -866,6 +933,7 @@ def spawn_children_thread() -> None: self._processing_pool_name, self._process_type, self._skip_awaiting_futures, + self._ready_counter, ), ) process.start() diff --git a/clients/python/src/taskbroker_client/worker/workerchild.py b/clients/python/src/taskbroker_client/worker/workerchild.py index c2c658b3..5395920c 100644 --- a/clients/python/src/taskbroker_client/worker/workerchild.py +++ b/clients/python/src/taskbroker_client/worker/workerchild.py @@ -11,7 +11,10 @@ from functools import partial from multiprocessing.synchronize import Event from types import FrameType -from typing import Any +from typing import TYPE_CHECKING, Any + +if TYPE_CHECKING: + from multiprocessing.sharedctypes import Synchronized # XXX: Don't import any modules that will import django here, do those within child_process import msgpack @@ -172,6 +175,7 @@ def child_process( processing_pool_name: str, process_type: str, skip_awaiting_futures: bool, + ready_counter: "Synchronized[int] | None" = None, ) -> None: """ The entrypoint for spawned worker children. @@ -802,6 +806,12 @@ 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 + # Run the worker loop run_worker( child_tasks, diff --git a/clients/python/tests/worker/test_worker.py b/clients/python/tests/worker/test_worker.py index 38940875..d1335aee 100644 --- a/clients/python/tests/worker/test_worker.py +++ b/clients/python/tests/worker/test_worker.py @@ -653,6 +653,117 @@ def test_batch_push_worker_health_check_touches_while_idle(tmp_path: Path) -> No assert taskworker._health_check_thread is None +def _make_push_worker(**kwargs: Any) -> PushTaskWorker: + return PushTaskWorker( + app_module="examples.app:app", + broker_service="127.0.0.1:50051", + max_child_task_count=100, + process_type="fork", + **kwargs, + ) + + +def test_min_ready_defaults_to_concurrency() -> None: + taskworker = _make_push_worker(concurrency=8) + assert taskworker._min_ready == 8 + + +def test_await_children_warm_returns_when_ready() -> None: + taskworker = _make_push_worker(concurrency=4, min_ready=4, warmup_timeout=5) + taskworker.worker_pool._ready_counter.value = 4 + + with mock.patch.object(taskworker, "_metrics") as mock_metrics: + start = time.time() + taskworker._await_children_warm() + elapsed = time.time() - start + + assert elapsed < 1 + # Records warmup duration, but no timeout. + timeout_calls = [ + c + for c in mock_metrics.incr.call_args_list + if c.args[0] == "taskworker.worker.warmup_timeout" + ] + assert timeout_calls == [] + mock_metrics.distribution.assert_any_call( + "taskworker.worker.warmup_duration", mock.ANY, tags=mock.ANY + ) + + +def test_await_children_warm_times_out() -> None: + taskworker = _make_push_worker(concurrency=4, min_ready=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() + taskworker._await_children_warm() + elapsed = time.time() - start + + assert elapsed >= 0.1 + mock_metrics.incr.assert_any_call( + "taskworker.worker.warmup_timeout", tags={"processing_pool": "unknown"} + ) + + +def test_await_children_warm_clamps_min_ready_to_concurrency() -> None: + # min_ready exceeds concurrency; without clamping this would always time out. + taskworker = _make_push_worker(concurrency=4, min_ready=1000, warmup_timeout=0.1) + taskworker.worker_pool._ready_counter.value = 4 + + with mock.patch.object(taskworker, "_metrics") as mock_metrics: + start = time.time() + taskworker._await_children_warm() + elapsed = time.time() - start + + assert elapsed < 1 + timeout_calls = [ + c + for c in mock_metrics.incr.call_args_list + if c.args[0] == "taskworker.worker.warmup_timeout" + ] + assert timeout_calls == [] + + +def test_await_children_warm_disabled_when_min_ready_zero() -> None: + taskworker = _make_push_worker(concurrency=4, min_ready=0, warmup_timeout=10) + # Counter stays at 0; gate is disabled so this must return immediately. + taskworker.worker_pool._ready_counter.value = 0 + + start = time.time() + taskworker._await_children_warm() + elapsed = time.time() - start + + assert elapsed < 1 + + +def test_await_children_warm_unblocks_when_children_warm() -> None: + taskworker = _make_push_worker(concurrency=2, min_ready=2, warmup_timeout=5) + taskworker.worker_pool._ready_counter.value = 0 + + def warm_up() -> None: + time.sleep(0.2) + taskworker.worker_pool._ready_counter.value = 2 + + warmer = threading.Thread(target=warm_up) + warmer.start() + try: + with mock.patch.object(taskworker, "_metrics") as mock_metrics: + start = time.time() + taskworker._await_children_warm() + elapsed = time.time() - start + finally: + warmer.join() + + assert 0.2 <= elapsed < 5 + timeout_calls = [ + c + for c in mock_metrics.incr.call_args_list + if c.args[0] == "taskworker.worker.warmup_timeout" + ] + assert timeout_calls == [] + + class TestWorkerServicer(TestCase): def test_push_task_success(self) -> None: taskworker = PushTaskWorker( @@ -778,6 +889,30 @@ 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: + todo: queue.Queue[InflightTaskActivation] = queue.Queue() + processed: queue.Queue[ProcessingResult] = queue.Queue() + shutdown = Event() + ctx = get_context("fork") + ready_counter = ctx.Value("i", 0) + + todo.put(SIMPLE_TASK) + child_process( + "examples.app:app", + todo, + processed, + shutdown, + max_task_count=1, + processing_pool_name="test", + process_type="fork", + skip_awaiting_futures=False, + ready_counter=ready_counter, + ) + + # The child increments the counter once warmup is done, before consuming. + assert ready_counter.value == 1 + + def test_child_process_remove_start_time_kwargs() -> None: activation = InflightTaskActivation( host="localhost:50051", From a27b15e3d6a97f29250f84560446e2608c1ffec8 Mon Sep 17 00:00:00 2001 From: Enoch Tang Date: Thu, 25 Jun 2026 15:04:53 -0400 Subject: [PATCH 2/3] dont advertise SERVING after shutdown event --- clients/python/src/taskbroker_client/worker/worker.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/clients/python/src/taskbroker_client/worker/worker.py b/clients/python/src/taskbroker_client/worker/worker.py index 6ce61222..95a6a42e 100644 --- a/clients/python/src/taskbroker_client/worker/worker.py +++ b/clients/python/src/taskbroker_client/worker/worker.py @@ -425,6 +425,11 @@ def signal_handler(*args: Any) -> None: # the NEG/readiness set while its child processes are still loading. self._await_children_warm() + # If shutdown was requested during warmup, don't advertise SERVING. + # Bail to the finally below, which sets NOT_SERVING and tears everything down. + if self._grpc_sync_event.is_set(): + return 0 + # Indicate that the server is ready health_servicer.set("", health_pb2.HealthCheckResponse.SERVING) health_servicer.set(WORKER_SERVICE_NAME, health_pb2.HealthCheckResponse.SERVING) From 6841f94016d13808c691567c3cc02cb92b2fe174 Mon Sep 17 00:00:00 2001 From: Enoch Tang Date: Thu, 25 Jun 2026 16:15:02 -0400 Subject: [PATCH 3/3] remove min_ready --- .../python/src/taskbroker_client/constants.py | 4 +- .../src/taskbroker_client/worker/worker.py | 17 ++-- clients/python/tests/worker/test_worker.py | 82 ++++++++++--------- 3 files changed, 52 insertions(+), 51 deletions(-) diff --git a/clients/python/src/taskbroker_client/constants.py b/clients/python/src/taskbroker_client/constants.py index 7d5cd08c..ad157364 100644 --- a/clients/python/src/taskbroker_client/constants.py +++ b/clients/python/src/taskbroker_client/constants.py @@ -59,8 +59,8 @@ DEFAULT_WORKER_WARMUP_TIMEOUT_SEC = 90.0 """ -Max seconds PushTaskWorker waits for >= min_ready children to warm up -before flipping gRPC health to SERVING anyway. +Max seconds PushTaskWorker waits for all children to warm up before +flipping gRPC health to SERVING anyway. """ diff --git a/clients/python/src/taskbroker_client/worker/worker.py b/clients/python/src/taskbroker_client/worker/worker.py index 95a6a42e..690b6514 100644 --- a/clients/python/src/taskbroker_client/worker/worker.py +++ b/clients/python/src/taskbroker_client/worker/worker.py @@ -144,7 +144,6 @@ def __init__( push_task_timeout: float = 5, update_in_batches: bool = False, skip_awaiting_futures: bool = True, - min_ready: int | None = None, warmup_timeout: float = DEFAULT_WORKER_WARMUP_TIMEOUT_SEC, ) -> None: app = import_app(app_module) @@ -205,7 +204,6 @@ def __init__( self._grpc_secrets = parse_rpc_secret_list(app.config["rpc_secret"]) self._push_task_timeout = push_task_timeout - self._min_ready = concurrency if min_ready is None else min_ready self._warmup_timeout = warmup_timeout def _create_client( @@ -318,20 +316,19 @@ def _stop_health_check_thread(self) -> None: def _await_children_warm(self) -> None: """ - Block until at least min_ready children have warmed up or warmup_timeout elapses. + Block until all children have warmed up or warmup_timeout elapses. On timeout we fall through and serve anyway, a degraded-but-routable pod - beats one that never becomes ready. The min_ready is clamped to concurrency and a value - <= 0 disables the gate entirely. + beats one that never becomes ready. """ - min_ready = min(self._min_ready, self._concurrency) - if min_ready <= 0: + required = self._concurrency + if required <= 0: return warmup_start = time.monotonic() deadline = warmup_start + self._warmup_timeout timed_out = False - while self.worker_pool.ready_count < min_ready: + while self.worker_pool.ready_count < required: if time.monotonic() >= deadline: timed_out = True self._metrics.incr( @@ -343,7 +340,7 @@ def _await_children_warm(self) -> None: extra={ "processing_pool": self._processing_pool_name, "ready_count": self.worker_pool.ready_count, - "min_ready": min_ready, + "required": required, "warmup_timeout": self._warmup_timeout, }, ) @@ -362,7 +359,7 @@ def _await_children_warm(self) -> None: extra={ "processing_pool": self._processing_pool_name, "ready_count": self.worker_pool.ready_count, - "min_ready": min_ready, + "required": required, "timed_out": timed_out, }, ) diff --git a/clients/python/tests/worker/test_worker.py b/clients/python/tests/worker/test_worker.py index d1335aee..bcf6186a 100644 --- a/clients/python/tests/worker/test_worker.py +++ b/clients/python/tests/worker/test_worker.py @@ -663,13 +663,8 @@ def _make_push_worker(**kwargs: Any) -> PushTaskWorker: ) -def test_min_ready_defaults_to_concurrency() -> None: - taskworker = _make_push_worker(concurrency=8) - assert taskworker._min_ready == 8 - - def test_await_children_warm_returns_when_ready() -> None: - taskworker = _make_push_worker(concurrency=4, min_ready=4, warmup_timeout=5) + 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: @@ -691,7 +686,7 @@ def test_await_children_warm_returns_when_ready() -> None: def test_await_children_warm_times_out() -> None: - taskworker = _make_push_worker(concurrency=4, min_ready=4, warmup_timeout=0.1) + taskworker = _make_push_worker(concurrency=4, warmup_timeout=0.1) # Never becomes ready. taskworker.worker_pool._ready_counter.value = 0 @@ -706,39 +701,8 @@ def test_await_children_warm_times_out() -> None: ) -def test_await_children_warm_clamps_min_ready_to_concurrency() -> None: - # min_ready exceeds concurrency; without clamping this would always time out. - taskworker = _make_push_worker(concurrency=4, min_ready=1000, warmup_timeout=0.1) - taskworker.worker_pool._ready_counter.value = 4 - - with mock.patch.object(taskworker, "_metrics") as mock_metrics: - start = time.time() - taskworker._await_children_warm() - elapsed = time.time() - start - - assert elapsed < 1 - timeout_calls = [ - c - for c in mock_metrics.incr.call_args_list - if c.args[0] == "taskworker.worker.warmup_timeout" - ] - assert timeout_calls == [] - - -def test_await_children_warm_disabled_when_min_ready_zero() -> None: - taskworker = _make_push_worker(concurrency=4, min_ready=0, warmup_timeout=10) - # Counter stays at 0; gate is disabled so this must return immediately. - taskworker.worker_pool._ready_counter.value = 0 - - start = time.time() - taskworker._await_children_warm() - elapsed = time.time() - start - - assert elapsed < 1 - - def test_await_children_warm_unblocks_when_children_warm() -> None: - taskworker = _make_push_worker(concurrency=2, min_ready=2, warmup_timeout=5) + taskworker = _make_push_worker(concurrency=2, warmup_timeout=5) taskworker.worker_pool._ready_counter.value = 0 def warm_up() -> None: @@ -764,6 +728,46 @@ def warm_up() -> None: assert timeout_calls == [] +def test_start_does_not_serve_when_shutdown_during_warmup() -> None: + from grpc_health.v1 import health_pb2 + + 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() + fake_server = mock.MagicMock() + + with ( + mock.patch.object(taskworker.worker_pool, "start_metrics_thread"), + mock.patch.object(taskworker.worker_pool, "start_result_thread"), + mock.patch.object(taskworker.worker_pool, "start_spawn_children_thread"), + mock.patch.object(taskworker.worker_pool, "shutdown"), + mock.patch("taskbroker_client.worker.worker.grpc.server", return_value=fake_server), + mock.patch( + "taskbroker_client.worker.worker.health.HealthServicer", return_value=fake_health + ), + mock.patch("taskbroker_client.worker.worker.health_pb2_grpc.add_HealthServicer_to_server"), + mock.patch( + "taskbroker_client.worker.worker.taskbroker_pb2_grpc" + ".add_WorkerServiceServicer_to_server" + ), + ): + exitcode = taskworker.start() + + assert exitcode == 0 + # Health must never have been flipped to SERVING. + serving_calls = [ + c + for c in fake_health.set.call_args_list + if c.args[1] == health_pb2.HealthCheckResponse.SERVING + ] + assert serving_calls == [] + # We never reached server.wait_for_termination() (returned before it). + fake_server.wait_for_termination.assert_not_called() + + class TestWorkerServicer(TestCase): def test_push_task_success(self) -> None: taskworker = PushTaskWorker(