Arax pathfinder

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_<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** 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
 

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

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

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,
+    )