diff --git a/README.md b/README.md index f601a01..066b77d 100644 --- a/README.md +++ b/README.md @@ -14,33 +14,46 @@ The main entrypoint is `./compose.yml` and will spin everything up. If you want to add a new operation/worker, add a new service in `compose.yml` under `services`. -### Worker data (LMDB) downloads +### Worker data (LMDB / sqlite) downloads -A couple of workers read from large, read-only LMDB datasets that are too big to +A couple of workers read from large, read-only sqlite databases and LMDB datasets that are too big to commit to git (they're gitignored and volume-mounted from the host): - **`aragorn_omnicorp`** → `./omnicorp_lmdb/` (`curies.lmdb`, `shared_counts.lmdb`) - **`score_paths`** → `./pathfinder_embeddings/` (a directory-style LMDB) +- **`arax_pathfinder`** → `./arax_pathfinder_dbs/` (`curie_ngd_v1.0_.sqlite`, `tier0-info-for-overlay_v1.0_.sqlite`) -So a new developer doesn't have to source these by hand, each worker can fetch -its dataset on first startup. Point it at a `.tar.gz` on an external server by -adding the matching variable to your root `.env` file: +So a new developer doesn't have to source these by hand, each worker can fetch its dataset on first +startup. Two download mechanisms are supported, depending on where the dataset lives: + +**LMDB datasets (`aragorn_omnicorp`, `score_paths`)** are fetched as a `.tar.gz` from a plain HTTP(S) +URL and extracted in place. Add the matching variable to your root `.env` file: ```dotenv OMNICORP_LMDB_URL=https://example.org/path/omnicorp_lmdb.tar.gz PATHFINDER_EMBEDDINGS_URL=https://example.org/path/pathfinder_embeddings.tar.gz ``` -On startup the worker checks whether its LMDB files already exist in the -volume-mounted directory. If they're missing and a URL is set, it downloads the -archive and extracts it into that directory — which lives on the host, so the -data persists across restarts and is only downloaded once. If the files are -already present, or no URL is configured, the download is skipped (production -mounts this data out of band, so it's unaffected). +The archive for each dataset should contain the expected files at its top level: `curies.lmdb` and +`shared_counts.lmdb` for omnicorp, `data.mdb` (and `lock.mdb`) for the embeddings. + +**`arax_pathfinder`'s sqlite databases** don't live behind a URL — they're on a private, +SSH-accessible host (`arax-databases.rtx.ai`), so each file is fetched individually via `scp` +instead. The filenames and remote directory both embed a data-tier version that changes +periodically, so only one variable needs updating when a new tier ships: + +```dotenv +ARAX_PATHFINDER_TIER_VERSION=tier0-20260621 +``` + +This requires an SSH key with access to that host, mounted read-only into the container +(`~/.ssh:/root/.ssh:ro` in docker-compose.yml). -The archive for each dataset should contain the expected files at its top level: -`curies.lmdb` and `shared_counts.lmdb` for omnicorp, `data.mdb` (and -`lock.mdb`) for the embeddings. +On startup, each worker checks whether its files already exist in the volume-mounted directory. If +they're missing and a source is configured (URL or scp path), it fetches them into that directory — +which lives on the host, so the data persists across restarts and is only downloaded once. If the +files are already present, or no source is configured, the download is skipped (production mounts +this data out of band, so it's unaffected). ### Worker diff --git a/compose.yml b/compose.yml index 8847600..5fbf34b 100644 --- a/compose.yml +++ b/compose.yml @@ -204,7 +204,7 @@ services: - ./.env:/app/.env # First run? Set PATHFINDER_EMBEDDINGS_URL in your .env to a .tar.gz and # the worker downloads + extracts the embeddings LMDB into this mount on - # startup (see README "Worker data (LMDB) downloads"). + # startup (see README "Worker data (LMDB / sqlite) downloads"). - ./pathfinder_embeddings:/app/pathfinder_embeddings ######### Example ARA @@ -311,7 +311,7 @@ services: - ./.env:/app/.env # First run? Set OMNICORP_LMDB_URL in your .env to a .tar.gz and the worker # downloads + extracts curies.lmdb / shared_counts.lmdb into this mount on - # startup (see README "Worker data (LMDB) downloads"). + # startup (see README "Worker data (LMDB / sqlite) downloads"). - ./omnicorp_lmdb:/app/omnicorp_lmdb aragorn_score: container_name: aragorn_score @@ -344,6 +344,28 @@ services: volumes: - ./logs:/app/logs - ./.env:/app/.env + arax_pathfinder: + container_name: arax_pathfinder + build: + context: . + dockerfile: workers/arax_pathfinder/Dockerfile + restart: unless-stopped + depends_on: + shepherd_db: + condition: service_healthy + shepherd_broker: + condition: service_healthy + volumes: + - ./logs:/app/logs + - ./.env:/app/.env + # First run? The worker scp's its two sqlite dbs down from + # arax-databases.rtx.ai on startup. Set ARAX_PATHFINDER_TIER_VERSION in + # your .env if you need a tier other than the default (see README "Worker + # data (LMDB / sqlite) downloads"). Requires your SSH key to have access + # to that host. + - ./arax_pathfinder_dbs:/app/arax_pathfinder_dbs + # Replace '/Users/facadmin' with your local home directory path + - /Users/facadmin/.ssh:/home/nru/.ssh:ro arax_rank: container_name: arax_rank diff --git a/shepherd_utils/config.py b/shepherd_utils/config.py index 67f4dce..c2b8237 100644 --- a/shepherd_utils/config.py +++ b/shepherd_utils/config.py @@ -84,7 +84,24 @@ class Settings(BaseSettings): sync_kg_retrieval_url: str = "http://host.docker.internal:8080/query" kg_rehydrate_url: str = "http://host.docker.internal:8080/rehydrate" default_data_tier: int = 0 + + + # ARAX configs arax_url: str = "https://arax.ncats.io/shepherd/api/arax/v1.4/query" + plover_url: str = "https://kg2cploverdb.ci.transltr.io" + arax_biolink_version: str = "4.2.5" + arax_blocked_list_url: str = ( + "https://raw.githubusercontent.com/RTXteam/RTX/master/" + "code/ARAX/KnowledgeSources/general_concepts.json" + ) + + arax_pathfinder_dbs_dir: str = "arax_pathfinder_dbs" + arax_pathfinder_tier_version: str = "tier0-20260621" + arax_pathfinder_curie_ngd_sqlite_filename: str = "curie_ngd_v1.0_{version}.sqlite" + arax_pathfinder_tier0_overlay_sqlite_filename: str = "tier0-info-for-overlay_v1.0_{version}.sqlite" + arax_pathfinder_sqlite_host: str = "rtxconfig@arax-databases.rtx.ai" + arax_pathfinder_sqlite_remote_dir: str = "~/{version}" + # End of ARAX configs pathfinder_redis_host: str = "host.docker.internal" pathfinder_redis_port: int = 6383 diff --git a/shepherd_utils/data_download.py b/shepherd_utils/data_download.py index 6df753f..dce8282 100644 --- a/shepherd_utils/data_download.py +++ b/shepherd_utils/data_download.py @@ -1,32 +1,42 @@ -"""Ensure large read-only LMDB datasets are present, downloading them on first +"""Ensure large read-only datasets are present, downloading them on first run so new developers can spin the stack up locally. -The ``aragorn_omnicorp`` and ``score_paths`` workers read from LMDB datasets -that are far too large to commit to git -- they're gitignored and volume-mounted -from the host (``./omnicorp_lmdb`` and ``./pathfinder_embeddings``). 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 a missing LMDB. +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``). 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. -When a download URL is configured (``OMNICORP_LMDB_URL`` / -``PATHFINDER_EMBEDDINGS_URL``, read via :mod:`shepherd_utils.config`), each of -those workers calls the matching ``ensure_*`` helper at startup: +Two flavors of remote source are supported: + +* **HTTP** -- 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). +* **SCP** -- individual files fetched from a private, SSH-accessible host via + the system ``scp`` binary (used by ``arax_pathfinder`` below, whose two + sqlite databases live on ``arax-databases.rtx.ai`` rather than behind a + plain URL -- there's no bucket/CDN in front of them, just SSH access). + +When a download source is configured (read via :mod:`shepherd_utils.config`), +each worker calls its matching ``ensure_*`` helper at startup: * if the expected files are already present it's a no-op; -* otherwise the dataset is fetched as a ``.tar.gz`` from the external server and - extracted into the (volume-mounted) target directory, so it persists on the - host and is only downloaded once. +* otherwise the dataset is fetched and written into the (volume-mounted) + target directory, so it persists on the host and is only downloaded once. -With no URL configured the call is a no-op that logs how to enable the download, -so production -- where the data is already mounted -- is unaffected. +With no source configured the call is a no-op that logs how to enable the +download, so production -- where the data is already mounted -- is +unaffected. """ import logging import os +import subprocess import tarfile import tempfile import urllib.request -from typing import List, Optional +from typing import Dict, List, Optional, Tuple from shepherd_utils.config import settings @@ -153,6 +163,129 @@ def ensure_lmdb_dataset( logger.info(f"{name}: dataset ready in {target_dir}.") +def _scp_download(remote_path: str, dest_path: str, logger: logging.Logger) -> None: + """Copy a single file from a remote host to ``dest_path`` via ``scp``. + + Unlike ``_download`` above, these sqlite files aren't behind a plain URL -- + they live on a private, SSH-accessible host (see README), so this shells + out to the system ``scp`` binary and relies on the caller's SSH key (or + agent) for auth rather than any credential this code holds. + + ``BatchMode=yes`` makes scp fail fast instead of hanging on an interactive + password/passphrase prompt if the key isn't set up. The known-hosts file is + redirected to a scratch path so this still works even when ``~/.ssh`` is + mounted read-only -- a fresh container has nothing pinned there yet, and + ``accept-new`` trusts the host key on first connect without prompting. + ``-C`` enables compression, which helps for a database-sized transfer. + """ + logger.info(f"Downloading {remote_path} via scp ...") + cmd = [ + "scp", + "-C", + "-o", "BatchMode=yes", + "-o", "StrictHostKeyChecking=accept-new", + "-o", "UserKnownHostsFile=/tmp/known_hosts", + remote_path, + dest_path, + ] + try: + subprocess.run(cmd, check=True, capture_output=True, text=True) + except FileNotFoundError as e: + raise RuntimeError( + "scp binary not found in this image -- install openssh-client." + ) from e + except subprocess.CalledProcessError as e: + raise RuntimeError( + f"scp failed for {remote_path} (exit {e.returncode}): " + f"{e.stderr.strip()}. Confirm your SSH key has access to the " + f"source host and is mounted into the container (see README)." + ) from e + size_mb = os.path.getsize(dest_path) / 1e6 + logger.info(f"Download complete: {dest_path} ({size_mb:.0f} MB)") + + +def ensure_scp_dataset( + name: str, + target_dir: str, + file_sources: Dict[str, str], + logger: Optional[logging.Logger] = None, +) -> None: + """Ensure each file in ``file_sources`` exists under ``target_dir``. + + Unlike ``ensure_lmdb_dataset`` (one ``.tar.gz`` archive fetched over HTTP + and extracted), each of these files is fetched individually via ``scp`` + from a private, SSH-accessible host. ``file_sources`` maps the expected + local filename to its ``user@host:path`` remote source; a filename whose + source is empty is skipped (warned about) rather than downloaded, same as + an unset ``url`` in ``ensure_lmdb_dataset``. + + Idempotent: once a file is present it's left alone, so it's safe to call + unconditionally on every worker startup. + """ + logger = logger or logging.getLogger(__name__) + required_files = list(file_sources.keys()) + + missing = _missing_files(target_dir, required_files) + if not missing: + logger.info( + f"{name}: dataset already present in {target_dir}; skipping download." + ) + return + + os.makedirs(target_dir, exist_ok=True) + logger.info(f"{name}: dataset missing from {target_dir} (missing: {missing}).") + + attempted = [] + for filename in missing: + remote = file_sources.get(filename) + if not remote: + logger.warning( + f"{name}: {filename} missing from {target_dir} and no source " + f"configured for it. Set the corresponding *_SOURCE env var (see " + f"the README) to download it automatically, or provide the file " + f"manually. The worker will fail to start without it." + ) + continue + attempted.append(filename) + + dest_path = os.path.join(target_dir, filename) + # Download to a temp file in the same dir first, then atomically rename, + # so a partial/interrupted transfer is never mistaken for a complete + # file (same reasoning as the tar.gz download above). + tmp_fd, tmp_path = tempfile.mkstemp(suffix=".part", dir=target_dir) + os.close(tmp_fd) + try: + _scp_download(remote, tmp_path, logger) + os.replace(tmp_path, dest_path) + except Exception: + try: + os.remove(tmp_path) + except OSError: + pass + raise + + # Only files we actually attempted (had a source) count toward failure -- + # a file with no source configured was already warned about above and is + # expected to still be missing, same as an unset url in + # ensure_lmdb_dataset. Checking against `required_files` here would raise + # even when nothing went wrong. + still_missing = _missing_files(target_dir, attempted) + if still_missing: + raise RuntimeError( + f"{name}: still missing expected files after download attempt: " + f"{still_missing}. Check that the *_SOURCE env vars are set and that " + f"your SSH key has access to the source host." + ) + if _missing_files(target_dir, required_files): + logger.warning( + f"{name}: dataset partially ready in {target_dir} -- some files have " + f"no source configured (see warnings above). The worker will fail " + f"when it tries to open them." + ) + else: + logger.info(f"{name}: dataset ready in {target_dir}.") + + def ensure_omnicorp_lmdb(logger: Optional[logging.Logger] = None) -> None: """Ensure the omnicorp curies + shared-counts LMDBs are present. @@ -187,3 +320,65 @@ def ensure_pathfinder_embeddings(logger: Optional[logging.Logger] = None) -> Non url=settings.pathfinder_embeddings_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 + templates + the current ``arax_pathfinder_tier_version``. + + Single source of truth for these two paths: ``ensure_arax_pathfinder_dbs`` + (below) uses it to know what to download and where, and worker.py's + ``execute_pathfinding_sync`` uses it to know what to open, so the two can + never disagree about a file's location the way two independently-defined + settings could. + """ + version = settings.arax_pathfinder_tier_version + curie_ngd_path = os.path.join( + settings.arax_pathfinder_dbs_dir, + settings.arax_pathfinder_curie_ngd_sqlite_filename.format(version=version), + ) + node_degree_path = os.path.join( + settings.arax_pathfinder_dbs_dir, + settings.arax_pathfinder_tier0_overlay_sqlite_filename.format(version=version), + ) + return curie_ngd_path, node_degree_path + + +def ensure_arax_pathfinder_dbs(logger: Optional[logging.Logger] = None) -> None: + """Ensure the arax_pathfinder worker's two sqlite databases are present. + + Both live on a private, SSH-accessible host (``arax-databases.rtx.ai``) + rather than behind a plain download URL, so each is fetched individually + via ``scp`` instead of the tar.gz + extract flow used for the LMDB + datasets above. Both are expected in the same directory (see the + ``arax_pathfinder`` volume mount in docker-compose.yml). + + The local filenames and the remote directory both embed a data-tier + version (e.g. ``tier0-20260621``) that changes periodically as new tiers + ship. Rather than duplicate that string across separate path/source + settings -- which can drift out of sync if only one is updated -- the + filenames and remote dir are templates with a ``{version}`` placeholder, + filled in from the single ``arax_pathfinder_tier_version`` setting. + Bumping to a new tier is then one env var change + (``ARAX_PATHFINDER_TIER_VERSION``) rather than several. + """ + curie_ngd_path, node_degree_path = arax_pathfinder_sqlite_paths() + target_dir = settings.arax_pathfinder_dbs_dir + + version = settings.arax_pathfinder_tier_version + remote_dir = settings.arax_pathfinder_sqlite_remote_dir.format(version=version) + host = settings.arax_pathfinder_sqlite_host + + curie_ngd_filename = os.path.basename(curie_ngd_path) + node_degree_filename = os.path.basename(node_degree_path) + + ensure_scp_dataset( + name="arax_pathfinder", + target_dir=target_dir, + file_sources={ + curie_ngd_filename: f"{host}:{remote_dir}/{curie_ngd_filename}", + node_degree_filename: f"{host}:{remote_dir}/{node_degree_filename}", + }, + logger=logger, + ) diff --git a/workers/arax/inject_shepherd_arax_provenance.py b/shepherd_utils/inject_shepherd_arax_provenance.py similarity index 100% rename from workers/arax/inject_shepherd_arax_provenance.py rename to shepherd_utils/inject_shepherd_arax_provenance.py diff --git a/workers/arax/worker.py b/workers/arax/worker.py index 3d26a44..84ab9e4 100644 --- a/workers/arax/worker.py +++ b/workers/arax/worker.py @@ -4,9 +4,10 @@ import json import logging import uuid - import httpx -from inject_shepherd_arax_provenance import add_shepherd_arax_to_edge_sources +from shepherd_utils.inject_shepherd_arax_provenance import ( + add_shepherd_arax_to_edge_sources, +) from shepherd_utils.config import settings from shepherd_utils.db import get_message, save_message @@ -21,33 +22,50 @@ tracer = setup_tracer(STREAM) -async def arax(task, logger: logging.Logger): +def is_pathfinder_query(message): try: - query_id = task[1]["query_id"] - logger.info(f"Getting message from db for query id {query_id}") - message = await get_message(query_id, logger) - message["submitter"] = "Shepherd" - logger.info(f"Get the message from db {message}") - - headers = {"Content-Type": "application/json"} - async with httpx.AsyncClient(timeout=270) as client: - response = await client.post( - settings.arax_url, json=message, headers=headers - ) - - logger.info(f"Status Code from ARAX response: {response.status_code}") - result = response.json() - result = add_shepherd_arax_to_edge_sources(result) - - except Exception as e: - logger.error(f"Error occurred in ARAX entry module: {e}") - result = {"status": "error", "error": str(e)} - - response_id = task[1]["response_id"] + # this can still fail if the input looks like e.g.: + # "query_graph": None + qedges = message.get("message", {}).get("query_graph", {}).get("edges", {}) + except: + qedges = {} + try: + # this can still fail if the input looks like e.g.: + # "query_graph": None + qpaths = message.get("message", {}).get("query_graph", {}).get("paths", {}) + except: + qpaths = {} + if len(qpaths) > 1: + raise Exception("Only a single path is supported", 400) + if (len(qpaths) > 0) and (len(qedges) > 0): + raise Exception("Mixed mode pathfinder queries are not supported", 400) + return len(qpaths) == 1 - await save_message(response_id, result, logger) - task[1]["workflow"] = json.dumps([{"id": "arax"}]) +async def arax(task, logger: logging.Logger): + query_id = task[1]["query_id"] + logger.info(f"Getting message from db for query id {query_id}") + message = await get_message(query_id, logger) + if is_pathfinder_query(message): + task[1]["workflow"] = json.dumps([{"id": "arax.pathfinder"}]) + else: + try: + message["submitter"] = "Shepherd" + logger.info(f"Get the message from db {message}") + headers = {"Content-Type": "application/json"} + async with httpx.AsyncClient(timeout=270) as client: + response = await client.post( + settings.arax_url, json=message, headers=headers + ) + logger.info(f"Status Code from ARAX response: {response.status_code}") + result = response.json() + result = add_shepherd_arax_to_edge_sources(result) + except Exception as e: + logger.error(f"Error occurred calling ARAX service: {e}") + result = {"status": "error", "error": str(e)} + response_id = task[1]["response_id"] + await save_message(response_id, result, logger) + task[1]["workflow"] = json.dumps([{"id": "arax"}]) async def process_task(task, parent_ctx, logger: logging.Logger, limiter): diff --git a/workers/arax_pathfinder/Dockerfile b/workers/arax_pathfinder/Dockerfile new file mode 100644 index 0000000..68ae2b6 --- /dev/null +++ b/workers/arax_pathfinder/Dockerfile @@ -0,0 +1,34 @@ +# Use RENCI python base image +FROM ghcr.io/translatorsri/renci-python-image:3.12.13 + +# Add image info +LABEL org.opencontainers.image.source https://github.com/BioPack-team/shepherd + +ENV PYTHONHASHSEED=0 + +# set up requirements +WORKDIR /app + +# make sure all is writeable for the nru USER later on +RUN chmod -R 777 . + +# Install requirements +COPY ./shepherd_utils ./shepherd_utils +COPY ./pyproject.toml . +RUN pip install . + +COPY ./workers/arax_pathfinder/requirements.txt . +RUN pip install -r requirements.txt + +# switch to the non-root user (nru). defined in the base image +USER nru + +# Copy in files +COPY ./workers/arax_pathfinder ./ + +# Set up base for command and any variables +# that shouldn't be modified +# ENTRYPOINT ["uvicorn", "shepherd_server.server:APP"] + +# Variables that can be overriden +CMD ["python", "worker.py"] diff --git a/workers/arax_pathfinder/__init__.py b/workers/arax_pathfinder/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/workers/arax_pathfinder/requirements.txt b/workers/arax_pathfinder/requirements.txt new file mode 100644 index 0000000..167e022 --- /dev/null +++ b/workers/arax_pathfinder/requirements.txt @@ -0,0 +1,3 @@ +catrax-pathfinder==2.4.3 +biolink-helper-pkg==1.0.0 +httpx==0.28.1 \ No newline at end of file diff --git a/workers/arax_pathfinder/worker.py b/workers/arax_pathfinder/worker.py new file mode 100644 index 0000000..d81a524 --- /dev/null +++ b/workers/arax_pathfinder/worker.py @@ -0,0 +1,289 @@ +"""Arax ARA Pathfinder module.""" + +import asyncio +import json +import logging +import time +import uuid +from pathlib import Path + +import httpx +import requests +from biolink_helper_pkg import BiolinkHelper +from pathfinder.Pathfinder import Pathfinder + +from shepherd_utils.config import settings +from shepherd_utils.data_download import ( + arax_pathfinder_sqlite_paths, + ensure_arax_pathfinder_dbs, +) +from shepherd_utils.db import ( + get_message, + save_message, +) +from shepherd_utils.inject_shepherd_arax_provenance import ( + add_shepherd_arax_to_edge_sources, +) +from shepherd_utils.otel import setup_tracer +from shepherd_utils.shared import ( + get_tasks, + handle_task_failure, + wrap_up_task, +) + +# Queue name +STREAM = "arax.pathfinder" +# Consumer group, most likely you don't need to change this. +GROUP = "consumer" +CONSUMER = str(uuid.uuid4())[:8] +TASK_LIMIT = 100 +tracer = setup_tracer(STREAM) + +NUM_TOTAL_HOPS = 4 +MAX_HOPS_TO_EXPLORE = 4 +MAX_PATHFINDER_PATHS = 500 +PRUNE_TOP_K = 75 +NODE_DEGREE_THRESHOLD = 10000 + +OUT_PATH = Path("general_concepts.json") + + +def download_file(url: str, out_path: Path, overwrite: bool = False) -> Path: + out_path = Path(out_path) + + if out_path.exists() and not overwrite: + return out_path + + out_path.parent.mkdir(parents=True, exist_ok=True) + + r = requests.get(url, timeout=60) + r.raise_for_status() + + out_path.write_bytes(r.content) + return out_path + + +def get_blocked_list(): + download_file(settings.arax_blocked_list_url, OUT_PATH, False) + + with open(OUT_PATH, "r") as file: + json_block_list = json.load(file) + synonyms = set(s.lower() for s in json_block_list["synonyms"]) + return set(json_block_list["curies"]), synonyms + +async def rehydrate(kg, retriever_url, logger): + headers = {"Content-Type": "application/json", "Accept": "application/json"} + payload = { + "message": { + "knowledge_graph": kg + }, + "parameters": { + "rehydrate": True, + "tier": 0 + } + } + + try: + async with httpx.AsyncClient(timeout=30.0) as client: + res = await client.post( + retriever_url.replace("query", "rehydrate"), headers=headers, json=payload + ) + res.raise_for_status() + return res.json()["message"]["knowledge_graph"] + + except httpx.HTTPStatusError as http_err: + logger.error(f"HTTP error occurred: {http_err}") + if res.text: + logger.error(f"Error details: {res.text}") + raise http_err + except httpx.ConnectError as conn_err: + logger.error(f"Connection error occurred: {conn_err}") + raise conn_err + except httpx.TimeoutException as timeout_err: + logger.error(f"Timeout error occurred: {timeout_err}") + raise timeout_err + except httpx.RequestError as req_err: + logger.error(f"An unexpected error occurred: {req_err}") + raise req_err + except json.JSONDecodeError: + logger.error("Failed to parse the response as JSON.") + logger.error(f"Raw response: {res.text}") + raise + except Exception as e: + logger.error(f"An unexpected error occurred: {e}") + raise e + +def execute_pathfinding_sync( + pinned_node_ids, pinned_node_keys, intermediate_categories, logger +): + + blocked_curies, blocked_synonyms = get_blocked_list() + + curie_ngd_path, node_degree_path = arax_pathfinder_sqlite_paths() + pathfinder_instance = Pathfinder( + f"retriever:{settings.sync_kg_retrieval_url}", + f"sqlite:{curie_ngd_path}", + f"sqlite:{node_degree_path}", + blocked_curies, + blocked_synonyms, + logger, + ) + + biolink_cache_dir = "/tmp/biolink" + Path(biolink_cache_dir).mkdir(parents=True, exist_ok=True) + biolink_helper = BiolinkHelper(settings.arax_biolink_version, biolink_cache_dir) + descendants = set(biolink_helper.get_descendants(intermediate_categories[0])) + + start = time.perf_counter() + logger.info("Starting pathfinder.get_paths() in worker thread") + + result, aux_graphs, knowledge_graph = pathfinder_instance.get_paths( + pinned_node_ids[0], + pinned_node_ids[1], + pinned_node_keys[0], + pinned_node_keys[1], + NUM_TOTAL_HOPS, + MAX_HOPS_TO_EXPLORE, + MAX_PATHFINDER_PATHS, + PRUNE_TOP_K, + NODE_DEGREE_THRESHOLD, + descendants, + ) + + elapsed = time.perf_counter() - start + logger.info(f"pathfinder.get_paths() finished in {elapsed:.3f} seconds") + + return result, aux_graphs, knowledge_graph + + +async def pathfinder(task, logger: logging.Logger): + start = time.time() + query_id = task[1]["query_id"] + response_id = task[1]["response_id"] + message = await get_message(query_id, logger) + parameters = message.get("parameters") or {} + parameters["timeout"] = parameters.get("timeout", settings.lookup_timeout) + parameters["tiers"] = parameters.get("tiers") or [0] + message["parameters"] = parameters + + qgraph = message["message"]["query_graph"] + pinned_node_keys = [] + pinned_node_ids = [] + for node_key, node in qgraph["nodes"].items(): + pinned_node_keys.append(node_key) + if node.get("ids", None) is not None: + pinned_node_ids.append(node["ids"][0]) + if len(set(pinned_node_ids)) != 2: + logger.error("Pathfinder queries require two pinned nodes.") + return message, 500 + + intermediate_categories = [] + path_key = next(iter(qgraph["paths"].keys())) + qpath = qgraph["paths"][path_key] + if ( + qpath.get("constraints", None) is not None + and len(qpath.get("constraints", [])) > 0 + ): + constraints = qpath["constraints"] + if len(constraints) > 1: + logger.error("Pathfinder queries do not support multiple constraints.") + return message, 500 + if len(constraints) > 0: + intermediate_categories = ( + constraints[0].get("intermediate_categories", None) or [] + ) + if len(intermediate_categories) > 1: + logger.error( + "Pathfinder queries do not support multiple intermediate categories" + ) + return message, 500 + else: + intermediate_categories = ["biolink:NamedThing"] + + try: + result, aux_graphs, knowledge_graph = await asyncio.to_thread( + execute_pathfinding_sync, + pinned_node_ids, + pinned_node_keys, + intermediate_categories, + logger, + ) + logger.info(f"Rehydrating knowledge graph with retriever") + knowledge_graph = await rehydrate(knowledge_graph, settings.kg_rehydrate_url, logger) + + res = [] + if result is not None: + res.append( + { + "id": result["id"], + "analyses": result["analyses"], + "node_bindings": result["node_bindings"], + "essence": "result", + } + ) + if aux_graphs is None: + aux_graphs = {} + if knowledge_graph is None: + knowledge_graph = {} + message["message"]["knowledge_graph"] = knowledge_graph + message["message"]["auxiliary_graphs"] = aux_graphs + message["message"]["results"] = res + + message = add_shepherd_arax_to_edge_sources(message) + + await save_message(response_id, message, logger) + except Exception as e: + logger.error( + f"PathFinder failed to find paths between {pinned_node_keys[0]} and {pinned_node_keys[1]}. " + f"Error message is: {e}" + ) + message = {"status": "error", "error": str(e)} + await save_message(response_id, message, logger) + + logger.info(f"Task took {time.time() - start}") + + +async def process_task(task, parent_ctx, logger: logging.Logger, limiter): + """Process a given task and ACK in redis.""" + start = time.time() + span = tracer.start_span(STREAM, context=parent_ctx) + try: + await pathfinder(task, logger) + # Always wrap up the task to ACK it in the broker + try: + await wrap_up_task(STREAM, GROUP, task, logger) + except Exception as e: + logger.error(f"Task {task[0]}: Failed to wrap up task: {e}") + except asyncio.CancelledError: + logger.warning(f"Task {task[0]} was cancelled") + except Exception as e: + logger.error(f"Task {task[0]} failed with unhandled error: {e}", exc_info=True) + await handle_task_failure(STREAM, GROUP, task, logger) + finally: + span.end() + limiter.release() + logger.info(f"Finished task {task[0]} in {time.time() - start}") + + +async def poll_for_tasks(): + """On initialization, poll indefinitely for available tasks.""" + # Ensure the two sqlite databases exist before any task tries to open them + # (a first-run local `docker compose up` starts with the volume-mounted + # directory empty). No-op once present or when no scp source is configured + # (e.g. production, where the data is mounted out of band). + ensure_arax_pathfinder_dbs(logging.getLogger(STREAM)) + while True: + try: + async for task, parent_ctx, logger, limiter in get_tasks( + STREAM, GROUP, CONSUMER, TASK_LIMIT + ): + asyncio.create_task(process_task(task, parent_ctx, logger, limiter)) + except asyncio.CancelledError: + logging.info("Poll loop cancelled, shutting down.") + except Exception as e: + logging.error(f"Error in task polling loop: {e}", exc_info=True) + await asyncio.sleep(5) # back off before retrying + + +if __name__ == "__main__": + asyncio.run(poll_for_tasks())