From ce1ab998875a2eb61df8b348998b2a075d7bda0d Mon Sep 17 00:00:00 2001 From: Rahul Johny Date: Mon, 5 Oct 2026 12:34:53 +0530 Subject: [PATCH 1/2] UN-4223 [FIX] Kill a stuck PG consumer child instead of restarting the whole pod One hung task (an LLM call that never returns) froze its child's heartbeat. The fleet probe reported the OLDEST child's age, so that single child failed liveness, and the container restart then waited out the full termination grace for the hung call, leaving the pod consuming nothing for up to ~2h. - The supervisor now SIGKILLs a child whose heartbeat stays frozen past the stuck-child cap and re-forks it. Its lease stops renewing, so the reaper redelivers the message (bounded by max_attempts). The cap defaults to HEALTH_STALE_SECONDS (the existing per-task bound) and is overridable via WORKER_PG_QUEUE_CONSUMER_STUCK_CHILD_SECONDS; with no health port it is off unless set explicitly. - /health now goes 503 only once at least half the children are stale. The oldest child's age and the stale-child count remain in the body. Co-Authored-By: Claude Opus 5.5 --- workers/pg_queue_consumer/supervisor.py | 152 +++++++++++++++-- workers/queue_backend/pg_queue/metrics.py | 5 +- workers/tests/test_pg_consumer_supervisor.py | 161 ++++++++++++++++++- 3 files changed, 302 insertions(+), 16 deletions(-) diff --git a/workers/pg_queue_consumer/supervisor.py b/workers/pg_queue_consumer/supervisor.py index 2474983083..e99b5a5345 100644 --- a/workers/pg_queue_consumer/supervisor.py +++ b/workers/pg_queue_consumer/supervisor.py @@ -17,12 +17,20 @@ path (byte-identical to before this module existed). **Health**: the supervisor owns the single liveness port and reports the *fleet's* -freshness — the staleness of the oldest-polling child (each child publishes its -last-poll wall-time into a shared array). A child that dies is re-forked +freshness — the age that at least half the children have gone without polling +(each child publishes its last-poll wall-time into a shared array), so the probe +trips only when half the fleet is wedged. A child that dies is re-forked internally (transient); a child that **crash-loops** (dies immediately N times in a row, never reaching a real poll) forces the probe to 503 so k8s restarts the pod rather than the supervisor masking a wedged fleet with fresh-looking re-forks. +**Stuck children** (UN-4223): a child whose heartbeat stays frozen past the +stuck-child cap (default: the health stale window) is stuck on one task, e.g. an +LLM call that never returns. The supervisor SIGKILLs and re-forks just that +child. Before this, the probe tracked the OLDEST child, so one hung call failed +liveness and the container restart then waited out the full termination grace +for that call, leaving the pod consuming nothing for up to ~2h. + **Readiness** (UN-4136): the same port serves ``/ready``, which answers 200 only once EVERY child has finished its ``import worker`` bootstrap and built its consumer. ``/health`` cannot say this — the heartbeats are seeded fresh at @@ -72,6 +80,10 @@ _MIN_HEALTHY_UPTIME_SECONDS = 10.0 # Consecutive immediate crashes after which the fleet probe is forced unhealthy. _CRASH_LOOP_THRESHOLD = 3 +# Fraction of children that must be stale before the fleet probe goes 503. One +# stuck child is the supervisor's to kill (see _kill_stuck_children), not a reason +# to restart the whole container and drain every healthy sibling with it (UN-4223). +_UNHEALTHY_STALE_FRACTION = 0.5 # Fallback graceful-drain budget (s, shared across all children) on shutdown, used # only when neither an explicit override nor the consumer VT is set — see # shutdown_grace_from_env(). @@ -142,6 +154,40 @@ def shutdown_grace_from_env() -> float: return max(_DEFAULT_SHUTDOWN_GRACE_SECONDS, float(vt)) +def stuck_child_seconds_from_env() -> float | None: + """Heartbeat age (seconds) past which the supervisor SIGKILLs a child as stuck, + or ``None`` to never kill. + + A child's heartbeat is frozen while it runs a task, so this is a per-task + wall-clock cap. It defaults to ``WORKER_PG_QUEUE_CONSUMER_HEALTH_STALE_SECONDS``, + the threshold that is already documented as the upper bound on one task and + that, before UN-4223, restarted the WHOLE container when a single child crossed + it. Killing just that child keeps the same bound without draining its healthy + siblings. ``WORKER_PG_QUEUE_CONSUMER_STUCK_CHILD_SECONDS`` overrides it. + + With no health port configured there was no probe to enforce the bound, so + nothing is killed unless the override is set explicitly. + """ + from queue_backend.pg_queue.consumer import ( + _DEFAULT_HEALTH_STALE_SECONDS, + consumer_env, + ) + + override: float | None = consumer_env("STUCK_CHILD_SECONDS", None, float) + if override is not None: + # 0 / negative would kill every child mid-task on the next tick; inf / nan + # would silently disable the cap. Fail fast at startup instead. + if not math.isfinite(override) or override <= 0: + raise ValueError( + "WORKER_PG_QUEUE_CONSUMER_STUCK_CHILD_SECONDS must be a finite " + f"number > 0, got {override!r}" + ) + return override + if consumer_env("HEALTH_PORT", None, int) is None: + return None + return consumer_env("HEALTH_STALE_SECONDS", _DEFAULT_HEALTH_STALE_SECONDS, float) + + class _Fleet: """Owns the per-slot child state — pid, last-fork, heartbeat, crash count and pending-restart schedule — keeping them mutually consistent. Slots are @@ -169,6 +215,9 @@ def __init__(self, concurrency: int) -> None: self._last_fork: dict[int, float] = {} self._consecutive_crashes: dict[int, int] = {} self._restart_due: dict[int, float] = {} # slot -> monotonic not-before + # Slots already SIGKILLed as stuck, so a child that takes a tick to die + # is not re-killed and re-logged on every monitor pass. Cleared in reap(). + self._stuck_killed: set[int] = set() @property def concurrency(self) -> int: @@ -212,6 +261,7 @@ def reap(self, slot: int) -> float: forked_at = self._last_fork.pop(slot, time.monotonic()) self._pids.pop(slot, None) self._loaded[slot] = 0 + self._stuck_killed.discard(slot) return time.monotonic() - forked_at def schedule_restart(self, slot: int, uptime: float) -> int: @@ -264,17 +314,60 @@ def all_loaded(self) -> bool: """Readiness verdict source: True once every slot's child has loaded.""" return self.loaded_count() == self._n - def oldest_age(self) -> float: + def _ages(self) -> list[float]: + """Every slot's heartbeat age (seconds), oldest first.""" now = time.time() - return max((now - hb for hb in self._heartbeats), default=0.0) + return sorted((now - hb for hb in self._heartbeats), reverse=True) + + def oldest_age(self) -> float: + return max(self._ages(), default=0.0) + + def quorum_age(self) -> float: + """The age that at least ``_UNHEALTHY_STALE_FRACTION`` of the children have + reached: past the stale window only once that share of the fleet is stale. + With one child this is that child's age, as before. + """ + ages = self._ages() + if not ages: + return 0.0 + quorum = max(1, math.ceil(len(ages) * _UNHEALTHY_STALE_FRACTION)) + return ages[quorum - 1] + + def stale_count(self, stale_after: float) -> int: + return sum(1 for age in self._ages() if age > stale_after) + + def stuck_children(self, stuck_after: float) -> list[tuple[int, int, float]]: + """``(slot, pid, heartbeat_age)`` for each live child stuck past + ``stuck_after`` and not yet killed. + + A re-forked child's heartbeat is NOT reseeded (see :meth:`record_fork`), so + it carries its predecessor's age until it first polls. Requiring the child + itself to have been up for ``stuck_after`` stops a fresh replacement from + being killed mid-bootstrap for its predecessor's staleness. + """ + now_wall, now_mono = time.time(), time.monotonic() + stuck = [] + for slot, pid in self.alive_items(): + if slot in self._stuck_killed: + continue + age = now_wall - self._heartbeats[slot] + uptime = now_mono - self._last_fork.get(slot, now_mono) + if age > stuck_after and uptime > stuck_after: + stuck.append((slot, pid, age)) + return stuck + + def mark_stuck_killed(self, slot: int) -> None: + self._validate(slot) + self._stuck_killed.add(slot) def freshness(self) -> float: """Liveness verdict source: a crash-looping fleet is force-stale (``inf``) so the probe trips 503 even if a just-constructed child briefly looked - fresh; otherwise the oldest child's staleness (catches a wedged-alive - child). + fresh; otherwise :meth:`quorum_age`, so the pod is restarted only when at + least half its children are wedged. A single wedged child is killed and + re-forked by the supervisor instead (:func:`_kill_stuck_children`). """ - return float("inf") if self.is_crash_looping() else self.oldest_age() + return float("inf") if self.is_crash_looping() else self.quorum_age() def _run_child(slot: int, heartbeats, loaded) -> None: # noqa: ANN001 (ctypes arrays) @@ -393,6 +486,36 @@ def _reap_dead(fleet: _Fleet, stopping: threading.Event) -> None: ) +def _kill_stuck_children( + fleet: _Fleet, stuck_after: float | None, stopping: threading.Event +) -> None: + """SIGKILL each child whose heartbeat has been frozen past ``stuck_after``. + + A child's heartbeat only advances between tasks, so a frozen one means a single + task has run past its cap (e.g. an LLM call that never returns). Killing it + frees the slot: ``_reap_dead`` reaps it on the next tick and re-forks it, and + its message's lease is no longer renewed, so it lapses and the reaper + redelivers it (bounded by ``max_attempts``). SIGTERM would not do: the child + only acts on it between tasks, which is exactly what a stuck one never reaches. + + Skipped during shutdown, where ``_join_children`` owns the children. + """ + if stuck_after is None or stopping.is_set(): + return + for slot, pid, age in fleet.stuck_children(stuck_after): + logger.error( + "PG-queue consumer: child slot=%s pid=%s has been on one task for " + "%.0fs (cap %.0fs) — SIGKILL; its message will be redelivered", + slot, + pid, + age, + stuck_after, + ) + with contextlib.suppress(ProcessLookupError): + os.kill(pid, signal.SIGKILL) + fleet.mark_stuck_killed(slot) + + def _restart_due_children(fleet: _Fleet, stopping: threading.Event) -> None: """Re-fork the slots whose backoff has elapsed — non-blocking (the backoff is a scheduled not-before, not an in-loop sleep), and re-checking ``stopping`` @@ -412,10 +535,12 @@ def run_supervised(concurrency: int) -> None: fleet = _Fleet(concurrency) grace_seconds = shutdown_grace_from_env() + stuck_after = stuck_child_seconds_from_env() logger.info( "PG-queue consumer supervisor: shutdown drain grace = %.0fs (shared across " - "children)", + "children); stuck-child cap = %s", grace_seconds, + "off" if stuck_after is None else f"{stuck_after:.0f}s", ) stopping = threading.Event() @@ -453,6 +578,7 @@ def _on_term(signum: int, _frame: object) -> None: try: while not stopping.is_set(): _reap_dead(fleet, stopping) + _kill_stuck_children(fleet, stuck_after, stopping) _restart_due_children(fleet, stopping) stopping.wait(_MONITOR_INTERVAL_SECONDS) # responsive to SIGTERM finally: @@ -518,8 +644,10 @@ def _maybe_start_supervisor_health(fleet: _Fleet) -> LivenessServer | None: Reuses the single-process consumer's env knobs (``..._HEALTH_PORT`` / ``..._HEALTH_STALE_SECONDS``) and the same HTTP contract (``/health`` → 200/503), so the k8s probe config is unchanged. The JSON body differs - (``check="pg_queue_fleet"``, age key ``oldest_child_seconds_since_poll``) since - the freshness source is the fleet's oldest child, not one poll loop. + (``check="pg_queue_fleet"``, age key ``quorum_child_seconds_since_poll``) since + the freshness source is the fleet's :meth:`_Fleet.quorum_age`, not one poll + loop. The oldest child's age and the stale-child count stay in the body for + diagnosis. ``/ready`` → 200 only once every child has loaded (:meth:`_Fleet.all_loaded`), for the chart's opt-in ``startupProbe``. @@ -549,6 +677,8 @@ def _extra_status() -> dict[str, object]: "loaded_children": fleet.loaded_count(), "concurrency": fleet.concurrency, "crash_looping": fleet.is_crash_looping(), + "oldest_child_seconds_since_poll": fleet.oldest_age(), + "stale_children": fleet.stale_count(stale_after), "liveness_probe_bound": True, } @@ -564,7 +694,7 @@ def _extra_status() -> dict[str, object]: stale_after=stale_after, port=port, check_name="pg_queue_fleet", - age_key="oldest_child_seconds_since_poll", + age_key="quorum_child_seconds_since_poll", extra_status_fn=_extra_status, metrics_fn=metrics.render, ready_fn=fleet.all_loaded, diff --git a/workers/queue_backend/pg_queue/metrics.py b/workers/queue_backend/pg_queue/metrics.py index d34c7d885d..7a5a5487d8 100644 --- a/workers/queue_backend/pg_queue/metrics.py +++ b/workers/queue_backend/pg_queue/metrics.py @@ -75,8 +75,9 @@ class ConsumerMetrics(_Exporter): """Per-pod metrics for a PG-queue consumer (or the fleet supervisor). ``freshness_fn`` is the same heartbeat the liveness probe reads — - seconds since the poll loop last made progress (for the supervisor, the - OLDEST child's, so one wedged child surfaces). The optional fleet hooks + seconds since the poll loop last made progress (for the supervisor, the age + at least half its children have reached; a single wedged child is killed by + the supervisor rather than surfaced here). The optional fleet hooks exist because the supervisor's ``/health`` JSON already reports them and an operator graphing the fleet needs them as numbers, not JSON. """ diff --git a/workers/tests/test_pg_consumer_supervisor.py b/workers/tests/test_pg_consumer_supervisor.py index e4d88f458d..5eccaacb77 100644 --- a/workers/tests/test_pg_consumer_supervisor.py +++ b/workers/tests/test_pg_consumer_supervisor.py @@ -22,6 +22,7 @@ _Fleet, _child_after_fork, _join_children, + _kill_stuck_children, _reap_dead, _restart_due_children, _run_child, @@ -29,6 +30,7 @@ _wait_for_exit, concurrency_from_env, shutdown_grace_from_env, + stuck_child_seconds_from_env, ) _MOD = "pg_queue_consumer.supervisor" @@ -171,11 +173,41 @@ def test_freshness_is_inf_when_crash_looping(self): f.schedule_restart(0, uptime=0.1) assert math.isinf(f.freshness()) - def test_freshness_is_oldest_age_when_healthy(self): - f = _Fleet(2) - f._heartbeats[1] = time.time() - 100 + def test_freshness_with_one_child_is_its_age(self): + f = _Fleet(1) + f._heartbeats[0] = time.time() - 100 assert 99 < f.freshness() < 102 + def test_one_stale_child_of_many_does_not_age_the_fleet(self): + # UN-4223: one hung child must not fail liveness for the whole pod. + f = _Fleet(20) + f._heartbeats[3] = time.time() - 9999 + assert f.freshness() < 5 + assert f.oldest_age() > 9990 # still visible for diagnosis + + @pytest.mark.parametrize( + ("concurrency", "stale", "trips"), + [ + (20, 9, False), + (20, 10, True), # exactly half + (3, 1, False), + (3, 2, True), # ceil(1.5) + (2, 1, True), # half of two + ], + ) + def test_freshness_trips_once_half_the_fleet_is_stale( + self, concurrency, stale, trips + ): + f = _Fleet(concurrency) + for slot in range(stale): + f._heartbeats[slot] = time.time() - 1000 + assert (f.freshness() > 500) is trips + + def test_stale_count(self): + f = _Fleet(4) + f._heartbeats[0] = f._heartbeats[2] = time.time() - 1000 + assert f.stale_count(500) == 2 + def test_due_restarts_respects_backoff(self, monkeypatch): f = _Fleet(1) clock = [1000.0] @@ -250,6 +282,125 @@ def test_already_reaped_is_treated_gone(self): assert f.alive_count() == 0 +_STUCK = "WORKER_PG_QUEUE_CONSUMER_STUCK_CHILD_SECONDS" +_PORT = "WORKER_PG_QUEUE_CONSUMER_HEALTH_PORT" +_STALE = "WORKER_PG_QUEUE_CONSUMER_HEALTH_STALE_SECONDS" + + +class TestStuckChildSecondsFromEnv: + @pytest.fixture(autouse=True) + def _clean(self, monkeypatch): + for var in (_STUCK, _PORT, _STALE): + monkeypatch.delenv(var, raising=False) + + def test_defaults_to_health_stale_window(self, monkeypatch): + # The same bound that used to restart the whole container. + monkeypatch.setenv(_PORT, "8090") + monkeypatch.setenv(_STALE, "7260") + assert stuck_child_seconds_from_env() == pytest.approx(7260.0) + + def test_off_without_a_health_port(self): + # No probe ever enforced a cap here, so none is introduced implicitly. + assert stuck_child_seconds_from_env() is None + + def test_override_wins_even_without_a_port(self, monkeypatch): + monkeypatch.setenv(_PORT, "8090") + monkeypatch.setenv(_STALE, "7260") + monkeypatch.setenv(_STUCK, "3660") + assert stuck_child_seconds_from_env() == pytest.approx(3660.0) + monkeypatch.delenv(_PORT) + assert stuck_child_seconds_from_env() == pytest.approx(3660.0) + + @pytest.mark.parametrize("bad", ["0", "-5", "inf", "nan"]) + def test_invalid_override_raises(self, monkeypatch, bad): + monkeypatch.setenv(_STUCK, bad) + with pytest.raises(ValueError, match="STUCK_CHILD_SECONDS"): + stuck_child_seconds_from_env() + + def test_malformed_override_raises(self, monkeypatch): + monkeypatch.setenv(_STUCK, "1h") + with pytest.raises(ValueError, match="STUCK_CHILD_SECONDS"): + stuck_child_seconds_from_env() + + +class TestKillStuckChildren: + @staticmethod + def _fleet_with_stuck_slot_1(monkeypatch) -> _Fleet: # noqa: ANN001 + clock = [1000.0] + monkeypatch.setattr(f"{_MOD}.time.monotonic", lambda: clock[0]) + f = _Fleet(3) + for slot, pid in enumerate((111, 222, 333)): + f.record_fork(slot, pid) + clock[0] += 5000.0 # every child has been up 5000s + f._heartbeats[1] = time.time() - 4000 # slot 1 frozen on one task + return f + + def test_only_the_stuck_child_is_sigkilled(self, monkeypatch): + import signal as _signal + + f = self._fleet_with_stuck_slot_1(monkeypatch) + with patch(f"{_MOD}.os.kill") as kill: + _kill_stuck_children(f, 3600.0, threading.Event()) + kill.assert_called_once_with(222, _signal.SIGKILL) + + def test_not_killed_again_before_it_is_reaped(self, monkeypatch): + f = self._fleet_with_stuck_slot_1(monkeypatch) + with patch(f"{_MOD}.os.kill") as kill: + _kill_stuck_children(f, 3600.0, threading.Event()) + _kill_stuck_children(f, 3600.0, threading.Event()) + assert kill.call_count == 1 + + def test_killed_child_is_reaped_and_reforked_without_counting_a_crash( + self, monkeypatch + ): + f = self._fleet_with_stuck_slot_1(monkeypatch) + with patch(f"{_MOD}.os.kill"): + _kill_stuck_children(f, 3600.0, threading.Event()) + + def _waitpid(pid, _flags): # noqa: ANN001, ANN202 + return (pid, 9) if pid == 222 else (0, 0) + + with patch(f"{_MOD}.os.waitpid", side_effect=_waitpid): + _reap_dead(f, threading.Event()) + assert sorted(f.alive_items()) == [(0, 111), (2, 333)] + assert f._consecutive_crashes[1] == 0 # ran long: not a crash loop + assert 1 in f._restart_due + assert 1 not in f._stuck_killed # its replacement can be killed later + + def test_fresh_replacement_is_not_killed_for_its_predecessors_age(self, monkeypatch): + # A re-fork keeps the slot's old heartbeat until it first polls; it must + # not be killed mid-bootstrap for that. + clock = [1000.0] + monkeypatch.setattr(f"{_MOD}.time.monotonic", lambda: clock[0]) + f = _Fleet(1) + f._heartbeats[0] = time.time() - 4000 + f.record_fork(0, 444) + clock[0] += 30.0 + with patch(f"{_MOD}.os.kill") as kill: + _kill_stuck_children(f, 3600.0, threading.Event()) + kill.assert_not_called() + + def test_disabled_cap_kills_nothing(self, monkeypatch): + f = self._fleet_with_stuck_slot_1(monkeypatch) + with patch(f"{_MOD}.os.kill") as kill: + _kill_stuck_children(f, None, threading.Event()) + kill.assert_not_called() + + def test_nothing_killed_while_stopping(self, monkeypatch): + f = self._fleet_with_stuck_slot_1(monkeypatch) + stopping = threading.Event() + stopping.set() + with patch(f"{_MOD}.os.kill") as kill: + _kill_stuck_children(f, 3600.0, stopping) + kill.assert_not_called() + + def test_already_gone_child_is_tolerated(self, monkeypatch): + f = self._fleet_with_stuck_slot_1(monkeypatch) + with patch(f"{_MOD}.os.kill", side_effect=ProcessLookupError()): + _kill_stuck_children(f, 3600.0, threading.Event()) + assert 1 in f._stuck_killed + + class TestRestartDueChildren: def test_due_slot_is_reforked(self): f = _Fleet(1) @@ -493,6 +644,10 @@ def _ready() -> tuple[int, dict]: try: with urllib.request.urlopen(f"{base}/health", timeout=5) as resp: assert resp.status == 200 # alive, but... + body = json.loads(resp.read()) + assert body["stale_children"] == 0 + assert "quorum_child_seconds_since_poll" in body + assert "oldest_child_seconds_since_poll" in body code, body = _ready() assert (code, body["status"], body["loaded_children"]) == (503, "starting", 0) From 03a3333a6c435896a9fc0fdf25630166fcbb84d7 Mon Sep 17 00:00:00 2001 From: Rahul Johny Date: Mon, 5 Oct 2026 12:52:33 +0530 Subject: [PATCH 2/2] UN-4223 [FIX] Address review: two-child liveness, bootstrap kills, kill metrics - Reseed a stuck-killed slot's heartbeat when it is reaped. In a two-child fleet one slot is the quorum, so the frozen age kept /health at 503 through the replacement's bootstrap and could still restart the pod. Crash exits are not reseeded, so crash-loop detection is unchanged. - Only kill children that have finished loading. With a cap shorter than the bootstrap, a child that had not polled yet could be killed before starting. - Warn at startup when STUCK_CHILD_SECONDS exceeds HEALTH_STALE_SECONDS, where a stuck child can fail the probe before it is killed. - Export pg_consumer_oldest_child_age_seconds and pg_consumer_stuck_child_kills_total, so a stuck child and each kill are visible while the half-of-fleet heartbeat stays healthy. Co-Authored-By: Claude Opus 5.5 --- workers/pg_queue_consumer/supervisor.py | 61 +++++++++++++++++--- workers/queue_backend/pg_queue/metrics.py | 42 ++++++++++++++ workers/tests/test_pg_consumer_supervisor.py | 54 +++++++++++++++++ workers/tests/test_pg_metrics.py | 15 +++++ 4 files changed, 165 insertions(+), 7 deletions(-) diff --git a/workers/pg_queue_consumer/supervisor.py b/workers/pg_queue_consumer/supervisor.py index e99b5a5345..5932f54970 100644 --- a/workers/pg_queue_consumer/supervisor.py +++ b/workers/pg_queue_consumer/supervisor.py @@ -188,6 +188,31 @@ def stuck_child_seconds_from_env() -> float | None: return consumer_env("HEALTH_STALE_SECONDS", _DEFAULT_HEALTH_STALE_SECONDS, float) +def _warn_if_cap_exceeds_stale_window(stuck_after: float | None) -> None: + """A cap above the health stale window lets a stuck child count as stale + before it is killed. In a small fleet one child can be half of it, so the + probe would restart the pod first, the outage the cap exists to prevent. + """ + from queue_backend.pg_queue.consumer import ( + _DEFAULT_HEALTH_STALE_SECONDS, + consumer_env, + ) + + if stuck_after is None or consumer_env("HEALTH_PORT", None, int) is None: + return + stale_after = consumer_env( + "HEALTH_STALE_SECONDS", _DEFAULT_HEALTH_STALE_SECONDS, float + ) + if stuck_after > stale_after: + logger.warning( + "PG-queue consumer supervisor: STUCK_CHILD_SECONDS (%.0fs) exceeds " + "HEALTH_STALE_SECONDS (%.0fs) — a stuck child can fail the probe " + "before it is killed; keep the cap at or below the stale window", + stuck_after, + stale_after, + ) + + class _Fleet: """Owns the per-slot child state — pid, last-fork, heartbeat, crash count and pending-restart schedule — keeping them mutually consistent. Slots are @@ -218,6 +243,7 @@ def __init__(self, concurrency: int) -> None: # Slots already SIGKILLed as stuck, so a child that takes a tick to die # is not re-killed and re-logged on every monitor pass. Cleared in reap(). self._stuck_killed: set[int] = set() + self._stuck_kill_count = 0 @property def concurrency(self) -> int: @@ -245,7 +271,8 @@ def record_fork(self, slot: int, pid: int) -> None: """Mark ``slot`` alive under ``pid``; clears any pending restart. Note the heartbeat is deliberately NOT reseeded here — a re-forked child must earn freshness by actually polling, so a crash-looping slot ages instead of - looking perpetually fresh. + looking perpetually fresh. (A slot the supervisor killed as stuck is the + one exception; :meth:`reap` reseeds it.) """ self._validate(slot) self._pids[slot] = pid @@ -257,11 +284,20 @@ def reap(self, slot: int) -> float: Also clears the slot's loaded flag: its replacement must finish its own bootstrap before the fleet counts as loaded again. + + A slot the supervisor killed as stuck gets its heartbeat reseeded: the + stuck task is gone, and leaving its frozen age in place would keep the + slot stale for the replacement's whole bootstrap. In a two-child fleet + that one slot is the quorum, so the probe would restart the pod anyway. + A crash loop is still caught, by :meth:`is_crash_looping`. The child is + dead by now, so the parent is the slot's only writer. """ forked_at = self._last_fork.pop(slot, time.monotonic()) self._pids.pop(slot, None) self._loaded[slot] = 0 - self._stuck_killed.discard(slot) + if slot in self._stuck_killed: + self._stuck_killed.discard(slot) + self._heartbeats[slot] = time.time() return time.monotonic() - forked_at def schedule_restart(self, slot: int, uptime: float) -> int: @@ -340,15 +376,17 @@ def stuck_children(self, stuck_after: float) -> list[tuple[int, int, float]]: """``(slot, pid, heartbeat_age)`` for each live child stuck past ``stuck_after`` and not yet killed. - A re-forked child's heartbeat is NOT reseeded (see :meth:`record_fork`), so - it carries its predecessor's age until it first polls. Requiring the child - itself to have been up for ``stuck_after`` stops a fresh replacement from - being killed mid-bootstrap for its predecessor's staleness. + Only children that have finished loading count: before that the slot + holds the parent's seed or a predecessor's heartbeat, not this child's, + and a slow bootstrap is not a stuck task. A loaded child's heartbeat + starts fresh, because the consumer stamps its last poll when it is built. + Requiring the child to have been up for ``stuck_after`` as well covers + the moment between it setting its loaded flag and first publishing. """ now_wall, now_mono = time.time(), time.monotonic() stuck = [] for slot, pid in self.alive_items(): - if slot in self._stuck_killed: + if slot in self._stuck_killed or not self._loaded[slot]: continue age = now_wall - self._heartbeats[slot] uptime = now_mono - self._last_fork.get(slot, now_mono) @@ -359,6 +397,12 @@ def stuck_children(self, stuck_after: float) -> list[tuple[int, int, float]]: def mark_stuck_killed(self, slot: int) -> None: self._validate(slot) self._stuck_killed.add(slot) + self._stuck_kill_count += 1 + + @property + def stuck_kill_count(self) -> int: + """Children killed as stuck since the supervisor started.""" + return self._stuck_kill_count def freshness(self) -> float: """Liveness verdict source: a crash-looping fleet is force-stale (``inf``) @@ -536,6 +580,7 @@ def run_supervised(concurrency: int) -> None: fleet = _Fleet(concurrency) grace_seconds = shutdown_grace_from_env() stuck_after = stuck_child_seconds_from_env() + _warn_if_cap_exceeds_stale_window(stuck_after) logger.info( "PG-queue consumer supervisor: shutdown drain grace = %.0fs (shared across " "children); stuck-child cap = %s", @@ -688,6 +733,8 @@ def _extra_status() -> dict[str, object]: freshness_fn=fleet.freshness, alive_children_fn=lambda: float(fleet.alive_count()), concurrency_fn=lambda: float(fleet.concurrency), + oldest_child_age_fn=fleet.oldest_age, + stuck_child_kills_fn=lambda: float(fleet.stuck_kill_count), ) server = LivenessServer( freshness_fn=fleet.freshness, diff --git a/workers/queue_backend/pg_queue/metrics.py b/workers/queue_backend/pg_queue/metrics.py index 7a5a5487d8..1e78d28cd0 100644 --- a/workers/queue_backend/pg_queue/metrics.py +++ b/workers/queue_backend/pg_queue/metrics.py @@ -65,12 +65,37 @@ def _function_gauge(self, name: str, doc: str, fn: Callable[[], float]) -> None: Gauge(name, doc, registry=self.registry).set_function(fn) + def _function_counter(self, name: str, doc: str, fn: Callable[[], float]) -> None: + self.registry.register(_FunctionCounter(name, doc, fn)) + def render(self) -> bytes: from prometheus_client import generate_latest return generate_latest(self.registry) +class _FunctionCounter: + """A counter whose value is read from ``fn`` at scrape time. + + ``prometheus_client.Counter`` can only be incremented, not backed by a + callback; this suits a count kept by its owner (e.g. the supervisor's + ``_Fleet``) rather than by the metrics layer. + """ + + def __init__(self, name: str, doc: str, fn: Callable[[], float]) -> None: + self._name, self._doc, self._fn = name, doc, fn + + def describe(self) -> Iterable[Metric]: + from prometheus_client.core import CounterMetricFamily + + return (CounterMetricFamily(self._name, self._doc),) + + def collect(self) -> Iterable[Metric]: + from prometheus_client.core import CounterMetricFamily + + return (CounterMetricFamily(self._name, self._doc, value=self._fn()),) + + class ConsumerMetrics(_Exporter): """Per-pod metrics for a PG-queue consumer (or the fleet supervisor). @@ -88,6 +113,8 @@ def __init__( freshness_fn: Callable[[], float], alive_children_fn: Callable[[], float] | None = None, concurrency_fn: Callable[[], float] | None = None, + oldest_child_age_fn: Callable[[], float] | None = None, + stuck_child_kills_fn: Callable[[], float] | None = None, ) -> None: super().__init__() self._function_gauge( @@ -108,6 +135,21 @@ def __init__( "Configured child-process concurrency of the supervisor fleet", concurrency_fn, ) + if oldest_child_age_fn is not None: + self._function_gauge( + "pg_consumer_oldest_child_age_seconds", + "Seconds since the most-stale child in the supervisor fleet last " + "polled (one child stuck on a task shows here, not in the heartbeat " + "age, until it is killed)", + oldest_child_age_fn, + ) + if stuck_child_kills_fn is not None: + self._function_counter( + "pg_consumer_stuck_child_kills", + "Children the supervisor SIGKILLed for running one task past the " + "stuck-child cap (each one's message is redelivered)", + stuck_child_kills_fn, + ) @dataclass(frozen=True) diff --git a/workers/tests/test_pg_consumer_supervisor.py b/workers/tests/test_pg_consumer_supervisor.py index 5eccaacb77..a427b8e499 100644 --- a/workers/tests/test_pg_consumer_supervisor.py +++ b/workers/tests/test_pg_consumer_supervisor.py @@ -317,6 +317,17 @@ def test_invalid_override_raises(self, monkeypatch, bad): with pytest.raises(ValueError, match="STUCK_CHILD_SECONDS"): stuck_child_seconds_from_env() + def test_warns_when_cap_exceeds_stale_window(self, monkeypatch, caplog): + monkeypatch.setenv(_PORT, "8090") + monkeypatch.setenv(_STALE, "180") + with caplog.at_level("WARNING", logger=_MOD): + sup._warn_if_cap_exceeds_stale_window(3600.0) + assert "exceeds" in caplog.text + caplog.clear() + sup._warn_if_cap_exceeds_stale_window(180.0) + sup._warn_if_cap_exceeds_stale_window(None) + assert caplog.text == "" + def test_malformed_override_raises(self, monkeypatch): monkeypatch.setenv(_STUCK, "1h") with pytest.raises(ValueError, match="STUCK_CHILD_SECONDS"): @@ -331,6 +342,7 @@ def _fleet_with_stuck_slot_1(monkeypatch) -> _Fleet: # noqa: ANN001 f = _Fleet(3) for slot, pid in enumerate((111, 222, 333)): f.record_fork(slot, pid) + f.loaded[slot] = 1 clock[0] += 5000.0 # every child has been up 5000s f._heartbeats[1] = time.time() - 4000 # slot 1 frozen on one task return f @@ -366,6 +378,47 @@ def _waitpid(pid, _flags): # noqa: ANN001, ANN202 assert f._consecutive_crashes[1] == 0 # ran long: not a crash loop assert 1 in f._restart_due assert 1 not in f._stuck_killed # its replacement can be killed later + assert f.stuck_kill_count == 1 + + def test_reaping_a_killed_child_reseeds_its_heartbeat(self, monkeypatch): + # Review finding: in a two-child fleet one slot is the quorum, so leaving + # the killed child's frozen age in place kept /health at 503 through the + # replacement's bootstrap. + clock = [1000.0] + monkeypatch.setattr(f"{_MOD}.time.monotonic", lambda: clock[0]) + f = _Fleet(2) + for slot, pid in enumerate((111, 222)): + f.record_fork(slot, pid) + f.loaded[slot] = 1 + clock[0] += 5000.0 + f._heartbeats[1] = time.time() - 4000 + assert f.freshness() > 3600 # stale before the kill + with patch(f"{_MOD}.os.kill"): + _kill_stuck_children(f, 3600.0, threading.Event()) + with patch( + f"{_MOD}.os.waitpid", + side_effect=lambda pid, _flags: (pid, 9) if pid == 222 else (0, 0), + ): + _reap_dead(f, threading.Event()) + assert f.freshness() < 5 # healthy while the replacement boots + + def test_a_crashed_child_is_not_reseeded(self): + # The reseed is only for stuck kills; a crash must still age the slot. + f = _Fleet(1) + f._heartbeats[0] = time.time() - 500 + f.record_fork(0, 111) + with patch(f"{_MOD}.os.waitpid", return_value=(111, 0)): + _reap_dead(f, threading.Event()) + assert f.oldest_age() > 400 + + def test_a_child_still_loading_is_never_killed(self, monkeypatch): + # Review finding: with a cap shorter than the bootstrap, a child that has + # not polled yet passed both checks and was killed before it could start. + f = self._fleet_with_stuck_slot_1(monkeypatch) + f.loaded[1] = 0 + with patch(f"{_MOD}.os.kill") as kill: + _kill_stuck_children(f, 3600.0, threading.Event()) + kill.assert_not_called() def test_fresh_replacement_is_not_killed_for_its_predecessors_age(self, monkeypatch): # A re-fork keeps the slot's old heartbeat until it first polls; it must @@ -375,6 +428,7 @@ def test_fresh_replacement_is_not_killed_for_its_predecessors_age(self, monkeypa f = _Fleet(1) f._heartbeats[0] = time.time() - 4000 f.record_fork(0, 444) + f.loaded[0] = 1 # loaded, but its heartbeat not yet published clock[0] += 30.0 with patch(f"{_MOD}.os.kill") as kill: _kill_stuck_children(f, 3600.0, threading.Event()) diff --git a/workers/tests/test_pg_metrics.py b/workers/tests/test_pg_metrics.py index e21bc2b67f..da4a8ce0d0 100644 --- a/workers/tests/test_pg_metrics.py +++ b/workers/tests/test_pg_metrics.py @@ -143,6 +143,19 @@ def test_fleet_hooks_are_optional(self): 4.0 ) + def test_stuck_child_hooks(self): + kills = 0.0 + metrics = ConsumerMetrics( + freshness_fn=lambda: 0.0, + oldest_child_age_fn=lambda: 99.0, + stuck_child_kills_fn=lambda: kills, + ) + oldest = _sample(metrics, "pg_consumer_oldest_child_age_seconds") + assert oldest == pytest.approx(99.0) + assert _sample(metrics, "pg_consumer_stuck_child_kills_total") == 0.0 + kills = 2.0 # read at scrape time + assert _sample(metrics, "pg_consumer_stuck_child_kills_total") == 2.0 + def test_render_is_prometheus_exposition(self): body = ConsumerMetrics(freshness_fn=lambda: 1.0).render() assert b"pg_consumer_heartbeat_age_seconds" in body @@ -596,5 +609,7 @@ def test_supervisor_fleet_metrics_served(self, monkeypatch): assert b"pg_consumer_heartbeat_age_seconds" in body assert b"pg_consumer_alive_children" in body assert b"pg_consumer_configured_concurrency 2.0" in body + assert b"pg_consumer_oldest_child_age_seconds" in body + assert b"pg_consumer_stuck_child_kills_total 0.0" in body finally: server.stop()