Claude/score paths worker cpu xxdxvp

README.md

@@ -127,9 +127,9 @@ single large pod. On Kubernetes the memory `limit` (OOMKilled + restart) plus
 regular rollouts already recycle pods, so leaked-resource cleanup comes for free
 — add an RSS-based `livenessProbe` only if the monitor shows OOMKills in
 practice. CPU-bound pool workers (`merge_message`, `score_paths`, `arax_rank`,
-`aragorn_score`, `aragorn_omnicorp`) size their process/thread pools from the
-in-code default, so raising `TASK_LIMIT` for those only deepens the intake queue
-rather than adding parallelism.
+`aragorn_score`, `aragorn_omnicorp`) size their process pools from the in-code
+default, so raising `TASK_LIMIT` for those only deepens the intake queue rather
+than adding parallelism.
 
 ### Message Broker Streams
 

shepherd_utils/config.py

@@ -230,6 +230,14 @@ class Settings(BaseSettings):
     # POOL_TASK_TIMEOUT_SEC; 0 disables the timeout.
     pool_task_timeout_sec: float = 300.0
 
+    # score_paths gets a tighter ceiling than the shared default: scoring runs
+    # at the tail of a query whose lookups were already bounded by
+    # lookup_timeout, so a scoring that outlives that budget is past the point
+    # of being useful to the client. Matching the two keeps the worst case for
+    # a single query's scoring stage predictable. Per-Deployment override via
+    # SCORE_PATHS_TASK_TIMEOUT_SEC; 0 disables the timeout.
+    score_paths_task_timeout_sec: float = 210.0
+
     # Recycle each process-pool child after this many tasks. A child that once
     # processed a very large message keeps that peak RSS for its whole life
     # (freed memory isn't fully returned to the OS), so long-lived children

shepherd_utils/reclaim.py

@@ -52,10 +52,13 @@
     # Pathfinding runs in a process pool bounded by pool_task_timeout_sec
     # (300s), so no legitimate task can outlive that; the floor sits just above.
     "arax.pathfinder": 360,
+    # Scoring runs in a process pool bounded by score_paths_task_timeout_sec
+    # (210s), so the floor sits just above that -- same shape as the lookup
+    # workers, whose ceiling is the identical 210s lookup_timeout.
+    "score_paths": 240,
     # Medium-duration workers.
     "arax.rank": 60,
     "merge_message": 60,
-    "score_paths": 60,
     "example.score": 30,
     # finish_query sends the async callback, which retries with backoff and can
     # legitimately run for minutes against a slow callback endpoint (httpx

tests/unit/test_worker_dispatch_concurrency.py

@@ -1,18 +1,27 @@
-"""Regression guard: CPU-bound workers must dispatch tasks concurrently.
+"""Regression guards: how CPU-bound workers must run their heavy work.
 
-``arax_rank`` and ``aragorn_score`` previously awaited ``process_task`` inline
-inside their ``poll_for_tasks`` ``async for`` body, which serialized every task
-through their ProcessPoolExecutor (only one task ever ran at a time -- the same
-bug fixed in ``merge_message``). The corrected code dispatches each task with
-``asyncio.create_task(process_task(...))``, matching the ``filter_results_top_n``
-template.
+Two invariants, both learned the hard way, both pinned statically here.
+
+**Dispatch concurrently.** ``arax_rank`` and ``aragorn_score`` previously
+awaited ``process_task`` inline inside their ``poll_for_tasks`` ``async for``
+body, which serialized every task through their ProcessPoolExecutor (only one
+task ever ran at a time -- the same bug fixed in ``merge_message``). The
+corrected code dispatches each task with ``asyncio.create_task(process_task
+(...))``, matching the ``filter_results_top_n`` template.
+
+**Offload to a process pool, not a thread pool.** These workers' hot loops are
+mostly pure Python, so in a ``ThreadPoolExecutor`` they hold the GIL and starve
+the event loop the heartbeat pings from. Once the heartbeat goes stale past
+``HEARTBEAT_TTL_SEC`` a peer stops treating the worker as alive and reclaims its
+in-flight tasks; past ``worker_loop_stall_exit_sec`` the loop watchdog restarts
+the pod outright. ``arax_pathfinder`` hit exactly this on ``asyncio.to_thread``
+and ``score_paths`` on a ``ThreadPoolExecutor``; both now use
+``ProcessPoolManager``, which also brings the OOM self-heal and per-task timeout.
 
 ``poll_for_tasks`` is an unbounded ``while True`` loop whose ``CancelledError``
 handler intentionally does not return (the shared worker template), so it can't
-be run to completion in a unit test; and ``arax_rank`` isn't importable outside
-its container. So we pin the dispatch shape statically on the source of
-``poll_for_tasks``: it must wrap ``process_task`` in ``asyncio.create_task`` and
-must not ``await`` it inline.
+be run to completion in a unit test; and these workers aren't importable outside
+their containers. So both invariants are pinned on the worker sources instead.
 """
 
 from pathlib import Path
@@ -25,11 +34,16 @@
     "workers/arax_rank/worker.py",
     "workers/aragorn_score/worker.py",
     "workers/arax_pathfinder/worker.py",
+    "workers/score_paths/worker.py",
 ]
 
 
+def _worker_source(worker_file: str) -> str:
+    return (REPO_ROOT / worker_file).read_text()
+
+
 def _poll_for_tasks_source(worker_file: str) -> str:
-    text = (REPO_ROOT / worker_file).read_text()
+    text = _worker_source(worker_file)
     start = text.index("async def poll_for_tasks")
     # poll_for_tasks is the last definition before the __main__ guard.
     end = text.index('if __name__ == "__main__"', start)
@@ -51,3 +65,52 @@ def test_poll_for_tasks_dispatches_concurrently(worker_file):
         "serializes every task through the process pool (the merge_message bug). "
         "Dispatch with asyncio.create_task(process_task(...)) instead."
     )
+
+
+@pytest.mark.parametrize("worker_file", WORKER_FILES)
+def test_cpu_bound_work_runs_in_a_process_pool(worker_file):
+    src = _worker_source(worker_file)
+
+    assert "ProcessPoolManager" in src, (
+        f"{worker_file}: CPU-bound work must be offloaded with "
+        "ProcessPoolManager; found no reference to it."
+    )
+    # Thread-pool bug shape: heavy pure-Python work sharing the GIL with the
+    # event loop, which starves the heartbeat and gets the worker's tasks
+    # reclaimed while it is still very much alive.
+    # Matched on the call, not the bare name, so the docstrings explaining why
+    # these workers moved off a thread pool don't trip their own guard.
+    assert "ThreadPoolExecutor(" not in src, (
+        f"{worker_file}: offloads to a ThreadPoolExecutor, whose threads hold "
+        "the GIL against the event loop and can starve the heartbeat past "
+        "HEARTBEAT_TTL_SEC. Use ProcessPoolManager instead."
+    )
+    assert "asyncio.to_thread(" not in src, (
+        f"{worker_file}: offloads with asyncio.to_thread, which has the same "
+        "GIL/heartbeat problem as a ThreadPoolExecutor. Use ProcessPoolManager."
+    )
+
+
+def test_score_paths_scores_in_bounded_chunks():
+    """score_paths must not materialise a whole message's features at once.
+
+    Scoring every analysis in one batch held the float16 row list, its stacked
+    copy and the float32 cast live simultaneously -- 24 GiB for an observed
+    387k-analysis message, which cgroup-OOM-killed the pod (SIGKILL, so no
+    traceback: the logs simply stopped mid-task). Chunking bounds that at
+    SCORE_CHUNK_SIZE rows regardless of message size.
+
+    This is a string check because CI cannot import the module (torch, lmdb and
+    bmt live only in the worker's image), so there is no other guard on it.
+    """
+    src = _worker_source("workers/score_paths/worker.py")
+
+    assert (
+        "SCORE_CHUNK_SIZE" in src
+    ), "score_paths must score in bounded chunks; SCORE_CHUNK_SIZE is gone."
+    # The one-shot shape: stacking every feature row into a single array.
+    assert "np.stack(" not in src, (
+        "score_paths stacks all feature rows into one array again. That scales "
+        "peak memory with the message and OOM-killed the pod on large ones; "
+        "score in SCORE_CHUNK_SIZE batches instead."
+    )