diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index e38c517..307f9eb 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -49,6 +49,7 @@ jobs: - aragorn_pathfinder - aragorn_score - arax + - arax_pathfinder - arax_rank - bte - bte_lookup diff --git a/.gitignore b/.gitignore index 4fb74c6..cf04e5c 100644 --- a/.gitignore +++ b/.gitignore @@ -5,6 +5,7 @@ debug/ gandalf_mmap/ omnicorp_lmdb/ pathfinder_embeddings/ +arax_pathfinder_dbs/ # Project specific scripts/ diff --git a/README.md b/README.md index f601a01..4921dfc 100644 --- a/README.md +++ b/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_.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** 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 diff --git a/compose.test.yml b/compose.test.yml index 0bb7690..1db188b 100644 --- a/compose.test.yml +++ b/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 diff --git a/compose.yml b/compose.yml index 8847600..0e1a39a 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,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 diff --git a/scripts/test_pathfinder.py b/scripts/test_pathfinder.py new file mode 100644 index 0000000..1e574b0 --- /dev/null +++ b/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()) diff --git a/shepherd_utils/config.py b/shepherd_utils/config.py index 67f4dce..7f7e65c 100644 --- a/shepherd_utils/config.py +++ b/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 diff --git a/shepherd_utils/data_download.py b/shepherd_utils/data_download.py index 6df753f..610a948 100644 --- a/shepherd_utils/data_download.py +++ b/shepherd_utils/data_download.py @@ -1,32 +1,41 @@ -"""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 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). +* **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). + +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 tarfile import tempfile +import urllib.error import urllib.request -from typing import List, Optional +from typing import Dict, List, Optional, Tuple from shepherd_utils.config import settings @@ -43,29 +52,40 @@ def _missing_files(target_dir: str, required_files: List[str]) -> List[str]: def _download(url: str, dest_path: str, logger: logging.Logger) -> None: """Stream ``url`` to ``dest_path``, logging progress periodically.""" logger.info(f"Downloading dataset from {url} ...") - # nosec B310: the URL is operator-configured (an env var), not user input. - with urllib.request.urlopen(url) as resp: # noqa: S310 - header = resp.headers.get("Content-Length") - total = int(header) if header and header.isdigit() else None - read = 0 - step = 50 * 1024 * 1024 # log roughly every 50 MB - next_log = step - with open(dest_path, "wb") as out: - while True: - chunk = resp.read(1024 * 1024) - if not chunk: - break - out.write(chunk) - read += len(chunk) - if read >= next_log: - if total: - logger.info( - f" ... {read / 1e6:.0f}/{total / 1e6:.0f} MB " - f"({100 * read / total:.0f}%)" - ) - else: - logger.info(f" ... {read / 1e6:.0f} MB") - next_log += step + try: + # nosec B310: the URL is operator-configured (an env var), not user input. + with urllib.request.urlopen(url) as resp: # noqa: S310 + header = resp.headers.get("Content-Length") + total = int(header) if header and header.isdigit() else None + read = 0 + step = 50 * 1024 * 1024 # log roughly every 50 MB + next_log = step + with open(dest_path, "wb") as out: + while True: + chunk = resp.read(1024 * 1024) + if not chunk: + break + out.write(chunk) + read += len(chunk) + if read >= next_log: + if total: + logger.info( + f" ... {read / 1e6:.0f}/{total / 1e6:.0f} MB " + f"({100 * read / total:.0f}%)" + ) + else: + logger.info(f" ... {read / 1e6:.0f} MB") + next_log += step + except urllib.error.HTTPError as e: + raise RuntimeError( + f"download failed for {url}: HTTP {e.code} {e.reason}. Confirm the " + f"URL is correct and reachable." + ) from e + except urllib.error.URLError as e: + raise RuntimeError( + f"download failed for {url}: {e.reason}. Confirm the URL is correct " + f"and reachable from inside the container." + ) from e logger.info(f"Download complete: {read / 1e6:.0f} MB") @@ -153,6 +173,87 @@ def ensure_lmdb_dataset( logger.info(f"{name}: dataset ready in {target_dir}.") +def ensure_http_files_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``, + fetching any missing ones directly over HTTP(S) -- one URL per file, no + archive/extract step (contrast with ``ensure_lmdb_dataset``'s single + ``.tar.gz``). + + ``file_sources`` maps the expected local filename to its URL; a filename + whose URL 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: + url = file_sources.get(filename) + if not url: + logger.warning( + f"{name}: {filename} missing from {target_dir} and no URL " + f"configured for it. Set the corresponding *_URL 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: + _download(url, 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 URL) count toward failure -- a + # file with no URL 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}." + ) + if _missing_files(target_dir, required_files): + logger.warning( + f"{name}: dataset partially ready in {target_dir} -- some files have " + f"no URL 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 +288,59 @@ 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 are served as plain files over HTTPS, so each is fetched individually + with a normal GET -- no archive/extract step. Both are expected in the + same directory (see the ``arax_pathfinder`` volume mount in + docker-compose.yml). + + The version tag shows up in the filename (e.g. + ``curie_ngd_v1.0_tier0-20260621.sqlite``) but not in the URL path -- the + ``tier0`` segment in ``arax_pathfinder_sqlite_base_url`` is fixed, not the + tier version. Only the filename templates are filled in from + ``arax_pathfinder_tier_version``. Bumping to a new tier is one env var + change (``ARAX_PATHFINDER_TIER_VERSION``). + """ + curie_ngd_path, node_degree_path = arax_pathfinder_sqlite_paths() + target_dir = settings.arax_pathfinder_dbs_dir + base_url = settings.arax_pathfinder_sqlite_base_url + + curie_ngd_filename = os.path.basename(curie_ngd_path) + node_degree_filename = os.path.basename(node_degree_path) + + ensure_http_files_dataset( + name="arax_pathfinder", + target_dir=target_dir, + file_sources={ + curie_ngd_filename: f"{base_url}/{curie_ngd_filename}", + node_degree_filename: f"{base_url}/{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/shepherd_utils/reclaim.py b/shepherd_utils/reclaim.py index 54121c6..e858878 100644 --- a/shepherd_utils/reclaim.py +++ b/shepherd_utils/reclaim.py @@ -49,6 +49,9 @@ "aragorn.score": 240, "bte.lookup": 240, "example.lookup": 240, + # Pathfinding runs in a process pool bounded by pool_task_timeout_sec + # (300s), so no legitimate task can outlive that; the floor sits just above. + "arax.pathfinder": 360, # Medium-duration workers. "arax.rank": 60, "merge_message": 60, diff --git a/tests/unit/test_arax_pathfinder.py b/tests/unit/test_arax_pathfinder.py new file mode 100644 index 0000000..f572207 --- /dev/null +++ b/tests/unit/test_arax_pathfinder.py @@ -0,0 +1,254 @@ +"""Tests for the arax_pathfinder process-pool entrypoint. + +The worker imports ``pathfinder.Pathfinder`` and ``biolink_helper_pkg``, which +are installed only inside its container (see +``workers/arax_pathfinder/requirements.txt``) and are not in +``test-requirements.txt``. We stub both into ``sys.modules`` before importing +the worker, the same trick ``tests/conftest.py`` uses for ``shepherd_utils.otel``. +""" + +import copy +import logging +import sys +import types +from unittest.mock import MagicMock + +import pytest + +if "pathfinder" not in sys.modules: + _pathfinder_pkg = types.ModuleType("pathfinder") + _pathfinder_mod = types.ModuleType("pathfinder.Pathfinder") + _pathfinder_mod.Pathfinder = MagicMock(name="Pathfinder") + _pathfinder_pkg.Pathfinder = _pathfinder_mod + sys.modules["pathfinder"] = _pathfinder_pkg + sys.modules["pathfinder.Pathfinder"] = _pathfinder_mod + +if "biolink_helper_pkg" not in sys.modules: + _biolink_mod = types.ModuleType("biolink_helper_pkg") + _biolink_mod.BiolinkHelper = MagicMock(name="BiolinkHelper") + sys.modules["biolink_helper_pkg"] = _biolink_mod + +from workers.arax_pathfinder import worker as pf_worker # noqa: E402 + +LOGGER = logging.getLogger(__name__) + +QUERY = { + "message": { + "query_graph": { + "nodes": { + "n0": {"ids": ["MONDO:0005148"]}, + "n1": {"ids": ["CHEBI:15365"]}, + }, + "paths": {"p0": {"subject": "n0", "object": "n1", "constraints": []}}, + } + } +} + +PATHS_RESULT = ( + {"id": "r0", "analyses": [{"score": 1.0}], "node_bindings": {"n0": []}}, + {"aux0": {"edges": ["e0"]}}, + { + "nodes": {"MONDO:0005148": {}}, + "edges": {"e0": {"predicate": "biolink:related_to"}}, + }, +) + + +def _patch_query(mocker, query=None): + return mocker.patch( + "workers.arax_pathfinder.worker.get_message_sync", + return_value=copy.deepcopy(query if query is not None else QUERY), + ) + + +def test_pathfinder_task_searches_rehydrates_and_saves(mocker): + """The process-pool entrypoint reads by id, searches, and writes back. + + Only the two ids cross into the child: the message is loaded with + ``get_message_sync``, assembled, and persisted with ``save_message_sync`` -- + the knowledge graph never has to be pickled back to the parent. + """ + _patch_query(mocker) + search = mocker.patch( + "workers.arax_pathfinder.worker.execute_pathfinding", + return_value=copy.deepcopy(PATHS_RESULT), + ) + rehydrated_kg = {"nodes": {}, "edges": {"e0": {"predicate": "biolink:related_to"}}} + rehydrate = mocker.patch( + "workers.arax_pathfinder.worker.rehydrate", return_value=rehydrated_kg + ) + save = mocker.patch("workers.arax_pathfinder.worker.save_message_sync") + + pf_worker.arax_pathfinder_task("query-1", "resp-1", LOGGER) + + search.assert_called_once() + rehydrate.assert_called_once() + save.assert_called_once() + saved_id, message = save.call_args.args + assert saved_id == "resp-1" + assert message["message"]["knowledge_graph"] is rehydrated_kg + assert message["message"]["auxiliary_graphs"] == {"aux0": {"edges": ["e0"]}} + assert len(message["message"]["results"]) == 1 + assert message["message"]["results"][0]["essence"] == "result" + # Provenance is injected before saving. + assert message["message"]["knowledge_graph"]["edges"]["e0"]["sources"] == [ + { + "resource_id": "infores:shepherd-arax", + "resource_role": "aggregator_knowledge_source", + "source_record_urls": None, + "upstream_resource_ids": ["infores:arax"], + } + ] + # Defaults are filled in on the message that gets saved. + assert message["parameters"]["tiers"] == [0] + + +def test_pathfinder_task_saves_empty_graphs_when_no_paths_found(mocker): + """A search that finds nothing still writes a well-formed TRAPI message.""" + _patch_query(mocker) + mocker.patch( + "workers.arax_pathfinder.worker.execute_pathfinding", + return_value=(None, None, None), + ) + mocker.patch("workers.arax_pathfinder.worker.rehydrate", return_value=None) + save = mocker.patch("workers.arax_pathfinder.worker.save_message_sync") + + pf_worker.arax_pathfinder_task("query-2", "resp-2", LOGGER) + + _, message = save.call_args.args + assert message["message"]["results"] == [] + assert message["message"]["auxiliary_graphs"] == {} + assert message["message"]["knowledge_graph"] == {} + + +@pytest.mark.parametrize( + "qgraph, expected", + [ + pytest.param( + { + "nodes": {"n0": {"ids": ["MONDO:0005148"]}, "n1": {}}, + "paths": {"p0": {"constraints": []}}, + }, + "two pinned nodes", + id="one_pinned_node", + ), + pytest.param( + { + "nodes": { + "n0": {"ids": ["MONDO:0005148"]}, + "n1": {"ids": ["CHEBI:15365"]}, + }, + "paths": { + "p0": { + "constraints": [ + {"intermediate_categories": ["biolink:Gene"]}, + {"intermediate_categories": ["biolink:Drug"]}, + ] + } + }, + }, + "multiple constraints", + id="multiple_constraints", + ), + pytest.param( + { + "nodes": { + "n0": {"ids": ["MONDO:0005148"]}, + "n1": {"ids": ["CHEBI:15365"]}, + }, + "paths": { + "p0": { + "constraints": [ + { + "intermediate_categories": [ + "biolink:Gene", + "biolink:Drug", + ] + } + ] + } + }, + }, + "multiple intermediate categories", + id="multiple_intermediate_categories", + ), + ], +) +def test_unanswerable_query_graph_raises(mocker, qgraph, expected): + """Validation failures raise so run_task_lifecycle can fail the query. + + They used to ``return message, 500``, a value the caller discarded -- so the + task was wrapped up as a success and no message was ever written to + ``response_id``, leaving the next worker to ``KeyError`` on it. Raising + routes the query to ``finish_query`` with an ERROR status instead. + """ + _patch_query(mocker, {"message": {"query_graph": qgraph}}) + save = mocker.patch("workers.arax_pathfinder.worker.save_message_sync") + search = mocker.patch("workers.arax_pathfinder.worker.execute_pathfinding") + + with pytest.raises(ValueError, match=expected): + pf_worker.arax_pathfinder_task("query-3", "resp-3", LOGGER) + + search.assert_not_called() + save.assert_not_called() + + +def test_search_failure_propagates_instead_of_saving_error_blob(mocker): + """A failed search raises rather than persisting a non-TRAPI error blob. + + The old code caught everything and saved ``{"status": "error", ...}`` to the + response id, then reported success -- so that blob flowed on through the + workflow. Now the exception reaches ``run_task_lifecycle``. + """ + _patch_query(mocker) + mocker.patch( + "workers.arax_pathfinder.worker.execute_pathfinding", + side_effect=RuntimeError("sqlite is on fire"), + ) + save = mocker.patch("workers.arax_pathfinder.worker.save_message_sync") + + with pytest.raises(RuntimeError, match="sqlite is on fire"): + pf_worker.arax_pathfinder_task("query-4", "resp-4", LOGGER) + + save.assert_not_called() + + +def test_rehydrate_failure_propagates(mocker): + """A rehydrate failure fails the task too, rather than saving a partial KG.""" + _patch_query(mocker) + mocker.patch( + "workers.arax_pathfinder.worker.execute_pathfinding", + return_value=copy.deepcopy(PATHS_RESULT), + ) + mocker.patch( + "workers.arax_pathfinder.worker.rehydrate", + side_effect=RuntimeError("retriever unreachable"), + ) + save = mocker.patch("workers.arax_pathfinder.worker.save_message_sync") + + with pytest.raises(RuntimeError, match="retriever unreachable"): + pf_worker.arax_pathfinder_task("query-5", "resp-5", LOGGER) + + save.assert_not_called() + + +def test_descendants_are_memoized_per_child(mocker): + """The Biolink model is built once per pool child, not once per task. + + ``BiolinkHelper`` construction parses the model, and the old code did it + inside every task alongside a fresh blocked-list HTTP fetch. + """ + mocker.patch.object(pf_worker, "_biolink_helper", None) + mocker.patch.object(pf_worker, "_descendants_cache", {}) + helper = MagicMock() + helper.get_descendants.return_value = ["biolink:Gene", "biolink:Protein"] + helper_cls = mocker.patch( + "workers.arax_pathfinder.worker.BiolinkHelper", return_value=helper + ) + + first = pf_worker.get_descendants("biolink:NamedThing") + second = pf_worker.get_descendants("biolink:NamedThing") + + assert first == second == {"biolink:Gene", "biolink:Protein"} + helper_cls.assert_called_once() + helper.get_descendants.assert_called_once_with("biolink:NamedThing") diff --git a/tests/unit/test_inject_shepherd_arax_provenance.py b/tests/unit/test_inject_shepherd_arax_provenance.py index 6c50dc0..0fbdf68 100644 --- a/tests/unit/test_inject_shepherd_arax_provenance.py +++ b/tests/unit/test_inject_shepherd_arax_provenance.py @@ -1,4 +1,4 @@ -"""Tests for ``workers.arax.inject_shepherd_arax_provenance``. +"""Tests for ``shepherd_utils.inject_shepherd_arax_provenance``. The shepherd-arax injector tags every kgraph edge with an aggregator ``infores:shepherd-arax`` source so downstream consumers can attribute @@ -12,7 +12,7 @@ import copy -from workers.arax.inject_shepherd_arax_provenance import ( +from shepherd_utils.inject_shepherd_arax_provenance import ( SHEPHERD_ARAX_SOURCE, add_shepherd_arax_to_edge_sources, ) diff --git a/tests/unit/test_worker_dispatch_concurrency.py b/tests/unit/test_worker_dispatch_concurrency.py index 5dbead9..faec548 100644 --- a/tests/unit/test_worker_dispatch_concurrency.py +++ b/tests/unit/test_worker_dispatch_concurrency.py @@ -24,6 +24,7 @@ WORKER_FILES = [ "workers/arax_rank/worker.py", "workers/aragorn_score/worker.py", + "workers/arax_pathfinder/worker.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..890d662 --- /dev/null +++ b/workers/arax_pathfinder/Dockerfile @@ -0,0 +1,39 @@ +# 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 +# Enable faulthandler from interpreter startup so a SIGABRT dumps every thread's +# traceback -- including during a spawned pool child's import phase, before the +# pool initializer's faulthandler.enable() would run. Lets ProcessPoolManager's +# SIGABRT-on-timeout capture a stuck child's stack even if it hangs in startup. +ENV PYTHONFAULTHANDLER=1 + +# 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..e1e8e9f --- /dev/null +++ b/workers/arax_pathfinder/worker.py @@ -0,0 +1,366 @@ +"""Arax ARA Pathfinder module.""" + +import asyncio +import json +import logging +import time +import uuid +from pathlib import Path + +import httpx +from biolink_helper_pkg import BiolinkHelper +from pathfinder.Pathfinder import Pathfinder + +from shepherd_utils.config import settings +from shepherd_utils.cpu import resolve_pool_workers +from shepherd_utils.data_download import ( + arax_pathfinder_sqlite_paths, + ensure_arax_pathfinder_dbs, + ensure_http_files_dataset, +) +from shepherd_utils.db import ( + get_message_sync, + save_message_sync, +) +from shepherd_utils.inject_shepherd_arax_provenance import ( + add_shepherd_arax_to_edge_sources, +) +from shepherd_utils.otel import setup_tracer +from shepherd_utils.process_pool import ProcessPoolManager +from shepherd_utils.shared import get_tasks, run_task_lifecycle + +# 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 = 10 +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 + +# The ARAX blocked-concept list, fetched once at worker startup (see +# poll_for_tasks) and read by each pool child. Kept in the working directory +# (/app in the image) so it needs no extra volume mount. +BLOCKED_LIST_DIR = "." +BLOCKED_LIST_FILENAME = "general_concepts.json" +BLOCKED_LIST_PATH = Path(BLOCKED_LIST_DIR) / BLOCKED_LIST_FILENAME + +BIOLINK_CACHE_DIR = "/tmp/biolink" + +REHYDRATE_TIMEOUT_SEC = 30.0 + +# Per-child caches. Children are spawned once and reused for up to +# ``settings.pool_max_tasks_per_child`` tasks, so the blocked list and the +# Biolink model are parsed once per child instead of once per task (previously +# every task re-read the JSON and rebuilt the BiolinkHelper, and re-fetched the +# blocked list over HTTP). Only ``Pathfinder`` is still built per task, because +# it takes the task's logger. +_blocked_list_cache = None +_biolink_helper = None +_descendants_cache: dict = {} + + +def ensure_blocked_list(logger: logging.Logger) -> None: + """Fetch the ARAX blocked-concept list if it isn't on disk yet. + + Uses the shared downloader so the file lands via a temp file + atomic + rename; the previous per-task ``requests.get`` wrote the destination + directly, so concurrent tasks could race on a half-written file. + Idempotent, so it is safe to call at startup and again lazily in a child. + """ + ensure_http_files_dataset( + name="arax_blocked_list", + target_dir=BLOCKED_LIST_DIR, + file_sources={BLOCKED_LIST_FILENAME: settings.arax_blocked_list_url}, + logger=logger, + ) + + +def get_blocked_list(logger: logging.Logger): + """``(blocked_curies, blocked_synonyms)``, parsed once per pool child.""" + global _blocked_list_cache + if _blocked_list_cache is None: + if not BLOCKED_LIST_PATH.exists(): + # Startup fetch failed or this child outlived a wiped working dir. + ensure_blocked_list(logger) + with open(BLOCKED_LIST_PATH, "r") as file: + json_block_list = json.load(file) + synonyms = set(s.lower() for s in json_block_list["synonyms"]) + _blocked_list_cache = (set(json_block_list["curies"]), synonyms) + return _blocked_list_cache + + +def get_descendants(category: str): + """Biolink descendants of ``category``, memoized per pool child.""" + global _biolink_helper + if category not in _descendants_cache: + if _biolink_helper is None: + Path(BIOLINK_CACHE_DIR).mkdir(parents=True, exist_ok=True) + _biolink_helper = BiolinkHelper( + settings.arax_biolink_version, BIOLINK_CACHE_DIR + ) + _descendants_cache[category] = set(_biolink_helper.get_descendants(category)) + return _descendants_cache[category] + + +def rehydrate(kg, rehydrate_url, logger): + """POST the knowledge graph to the retriever and return the rehydrated one. + + Synchronous because it runs inside the process-pool child alongside the + pathfinding, so the (potentially very large) knowledge graph never has to + be pickled back to the parent or re-encoded on its event loop. + """ + headers = {"Content-Type": "application/json", "Accept": "application/json"} + payload = { + "message": {"knowledge_graph": kg}, + "parameters": {"rehydrate": True, "tier": 0}, + } + + try: + with httpx.Client(timeout=REHYDRATE_TIMEOUT_SEC) as client: + res = client.post(rehydrate_url, 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 http_err.response.text: + logger.error(f"Error details: {http_err.response.text}") + raise + except httpx.ConnectError as conn_err: + logger.error(f"Connection error occurred: {conn_err}") + raise + except httpx.TimeoutException as timeout_err: + logger.error(f"Timeout error occurred: {timeout_err}") + raise + except httpx.RequestError as req_err: + logger.error(f"An unexpected error occurred: {req_err}") + raise + except json.JSONDecodeError: + logger.error("Failed to parse the rehydrate response as JSON.") + raise + except Exception as e: + logger.error(f"An unexpected error occurred: {e}") + raise e + + +def parse_query_graph(qgraph): + """Pull the pinned nodes and intermediate category out of the query graph. + + Raises ``ValueError`` on anything Pathfinder can't answer. Raising (rather + than returning a status code, which the old caller discarded) lets + ``run_task_lifecycle`` record the failure on the span and route the query to + ``finish_query`` with an ERROR status, instead of continuing the workflow + with no response message ever written. + """ + 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: + raise ValueError("Pathfinder queries require two pinned nodes.") + + 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: + raise ValueError("Pathfinder queries do not support multiple constraints.") + if len(constraints) > 0: + intermediate_categories = ( + constraints[0].get("intermediate_categories", None) or [] + ) + if len(intermediate_categories) > 1: + raise ValueError( + "Pathfinder queries do not support multiple intermediate categories" + ) + else: + intermediate_categories = ["biolink:NamedThing"] + + return pinned_node_keys, pinned_node_ids, intermediate_categories + + +def execute_pathfinding( + pinned_node_ids, pinned_node_keys, intermediate_categories, logger +): + blocked_curies, blocked_synonyms = get_blocked_list(logger) + + 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, + ) + + descendants = get_descendants(intermediate_categories[0]) + + start = time.perf_counter() + logger.info("Starting pathfinder.get_paths()") + + 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 + + +def arax_pathfinder_task( + query_id: str, response_id: str, logger: logging.Logger +) -> None: + """Process-pool entrypoint: load, search, rehydrate, and save in the child. + + Only the two small ids cross the process-pool boundary; the (potentially + very large) message is read from Redis, searched over, and written back + inside the child. That keeps the payload off the parent's heap and, more + importantly, keeps the graph search off the parent's event loop -- it used + to run via ``asyncio.to_thread``, where the GIL meant a handful of + concurrent searches could starve the heartbeat past HEARTBEAT_TTL_SEC and + get a live worker's tasks reclaimed out from under it (matching + aragorn_score / arax_rank). + """ + start = time.time() + message = get_message_sync(query_id) + parameters = message.get("parameters") or {} + parameters["timeout"] = parameters.get("timeout", settings.lookup_timeout) + parameters["tiers"] = parameters.get("tiers") or [0] + message["parameters"] = parameters + + pinned_node_keys, pinned_node_ids, intermediate_categories = parse_query_graph( + message["message"]["query_graph"] + ) + + try: + result, aux_graphs, knowledge_graph = execute_pathfinding( + pinned_node_ids, + pinned_node_keys, + intermediate_categories, + logger, + ) + logger.info("Rehydrating knowledge graph with retriever") + knowledge_graph = rehydrate(knowledge_graph, settings.kg_rehydrate_url, logger) + except Exception as e: + # Let the failure reach run_task_lifecycle, which records it on the span + # and routes the query to finish_query with an ERROR status. Previously + # this saved a non-TRAPI {"status": "error"} blob and reported success, + # so the bogus payload flowed on down the workflow. + logger.error( + f"PathFinder failed to find paths between {pinned_node_keys[0]} and " + f"{pinned_node_keys[1]}. Error message is: {e}" + ) + raise + + 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) + + save_message_sync(response_id, message) + logger.info(f"Task took {time.time() - start}") + + +async def process_task(task, parent_ctx, logger, limiter, loop, pool): + """Process a given task and ACK in redis. + + Pathfinding is CPU-bound, so it is dispatched to a process pool while the + span, wrap-up, and error handling are shared with every worker. Only the + ids are handed to the child; the message load/save happen there (see + ``arax_pathfinder_task``) so the payload never crosses the process boundary. + + Dispatch goes through ``pool`` (a ProcessPoolManager) so a child dying on an + oversized message replaces the pool instead of poisoning it for good, and a + search that runs away is time-bounded rather than holding its slot forever. + """ + + async def _run(task, logger): + await pool.run( + loop, + arax_pathfinder_task, + task[1]["query_id"], + task[1]["response_id"], + logger, + ) + + await run_task_lifecycle(STREAM, GROUP, task, parent_ctx, logger, limiter, _run) + + +async def poll_for_tasks(): + """On initialization, poll indefinitely for available tasks.""" + startup_logger = logging.getLogger(STREAM) + # 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(startup_logger) + # Fetch the blocked-concept list once here rather than per task, so the pool + # children only ever read it. + ensure_blocked_list(startup_logger) + loop = asyncio.get_running_loop() + # Size the pool by the pod's actual CPU allocation (cgroup limit), not + # os.cpu_count() -- see aragorn_score.poll_for_tasks. Each child loads a full + # message, so this also bounds peak memory. POOL_MAX_WORKERS overrides. + max_workers = resolve_pool_workers(TASK_LIMIT, startup_logger) + logging.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="arax.pathfinder process pool", + task_timeout=settings.pool_task_timeout_sec, + ) + while True: + try: + async for task, parent_ctx, logger, limiter in get_tasks( + STREAM, GROUP, CONSUMER, max_workers + ): + asyncio.create_task( + process_task(task, parent_ctx, logger, limiter, loop, pool) + ) + 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())