Arax pathfinder

.github/workflows/release.yml

@@ -49,6 +49,7 @@ jobs:
           - aragorn_pathfinder
           - aragorn_score
           - arax
+          - arax_pathfinder
           - arax_rank
           - bte
           - bte_lookup

.gitignore

@@ -5,6 +5,7 @@ debug/
 gandalf_mmap/
 omnicorp_lmdb/
 pathfinder_embeddings/
+arax_pathfinder_dbs/
 
 # Project specific
 scripts/

README.md

@@ -14,33 +14,45 @@ 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_<tier-version>.sqlite`, `tier0-info-for-overlay_v1.0_<tier-version>.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** are served as plain files over HTTPS. The
+filenames embed a Knowledge Graph version that changes periodically, so only one variable needs updating
+when a new Knowledge Graph 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:/home/nru/.ssh:ro` in 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
 

compose.test.yml

@@ -64,6 +64,9 @@ services:
   arax:
     cpus: 1
     mem_limit: 3g
+  arax_pathfinder:
+    cpus: 6
+    mem_limit: 22g
   arax_rank:
     cpus: 4
     mem_limit: 10g

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,25 @@ 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 downloads its two sqlite dbs from
+      # kg2webhost.rtx.ai on startup, no credentials needed. Set
+      # ARAX_PATHFINDER_TIER_VERSION in your .env if you need a tier other than
+      # the default (see README "Worker data (LMDB / sqlite) downloads").
+      - ./arax_pathfinder_dbs:/app/arax_pathfinder_dbs
 
   arax_rank:
     container_name: arax_rank

scripts/test_pathfinder.py

@@ -0,0 +1,140 @@
+import asyncio
+import json
+import time
+from datetime import datetime
+from pathlib import Path
+
+import httpx
+
+target_urls = {
+    "aragorn-ci": "https://shepherd.ci.transltr.io/aragorn",
+    "arax-ci": "https://shepherd.ci.transltr.io/arax",
+    "aragorn-dev": "https://shepherd.renci.org/aragorn",
+    "arax-dev": "https://shepherd.renci.org/arax",
+    "aragorn-local": "http://localhost:5439/aragorn",
+    "arax-local": "http://localhost:5439/arax",
+    "bte-local": "http://localhost:5439/bte",
+}
+
+RESPONSES_DIR = "responses"
+
+
+def generate_query(curie1: str, curie2: str) -> dict:
+    """Given a curie, return a TRAPI message."""
+    parameters = {
+        # "timeout": 300,
+        # "tiers": [0],
+    }
+    return {
+      "message": {
+          "query_graph": {
+              "nodes": {
+                  "on": {
+                      "constraints": [],
+                      "ids": [
+                          curie1
+                      ],
+                  },
+                  "sn": {
+                      "constraints": [],
+                      "ids": [
+                          curie2
+                      ],
+                  }
+              },
+              "paths": {
+                  "p0": {
+                      "object": "on",
+                      "subject": "sn",
+                      "predicates": [
+                          "biolink:related_to",
+                      ]
+                  }
+              }
+          }
+      },
+      "parameters": parameters,
+      "log_level": "DEBUG",
+  }
+
+
+async def single_lookup(curies: tuple[str, str], target: str):
+    """Run a single query lookup synchronously."""
+    query = generate_query(curies[0], curies[1])
+    start_time = datetime.now()
+    try:
+        async with httpx.AsyncClient(timeout=600000) as client:
+            response = await client.post(
+                f"{target_urls[target]}/query",
+                json=query,
+            )
+            response.raise_for_status()
+            response_json = response.json()
+            results = (response_json.get("message") or {}).get("results") or []
+            num_results = len(results)
+            assert num_results == 1
+            num_analyses = len(results[0]["analyses"])
+    except Exception as e:
+        num_results = 0
+        num_analyses = 0
+        response_json = {
+            "Error": str(e),
+        }
+
+    stop_time = datetime.now()
+    print(f"{curies[0]}->{curies[1]} took {stop_time - start_time} seconds and gave {num_analyses} results")
+    out_dir = Path(RESPONSES_DIR) / "pathfinder" / target
+    out_dir.mkdir(parents=True, exist_ok=True)
+    response_path = out_dir / f"{('_').join(curies[0].split(':'))}_{('_').join(curies[1].split(':'))}_response.json"
+    with response_path.open("w", encoding="utf-8") as f:
+        json.dump(response_json, f, indent=2)
+
+
+query_list = [
+    ('MONDO:0021095', 'MONDO:0005105'),
+    ('CHEBI:9139', 'MONDO:0004975'),
+    ('CHEBI:5118', 'MONDO:0100233'),
+    ('MONDO:0005180', 'MONDO:0005105'),
+    ('MONDO:0019632', 'MONDO:0005340'),
+    ('CHEBI:27881', 'NCBIGene:2739'),
+    ('CHEBI:45783', 'MONDO:0004979'),  # Imatinib -> Asthma
+    ('GO:0006914', 'MONDO:0005265'),
+    ('NCBIGene:3458', 'CHEBI:16828'),
+    ('MONDO:0005532', 'MONDO:0005180'),
+    ('CHEBI:15647', 'UNII:31YO63LBSN'),
+    ('CHEBI:28364', 'MONDO:0005311'),
+    ('NCBIGene:3458', 'MONDO:0100096'),
+    ('NCBIGene:27240', 'MONDO:0100096'),
+    ('CHEBI:3750', 'MONDO:0013209'),
+    ('CHEBI:83766', 'MONDO:0008170'),
+    ('CHEBI:45783', 'MONDO:0004784'),
+    ('UNII:7SE5582Q2P', 'MONDO:0007037'),
+    ('MONDO:0005011', 'MONDO:0005180'),
+    ('CHEBI:15365', 'MONDO:0005575'),
+    ('CHEBI:50924', 'MONDO:0007256'),
+    ('CHEBI:45713', 'NCBIGene:2739'),
+    ('NCBIGene:54716', 'MONDO:0100096'),
+    ('CHEBI:7465', 'MONDO:0008218'),
+    # ('CHEBI:10033', 'MONDO:0004992'),  # Warfarin -> Cancer, DON'T RUN
+]
+
+
+async def main():
+    """Run the given query and time it."""
+    targets = ["arax-local"]
+    runs_per_target = 1
+
+    start = time.time()
+    queries = []
+    for curies in query_list:
+        queries.extend([
+            single_lookup(curies, target)
+            for target in targets
+            for _ in range(runs_per_target)
+        ])
+    await asyncio.gather(*queries)
+    print(f"\nAll queries took {time.time() - start:.2f} seconds")
+
+
+if __name__ == "__main__":
+    asyncio.run(main())

shepherd_utils/config.py

@@ -84,7 +84,23 @@ 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"
+    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_base_url: str = "https://kg2webhost.rtx.ai/tier0"
+    # End of ARAX configs
 
     pathfinder_redis_host: str = "host.docker.internal"
     pathfinder_redis_port: int = 6383