From 67bcf3dc3f840cc32a57415d742d40b879b73b71 Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 1 Sep 2026 01:56:11 +0000 Subject: [PATCH 1/4] Run score_paths scoring in a process pool, not a thread pool score_paths was the last CPU-bound worker still offloading to a ThreadPoolExecutor, so its four concurrent scorings ran in the same process as the event loop the heartbeat pings from. Most of a scoring's time is the feature build -- convert_path_to_components walks every edge of every path in pure Python -- which holds the GIL; only the LMDB reads and the torch forward release it. Under load the loop thread's turnaround degrades from milliseconds to seconds, and once the heartbeat goes unrefreshed past HEARTBEAT_TTL_SEC (15s) peers stop counting the worker as alive and can XCLAIM its in-flight tasks out from under it. Past worker_loop_stall_exit_sec (60s) the loop watchdog force-exits the pod. This is the same failure arax_pathfinder hit on asyncio.to_thread and aragorn_score / arax_rank hit before them; nothing here was different except that it hadn't been migrated yet. Scoring now goes through ProcessPoolManager like its siblings: - score_paths_task is the child entrypoint. Only the response_id and the task's log level cross the boundary -- the message is loaded, scored and saved inside the child, so the payload never lands on the parent's heap either. - The child attaches its own QueryLogHandler and hands the formatted records back with the result; the parent folds them into the task's query logger, so the per-query scoring lines (feature build stats, score ranges) still reach the query's log list rather than only container stderr. - Per-child state (biolink Toolkit, embeddings LMDB, the MLP) is built on first use rather than in the pool initializer, so a bad checkpoint or an unreadable LMDB fails one task with a traceback instead of killing children at startup and leaving the pool rebuilding itself in a loop. Each child caps torch to one intra-op thread: the pool is already sized to the pod's CPU allocation, so a full thread pool per child just oversubscribes it. - The read-only LMDB is opened with lock=False in every child, so the pages are shared through the page cache instead of copied per child. - The parent validates the embeddings cache and the weights file at startup, before any child spawns, so a bad mount still fails fast instead of surfacing as every task failing one at a time. This also brings the OOM self-heal and the pool_task_timeout_sec (300s) per-task ceiling to this stream; scoring previously had no timeout at all, so a pathological message could hold its slot indefinitely. Pool size comes from resolve_pool_workers (cgroup-aware, POOL_MAX_WORKERS overrides) and doubles as the in-flight task limit, matching arax_rank. The dispatch-concurrency regression test now covers score_paths, and gains a second invariant: these workers must offload through ProcessPoolManager, and must not instantiate a ThreadPoolExecutor or use asyncio.to_thread. Co-Authored-By: Claude Claude-Session: https://claude.ai/code/session_01XpZWD8DD34dqV3mMkJHzky --- .../unit/test_worker_dispatch_concurrency.py | 62 ++++-- workers/score_paths/worker.py | 183 +++++++++++++++--- 2 files changed, 202 insertions(+), 43 deletions(-) diff --git a/tests/unit/test_worker_dispatch_concurrency.py b/tests/unit/test_worker_dispatch_concurrency.py index faec548..0e2ce84 100644 --- a/tests/unit/test_worker_dispatch_concurrency.py +++ b/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,27 @@ 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." + ) diff --git a/workers/score_paths/worker.py b/workers/score_paths/worker.py index d23a89a..3e2a9ce 100644 --- a/workers/score_paths/worker.py +++ b/workers/score_paths/worker.py @@ -1,10 +1,10 @@ """Path scoring module""" import asyncio +import logging +import os import time import uuid -from concurrent.futures import ThreadPoolExecutor -from functools import partial import lmdb import numpy as np @@ -13,10 +13,12 @@ from torch import nn from shepherd_utils.config import settings +from shepherd_utils.cpu import resolve_pool_workers from shepherd_utils.data_download import ensure_pathfinder_embeddings from shepherd_utils.db import get_message_sync, save_message_sync -from shepherd_utils.logger import get_worker_logger +from shepherd_utils.logger import QueryLogger, get_query_handler, get_worker_logger from shepherd_utils.otel import setup_tracer +from shepherd_utils.process_pool import ProcessPoolManager from shepherd_utils.shared import get_tasks, run_task_lifecycle STREAM = "score_paths" @@ -24,9 +26,17 @@ CONSUMER = str(uuid.uuid4())[:8] TASK_LIMIT = 4 EMBEDDING_DIR = settings.pathfinder_embeddings_dir +MODEL_WEIGHTS = "model_weights/squashbert_direct_3hop.pt" tracer = setup_tracer(STREAM) LOGGER = get_worker_logger(STREAM) +# Per-child scoring state, built on first use by ``_ensure_scoring_state``. +# These live in the process-pool children, not the parent: the parent only +# validates the data at startup and never scores anything itself. +bmt = None +embedding_env = None +mlp = None + def convert_path_to_components(source, target, path, knowledge_graph, logger): try: @@ -125,8 +135,76 @@ def _probe_cache(env): return n, key.decode("utf-8", errors="replace") -def score_paths(task, logger): - response_id = task[1]["response_id"] +def _open_embeddings(): + """Open the embeddings LMDB read-only. + + ``lock=False`` on a read-only env is what lets every pool child map the same + database concurrently; the pages are shared through the OS page cache rather + than copied per child. + """ + return lmdb.open( + EMBEDDING_DIR, readonly=True, lock=False, readahead=False, subdir=True + ) + + +def _build_mlp(): + """Build the scoring MLP and load its trained weights.""" + model = nn.Sequential( + nn.Linear(11 * 768, 1536), + nn.GELU(), + nn.LayerNorm(1536), + nn.Linear(1536, 1536), + nn.GELU(), + nn.LayerNorm(1536), + nn.Linear(1536, 1), + ) + ckpt = torch.load(MODEL_WEIGHTS, map_location="cpu") + model.load_state_dict({k.removeprefix("net."): v for k, v in ckpt["model"].items()}) + model.eval() + return model + + +def _validate_scoring_data(logger) -> None: + """Fail fast at startup if the data the pool children need is missing. + + The children do the actual loading, so without this an empty volume mount or + a missing checkpoint would surface only as every task failing individually. + The env opened here is closed again immediately -- the parent never scores. + """ + env = _open_embeddings() + try: + count, sample = _probe_cache(env) + finally: + env.close() + logger.info(f"embeddings cache: {count} entries (sample key: {sample!r})") + if not os.path.exists(MODEL_WEIGHTS): + raise RuntimeError(f"model weights not found at {MODEL_WEIGHTS}") + + +def _ensure_scoring_state(logger) -> None: + """Build this child's scoring state on first use, then reuse it. + + Loaded lazily rather than through the pool's ``initializer`` so a failure + here (an unreadable LMDB, a corrupt checkpoint) surfaces as an ordinary task + failure with a traceback, instead of killing the child before it takes any + work and leaving the pool to rebuild itself in a loop. Each child pays this + once and amortizes it over ``pool_max_tasks_per_child`` tasks. + """ + global bmt, embedding_env, mlp + if mlp is not None: + return + # One intra-op thread per child. The pool is already sized to the pod's CPU + # allocation, so letting each child spin up a full torch thread pool + # oversubscribes that quota several times over and the children mostly end + # up contending with each other. + torch.set_num_threads(1) + bmt = Toolkit() + embedding_env = _open_embeddings() + mlp = _build_mlp() + logger.debug(f"score_paths child {os.getpid()} loaded its scoring state.") + + +def score_paths(response_id, logger): message = get_message_sync(response_id) try: paths = message["message"]["query_graph"]["paths"] @@ -251,46 +329,89 @@ def score_paths(task, logger): logger.error(f"Failed to save a message into redis: {e}") -async def process_task(task, parent_ctx, logger, limiter): +def score_paths_task(response_id: str, log_level: int = logging.INFO) -> list[dict]: + """Process-pool entrypoint: load, score, and save entirely in the child. + + Only the small ``response_id`` and the task's log level cross the process + boundary; the (potentially very large) message is read from Redis, scored, + and written back inside the child. That keeps the payload off the parent's + heap and -- more importantly -- keeps the feature build off the parent's + event loop. It used to run in a ``ThreadPoolExecutor``, where the path walk + is mostly pure Python and so holds the GIL: a few concurrent scorings could + starve the heartbeat past ``HEARTBEAT_TTL_SEC`` and get a live worker's + tasks reclaimed out from under it (matching the fix already applied to + arax_pathfinder / aragorn_score / arax_rank). + + Returns this child's log records, already formatted and oldest-first, for + the parent to fold into the query's logs -- the child can't reach the + parent's query log handler itself. + """ + # logging.getLogger hands back the same object for the whole life of the + # child, so attach a call-scoped handler and remove it in finally -- + # otherwise handlers accumulate across the child's successive tasks and one + # query's logs leak into the next. + query_log_handler = QueryLogger().log_handler + logger = get_worker_logger(f"{STREAM}.worker.{os.getpid()}") + logger.setLevel(log_level) + logger.addHandler(query_log_handler) + try: + _ensure_scoring_state(logger) + score_paths(response_id, logger) + return query_log_handler.drain() + finally: + logger.removeHandler(query_log_handler) + + +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 stay shared with every other worker. The child's + log records come back with the result and are folded into this task's query + logger so they still reach the query's log list. + """ + async def _run(task, logger): - loop = asyncio.get_event_loop() - await loop.run_in_executor(executor, partial(score_paths, task, logger)) + response_id = task[1]["response_id"] + entries = await pool.run( + loop, score_paths_task, response_id, logger.getEffectiveLevel() + ) + handler = get_query_handler(logger) + if handler is not None and entries: + handler.ingest(entries) await run_task_lifecycle(STREAM, GROUP, task, parent_ctx, logger, limiter, _run) async def poll_for_tasks(): - global bmt, mlp, embedding_env, executor - # Ensure the embeddings LMDB exists before we open it below (a first-run + loop = asyncio.get_running_loop() + # Ensure the embeddings LMDB exists before any child opens it (a first-run # local `docker compose up` starts with the volume-mounted directory empty). # No-op once present or when no download URL is configured (e.g. production, - # where the data is mounted out of band). + # where the data is mounted out of band). Downloading in the parent also + # keeps the children from racing each other for the same archive. ensure_pathfinder_embeddings(LOGGER) - bmt = Toolkit() - embedding_env = lmdb.open( - EMBEDDING_DIR, readonly=True, lock=False, readahead=False, subdir=True - ) - count, sample = _probe_cache(embedding_env) - LOGGER.info(f"embeddings cache: {count} entries (sample key: {sample!r})") - mlp = nn.Sequential( - nn.Linear(11 * 768, 1536), - nn.GELU(), - nn.LayerNorm(1536), - nn.Linear(1536, 1536), - nn.GELU(), - nn.LayerNorm(1536), - nn.Linear(1536, 1), + _validate_scoring_data(LOGGER) + # Size the pool by the pod's actual CPU allocation (cgroup limit), not + # os.cpu_count() -- see aragorn_omnicorp.poll_for_tasks. Each child holds its + # own copy of the MLP plus a full message, so this also bounds peak memory. + # POOL_MAX_WORKERS overrides. + max_workers = resolve_pool_workers(TASK_LIMIT, LOGGER) + LOGGER.info(f"{STREAM}: process pool sized to {max_workers} worker(s).") + pool = ProcessPoolManager( + max_workers, + max_tasks_per_child=settings.pool_max_tasks_per_child, + name="score_paths process pool", + task_timeout=settings.pool_task_timeout_sec, ) - ckpt = torch.load("model_weights/squashbert_direct_3hop.pt", map_location="cpu") - mlp.load_state_dict({k.removeprefix("net."): v for k, v in ckpt["model"].items()}) - mlp.eval() - executor = ThreadPoolExecutor(max_workers=TASK_LIMIT) while True: try: async for task, parent_ctx, logger, limiter in get_tasks( - STREAM, GROUP, CONSUMER, TASK_LIMIT + STREAM, GROUP, CONSUMER, max_workers ): - asyncio.create_task(process_task(task, parent_ctx, logger, limiter)) + asyncio.create_task( + process_task(task, parent_ctx, logger, limiter, loop, pool) + ) except asyncio.CancelledError: LOGGER.info("Poll loop cancelled, shutting down.") except Exception as e: From e9f91dc8289383ffa01a4ef365f393495eae6e29 Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 1 Sep 2026 02:38:06 +0000 Subject: [PATCH 2/4] Cap score_paths tasks at 210s score_paths now bounds a scoring run with its own score_paths_task_timeout_sec (210s) rather than the shared 300s pool_task_timeout_sec. Scoring runs at the tail of a query whose lookups were already bounded by the identical 210s lookup_timeout, so a scoring that outlives that budget is past the point of being useful to the client. Per-Deployment override via SCORE_PATHS_TASK_TIMEOUT_SEC. The reclaim idle floor for the stream moves 60s -> 240s to match. That floor has to exceed the worst-case legitimate task duration or a peer can XCLAIM a message out from under a worker that is simply slow, and the ceiling now makes that worst case an explicit 210s. This is the same shape the lookup workers already use: a 210s internal timeout with the floor just above it. Also drops the "process/thread pools" wording from the README's pool-worker note, since score_paths was the thread-pool one. --- README.md | 6 +++--- shepherd_utils/config.py | 8 ++++++++ shepherd_utils/reclaim.py | 5 ++++- workers/score_paths/worker.py | 2 +- 4 files changed, 16 insertions(+), 5 deletions(-) diff --git a/README.md b/README.md index aa2f94e..01c6554 100644 --- a/README.md +++ b/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 diff --git a/shepherd_utils/config.py b/shepherd_utils/config.py index 6f69970..3aa86c7 100644 --- a/shepherd_utils/config.py +++ b/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 diff --git a/shepherd_utils/reclaim.py b/shepherd_utils/reclaim.py index e858878..31d523f 100644 --- a/shepherd_utils/reclaim.py +++ b/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 diff --git a/workers/score_paths/worker.py b/workers/score_paths/worker.py index 3e2a9ce..f7b00c1 100644 --- a/workers/score_paths/worker.py +++ b/workers/score_paths/worker.py @@ -402,7 +402,7 @@ async def poll_for_tasks(): max_workers, max_tasks_per_child=settings.pool_max_tasks_per_child, name="score_paths process pool", - task_timeout=settings.pool_task_timeout_sec, + task_timeout=settings.score_paths_task_timeout_sec, ) while True: try: From a2a1414e1f8c2377267decd69391cbddecd27dd0 Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 1 Sep 2026 03:15:36 +0000 Subject: [PATCH 3/4] Share the score_paths weights across pool children via mmap Moving scoring into a process pool gave each child its own copy of the 61 MB checkpoint, where the thread pool had loaded it once. Memory-mapping the weights gets that back: torch.load(mmap=True) returns file-backed MAP_PRIVATE tensors, and load_state_dict(assign=True) makes those tensors be the module's parameters rather than a destination to copy into -- the default allocates fresh storage per child and undoes the sharing. Scoring only reads them (eval mode, inference_mode), so nothing triggers a copy-on-write fault and the pages stay shared for the life of the pod. Verified against the real checkpoint: parameters and the forward pass are bit-identical to the plain load, the file shows up in /proc/self/maps, and two spawned children that each run a full forward pass report Shared_Clean 58.5 MB / Private_Dirty 0.0 MB for the mapping -- the pages are shared, not copied. mmap needs torch's zipfile checkpoint format (its default since 1.6). A checkpoint re-saved in the legacy format raises RuntimeError, so that falls back to a plain private-copy load with a warning rather than failing every task; the fallback path is exercised and the exception type confirmed. The biolink Toolkit remains genuinely per-child -- live Python objects, no equivalent trick -- which is now noted in _ensure_scoring_state alongside the two things that are shared, since POOL_MAX_WORKERS is the lever for it. --- workers/score_paths/worker.py | 43 +++++++++++++++++++++++++++++++---- 1 file changed, 38 insertions(+), 5 deletions(-) diff --git a/workers/score_paths/worker.py b/workers/score_paths/worker.py index f7b00c1..7f43046 100644 --- a/workers/score_paths/worker.py +++ b/workers/score_paths/worker.py @@ -147,8 +147,23 @@ def _open_embeddings(): ) -def _build_mlp(): - """Build the scoring MLP and load its trained weights.""" +def _build_mlp(logger): + """Build the scoring MLP and load its trained weights. + + The weights are memory-mapped and assigned rather than copied, so the pool's + children share one 61 MB mapping instead of each allocating its own copy. + ``mmap=True`` hands back file-backed ``MAP_PRIVATE`` tensors; ``assign=True`` + makes those tensors *be* the module's parameters instead of a destination to + copy into (the default would allocate fresh storage per child and undo the + sharing). Scoring only ever reads them -- the model is in ``eval`` mode under + ``inference_mode`` -- so nothing triggers a copy-on-write fault and the pages + stay shared for the life of the pod. + + ``mmap=True`` needs a checkpoint in torch's zipfile format (the default since + torch 1.6). A checkpoint re-saved in the legacy format would raise, so fall + back to a plain load: that child then pays for its own copy, which is the + pre-mmap behaviour and strictly better than failing every task. + """ model = nn.Sequential( nn.Linear(11 * 768, 1536), nn.GELU(), @@ -158,8 +173,19 @@ def _build_mlp(): nn.LayerNorm(1536), nn.Linear(1536, 1), ) - ckpt = torch.load(MODEL_WEIGHTS, map_location="cpu") - model.load_state_dict({k.removeprefix("net."): v for k, v in ckpt["model"].items()}) + try: + ckpt = torch.load(MODEL_WEIGHTS, map_location="cpu", mmap=True) + assign = True + except (RuntimeError, ValueError) as e: + logger.warning( + f"Could not memory-map {MODEL_WEIGHTS} ({e}); loading a private copy " + "of the weights instead. Every pool child will hold its own." + ) + ckpt = torch.load(MODEL_WEIGHTS, map_location="cpu") + assign = False + model.load_state_dict( + {k.removeprefix("net."): v for k, v in ckpt["model"].items()}, assign=assign + ) model.eval() return model @@ -189,6 +215,13 @@ def _ensure_scoring_state(logger) -> None: failure with a traceback, instead of killing the child before it takes any work and leaving the pool to rebuild itself in a loop. Each child pays this once and amortizes it over ``pool_max_tasks_per_child`` tasks. + + Two of the three are shared across children rather than duplicated: the + embeddings LMDB and the model weights are both file-backed mappings, so the + OS page cache serves every child from one copy (see ``_open_embeddings`` and + ``_build_mlp``). The biolink ``Toolkit`` is live Python objects and so is + genuinely per-child -- the one place pool size costs real memory, which is + why ``POOL_MAX_WORKERS`` exists for a memory-tight deployment. """ global bmt, embedding_env, mlp if mlp is not None: @@ -200,7 +233,7 @@ def _ensure_scoring_state(logger) -> None: torch.set_num_threads(1) bmt = Toolkit() embedding_env = _open_embeddings() - mlp = _build_mlp() + mlp = _build_mlp(logger) logger.debug(f"score_paths child {os.getpid()} loaded its scoring state.") From 7d107757e58da6a7fb1ef5e54c74344944693046 Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 1 Sep 2026 10:52:26 +0000 Subject: [PATCH 4/4] Score paths in bounded chunks instead of one batch per message Production logs show a pod dying mid-task with no traceback: a cgroup OOM SIGKILL, which no in-process handler can catch (a Python MemoryError would have hit the "Error scoring paths" handler and logged). Four tasks were in flight, three of them 387,583 analyses each. Scoring a message in one pass held three representations of the same data at once -- the float16 row list, np.stack's copy of it, and the float32 cast: analyses rows stack float32 peak 138,824 2.18 2.18 4.37 8.74 GiB 387,583 6.10 6.10 12.20 24.40 GiB 79,906 1.26 1.26 2.51 5.03 GiB The last line logged was a 79,906-analysis feature build completing; the very next statement was its np.stack(...).astype(np.float32), a 1.26 GiB stack plus a 2.51 GiB allocation on top of three 387k row lists that had been accumulating for 138s. Those messages then outlived their worker, got reclaimed, killed the next pod the same way, and were finally dead-lettered by the poison-pill breaker after three deliveries. Scoring now runs every SCORE_CHUNK_SIZE (4096) rows and frees the chunk, so peak is ~208 MB whatever the message size instead of scaling with it. Rows are copied straight into a float32 batch rather than stacked as float16 and cast, so only one array of the batch exists at a time; float16 converts exactly, so the input matrix is unchanged. The summary log lines keep their shape, with build and MLP time accumulated separately now that the two interleave and the score range tracked as running aggregates rather than a per-analysis list. Verified against the real checkpoint at 1, 100, 4095, 4096, 4097, 8192 and 10000 rows: every score is bit-identical to the one-shot path, the reported count/min/max/mean match, and the largest float32 batch allocated for a 60k-row message is 4096 rows (132 MiB) rather than 60000 (1.89 GiB). Note this bounds the scoring memory, not the message itself: a 387k-analysis TRAPI payload still has to be decoded into the child's heap. What changes is that the multi-GiB feature arrays no longer sit on top of it, the cost is one child's rather than the parent's, and it is returned to the OS when the child recycles. The chunking is pinned by a static test, since CI cannot import this module -- torch, lmdb and bmt live only in the worker image. --- .../unit/test_worker_dispatch_concurrency.py | 25 +++++ workers/score_paths/worker.py | 103 ++++++++++++++---- 2 files changed, 104 insertions(+), 24 deletions(-) diff --git a/tests/unit/test_worker_dispatch_concurrency.py b/tests/unit/test_worker_dispatch_concurrency.py index 0e2ce84..c9b28ce 100644 --- a/tests/unit/test_worker_dispatch_concurrency.py +++ b/tests/unit/test_worker_dispatch_concurrency.py @@ -89,3 +89,28 @@ def test_cpu_bound_work_runs_in_a_process_pool(worker_file): 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." + ) diff --git a/workers/score_paths/worker.py b/workers/score_paths/worker.py index 7f43046..0169554 100644 --- a/workers/score_paths/worker.py +++ b/workers/score_paths/worker.py @@ -27,6 +27,17 @@ TASK_LIMIT = 4 EMBEDDING_DIR = settings.pathfinder_embeddings_dir MODEL_WEIGHTS = "model_weights/squashbert_direct_3hop.pt" +# 11 embeddings of 768 dims each (4 node names, 4 categories, 3 hop phrases) -- +# the MLP's input width, and the width of every feature row. +FEATURE_DIM = 11 * 768 +# Analyses scored per forward pass. Peak memory is bounded by this rather than +# by the message's size: one float32 batch (4096 x 8448 x 4 = 138 MB) plus the +# float16 rows still pending in the chunk (69 MB), and both are freed at the end +# of each chunk. Scoring the whole message in one pass instead is what +# OOM-killed the pod on large messages -- a 387k-analysis message needed the +# row list, its stacked copy and the float32 cast all live at once, 24 GiB in +# total. The MLP is row-independent, so chunk size changes only the batching. +SCORE_CHUNK_SIZE = 4096 tracer = setup_tracer(STREAM) LOGGER = get_worker_logger(STREAM) @@ -165,7 +176,7 @@ def _build_mlp(logger): pre-mmap behaviour and strictly better than failing every task. """ model = nn.Sequential( - nn.Linear(11 * 768, 1536), + nn.Linear(FEATURE_DIM, 1536), nn.GELU(), nn.LayerNorm(1536), nn.Linear(1536, 1536), @@ -237,6 +248,34 @@ def _ensure_scoring_state(logger) -> None: logger.debug(f"score_paths child {os.getpid()} loaded its scoring state.") +def _score_chunk(rows, index, results): + """Score one chunk of feature rows and write the scores onto the analyses. + + The rows are copied straight into a float32 batch rather than stacked as + float16 and cast afterwards, so only one array of the batch exists at a + time. float16 converts to float32 exactly, so this is the same input the + stack-then-cast path produced. + + Returns ``(count, minimum, maximum, total)`` for the chunk, letting the + caller keep running statistics for the summary log line without holding + every score of a large message in a list. + """ + features = np.empty((len(rows), FEATURE_DIM), dtype=np.float32) + for i, row in enumerate(rows): + features[i] = row + with torch.inference_mode(): + logits = mlp(torch.from_numpy(features)).squeeze(-1) + scores = torch.sigmoid(logits).numpy() + for (result_ind, analysis_ind), score in zip(index, scores): + results[result_ind]["analyses"][analysis_ind]["score"] = float(score) + return ( + len(scores), + float(scores.min()), + float(scores.max()), + float(scores.sum(dtype=np.float64)), + ) + + def score_paths(response_id, logger): message = get_message_sync(response_id) try: @@ -252,12 +291,38 @@ def score_paths(response_id, logger): f"Scoring {response_id}: {len(results)} results, " f"{total_analyses} analyses, {len(auxiliary_graphs)} aux graphs" ) - feature_rows = [] - embedding_index = [] + chunk_rows = [] + chunk_index = [] skip_no_binding = 0 skip_bad_path = 0 skip_missing_emb = 0 missing_samples = [] + # Running totals for the summary lines. Scoring is interleaved with the + # feature build now, so the two timings are accumulated separately + # rather than measured as consecutive phases. + mlp_time = 0.0 + scored = 0 + score_min = float("inf") + score_max = float("-inf") + score_sum = 0.0 + + def flush_chunk(): + """Score the pending rows, write them back, and free the chunk.""" + nonlocal mlp_time, scored, score_min, score_max, score_sum + if not chunk_rows: + return + started = time.time() + count, lowest, highest, total = _score_chunk( + chunk_rows, chunk_index, results + ) + mlp_time += time.time() - started + scored += count + score_min = min(score_min, lowest) + score_max = max(score_max, highest) + score_sum += total + chunk_rows.clear() + chunk_index.clear() + t0 = time.time() with embedding_env.begin() as txn: for result_ind, result in enumerate(results): @@ -307,11 +372,14 @@ def score_paths(response_id, logger): analysis["score"] = 0.0 skip_missing_emb += 1 continue - feature_rows.append(features) - embedding_index.append((result_ind, analysis_ind)) - build_time = time.time() - t0 + chunk_rows.append(features) + chunk_index.append((result_ind, analysis_ind)) + if len(chunk_rows) >= SCORE_CHUNK_SIZE: + flush_chunk() + flush_chunk() + build_time = time.time() - t0 - mlp_time skipped = skip_no_binding + skip_bad_path + skip_missing_emb - msg = f"Feature build: {len(feature_rows)}/{total_analyses} ready in {build_time:.1f}s" + msg = f"Feature build: {scored}/{total_analyses} ready in {build_time:.1f}s" if skipped: msg += ( f"; skipped {skipped} " @@ -322,24 +390,11 @@ def score_paths(response_id, logger): if missing_samples: msg += f"; missing keys e.g. {missing_samples}" logger.info(msg) - if feature_rows: - features = np.stack(feature_rows).astype(np.float32) - t0 = time.time() - with torch.inference_mode(): - logits = mlp(torch.from_numpy(features)).squeeze(-1) - all_scores = torch.sigmoid(logits).numpy() - mlp_time = time.time() - t0 - - scores = [] - for (r_idx, a_idx), s in zip(embedding_index, all_scores): - s = float(s) - results[r_idx]["analyses"][a_idx]["score"] = s - scores.append(s) - + if scored: logger.info( - f"Scored {len(scores)} paths in {mlp_time:.1f}s; " - f"scores [{min(scores):.3f}, {max(scores):.3f}] " - f"mean {sum(scores) / len(scores):.3f}" + f"Scored {scored} paths in {mlp_time:.1f}s; " + f"scores [{score_min:.3f}, {score_max:.3f}] " + f"mean {score_sum / scored:.3f}" ) else: logger.info("No paths to score")