From 8c26f47ea60501f5d69d80c3fb48ea79dd473ded Mon Sep 17 00:00:00 2001 From: Claude Date: Mon, 10 Aug 2026 16:14:20 +0000 Subject: [PATCH 1/3] Move arax_pathfinder onto a process pool and run_task_lifecycle The pathfinding search ran via asyncio.to_thread, so its GIL-bound work competed with the event loop -- and TASK_LIMIT was 100, meaning up to 100 concurrent searches. The heartbeat is a coroutine on that same loop pinging every 5s against a 15s freshness window, so under load a live worker read as dead: reclaim_orphaned could XCLAIM a still-running task to a peer and duplicate it, and a stall past 60s tripped LoopWatchdog into os._exit(1). Restructure to match aragorn_score / arax_rank: - arax_pathfinder_task is a sync process-pool entrypoint taking only the query/response ids. It loads, searches, rehydrates, and saves entirely in a spawned child, so neither the search nor the message encode touches the parent's loop or heap. rehydrate becomes sync (httpx.Client) since it now runs in the child. - process_task delegates to the shared run_task_lifecycle via a pool closure, replacing the hand-rolled copy that used tracer.start_span instead of start_as_current_span (so httpx spans never nested under the task span) and never recorded exceptions on the span. - poll_for_tasks builds a ProcessPoolManager sized by resolve_pool_workers and passes max_workers to get_tasks; TASK_LIMIT drops 100 -> 10. Failures now reach the lifecycle instead of being swallowed. The three query-graph validation failures returned `message, 500`, a value the caller discarded, so the task was wrapped up as a success with no message ever written to response_id -- leaving the next worker to KeyError on it. Search failures were saved as a non-TRAPI {"status": "error"} blob that then flowed down the workflow. Both now raise and route to finish_query with ERROR. Also: - The blocked-concept list is fetched once at startup via ensure_http_files_dataset (temp file + atomic rename) rather than with a per-task requests.get racing on the same relative path; children parse it once and memoize, as they do the BiolinkHelper descendants. - Add "arax.pathfinder" to PER_STREAM_MIN_IDLE_SEC. It was absent, so its reclaim floor was the 30s fast default; 360s sits above pool_task_timeout_sec. - Dockerfile gets PYTHONFAULTHANDLER=1 so a child hung during import still dumps a stack when ProcessPoolManager SIGABRTs it. - Fix rehydrate's no-op `.replace("query", "rehydrate")` on an already-correct URL, and its references to a possibly-unbound `res` in the error handlers. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01H51TTxq1F6o6yZSP1C86MZ --- shepherd_utils/reclaim.py | 3 + tests/unit/test_arax_pathfinder.py | 254 ++++++++++++ .../unit/test_worker_dispatch_concurrency.py | 1 + workers/arax_pathfinder/Dockerfile | 5 + workers/arax_pathfinder/worker.py | 368 +++++++++++------- 5 files changed, 485 insertions(+), 146 deletions(-) create mode 100644 tests/unit/test_arax_pathfinder.py diff --git a/shepherd_utils/reclaim.py b/shepherd_utils/reclaim.py index 54121c6..e858878 100644 --- a/shepherd_utils/reclaim.py +++ b/shepherd_utils/reclaim.py @@ -49,6 +49,9 @@ "aragorn.score": 240, "bte.lookup": 240, "example.lookup": 240, + # 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, # Medium-duration workers. "arax.rank": 60, "merge_message": 60, diff --git a/tests/unit/test_arax_pathfinder.py b/tests/unit/test_arax_pathfinder.py new file mode 100644 index 0000000..f572207 --- /dev/null +++ b/tests/unit/test_arax_pathfinder.py @@ -0,0 +1,254 @@ +"""Tests for the arax_pathfinder process-pool entrypoint. + +The worker imports ``pathfinder.Pathfinder`` and ``biolink_helper_pkg``, which +are installed only inside its container (see +``workers/arax_pathfinder/requirements.txt``) and are not in +``test-requirements.txt``. We stub both into ``sys.modules`` before importing +the worker, the same trick ``tests/conftest.py`` uses for ``shepherd_utils.otel``. +""" + +import copy +import logging +import sys +import types +from unittest.mock import MagicMock + +import pytest + +if "pathfinder" not in sys.modules: + _pathfinder_pkg = types.ModuleType("pathfinder") + _pathfinder_mod = types.ModuleType("pathfinder.Pathfinder") + _pathfinder_mod.Pathfinder = MagicMock(name="Pathfinder") + _pathfinder_pkg.Pathfinder = _pathfinder_mod + sys.modules["pathfinder"] = _pathfinder_pkg + sys.modules["pathfinder.Pathfinder"] = _pathfinder_mod + +if "biolink_helper_pkg" not in sys.modules: + _biolink_mod = types.ModuleType("biolink_helper_pkg") + _biolink_mod.BiolinkHelper = MagicMock(name="BiolinkHelper") + sys.modules["biolink_helper_pkg"] = _biolink_mod + +from workers.arax_pathfinder import worker as pf_worker # noqa: E402 + +LOGGER = logging.getLogger(__name__) + +QUERY = { + "message": { + "query_graph": { + "nodes": { + "n0": {"ids": ["MONDO:0005148"]}, + "n1": {"ids": ["CHEBI:15365"]}, + }, + "paths": {"p0": {"subject": "n0", "object": "n1", "constraints": []}}, + } + } +} + +PATHS_RESULT = ( + {"id": "r0", "analyses": [{"score": 1.0}], "node_bindings": {"n0": []}}, + {"aux0": {"edges": ["e0"]}}, + { + "nodes": {"MONDO:0005148": {}}, + "edges": {"e0": {"predicate": "biolink:related_to"}}, + }, +) + + +def _patch_query(mocker, query=None): + return mocker.patch( + "workers.arax_pathfinder.worker.get_message_sync", + return_value=copy.deepcopy(query if query is not None else QUERY), + ) + + +def test_pathfinder_task_searches_rehydrates_and_saves(mocker): + """The process-pool entrypoint reads by id, searches, and writes back. + + Only the two ids cross into the child: the message is loaded with + ``get_message_sync``, assembled, and persisted with ``save_message_sync`` -- + the knowledge graph never has to be pickled back to the parent. + """ + _patch_query(mocker) + search = mocker.patch( + "workers.arax_pathfinder.worker.execute_pathfinding", + return_value=copy.deepcopy(PATHS_RESULT), + ) + rehydrated_kg = {"nodes": {}, "edges": {"e0": {"predicate": "biolink:related_to"}}} + rehydrate = mocker.patch( + "workers.arax_pathfinder.worker.rehydrate", return_value=rehydrated_kg + ) + save = mocker.patch("workers.arax_pathfinder.worker.save_message_sync") + + pf_worker.arax_pathfinder_task("query-1", "resp-1", LOGGER) + + search.assert_called_once() + rehydrate.assert_called_once() + save.assert_called_once() + saved_id, message = save.call_args.args + assert saved_id == "resp-1" + assert message["message"]["knowledge_graph"] is rehydrated_kg + assert message["message"]["auxiliary_graphs"] == {"aux0": {"edges": ["e0"]}} + assert len(message["message"]["results"]) == 1 + assert message["message"]["results"][0]["essence"] == "result" + # Provenance is injected before saving. + assert message["message"]["knowledge_graph"]["edges"]["e0"]["sources"] == [ + { + "resource_id": "infores:shepherd-arax", + "resource_role": "aggregator_knowledge_source", + "source_record_urls": None, + "upstream_resource_ids": ["infores:arax"], + } + ] + # Defaults are filled in on the message that gets saved. + assert message["parameters"]["tiers"] == [0] + + +def test_pathfinder_task_saves_empty_graphs_when_no_paths_found(mocker): + """A search that finds nothing still writes a well-formed TRAPI message.""" + _patch_query(mocker) + mocker.patch( + "workers.arax_pathfinder.worker.execute_pathfinding", + return_value=(None, None, None), + ) + mocker.patch("workers.arax_pathfinder.worker.rehydrate", return_value=None) + save = mocker.patch("workers.arax_pathfinder.worker.save_message_sync") + + pf_worker.arax_pathfinder_task("query-2", "resp-2", LOGGER) + + _, message = save.call_args.args + assert message["message"]["results"] == [] + assert message["message"]["auxiliary_graphs"] == {} + assert message["message"]["knowledge_graph"] == {} + + +@pytest.mark.parametrize( + "qgraph, expected", + [ + pytest.param( + { + "nodes": {"n0": {"ids": ["MONDO:0005148"]}, "n1": {}}, + "paths": {"p0": {"constraints": []}}, + }, + "two pinned nodes", + id="one_pinned_node", + ), + pytest.param( + { + "nodes": { + "n0": {"ids": ["MONDO:0005148"]}, + "n1": {"ids": ["CHEBI:15365"]}, + }, + "paths": { + "p0": { + "constraints": [ + {"intermediate_categories": ["biolink:Gene"]}, + {"intermediate_categories": ["biolink:Drug"]}, + ] + } + }, + }, + "multiple constraints", + id="multiple_constraints", + ), + pytest.param( + { + "nodes": { + "n0": {"ids": ["MONDO:0005148"]}, + "n1": {"ids": ["CHEBI:15365"]}, + }, + "paths": { + "p0": { + "constraints": [ + { + "intermediate_categories": [ + "biolink:Gene", + "biolink:Drug", + ] + } + ] + } + }, + }, + "multiple intermediate categories", + id="multiple_intermediate_categories", + ), + ], +) +def test_unanswerable_query_graph_raises(mocker, qgraph, expected): + """Validation failures raise so run_task_lifecycle can fail the query. + + They used to ``return message, 500``, a value the caller discarded -- so the + task was wrapped up as a success and no message was ever written to + ``response_id``, leaving the next worker to ``KeyError`` on it. Raising + routes the query to ``finish_query`` with an ERROR status instead. + """ + _patch_query(mocker, {"message": {"query_graph": qgraph}}) + save = mocker.patch("workers.arax_pathfinder.worker.save_message_sync") + search = mocker.patch("workers.arax_pathfinder.worker.execute_pathfinding") + + with pytest.raises(ValueError, match=expected): + pf_worker.arax_pathfinder_task("query-3", "resp-3", LOGGER) + + search.assert_not_called() + save.assert_not_called() + + +def test_search_failure_propagates_instead_of_saving_error_blob(mocker): + """A failed search raises rather than persisting a non-TRAPI error blob. + + The old code caught everything and saved ``{"status": "error", ...}`` to the + response id, then reported success -- so that blob flowed on through the + workflow. Now the exception reaches ``run_task_lifecycle``. + """ + _patch_query(mocker) + mocker.patch( + "workers.arax_pathfinder.worker.execute_pathfinding", + side_effect=RuntimeError("sqlite is on fire"), + ) + save = mocker.patch("workers.arax_pathfinder.worker.save_message_sync") + + with pytest.raises(RuntimeError, match="sqlite is on fire"): + pf_worker.arax_pathfinder_task("query-4", "resp-4", LOGGER) + + save.assert_not_called() + + +def test_rehydrate_failure_propagates(mocker): + """A rehydrate failure fails the task too, rather than saving a partial KG.""" + _patch_query(mocker) + mocker.patch( + "workers.arax_pathfinder.worker.execute_pathfinding", + return_value=copy.deepcopy(PATHS_RESULT), + ) + mocker.patch( + "workers.arax_pathfinder.worker.rehydrate", + side_effect=RuntimeError("retriever unreachable"), + ) + save = mocker.patch("workers.arax_pathfinder.worker.save_message_sync") + + with pytest.raises(RuntimeError, match="retriever unreachable"): + pf_worker.arax_pathfinder_task("query-5", "resp-5", LOGGER) + + save.assert_not_called() + + +def test_descendants_are_memoized_per_child(mocker): + """The Biolink model is built once per pool child, not once per task. + + ``BiolinkHelper`` construction parses the model, and the old code did it + inside every task alongside a fresh blocked-list HTTP fetch. + """ + mocker.patch.object(pf_worker, "_biolink_helper", None) + mocker.patch.object(pf_worker, "_descendants_cache", {}) + helper = MagicMock() + helper.get_descendants.return_value = ["biolink:Gene", "biolink:Protein"] + helper_cls = mocker.patch( + "workers.arax_pathfinder.worker.BiolinkHelper", return_value=helper + ) + + first = pf_worker.get_descendants("biolink:NamedThing") + second = pf_worker.get_descendants("biolink:NamedThing") + + assert first == second == {"biolink:Gene", "biolink:Protein"} + helper_cls.assert_called_once() + helper.get_descendants.assert_called_once_with("biolink:NamedThing") diff --git a/tests/unit/test_worker_dispatch_concurrency.py b/tests/unit/test_worker_dispatch_concurrency.py index 5dbead9..faec548 100644 --- a/tests/unit/test_worker_dispatch_concurrency.py +++ b/tests/unit/test_worker_dispatch_concurrency.py @@ -24,6 +24,7 @@ WORKER_FILES = [ "workers/arax_rank/worker.py", "workers/aragorn_score/worker.py", + "workers/arax_pathfinder/worker.py", ] diff --git a/workers/arax_pathfinder/Dockerfile b/workers/arax_pathfinder/Dockerfile index 68ae2b6..890d662 100644 --- a/workers/arax_pathfinder/Dockerfile +++ b/workers/arax_pathfinder/Dockerfile @@ -5,6 +5,11 @@ FROM ghcr.io/translatorsri/renci-python-image:3.12.13 LABEL org.opencontainers.image.source https://github.com/BioPack-team/shepherd ENV PYTHONHASHSEED=0 +# Enable faulthandler from interpreter startup so a SIGABRT dumps every thread's +# traceback -- including during a spawned pool child's import phase, before the +# pool initializer's faulthandler.enable() would run. Lets ProcessPoolManager's +# SIGABRT-on-timeout capture a stuck child's stack even if it hangs in startup. +ENV PYTHONFAULTHANDLER=1 # set up requirements WORKDIR /app diff --git a/workers/arax_pathfinder/worker.py b/workers/arax_pathfinder/worker.py index b9fc501..e1e8e9f 100644 --- a/workers/arax_pathfinder/worker.py +++ b/workers/arax_pathfinder/worker.py @@ -8,35 +8,33 @@ from pathlib import Path import httpx -import requests from biolink_helper_pkg import BiolinkHelper from pathfinder.Pathfinder import Pathfinder from shepherd_utils.config import settings +from shepherd_utils.cpu import resolve_pool_workers from shepherd_utils.data_download import ( arax_pathfinder_sqlite_paths, ensure_arax_pathfinder_dbs, + ensure_http_files_dataset, ) from shepherd_utils.db import ( - get_message, - save_message, + get_message_sync, + save_message_sync, ) from shepherd_utils.inject_shepherd_arax_provenance import ( add_shepherd_arax_to_edge_sources, ) from shepherd_utils.otel import setup_tracer -from shepherd_utils.shared import ( - get_tasks, - handle_task_failure, - wrap_up_task, -) +from shepherd_utils.process_pool import ProcessPoolManager +from shepherd_utils.shared import get_tasks, run_task_lifecycle # Queue name STREAM = "arax.pathfinder" # Consumer group, most likely you don't need to change this. GROUP = "consumer" CONSUMER = str(uuid.uuid4())[:8] -TASK_LIMIT = 100 +TASK_LIMIT = 10 tracer = setup_tracer(STREAM) NUM_TOTAL_HOPS = 4 @@ -45,34 +43,78 @@ PRUNE_TOP_K = 75 NODE_DEGREE_THRESHOLD = 10000 -OUT_PATH = Path("general_concepts.json") - - -def download_file(url: str, out_path: Path, overwrite: bool = False) -> Path: - out_path = Path(out_path) - - if out_path.exists() and not overwrite: - return out_path - - out_path.parent.mkdir(parents=True, exist_ok=True) - - r = requests.get(url, timeout=60) - r.raise_for_status() - - out_path.write_bytes(r.content) - return out_path +# The ARAX blocked-concept list, fetched once at worker startup (see +# poll_for_tasks) and read by each pool child. Kept in the working directory +# (/app in the image) so it needs no extra volume mount. +BLOCKED_LIST_DIR = "." +BLOCKED_LIST_FILENAME = "general_concepts.json" +BLOCKED_LIST_PATH = Path(BLOCKED_LIST_DIR) / BLOCKED_LIST_FILENAME + +BIOLINK_CACHE_DIR = "/tmp/biolink" + +REHYDRATE_TIMEOUT_SEC = 30.0 + +# Per-child caches. Children are spawned once and reused for up to +# ``settings.pool_max_tasks_per_child`` tasks, so the blocked list and the +# Biolink model are parsed once per child instead of once per task (previously +# every task re-read the JSON and rebuilt the BiolinkHelper, and re-fetched the +# blocked list over HTTP). Only ``Pathfinder`` is still built per task, because +# it takes the task's logger. +_blocked_list_cache = None +_biolink_helper = None +_descendants_cache: dict = {} + + +def ensure_blocked_list(logger: logging.Logger) -> None: + """Fetch the ARAX blocked-concept list if it isn't on disk yet. + + Uses the shared downloader so the file lands via a temp file + atomic + rename; the previous per-task ``requests.get`` wrote the destination + directly, so concurrent tasks could race on a half-written file. + Idempotent, so it is safe to call at startup and again lazily in a child. + """ + ensure_http_files_dataset( + name="arax_blocked_list", + target_dir=BLOCKED_LIST_DIR, + file_sources={BLOCKED_LIST_FILENAME: settings.arax_blocked_list_url}, + logger=logger, + ) -def get_blocked_list(): - download_file(settings.arax_blocked_list_url, OUT_PATH, False) +def get_blocked_list(logger: logging.Logger): + """``(blocked_curies, blocked_synonyms)``, parsed once per pool child.""" + global _blocked_list_cache + if _blocked_list_cache is None: + if not BLOCKED_LIST_PATH.exists(): + # Startup fetch failed or this child outlived a wiped working dir. + ensure_blocked_list(logger) + with open(BLOCKED_LIST_PATH, "r") as file: + json_block_list = json.load(file) + synonyms = set(s.lower() for s in json_block_list["synonyms"]) + _blocked_list_cache = (set(json_block_list["curies"]), synonyms) + return _blocked_list_cache + + +def get_descendants(category: str): + """Biolink descendants of ``category``, memoized per pool child.""" + global _biolink_helper + if category not in _descendants_cache: + if _biolink_helper is None: + Path(BIOLINK_CACHE_DIR).mkdir(parents=True, exist_ok=True) + _biolink_helper = BiolinkHelper( + settings.arax_biolink_version, BIOLINK_CACHE_DIR + ) + _descendants_cache[category] = set(_biolink_helper.get_descendants(category)) + return _descendants_cache[category] - with open(OUT_PATH, "r") as file: - json_block_list = json.load(file) - synonyms = set(s.lower() for s in json_block_list["synonyms"]) - return set(json_block_list["curies"]), synonyms +def rehydrate(kg, rehydrate_url, logger): + """POST the knowledge graph to the retriever and return the rehydrated one. -async def rehydrate(kg, retriever_url, logger): + Synchronous because it runs inside the process-pool child alongside the + pathfinding, so the (potentially very large) knowledge graph never has to + be pickled back to the parent or re-encoded on its event loop. + """ headers = {"Content-Type": "application/json", "Accept": "application/json"} payload = { "message": {"knowledge_graph": kg}, @@ -80,43 +122,79 @@ async def rehydrate(kg, retriever_url, logger): } try: - async with httpx.AsyncClient(timeout=30.0) as client: - res = await client.post( - retriever_url.replace("query", "rehydrate"), - headers=headers, - json=payload, - ) + with httpx.Client(timeout=REHYDRATE_TIMEOUT_SEC) as client: + res = client.post(rehydrate_url, headers=headers, json=payload) res.raise_for_status() return res.json()["message"]["knowledge_graph"] except httpx.HTTPStatusError as http_err: logger.error(f"HTTP error occurred: {http_err}") - if res.text: - logger.error(f"Error details: {res.text}") - raise http_err + if http_err.response.text: + logger.error(f"Error details: {http_err.response.text}") + raise except httpx.ConnectError as conn_err: logger.error(f"Connection error occurred: {conn_err}") - raise conn_err + raise except httpx.TimeoutException as timeout_err: logger.error(f"Timeout error occurred: {timeout_err}") - raise timeout_err + raise except httpx.RequestError as req_err: logger.error(f"An unexpected error occurred: {req_err}") - raise req_err + raise except json.JSONDecodeError: - logger.error("Failed to parse the response as JSON.") - logger.error(f"Raw response: {res.text}") + logger.error("Failed to parse the rehydrate response as JSON.") raise except Exception as e: logger.error(f"An unexpected error occurred: {e}") raise e -def execute_pathfinding_sync( +def parse_query_graph(qgraph): + """Pull the pinned nodes and intermediate category out of the query graph. + + Raises ``ValueError`` on anything Pathfinder can't answer. Raising (rather + than returning a status code, which the old caller discarded) lets + ``run_task_lifecycle`` record the failure on the span and route the query to + ``finish_query`` with an ERROR status, instead of continuing the workflow + with no response message ever written. + """ + pinned_node_keys = [] + pinned_node_ids = [] + for node_key, node in qgraph["nodes"].items(): + pinned_node_keys.append(node_key) + if node.get("ids", None) is not None: + pinned_node_ids.append(node["ids"][0]) + if len(set(pinned_node_ids)) != 2: + raise ValueError("Pathfinder queries require two pinned nodes.") + + intermediate_categories = [] + path_key = next(iter(qgraph["paths"].keys())) + qpath = qgraph["paths"][path_key] + if ( + qpath.get("constraints", None) is not None + and len(qpath.get("constraints", [])) > 0 + ): + constraints = qpath["constraints"] + if len(constraints) > 1: + raise ValueError("Pathfinder queries do not support multiple constraints.") + if len(constraints) > 0: + intermediate_categories = ( + constraints[0].get("intermediate_categories", None) or [] + ) + if len(intermediate_categories) > 1: + raise ValueError( + "Pathfinder queries do not support multiple intermediate categories" + ) + else: + intermediate_categories = ["biolink:NamedThing"] + + return pinned_node_keys, pinned_node_ids, intermediate_categories + + +def execute_pathfinding( pinned_node_ids, pinned_node_keys, intermediate_categories, logger ): - - blocked_curies, blocked_synonyms = get_blocked_list() + blocked_curies, blocked_synonyms = get_blocked_list(logger) curie_ngd_path, node_degree_path = arax_pathfinder_sqlite_paths() pathfinder_instance = Pathfinder( @@ -128,13 +206,10 @@ def execute_pathfinding_sync( logger, ) - biolink_cache_dir = "/tmp/biolink" - Path(biolink_cache_dir).mkdir(parents=True, exist_ok=True) - biolink_helper = BiolinkHelper(settings.arax_biolink_version, biolink_cache_dir) - descendants = set(biolink_helper.get_descendants(intermediate_categories[0])) + descendants = get_descendants(intermediate_categories[0]) start = time.perf_counter() - logger.info("Starting pathfinder.get_paths() in worker thread") + logger.info("Starting pathfinder.get_paths()") result, aux_graphs, knowledge_graph = pathfinder_instance.get_paths( pinned_node_ids[0], @@ -155,130 +230,131 @@ def execute_pathfinding_sync( return result, aux_graphs, knowledge_graph -async def pathfinder(task, logger: logging.Logger): +def arax_pathfinder_task( + query_id: str, response_id: str, logger: logging.Logger +) -> None: + """Process-pool entrypoint: load, search, rehydrate, and save in the child. + + Only the two small ids cross the process-pool boundary; the (potentially + very large) message is read from Redis, searched over, and written back + inside the child. That keeps the payload off the parent's heap and, more + importantly, keeps the graph search off the parent's event loop -- it used + to run via ``asyncio.to_thread``, where the GIL meant a handful of + concurrent searches could starve the heartbeat past HEARTBEAT_TTL_SEC and + get a live worker's tasks reclaimed out from under it (matching + aragorn_score / arax_rank). + """ start = time.time() - query_id = task[1]["query_id"] - response_id = task[1]["response_id"] - message = await get_message(query_id, logger) + message = get_message_sync(query_id) parameters = message.get("parameters") or {} parameters["timeout"] = parameters.get("timeout", settings.lookup_timeout) parameters["tiers"] = parameters.get("tiers") or [0] message["parameters"] = parameters - qgraph = message["message"]["query_graph"] - pinned_node_keys = [] - pinned_node_ids = [] - for node_key, node in qgraph["nodes"].items(): - pinned_node_keys.append(node_key) - if node.get("ids", None) is not None: - pinned_node_ids.append(node["ids"][0]) - if len(set(pinned_node_ids)) != 2: - logger.error("Pathfinder queries require two pinned nodes.") - return message, 500 - - intermediate_categories = [] - path_key = next(iter(qgraph["paths"].keys())) - qpath = qgraph["paths"][path_key] - if ( - qpath.get("constraints", None) is not None - and len(qpath.get("constraints", [])) > 0 - ): - constraints = qpath["constraints"] - if len(constraints) > 1: - logger.error("Pathfinder queries do not support multiple constraints.") - return message, 500 - if len(constraints) > 0: - intermediate_categories = ( - constraints[0].get("intermediate_categories", None) or [] - ) - if len(intermediate_categories) > 1: - logger.error( - "Pathfinder queries do not support multiple intermediate categories" - ) - return message, 500 - else: - intermediate_categories = ["biolink:NamedThing"] + pinned_node_keys, pinned_node_ids, intermediate_categories = parse_query_graph( + message["message"]["query_graph"] + ) try: - result, aux_graphs, knowledge_graph = await asyncio.to_thread( - execute_pathfinding_sync, + result, aux_graphs, knowledge_graph = execute_pathfinding( pinned_node_ids, pinned_node_keys, intermediate_categories, logger, ) - logger.info(f"Rehydrating knowledge graph with retriever") - knowledge_graph = await rehydrate( - knowledge_graph, settings.kg_rehydrate_url, logger - ) - - res = [] - if result is not None: - res.append( - { - "id": result["id"], - "analyses": result["analyses"], - "node_bindings": result["node_bindings"], - "essence": "result", - } - ) - if aux_graphs is None: - aux_graphs = {} - if knowledge_graph is None: - knowledge_graph = {} - message["message"]["knowledge_graph"] = knowledge_graph - message["message"]["auxiliary_graphs"] = aux_graphs - message["message"]["results"] = res - - message = add_shepherd_arax_to_edge_sources(message) - - await save_message(response_id, message, logger) + logger.info("Rehydrating knowledge graph with retriever") + knowledge_graph = rehydrate(knowledge_graph, settings.kg_rehydrate_url, logger) except Exception as e: + # Let the failure reach run_task_lifecycle, which records it on the span + # and routes the query to finish_query with an ERROR status. Previously + # this saved a non-TRAPI {"status": "error"} blob and reported success, + # so the bogus payload flowed on down the workflow. logger.error( - f"PathFinder failed to find paths between {pinned_node_keys[0]} and {pinned_node_keys[1]}. " - f"Error message is: {e}" + f"PathFinder failed to find paths between {pinned_node_keys[0]} and " + f"{pinned_node_keys[1]}. Error message is: {e}" + ) + raise + + res = [] + if result is not None: + res.append( + { + "id": result["id"], + "analyses": result["analyses"], + "node_bindings": result["node_bindings"], + "essence": "result", + } ) - message = {"status": "error", "error": str(e)} - await save_message(response_id, message, logger) + if aux_graphs is None: + aux_graphs = {} + if knowledge_graph is None: + knowledge_graph = {} + message["message"]["knowledge_graph"] = knowledge_graph + message["message"]["auxiliary_graphs"] = aux_graphs + message["message"]["results"] = res + + message = add_shepherd_arax_to_edge_sources(message) + save_message_sync(response_id, message) logger.info(f"Task took {time.time() - start}") -async def process_task(task, parent_ctx, logger: logging.Logger, limiter): - """Process a given task and ACK in redis.""" - start = time.time() - span = tracer.start_span(STREAM, context=parent_ctx) - try: - await pathfinder(task, logger) - # Always wrap up the task to ACK it in the broker - try: - await wrap_up_task(STREAM, GROUP, task, logger) - except Exception as e: - logger.error(f"Task {task[0]}: Failed to wrap up task: {e}") - except asyncio.CancelledError: - logger.warning(f"Task {task[0]} was cancelled") - except Exception as e: - logger.error(f"Task {task[0]} failed with unhandled error: {e}", exc_info=True) - await handle_task_failure(STREAM, GROUP, task, logger) - finally: - span.end() - limiter.release() - logger.info(f"Finished task {task[0]} in {time.time() - start}") +async def process_task(task, parent_ctx, logger, limiter, loop, pool): + """Process a given task and ACK in redis. + + Pathfinding is CPU-bound, so it is dispatched to a process pool while the + span, wrap-up, and error handling are shared with every worker. Only the + ids are handed to the child; the message load/save happen there (see + ``arax_pathfinder_task``) so the payload never crosses the process boundary. + + Dispatch goes through ``pool`` (a ProcessPoolManager) so a child dying on an + oversized message replaces the pool instead of poisoning it for good, and a + search that runs away is time-bounded rather than holding its slot forever. + """ + + async def _run(task, logger): + await pool.run( + loop, + arax_pathfinder_task, + task[1]["query_id"], + task[1]["response_id"], + logger, + ) + + await run_task_lifecycle(STREAM, GROUP, task, parent_ctx, logger, limiter, _run) async def poll_for_tasks(): """On initialization, poll indefinitely for available tasks.""" + startup_logger = logging.getLogger(STREAM) # Ensure the two sqlite databases exist before any task tries to open them # (a first-run local `docker compose up` starts with the volume-mounted # directory empty). No-op once present or when no scp source is configured # (e.g. production, where the data is mounted out of band). - ensure_arax_pathfinder_dbs(logging.getLogger(STREAM)) + ensure_arax_pathfinder_dbs(startup_logger) + # Fetch the blocked-concept list once here rather than per task, so the pool + # children only ever read it. + ensure_blocked_list(startup_logger) + loop = asyncio.get_running_loop() + # Size the pool by the pod's actual CPU allocation (cgroup limit), not + # os.cpu_count() -- see aragorn_score.poll_for_tasks. Each child loads a full + # message, so this also bounds peak memory. POOL_MAX_WORKERS overrides. + max_workers = resolve_pool_workers(TASK_LIMIT, startup_logger) + logging.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="arax.pathfinder process pool", + task_timeout=settings.pool_task_timeout_sec, + ) 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: logging.info("Poll loop cancelled, shutting down.") except Exception as e: From 78d17afc331c7dda16244beed3405477dc4884c3 Mon Sep 17 00:00:00 2001 From: Max Wang Date: Mon, 10 Aug 2026 15:17:50 -0400 Subject: [PATCH 2/3] Add script for testing against pathfinder --- scripts/test_pathfinder.py | 140 +++++++++++++++++++++++++++++++++++++ 1 file changed, 140 insertions(+) create mode 100644 scripts/test_pathfinder.py diff --git a/scripts/test_pathfinder.py b/scripts/test_pathfinder.py new file mode 100644 index 0000000..1e574b0 --- /dev/null +++ b/scripts/test_pathfinder.py @@ -0,0 +1,140 @@ +import asyncio +import json +import time +from datetime import datetime +from pathlib import Path + +import httpx + +target_urls = { + "aragorn-ci": "https://shepherd.ci.transltr.io/aragorn", + "arax-ci": "https://shepherd.ci.transltr.io/arax", + "aragorn-dev": "https://shepherd.renci.org/aragorn", + "arax-dev": "https://shepherd.renci.org/arax", + "aragorn-local": "http://localhost:5439/aragorn", + "arax-local": "http://localhost:5439/arax", + "bte-local": "http://localhost:5439/bte", +} + +RESPONSES_DIR = "responses" + + +def generate_query(curie1: str, curie2: str) -> dict: + """Given a curie, return a TRAPI message.""" + parameters = { + # "timeout": 300, + # "tiers": [0], + } + return { + "message": { + "query_graph": { + "nodes": { + "on": { + "constraints": [], + "ids": [ + curie1 + ], + }, + "sn": { + "constraints": [], + "ids": [ + curie2 + ], + } + }, + "paths": { + "p0": { + "object": "on", + "subject": "sn", + "predicates": [ + "biolink:related_to", + ] + } + } + } + }, + "parameters": parameters, + "log_level": "DEBUG", + } + + +async def single_lookup(curies: tuple[str, str], target: str): + """Run a single query lookup synchronously.""" + query = generate_query(curies[0], curies[1]) + start_time = datetime.now() + try: + async with httpx.AsyncClient(timeout=600000) as client: + response = await client.post( + f"{target_urls[target]}/query", + json=query, + ) + response.raise_for_status() + response_json = response.json() + results = (response_json.get("message") or {}).get("results") or [] + num_results = len(results) + assert num_results == 1 + num_analyses = len(results[0]["analyses"]) + except Exception as e: + num_results = 0 + num_analyses = 0 + response_json = { + "Error": str(e), + } + + stop_time = datetime.now() + print(f"{curies[0]}->{curies[1]} took {stop_time - start_time} seconds and gave {num_analyses} results") + out_dir = Path(RESPONSES_DIR) / "pathfinder" / target + out_dir.mkdir(parents=True, exist_ok=True) + response_path = out_dir / f"{('_').join(curies[0].split(':'))}_{('_').join(curies[1].split(':'))}_response.json" + with response_path.open("w", encoding="utf-8") as f: + json.dump(response_json, f, indent=2) + + +query_list = [ + ('MONDO:0021095', 'MONDO:0005105'), + ('CHEBI:9139', 'MONDO:0004975'), + ('CHEBI:5118', 'MONDO:0100233'), + ('MONDO:0005180', 'MONDO:0005105'), + ('MONDO:0019632', 'MONDO:0005340'), + ('CHEBI:27881', 'NCBIGene:2739'), + ('CHEBI:45783', 'MONDO:0004979'), # Imatinib -> Asthma + ('GO:0006914', 'MONDO:0005265'), + ('NCBIGene:3458', 'CHEBI:16828'), + ('MONDO:0005532', 'MONDO:0005180'), + ('CHEBI:15647', 'UNII:31YO63LBSN'), + ('CHEBI:28364', 'MONDO:0005311'), + ('NCBIGene:3458', 'MONDO:0100096'), + ('NCBIGene:27240', 'MONDO:0100096'), + ('CHEBI:3750', 'MONDO:0013209'), + ('CHEBI:83766', 'MONDO:0008170'), + ('CHEBI:45783', 'MONDO:0004784'), + ('UNII:7SE5582Q2P', 'MONDO:0007037'), + ('MONDO:0005011', 'MONDO:0005180'), + ('CHEBI:15365', 'MONDO:0005575'), + ('CHEBI:50924', 'MONDO:0007256'), + ('CHEBI:45713', 'NCBIGene:2739'), + ('NCBIGene:54716', 'MONDO:0100096'), + ('CHEBI:7465', 'MONDO:0008218'), + # ('CHEBI:10033', 'MONDO:0004992'), # Warfarin -> Cancer, DON'T RUN +] + + +async def main(): + """Run the given query and time it.""" + targets = ["arax-local"] + runs_per_target = 1 + + start = time.time() + queries = [] + for curies in query_list: + queries.extend([ + single_lookup(curies, target) + for target in targets + for _ in range(runs_per_target) + ]) + await asyncio.gather(*queries) + print(f"\nAll queries took {time.time() - start:.2f} seconds") + + +if __name__ == "__main__": + asyncio.run(main()) From a782c5fd514f6aa5929aee9eabaeed4614208c2e Mon Sep 17 00:00:00 2001 From: Max Wang Date: Mon, 10 Aug 2026 15:20:19 -0400 Subject: [PATCH 3/3] Set arax pathfinder resources --- compose.test.yml | 3 +++ 1 file changed, 3 insertions(+) diff --git a/compose.test.yml b/compose.test.yml index 0bb7690..1db188b 100644 --- a/compose.test.yml +++ b/compose.test.yml @@ -64,6 +64,9 @@ services: arax: cpus: 1 mem_limit: 3g + arax_pathfinder: + cpus: 6 + mem_limit: 22g arax_rank: cpus: 4 mem_limit: 10g