Aragorn pathfinder update

.gitignore

@@ -5,6 +5,7 @@ debug/
 gandalf_mmap/
 omnicorp_lmdb/
 pathfinder_embeddings/
+pathfinder_search_graph/
 arax_pathfinder_dbs/
 arax_dbs/
 

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

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

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

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

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,27 +604,42 @@ 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
         if num_tries < 4:
             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
 

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}