Claude/omnicorp worker pool error ytpcn2

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
@@ -247,13 +248,66 @@ 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:
+        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": (
+                    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)

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"
 

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

tests/unit/test_callback_size_limit.py

@@ -0,0 +1,104 @@
+"""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 import base_routes
+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_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"]

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