diff --git a/.gitignore b/.gitignore index ea5de91a..2537f1f1 100644 --- a/.gitignore +++ b/.gitignore @@ -5,6 +5,7 @@ debug/ gandalf_mmap/ omnicorp_lmdb/ pathfinder_embeddings/ +pathfinder_search_graph/ arax_pathfinder_dbs/ arax_dbs/ diff --git a/compose.yml b/compose.yml index 16ebb6f9..5c1879ea 100644 --- a/compose.yml +++ b/compose.yml @@ -343,6 +343,10 @@ services: - ./workers/aragorn_pathfinder/debug:/app/debug - ./logs:/app/logs - ./.env:/app/.env + # First run? Set PATHFINDER_SEARCH_GRAPH_URL in your .env to a .tar.gz and + # the worker downloads + extracts the search graph into this mount on + # startup (see README "Worker data (LMDB / sqlite) downloads"). + - ./pathfinder_search_graph:/app/pathfinder_search_graph aragorn_omnicorp: container_name: aragorn_omnicorp platform: linux/amd64 diff --git a/shepherd_server/base_routes.py b/shepherd_server/base_routes.py index bf65b14f..e81abdbb 100644 --- a/shepherd_server/base_routes.py +++ b/shepherd_server/base_routes.py @@ -22,7 +22,7 @@ DecompressedTooLargeError, add_query, add_ready_callback, - decompress_zstd, + NotAJSONObjectError, get_callback_query_id, get_logs, get_message, @@ -31,7 +31,8 @@ get_response_query_parameters, remove_callback_id, save_logs, - save_message, + encode_callback_body, + save_encoded_message, ) from shepherd_utils.logger import ( attach_query_handler, @@ -738,12 +739,15 @@ async def callback( }, status_code=413, ) + # Stored as it came, never parsed here: parsing a large message holds the + # GIL long enough for uvicorn's supervisor to take this worker for hung and + # kill it. merge_message parses it. The cap applies to what that has to + # hold in memory, so a compressed body is bounded on its decompressed size. + zstd_encoded = "zstd" in request.headers.get("content-encoding", "").lower() try: - if "zstd" in request.headers.get("content-encoding", "").lower(): - # The cap applies to what we have to hold in memory, so a - # compressed body is bounded on its decompressed size too. - raw = decompress_zstd(raw, max_bytes) - response = orjson.loads(raw) + blob, payload_size = await asyncio.to_thread( + encode_callback_body, raw, zstd_encoded, max_bytes + ) except DecompressedTooLargeError as e: await _fail_query_for_oversized_callback( callback_id, @@ -761,7 +765,7 @@ async def callback( }, status_code=413, ) - except (orjson.JSONDecodeError, zstandard.ZstdError): + except (NotAJSONObjectError, zstandard.ZstdError): logger.warning(f"Rejecting callback {callback_id}: invalid request body.") await _save_callback_error_logs(callback_id, logger) return JSONResponse( @@ -785,31 +789,7 @@ async def callback( level_number = await get_query_log_level(original_query[0], logger) logger.setLevel(level_number) logger.debug(f"Got original query: {original_query}") - if not isinstance(response.get("message"), dict): - response["message"] = {} - results = response["message"].get("results") - if results is None: - response["message"]["results"] = [] - kgraph = response["message"].get("knowledge_graph") - if kgraph is None: - response["message"]["knowledge_graph"] = { - "nodes": {}, - "edges": {}, - } - - logger.info( - f"[{callback_id}] Got back {len(response['message']['results'])} results." - ) - logger.debug( - f"[{callback_id}] for query graph: {response['message'].get('query_graph')}" - ) - # if len(response["message"]["results"]) > 0: - # with open( - # f"shepherd_server/debug/{query_id}_{callback_id}_response.json", - # "w", - # encoding="utf-8", - # ) as f: - # json.dump(response, f, indent=2) + logger.info(f"[{callback_id}] Got back a {payload_size}-byte response.") query_state = await get_query_state(original_query[0], logger) if query_state is None: logger.warning( @@ -826,7 +806,7 @@ async def callback( # to) could still record the callback as ready once Redis was back -- # e.g. across a Redis restart -- and merge_message would then skip it # as missing, silently dropping results the sender was told arrived. - await save_message(callback_id, response, logger, raise_on_failure=True) + await save_encoded_message(callback_id, blob, logger, raise_on_failure=True) except Exception: logger.error( f"[{callback_id}] could not save callback payload after retries; " @@ -855,12 +835,9 @@ async def callback( # adds otel trace to carrier for next worker parent_ctx = extract(json.loads(original_query[1])) with tracer.start_as_current_span("callback", context=parent_ctx) as span: - kgraph = response["message"]["knowledge_graph"] span.set_attribute("callback.id", callback_id) - span.set_attribute("callback.results", len(response["message"]["results"])) - span.set_attribute("callback.kg_nodes", len(kgraph.get("nodes", {}))) - span.set_attribute("callback.kg_edges", len(kgraph.get("edges", {}))) span.set_attribute("callback.payload_bytes", len(raw)) + span.set_attribute("callback.decompressed_bytes", payload_size) span_carrier = {} inject(span_carrier) # add new task to merge callback response into original message diff --git a/shepherd_utils/config.py b/shepherd_utils/config.py index 4b146b23..75734575 100644 --- a/shepherd_utils/config.py +++ b/shepherd_utils/config.py @@ -201,6 +201,7 @@ class Settings(BaseSettings): omnicorp_shared_counts_lmdb_path: str = "/app/omnicorp_lmdb/shared_counts.lmdb" pathfinder_embeddings_dir: str = "pathfinder_embeddings" + pathfinder_search_graph_dir: str = "pathfinder_search_graph" # First-run data downloads for local docker compose development. When set, # the worker fetches a ``.tar.gz`` from this URL on startup (if the LMDB @@ -210,6 +211,7 @@ class Settings(BaseSettings): # download is skipped there. See shepherd_utils/data_download.py. omnicorp_lmdb_url: str = "" pathfinder_embeddings_url: str = "" + pathfinder_search_graph_url: str = "" # Socket timeout (seconds) for those startup dataset downloads. This bounds # each individual blocking socket operation -- the connect and every diff --git a/shepherd_utils/data_download.py b/shepherd_utils/data_download.py index eb9e902a..8a5abc9a 100644 --- a/shepherd_utils/data_download.py +++ b/shepherd_utils/data_download.py @@ -3,16 +3,18 @@ Several workers read from datasets that are far too large to commit to git -- they're gitignored and volume-mounted from the host (``./omnicorp_lmdb``, -``./pathfinder_embeddings``, ``./arax_pathfinder_dbs``, ``./arax_dbs``). In -production these volumes are provisioned out of band, but a developer running -``docker compose up`` for the first time has empty directories, and the -workers crash on startup trying to open missing files. +``./pathfinder_embeddings``, ``./pathfinder_search_graph``, +``./arax_pathfinder_dbs``, ``./arax_dbs``). In production these volumes are +provisioned out of band, but a developer running ``docker compose up`` for the +first time has empty directories, and the workers crash on startup trying to +open missing files. Two flavors of HTTP source are supported: * **Archive** -- a single ``.tar.gz`` fetched via ``urllib`` and extracted in - place (``OMNICORP_LMDB_URL`` / ``PATHFINDER_EMBEDDINGS_URL``, used by - ``aragorn_omnicorp`` and ``score_paths`` below). + place (``OMNICORP_LMDB_URL`` / ``PATHFINDER_EMBEDDINGS_URL`` / + ``PATHFINDER_SEARCH_GRAPH_URL``, used by ``aragorn_omnicorp``, + ``score_paths`` and ``aragorn_pathfinder`` below). * **Per-file** -- individual files fetched directly, no archive/extract step (``arax_pathfinder`` below, whose two sqlite databases are served as plain files rather than bundled into one archive, plus the ARAX blocked-concept @@ -315,6 +317,21 @@ def ensure_pathfinder_embeddings(logger: Optional[logging.Logger] = None) -> Non ) +def ensure_pathfinder_search_graph(logger: Optional[logging.Logger] = None) -> None: + """Ensure the aragorn_pathfinder search graph is present. + + The backbone's pair graph as a CSR plus its node CURIEs in index order, + written by squashbert's ``scripts/export_search_graph.py``. + """ + ensure_lmdb_dataset( + name="aragorn_pathfinder", + target_dir=settings.pathfinder_search_graph_dir, + required_files=["offsets.npy", "neighbors.npy", "weights.npy", "nodes.txt"], + url=settings.pathfinder_search_graph_url, + logger=logger, + ) + + def arax_pathfinder_sqlite_paths() -> Tuple[str, str]: """Return ``(curie_ngd_path, node_degree_path)`` for the arax_pathfinder sqlite databases, built from ``arax_pathfinder_dbs_dir`` + the filename diff --git a/shepherd_utils/db.py b/shepherd_utils/db.py index 4ca94a08..7c48f5ff 100644 --- a/shepherd_utils/db.py +++ b/shepherd_utils/db.py @@ -225,9 +225,18 @@ def encode_message(obj: Any) -> bytes: return zstandard.compress(orjson.dumps(obj)) +def decompress_frame(blob: bytes) -> bytes: + """Decompress a stored blob. A frame that omits its content size (a + streaming one, as a callback's sender may post) can't be decompressed in + one shot, so it is streamed instead.""" + if zstandard.frame_content_size(blob) == -1: + return decompress_zstd(blob) + return zstandard.decompress(blob) + + def decode_message(blob: bytes) -> Any: """Deserialize a stored message blob back into a Python object.""" - return orjson.loads(zstandard.decompress(blob)) + return orjson.loads(decompress_frame(blob)) class DecompressedTooLargeError(Exception): @@ -281,6 +290,55 @@ def decompress_zstd(blob: bytes, max_bytes: int = 0) -> bytes: return b"".join(chunks) +class NotAJSONObjectError(ValueError): + """Raised by ``encode_callback_body`` for a body that can't be a TRAPI + message: it doesn't start like a JSON object.""" + + +def _starts_like_json_object(head: bytes) -> bool: + return head.lstrip()[:1] == b"{" + + +def encode_callback_body( + body: bytes, zstd_encoded: bool, max_bytes: int = 0 +) -> tuple[bytes, int]: + """A callback body as a stored blob, without parsing it. + + Returns ``(blob, decompressed size)``. Parsing a large TRAPI message holds + the GIL for seconds, long enough for the server's process supervisor to + take a worker for hung and kill it, so the callback route stores the body + as it came and leaves parsing to merge_message. A zstd body whose frame + carries its content size already is a stored blob; any other body is + (re)compressed into one, so ``get_response_size`` can still read its size + from the header. zstd releases the GIL throughout, so run this in a thread. + + ``max_bytes`` bounds the decompressed size (0 disables it), raising + ``DecompressedTooLargeError``. A body that doesn't start like a JSON + object raises ``NotAJSONObjectError``, and a bad zstd frame + ``zstandard.ZstdError``; anything wrong further in is merge_message's to + find. + """ + if not zstd_encoded: + if not _starts_like_json_object(body[:64]): + raise NotAJSONObjectError("request body is not a JSON object") + return zstandard.compress(body), len(body) + size = zstandard.frame_content_size(body) + if size == -1: + data = decompress_zstd(body, max_bytes) + if not _starts_like_json_object(data[:64]): + raise NotAJSONObjectError("request body is not a JSON object") + return zstandard.compress(data), len(data) + if max_bytes > 0 and size > max_bytes: + raise DecompressedTooLargeError( + f"zstd frame declares {size} decompressed bytes, over the " + f"{max_bytes}-byte limit" + ) + head = zstandard.ZstdDecompressor().stream_reader(io.BytesIO(body)).read(64) + if not _starts_like_json_object(head): + raise NotAJSONObjectError("request body is not a JSON object") + return body, size + + # Idempotent DDL applied on every startup. The image's init_db.sql only runs # when the Postgres data volume is brand new, so deployments whose volume # predates a schema addition never pick it up from there; re-running these @@ -546,18 +604,33 @@ async def save_message( callers whose success response promises the payload is stored (e.g. ARS submit); default keeps the historical swallow-and-log. """ - start = time.time() + start_comp = time.time() + # a thread: orjson+zstd over a multi-MB message is pure CPU try: - start_comp = time.time() - # a thread: orjson+zstd over a multi-MB message is pure CPU, and the - # server's callback path runs this on its event loop compressed = await asyncio.to_thread(encode_message, response) - logger.info(f"Compression took {time.time() - start_comp}") - await data_db_client.set( - callback_id, - compressed, - ex=settings.redis_ttl, - ) + except Exception as e: + logger.error(f"Failed to encode a message for redis: {e}") + if raise_on_failure: + raise + return + logger.info(f"Compression took {time.time() - start_comp}") + await save_encoded_message( + callback_id, compressed, logger, num_tries, raise_on_failure + ) + + +async def save_encoded_message( + message_id: str, + blob: bytes, + logger: logging.Logger, + num_tries: int = 0, + raise_on_failure: bool = False, +): + """Store an already-encoded message blob, retrying as ``save_message`` + does.""" + start = time.time() + try: + await data_db_client.set(message_id, blob, ex=settings.redis_ttl) logger.debug(f"Saving message took {time.time() - start} seconds") except Exception as e: # failed to put message in db @@ -565,8 +638,8 @@ async def save_message( num_tries += 1 logger.warning(f"Failed to save message {num_tries} times. Trying again...") await asyncio.sleep(0.5) - await save_message( - callback_id, response, logger, num_tries, raise_on_failure + await save_encoded_message( + message_id, blob, logger, num_tries, raise_on_failure ) else: logger.error(f"Failed to save a message into redis: {e}") @@ -728,7 +801,7 @@ async def get_message( raise KeyError(f"Failed to get {message_id} from db") if raw: - result = zstandard.decompress(blob) + result = decompress_frame(blob) logger.debug(f"Getting raw message took {time.time() - start} seconds") return result diff --git a/shepherd_utils/trapi.py b/shepherd_utils/trapi.py index 64737503..121668b3 100644 --- a/shepherd_utils/trapi.py +++ b/shepherd_utils/trapi.py @@ -590,3 +590,23 @@ def empty_message() -> MessageDict: "MessageDict", {"knowledge_graph": {"nodes": {}, "edges": {}}, "results": []}, ) + + +# A query node pinned to more ids than this logs as "N curies" instead. +MAX_LOGGED_NODE_IDS = 10 + + +def summarize_query_graph(query_graph): + """The query graph to log: a node's long id list becomes its count, so a + lookup pinning thousands of curies logs in a line.""" + if not isinstance(query_graph, dict) or not isinstance( + query_graph.get("nodes"), dict + ): + return query_graph + nodes = {} + for node_id, node in query_graph["nodes"].items(): + ids = node.get("ids") if isinstance(node, dict) else None + if isinstance(ids, list) and len(ids) > MAX_LOGGED_NODE_IDS: + node = {**node, "ids": f"{len(ids)} curies"} + nodes[node_id] = node + return {**query_graph, "nodes": nodes} diff --git a/tests/unit/test_aragorn_pathfinder.py b/tests/unit/test_aragorn_pathfinder.py index 519ae32e..19c04c38 100644 --- a/tests/unit/test_aragorn_pathfinder.py +++ b/tests/unit/test_aragorn_pathfinder.py @@ -1,8 +1,9 @@ """Tests for ``workers.aragorn_pathfinder.worker.shadowfax``. -The pathfinder entry point validates the input pathfinder query, builds a -3-hop expanded query, and delegates to the gandalf stream. We exercise the -validation paths and the happy path (mocking out postgres + redis). +The pathfinder entry point validates the input pathfinder query, sends +one-hop edge lookups to Retriever, waits for their callbacks, and builds the +paths in a process pool. We exercise the validation paths and the happy path +(mocking out postgres, redis and the pool). """ import json @@ -10,11 +11,22 @@ import pytest +from workers.aragorn_pathfinder import worker as pathfinder_worker from workers.aragorn_pathfinder.worker import shadowfax logger = logging.getLogger(__name__) +@pytest.fixture(autouse=True) +def pool(mocker): + """Stands in for the process pool poll_for_tasks sets up: the path build + is not run, only the lookups that feed it.""" + pool = mocker.Mock() + pool.run = mocker.AsyncMock(return_value=[]) + mocker.patch.object(pathfinder_worker, "POOL", pool) + return pool + + def _make_task(workflow=None): return [ "test", @@ -50,9 +62,9 @@ def _pathfinder_message(constraints=None): @pytest.mark.asyncio -async def test_shadowfax_happy_path_dispatches_to_gandalf(redis_mock, mocker): - """A valid pathfinder query: callback registered, threehop saved, gandalf - task enqueued, polling loop exits when no callbacks remain.""" +async def test_shadowfax_happy_path_builds_paths(redis_mock, mocker, pool): + """A valid pathfinder query: callback registered, edge lookup sent, the + wait exits when no callbacks remain, and the paths are built in the pool.""" mocker.patch( "workers.aragorn_pathfinder.worker.get_message", new_callable=mocker.AsyncMock, @@ -79,6 +91,7 @@ async def test_shadowfax_happy_path_dispatches_to_gandalf(redis_mock, mocker): assert mock_add_cb.called assert mock_httpx.called + pool.run.assert_awaited_once() @pytest.mark.asyncio @@ -131,9 +144,9 @@ async def test_shadowfax_rejects_multiple_intermediate_categories(redis_mock, mo @pytest.mark.asyncio -async def test_shadowfax_propagates_gandalf_parameters(redis_mock, mocker): - """Custom gandalf_parameters in the input should ride along into the - saved threehop's parameters.""" +async def test_shadowfax_sends_its_own_lookup_parameters(redis_mock, mocker): + """The edge lookups carry the worker's own retrieval parameters (one + tier, no tier fallback, dehydrated), not the caller's 3-hop filters.""" msg = _pathfinder_message() msg["parameters"]["filter_config"] = { "min_information_content": 1, @@ -168,15 +181,17 @@ async def test_shadowfax_propagates_gandalf_parameters(redis_mock, mocker): args, kwargs = mock_httpx.call_args query_parameters = kwargs["json"]["parameters"] - assert query_parameters["filter_config"]["min_information_content"] == 1 - assert query_parameters["filter_config"]["max_node_degree"] == 10 + assert "filter_config" not in query_parameters + assert query_parameters["timeout"] == 5 + assert query_parameters["tier_fallback"] is False + assert query_parameters["dehydrated"] is True @pytest.mark.asyncio async def test_shadowfax_sends_valid_trapi_2_query(redis_mock, mocker): """A 2.0 path constraint (required_intermediate_categories) is accepted, - and the expanded 3-hop query sent to Retriever is a valid TRAPI 2.0 query - with no top-level log_level.""" + and the edge lookup sent to Retriever is a valid TRAPI 2.0 query with no + top-level log_level.""" from translator_tom import Query msg = _pathfinder_message( @@ -209,5 +224,4 @@ async def test_shadowfax_sends_valid_trapi_2_query(redis_mock, mocker): sent = mock_httpx.call_args.kwargs["json"] assert "log_level" not in sent - assert sent["parameters"]["log_level"] == "DEBUG" Query.from_dict(sent) diff --git a/tests/unit/test_aragorn_pathfinder_paths.py b/tests/unit/test_aragorn_pathfinder_paths.py new file mode 100644 index 00000000..fd4edecd --- /dev/null +++ b/tests/unit/test_aragorn_pathfinder_paths.py @@ -0,0 +1,733 @@ +"""Tests for aragorn.pathfinder's path finding and building. + +``search_paths`` is checked against a brute-force enumeration over small +graphs written as the same CSR the worker memory-maps; ``edge_lookups`` and +``build_paths`` are pure; ``build_paths_task`` and an on-backbone +``shadowfax`` run with Redis and the outbound calls mocked. +""" + +import itertools +import json +import logging +import random +from math import inf, isclose + +import numpy as np +import pytest +from translator_tom.v2_0 import AsyncQuery + +from shepherd_utils import data_download +from shepherd_utils.config import settings +from shepherd_utils.data_download import ensure_pathfinder_search_graph +from workers.aragorn_pathfinder import worker +from workers.aragorn_pathfinder.worker import ( + build_paths, + build_paths_task, + cheapest_legs, + edge_lookups, + load_search_graph, + one_hop, + search_paths, + shadowfax, +) + +logger = logging.getLogger(__name__) + + +# --- helpers ----------------------------------------------------------------- + + +def write_graph(directory, num_nodes, edges): + """Write an undirected weighted graph as the worker's CSR .npy files and + return the adjacency it encodes.""" + adjacency = {node: {} for node in range(num_nodes)} + for u, v, w in edges: + adjacency[u][v] = w + adjacency[v][u] = w + offsets, neighbors, weights = [0], [], [] + for node in range(num_nodes): + for other, w in sorted(adjacency[node].items()): + neighbors.append(other) + weights.append(w) + offsets.append(len(neighbors)) + directory.mkdir(parents=True, exist_ok=True) + np.save(directory / "offsets.npy", np.array(offsets, dtype=np.int64)) + np.save(directory / "neighbors.npy", np.array(neighbors, dtype=np.int32)) + np.save(directory / "weights.npy", np.array(weights, dtype=np.float32)) + return adjacency + + +def random_graph(directory, num_nodes=30, num_edges=110, seed=0): + rng = random.Random(seed) + edges = [] + for _ in range(num_edges): + u, v = rng.sample(range(num_nodes), 2) + edges.append((u, v, rng.random())) + return write_graph(directory, num_nodes, edges) + + +def weight(adjacency, u, v): + return float(np.float32(adjacency[u][v])) + + +def path_cost(adjacency, path): + return sum(weight(adjacency, u, v) for u, v in zip(path, path[1:])) + + +def centres(path): + """The nodes a path is centred on: its middle, or both intermediates of a + three-hop path.""" + if len(path) == 4: + return [path[1], path[2]] + return [path[len(path) // 2]] + + +def brute_force(adjacency, source, target): + """Every centre's cheapest loopless path of two to four hops.""" + best = {} + + def walk(path): + if path[-1] == target: + if len(path) >= 3: + cost = path_cost(adjacency, path) + for centre in centres(path): + if centre not in best or cost < best[centre]: + best[centre] = cost + return + if len(path) == 5: + return + for node in adjacency[path[-1]]: + if node not in path: + walk(path + [node]) + + walk([source]) + return best + + +@pytest.fixture +def graph_dir(tmp_path): + yield tmp_path / "graph" + worker._search_graphs.clear() + + +# --- search ------------------------------------------------------------------ + + +def test_load_search_graph_maps_the_csr_once(graph_dir): + write_graph(graph_dir, 3, [(0, 1, 0.5), (1, 2, 0.25)]) + + graph = load_search_graph(str(graph_dir)) + + assert load_search_graph(str(graph_dir)) is graph + offsets, neighbors, weights = graph + assert offsets.tolist() == [0, 1, 3, 4] + assert neighbors.tolist() == [1, 0, 2, 1] + assert isinstance(weights, np.memmap) + + +def test_cheapest_legs(graph_dir): + # 0 -1- 1 -1- 3 + # 0 -2- 2 -2- 3 0 -9- 4 (the far end) + # 1 -5- 2 3 -1- 4 + write_graph( + graph_dir, + 5, + [(0, 1, 1), (1, 3, 1), (0, 2, 2), (2, 3, 2), (1, 2, 5), (0, 4, 9), (3, 4, 1)], + ) + one, two, via = cheapest_legs(load_search_graph(str(graph_dir)), 0, 4) + + # one-hop legs: never to the start or the far end + assert one.tolist() == [inf, 1, 2, inf, inf] + # two-hop legs to 3: through 1 (cost 2) then through 2 (cost 4) + assert two[0][3] == 2 and via[0][3] == 1 + assert two[1][3] == 4 and via[1][3] == 2 + # 0 -> 1 -> 2 (cost 6) and 0 -> 2 -> 1 (cost 7) + assert two[0][2] == 6 and via[0][2] == 1 + assert two[0][1] == 7 and via[0][1] == 2 + # no leg passes through or ends at the far end, or returns to the start + assert two[0][4] == inf and two[0][0] == inf + assert via[1][1] == -1 + + +def test_search_paths_matches_brute_force(graph_dir): + adjacency = random_graph(graph_dir) + + checked = 0 + for source, target in itertools.permutations(range(8), 2): + paths, _ = search_paths(str(graph_dir), source, target, 10**6) + expected = brute_force(adjacency, source, target) + + found = {} + for path in paths: + assert path[0] == source and path[-1] == target + assert 3 <= len(path) <= 5 + assert len(set(path)) == len(path), f"{path} has a loop" + assert all(v in adjacency[u] for u, v in zip(path, path[1:])) + for centre in centres(path): + found.setdefault(centre, path_cost(adjacency, path)) + # every centre a loopless path has, at the cheapest such path's cost + assert found.keys() == expected.keys() + for centre, cost in expected.items(): + assert isclose(found[centre], cost, rel_tol=1e-6), centre + checked += len(paths) + assert checked > 100 + + +def test_search_paths_come_out_cheapest_first_and_once(graph_dir): + adjacency = random_graph(graph_dir, seed=1) + + paths, middles = search_paths(str(graph_dir), 0, 1, 10**6) + + costs = [path_cost(adjacency, path) for path in paths] + assert all(a <= b + 1e-6 for a, b in zip(costs, costs[1:])) + assert len(set(paths)) == len(paths) + assert middles >= len(paths) + + +def test_search_paths_stops_at_max_paths(graph_dir): + random_graph(graph_dir, seed=2) + + everything, _ = search_paths(str(graph_dir), 0, 1, 10**6) + first, _ = search_paths(str(graph_dir), 0, 1, 5) + + assert len(everything) > 5 + assert first == everything[:5] + + +def test_search_paths_between_unconnected_nodes_is_empty(graph_dir): + write_graph(graph_dir, 4, [(0, 1, 1.0), (2, 3, 1.0)]) + + assert search_paths(str(graph_dir), 0, 3, 10)[0] == [] + + +def test_search_paths_skips_unusable_edges(graph_dir): + # the only route through 2 has an infinite weight + write_graph(graph_dir, 4, [(0, 1, 1), (1, 3, 1), (0, 2, inf), (2, 3, 1)]) + + paths, _ = search_paths(str(graph_dir), 0, 3, 10) + + assert paths == [(0, 1, 3)] + + +def test_search_paths_gives_up_at_the_deadline(graph_dir, mocker): + random_graph(graph_dir, seed=3) + mocker.patch.object(worker, "SEARCH_BUDGET_SEC", -1.0) + + paths, middles = search_paths(str(graph_dir), 0, 1, 10) + + assert paths == [] + assert middles > 0 + + +# --- lookups ----------------------------------------------------------------- + + +def query_ids(lookup): + nodes = lookup["message"]["query_graph"]["nodes"] + return nodes["n0"].get("ids"), nodes["n1"].get("ids") + + +def test_one_hop(): + lookup = one_hop(["A"], None, {"timeout": 5}, True) + + assert lookup == { + "message": { + "query_graph": { + "nodes": {"n0": {"ids": ["A"]}, "n1": {}}, + "edges": { + "e0": { + "subject": "n0", + "object": "n1", + "predicates": ["biolink:related_to"], + } + }, + } + }, + "parameters": {"timeout": 5, "subclass": True}, + } + + +def test_edge_lookups_without_paths_is_the_direct_edge(): + lookups = edge_lookups([], "S", "T", {"timeout": 5}) + + assert len(lookups) == 1 + assert query_ids(lookups[0]) == (["S"], ["T"]) + assert lookups[0]["parameters"] == {"timeout": 5, "subclass": True} + + +def test_edge_lookups_cover_every_hop(): + paths = [ + ["S", "M", "T"], + ["S", "A", "B", "T"], + ["S", "A", "M", "C", "T"], + ["S", "D", "M", "C", "T"], + ] + + lookups = edge_lookups(paths, "S", "T", {"timeout": 5}) + + # the endpoints are expanded to their subclasses, against anything + assert [query_ids(lookup) for lookup in lookups[:3]] == [ + (["S"], ["T"]), + (["S"], None), + (None, ["T"]), + ] + assert all(lookup["parameters"]["subclass"] for lookup in lookups[:3]) + # the hops between intermediates are one exact lookup + assert len(lookups) == 4 + near, far = query_ids(lookups[3]) + assert near == ["A", "C", "D"] + assert far == ["B", "M"] + assert lookups[3]["parameters"] == {"timeout": 5, "subclass": False} + pairs = set(itertools.product(near, far)) + for u, v in [("A", "B"), ("A", "M"), ("C", "M"), ("D", "M")]: + assert (u, v) in pairs + + +def test_edge_lookups_for_two_hop_paths_need_no_interior_lookup(): + lookups = edge_lookups([["S", "M", "T"], ["S", "N", "T"]], "S", "T", {}) + + assert [query_ids(lookup) for lookup in lookups] == [ + (["S"], ["T"]), + (["S"], None), + (None, ["T"]), + ] + + +# --- building paths ---------------------------------------------------------- + + +def edge(subject, obj, predicate="biolink:related_to"): + return {"subject": subject, "object": obj, "predicate": predicate} + + +def node(*categories): + return {"categories": list(categories)} + + +def make_plan(node_paths, category=None): + return { + "source": "S", + "target": "T", + "node_paths": node_paths, + "path_key": "p0", + "subject_key": "n0", + "object_key": "n1", + "intermediate_category": category, + } + + +def make_message(): + return { + "query_graph": {}, + "knowledge_graph": { + "nodes": { + "S": node("biolink:Disease"), + "T": node("biolink:Disease"), + "A": node("biolink:Gene"), + "B": node("biolink:ChemicalEntity"), + "M": node("biolink:Gene"), + }, + "edges": { + "st": edge("S", "T"), + "sa": edge("S", "A"), + "as": edge("A", "S", "biolink:affects"), + "ab": edge("B", "A"), + "bt": edge("B", "T"), + "sm": edge("S", "M"), + "m_sub": edge("M", "T", "biolink:subclass_of"), + "sb": edge("S", "B"), + }, + }, + "auxiliary_graphs": {"support": {"edges": ["sa"]}}, + "results": [{"stale": True}], + } + + +def bound_paths(message): + """{aux id: sorted edge ids} for each path analysis.""" + (result,) = message["results"] + out = {} + for analysis in result["analyses"]: + assert analysis["resource_id"] == "infores:shepherd-aragorn" + (aux_id,) = analysis["path_bindings"]["p0"]["ids"] + out[aux_id] = sorted(message["auxiliary_graphs"][aux_id]["edges"]) + return out + + +def test_build_paths(): + message = make_message() + plan = make_plan([("S", "A", "B", "T"), ("S", "M", "T"), ("S", "B", "T")]) + + build_paths(message, plan, logger) + + (result,) = message["results"] + assert result["node_bindings"] == { + "n0": {"ids": ["S"]}, + "n1": {"ids": ["T"]}, + } + # the direct edge first, then every edge of each hop, either way round; the + # path through M has only a subclass_of edge to T and is dropped + assert bound_paths(message) == { + "p_0": ["st"], + "p_1": ["ab", "as", "bt", "sa"], + "p_3": ["bt", "sb"], + } + # the lookups' own aux graphs stay + assert message["auxiliary_graphs"]["support"] == {"edges": ["sa"]} + + +def test_build_paths_keeps_one_path_per_set_of_intermediates(): + message = make_message() + + build_paths(message, make_plan([("S", "B", "T"), ("S", "B", "T")]), logger) + + assert list(bound_paths(message)) == ["p_0", "p_1"] + + +def test_build_paths_without_a_direct_edge(): + message = make_message() + del message["knowledge_graph"]["edges"]["st"] + + build_paths(message, make_plan([("S", "B", "T")]), logger) + + assert bound_paths(message) == {"p_1": ["bt", "sb"]} + + +@pytest.mark.parametrize( + "category, expected", + [ + ("biolink:Gene", ["p_1"]), + ("biolink:ChemicalEntity", ["p_1", "p_2"]), + ("biolink:NamedThing", ["p_0", "p_1", "p_2"]), + (None, ["p_0", "p_1", "p_2"]), + ], +) +def test_build_paths_requires_the_intermediate_category(category, expected): + message = make_message() + plan = make_plan([("S", "A", "B", "T"), ("S", "B", "T")], category) + + build_paths(message, plan, logger) + + assert list(bound_paths(message)) == expected + + +def test_build_paths_with_nothing_to_bind(): + message = make_message() + message["knowledge_graph"]["edges"] = {} + + build_paths(message, make_plan([("S", "A", "B", "T")]), logger) + + assert message["results"] == [] + assert message["auxiliary_graphs"] == {"support": {"edges": ["sa"]}} + + +def test_build_paths_on_an_empty_message(): + message = {} + + build_paths(message, make_plan([("S", "A", "T")]), logger) + + assert message == {"knowledge_graph": {"nodes": {}, "edges": {}}, "results": []} + + +def test_build_paths_task_saves_the_paths_without_orphans(mocker): + stored = {"message": make_message()} + stored["message"]["auxiliary_graphs"] = {} + mocker.patch.object(worker, "get_message_sync", return_value=stored) + save = mocker.patch.object(worker, "save_response_sync") + + entries = build_paths_task("rid", make_plan([("S", "B", "T")]), logging.INFO) + + save.assert_called_once() + response_id, response = save.call_args.args + assert response_id == "rid" + assert bound_paths(response["message"]) == {"p_0": ["st"], "p_1": ["bt", "sb"]} + # edges and nodes no path uses are gone + assert set(response["message"]["knowledge_graph"]["edges"]) == {"st", "sb", "bt"} + assert set(response["message"]["knowledge_graph"]["nodes"]) == {"S", "B", "T"} + assert any("Built 2 paths" in json.dumps(entry) for entry in entries) + + +def test_build_paths_takes_an_inferred_endpoint_edge_and_keeps_its_support(mocker): + """Gandalf expands an endpoint to its subclass children and returns a + child's edge as an inferred edge on the endpoint, supported by the child's + own edge and its subclass_of edge. The inferred edge is the hop; the + child's edge is not, and the support survives the orphan filter.""" + message = { + "knowledge_graph": { + "nodes": { + "S": node("biolink:Disease"), + "C": node("biolink:Disease"), + "M": node("biolink:Gene"), + "T": node("biolink:Disease"), + }, + "edges": { + "cm": edge("C", "M", "biolink:affects"), + "cs": edge("C", "S", "biolink:subclass_of"), + "cm_cs": { + **edge("S", "M", "biolink:affects"), + "attributes": [ + { + "attribute_type_id": "biolink:support_graphs", + "value": ["support_cm_cs"], + } + ], + }, + "mt": edge("M", "T"), + }, + }, + "auxiliary_graphs": {"support_cm_cs": {"edges": ["cm", "cs"]}}, + "results": [], + } + mocker.patch.object(worker, "get_message_sync", return_value={"message": message}) + save = mocker.patch.object(worker, "save_response_sync") + + build_paths_task("rid", make_plan([("S", "M", "T")]), logging.INFO) + + saved = save.call_args.args[1]["message"] + assert bound_paths(saved) == {"p_1": ["cm_cs", "mt"]} + assert saved["auxiliary_graphs"]["support_cm_cs"] == {"edges": ["cm", "cs"]} + assert set(saved["knowledge_graph"]["edges"]) == {"cm_cs", "mt", "cm", "cs"} + assert set(saved["knowledge_graph"]["nodes"]) == {"S", "C", "M", "T"} + + +def test_build_paths_says_where_the_dropped_paths_lack_a_hop(caplog): + """A path whose hop came back with no edges is dropped, counted by where + the hop is, and the first such hop is named with whether its nodes came + back at all.""" + message = make_message() + plan = make_plan( + [ + ("S", "X", "T"), # no S - X edge: at the source + ("S", "A", "X", "T"), # no A - X edge: between intermediates + ("S", "M", "T"), # M - T only subclass_of: at the target + ] + ) + + with caplog.at_level(logging.INFO): + build_paths(message, plan, logger) + + assert bound_paths(message) == {"p_0": ["st"]} + assert ( + "dropped 3 with a hop that has no edges {'at the source': 1, " + "'between intermediates': 1, 'at the target': 1}" + ) in caplog.text + assert "First hop with no edges: S (in the edges) - X (not in the edges)" in ( + caplog.text + ) + + +# --- shadowfax over the backbone --------------------------------------------- + + +class InlinePool: + """Runs the search for real and records the path build it is handed, + running that too when asked.""" + + def __init__(self, build=False): + self.build = build + self.builds = [] + + async def run(self, loop, fn, *args): + if fn is search_paths: + return fn(*args) + self.builds.append(args) + return fn(*args) if self.build else [] + + +def backbone_query(constraints=None): + query = { + "message": { + "query_graph": { + "nodes": {"n0": {"ids": ["X:S"]}, "n1": {"ids": ["X:T"]}}, + "paths": {"p0": {"subject": "n0", "object": "n1"}}, + } + }, + "parameters": {"timeout": 5}, + } + if constraints is not None: + query["message"]["query_graph"]["paths"]["p0"]["constraints"] = constraints + return query + + +def backbone_merged_message(): + """The response once every lookup over the backbone is merged: A is the + only gene, and there is no direct edge.""" + return { + "knowledge_graph": { + "nodes": { + "X:S": node("biolink:Disease"), + "X:A": node("biolink:Gene"), + "X:B": node("biolink:ChemicalEntity"), + "X:M": node("biolink:Disease"), + "X:T": node("biolink:Disease"), + }, + "edges": { + "sa": edge("X:S", "X:A", "biolink:affects"), + "ab": edge("X:A", "X:B"), + "bt": edge("X:B", "X:T"), + "sm": edge("X:S", "X:M"), + "mt": edge("X:M", "X:T"), + }, + }, + "results": [], + } + + +async def run_on_backbone(mocker, graph_dir, query, pool): + """``shadowfax`` on ``query`` over the backbone S - A - B - T, S - M - T, + with Redis and the outbound calls mocked. Returns its metadata, the + lookups it sent, the callback registrations and the response it saved + (None unless the pool runs the build).""" + write_graph( + graph_dir, + 5, + [(0, 1, 1), (1, 2, 1), (2, 3, 1), (0, 4, 1), (4, 3, 1)], + ) + curies = ["X:S", "X:A", "X:B", "X:T", "X:M"] + mocker.patch.object(worker, "SEARCH_GRAPH_DIR", str(graph_dir)) + mocker.patch.object(worker, "NODE_CURIES", curies) + mocker.patch.object(worker, "NODE_INDEX", {c: i for i, c in enumerate(curies)}) + mocker.patch.object(worker, "POOL", pool) + mocker.patch.object( + worker, "get_message", new_callable=mocker.AsyncMock, return_value=query + ) + add_callback = mocker.patch.object( + worker, "add_callback_id", new_callable=mocker.AsyncMock + ) + mocker.patch.object( + worker, + "get_running_callbacks", + new_callable=mocker.AsyncMock, + return_value=[], + ) + mocker.patch.object( + worker, + "get_message_sync", + return_value={"message": backbone_merged_message()}, + ) + save = mocker.patch.object(worker, "save_response_sync") + post = mocker.patch( + "httpx.AsyncClient.post", + new_callable=mocker.AsyncMock, + return_value=mocker.Mock(status_code=200), + ) + task = [ + "1-0", + { + "query_id": "qid", + "response_id": "rid", + "workflow": "[]", + "log_level": "20", + "otel": "{}", + "metadata": "{}", + }, + ] + + metadata = await shadowfax(task, logger) + + lookups = [call.kwargs["json"] for call in post.call_args_list] + saved = save.call_args.args[1]["message"] if save.called else None + return metadata, lookups, add_callback, saved + + +async def test_shadowfax_searches_the_backbone(redis_mock, mocker, graph_dir): + pool = InlinePool() + + metadata, lookups, _, _ = await run_on_backbone( + mocker, graph_dir, backbone_query(), pool + ) + + assert json.loads(metadata) == {"rehydrate": True} + sent = [query_ids(lookup) for lookup in lookups] + assert sorted(sent, key=str) == sorted( + [ + (["X:S"], ["X:T"]), + (["X:S"], None), + (None, ["X:T"]), + (["X:A"], ["X:B"]), + ], + key=str, + ) + (build,) = pool.builds + response_id, plan, _ = build + assert response_id == "rid" + assert sorted(plan["node_paths"]) == [ + ["X:S", "X:A", "X:B", "X:T"], + ["X:S", "X:M", "X:T"], + ] + assert plan["source"] == "X:S" and plan["target"] == "X:T" + + +async def test_shadowfax_on_the_backbone_sends_valid_lookups_and_builds_the_paths( + redis_mock, mocker, graph_dir +): + """Every lookup a query on the backbone sends -- the open endpoint lookups + and the interior one pinning many ids -- is a valid TRAPI 2.0 query with + the worker's parameters and its own callback, and once they are merged + the paths are built and saved.""" + metadata, lookups, add_callback, saved = await run_on_backbone( + mocker, graph_dir, backbone_query(), InlinePool(build=True) + ) + + assert json.loads(metadata) == {"rehydrate": True} + assert len(lookups) == 4 + for lookup in lookups: + AsyncQuery.from_dict(lookup) + assert "log_level" not in lookup + parameters = {k: v for k, v in lookup["parameters"].items() if k != "subclass"} + assert parameters == { + "timeout": 5, + "tiers": [settings.default_data_tier], + "tier_fallback": False, + "dehydrated": True, + } + assert add_callback.await_count == len(lookups) + callbacks = {lookup["callback"] for lookup in lookups} + assert len(callbacks) == len(lookups) + assert sorted(bound_paths(saved).values()) == [ + ["ab", "bt", "sa"], + ["mt", "sm"], + ] + + +async def test_shadowfax_keeps_to_the_intermediate_category( + redis_mock, mocker, graph_dir +): + """A required_intermediate_categories constraint reaches the path build, + and only paths through a node of that category are kept.""" + pool = InlinePool(build=True) + query = backbone_query( + constraints=[{"required_intermediate_categories": ["biolink:Gene"]}] + ) + + _, lookups, _, saved = await run_on_backbone(mocker, graph_dir, query, pool) + + for lookup in lookups: + AsyncQuery.from_dict(lookup) + (build,) = pool.builds + assert build[1]["intermediate_category"] == "biolink:Gene" + # the path through A, the gene; not the one through M + assert list(bound_paths(saved).values()) == [["ab", "bt", "sa"]] + + +# --- the search graph download ---------------------------------------------- + + +def test_ensure_pathfinder_search_graph_wires_settings(tmp_path, mocker): + ensure = mocker.patch.object(data_download, "ensure_lmdb_dataset") + mocker.patch.object( + data_download.settings, "pathfinder_search_graph_dir", str(tmp_path) + ) + mocker.patch.object( + data_download.settings, "pathfinder_search_graph_url", "http://x/g.tar.gz" + ) + + ensure_pathfinder_search_graph(logger) + + ensure.assert_called_once_with( + name="aragorn_pathfinder", + target_dir=str(tmp_path), + required_files=["offsets.npy", "neighbors.npy", "weights.npy", "nodes.txt"], + url="http://x/g.tar.gz", + logger=logger, + ) diff --git a/tests/unit/test_callback_log_level.py b/tests/unit/test_callback_log_level.py index 44a67e0d..5f5ced8b 100644 --- a/tests/unit/test_callback_log_level.py +++ b/tests/unit/test_callback_log_level.py @@ -231,12 +231,12 @@ async def test_callback_that_cannot_be_stored_asks_for_redelivery( async def _add_ready_callback(response_id, callback_id, logger): ready.append(callback_id) - async def _save_message(callback_id, response, logger, **kwargs): + async def _save_encoded_message(callback_id, blob, logger, **kwargs): assert kwargs.get("raise_on_failure") is True raise ConnectionError("Error 111 connecting to shepherd_broker:6379") monkeypatch.setattr(base_routes, "add_ready_callback", _add_ready_callback) - monkeypatch.setattr(base_routes, "save_message", _save_message) + monkeypatch.setattr(base_routes, "save_encoded_message", _save_encoded_message) response = await callback( ARATargetEnum.ARAGORN, "cb-8", _make_request(_callback_body()) diff --git a/tests/unit/test_callback_unparsed.py b/tests/unit/test_callback_unparsed.py new file mode 100644 index 00000000..f192ff1d --- /dev/null +++ b/tests/unit/test_callback_unparsed.py @@ -0,0 +1,292 @@ +"""The /callback route stores what it was posted without parsing it. + +Parsing a large TRAPI message holds the GIL for seconds, which is long enough +for uvicorn's supervisor to take the server worker for hung and kill it. So the +route only stores the body (``encode_callback_body``) and merge_message parses +it (``read_callback_response``), finding there any body that isn't a TRAPI +response. +""" + +import copy +import io +import logging + +import orjson +import pytest +import zstandard +from starlette.requests import Request + +from shepherd_server import base_routes +from shepherd_utils import db +from shepherd_server.base_routes import ARATargetEnum, callback +from shepherd_utils.db import ( + DecompressedTooLargeError, + NotAJSONObjectError, + decode_message, + encode_callback_body, + get_message, + get_message_sync, + save_message, + save_message_sync, +) +from tests.helpers.generate_messages import generate_response, response_1, response_2 +from workers.merge_message.worker import merge_messages_by_ids, read_callback_response + +logger = logging.getLogger(__name__) + +BODY = {"message": {"results": [{"score": 0.5}], "knowledge_graph": None}} + + +def streaming_frame(data: bytes) -> bytes: + """A zstd frame that omits its content size, as a streaming sender writes.""" + buf = io.BytesIO() + with zstandard.ZstdCompressor().stream_writer(buf, closefd=False) as writer: + writer.write(data) + return buf.getvalue() + + +# --- encode_callback_body ---------------------------------------------------- + + +def test_a_json_body_is_compressed_with_its_size(): + raw = orjson.dumps(BODY) + + blob, size = encode_callback_body(raw, zstd_encoded=False) + + assert size == len(raw) + assert zstandard.frame_content_size(blob) == len(raw) + assert decode_message(blob) == BODY + + +def test_a_zstd_body_with_its_size_is_stored_as_it_came(): + raw = orjson.dumps(BODY) + frame = zstandard.compress(raw) + + blob, size = encode_callback_body(frame, zstd_encoded=True, max_bytes=10_000) + + assert blob is frame + assert size == len(raw) + + +def test_a_streaming_zstd_body_is_recompressed_with_its_size(): + raw = orjson.dumps(BODY) + + blob, size = encode_callback_body(streaming_frame(raw), zstd_encoded=True) + + assert size == len(raw) + assert zstandard.frame_content_size(blob) == len(raw) + assert decode_message(blob) == BODY + + +def test_a_stored_streaming_frame_still_decodes(): + raw = orjson.dumps(BODY) + + assert decode_message(streaming_frame(raw)) == BODY + + +@pytest.mark.parametrize("streaming", [False, True]) +def test_the_cap_is_on_the_decompressed_size(streaming): + raw = orjson.dumps({"message": {"pad": "x" * 5000}}) + frame = streaming_frame(raw) if streaming else zstandard.compress(raw) + + with pytest.raises(DecompressedTooLargeError): + encode_callback_body(frame, zstd_encoded=True, max_bytes=1000) + # no cap + assert encode_callback_body(frame, zstd_encoded=True)[1] == len(raw) + + +@pytest.mark.parametrize( + "body, zstd_encoded", + [ + (b"this is not json", False), + (b" [1, 2]", False), + (b"", False), + (zstandard.compress(b"not json"), True), + (streaming_frame(b"[]"), True), + ], +) +def test_a_body_that_is_not_a_json_object_is_refused(body, zstd_encoded): + with pytest.raises(NotAJSONObjectError): + encode_callback_body(body, zstd_encoded) + + +def test_a_bad_zstd_frame_is_refused(): + with pytest.raises(zstandard.ZstdError): + encode_callback_body(b"not a zstd frame", zstd_encoded=True) + + +# --- the route --------------------------------------------------------------- + + +def make_request(body: bytes, headers=None) -> Request: + headers = {"content-length": len(body), **(headers or {})} + scope = { + "type": "http", + "method": "POST", + "path": "/aragorn/callback/cb-1", + "headers": [(k.encode(), str(v).encode()) for k, v in headers.items()], + } + sent = False + + async def receive(): + nonlocal sent + if sent: + return {"type": "http.disconnect"} + sent = True + return {"type": "http.request", "body": body, "more_body": False} + + return Request(scope, receive) + + +@pytest.fixture +def route(redis_mock, monkeypatch, mocker): + """The route with its Postgres lookups stubbed; returns the enqueued tasks + and the callbacks recorded as ready.""" + tasks, ready = [], [] + + async def _get_callback_query_id(callback_id, logger): + return ("q-1", "{}") + + async def _get_query_state(query_id, logger): + return [None] * 7 + ["resp-1"] + + async def _add_ready_callback(response_id, callback_id, logger): + ready.append((response_id, callback_id)) + + async def _add_task(stream, fields, logger): + tasks.append((stream, fields)) + + async def _save_logs(response_id, logger): + return None + + monkeypatch.setattr(base_routes, "get_callback_query_id", _get_callback_query_id) + monkeypatch.setattr(base_routes, "get_query_state", _get_query_state) + monkeypatch.setattr(base_routes, "add_ready_callback", _add_ready_callback) + monkeypatch.setattr(base_routes, "add_task", _add_task) + monkeypatch.setattr(base_routes, "save_logs", _save_logs) + # what the route parses: the stored query, but never the posted body + loads = mocker.patch("orjson.loads", wraps=orjson.loads) + return tasks, ready, loads + + +async def test_the_route_stores_the_body_without_parsing_it(route): + tasks, ready, loads = route + await save_message("q-1", {"message": {}}, logger) + raw = orjson.dumps(BODY) + + response = await callback(ARATargetEnum.ARAGORN, "cb-1", make_request(raw)) + + assert response.status_code == 200 + assert raw not in [call.args[0] for call in loads.call_args_list] + assert await get_message("cb-1", logger) == BODY + assert ready == [("resp-1", "cb-1")] + assert tasks[0][0] == "merge_message" + + +async def test_the_route_stores_a_zstd_body_as_it_came(route): + await save_message("q-1", {"message": {}}, logger) + frame = zstandard.compress(orjson.dumps(BODY)) + + response = await callback( + ARATargetEnum.ARAGORN, + "cb-1", + make_request(frame, {"content-encoding": "zstd"}), + ) + + assert response.status_code == 200 + assert await db.data_db_client.get("cb-1") == frame + + +async def test_the_route_refuses_a_body_that_is_not_json(route): + tasks, ready, _ = route + + response = await callback( + ARATargetEnum.ARAGORN, "cb-1", make_request(b"oops") + ) + + assert response.status_code == 422 + assert ready == [] and tasks == [] + + +# --- merge_message reads it -------------------------------------------------- + + +@pytest.fixture +def sync_store(mocker): + storage = {} + client = mocker.Mock() + client.set.side_effect = lambda key, blob, ex=None: storage.__setitem__(key, blob) + client.get.side_effect = lambda key: storage.get(key) + mocker.patch("shepherd_utils.db._get_sync_data_db", return_value=client) + return storage + + +def test_read_callback_response_fills_in_what_is_missing(sync_store, caplog): + sync_store["cb"] = encode_callback_body( + orjson.dumps({"message": {"query_graph": {"nodes": {}}}}), False + )[0] + + with caplog.at_level(logging.INFO): + response = read_callback_response("cb", logger) + + assert response["message"]["results"] == [] + assert response["message"]["knowledge_graph"] == {"nodes": {}, "edges": {}} + assert "[cb] Got back 0 results, 0 nodes and 0 edges." in caplog.text + + +@pytest.mark.parametrize("message", [None, [], "nope"]) +def test_read_callback_response_without_a_message(sync_store, message): + save_message_sync("cb", {"message": message}) + + response = read_callback_response("cb", logger) + + assert response["message"] == { + "results": [], + "knowledge_graph": {"nodes": {}, "edges": {}}, + } + + +def test_read_callback_response_logs_a_long_id_list_as_its_count(sync_store, caplog): + ids = [f"X:{i}" for i in range(500)] + save_message_sync( + "cb", {"message": {"query_graph": {"nodes": {"n0": {"ids": ids}}}}} + ) + + with caplog.at_level(logging.DEBUG): + read_callback_response("cb", logger) + + assert "500 curies" in caplog.text + assert "X:499" not in caplog.text + + +@pytest.mark.parametrize( + "blob", + [ + zstandard.compress(b'{"message": '), + zstandard.compress(b"[1, 2]"), + b"not zstd", + ], +) +def test_read_callback_response_skips_what_is_not_a_response(sync_store, blob, caplog): + sync_store["cb"] = blob + + assert read_callback_response("cb", logger) is None + assert "cb" in caplog.text + + +def test_read_callback_response_skips_a_missing_callback(sync_store): + assert read_callback_response("gone", logger) is None + + +def test_a_bad_callback_does_not_stop_the_rest_merging(sync_store): + query_graph = response_1["message"]["query_graph"] + save_message_sync("qid", {"message": {"query_graph": copy.deepcopy(query_graph)}}) + save_message_sync("rid", generate_response()) + sync_store["bad"] = zstandard.compress(b'{"message": {"results": [') + save_message_sync("good", copy.deepcopy(response_2)) + + merged, entries = merge_messages_by_ids("test_ara", "qid", "rid", ["bad", "good"]) + + assert merged == ["good"] + assert get_message_sync("rid")["message"]["results"] + assert any("bad is not valid JSON" in str(entry) for entry in entries) diff --git a/tests/unit/test_data_download.py b/tests/unit/test_data_download.py index d227113a..4009afe5 100644 --- a/tests/unit/test_data_download.py +++ b/tests/unit/test_data_download.py @@ -16,6 +16,7 @@ ensure_lmdb_dataset, ensure_omnicorp_lmdb, ensure_pathfinder_embeddings, + ensure_pathfinder_search_graph, ) logger = logging.getLogger(__name__) @@ -159,6 +160,38 @@ def test_ensure_pathfinder_embeddings_wires_settings(tmp_path, mocker): assert (target / "lock.mdb").exists() +def test_ensure_pathfinder_search_graph_downloads_the_csr_and_node_list( + tmp_path, mocker +): + files = ["offsets.npy", "neighbors.npy", "weights.npy", "nodes.txt"] + url = _make_archive(tmp_path, files) + target = tmp_path / "pathfinder_search_graph" + mocker.patch.object( + data_download.settings, "pathfinder_search_graph_dir", str(target) + ) + mocker.patch.object(data_download.settings, "pathfinder_search_graph_url", url) + + ensure_pathfinder_search_graph(logger) + + assert all((target / name).exists() for name in files) + + +def test_ensure_pathfinder_search_graph_needs_every_array_and_the_node_list( + tmp_path, mocker +): + """An archive without nodes.txt would leave the worker unable to map its + search graph's indices back to CURIEs: it is refused, naming the file.""" + url = _make_archive(tmp_path, ["offsets.npy", "neighbors.npy", "weights.npy"]) + target = tmp_path / "pathfinder_search_graph" + mocker.patch.object( + data_download.settings, "pathfinder_search_graph_dir", str(target) + ) + mocker.patch.object(data_download.settings, "pathfinder_search_graph_url", url) + + with pytest.raises(RuntimeError, match="nodes.txt"): + ensure_pathfinder_search_graph(logger) + + # --- download timeout -------------------------------------------------------- # # urllib defaults to no timeout, so a connection that opens and then stalls hung diff --git a/tests/unit/test_lookup_query_data_lost.py b/tests/unit/test_lookup_query_data_lost.py index c8f2232a..d6b78339 100644 --- a/tests/unit/test_lookup_query_data_lost.py +++ b/tests/unit/test_lookup_query_data_lost.py @@ -146,6 +146,12 @@ def lookup(request, redis_mock, cleanup, mocker): return_value=mocker.Mock(status_code=200), ) mocker.patch.object(module.asyncio, "sleep", new_callable=mocker.AsyncMock) + if hasattr(module, "POOL"): + # aragorn.pathfinder builds its paths in a process pool that only + # poll_for_tasks sets up; the build itself is not under test here. + pool = mocker.Mock() + pool.run = mocker.AsyncMock(return_value=[]) + mocker.patch.object(module, "POOL", pool) running = mocker.patch.object( module, "get_running_callbacks", diff --git a/tests/unit/test_merge_message_helpers.py b/tests/unit/test_merge_message_helpers.py index 8466d3b2..9d810796 100644 --- a/tests/unit/test_merge_message_helpers.py +++ b/tests/unit/test_merge_message_helpers.py @@ -458,9 +458,11 @@ def test_merge_messages_combines_aux_graphs(): assert "aux_b" in aux -def test_merge_messages_pathfinder_query_returns_single_result(): - """A pathfinder query (paths instead of edges) should return a single - pathfinder result with the start/end node IDs threaded through.""" +def test_merge_messages_pathfinder_query_only_accumulates_the_kgraph(): + """A pathfinder query's callbacks are aragorn.pathfinder's edge lookups: + their edges are merged into the knowledge graph, but their results are + not turned into paths here (the pathfinder worker builds those once every + lookup is in), so the response keeps the results it already had.""" pathfinder_qg = { "nodes": { "n0": {"ids": ["MONDO:0001"]}, @@ -473,7 +475,20 @@ def test_merge_messages_pathfinder_query_returns_single_result(): new_response = { "message": { "query_graph": pathfinder_qg, - "knowledge_graph": {"nodes": {}, "edges": {}}, + "knowledge_graph": { + "nodes": { + "MONDO:0001": {"categories": ["biolink:Disease"]}, + "MONDO:0002": {"categories": ["biolink:Disease"]}, + }, + "edges": { + "kedge_lookup": { + "subject": "MONDO:0001", + "predicate": "biolink:related_to", + "object": "MONDO:0002", + "sources": [], + } + }, + }, "auxiliary_graphs": {}, "results": [ { @@ -483,7 +498,7 @@ def test_merge_messages_pathfinder_query_returns_single_result(): }, "analyses": [ { - "edge_bindings": {"p1": {"ids": ["kedge_pathfinder"]}}, + "edge_bindings": {"e0": {"ids": ["kedge_lookup"]}}, } ], "score": 0.42, @@ -506,14 +521,91 @@ def test_merge_messages_pathfinder_query_returns_single_result(): new_response=new_response, logger=logger, ) - assert len(out["message"]["results"]) == 1 - pf = out["message"]["results"][0] - assert pf["node_bindings"]["n0"]["ids"][0] == "MONDO:0001" - assert pf["node_bindings"]["n1"]["ids"][0] == "MONDO:0002" - # An auxiliary graph should have been created and associated with an - # analysis path binding. - assert pf["analyses"][0]["score"] == 0.42 - assert "p1" in pf["analyses"][0]["path_bindings"] + assert out["message"]["results"] == [] + assert set(out["message"]["knowledge_graph"]["nodes"]) == { + "MONDO:0001", + "MONDO:0002", + } + assert len(out["message"]["knowledge_graph"]["edges"]) == 1 + + +def test_merge_messages_pathfinder_callbacks_accumulate_under_its_query_graph(): + """Every edge lookup's callback adds to the same knowledge graph, and the + response keeps the pathfinder's own query graph (its path), not a + lookup's one-hop one.""" + pathfinder_qg = { + "nodes": { + "n0": {"ids": ["MONDO:0001"]}, + "n1": {"ids": ["MONDO:0002"]}, + }, + "paths": { + "p0": {"subject": "n0", "object": "n1"}, + }, + } + + def lookup_callback(edge_id, subject, obj): + return { + "message": { + "query_graph": { + "nodes": {"n0": {"ids": [subject]}, "n1": {"ids": [obj]}}, + "edges": { + "e0": { + "subject": "n0", + "object": "n1", + "predicates": ["biolink:related_to"], + } + }, + }, + "knowledge_graph": { + "nodes": { + subject: {"categories": ["biolink:Gene"]}, + obj: {"categories": ["biolink:Gene"]}, + }, + "edges": { + edge_id: { + "subject": subject, + "predicate": "biolink:affects", + "object": obj, + "sources": [], + } + }, + }, + "results": [ + { + "node_bindings": { + "n0": {"ids": [subject]}, + "n1": {"ids": [obj]}, + }, + "analyses": [{"edge_bindings": {"e0": {"ids": [edge_id]}}}], + } + ], + } + } + + response = { + "message": { + "query_graph": pathfinder_qg, + "knowledge_graph": {"nodes": {}, "edges": {}}, + "results": [], + } + } + for callback in ( + lookup_callback("kedge_1", "MONDO:0001", "NCBIGene:1"), + lookup_callback("kedge_2", "NCBIGene:1", "MONDO:0002"), + ): + response = merge_messages( + target="aragorn", + original_query_graph=pathfinder_qg, + response=response, + new_response=callback, + logger=logger, + ) + + kgraph = response["message"]["knowledge_graph"] + assert set(kgraph["edges"]) == {"kedge_1", "kedge_2"} + assert set(kgraph["nodes"]) == {"MONDO:0001", "NCBIGene:1", "MONDO:0002"} + assert response["message"]["results"] == [] + assert response["message"]["query_graph"] == pathfinder_qg def test_merge_messages_pathfinder_missing_endpoint_raises(): diff --git a/tests/unit/test_score_paths.py b/tests/unit/test_score_paths.py new file mode 100644 index 00000000..e8bcfa63 --- /dev/null +++ b/tests/unit/test_score_paths.py @@ -0,0 +1,799 @@ +"""Tests for the score_paths worker and its sequence-reranker ``scorer``. + +The reranker itself is a TorchScript model squashbert exports; here a fake +model stands in, so what is checked is everything around it: how a TRAPI path +is read into nodes, categories and hop phrases, the features looked up for it, +and how the model's output becomes each analysis' score. + +torch is not a test requirement (its default Linux wheel is gigabytes), so when +it is missing a stand-in with the handful of calls the scorer makes is +installed in its place. +""" + +import contextlib +import hashlib +import io +import json +import logging +import sys +import types +from pathlib import Path + +import lmdb +import numpy as np +import pytest + +try: + import torch # noqa: F401 +except ImportError: + _torch = types.ModuleType("torch") + _torch.int64 = np.int64 + _torch.from_numpy = np.asarray + _torch.full = lambda shape, value, dtype=None: np.full(shape, value, dtype=dtype) + _torch.zeros = lambda shape, dtype=None: np.zeros(shape, dtype=dtype) + _torch.inference_mode = contextlib.nullcontext + _torch.set_num_threads = lambda n: None + + def _no_jit(*args, **kwargs): + raise RuntimeError("torch is not installed") + + _torch.jit = types.SimpleNamespace(load=_no_jit) + sys.modules["torch"] = _torch + +# The worker imports its scorer as a top-level module, as it runs in its image. +sys.path.insert(0, str(Path(__file__).parents[2] / "workers" / "score_paths")) +import scorer # noqa: E402 +from workers.score_paths import worker as score_worker # noqa: E402 + +logger = logging.getLogger(__name__) + + +# --- fakes ------------------------------------------------------------------- + + +class FakeBmt: + """The few biolink Toolkit calls the scorer makes.""" + + predicates = { + "related to", + "interacts with", + "affects", + "affected by", + "causes", + "caused by", + "treats", + "treated by", + "broken", + } + symmetric = {"related to", "interacts with"} + inverses = { + "affects": "affected by", + "affected by": "affects", + "causes": "caused by", + "caused by": "causes", + "treats": "treated by", + "treated by": "treats", + } + depths = { + "biolink:NamedThing": 0, + "biolink:BiologicalEntity": 2, + "biolink:Disease": 4, + "biolink:Gene": 5, + "biolink:SmallMolecule": 6, + } + mixins = {"biolink:GeneOrGeneProduct", "biolink:ChemicalOrDrugOrTreatment"} + + def is_predicate(self, name): + return name in self.predicates + + def is_symmetric(self, name): + return name in self.symmetric + + def get_inverse(self, name): + if name == "broken": + raise AttributeError("declared inverse is not in the schema") + return self.inverses.get(name) + + def is_category(self, name): + return name in self.depths + + def get_element_depth(self, name): + return self.depths[name] + + def is_mixin(self, name): + return name in self.mixins + + +BMT = FakeBmt() + + +class FakeTxn: + def __init__(self, data): + self.data = data + + def get(self, key): + return self.data.get(key) + + +class _Output: + def __init__(self, values): + self.values = values + + def float(self): + return self + + def numpy(self): + return self.values + + +class FakeModel: + """Scores a path by the first value of its second node's vector, and + records what it was given.""" + + def __init__(self): + self.calls = [] + + def __call__(self, node_x, edge_x, n_hops, scalars, extra): + node_x, edge_x, n_hops, scalars = ( + np.asarray(x) for x in (node_x, edge_x, n_hops, scalars) + ) + self.calls.append( + { + "node_x": node_x, + "edge_x": edge_x, + "n_hops": n_hops, + "scalars": scalars, + } + ) + return _Output(node_x[:, 1, 0].astype(np.float32)) + + +DIM = 2 +SCALAR_NAMES = ( + ["n_hops_scaled"] + + [name for name, _, _ in scorer._NODE_COLUMNS] + + [f"{base}_{stat}" for base in scorer._EDGE_BASES for stat in ("mean", "min")] +) + + +def make_config(**overrides): + config = { + "embed_dim": DIM, + "scalar_names": SCALAR_NAMES, + "node_context": True, + "edge_context": True, + "length_scale": {2: (0.1, 2.0)}, + "biolink_version": "4.2.5", + } + config.update(overrides) + return config + + +def store(**vectors): + """An LMDB stand-in holding the given vectors, keyed as squashbert's + export writes them.""" + data = {} + for kind, entries in vectors.items(): + for text, values in entries.items(): + data[scorer.lmdb_key(kind, text)] = np.asarray( + values, dtype=scorer.DTYPES[kind] + ).tobytes() + return data + + +def edge(subject, obj, predicate="biolink:related_to", **extra): + return {"subject": subject, "object": obj, "predicate": predicate, **extra} + + +def node(*categories): + return {"categories": list(categories)} + + +def sigmoid(z): + return 1.0 / (1.0 + np.exp(-z)) + + +# --- keys -------------------------------------------------------------------- + + +def test_lmdb_key(): + assert scorer.lmdb_key("node", "MONDO:1") == b"node:MONDO:1" + + +def test_lmdb_key_hashes_what_lmdb_cannot_hold(): + text = "x" * 600 + + key = scorer.lmdb_key("phrase", text) + + expected = hashlib.sha256(f"phrase:{text}".encode()).hexdigest() + assert key == f"phrase#{expected}".encode() + assert len(key) < scorer.MAX_KEY_BYTES + + +def test_edge_context_key_ignores_the_order_of_the_pair(): + key = scorer.edge_context_key("B:2", "A:1") + + digest = hashlib.blake2b(b"A:1|B:2", digest_size=8).digest() + assert key == b"edge_context#" + digest + assert scorer.edge_context_key("A:1", "B:2") == key + + +# --- loading ----------------------------------------------------------------- + + +def test_load_reranker(mocker): + model = mocker.Mock() + load = mocker.patch.object(scorer.torch.jit, "load", return_value=model) + txn = FakeTxn( + { + scorer.MODEL_KEY: b"model bytes", + scorer.CONFIG_KEY: json.dumps( + {"length_scale": {"2": [0.5, 2], "3": [1, 3]}, "embed_dim": 4} + ).encode(), + } + ) + + loaded, config = scorer.load_reranker(txn) + + assert loaded is model + model.eval.assert_called_once() + blob = load.call_args.args[0] + assert isinstance(blob, io.BytesIO) and blob.getvalue() == b"model bytes" + assert config == {"length_scale": {2: (0.5, 2.0), 3: (1.0, 3.0)}, "embed_dim": 4} + + +@pytest.mark.parametrize("missing", [scorer.MODEL_KEY, scorer.CONFIG_KEY]) +def test_load_reranker_without_one(missing): + data = {scorer.MODEL_KEY: b"model", scorer.CONFIG_KEY: b"{}"} + del data[missing] + + with pytest.raises(RuntimeError, match="no reranker"): + scorer.load_reranker(FakeTxn(data)) + + +def test_load_toolkit_reads_the_trained_biolink_release(mocker): + toolkit = mocker.patch.object(scorer, "Toolkit") + + scorer.load_toolkit({"biolink_version": "4.2.5"}) + + toolkit.assert_called_once_with( + schema="https://raw.githubusercontent.com/biolink/biolink-model/" + "refs/tags/v4.2.5/biolink-model.yaml", + predicate_map="https://raw.githubusercontent.com/biolink/biolink-model/" + "refs/tags/v4.2.5/predicate_mapping.yaml", + ) + + +# --- reading a path ---------------------------------------------------------- + + +def test_predicate_names(): + assert scorer.predicate_name("biolink:affected_by") == "affected by" + assert scorer.predicate_curie("affected by") == "biolink:affected_by" + + +@pytest.mark.parametrize( + "predicate, inverse", + [ + ("biolink:affects", "biolink:affected_by"), + ("biolink:interacts_with", "biolink:interacts_with"), + ("biolink:has_chemical_role", "biolink:is_chemical_role_of"), + ("is_chemical_role_of", "biolink:has_chemical_role"), + ("biolink:not_a_predicate", None), + ("biolink:broken", None), + ], +) +def test_inverse_predicate(predicate, inverse): + assert scorer.inverse_predicate(BMT, predicate) == inverse + + +@pytest.mark.parametrize( + "categories, expected", + [ + ( + ("biolink:BiologicalEntity", "biolink:Gene", "biolink:NamedThing"), + "biolink:Gene", + ), + (("biolink:GeneOrGeneProduct", "biolink:Disease"), "biolink:Disease"), + (("biolink:GeneOrGeneProduct",), "biolink:GeneOrGeneProduct"), + ], +) +def test_most_specific_category(categories, expected): + assert scorer.most_specific_category(BMT, categories) == expected + + +def test_most_specific_category_needs_a_biolink_one(): + with pytest.raises(ValueError, match="No biolink categories"): + scorer.most_specific_category(BMT, ("not:a_category",)) + + +def qualifiers(**slots): + return [ + {"qualifier_type_id": f"biolink:{slot}", "qualifier_value": value} + for slot, value in slots.items() + ] + + +@pytest.mark.parametrize( + "slots, forward, backward", + [ + ({}, "affects", "affected by"), + ( + { + "qualified_predicate": "biolink:causes", + "object_aspect_qualifier": "activity", + "object_direction_qualifier": "increased", + }, + "causes increased activity of", + "has increased activity caused by", + ), + ( + {"object_aspect_qualifier": "activity_or_abundance"}, + "affects activity or abundance of", + "has activity or abundance affected by", + ), + # a qualified predicate only reads with a direction or without an aspect + ( + { + "qualified_predicate": "biolink:causes", + "object_aspect_qualifier": "activity", + }, + "affects activity of", + "has activity affected by", + ), + # a qualified predicate with no inverse reads backwards as the bare one + ( + {"qualified_predicate": "biolink:regulates"}, + "regulates", + "affected by", + ), + # these have no wording, so the edge reads as its bare predicate + ({"object_direction_qualifier": "increased"}, "affects", "affected by"), + ({"subject_aspect_qualifier": "activity"}, "affects", "affected by"), + ], +) +def test_edge_phrase(slots, forward, backward): + affects = edge("A", "B", "biolink:affects", qualifiers=qualifiers(**slots)) + + assert scorer.edge_phrase(BMT, affects, backward=False) == forward + assert scorer.edge_phrase(BMT, affects, backward=True) == backward + + +def test_edge_phrase_cannot_walk_back_an_edge_with_no_inverse(): + one_way = edge("A", "B", "biolink:not_a_predicate") + + assert scorer.edge_phrase(BMT, one_way, backward=False) == "not a predicate" + assert scorer.edge_phrase(BMT, one_way, backward=True) is None + + +def test_edge_phrase_rejects_a_qualifier_it_cannot_read(): + bad = edge( + "A", + "B", + "biolink:affects", + qualifiers=qualifiers(object_aspect_qualifier="activity; drop table"), + ) + + with pytest.raises(ValueError, match="cannot be read"): + scorer.edge_phrase(BMT, bad, backward=False) + + +def negated(value): + return {"attributes": [{"attribute_type_id": "biolink:negated", "value": value}]} + + +def test_hop_phrase_reads_every_edge_of_the_hop(): + edges = { + "e1": edge("A", "B", "biolink:affects"), + "e2": edge("B", "A", "biolink:treats"), + "e3": edge("A", "B", "biolink:affects"), + "sub": edge("A", "B", "biolink:subclass_of"), + "other": edge("A", "C", "biolink:causes"), + } + + phrase = scorer.hop_phrase(BMT, "A", "B", list(edges), edges) + + assert phrase == "[affects or treated by]" + assert scorer.hop_phrase(BMT, "A", "B", ["e1", "missing"], edges) == "affects" + + +def test_hop_phrase_drops_a_phrase_only_negated_edges_carry(): + edges = { + "not": edge("A", "B", "biolink:affects", **negated(True)), + "yes": edge("A", "B", "biolink:treats", **negated(False)), + "also_not": edge( + "A", + "B", + "biolink:causes", + attributes=[{"original_attribute_name": "negated", "value": True}], + ), + "affirmed": edge("A", "B", "biolink:causes"), + } + + assert scorer.hop_phrase(BMT, "A", "B", list(edges), edges) == "[causes or treats]" + assert scorer.hop_phrase(BMT, "A", "B", ["not"], edges) is None + + +def test_hop_phrase_rejects_a_negated_that_is_not_a_boolean(): + edges = {"e": edge("A", "B", "biolink:affects", **negated("yes"))} + + with pytest.raises(ValueError, match="expected a boolean"): + scorer.hop_phrase(BMT, "A", "B", ["e"], edges) + + +def test_hop_phrase_with_nothing_walkable(): + edges = {"e": edge("B", "A", "biolink:not_a_predicate")} + + assert scorer.hop_phrase(BMT, "A", "B", ["e"], edges) is None + + +def test_order_path_nodes(): + edges = { + "bt": edge("T", "B"), + "ab": edge("A", "B"), + "sa": edge("S", "A"), + "sa2": edge("A", "S"), + } + + assert scorer.order_path_nodes("S", "T", list(edges), edges) == [ + "S", + "A", + "B", + "T", + ] + + +def test_order_path_nodes_needs_a_connected_path(): + edges = {"sa": edge("S", "A"), "bt": edge("B", "T")} + + with pytest.raises(ValueError, match="Could not construct path"): + scorer.order_path_nodes("S", "T", ["sa", "bt", "gone"], edges) + + +def test_read_path(): + kgraph = { + "nodes": { + "S": node("biolink:BiologicalEntity", "biolink:Disease"), + "A": node(), + "T": node("biolink:SmallMolecule"), + }, + "edges": { + "sa": edge("S", "A", "biolink:affects"), + "ta": edge("T", "A", "biolink:treats"), + }, + } + + assert scorer.read_path(BMT, "S", "T", ["sa", "ta"], kgraph) == ( + ["S", "A", "T"], + ["biolink:Disease", "biolink:NamedThing", "biolink:SmallMolecule"], + ["affects", "treated by"], + ) + + +def test_read_path_needs_every_node(): + kgraph = {"nodes": {"S": node()}, "edges": {"st": edge("S", "T")}} + + with pytest.raises(ValueError, match="T is not in the knowledge graph"): + scorer.read_path(BMT, "S", "T", ["st"], kgraph) + + +def test_read_path_needs_every_hop_walkable(): + kgraph = { + "nodes": {"S": node(), "T": node()}, + "edges": {"ts": edge("T", "S", "biolink:not_a_predicate")}, + } + + with pytest.raises(ValueError, match="no edge of the path can be walked"): + scorer.read_path(BMT, "S", "T", ["ts"], kgraph) + + +def test_is_direct_path(): + kgraph = { + "edges": { + "st": edge("S", "T"), + "ts": edge("T", "S"), + "sa": edge("S", "A"), + } + } + + assert scorer.is_direct_path("S", "T", ["st", "ts"], kgraph) + assert not scorer.is_direct_path("S", "T", ["st", "sa"], kgraph) + assert not scorer.is_direct_path("S", "T", ["missing"], kgraph) + assert not scorer.is_direct_path("S", "T", [], kgraph) + + +# --- features ---------------------------------------------------------------- + + +def test_path_scalars(): + node_table = { + "S": [1, 2, 0.5, 1], + "A": [3, 4, 0.0, 0], + "T": [5, 0, 1.0, 1], + } + edge_table = { + frozenset("SA"): [1, 10, 1, 1, 4], + frozenset("AT"): [3, 0, 3, 2, 6], + } + + scalars = scorer.path_scalars( + SCALAR_NAMES, + ["S", "A", "T"], + lambda c: np.array(node_table[c], dtype=np.float32), + lambda a, b: np.array(edge_table[frozenset((a, b))], dtype=np.float32), + ) + + got = dict(zip(SCALAR_NAMES, scalars.tolist())) + assert got["n_hops_scaled"] == 2 / scorer.N_HOPS_SCALE + assert got["deg_out_max"] == 5 and got["deg_out_mean"] == 3 + assert got["deg_in_max"] == 4 and got["deg_in_mean"] == 2 + assert got["ic_mean"] == 0.5 + assert got["ic_present_frac"] == pytest.approx(2 / 3) + assert got["edge_count_mean"] == 2 and got["edge_count_min"] == 1 + assert got["pubs_mean"] == 5 and got["pubs_min"] == 0 + assert got["attributes_seen_mean"] == 5 and got["attributes_seen_min"] == 4 + + +def test_path_scalars_without_the_context_tables(): + scalars = scorer.path_scalars(SCALAR_NAMES, ["S", "A", "B", "T"]) + + assert scalars[0] == 3 / scorer.N_HOPS_SCALE + assert not scalars[1:].any() + + +def test_path_features(): + data = store( + node={"S": [1, 2], "T": [3, 4]}, + category={"biolink:Disease": [5, 6], "biolink:Gene": [7, 8]}, + phrase={"affects": [0.25, 0.5]}, + node_context={"S": [2, 0, 0, 0]}, + ) + data[scorer.edge_context_key("S", "T")] = np.array( + [4, 0, 0, 0, 0], dtype=np.float32 + ).tobytes() + kept = {} + + node_x, edge_x, scalars = scorer.path_features( + FakeTxn(data), + make_config(), + ["S", "T"], + ["biolink:Disease", "biolink:Gene"], + ["affects"], + kept, + ) + + assert node_x.tolist() == [[1, 2, 5, 6], [3, 4, 7, 8]] + assert edge_x.tolist() == [[0.25, 0.5]] + got = dict(zip(SCALAR_NAMES, scalars.tolist())) + # T has no node context: zeros + assert got["deg_out_max"] == 2 and got["deg_out_mean"] == 1 + assert got["edge_count_mean"] == 4 + # the categories and phrase are kept for the next path, the nodes are not + assert set(kept) == { + ("category", "biolink:Disease"), + ("category", "biolink:Gene"), + ("phrase", "affects"), + } + + +def test_path_features_names_the_missing_vector(): + data = store(node={"S": [1, 2]}, category={"biolink:Disease": [0, 0]}) + + with pytest.raises(KeyError, match="node:T"): + scorer.path_features( + FakeTxn(data), + make_config(), + ["S", "T"], + ["biolink:Disease", "biolink:Disease"], + ["related to"], + ) + + +# --- scoring a message ------------------------------------------------------- + + +def scoring_message(): + """One result with a path of every kind, and one with no bindings.""" + kgraph = { + "nodes": {name: node("biolink:Disease") for name in ("S", "A", "B", "C", "T")}, + "edges": { + "st": edge("S", "T"), + "sa": edge("S", "A", "biolink:affects"), + "at": edge("A", "T", "biolink:affects"), + "sb": edge("S", "B"), + "bt": edge("B", "T"), + "ac": edge("A", "C"), + "ct": edge("C", "T"), + }, + } + aux = { + "direct": {"edges": ["st"]}, + "via_a": {"edges": ["sa", "at"]}, + "via_b": {"edges": ["sb", "bt"]}, + "broken": {"edges": ["sa"]}, + "three_hops": {"edges": ["sa", "ac", "ct"]}, + } + + def analysis(aux_id): + return {"path_bindings": {"p0": {"ids": [aux_id]}}} + + return { + "message": { + "query_graph": {"paths": {"p0": {"subject": "n0", "object": "n1"}}}, + "knowledge_graph": kgraph, + "auxiliary_graphs": aux, + "results": [ + { + "node_bindings": {"n0": {"ids": ["S"]}, "n1": {"ids": ["T"]}}, + "analyses": [ + analysis("direct"), + analysis("via_a"), + analysis("via_b"), + analysis("no_such_aux"), + {"path_bindings": {}}, + analysis("broken"), + analysis("three_hops"), + {**analysis("via_a"), "score": 0.9}, + ], + }, + {"node_bindings": {}, "analyses": [analysis("direct")]}, + ], + } + } + + +def scoring_store(): + # B has no vector: its path cannot be scored + return FakeTxn( + store( + node={"S": [0, 0], "A": [0.5, 0], "C": [1, 0], "T": [0, 0]}, + category={"biolink:Disease": [0, 0]}, + phrase={"affects": [0, 0], "related to": [0, 0]}, + ) + ) + + +@pytest.mark.parametrize("chunk_size", [1, 1024]) +def test_score_message(chunk_size): + message = scoring_message() + model = FakeModel() + + stats = scorer.score_message( + message, scoring_store(), model, make_config(), BMT, chunk_size + ) + + scored, unbound = message["message"]["results"] + via_a = sigmoid((0.5 - 0.1) / 2.0) + assert [a["score"] for a in scored["analyses"]] == pytest.approx( + [1.0, via_a, 0.0, 0.0, 0.0, 0.0, 0.0, via_a] + ) + assert unbound["analyses"][0]["score"] == 0.0 + assert {k: v for k, v in stats.items() if k != "model_seconds"} == { + "analyses": 9, + "scored": 2, + "direct": 1, + "no_binding": 3, + "bad_path": 1, + "no_scale": 1, + "missing": 1, + "missing_samples": ["node:B"], + } + assert len(model.calls) == (2 if chunk_size == 1 else 1) + for call in model.calls: + assert call["n_hops"].tolist() == [2] * len(call["n_hops"]) + assert call["node_x"].shape[1:] == (3, 2 * DIM) + assert call["edge_x"].shape[1:] == (2, DIM) + assert call["scalars"].shape[1:] == (len(SCALAR_NAMES),) + + +def test_score_message_ranks_lengths_on_their_own_scales(): + message = scoring_message() + config = make_config(length_scale={2: (0.1, 2.0), 3: (0.25, 0.5)}) + model = FakeModel() + + stats = scorer.score_message(message, scoring_store(), model, config, BMT, 1024) + + scores = [a["score"] for a in message["message"]["results"][0]["analyses"]] + assert scores[1] == pytest.approx(sigmoid((0.5 - 0.1) / 2.0)) + # the three-hop path's second node is A as well, on the three-hop scale + assert scores[6] == pytest.approx(sigmoid((0.5 - 0.25) / 0.5)) + assert stats["no_scale"] == 0 + assert sorted(call["n_hops"][0] for call in model.calls) == [2, 3] + + +# --- the worker -------------------------------------------------------------- + + +class FakeEnv: + @contextlib.contextmanager + def begin(self): + yield scoring_store() + + +@pytest.fixture +def scoring_state(mocker): + mocker.patch.object(score_worker, "embedding_env", FakeEnv()) + mocker.patch.object(score_worker, "model", FakeModel()) + mocker.patch.object(score_worker, "config", make_config()) + mocker.patch.object(score_worker, "bmt", BMT) + + +def test_score_paths_scores_and_saves(mocker, scoring_state, caplog): + message = scoring_message() + mocker.patch.object(score_worker, "get_message_sync", return_value=message) + save = mocker.patch.object(score_worker, "save_response_sync") + + with caplog.at_level(logging.INFO): + score_worker.score_paths("rid", logger) + + save.assert_called_once_with("rid", message) + assert message["message"]["results"][0]["analyses"][0]["score"] == 1.0 + assert "Scored 2/9 analyses" in caplog.text + assert "missing embedding: 1" in caplog.text + + +def test_score_paths_still_saves_when_scoring_fails(mocker, scoring_state): + message = scoring_message() + mocker.patch.object(score_worker, "get_message_sync", return_value=message) + mocker.patch.object(scorer, "score_message", side_effect=RuntimeError("boom")) + save = mocker.patch.object(score_worker, "save_response_sync") + + score_worker.score_paths("rid", logger) + + save.assert_called_once() + analyses = message["message"]["results"][0]["analyses"] + # an analysis that had a score keeps it; the rest get 0 + assert [a["score"] for a in analyses] == [0.0] * 7 + [0.9] + + +def test_score_paths_retries_the_save(mocker, scoring_state): + mocker.patch.object( + score_worker, "get_message_sync", return_value=scoring_message() + ) + mocker.patch.object(score_worker.time, "sleep") + save = mocker.patch.object( + score_worker, + "save_response_sync", + side_effect=[ConnectionError("down"), ConnectionError("down"), None], + ) + + score_worker.score_paths("rid", logger) + + assert save.call_count == 3 + + +def test_validate_scoring_data_needs_a_reranker(tmp_path, mocker): + mocker.patch.object(score_worker, "EMBEDDING_DIR", str(tmp_path)) + env = lmdb.open(str(tmp_path)) + with env.begin(write=True) as txn: + txn.put(b"node:S", b"\x00\x00") + env.close() + + with pytest.raises(RuntimeError, match="no reranker"): + score_worker._validate_scoring_data(logger) + + +def test_validate_scoring_data(tmp_path, mocker, caplog): + mocker.patch.object(score_worker, "EMBEDDING_DIR", str(tmp_path)) + mocker.patch.object(scorer.torch.jit, "load", return_value=mocker.Mock()) + env = lmdb.open(str(tmp_path)) + with env.begin(write=True) as txn: + txn.put(scorer.MODEL_KEY, b"model") + txn.put( + scorer.CONFIG_KEY, + json.dumps( + { + "length_scale": {"3": [0, 1], "2": [0, 1]}, + "biolink_version": "4.2.5", + "node_context": True, + "edge_context": False, + } + ).encode(), + ) + env.close() + + with caplog.at_level(logging.INFO): + score_worker._validate_scoring_data(logger) + + assert "2 entries" in caplog.text + assert "reranker scores [2, 3]-hop paths, biolink 4.2.5" in caplog.text diff --git a/tests/unit/test_summarize_query_graph.py b/tests/unit/test_summarize_query_graph.py new file mode 100644 index 00000000..bde2584e --- /dev/null +++ b/tests/unit/test_summarize_query_graph.py @@ -0,0 +1,48 @@ +"""Tests for how a callback's query graph is logged.""" + +from shepherd_utils.trapi import MAX_LOGGED_NODE_IDS, summarize_query_graph + + +def _query_graph(n0_ids, n1_ids): + return { + "nodes": { + "n0": {"ids": n0_ids, "categories": ["biolink:Gene"]}, + "n1": {"ids": n1_ids}, + "n2": {"categories": ["biolink:Disease"]}, + }, + "edges": {"e0": {"subject": "n0", "object": "n1"}}, + } + + +def test_a_long_id_list_logs_as_its_count(): + many = [f"X:{i}" for i in range(5000)] + few = ["Y:1", "Y:2"] + query_graph = _query_graph(many, few) + + summary = summarize_query_graph(query_graph) + + assert summary == { + "nodes": { + "n0": {"ids": "5000 curies", "categories": ["biolink:Gene"]}, + "n1": {"ids": ["Y:1", "Y:2"]}, + "n2": {"categories": ["biolink:Disease"]}, + }, + "edges": {"e0": {"subject": "n0", "object": "n1"}}, + } + # the callback goes on to merge the graph it was sent, unchanged + assert query_graph["nodes"]["n0"]["ids"] is many + + +def test_ids_up_to_the_threshold_are_logged(): + ids = [f"X:{i}" for i in range(MAX_LOGGED_NODE_IDS)] + + summary = summarize_query_graph(_query_graph(ids, ids + ["X:extra"])) + + assert summary["nodes"]["n0"]["ids"] == ids + assert summary["nodes"]["n1"]["ids"] == f"{MAX_LOGGED_NODE_IDS + 1} curies" + + +def test_anything_but_a_query_graph_is_logged_as_is(): + assert summarize_query_graph(None) is None + assert summarize_query_graph({"edges": {}}) == {"edges": {}} + assert summarize_query_graph({"nodes": {"n0": None}}) == {"nodes": {"n0": None}} diff --git a/workers/aragorn_pathfinder/requirements.txt b/workers/aragorn_pathfinder/requirements.txt index 63370b69..cc634450 100644 --- a/workers/aragorn_pathfinder/requirements.txt +++ b/workers/aragorn_pathfinder/requirements.txt @@ -1,2 +1,2 @@ httpx -networkx +numpy diff --git a/workers/aragorn_pathfinder/worker.py b/workers/aragorn_pathfinder/worker.py index ecf332e5..4932a590 100644 --- a/workers/aragorn_pathfinder/worker.py +++ b/workers/aragorn_pathfinder/worker.py @@ -3,26 +3,43 @@ import asyncio import json import logging +import os import time import uuid +from collections import Counter, defaultdict +from dataclasses import dataclass +from itertools import pairwise +from math import isfinite +from pathlib import Path +from typing import Optional import httpx +import numpy as np +from opentelemetry.propagate import inject from shepherd_utils.config import settings +from shepherd_utils.cpu import resolve_pool_workers +from shepherd_utils.data_download import ensure_pathfinder_search_graph from shepherd_utils.db import ( abandon_lookup_if_data_lost, add_callback_id, cleanup_callbacks, get_message, + get_message_sync, get_running_callbacks, + remove_callback_id, save_response, + save_response_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 ( + filter_kgraph_orphans, get_tasks, run_task_lifecycle, ) +from shepherd_utils.trapi import make_binding # Queue name STREAM = "aragorn.pathfinder" @@ -33,24 +50,450 @@ tracer = setup_tracer(STREAM) LOGGER = get_worker_logger(STREAM) +# The backbone's pair graph as a CSR (offsets, neighbors, weights) plus its node +# CURIEs in index order, written by squashbert's scripts/export_search_graph.py. +SEARCH_GRAPH_DIR = settings.pathfinder_search_graph_dir +MAX_PATHS = 12000 +SEARCH_BUDGET_SEC = 5.0 +MIDDLE_CHUNK = 4096 +# Each lookup holds a Postgres connection while it registers its callback; half +# the pool leaves the rest for the callback wait and the other queries. +MAX_LOOKUPS_IN_FLIGHT = max(1, settings.postgres_pool_max_size // 2) +SOURCE = "infores:shepherd-aragorn" +ONTOLOGY_PREDICATES = {"biolink:subclass_of", "biolink:superclass_of"} + +# Set up in poll_for_tasks. +NODE_CURIES: list[str] = [] +NODE_INDEX: dict[str, int] = {} +POOL: Optional[ProcessPoolManager] = None +# Per-child CSR, memory-mapped on first use so the pool children share one copy. +_search_graphs: dict = {} + + +@dataclass +class AsyncResponse: + status_code: int + success: bool + callback_id: str + error: Optional[str] = None + # True when the submission was sent but we never heard the ACK. The + # retrieval service may well be running the lookup and will POST the + # callback later, so the callback id must be kept. + in_flight: bool = False + + +def load_search_graph(graph_dir: str): + """The search graph's CSR arrays, memory-mapped read-only.""" + if graph_dir not in _search_graphs: + _search_graphs[graph_dir] = tuple( + np.load(Path(graph_dir) / f"{name}.npy", mmap_mode="r") + for name in ("offsets", "neighbors", "weights") + ) + return _search_graphs[graph_dir] + + +def cheapest_legs(graph, start, far): + """The one-hop leg and the two cheapest two-hop legs, through different + intermediates, from ``start`` to every node. + + A leg never passes through or ends at ``start`` or ``far``. ``via`` is + each two-hop leg's intermediate. + """ + offsets, neighbors, weights = graph + num_nodes = len(offsets) - 1 + + def gather(rows): + lo = np.asarray(offsets[rows], dtype=np.int64) + counts = np.asarray(offsets[rows + 1], dtype=np.int64) - lo + ends = np.cumsum(counts) + slots = np.arange(int(counts.sum())) + np.repeat(lo - (ends - counts), counts) + return slots, counts + + slots, _ = gather(np.array([start])) + heads = np.asarray(neighbors[slots], dtype=np.int64) + cost = np.asarray(weights[slots], dtype=np.float64) + kept = np.isfinite(cost) & (heads != far) & (heads != start) + heads, cost = heads[kept], cost[kept] + one = np.full(num_nodes, np.inf) + np.minimum.at(one, heads, cost) + + first = np.unique(heads) + slots, counts = gather(first) + via_all = np.repeat(first, counts) + heads = np.asarray(neighbors[slots], dtype=np.int64) + cost = np.repeat(one[first], counts) + np.asarray(weights[slots], dtype=np.float64) + kept = (heads != start) & (heads != far) & (heads != via_all) & np.isfinite(cost) + via_all, heads, cost = via_all[kept], heads[kept], cost[kept] + + two = np.full((2, num_nodes), np.inf) + via = np.full((2, num_nodes), -1, dtype=np.int64) + for rank in range(2): + np.minimum.at(two[rank], heads, cost) + at = np.flatnonzero(np.isfinite(cost) & (cost == two[rank][heads])) + if not len(at): + break + _, first_at = np.unique(heads[at], return_index=True) + chosen = at[first_at] + via[rank][heads[chosen]] = via_all[chosen] + cost[chosen] = np.inf + return one, two, via + + +def search_paths(graph_dir, source, target, max_paths): + """Process-pool entrypoint: middle-first paths between two backbone nodes. + + Every node within two hops of both ends is priced at its cheapest loopless + path with that node at the centre, and the middles come out in that order, + each with that one path. A three-hop path is centred on both its + intermediates and comes out once. + """ + deadline = time.monotonic() + SEARCH_BUDGET_SEC + graph = load_search_graph(graph_dir) + one_s, two_s, via_s = cheapest_legs(graph, source, target) + one_t, two_t, via_t = cheapest_legs(graph, target, source) + # The cheapest leg out and back can share their intermediate; the runner-up + # on either side prices the middle loopless. + both_two = np.where( + via_s[0] != via_t[0], + two_s[0] + two_t[0], + np.minimum(two_s[0] + two_t[1], two_s[1] + two_t[0]), + ) + score = np.minimum.reduce( + [one_s + one_t, one_s + two_t[0], two_s[0] + one_t, both_two] + ) + middles = np.flatnonzero(np.isfinite(score)) + order = middles[np.argsort(score[middles], kind="stable")] + + paths = [] + seen = set() + for chunk in range(0, len(order), MIDDLE_CHUNK): + if time.monotonic() > deadline: + break + ms = order[chunk : chunk + MIDDLE_CHUNK] + no_via = np.full(len(ms), -1) + f_cost = np.column_stack([one_s[ms], two_s[0][ms], two_s[1][ms]]) + f_via = np.column_stack([no_via, via_s[0][ms], via_s[1][ms]]) + b_cost = np.column_stack([one_t[ms], two_t[0][ms], two_t[1][ms]]) + b_via = np.column_stack([no_via, via_t[0][ms], via_t[1][ms]]) + pairs = f_cost[:, :, None] + b_cost[:, None, :] + pairs[(f_via[:, :, None] == b_via[:, None, :]) & (f_via[:, :, None] >= 0)] = ( + np.inf + ) + pairs = pairs.reshape(len(ms), 9) + picks = pairs.argmin(axis=1) + costs = pairs[np.arange(len(ms)), picks] + for i, m in enumerate(ms.tolist()): + if not isfinite(costs[i]): + continue + a = int(f_via[i, picks[i] // 3]) + b = int(b_via[i, picks[i] % 3]) + nodes = ( + source, + *((a,) if a >= 0 else ()), + m, + *((b,) if b >= 0 else ()), + target, + ) + if len(nodes) == 4: + if nodes in seen: + continue + seen.add(nodes) + paths.append(nodes) + if len(paths) >= max_paths: + return paths, len(order) + return paths, len(order) + + +def one_hop(left, right, parameters, subclass): + """A one-hop lookup between two query nodes, each pinned to ids or open.""" + return { + "message": { + "query_graph": { + "nodes": { + "n0": {} if left is None else {"ids": list(left)}, + "n1": {} if right is None else {"ids": list(right)}, + }, + "edges": { + "e0": { + "subject": "n0", + "object": "n1", + "predicates": ["biolink:related_to"], + } + }, + } + }, + "parameters": {**parameters, "subclass": subclass}, + } + + +def edge_lookups(node_paths, source, target, parameters): + """One-hop lookups for every hop the paths take, and for the direct edge. + + Gandalf expands the endpoints to their subclass children: the direct edge, + and each endpoint against anything, since pinning the other side would + expand that too. The hops between intermediates are exact, pinned in one + lookup with every near node on one side and every far node on the other: + one lookup per hop would be thousands of lookups, each with its callback + to register and merge. It also returns the edges between the other pairs, + which no path uses and the orphan filter drops. + """ + lookups = [one_hop([source], [target], parameters, True)] + if node_paths: + lookups.append(one_hop([source], None, parameters, True)) + lookups.append(one_hop(None, [target], parameters, True)) + hops = set() + for path in node_paths: + if len(path) > 3: + hops.add((path[1], path[2])) + if len(path) > 4: + hops.add((path[-2], path[-3])) + if hops: + near = sorted({u for u, _ in hops}) + far = sorted({v for _, v in hops}) + lookups.append(one_hop(near, far, parameters, False)) + return lookups + + +async def run_async_lookup( + client: httpx.AsyncClient, + message: dict, + query_id: str, + logger: logging.Logger, +) -> AsyncResponse: + """Return an async lookup response with callback id.""" + callback_id = str(uuid.uuid4())[:8] + with tracer.start_as_current_span("aragorn.pathfinder.lookup") as span: + span.set_attribute("callback.id", callback_id) + lookup_carrier = {} + inject(lookup_carrier) + await add_callback_id(query_id, callback_id, json.dumps(lookup_carrier), logger) + + message["callback"] = f"{settings.callback_host}/aragorn/callback/{callback_id}" + logger.debug( + f"[{callback_id}] Sending edge lookup to {settings.kg_retrieval_url}" + ) + try: + response = await client.post(settings.kg_retrieval_url, json=message) + return AsyncResponse( + status_code=response.status_code, + success=response.status_code == 200, + callback_id=callback_id, + ) + except httpx.ReadTimeout as e: + span.record_exception(e) + return AsyncResponse( + status_code=500, + success=False, + callback_id=callback_id, + error=f"{type(e).__name__}: no ACK within {client.timeout.read}s", + in_flight=True, + ) + except Exception as e: + span.record_exception(e) + return AsyncResponse( + status_code=500, + success=False, + callback_id=callback_id, + error=f"{type(e).__name__}: {e}", + ) + + +async def send_lookups(lookups, query_id, logger: logging.Logger): + """Send the lookups, ``MAX_LOOKUPS_IN_FLIGHT`` at a time, dropping the + callbacks of any that failed.""" + in_flight = asyncio.Semaphore(MAX_LOOKUPS_IN_FLIGHT) + + async def send(client, lookup): + async with in_flight: + return await run_async_lookup(client, lookup, query_id, logger) + + async with httpx.AsyncClient( + timeout=settings.kg_retrieval_submit_timeout + ) as client: + responses = await asyncio.gather( + *(send(client, lookup) for lookup in lookups), + return_exceptions=True, + ) + for response in responses: + if isinstance(response, Exception): + logger.error( + f"Failed to do lookup and unable to remove callback id: {response}" + ) + elif not response.success: + if response.in_flight: + logger.warning( + f"[{response.callback_id}] Lookup submitted but not " + f"acknowledged, waiting for its callback anyway: {response.error}" + ) + continue + logger.error( + f"[{response.callback_id}] Failed to do lookup, removing callback " + f"id: {response.error}" + ) + await remove_callback_id(response.callback_id, logger) + + +async def wait_for_callbacks( + query_id, response_id, max_query_time, logger: logging.Logger +): + """Wait until every lookup's callback is merged, or the query times out.""" + start_time = time.time() + running_callback_ids = [""] + try: + while time.time() - start_time < max_query_time: + try: + running_callback_ids = await get_running_callbacks(query_id, logger) + except Exception: + # Brief backoff then retry the check rather than giving up + await asyncio.sleep(5) + continue + # fail now if the query data the callbacks merge into is gone + await abandon_lookup_if_data_lost( + query_id, response_id, running_callback_ids, logger + ) + if len(running_callback_ids) > 0: + await asyncio.sleep(1) + continue + logger.debug("Got all lookups back. Continuing...") + break + + if time.time() - start_time > max_query_time: + logger.warning( + f"Timed out getting lookup callbacks. {len(running_callback_ids)} queries were still running..." + ) + logger.warning(f"Running callbacks: {running_callback_ids}") + await cleanup_callbacks(query_id, logger) + except asyncio.CancelledError: + logger.warning("Cancelled while waiting for callbacks") + + +def build_paths(message, plan, logger: logging.Logger): + """Turn the search's node paths into TRAPI paths over the looked-up edges. + + A hop is every edge between its two nodes. At an endpoint that includes + gandalf's inferred edges from the endpoint's subclass children, each + supported by the child's edge and its subclass_of edge. A path with a hop + that came back empty is dropped. + """ + kgraph = message.get("knowledge_graph") or {"nodes": {}, "edges": {}} + nodes, edges = kgraph["nodes"], kgraph["edges"] + aux_graphs = message.get("auxiliary_graphs") or {} + source, target = plan["source"], plan["target"] + + by_pair = defaultdict(list) + for edge_id, edge in edges.items(): + if edge["predicate"] not in ONTOLOGY_PREDICATES: + by_pair[frozenset((edge["subject"], edge["object"]))].append(edge_id) + + def path_edges(path): + """The path's edges, or None and the first hop that has none.""" + found = [] + for hop, (u, v) in enumerate(pairwise(path)): + hop_edges = by_pair.get(frozenset((u, v))) + if not hop_edges: + return None, hop + found.extend(hop_edges) + return list(dict.fromkeys(found)), None + + category = plan["intermediate_category"] + analyses = [] + seen = set() + dropped = Counter() + first_missing = None + # The direct edge is the one-hop path, first. + for index, path in enumerate([[source, target], *plan["node_paths"]]): + intermediates = tuple(path[1:-1]) + if intermediates in seen: + continue + seen.add(intermediates) + if category is not None and category != "biolink:NamedThing": + if not any( + category in (nodes.get(node, {}).get("categories") or []) + for node in intermediates + ): + continue + edge_ids, missing = path_edges(path) + if edge_ids is None: + if len(path) == 2: + dropped["direct"] += 1 + elif missing == 0: + dropped["at the source"] += 1 + elif missing == len(path) - 2: + dropped["at the target"] += 1 + else: + dropped["between intermediates"] += 1 + first_missing = first_missing or (path[missing], path[missing + 1]) + continue + aux_id = f"p_{index}" + aux_graphs[aux_id] = {"edges": edge_ids} + analyses.append( + { + "resource_id": SOURCE, + "path_bindings": {plan["path_key"]: make_binding([aux_id])}, + } + ) + logger.info( + f"Built {len(analyses)} paths from {len(edges)} edges; dropped " + f"{dropped.total()} with a hop that has no edges {dict(dropped)}." + ) + if first_missing: + # Whether the hop's nodes are in the graph at all tells ids that do not + # match it apart from a pair it holds no edge between. + linked = set().union(*by_pair) + u, v = first_missing + logger.info( + f"First hop with no edges: {u} ({'in' if u in linked else 'not in'} " + f"the edges) - {v} ({'in' if v in linked else 'not in'} the edges)." + ) + + message["knowledge_graph"] = kgraph + message["results"] = [] + if analyses: + message["auxiliary_graphs"] = aux_graphs + message["results"] = [ + { + "node_bindings": { + plan["subject_key"]: make_binding([source]), + plan["object_key"]: make_binding([target]), + }, + "analyses": analyses, + } + ] + + +def build_paths_task(response_id: str, plan: dict, log_level: int) -> list[dict]: + """Process-pool entrypoint: build the paths into the merged response and + save it, returning this child's log records for the parent.""" + 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: + response = get_message_sync(response_id) + build_paths(response["message"], plan, logger) + filter_kgraph_orphans(response, logger) + save_response_sync(response_id, response) + return query_log_handler.drain() + finally: + logger.removeHandler(query_log_handler) + async def shadowfax(task, logger: logging.Logger) -> str: - """Processes pathfinder queries. This is done by using literature - co-occurrence to find nodes that occur in publications with our input - nodes, then finding paths that connect our input nodes through these - intermediate nodes.""" + """Processes pathfinder queries. Middle-first search over the backbone + finds the paths as nodes, retriever returns the edges for every hop they + take, and once all of those lookups are merged the paths are built from + them.""" # given a task, get the message from the db query_id = task[1]["query_id"] response_id = task[1]["response_id"] - otel = task[1]["otel"] metadata = json.loads(task[1]["metadata"]) message = await get_message(query_id, logger) - response = await get_message(response_id, logger) parameters = message.get("parameters") or {} parameters["timeout"] = parameters.get("timeout", settings.lookup_timeout) if metadata.get("rehydrate"): # do rehydration + response = await get_message(response_id, logger) parameters["rehydrate"] = parameters.get("rehydrate", True) response["parameters"] = parameters logger.debug( @@ -69,231 +512,85 @@ async def shadowfax(task, logger: logging.Logger) -> str: await save_response(response_id, response_json, logger) return json.dumps({}) - filter_config = parameters.get("filter_config", {}) - parameters["filter_config"] = { - "min_information_content": filter_config.get("min_information_content", 60), - "max_node_degree": filter_config.get("max_node_degree", 10000), - } - parameters["dehydrated"] = parameters.get("dehydrated", True) - message["parameters"] = parameters - qgraph = message["message"]["query_graph"] - pinned_node_keys = [] - pinned_node_ids = [] - retriever_query = {"message": message["message"], "parameters": parameters} - for node_key, node in qgraph["nodes"].items(): - pinned_node_keys.append(node_key) - if node.get("ids", None) is not None: - # TODO: silently only grabbing the first id - pinned_node_ids.append(node["ids"][0]) + pinned_node_ids = [ + node["ids"][0] for node in qgraph["nodes"].values() if node.get("ids") + ] if len(set(pinned_node_ids)) != 2: raise Exception("Pathfinder queries require two pinned nodes.") + path_key, qpath = next(iter(qgraph["paths"].items())) + constraints = qpath.get("constraints") or [] + if len(constraints) > 1: + raise Exception("Pathfinder queries do not support multiple constraints.") + intermediate_categories = ( + (constraints[0].get("required_intermediate_categories") or []) + if constraints + else [] + ) + if len(intermediate_categories) > 1: + raise Exception( + "Pathfinder queries do not support multiple intermediate categories" + ) + source = qgraph["nodes"][qpath["subject"]]["ids"][0] + target = qgraph["nodes"][qpath["object"]]["ids"][0] - path_key = next(iter(qgraph["paths"].keys())) - qpath = qgraph["paths"][path_key] - if qpath.get("constraints", None) is not None: - constraints = qpath["constraints"] - # TODO: need to wrap up tasks - if len(constraints) > 1: - raise Exception("Pathfinder queries do not support multiple constraints.") - if len(constraints) > 0: - # TRAPI 2.0 renamed PathConstraint.intermediate_categories. - intermediate_categories = ( - constraints[0].get("required_intermediate_categories", None) or [] - ) - if len(intermediate_categories) > 1: - raise Exception( - "Pathfinder queries do not support multiple intermediate categories" - ) - - # Create 3-hop query - - retriever_query["message"]["query_graph"] = { - "nodes": { - pinned_node_keys[0]: {"ids": [pinned_node_ids[0]]}, - "intermediate_0": { - "categories": ["biolink:NamedThing"], - }, - "intermediate_1": { - "categories": ["biolink:NamedThing"], - }, - pinned_node_keys[1]: {"ids": [pinned_node_ids[1]]}, - }, - "edges": { - "e0": { - "subject": pinned_node_keys[0], - "object": "intermediate_0", - "predicates": [ - "biolink:physically_interacts_with", - "biolink:genetically_interacts_with", - "biolink:contributes_to", - "biolink:contribution_from", - "biolink:affects", - "biolink:affected_by", - "biolink:acts_upstream_of", - "biolink:has_upstream_actor", - "biolink:enables", - "biolink:enabled_by", - "biolink:produces", - "biolink:produced_by", - "biolink:has_participant", - "biolink:participates_in", - "biolink:derives_from", - "biolink:derives_into", - "biolink:transcribed_to", - "biolink:transcribed_from", - "biolink:translates_to", - "biolink:translation_of", - "biolink:has_gene_product", - "biolink:gene_product_of", - "biolink:genetically_associated_with", - "biolink:located_in", - "biolink:location_of", - "biolink:contains_process", - "biolink:occurs_in", - "biolink:affects_likelihood_of", - "biolink:likelihood_affected_by", - "biolink:active_in", - "biolink:has_active_component", - ], - }, - "e1": { - "subject": "intermediate_0", - "object": "intermediate_1", - "predicates": [ - "biolink:physically_interacts_with", - "biolink:genetically_interacts_with", - "biolink:contributes_to", - "biolink:contribution_from", - "biolink:affects", - "biolink:affected_by", - "biolink:acts_upstream_of", - "biolink:has_upstream_actor", - "biolink:enables", - "biolink:enabled_by", - "biolink:produces", - "biolink:produced_by", - "biolink:has_participant", - "biolink:participates_in", - "biolink:derives_from", - "biolink:derives_into", - "biolink:transcribed_to", - "biolink:transcribed_from", - "biolink:translates_to", - "biolink:translation_of", - "biolink:has_gene_product", - "biolink:gene_product_of", - "biolink:genetically_associated_with", - "biolink:located_in", - "biolink:location_of", - "biolink:contains_process", - "biolink:occurs_in", - "biolink:affects_likelihood_of", - "biolink:likelihood_affected_by", - "biolink:active_in", - "biolink:has_active_component", - ], - }, - "e2": { - "subject": "intermediate_1", - "object": pinned_node_keys[1], - "predicates": [ - "biolink:physically_interacts_with", - "biolink:genetically_interacts_with", - "biolink:contributes_to", - "biolink:contribution_from", - "biolink:affects", - "biolink:affected_by", - "biolink:acts_upstream_of", - "biolink:has_upstream_actor", - "biolink:enables", - "biolink:enabled_by", - "biolink:produces", - "biolink:produced_by", - "biolink:has_participant", - "biolink:participates_in", - "biolink:derives_from", - "biolink:derives_into", - "biolink:transcribed_to", - "biolink:transcribed_from", - "biolink:translates_to", - "biolink:translation_of", - "biolink:has_gene_product", - "biolink:gene_product_of", - "biolink:genetically_associated_with", - "biolink:located_in", - "biolink:location_of", - "biolink:contains_process", - "biolink:occurs_in", - "biolink:affects_likelihood_of", - "biolink:likelihood_affected_by", - "biolink:active_in", - "biolink:has_active_component", - ], - }, - }, + lookup_parameters = { + "timeout": parameters["timeout"], + "tiers": [settings.default_data_tier], + # tier 1 would expand every pinned node and cannot dehydrate + "tier_fallback": False, + "dehydrated": True, } - callback_id = str(uuid.uuid4())[:8] - # Put callback UID and query ID in postgres - await add_callback_id(query_id, callback_id, otel, logger) - - retriever_query["callback"] = ( - f"{settings.callback_host}/aragorn/callback/{callback_id}" - ) - logger.info( - f"[{callback_id}] Sending pathfinder query to {settings.kg_retrieval_url}" - ) - with tracer.start_as_current_span("aragorn.pathfinder") as span: - span.set_attribute("callback.id", callback_id) - async with httpx.AsyncClient(timeout=100) as client: - retriever_async_response = await client.post( - settings.kg_retrieval_url, - json=retriever_query, - ) - try: - retriever_async_response.raise_for_status() - except Exception as e: - logger.error(f"Error contacting retriever: {e}") - logger.debug(f"Error details: {retriever_async_response.json()}") - raise e - - # this worker might have a timeout set for if the lookups don't finish within a certain - # amount of time - MAX_QUERY_TIME = message["parameters"]["timeout"] - start_time = time.time() - running_callback_ids = [""] - try: - while time.time() - start_time < MAX_QUERY_TIME: - try: - # see if there are existing lookups going - running_callback_ids = await get_running_callbacks(query_id, logger) - except Exception: - # Brief backoff then retry the check rather than giving up - await asyncio.sleep(5) - continue - # fail now if the query data the callbacks merge into is gone - await abandon_lookup_if_data_lost( - query_id, response_id, running_callback_ids, logger + node_paths = [] + if source in NODE_INDEX and target in NODE_INDEX: + with tracer.start_as_current_span("aragorn.pathfinder.search") as span: + paths, middles = await POOL.run( + asyncio.get_running_loop(), + search_paths, + SEARCH_GRAPH_DIR, + NODE_INDEX[source], + NODE_INDEX[target], + MAX_PATHS, ) - # logger.info(f"Got back {len(running_callback_ids)} running lookups") - # if there are, continue to wait - if len(running_callback_ids) > 0: - await asyncio.sleep(1) - continue - # if there aren't, lookup is complete and we need to pass on to next workflow operation - if len(running_callback_ids) == 0: - logger.debug("Got all lookups back. Continuing...") - break + span.set_attribute("paths", len(paths)) + node_paths = [[NODE_CURIES[node] for node in path] for path in paths] + logger.info( + f"Backbone search found {len(node_paths)} paths over {middles} middles." + ) + else: + logger.warning( + f"{source} or {target} is not on the backbone; " + "looking up the direct edge only." + ) - if time.time() - start_time > MAX_QUERY_TIME: - logger.warning( - f"Timed out getting lookup callbacks. {len(running_callback_ids)} queries were still running..." - ) - logger.warning(f"Running callbacks: {running_callback_ids}") - await cleanup_callbacks(query_id, logger) - except asyncio.CancelledError: - logger.warning(f"Task {task[0]}: Cancelled while waiting for callbacks") + lookups = edge_lookups(node_paths, source, target, lookup_parameters) + logger.info(f"Sending {len(lookups)} edge lookups to {settings.kg_retrieval_url}") + await send_lookups(lookups, query_id, logger) + await wait_for_callbacks(query_id, response_id, parameters["timeout"], logger) + + plan = { + "source": source, + "target": target, + "node_paths": node_paths, + "path_key": path_key, + "subject_key": qpath["subject"], + "object_key": qpath["object"], + "intermediate_category": ( + intermediate_categories[0] if intermediate_categories else None + ), + } + with tracer.start_as_current_span("aragorn.pathfinder.build_paths"): + entries = await POOL.run( + asyncio.get_running_loop(), + build_paths_task, + response_id, + plan, + logger.getEffectiveLevel(), + ) + handler = get_query_handler(logger) + if handler is not None and entries: + handler.ingest(entries) return json.dumps({"rehydrate": True}) @@ -310,6 +607,28 @@ async def _run(task, logger): async def poll_for_tasks(): """On initialization, poll indefinitely for available tasks.""" + global NODE_CURIES, NODE_INDEX, POOL + # In the parent, before any pool child maps the files. + ensure_pathfinder_search_graph(LOGGER) + NODE_CURIES = (Path(SEARCH_GRAPH_DIR) / "nodes.txt").read_text().splitlines() + NODE_INDEX = {curie: index for index, curie in enumerate(NODE_CURIES)} + offsets, _, _ = load_search_graph(SEARCH_GRAPH_DIR) + if len(offsets) != len(NODE_CURIES) + 1: + raise RuntimeError( + f"{SEARCH_GRAPH_DIR} holds {len(offsets) - 1} nodes in its CSR but " + f"{len(NODE_CURIES)} in nodes.txt" + ) + LOGGER.info(f"{STREAM}: search graph over {len(NODE_CURIES)} nodes.") + # The search and the path building are CPU-bound; the waits between them + # are not, so more tasks run than there are pool workers. + 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="aragorn.pathfinder process pool", + task_timeout=settings.pool_task_timeout_sec, + ) while True: try: async for task, parent_ctx, logger, limiter in get_tasks( diff --git a/workers/merge_message/worker.py b/workers/merge_message/worker.py index 0f7015a9..c912cb5c 100644 --- a/workers/merge_message/worker.py +++ b/workers/merge_message/worker.py @@ -11,6 +11,9 @@ from itertools import combinations from typing import Any, Optional, Union +import orjson +import zstandard + from shepherd_utils.broker import ( add_task, mark_task_as_complete, @@ -56,6 +59,7 @@ node_binding_ids, qedge_qualifier_sets, query_log_level, + summarize_query_graph, ) # Queue name @@ -560,107 +564,13 @@ def merge_messages( return merged_messages elif "paths" in original_query_graph: - # Pathfinder query - path_id, og_path = next(iter(original_query_graph["paths"].items())) - subject_node_id = og_path.get("subject") - object_node_id = og_path.get("object") - if subject_node_id is None or object_node_id is None: + # Pathfinder query: each callback is one of aragorn.pathfinder's edge + # lookups, so only the knowledge graph accumulates here. The paths are + # built from it once every lookup is in. + _, og_path = next(iter(original_query_graph["paths"].items())) + if og_path.get("subject") is None or og_path.get("object") is None: raise KeyError("Missing either subject or object from path.") - - intermediate_category = None - constraints = og_path.get("constraints") or [] - if len(constraints) > 0: - intermediate_categories = ( - constraints[0].get("required_intermediate_categories") or [] - ) - if len(intermediate_categories) > 0: - intermediate_category = intermediate_categories[0] - - kg_nodes = result["message"]["knowledge_graph"].get("nodes", {}) - kg_edges = result["message"]["knowledge_graph"].get("edges", {}) - - aux_counter = 0 - score = 0 - analyses = [] - for new_result in new_response["message"]["results"]: - path_edge_ids = set() - for analysis in new_result.get("analyses") or []: - path_edge_ids.update(edge_binding_ids(analysis)) - score = new_result.get("score") - if not path_edge_ids: - continue - - if ( - intermediate_category is not None - and intermediate_category != "biolink:NamedThing" - ): - pinned_ids = set() - for pinned in (subject_node_id, object_node_id): - pinned_ids.update(node_binding_ids(new_result, pinned)) - intermediate_node_ids = set() - for edge_id in path_edge_ids: - edge = kg_edges.get(edge_id) - if edge is None: - continue - for node_id in (edge.get("subject"), edge.get("object")): - if node_id and node_id not in pinned_ids: - intermediate_node_ids.add(node_id) - if not any( - intermediate_category - in (kg_nodes.get(nid, {}).get("categories") or []) - for nid in intermediate_node_ids - ): - continue - - aux_id = f"a_{aux_counter}" - aux_counter += 1 - - # Add new aux graph to message - result["message"].setdefault("auxiliary_graphs", {})[aux_id] = { - "edges": list(path_edge_ids), - } - - analysis = { - "resource_id": source, - "path_bindings": {path_id: make_binding([aux_id])}, - } - if score is not None: - analysis["score"] = score - analyses.append(analysis) - - # --- Resolve the pinned node IDs from any old result's node_bindings --- - # (They should all bind the same start/end IDs since they're pinned) - if not analyses: - result["message"]["results"] = [] - return result - - start_kg_id = None - end_kg_id = None - for new_result in new_response["message"]["results"]: - subject_ids = node_binding_ids(new_result, subject_node_id) - if subject_ids: - start_kg_id = subject_ids[0] - object_ids = node_binding_ids(new_result, object_node_id) - if object_ids: - end_kg_id = object_ids[0] - if start_kg_id and end_kg_id: - break - - if not start_kg_id or not end_kg_id: - result["message"]["results"] = [] - return result - - # --- Assemble the single Pathfinder result --- - pathfinder_result = { - "node_bindings": { - subject_node_id: make_binding([start_kg_id]), - object_node_id: make_binding([end_kg_id]), - }, - "analyses": analyses, - } - - result["message"]["results"] = [pathfinder_result] - + result["message"]["results"] = response["message"].get("results") or [] return result raise TypeError("Unsupported query type.") @@ -738,6 +648,48 @@ def take_callback_logs( # --------------------------------------------------------------------------- +def read_callback_response( + callback_id: str, logger: logging.Logger +) -> Optional[dict[str, Any]]: + """A callback's stored payload as a TRAPI response, or None when it is + gone or can't be read. + + The callback route stores what was posted without parsing it (a large one + would stall the server), so this is where a body that isn't a TRAPI + response is found, and where a response missing its results or knowledge + graph is given empty ones. + """ + try: + response = get_message_sync(callback_id) + except KeyError: + logger.error(f"Missing callback {callback_id} while folding; skipping.") + return None + except (orjson.JSONDecodeError, zstandard.ZstdError) as e: + logger.error(f"Callback {callback_id} is not valid JSON; skipping. {e}") + return None + if not isinstance(response, dict): + logger.error(f"Callback {callback_id} is not a JSON object; skipping.") + return None + if not isinstance(response.get("message"), dict): + response["message"] = {} + message = response["message"] + if message.get("results") is None: + message["results"] = [] + if message.get("knowledge_graph") is None: + message["knowledge_graph"] = {"nodes": {}, "edges": {}} + kgraph = message["knowledge_graph"] + logger.info( + f"[{callback_id}] Got back {len(message['results'])} results, " + f"{len(kgraph.get('nodes') or {})} nodes and " + f"{len(kgraph.get('edges') or {})} edges." + ) + logger.debug( + f"[{callback_id}] for query graph: " + f"{summarize_query_graph(message.get('query_graph'))}" + ) + return response + + def merge_messages_by_ids( target: str, query_id: str, @@ -790,12 +742,8 @@ def merge_messages_by_ids( merged: list[str] = [] callback_log_entries: list[dict] = [] for callback_id in callback_ids: - try: - callback_response = get_message_sync(callback_id) - except KeyError: - worker_logger.error( - f"Missing callback {callback_id} while folding; skipping." - ) + callback_response = read_callback_response(callback_id, worker_logger) + if callback_response is None: continue callback_log_entries.extend( take_callback_logs( diff --git a/workers/score_paths/model_weights/squashbert_direct_3hop.pt b/workers/score_paths/model_weights/squashbert_direct_3hop.pt deleted file mode 100644 index 2774d4dc..00000000 Binary files a/workers/score_paths/model_weights/squashbert_direct_3hop.pt and /dev/null differ diff --git a/workers/score_paths/requirements.txt b/workers/score_paths/requirements.txt index f6416df6..751ff013 100644 --- a/workers/score_paths/requirements.txt +++ b/workers/score_paths/requirements.txt @@ -1,5 +1,4 @@ -bmt>=1.4.6 +bmt==1.5.0 lmdb numpy -torch>=2.10.0 -scikit-learn +torch>=2.11.0 diff --git a/workers/score_paths/scorer.py b/workers/score_paths/scorer.py new file mode 100644 index 00000000..770cdf5a --- /dev/null +++ b/workers/score_paths/scorer.py @@ -0,0 +1,489 @@ +"""Score a TRAPI 2.0 message's paths with the sequence reranker. + +Everything the reranker reads is in one LMDB, written by squashbert's +scripts/export_reranker.py under the keys below: the TorchScript model, its +config, and the vectors and context tables for the graph's nodes, categories +and hop phrases. A path is read the way squashbert's training paths were +written, so it looks up the vectors it was trained on; squashbert's tests +check that reading and the scores, which is why this module imports nothing +from shepherd_utils. +""" + +import functools +import hashlib +import io +import json +import re +import time + +import numpy as np +import torch +from bmt import Toolkit + +NODE, CATEGORY, PHRASE = "node", "category", "phrase" +NODE_CONTEXT, EDGE_CONTEXT = "node_context", "edge_context" +# Nodes and categories keep the SAPBERT caches' float16 +DTYPES = { + NODE: np.float16, + CATEGORY: np.float16, + PHRASE: np.float32, + NODE_CONTEXT: np.float32, + EDGE_CONTEXT: np.float32, +} +MODEL_KEY = b"reranker:model.pt" +CONFIG_KEY = b"reranker:config.json" +# LMDB refuses keys over 511 bytes +MAX_KEY_BYTES = 480 + +SCHEMA_URL = ( + "https://raw.githubusercontent.com/biolink/biolink-model/" + "refs/tags/v{version}/biolink-model.yaml" +) +PREDICATE_MAP_URL = ( + "https://raw.githubusercontent.com/biolink/biolink-model/" + "refs/tags/v{version}/predicate_mapping.yaml" +) + +ONTOLOGY_PREDICATES = {"biolink:subclass_of", "biolink:superclass_of"} +# The qualifiers that decide how an edge reads +READING_SLOTS = ( + "qualified_predicate", + "object_aspect_qualifier", + "object_direction_qualifier", + "subject_aspect_qualifier", + "subject_direction_qualifier", +) +# A qualifier value a hop label can hold +_VALUE = re.compile(r"[\w.:-]+( [\w.:-]+)*") + +N_HOPS_SCALE = 8.0 +# The node table's columns: out-degree, in-degree, IC, IC present +NODE_CONTEXT_DIM = 4 +_NODE_COLUMNS = ( + ("deg_out_max", 0, np.max), + ("deg_out_mean", 0, np.mean), + ("deg_in_max", 1, np.max), + ("deg_in_mean", 1, np.mean), + ("ic_mean", 2, np.mean), + ("ic_present_frac", 3, np.mean), +) +_EDGE_BASES = ( + "edge_count", + "pubs", + "primary_sources", + "distinct_sources", + "attributes_seen", +) + + +def lmdb_key(kind: str, text: str) -> bytes: + key = f"{kind}:{text}".encode() + if len(key) > MAX_KEY_BYTES: + key = f"{kind}#".encode() + hashlib.sha256(key).hexdigest().encode() + return key + + +def edge_context_key(a: str, b: str) -> bytes: + """Keyed by the pair's 64-bit blake2b digest, as squashbert's PairTable + stores every pair of the graph.""" + pair = "|".join(sorted((a, b))) + digest = hashlib.blake2b(pair.encode(), digest_size=8).digest() + return f"{EDGE_CONTEXT}#".encode() + digest + + +def load_reranker(txn): + """`(TorchScript model, config)` from the LMDB.""" + blob, config = txn.get(MODEL_KEY), txn.get(CONFIG_KEY) + if blob is None or config is None: + raise RuntimeError( + "the embeddings LMDB holds no reranker; write one with squashbert's " + "scripts/export_reranker.py" + ) + model = torch.jit.load(io.BytesIO(blob), map_location="cpu") + model.eval() + config = json.loads(config) + config["length_scale"] = { + int(length): (float(mean), float(spread)) + for length, (mean, spread) in config["length_scale"].items() + } + return model, config + + +def load_toolkit(config) -> Toolkit: + """The biolink release the training paths were read by.""" + version = config["biolink_version"] + return Toolkit( + schema=SCHEMA_URL.format(version=version), + predicate_map=PREDICATE_MAP_URL.format(version=version), + ) + + +def predicate_name(predicate: str) -> str: + return predicate.removeprefix("biolink:").replace("_", " ") + + +def predicate_curie(name: str) -> str: + return "biolink:" + name.replace(" ", "_") + + +# A bmt lookup takes about a millisecond and a message asks the same few +# thousands of times, so answers are kept for the life of the process. +@functools.lru_cache(maxsize=None) +def inverse_predicate(bmt, predicate: str) -> str | None: + # squashbert's stand-in for an inverse the model lacks + if predicate in ("biolink:has_chemical_role", "has_chemical_role"): + return "biolink:is_chemical_role_of" + if predicate in ("biolink:is_chemical_role_of", "is_chemical_role_of"): + return "biolink:has_chemical_role" + name = predicate_name(predicate) + if not bmt.is_predicate(name): + return None + if bmt.is_symmetric(name): + return predicate_curie(name) + try: + inverse = bmt.get_inverse(name) + except AttributeError: + # bmt raises on a declared inverse the schema does not hold + return None + return None if inverse is None else predicate_curie(inverse) + + +@functools.lru_cache(maxsize=None) +def most_specific_category(bmt, categories: tuple[str, ...]) -> str: + """The deepest biolink class, else the first mixin; not the first listed, + which the translator KGX's alphabetical lists make BiologicalEntity.""" + best = None + for category in categories: + depth = bmt.get_element_depth(category) if bmt.is_category(category) else None + if depth is not None and (best is None or depth > best[0]): + best = (depth, category) + if best is not None: + chosen = best[1] + else: + mixins = [category for category in categories if bmt.is_mixin(category)] + if not mixins: + raise ValueError(f"No biolink categories found in {categories!r}") + chosen = mixins[0] + if chosen.startswith("biolink:"): + return chosen + return "biolink:" + chosen.replace(" ", "") + + +def _worded(verb: str, aspect, direction, walked_back: bool) -> str: + word = predicate_name(verb) + if aspect is None: + return word + change = aspect.replace("_", " ") + if direction is not None: + change = f"{direction.replace('_', ' ')} {change}" + return f"has {change} {word}" if walked_back else f"{word} {change} of" + + +def edge_phrase(bmt, edge: dict, backward: bool) -> str | None: + """The phrase a TRAPI edge reads as, walked as stored or backwards; None + where it cannot be walked that way.""" + predicate = edge["predicate"] + found = { + item["qualifier_type_id"].removeprefix("biolink:"): item["qualifier_value"] + for item in edge.get("qualifiers") or [] + } + slots = { + slot: str(found[slot]) for slot in READING_SLOTS if found.get(slot) is not None + } + verb, aspect, direction = predicate, None, None + # Subject-side qualifiers, and a direction without an aspect, have no + # wording: such an edge reads as its bare predicate + unworded = ( + slots.get("subject_aspect_qualifier") is not None + or slots.get("subject_direction_qualifier") is not None + or ( + slots.get("object_direction_qualifier") is not None + and slots.get("object_aspect_qualifier") is None + ) + ) + if not unworded: + aspect = slots.get("object_aspect_qualifier") + direction = slots.get("object_direction_qualifier") + qualified = slots.get("qualified_predicate") + if qualified is not None and (direction is not None or aspect is None): + verb = predicate_curie(predicate_name(qualified)) + for value in (aspect, direction): + if value is not None and not _VALUE.fullmatch(value): + raise ValueError(f"qualifier value {value!r} cannot be read") + if not backward: + return _worded(verb, aspect, direction, False) + # Only the predicate decides whether an edge can be walked backwards; a + # qualified predicate with no inverse reads as the bare inverse + plain = inverse_predicate(bmt, predicate) + if plain is None: + return None + inverse = inverse_predicate(bmt, verb) + if inverse is None: + return _worded(plain, None, None, False) + return _worded(inverse, aspect, direction, True) + + +def hop_phrase(bmt, current: str, following: str, edge_ids, edges: dict) -> str | None: + """Every edge of the path between two nodes, read from `current`; None + where none can be walked across.""" + carried = {} + for edge_id in edge_ids: + edge = edges.get(edge_id) + if not edge or edge["predicate"] in ONTOLOGY_PREDICATES: + continue + if edge["subject"] == current and edge["object"] == following: + backward = False + elif edge["object"] == current and edge["subject"] == following: + backward = True + else: + continue + phrase = edge_phrase(bmt, edge, backward) + if phrase is None: + continue + negated = False + for attribute in edge.get("attributes") or (): + if ( + attribute.get("original_attribute_name") == "negated" + or attribute.get("attribute_type_id") == "biolink:negated" + ): + value = attribute.get("value") + if value is not None and not isinstance(value, bool): + raise ValueError(f"negated is {value!r}; expected a boolean") + negated = bool(value) + break + # A label only NOT annotations carry is dropped + carried[phrase] = carried.get(phrase, False) or not negated + phrases = sorted(phrase for phrase, kept in carried.items() if kept) + if not phrases: + return None + if len(phrases) == 1: + return phrases[0] + return "[" + " or ".join(phrases) + "]" + + +def order_path_nodes(source: str, target: str, edge_ids, edges: dict) -> list[str]: + ordered = [source] + while target not in ordered: + progress = False + for edge_id in edge_ids: + edge = edges.get(edge_id) + if not edge: + continue + tail = ordered[-1] + if edge["subject"] == tail and edge["object"] not in ordered: + ordered.append(edge["object"]) + progress = True + elif edge["object"] == tail and edge["subject"] not in ordered: + ordered.append(edge["subject"]) + progress = True + if not progress: + raise ValueError( + f"Could not construct path from {source} to {target} from given edges" + ) + return ordered + + +def read_path(bmt, source: str, target: str, edge_ids, knowledge_graph: dict): + """`(node CURIEs, their categories, hop phrases)`; ValueError for a path + that cannot be read.""" + nodes, edges = knowledge_graph["nodes"], knowledge_graph["edges"] + ordered = order_path_nodes(source, target, edge_ids, edges) + categories = [] + for curie in ordered: + node = nodes.get(curie) + if node is None: + raise ValueError(f"{curie} is not in the knowledge graph") + found = node.get("categories") or [] + categories.append( + most_specific_category(bmt, tuple(found)) if found else "biolink:NamedThing" + ) + phrases = [] + for current, following in zip(ordered[:-1], ordered[1:]): + phrase = hop_phrase(bmt, current, following, edge_ids, edges) + if phrase is None: + raise ValueError( + f"no edge of the path can be walked from {current} to {following}" + ) + phrases.append(phrase) + return ordered, categories, phrases + + +def path_scalars(scalar_names, node_curies, node_context=None, edge_context=None): + """One path's scalars in `scalar_names` order. `node_context(curie)` and + `edge_context(a, b)` give a table's vector, zeros where it has none; a + table not given leaves its scalars zero, as is the retrieval cost g.""" + out = np.zeros(len(scalar_names), dtype=np.float32) + col = {name: j for j, name in enumerate(scalar_names)} + out[col["n_hops_scaled"]] = (len(node_curies) - 1) / N_HOPS_SCALE + if node_context is not None: + vecs = np.stack([node_context(c) for c in node_curies]) + for name, j, reduce in _NODE_COLUMNS: + out[col[name]] = reduce(vecs[:, j]) + if edge_context is not None: + hops = np.stack( + [edge_context(a, b) for a, b in zip(node_curies[:-1], node_curies[1:])] + ) + for j, base in enumerate(_EDGE_BASES): + out[col[f"{base}_mean"]] = hops[:, j].mean() + out[col[f"{base}_min"]] = hops[:, j].min() + return out + + +def vector(txn, kind: str, text: str) -> np.ndarray: + raw = txn.get(lmdb_key(kind, text)) + if raw is None: + raise KeyError(f"{kind}:{text}") + return np.frombuffer(raw, dtype=DTYPES[kind]) + + +def path_features(txn, config, nodes, categories, phrases, kept=None): + """One path's `(node_x, edge_x, scalars)`; KeyError naming the first + vector the LMDB lacks. `kept` carries the category and phrase vectors, + as float32, from one path of a message to the next: there are few, read + over and over. Nodes are not kept, so it stays small.""" + dim = config["embed_dim"] + kept = {} if kept is None else kept + + def repeated(kind, text): + if (kind, text) not in kept: + kept[kind, text] = vector(txn, kind, text).astype(np.float32) + return kept[kind, text] + + node_x = np.empty((len(nodes), 2 * dim), dtype=np.float32) + for j, (curie, category) in enumerate(zip(nodes, categories)): + node_x[j, :dim] = vector(txn, NODE, curie) + node_x[j, dim:] = repeated(CATEGORY, category) + edge_x = np.stack([repeated(PHRASE, phrase) for phrase in phrases]) + + def context(key, width): + raw = txn.get(key) + if raw is None: + return np.zeros(width, dtype=np.float32) + return np.frombuffer(raw, dtype=np.float32) + + scalars = path_scalars( + config["scalar_names"], + nodes, + ( + (lambda c: context(lmdb_key(NODE_CONTEXT, c), NODE_CONTEXT_DIM)) + if config["node_context"] + else None + ), + ( + (lambda a, b: context(edge_context_key(a, b), len(_EDGE_BASES))) + if config["edge_context"] + else None + ), + ) + return node_x, edge_x, scalars + + +def is_direct_path(source, target, path, knowledge_graph): + """Whether every edge of the path joins the two pinned nodes.""" + edges = knowledge_graph["edges"] + return bool(path) and all( + eid in edges + and {edges[eid]["subject"], edges[eid]["object"]} == {source, target} + for eid in path + ) + + +def score_message(message, txn, model, config, bmt, chunk_size): + """Give every analysis a `score` in place, and return counts for the log. + + The direct edge scores 1. Other paths are scored a length at a time, at + most `chunk_size` at once, on the checkpoint's length scale -- how paths + of different lengths are ranked together -- through a logistic into + (0, 1). A path that cannot be read, has no length scale, or lacks a + vector keeps 0. + """ + stats = { + "analyses": 0, + "scored": 0, + "direct": 0, + "no_binding": 0, + "bad_path": 0, + "no_scale": 0, + "missing": 0, + "missing_samples": [], + "model_seconds": 0.0, + } + msg = message["message"] + knowledge_graph = msg["knowledge_graph"] + auxiliary_graphs = msg.get("auxiliary_graphs") or {} + qpath_id, qpath = next(iter(msg["query_graph"]["paths"].items())) + pending, kept = {}, {} + + def flush(n_hops): + batch = pending.pop(n_hops, []) + if not batch: + return + started = time.time() + node_x, edge_x, scalars = ( + np.stack(column) for column in zip(*(features for _, features in batch)) + ) + with torch.inference_mode(): + raw = model( + torch.from_numpy(node_x), + torch.from_numpy(edge_x), + torch.full((len(batch),), n_hops, dtype=torch.int64), + torch.from_numpy(scalars), + torch.zeros(len(batch), dtype=torch.int64), + ) + raw = raw.float().numpy() + stats["model_seconds"] += time.time() - started + mean, spread = config["length_scale"][n_hops] + z = ((raw.astype(np.float64) - mean) / spread).astype(np.float32) + scores = 1.0 / (1.0 + np.exp(-z.astype(np.float64))) + for (analysis, _), score in zip(batch, scores): + analysis["score"] = float(score) + stats["scored"] += len(batch) + + for result in msg.get("results") or []: + analyses = result.get("analyses") or [] + for analysis in analyses: + analysis["score"] = 0.0 + stats["analyses"] += 1 + try: + source = result["node_bindings"][qpath["subject"]]["ids"][0] + target = result["node_bindings"][qpath["object"]]["ids"][0] + except (KeyError, IndexError, TypeError): + stats["no_binding"] += len(analyses) + continue + for analysis in analyses: + try: + path_binding = (analysis.get("path_bindings") or {}).get(qpath_id) + aux_id = path_binding["ids"][0] + edge_ids = auxiliary_graphs[aux_id]["edges"] + except (KeyError, IndexError, TypeError): + stats["no_binding"] += 1 + continue + if is_direct_path(source, target, edge_ids, knowledge_graph): + analysis["score"] = 1.0 + stats["direct"] += 1 + continue + try: + nodes, categories, phrases = read_path( + bmt, source, target, edge_ids, knowledge_graph + ) + except ValueError: + stats["bad_path"] += 1 + continue + n_hops = len(phrases) + if n_hops not in config["length_scale"]: + stats["no_scale"] += 1 + continue + try: + features = path_features(txn, config, nodes, categories, phrases, kept) + except KeyError as e: + stats["missing"] += 1 + samples = stats["missing_samples"] + if len(samples) < 5 and e.args[0] not in samples: + samples.append(e.args[0]) + continue + pending.setdefault(n_hops, []).append((analysis, features)) + if len(pending[n_hops]) >= chunk_size: + flush(n_hops) + for n_hops in sorted(pending): + flush(n_hops) + return stats diff --git a/workers/score_paths/worker.py b/workers/score_paths/worker.py index 6b9b7f8d..90c37545 100644 --- a/workers/score_paths/worker.py +++ b/workers/score_paths/worker.py @@ -1,4 +1,5 @@ -"""Path scoring module""" +"""Path scoring module: every analysis scored by the sequence reranker +(``scorer``), from the LMDB squashbert's scripts/export_reranker.py writes.""" import asyncio import logging @@ -7,10 +8,7 @@ import uuid import lmdb -import numpy as np import torch -from bmt import Toolkit -from torch import nn from shepherd_utils.config import settings from shepherd_utils.cpu import resolve_pool_workers @@ -21,23 +19,17 @@ from shepherd_utils.process_pool import ProcessPoolManager from shepherd_utils.shared import get_tasks, run_task_lifecycle +import scorer + STREAM = "score_paths" GROUP = "consumer" CONSUMER = str(uuid.uuid4())[:8] TASK_LIMIT = 6 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 +# Paths scored per forward pass, per length. It bounds peak memory (about 40 KB +# of inputs a path) whatever the message's size: scoring a whole message in +# one pass is what OOM-killed the pod on large messages. +SCORE_CHUNK_SIZE = 1024 tracer = setup_tracer(STREAM) LOGGER = get_worker_logger(STREAM) @@ -46,104 +38,8 @@ # 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: - edges = knowledge_graph["edges"] - nodes = knowledge_graph["nodes"] - ordered = [source] - while target not in ordered: - noncylcic = False - for eid in path: - edge = edges.get(eid) - if not edge: - continue - tail = ordered[-1] - if edge["subject"] == tail and edge["object"] not in ordered: - ordered.append(edge["object"]) - noncylcic = True - elif edge["object"] == tail and edge["subject"] not in ordered: - ordered.append(edge["subject"]) - noncylcic = True - if not noncylcic: - return None # Cyclic paths are not supported - if len(ordered) != 4: - return None - names = [] - cat_phrases = [] - for nid in ordered: - node = nodes.get(nid) - if node is None: - return None - name = node.get("name") - cats = node.get("categories") or [] - if not name or not cats or not cats[0]: - return None - names.append(name) - cat_phrases.append(cats[0].removeprefix("biolink:")) - hop_phrases = [] - for cur, nxt in zip(ordered[:-1], ordered[1:]): - preds = [] - seen = set() - for eid in path: - edge = edges.get(eid) - if not edge: - continue - pname = edge["predicate"].removeprefix("biolink:").replace("_", " ") - if not bmt.is_predicate(pname): - continue # mirror cache: drop non-biolink predicates - if edge["subject"] == cur and edge["object"] == nxt: - name = pname - elif edge["object"] == cur and edge["subject"] == nxt: - if bmt.is_symmetric(pname): - name = pname - else: - inv = bmt.get_inverse(pname) - if inv is None: - continue # mirror cache: silently drop non-invertibles - name = inv - else: - continue - if name not in seen: - seen.add(name) - preds.append(name) - if not preds: - return None - preds.sort() - hop_phrases.append( - preds[0] if len(preds) == 1 else "[" + " or ".join(preds) + "]" - ) - return names, cat_phrases, hop_phrases - except Exception as e: - logger.error(f"Failed to convert path to components: {e}") - return None - - -def _lookup(txn, key): - raw = txn.get(key.encode("utf-8")) - if raw is None: - raise KeyError(key) - return np.frombuffer(raw, dtype=np.float16) - - -def _probe_cache(env): - """Confirm the LMDB cache is non-empty and decodes to (768,) float16.""" - with env.begin() as txn: - n = txn.stat()["entries"] - if n == 0: - raise RuntimeError("embeddings cache is empty") - cursor = txn.cursor() - cursor.first() - key, value = cursor.item() - expected_bytes = 768 * np.dtype(np.float16).itemsize - if len(value) != expected_bytes: - raise RuntimeError( - f"embeddings cache has wrong value size: got {len(value)} bytes, " - f"expected {expected_bytes} (768-dim float16)" - ) - return n, key.decode("utf-8", errors="replace") +model = None +config = None def _open_embeddings(): @@ -158,251 +54,101 @@ def _open_embeddings(): ) -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. + an LMDB without a reranker 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) + with env.begin() as txn: + entries = txn.stat()["entries"] + _, found = scorer.load_reranker(txn) 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}") + logger.info( + f"embeddings LMDB: {entries} entries; reranker scores " + f"{sorted(found['length_scale'])}-hop paths, biolink " + f"{found['biolink_version']}, node context " + f"{found['node_context']}, edge context {found['edge_context']}" + ) 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 + here (an unreadable LMDB, a corrupt model) 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. + The embeddings LMDB is a file-backed mapping, so the OS page cache serves + every child from one copy (see ``_open_embeddings``). The reranker (about + 16 MB) and the biolink ``Toolkit`` are 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: + global bmt, embedding_env, model, config + if model 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) + with embedding_env.begin() as txn: + model, config = scorer.load_reranker(txn) + bmt = scorer.load_toolkit(config) 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"] - results = message["message"]["results"] - knowledge_graph = message["message"]["knowledge_graph"] + results = message["message"].get("results") or [] auxiliary_graphs = message["message"].get("auxiliary_graphs") or {} - qpath_id, qpath = next(iter(paths.items())) - subject_qnode = qpath["subject"] - object_qnode = qpath["object"] total_analyses = sum(len(r.get("analyses") or []) for r in results) logger.info( f"Scoring {response_id}: {len(results)} results, " f"{total_analyses} analyses, {len(auxiliary_graphs)} aux graphs" ) - 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() + started = time.time() with embedding_env.begin() as txn: - for result_ind, result in enumerate(results): - try: - # TRAPI 2.0: one binding {"ids": [...]} per qnode/qpath. - source = result["node_bindings"][subject_qnode]["ids"][0] - target = result["node_bindings"][object_qnode]["ids"][0] - except (KeyError, IndexError, TypeError): - continue - analyses = result.get("analyses") or [] - for analysis_ind, analysis in enumerate(analyses): - path_binding = (analysis.get("path_bindings") or {}).get(qpath_id) - try: - aux_id = path_binding["ids"][0] - edge_ids = auxiliary_graphs[aux_id]["edges"] - except (KeyError, IndexError, TypeError): - analysis["score"] = 0.0 - skip_no_binding += 1 - continue - components = convert_path_to_components( - source, target, edge_ids, knowledge_graph, logger - ) - if components is None: - analysis["score"] = 0.0 - skip_bad_path += 1 - continue - names, cats, hops = components - try: - features = np.concatenate( - [ - _lookup(txn, names[0]), - _lookup(txn, cats[0]), - _lookup(txn, hops[0]), - _lookup(txn, names[1]), - _lookup(txn, cats[1]), - _lookup(txn, hops[1]), - _lookup(txn, names[2]), - _lookup(txn, cats[2]), - _lookup(txn, hops[2]), - _lookup(txn, names[3]), - _lookup(txn, cats[3]), - ] - ) - except KeyError as e: - key = e.args[0] - if len(missing_samples) < 5 and key not in missing_samples: - missing_samples.append(key) - analysis["score"] = 0.0 - skip_missing_emb += 1 - continue - 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: {scored}/{total_analyses} ready in {build_time:.1f}s" + stats = scorer.score_message( + message, txn, model, config, bmt, SCORE_CHUNK_SIZE + ) + elapsed = time.time() - started + skipped = ( + stats["no_binding"] + + stats["bad_path"] + + stats["no_scale"] + + stats["missing"] + ) + msg = ( + f"Scored {stats['scored']}/{stats['analyses']} analyses in " + f"{elapsed:.1f}s ({stats['model_seconds']:.1f}s in the reranker); " + f"{stats['direct']} direct" + ) if skipped: msg += ( f"; skipped {skipped} " - f"(no binding: {skip_no_binding}, " - f"bad path: {skip_bad_path}, " - f"missing embedding: {skip_missing_emb})" + f"(no binding: {stats['no_binding']}, " + f"bad path: {stats['bad_path']}, " + f"no length scale: {stats['no_scale']}, " + f"missing embedding: {stats['missing']})" ) - if missing_samples: - msg += f"; missing keys e.g. {missing_samples}" + if stats["missing_samples"]: + msg += f"; missing keys e.g. {stats['missing_samples']}" logger.info(msg) - if scored: - logger.info( - 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") except Exception as e: logger.error(f"Error scoring paths: {e}", exc_info=True) for result in message["message"].get("results", []) or []: - for analysis in result.get("analyses", []) or []: + for analysis in result.get("analyses") or []: analysis.setdefault("score", 0.0) for attempt in range(5): try: @@ -483,8 +229,8 @@ async def poll_for_tasks(): _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. + # own copy of the reranker 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(