feat(worker): defer SERVING until child processes are warm

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
 """

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()

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,

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",