diff --git a/clients/python/src/taskbroker_client/constants.py b/clients/python/src/taskbroker_client/constants.py index 817a70fb..ad157364 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 all 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..690b6514 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,7 @@ def __init__( push_task_timeout: float = 5, update_in_batches: bool = False, skip_awaiting_futures: bool = True, + warmup_timeout: float = DEFAULT_WORKER_WARMUP_TIMEOUT_SEC, ) -> None: app = import_app(app_module) @@ -202,6 +204,8 @@ def __init__( self._grpc_secrets = parse_rpc_secret_list(app.config["rpc_secret"]) self._push_task_timeout = push_task_timeout + self._warmup_timeout = warmup_timeout + def _create_client( self, service: str, @@ -310,6 +314,56 @@ 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 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. + """ + 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 < required: + 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, + "required": required, + "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, + "required": required, + "timed_out": timed_out, + }, + ) + def start(self) -> int: """ This starts the worker gRPC server. @@ -364,6 +418,15 @@ 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() + + # 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) @@ -739,10 +802,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 +935,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..bcf6186a 100644 --- a/clients/python/tests/worker/test_worker.py +++ b/clients/python/tests/worker/test_worker.py @@ -653,6 +653,121 @@ 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_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: + 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, 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_unblocks_when_children_warm() -> None: + taskworker = _make_push_worker(concurrency=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 == [] + + +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( @@ -778,6 +893,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",