From 0cdfd285ad8989ef719faf5476eccb89832b5fcc Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 10 Jul 2026 16:32:58 +0000 Subject: [PATCH 1/5] Recover aragorn.omnicorp worker from a broken process pool A single pool child dying abruptly -- almost always an OOM kill while overlaying a very large message -- puts the shared ProcessPoolExecutor into a permanently broken state. From then on every run_in_executor call raises BrokenProcessPool ("A process in the process pool was terminated abruptly while the future was running or pending"), not just the task whose child died. Because poll_for_tasks reused one executor for the worker's whole lifetime, one oversized query wedged the worker into failing every subsequent task until the pod restarted. Wrap the executor in a PoolManager that catches the breakage, tears the dead pool down, and stands up a fresh one so the worker recovers on its next task. Replacement is guarded so the burst of concurrent BrokenProcessPool errors (every in-flight future fails at once) swaps the pool only once. The triggering task still fails cleanly via run_task_lifecycle; only the poisoning of later tasks is fixed. Also add an opt-in OMNICORP_MAX_TASKS_PER_CHILD knob to recycle pool children under memory pressure so peak resident size from a big message is returned to the OS instead of accumulating across tasks. Left unset by default since it forces the spawn start method. Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_01TiiFCcToq8LgiEeaQhQM5e --- tests/unit/aragorn/test_aragorn_omnicorp.py | 44 +++++++++ workers/aragorn_omnicorp/worker.py | 102 ++++++++++++++++++-- 2 files changed, 137 insertions(+), 9 deletions(-) diff --git a/tests/unit/aragorn/test_aragorn_omnicorp.py b/tests/unit/aragorn/test_aragorn_omnicorp.py index bf47bf9..cbbc0d9 100644 --- a/tests/unit/aragorn/test_aragorn_omnicorp.py +++ b/tests/unit/aragorn/test_aragorn_omnicorp.py @@ -5,17 +5,32 @@ LMDB shims and the full overlay on a small TRAPI message. """ +import asyncio import copy import json import logging +import os import lmdb import pytest +from concurrent.futures.process import BrokenProcessPool + from shepherd_utils.config import settings from workers.aragorn_omnicorp import worker +# Top-level (picklable) callables for the process-pool recovery test. +def _pool_echo(value): + return value + + +def _pool_suicide(_): + # Abruptly kill the child so the executor enters its broken state, mimicking + # a kernel OOM kill of a worker processing an oversized message. + os._exit(1) + + def _build_curies_lmdb(path, entries): """Write ``{curie: {"pmc": int, "index": int}}`` entries to a single-file LMDB.""" env = lmdb.open(str(path), subdir=False, map_size=10 * 1024 * 1024) @@ -426,6 +441,35 @@ def test_aragorn_omnicorp_loads_overlays_saves_and_preserves_workflow( assert article_attrs and article_attrs[0]["value"] == 100 +async def test_pool_manager_recovers_after_child_dies(): + """A child dying breaks the pool for one task, then the pool self-heals. + + Regression test for the wedged-worker bug: previously a single OOM-killed + child poisoned the shared ProcessPoolExecutor so every subsequent task + failed with "A process in the process pool was terminated abruptly...". The + PoolManager must replace the broken pool so the very next task succeeds. + """ + loop = asyncio.get_running_loop() + pool = worker.PoolManager(max_workers=1) + try: + # Healthy pool runs a task fine. + assert await pool.run(loop, _pool_echo, "before") == "before" + + first = pool._executor + + # A child dying surfaces as BrokenProcessPool for the triggering task... + with pytest.raises(BrokenProcessPool): + await pool.run(loop, _pool_suicide, None) + + # ...but the manager swaps in a fresh executor... + assert pool._executor is not first + + # ...so the next task runs on the healthy pool instead of re-raising. + assert await pool.run(loop, _pool_echo, "after") == "after" + finally: + pool._executor.shutdown(wait=False, cancel_futures=True) + + def test_generate_curie_pairs_stops_at_max_pairs(): """``max_pairs`` bounds the mapping instead of materializing every pair.""" from workers.aragorn_omnicorp.worker import generate_curie_pairs diff --git a/workers/aragorn_omnicorp/worker.py b/workers/aragorn_omnicorp/worker.py index 454fb17..a2ed350 100644 --- a/workers/aragorn_omnicorp/worker.py +++ b/workers/aragorn_omnicorp/worker.py @@ -19,6 +19,7 @@ import uuid from collections import defaultdict from concurrent.futures import ProcessPoolExecutor +from concurrent.futures.process import BrokenProcessPool from datetime import datetime from itertools import combinations from typing import Dict, List @@ -51,6 +52,35 @@ # seconds to 45 minutes-3 hours. Override via the env var if needed. OMNICORP_MAX_CURIE_PAIRS = int(os.environ.get("OMNICORP_MAX_CURIE_PAIRS", 1_000_000)) + +def _parse_max_tasks_per_child(): + """Optional ceiling on overlays a single pool child runs before it recycles. + + Recycling a child returns the memory a very large message forced it to grow + (peak resident size persists for the process's lifetime) back to the OS + instead of letting it accumulate across tasks, so the child stays well below + the OOM line. Left unset (``None``) by default because enabling it makes the + ``ProcessPoolExecutor`` fall back to the ``spawn`` start method -- a Python + requirement of ``max_tasks_per_child`` -- which re-imports this worker in + each fresh child. Opt in via ``OMNICORP_MAX_TASKS_PER_CHILD`` on a + memory-pressured deployment. + """ + raw = os.environ.get("OMNICORP_MAX_TASKS_PER_CHILD") + if not raw: + return None + try: + value = int(raw) + except ValueError: + logging.warning( + f"Ignoring invalid OMNICORP_MAX_TASKS_PER_CHILD={raw!r}; leaving pool " + "children un-recycled." + ) + return None + return value if value > 0 else None + + +OMNICORP_MAX_TASKS_PER_CHILD = _parse_max_tasks_per_child() + # Both LMDBs are opened lazily on first use so importing the worker (e.g. in # tests) does not require the live data files. Static datasets, so we open # read-only and skip the writer lock. @@ -502,8 +532,64 @@ def aragorn_omnicorp(response_id: str, logger: logging.Logger) -> None: save_message_sync(response_id, response) +class PoolManager: + """Owns the process pool and transparently replaces it once it breaks. + + A single child dying abruptly -- most often an OOM kill while overlaying a + very large message -- puts a ``ProcessPoolExecutor`` into a *permanently* + broken state: every subsequent submission raises ``BrokenProcessPool``, not + just the one whose child died. Because ``poll_for_tasks`` reused one executor + for the worker's whole lifetime, that turned a single oversized query into a + wedged worker that failed every following task ("A process in the process + pool was terminated abruptly...") until the pod restarted. This wrapper + catches the breakage, tears the dead pool down, and stands up a fresh one so + the worker recovers on its next task. + """ + + def __init__(self, max_workers: int): + self._max_workers = max_workers + self._lock = asyncio.Lock() + self._executor = self._new_executor() + + def _new_executor(self) -> ProcessPoolExecutor: + kwargs = {"max_workers": self._max_workers} + if OMNICORP_MAX_TASKS_PER_CHILD is not None: + kwargs["max_tasks_per_child"] = OMNICORP_MAX_TASKS_PER_CHILD + return ProcessPoolExecutor(**kwargs) + + async def run(self, loop, fn, *args): + """Run ``fn`` in the pool, recreating it if it was broken by the call.""" + executor = self._executor + try: + return await loop.run_in_executor(executor, fn, *args) + except BrokenProcessPool: + # A child died; the whole executor is now unusable. Replace it before + # re-raising so this task fails cleanly (handled by run_task_lifecycle) + # while the next task gets a healthy pool. + await self._replace(executor) + raise + + async def _replace(self, broken: ProcessPoolExecutor) -> None: + async with self._lock: + # When a child dies, every future in flight on that pool raises + # BrokenProcessPool at once, so several tasks may race here. Only the + # first to see this particular pool swaps it; the rest are no-ops. + if self._executor is not broken: + return + logging.error( + "aragorn.omnicorp process pool broke (a child likely died or was " + "OOM-killed on a large message); replacing it so the worker keeps " + "processing tasks." + ) + try: + broken.shutdown(wait=False, cancel_futures=True) + except Exception: + pass + self._executor = self._new_executor() + + async def process_task( - task, parent_ctx, logger: logging.Logger, limiter, loop, executor + task, parent_ctx, logger: logging.Logger, limiter, loop, pool: "PoolManager" ): """Process a given task and ACK in redis. @@ -514,16 +600,14 @@ async def process_task( everything else until it finished. Only the ``response_id`` is handed to the child; the message load/save happen there (see ``aragorn_omnicorp``) so the payload never crosses the process boundary. + + Dispatch goes through ``pool`` (not a raw executor) so a child dying on an + oversized message replaces the pool instead of poisoning it for good. """ async def _run(task, logger): response_id = task[1]["response_id"] - await loop.run_in_executor( - executor, - aragorn_omnicorp, - response_id, - logger, - ) + await pool.run(loop, aragorn_omnicorp, response_id, logger) await run_task_lifecycle(STREAM, GROUP, task, parent_ctx, logger, limiter, _run) @@ -537,14 +621,14 @@ async def poll_for_tasks(): cpu_count = os.cpu_count() cpu_count = cpu_count if cpu_count is not None else 1 cpu_count = min(cpu_count, TASK_LIMIT) - executor = ProcessPoolExecutor(max_workers=cpu_count) + pool = PoolManager(cpu_count) while True: try: async for task, parent_ctx, logger, limiter in get_tasks( STREAM, GROUP, CONSUMER, TASK_LIMIT ): asyncio.create_task( - process_task(task, parent_ctx, logger, limiter, loop, executor) + process_task(task, parent_ctx, logger, limiter, loop, pool) ) except asyncio.CancelledError: logging.info("Poll loop cancelled, shutting down.") From 51e9d1880d051153677aa1f862b87b774066931f Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 10 Jul 2026 16:40:42 +0000 Subject: [PATCH 2/5] Extend broken-pool recovery to aragorn.score and arax.rank aragorn_score and arax_rank had the same wedged-worker vulnerability just fixed in aragorn.omnicorp: poll_for_tasks created a single ProcessPoolExecutor reused for the worker's whole lifetime, so one child dying abruptly (a kernel OOM kill on a large message) put the executor into a permanently broken state and every subsequent task failed with "A process in the process pool was terminated abruptly while the future was running or pending" until the pod restarted. Extract the self-healing wrapper introduced for omnicorp into a shared shepherd_utils.process_pool.ProcessPoolManager and route all three CPU- bound workers through it. The manager catches BrokenProcessPool, tears the dead pool down, and stands up a fresh one so the worker recovers on its next task; the triggering task still fails cleanly via run_task_lifecycle. Surveyed the remaining workers: merge_message already recovers via its own holder + _recreate_executor path; score_paths and monitor/alerts use thread pools, which a task crash cannot permanently poison. No other worker is affected. Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_01TiiFCcToq8LgiEeaQhQM5e --- shepherd_utils/process_pool.py | 90 +++++++++++++++++++++ tests/unit/aragorn/test_aragorn_omnicorp.py | 44 ---------- tests/unit/test_process_pool.py | 69 ++++++++++++++++ workers/aragorn_omnicorp/worker.py | 67 ++------------- workers/aragorn_score/worker.py | 18 ++--- workers/arax_rank/worker.py | 20 +++-- 6 files changed, 183 insertions(+), 125 deletions(-) create mode 100644 shepherd_utils/process_pool.py create mode 100644 tests/unit/test_process_pool.py diff --git a/shepherd_utils/process_pool.py b/shepherd_utils/process_pool.py new file mode 100644 index 0000000..76a6773 --- /dev/null +++ b/shepherd_utils/process_pool.py @@ -0,0 +1,90 @@ +"""Self-healing process pool for CPU-bound worker overlays. + +A ``ProcessPoolExecutor`` that loses a child abruptly -- almost always a kernel +OOM kill while a child processes a very large message -- becomes *permanently* +broken: every subsequent submission raises ``BrokenProcessPool``, not only the +task whose child died. A worker that creates one executor for its whole +lifetime (as ``poll_for_tasks`` does) would therefore fail every task after the +first such death -- with "A process in the process pool was terminated abruptly +while the future was running or pending" -- until the pod was restarted. + +``ProcessPoolManager`` wraps the executor, catches the breakage, tears the dead +pool down, and stands up a fresh one so the worker recovers on its next task. +The task whose child died still fails cleanly (its ``run`` call re-raises, and +``run_task_lifecycle`` routes it to ``finish_query`` with an error); only the +poisoning of *later* tasks is fixed. +""" + +import asyncio +import logging +from concurrent.futures import ProcessPoolExecutor +from concurrent.futures.process import BrokenProcessPool +from typing import Optional + + +class ProcessPoolManager: + """Owns a ``ProcessPoolExecutor`` and transparently replaces it once broken. + + Args: + max_workers: pool size handed to ``ProcessPoolExecutor``. + max_tasks_per_child: optional ceiling on tasks a single child runs + before it recycles, returning the memory a large message forced it + to grow back to the OS. ``None`` (default) leaves children alive for + the pool's lifetime; setting it makes the executor fall back to the + ``spawn`` start method (a Python requirement of the parameter). + name: label used in the log line emitted when the pool is replaced. + """ + + def __init__( + self, + max_workers: int, + max_tasks_per_child: Optional[int] = None, + name: str = "process pool", + ): + self._max_workers = max_workers + self._max_tasks_per_child = max_tasks_per_child + self._name = name + self._lock = asyncio.Lock() + self._executor = self._new_executor() + + def _new_executor(self) -> ProcessPoolExecutor: + kwargs = {"max_workers": self._max_workers} + if self._max_tasks_per_child is not None: + kwargs["max_tasks_per_child"] = self._max_tasks_per_child + return ProcessPoolExecutor(**kwargs) + + async def run(self, loop, fn, *args): + """Run ``fn(*args)`` in the pool, recreating it if the call broke it.""" + executor = self._executor + try: + return await loop.run_in_executor(executor, fn, *args) + except BrokenProcessPool: + # A child died; the whole executor is now unusable. Replace it before + # re-raising so this task fails cleanly (handled upstream) while the + # next task gets a healthy pool. + await self._replace(executor) + raise + + async def _replace(self, broken: ProcessPoolExecutor) -> None: + async with self._lock: + # When a child dies, every future in flight on that pool raises + # BrokenProcessPool at once, so several tasks may race here. Only the + # first to see this particular pool swaps it; the rest are no-ops. + if self._executor is not broken: + return + logging.error( + f"{self._name} broke (a child likely died or was OOM-killed on a " + "large message); replacing it so the worker keeps processing tasks." + ) + try: + broken.shutdown(wait=False, cancel_futures=True) + except Exception: + pass + self._executor = self._new_executor() + + def shutdown(self) -> None: + """Tear the current executor down (best effort).""" + try: + self._executor.shutdown(wait=False, cancel_futures=True) + except Exception: + pass diff --git a/tests/unit/aragorn/test_aragorn_omnicorp.py b/tests/unit/aragorn/test_aragorn_omnicorp.py index cbbc0d9..bf47bf9 100644 --- a/tests/unit/aragorn/test_aragorn_omnicorp.py +++ b/tests/unit/aragorn/test_aragorn_omnicorp.py @@ -5,32 +5,17 @@ LMDB shims and the full overlay on a small TRAPI message. """ -import asyncio import copy import json import logging -import os import lmdb import pytest -from concurrent.futures.process import BrokenProcessPool - from shepherd_utils.config import settings from workers.aragorn_omnicorp import worker -# Top-level (picklable) callables for the process-pool recovery test. -def _pool_echo(value): - return value - - -def _pool_suicide(_): - # Abruptly kill the child so the executor enters its broken state, mimicking - # a kernel OOM kill of a worker processing an oversized message. - os._exit(1) - - def _build_curies_lmdb(path, entries): """Write ``{curie: {"pmc": int, "index": int}}`` entries to a single-file LMDB.""" env = lmdb.open(str(path), subdir=False, map_size=10 * 1024 * 1024) @@ -441,35 +426,6 @@ def test_aragorn_omnicorp_loads_overlays_saves_and_preserves_workflow( assert article_attrs and article_attrs[0]["value"] == 100 -async def test_pool_manager_recovers_after_child_dies(): - """A child dying breaks the pool for one task, then the pool self-heals. - - Regression test for the wedged-worker bug: previously a single OOM-killed - child poisoned the shared ProcessPoolExecutor so every subsequent task - failed with "A process in the process pool was terminated abruptly...". The - PoolManager must replace the broken pool so the very next task succeeds. - """ - loop = asyncio.get_running_loop() - pool = worker.PoolManager(max_workers=1) - try: - # Healthy pool runs a task fine. - assert await pool.run(loop, _pool_echo, "before") == "before" - - first = pool._executor - - # A child dying surfaces as BrokenProcessPool for the triggering task... - with pytest.raises(BrokenProcessPool): - await pool.run(loop, _pool_suicide, None) - - # ...but the manager swaps in a fresh executor... - assert pool._executor is not first - - # ...so the next task runs on the healthy pool instead of re-raising. - assert await pool.run(loop, _pool_echo, "after") == "after" - finally: - pool._executor.shutdown(wait=False, cancel_futures=True) - - def test_generate_curie_pairs_stops_at_max_pairs(): """``max_pairs`` bounds the mapping instead of materializing every pair.""" from workers.aragorn_omnicorp.worker import generate_curie_pairs diff --git a/tests/unit/test_process_pool.py b/tests/unit/test_process_pool.py new file mode 100644 index 0000000..7c70d15 --- /dev/null +++ b/tests/unit/test_process_pool.py @@ -0,0 +1,69 @@ +"""Tests for the self-healing ProcessPoolManager. + +Regression coverage for the wedged-worker bug: a single OOM-killed child used +to poison the shared ProcessPoolExecutor so every subsequent task failed with +"A process in the process pool was terminated abruptly...". The manager must +replace the broken pool so the very next task succeeds. +""" + +import asyncio +import os + +from concurrent.futures.process import BrokenProcessPool + +import pytest + +from shepherd_utils.process_pool import ProcessPoolManager + + +# Top-level (picklable) callables so the process pool can dispatch them. +def _echo(value): + return value + + +def _suicide(_): + # Abruptly kill the child so the executor enters its broken state, mimicking + # a kernel OOM kill of a worker processing an oversized message. + os._exit(1) + + +async def test_manager_recovers_after_child_dies(): + loop = asyncio.get_running_loop() + pool = ProcessPoolManager(max_workers=1, name="test pool") + try: + # Healthy pool runs a task fine. + assert await pool.run(loop, _echo, "before") == "before" + + first = pool._executor + + # A child dying surfaces as BrokenProcessPool for the triggering task... + with pytest.raises(BrokenProcessPool): + await pool.run(loop, _suicide, None) + + # ...but the manager swaps in a fresh executor... + assert pool._executor is not first + + # ...so the next task runs on the healthy pool instead of re-raising. + assert await pool.run(loop, _echo, "after") == "after" + finally: + pool.shutdown() + + +async def test_replace_is_idempotent_for_the_same_broken_pool(): + """Concurrent tasks that all saw the same dead pool replace it only once.""" + loop = asyncio.get_running_loop() + pool = ProcessPoolManager(max_workers=2, name="test pool") + try: + await pool.run(loop, _echo, "warmup") + broken = pool._executor + + # Two callers both observed `broken`; the first swaps it, the second is + # a no-op (identity guard) rather than churning a second fresh pool. + await pool._replace(broken) + replaced_once = pool._executor + await pool._replace(broken) + + assert replaced_once is not broken + assert pool._executor is replaced_once + finally: + pool.shutdown() diff --git a/workers/aragorn_omnicorp/worker.py b/workers/aragorn_omnicorp/worker.py index a2ed350..e944989 100644 --- a/workers/aragorn_omnicorp/worker.py +++ b/workers/aragorn_omnicorp/worker.py @@ -18,8 +18,6 @@ import os import uuid from collections import defaultdict -from concurrent.futures import ProcessPoolExecutor -from concurrent.futures.process import BrokenProcessPool from datetime import datetime from itertools import combinations from typing import Dict, List @@ -30,6 +28,7 @@ from shepherd_utils.config import settings from shepherd_utils.db import get_message_sync, save_message_sync from shepherd_utils.otel import setup_tracer +from shepherd_utils.process_pool import ProcessPoolManager from shepherd_utils.shared import get_tasks, run_task_lifecycle # Queue name @@ -532,64 +531,8 @@ def aragorn_omnicorp(response_id: str, logger: logging.Logger) -> None: save_message_sync(response_id, response) -class PoolManager: - """Owns the process pool and transparently replaces it once it breaks. - - A single child dying abruptly -- most often an OOM kill while overlaying a - very large message -- puts a ``ProcessPoolExecutor`` into a *permanently* - broken state: every subsequent submission raises ``BrokenProcessPool``, not - just the one whose child died. Because ``poll_for_tasks`` reused one executor - for the worker's whole lifetime, that turned a single oversized query into a - wedged worker that failed every following task ("A process in the process - pool was terminated abruptly...") until the pod restarted. This wrapper - catches the breakage, tears the dead pool down, and stands up a fresh one so - the worker recovers on its next task. - """ - - def __init__(self, max_workers: int): - self._max_workers = max_workers - self._lock = asyncio.Lock() - self._executor = self._new_executor() - - def _new_executor(self) -> ProcessPoolExecutor: - kwargs = {"max_workers": self._max_workers} - if OMNICORP_MAX_TASKS_PER_CHILD is not None: - kwargs["max_tasks_per_child"] = OMNICORP_MAX_TASKS_PER_CHILD - return ProcessPoolExecutor(**kwargs) - - async def run(self, loop, fn, *args): - """Run ``fn`` in the pool, recreating it if it was broken by the call.""" - executor = self._executor - try: - return await loop.run_in_executor(executor, fn, *args) - except BrokenProcessPool: - # A child died; the whole executor is now unusable. Replace it before - # re-raising so this task fails cleanly (handled by run_task_lifecycle) - # while the next task gets a healthy pool. - await self._replace(executor) - raise - - async def _replace(self, broken: ProcessPoolExecutor) -> None: - async with self._lock: - # When a child dies, every future in flight on that pool raises - # BrokenProcessPool at once, so several tasks may race here. Only the - # first to see this particular pool swaps it; the rest are no-ops. - if self._executor is not broken: - return - logging.error( - "aragorn.omnicorp process pool broke (a child likely died or was " - "OOM-killed on a large message); replacing it so the worker keeps " - "processing tasks." - ) - try: - broken.shutdown(wait=False, cancel_futures=True) - except Exception: - pass - self._executor = self._new_executor() - - async def process_task( - task, parent_ctx, logger: logging.Logger, limiter, loop, pool: "PoolManager" + task, parent_ctx, logger: logging.Logger, limiter, loop, pool: ProcessPoolManager ): """Process a given task and ACK in redis. @@ -621,7 +564,11 @@ async def poll_for_tasks(): cpu_count = os.cpu_count() cpu_count = cpu_count if cpu_count is not None else 1 cpu_count = min(cpu_count, TASK_LIMIT) - pool = PoolManager(cpu_count) + pool = ProcessPoolManager( + cpu_count, + max_tasks_per_child=OMNICORP_MAX_TASKS_PER_CHILD, + name="aragorn.omnicorp process pool", + ) while True: try: async for task, parent_ctx, logger, limiter in get_tasks( diff --git a/workers/aragorn_score/worker.py b/workers/aragorn_score/worker.py index 4ff512d..190f844 100644 --- a/workers/aragorn_score/worker.py +++ b/workers/aragorn_score/worker.py @@ -8,13 +8,13 @@ import os import uuid from collections import defaultdict -from concurrent.futures import ProcessPoolExecutor from itertools import combinations import numpy as np from shepherd_utils.db import get_message_sync, save_message_sync from shepherd_utils.otel import setup_tracer +from shepherd_utils.process_pool import ProcessPoolManager from shepherd_utils.shared import get_tasks, run_task_lifecycle # Queue name @@ -1229,23 +1229,21 @@ def aragorn_score(response_id: str, logger: logging.Logger) -> None: save_message_sync(response_id, in_message) -async def process_task(task, parent_ctx, logger, limiter, loop, executor): +async def process_task(task, parent_ctx, logger, limiter, loop, pool): """Process a given task and ACK in redis. Scoring is CPU-bound, so it is dispatched to a process pool while the span, wrap-up, and error handling are shared with every worker. Only the ``response_id`` is handed to the child; the message load/save happen there (see ``aragorn_score``) so the payload never crosses the process boundary. + + Dispatch goes through ``pool`` (a ProcessPoolManager) so a child dying on an + oversized message replaces the pool instead of poisoning it for good. """ async def _run(task, logger): response_id = task[1]["response_id"] - await loop.run_in_executor( - executor, - aragorn_score, - response_id, - logger, - ) + await pool.run(loop, aragorn_score, response_id, logger) await run_task_lifecycle(STREAM, GROUP, task, parent_ctx, logger, limiter, _run) @@ -1256,14 +1254,14 @@ async def poll_for_tasks(): cpu_count = os.cpu_count() cpu_count = cpu_count if cpu_count is not None else 1 cpu_count = min(cpu_count, TASK_LIMIT) - executor = ProcessPoolExecutor(max_workers=cpu_count) + pool = ProcessPoolManager(cpu_count, name="aragorn.score process pool") while True: try: async for task, parent_ctx, logger, limiter in get_tasks( STREAM, GROUP, CONSUMER, TASK_LIMIT ): asyncio.create_task( - process_task(task, parent_ctx, logger, limiter, loop, executor) + process_task(task, parent_ctx, logger, limiter, loop, pool) ) except asyncio.CancelledError: logging.info("Poll loop cancelled, shutting down.") diff --git a/workers/arax_rank/worker.py b/workers/arax_rank/worker.py index d9a03ac..401a773 100644 --- a/workers/arax_rank/worker.py +++ b/workers/arax_rank/worker.py @@ -15,10 +15,10 @@ import logging import os import uuid -from concurrent.futures import ProcessPoolExecutor from shepherd_utils.db import get_message, save_message from shepherd_utils.otel import setup_tracer +from shepherd_utils.process_pool import ProcessPoolManager from shepherd_utils.shared import get_tasks, run_task_lifecycle from ranker import arax_rank @@ -78,11 +78,14 @@ def rank_message(in_message: dict, logger: logging.Logger) -> dict: return in_message -async def process_task(task, parent_ctx, logger, limiter, loop, executor): +async def process_task(task, parent_ctx, logger, limiter, loop, pool): """Process a given task and ACK in redis. Ranking is CPU-bound, so it is dispatched to a process pool while the span, wrap-up, and error handling are shared with every worker. + + Dispatch goes through ``pool`` (a ProcessPoolManager) so a child dying on an + oversized message replaces the pool instead of poisoning it for good. """ async def _run(task, logger): @@ -90,12 +93,7 @@ async def _run(task, logger): message = await get_message(response_id, logger) if message is not None: # Run ranking in process pool for CPU-intensive operations - ranked_message = await loop.run_in_executor( - executor, - rank_message, - message, - logger, - ) + ranked_message = await pool.run(loop, rank_message, message, logger) if ranked_message is None: logger.error("Ranking returned None. Returning original message.") ranked_message = message @@ -110,14 +108,14 @@ async def poll_for_tasks() -> None: """ Main loop to poll for and process ranking tasks. - Creates a single ProcessPoolExecutor that is reused across all tasks + Creates a single self-healing process pool that is reused across all tasks for better performance. """ loop = asyncio.get_running_loop() cpu_count = os.cpu_count() cpu_count = cpu_count if cpu_count is not None else 1 cpu_count = min(cpu_count, TASK_LIMIT) - executor = ProcessPoolExecutor(max_workers=cpu_count) + pool = ProcessPoolManager(cpu_count, name="arax.rank process pool") while True: try: @@ -125,7 +123,7 @@ async def poll_for_tasks() -> None: STREAM, GROUP, CONSUMER, TASK_LIMIT ): asyncio.create_task( - process_task(task, parent_ctx, logger, limiter, loop, executor) + process_task(task, parent_ctx, logger, limiter, loop, pool) ) except asyncio.CancelledError: logging.info("Poll loop cancelled, shutting down.") From eaf45db7cba3f0324efde7da793adbaa6580061e Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 10 Jul 2026 17:12:34 +0000 Subject: [PATCH 3/5] Add configurable request-size limit to /callback endpoint The callback endpoint buffers each subservice's posted TRAPI response into memory before decoding it. An oversized payload can OOM the server (and, once merged, the downstream CPU-bound workers). Reject bodies larger than a configurable limit with 413 before the whole body is read. _read_body_within_limit checks the declared Content-Length first (so well-behaved clients are rejected without buffering anything), then reads the stream in chunks and aborts the moment the running total crosses the limit, so a missing or dishonest Content-Length can't force an unbounded buffer. Enforced in the shared callback() so every ARA's /callback route is covered. The limit is set via callback_max_request_size (Kubernetes-style sizes like "100MB"/"512Mi" or a plain byte count), default 100MB; 0 disables it. Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_01TiiFCcToq8LgiEeaQhQM5e --- shepherd_server/base_routes.py | 45 ++++++++++++- shepherd_utils/config.py | 12 ++++ tests/unit/test_callback_size_limit.py | 90 ++++++++++++++++++++++++++ 3 files changed, 146 insertions(+), 1 deletion(-) create mode 100644 tests/unit/test_callback_size_limit.py diff --git a/shepherd_server/base_routes.py b/shepherd_server/base_routes.py index 64d2e14..eb839f0 100644 --- a/shepherd_server/base_routes.py +++ b/shepherd_server/base_routes.py @@ -247,13 +247,56 @@ async def run_async_query( ) +async def _read_body_within_limit(request: Request, max_bytes: int): + """Read the request body, aborting if it exceeds ``max_bytes``. + + Returns the body bytes, or ``None`` if the limit was exceeded. A + ``max_bytes`` of 0 disables the limit and reads the whole body. + + The declared ``Content-Length`` is checked first so well-behaved clients are + rejected without buffering anything; the stream is then read in chunks and + aborted the moment the running total crosses the limit, so a missing or + dishonest ``Content-Length`` can't force us to buffer an unbounded body. + """ + if max_bytes <= 0: + return await request.body() + + content_length = request.headers.get("content-length") + if content_length is not None: + try: + if int(content_length) > max_bytes: + return None + except ValueError: + pass + + chunks = [] + total = 0 + async for chunk in request.stream(): + total += len(chunk) + if total > max_bytes: + return None + chunks.append(chunk) + return b"".join(chunks) + + async def callback( target: ARATargetEnum, callback_id: str, request: Request, ) -> Response: """Handle asynchronous callback queries from subservices.""" - raw = await request.body() + max_bytes = settings.callback_max_request_size_bytes + raw = await _read_body_within_limit(request, max_bytes) + if raw is None: + return JSONResponse( + content={ + "detail": ( + f"Request body exceeds the maximum allowed size of {max_bytes} " + "bytes." + ) + }, + status_code=413, + ) try: if "zstd" in request.headers.get("content-encoding", "").lower(): raw = decompress_zstd(raw) diff --git a/shepherd_utils/config.py b/shepherd_utils/config.py index f12f14d..2ada699 100644 --- a/shepherd_utils/config.py +++ b/shepherd_utils/config.py @@ -64,6 +64,13 @@ class Settings(BaseSettings): lookup_timeout: int = 210 callback_host: str = "http://host.docker.internal:5439" + # Maximum accepted request body size for the ``/callback`` endpoint. + # Subservices POST their TRAPI responses there; an oversized payload risks + # OOMing the server buffering it and, once merged, the downstream CPU-bound + # workers. Anything larger is rejected with 413 before it is read into + # memory. Accepts Kubernetes-style sizes ("100MB", "512Mi", ...) or a plain + # byte count; 0 (or unparseable) disables the limit. + callback_max_request_size: str = "100MB" kg_retrieval_url: str = "http://host.docker.internal:8080/asyncquery" sync_kg_retrieval_url: str = "http://host.docker.internal:8080/query" kg_rehydrate_url: str = "http://host.docker.internal:8080/rehydrate" @@ -187,6 +194,11 @@ def pg_volume_capacity_bytes(self) -> int: """Configured Postgres volume size in bytes (0 if unset).""" return parse_size_to_bytes(self.pg_volume_capacity) + @property + def callback_max_request_size_bytes(self) -> int: + """Max accepted ``/callback`` body in bytes (0 disables the limit).""" + return parse_size_to_bytes(self.callback_max_request_size) + class Config: env_file = ".env" diff --git a/tests/unit/test_callback_size_limit.py b/tests/unit/test_callback_size_limit.py new file mode 100644 index 0000000..d3e24bb --- /dev/null +++ b/tests/unit/test_callback_size_limit.py @@ -0,0 +1,90 @@ +"""Tests for the /callback request-size limit. + +The callback endpoint buffers the posted TRAPI response into memory, so an +oversized payload can OOM the server (and, once merged, the downstream workers). +``callback`` rejects anything larger than ``callback_max_request_size`` with a +413 before the whole body is read. +""" + +import pytest +from starlette.requests import Request + +from shepherd_server.base_routes import ARATargetEnum, _read_body_within_limit, callback +from shepherd_utils.config import settings + + +def _make_request(body: bytes, headers: dict, chunk_size: int | None = None) -> Request: + """Build a Starlette Request that streams ``body`` from an ASGI receive. + + ``chunk_size`` splits the body across multiple ``http.request`` messages so + the streaming-abort path can be exercised; ``None`` sends it in one chunk. + """ + header_list = [(k.lower().encode(), str(v).encode()) for k, v in headers.items()] + scope = { + "type": "http", + "method": "POST", + "path": "/callback/abc", + "headers": header_list, + } + if chunk_size: + chunks = [body[i : i + chunk_size] for i in range(0, len(body), chunk_size)] + else: + chunks = [body] + if not chunks: + chunks = [b""] + messages = [ + {"type": "http.request", "body": c, "more_body": i < len(chunks) - 1} + for i, c in enumerate(chunks) + ] + it = iter(messages) + + async def receive(): + try: + return next(it) + except StopIteration: + return {"type": "http.disconnect"} + + return Request(scope, receive) + + +async def test_reads_full_body_when_under_limit(): + body = b'{"message": {}}' + request = _make_request(body, {"content-length": len(body)}) + assert await _read_body_within_limit(request, 1000) == body + + +async def test_rejects_on_declared_content_length(): + # Fast path: the declared Content-Length alone is over the limit, so we + # reject without reading the (here intentionally absent) body. + request = _make_request(b"", {"content-length": 5000}) + assert await _read_body_within_limit(request, 1000) is None + + +async def test_rejects_when_stream_exceeds_limit_without_content_length(): + # No Content-Length header: the cap must still trip mid-stream. + body = b"x" * 5000 + request = _make_request(body, {}, chunk_size=512) + assert await _read_body_within_limit(request, 1000) is None + + +async def test_lying_content_length_is_caught_by_stream_cap(): + # Content-Length under-reports; the streaming total still enforces the cap. + body = b"x" * 5000 + request = _make_request(body, {"content-length": 10}, chunk_size=512) + assert await _read_body_within_limit(request, 1000) is None + + +async def test_zero_limit_disables_the_cap(): + body = b"x" * 5000 + request = _make_request(body, {"content-length": len(body)}) + assert await _read_body_within_limit(request, 0) == body + + +async def test_callback_returns_413_for_oversized_payload(monkeypatch): + monkeypatch.setattr(settings, "callback_max_request_size", "1000") + body = b"x" * 5000 + request = _make_request(body, {"content-length": len(body)}) + + response = await callback(ARATargetEnum.ARAGORN, "cb-1", request) + + assert response.status_code == 413 From 33d757ea671b5487d8207dd96fc5d287dc85b41f Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 10 Jul 2026 17:50:50 +0000 Subject: [PATCH 4/5] Drop rejected oversized callback so the lookup worker doesn't hang A callback only leaves the per-query running set (the postgres callbacks table that get_running_callbacks reads) once merge_message has processed it. When /callback rejects an oversized payload with 413 it never reaches merge_message, so the callback stayed "running" and the lookup worker blocked waiting for it until its whole-query timeout. Call remove_callback_id on the 413 path so the rejected callback is dropped from the running set immediately and the lookup proceeds with whatever else came back, exactly as it already does when a lookup POST fails. Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_01TiiFCcToq8LgiEeaQhQM5e --- shepherd_server/base_routes.py | 11 +++++++++++ tests/unit/test_callback_size_limit.py | 16 +++++++++++++++- 2 files changed, 26 insertions(+), 1 deletion(-) diff --git a/shepherd_server/base_routes.py b/shepherd_server/base_routes.py index eb839f0..e2e42de 100644 --- a/shepherd_server/base_routes.py +++ b/shepherd_server/base_routes.py @@ -24,6 +24,7 @@ get_logs, get_message, get_query_state, + remove_callback_id, save_message, ) from shepherd_utils.logger import QueryLogger, setup_logging @@ -288,6 +289,16 @@ async def callback( max_bytes = settings.callback_max_request_size_bytes raw = await _read_body_within_limit(request, max_bytes) if raw is None: + logger = logging.getLogger(f"shepherd.{callback_id}") + logger.warning( + f"Rejecting callback {callback_id}: request body exceeds the maximum " + f"allowed size of {max_bytes} bytes." + ) + # Drop this callback from the running set so the lookup worker stops + # waiting on it. Without this the lookup blocks until its whole-query + # timeout, since a callback only leaves the set once merge_message has + # processed it -- which never happens for a payload we refused to read. + await remove_callback_id(callback_id, logger) return JSONResponse( content={ "detail": ( diff --git a/tests/unit/test_callback_size_limit.py b/tests/unit/test_callback_size_limit.py index d3e24bb..2e219c3 100644 --- a/tests/unit/test_callback_size_limit.py +++ b/tests/unit/test_callback_size_limit.py @@ -9,6 +9,7 @@ import pytest from starlette.requests import Request +from shepherd_server import base_routes from shepherd_server.base_routes import ARATargetEnum, _read_body_within_limit, callback from shepherd_utils.config import settings @@ -80,11 +81,24 @@ async def test_zero_limit_disables_the_cap(): assert await _read_body_within_limit(request, 0) == body -async def test_callback_returns_413_for_oversized_payload(monkeypatch): +async def test_callback_returns_413_and_drops_callback_for_oversized_payload( + monkeypatch, +): monkeypatch.setattr(settings, "callback_max_request_size", "1000") + removed = [] + + async def _fake_remove(callback_id, logger): + removed.append(callback_id) + + # Patch the name as imported into base_routes. + monkeypatch.setattr(base_routes, "remove_callback_id", _fake_remove) + body = b"x" * 5000 request = _make_request(body, {"content-length": len(body)}) response = await callback(ARATargetEnum.ARAGORN, "cb-1", request) assert response.status_code == 413 + # The rejected callback must be dropped from the running set so the lookup + # worker doesn't hang waiting for it until its timeout. + assert removed == ["cb-1"] From ee4593fd85427b948f5338d1483ef48028798114 Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 10 Jul 2026 18:15:32 +0000 Subject: [PATCH 5/5] Stop merge_message lock-losers from re-enqueuing (log spew + churn) /callback enqueues one wake task per callback and adds the callback to the query's ready set. The first worker to grab the query lock becomes the holder and drains the whole ready set; every other wake task for that query is a loser. A loser used to re-enqueue itself after merge_contention_backoff (0.1s) whenever its callback was still pending, so during a single burst merge each contended callback bounced through the stream ~10x/second while the holder drained. The shared get_tasks helper logs "Doing task " at INFO on every redelivery, and re-enqueued wakes carry the same callback_id/query_id, so the logs filled with hundreds of near-identical "Doing task" lines for what was really one callback being retried -- no duplicate merging, just wake-task churn. Have losers simply ack and drop: the holder already drains the ready set to empty, so a loser has nothing to add. To cover the narrow race where a callback lands after the holder's final drain pass but before it releases the lock, the holder now re-checks the ready set once after unlocking and kicks a single wake if anything remains. Holder crashes are still recovered by reclaim (60s min-idle) once the 45s lock TTL lapses. Net effect: at most one re-enqueue per drain instead of a storm, and one "Doing task" per callback instead of hundreds. Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_01TiiFCcToq8LgiEeaQhQM5e --- workers/merge_message/worker.py | 45 +++++++++++++++++++++------------ 1 file changed, 29 insertions(+), 16 deletions(-) diff --git a/workers/merge_message/worker.py b/workers/merge_message/worker.py index 83c6af8..e0b1eac 100644 --- a/workers/merge_message/worker.py +++ b/workers/merge_message/worker.py @@ -27,7 +27,6 @@ get_message, get_message_sync, get_ready_callbacks, - is_ready_callback, remove_callback_id, save_message_sync, ) @@ -801,22 +800,19 @@ async def process_query(task, parent_ctx, logger, limiter): # drains the whole query, so a loser has nothing useful to add. got_lock = await try_lock(response_id, CONSUMER, logger) if not got_lock: - still_pending = await is_ready_callback( - response_id, callback_id, logger + # Someone else holds this query's lock and drains its ready + # set to empty, so our callback (added to the set before this + # wake task was enqueued) will be merged by them. Just ack and + # move on -- no re-enqueue. The holder does one final ready-set + # check after releasing the lock (below) to catch a callback + # that lands in the narrow window past its last drain pass, so + # nothing is stranded. This replaces the old per-loser + # re-enqueue, which spun a wake task per contended callback + # every merge_contention_backoff and flooded the logs with + # near-identical "Doing task" lines during a single merge. + logger.debug( + f"[{callback_id}] Lock busy; holder will drain it. Acking." ) - if still_pending: - # The holder may finish just before it would have swept - # us; re-enqueue (after a short backoff) so we're retried. - await asyncio.sleep(settings.merge_contention_backoff) - await _reenqueue_wake_task(task, logger) - logger.debug( - f"[{callback_id}] Lock busy; re-enqueued wake task." - ) - else: - # Already merged by the lock holder -- nothing to do. - logger.debug( - f"[{callback_id}] Already merged by holder; acking." - ) return logger.info(f"[{callback_id}] Obtained lock for {response_id}.") @@ -884,6 +880,23 @@ async def process_query(task, parent_ctx, logger, limiter): f"{time.time() - lock_time:.2f}s" ) await remove_lock(response_id, CONSUMER, logger) + # Close the race where a callback landed in the ready set after + # our final drain pass read it empty but before we released the + # lock: that callback's own wake task would have found the lock + # held and dropped (losers no longer re-enqueue). Now that the + # lock is free, re-check and kick exactly one wake if anything + # remains so the late arrival still gets merged. One conditional + # re-enqueue replaces the old storm of per-loser re-enqueues. + try: + leftover = await get_ready_callbacks(response_id, logger) + except Exception: + leftover = [] + if leftover: + logger.debug( + f"[{callback_id}] {len(leftover)} callback(s) arrived " + "post-drain; kicking one wake task." + ) + await _reenqueue_wake_task(task, logger) except Exception as e: logger.error( f"Task {task[0]} failed with unhandled error: {e}", exc_info=True