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/tests/unit/test_worker_dispatch_concurrency.py b/tests/unit/test_worker_dispatch_concurrency.py index faec548..c9b28ce 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,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." + ) diff --git a/workers/score_paths/worker.py b/workers/score_paths/worker.py index d23a89a..0169554 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,28 @@ CONSUMER = str(uuid.uuid4())[:8] 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) +# 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 +146,137 @@ 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(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(FEATURE_DIM, 1536), + nn.GELU(), + nn.LayerNorm(1536), + nn.Linear(1536, 1536), + nn.GELU(), + nn.LayerNorm(1536), + nn.Linear(1536, 1), + ) + 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 + + +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. + + 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: + 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) + 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: paths = message["message"]["query_graph"]["paths"] @@ -141,12 +291,38 @@ def score_paths(task, 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): @@ -196,11 +372,14 @@ def score_paths(task, 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} " @@ -211,24 +390,11 @@ def score_paths(task, 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") @@ -251,46 +417,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.score_paths_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: