Claude/score paths worker cpu xxdxvp - #147
Merged
Merged
Conversation
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 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01XpZWD8DD34dqV3mMkJHzky
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.
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.
Codecov Report❌ Patch coverage is
... and 3 files with indirect coverage changes Continue to review full report in Codecov by Harness.
🚀 New features to boost your workflow:
|
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.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
No description provided.