diff --git a/.env b/.env new file mode 100644 index 00000000..28994cc6 --- /dev/null +++ b/.env @@ -0,0 +1,14 @@ + +KGX_DATA_SETS="bdc-studies-kgx:v1.0" +INPUT_DATA_SETS="heal-mds-studies:v1.0" + +LAKEFS_ACCESS_KEY="" +LAKEFS_SECRET_KEY="" +LAKEFS_REPO="" +LAKEFS_BRANCH="" +LAKEFS_URL="https://lakefs.apps.renci.org" + +BIOMEGATRON_URL="https://med-nemo.apps.renci.org/annotate" +SAPBERT_URL="https://sap-qdrant.apps.renci.org/annotate" +NODE_NORM_URL="https://nodenormalization-sri.renci.org/get_normalized_nodes?conflate=false&description=true&curie=" +NAME_RES_URL="https://name-resolution-sri.renci.org/reverse_lookup" \ No newline at end of file diff --git a/.github/workflows/build-push-dev-image.yml b/.github/workflows/build-push-dev-image.yml new file mode 100644 index 00000000..ddc949e5 --- /dev/null +++ b/.github/workflows/build-push-dev-image.yml @@ -0,0 +1,86 @@ +# Workflow responsible for the +# development release processes. +# +name: Build-Push-Dev-Image +on: + push: + branches: + - develop + paths-ignore: + - README.md + - .old_cicd/* + - .github/* + - .github/workflows/* + - LICENSE + - .gitignore + - .dockerignore + - .githooks + # Do not build another image on a pull request. + # Any push to develop will trigger a new build however. + pull_request: + branches-ignore: + - '*' + +jobs: + build-push-dev-image: + runs-on: ubuntu-latest + steps: + + - name: Checkout Code + uses: actions/checkout@v3 + with: + ref: ${{ github.head_ref }} + # fetch-depth: 0 means, get all branches and commits + fetch-depth: 0 + + - name: Set short git commit SHA + id: vars + run: | + echo "short_sha=$(git rev-parse --short ${{ github.sha }})" >> $GITHUB_OUTPUT + # https://github.blog/changelog/2022-10-11-github-actions-deprecating-save-state-and-set-output-commands/ + + - name: Confirm git commit SHA output + run: echo ${{ steps.vars.outputs.short_sha }} + + # Docker Buildx is important to caching in the Build And Push Container + # step + # https://github.com/marketplace/actions/build-and-push-docker-images + + - name: Set up Docker Buildx + uses: docker/setup-buildx-action@v3 + with: + driver-opts: | + network=host + + - name: Login to DockerHub + uses: docker/login-action@v3 + with: + username: ${{ secrets.DOCKERHUB_USERNAME }} + password: ${{ secrets.DOCKERHUB_TOKEN }} + logout: true + + - name: Login to Container Registry + uses: docker/login-action@v3 + with: + registry: containers.renci.org + username: ${{ secrets.CONTAINERHUB_USERNAME }} + password: ${{ secrets.CONTAINERHUB_TOKEN }} + logout: true + + + # Notes on Cache: + # https://docs.docker.com/build/ci/github-actions/examples/#inline-cache + - name: Build Push Container + uses: docker/build-push-action@v5 + with: + context: . + push: true + # Push to renci-registry and dockerhub here. + # cache comes from dockerhub. + tags: | + ${{ github.repository }}:develop + ${{ github.repository }}:${{ steps.vars.outputs.short_sha }} + containers.renci.org/${{ github.repository }}:develop + containers.renci.org/${{ github.repository }}:${{ steps.vars.outputs.short_sha }} + cache-from: type=registry,ref=${{ github.repository }}:buildcache-dev + cache-to: type=registry,ref=${{ github.repository }}:buildcache-dev,mode=max diff --git a/.github/workflows/build-push-release.yml b/.github/workflows/build-push-release.yml new file mode 100644 index 00000000..07b22d21 --- /dev/null +++ b/.github/workflows/build-push-release.yml @@ -0,0 +1,131 @@ +# Workflow responsible for the +# major release processes. +# + +name: Build-Push-Release +on: + push: + branches: + - master + - main + paths-ignore: + - README.md + - .old_cicd/* + - .github/* + - .github/workflows/* + - LICENSE + - .gitignore + - .dockerignore + - .githooks + tags-ignore: + - '*' +jobs: + build-push-release: + runs-on: ubuntu-latest + steps: + - name: Checkout Code + uses: actions/checkout@v3 + with: + ref: ${{ github.head_ref }} + fetch-depth: 0 + + - name: Set short git commit SHA + id: vars + run: | + echo "short_sha=$(git rev-parse --short ${{ github.sha }})" >> $GITHUB_OUTPUT + # https://github.blog/changelog/2022-10-11-github-actions-deprecating-save-state-and-set-output-commands/ + + - name: Confirm git commit SHA output + run: echo ${{ steps.vars.outputs.short_sha }} + + # https://github.com/marketplace/actions/git-semantic-version + - name: Semver Check + uses: paulhatch/semantic-version@v5.0.3 + id: version + with: + # The prefix to use to identify tags + tag_prefix: "v" + # A string which, if present in a git commit, indicates that a change represents a + # major (breaking) change, supports regular expressions wrapped with '/' + major_pattern: "/breaking:|major:/" + # A string which indicates the flags used by the `major_pattern` regular expression. Supported flags: idgs + major_regexp_flags: "ig" + # Same as above except indicating a minor change, supports regular expressions wrapped with '/' + minor_pattern: "/feat:|feature:|minor:/" + # A string which indicates the flags used by the `minor_pattern` regular expression. Supported flags: idgs + minor_regexp_flags: "ig" + # A string to determine the format of the version output + # version_format: "${major}.${minor}.${patch}-prerelease${increment}" + version_format: "${major}.${minor}.${patch}" + search_commit_body: false + + # Docker Buildx is important to caching in the Build And Push Container + # step + # https://github.com/marketplace/actions/build-and-push-docker-images + - name: Set up Docker Buildx + uses: docker/setup-buildx-action@v3 + with: + driver-opts: | + network=host + + - name: Login to DockerHub + uses: docker/login-action@v3 + with: + username: ${{ secrets.DOCKERHUB_USERNAME }} + password: ${{ secrets.DOCKERHUB_TOKEN }} + logout: true + + - name: Login to Container Registry + uses: docker/login-action@v3 + with: + registry: containers.renci.org + username: ${{ secrets.CONTAINERHUB_USERNAME }} + password: ${{ secrets.CONTAINERHUB_TOKEN }} + logout: true + + # Notes on Cache: + # https://docs.docker.com/build/ci/github-actions/examples/#inline-cache + - name: Build Push Container + uses: docker/build-push-action@v5 + with: + push: true + # Push to renci-registry and dockerhub here. + # cache comes from dockerhub. + tags: | + containers.renci.org/${{ github.repository }}:v${{ steps.version.outputs.version }} + containers.renci.org/${{ github.repository }}:latest + containers.renci.org/${{ github.repository }}:${{ steps.vars.outputs.short_sha }} + ${{ github.repository }}:v${{ steps.version.outputs.version }} + ${{ github.repository }}:latest + ${{ github.repository }}:${{ steps.vars.outputs.short_sha }} + cache-from: type=registry,ref=${{ github.repository }}:buildcache-release + cache-to: type=registry,ref=${{ github.repository }}:buildcache-release,mode=max + +#==========================TAG & RELEASE W/ NOTES ========================= + + # Note: GITHUB_TOKEN is autogenerated feature of github app + # which is auto-enabled when using github actions. + # https://docs.github.com/en/actions/security-guides/automatic-token-authentication + # https://docs.github.com/en/rest/git/tags?apiVersion=2022-11-28#create-a-tag-object + # https://docs.github.com/en/rest/git/refs?apiVersion=2022-11-28#create-a-reference + # This creates a "lightweight" ref tag. + - name: Create Tag for Release + run: | + curl \ + -s --fail -X POST \ + -H "Accept: application/vnd.github+json" \ + -H "Authorization: Bearer ${{ secrets.GITHUB_TOKEN }}" \ + -H "X-GitHub-Api-Version: 2022-11-28" \ + https://api.github.com/repos/${{ github.repository }}/git/refs \ + -d '{"ref":"refs/tags/v${{ steps.version.outputs.version }}","sha":"${{ github.sha }}"}' + +# https://cli.github.com/manual/gh_release_create + - name: Create Release + env: + RELEASE_VERSION: ${{ steps.version.outputs.version }} + GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} + run: | + gh release create ${{ env.RELEASE_VERSION }} \ + -t "${{ env.RELEASE_VERSION }}" \ + --generate-notes \ + --latest \ No newline at end of file diff --git a/.github/workflows/code-checks.yml b/.github/workflows/code-checks.yml new file mode 100644 index 00000000..1d448aa3 --- /dev/null +++ b/.github/workflows/code-checks.yml @@ -0,0 +1,119 @@ +# Workflow responsible for core acceptance testing. +# Tests Currently Run: +# - flake8-linter +# - PYTest +# - Bandit +# For PR Vulnerability Scanning a separate workflow will run. +# The build-push-dev-image and build-push-release workflows +# handle the develop and release image storage respectively. +# +# + +name: Code-Checks +on: + push: + branches-ignore: + - master + - main + - develop + pull_request: + branches: + - develop + - master + - main + types: [opened, synchronize] + paths-ignore: + - README.md + - .old_cicd/* + - .github/* + - .github/workflows/* + - LICENSE + - .gitignore + - .dockerignore + - .githooks + +jobs: + ############################## flake8-linter ############################## + flake8-linter: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v3 + + - name: Set up Python + uses: actions/setup-python@v4 + with: + python-version: "3.12" + + # flake8 parses sources, it never imports them, so this job needs + # neither requirements.txt nor a cached interpreter. Caching + # ${{ env.pythonLocation }} also restores a stale site-packages/pip + # over the runner's own, which breaks pip outright. + - name: Lint with flake8 + run: | + pip install flake8 + flake8 --ignore=E,W dags + # We continue on error here until the code is clean + # flake8 --ignore=E,W --exit-zero . + continue-on-error: true + + ################################### PYTEST ################################### + # pytest: + # runs-on: ubuntu-latest + # steps: + # - uses: actions/checkout@v3 + # - name: Set up Python + # uses: actions/setup-python@v4 + # with: + # python-version: '3.12' + + # - name: Install Requirements + # run: | + # pip install -r requirements.txt + # pip install coverage + # pip install ./tests + + # - name: Test with pytest + # run: | + # make test + ############################## test-image-build ############################## + test-image-build: + runs-on: ubuntu-latest + # if: ${{ github.actor == 'dependabot[bot]' }} + steps: + - uses: actions/checkout@v3 + + - name: Set short git commit SHA + id: vars + run: | + echo "short_sha=$(git rev-parse --short ${{ github.sha }})" >> $GITHUB_OUTPUT + # https://github.blog/changelog/2022-10-11-github-actions-deprecating-save-state-and-set-output-commands/ + - name: Confirm git commit SHA output + run: echo ${{ steps.vars.outputs.short_sha }} + + - name: Set up Docker Buildx + uses: docker/setup-buildx-action@v3 + + - name: Login to DockerHub + uses: docker/login-action@v3 + with: + username: ${{ secrets.DOCKERHUB_USERNAME }} + password: ${{ secrets.DOCKERHUB_TOKEN }} + logout: true + + - name: Parse Github Reference Name + id: branch + run: | + REF=${{ github.ref_name }} + echo "GHR=${REF%/*}" >> $GITHUB_OUTPUT + + # Notes on Cache: + # https://docs.docker.com/build/ci/github-actions/examples/#inline-cache + - name: Build Container + uses: docker/build-push-action@v5 + with: + context: . + push: true + tags: | + ${{ github.repository }}:test_${{ steps.branch.outputs.GHR }} + cache-from: type=registry,ref=${{ github.repository }}:buildcache + cache-to: type=registry,ref=${{ github.repository }}:buildcache,mode=max diff --git a/.github/workflows/trivy-pr-scan.yml b/.github/workflows/trivy-pr-scan.yml new file mode 100644 index 00000000..3f85d19b --- /dev/null +++ b/.github/workflows/trivy-pr-scan.yml @@ -0,0 +1,67 @@ +name: trivy-pr-scan +on: + pull_request: + branches: + - develop + - master + - main + types: [ opened, synchronize ] + paths-ignore: + - README.md + - .old_cicd/* + - .github/* + - .github/workflows/* + - LICENSE + - .gitignore + - .dockerignore + - .githooks + +jobs: + trivy-pr-scan: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v3 + + - name: Set up Docker Buildx + uses: docker/setup-buildx-action@v3 + with: + driver-opts: | + network=host + + - name: Login to DockerHub + uses: docker/login-action@v3 + with: + username: ${{ secrets.DOCKERHUB_USERNAME }} + password: ${{ secrets.DOCKERHUB_TOKEN }} + logout: true + + # Notes on Cache: + # https://docs.docker.com/build/ci/github-actions/examples/#inline-cache + - name: Build Container + uses: docker/build-push-action@v5 + with: + context: . + push: false + load: true + tags: ${{ github.repository }}:vuln-test + cache-from: type=registry,ref=${{ github.repository }}:buildcache + cache-to: type=registry,ref=${{ github.repository }}:buildcache,mode=max + + # We will not be concerned with Medium and Low vulnerabilities + - name: Run Trivy vulnerability scanner + uses: aquasecurity/trivy-action@0.36.0 + with: + image-ref: '${{ github.repository }}:vuln-test' + format: 'sarif' + severity: 'CRITICAL,HIGH' + ignore-unfixed: true + output: 'trivy-results.sarif' + exit-code: '1' + # Scan results should be viewable in GitHub Security Dashboard + # We still fail the job if results are found, so below will always run + # unless manually canceled. + - name: Upload Trivy scan results to GitHub Security tab + uses: github/codeql-action/upload-sarif@v3 + if: '!cancelled()' + with: + sarif_file: 'trivy-results.sarif' diff --git a/.gitignore b/.gitignore new file mode 100644 index 00000000..6c46fe7f --- /dev/null +++ b/.gitignore @@ -0,0 +1,154 @@ +# Git ignore bioler plate from https://github.com/github/gitignore/blob/master/Python.gitignore +.secret-env +.vscode/ + +# Byte-compiled / optimized / DLL files +__pycache__/ +*.py[cod] +*$py.class + +# C extensions +*.so + +# Distribution / packaging +.Python +build/ +develop-eggs/ +dist/ +downloads/ +eggs/ +.eggs/ +lib/ +lib64/ +parts/ +sdist/ +var/ +wheels/ +share/python-wheels/ +*.egg-info/ +.installed.cfg +*.egg +MANIFEST + +# PyInstaller +# Usually these files are written by a python script from a template +# before PyInstaller builds the exe, so as to inject date/other infos into it. +*.manifest +*.spec + +# Installer logs +pip-log.txt +pip-delete-this-directory.txt + +# Unit test / coverage reports +htmlcov/ +.tox/ +.nox/ +.coverage +.coverage.* +.cache +nosetests.xml +coverage.xml +*.cover +*.py,cover +.hypothesis/ +.pytest_cache/ +cover/ + +# Translations +*.mo +*.pot + +# Django stuff: +*.log +local_settings.py +db.sqlite3 +db.sqlite3-journal + +# Flask stuff: +instance/ +.webassets-cache + +# Scrapy stuff: +.scrapy + +# Sphinx documentation +docs/_build/ + +# PyBuilder +.pybuilder/ +target/ + +# Jupyter Notebook +.ipynb_checkpoints + +# IPython +profile_default/ +ipython_config.py + +# pyenv +# For a library or package, you might want to ignore these files since the code is +# intended to run in multiple environments; otherwise, check them in: +.python-version + +# pipenv +# According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control. +# However, in case of collaboration, if having platform-specific dependencies or dependencies +# having no cross-platform support, pipenv may install dependencies that don't work, or not +# install all needed dependencies. +#Pipfile.lock + +# PEP 582; used by e.g. github.com/David-OConnor/pyflow +__pypackages__/ + +# Celery stuff +celerybeat-schedule +celerybeat.pid + +# SageMath parsed files +*.sage.py + +# Environments +.secrets-env +.venv +env/ +venv/ +ENV/ +env.bak/ +venv.bak/ + +# Spyder project settings +.spyderproject +.spyproject + +# PyCharm +.idea + +# Rope project settings +.ropeproject + +# Mac +.DS_Store + +# mkdocs documentation +/site + +# mypy +.mypy_cache/ +.dmypy.json +dmypy.json + +# Pyre type checker +.pyre/ + +# pytype static type analyzer +.pytype/ + +# Cython debug symbols +cython_debug/ + +# Local output directories +dags/roger/data +local_storage +logs +tests/integration/data/bulk/ diff --git a/.pylintrc b/.pylintrc new file mode 100644 index 00000000..e4e7438b --- /dev/null +++ b/.pylintrc @@ -0,0 +1,4 @@ +[MAIN] +disable=invalid-name, + no-member, + no-value-for-parameter diff --git a/CLAUDE.md b/CLAUDE.md new file mode 100644 index 00000000..a44b43fa --- /dev/null +++ b/CLAUDE.md @@ -0,0 +1,340 @@ +# CLAUDE.md + +This file provides guidance to Claude Code (claude.ai/code) when working with code in this repository. + +## What Roger Is + +Roger is an automated graph-data curation pipeline that transforms KGX (Knowledge Graph Exchange) files into a FalkorDB/RedisGraph database, orchestrated by Apache Airflow (3.x). It processes biomedical datasets (TopMed, BDC, AnVIL, dbGaP, HEAL, RADx, SPARC, NIDA, KFDRC, CRDC, CTN, BACPAC, etc.) through annotation (via Dug), Elasticsearch indexing, KGX merge/normalization, schema inference, and bulk loading. + +Two largely independent workloads live here: + +1. **`annotate_and_index`** — per-dataset Dug annotation, TranQL concept expansion, KGX conversion, and Elasticsearch indexing. This is the incremental one (see below). +2. **`knowledge_graph_build`** — the KGX merge → schema → bulk-load → validate chain into FalkorDB. + +## Commands + +```bash +make install # pip install -r requirements.txt (Python 3.12 in CI/Docker) +make test.lint # flake8 dags (CI uses: flake8 --ignore=E,W dags) +make test.unit # pytest tests/unit +make test.integration # pytest tests/integration +make test # unit + integration + +# Single test +python3 -m pytest tests/unit/test_config.py::test_merge -v + +# Local stack (Airflow UI on :8080, Postgres, Elasticsearch :9200, Redis Stack/FalkorDB :6379) +make stack.init # docker-compose up airflow-init (first time) +make stack # docker-compose up + +make clean # wipe logs and local_storage +``` + +The package source lives in **`src/roger/`**; the Docker image sets `PYTHONPATH=/opt/airflow/dags/repo/src/` so DAGs can import it. + +### Broken/stale Makefile targets + +- `make build` and anything using `VERSION` read `./dags/_version.py`, **which does not exist**. The real version is `version = 0.10.4.2` in `setup.cfg`. Bump that and build/tag the image manually, or restore `dags/_version.py` (note `dags/__init__.py` still does `from ._version import version as __version__`, so that import is also dead). +- `test.doc` and parts of `rm_dirs` reference a removed `dags/roger/` copy. + +### Running tests without a local environment + +There are **no roger dependencies installed on the host** — no Airflow, no Dug, no avalon. The system `python3` is 3.11 and Dug requires 3.12 (`typing.override`). Run everything in a container instead. + +Fastest path is the already-built roger image, which has every dependency including the exact Dug/`dug_data_model` versions in production: + +```bash +docker run --rm -v $PWD:/repo:ro -w /repo --user root \ + --entrypoint bash containers.renci.org/helxplatform/roger: -c \ + "pip install -q pytest flake8; PYTHONPATH=/repo/src python -m pytest tests/unit -q; \ + flake8 --ignore=E,W dags" +``` + +Alternative (no image handy): a `python:3.12-slim-trixie` container plus a full `pip install` of `requirements.txt`'s git dependencies — works, but takes ~5 minutes with a cold pip cache. + +Notes: +- `dags/__pycache__` is root-owned by docker, so host `py_compile` of `dags/` fails with EACCES. Use `ast.parse` for a syntax-only check. +- Mount `:ro` for reading/tests; mount read-write and pass `--user $(id -u):$(id -g)` for anything that writes, so files stay owned by you. + +## Architecture + +### Knowledge graph build data flow + +``` +GET KGX files → MERGE nodes (dedup) → CREATE SCHEMA (type inference) + → CREATE BULK CSVs → BULK LOAD into FalkorDB → VALIDATE (test queries) → CHECK TRANQL +``` + +### Layers + +- **`src/roger/core/`** — graph operations. `base.py` has the `Roger` orchestrator class and top-level functions (`get_kgx`, `merge_nodes`, `create_schema`, `create_bulk_load`, `bulk_load`, `validate`). `bulkload.py` generates CSVs for the FalkorDB bulk loader; `redis_graph.py` wraps the `falkordb` SDK; `storage.py` owns all path conventions and file globbing. +- **`src/roger/models/kgx.py`** — KGX merge and schema inference. Type-conflict rules: bool/float/int conflicts → string; any list value → list. +- **`src/roger/pipelines/`** — one class per dataset, all inheriting `DugPipeline` (`base.py`). A pipeline defines `pipeline_name`, optional `parser_name`, `input_version`, file discovery, and the annotate/crawl/index/KGX-convert steps. See `src/roger/pipelines/README.md` for adding a new dataset. +- **`src/roger/tasks.py`** — all Airflow glue: config injection, LakeFS I/O, incremental state, and the DAG-building helpers. The biggest and most subtle file in the repo. +- **`dags/`** — `knowledge_graph_build.py`, `annotate_and_index.py`, `index_only.py`. +- **`src/roger/cli.py`** — non-Airflow CLI: `python3 -m roger --get-kgx --merge-kgx --create-schema --create-bulk --insert --validate -d `. +- **`scripts/`** — one-off maintenance tooling; see the data-model migration section. + +### `annotate_and_index` DAG shape + +File-producing work and Elasticsearch work are deliberately **separate task groups**, because lakefs is the source of truth and ES is a derived index that gets rebuilt wholesale: + +``` +init + └─> {name}_dataset_pipeline_task_group (one per dataset, incremental) + annotate_{name}_files + ├─> make_kgx_{name} + └─> crawl_{name} + └─> complete_{name} (trigger_rule="none_failed") + └─> wipe_es_indexes (ONE global task, all datasets) + └─> {name}_es_index_task_group (one per dataset, always full) + index_{name}_variables ─> validate_{name}_index_variables + index_{name}_concepts ─> validate_{name}_index_concepts + └─> complete_{name} + └─> finish (trigger_rule="none_failed") +``` + +Why it is split this way: + +- ES index names are **global and shared across every dataset** (see `indexing.*_index` in `config.yaml`). A per-dataset wipe would destroy sibling datasets' documents, so there is exactly one `wipe_es_indexes` task sitting between the file groups and the ES groups. +- The file groups run incrementally. The ES groups are built with `incremental_pull=False`, so they always pull the *complete* set of annotate/crawl outputs from lakefs and reindex from scratch. That is what makes upstream **deletions** work: a file removed from a source repo is removed from lakefs, so it simply isn't there during the rebuild and its documents never come back. No per-document ES bookkeeping. +- ES document ids are stable (`element.id`, `concept_id`), so reindexing is an upsert and retries are safe. + +**Caveat:** the wipe is unconditional. Every run that reaches it empties all indexes and then repopulates them. If a file group fails and its ES group is skipped, the indexes end up partially populated until the next successful run. + +### `index_only` DAG + +`dags/index_only.py` re-indexes Elasticsearch from `annotate_and_index` outputs **already committed to the runtime repo** — for example after merging the dev runtime branch into prod. No annotate or crawl re-runs. `create_index_only_taskgroup` pulls inputs by explicit path (`{ANNOTATE_DAG_ID}/{group_id}.{task_id}/`) from the configured runtime repo/branch, with `params={"incremental": False}`. + +Note the input pairing, which is easy to get wrong: `validate_indexed_variables` reads only annotated elements, while `validate_indexed_concepts` pairs expanded concepts (crawl output) against annotated elements (annotate output) and asserts matching counts — so it needs **both** prefixes. + +## Incremental ingestion + +The problem: `annotate_{slug}_files` downloads from external LakeFS repos that keep growing, and re-annotating everything on every run is expensive for no gain. Since lakefs is version-oriented, each task instead diffs against the last commit it successfully consumed. + +### State storage + +Per-task Airflow Variables, keyed by `incremental_state_key()`: + +``` +roger_incr::{dag_id}::{task_id}::{repo}@{branch} +``` + +Use `from airflow.sdk import Variable` — `airflow.models.Variable` does direct DB access and is blocked on Airflow 3.x workers (`tasks.py` imports the SDK one with a fallback). + +State is **only advanced after a successful commit+merge**, inside +`avalon_commit_callback`'s try block, so a failed merge can never mark +unprocessed commits as consumed. That merge failure is re-raised rather than +logged and dropped: the `clean_up` at the end of the callback deletes the +local output, so swallowing it destroyed the work *and* reported success. + +### Flow through one task + +1. **`setup_input_data`** (the `pre_execute` hook) groups the task's configured repos by `(repo, branch)`. For each group it resolves the ref tip, reads the last-consumed commit from the Variable, and diffs. +2. **`resolve_ref_tip`** uses `refs_api.log_commits(ref, amount=1)` — **not** `branches_api.get_branch`, because dataset version refs (e.g. `topmed:v2.0`) are frequently **tags, not branches**. +3. **`get_changed_files`** does the diff roger-side via `refs_api.diff_refs` + `pagination_helper`, keeping `removed` entries (avalon's own `get_changes` drops them) and prefix-filtering with trailing-slash normalization. `'*'` or empty means no filter; `conflict`/`prefix_changed` are ignored. +4. Downloads are pinned to the **resolved tip commit**, not the branch name, so a push mid-run can't produce a torn read. +5. First run, `incremental=False`, `incremental_pull=False`, or a `NotFoundException` all fall back to a full `get_files(changes_only=False)`. +6. If nothing changed in any group, the task raises `AirflowSkipException`, which skips it and cascades through the group. `complete_{name}` and `finish` use `trigger_rule="none_failed"` so a skipped group still finishes green. +7. `write_state_file` persists the resolved tips to a JSON file next to the task dir (`generate_dir_name_from_task_instance(..., suffix='state')`). This file is the channel from `pre_execute` to the success callback — the callback must **never** re-resolve the branch itself, or commits that landed mid-run would be marked consumed without being processed. + +Downstream tasks are incremental the same way, diffing roger's own output-repo prefixes (`{dag_id}/{upstream_task_id}/`). + +### Deletion propagation + +When the diff reports removals, `avalon_commit_callback` maps them to derived outputs and deletes those on the temp branch before the merge, so removals land atomically with the rest of the commit: + +- `removed_bases()` maps a removed source path to a base name — first path segment under `{dag}/{task}/` for roger-repo paths, basename minus extension for external repos. +- `stale_output_paths()` matches existing outputs as `rel.startswith(base + '/')` or `rel == f"{base}_kgx.json"`. +- Deletion goes through `objects_api.delete_objects(PathList(...))` + a `commits_api.commit`. + +This is a **heuristic**, deliberately fail-safe: archive-style sources (one tarball expanding into many files) won't map 1:1, so nothing is deleted for them rather than the wrong thing being deleted. A `.removed_files.json` manifest is also written on every run with state — as a dotfile, because `storage.py` readers glob `*.json`/`**/*.json` over pulled task outputs and glob skips dotfiles. Writing it every run (not just on removals) guarantees `put_files` has content even on a removal-only run. + +### dbGaP sibling files + +`find_sibling_files()` handles a wrinkle specific to dbGaP data dicts: the parser needs a sibling `GapExchange_` file in the same lakefs directory for study name/description, but an incremental diff only carries the changed data dicts. So each affected directory is listed and any marker file not already downloaded is pulled in. + +### Turning it off + +`params={"incremental": false}` on a DAG run forces a full pull everywhere. Per-task, `create_python_task(..., incremental_pull=False)` opts a single task out permanently (this is what the ES groups use). + +## Annotation cost, caching, and resume + +Annotation is the dominant cost in the whole repo, and two structural facts +explain nearly all of it. + +**dbGaP parsers emit the study element into every data-dict file.** So a study +with 55,000 data dicts annotates its study description 55,000 times. +Measured on `bdc-parent` (61,597 XML files, **24 distinct studies**; Framingham +`phs000007.v35.p16` alone is 54,986 files): + +| | per file | +|---|---| +| study element (the same one every time) | ~53.5 s | +| variable element (what the file actually contributes) | ~1.4 s | + +That is 55 s/file, 39 days for the dataset, **96% of it recomputing 24 +answers**. + +**Dug's cached session cached only one of its four calls.** +`DugFactory.build_http_session` returns `requests_cache.CachedSession`, whose +`allowable_methods` defaults to `('GET', 'HEAD')`. Of the four annotation +calls, only node normalization is a GET (`DefaultNormalizer.make_request`, +`dug/core/annotators/_base.py`); nemo token classification, sapbert, and +name-resolution synonyms are all **POST** and so were never cached. + +`roger.utils.http_utils.enable_post_caching` fixes it +(`annotation.cache_post_requests`, on by default). The request body is part of +requests_cache's key for POST, so this is correct, not a heuristic. Faster +annotator endpoints do not help here: the bottleneck is call *count*. + +The same function also sets `expire_after`, which dug never did — so the +normalizer GETs that *were* being cached had no expiry and grew unbounded. +See the eviction note below. + +Keep `annotation.http_cache_expire_seconds` **nonzero** (default 30 days). +requests_cache's redis backend writes entries with `SETEX` only when an expiry +is set. That makes annotation cache keys volatile while the FalkorDB graph keys +in the same redis stay permanent — so redis can be given a `maxmemory` with +`volatile-lru` and will evict cache before it ever touches the graph. With no +expiry the cache is permanent and unbounded, and under `noeviction` (the +deployed default, with `maxmemory 0`) it grows until the pod is OOMKilled, +taking the loaded graph with it. + +`annotation.annotate_workers` (default 4) threads `annotate_files` over input +files. Files are wholly independent — own parse, own `Crawler`, own output dir +— and the work is nearly all HTTP wait, so this scales despite the GIL. Each +worker gets its own session and annotator via +`DugPipeline.thread_annotation_context`; the response cache is shared, so +workers still see each other's annotations. Element-level concurrency is not +possible without changing dug: `Crawler.annotate_elements` is a serial loop, +and inside it `AnnotateSapbert.__call__` does one classify call, then a sapbert +call per entity, then a normalize *and* a synonym call per identifier, all +sequentially. Log volume scales with worker count — see the ephemeral-storage +history in `roger.logger`. + +### Resume + +Task output only reaches lakefs on task *success*, so a 39-day task that died +at file 40,000 used to discard all of it and restart at zero. Three pieces make +retries resume: + +- `DugPipeline.annotation_is_complete` skips input files whose `elements.txt` + **and** `concepts.txt` both exist and are non-empty. Both are required: they + are written in sequence, so a kill between them leaves a directory that looks + started but is unusable. +- `clean_up(..., keep_output=True)` is the failure callback, so the dead try's + output survives. +- `reuse_prior_try_outputs` (from `setup_input_data`) hard-links earlier tries' + outputs into the current try's dir — necessary because + `generate_dir_name_from_task_instance` stamps the try number into the path, so + a retry otherwise starts in an empty directory. The successful commit then + includes everything, and `clean_up(..., all_tries=True)` clears every try dir. + +This covers retries within a dag run. It does **not** checkpoint mid-task: a +single try that never succeeds commits nothing, and a fresh dag run gets a new +`run_id` and therefore new dirs. Sharding a dataset across mapped tasks is the +next step if that becomes the binding constraint. + +## LakeFS integration (via the `avalon` library) + +When `ROGER_LAKEFS__CONFIG_ENABLED=true`, each task pulls inputs from a LakeFS repo/branch (`get_files()`), works in a task-specific local dir (named by `generate_dir_name_from_task_instance`), writes outputs back (`put_files()`), and commits via a temp branch merged after task success (`Merge(strategy="source-wins")`). Without LakeFS, tasks read/write a shared local data root. + +`create_python_task()` wires the callbacks: + +| flag | effect | +|---|---| +| `no_input_files=True` | skip the `pre_execute` download entirely (used by `wipe_es_indexes`) | +| `no_output_files=True` | `post_execute` runs `record_state_callback` (advances state only) instead of `avalon_commit_callback` (commits output) | +| `incremental_pull=False` | always download full inputs, even when the DAG run is incremental | +| `pass_conf` | whether the DAG run conf is forwarded into the callable | + +Output is committed from **`post_execute`**, not `on_success_callback`. +Airflow runs `post_execute` inside `_execute_task`, before it records +`end_date` and releases downstream; success callbacks run later, in +`finalize()`. Committing there let downstream tasks read the branch before the +output landed — `BulkLoad` loaded an edgeless graph 105s early, and +`make_kgx` built KGX from annotations 4.8 hours stale, both green. Airflow +also only *logs* exceptions raised by state-change callbacks +(`_run_task_state_change_callbacks`), so a failed upload or merge left the +task successful; from `post_execute` it fails the task. + +`on_failure_callback` and `on_skipped_callback` both run `clean_up` — the skip variant matters because `pre_execute` creates the input dir *before* it can raise `AirflowSkipException`. + +## Dug data model and jsonpickle + +Dug's data classes were extracted into a separate **`dug_data_model`** library: `DugElement` (`v2.base`), `DugConcept` (`v2.concept`), `DugVariable` (`v2.variable`), `DugStudy` (`v2.study`), `DugSection` (`v2.section`). `DugIdentifier` still lives in `dug.core.annotators._base`. + +Intermediate artifacts (`elements.txt`, `concepts.txt`, `expanded_concepts.txt`) are **jsonpickle-encoded**, meaning they embed fully-qualified class paths like `"py/object": "dug_data_model.v2.concept.DugConcept"`. + +**The trap:** when a stored class's module fails to import, jsonpickle does *not* raise — it silently returns the raw dict. The failure surfaces much later and much more confusingly, as `'dict' object has no attribute 'id'` inside ES indexing. The old `dug.core.parsers._base` module is exactly this case: it no longer imports at all (circular import via `dug.core.loaders.InputFile`). + +This interacts badly with incremental ingestion: artifacts are only rewritten when their source changes, so they can sit in lakefs across multiple Dug upgrades, pinned to long-dead class paths. + +**Do not fix this by re-annotating** — that is a month of compute. Use the migration: + +```bash +lakectl local clone lakefs://// ./roger-out + +# what's in there, and does any stored field lack a home in the current model? +docker run --rm -v $PWD:/w -w /w --entrypoint python \ + containers.renci.org/helxplatform/roger: \ + /w/scripts/migrate_pickled_classes.py --scan ./roger-out + +# rewrite in place (note the --user flag; the image user does not own your files) +docker run --rm -v $PWD:/w -w /w --user $(id -u):$(id -g) --entrypoint python \ + containers.renci.org/helxplatform/roger: \ + /w/scripts/migrate_pickled_classes.py --fix ./roger-out + +lakectl local commit ./roger-out -m "restamp pickled dug classes" +``` + +`--scan` prints each stored class as `ok`/`STALE`, then a field-drift table with the count of objects inspected per class. Two lines gate the migration: + +- `!! no current class named: [...]` — a class was renamed, not just moved; needs an explicit mapping. +- `!! stored but not declared -> [...]` — a *field* was renamed or dropped; migrating would silently lose those values. + +`--fix` aliases the dead modules, decodes, fills in fields added since (via `model_construct()` defaults, because `__setstate__` assigns `__dict__` wholesale and new fields would otherwise just be absent), and re-encodes with current class paths. It is idempotent and skips already-current files. `--self-check` round-trips a synthetic legacy payload with no lakefs data needed. + +Repeat per dataset prefix. Any future Dug release that moves data classes needs the same pass. + +## Configuration + +`RogerConfig` loads `src/roger/config/config.yaml` (override file via `ROGER_CONFIG_FILE`), then applies env vars prefixed `ROGER_`. Dots in the config path map to `_`, and literal underscores in key names are escaped as `__`: + +```bash +ROGER_REDISGRAPH_HOST=... # redisgraph.host +ROGER_KGX_DATA__SETS=topmed:v1.0 # kgx.data_sets +ROGER_LAKEFS__CONFIG_ENABLED=true # lakefs_config.enabled +ROGER_DUG__INPUTS_DATA__SETS=topmed:v2.0,anvil:v1.0 # which pipelines the DAGs build +``` + +`ROGER_DUG__INPUTS_DATA__SETS` is read directly by the DAG files as `name:version` pairs; the version becomes each pipeline class's `input_version`, i.e. the external lakefs ref (branch **or tag**) it pulls from. + +Main sections: `redisgraph`, `kgx`, `dug_inputs`, `bulk_loader`, `annotation` (annotator type + normalizer/synonym service URLs), `indexing` (ES index names, TranQL queries, `element_mapping`), `elasticsearch`, `lakefs_config`, `validation` (test queries). `dev-config.yaml` and `test-config.yaml` sit alongside the default. `RogerConfig.to_dug_conf()` bridges Roger config to Dug. + +## Bulk load CSV quirks + +The FalkorDB bulk loader requires every column populated, so `bulkload.py` groups entities by which attributes have values — output files look like `data/bulk/nodes/.csv--`. Values are cast per the inferred schema; the column separator is `0x1E` (configurable as `bulk_loader.separator`). + +## Key External Dependencies + +- **Dug** (helxplatform/dug, `DugModel2.0` branch) — annotation, indexing, concept expansion +- **dug_data_model** — the extracted Dug data classes (see the jsonpickle section) +- **avalon** (`lakefs-1.71.0` fork) — LakeFS client, wrapping `lakefs_sdk` 1.12 +- **falkordb** + `falkordb-bulk-loader` — graph database (not the old redisgraph package) +- **bmt** — Biolink Model Tools (model version pinned in `kgx.biolink_model_version`) +- **apache-airflow 3.2.0** +- **jsonpickle** — unpinned; 4.x deprecation warnings about `keys` defaulting to True in 5.0 are expected noise + +## Gotchas worth remembering + +- `DugPipeline.search_obj` / `.index_obj` are lazily built. Any method that touches Elasticsearch first in a task must initialize them (`clear_index` does this) or you get `'NoneType' object has no attribute 'es'`. +- Dataset version refs are often **tags**. Never assume branch APIs work on them. +- Incremental state Variables have no compare-and-swap, so `annotate_and_index` and `index_only` both set `max_active_runs=1`. +- Task ids encode the group (`{name}_dataset_pipeline_task_group.annotate_{name}_files`), and lakefs output paths are built from those ids — renaming a task or group orphans its previous outputs and resets its incremental state. +- Pre-existing lint noise that is not worth "fixing" blind: `F401` unused imports in `dags/__init__.py` and `tasks.py`, `F824` unused `nonlocal` in `storage.py`. + +## CI + +GitHub Actions (`.github/workflows/`): `code-checks.yml` runs flake8 + pytest on Python 3.12 and a Docker build test; dev images push on `develop` branch; Trivy scans PRs. Default PR target branch is `develop`. diff --git a/Dockerfile b/Dockerfile new file mode 100644 index 00000000..efd301d5 --- /dev/null +++ b/Dockerfile @@ -0,0 +1,89 @@ +# Use a Debian-based image for better compatibility +FROM python:3.12.13-slim-trixie +# FROM dhi.io/python:3.12-debian13-dev +# Set Airflow version and home directory + +ARG AIRFLOW_VERSION=3.2.2 + +ARG AIRFLOW_HOME=/opt/airflow + +# Environment variables +ENV AIRFLOW_HOME=${AIRFLOW_HOME} +ENV AIRFLOW__CORE__LOAD_EXAMPLES=False +ENV AIRFLOW__CORE__EXECUTOR=LocalExecutor +ENV AIRFLOW__DATABASE__SQL_ALCHEMY_CONN=postgresql+psycopg2://airflow:airflow@postgres:5432/airflow +ENV PYTHONUNBUFFERED=1 + +# Create airflow user and directories +RUN groupadd -g 50000 airflow +RUN useradd --uid 50000 --home-dir ${AIRFLOW_HOME} -g 50000 --create-home airflow && \ + mkdir -p ${AIRFLOW_HOME}/dags ${AIRFLOW_HOME}/logs ${AIRFLOW_HOME}/plugins ${AIRFLOW_HOME}/config + +# Install system dependencies +RUN apt-get update && apt-get install -y --no-install-recommends \ + build-essential \ + libpq-dev \ + libffi-dev \ + libssl-dev \ + curl \ + tini \ + tzdata \ + git \ + && rm -rf /var/lib/apt/lists/* + +# Upgrade pip tools +RUN pip install --no-cache-dir --upgrade pip setuptools wheel + +# Install Airflow (with PostgreSQL, Celery, Redis support) +RUN pip install --no-cache-dir \ + "apache-airflow[postgres,celery,redis,fab]==${AIRFLOW_VERSION}" \ + "apache-airflow-providers-cncf-kubernetes" \ + --constraint "https://raw.githubusercontent.com/apache/airflow/constraints-${AIRFLOW_VERSION}/constraints-3.12.txt" + +# Fix auth rollback bug. +RUN pip install --no-cache-dir \ + "apache-airflow-providers-fab==3.3.0rc1" + +# Optional: install extra packages +RUN pip install --no-cache-dir psycopg2-binary redis + +COPY ./requirements.txt /tmp/requirements.txt + +RUN pip install -r /tmp/requirements.txt + +RUN rm /tmp/requirements.txt + +# COPY . /opt/roger +# RUN pip install /opt/roger + +RUN apt-get purge -y --auto-remove \ + build-essential \ + libpq-dev \ + libffi-dev \ + libssl-dev \ + curl \ + git && \ + apt-get clean + +RUN if [ -n "$ROGER_SOURCE" ]; then pip install -e $ROGER_SOURCE; fi + +# Set ownership +RUN chown -R airflow:airflow ${AIRFLOW_HOME} + +# Vulnerability cleanup +RUN apt-get purge -y --auto-remove --allow-remove-essential perl-base + +# Switch to airflow user +USER airflow +WORKDIR ${AIRFLOW_HOME} + +ENV PYTHONPATH=/opt/airflow/dags/repo/src/ + +# Expose Airflow webserver port +EXPOSE 8080 + +# Use tini for signal handling +ENTRYPOINT ["/usr/bin/tini", "--"] + +# Default command +CMD ["airflow", "webserver"] diff --git a/Makefile b/Makefile new file mode 100644 index 00000000..ef227aa4 --- /dev/null +++ b/Makefile @@ -0,0 +1,76 @@ +PYTHON = $(shell which python3) +PYTHONPATH = dags +VERSION_FILE = ./dags/_version.py +VERSION = $(shell cut -d " " -f 3 ${VERSION_FILE}) +DOCKER_REPO = docker.io +DOCKER_OWNER = helxplatform +DOCKER_APP = roger +DOCKER_TAG = ${VERSION} +DOCKER_IMAGE = ${DOCKER_OWNER}/${DOCKER_APP}:$(DOCKER_TAG) + +.DEFAULT_GOAL = help + +.PHONY: help clean install test build image publish + +help: + @grep -E '^#[a-zA-Z\.\-]+:.*$$' $(MAKEFILE_LIST) | tr -d '#' | awk 'BEGIN {FS = ": "}; {printf "\033[36m%-20s\033[0m %s\n", $$1, $$2}' + +mk_dirs: + mkdir -p {logs,plugins} + mkdir -p local_storage/elastic + mkdir -p local_storage/redis + +rm_dirs: + rm -rf logs/* + rm -rf local_storage/elastic/* + rm -rf local_storage/redis/* + rm -rf ./dags/roger/data/* + +#install: Install application along with required packages to local environment +install: + ${PYTHON} -m pip install --upgrade pip + ${PYTHON} -m pip install -r requirements.txt + +#test.lint: Run flake8 on the source code +test.lint: + ${PYTHON} -m flake8 dags + +#test.doc: Run doctests in the source code +test.doc: + echo "Running doc tests..." + ${PYTHON} -m pytest --doctest-modules dags/roger + +#test.unit: Run unit tests +test.unit: + ${PYTHON} --version + ${PYTHON} -m pytest tests/unit + +#test.integration: Run unit tests +test.integration: + echo "Running integration tests..." + ${PYTHON} -m pytest tests/integration + +#test: Run all tests +test: test.unit test.integration + +#build: Build the Docker image +build: + echo "Building docker image: ${DOCKER_IMAGE}" + docker build --no-cache -t ${DOCKER_IMAGE} -f Dockerfile . + echo "Successfully built: ${DOCKER_IMAGE}" + +#publish: Push the Docker image +publish: + docker tag ${DOCKER_IMAGE} ${DOCKER_REPO}/${DOCKER_IMAGE} + docker push ${DOCKER_REPO}/${DOCKER_IMAGE} + +#clean: Remove old data +clean: rm_dirs mk_dirs + +#stack.init: Initialize the airflow DB +stack.init: mk_dirs + docker-compose up airflow-init + +#stack: Bring up Airflow and all backend services +stack: stack.init + docker-compose up diff --git a/README.md b/README.md index 0e7038aa..84d316c5 100644 --- a/README.md +++ b/README.md @@ -1,505 +1,235 @@ -# roger +# Roger -Roger is an automated graph data curation pipeline. -![image](https://user-images.githubusercontent.com/306971/97894880-0c17ef00-1d01-11eb-8162-0b6bd977769d.png) +Roger is an automated graph-data curation pipeline. It takes biomedical dataset metadata (dbGaP data dictionaries, study descriptions, common data elements) and turns it into two searchable products: -The first workflow transforms Knowledge Graph eXchange ([KGX](https://github.com/biolink/kgx)) files into a graph database in phases: -* **get**: Fetch KGX files from a repository. -* **merge**: Merge duplicate nodes accross multiple KGX files. -* **schema**: Infer the schema properties of nodes and edges. -* **bulk create**: Format for [bulk load to Redisgraph](https://github.com/RedisGraph/redisgraph-bulk-loader). -* **bulk load**: Load into Redisgraph -* **validate**: Execute test queries to validate the bulk load. +1. **Elasticsearch indexes** of variables, studies, sections, concepts, and knowledge-graph answers — what the [Dug](https://github.com/helxplatform/dug) semantic search UI queries. +2. **A FalkorDB/RedisGraph knowledge graph** built from [KGX](https://github.com/biolink/kgx) files, queried via TranQL. -## Installation +Everything runs as [Apache Airflow](https://airflow.apache.org/) 3.x DAGs. Roger is part of the [HeLx](https://helx.renci.org/) platform. -Requires Python 3.7+, Docker, and Make. +--- -Also requires KGX fork with Redisgraph Transformer. +## Start here: the two workflows + +Roger has two largely independent pipelines. Knowing which one you're looking at is the fastest way to orient yourself. + +### 1. `annotate_and_index` — metadata → Elasticsearch + +Per dataset (TopMed, BDC, AnVIL, HEAL, RADx, …): ``` -$ git clone https://github.com/stevencox/kgx -$ git clone -$ cd -$ pip install requirements.txt -$ bin/roger all +annotate ──> crawl_tranql ──┐ + │ ├──> [wipe ES] ──> index_variables ──> validate + └──> make_kgx │ index_concepts ──> validate + └─────────────────────────────────────────────> ``` -Roger can also be run via a Makefile: +- **annotate** — parse source files into Dug elements, annotate free text against ontologies (via a normalizer/synonym service), producing concepts. +- **crawl_tranql** — expand each concept through the knowledge graph via TranQL, attaching "kg answers". +- **make_kgx** — emit KGX nodes/edges so annotations can flow into the graph build. +- **index / validate** — push into Elasticsearch, then run search queries asserting the documents are actually findable. + +### 2. `knowledge_graph_build` — KGX → FalkorDB + ``` -cd bin -make clean install validate +get ──> merge ──> schema ──> bulk create ──> bulk load ──> validate ──> check tranql ``` -## Design +- **get** — fetch KGX files for a configured data version. +- **merge** — deduplicate nodes across files, unioning their properties. +- **schema** — infer the property set and type of every node/edge category. +- **bulk create** — write CSVs for the FalkorDB bulk loader. +- **bulk load** — load the graph. +- **validate** — run timed test queries as a sanity check. -Roger's is designed to transform data through well defined and transparent phases. +A third DAG, `index_only`, re-runs just the Elasticsearch half from artifacts already stored in LakeFS — useful after promoting data from dev to prod without re-annotating anything. -In general, each phase -* Reads and writes a set of files. -* Manages data in a single, configurable, root directory. +--- -Configuration is at roger/config.yaml. +## Repo tour -Roger can load Redisgraph -* By running the RedisgraphTransformer (currently on a fork of KGX) - * Currently, this is very slow. -* By bulk loading Redisgraph +``` +src/roger/ the actual package (Docker sets PYTHONPATH here) +├── tasks.py ALL Airflow glue: config injection, LakeFS I/O, +│ incremental state, DAG/task-group builders. +│ The densest file in the repo — read it second. +├── pipelines/ one class per dataset, all subclassing DugPipeline +│ ├── base.py DugPipeline: annotate, crawl_tranql, index_*, validate_* +│ └── README.md how to add a new dataset <-- read this to contribute +├── core/ graph side: base.py orchestrator, bulkload.py, +│ redis_graph.py, storage.py (all path conventions) +├── models/kgx.py KGX merge + schema inference rules +├── config/ config.yaml (+ dev-config.yaml, test-config.yaml) +└── cli.py run the graph build without Airflow -To build a bulk load, we -* Ensure no duplicate nodes exist -* Preserve all properties present across duplicate nodes -* Ensure all nodes of the same type have exactly the same properties -* Generate a comprehensive header (schema) for all nodes and edges -These constraints are managed in the steps below. +dags/ DAG definitions only; logic lives in src/roger +├── annotate_and_index.py +├── knowledge_graph_build.py +└── index_only.py -### Get -Fetches KGX files according to a data version selecting the set of files to use. -### Merge -Merges nodes duplicated across files aggregating properties from all nodes -### Schema -Identify and record the schema (properties) of every edge and node type. -### Bulk Create -Create bulk load CSV files conforming to the Redisgraph Bulk Loader's requirements. -### Bulk Load -Use the bulk loader to load Redisgraph logging statistics on each type of loaded object. -### Validate -Runs a configurable list of queries with timing information to quality check the generated graph database. +scripts/ one-off maintenance tooling +tests/unit/ the fast tests; start here +bin/ Helm/k8s deployment helpers (bin/roger init|start|stop) +``` -## Execution +Reading order for a new contributor: this file → `src/roger/pipelines/README.md` → `src/roger/tasks.py` → the dataset pipeline you care about. -### Redisgraph +--- -Roger uses Redisgraph's new bulk loader which is available in the 'edge' tagged Docker image. +## Quickstart -You can run the container like this and use it immediately -``` -docker run -p 6379:6379 -it --rm --name redisgraph redislabs/redisgraph:edge -``` -or run it with `/bin/bash` at the end to get a shell like this: -``` -docker run -p 6379:6379 -it --rm --name redisgraph redislabs/redisgraph:edge /bin/bash -``` -This lets you have a look around inside the container. To start Redis with the graph database plugin: -``` -# redis-server --loadmodule /usr/lib/redis/modules/redisgraph.so +Requires Docker, Docker Compose, Make, and Python 3.12 if you want to run anything on the host. + +```shell +make stack.init # one-time: docker-compose up airflow-init +make stack # bring everything up ``` -A clean Roger build looks like this. Times below are on a Macbook Air. +That gives you Airflow UI on `:8080`, Postgres, Elasticsearch on `:9200`, and Redis Stack / FalkorDB on `:6379`. Open the UI, pick a DAG, and hit trigger. -This can be run in the bin directory as -``` -$ make clean install validate -``` -Or via the roger CLI -``` -$ ../bin/roger all -[roger][core.py][ get] DEBUG: wrote data/kgx/chembio_kgx-v0.1.json: edges: 21637 nodes: 8725 time: 13870 -[roger][core.py][ get] DEBUG: wrote data/kgx/chemical_normalization-v0.1.json: edges: 277030 nodes: 72963 time: 15455 -[roger][core.py][ get] DEBUG: wrote data/kgx/cord19-phenotypes-v0.1.json: edges: 24 nodes: 25 time: 392 -[roger][core.py][ get] DEBUG: wrote data/kgx/ctd-v0.1.json: edges: 48363 nodes: 24008 time: 7143 -[roger][core.py][ get] DEBUG: wrote data/kgx/foodb-v0.1.json: edges: 5429 nodes: 4536 time: 1974 -[roger][core.py][ get] DEBUG: wrote data/kgx/mychem-v0.1.json: edges: 123119 nodes: 5496 time: 12271 -[roger][core.py][ get] DEBUG: wrote data/kgx/pharos-v0.1.json: edges: 287750 nodes: 224349 time: 40150 -[roger][core.py][ get] DEBUG: wrote data/kgx/topmed-v0.1.json: edges: 63860 nodes: 15870 time: 10901 - -real 1m58.722s -user 1m4.472s -sys 0m4.625s -[roger][core.py][ merge] INFO: merging data/kgx/chembio_kgx-v0.1.json -[roger][core.py][ merge] DEBUG: merged data/kgx/chemical_normalization-v0.1.json load: 1377 scope: 60 merge: 39 -[roger][core.py][ merge] DEBUG: merged data/kgx/cord19-phenotypes-v0.1.json load: 118 scope: 38 merge: 0 -[roger][core.py][ merge] DEBUG: merged data/kgx/ctd-v0.1.json load: 1151 scope: 26 merge: 19 -[roger][core.py][ merge] DEBUG: merged data/kgx/foodb-v0.1.json load: 141 scope: 24 merge: 1 -[roger][core.py][ merge] DEBUG: merged data/kgx/mychem-v0.1.json load: 1763 scope: 9 merge: 5 -[roger][core.py][ merge] DEBUG: merged data/kgx/pharos-v0.1.json load: 8426 scope: 218 merge:126 -[roger][core.py][ merge] DEBUG: merged data/kgx/topmed-v0.1.json load: 873 scope: 194 merge: 4 -[roger][core.py][ merge] INFO: data/kgx/chembio_kgx-v0.1.json rewrite: 1323. total merge time: 62921 -[roger][core.py][ merge] INFO: merge data/merge/chemical_normalization-v0.1.json is up to date. -[roger][core.py][ merge] INFO: merge data/merge/cord19-phenotypes-v0.1.json is up to date. -[roger][core.py][ merge] INFO: merge data/merge/ctd-v0.1.json is up to date. -[roger][core.py][ merge] INFO: merge data/merge/foodb-v0.1.json is up to date. -[roger][core.py][ merge] INFO: merge data/merge/mychem-v0.1.json is up to date. -[roger][core.py][ merge] INFO: merge data/merge/pharos-v0.1.json is up to date. -[roger][core.py][ merge] INFO: merge data/merge/topmed-v0.1.json is up to date. - -real 1m8.211s -user 0m53.546s -sys 0m3.894s -[roger][core.py][ is_up_to_date] DEBUG: no targets found -[roger][core.py][ create_schema] DEBUG: analyzing schema of data/kgx/chembio_kgx-v0.1.json. -[roger][core.py][ create_schema] DEBUG: analyzing schema of data/kgx/chemical_normalization-v0.1.json. -[roger][core.py][ create_schema] DEBUG: analyzing schema of data/kgx/cord19-phenotypes-v0.1.json. -[roger][core.py][ create_schema] DEBUG: analyzing schema of data/kgx/ctd-v0.1.json. -[roger][core.py][ create_schema] DEBUG: analyzing schema of data/kgx/foodb-v0.1.json. -[roger][core.py][ create_schema] DEBUG: analyzing schema of data/kgx/mychem-v0.1.json. -[roger][core.py][ create_schema] DEBUG: analyzing schema of data/kgx/pharos-v0.1.json. -[roger][core.py][ create_schema] DEBUG: analyzing schema of data/kgx/topmed-v0.1.json. -[roger][core.py][ write_schema] INFO: writing schema: data/schema/predicate-schema.json -[roger][core.py][ write_schema] INFO: writing schema: data/schema/category-schema.json - -real 0m46.205s -user 0m34.701s -sys 0m3.237s -[roger][core.py][ is_up_to_date] DEBUG: no targets found -[roger][core.py][ create] INFO: processing data/merge/chembio_kgx-v0.1.json -[roger][core.py][ write_bulk] INFO: --creating data/bulk/nodes/chemical_substance.csv -[roger][core.py][ write_bulk] INFO: --creating data/bulk/nodes/gene.csv -[roger][core.py][ write_bulk] INFO: --creating data/bulk/nodes/named_thing.csv -[roger][core.py][ write_bulk] INFO: --creating data/bulk/edges/directly_interacts_with.csv -[roger][core.py][ create] INFO: processing data/merge/chemical_normalization-v0.1.json -[roger][core.py][ write_bulk] INFO: --creating data/bulk/edges/similar_to.csv -[roger][core.py][ create] INFO: processing data/merge/cord19-phenotypes-v0.1.json -[roger][core.py][ write_bulk] INFO: --creating data/bulk/nodes/disease.csv -[roger][core.py][ write_bulk] INFO: --creating data/bulk/nodes/phenotypic_feature.csv -[roger][core.py][ write_bulk] INFO: --creating data/bulk/edges/has_phenotype.csv -[roger][core.py][ create] INFO: processing data/merge/ctd-v0.1.json -[roger][core.py][ write_bulk] INFO: --creating data/bulk/edges/treats.csv -[roger][core.py][ write_bulk] INFO: --creating data/bulk/edges/contributes_to.csv -[roger][core.py][ write_bulk] INFO: --creating data/bulk/edges/decreases_activity_of.csv -[roger][core.py][ write_bulk] INFO: --creating data/bulk/edges/decreases_molecular_interaction.csv -[roger][core.py][ write_bulk] INFO: --creating data/bulk/edges/increases_activity_of.csv -[roger][core.py][ write_bulk] INFO: --creating data/bulk/edges/increases_localization_of.csv -[roger][core.py][ write_bulk] INFO: --creating data/bulk/edges/increases_expression_of.csv -[roger][core.py][ write_bulk] INFO: --creating data/bulk/edges/decreases_response_to.csv -[roger][core.py][ write_bulk] INFO: --creating data/bulk/edges/increases_molecular_interaction.csv -[roger][core.py][ write_bulk] INFO: --creating data/bulk/edges/increases_degradation_of.csv -[roger][core.py][ write_bulk] INFO: --creating data/bulk/edges/affects_activity_of.csv -[roger][core.py][ write_bulk] INFO: --creating data/bulk/edges/decreases_localization_of.csv -[roger][core.py][ write_bulk] INFO: --creating data/bulk/edges/affects_localization_of.csv -[roger][core.py][ write_bulk] INFO: --creating data/bulk/edges/decreases_secretion_of.csv -[roger][core.py][ write_bulk] INFO: --creating data/bulk/edges/increases_secretion_of.csv -[roger][core.py][ write_bulk] INFO: --creating data/bulk/edges/affects_response_to.csv -[roger][core.py][ write_bulk] INFO: --creating data/bulk/edges/increases_response_to.csv -[roger][core.py][ write_bulk] INFO: --creating data/bulk/edges/increases_synthesis_of.csv -[roger][core.py][ write_bulk] INFO: --creating data/bulk/edges/increases_transport_of.csv -[roger][core.py][ write_bulk] INFO: --creating data/bulk/edges/increases_mutation_rate_of.csv -[roger][core.py][ write_bulk] INFO: --creating data/bulk/edges/affects_metabolic_processing_of.csv -[roger][core.py][ write_bulk] INFO: --creating data/bulk/edges/decreases_metabolic_processing_of.csv -[roger][core.py][ write_bulk] INFO: --creating data/bulk/edges/increases_metabolic_processing_of.csv -[roger][core.py][ write_bulk] INFO: --creating data/bulk/edges/decreases_degradation_of.csv -[roger][core.py][ write_bulk] INFO: --creating data/bulk/edges/affects_synthesis_of.csv -[roger][core.py][ write_bulk] INFO: --creating data/bulk/edges/decreases_molecular_modification_of.csv -[roger][core.py][ write_bulk] INFO: --creating data/bulk/edges/increases_molecular_modification_of.csv -[roger][core.py][ write_bulk] INFO: --creating data/bulk/edges/decreases_synthesis_of.csv -[roger][core.py][ write_bulk] INFO: --creating data/bulk/edges/decreases_expression_of.csv -[roger][core.py][ write_bulk] INFO: --creating data/bulk/edges/affects.csv -[roger][core.py][ write_bulk] INFO: --creating data/bulk/edges/decreases_stability_of.csv -[roger][core.py][ write_bulk] INFO: --creating data/bulk/edges/molecularly_interacts_with.csv -[roger][core.py][ write_bulk] INFO: --creating data/bulk/edges/affects_degradation_of.csv -[roger][core.py][ write_bulk] INFO: --creating data/bulk/edges/increases_uptake_of.csv -[roger][core.py][ write_bulk] INFO: --creating data/bulk/edges/decreases_mutation_rate_of.csv -[roger][core.py][ write_bulk] INFO: --creating data/bulk/edges/increases_stability_of.csv -[roger][core.py][ write_bulk] INFO: --creating data/bulk/edges/affects_expression_of.csv -[roger][core.py][ write_bulk] INFO: --creating data/bulk/edges/affects_secretion_of.csv -[roger][core.py][ write_bulk] INFO: --creating data/bulk/edges/decreases_uptake_of.csv -[roger][core.py][ write_bulk] INFO: --creating data/bulk/edges/affects_transport_of.csv -[roger][core.py][ write_bulk] INFO: --creating data/bulk/edges/decreases_transport_of.csv -[roger][core.py][ write_bulk] INFO: --creating data/bulk/edges/affects_uptake_of.csv -[roger][core.py][ create] INFO: processing data/merge/foodb-v0.1.json -[roger][core.py][ write_bulk] INFO: --creating data/bulk/edges/related_to.csv -[roger][core.py][ create] INFO: processing data/merge/mychem-v0.1.json -[roger][core.py][ write_bulk] INFO: --creating data/bulk/edges/causes_adverse_event.csv -[roger][core.py][ write_bulk] INFO: --creating data/bulk/edges/causes.csv -[roger][core.py][ write_bulk] INFO: --creating data/bulk/edges/Unmapped_Relation.csv -[roger][core.py][ create] INFO: processing data/merge/pharos-v0.1.json -[roger][core.py][ write_bulk] INFO: --creating data/bulk/edges/gene_associated_with_condition.csv -[roger][core.py][ create] INFO: processing data/merge/topmed-v0.1.json -[roger][core.py][ write_bulk] INFO: --creating data/bulk/nodes/cell.csv -[roger][core.py][ write_bulk] INFO: --creating data/bulk/nodes/molecular_activity.csv -[roger][core.py][ write_bulk] INFO: --creating data/bulk/nodes/anatomical_entity.csv -[roger][core.py][ write_bulk] INFO: --creating data/bulk/nodes/cellular_component.csv -[roger][core.py][ write_bulk] INFO: --creating data/bulk/nodes/biological_process.csv -[roger][core.py][ write_bulk] INFO: --creating data/bulk/edges/association.csv -[roger][core.py][ write_bulk] INFO: --creating data/bulk/edges/has_part.csv -[roger][core.py][ write_bulk] INFO: --creating data/bulk/edges/part_of.csv - -real 1m7.897s -user 0m58.467s -sys 0m2.791s -[roger][core.py][ insert] INFO: bulk loading - nodes: ['data/bulk/nodes/gene.csv', 'data/bulk/nodes/molecular_activity.csv', 'data/bulk/nodes/phenotypic_feature.csv', 'data/bulk/nodes/cell.csv', 'data/bulk/nodes/biological_process.csv', 'data/bulk/nodes/chemical_substance.csv', 'data/bulk/nodes/cellular_component.csv', 'data/bulk/nodes/anatomical_entity.csv', 'data/bulk/nodes/named_thing.csv', 'data/bulk/nodes/disease.csv'] - edges: ['data/bulk/edges/part_of.csv', 'data/bulk/edges/decreases_metabolic_processing_of.csv', 'data/bulk/edges/decreases_uptake_of.csv', 'data/bulk/edges/decreases_secretion_of.csv', 'data/bulk/edges/decreases_molecular_modification_of.csv', 'data/bulk/edges/increases_synthesis_of.csv', 'data/bulk/edges/causes_adverse_event.csv', 'data/bulk/edges/decreases_localization_of.csv', 'data/bulk/edges/decreases_stability_of.csv', 'data/bulk/edges/treats.csv', 'data/bulk/edges/affects_activity_of.csv', 'data/bulk/edges/increases_secretion_of.csv', 'data/bulk/edges/decreases_expression_of.csv', 'data/bulk/edges/affects_transport_of.csv', 'data/bulk/edges/Unmapped_Relation.csv', 'data/bulk/edges/affects_localization_of.csv', 'data/bulk/edges/increases_stability_of.csv', 'data/bulk/edges/decreases_activity_of.csv', 'data/bulk/edges/increases_response_to.csv', 'data/bulk/edges/causes.csv', 'data/bulk/edges/decreases_degradation_of.csv', 'data/bulk/edges/similar_to.csv', 'data/bulk/edges/decreases_synthesis_of.csv', 'data/bulk/edges/affects_expression_of.csv', 'data/bulk/edges/affects_uptake_of.csv', 'data/bulk/edges/has_part.csv', 'data/bulk/edges/affects_synthesis_of.csv', 'data/bulk/edges/affects_response_to.csv', 'data/bulk/edges/increases_molecular_interaction.csv', 'data/bulk/edges/increases_localization_of.csv', 'data/bulk/edges/increases_expression_of.csv', 'data/bulk/edges/increases_uptake_of.csv', 'data/bulk/edges/related_to.csv', 'data/bulk/edges/increases_mutation_rate_of.csv', 'data/bulk/edges/affects.csv', 'data/bulk/edges/decreases_transport_of.csv', 'data/bulk/edges/gene_associated_with_condition.csv', 'data/bulk/edges/directly_interacts_with.csv', 'data/bulk/edges/increases_metabolic_processing_of.csv', 'data/bulk/edges/molecularly_interacts_with.csv', 'data/bulk/edges/increases_degradation_of.csv', 'data/bulk/edges/affects_metabolic_processing_of.csv', 'data/bulk/edges/has_phenotype.csv', 'data/bulk/edges/decreases_response_to.csv', 'data/bulk/edges/decreases_molecular_interaction.csv', 'data/bulk/edges/increases_activity_of.csv', 'data/bulk/edges/association.csv', 'data/bulk/edges/affects_secretion_of.csv', 'data/bulk/edges/decreases_mutation_rate_of.csv', 'data/bulk/edges/contributes_to.csv', 'data/bulk/edges/increases_transport_of.csv', 'data/bulk/edges/increases_molecular_modification_of.csv', 'data/bulk/edges/affects_degradation_of.csv'] -[roger][core.py][ insert] INFO: deleting graph test in preparation for bulk load. -[roger][core.py][ insert] INFO: no graph to delete -[roger][core.py][ insert] INFO: bulk loading graph: test -gene [####################################] 100% -17868 nodes created with label 'gene' -3 nodes created with label 'molecular_activity' -phenotypic_feature [####################################] 100% -3723 nodes created with label 'phenotypic_feature' -8 nodes created with label 'cell' -2 nodes created with label 'biological_process' -chemical_substance [####################################] 100% -252966 nodes created with label 'chemical_substance' -2 nodes created with label 'cellular_component' -12 nodes created with label 'anatomical_entity' -named_thing [####################################] 100% -25903 nodes created with label 'named_thing' -disease [####################################] 100% -9777 nodes created with label 'disease' -part_of [####################################] 100% -31532 relations created for type 'part_of' -24 relations created for type 'decreases_metabolic_processing_of' -26 relations created for type 'decreases_uptake_of' -192 relations created for type 'decreases_secretion_of' -13 relations created for type 'decreases_molecular_modification_of' -186 relations created for type 'increases_synthesis_of' -causes_adverse_event [####################################] 100% -66461 relations created for type 'causes_adverse_event' -39 relations created for type 'decreases_localization_of' -12 relations created for type 'decreases_stability_of' -treats [####################################] 100% -11485 relations created for type 'treats' -307 relations created for type 'affects_activity_of' -527 relations created for type 'increases_secretion_of' -decreases_expression_of [####################################] 100% -2791 relations created for type 'decreases_expression_of' -91 relations created for type 'affects_transport_of' -28 relations created for type 'Unmapped_Relation' -506 relations created for type 'affects_localization_of' -48 relations created for type 'increases_stability_of' -decreases_activity_of [####################################] 100% -240317 relations created for type 'decreases_activity_of' -762 relations created for type 'increases_response_to' -causes [####################################] 100% -46277 relations created for type 'causes' -69 relations created for type 'decreases_degradation_of' -similar_to [####################################] 100% -277030 relations created for type 'similar_to' -21 relations created for type 'decreases_synthesis_of' -259 relations created for type 'affects_expression_of' -19 relations created for type 'affects_uptake_of' -has_part [####################################] 100% -31532 relations created for type 'has_part' -42 relations created for type 'affects_synthesis_of' -1804 relations created for type 'affects_response_to' -1495 relations created for type 'increases_molecular_interaction' -119 relations created for type 'increases_localization_of' -increases_expression_of [####################################] 100% -4178 relations created for type 'increases_expression_of' -118 relations created for type 'increases_uptake_of' -related_to [####################################] 100% -5429 relations created for type 'related_to' -564 relations created for type 'increases_mutation_rate_of' -116 relations created for type 'affects' -17 relations created for type 'decreases_transport_of' -gene_associated_with_condition [####################################] 100% -36017 relations created for type 'gene_associated_with_condition' -directly_interacts_with [####################################] 100% -30826 relations created for type 'directly_interacts_with' -467 relations created for type 'increases_metabolic_processing_of' -49 relations created for type 'molecularly_interacts_with' -increases_degradation_of [####################################] 100% -3394 relations created for type 'increases_degradation_of' -337 relations created for type 'affects_metabolic_processing_of' -24 relations created for type 'has_phenotype' -904 relations created for type 'decreases_response_to' -513 relations created for type 'decreases_molecular_interaction' -increases_activity_of [####################################] 100% -12061 relations created for type 'increases_activity_of' -796 relations created for type 'association' -242 relations created for type 'affects_secretion_of' -1 relations created for type 'decreases_mutation_rate_of' -contributes_to [####################################] 100% -16172 relations created for type 'contributes_to' -153 relations created for type 'increases_transport_of' -54 relations created for type 'increases_molecular_modification_of' -24 relations created for type 'affects_degradation_of' -Construction of graph 'test' complete: 310264 nodes created, 826470 relations created in 268.800857 seconds - -real 4m31.889s -user 2m50.070s -sys 0m4.201s -config:{ - "username": "", - "password": "", - "host": "localhost", - "graph": "test", - "ports": { - "http": 6379 - } -} -+-------------+ -| b'COUNT(a)' | -+-------------+ -| 310264 | -+-------------+ - -Cached execution 0.0 -internal execution time 19.4111 -[roger][core.py][ validate] INFO: Query count_nodes:Count Nodes ran in 39ms: MATCH (a) RETURN COUNT(a) -+-------------+ -| b'COUNT(e)' | -+-------------+ -| 826470 | -+-------------+ - -Cached execution 0.0 -internal execution time 6.2872 -[roger][core.py][ validate] INFO: Query count_edges:Count Edges ran in 14ms: MATCH (a)-[e]-(b) RETURN COUNT(e) -+-----------------+-------------------------------+ -| b'a.category' | b'b.id' | -+-----------------+-------------------------------+ -| ['named_thing'] | NCBIGene:5978 | -| ['named_thing'] | GO:0043336 | -| ['named_thing'] | CHEBI:24433 | -| ['named_thing'] | UBERON:0000178 | -| ['named_thing'] | TOPMED.VAR:phv00177354.v2.p10 | -| ['named_thing'] | TOPMED.VAR:phv00003307.v1.p10 | -| ['named_thing'] | TOPMED.VAR:phv00010123.v5.p10 | -...about 400 lines elided here... -| ['named_thing'] | TOPMED.VAR:phv00001046.v1.p10 | -| ['named_thing'] | TOPMED.VAR:phv00116572.v2.p2 | -| ['named_thing'] | TOPMED.VAR:phv00210333.v1.p1 | -| ['named_thing'] | TOPMED.VAR:phv00307964.v1.p1 | -| ['named_thing'] | TOPMED.VAR:phv00083411.v1.p3 | -+-----------------+-------------------------------+ - -Cached execution 0.0 -internal execution time 656.6031 -[roger][core.py][ validate] INFO: Query connectivity:TOPMED Connectivity ran in 725ms: MATCH (a { id : 'TOPMED.TAG:8' })--(b) RETURN a.category, b.id -+-----------------+--------------------------------+ -| b'a.category' | b'b.id' | -+-----------------+--------------------------------+ -| ['named_thing'] | TOPMED.TAG:64 | -| ['named_thing'] | TOPMED.STUDY:phs000007.v29.p10 | -+-----------------+--------------------------------+ - -Cached execution 0.0 -internal execution time 123.0736 -[roger][core.py][ validate] INFO: Query connectivity:TOPMED Connectivity ran in 129ms: MATCH (a { id : 'TOPMED.VAR:phv00000484.v1.p10' })--(b) RETURN a.category, b.id -+-----------------+--------------------------------+ -| b'a.category' | b'b.id' | -+-----------------+--------------------------------+ -| ['named_thing'] | TOPMED.TAG:30 | -| ['named_thing'] | TOPMED.STUDY:phs000007.v29.p10 | -+-----------------+--------------------------------+ - -Cached execution 0.0 -internal execution time 111.434 -[roger][core.py][ validate] INFO: Query connectivity:TOPMED Connectivity ran in 116ms: MATCH (a { id : 'TOPMED.VAR:phv00000487.v1.p10' })--(b) RETURN a.category, b.id -+-----------------+--------------------------------+ -| b'a.category' | b'b.id' | -+-----------------+--------------------------------+ -| ['named_thing'] | TOPMED.TAG:74 | -| ['named_thing'] | TOPMED.STUDY:phs000007.v29.p10 | -+-----------------+--------------------------------+ - -Cached execution 0.0 -internal execution time 110.0168 -[roger][core.py][ validate] INFO: Query connectivity:TOPMED Connectivity ran in 113ms: MATCH (a { id : 'TOPMED.VAR:phv00000496.v1.p10' })--(b) RETURN a.category, b.id -+-----------------+--------------------------------+ -| b'a.category' | b'b.id' | -+-----------------+--------------------------------+ -| ['named_thing'] | TOPMED.TAG:26 | -| ['named_thing'] | TOPMED.STUDY:phs000007.v29.p10 | -+-----------------+--------------------------------+ - -Cached execution 0.0 -internal execution time 118.366 -[roger][core.py][ validate] INFO: Query connectivity:TOPMED Connectivity ran in 122ms: MATCH (a { id : 'TOPMED.VAR:phv00000517.v1.p10' })--(b) RETURN a.category, b.id -+-----------------+--------------------------------+ -| b'a.category' | b'b.id' | -+-----------------+--------------------------------+ -| ['named_thing'] | TOPMED.TAG:40 | -| ['named_thing'] | TOPMED.STUDY:phs000007.v29.p10 | -+-----------------+--------------------------------+ - -Cached execution 0.0 -internal execution time 120.7783 -[roger][core.py][ validate] INFO: Query connectivity:TOPMED Connectivity ran in 128ms: MATCH (a { id : 'TOPMED.VAR:phv00000518.v1.p10' })--(b) RETURN a.category, b.id -+-----------------+--------------------------------+ -| b'a.category' | b'b.id' | -+-----------------+--------------------------------+ -| ['named_thing'] | TOPMED.STUDY:phs000007.v29.p10 | -| ['named_thing'] | TOPMED.TAG:7 | -+-----------------+--------------------------------+ - -Cached execution 0.0 -internal execution time 115.9252 -[roger][core.py][ validate] INFO: Query connectivity:TOPMED Connectivity ran in 120ms: MATCH (a { id : 'TOPMED.VAR:phv00000528.v1.p10' })--(b) RETURN a.category, b.id -+-----------------+--------------------------------+ -| b'a.category' | b'b.id' | -+-----------------+--------------------------------+ -| ['named_thing'] | TOPMED.TAG:8 | -| ['named_thing'] | TOPMED.STUDY:phs000007.v29.p10 | -+-----------------+--------------------------------+ - -Cached execution 0.0 -internal execution time 164.6341 -[roger][core.py][ validate] INFO: Query connectivity:TOPMED Connectivity ran in 170ms: MATCH (a { id : 'TOPMED.VAR:phv00000529.v1.p10' })--(b) RETURN a.category, b.id -+-----------------+--------------------------------+ -| b'a.category' | b'b.id' | -+-----------------+--------------------------------+ -| ['named_thing'] | TOPMED.STUDY:phs000007.v29.p10 | -| ['named_thing'] | TOPMED.TAG:7 | -+-----------------+--------------------------------+ - -Cached execution 0.0 -internal execution time 137.6454 -[roger][core.py][ validate] INFO: Query connectivity:TOPMED Connectivity ran in 144ms: MATCH (a { id : 'TOPMED.VAR:phv00000530.v1.p10' })--(b) RETURN a.category, b.id -+-----------------+--------------------------------+ -| b'a.category' | b'b.id' | -+-----------------+--------------------------------+ -| ['named_thing'] | TOPMED.TAG:8 | -| ['named_thing'] | TOPMED.STUDY:phs000007.v29.p10 | -+-----------------+--------------------------------+ - -Cached execution 0.0 -internal execution time 138.8376 -[roger][core.py][ validate] INFO: Query connectivity:TOPMED Connectivity ran in 143ms: MATCH (a { id : 'TOPMED.VAR:phv00000531.v1.p10' })--(b) RETURN a.category, b.id -+-------------+-------------+ -| b'count(a)' | b'count(b)' | -+-------------+-------------+ -| 1295945 | 1295945 | -+-------------+-------------+ - -Cached execution 0.0 -internal execution time 1661.8929 -[roger][core.py][ validate] INFO: Query count_connected_nodes:Count Connected Nodes ran in 1666ms: MATCH (a)-[e]-(b) RETURN count(a), count(b) -+-----------------------+-----------------------+ -| b'count(distinct(a))' | b'count(distinct(b))' | -+-----------------------+-----------------------+ -| 12156 | 196144 | -+-----------------------+-----------------------+ - -Cached execution 0.0 -internal execution time 3538.9259 -[roger][core.py][ validate] INFO: Query query_by_type:Query by Type ran in 3543ms: MATCH (a:gene)-[e]-(b) WHERE 'chemical_substance' IN b.category RETURN count(distinct(a)), count(distinct(b)) +Without Make: + +```shell +mkdir -p logs plugins local_storage/elastic +docker-compose up airflow-init +docker-compose up ``` -## Airflow +To wipe local state between runs: `make clean` (removes logs and `local_storage`). -This is a local run of Roger in Airflow. Next steps: Kubernetes. +For a no-Airflow local run driven entirely by the CLI and Makefiles, see **`roger-cli-steps.md`**. -![image](https://user-images.githubusercontent.com/306971/97792736-9acd2480-1bb8-11eb-8052-371a3188f3e4.png) +--- -Detailed feedback for each task is available including output logs +## Running the tests -![image](https://user-images.githubusercontent.com/306971/97792727-5fcaf100-1bb8-11eb-85b5-03cad151e0a0.png) +There are **no Roger dependencies installed on your host** and the dependency set is awkward (Dug needs Python 3.12; several deps install from git). Don't fight it — use a container. -### Running in Airflow -In one window: -``` -airflow scheduler -``` -In another: +If you have a built Roger image: + +```shell +docker run --rm -v $PWD:/repo:ro -w /repo --user root \ + --entrypoint bash containers.renci.org/helxplatform/roger: -c \ + "pip install -q pytest flake8; \ + PYTHONPATH=/repo/src python -m pytest tests/unit -q; \ + flake8 --ignore=E,W dags" ``` -airflow webserver -p 8080 + +Otherwise a `python:3.12-slim-trixie` container plus `pip install -r requirements.txt` works; budget ~5 minutes for the git dependencies. + +Make targets (`make test.unit`, `make test.lint`, `make test`) assume the deps are already importable, so they're really for CI and for inside the image. + +--- + +## Configuration + +`RogerConfig` loads `src/roger/config/config.yaml`, then overlays environment variables prefixed `ROGER_`. Dots in the config path become `_`, and a literal underscore in a key name is escaped by doubling it: + +```shell +ROGER_REDISGRAPH_HOST=localhost # redisgraph.host +ROGER_ELASTICSEARCH_HOST=localhost # elasticsearch.host +ROGER_LAKEFS__CONFIG_ENABLED=true # lakefs_config.enabled +ROGER_DUG__INPUTS_DATA__SETS=topmed:v2.0,anvil:v1.0 # which pipelines to build ``` -Open localhost:8080 in a browser. -Then run: +That last one is the important knob: it's a comma-separated list of `pipeline_name:version`, read directly by the DAG files. The name selects a class from `src/roger/pipelines/`; the version is the external LakeFS ref that pipeline reads from (often a **tag**, not a branch). Point `ROGER_CONFIG_FILE` at a different file to swap the whole config. + +Main sections: `redisgraph`, `kgx`, `dug_inputs`, `bulk_loader`, `annotation`, `indexing`, `elasticsearch`, `lakefs_config`, `validation`. + +--- + +## Concepts you'll hit early + +### LakeFS is the source of truth + +When `lakefs_config.enabled` is on, every task pulls its inputs from a LakeFS repo/branch, works in a task-specific local directory, writes its outputs back, and commits them on a temp branch that gets merged on success. Outputs are addressed by task id, so `annotate_and_index/{group}.{task}/…` is where a task's results live. Nothing is passed between tasks in memory. + +### Runs are incremental + +Re-annotating every dataset on every run is prohibitively expensive, so each task records the last source commit it successfully consumed (in an Airflow Variable) and only processes what changed since. If nothing changed, the task skips and the group still completes green. + +Force a full run with DAG params: `{"incremental": false}`. + +### Elasticsearch is derived, not authoritative + +The ES indexes are shared across all datasets, and they're wiped and rebuilt from whatever files remain in LakeFS. That's deliberate: it means a source file deleted upstream disappears from search without any per-document bookkeeping. It also means a run that reaches the wipe but fails partway leaves the indexes incomplete until the next good run. + +### Artifacts are jsonpickle, and that has teeth + +Intermediate files (`elements.txt`, `concepts.txt`, `expanded_concepts.txt`) are jsonpickle-encoded Python objects, so they embed fully-qualified class paths. If a Dug release moves those classes, old artifacts silently decode to plain dicts instead of raising, and you get a confusing `'dict' object has no attribute 'id'` during indexing. Fix is `scripts/migrate_pickled_classes.py`, not re-annotation. See `CLAUDE.md` for the runbook. + +--- + +## KGX merge and schema rules + +Worth knowing before debugging a weird graph load: + +- Duplicate nodes across files are merged, keeping the union of their properties. +- Every node of a given type must end up with exactly the same property set, so the schema step resolves conflicts: + - a property flip-flopping between bool / float / int becomes **string** + - a property that is ever a string and never a list becomes **string** + - a property that is ever a list becomes a **list** +- The FalkorDB bulk loader requires every column populated, so entities are grouped by which attributes actually have values and written to separate CSVs: `data/bulk/nodes/.csv--`. Column separator is `0x1E`. + +--- + +## Adding a dataset + +Add a subclass of `DugPipeline` in `src/roger/pipelines/`, set `pipeline_name` (and usually `parser_name` and `input_version`), then add it to `ROGER_DUG__INPUTS_DATA__SETS`. The DAGs build task groups for whatever is listed there — no DAG edits needed. + +Full walkthrough with the customization hooks: **`src/roger/pipelines/README.md`**. + +--- + +## Deploying to Kubernetes + +Roger installs via [Helm](https://helm.sh). Prerequisites: + +1. **A persistent volume** — create a `ReadWriteMany` PVC named `roger-data-pvc` for Roger's data directory. +2. **Git SSH secrets** — `airflow-secrets` (key `gitSshKey`, used by `AIRFLOW__KUBERNETES__GIT_SSH_KEY_SECRET_NAME`) and `airflow-git-keys` (`id_rsa`, `id_rsa.pub`, `known_hosts`, used by `airflow.dags.git.secret`), both base64-encoded. + +Then: + +```shell +cd bin/ +export NAMESPACE= +export RELEASE_NAME= +export CLUSTER_DOMAIN=cluster.local +./roger init # initialize helm dependencies (airflow + redis charts) +./roger start # install; follow the printed notes for port-forwarding ``` -python tranql_translator.py + +`./roger stop` tears it down, `./roger restart` cycles it. Trigger config for a run targeting a specific graph: + +```json +{"redisgraph": {"host": "", "port": 6379, "graph": "graph-name"}} ``` -The Airflow interface shows the workflow: -![image](https://user-images.githubusercontent.com/306971/97787955-b968f680-1b8b-11eb-86cc-4d93842eafd3.png) -Use the Trigger icon to run the workflow immediatley. +--- + +## Troubleshooting + +| Symptom | Likely cause | +|---|---| +| `'dict' object has no attribute 'id'` during indexing | LakeFS artifacts written by an older Dug; run `scripts/migrate_pickled_classes.py` | +| `'NoneType' object has no attribute 'es'` | An ES method ran before `search_obj`/`index_obj` were lazily built | +| Tasks skip with "No changes in source refs" | Working as intended — incremental found nothing new. Re-run with `{"incremental": false}` to force | +| Search results missing after a partial run | The ES wipe ran but a rebuild task failed; re-run the DAG | +| `make build` fails on a missing version file | Known: the Makefile reads `dags/_version.py`, which no longer exists. Version lives in `setup.cfg` | + +--- + +## Further reading +- **`CLAUDE.md`** — deep architecture notes: incremental state machine, deletion propagation, LakeFS task wiring, the Dug data-model migration runbook, and accumulated gotchas. +- **`src/roger/pipelines/README.md`** — adding and customizing dataset pipelines. +- **`roger-cli-steps.md`** — local deployment driven by the CLI instead of Airflow. +- **`bin/Readme.md`** — deployment helper scripts. +## Key dependencies +[Dug](https://github.com/helxplatform/dug) (annotation and search) · `dug_data_model` (the shared data classes) · [avalon](https://github.com/helxplatform/avalon) (LakeFS client) · `falkordb` + `falkordb-bulk-loader` · `bmt` (Biolink Model Toolkit) · Apache Airflow 3.2.0 +## License +See [LICENSE](LICENSE). diff --git a/bin/Makefile b/bin/Makefile index 21ba47bb..a9833163 100644 --- a/bin/Makefile +++ b/bin/Makefile @@ -1,61 +1,24 @@ -########################################################## -## -## -## Make the Roger database in phases. -## -## Opertions -## -## get: Fetch versioned knowledge graph exchange -## (KGX) formatted data files. -## -## merge: Merge nodes, consolidating duplicates -## and preserving fields. -## -## schema: Identify the all properties in each -## predicate and node type. -## -## tables: Write tabular formatted data for all -## edges and nodes. -## -## install: Bulk load a Redisgraph instance. -## -## validate: Validate database contents. -## -## clean: Delete all data artifacts. -## -## -########################################################## - -# Root of Roger -ROGER_HOME=$(PWD)/.. - -# Path to Roger executable -ROGER=${ROGER_HOME}/bin/roger - -# Location of data -DATA_ROOT=${ROGER_HOME}/roger/data +ROGER_MAKE_DIR=./roger_graph_build +ANNOTATE_MAKE_DIR=./dug_annotate +INDEXING_MAKE_DIR=./dug_indexing + RM=/bin/rm -TIME=/usr/bin/time -clean: - $(RM) -rf $(DATA_ROOT) +DATA_ROOT=${ROGERENV_DATA__ROOT} -get: - $(TIME) $(ROGER) kgx get --data-root $(DATA_ROOT) -merge: get - $(TIME) $(ROGER) kgx merge --data-root $(DATA_ROOT) +clean: + $(RM) -rf $(DATA_ROOT) -schema: merge - $(TIME) $(ROGER) kgx schema --data-root $(DATA_ROOT) -tables: schema - $(TIME) $(ROGER) bulk create --data-root $(DATA_ROOT) +annotate: + make -C ${ANNOTATE_MAKE_DIR} all -install: tables - $(TIME) $(ROGER) bulk load --data-root $(DATA_ROOT) +graph: + make -C ${ROGER_MAKE_DIR} all -validate: - $(TIME) $(ROGER) bulk validate --data-root $(DATA_ROOT) +index: + make -C ${INDEXING_MAKE_DIR} all +all: annotate graph index \ No newline at end of file diff --git a/bin/Readme.md b/bin/Readme.md new file mode 100644 index 00000000..bcade36a --- /dev/null +++ b/bin/Readme.md @@ -0,0 +1,136 @@ +### Running Roger + +This document outlines some of the ways that Roger can be run. + +### Roger Configuration + +Configuration is mainly managed through `roger/roger/config.yaml`. +Each values in this config file can be overridden by shell environment +variables. For instance to override the following : + +``` + kgx: + biolink_model_version: 1.5.0 + dataset_version: v1.0 +``` + +Overridding variables can be exported as: + +```shell script +export ROGERENV_KGX_BIOLINK__MODEL__VERSION=1.6 +export ROGERENV_KGX_DATASET__VERSION=v1.1 +``` +Some things to note are: +* Environment variables should be prefixed by `ROGERENV_` +* Single Underscore `_` character denotes sub-key in the yaml +* Double Underscores `__` are treated as regular underscore +* Keys in yaml are in lower and environment variables that override them should be in upper case. + +### Deploy Script + +`roger/bin/deploy` script can be used to deploy Roger's dependencies in either docker or kubernetes. +For full capabilities use: +```shell script +cd roger/bin +./deploy help +``` + +##### Docker + +For local development we can use docker containers to run backend services that roger depends on. +These are Redis store, Elastic search and Tranql web service. + +Eg: +```shell script +cd roger/bin +./deploy docker config # to display the configuration (port address and passwords) +./deploy docker start # to start +./deploy help # for help on commands +``` + +##### Kubernetes + +For running on k8s we can configure git branch and docker images by exporting: +```shell script +export NAMESPACE=your-namespace +export RELEASE=roger +export CLUSTER_DOMAIN=cluster.local +export WORKING_GIT_BRANCH=develop +``` +deploy using : + +```shell script +cd roger/bin +./deploy k8s config # to display the configuration +./deploy k8s start # to start +./deploy k8s help # for help on commands +``` + +### Local Development + +##### Setup python virtual env + +```shell script +cd roger +python -m venv venv +source venv/bin/activate +pip install -r requirements.txt +``` + +##### Configuration + +Refer to configuration section to override server names and passwords to +passwords etc.. to the backend servers. + +For development there is a dev.env file in `roger/bin/` directory with some start +up variables. Modify as needed. The following command can be used to export them into +shell. +```shell script +export $(grep -v'^#' bin/dev.env | xargs 0) +``` + + +##### Run a task + +To run a single task : + +```shell script +python cli.py -l # runs annotatation task +python cli.py -h # see the full list of available arguments. +``` + +##### Using the Makefiles + +Another way to run roger is as a pipeline, where each task is +In `roger/roger/bin/` there is a root make file and in the `roger/roger/bin/dug_annotate`, +`roger/roger/bin/dug_indexing` and `roger/roger/bin/roger_graph_build`. + +Running all pipelines end to end: + +```shell script +cd roger/roger/bin/ +make all +``` + +Running annotation pipeline: + +```shell script +cd roger/roger/bin/ +make annotate +``` + +Running graph pipeline: + +```shell script +cd roger/roger/bin/ +make graph +``` + +Running index pipeline: + +```shell script +cd roger/roger/bin/ +make index +``` + + diff --git a/bin/airk8s b/bin/airk8s deleted file mode 100755 index c6b32aae..00000000 --- a/bin/airk8s +++ /dev/null @@ -1,53 +0,0 @@ -#!/bin/bash - -set -x -set -e - -namespace=${NAMESPACE:-scox} -version=v7.11.0 - -init () { - helm repo add stable https://kubernetes-charts.storage.googleapis.com - helm repo update -} -start () { - helm install "airflow" stable/airflow \ - --version "$version" \ - --namespace "$namespace" \ - --values ./custom-values.yaml -} -status () { - helm status "airflow" --namespace $namespace - echo Scheduler: - kubectl -n $namespace logs $(kubectl get pods | grep airflow-scheduler | awk '{ print $1 }') -c git-sync - echo Worker: - kubectl -n $namespace logs $(kubectl get pods | grep airflow-worker | awk '{ print $1 }') -c git-sync -} -stop () { - helm delete "airflow" --namespace $namespace -} -connect () { - kubectl exec -it \ - --namespace $namespace \ - --container airflow-web \ - Deployment/airflow-web \ - /bin/bash -} -web () { - export NODE_PORT=$(kubectl get --namespace $namespace -o jsonpath="{.spec.ports[0].nodePort}" services airflow-web) - export NODE_IP=$(kubectl get nodes --namespace $namespace -o jsonpath="{.items[0].status.addresses[0].address}") - echo http://$NODE_IP:$NODE_PORT/ - export AIRFLOW_UI=http://$NODE_IP:$NODE_PORT/ -} -gitsecret () { - kubectl create secret generic \ - airflow-git-keys \ - --from-file=id_rsa=$HOME/.ssh/id_rsa \ - --from-file=id_rsa.pub=$HOME/.ssh/id_rsa.pub \ - --from-file=known_hosts=$HOME/.ssh/known_hosts \ - --namespace $namespace -} - -$* - -exit 0 diff --git a/bin/custom-values.yaml b/bin/custom-values.yaml deleted file mode 100644 index e60f3297..00000000 --- a/bin/custom-values.yaml +++ /dev/null @@ -1,157 +0,0 @@ -# -# NOTE: -# - This is intended to be a `custom-values.yaml` starting point for non-production deployment (like minikube) - -# External Dependencies: -# - A PUBLIC git repo for DAGs: ssh://git@repo.example.com:my-airflow-dags.git -# - -################################### -# Airflow - Common Configs -################################### -airflow: - ## the airflow executor type to use - ## - executor: CeleryExecutor -# executor: KubernetesExecutor - - ## the fernet key used to encrypt the connections in the database - ## - fernetKey: "7T512UXSSmBOkpWimFHIVb8jK6lfmSAvx4mO6Arehnc=" - - ## environment variables for the web/scheduler/worker Pods (for airflow configs) - ## - config: - # Security - AIRFLOW__CORE__SECURE_MODE: "True" - AIRFLOW__API__AUTH_BACKEND: "airflow.api.auth.backend.deny_all" - AIRFLOW__WEBSERVER__EXPOSE_CONFIG: "False" - AIRFLOW__WEBSERVER__RBAC: "False" - - # DAGS - AIRFLOW__CORE__LOAD_EXAMPLES: "False" - - ## Disable noisy "Handling signal: ttou" Gunicorn log messages - GUNICORN_CMD_ARGS: "--log-level WARNING" - -################################### -# Airflow - Scheduler Configs -################################### -scheduler: - - ## custom airflow connections for the airflow scheduler - ## -# connections: -# - id: my_aws -# type: aws -# extra: | -# { -# "aws_access_key_id": "XXXXXXXXXXXXXXXXXXX", -# "aws_secret_access_key": "XXXXXXXXXXXXXXX", -# "region_name":"eu-central-1" -# } - - ## custom airflow variables for the airflow scheduler - ## - variables: | - { "environment": "dev" } - - ## custom airflow pools for the airflow scheduler - ## - pools: | - { - "example": { - "description": "This is an example pool with 2 slots.", - "slots": 2 - } - } - -################################### -# Airflow - WebUI Configs -################################### -web: - ## configs for the Service of the web Pods - ## - service: - type: NodePort - -################################### -# Airflow - Worker Configs -################################### -workers: - ## the number of workers Pods to run - ## - replicas: 1 - -################################### -# Airflow - DAGs Configs -################################### -dags: - ## configs for the DAG git repository & sync container - ## - git: - ## url of the git repository - ## - #url: "ssh://git@repo.example.com/my-airflow-dags.git" - #url: "ssh://git@github.com/stevencox/airflow.git" - url: "ssh://git@github.com/stevencox/roger.git" - - ## the branch/tag/sha1 which we clone - ## - ref: main - - ## the name of a pre-created secret containing files for ~/.ssh/ - ## - ## NOTE: - ## - this is ONLY RELEVANT for SSH git repos - ## - the secret commonly includes files: id_rsa, id_rsa.pub, known_hosts - ## - known_hosts is NOT NEEDED if `git.sshKeyscan` is true - ## - secret: airflow-git-keys - - ## the name of the private key file in your `git.secret` - ## - ## NOTE: - ## - this is ONLY RELEVANT for PRIVATE SSH git repos - ## - privateKeyName: id_rsa - - ## the host name of the git repo - ## - ## NOTE: - ## - this is ONLY REQUIRED for SSH git repos - ## - ## EXAMPLE: - ## repoHost: "github.com" - ## - repoHost: "github.com" - - ## the port of the git repo - ## - ## NOTE: - ## - this is ONLY REQUIRED for SSH git repos - ## - repoPort: 22 - - ## configs for the git-sync container - ## - gitSync: - ## enable the git-sync sidecar container - ## - enabled: true - - ## the git sync interval in seconds - ## - refreshTime: 60 - -################################### -# Database - PostgreSQL Chart -################################### -postgresql: - enabled: true - -################################### -# Database - Redis Chart -################################### -redis: - enabled: true diff --git a/bin/deploy b/bin/deploy new file mode 100644 index 00000000..47dd6cba --- /dev/null +++ b/bin/deploy @@ -0,0 +1,119 @@ +#!/usr/bin/env bash +# ---- Kubernetes ------ + +k8s () { + + namespace=${NAMESPACE:-} + release=${RELEASE:-roger} + cluster_domain=${CLUSTER_DOMAIN:-cluster.local} + branch=${WORKING_GIT_BRANCH:-develop} + + help () { + echo " + Usage : ./deploy k8s [sub-command] + + Deploys Roger pipeline on kubernetes along airflow. + + Available sub-commands: + - config : view configuration + - init : Initializes helm dependencies for install. + - start : Runs helm upgrade/install. + - stop : Stops running instance. + - restart : Restarts running instance. + - client : If redis is installed on the system, it will try to connect to + " + + } + config() { + echo " + Configuration for k8s instance. + To modify this values export variables with new values. + eg: export NAMESPACE=my-namespace + + NAMESPACE: ${namespace} + RELEASE: ${release} + CLUSTER_DOMAIN: ${cluster_domain} + WORKING_GIT_BRANCH: ${branch} + " + } + init () { + helm dependency update ../helm + } + start () { + init + helm upgrade --install $release \ + --set redis.clusterDomain=$cluster_domain \ + --set airflow.airflow.config.AIRFLOW__KUBERNETES__GIT_BRANCH=$branch \ + --set airflow.dags.git.ref=$branch \ + --namespace=$namespace \ + ../helm + } + stop () { + helm delete $release \ + --namespace=$namespace + } + restart () { + stop + start + } + status () { + helm --namespace=$namespace status $release + } + client () { + redis-cli -h 127.0.0.1 -p 6379 -a $REDIS_PASSWORD + } + $* +} +#---------End Kubernetes------------------- + +#---------Docker-compose ------------------ + +docker() { + COMPOSE_FILE=./docker_backend/docker-compose.yaml + help () { + echo " + Usage: ./deploy docker [subcommand] + + Run docker based backends. + + Available sub-commands: + config: Print contents of ./.env file + init: Export ./.env file contents as shell variables. + start: Runs docker containers up using ${COMPOSE_FILE}. + stop: Stops running docker containers. + restart: Restarts containers. + " + } + config() { + grep -v "^#" dev.env + } + init() { + export $(config | xargs -0) + } + start() { + init + docker-compose -f ${COMPOSE_FILE} up -d + } + stop() { + init + docker-compose -f ${COMPOSE_FILE} down + } + $* +} + +help () { + echo " + Usage : ./deploy [env-type] [subcommand] + + Deploys roger dependencies in docker / k8s + + env-type: either k8s or docker + + Read below for the subcommands avaible or use + ./deploy [env-type] help . + + " + docker help + k8s help +} +$* \ No newline at end of file diff --git a/bin/dev.env b/bin/dev.env new file mode 100644 index 00000000..1653a62d --- /dev/null +++ b/bin/dev.env @@ -0,0 +1,4 @@ +ROGERENV_DATA__ROOT=~/roger-data +ROGERENV_KGX_DATASET__VERSION=test +ROGERENV_ELASTIC__SEARCH_PASSWORD=changeme +ROGERENV_REDISGRAPH_PASSWORD=changeme \ No newline at end of file diff --git a/bin/docker_backend/docker-compose.yaml b/bin/docker_backend/docker-compose.yaml new file mode 100644 index 00000000..d87c6ae6 --- /dev/null +++ b/bin/docker_backend/docker-compose.yaml @@ -0,0 +1,68 @@ +version: '3.0' + +################################################################################# +## +## A service stack for the Roger pipeline. +## +################################################################################# +services: + + ################################################################################# + ## + ## The OpenAPI endpoint for search. This is the only service to be + ## exposed beyond the internal network. + ## + ################################################################################# + tranql: + image: renciorg/tranql-app:0.35 + depends_on: + - redis + restart: always + networks: + - roger-network + environment: + - REDIS_PASSWORD=$ROGERENV_REDISGRAPH_PASSWORD + entrypoint: /usr/local/bin/gunicorn --workers=2 --bind=0.0.0.0:8001 --name=tranql --timeout=600 tranql.api:app + ports: + - 8001:8001 + volumes: + - ./tranql-schema.yaml:/tranql/tranql/conf/schema.yaml + ################################################################################# + ## + ## A search engine providing scalable indexing and full text search. + ## + ################################################################################# + elasticsearch: + image: docker.elastic.co/elasticsearch/elasticsearch:7.6.1 + networks: + - roger-network + environment: + - ELASTIC_PASSWORD=$ROGERENV_ELASTIC__SEARCH_PASSWORD + - discovery.type=single-node + - xpack.security.enabled=true + volumes: + - ./data/elastic:/bitnami/elasticsearch/data + ports: + - '9200:9200' + - '9300:9300' + + ################################################################################# + ## + ## A memory cache for results of high volume service requests. + ## + ################################################################################# + redis: + image: 'redislabs/redisgraph' + networks: + - roger-network + command: redis-server --requirepass ${ROGERENV_REDISGRAPH_PASSWORD} --loadmodule /usr/lib/redis/modules/redisgraph.so + environment: + - REDIS_DISABLE_COMMANDS=FLUSHDB,FLUSHALL + volumes: + - ./data/redis:/data + ports: + - '6379:6379' + +networks: + roger-network: + driver: bridge diff --git a/bin/docker_backend/tranql-schema.yaml b/bin/docker_backend/tranql-schema.yaml new file mode 100644 index 00000000..965d12c3 --- /dev/null +++ b/bin/docker_backend/tranql-schema.yaml @@ -0,0 +1,12 @@ +schema: + redis: + doc: | + Roger is a knowledge graph built by aggregeting several kgx formatted knowledge graphs from several sources. + url: "redis:" + redis: true + redis_connection_params: + # Host here is the service name in the docker composed container. + host: redis + port: 6379 + # SET USERNAME and PASSWORD + # via ROGER_USERNAME , ROGER_PASSWORD Env vars (i.e capitialize service name) diff --git a/bin/dug_annotate/Makefile b/bin/dug_annotate/Makefile new file mode 100644 index 00000000..745c01de --- /dev/null +++ b/bin/dug_annotate/Makefile @@ -0,0 +1,44 @@ +########################################################## +## +## +## Annotate files using Dug. +## +## Operations +## +## annotate_and_normalize: Annotates Variable files using entity name resolution service with curies. +## +## create_kgx_files: Creates KGX formatted knowledge graphs from annotation result set. +## +## clean: Delete all data artifacts. +## +## +########################################################## + +# Root +THIS_MAKEFILE_PATH:=$(word $(words $(MAKEFILE_LIST)),$(MAKEFILE_LIST)) +THIS_DIR:=$(shell cd $(dir $(THIS_MAKEFILE_PATH));pwd) + +ROGER_HOME=${THIS_DIR}/../.. +CLI_WRAPPER=${ROGER_HOME}/cli.py + +# Override Roger data dir ENV +ANNOTATE_DIR=${ROGERENV_DATA__ROOT}/dug/annotations +KGX_DIR=${ROGERENV_DATA__ROOT}/dug/kgx + +RM=/bin/rm +TIME=/usr/bin/time + +clean: + $(RM) -rf ${ANNOTATE_DIR} + $(RM) -rf ${KGX_DIR} + +get_input_files: + $(TIME) roger -gd + +annotate_and_normalize: + $(TIME) roger -l + +create_kgx_files: + $(TIME) roger -t + +all: get_input_files annotate_and_normalize create_kgx_files diff --git a/bin/dug_indexing/Makefile b/bin/dug_indexing/Makefile new file mode 100644 index 00000000..479f2b2c --- /dev/null +++ b/bin/dug_indexing/Makefile @@ -0,0 +1,52 @@ +########################################################## +## +## +## Annotate files using Dug. +## +## Operations +## +## annotate_and_normalize: Annotates Variable files using entity name resolution service with curies. +## +## create_kgx_files: Creates KGX formatted knowledge graphs from annotation result set. +## +## clean: Delete all data artifacts. +## +## +########################################################## + +# Root +THIS_MAKEFILE_PATH:=$(word $(words $(MAKEFILE_LIST)),$(MAKEFILE_LIST)) +THIS_DIR:=$(shell cd $(dir $(THIS_MAKEFILE_PATH));pwd) + +ROGER_HOME=${THIS_DIR}/../.. + +# Override Roger data dir ENV +INDEXING_DIR=${ROGERENV_DATA__ROOT}/dug/expanded_concepts +CRAWL_DIR=${ROGERENV_DATA__ROOT}/dug/crawl + + +RM=/bin/rm +TIME=/usr/bin/time + +clean: + $(RM) -rf ${INDEXING_DIR} + $(RM) -rf ${CRAWL_DIR} + +crawl_concepts: + $(TIME) roger -C + +index_concepts: crawl_concepts + $(TIME) roger -ic + +index_variables: + $(TIME) roger -iv + +validate_indexed_concepts: index_concepts + $(TIME) roger -vc + +validate_indexed_variables: index_variables + $(TIME) roger -vv + +all: validate_indexed_concepts validate_indexed_variables + + diff --git a/bin/roger b/bin/roger index 5d96cde8..4626df88 100755 --- a/bin/roger +++ b/bin/roger @@ -1,53 +1,7 @@ +#!/usr/bin/env bash #set -x set -e -namespace=${NAMESPACE:-scox} -release=redisgraph -image_repository=redislabs/redisgraph -image_tag=edge - -# https://github.com/bitnami/charts/tree/master/bitnami/redis -init () { - helm repo add bitnami https://charts.bitnami.com/bitnami -} -start () { - helm install $release \ - --set image.repository=$image_repository \ - --set image.tag=$image_tag \ - --set redis.command="redis-server" \ - --set redis.args="--loadmodule /usr/lib/redis/modules/redisgraph.so" \ - --set master.command="redis-server --loadmodule /usr/lib/redis/modules/redisgraph.so" \ - --set slave.command="redis-server --loadmodule /usr/lib/redis/modules/redisgraph.so" \ - --namespace=$namespace \ - bitnami/redis -} -start () { - helm install $release \ - --set image.repository=$image_repository \ - --set image.tag=$image_tag \ - --namespace=$namespace \ - bitnami/redis -} -stop () { - helm delete $release \ - --namespace=$namespace -} -restart () { - stop - start -} -status () { - kubectl --namespace=$namespace get pods | grep $release - export REDIS_PASSWORD=$(kubectl get secret --namespace $namespace redisgraph -o jsonpath="{.data.redis-password}" | base64 --decode) -} -client () { - #kubectl port-forward --namespace $namespace svc/redisgraph-master 6380:6379 & - redis-cli -h 127.0.0.1 -p 6380 -a $REDIS_PASSWORD -} -#---------------------------- - - - DIR="$( cd "$( dirname "${BASH_SOURCE[0]}" )" >/dev/null 2>&1 && pwd )" export ROGER_HOME=$( dirname $DIR ) @@ -56,7 +10,7 @@ export PYTHONPATH=$ROGER_HOME:$ROGER_HOME/../kgx export DB_NAME=test roger () { - python $ROGER_HOME/roger/core.py $* + python $ROGER_HOME/dags/roger/core.py $* } kgx () { diff --git a/bin/roger_graph_build/Makefile b/bin/roger_graph_build/Makefile new file mode 100644 index 00000000..f43c9c8c --- /dev/null +++ b/bin/roger_graph_build/Makefile @@ -0,0 +1,66 @@ +########################################################## +## +## +## Make the Roger database in phases. +## +## Opertions +## +## get: Fetch versioned knowledge graph exchange +## (KGX) formatted data files. +## +## merge: Merge nodes, consolidating duplicates +## and preserving fields. +## +## schema: Identify the all properties in each +## predicate and node type. +## +## tables: Write tabular formatted data for all +## edges and nodes. +## +## install: Bulk load a Redisgraph instance. +## +## validate: Validate database contents. +## +## clean: Delete all data artifacts. +## +## +########################################################## + +# Root of Roger +# Root +THIS_MAKEFILE_PATH:=$(word $(words $(MAKEFILE_LIST)),$(MAKEFILE_LIST)) +THIS_DIR:=$(shell cd $(dir $(THIS_MAKEFILE_PATH));pwd) + +ROGER_HOME=${THIS_DIR}/../.. + +# Path to Roger executable +CLI_WRAPPER=${ROGER_HOME}/cli.py + +# Location of data +DATA_ROOT=${ROGER_HOME}/roger/data + +RM=/bin/rm +TIME=/usr/bin/time + +clean: + $(RM) -rf $(DATA_ROOT) + +get: + $(TIME) python ${CLI_WRAPPER} -g + +merge: get + $(TIME) python ${CLI_WRAPPER} -m + +schema: merge + $(TIME) python ${CLI_WRAPPER} -s + +tables: schema + $(TIME) python ${CLI_WRAPPER} -b + +install: tables + $(TIME) python ${CLI_WRAPPER} -i + +validate: + $(TIME) python ${CLI_WRAPPER} -a + +all: install validate \ No newline at end of file diff --git a/dags/__init__.py b/dags/__init__.py new file mode 100644 index 00000000..f0aee1ff --- /dev/null +++ b/dags/__init__.py @@ -0,0 +1 @@ +from ._version import version as __version__ diff --git a/dags/annotate_and_index.py b/dags/annotate_and_index.py new file mode 100644 index 00000000..4e7324d5 --- /dev/null +++ b/dags/annotate_and_index.py @@ -0,0 +1,62 @@ +"""DAG which performs Dug annotate and index operations + +This DAG differes slightly from prior versions of the same functionality in +Roger not only in that the annotation and indexing happen in the same DAG, but +also those tasks are broken out into sub-DAGs organized by dataset. Each dataset +has a subdag for all tasks. +""" + +import os + +from airflow.models import DAG +from airflow.providers.standard.operators.empty import EmptyOperator +from roger.tasks import (default_args, create_pipeline_taskgroup, + create_es_taskgroup, create_es_wipe_task) + +env_enabled_datasets = os.getenv( + "ROGER_DUG__INPUTS_DATA__SETS", "topmed,anvil").split(",") + +with DAG( + dag_id='annotate_and_index', + default_args=default_args, + # incremental state Variables have no compare-and-swap; serialize runs + max_active_runs=1, + params= + { + "repository_id": None, + "branch_name": None, + "commitid_from": None, + "commitid_to": None, + # diff source refs against the last ingested commit and only + # process new/changed files; set false to force a full run + "incremental": True + }, + # schedule_interval=None +) as dag: + init = EmptyOperator(task_id="init", dag=dag) + finish = EmptyOperator(task_id="finish", dag=dag, + trigger_rule="none_failed") + + + from roger import pipelines + from roger.config import config + envspec = os.getenv("ROGER_DUG__INPUTS_DATA__SETS","topmed:v2.0") + data_sets = envspec.split(",") + pipeline_names = {x.split(':')[0]: x.split(':')[1] for x in data_sets} + pipeline_classes = list(pipelines.get_pipeline_classes(pipeline_names)) + + if pipeline_classes: + # file-based tasks run incrementally per dataset; then one global + # index wipe; then elastic rebuilds from the full file set left in + # lakefs (so upstream deletions simply vanish from the indexes) + wipe_es = create_es_wipe_task(dag, pipeline_classes[0], config) + for pipeline_class in pipeline_classes: + init >> create_pipeline_taskgroup(dag, pipeline_class, config) \ + >> wipe_es + wipe_es >> create_es_taskgroup(dag, pipeline_class, config) \ + >> finish + else: + init >> finish + +if __name__ == "__main__": + dag.test() diff --git a/dags/index_only.py b/dags/index_only.py new file mode 100644 index 00000000..d33a5c84 --- /dev/null +++ b/dags/index_only.py @@ -0,0 +1,38 @@ +"""DAG which only runs the Dug ES indexing steps. + +Re-indexes Elasticsearch (concepts + variables) from annotate_and_index +outputs already present in the runtime repo on the configured branch -- e.g. +after merging the dev runtime branch into prod. No annotate or crawl is +re-run; inputs are pulled by explicit path from the runtime repo. +""" + +import os + +from airflow.models import DAG +from airflow.providers.standard.operators.empty import EmptyOperator +from roger.tasks import default_args, create_index_only_taskgroup + +with DAG( + dag_id='index_only', + default_args=default_args, + max_active_runs=1, + params={ + # re-index everything present on the branch; not a diff + "incremental": False, + }, +) as dag: + init = EmptyOperator(task_id="init", dag=dag) + finish = EmptyOperator(task_id="finish", dag=dag, + trigger_rule="none_failed") + + from roger import pipelines + from roger.config import config + envspec = os.getenv("ROGER_DUG__INPUTS_DATA__SETS", "topmed:v2.0") + data_sets = envspec.split(",") + pipeline_names = {x.split(':')[0]: x.split(':')[1] for x in data_sets} + for pipeline_class in pipelines.get_pipeline_classes(pipeline_names): + init >> create_index_only_taskgroup(dag, pipeline_class, config) \ + >> finish + +if __name__ == "__main__": + dag.test() diff --git a/dags/knowledge_graph_build.py b/dags/knowledge_graph_build.py new file mode 100644 index 00000000..b34aefc5 --- /dev/null +++ b/dags/knowledge_graph_build.py @@ -0,0 +1,124 @@ +# -*- coding: utf-8 -*- +# + +""" +An Airflow workflow for the Roger Translator KGX data pipeline. +""" + +from airflow.models import DAG +from airflow.providers.standard.operators.empty import EmptyOperator +import roger +from roger.tasks import default_args, create_python_task +from roger.config import config + +""" Build the workflow's tasks and DAG. """ +with DAG( + dag_id='knowledge_graph_build', + default_args=default_args, + # schedule_interval=None +) as dag: + + """ Build the workflow tasks. """ + intro = EmptyOperator(task_id='Intro') + + # Merge nodes needs inputs from two sources + # 1. baseline and/or CDE KGX files from LakeFS (External repo) + # 2. Infer which local kgx files are needed based on dug_inputs and grab them from the current repo + + # build the annotate and index pipeline output locations + #lakefs://yk-heal/main/annotate_and_index/crdc_dataset_pipeline_task_group.make_kgx_crdc/ + working_repo = config.lakefs_config.repo + branch = config.lakefs_config.branch + kgx_repos = config.kgx.data_sets + input_repos = [{ + 'name': repo.split(':')[0], + 'branch': repo.split(':')[1], + 'path': '*' + } for repo in kgx_repos] + + # Figure out a way to extract paths + get_path_on_lakefs = lambda d: f"annotate_and_index/{d}_dataset_pipeline_task_group.make_kgx_{d}/" + + + for dataset in config.dug_inputs.data_sets: + dataset_name = dataset.split(":")[0] + # add datasets from the other pipeline + input_repos.append( + { + 'name': working_repo, + 'branch': branch, + 'path': get_path_on_lakefs(dataset_name) + } + ) + + merge_nodes = create_python_task (dag, name="MergeNodes", + a_callable=roger.merge_nodes, + external_repos=input_repos, + # merges every KGX node across the + # baseline graph and all dataset + # outputs in one task; OOMed at the + # chart default + memory="15Gi", + ) + + # The rest of these guys can just operate on the local lakefs repo/branch + # we need to add input dir and output dir similar to what we did for dug tasks + + create_nodes_schema = create_python_task(dag, + name="CreateNodesSchema", + a_callable=roger.create_nodes_schema + ) + create_edges_schema = create_python_task(dag, + name="CreateEdgesSchema", + a_callable=roger.create_edges_schema) + + create_bulk_load_nodes = create_python_task(dag, + name="CreateBulkLoadNodes", + a_callable=roger.create_bulk_nodes, + clear_output_prefix=True) + create_bulk_load_edges = create_python_task(dag, + name="CreateBulkLoadEdges", + a_callable=roger.create_bulk_edges, + clear_output_prefix=True, + # edges dominate memory; the + # rest run at the chart default + memory="15Gi") + bulk_load = create_python_task(dag, + name="BulkLoad", + a_callable=roger.bulk_load, + no_output_files=True, + # deletes the graph and reloads it whole, + # so it always needs the complete node and + # edge csv set -- an incremental pull hands + # it only what changed and it would rebuild + # the graph from that fragment + incremental_pull=False, + # the loader holds a node-id -> internal + # id map for every node in the graph so it + # can resolve edge endpoints; at 3.9M nodes + # and 78M edges the chart default 2Gi is + # not enough + memory="15Gi") + check_tranql = create_python_task(dag, + name="CheckTranql", + a_callable=roger.check_tranql, + no_output_files=True) + validate = create_python_task(dag, + name="Validate", + a_callable=roger.validate, + no_output_files=True) + + + """ Build the DAG. """ + merge_nodes.set_upstream(intro) + create_nodes_schema.set_upstream(merge_nodes) + create_edges_schema.set_upstream(merge_nodes) + create_bulk_load_nodes.set_upstream(create_nodes_schema) + create_bulk_load_nodes.set_upstream(merge_nodes) + create_bulk_load_edges.set_upstream(create_edges_schema) + create_bulk_load_edges.set_upstream(merge_nodes) + bulk_load.set_upstream(create_bulk_load_nodes) + bulk_load.set_upstream(create_bulk_load_edges) + validate.set_upstream(bulk_load) + check_tranql.set_upstream(bulk_load) + diff --git a/dags/metadata.yaml b/dags/metadata.yaml new file mode 100644 index 00000000..0cedb6a0 --- /dev/null +++ b/dags/metadata.yaml @@ -0,0 +1,206 @@ +kgx: + versions: + - files: + - biolink-v1.0.json + - ctd-v1.0.json + - gtopdb-v1.0.json + - hetio-v1.0.json + - hgnc-v1.0.json + - hmdb-v1.0.json + - kegg-v1.0.json + - mychem-v1.0.json + - ontological-hierarchy-v1.0.json + - panther-v1.0.json + - foodb-v1.0.json + - pharos-v1.0.json + - intact-v1.0.json + - human-goa-v1.0.json + - uberongraph-v1.0.json + - viral-proteome-v1.0.json + version: v1.0 + name: baseline-graph + format: json + - files: + - biolink-v2.0.json + - ctd-v2.0.json + - gtopdb-v2.0.json + - hetio-v2.0.json + - hgnc-v2.0.json + - hmdb-v2.0.json + - kegg-v2.0.json + - mychem-v2.0.json + - ontological-hierarchy-v2.0.json + - panther-v2.0.json + - foodb-v2.0.json + - pharos-v2.0.json + - intact-v2.0.json + - human-goa-v2.0.json + - uberongraph-v2.0.json + - viral-proteome-v2.0.json + version: v2.0 + name: baseline-graph + format: json + - files: + - heal/sparc/curation-export-processed.json + version: v2.0 + name: sparc-kgx + format: json + - files: + - Biolink_edges_v3.0.jsonl + - Biolink_nodes_v3.0.jsonl + - CTD_edges_v3.0.jsonl + - CTD_nodes_v3.0.jsonl + - DrugCentral_edges_v3.0.jsonl + - DrugCentral_nodes_v3.0.jsonl + - GtoPdb_edges_v3.0.jsonl + - GtoPdb_nodes_v3.0.jsonl + - Hetio_edges_v3.0.jsonl + - Hetio_nodes_v3.0.jsonl + - HGNC_edges_v3.0.jsonl + - HGNC_nodes_v3.0.jsonl + - HMDB_edges_v3.0.jsonl + - HMDB_nodes_v3.0.jsonl + - HumanGOA_edges_v3.0.jsonl + - HumanGOA_nodes_v3.0.jsonl + - IntAct_edges_v3.0.jsonl + - IntAct_nodes_v3.0.jsonl + - OntologicalHierarchy_edges_v3.0.jsonl + - OntologicalHierarchy_nodes_v3.0.jsonl + - PANTHER_edges_v3.0.jsonl + - PANTHER_nodes_v3.0.jsonl + - PHAROS_edges_v3.0.jsonl + - PHAROS_nodes_v3.0.jsonl + - UberGraph_edges_v3.0.jsonl + - UberGraph_nodes_v3.0.jsonl + version: v3.0 + name: baseline-graph + format: jsonl + - version: test + files: + - panther.json + name: test + - version: v3.0 + name: cde-graph + format: jsonl + files: + - cde/annotated_edges_v3.0.jsonl + - cde/annotated_nodes_v3.0.jsonl + - version: v4.0 + name: baseline-graph + format: jsonl + files: + - baseline-4.0/edges_v4.0.jsonl + - baseline-4.0/nodes_v4.0.jsonl + - version: v4.0 + name: cde-graph + format: jsonl + files: + - cde/annotated_edges_v4.0.jsonl + - cde/annotated_nodes_v4.0.jsonl + - version: v5.0 + name: baseline-graph + format: jsonl + files: + - baseline-5.0/edges_v5.0.jsonl + - baseline-5.0/nodes_v5.0.jsonl + - version: v5.0 + name: cde-graph + format: jsonl + files: + - cde/annotated_edges_v5.0.jsonl + - cde/annotated_nodes_v5.0.jsonl +dug_inputs: + versions: + - name: bdc + version: v1.0 + files: + s3: + - "bdc/v1.0/bdc_dbgap_data_dicts.tar.gz" + stars: + - "bdc_dbgap_data_dicts.tar.gz" + format: dbGaP + - name: bdc + version: v2.0 + files: + s3: + - "bdc/v2.0/bdc_dbgap_data_dicts.tar.gz" + stars: + - "bdc_dbgap_data_dicts.tar.gz" + format: dbGaP + - name: bdc + version: v3.0 + files: + s3: + - "bdc/v3.0/bdc_dbgap_data_dicts.tar.gz" + format: dbGaP + - name: nida + version: v1.0 + files: + s3: + - "nida/v1.0/nida-12studies.tar.gz" + stars: + - "nida-12studies.tar.gz" + format: nida + - name: sparc + version: v1.0 + files: + s3: + - "sparc/v1.0/sparc-dbgap-xml-formatted.tar.gz" + stars: + - "sparc-dbgap-xml-formatted.tar.gz" + format: sparc + - name: topmed + version: v2.0 + files: + s3: + - "topmed/v2.0/topmed_tags_v2.0.json" + - "topmed/v2.0/topmed_variables_v2.0.csv" + stars: + - topmed_variables_v2.0.csv + - topmed_tags_v2.0.json + format: topmed + - name: anvil + version: v1.0 + files: + s3: + - "bdc/v1.0/anvil_dbgap_data_dicts.tar.gz" + stars: + - "anvil_dbgap_data_dicts.tar.gz" + format: anvil + - name: kfdrc + version: v1.0 + files: + s3: + - "bdc/v1.0/KFDRC.tar.gz" + format: kfdrc + - name: crdc + version: v1.0 + files: + s3: + - "bdc/v1.0/CRDC.tar.gz" + format: crdc + - name: sprint + version: v1.0 + files: + s3: + - "sprint/v1.0/StanfordSPRINT_DataDictionary_2020-12-16.tar.gz" + format: sprint + - name: bacpac + version: v1.0 + files: + s3: + - "heal-datasets/bacpac/bacpac_baseline_do_measures.tar.gz" + format: bacpac + - name: heal-studies + version: v1.0 + files: + s3: + - heal-datasets/ingest-8-23/heal_studies.tar.gz + - heal-datasets/ingest-8-23/heal_mds_import.tar.gz + format: heal-studies + - name: heal-research-programs + version: v1.0 + files: + s3: + - heal-datasets/ingest-8-23/heal_research_programs.tar.gz + format: heal-research diff --git a/dags/test_dag.py b/dags/test_dag.py new file mode 100644 index 00000000..3073b30e --- /dev/null +++ b/dags/test_dag.py @@ -0,0 +1,37 @@ +"""Just a test dag to see if all the wrappers are working correctly. +""" + +from airflow.models import DAG +from airflow.providers.standard.operators.empty import EmptyOperator +from roger.tasks import default_args, create_python_task + +with DAG( + dag_id='test_dag', + default_args=default_args, + params= + { + "repository_id": None, + "branch_name": None, + "commitid_from": None, + "commitid_to": None + }, + # schedule_interval=None +) as dag: + + init = EmptyOperator(task_id="init", dag=dag) + finish = EmptyOperator(task_id="finish", dag=dag) + + def print_context(ds=None, **kwargs): + print(">>>All kwargs") + print(kwargs) + print(">>>All ds") + print(ds) + + (init >> + create_python_task(dag, "print_context", print_context) >> + finish) + + #run_this = PythonOperator(task_id="print_the_context", python_callable=print_context) + +if __name__ == "__main__": + dag.test() diff --git a/dags/test_metadata.yaml b/dags/test_metadata.yaml new file mode 100644 index 00000000..54d508c4 --- /dev/null +++ b/dags/test_metadata.yaml @@ -0,0 +1,124 @@ +# This is a file that lists the data to be used for testing purposes +# It contains a reduced set of the metadata.yaml file +kgx: + versions: + - files: + - biolink-v1.0.json + - ctd-v1.0.json + - gtopdb-v1.0.json + - hetio-v1.0.json + - hgnc-v1.0.json + - hmdb-v1.0.json + - kegg-v1.0.json + - mychem-v1.0.json + - ontological-hierarchy-v1.0.json + - panther-v1.0.json + - foodb-v1.0.json + - pharos-v1.0.json + - intact-v1.0.json + - human-goa-v1.0.json + - uberongraph-v1.0.json + - viral-proteome-v1.0.json + version: v1.0 + name: baseline-graph + format: json + - files: + - biolink-v2.0.json + - ctd-v2.0.json + - gtopdb-v2.0.json + - hetio-v2.0.json + - hgnc-v2.0.json + - hmdb-v2.0.json + - kegg-v2.0.json + - mychem-v2.0.json + - ontological-hierarchy-v2.0.json + - panther-v2.0.json + - foodb-v2.0.json + - pharos-v2.0.json + - intact-v2.0.json + - human-goa-v2.0.json + - uberongraph-v2.0.json + - viral-proteome-v2.0.json + version: v2.0 + name: baseline-graph + format: json + - files: + - heal/sparc/curation-export-processed.json + version: v2.0 + name: sparc-kgx + format: json + - files: + - Biolink_edges_v3.0.jsonl + - Biolink_nodes_v3.0.jsonl + - CTD_edges_v3.0.jsonl + - CTD_nodes_v3.0.jsonl + - DrugCentral_edges_v3.0.jsonl + - DrugCentral_nodes_v3.0.jsonl + - GtoPdb_edges_v3.0.jsonl + - GtoPdb_nodes_v3.0.jsonl + - Hetio_edges_v3.0.jsonl + - Hetio_nodes_v3.0.jsonl + - HGNC_edges_v3.0.jsonl + - HGNC_nodes_v3.0.jsonl + - HMDB_edges_v3.0.jsonl + - HMDB_nodes_v3.0.jsonl + - HumanGOA_edges_v3.0.jsonl + - HumanGOA_nodes_v3.0.jsonl + - IntAct_edges_v3.0.jsonl + - IntAct_nodes_v3.0.jsonl + - OntologicalHierarchy_edges_v3.0.jsonl + - OntologicalHierarchy_nodes_v3.0.jsonl + - PANTHER_edges_v3.0.jsonl + - PANTHER_nodes_v3.0.jsonl + - PHAROS_edges_v3.0.jsonl + - PHAROS_nodes_v3.0.jsonl + - UberGraph_edges_v3.0.jsonl + - UberGraph_nodes_v3.0.jsonl + version: v3.0 + name: baseline-graph + format: jsonl + - version: test + files: + - hgnc_nodes.jsonl + - hgnc_edges.jsonl + name: test + - version: v3.0 + name: cde-graph + format: jsonl + files: + - cde/annotated_edges_v3.0.jsonl + - cde/annotated_nodes_v3.0.jsonl +dug_inputs: + versions: + - name: bdc + version: v1.0 + files: + s3: + - "bdc/v1.0/bdc_dbgap_data_dicts.tar.gz" + stars: + - "bdc_dbgap_data_dicts.tar.gz" + format: dbGaP + - name: nida + version: v1.0 + files: + s3: + - "nida/v1.0/nida-12studies.tar.gz" + stars: + - "nida-12studies.tar.gz" + format: nida + - name: sparc + version: v1.0 + files: + s3: + - "sparc/v1.0/sparc-dbgap-xml-formatted.tar.gz" + stars: + - "sparc-dbgap-xml-formatted.tar.gz" + format: sparc + - name: anvil + version: v1.0 + files: + s3: + - "bdc/v1.0/anvil_dbgap_data_dicts.tar.gz" + stars: + - "anvil_dbgap_data_dicts.tar.gz" + format: anvil \ No newline at end of file diff --git a/docker-compose.yaml b/docker-compose.yaml new file mode 100644 index 00000000..9a52c670 --- /dev/null +++ b/docker-compose.yaml @@ -0,0 +1,167 @@ + +x-airflow-common: &airflow-common + build: . + environment: &airflow_common_environment + AIRFLOW__CORE__EXECUTOR: LocalExecutor + AIRFLOW__DATABASE__SQL_ALCHEMY_CONN: postgresql+psycopg2://airflow:airflow@postgres:5432/airflow + AIRFLOW__CORE__FERNET_KEY: '' + AIRFLOW__CORE__DAGS_ARE_PAUSED_AT_CREATION: 'true' + AIRFLOW__CORE__LOAD_EXAMPLES: 'false' + AIRFLOW__API__AUTH_BACKENDS: 'airflow.api.auth.backend.basic_auth,airflow.api.auth.backend.session' + AIRFLOW__SCHEDULER__ENABLE_HEALTH_CHECK: 'true' + AIRFLOW__CORE__SIMPLE_AUTH_MANAGER_USERS: 'admin:Admin' + ROGER_ELASTICSEARCH_HOST: "elasticsearch" + ROGER_ELASTICSEARCH_PASSWORD: "" + ROGER_ELASTICSEARCH_SCHEME: "http" + ROGER_ELASTICSEARCH_USERNAME: "elastic" + ROGER_REDISGRAPH_GRAPH: "test" + ROGER_REDISGRAPH_HOST: "redis-stack" + ROGER_REDISGRAPH_PASSWORD: "" + ROGER_REDISGRAPH_PORT: "6379" + ROGER_KGX_DATA__SETS: ${KGX_DATA_SETS} + ROGER_LAKEFS__CONFIG_ACCESS__KEY__ID: ${LAKEFS_ACCESS_KEY} + ROGER_LAKEFS__CONFIG_BRANCH: ${LAKEFS_BRANCH} + ROGER_LAKEFS__CONFIG_ENABLED: "true" + ROGER_LAKEFS__CONFIG_HOST: ${LAKEFS_URL} + ROGER_LAKEFS__CONFIG_REPO: ${LAKEFS_REPO} + ROGER_LAKEFS__CONFIG_SECRET__ACCESS__KEY: ${LAKEFS_SECRET_KEY} + ROGER_DUG__INPUTS_DATA__SETS: ${INPUT_DATA_SETS} + ROGER_ANNOTATION_ANNOTATOR__ARGS_SAPBERT_CLASSIFICATION__URL: ${BIOMEGATRON_URL} + ROGER_ANNOTATION_ANNOTATOR__ARGS_SAPBERT_ANNOTATOR__URL : ${SAPBERT_URL} + ROGER_ANNOTATION_NORMALIZER : ${NODE_NORM_URL} + ROGER_ANNOTATION_SYNONYM__SERVICE : ${NAME_RES_URL} + volumes: + - ./dags:/opt/airflow/dags + - ./logs:/opt/airflow/logs + - ./plugins:/opt/airflow/plugins + - ./config:/opt/airflow/config + - .:/opt/roger + depends_on: + postgres: + condition: service_healthy + networks: + - airflow-network + +services: + postgres: + image: postgres:15-alpine + environment: + POSTGRES_USER: airflow + POSTGRES_PASSWORD: airflow + POSTGRES_DB: airflow + volumes: + - postgres-db-volume:/var/lib/postgresql/data + healthcheck: + test: ["CMD", "pg_isready", "-U", "airflow"] + interval: 10s + retries: 5 + start_period: 5s + ports: + - "5432:5432" + networks: + - airflow-network + + # --- NEW: Elasticsearch Service --- + elasticsearch: + image: docker.elastic.co/elasticsearch/elasticsearch:8.11.1 + environment: + - discovery.type=single-node + - xpack.security.enabled=false + - ES_JAVA_OPTS=-Xms512m -Xmx512m # Limit RAM usage for dev + volumes: + - elasticsearch-data:/usr/share/elasticsearch/data + ports: + - "9200:9200" # REST API + - "9300:9300" # Internal transport + networks: + - airflow-network + + # --- NEW: Redis Stack (includes RedisGraph/FalkorDB) --- + redis-stack: + image: redis/redis-stack:latest + volumes: + - redis-stack-data:/data + ports: + - "6379:6379" # Redis port + - "8001:8001" # RedisInsight UI + healthcheck: + test: [ "CMD", "redis-cli", "ping" ] + interval: 10s + retries: 5 + start_period: 5s + networks: + - airflow-network + + airflow-init: + <<: *airflow-common + entrypoint: /bin/bash + command: + - -c + - | + mkdir -p /opt/airflow/logs /opt/airflow/dags /opt/airflow/plugins + airflow db migrate + depends_on: + postgres: + condition: service_healthy + + airflow-webserver: + <<: *airflow-common + command: airflow api-server + ports: + - "8080:8080" + healthcheck: + test: ["CMD", "curl", "--fail", "http://localhost:8080/api/v2/monitor/health"] + interval: 30s + timeout: 10s + retries: 5 + start_period: 30s + depends_on: + airflow-init: + condition: service_completed_successfully + + airflow-scheduler: + <<: *airflow-common + command: airflow scheduler + healthcheck: + test: ["CMD", "airflow", "jobs", "check", "--job-type", "SchedulerJob", "--hostname", "$${HOSTNAME}"] + interval: 30s + timeout: 10s + retries: 5 + start_period: 30s + depends_on: + airflow-init: + condition: service_completed_successfully + + airflow-triggerer: + <<: *airflow-common + command: airflow triggerer + healthcheck: + test: ["CMD-SHELL", 'airflow jobs check --job-type TriggererJob --hostname "$${HOSTNAME}"'] + interval: 30s + timeout: 10s + retries: 5 + start_period: 30s + depends_on: + airflow-init: + condition: service_completed_successfully + + airflow-dag-processor: + <<: *airflow-common + command: airflow dag-processor + healthcheck: + test: [ "CMD", "airflow", "jobs", "check", "--job-type", "DagProcessorJob", "--hostname", "$${HOSTNAME}" ] + interval: 30s + timeout: 10s + retries: 5 + start_period: 30s + depends_on: + airflow-init: + condition: service_completed_successfully + +volumes: + postgres-db-volume: + elasticsearch-data: # <-- New volume for Elasticsearch + redis-stack-data: +networks: + airflow-network: + driver: bridge diff --git a/ext/containers/FalkorDB/Dockerfile b/ext/containers/FalkorDB/Dockerfile new file mode 100644 index 00000000..e929a199 --- /dev/null +++ b/ext/containers/FalkorDB/Dockerfile @@ -0,0 +1,5 @@ +FROM falkordb/falkordb-server:v4.14.4-alpine +# Remove Go if present (apk-installed or tarball) +RUN apk del --no-network go 2>/dev/null || true \ + && rm -rf /usr/local/go /usr/lib/go /root/go /go 2>/dev/null || true +RUN apk add bash diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 00000000..862ab8f0 --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,11 @@ +[build-system] +requires = [ + "setuptools>=42", + "wheel" +] +build-backend = "setuptools.build_meta" + +[tool.pytest.ini_options] +testpaths = [ + "tests", +] diff --git a/requirements.txt b/requirements.txt index 0fe342a6..817aed9e 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,189 +1,32 @@ -alabaster==0.7.12 -alembic==1.4.2 -antlr4-python3-runtime==4.8 -apache-airflow==1.10.12 -apispec==1.3.3 -appnope==0.1.0 -argcomplete==1.12.0 -argon2-cffi==20.1.0 -async-generator==1.10 -attrs==19.3.0 -Babel==2.8.0 -backcall==0.2.0 -biolink-model==1.2.5 -biolinkml==1.5.8 -bleach==3.2.1 -bmt==0.1.1 -cached-property==1.5.1 -cachetools==4.1.1 -cattrs==1.0.0 -certifi==2020.6.20 -cffi==1.14.3 -CFGraph==0.2.1 -chardet==3.0.4 -click==7.1.2 -colorama==0.4.3 -colorlog==4.0.2 -configparser==3.5.3 -croniter==0.3.34 -decorator==4.4.2 -defusedxml==0.6.0 -dill==0.3.2 -dnspython==1.16.0 -docker==4.3.1 -docutils==0.16 -email-validator==1.1.1 -entrypoints==0.3 -env==0.1.0 -Flask==1.1.2 -Flask-Admin==1.5.4 -Flask-AppBuilder==2.3.4 -Flask-Babel==1.0.0 -Flask-Caching==1.3.3 -Flask-JWT-Extended==3.24.1 -Flask-Login==0.4.1 -Flask-OpenID==1.2.5 -Flask-SQLAlchemy==2.4.4 -flask-swagger==0.2.14 -Flask-WTF==0.14.3 -funcsigs==1.0.2 -future==0.18.2 -graphviz==0.14.1 -gunicorn==20.0.4 -idna==2.10 -imagesize==1.2.0 -importlib-metadata==1.7.0 -iniconfig==1.1.1 -ipykernel==5.3.4 -ipython==7.18.1 -ipython-genutils==0.2.0 -ipywidgets==7.5.1 -iso8601==0.1.12 -isodate==0.6.0 -itsdangerous==1.1.0 -jedi==0.17.2 -Jinja2==2.11.2 -json-merge-patch==0.2 -jsonasobj==1.2.1 -jsonlines==1.2.0 -jsonschema==3.2.0 -jupyter==1.0.0 -jupyter-client==6.1.7 -jupyter-console==6.2.0 -jupyter-core==4.6.3 -jupyterlab-pygments==0.1.2 -kgx==0.1.0 -lazy-object-proxy==1.5.1 -lockfile==0.12.2 -Mako==1.1.3 -Markdown==2.6.11 -MarkupSafe==1.1.1 -marshmallow==2.21.0 -marshmallow-enum==1.5.1 -marshmallow-sqlalchemy==0.23.1 -mistune==0.8.4 -mypy==0.790 -mypy-extensions==0.4.3 -natsort==7.0.1 -nbclient==0.5.1 -nbconvert==6.0.7 -nbformat==5.0.8 -neo4jrestclient==2.1.1 -nest-asyncio==1.4.1 -networkx==2.5 -notebook==6.1.4 -numpy==1.19.1 -ordered-set==4.0.2 -packaging==20.4 -pandas==1.1.0 -pandocfilters==1.4.2 -parso==0.7.1 -pathlib==1.0.1 -pathtools==0.1.2 -pbr==5.5.0 -pendulum==1.4.4 -pexpect==4.8.0 -pickleshare==0.7.5 -pluggy==0.13.1 -prefixcommons==0.1.9 -prison==0.1.3 -prologterms==0.0.6 -prometheus-client==0.8.0 -prompt-toolkit==3.0.8 -psutil==5.7.2 -PTable==0.9.2 -ptyprocess==0.6.0 -py==1.9.0 -pycparser==2.20 -Pygments==2.6.1 -PyJSG==0.10.0 -PyJWT==1.7.1 -pyparsing==2.4.7 -pyrsistent==0.16.0 -PyShEx==0.7.14 -PyShExC==0.8.2 -pystache==0.5.4 -pytest==6.1.1 -python-daemon==2.2.4 -python-dateutil==2.8.1 -python-editor==1.0.4 -python-nvd3==0.15.0 -python-slugify==4.0.1 -python3-openid==3.2.0 -pytz==2020.1 -pytzdata==2020.1 -PyYAML==5.3.1 -pyzmq==19.0.2 -qtconsole==4.7.7 -QtPy==1.9.0 -rdflib==5.0.0 -rdflib-jsonld==0.5.0 -redis==3.5.3 -redisgraph==2.1.5 -redisgraph-bulk-loader==0.9.3 -requests==2.24.0 -Send2Trash==1.5.0 -setproctitle==1.1.10 -ShExJSG==0.7.0 -six==1.15.0 -snowballstemmer==2.0.0 -sparql-slurper==0.3.4 -SPARQLWrapper==1.8.5 -Sphinx==3.2.1 -sphinx-click==2.5.0 -sphinx-rtd-theme==0.5.0 -sphinxcontrib-applehelp==1.0.2 -sphinxcontrib-devhelp==1.0.2 -sphinxcontrib-htmlhelp==1.0.3 -sphinxcontrib-jsmath==1.0.1 -sphinxcontrib-qthelp==1.0.3 -sphinxcontrib-serializinghtml==1.1.4 -SQLAlchemy==1.3.18 -SQLAlchemy-JSONField==0.9.0 -SQLAlchemy-Utils==0.36.8 -stringcase==1.2.0 -tabulate==0.8.7 -tenacity==4.12.0 -terminado==0.9.1 -terminaltables==3.1.0 -testpath==0.4.4 -text-unidecode==1.3 -thrift==0.13.0 -toml==0.10.1 -tornado==6.0.4 -traitlets==5.0.5 -typed-ast==1.4.1 -typing-extensions==3.7.4.2 -tzlocal==1.5.1 -unicodecsv==0.14.1 -urllib3==1.25.10 -validators==0.18.1 -watchdog==0.10.3 -wcwidth==0.2.5 -webencodings==0.5.1 -websocket-client==0.57.0 -Werkzeug==0.16.1 -widgetsnbextension==3.5.1 -WTForms==2.3.3 -zipp==3.1.0 -zope.deprecation==4.4.0 +elasticsearch>=8.5.2 +flatten-dict +jsonpickle +git+https://github.com/falkordb/falkordb-bulk-loader.git@v1.0.6 +setuptools>=66 +pytest +PyYAML +git+https://github.com/helxplatform/dug@develop +orjson>=3.11 +git+https://github.com/helxplatform/kg_utils.git@v0.0.10.1 +# kg_utils/merging.py hashes an f-string; xxhash 4.0 dropped the +# implicit str->bytes encoding, so every edge raises "Strings must +# be encoded before hashing" and MergeNodes dies on the edge sort. +# roger/models/kgx.py hashes a str the same way. +xxhash<4 +git+https://github.com/helxplatform/python-stringcase@1.2.1 +bmt==1.4.4 +git+https://github.com/helxplatform/avalon.git@lakefs-1.71.0 +h11>=0.16.0 +starlette<1.0.0 +datetime +# Latest redisgraph version +redis==5.3.1 +falkordb +#--- patch +aiohttp>=3.13.3 +werkzeug==3.1.5 +cryptography>=44.0.1 +urllib3>=2.6.3 +jaraco.context==6.1.0 +marshmallow==3.26.2 # upgrade from 3.26.1 specified in airflow constraints file + diff --git a/roger-cli-steps.md b/roger-cli-steps.md new file mode 100644 index 00000000..8e132746 --- /dev/null +++ b/roger-cli-steps.md @@ -0,0 +1,27 @@ +# Deployment with Roger CLI + +## QUICK Local Set Up + +This is list steps to produce a local deployment of Roger. This set up does NOT use airflow and instead only uses the Roger CLI via **Makefile** commands. + +### Prerequsite Steps + +- Set up Roger dependencies by ensuring that the `.env` has all the correct information. +- Run the following docker compose commands + - `docker compose up tranql -d`: starts up tranql which is the API handlerfor redis graph in the `graph` stage + - `docker compose up redis -d`: starts up redis which will be used via redis graph for the `graph` stage + - `docker compose up dug -d`: starts up dug API to work as the API handler for elastic search in the `index` stage + - `docker compose up elasticsearch -d`: starts up elastic search for the `index` stage + +### Roger CLI Steps + +1) `python3 -m venv ~/.environments/roger` +2) `source ~/.environments/roger/bin/activate` +3) `pip install -r requirements.txt` +4) `export PYTHONPATH=$PWD/dags` +5) Change the elasticsearch and redisgraph `host` values to localhost in `dags/roger/config/config.yaml` +6) Get the S3 Bucket credentials (access_key, bucket, host, secret_key) and export them as environment variables with ROGER_S3_ in the front of the value like: `ROGER_S3_ACCESS__KEY=XXXXKEYXXXX` +7) `cd bin/` and here either run `make all` OR separate the commands into three steps: + 1) `make annotate`: executes the CLI related commands found in `bin/dug_annotate/Makefile` + 2) `make graph`: executes the CLI related commands found in `bin/roger_graph_build/Makefile` + 3) `make index`: executes the CLI related commands found in `bin/dug_index/Makefile` diff --git a/roger/config.yaml b/roger/config.yaml deleted file mode 100644 index f1377a3c..00000000 --- a/roger/config.yaml +++ /dev/null @@ -1,56 +0,0 @@ -redisgraph: - username: "" - password: "" - host: localhost - graph: test - ports: - http: 6379 - -logging: - level: DEBUG - format: '[%(name)s][%(filename)s][%(funcName)20s] %(levelname)s: %(message)s' - -data_root: roger/data -base_data_uri: https://stars.renci.org/var/kgx_data - -#https://github.com/RedisGraph/redisgraph-bulk-loader/blob/master/redisgraph_bulk_loader/bulk_insert.py#L43 -bulk_loader: - separator: "|" - enforce_schema: False - skip_invalid_nodes: False - skip_invalid_edges: False - quote: 0 - max_token_count: 1024 - max_buffer_size: 2048 - max_token_size: 500 - index: [] - full_text_index: [] - -validation: - queries: - count_nodes: - name: "Count Nodes" - query: "MATCH (a) RETURN COUNT(a)" - count_edges: - name: "Count Edges" - query: "MATCH (a)-[e]-(b) RETURN COUNT(e)" - connectivity: - name: TOPMED Connectivity - query: "MATCH (a { id : '$var' })--(b) RETURN a.category, b.id" - args: - - var: TOPMED.TAG:8 - - var: TOPMED.VAR:phv00000484.v1.p10 - - var: TOPMED.VAR:phv00000487.v1.p10 - - var: TOPMED.VAR:phv00000496.v1.p10 - - var: TOPMED.VAR:phv00000517.v1.p10 - - var: TOPMED.VAR:phv00000518.v1.p10 - - var: TOPMED.VAR:phv00000528.v1.p10 - - var: TOPMED.VAR:phv00000529.v1.p10 - - var: TOPMED.VAR:phv00000530.v1.p10 - - var: TOPMED.VAR:phv00000531.v1.p10 - count_connected_nodes: - name: Count Connected Nodes - query: "MATCH (a)-[e]-(b) RETURN count(a), count(b)" - query_by_type: - name: Query by Type - query: "MATCH (a:gene)-[e]-(b) WHERE 'chemical_substance' IN b.category RETURN count(distinct(a)), count(distinct(b))" diff --git a/roger/core.py b/roger/core.py deleted file mode 100644 index 3d354071..00000000 --- a/roger/core.py +++ /dev/null @@ -1,688 +0,0 @@ -import argparse -import glob -import json -import os -import redis -import requests -import shutil -import time -import yaml -import sys -import traceback -from biolink import model -from collections import defaultdict -from enum import Enum -from io import StringIO -from kgx.cli import redisgraph_upload -from roger.roger_util import get_logger, get_config -from redisgraph_bulk_loader.bulk_insert import bulk_insert -from roger.roger_db import RedisGraph -from string import Template - -log = get_logger () -config = get_config () -data_root = config['data_root'] - -class SchemaType(Enum): - """ High level semantic metatdata concepts. - Categories are classes in an ontological model like Biolink. - Predicates are links between nodes. """ - CATEGORY = "category" - PREDICATE = "predicate" - -class FileFormat(Enum): - """ File formats this module knows about. """ - JSON = "json" - YAML = "yaml" - -class Util: - - @staticmethod - def current_time_in_millis(): - """ - Get current time in milliseconds. - - Returns - ------- - int - Time in milliseconds - - """ - return int(round(time.time() * 1000)) - - """ A just do it approach to getting data. """ - @staticmethod - def read_file(path): - """ Read a file. - :param path: Path to a file. - """ - text = None - with open(path, "r") as stream: - text = stream.read () - return text - - @staticmethod - def read_url(url): - """ Read data from a URL. - :param url: The URL to read. """ - return requests.get (url).text - - @staticmethod - def read_data(path): - """ Read data from a URL or File. HTTP(S) is the only supported protocol. - :param path: A URL or file path. """ - text = None - if Util.is_web(path): - text = Util.read_url (path) - else: - text = Util.read_file (path) - return text - - @staticmethod - def read_object(path, key=None): - """ Read on object from a path. - :param path: A URL or file path. Supports YAML and JSON depending on extension. - :param key: A configuration key. This is prepended to the path if present. - :raises ValueError: If the key is not in the configuration. """ - if key is not None: - prefix = config[key] - path = f"{prefix}/{path}" if Util.is_web(prefix) \ - else os.path.join (prefix, path) - obj = None - if path.endswith (".yaml") or path.endswith (".yml"): - obj = yaml.safe_load (Util.read_data (path)) - elif path.endswith (".json"): - obj = json.loads (Util.read_data (path)) - return obj - - @staticmethod - def is_web (uri): - """ The URI is a web URI (starts with http or https). - :param uri: A URI """ - return uri.startswith("http://") or uri.startswith ("https://") - - @staticmethod - def write_object (obj, path, key=None): - """ Write an object to a path. YAML and JSON supported based on extension. - :param obj: The object to write. - :param path: The path to write to. - :param key: The configuration key to prepend to the path. - """ - """ Prepend a prefix from the configuration file if a key is given. """ - if key is not None: - prefix = config[key] - path = f"{prefix}/{path}" if Util.is_web(prefix) \ - else os.path.join (prefix, path) - """ Ensure the directory to be written to exists. """ - dirname = os.path.dirname (path) - if not os.path.exists (dirname): - os.makedirs (dirname, exist_ok=True) - """ Write the file in the specified format. """ - if path.endswith (".yaml") or path.endswith (".yml"): - with open(path, 'w') as outfile: - yaml.dump (obj, stream) - elif path.endswith (".json"): - with open (path, "w") as stream: - json.dump (obj, stream, indent=2) - else: - """ Raise an exception if invalid. """ - raise ValueError (f"Unrecognized extension: {path}") - - @staticmethod - def kgx_path (name): - """ Form a KGX object path. - :path name: Name of the KGX object. """ - return os.path.join (data_root, "kgx", name) - - @staticmethod - def kgx_objects (): - """ A list of KGX objects. """ - kgx_pattern = Util.kgx_path("**.json") - return sorted(glob.glob (kgx_pattern)) - - @staticmethod - def merge_path (name): - """ Form a merged KGX object path. - :path name: Name of the merged KGX object. """ - return os.path.join (data_root, "merge", name) - - @staticmethod - def merged_objects (): - """ A list of merged KGX objects. """ - merged_pattern = Util.merge_path("**.json") - return sorted(glob.glob (merged_pattern)) - - @staticmethod - def schema_path (name): - """ Path to a schema object. - :param name: Name of the object to get a path for. """ - return os.path.join (data_root, "schema", name) - - @staticmethod - def bulk_path (name): - """ Path to a bulk load object. - :param name: Name of the object. """ - return os.path.join (data_root, "bulk", name) - - @staticmethod - def read_schema (schema_type: SchemaType): - """ Read a schema object. - :param schema_type: Schema type of the object to read. """ - path = Util.schema_path (f"{schema_type.value}-schema.json") - return Util.read_object (path) - - @staticmethod - def get_uri (path, key): - """ Build a URI. - :param path: The path of an object. - :param key: The key of a configuration value to prepend to the object. """ - return f"{config[key]}/{path}" - - @staticmethod - def get_relative_path (path): - return os.path.join (os.path.dirname (__file__), path) - - @staticmethod - def read_relative_object (path): - return Util.read_object (Util.get_relative_path(path)) - - @staticmethod - def trunc(text, limit): - return ('..' + text[-limit-2:]) if len(text) > limit else text - - @staticmethod - def is_up_to_date (source, targets): - target_time_list = [ os.stat (f).st_mtime for f in targets if os.path.exists(f) ] - if len(target_time_list) == 0: - log.debug (f"no targets found") - return False - source = [ os.stat (f).st_mtime for f in source if os.path.exists (f) ] - if len(source) == 0: - log.debug ("no source found. up to date") - return True - return max(source) < min(target_time_list) - -class KGXModel: - """ Abstractions for transforming Knowledge Graph Exchange formatted data. """ - def __init__(self, biolink): - self.biolink = biolink - - def get (self, dataset_version = "v0.1"): - """ Read metadata for edge and node files, then join them into whole KGX objects - containing both nodes and edges. - :param dataset_version: Data version to operate on. - """ - metadata = Util.read_relative_object ("metadata.yaml") - for item in metadata['versions']: - if item['version'] == dataset_version: - for edge_url in item['edgeFiles']: - start = Util.current_time_in_millis () - edge_url = Util.get_uri (edge_url, "base_data_uri") - node_url = edge_url.replace ("-edge-", "-node-") - subgraph_basename = os.path.basename (edge_url.replace ("-edge", "")) - subgraph_path = Util.kgx_path (subgraph_basename) - if os.path.exists (subgraph_path): - log.info (f"cached kgx: {subgraph_path}") - continue - subgraph = { - "edges" : Util.read_object (edge_url), - "nodes" : Util.read_object (node_url) - } - Util.write_object (subgraph, subgraph_path) - total_time = Util.current_time_in_millis () - start - - edges = len(subgraph['edges']) - nodes = len(subgraph['nodes']) - log.debug ("wrote {:>45}: edges:{:>7} nodes: {:>7} time:{:>8}".format ( - Util.trunc(subgraph_path, 45), edges, nodes, total_time)) - - def create_schema (self): - """ - Determine the schema of each type of object. We have to do this to make it possible - to write tabular data. Need to know all possible columns in advance and correct missing - fields. - """ - if self.schema_up_to_date(): - log.info (f"schema is up to date.") - return - - predicate_schemas = defaultdict(lambda:None) - category_schemas = defaultdict(lambda:None) - for subgraph in Util.kgx_objects (): - """ Read a kgx data file. """ - log.debug (f"analyzing schema of {subgraph}.") - basename = os.path.basename (subgraph).replace (".json", "") - graph = Util.read_object (subgraph) - """ Infer predicate schemas. """ - for edge in graph['edges']: - predicate = edge['edge_label'] - if not predicate in predicate_schemas: - predicate_schemas[predicate] = edge - for k in edge.keys (): - edge[k] = '' - else: - for k in edge.keys (): - if not k in predicate_schemas[predicate]: - predicate_schemas[predicate][k] = '' - """ Infer node schemas. """ - for node in graph['nodes']: - node_type = self.biolink.get_leaf_class (node['category']) - if not node_type in category_schemas: - category_schemas[node_type] = node - for k in node.keys (): - node[k] = '' - else: - for k in node.keys (): - if not k in category_schemas[node_type]: - category_schemas[node_type][k] = '' - """ Write node and predicate schemas. """ - self.write_schema (predicate_schemas, SchemaType.PREDICATE) - self.write_schema (category_schemas, SchemaType.CATEGORY) - - def schema_up_to_date (self): - return Util.is_up_to_date ( - source=Util.kgx_objects (), - targets=[ - Util.schema_path (f"{SchemaType.PREDICATE.value}-schema.json"), - Util.schema_path (f"{SchemaType.PREDICATE.value}-schema.json") - ]) - - def write_schema (self, schema, schema_type: SchemaType): - """ Output the schema file. - :param schema: Schema to get keys from. - :param schema_type: Type of schema to write. """ - file_name = Util.schema_path (f"{schema_type.value}-schema.json") - log.info (f"writing schema: {file_name}") - dictionary = { k : self.format_keys(v.keys(), schema_type) for k, v in schema.items () } - Util.write_object (dictionary, file_name) - - def merge_nodes (self, L, R): - for k in L.keys (): - R_v = R.get (k, None) - if R_v == '' or R_v == None: - L[k] = R_v - - def diff_lists (self, L, R): - return list(list(set(L)-set(R)) + list(set(R)-set(L))) - - def merge (self): - """ Merge nodes. Would be good to have something less computationally intensive. """ - for path in Util.kgx_objects (): - new_path = path.replace ('/kgx/', '/merge/') - - source_stats = os.stat (path) - if os.path.exists (new_path): - dest_stats = os.stat (new_path) - if dest_stats.st_mtime > source_stats.st_mtime: - log.info (f"merge {new_path} is up to date.") - continue - - log.info (f"merging {path}") - graph = Util.read_object (path) - graph_nodes = graph.get ('nodes', []) - graph_map = { n['id'] : n for n in graph_nodes } - graph_keys = graph_map.keys () - total_merge_time = 0 - for path_2 in Util.kgx_objects (): - if path_2 == path: - continue - start = Util.current_time_in_millis () - other_graph = Util.read_object (path_2) - load_time = Util.current_time_in_millis () - start - - start = Util.current_time_in_millis () - other_nodes = other_graph.get('nodes', []) - other_map = { n['id'] : n for n in other_nodes } - other_keys = set(other_map.keys()) - intersection = [ v for v in graph_keys if v in other_keys ] - difference = list(set(other_keys) - set(graph_keys)) - scope_time = Util.current_time_in_millis () - start - - start = Util.current_time_in_millis () - for i in intersection: - self.merge_nodes (graph_map[i], other_map[i]) - other_graph['nodes'] = [ other_map[i] for i in difference ] - merge_time = Util.current_time_in_millis () - start - - start = Util.current_time_in_millis () - Util.write_object (other_graph, path_2.replace ('kgx', 'merge')) - write_time = Util.current_time_in_millis () - start - log.debug ("merged {:>45} load:{:>5} scope:{:>7} merge:{:>3}".format( - Util.trunc(path_2, 45), load_time, scope_time, merge_time)) - total_merge_time += load_time + scope_time + merge_time + write_time - - start = Util.current_time_in_millis () - Util.write_object (graph, new_path) - rewrite_time = Util.current_time_in_millis () - start - log.info (f"{path} rewrite: {rewrite_time}. total merge time: {total_merge_time}") - - def format_keys (self, keys, schema_type : SchemaType): - """ Format schema keys. Make source and destination first in edges. Make - id first in nodes. Remove keys for fields we can't yet represent. - :param keys: List of keys. - :param schema_type: Type of schema to conform to. - """ - """ Sort keys. """ - k_list = sorted(keys) - if schema_type == SchemaType.PREDICATE: - """ Rename subject and object to src and dest """ - k_list.remove ('subject') - k_list.remove ('object') - k_list.insert (0, 'src') - k_list.insert (1, 'dest') - elif schema_type == SchemaType.CATEGORY: - """ Make id the first field. Remove smiles. It causes ast parse errors. - TODO: update bulk loader to ignore AST on selected fields. - """ - k_list.remove ('id') - if 'simple_smiles' in k_list: - k_list.remove ('simple_smiles') - k_list.insert (0, 'id') - return k_list - - def load (self): - """ Use KGX to load a data set into Redisgraph """ - input_format = "json" - uri = f"redis://{config['redisgraph']['host']}:{config['redisgraph']['ports']['http']}/" - username = config['redisgraph']['username'] - password = config['redisgraph']['password'] - log.info (f"connecting to redisgraph: {uri}") - for subgraph in glob.glob (f"{kgx_repo}/**.json"): - redisgraph_upload(inputs=[ subgraph ], - input_format=input_format, - input_compression=None, - uri=uri, - username=username, - password=password, - node_filters=[], - edge_filters=[]) - -class BiolinkModel: - """ Programmatic model of Biolink. """ - def to_camel_case(self, snake_str): - """ Convert a snake case string to camel case. """ - components = snake_str.split('_') - return ''.join(x.title() for x in components) - - def get_class(self, name): - """ Get a Python class from a string name. """ - return getattr(sys.modules["biolink.model"], name) - - def is_derived (self, a_class_name, classes): - """ Return true if the class derives from any of the provided classes. """ - for c in classes: - if isinstance (self.get_class(self.to_camel_case(a_class_name)), c): - return True - return False - - def get_leaf_class (self, names): - """ Return the leaf classes in the provided list of names. """ - classes = [ self.get_class(self.to_camel_case(n)) for n in names ] - leaves = [ n for n in names if not self.is_derived (n, classes) ] - return leaves [0] - -class BulkLoad: - """ Tools for creating a Redisgraph bulk load dataset. """ - def __init__(self, biolink): - self.biolink = biolink - - def tables_up_to_date (self): - return Util.is_up_to_date ( - source=[ - Util.schema_path (f"{SchemaType.PREDICATE.value}-schema.json"), - Util.schema_path (f"{SchemaType.PREDICATE.value}-schema.json") - ] + Util.merged_objects (), - targets=glob.glob (Util.bulk_path ("nodes/**.csv")) + \ - glob.glob (Util.bulk_path ("edges/**.csv"))) - - def create (self): - """ Check source times. """ - if self.tables_up_to_date (): - log.info ("up to date.") - return - - """ Format the data for bulk load. """ - predicates_schema = Util.read_schema (SchemaType.PREDICATE) - categories_schema = Util.read_schema (SchemaType.CATEGORY) - bulk_path = Util.bulk_path("") - if os.path.exists(bulk_path): - shutil.rmtree(bulk_path) - - state = defaultdict(lambda:None) - for subgraph in Util.merged_objects (): - log.info (f"processing {subgraph}") - graph = Util.read_object (subgraph) - - """ Write node data for bulk load. """ - categories = defaultdict(lambda: []) - for node in graph['nodes']: - index = self.biolink.get_leaf_class (node['category']) - categories[index].append (node) - self.write_bulk (Util.bulk_path("nodes"), categories, categories_schema, - state=state, f=subgraph) - - """ Write predicate data for bulk load. """ - predicates = defaultdict(lambda: []) - for edge in graph['edges']: - predicates[edge['edge_label']].append (edge) - edge['src'] = edge.pop ('subject') - edge['dest'] = edge.pop ('object') - self.write_bulk (Util.bulk_path("edges"), predicates, predicates_schema) - - def cleanup (self, v): - """ Filter problematic text. - :param v: A value to filter and clean. - """ - if isinstance(v, list): - v = [ self.cleanup(val) for val in v ] - elif isinstance (v, str): - """ Some values contain the CSV separator character. 'fix' that. """ - if len(v) > 1 and v[0] == '[' and v[-1] == ']': - v = v.replace ("[", "@").replace ("]", "@") #f" {v}" - v = v.replace ("|","^") - return v - - def write_bulk (self, bulk_path, obj_map, schema, state={}, f=None): - """ Write a bulk load group of objects. - :param bulk_path: Path to the bulk loader object to write. - :param obj_map: A map of biolink type to list of objects. - :param schema: The schema (nodes or predicates) containing identifiers. - :param state: Track state of already written objects to avoid duplicates. - """ - os.makedirs (bulk_path, exist_ok=True) - for key, objects in obj_map.items (): - out_file = f"{bulk_path}/{key}.csv" - if len(objects) == 0: - continue - new_file = not os.path.exists (out_file) - all_keys = schema[key] - with open (out_file, "a") as stream: - if new_file: - log.info (f" --creating {out_file}") - stream.write ("|".join (all_keys)) - stream.write ("\n") - """ Make all objects conform to the schema. """ - for obj in objects: - for akey in all_keys: - if not akey in obj: - obj[akey] = "" - """ Write fields, skipping duplicate objects. """ - for obj in objects: - oid = str(obj['id']) - if oid in state: - continue - state[oid] = oid - values = [ self.cleanup(obj[k]) for k in all_keys if not 'smiles' in k ] - clean = list(map(str, values)) - s = "|".join (clean) - stream.write (s) - stream.write ("\n") - - def insert (self): - redisgraph = config.get('redisgraph', {}) - bulk_loader = config.get('bulk_loader', {}) - nodes = sorted(glob.glob (Util.bulk_path ("nodes/**.csv"))) - edges = sorted(glob.glob (Util.bulk_path ("edges/**.csv"))) - graph = redisgraph['graph'] - log.info (f"bulk loading \n nodes: {nodes} \n edges: {edges}") - print (f"bulk loading \n nodes: {nodes} \n edges: {edges}") - - try: - log.info (f"deleting graph {graph} in preparation for bulk load.") - db = self.get_redisgraph (redisgraph) - db.redis_graph.delete () - except redis.exceptions.ResponseError: - log.info ("no graph to delete") - - log.info (f"bulk loading graph: {graph}") - args = [] - if len(nodes) > 0: - args.extend (("-n " + " -n ".join (nodes)).split ()) - if len(edges) > 0: - args.extend (("-r " + " -r ".join (edges)).split ()) - args.extend ([ "--separator=|" ]) - args.extend ([ redisgraph['graph'] ]) - """ standalone_mode=False tells click not to sys.exit() """ - bulk_insert (args, standalone_mode=False) - - def get_redisgraph (self, redisgraph): - return RedisGraph (host=redisgraph['host'], - port=redisgraph['ports']['http'], - graph=redisgraph['graph']) - - def validate (self): - redisgraph = config.get('redisgraph', {}) - print (f"config:{json.dumps(redisgraph, indent=2)}") - db = self.get_redisgraph (redisgraph) - validation_queries = config.get('validation', {}).get('queries', []) - for key, query in validation_queries.items (): - text = query['query'] - name = query['name'] - args = query.get('args', [{}]) - for arg in args: - start = Util.current_time_in_millis () - instance = Template (text).safe_substitute (arg) - db.query (instance) - duration = Util.current_time_in_millis () - start - log.info (f"Query {key}:{name} ran in {duration}ms: {instance}") - -class Roger: - """ Consolidate Roger functionality for a cleaner interface. """ - - def __init__(self, to_string=False): - """ Initialize. - :param to_string: Log messages to a string, available as self.log_stream.getvalue() - after execution completes. - """ - import logging - if to_string: - """ Add a stream handler to enable to_string. """ - self.log_stream = StringIO() - self.string_handler = logging.StreamHandler (self.log_stream) - log.addHandler (self.string_handler) - self.biolink = BiolinkModel () - self.kgx = KGXModel (self.biolink) - self.bulk = BulkLoad (self.biolink) - - def __enter__(self): - """ Implement Python's Context Manager interface. """ - return self - - def __exit__(self, exception_type, exception_value, traceback): - """ Implement Python's Context Manager interface. We use this finalizer - to detach the stream handler appended in the constructor. - :param exception_type: Type of exception, if one occurred. - :param exception_value: The exception, if one occurred. - :param traceback: The stack trace explaining the exception. - """ - if exception_type or exception_value or traceback: - log.error ("{} {} {}".format (exception_type, exception_value, traceback)) - log.removeHandler (self.string_handler) - -class RogerUtil: - """ An interface abstracting Roger's inner workings to make it easier to - incorporate into external tools like workflow engines. """ - @staticmethod - def get_kgx (to_string=False): - output = None - with Roger (to_string) as roger: - roger.kgx.get () - output = roger.log_stream.getvalue () if to_string else None - return output - - @staticmethod - def create_schema (to_string=False): - output = None - with Roger (to_string) as roger: - roger.kgx.create_schema () - output = roger.log_stream.getvalue () if to_string else None - return output - - @staticmethod - def merge_nodes (to_string=False): - output = None - with Roger (to_string) as roger: - roger.kgx.merge () - output = roger.log_stream.getvalue () if to_string else None - return output - - @staticmethod - def create_bulk_load (to_string=False): - output = None - with Roger (to_string) as roger: - roger.bulk.create () - output = roger.log_stream.getvalue () if to_string else None - return output - - @staticmethod - def bulk_load (to_string=False): - output = None - with Roger (to_string) as roger: - roger.bulk.insert () - output = roger.log_stream.getvalue () if to_string else None - return output - - @staticmethod - def validate (to_string=False): - output = None - with Roger (to_string) as roger: - roger.bulk.validate () - output = roger.log_stream.getvalue () if to_string else None - return output - -if __name__ == "__main__": - """ Roger CLI. """ - parser = argparse.ArgumentParser(description='Roger') - parser.add_argument('-v', '--dataset-version', help="Dataset version.", default="v0.1") - parser.add_argument('-d', '--data-root', help="Root of data hierarchy", default=None) - parser.add_argument('-g', '--get-kgx', help="Get KGX objects", action='store_true') - parser.add_argument('-l', '--load-kgx', help="Load via KGX", action='store_true') - parser.add_argument('-s', '--create-schema', help="Infer schema", action='store_true') - parser.add_argument('-m', '--merge-kgx', help="Merge KGX nodes", action='store_true') - parser.add_argument('-b', '--create-bulk', help="Create bulk load", action='store_true') - parser.add_argument('-i', '--insert', help="Do the bulk insert", action='store_true') - parser.add_argument('-a', '--validate', help="Validate the insert", action='store_true') - args = parser.parse_args () - - biolink = BiolinkModel () - kgx = KGXModel (biolink) - bulk = BulkLoad (biolink) - if args.data_root is not None: - data_root = get_config()['data_root'] = args.data_root - log.info (f"data root:{data_root}") - if args.get_kgx: - kgx.get (dataset_version=args.dataset_version) - if args.load_kgx: - kgx.load () - if args.merge_kgx: - kgx.merge () - if args.create_schema: - kgx.create_schema () - if args.create_bulk: - bulk.create () - if args.insert: - bulk.insert () - if args.validate: - bulk.validate () - - sys.exit (0) diff --git a/roger/metadata.yaml b/roger/metadata.yaml deleted file mode 100644 index 1b56206e..00000000 --- a/roger/metadata.yaml +++ /dev/null @@ -1,37 +0,0 @@ -versions: -- edgeFiles: - # - biolink_kgx-edge-v0.1.json - - chembio_kgx-edge-v0.1.json - - chemical_normalization-edge-v0.1.json - - cord19-phenotypes-edge-v0.1.json -# - cord19-scibite-edge-v0.1.json -# - cord19-scigraph-edge-v0.1.json - - ctd-edge-v0.1.json - - foodb-edge-v0.1.json -# - kegg-edge-v0.1.json - - mychem-edge-v0.1.json -# - panther-edge-v0.1.json - - pharos-edge-v0.1.json - - topmed-edge-v0.1.json - nodeFiles: - - biolink_kgx-node-v0.1.json - - chembio_kgx-node-v0.1.json - - chemical_normalization-node-v0.1.json - - cord19-phenotypes-node-v0.1.json - - cord19-scibite-node-v0.1.json - - cord19-scigraph-node-v0.1.json - - ctd-node-v0.1.json - - foodb-node-v0.1.json - # - kegg-node-v0.1.json - - mychem-node-v0.1.json - - panther-node-v0.1.json - - pharos-node-v0.1.json - - topmed-node-v0.1.json - version: v0.1 -- version: test - edgeFiles: - - cord19-phenotypes-edge-v0.1.json - - chembio_kgx-edge-v0.1.json - nodeFiles: - - cord19-phenotypes-node-v0.1.json - - chembio_kgx-node-v0.1.json diff --git a/roger/roger_util.py b/roger/roger_util.py deleted file mode 100644 index 35c5f3f7..00000000 --- a/roger/roger_util.py +++ /dev/null @@ -1,58 +0,0 @@ -import logging -import requests -import sys -import yaml -from os import path -from typing import Dict, Any, Optional - -config: Optional[Dict[str, Any]] = None -logger: Optional[logging.Logger] = None - -CONFIG_FILENAME = path.join(path.dirname(path.abspath(__file__)), 'config.yaml') - -def get_config(filename: str = CONFIG_FILENAME) -> dict: - """ - Get config as a dictionary - - Parameters - ---------- - filename: str - The filename with all the configuration - - Returns - ------- - dict - A dictionary containing all the entries from the config YAML - - """ - global config - if config is None: - config = yaml.load(open(filename), Loader=yaml.FullLoader) - return config - -def get_logger(name: str = 'roger') -> logging.Logger: - """ - Get an instance of logger. - - Parameters - ---------- - name: str - The name of logger - - Returns - ------- - logging.Logger - An instance of logging.Logger - - """ - global logger - if logger is None: - config = get_config() - logger = logging.getLogger(name) - handler = logging.StreamHandler(sys.stdout) - formatter = logging.Formatter(config['logging']['format']) - handler.setFormatter(formatter) - logger.addHandler(handler) - logger.setLevel(config['logging']['level']) - logger.propagate = False - return logger diff --git a/scripts/migrate_pickled_classes.py b/scripts/migrate_pickled_classes.py new file mode 100644 index 00000000..a6e06b11 --- /dev/null +++ b/scripts/migrate_pickled_classes.py @@ -0,0 +1,372 @@ +#!/usr/bin/env python +"""Bring jsonpickle artifacts written by an older dug up to the current classes. + +Dug's data classes moved into the dug_data_model library, and the module they +used to live in (dug.core.parsers._base) no longer imports. jsonpickle does not +raise on that -- it hands back the raw dict -- so stale files index as dicts and +blow up with "'dict' object has no attribute 'id'". + + python scripts/migrate_pickled_classes.py --scan # what's in there + python scripts/migrate_pickled_classes.py --fix # rewrite in place + python scripts/migrate_pickled_classes.py --restamp # class paths only, + # for huge trees + python scripts/migrate_pickled_classes.py --self-check # no dir needed + +Run inside the roger image so the current dug/dug_data_model are importable. +""" + +import argparse +import importlib +import os +import json +import re +import sys +import types +from pathlib import Path + +import jsonpickle +from pydantic import BaseModel + +from dug.core import DugConcept, DugVariable, DugStudy, DugSection +from dug.core.annotators import DugIdentifier +from dug.core.parsers import DugElement + +# Class name -> the class as it exists today. Legacy module paths resolve +# against this by name, so a class moving again needs no new mapping here. +CURRENT = {c.__name__: c for c in ( + DugElement, DugConcept, DugVariable, DugStudy, DugSection, DugIdentifier)} + +PY_OBJECT = re.compile(r'"py/object":\s*"([^"]+)"') +ARTIFACTS = ('elements.txt', 'concepts.txt', 'expanded_concepts.txt') + + +def install_alias(module_path): + """Stand in for a legacy module, resolving class names against CURRENT.""" + shim = types.ModuleType(module_path) + # PEP 562: unknown attribute lookups on the module land here + shim.__getattr__ = lambda name: CURRENT[name] + sys.modules[module_path] = shim + + +def classes_in(text): + return set(PY_OBJECT.findall(text)) + + +def broken_modules(class_paths): + """Of the modules these classes claim to live in, which no longer import.""" + broken = {} + for path in class_paths: + module_path, _, class_name = path.rpartition('.') + if module_path in broken or module_path in sys.modules: + continue + try: + importlib.import_module(module_path) + except Exception as exc: # noqa: BLE001 - any import failure counts + broken[module_path] = f"{type(exc).__name__}: {exc}" + return broken + + +def fill_defaults(obj, seen=None): + """Add fields the current model has that the old state never wrote. + + jsonpickle restores via __setstate__, which assigns __dict__ wholesale, so + fields added since the file was written are simply absent. model_construct() + gives us the defaults without running validation over the whole graph. + """ + if seen is None: + seen = set() + if id(obj) in seen: + return + seen.add(id(obj)) + + if isinstance(obj, BaseModel): + defaults = type(obj).model_construct().__dict__ + for key, value in defaults.items(): + if key not in obj.__dict__: + obj.__dict__[key] = value + for value in list(obj.__dict__.values()): + fill_defaults(value, seen) + elif isinstance(obj, dict): + for value in obj.values(): + fill_defaults(value, seen) + elif isinstance(obj, (list, tuple, set)): + for value in obj: + fill_defaults(value, seen) + + +def artifact_files(root): + return sorted(p for p in Path(root).rglob('*.txt') if p.name in ARTIFACTS) + + +def field_drift(root): + return field_drift_paths(artifact_files(root)) + + +def field_drift_paths(paths): + """Per class, which stored fields the current model no longer declares. + + Class names matching is not enough: a renamed field would migrate into a + model that has no home for it, and the value would quietly vanish. + """ + drift = {} + for path in paths: + obj = jsonpickle.decode(path.read_text()) + stack, seen = [obj], set() + while stack: + item = stack.pop() + if id(item) in seen: + continue + seen.add(id(item)) + if isinstance(item, BaseModel): + name = type(item).__name__ + stored, declared = set(item.__dict__), set(type(item).model_fields) + counted = drift.setdefault(name, [set(), set(), 0]) + counted[0] |= stored - declared + counted[1] |= declared - stored + counted[2] += 1 + stack.extend(item.__dict__.values()) + elif isinstance(item, dict): + stack.extend(item.values()) + elif isinstance(item, (list, tuple, set)): + stack.extend(item) + return drift + + +def scan(root): + files = artifact_files(root) + print(f"{len(files)} artifact file(s) under {root}") + found = set() + for path in files: + found |= classes_in(path.read_text()) + broken = broken_modules(found) + for cls in sorted(found): + module_path = cls.rpartition('.')[0] + mark = 'STALE' if module_path in broken else 'ok' + print(f" [{mark:5}] {cls}") + for module_path, why in broken.items(): + print(f"\n{module_path} does not import -> {why}") + missing = [c.rpartition('.')[2] for c in found + if c.startswith(module_path + '.') + and c.rpartition('.')[2] not in CURRENT] + if missing: + print(f" !! no current class named: {sorted(set(missing))}") + install_alias(module_path) + + if broken: + print("\nstored fields vs current models:") + drift = field_drift(root) + if not drift: + print(" !! no model objects decoded -- scan proved nothing") + for name, (unknown, absent, count) in sorted(drift.items()): + print(f" {name}: {count} object(s) inspected") + if unknown: + print(f" !! stored but not declared -> " + f"{sorted(unknown)} (renamed? migration drops these)") + if absent: + print(f" new field, will take its default -> " + f"{sorted(absent)}") + return broken + + + +def dead_modules(sample_paths): + """Which stored modules no longer import, from a sample of files. + + Only the module set comes from the sample -- never the class list. A + sample of concepts.txt files sees DugConcept and nothing else, and a + mapping built from that would restamp DugConcept while leaving + DugVariable pointing at the dead module: a half-migrated file that still + decodes to plain dicts. + """ + dead, found = set(), set() + for path in sample_paths: + found |= classes_in(path.read_text()) + for cls in found: + module_path = cls.rpartition('.')[0] + if module_path in dead or module_path in sys.modules: + continue + try: + importlib.import_module(module_path) + except Exception: # noqa: BLE001 - module is gone; restamp it + dead.add(module_path) + return dead + + +def restamp_text(text, dead): + """Repoint every py/object under a dead module at its current home. + + Resolves by class name, so classes that never appeared in the sample are + still rewritten. Returns (new_text, unmapped stored paths). + """ + unmapped = set() + + def repoint(match): + stored = match.group(1) + module_path, _, name = stored.rpartition('.') + if module_path not in dead: + return match.group(0) + current = CURRENT.get(name) + if current is None: + unmapped.add(stored) + return match.group(0) + return match.group(0).replace( + stored, f"{current.__module__}.{current.__name__}") + + return PY_OBJECT.sub(repoint, text), unmapped + + +def restamp(root, sample=20, dry_run=False): + """Rewrite stored class paths as text, without decoding. + + --fix decodes and re-encodes every file, which is right when fields have + to be backfilled but costs seconds per file; at 150k artifacts that is + days. When the stored and current models declare the same fields, the + only thing that needs to change is the dotted path in "py/object", so a + streaming string replace does the whole job at I/O speed. + + Field equivalence is not assumed: a sample is decoded through the shim + and checked for drift first, and any drift aborts the run. + """ + files = artifact_files(root) + print(f"{len(files)} artifact file(s) under {root}") + if not files: + return 0 + + # stratify: concepts.txt holds only concepts and elements.txt only + # elements, so an unstratified slice can miss whole classes + sample_paths = [] + per_kind = max(1, sample // len(ARTIFACTS)) + for kind in ARTIFACTS: + of_kind = [f for f in files + if f.name == kind and f.stat().st_size > 64] + step = max(1, len(of_kind) // per_kind) + sample_paths += of_kind[::step][:per_kind] + if not sample_paths: + raise SystemExit("no non-empty artifact files to sample") + + dead = dead_modules(sample_paths) + if not dead: + print("nothing stale in the sample; already current") + return 0 + print("dead module(s):", ", ".join(sorted(dead))) + for module_path in dead: + install_alias(module_path) + print(f"\nchecking {len(sample_paths)} sampled file(s) for field drift") + drift = field_drift_paths(sample_paths) + blocked = False + for name, (unknown, absent, count) in sorted(drift.items()): + print(f" {name}: {count} object(s) inspected") + if unknown: + blocked = True + print(f" !! stored but not declared -> {sorted(unknown)}") + if absent: + print(f" new field, NOT backfilled by restamp -> " + f"{sorted(absent)}") + if blocked: + raise SystemExit( + "aborting: restamp only rewrites class paths, so renamed or " + "dropped fields would be silently lost. Use --fix for these.") + if not drift: + raise SystemExit("aborting: decoded no model objects from the sample") + + changed = 0 + for i, path in enumerate(files, 1): + text = path.read_text() + new_text, unmapped = restamp_text(text, dead) + if unmapped: + raise SystemExit( + f"!! no current class for {sorted(unmapped)} in {path}; " + f"add an explicit mapping before restamping") + if new_text == text: + continue + changed += 1 + if not dry_run: + # write-then-rename: a pod killed mid-write must not leave a + # truncated artifact behind, and there are 150k of them + tmp = path.with_name(path.name + '.restamp-tmp') + tmp.write_text(new_text) + os.replace(tmp, path) + if changed % 5000 == 0: + print(f" {changed} rewritten ({i}/{len(files)} scanned)") + print(f"{'would rewrite' if dry_run else 'rewrote'} {changed} " + f"of {len(files)} file(s)") + return changed + + +def fix(root, dry_run=False): + files = artifact_files(root) + for module_path in broken_modules( + {c for p in files for c in classes_in(p.read_text())}): + print(f"aliasing legacy module {module_path}") + install_alias(module_path) + + changed = 0 + for path in files: + text = path.read_text() + obj = jsonpickle.decode(text) + fill_defaults(obj) + rewritten = jsonpickle.encode(obj, indent=2) + if classes_in(rewritten) == classes_in(text): + continue + changed += 1 + print(f"{'would rewrite' if dry_run else 'rewrote'} {path}") + if not dry_run: + path.write_text(rewritten) + print(f"{changed} of {len(files)} file(s) needed migration") + return changed + + +def self_check(): + "Round-trip a synthetic legacy payload; fails loudly if the shim regresses." + legacy = json.dumps({"UMLS:C1": { + "py/object": "dug.core.parsers._base.DugConcept", + "py/state": { + "__dict__": {"id": "UMLS:C1", "name": "n", "description": "", + "type": "concept", "search_terms": ["s"], + "identifiers": {}, "kg_answers": {}}, + "__pydantic_extra__": None, + "__pydantic_fields_set__": {"py/set": ["id", "name"]}, + "__pydantic_private__": None}}}) + + assert isinstance(jsonpickle.decode(legacy)["UMLS:C1"], dict), \ + "legacy payload decoded without the shim -- test no longer meaningful" + + install_alias("dug.core.parsers._base") + concept = jsonpickle.decode(legacy)["UMLS:C1"] + assert isinstance(concept, DugConcept), type(concept) + assert concept.id == "UMLS:C1" + + fill_defaults(concept) + reencoded = jsonpickle.encode(concept) + assert "dug.core.parsers._base" not in reencoded, reencoded[:200] + assert isinstance(jsonpickle.decode(reencoded), DugConcept) + print("self-check ok") + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument('--scan', metavar='DIR') + parser.add_argument('--fix', metavar='DIR') + parser.add_argument('--restamp', metavar='DIR', + help='rewrite class paths as text (fast path for\n' + 'very large trees); verifies a sample first') + parser.add_argument('--sample', type=int, default=20, + help='files to decode for the drift check') + parser.add_argument('--dry-run', action='store_true') + parser.add_argument('--self-check', action='store_true') + args = parser.parse_args() + + if args.self_check: + self_check() + elif args.scan: + scan(args.scan) + elif args.restamp: + restamp(args.restamp, sample=args.sample, dry_run=args.dry_run) + elif args.fix: + fix(args.fix, dry_run=args.dry_run) + else: + parser.error("one of --scan, --restamp, --fix, --self-check is required") + + +if __name__ == '__main__': + main() diff --git a/scripts/seed_incremental_state.py b/scripts/seed_incremental_state.py new file mode 100644 index 00000000..726f5d76 --- /dev/null +++ b/scripts/seed_incremental_state.py @@ -0,0 +1,126 @@ +#!/usr/bin/env python +"""Seed the incremental-ingest Airflow Variables for datasets already ingested. + +A fresh instance has no state, so the first run re-annotates everything even +though the data is already in lakefs. This writes the "last consumed commit" +Variables so the first run diffs instead. + + # print the commands, change nothing + python scripts/seed_incremental_state.py \ + --dataset heal-mds-studies:main=a352eccb... \ + --dataset heal-cdes:main=f40bd330... + + # ...and the downstream tasks, so crawl/make_kgx don't redo everything + python scripts/seed_incremental_state.py \ + --dataset heal-mds-studies:main=a352eccb... \ + --downstream := \ + --apply + +Only pass --downstream when the runtime repo ALREADY holds crawl and KGX +outputs for these datasets; it claims everything up to that commit is +consumed, and work that never ran would be skipped permanently. + +Run inside the roger image / an airflow pod. --apply shells out to the airflow +CLI, so it needs a pod with metadata-DB access (scheduler, not a worker). +""" + +import argparse +import sys + +from roger.tasks import ANNOTATE_DAG_ID, file_task_group_id, \ + incremental_state_key + + +def parse_spec(spec): + "name:branch=commit -> (name, branch, commit)" + try: + target, commit = spec.split('=', 1) + name, branch = target.rsplit(':', 1) + except ValueError: + raise argparse.ArgumentTypeError( + f"expected name:branch=commit, got {spec!r}") + if not (name and branch and commit): + raise argparse.ArgumentTypeError( + f"expected name:branch=commit, got {spec!r}") + return name, branch, commit + + +def build(datasets, downstream=None): + """[(key, commit)] for the annotate tasks, plus crawl/make_kgx when a + downstream (runtime repo) target is given.""" + pairs = [] + for name, branch, commit in datasets: + group = file_task_group_id(name) + pairs.append((incremental_state_key( + ANNOTATE_DAG_ID, f"{group}.annotate_{name}_files", + name, branch), commit)) + if downstream: + repo, ds_branch, ds_commit = downstream + for task in (f"crawl_{name}", f"make_kgx_{name}"): + pairs.append((incremental_state_key( + ANNOTATE_DAG_ID, f"{group}.{task}", + repo, ds_branch), ds_commit)) + return pairs + + +def self_check(): + pairs = build([parse_spec("heal-cdes:main=abc123")]) + assert pairs == [( + "roger_incr::annotate_and_index::" + "heal-cdes_dataset_pipeline_task_group.annotate_heal-cdes_files::" + "heal-cdes@main", "abc123")], pairs + + pairs = build([parse_spec("d:main=c1")], parse_spec("out:prod=c2")) + assert [k for k, _ in pairs] == [ + "roger_incr::annotate_and_index::" + "d_dataset_pipeline_task_group.annotate_d_files::d@main", + "roger_incr::annotate_and_index::" + "d_dataset_pipeline_task_group.crawl_d::out@prod", + "roger_incr::annotate_and_index::" + "d_dataset_pipeline_task_group.make_kgx_d::out@prod"], pairs + assert [v for _, v in pairs] == ["c1", "c2", "c2"], pairs + print("self-check ok") + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument('--dataset', action='append', type=parse_spec, + metavar='NAME:BRANCH=COMMIT', default=[], + help='pipeline name, its input_version ref, and the ' + 'source commit already ingested') + parser.add_argument('--downstream', type=parse_spec, default=None, + metavar='REPO:BRANCH=COMMIT', + help='runtime repo/branch and its current tip, to ' + 'also seed crawl_* and make_kgx_*') + parser.add_argument('--apply', action='store_true', + help='set the Variables instead of printing commands') + parser.add_argument('--self-check', action='store_true') + args = parser.parse_args() + + if args.self_check: + self_check() + return + if not args.dataset: + parser.error("at least one --dataset is required") + + pairs = build(args.dataset, args.downstream) + + if not args.apply: + for key, commit in pairs: + print(f"airflow variables set '{key}' {commit}") + print(f"\n# {len(pairs)} variable(s); re-run with --apply to set them", + file=sys.stderr) + return + + # airflow.sdk.Variable only works inside a running task (it needs + # SUPERVISOR_COMMS), so shell out to the CLI, which is what the printed + # commands do anyway. Needs a pod with metadata-DB access. + import subprocess + for key, commit in pairs: + subprocess.run(['airflow', 'variables', 'set', key, commit], + check=True) + print(f"set {key} = {commit}") + + +if __name__ == '__main__': + main() diff --git a/setup.cfg b/setup.cfg new file mode 100644 index 00000000..fa232c8f --- /dev/null +++ b/setup.cfg @@ -0,0 +1,39 @@ +[metadata] +name = roger +version = 0.10.4.2 +author = Renaissance Computing Institute + RTI +description = Data pipeline automation for dug +long_description = file: README.md +long_description_content_type = text/markdown +url = https://github.com/helxplatform/roger +project_urls = + Bug Tracker = https://github.com/helxplatform/roger/issues +classifiers = + Programming Language :: Python :: 3 + License :: OSI Approved :: MIT License + Operating System :: OS Independent + +[options] +package_dir = + = src +packages = find: +python_requires = >=3.10 +include_package_data = true +install_requires = + orjson + requests + requests_cache + redis + +[options.entry_points] +console_scripts = + dug = dug.cli:main + roger = roger.cli:main + +[options.extras_require] +rest = + jsonschema + apache-airflow + +[options.packages.find] +where = src diff --git a/setup.py b/setup.py new file mode 100644 index 00000000..45f160da --- /dev/null +++ b/setup.py @@ -0,0 +1,6 @@ +#!/usr/bin/env python + +import setuptools + +if __name__ == "__main__": + setuptools.setup() \ No newline at end of file diff --git a/src/roger/__init__.py b/src/roger/__init__.py new file mode 100644 index 00000000..e950b109 --- /dev/null +++ b/src/roger/__init__.py @@ -0,0 +1,20 @@ +"Roger: an automated graph data curation pipeline." + +from roger.core.base import ( + Roger, + roger_cli, + get_kgx, + create_schema, + create_edges_schema, + create_nodes_schema, + merge_nodes, + create_bulk_load, + create_bulk_nodes, + create_bulk_edges, + bulk_load, + validate, + check_tranql, +) + +if __name__ == "__main__": + roger_cli() diff --git a/src/roger/_version.py b/src/roger/_version.py new file mode 100644 index 00000000..adcf54c7 --- /dev/null +++ b/src/roger/_version.py @@ -0,0 +1,2 @@ +version = "0.10.4" + diff --git a/src/roger/cli.py b/src/roger/cli.py new file mode 100644 index 00000000..4c0b3f02 --- /dev/null +++ b/src/roger/cli.py @@ -0,0 +1,155 @@ +"""CLI interface for roger +""" + +import sys +import argparse +import os +import time +import pathlib + +import roger +from roger.config import config +from roger.logger import get_logger +from roger.pipelines import get_pipeline_classes + +log = get_logger() + +def get_arguments(): + "Parse argv" + + parser = argparse.ArgumentParser(description='Roger common cli tool.') + """ Common CLI. """ + parser.add_argument('-d', '--data-root', default=None, + help="Root of data hierarchy") + parser.add_argument('-ip', '--input-path', default=None, type=pathlib.Path, + help="Input files path") + parser.add_argument('-op', '--output-path', default=None, type=pathlib.Path, + help="Output files path") + + """ Roger CLI. """ + parser.add_argument('-v', '--dataset-version', help="Dataset version.", + default="v1.0") + parser.add_argument('-g', '--get-kgx', help="Get KGX objects", + action='store_true') + parser.add_argument('-s', '--create-schema', help="Infer schema", + action='store_true') + parser.add_argument('-m', '--merge-kgx', help="Merge KGX nodes", + action='store_true') + parser.add_argument('-b', '--create-bulk', help="Create bulk load", + action='store_true') + parser.add_argument('-i', '--insert', help="Do the bulk insert", + action='store_true') + parser.add_argument('-a', '--validate', help="Validate the insert", + action='store_true') + + dataset_envspec = os.getenv("ROGER_DUG__INPUTS_DATA__SETS", + "topmed:v2.0,dbGaP:v1.0,anvil:v1.0") + data_sets = dataset_envspec.split(",") + parser.add_argument('-D', '--datasets', action="append", + default=None, + help="Dataset pipelines name:vers to run. " + f"(default: f{str(data_sets)}") + + """ Dug Annotation CLI. """ + parser.add_argument('-gd', '--get_dug_input_files', action="store_true", + help="Gets input files for annotation") + parser.add_argument('-l', '--load-and-annotate', action="store_true", + help="Annotates and normalizes datasets of varaibles.") + parser.add_argument('-t', '--make-tagged-kg', action="store_true", + help="Creates KGX files from annotated variable " + "datasets.") + + """ Dug indexing CLI . """ + parser.add_argument('-iv', '--index-variables', action="store_true", + help="Index annotated variables to elastic search.") + parser.add_argument('-C', '--crawl-concepts', action="store_true", + help="Crawl tranql and index concepts") + parser.add_argument('-ic', '--index-concepts', action="store_true", + help="Index expanded concepts to elastic search.") + + parser.add_argument('-vc', '--validate-concepts', action="store_true", + help="Validates indexing of concepts") + + parser.add_argument('-vv', '--validate-variables', action="store_true", + help="Validates indexing of variables") + + args = parser.parse_args () + + if args.data_root is not None: + data_root = args.data_root + config.data_root = data_root + log.info (f"data root:{data_root}") + + if not args.datasets: + args.datasets = data_sets + + return args + +def main(): + start = time.time() + log.info(f"Start TIME:{start}") + + args = get_arguments() + # When all lights are on... + + # Instantiate the pipeline classes + pipeline_names = {x.split(':')[0]: x.split(':')[1] for x in args.datasets} + log.info("Working on dataset list %s", str(pipeline_names)) + pipeline_classes = get_pipeline_classes(pipeline_names) + pipelines = [pipeclass(config) for pipeclass in pipeline_classes] + + for pipe in pipelines: + # Do all actions for one pipeline first, then move on to the next: + log.info("Running pipeline %s", pipe.pipeline_name) + + # Annotation comes first + if args.get_dug_input_files: + pipe.get_versioned_files() + + if args.load_and_annotate: + pipe.clear_annotation_cached(output_data_path=args.output_path) + pipe.annotate(input_data_path=args.input_path, + output_data_path=args.output_path) + + if args.make_tagged_kg: + pipe.make_kg_tagged() + + # Roger things + if args.get_kgx: + roger.get_kgx(config=config) + if args.merge_kgx: + roger.merge_nodes(config=config) + if args.create_schema: + roger.create_schema(config=config) + if args.create_bulk: + roger.create_bulk_load(config=config) + if args.insert: + roger.bulk_load(config=config) + if args.validate: + roger.validate(config=config) + roger.check_tranql(config=config) + + # Back to dug indexing + if args.index_variables: + pipe.index_variables() + + if args.validate_variables: + pipe.validate_indexed_variables() + + if args.crawl_concepts: + pipe.crawl_tranql() + + if args.index_concepts: + pipe.index_concepts() + + if args.validate_concepts: + pipe.validate_indexed_concepts() + + end = time.time() + time_elapsed = end - start + log.info(f"Completion TIME:{time_elapsed}") + + sys.exit (0) + +if __name__ == "__main__": + main() diff --git a/src/roger/components/__init__.py b/src/roger/components/__init__.py new file mode 100644 index 00000000..49314f47 --- /dev/null +++ b/src/roger/components/__init__.py @@ -0,0 +1 @@ +"Data conversion utilities" diff --git a/src/roger/components/data_conversion.py b/src/roger/components/data_conversion.py new file mode 100644 index 00000000..46dd61fe --- /dev/null +++ b/src/roger/components/data_conversion.py @@ -0,0 +1,71 @@ +"Data conversion utility methods" + +from typing import Any + + +_type_map = { + list.__name__: { + 'priority': 0, + 'constructor': lambda x: list([x]) + }, + str.__name__: { + 'priority': 1, + 'constructor': lambda x: str(x) + }, + bool.__name__: { + 'priority': 2, + 'constructor': lambda x: True if x else False + }, + float.__name__: { + 'priority': 2, + 'constructor': lambda x: float(x), + }, + int.__name__: { + 'priority': 2, + 'constructor': lambda x: int(x) + }, + type(None).__name__: { + 'priority': 3, + 'constructor': lambda x: '', + } +} + +def cast(value: Any, to_type: str): + """ + Parses a value to dest type. + :param value: value to parse + :param to_type: destination type + :return: parsed value + """ + if to_type not in _type_map: + raise TypeError( + f'Type {to_type} not found in conversion map. ' + f'Available types are {_type_map.keys()}') + dest_type_constructor = _type_map[to_type]['constructor'] + return dest_type_constructor(value) + +def compare_types(data_type: str, data_type_2: str): + """ + Of two python types selects the one we would like to upcast to. + :param data_type: + :param data_type_2: + :return: + """ + assert data_type in _type_map, ( + f"Unrecognised type {data_type} From types:" + f"{list(_type_map.keys())}") + + assert data_type_2 in _type_map, ( + f"Unrecognised type {data_type} From types: " + f"{list(_type_map.keys())}") + + d1_val = _type_map[data_type]['priority'] + d2_val = _type_map[data_type_2]['priority'] + + if data_type != data_type_2 and d1_val == d2_val: + # For float int and bool have same priority + # treat them as strings. + d1_val = (d1_val - 1) + data_type = str.__name__ + + return data_type if d1_val < d2_val else data_type_2 diff --git a/src/roger/components/data_conversion_utils.py b/src/roger/components/data_conversion_utils.py new file mode 100644 index 00000000..f6f60eb0 --- /dev/null +++ b/src/roger/components/data_conversion_utils.py @@ -0,0 +1,69 @@ +from typing import Any + + +class TypeConversionUtil: + + type_map = { + list.__name__: { + 'priority': 0, + 'constructor': lambda x: list([x]) + }, + str.__name__: { + 'priority': 1, + 'constructor': lambda x: str(x) + }, + bool.__name__: { + 'priority': 2, + 'constructor': lambda x: True if x else False + }, + float.__name__: { + 'priority': 2, + 'constructor': lambda x: float(x), + }, + int.__name__: { + 'priority': 2, + 'constructor': lambda x: int(x) + }, + type(None).__name__: { + 'priority': 3, + 'constructor': lambda x: '', + } + } + + @staticmethod + def cast(value: Any, to_type: str): + """ + Parses a value to dest type. + :param value: value to parse + :param to_type: destination type + :return: parsed value + """ + if to_type not in TypeConversionUtil.type_map: + raise TypeError(f'Type {to_type} not found in conversion map. Available types are {TypeConversionUtil.type_map.keys()}') + dest_type_constructor = TypeConversionUtil.type_map[to_type]['constructor'] + return dest_type_constructor(value) + + @staticmethod + def compare_types(data_type: str, data_type_2: str): + """ + Of two python types selects the one we would like to upcast to. + :param data_type: + :param data_type_2: + :return: + """ + assert data_type in TypeConversionUtil.type_map, f"Unrecognised type {data_type} From types:" \ + f"{list(TypeConversionUtil.type_map.keys())}" + + assert data_type_2 in TypeConversionUtil.type_map, f"Unrecognised type {data_type} From types: " \ + f"{list(TypeConversionUtil.type_map.keys())}" + + d1_val = TypeConversionUtil.type_map[data_type]['priority'] + d2_val = TypeConversionUtil.type_map[data_type_2]['priority'] + + if data_type != data_type_2 and d1_val == d2_val: + # For float int and bool have same priority + # treat them as strings. + d1_val = (d1_val - 1) + data_type = str.__name__ + + return data_type if d1_val < d2_val else data_type_2 diff --git a/src/roger/config/__init__.py b/src/roger/config/__init__.py new file mode 100644 index 00000000..a7732b0a --- /dev/null +++ b/src/roger/config/__init__.py @@ -0,0 +1,478 @@ +import json +import os +import warnings +from dataclasses import dataclass, field +from pathlib import Path +from typing import Dict, Optional, List, Union + +import yaml +from dug.config import Config as DugConfig +from flatten_dict import flatten, unflatten + +from ._base import DictLike +from .s3_config import S3Config + +if os.environ.get('ROGER_CONFIG_FILE', None): + CONFIG_FILENAME = Path(os.environ.get('ROGER_CONFIG_FILE')) +else: + CONFIG_FILENAME = Path(__file__).parent.resolve() / "config.yaml" + +@dataclass +class RedisConfig(DictLike): + username: str = "" + password: str = "" + host: str = "redis" + graph: str = "test" + port: int = 6379 + use_redis_cache: bool = True + + def __post_init__(self): + self.port = int(self.port) + + +@dataclass +class LakefsConfig(DictLike): + host: str + access_key_id: str + secret_access_key: str + branch: str + repo: str + enabled: Union[bool, str] = False + + def __post_init__(self): + if isinstance(self.enabled, str): + self.enabled = self.enabled.lower() == "true" + + + +@dataclass +class LoggingConfig(DictLike): + level: str = "DEBUG" + format: str = '[%(name)s][%(filename)s][%(lineno)d][%(funcName)20s] %(levelname)s: %(message)s' + + +@dataclass +class KgxConfig(DictLike): + biolink_model_version: str = "1.5.0" + merge_db_temp_dir: str = "workspace" + data_sets: List = field(default_factory=lambda: ['baseline-graph:v5.0']) + + def __post_init__(self): + # Convert strings to list. In cases where this is passed as env variable with a single value + # cast it to a list. eg ROGER_KGX_DATA__SET="spark,baseline-data" could be converted to + # config.kgx.data_set = ["spark", "baseline-data"] + self.data_sets = [data_set.strip(" ") for data_set in self.data_sets.split(",")] \ + if isinstance(self.data_sets, str) else self.data_sets + + +@dataclass +class DugInputsConfig(DictLike): + data_source: str = 'stars' + data_sets: List = field(default_factory=lambda: ['topmed', 'bdc']) + + def __post_init__(self): + # Convert strings to list. In cases where this is passed as env variable with a single value + # cast it to a list. eg ROGER_KGX_DATA__SET="spark,baseline-data" could be converted to + # config.kgx.data_set = ["spark", "baseline-data"] + self.data_sets = [data_set.strip(" ") for data_set in self.data_sets.split(",")] \ + if isinstance(self.data_sets, str) else self.data_sets + + +@dataclass +class BulkLoaderConfig(DictLike): + separator: str = "0x1E" + enforce_schema: bool = False + skip_invalid_nodes: bool = False + skip_invalid_edges: bool = False + quote: int = 0 + max_token_count: int = 1024 + max_buffer_size: int = 2048 + max_token_size: int = 500 + index: list = field(default_factory=list) + full_text_index: list = field(default_factory=list) + + +@dataclass +class AnnotationConfig(DictLike): + annotator_type: str = "monarch" + annotator_args: dict = field( + default_factory=lambda: { + "monarch": { + "url": "https://api.monarchinitiative.org/api/nlp/annotate/entities?min_length=4&longest_only=false&include_abbreviation=false&include_acronym=false&include_numbers=false&content=" + }, + "sapbert": { + "classification_url": "https://med-nemo.apps.renci.org/annotate/", + "annotator_url": "https://babel-sapbert.apps.renci.org/annotate/", + "score_threshold": 0.8, + "bagel": { + "enabled": False, + "url": "https://bagel.apps.renci.org/group_synonyms_openai", + "prompt": "bagel/ask_classes", + "llm_args": { + "llm_model_name": "gpt-4o-2024-05-13", + "organization": "", + "access_key": "", + "llm_model_args": { + "top_p": 0, + "temperature": 0.1 + } + } + } + }, + } + ) + normalizer: str = "https://nodenormalization-sri.renci.org/get_normalized_nodes?curie=" + synonym_service: str = "https://onto.renci.org/synonyms/" + ontology_metadata: str = "https://api.monarchinitiative.org/api/bioentity/" + clear_http_cache: bool = False + # Bounds on the annotation service calls. Without a read timeout a broken + # connection stalls the pipeline indefinitely rather than failing. + http_connect_timeout: float = 10.0 + http_read_timeout: float = 120.0 + http_retries: int = 3 + http_retry_backoff: float = 1.0 + # dug hands the annotator a requests_cache CachedSession, but + # requests_cache only caches GET/HEAD by default. Of the four annotation + # calls only node normalization is a GET; nemo token classification, + # sapbert, and name-resolution synonyms are POSTs and so were never + # cached. That matters enormously for dbGaP-shaped data: each data-dict + # file re-annotates its parent study element, so bdc-parent annotated the + # same 24 study descriptions 61,597 times, at ~53s each against ~1.4s for + # the variable that file actually contributes. All these services are + # read-only lookups keyed entirely by the request body, so caching POSTs + # is safe. + cache_post_requests: bool = True + # Nonzero for two reasons. The normalizer and synonym services have + # stable urls but drifting content, so entries should not live forever. + # (dug set no expiry at all, so the normalizer GETs it did cache were + # permanent; setting this bounds those too.) + # More importantly, requests_cache's redis backend writes entries with + # SETEX when an expiry is set, which makes cache keys volatile while the + # graph keys sharing that redis stay permanent -- so redis can be given + # a maxmemory with volatile-lru and will evict annotation cache before it + # ever touches the loaded graph. With no expiry the cache is permanent, + # unbounded, and under noeviction grows until the pod is OOMKilled, + # taking the graph with it. 0 disables expiry. + http_cache_expire_seconds: int = 30 * 24 * 3600 + # Files annotate independently, and the work is almost entirely waiting + # on http, so threads help even under the GIL. Each worker gets its own + # session and annotator. + annotate_workers: int = 4 + # dug resolves identifiers one at a time; both the normalizer and the + # name resolution service take a list and answer it in about the time + # they take to answer one (0.13ms/curie at n=200 vs 10-26ms at n=1). + # See roger.utils.batched_annotator. Off puts the serial path back. + batch_identifier_lookups: bool = True + # Memory limit for the annotate task pods. The chart default (2Gi) was + # sized for the serial annotator; annotate_workers threads hold one + # parsed file each, and some files are large -- bdc-biolincc has files + # with 13,299 elements, and jsonpickle.encode builds the whole output + # string in memory before it is written. Four of those at once OOMKilled + # the task at 2Gi. + annotate_memory: str = "6Gi" + preprocessor: dict = field(default_factory=lambda: + { + "debreviator": { + "BMI": "body mass index" + }, + "stopwords": "the", + } + ) + + ontology_greenlist: List[str] = field(default_factory=lambda: [ + "PATO", "CHEBI", "MONDO", "UBERON", "HP", "MESH", "UMLS" + ]) + + def __post_init__(self): + self.annotator_args["sapbert"]["bagel"]["enabled"] = str(self.annotator_args["sapbert"]["bagel"][ + "enabled"]).lower() == "true" + # These can arrive from environment variables, where every value is a + # string. urllib3 rejects a string timeout outright, and a string retry + # count fails on the first retry, so coerce them here. + self.http_connect_timeout = float(self.http_connect_timeout) + self.http_read_timeout = float(self.http_read_timeout) + self.http_retries = int(self.http_retries) + self.http_retry_backoff = float(self.http_retry_backoff) + # ROGER_ANNOTATION_CACHE__POST__REQUESTS=false would otherwise be a + # truthy string, silently leaving the cache on + if isinstance(self.cache_post_requests, str): + self.cache_post_requests = ( + self.cache_post_requests.strip().lower() == "true") + else: + self.cache_post_requests = bool(self.cache_post_requests) + self.http_cache_expire_seconds = int(self.http_cache_expire_seconds) + self.annotate_workers = max(1, int(self.annotate_workers)) + if isinstance(self.batch_identifier_lookups, str): + self.batch_identifier_lookups = ( + self.batch_identifier_lookups.strip().lower() == "true") + else: + self.batch_identifier_lookups = bool( + self.batch_identifier_lookups) + + +@dataclass +class IndexingConfig(DictLike): + variables_index: str = "variables_index" + concepts_index: str = "concepts_index" + kg_index: str = "kg_index" + studies_index: str = "studies_index" + sections_index: str = "sections_index" + tranql_min_score: float = 0.2 + excluded_identifiers: List[str] = field(default_factory=lambda: [ + "CHEBI:17336" + ]) + + queries: dict = field(default_factory=lambda: { + "disease": ["disease", "phenotypic_feature"], + "pheno": ["phenotypic_feature", "disease"], + "anat": ["disease", "anatomical_entity"], + "chem_to_disease": ["chemical_substance", "disease"], + "phen_to_anat": ["phenotypic_feature", "anatomical_entity"], + "anat_to_disease": ["anatomical_entity", "disease"], + "anat_to_pheno": ["anatomical_entity", "phenotypic_feature"], + }) + tranql_endpoint: str = "http://tranql-service/tranql/query?dynamic_id_resolution=true&asynchronous=false" + # by default skips node to element queries + node_to_element_queries: dict = field(default_factory=lambda: {}) + element_mapping: str = "" + def __post_init__(self): + # convert element mapping to dict + if self.element_mapping and len(self.element_mapping.split(',')): + final_element_mapping = {} + for mapping in self.element_mapping.split(','): + if not mapping: + continue + original_name = mapping.split(':')[0].lower().strip() + final_name = mapping.split(':')[1].strip() + final_element_mapping[original_name] = final_name + self.element_mapping = final_element_mapping + node_to_el_enabled = True if str(self.node_to_element_queries.get("enabled")).lower() == "true" else False + final_node_to_element_queries = {} + if node_to_el_enabled: + for key in filter(lambda k: k != "enabled", self.node_to_element_queries.keys()): + final_node_to_element_queries[key] = self.node_to_element_queries[key] + self.node_to_element_queries = final_node_to_element_queries + +@dataclass +class ElasticsearchConfig(DictLike): + host: str = "elasticsearch" + username: str = "elastic" + password: str = "" + nboost_host: str = "" + scheme: str = "http" + ca_path: str = "" + + + +class RogerConfig(DictLike): + + OS_VAR_PREFIX = "ROGER_" + + def __init__(self, **kwargs): + self.redisgraph = RedisConfig(**kwargs.pop('redisgraph', {})) + self.logging = LoggingConfig(**kwargs.pop('logging', {})) + self.kgx = KgxConfig(**kwargs.pop('kgx', {})) + self.dug_inputs = DugInputsConfig(**kwargs.pop('dug_inputs', {})) + self.bulk_loader = BulkLoaderConfig(**kwargs.pop('bulk_loader', {})) + self.annotation = AnnotationConfig(**kwargs.pop('annotation', {})) + self.indexing = IndexingConfig(**kwargs.pop('indexing', {})) + self.elasticsearch = ElasticsearchConfig(**kwargs.pop('elasticsearch')) + self.s3_config = S3Config(**kwargs.pop('s3', {})) + + self.data_root: str = kwargs.pop("data_root", "") + self.dug_data_root: str = kwargs.pop("dug_data_root", "") + self.kgx_base_data_uri: str = kwargs.pop("kgx_base_data_uri", "") + self.annotation_base_data_uri: str = kwargs.pop("annotation_base_data_uri", "") + self.validation = kwargs.pop("validation") + self.dag_run = kwargs.pop('dag_run', None) + self.lakefs_config = LakefsConfig(**kwargs.pop("lakefs_config")) + + def to_dug_conf(self) -> DugConfig: + return DugConfig( + elastic_host=self.elasticsearch.host, + elastic_password=self.elasticsearch.password, + elastic_username=self.elasticsearch.username, + elastic_scheme=self.elasticsearch.scheme, + elastic_ca_path=self.elasticsearch.ca_path, + redis_host=self.redisgraph.host, + redis_password=self.redisgraph.password, + redis_port=self.redisgraph.port, + use_redis_cache=self.redisgraph.use_redis_cache, + nboost_host=self.elasticsearch.nboost_host, + preprocessor=self.annotation.preprocessor, + annotator_type=self.annotation.annotator_type, + annotator_args=self.annotation.annotator_args, + concepts_index_name=self.indexing.get('concepts_index'), + variables_index_name=self.indexing.get('variables_index'), + studies_index_name=self.indexing.get('studies_index'), + sections_index_name=self.indexing.get('sections_index'), + kg_index_name=self.indexing.get('kg_index'), + normalizer={ + 'url': self.annotation.normalizer, + }, + synonym_service={ + 'url': self.annotation.synonym_service, + }, + ontology_helper={ + 'url': self.annotation.ontology_metadata, + }, + tranql_exclude_identifiers=self.indexing.excluded_identifiers, + tranql_queries=self.indexing.queries, + concept_expander={ + 'url': self.indexing.tranql_endpoint, + 'min_tranql_score': self.indexing.tranql_min_score, + }, + ontology_greenlist=self.annotation.ontology_greenlist, + node_to_element_queries=self.indexing.node_to_element_queries, + ) + + @property + def dict(self): + output = {} + for key, value in self.__dict__.items(): + if hasattr(value, '__dict__'): + output[key] = value.__dict__ + else: + output[key] = value + return output + + @classmethod + def factory(cls, file_path: str): + file_path = Path(file_path).resolve() + with file_path.open() as config_file: + file_data = yaml.load(config_file, Loader=yaml.FullLoader) + + override_data = cls.get_override_data(cls.OS_VAR_PREFIX) + + combined_data = cls.merge_dicts(file_data, override_data) + + return RogerConfig(**combined_data) + + @staticmethod + def merge_dicts(dict_a, dict_b): + flat_a = flatten(dict_a, reducer='dot') + flat_b = flatten(dict_b, reducer='dot') + flat_a.update(flat_b) + return unflatten(flat_a, 'dot') + + @staticmethod + def get_override_data(prefix): + override_data = {} + os_var_keys = os.environ.keys() + keys_of_interest = filter(lambda x: x.startswith(prefix), os_var_keys) + for key in keys_of_interest: + value = os.environ.get(key) + var_name = key.replace(prefix, "", 1) + var_name = var_name.lstrip("_") + var_name = var_name.replace("__", "~") + var_name = var_name.replace("_", ".") + var_name = var_name.replace("~", "_") + var_name = var_name.lower() + override_data[var_name] = value + return unflatten(override_data, 'dot') + + +class Config: + """ + Singleton config wrapper + """ + __instance__: Optional[Dict] = None + os_var_prefix = "ROGERENV_" + + def __init__(self, file_name: str): + if not Config.__instance__: + Config.__instance__ = Config.read_config_file(file_name=file_name) + os_var_keys = os.environ.keys() + keys_of_interest = [x for x in os_var_keys if x.startswith(Config.os_var_prefix)] + for key in keys_of_interest: + new_key = key.replace(Config.os_var_prefix, "") + value = os.environ[key] + new_dict = Config.os_var_to_dict(new_key, value) + try: + Config.update(new_dict) + except ValueError as e: + warnings.warn(f"{e} encountered trying to assign string from " + f"OS variable `{key}` to a dictionary object." + f"Please specify inner keys.") + + @staticmethod + def os_var_to_dict(var_name, value): + var_name = var_name.replace("__", "~") + var_name = var_name.replace("_", ".") + var_name = var_name.replace("~", "_") + var_name = var_name.lower() + m = {var_name: value} + result = unflatten(m, "dot") + return result + + @staticmethod + def read_config_file(file_name: str): + return yaml.load(open(file_name), Loader=yaml.FullLoader) + + def __getattr__(self, item): + """ + Proxies calls to instance dict. + Note: dict.update is overridden to do partial updates. + Refer to Config.update method. + :param item: method called + :return: proxied method + """ + if item == 'update': + # overrides default dict update method + return self.update + return getattr(Config.__instance__, item) + + def __getitem__(self, item): + """ + Makes config object subscriptable + :param item: key to lookup + :return: value stored in key + """ + return self.__instance__.get(item) + + @staticmethod + def update(new_value: Dict): + """ + Updates dictionary partially. + Given a config {'name': {'first': 'name', 'last': 'name'}} + and a partial update {'name': {'first': 'new name'} } + result would be {'name': {'first': 'new name', 'last': 'name'}} + :param new_value: parts to update + :return: updated dict + """ + config_flat = flatten(Config.__instance__) + new_value_flat = flatten(new_value) + config_flat.update(new_value_flat) + Config.__instance__ = unflatten(config_flat) + return Config.__instance__ + + def __str__(self): + flat = flatten(Config.__instance__) + for k in flat: + if 'PASSWORD' in k or 'password' in k or 'key' in k.lower(): + flat[k] = '******' + flat = unflatten(flat) + result = json.dumps(flat) + return f"""{result}""" + + +def get_default_config(file_name: str = CONFIG_FILENAME) -> RogerConfig: + """ + Get config as a dictionary + + Parameters + ---------- + file_name: str + The filename with all the configuration + + Returns + ------- + dict + A dictionary containing all the entries from the config YAML + + """ + config_instance = RogerConfig.factory(file_name) + return config_instance + + +config: RogerConfig = get_default_config() diff --git a/src/roger/config/_base.py b/src/roger/config/_base.py new file mode 100644 index 00000000..77309666 --- /dev/null +++ b/src/roger/config/_base.py @@ -0,0 +1,11 @@ +class DictLike: + def __getitem__(self, item): + if not hasattr(self, item): + raise KeyError(item) + return getattr(self, item) + + def __setitem__(self, key, value): + setattr(self, key, value) + + def get(self, key, default=None): + return getattr(self, key, default) \ No newline at end of file diff --git a/src/roger/config/config.yaml b/src/roger/config/config.yaml new file mode 100644 index 00000000..3d883d13 --- /dev/null +++ b/src/roger/config/config.yaml @@ -0,0 +1,208 @@ +redisgraph: + username: "" + password: "weak" + host: localhost + graph: test + port: 6379 + +logging: + # DEBUG here is roger's own logger; third-party HTTP chatter is capped + # separately in roger.logger.quiet_noisy_loggers. + level: INFO + format: '[%(name)s][%(filename)s][%(lineno)d][%(funcName)20s] %(levelname)s: %(message)s' + +data_root: roger/data + +kgx_base_data_uri: https://stars.renci.org/var/kgx_data/ +annotation_base_data_uri: https://stars.renci.org/var/dug/ + +kgx: + biolink_model_version: v3.1.2 + merge_db_temp_dir: workspace + data_sets: + - baseline-graph:v5.0 + +dug_inputs: + data_source: s3 + data_sets: + - topmed:v1.0 + - bdc:v1.0 + - anvil:v1.0 + +#https://github.com/RedisGraph/redisgraph-bulk-loader/blob/master/redisgraph_bulk_loader/bulk_insert.py#L43 +bulk_loader: + separator: 0x1E + enforce_schema: False + skip_invalid_nodes: False + skip_invalid_edges: False + quote: 0 + max_token_count: 1024 + max_buffer_size: 2048 + max_token_size: 500 + index: [] + full_text_index: [] + +annotation: + clear_http_cache: false + # Bounds on the annotation service calls. Without a read timeout a broken + # connection stalls the pipeline indefinitely rather than failing. + http_connect_timeout: 10.0 + http_read_timeout: 120.0 + http_retries: 3 + http_retry_backoff: 1.0 + # requests_cache caches only GET/HEAD by default. Node normalization is + # the one annotation call that is a GET; classification, sapbert and + # synonym lookup are POSTs and were never cached. See AnnotationConfig for + # why this dominates dbGaP-shaped datasets. + cache_post_requests: true + # 30 days. Keep this nonzero: it makes requests_cache write redis entries + # with SETEX, so the annotation cache is evictable (set redis maxmemory + + # volatile-lru) while the graph keys in the same redis are not. dug set no + # expiry, so the normalizer GETs it cached were permanent. 0 disables. + http_cache_expire_seconds: 2592000 + # threads over input files inside one annotate task + annotate_workers: 4 + # Batch the normalize/synonym lookups dug makes per identifier into one + # request each per element. Both services are flat in batch size, so this + # turns ~2 requests per identifier into 2 per element. + batch_identifier_lookups: true + # annotate task pod memory. Threads hold a parsed file each, so this scales + # with annotate_workers; 2Gi (the chart default) OOMKilled bdc-biolincc, + # whose largest files carry 13,299 elements. + annotate_memory: 6Gi + annotator_type: sapbert + annotator_args: + monarch: + url: "https://api.monarchinitiative.org/api/nlp/annotate/entities?min_length=4&longest_only=false&include_abbreviation=false&include_acronym=false&include_numbers=false&content=" + sapbert: + classification_url: "https://med-nemo.apps.renci.org/annotate/" + annotator_url: "https://sap-qdrant.apps.renci.org/annotate/" + score_threshold: 0.8 + bagel: + enabled: false + url: "http://localhost:9099/group_synonyms_openai" + prompt: "bagel/ask_classes" + llm_args: + llm_model_name: "gpt-4o-2024-05-13" + organization: + access_key: + llm_model_args: + top_p: 0 + temperature: 0.1 + normalizer: "https://nodenormalization-dev.apps.renci.org/get_normalized_nodes?conflate=false&description=true&curie=" + synonym_service: "https://name-resolution-sri.renci.org/reverse_lookup" + ontology_metadata: "https://api.monarchinitiative.org/api/bioentity/" + + preprocessor: + debreviator: + BMI: "body mass index" + stopwords: "the" + ontology_greenlist: ["PATO", "CHEBI", "MONDO", "UBERON", "HP", "MESH", "UMLS"] + +indexing: + # colon seperated mappings list by comma + # eg : dbgap:Non-HEAL Studies,bacpac:HEAL Research Programs + element_mapping: "" + variables_index: "variables_index" + studies_index: "studies_index" + sections_index: "sections_index" + concepts_index: "concepts_index" + kg_index: "kg_index" + tranql_min_score: 0.2 + excluded_identifiers: + - "CHEBI:17336" + queries: + "disease": ["disease", "phenotypic_feature"] + "pheno": ["phenotypic_feature", "disease"] + "anat": ["disease", "anatomical_entity"] + "chem_to_disease": ["chemical_entity", "disease"] + "small_molecule_to_disease": ["small_molecule", "disease"] + "chemical_mixture_to_disease": ["chemical_mixture", "disease"] + "phen_to_anat": ["phenotypic_feature", "anatomical_entity"] + tranql_endpoint: "http://tranql-service/tranql/query?dynamic_id_resolution=true&asynchronous=false" + node_to_element_queries: + enabled: false + cde: + node_type: biolink:Publication + curie_prefix: "HEALCDE" + list_field_choose_first: + - "files" + attribute_mapping: + name: "name" + desc: "summary" + collection_name: "cde_category" + collection_id: "cde_category" + action: "files" + +elasticsearch: + host: localhost + username: elastic + password: "12345" + nboost_host: "" + scheme: "http" + ca_path: "" + +validation: + queries: + count_nodes: + name: "Count Nodes" + query: "MATCH (a) RETURN COUNT(a)" + count_edges: + name: "Count Edges" + query: "MATCH (a)-[e]-(b) RETURN COUNT(e)" + connectivity: + name: TOPMED Connectivity + query: "MATCH (a { id : '$var' })--(b) RETURN a.category, b.id" + args: + - var: TOPMED.TAG:8 + - var: TOPMED.VAR:phv00000484.v1.p10 + - var: TOPMED.VAR:phv00000487.v1.p10 + - var: TOPMED.VAR:phv00000496.v1.p10 + - var: TOPMED.VAR:phv00000517.v1.p10 + - var: TOPMED.VAR:phv00000518.v1.p10 + - var: TOPMED.VAR:phv00000528.v1.p10 + - var: TOPMED.VAR:phv00000529.v1.p10 + - var: TOPMED.VAR:phv00000530.v1.p10 + - var: TOPMED.VAR:phv00000531.v1.p10 + count_connected_nodes: + name: Count Connected Nodes + query: "MATCH (a)-[e]-(b) RETURN count(a), count(b)" + query_by_type: + name: Query by Type + query: "MATCH (a:gene)-[e]-(b) WHERE 'chemical_substance' IN b.category RETURN count(distinct(a)), count(distinct(b))" + smiles_values: + name: Query Chemicals with smiles that look like arrays + query: "Match (a: chemical_substance { simple_smiles: '$var' }) RETURN a.id" + args: + - var: "[Os+6]" + - var: "[SiH2]" + - var: "[CH]" + - var: "[S-2]" + - var: "[Ti+4]" + - var: "[P-3]" + - var: "[Ca+2]" + - var: "[Au+3]" + - var: "[TeH2]" + - var: "[Pb]" + - var: "[B+]" + - var: "[AsH]" + - var: "[O-][I+2]([O-])[O-]" + - var: "[He+]" + - var: "[Mo+6]" + - var: "[N-]=[N+]=[N-]" + - var: "[Ag+]" + - var: "[Zn+2]" + - var: "[C-]#[O+]" +s3: + host: "" + bucket: "" + access_key: "" + secret_key: "" + +lakefs_config: + enabled: false + access_key_id: "" + secret_access_key: "" + host: "" + branch: "" + repo: "" diff --git a/src/roger/config/dev-config.yaml b/src/roger/config/dev-config.yaml new file mode 100644 index 00000000..bece11a8 --- /dev/null +++ b/src/roger/config/dev-config.yaml @@ -0,0 +1,118 @@ +redisgraph: + username: "" + password: "" + host: redis + graph: test + port: 6379 + +logging: + level: DEBUG + format: '[%(name)s][%(filename)s][%(funcName)20s] %(levelname)s: %(message)s' + +data_root: "/Users/schreepc/Projects/helxplatform/roger/roger/test/data" +dug_data_root: dug_helpers/dug_data/topmed_data +base_data_uri: https://stars.renci.org/var/kgx_data/trapi-1.0/ +kgx: + biolink_model_version: test + +#https://github.com/RedisGraph/redisgraph-bulk-loader/blob/master/redisgraph_bulk_loader/bulk_insert.py#L43 +bulk_loader: + separator: 0x1E + enforce_schema: False + skip_invalid_nodes: False + skip_invalid_edges: False + quote: 0 + max_token_count: 1024 + max_buffer_size: 2048 + max_token_size: 500 + index: [] + full_text_index: [] + +annotation: + annotator: "https://api.monarchinitiative.org/api/nlp/annotate/entities?min_length=4&longest_only=false&include_abbreviation=false&include_acronym=false&include_numbers=false&content=" + normalizer: "https://nodenormalization-sri.renci.org/get_normalized_nodes?curie=" + synonym_service: "https://onto.renci.org/synonyms/" + ontology_metadata: "https://api.monarchinitiative.org/api/ontology/term/" + # The following are neo4j params that would not be used + # need to remove them from annotator constructor. + db_url: "" + username: "" + password: "" + +indexing: + variables_index: "variables_index" + concepts_index: "concepts_index" + kg_index: "kg_index" + tranql_min_score: 0.2 + excluded_identifiers: + - "CHEBI:17336" + queries: + "disease": ["disease", "phenotypic_feature"] + "pheno": ["phenotypic_feature", "disease"] + "anat": ["disease", "anatomical_entity"] + "chem_to_disease": ["chemical_substance", "disease"] + "phen_to_anat": ["phenotypic_feature", "anatomical_entity"] + "anat_to_disease": ["anatomical_entity", "disease"] + "anat_to_pheno": ["anatomical_entity", "phenotypic_feature"] + tranql_endpoint: "http://tranql-service/tranql/query?dynamic_id_resolution=true&asynchronous=false" + +elasticsearch: + host: elasticsearch + username: elastic + # temporary + password: "13431" + nboost_host: "" + + + +validation: + queries: + count_nodes: + name: "Count Nodes" + query: "MATCH (a) RETURN COUNT(a)" + count_edges: + name: "Count Edges" + query: "MATCH (a)-[e]-(b) RETURN COUNT(e)" + connectivity: + name: TOPMED Connectivity + query: "MATCH (a { id : '$var' })--(b) RETURN a.category, b.id" + args: + - var: TOPMED.TAG:8 + - var: TOPMED.VAR:phv00000484.v1.p10 + - var: TOPMED.VAR:phv00000487.v1.p10 + - var: TOPMED.VAR:phv00000496.v1.p10 + - var: TOPMED.VAR:phv00000517.v1.p10 + - var: TOPMED.VAR:phv00000518.v1.p10 + - var: TOPMED.VAR:phv00000528.v1.p10 + - var: TOPMED.VAR:phv00000529.v1.p10 + - var: TOPMED.VAR:phv00000530.v1.p10 + - var: TOPMED.VAR:phv00000531.v1.p10 + count_connected_nodes: + name: Count Connected Nodes + query: "MATCH (a)-[e]-(b) RETURN count(a), count(b)" + query_by_type: + name: Query by Type + query: "MATCH (a:gene)-[e]-(b) WHERE 'chemical_substance' IN b.category RETURN count(distinct(a)), count(distinct(b))" + smiles_values: + name: Query Chemicals with smiles that look like arrays + query: "Match (a: chemical_substance { simple_smiles: '$var' }) RETURN a.id" + args: + - var: "[Os+6]" + - var: "[SiH2]" + - var: "[CH]" + - var: "[S-2]" + - var: "[Ti+4]" + - var: "[P-3]" + - var: "[Ca+2]" + - var: "[Au+3]" + - var: "[TeH2]" + - var: "[Pb]" + - var: "[B+]" + - var: "[AsH]" + - var: "[O-][I+2]([O-])[O-]" + - var: "[He+]" + - var: "[Mo+6]" + - var: "[N-]=[N+]=[N-]" + - var: "[Ag+]" + - var: "[Zn+2]" + - var: "[C-]#[O+]" diff --git a/src/roger/config/s3_config.py b/src/roger/config/s3_config.py new file mode 100644 index 00000000..41fcccab --- /dev/null +++ b/src/roger/config/s3_config.py @@ -0,0 +1,11 @@ +from dataclasses import dataclass + +from ._base import DictLike + + +@dataclass +class S3Config(DictLike): + host: str = "" + bucket: str = "" + access_key: str = "" + secret_key: str = "" \ No newline at end of file diff --git a/src/roger/config/test-config.yaml b/src/roger/config/test-config.yaml new file mode 100644 index 00000000..3149ba15 --- /dev/null +++ b/src/roger/config/test-config.yaml @@ -0,0 +1,179 @@ +redisgraph: + username: "" + password: "weak" + host: localhost + graph: test + port: 6379 + use_redis_cache: False + +logging: + level: DEBUG + format: '[%(name)s][%(filename)s][%(lineno)d][%(funcName)20s] %(levelname)s: %(message)s' + +data_root: roger/data + +kgx_base_data_uri: https://stars.renci.org/var/kgx_data/ +annotation_base_data_uri: https://stars.renci.org/var/dug/ + +kgx: + biolink_model_version: v3.1.2 + merge_db_temp_dir: workspace + data_sets: + - baseline-graph:v5.0 + +dug_inputs: + data_source: s3 + data_sets: + - topmed:v1.0 + - bdc:v1.0 + - anvil:v1.0 + +#https://github.com/RedisGraph/redisgraph-bulk-loader/blob/master/redisgraph_bulk_loader/bulk_insert.py#L43 +bulk_loader: + separator: 0x1E + enforce_schema: False + skip_invalid_nodes: False + skip_invalid_edges: False + quote: 0 + max_token_count: 1024 + max_buffer_size: 2048 + max_token_size: 500 + index: [] + full_text_index: [] + +annotation: + clear_http_cache: false + annotator_type: sapbert + annotator_args: + monarch: + url: "https://api.monarchinitiative.org/api/nlp/annotate/entities?min_length=4&longest_only=false&include_abbreviation=false&include_acronym=false&include_numbers=false&content=" + sapbert: + classification_url: "https://med-nemo.apps.renci.org/annotate/" + annotator_url: "https://sap-qdrant.apps.renci.org/annotate/" + score_threshold: 0.8 + bagel: + enabled: false + url: "http://localhost:9099/group_synonyms_openai" + prompt: "bagel/ask_classes" + llm_args: + llm_model_name: "gpt-4o-2024-05-13" + organization: + access_key: + llm_model_args: + top_p: 0 + temperature: 0.1 + normalizer: "https://nodenormalization-exp.apps.renci.org/get_normalized_nodes?conflate=false&description=true&curie=" + synonym_service: "https://name-resolution-sri.renci.org/reverse_lookup" + ontology_metadata: "https://api.monarchinitiative.org/api/bioentity/" + + preprocessor: + debreviator: + BMI: "body mass index" + stopwords: "the" + ontology_greenlist: ["PATO", "CHEBI", "MONDO", "UBERON", "HP", "MESH", "UMLS"] + +indexing: + # colon seperated mappings list by comma + # eg : dbgap:Non-HEAL Studies,bacpac:HEAL Research Programs + element_mapping: "" + variables_index: "variables_index" + concepts_index: "concepts_index" + kg_index: "kg_index" + tranql_min_score: 0.2 + excluded_identifiers: + - "CHEBI:17336" + queries: + "disease": ["disease", "phenotypic_feature"] + "pheno": ["phenotypic_feature", "disease"] + "anat": ["disease", "anatomical_entity"] + "chem_to_disease": ["chemical_entity", "disease"] + "small_molecule_to_disease": ["small_molecule", "disease"] + "chemical_mixture_to_disease": ["chemical_mixture", "disease"] + "phen_to_anat": ["phenotypic_feature", "anatomical_entity"] + tranql_endpoint: "http://tranql-service/tranql/query?dynamic_id_resolution=true&asynchronous=false" + node_to_element_queries: + enabled: false + cde: + node_type: biolink:Publication + curie_prefix: "HEALCDE" + list_field_choose_first: + - "files" + attribute_mapping: + name: "name" + desc: "summary" + collection_name: "cde_category" + collection_id: "cde_category" + action: "files" + +elasticsearch: + host: localhost + username: elastic + password: "12345" + nboost_host: "" + scheme: "http" + ca_path: "" + +validation: + queries: + count_nodes: + name: "Count Nodes" + query: "MATCH (a) RETURN COUNT(a)" + count_edges: + name: "Count Edges" + query: "MATCH (a)-[e]-(b) RETURN COUNT(e)" + connectivity: + name: TOPMED Connectivity + query: "MATCH (a { id : '$var' })--(b) RETURN a.category, b.id" + args: + - var: TOPMED.TAG:8 + - var: TOPMED.VAR:phv00000484.v1.p10 + - var: TOPMED.VAR:phv00000487.v1.p10 + - var: TOPMED.VAR:phv00000496.v1.p10 + - var: TOPMED.VAR:phv00000517.v1.p10 + - var: TOPMED.VAR:phv00000518.v1.p10 + - var: TOPMED.VAR:phv00000528.v1.p10 + - var: TOPMED.VAR:phv00000529.v1.p10 + - var: TOPMED.VAR:phv00000530.v1.p10 + - var: TOPMED.VAR:phv00000531.v1.p10 + count_connected_nodes: + name: Count Connected Nodes + query: "MATCH (a)-[e]-(b) RETURN count(a), count(b)" + query_by_type: + name: Query by Type + query: "MATCH (a:gene)-[e]-(b) WHERE 'chemical_substance' IN b.category RETURN count(distinct(a)), count(distinct(b))" + smiles_values: + name: Query Chemicals with smiles that look like arrays + query: "Match (a: chemical_substance { simple_smiles: '$var' }) RETURN a.id" + args: + - var: "[Os+6]" + - var: "[SiH2]" + - var: "[CH]" + - var: "[S-2]" + - var: "[Ti+4]" + - var: "[P-3]" + - var: "[Ca+2]" + - var: "[Au+3]" + - var: "[TeH2]" + - var: "[Pb]" + - var: "[B+]" + - var: "[AsH]" + - var: "[O-][I+2]([O-])[O-]" + - var: "[He+]" + - var: "[Mo+6]" + - var: "[N-]=[N+]=[N-]" + - var: "[Ag+]" + - var: "[Zn+2]" + - var: "[C-]#[O+]" +s3: + host: "" + bucket: "" + access_key: "" + secret_key: "" + +lakefs_config: + enabled: false + access_key_id: "" + secret_access_key: "" + host: "" + branch: "" + repo: "" diff --git a/src/roger/core/__init__.py b/src/roger/core/__init__.py new file mode 100644 index 00000000..5a5c5e90 --- /dev/null +++ b/src/roger/core/__init__.py @@ -0,0 +1,4 @@ +"Core roger modules, now broken out into a submodule" + +from roger.core.enums import SchemaType, FileFormat +from roger.core.bulkload import BulkLoad diff --git a/src/roger/core/base.py b/src/roger/core/base.py new file mode 100644 index 00000000..7ba9409a --- /dev/null +++ b/src/roger/core/base.py @@ -0,0 +1,217 @@ +"Core Roger object and utilities" + +import argparse +import sys +from io import StringIO +import logging + +from roger.config import get_default_config as get_config +from roger.logger import get_logger +from roger.core.bulkload import BulkLoad +from roger.models.kgx import KGXModel +from roger.models.biolink import BiolinkModel + +log = get_logger() + +class Roger: + """ Consolidate Roger functionality for a cleaner interface. """ + + def __init__(self, to_string=False, config=None): + """ Initialize. + :param to_string: Log to str, available as self.log_stream.getvalue() + after execution completes. + """ + self.has_string_handler = to_string + if not config: + config = get_config() + self.config = config + if to_string: + # Add a stream handler to enable to_string. + self.log_stream = StringIO() + self.string_handler = logging.StreamHandler (self.log_stream) + log.addHandler (self.string_handler) + log.debug("config is %s", config.kgx.biolink_model_version) + self.biolink = BiolinkModel (config.kgx.biolink_model_version) + self.kgx = KGXModel (self.biolink, config=config) + self.bulk = BulkLoad (self.biolink, config=config) + + def __enter__(self): + """ Implement Python's Context Manager interface. """ + return self + + def __exit__(self, exception_type, exception_value, traceback): + """ Implement Python's Context Manager interface. We use this finalizer + to detach the stream handler appended in the constructor. + :param exception_type: Type of exception, if one occurred. + :param exception_value: The exception, if one occurred. + :param traceback: The stack trace explaining the exception. + """ + if exception_type or exception_value or traceback: + log.error (msg="Error:", + exc_info=(exception_type, exception_value, traceback)) + if self.has_string_handler: + log.removeHandler (self.string_handler) + +# interfaces abstracting Roger's inner workings to make it easier to +# incorporate into external tools like workflow engines. + +def get_kgx (to_string=False, config=None): + "get KGX dataset" + output = None + log.debug("Getting KGX method called.") + with Roger (to_string, config=config) as roger: + dataset_version=config.get('kgx', {}).get('dataset_version') + log.debug("dataset_version is %s", dataset_version) + roger.kgx.get(dataset_version=dataset_version) + output = roger.log_stream.getvalue() if to_string else None + return output + +def create_schema(to_string=False, config=None): + "Create noders and edges schemata" + o1 = create_nodes_schema(to_string=to_string, config=config) + o2 = create_edges_schema(to_string=to_string, config=config) + output = (o1 + o2 ) if to_string else None + return output + +def create_edges_schema(to_string=False, config=None, input_data_path=None, output_data_path=None): + "Create edges schema on KGX object" + output = None + with Roger(to_string, config=config) as roger: + roger.kgx.create_edges_schema( + input_data_path=input_data_path, + output_data_path=output_data_path + ) + output = roger.log_stream.getvalue() if to_string else None + return output + +def create_nodes_schema(to_string=False, config=None, input_data_path=None, output_data_path=None): + "Create nodes schema on KGX object" + output = None + with Roger(to_string, config=config) as roger: + roger.kgx.create_nodes_schema(input_data_path=input_data_path, + output_data_path=output_data_path) + output = roger.log_stream.getvalue() if to_string else None + return output + +def merge_nodes(to_string=False, config=None, input_data_path=None, output_data_path=None): + "Run KGX merge" + output = None + with Roger (to_string, config=config) as roger: + roger.kgx.merge(input_path=input_data_path, output_path=output_data_path) + output = roger.log_stream.getvalue () if to_string else None + return output + +def create_bulk_load(to_string=False, config=None, input_data_path=None, output_data_path=None): + "Generate bulk load files" + o1 = create_bulk_nodes(to_string=to_string, config=config) + o2 = create_bulk_edges(to_string=to_string, config=config) + output = (o1 + o2) if to_string else None + return output + +def create_bulk_nodes(to_string=False, config=None, input_data_path=None, output_data_path=None): + "Generate bulk node CSV file" + output = None + with Roger(to_string, config=config) as roger: + log.info("input path: %s", input_data_path) + log.info("output path: %s", output_data_path) + roger.bulk.create_nodes_csv_file(input_data_path, output_data_path) + output = roger.log_stream.getvalue() if to_string else None + return output + +def create_bulk_edges(to_string=False, config=None, input_data_path=None, output_data_path=None): + "Create bulk edges CSV file" + output = None + with Roger(to_string, config=config) as roger: + roger.bulk.create_edges_csv_file(input_data_path, output_data_path) + output = roger.log_stream.getvalue() if to_string else None + return output + +def bulk_load(to_string=False, config=None, input_data_path=None, output_data_path=None): + "Run bulk load insert process" + output = None + with Roger (to_string, config=config) as roger: + roger.bulk.insert(input_data_path=input_data_path) + output = roger.log_stream.getvalue () if to_string else None + return output + +def validate (to_string=False, config=None, input_data_path=None, output_data_path=None): + "Run bulk validate process" + output = None + with Roger (to_string, config=config) as roger: + roger.bulk.validate() + output = roger.log_stream.getvalue () if to_string else None + return output + +def check_tranql(to_string=False, config=None, input_data_path=None, output_data_path=None): + "Tranql server smoke check" + output = None + with Roger(to_string, config=config) as roger: + roger.bulk.wait_for_tranql() + output = roger.log_stream.getvalue() if to_string else None + return output + +def roger_cli(): + " Roger CLI. " + parser = argparse.ArgumentParser(description='Roger') + parser.add_argument('-v', + '--dataset-version', + help="Dataset version.", + default="v1.0") + parser.add_argument('-d', + '--data-root', + help="Root of data hierarchy", + default=None) + parser.add_argument('-g', + '--get-kgx', + help="Get KGX objects", + action='store_true') + parser.add_argument('-l', + '--load-kgx', + help="Load via KGX", + action='store_true') + parser.add_argument('-s', + '--create-schema', + help="Infer schema", + action='store_true') + parser.add_argument('-m', + '--merge-kgx', + help="Merge KGX nodes", + action='store_true') + parser.add_argument('-b', + '--create-bulk', + help="Create bulk load", + action='store_true') + parser.add_argument('-i', + '--insert', + help="Do the bulk insert", + action='store_true') + parser.add_argument('-a', + '--validate', + help="Validate the insert", + action='store_true') + args = parser.parse_args () + + biolink = BiolinkModel () + kgx = KGXModel (biolink) + bulk = BulkLoad (biolink) + if args.data_root is not None: + config = get_config() + data_root = args.data_root + config.update({'data_root': data_root}) + log.info("data root: %s", data_root) + if args.get_kgx: + kgx.get (dataset_version=args.dataset_version) + if args.load_kgx: + kgx.load () + if args.merge_kgx: + kgx.merge () + if args.create_schema: + kgx.create_schema () + if args.create_bulk: + bulk.create () + if args.insert: + bulk.insert () + if args.validate: + bulk.validate () + + sys.exit (0) diff --git a/src/roger/core/bulkload.py b/src/roger/core/bulkload.py new file mode 100644 index 00000000..4dc2736b --- /dev/null +++ b/src/roger/core/bulkload.py @@ -0,0 +1,506 @@ +"Bulk loader for Roger" + +import os +import glob +import shutil +from collections import defaultdict +from functools import reduce +from string import Template +import time + +import requests +import redis +from contextlib import contextmanager +from falkordb_bulk_loader.bulk_insert import bulk_insert + +from roger.config import get_default_config as get_config +from roger.logger import get_logger +from roger.core.redis_graph import RedisGraph +from roger.core.enums import SchemaType +from roger.models.biolink import BiolinkModel +from roger.components.data_conversion import cast +from roger.core import storage + +log = get_logger() + +class BulkLoad: + """ Tools for creating a Redisgraph bulk load dataset. """ + def __init__(self, biolink, config=None): + self.biolink = biolink + if not config: + config = get_config() + self.config = config + separator = self.config.get('bulk_loader',{}).get('separator', '|') + self.separator =(chr(separator) if isinstance(separator, int) + else separator) + + def create (self): + """Used in the CLI on args.create_bulk""" + self.create_nodes_csv_file() + self.create_edges_csv_file() + + def create_nodes_csv_file(self, input_data_path=None, output_data_path=None): + # clear out previous data + bulk_path = storage.bulk_path("nodes", output_data_path) + if os.path.exists(bulk_path): + shutil.rmtree(bulk_path) + categories_schema = storage.read_schema (SchemaType.CATEGORY, input_data_path) + state = defaultdict(lambda: None) + log.info(f"processing nodes") + """ Write node data for bulk load. """ + + categories = defaultdict(lambda: []) + category_error_nodes = set() + merged_nodes_file = storage.merged_objects('nodes', input_data_path) + counter = 1 + for node in storage.json_line_iter(merged_nodes_file): + if node.get('description'): + node['description'] = node['description'].replace('\n', + ' ') + if node.get('name'): + node['name'] = node['name'].replace('\n', + ' ') + if not node.get('category'): + category_error_nodes.add(node['id']) + node['category'] = [BiolinkModel.root_type] + index = self.biolink.get_leaf_class(node['category']) + categories[index].append(node) + if category_error_nodes: + log.error( + f"some nodes didn't have category assigned. " + f"KGX file has errors. " + f"Nodes {len(category_error_nodes)}. " + f"They will be typed {BiolinkModel.root_type}. " + f"Showing first 10: {list(category_error_nodes)[:10]}.") + # flush every 100K + if counter % 100_000 == 0: + self.write_bulk(storage.bulk_path("nodes", output_data_path), + categories, categories_schema, + state=state, is_relation=False) + # reset variables. + category_error_nodes = set() + categories = defaultdict(lambda: []) + counter += 1 + # write back if any thing left. + if len(categories): + self.write_bulk(storage.bulk_path("nodes", output_data_path), + categories, categories_schema, + state=state, is_relation=False) + + def create_edges_csv_file(self, input_data_path=None, output_data_path=None): + """ Write predicate data for bulk load. """ + # Clear out previous data + bulk_path = storage.bulk_path("edges", output_data_path) + if os.path.exists(bulk_path): + shutil.rmtree(bulk_path) + predicates_schema = storage.read_schema(SchemaType.PREDICATE, input_data_path) + predicates = defaultdict(lambda: []) + edges_file = storage.merged_objects('edges', input_data_path) + counter = 1 + state = {} + for edge in storage.json_line_iter(edges_file): + predicates[edge['predicate']].append(edge) + # write out every 100K , to avoid large predicate dict. + if counter % 100_000 == 0: + self.write_bulk( + storage.bulk_path("edges", output_data_path),predicates, predicates_schema, + state=state, is_relation=True) + predicates = defaultdict(lambda : []) + counter += 1 + # if there are some items left (if loop ended before counter reached the + # specified value) + if len(predicates): + self.write_bulk(storage.bulk_path("edges", output_data_path), predicates, + predicates_schema,state=state, is_relation=True) + + @staticmethod + def create_redis_schema_header(attributes: dict, is_relation=False): + """Creates col headers for csv to be used by redis bulk loader + + Column headers are generated by assigning redis types + :param attributes: dict of data labels with values as python type strs + :param separator: CSV separator + :return: list of attrs, each item is attributeLabel:redisGraphDataType + """ + redis_type_conversion_map = { + 'str': 'STRING', + 'float': 'FLOAT', # Do we need to handle double + 'int': 'INT', + 'bool': 'BOOL', + 'list': 'ARRAY' + } + col_headers = [] + def format_for_redis(label, typ): + return f'{label}:{typ}' + for attribute, attribute_type in attributes.items(): + col_headers.append(format_for_redis( + attribute, redis_type_conversion_map[attribute_type])) + # Note this two fields are only important to bulk loader + # they will not be members of the graph + # https://github.com/RedisGraph/redisgraph-bulk-loader/tree/master#input-schemas + if is_relation: + col_headers.append('internal_start_id:START_ID') + col_headers.append('internal_end_id:END_ID') + # replace id:STRING with id:ID + col_headers.append('id:ID') + col_headers = list(filter(lambda x: x != 'id:STRING', col_headers)) + return col_headers + + @staticmethod + def group_items_by_attributes_set(objects: list, processed_object_ids: set): + """ Groups items into a dictionary + + The keys the output dictionary are sets of attributes set for all + items accessed in that key. + + Eg.: + { set(id,name,category): [{id:'xx0',name:'bbb', 'category':['type']}.... + {id:'xx1', name:'bb2', category: ['type1']}] } + :param objects: list of nodes or edges + :param processed_object_ids: ids to skip since they are processed. + :return: dictionary grouping based on set attributes + """ + clustered_by_set_values = {} + improper_keys = set() + def value_set_test(val): + "Converted from lambda function, is this just 'if x:'?" + if (val is not None and val != [] and val != ''): + return True + return False + for obj in objects: + # redis bulk loader needs columns not to include ':' + # till backticks are implemented we should avoid these. + def key_filter(key): + # Make sure no colons in key names + return ':' not in key + keys_with_values = frozenset( + [k for k in obj.keys() + if value_set_test(obj[k]) and key_filter(k)]) + for key in [k for k in obj.keys() if obj[k] and not key_filter(k)]: + improper_keys.add(key) + # group by attributes that have values. # Why? + # Redis bulk loader has one issue + # imagine we have: + # + #{'name': 'x'} , {'name': 'y', 'is_metabolite': true} + # + # we have a common schema name:STRING,is_metabolite: + # + # BOOL values `x,` and `y,true` + # + # but x not having value for is_metabolite is not handled well, + # redis bulk loader says we should give it default if we were to + # enforce schema but due to the nature of the data assigning + # defaults is very not an option. hence grouping data into several + # csv's might be the right way (?) + if obj['id'] not in processed_object_ids: + val_list = clustered_by_set_values.get(keys_with_values, []) + val_list.append(obj) + clustered_by_set_values[keys_with_values] = val_list + return clustered_by_set_values, improper_keys + + def write_bulk(self, bulk_path, obj_map, schema, state={}, + is_relation=False): + """ Write a bulk load group of objects. + :param bulk_path: Path to the bulk loader object to write. + :param obj_map: A map of biolink type to list of objects. + :param schema: The schema (nodes or predicates) containing identifiers. + :param state: Track state of already written objs to avoid duplicates. + """ + + os.makedirs (bulk_path, exist_ok=True) + processed_objects_id = state.get('processed_id', set()) + called_x_times = state.get('called_times', 0) + called_x_times += 1 + for key, objects in obj_map.items (): + if len(objects) == 0: + continue + try: + all_keys = schema[key] + except Exception as e: + log.error(f"{key} not in {schema.keys()} " ) + raise Exception("error") from e + """ Make all objects conform to the schema. """ + clustered_by_set_values, improper_redis_keys = ( + self.group_items_by_attributes_set(objects, + processed_objects_id)) + + if improper_redis_keys: + log.warning( + "The following keys were skipped since they include " + "conflicting `:` that would cause errors while bulk " + "loading to redis. [%s]", str(improper_redis_keys)) + for index, set_attributes in enumerate( + clustered_by_set_values.keys()): + items = clustered_by_set_values[set_attributes] + # When parted files are saved let the file names be collected + # here + state['file_paths'] = state.get('file_paths', {}) + state['file_paths'][key] = state['file_paths'].get(key, {}) + out_file = state['file_paths'][key][set_attributes] = ( + state['file_paths'].get(key, {}).get(set_attributes, '')) + + # When calling write bulk , lets say we have processed some + # chemicals from file 1 and we start processing file 2 if we are + # using just index then we might (rather will) end up adding + # records to the wrong file so we need this to be unique as + # possible by adding called_x_times , if we already found + # out-file from state obj we are sure that the schemas match. + + # biolink: is not valid name so we need to remove : + file_key = key.replace(':', '~') + + out_file = ( + f"{bulk_path}/{file_key}.csv-{index}-{called_x_times}" + if not out_file + else out_file) + # store back file name + state['file_paths'][key][set_attributes] = out_file + new_file = not os.path.exists(out_file) + keys_for_header = {x: all_keys[x] for x in all_keys + if x in set_attributes} + redis_schema_header = self.create_redis_schema_header( + keys_for_header, is_relation) + with open(out_file, "a", encoding='utf-8') as stream: + if new_file: + state['file_paths'][key][set_attributes] = out_file + log.info(f" --creating {out_file}") + stream.write(self.separator.join(redis_schema_header)) + stream.write("\n") + else: + log.info(f" --appending to {out_file}") + + # Write fields, skipping duplicate objects. + for obj in items: + oid = str(obj['id']) + if oid in processed_objects_id: + continue + processed_objects_id.add(oid) + + # Add ID / START_ID / END_ID depending + internal_id_fields = { + 'internal_id': obj['id'] + } + if is_relation: + internal_id_fields.update({ + 'internal_start_id': obj['subject'], + 'internal_end_id': obj['object'] + }) + obj.update(internal_id_fields) + values = [] + + # uses redis schema header to preserve order when + # writing lines out. + for column_name in redis_schema_header: + # last key is the type + obj_key = ':'.join(column_name.split(':')[:-1]) + value = obj[obj_key] + + if obj_key not in internal_id_fields: + current_type = type(value).__name__ + expected_type = all_keys[obj_key] + # cast it if it doesn't match type in schema + # keys i.e all_keys + value = ( + cast(obj[obj_key], all_keys[obj_key]) + if expected_type != current_type + else value) + # escape quotes . + values.append(str(value).replace("\"", "\\\"")) + s = self.separator.join(values) + stream.write(s) + stream.write("\n") + state['processed_id'] = processed_objects_id + state['called_times'] = called_x_times + + def insert (self, input_data_path=None): + redisgraph = self.config.redisgraph + nodes = sorted(glob.glob (storage.bulk_path ("**/nodes/**.csv*", input_data_path), recursive=True)) + edges = sorted(glob.glob (storage.bulk_path ("**/edges/**.csv*", input_data_path), recursive=True)) + graph = redisgraph['graph'] + log.info(f"bulk loading \n nodes: {nodes} \n edges: {edges}") + + # An empty edge set is a valid loader invocation, so a build whose + # edge csvs never arrived loads nodes only and reports success. That + # is how the graph ended up with 3.9M nodes and no relationships: + # BulkLoad resolved the lakefs tip 100s before CreateBulkLoadEdges + # committed its 12GB of edges, so it read a commit where the prefix + # was still empty. + if nodes and not edges: + raise ValueError( + f"{len(nodes)} node csv(s) but no edge csv(s) under " + f"{storage.bulk_path('**/edges', input_data_path)}. " + "Refusing to bulk load an edgeless graph; check that " + "CreateBulkLoadEdges committed before this task ran.") + + try: + log.info (f"deleting graph {graph} in preparation for bulk load.") + db = self.get_redisgraph() + db.redis_graph.delete () + except redis.exceptions.ResponseError: + log.info("no graph to delete") + + log.info ("bulk loading graph: %s", str(graph)) + args = [] + collect_labels = set() + if len(nodes) > 0: + bulk_path_root = glob.glob(storage.bulk_path('**/nodes', path=input_data_path), recursive=True)[0] + os.path.sep + nodes_with_type = [] + collect_labels = set() + for x in nodes: + """ + These lines prep nodes bulk load by: + 1) appending to labels 'biolink.' + 2) combine labels to create a multilabel redis node i.e. "biolink.OrganismalEntity:biolink.SubjectOfInvestigation" + """ + file_name_type_part = x.replace(bulk_path_root, '').split('.')[0].split('~')[1] + all_labels = "biolink." + file_name_type_part + ":" + ":".join([f'biolink.{v.lstrip("biolink:")}' for v in self.biolink.toolkit.get_ancestors("biolink:" + file_name_type_part, reflexive=False, formatted=True )] ) + collect_labels.add("biolink." + file_name_type_part) + for v in self.biolink.toolkit.get_ancestors("biolink:" + file_name_type_part, reflexive=False, + formatted=True): + collect_labels.add(f'biolink.{v.lstrip("biolink:")}') + nodes_with_type.append(f"{all_labels} {x}") + args.extend(("-N " + " -N ".join(nodes_with_type)).split()) + if len(edges) > 0: + bulk_path_root = glob.glob(storage.bulk_path('**/edges', path=input_data_path), recursive=True)[0] + os.path.sep + edges_with_type = [f"biolink.{x.replace(bulk_path_root, '').strip(os.path.sep).split('.')[0].split('~')[1]} {x}" + for x in edges] + # Edge label now no longer has 'biolink:' + args.extend(("-R " + " -R ".join(edges_with_type)).split()) + args.extend([f"--separator={self.separator}"]) + args.extend([f"--server-url=redis://:{redisgraph['password']}@{redisgraph['host']}:{redisgraph['port']}"]) + args.extend(['--enforce-schema']) + args.extend(['-e']) + for lbl in collect_labels: + # Backtick every label. falkordb interpolates it straight into + # the pattern (`(e:{label})` in Graph._create_typed_index), and + # biolink labels contain a dot, so an unquoted one is a syntax + # error: "Invalid input '.': expected ')'". The loader catches + # that and only prints it, so every full text index silently + # failed to be created while the range indexes -- already + # quoted here -- succeeded. + args.extend([f'-i `{lbl}`:id', + f'-f `{lbl}`:name', + f'-f `{lbl}`:synonyms']) + args.extend([f"{redisgraph['graph']}"]) + """ standalone_mode=False tells click not to sys.exit() """ + log.debug(f"Calling bulk_insert with extended args: {args}") + with self.snapshots_paused(): + try: + bulk_insert(args, standalone_mode=False) + # self.add_indexes() + except Exception as e: + log.error(f"Unexpected {e.__class__.__name__}: {e}") + raise + + @contextmanager + def snapshots_paused(self): + """Turn off RDB snapshots for the body, then put them back. + + `save` triggers on write volume and a bulk load is nothing but + write volume, so bgsave forks over and over during the load. + Copy-on-write on a graph this size can add most of it again on top + of the resident set and OOMKill the pod -- the headroom between + redis maxmemory and the pod memory limit is not sized for it. A + snapshot taken partway through a load is worthless anyway: the + bulk csvs are the source of truth and the load starts by deleting + the graph. + + Failing to set this is not fatal. The load still runs, it just + runs with snapshots on, which is where we were before. + """ + client = self.get_redisgraph().r + prior = None + try: + # redis-py may hand back bytes depending on version/decoding + cfg = {(k.decode() if isinstance(k, bytes) else k): + (v.decode() if isinstance(v, bytes) else v) + for k, v in client.config_get("save").items()} + prior = cfg.get("save", "") + client.config_set("save", "") + log.info("Paused redis snapshots for bulk load (was save=%r)", + prior) + except Exception as e: + log.warning("Could not pause redis snapshots, loading with " + "them on: %s", e) + try: + yield + finally: + if prior is None: + return + try: + client.config_set("save", prior) + log.info("Restored redis save=%r", prior) + except Exception as e: + # loud: snapshots are now off until someone restores them + log.error("FAILED to restore redis save=%r, snapshots are " + "still disabled: %s", prior, e) + + def add_indexes(self): + redis_connection = self.get_redisgraph() + all_labels = redis_connection.query( + "Match (c) return distinct labels(c)").result_set + all_labels = reduce(lambda x, y: x + y, all_labels, []) + id_index_queries = [ + f'CREATE INDEX on :`{label}`(id)' for label in all_labels + ] + name_index_queries = ( + "CALL db.labels() YIELD label " + "CALL db.idx.fulltext.createNodeIndex(label, 'name', 'synonyms')") + + for query in id_index_queries: + redis_connection.query(query=query) + redis_connection.query(query=name_index_queries) + log.info(f"Indexes created for {len(all_labels)} labels.") + + def get_redisgraph(self): + return RedisGraph( + host=self.config.redisgraph.host, + port=self.config.redisgraph.port, + password=self.config.redisgraph.password, + graph=self.config.redisgraph.graph, + ) + + def validate(self): + + db = self.get_redisgraph() + validation_queries = self.config.get( + 'validation', {}).get('queries', []) + for key, query in validation_queries.items (): + text = query['query'] + name = query['name'] + args = query.get('args', [{}]) + for arg in args: + start = storage.current_time_in_millis () + instance = Template (text).safe_substitute (arg) + db.query (instance) + duration = storage.current_time_in_millis () - start + log.info (f"Query {key}:{name} ran in {duration}ms: {instance}") + + def wait_for_tranql(self): + retry_secs = 3 + tranql_endpoint = self.config.indexing.tranql_endpoint + log.info(f"Contacting {tranql_endpoint}") + graph_name = self.config["redisgraph"]["graph"] + test_query = "SELECT disease-> phenotypic_feature " \ + f"FROM 'redis:{graph_name}'" \ + f"WHERE disease='MONDO:0004979'" + is_done_loading = False + try: + while not is_done_loading: + response = requests.post(tranql_endpoint, data=test_query) + response_code = response.status_code + response = response.json() + is_done_loading = "message" in response and response_code == 200 + if is_done_loading: + break + else: + log.info(f"Tranql responsed with response: {response}") + log.info(f"Retrying in {retry_secs} secs...") + time.sleep(retry_secs) + except ConnectionError as e: + # convert exception to be more readable. + raise ConnectionError( + f"Attempting to contact {tranql_endpoint} " + f"failed due to connection error. " + f"Please check status of Tranql server.") from e diff --git a/src/roger/core/enums.py b/src/roger/core/enums.py new file mode 100644 index 00000000..b44323af --- /dev/null +++ b/src/roger/core/enums.py @@ -0,0 +1,15 @@ +"Enums for Roger" + +from enum import Enum + +class SchemaType(Enum): + """ High level semantic metatdata concepts. + Categories are classes in an ontological model like Biolink. + Predicates are links between nodes. """ + CATEGORY = "category" + PREDICATE = "predicate" + +class FileFormat(Enum): + """ File formats this module knows about. """ + JSON = "json" + YAML = "yaml" diff --git a/roger/roger_db.py b/src/roger/core/redis_graph.py similarity index 87% rename from roger/roger_db.py rename to src/roger/core/redis_graph.py index 75b94e6b..5d89ef74 100644 --- a/roger/roger_db.py +++ b/src/roger/core/redis_graph.py @@ -1,18 +1,22 @@ -import logging +import copy + import redis -from redisgraph import Node, Edge, Graph -from redis.exceptions import ResponseError -from roger.roger_util import get_config, get_logger +# from redisgraph import Node, Edge, Graph +# https://redis-py.readthedocs.io/en/v4.5.1/redismodules.html#redisgraph-commands +from falkordb.node import Node +from falkordb.edge import Edge + +from roger.logger import get_logger logger = get_logger () class RedisGraph: """ Graph abstraction over RedisGraph. A thin wrapper but provides us some options. """ - def __init__(self, host='localhost', port=6379, graph='default'): + def __init__(self, host='localhost', port=6379, graph='default', password=''): """ Construct a connection to Redis Graph. """ - self.r = redis.Redis(host=host, port=port) - self.redis_graph = Graph(graph, self.r) + self.r = redis.Redis(host=host, port=port, password=password) + self.redis_graph = self.r.graph(graph) def add_node (self, identifier=None, label=None, properties=None): """ Add a node with the given label and properties. """ @@ -57,7 +61,7 @@ def commit (self): def query (self, query): """ Query and return result set. """ result = self.redis_graph.query(query) - result.pretty_print() + print(result) return result def delete (self): diff --git a/src/roger/core/storage.py b/src/roger/core/storage.py new file mode 100644 index 00000000..d10d0d62 --- /dev/null +++ b/src/roger/core/storage.py @@ -0,0 +1,511 @@ +""" utils for roger + +This is home to the utilities that were formerly in dags/roger/core.py:Util +""" + +import os +import glob +import time +import pathlib +import pickle +import shutil +import yaml +import orjson as json +import requests +from urllib.request import urlretrieve +from pathlib import Path + +from roger.logger import get_logger +from roger.config import get_default_config as get_config +from roger.core import SchemaType + +log = get_logger() +config = get_config() + +data_dir_env_value = os.getenv("ROGER_DATA_DIR") + +if data_dir_env_value is None: + ROGER_DATA_DIR = Path(__file__).parent.resolve() / 'data' +else: + ROGER_DATA_DIR = Path(data_dir_env_value) + + +def current_time_in_millis(): + """ + Get current time in milliseconds. + + Returns + ------- + int + Time in milliseconds + + """ + return int(round(time.time() * 1000)) + +# A just do it approach to getting data. +def read_file(path): + """ Read a file. + :param path: Path to a file. + """ + text = None + with open(path, "r", encoding='utf-8') as stream: + text = stream.read() + return text + +def read_url(url): + """ Read data from a URL. + :param url: The URL to read. """ + return requests.get(url, timeout=60).text + +def read_data(path): + """ Read data from a URL or File. HTTP(S) is the only supported protocol. + :param path: A URL or file path. """ + text = None + if is_web(path): + text = read_url(path) + else: + text = read_file(path) + return text + +def read_object(path, key=None): + """ Read on object from a path. + :param path: A URL or file path. + Supports YAML and JSON depending on extension. + :param key: A configuration key. This is prepended to the path if present. + :raises ValueError: If the key is not in the configuration. """ + if key is not None: + prefix = config[key] + path = f"{prefix}/{path}" if is_web(prefix) \ + else os.path.join (prefix, path) + obj = None + if path.endswith(".yaml") or path.endswith (".yml"): + obj = yaml.safe_load (read_data (path)) + elif path.endswith(".json"): + obj = json.loads (read_data (path)) + elif path.endswith(".pickle"): + with open(file=path, mode="rb") as stream: + obj = pickle.load(stream) + elif path.endswith(".jsonl") or path.endswith('.txt'): + obj = read_data(path) + return obj + +def is_web (uri): + """ The URI is a web URI (starts with http or https). + :param uri: A URI """ + return uri.startswith("http://") or uri.startswith ("https://") + +def write_object (obj, path, key=None): + """ Write an object to a path. YAML and JSON supported based on extension. + :param obj: The object to write. + :param path: The path to write to. + :param key: The configuration key to prepend to the path. + """ + # Prepend a prefix from the configuration file if a key is given. + if key is not None: + prefix = config[key] + path = (f"{prefix}/{path}" if is_web(prefix) + else os.path.join (prefix, path)) + + # Ensure the directory to be written to exists. + dirname = os.path.dirname (path) + if not os.path.exists (dirname): + os.makedirs (dirname, exist_ok=True) + + # Write the file in the specified format. + if path.endswith (".yaml") or path.endswith (".yml"): + with open(path, 'w') as outfile: + yaml.dump (obj, outfile) + elif path.endswith (".json"): + with open (path, "w", encoding='utf-8') as stream: + stream.write(str(json.dumps (obj, option=json.OPT_INDENT_2).decode('utf-8'))) + elif path.endswith(".pickle"): + with open (path, "wb") as stream: + pickle.dump(obj, file=stream) + elif path.endswith(".jsonl") or path.endswith('.txt'): + with open (path, "w", encoding="utf-8") as stream: + stream.write(obj) + else: + # Raise an exception if invalid. + raise ValueError (f"Unrecognized extension: {path}") + +def mkdir(path, is_dir=False): + directory = os.path.dirname(path) if not is_dir else path + if not os.path.exists(directory): + os.makedirs(directory) + +def remove(path): + if os.path.exists(path): + if os.path.isdir(path): + shutil.rmtree(path) + else: + os.remove(path) + +def clear_dir(path): + remove(path) + mkdir(path, is_dir=True) + +###################### +# Path methods + +def kgx_path(name): + """ Form a KGX object path. + :path name: Name of the KGX object. """ + return str(ROGER_DATA_DIR / "kgx" / name) + +def kgx_objects(format_="json", path=None): + """ A list of KGX objects. """ + kgx_pattern = kgx_path(f"**.{format_}") + if path: + kgx_pattern = f"{path}/**/*.{format_}" + return sorted(glob.glob (kgx_pattern, recursive=True)) + +def merge_path(name, path: Path=None): + """ Form a merged KGX object path. + :path name: Name of the merged KGX object. """ + if path is None: + # create output dir + if not os.path.exists(ROGER_DATA_DIR / 'merge'): + os.makedirs(ROGER_DATA_DIR / 'merge') + return str(ROGER_DATA_DIR / 'merge' / name) + if not os.path.exists(path): + os.makedirs(path) + + return str(path.joinpath(name)) + +def merged_objects(file_type, path=None): + """ A list of merged KGX objects. """ + if not path: + merged_pattern = merge_path(f"**/{file_type}.jsonl") + else: + merged_pattern = merge_path(f"**/{file_type}.jsonl", path=path) + # this thing should always return one edges or nodes file (based on file_type) + try: + return sorted(glob.glob(merged_pattern, recursive=True))[0] + except IndexError: + raise ValueError(f"Could not find merged KGX of type {file_type} in {merged_pattern}") + + +def schema_path(name, path=None): + """ Path to a schema object. + :param name: Name of the object to get a path for. """ + if not path: + return str(ROGER_DATA_DIR / 'schema' / name) + return str (path / 'schema' / name) + +def bulk_path(name, path=None): + """ Path to a bulk load object. + :param name: Name of the object. """ + if not path: + return str(ROGER_DATA_DIR / 'bulk' / name) + else: + return str(path / name) + +def metrics_path(name): + """ + Path to write metrics to + :param name: + :return: + """ + return str(ROGER_DATA_DIR / "metrics" / name) + +def dug_kgx_path(name): + return str(ROGER_DATA_DIR / "dug" / "kgx" / name) + +def dug_annotation_path(name): + return str(ROGER_DATA_DIR / "dug" / "annotations" / name) + +def dug_expanded_concepts_path(name): + return str(ROGER_DATA_DIR / 'dug' / 'expanded_concepts' / name) + +def dug_expanded_concept_objects(data_path=None, format="pickle"): + "Return a list of files containing expaneded concept objects" + if data_path: + file_pattern = os.path.join(data_path, '**', f'expanded_concepts.{format}') + else: + file_pattern = dug_expanded_concepts_path( + os.path.join('*',f'expanded_concepts.{format}')) + return sorted(glob.glob(file_pattern, recursive=True)) + +def dug_expanded_elements_objects(data_path=None, format="txt"): + "Return a list of element files from the expanded concepts directory" + if data_path: + file_pattern = os.path.join(data_path, '**', f'elements.{format}') + else: + file_pattern = dug_expanded_concepts_path( + os.path.join('*', f'elements.{format}')) + return sorted(glob.glob(file_pattern, recursive=True)) + +def dug_extracted_elements_objects(data_path=None, format="txt"): + if data_path: + file_pattern = os.path.join(data_path, '**', f'extracted_graph_elements.{format}') + else: + file_pattern = dug_expanded_concepts_path( + os.path.join('*', f'extracted_graph_elements.{format}')) + return sorted(glob.glob(file_pattern, recursive=True)) + +def dug_crawl_path(name): + return str(ROGER_DATA_DIR / 'dug' / 'crawl' / name) + +def dug_kgx_objects(): + """ A list of dug KGX objects. """ + dug_kgx_pattern = dug_kgx_path("**.json") + return sorted(glob.glob(dug_kgx_pattern)) + +def dug_concepts_objects(data_path, format="pickle"): + """ A list of dug annotation Objects. """ + if not data_path: + concepts_file_path = dug_annotation_path( + os.path.join('*',f'concepts.{format}')) + else: + concepts_file_path = os.path.join( + data_path, '**', f'concepts.{format}') + return sorted(glob.glob(concepts_file_path, recursive=True)) + +def dug_elements_objects(data_path=None, format='pickle'): + """ A list of dug annotation Objects. """ + if not data_path: + concepts_file_pattern = dug_annotation_path( + os.path.join('*', f'elements.{format}')) + else: + concepts_file_pattern = os.path.join( + data_path, '**', f'elements.{format}') + return sorted(glob.glob(concepts_file_pattern, recursive=True)) + +def dug_input_files_path(name) -> pathlib.Path: + path = ROGER_DATA_DIR / "dug" / "input_files" / name + if not path.exists(): + log.info(f"Input file path: {path} does not exist, creating") + path.mkdir(parents=True, exist_ok=True) + else: + log.info(f"Input file path: {path} already exists") + return path + +def dug_topmed_objects(input_data_path=None): + "Return list of TOPMed source files" + if not input_data_path: + input_data_path = str(dug_input_files_path('topmed')) + topmed_file_pattern = os.path.join(input_data_path, "topmed_*.csv") + return sorted(glob.glob(topmed_file_pattern)) + +def dug_anvil_path(): + """Anvil source files""" + return dug_input_files_path('anvil') + +def dug_sprint_path(): + """Anvil source files""" + return dug_input_files_path('sprint') + +def dug_bacpac_path(): + """Anvil source files""" + return dug_input_files_path('bacpac') + +def dug_heal_mds_path(): + """HEAL MDS source files""" + return dug_input_files_path('heal-mds-imports') + +def dug_heal_research_program_path(): + """HEAL research programs source files""" + return dug_input_files_path('heal-research-programs') + +def dug_heal_study_path(): + """HEAL study source files""" + return dug_input_files_path('heal-study-imports') + +def dug_crdc_path(): + """Anvil source files""" + return dug_input_files_path('crdc') + +def dug_kfdrc_path(): + """Anvil source files""" + return dug_input_files_path('kfdrc') + +def dug_nida_objects(input_data_path=None): + "Return list of NIDA source files" + if not input_data_path: + input_data_path = str(dug_input_files_path('nida')) + nida_file_pattern = os.path.join(input_data_path, "NIDA-*.xml") + return sorted(glob.glob(nida_file_pattern)) + +def dug_sparc_objects(input_data_path=None): + if not input_data_path: + input_data_path = str(dug_input_files_path('sparc')) + file_pattern = os.path.join(input_data_path, "scicrunch/*.xml") + return sorted(glob.glob(file_pattern)) + +def dug_anvil_objects(input_data_path=None): + if not input_data_path: + input_data_path = dug_anvil_path() + files = get_files_recursive( + lambda file_name: ( + not file_name.startswith('GapExchange_') + and file_name.endswith('.xml')), + input_data_path) + return sorted([str(f) for f in files]) + +def dug_sprint_objects(input_data_path=None): + if not input_data_path: + input_data_path = dug_sprint_path() + files = get_files_recursive( + lambda file_name: file_name.endswith('.xml'), input_data_path) + return sorted([str(f) for f in files]) + +def dug_bacpac_objects(input_data_path=None): + "Return list of BACPAC source files" + if not input_data_path: + input_data_path = dug_bacpac_path() + files = get_files_recursive( + lambda file_name: file_name.endswith('.xml'), input_data_path) + return sorted([str(f) for f in files]) + +def dug_crdc_objects(input_data_path=None): + if not input_data_path: + input_data_path = dug_crdc_path() + files = get_files_recursive( + lambda file_name: ( + not file_name.startswith('GapExchange_') + and file_name.endswith('.xml')), + input_data_path) + return sorted([str(f) for f in files]) + +def dug_heal_study_objects(input_data_path=None): + "Return list of HEAL study source files" + if not input_data_path: + input_data_path = dug_heal_study_path() + files = get_files_recursive(lambda file_name : file_name.endswith('.xml'), + input_data_path) + return sorted([str(f) for f in files]) + +def dug_heal_research_program_objects(input_data_path=None): + "Return list of HEAL research program source files" + if not input_data_path: + input_data_path = dug_heal_research_program_path() + files = get_files_recursive(lambda file_name : file_name.endswith('.xml'), + input_data_path) + return sorted([str(f) for f in files]) + +def dug_kfdrc_objects(input_data_path=None): + if not input_data_path: + input_data_path = dug_kfdrc_path() + files = get_files_recursive( + lambda file_name: ( + not file_name.startswith('GapExchange_') + and file_name.endswith('.xml')), + input_data_path) + return sorted([str(f) for f in files]) + + +def dug_dd_xml_path(): + """ Topmed source files""" + return dug_input_files_path('db_gap') + +def get_files_recursive(file_name_filter, current_dir): + file_paths = [] + for child in current_dir.iterdir(): + if child.is_dir(): + file_paths += get_files_recursive(file_name_filter, child) + continue + if not file_name_filter(child.name): + continue + else: + file_paths += [child] + return file_paths + +def dug_dd_xml_objects(input_data_path=None): + if not input_data_path: + input_data_path = dug_dd_xml_path() + files = get_files_recursive( + lambda file_name: ( + not file_name.startswith('._') + and file_name.endswith('.xml')), + input_data_path) + return sorted([str(f) for f in files]) + +def copy_file_to_dir(file_location, dir_name): + return shutil.copy(file_location, dir_name) + +def read_schema (schema_type: SchemaType, path=None): + """ Read a schema object. + :param schema_type: Schema type of the object to read. """ + if path is not None: + path = path / '**' + location = glob.glob(schema_path (f"{schema_type.value}-schema.json", path=path), recursive=True)[0] + return read_object (location) + +def get_uri (path, key): + """ Build a URI. + :param path: The path of an object. + :param key: The key of a configuration value to prepend to the object. """ + # Incase config has http://..../ or http://... remove / and add back to + # avoid double http://...// + root_url = config[key].rstrip('/') + return f"{root_url}/{path}" + +def get_relative_path (path): + return os.path.join (os.path.dirname (__file__), path) + +def read_relative_object (path): + return read_object (get_relative_path(path)) + +def trunc(text, limit): + return ('..' + text[-limit-2:]) if len(text) > limit else text + + + +def json_line_iter(jsonl_file_path): + f = open(file=jsonl_file_path, mode='r', encoding='utf-8') + for line in f: + yield json.loads(line) + f.close() + +def jsonl_iter(file_name): + # iterating over jsonl files + with open(file_name) as stream: + for line in stream: + # yield on line at time + yield json.loads(line) + +def json_iter(json_file,entity_key): + with open(json_file) as stream: + data = json.loads(stream.read()) + return data[entity_key] + +def downloadfile(thread_num, inputq, doneq): + url = "" + t0 = 0 + pct = 0 + + def downloadprogress(blocknumber, readsize, totalfilesize): + nonlocal thread_num + nonlocal url, t0, pct + blocks_expected = ( + int(totalfilesize/readsize) + + (1 if totalfilesize%readsize != 0 else 0)) + t1 = int(current_time_in_millis()/1000) + elapsed_delta = t1 - t0 + pct = int(100 * blocknumber / blocks_expected) + if elapsed_delta >= 30: # every n seconds + log.info(f"thread-{thread_num} {pct}% of size:{totalfilesize} " + f"({blocknumber}/{blocks_expected}) url:{url}") + t0 = t1 + + num_files_processed = 0 + while inputq.empty() is False: + t0 = int(current_time_in_millis()/1000) + url, dst = inputq.get() + num_files_processed += 1 + log.info(f"thread-{thread_num} downloading {url}") + try: + path, httpMessage = urlretrieve( + url, dst, reporthook=downloadprogress) + if pct < 100: + httpMessageKeys = httpMessage.keys() + log.info(f"thread-{thread_num} urlretrieve path:'{path}' " + f"http-keys:{httpMessageKeys} " + f"httpMessage:'{httpMessage.as_string()}") + except Exception as e: + log.error(f"thread-{thread_num} downloadfile excepton: {e}") + continue + log.info(f"thread-{thread_num} downloaded {dst}") + doneq.put((thread_num,num_files_processed)) + log.info(f"thread-{thread_num} done!") + return diff --git a/src/roger/logger.py b/src/roger/logger.py new file mode 100644 index 00000000..fdf82274 --- /dev/null +++ b/src/roger/logger.py @@ -0,0 +1,96 @@ +import logging + +import sys + +from typing import Optional + +from roger.config import get_default_config + + + +logger: Optional[logging.Logger] = None + +# HTTP plumbing that logs per-request at DEBUG. requests_cache alone emitted +# ~84% of the lines in an annotate task (5 per HTTP call: cache directives, +# pre-read, post-read, pre-write, skipping-write), which is how a single +# long-running task produced a 2.9 GB log file -- big enough to evict the +# api-server, whose ephemeral-storage limit is 750Mi, when someone opened it +# in the UI. Silenced independently of roger's own level so that +# logging.level: DEBUG stays usable for debugging roger. +NOISY_LOGGERS = ( + 'requests_cache', + 'httpcore', + 'httpx', + 'urllib3.connectionpool', + 'elastic_transport', +) +NOISY_LOGGER_LEVEL = logging.WARNING + + +def quiet_noisy_loggers(level: int = NOISY_LOGGER_LEVEL) -> None: + """Cap the per-request DEBUG chatter from HTTP libraries. + + Children are set explicitly rather than left to inherit: a child that + already has a level of its own ignores the parent, and these libraries + log from submodules (requests_cache.policy.actions, httpcore.http11). + """ + existing = list(logging.root.manager.loggerDict) + for prefix in NOISY_LOGGERS: + logging.getLogger(prefix).setLevel(level) + for name in existing: + if name.startswith(prefix + '.'): + logging.getLogger(name).setLevel(level) + + +def get_logger(name: str = 'roger') -> logging.Logger: + + """ + + Get an instance of logger. + + + + Parameters + + ---------- + + name: str + + The name of logger + + + + Returns + + ------- + + logging.Logger + + An instance of logging.Logger + + + + """ + + global logger + + if logger is None: + + config = get_default_config() + + logger = logging.getLogger(name) + + handler = logging.StreamHandler(sys.stdout) + + formatter = logging.Formatter(config['logging']['format']) + + handler.setFormatter(formatter) + + logger.addHandler(handler) + + logger.setLevel(config['logging']['level']) + + logger.propagate = True + + return logger + diff --git a/src/roger/models/__init__.py b/src/roger/models/__init__.py new file mode 100644 index 00000000..3a994ff8 --- /dev/null +++ b/src/roger/models/__init__.py @@ -0,0 +1,4 @@ +"Data models for Roger" + +from roger.models.kgx import KGXModel +from roger.models.biolink import BiolinkModel diff --git a/src/roger/models/biolink.py b/src/roger/models/biolink.py new file mode 100644 index 00000000..43102188 --- /dev/null +++ b/src/roger/models/biolink.py @@ -0,0 +1,48 @@ +"Biolink data model for Roger" + +from bmt import Toolkit +from roger.logger import get_logger + +log = get_logger() + +class BiolinkModel: + "Biolink data model for Roger" + root_type = 'biolink:NamedThing' + + def __init__(self, bl_version='v3.1.2'): + self.bl_url = (f'https://raw.githubusercontent.com/biolink' + f'/biolink-model/{bl_version}/biolink-model.yaml') + log.info("bl_url is %s", self.bl_url) + self.toolkit = Toolkit() + + def find_biolink_leaves(self, biolink_concepts): + """Given list of concepts, returns leaves minus any parent concepts + :param biolink_concepts: list of biolink concepts + :return: leave concepts. + """ + ancestry_set = set() + all_concepts = set(biolink_concepts) + unknown_elements = set() + + for x in all_concepts: + current_element = self.toolkit.get_element(x) + if not current_element: + unknown_elements.add(x) + ancestors = set(self.toolkit.get_ancestors( + x, mixin=True, reflexive=False, formatted=True)) + ancestry_set = ancestry_set.union(ancestors) + leaf_set = all_concepts - ancestry_set - unknown_elements + return leaf_set + + def get_leaf_class (self, names): + """ Return the leaf classes in the provided list of names. """ + leaves = list(self.find_biolink_leaves(names)) + return leaves[0] + + def get_label(self, class_name): + "Return the label for the given class name" + element = self.toolkit.get_element(class_name) + if element: + name = element.name + return name + return class_name.replace("biolink:", "").replace("_", " ") diff --git a/src/roger/models/kgx.py b/src/roger/models/kgx.py new file mode 100644 index 00000000..e35ab07e --- /dev/null +++ b/src/roger/models/kgx.py @@ -0,0 +1,503 @@ +"KGX data model for Roger" + +import os +import time +import queue +from itertools import chain +import threading +from collections import defaultdict +from xxhash import xxh64_hexdigest +import orjson as json +import ntpath +from kg_utils.merging import DiskGraphMerger +from kg_utils.constants import * + +from roger.config import get_default_config +from roger.logger import get_logger +from roger.components.data_conversion import compare_types +from roger.core import storage +from roger.models.biolink import BiolinkModel +from roger.core.enums import SchemaType + +log = get_logger() + +class KGXModel: + """ Abstractions for transforming KGX formatted data. + + KGX stands for Knowledge Graph Exchange + """ + def __init__(self, biolink=None, config=None): + if not config: + config = get_default_config() + self.config = config + + # We need a temp director for the DiskGraphMerger + self.temp_directory = storage.merge_path( + self.config.kgx.merge_db_temp_dir) + log.debug(f"Setting temp_directory to : {self.temp_directory}") + isExist = os.path.exists(self.temp_directory) + if not isExist: + os.makedirs(self.temp_directory) + + self.merger = DiskGraphMerger(temp_directory=self.temp_directory, + chunk_size=5_000_000) + self.biolink_version = self.config.kgx.biolink_model_version + log.debug(f"Trying to get biolink version : {self.biolink_version}") + if biolink is None: + self.biolink = BiolinkModel(self.biolink_version) + else: + self.biolink = biolink + self.enable_metrics = self.config.get('enable_metrics', False) + + def get_kgx_json_format(self, files: list, dataset_version: str): + """Gets Json formatted kgx files. + + These files have a the following structure: + {"nodes": [{"id":"..."},...], "edges": [{"id":...},...}] } + + Parameters + ---------- + files : list of file names + dataset_version : dataset version from dataset meta-data information + + Returns None + ------- + + """ + file_tuple_q = queue.Queue() + thread_done_q = queue.Queue() + for nfile, file_name in enumerate(files): + # file_url or skip + file_name = dataset_version + "/" + file_name + file_url = storage.get_uri(file_name, "kgx_base_data_uri") + subgraph_basename = os.path.basename(file_name) + subgraph_path = storage.kgx_path(subgraph_basename) + if os.path.exists(subgraph_path): + log.info(f"cached kgx: {subgraph_path}") + continue + log.debug("#{}/{} to get: {}".format( + nfile+1, len(files), file_url)) + # folder + dirname = os.path.dirname (subgraph_path) + if not os.path.exists (dirname): + os.makedirs (dirname, exist_ok=True) + # add to queue + file_tuple_q.put((file_url,subgraph_path)) + + # start threads for each file download + threads = [] + for thread_num in range(len(files)): # len(files) + th = threading.Thread( + target=storage.downloadfile, + args=(thread_num, file_tuple_q, thread_done_q)) + th.start() + threads.append(th) + + # wait for each thread to complete + for nwait in range(len(threads)): + thread_num, num_files_processed = thread_done_q.get() + th = threads[thread_num] + th.join() + log.info(f"#{nwait+1}/{len(threads)} joined: " + f"thread-{thread_num} processed: " + f"{num_files_processed} file(s)") + + all_kgx_files = [] + for nfile, file_name in enumerate(files): + start = storage.current_time_in_millis() + file_name = dataset_version + "/" + file_name + file_url = storage.get_uri(file_name, "kgx_base_data_uri") + subgraph_basename = os.path.basename(file_name) + subgraph_path = storage.kgx_path(subgraph_basename) + all_kgx_files.append(subgraph_path) + if os.path.exists(subgraph_path): + log.info(f"cached kgx: {subgraph_path}") + continue + log.info ("#{}/{} read: {}".format(nfile+1, len(files), file_url)) + subgraph = storage.read_object(file_url) + storage.write_object(subgraph, subgraph_path) + total_time = storage.current_time_in_millis() - start + edges = len(subgraph['edges']) + nodes = len(subgraph['nodes']) + log.info( + "#{}/{} edges:{:>7} nodes: {:>7} time:{:>8} wrote: {}".format( + nfile+1, len(files), edges, nodes, + total_time/1000, subgraph_path)) + return all_kgx_files + + def get_kgx_jsonl_format(self, files, dataset_version): + """gets pairs of jsonl formatted kgx files. + + Files is expected to have all the pairs. + + I.e if kgx_1_nodes.jsonl exists its expected that kgx_1_edges.jsonl + exists in the same path. + File names should have strings *nodes*.jsonl and *edges*.jsonl. + Parameters + ---------- + files + dataset_version + + Returns + ------- + + """ + # make a paired list + paired_up = [] + log.info(f"getting {files}") + for file_name in files: + if "nodes" in file_name: + paired_up.append( + [file_name, file_name.replace('nodes', 'edges')]) + error = False + # validate that all pairs exist + if len(files) / 2 != len(paired_up): + log.error("Error paired up kgx jsonl files don't match " + "list of files specified in metadata.yaml") + error = True + for pairs in paired_up: + if pairs[0] not in files: + log.error( + f"{pairs[0]} not in original list " + f"of files from metadata.yaml") + error = True + if pairs[1] not in files: + error = True + log.error( + f"{pairs[1]} not in original list " + f"of files from metadata.yaml") + if error: + raise Exception("Metadata.yaml has inconsistent jsonl files") + + file_tuple_q = queue.Queue() + thread_done_q = queue.Queue() + for npairs, pairs in enumerate(paired_up): + for npair, p in enumerate(pairs): + file_name = dataset_version + "/" + p + file_url = storage.get_uri(file_name, "kgx_base_data_uri") + subgraph_basename = os.path.basename(file_name) + subgraph_path = storage.kgx_path(subgraph_basename) + if os.path.exists(subgraph_path): + log.info(f"skip cached kgx: {subgraph_path}") + continue + log.info ("#{}.{}/{} read: {}".format( + npairs+1, npair+1, len(paired_up), file_url)) + # folder + dirname = os.path.dirname (subgraph_path) + if not os.path.exists (dirname): + os.makedirs (dirname, exist_ok=True) + # add to queue + file_tuple_q.put((file_url,subgraph_path)) + + # start threads for each file download + threads = [] + for thread_num in range(file_tuple_q.qsize()): + th = threading.Thread( + target=storage.downloadfile, + args=(thread_num, file_tuple_q, thread_done_q)) + th.start() + threads.append(th) + + # wait for each thread to complete + for nwait in range(len(threads)): + thread_num, num_files_processed = thread_done_q.get() + th = threads[thread_num] + th.join() + log.info(f"#{nwait+1}/{len(threads)} joined: " + f"thread-{thread_num} processed: " + f"{num_files_processed} file(s)") + + all_kgx_files = [] + for pairs in paired_up: + nodes = 0 + edges = 0 + start = storage.current_time_in_millis() + for p in pairs: + file_name = dataset_version + "/" + p + file_url = storage.get_uri(file_name, "kgx_base_data_uri") + subgraph_basename = os.path.basename(file_name) + subgraph_path = storage.kgx_path(subgraph_basename) + all_kgx_files.append(subgraph_path) + if os.path.exists(subgraph_path): + log.info(f"cached kgx: {subgraph_path}") + continue + data = storage.read_object(file_url) + storage.write_object(data, subgraph_path) + if "edges" in p: + edges = len(data.split('\n')) + else: + nodes = len(data.split('\n')) + total_time = storage.current_time_in_millis() - start + log.info( + "wrote {:>45}: edges:{:>7} nodes: {:>7} time:{:>8}".format( + storage.trunc(subgraph_path, 45), edges, nodes, total_time)) + return all_kgx_files + + def get (self, dataset_version = "v1.0"): + """ Read metadata for KGX files and downloads them locally. + :param dataset_version: Data version to operate on. + """ + metadata = storage.read_relative_object ("../../metadata.yaml") + data_set_list = self.config.kgx.data_sets + kgx_files_remote = [] + for item in metadata['kgx']['versions']: + if (item['version'] == dataset_version and + item['name'] in data_set_list): + log.info(f"Getting KGX dataset {item['name']}, " + f"version {item['version']}") + if item['format'] == 'json': + kgx_files_remote += self.get_kgx_json_format( + item['files'], item['version']) + elif item['format'] == 'jsonl': + kgx_files_remote += self.get_kgx_jsonl_format( + item['files'], item['version']) + else: + raise ValueError( + f"Unrecognized format in metadata.yaml: " + f"{item['format']}, valid formats are `json` " + f"and `jsonl`.") + # Fetchs kgx generated from Dug Annotation workflow. + new_files = self.fetch_dug_kgx() + kgx_files_remote + all_files_in_dir = ( + storage.kgx_objects("json") + + storage.kgx_objects("jsonl")) + files_to_remove = [x for x in all_files_in_dir + if x not in new_files] + if len(files_to_remove): + log.info( + "Found some old files to remove from kgx dir : %s", + files_to_remove) + for file in files_to_remove: + storage.remove(file) + log.info("removed %s", file) + log.info("Done.") + + + + def fetch_dug_kgx(self): + """ + Copies files from dug output dir to roger kgx dir. + :return: + """ + dug_kgx_files = storage.dug_kgx_objects() + all_kgx_files = [] + log.info("Copying dug KGX files to %s. Found %d kgx files to copy.", + storage.kgx_path(''), len(dug_kgx_files)) + for file in dug_kgx_files: + file_name = ntpath.basename(file) + dest = storage.kgx_path(file_name) + all_kgx_files.append(dest) + storage.write_object({}, dest) + log.info(f"Copying from {file} to {dest}.") + storage.copy_file_to_dir(file, dest) + log.info("Done copying dug KGX files.") + return all_kgx_files + + def create_nodes_schema(self, input_data_path=None, output_data_path=None): + """ + Extracts schema for nodes based on biolink leaf types + :return: + """ + + category_schemas = defaultdict(lambda: None) + category_error_nodes = set() + merged_nodes_file = storage.merged_objects("nodes", input_data_path) + log.info(f"Processing : {merged_nodes_file}") + counter = 0 + for node in storage.json_line_iter(merged_nodes_file): + # Debuging code + if counter % 10000 == 0: + log.info(f"Processing node : {node} counter : {counter}") + counter += 1 + + if not node.get('category'): + category_error_nodes.add(node['id']) + node['category'] = [BiolinkModel.root_type] + + # Get all leaf types of this node + node_types = list( + self.biolink.find_biolink_leaves(node['category'])) + # pick the fist one to work on + node_type = node_types[0] + + + # make sure it is defined in the final dict + category_schemas[node_type] = category_schemas.get(node_type, {}) + + # compute full list of attributes and the value types of the + # attributes for that type. + for k in node.keys(): + current_type = type(node[k]).__name__ + if k not in category_schemas[node_type]: + category_schemas[node_type][k] = current_type + else: + previous_type = category_schemas[node_type][k] + category_schemas[node_type][k] = compare_types( + previous_type, current_type) + + # copy over final result to every other leaf type + for tp in node_types: + category_schemas[tp] = category_schemas[node_type] + + + if len(category_error_nodes): + log.warning(f"some nodes didn't have category assigned. " + f"KGX file has errors." + f"Nodes {len(category_error_nodes)}." + f"Showing first 10: {list(category_error_nodes)[:10]}." + f"These will be treated as {BiolinkModel.root_type}.") + + # Write node schemas. + self.write_schema(category_schemas, SchemaType.CATEGORY, output_path=output_data_path) + + def create_edges_schema(self, input_data_path=None, output_data_path=None): + """ + Create unified schema for all edges in an edges jsonl file. + :return: + """ + predicate_schemas = defaultdict(lambda: None) + merged_edges_file = storage.merged_objects("edges", input_data_path) + """ Infer predicate schemas. """ + for edge in storage.json_line_iter(merged_edges_file): + predicate = edge['predicate'] + predicate_schemas[predicate] = predicate_schemas.get(predicate, + {}) + for k in edge.keys(): + current_type = type(edge[k]).__name__ + if k not in predicate_schemas[predicate]: + predicate_schemas[predicate][k] = current_type + else: + previous_type = predicate_schemas[predicate][k] + predicate_schemas[predicate][k] = compare_types( + previous_type, current_type) + self.write_schema(predicate_schemas, SchemaType.PREDICATE, output_path=output_data_path) + + def create_schema (self): + """Determine the schema of each type of object. + + We have to do this to make it possible to write tabular data. Need to + know all possible columns in advance and correct missing fields. + """ + if self.schema_up_to_date(): + log.info (f"schema is up to date.") + return + + self.create_nodes_schema() + self.create_edges_schema() + + def schema_up_to_date (self): + return storage.is_up_to_date ( + source=storage.kgx_objects(), + targets=[ + storage.schema_path ( + f"{SchemaType.PREDICATE.value}-schema.json"), + storage.schema_path ( + f"{SchemaType.PREDICATE.value}-schema.json") + ]) + + def write_schema(self, schema, schema_type: SchemaType ,output_path=None): + """ Output the schema file. + + :param schema: Schema to get keys from. + :param schema_type: Type of schema to write. + """ + file_name = storage.schema_path (f"{schema_type.value}-schema.json", output_path) + log.info("writing schema: %s", file_name) + dictionary = { k : v for k, v in schema.items () } + storage.write_object (dictionary, file_name) + + def merge(self, input_path=None, output_path=None): + """ This version uses the disk merging from the kg_utils module """ + + metrics = {} + start = time.time() + + log.info(f"Input path = {input_path}, Output path = {output_path}") + + if input_path: + json_format_files = storage.kgx_objects("json", input_path) + jsonl_format_files = storage.kgx_objects("jsonl", input_path) + else: + json_format_files = storage.kgx_objects("json") + jsonl_format_files = storage.kgx_objects("jsonl") + + # Create lists of the nodes and edges files in both json and jsonl + # formats + jsonl_node_files = {file for file in jsonl_format_files + if "node" in file.split('/')[-1]} + jsonl_edge_files = {file for file in jsonl_format_files + if "edge" in file.split('/')[-1]} + log.info(f"Jsonl edge files : {jsonl_edge_files}") + log.info(f"Jsonl node files : {jsonl_node_files}") + + # Create all the needed iterators and sets thereof + jsonl_node_iterators = [storage.jsonl_iter(file_name) + for file_name in jsonl_node_files] + jsonl_edge_iterators = [storage.jsonl_iter(file_name) + for file_name in jsonl_edge_files] + json_node_iterators = [storage.json_iter(file_name, 'nodes') + for file_name in json_format_files] + json_edge_iterators = [storage.json_iter(file_name, 'edges') + for file_name in json_format_files] + all_node_iterators = json_node_iterators + jsonl_node_iterators + all_edge_iterators = json_edge_iterators + jsonl_edge_iterators + + # chain the iterators together + node_iterators = chain(*all_node_iterators) + edge_iterators = chain(*all_edge_iterators) + + # now do the merge + self.merger.merge_nodes(node_iterators) + merged_nodes = self.merger.get_merged_nodes_jsonl() + + + self.merger.merge_edges(edge_iterators) + merged_edges = self.merger.get_merged_edges_jsonl() + + write_merge_metric = {} + t = time.time() + start_nodes_jsonl = time.time() + + + nodes_file_path = storage.merge_path("nodes.jsonl", output_path) + + # stream out nodes to nodes.jsonl file + with open(nodes_file_path, 'w') as stream: + for nodes in merged_nodes: + stream.write(nodes) + + time_difference = time.time() - start_nodes_jsonl + log.info("writing nodes took : %s", str(time_difference)) + write_merge_metric['nodes_writing_time'] = time_difference + start_edge_jsonl = time.time() + + # stream out edges to edges.jsonl file + edges_file_path = storage.merge_path("edges.jsonl", output_path) + with open(edges_file_path, 'w') as stream: + for edges in merged_edges: + edges = json.loads(edges) + # Add an id field for the edges as some of the downstream + # processing expects it. + edges['id'] = xxh64_hexdigest( + edges['subject'] + edges['predicate'] + + edges['object'] + + edges.get("biolink:primary_knowledge_source", "")) + keys_to_del = set() + for key in edges: + if key.startswith('biolink:'): + keys_to_del.add(key) + for k in keys_to_del: + edges[k.replace('biolink:', '')] = edges[k] + del edges[k] + stream.write(json.dumps(edges).decode('utf-8') + '\n') + + write_merge_metric['edges_writing_time'] = time.time() - start_edge_jsonl + log.info(f"writing edges took: {time.time() - start_edge_jsonl}") + write_merge_metric['total_time'] = time.time() - t + metrics['write_jsonl'] = write_merge_metric + metrics['total_time'] = time.time() - start + log.info(f"total took: {time.time() - start}") + if self.enable_metrics: + metricsfile_path = storage.metrics_path('merge_metrics.yaml') + storage.write_object(metrics, metricsfile_path) + diff --git a/src/roger/pipelines/README.md b/src/roger/pipelines/README.md new file mode 100644 index 00000000..e77e6a29 --- /dev/null +++ b/src/roger/pipelines/README.md @@ -0,0 +1,99 @@ +# Building custom Dug data pipelines + +The pipelines submodule is where data pipelines can be defined for specific data +sets with specific, custom behaviors for each one. In previous versions of the +code, customizations for each pipeline were spread across several modules. With +this instantiation, the customizations for each data set pipeline are +consolidated into a single overridden subclass of the DataPipeline class. + +## What the base pipeline does + +The function `roger.tasks.create_pipeline_taskgroup`, when called with the given +data pipeline class, will emit an Airflow task group with the following +structure. If Airflow is not being used, another executor should use a similarly +structured set of calls and dependencies to ensure that the task pipeline +executes fully and in order. + +```mermaid +graph TD; + annotate-->index_variables; + annotate-->validate_index_variables; + index_variables-->validate_index_variables; + annotate-->make_kg_tagged; + annotate-->crawl_tranql; + annotate-->index_concepts; + crawl_tranql-->validate_index_concepts; + index_concepts-->validate_index_concepts; + annotate-->validate_index_concepts; +``` +The pipeline steps are briefly described below + +### annotate + +By default, `annotate` will call the `get_objects` method to collect a list of +parsable files. For each of these files, a Dug Crawler object will be created +which will apply the parser returned by the pipeline class's `get_parser_name` +method. (This by default will return `parser_name` if it's defined, or will fall +back to `pipeline_name`.) The results will be written to `elements.json` and +`concepts.json` as appropriate. + +### index_variables + +This will load the `elements.json` files from `annotate` and pass them to the +indexer built from a DugFactory object. (This is sending them to ElasticSearch +for indexing under the hood.) + +### make_kg_tagged + +All `elements.json` files will be loaded, and based on the annotations, a +Translator-compliant knowledge graph will be written to a `_kgx.json` file. + +### index_concepts + +The `concepts.json` files are read and submitted to ElasticSearch using the +indexer object derived from the embedded DugFactory object. + +### validate_index_concepts + +Concepts from `concepts.json` are double-checked to ensure that the ES indexing +process actually worked. + +## Defining a basic pipeline, with no customizations + +Simple pipelines, such as that for the BACPAC dataset, need very little +customization. All pipelines must define a `pipeline_name`, which will be used +as the default value for a number of other parameters if they are not +defined. In the case of BACPAC, a difference in case means that both the +`pipeline_name` and the `parser_name` need to be defined. + +```python +from roger.pipelines import DugPipeline + +class BacPacPipeline(DugPipeline): + "Pipeline for BACPAC data set" + pipeline_name = "bacpac" + parser_name = "BACPAC" +``` + +This is the full extent of the code needed to adapt the DugPipeline object to +BACPAC. Other data sets have more specific customizations that need more custom +code or variables defined. + +## More extensive customization + +Because the base pipeline (defined in `roger/pipelines/base.py:DugPipeline`) is +inherited as a subclass for customizing, effectively any part of the pipeline +that isn't part of Dug proper can be overriden. Here are some common +customizations that are expected to be necessary for many parts of the process: + +### get_objects + +The `get_objects` method by default looks in the `input_data_path` that is +passed to it, and if that is None, loads the default from the `ROGER_DATA_DIR` +environment variable. By default, it reads all files with the `.xml` extension +recursively anywhere in that directory or its subdirectories. + +One example customization is the anvil data pipeline, which additionally +excludes any file that starts with 'GapExchange_'. Any overriden method should +accept an optional `input_data_path` parameter and return a list of files, +sorted in the order that they should be processed. diff --git a/src/roger/pipelines/__init__.py b/src/roger/pipelines/__init__.py new file mode 100644 index 00000000..ff75b18d --- /dev/null +++ b/src/roger/pipelines/__init__.py @@ -0,0 +1,41 @@ +"Modules for individual datasets" + +import pkgutil +from pathlib import Path +import importlib + +from .base import DugPipeline, DDM2Pipeline + +def get_all_subclasses(cls): + """Recurse to get all subsubclasses, etc.""" + rval = [cls] + for sc in cls.__subclasses__(): + rval.extend(get_all_subclasses(sc)) + return rval + +def get_pipeline_classes(pipeline_names): + """Return a list of all defined pipeline classes + """ + + base_path = Path(__file__).resolve().parent + + for (_, mod_name, _) in pkgutil.iter_modules([base_path]): + if mod_name == 'base': + continue + + # No need to actuall get the module symbol, once it's imported, it will + # show up below in __subclasses__. + importlib.import_module(f"{__name__}.{mod_name}") + pipeline_list = [] + + for subclass in get_all_subclasses(DugPipeline): + if (getattr(subclass, 'pipeline_name') and + getattr(subclass, 'pipeline_name') in pipeline_names): + try: + subclass.input_version = pipeline_names[ + getattr(subclass, 'pipeline_name')] + except TypeError: + # If someone passed in the list, don't bother with the verison. + pass + pipeline_list.append(subclass) + return pipeline_list diff --git a/src/roger/pipelines/anvil.py b/src/roger/pipelines/anvil.py new file mode 100644 index 00000000..ec6d06f4 --- /dev/null +++ b/src/roger/pipelines/anvil.py @@ -0,0 +1,25 @@ +"Pipeline for anvil data" + +from roger.pipelines import DugPipeline +from roger.core import storage + +class AnvilPipeline(DugPipeline): + "Pipeline for Anvil data set" + pipeline_name = 'anvil' + parser_name = 'Anvil' + files_dir = 'anvil' + + def get_objects(self, input_data_path=None): + """Retrieve anvil objects + + This code is imported from roger.core.storage.dug_anvil_objects + """ + if not input_data_path: + input_data_path = storage.dug_input_files_path( + self.files_dir) + files = storage.get_files_recursive( + lambda file_name: ( + not file_name.startswith('GapExchange_') + and file_name.endswith('.xml')), + input_data_path) + return sorted([str(f) for f in files]) diff --git a/src/roger/pipelines/bacpac.py b/src/roger/pipelines/bacpac.py new file mode 100644 index 00000000..495ba3b9 --- /dev/null +++ b/src/roger/pipelines/bacpac.py @@ -0,0 +1,8 @@ +"Pipeline for BACPAC data" + +from roger.pipelines import DugPipeline + +class BacPacPipeline(DugPipeline): + "Pipeline for BACPAC data set" + pipeline_name = "bacpac" + parser_name = "BACPAC" diff --git a/src/roger/pipelines/base.py b/src/roger/pipelines/base.py new file mode 100644 index 00000000..f3af15f0 --- /dev/null +++ b/src/roger/pipelines/base.py @@ -0,0 +1,1333 @@ +"Base class for implementing a dataset annotate, crawl, and index pipeline" + +import os +import random +import threading +import time +import asyncio +from concurrent.futures import ThreadPoolExecutor +from io import StringIO +import logging +import re +import hashlib +import traceback +from datetime import datetime, timezone +from functools import reduce +from pathlib import Path +import tarfile +from typing import Union +import jsonpickle +from dug_data_model.v2 import dedupe_and_sort + +import requests + +from dug.core import get_parser, get_annotator, get_plugin_manager, DugConcept, DugVariable, DugStudy, DugSection +from dug.core.concept_expander import ConceptExpander +from dug.core.crawler import Crawler +from dug.core.factory import DugFactory +from dug.core.parsers import Parser, DugElement +from dug.core.annotators import Annotator +from dug.core.annotators.sapbert_annotator import AnnotateSapbert +from dug.core.async_search import Search +from dug.core.index import Index + +from roger.config import RogerConfig +from roger.core import storage +from roger.models.biolink import BiolinkModel +from roger.logger import get_logger + +from roger.utils.http_utils import enable_post_caching, harden_session +from roger.utils.batched_annotator import BatchedAnnotator +from roger.utils.s3_utils import S3Utils + +log = get_logger() + +class PipelineException(Exception): + "Exception raised from DugPipeline and related classes" + +def make_edge(subj, + obj, + predicate='biolink:related_to', + predicate_label='related to', + relation='biolink:related_to', + relation_label='related to' + ): + """Create an edge between two nodes. + + :param subj: The identifier of the subject. + :param pred: The predicate linking the subject and object. + :param obj: The object of the relation. + :param predicate: Biolink compatible edge type. + :param predicate_label: Edge label. + :param relation: Ontological edge type. + :param relation_label: Ontological edge type label. + :returns: Returns and edge. + """ + edge_id = hashlib.md5( + f'{subj}{predicate}{obj}'.encode('utf-8')).hexdigest() + return { + "subject": subj, + "predicate": predicate, + "predicate_label": predicate_label, + "id": edge_id, + "relation": relation, + "relation_label": relation_label, + "object": obj, + "provided_by": "renci.bdc.semanticsearch.annotator" + } + +# The six list fields dug's Index.index_element unions when a document id is +# already present. It matters: the same CDE variable id appears in several +# element files (e.g. BRTHDTC in both adult- and pediatric-demographic) with +# different concepts and parents, and search must find it by any of them. +MERGED_LIST_FIELDS = ('search_terms', 'optional_terms', 'parents', + 'programs', 'identifiers') + + +def merge_searchable_docs(prior: dict, new: dict) -> dict: + """Union the list fields of two searchable dicts for the same id. + + Scalars (name, description, data_type, ...) are taken from `new`; dug's + per-doc update path left them at whatever was written first, which let + edits upstream go stale. + """ + merged = dict(new) + for field in MERGED_LIST_FIELDS: + merged[field] = dedupe_and_sort( + (prior.get(field) or []) + (new.get(field) or [])) + tags = (prior.get('tags') or []) + (new.get('tags') or []) + merged['tags'] = [dict(t) for t in + {tuple(sorted(d.items())) for d in tags}] + return merged + + +class FileFetcher: + """A basic remote file fetcher class + """ + + def __init__( + self, + remote_host: str, + remote_dir: Union[str, Path], + local_dir: Union[str, Path] = "." + ): + self.remote_host = remote_host + if isinstance(remote_dir, str): + self.remote_dir = remote_dir.rstrip("/") + else: + self.remote_dir = str(remote_dir.as_posix()) + self.local_dir = Path(local_dir).resolve() + + def __call__(self, remote_file_path: Union[str, Path]) -> Path: + remote_path = self.remote_dir + "/" + remote_file_path + local_path = self.local_dir / remote_file_path + url = f"{self.remote_host}{remote_path}" + log.debug("Fetching %s", url) + try: + response = requests.get(url, allow_redirects=True, timeout=60) + except Exception as e: + log.error("Unexpected %s: %s", e.__class__.__name__, str(e)) + raise RuntimeError(f"Unable to fetch {url}") from e + + log.debug("Response: %d", response.status_code) + if response.status_code != 200: + log.debug("Unable to fetch %s: %d", url, response.status_code) + raise RuntimeError(f"Unable to fetch {url}") + + with local_path.open('wb') as file_obj: + file_obj.write(response.content) + return local_path + +class DugPipeline(): + "Base class for dataset pipelines" + + pipeline_name = None + unzip_source = True + input_version = "" + + def __init__(self, config: RogerConfig, to_string=False): + "Set instance variables and check to make sure we're overriden" + if not self.pipeline_name: + raise PipelineException( + "Subclass must at least define pipeline_name as class var") + self.config = config + self.bl_toolkit = BiolinkModel() + dug_conf = config.to_dug_conf() + self.element_mapping = config.indexing.element_mapping + self.factory = DugFactory(dug_conf) + # dug builds this session with no timeout, so bound it here before + # handing it to the Crawler. This is the session every annotation + # call goes through. + self.annotation_conf = config.annotation + self.cached_session = self.build_annotation_session() + # one session and annotator per worker thread; requests.Session is + # not documented as thread safe and a silently corrupted annotation + # run costs weeks + self._thread_local = threading.local() + self.event_loop = asyncio.new_event_loop() + self.log_stream = StringIO() + if to_string: + self.string_handler = logging.StreamHandler(self.log_stream) + log.addHandler(self.string_handler) + self.s3_utils = S3Utils(self.config.s3_config) + + self.tranqlizer: ConceptExpander = self.factory.build_tranqlizer() + + graph_name = self.config["redisgraph"]["graph"] + source = f"redis:{graph_name}" + self.tranql_queries: dict = self.factory.build_tranql_queries(source) + self.node_to_element_queries: list = ( + self.factory.build_element_extraction_parameters(source)) + + indexing_config = config.indexing + self.variables_index = indexing_config.get('variables_index') + self.studies_index = indexing_config.get('studies_index') + self.concepts_index = indexing_config.get('concepts_index') + self.sections_index = indexing_config.get('sections_index') + self.kg_index = indexing_config.get('kg_index') + + self.search_obj = None + self.index_obj = None + + + def build_annotation_session(self): + """A hardened, POST-caching http session for annotation calls. + + Built per call rather than shared so each annotate worker thread can + hold its own. + """ + session = harden_session( + self.factory.build_http_session(), + connect_timeout=self.annotation_conf.http_connect_timeout, + read_timeout=self.annotation_conf.http_read_timeout, + retries=self.annotation_conf.http_retries, + backoff_factor=self.annotation_conf.http_retry_backoff, + ) + if self.annotation_conf.cache_post_requests: + enable_post_caching( + session, + expire_seconds=self.annotation_conf.http_cache_expire_seconds) + return session + + def thread_annotation_context(self): + """(session, annotator) owned by the calling thread. + + The redis-backed response cache is shared, so worker threads still + see each other's annotations; only the client objects are private. + """ + local = self._thread_local + if getattr(local, 'session', None) is None: + local.session = self.build_annotation_session() + local.annotator = self.init_annotator() + return local.session, local.annotator + + def __enter__(self): + self.event_loop = asyncio.new_event_loop() + return self + + def __exit__(self, exc_type, exc_val, exc_tb): + if self.search_obj: + # close elastic search connection + self.event_loop.run_until_complete(self.search_obj.es.close()) + # close async loop + if self.event_loop.is_running() and not self.event_loop.is_closed(): + self.event_loop.close() + if exc_type or exc_val or exc_tb: + traceback.print_exc() + log.error("%s %s %s", exc_val, exc_val, exc_tb) + log.exception("Got an exception") + + def get_data_format(self): + """Access method for data_format parameter + + Defaults to pipeline_name unless self.data_format is set. This method + can also be overriden + """ + return getattr(self, 'data_format', self.pipeline_name) + + def get_files_dir(self): + """Access method for files_dir parameter + + Defaults to pipeline_name unless self.files_dir is set. This method can + also be overriden. + """ + return getattr(self, 'files_dir', self.pipeline_name) + + def get_parser_name(self): + """Access method for parser_name + + Defaults to pipeline_name unless self.parser_name is set. This method + can also be overriden. + """ + return getattr(self, 'parser_name', self.pipeline_name) + + def get_annotator_name(self): + """ Access method for annotator_name + + Defaults to annotator_monarch unless specified using + annotation.annotator_type in the configuration file. + """ + return self.config.annotation.annotator_type + + def get_parser(self): + dug_plugin_manager = get_plugin_manager() + parser: Parser = get_parser(dug_plugin_manager.hook, + self.get_parser_name()) + return parser + + def get_annotator(self): + dug_plugin_manager = get_plugin_manager() + annotator: Annotator = get_annotator( + dug_plugin_manager.hook, + self.get_annotator_name(), + self.config.to_dug_conf() + ) + return annotator + + def clear_annotation_cached(self, to_string=False, output_data_path=None): + if not output_data_path: + output_data_path = storage.dug_annotation_path("") + storage.clear_dir(output_data_path) + # Clear http session cache + if self.config.annotation.clear_http_cache: + self.cached_session.cache.clear() + + def init_annotator(self, max_retries=5, base_delay=1, max_delay=10): + attempt = 0 + while attempt < max_retries: + try: + log.info("Initializing annotator") + annotator = self.get_annotator() + # Only sapbert: the batching wrapper drives + # text_classification/annotate_classifiers, which the + # monarch annotator does not have. + if (self.annotation_conf.batch_identifier_lookups + and isinstance(annotator, AnnotateSapbert)): + annotator = BatchedAnnotator(annotator) + log.info("Batching identifier lookups") + return annotator # success + except Exception as e: + attempt += 1 + if attempt == max_retries: + log.error("Max retries reached when creating annotator. " + "Failing with error: %s", e) + raise + delay = min(base_delay * (2 ** (attempt - 1)), max_delay) + delay += random.uniform(0, 1) # add jitter + log.warning("Error occurred: %s. Retrying in %.2f seconds...", + e, delay) + time.sleep(delay) + + def annotate_files(self, parsable_files, output_data_path=None): + """ + Annotates a Data element file using a Dug parser. + :param parsable_files: Files to parse. + :param output_data_path: Where to write elements.txt/concepts.txt. + :return: None. + """ + if not output_data_path: + output_data_path = storage.dug_annotation_path('') + log.info("Intializing parser") + parser = self.get_parser() + log.info("Done intializing parser") + + pending = [f for f in parsable_files + if not self.annotation_is_complete(f, output_data_path)] + skipped = len(parsable_files) - len(pending) + if skipped: + log.info("Resuming: %d of %d files already annotated, %d to go", + skipped, len(parsable_files), len(pending)) + + workers = max(1, int(self.annotation_conf.annotate_workers)) + # A thread per file, not per element: files are wholly independent + # (own parse, own crawler, own output dir) and the work is nearly all + # http wait, so threads scale it despite the GIL. Element-level + # concurrency would have to live in dug's Crawler. + workers = min(workers, len(pending)) or 1 + log.info("Annotating %d files with %d worker(s)", len(pending), + workers) + + if workers == 1: + for ct, parse_file in enumerate(pending): + self.annotate_one_file(parse_file, parser, output_data_path, + ct, len(pending)) + return + + with ThreadPoolExecutor(max_workers=workers, + thread_name_prefix='annotate') as pool: + futures = [ + pool.submit(self.annotate_one_file, parse_file, parser, + output_data_path, ct, len(pending)) + for ct, parse_file in enumerate(pending) + ] + # surface the first failure rather than letting the pool swallow + # it; dug's annotate_elements has no per-element error handling, + # so a raised exception means that file produced nothing + for future in futures: + future.result() + + @staticmethod + def annotation_output_paths(parse_file, output_data_path): + "The (elements, concepts) files annotate_one_file writes for a source file" + stem = '.'.join(os.path.basename(parse_file).split('.')[:-1]) + element_dir = os.path.join(output_data_path, stem) + return (os.path.join(element_dir, 'elements.txt'), + os.path.join(element_dir, 'concepts.txt')) + + @classmethod + def annotation_is_complete(cls, parse_file, output_data_path): + """True if this file's annotation output is already fully written. + + Both files are required: they are written in sequence, so a task + killed between the two leaves a half-annotated directory that must be + redone. Annotation is the most expensive step in the pipeline -- tens + of seconds per file, weeks for a large dbGaP study -- so a retry that + started from scratch threw away everything the previous try paid for. + """ + return all(os.path.isfile(path) and os.path.getsize(path) > 0 + for path in cls.annotation_output_paths(parse_file, + output_data_path)) + + def annotate_one_file(self, parse_file, parser, output_data_path, + index=0, total=0): + "Parse and annotate a single input file, writing pickles for it" + session, annotator = self.thread_annotation_context() + log.info("Annotating %s (%d of %d)", parse_file, index + 1, total) + elements_file, concepts_file = self.annotation_output_paths( + parse_file, output_data_path) + crawler = Crawler( + crawl_file=parse_file, + parser=parser, + annotator=annotator, + tranqlizer='', + tranql_queries=[], + http_session=session + ) + + # Use the specified parser to parse the parse_file into elements. + elements = parser(parse_file) + log.debug("Parsed elements: %s", str(elements)) + + # This inserts the list of elements into the crawler where + # annotate_elements expects to find it. Maybe in some future version + # of Dug this could be a parameter instead of an attribute? + crawler.elements = elements + + # @TODO propose for Dug to make this a crawler class init param(??) + crawler.crawlspace = os.path.dirname(elements_file) + crawler.annotate_elements() + + # Extract out the concepts gotten out of annotation + # Extract out the elements + non_expanded_concepts = crawler.concepts + # The elements object will have been modified by annotate_elements, + # so we want to make sure to catch those modifications. + elements = crawler.elements + + # Write pickles of objects to file. elements first, then concepts; + # annotation_is_complete requires both, so a crash between them is + # correctly treated as unfinished on the next try. + log.info("Parsed and annotated: %s", parse_file) + storage.write_object(jsonpickle.encode(elements, indent=2), + elements_file) + storage.write_object( + jsonpickle.encode(non_expanded_concepts, indent=2), + concepts_file) + log.info("Serialized annotated elements and concepts to %s", + os.path.dirname(elements_file)) + + def convert_to_kgx_json(self, elements, written_nodes=None): + """ + Given an annotated and normalized set of study variables, + generate a KGX compliant graph given the normalized annotations. + Write that grpah to a graph database. + See BioLink Model for category descriptions. + https://biolink.github.io/biolink-model/notes.html + """ + if written_nodes is None: + written_nodes = set() + graph = { + "nodes": [], + "edges": [] + } + edges = graph['edges'] + nodes = graph['nodes'] + + for _, element in enumerate(elements): + # DugElement means a variable (Study variable...) + if not isinstance(element, DugElement): + continue + study_id = element.id + study_link = element.action + study_desc = element.description + study_name = element.name or element.id + + + if study_id not in written_nodes: + nodes.append({ + "id": study_id, + "category": ["biolink:Study"], + "name": study_name, + "url": study_link, + "description": study_desc + }) + written_nodes.add(study_id) + + # connect the study and the variable. + edges.append(make_edge( + subj=element.id, + relation_label='part of', + relation='BFO:0000050', + obj=study_id, + predicate='biolink:part_of', + predicate_label='part of')) + edges.append(make_edge( + subj=study_id, + relation_label='has part', + relation="BFO:0000051", + obj=element.id, + predicate='biolink:has_part', + predicate_label='has part')) + + # a node for the variable. Should be BL compatible + variable_node = { + "id": element.id, + "name": element.name, + "category": ["biolink:StudyVariable"], + # bulk loader parsing issue + "description": ( + element.description.replace("'", '`').replace('\n', ' ')) + } + if element.id not in written_nodes: + nodes.append(variable_node) + written_nodes.add(element.id) + + for identifier, metadata in element.concepts.items(): + identifier_object = metadata.identifiers.get(identifier) + # This logic is treating DBGap files. + # First item in current DBGap xml files is a topmed tag, + # This is treated as a DugConcept Object. But since its not + # a concept we get from annotation (?) its never added to + # variable.concepts.items (Where variable is a DugElement obj) + # The following logic is trying to extract types, and for the + # aformentioned topmed tag it adds + # `biolink:InfomrmationContentEntity` + # Maybe a better solution could be adding types on + # DugConcept objects + # More specifically Biolink compatible types (?) + # + if identifier_object: + category = identifier_object.types + elif identifier.startswith("TOPMED.TAG:"): + category = ["biolink:InformationContentEntity"] + else: + continue + if identifier not in written_nodes: + if isinstance(category, str): + bl_element = self.bl_toolkit.toolkit.get_element( + category) + category = [bl_element.class_uri or bl_element.slot_uri] + nodes.append({ + "id": identifier, + "category": category, + "name": metadata.name + }) + written_nodes.add(identifier) + # related to edge + edges.append(make_edge( + subj=element.id, + obj=identifier + )) + # related to edge + edges.append(make_edge( + subj=identifier, + obj=element.id)) + return graph + + def make_tagged_kg(self, elements): + """ Make a Translator standard knowledge graph representing + tagged study variables. + :param variables: The variables to model. + :param tags: The tags characterizing the variables. + :returns: dict with nodes and edges modeling a Translator/Biolink KG. + """ + graph = { + "nodes": [], + "edges": [] + } + edges = graph['edges'] + nodes = graph['nodes'] + + # Create graph elements to model tags and their + # links to identifiers gathered by semantic tagging + tag_map = {} + # @TODO extract this into config or maybe dug ?? + topmed_tag_concept_type = "TOPMed Phenotype Concept" + nodes_written = set() + for tag in elements: + if not (isinstance(tag, DugConcept) + and tag.type == topmed_tag_concept_type): + continue + tag_id = tag.id + tag_map[tag_id] = tag + nodes.append({ + "id": tag_id, + "name": tag.name, + "description": tag.description.replace("'", "`"), + "category": ["biolink:InformationContentEntity"] + }) + + # Link ontology identifiers we've found for this tag via nlp. + for identifier, metadata in tag.identifiers.items(): + if isinstance(metadata.types, str): + bl_element = self.bl_toolkit.toolkit.get_element( + metadata.types) + category = [bl_element.class_uri or bl_element.slot_uri] + else: + category = metadata.types + synonyms = metadata.synonyms if metadata.synonyms else [] + nodes.append({ + "id": identifier, + "name": metadata.label, + "category": category, + "synonyms": synonyms + }) + nodes_written.add(identifier) + edges.append(make_edge( + subj=tag_id, + obj=identifier)) + edges.append(make_edge( + subj=identifier, + obj=tag_id)) + + concepts_graph = self.convert_to_kgx_json(elements, + written_nodes=nodes_written) + graph['nodes'] += concepts_graph['nodes'] + graph['edges'] += concepts_graph['edges'] + + return graph + + def record_ingest_date(self, *indices): + if self.index_obj is None: + self.index_obj: Index = self.factory.build_indexer_obj() + timestamp = datetime.now(timezone.utc).isoformat() + for index in indices: + log.info("Recording ingest date %s on %s", timestamp, index) + self.index_obj.set_ingest_date(index, timestamp) + + def index_elements(self, elements_file): + """Submit one elements file to ElasticSearch, routed by element type. + + Bulk, for the same reason as _index_concepts: dug's index_element + costs 2 requests for a new document and 3 for an existing one + (exists + get + update), which on a 127k-document index is hours of + sequential round trips. + + Union semantics are preserved -- ids repeat across element files and + each occurrence can carry different concepts and parents -- but the + reads are batched: one mget per 1000 ids instead of two requests per + document. + """ + from elasticsearch.helpers import bulk + + if self.index_obj is None: + self.index_obj: Index = self.factory.build_indexer_obj() + + log.info("Indexing %s...", str(elements_file)) + elements = jsonpickle.decode(storage.read_object(elements_file)) + log.info("found %d from elements files.", len(elements)) + + # isinstance, not type(): pipelines may subclass these + index_for = ((DugVariable, self.variables_index), + (DugStudy, self.studies_index), + (DugSection, self.sections_index)) + + by_index = {} + for element in elements: + # concepts are indexed separately by _index_concepts + if isinstance(element, DugConcept) or not element.id: + continue + index = next((idx for cls, idx in index_for + if isinstance(element, cls)), None) + if index is None: + continue + # override data-type with mapping values + if element.type.lower() in self.element_mapping: + element.type = self.element_mapping[element.type.lower()] + docs = by_index.setdefault(index, {}) + doc_id = element.get_id() + doc = element.get_searchable_dict() + # same id twice inside one file + docs[doc_id] = (merge_searchable_docs(docs[doc_id], doc) + if doc_id in docs else doc) + + # fold in what earlier files (or datasets) already indexed + for index, docs in by_index.items(): + ids = list(docs) + for i in range(0, len(ids), 1000): + chunk = ids[i:i + 1000] + found = self.index_obj.es.mget( + index=index, body={'ids': chunk}) + for entry in found.get('docs', []): + if entry.get('found'): + doc_id = entry['_id'] + docs[doc_id] = merge_searchable_docs( + entry['_source'], docs[doc_id]) + + actions = [{'_op_type': 'index', '_index': index, + '_id': doc_id, '_source': doc} + for index, docs in by_index.items() + for doc_id, doc in docs.items()] + indexed, errors = bulk(self.index_obj.es, actions, + chunk_size=1000, raise_on_error=False, + request_timeout=120) + if errors: + log.error("%d document(s) failed to index; first: %s", + len(errors), errors[0]) + raise PipelineException( + f"Bulk indexing failed for {len(errors)} document(s) " + f"from {elements_file}") + log.info("Done indexing %s: %d document(s).", elements_file, indexed) + + def validate_indexed_element_file(self, elements_file): + "After submitting elements for indexing, verify that they're available" + elements = [x for x in jsonpickle.decode( + storage.read_object(elements_file)) + if not isinstance(x, DugConcept)] + # Pick ~ 10 % + sample_size = int(len(elements) * 0.1) + + # random.choices(elements, k=sample_size) + test_elements = elements[:sample_size] + log.info("Picked %d from %s for validation.", len(test_elements), + elements_file) + for element in test_elements: + # Pick a concept + concepts = [element.concepts[curie] for curie in element.concepts + if element.concepts[curie].name] + + if len(concepts): + # Pick the first concept + concept = concepts[0] + curie = concept.id + search_term = re.sub(r'-', ' ', concept.name) + search_term = re.sub(r'[^a-zA-Z0-9/\.<>_:\ ]+', '', search_term) + if len(search_term) < 3: + # anything less than 3 chars won't return hits + continue + log.debug("Searching for Concept: %s and Search term: %s", + str(curie), search_term) + all_elements_ids = self._search_elements(curie, search_term) + present = element.id in all_elements_ids + if not present: + log.error("Did not find expected variable %s in search " + "result.", str(element.id)) + log.error("Concept id : %s, Search term: %s", + str(concept.id), search_term) + raise PipelineException( + f"Validation exception - did not find variable " + f"{element.id} from {str(elements_file)}" + f"when searching variable index with Concept ID : " + f"{concept.id} using Search Term : {search_term} ") + else: + log.info( + "%s has no concepts annotated. Skipping validation for it.", + str(element.id)) + + def _search_elements(self, curie, search_term): + "Asynchronously call a search on the curie and search term" + if self.search_obj == None : + self.search_obj: Search = self.factory.build_search_obj() + + page_size = 10000 + offset = 0 + hits, total_items, _ = self.event_loop.run_until_complete( + self.search_obj.search_elements( + self.variables_index, + concept=curie, + query=search_term, + size=page_size, + offset=offset + )) + if total_items == 0: + log.error(f"No search elements returned for variable search: " + f"{self.variables_index}.") + log.error(f"Concept id : {curie}, Search term: {search_term}") + raise Exception(f"Validation error - Did not find {curie} for" + f"Search term: {search_term}") + + while len(hits) < total_items: + offset += page_size + new_hits, _, _ = self.event_loop.run_until_complete( + self.search_obj.search_elements( + self.variables_index, + concept=curie, + query=search_term, + size=page_size, + offset=offset + )) + hits += new_hits + id_list = [e['_source']['id'] for e in hits] + return id_list + + def crawl_concepts(self, concepts, data_set_name, output_path=None): + """Adds tranql KG to Concepts + + Terms grabbed from KG are also added as search terms + :param concepts: + :param data_set_name: + :return: + """ + # TODO crawl dir seems to be storaing crawling info to avoid + # re-crawling, but is that consting us much? , it was when tranql was + # slow, but might right to consider getting rid of it. + crawl_dir = storage.dug_crawl_path('crawl_output') + output_file_name = os.path.join(data_set_name, + 'expanded_concepts.txt') + extracted_dug_elements_file_name = os.path.join( + data_set_name, 'extracted_graph_elements.txt') + if not output_path: + output_file = storage.dug_expanded_concepts_path(output_file_name) + extracted_output_file = storage.dug_expanded_concepts_path( + extracted_dug_elements_file_name + ) + else: + output_file = os.path.join(output_path, output_file_name) + extracted_output_file = os.path.join( + output_path, extracted_dug_elements_file_name) + + Path(crawl_dir).mkdir(parents=True, exist_ok=True) + extracted_dug_elements = [] + log.debug("Creating Dug Crawler object") + crawler = Crawler( + crawl_file="", + parser=None, + annotator=None, + tranqlizer=self.tranqlizer, + tranql_queries=self.tranql_queries, + http_session=self.cached_session, + ) + crawler.crawlspace = crawl_dir + counter = 0 + total = len(concepts) + for concept in concepts.values(): + counter += 1 + try: + crawler.expand_concept(concept) + concept.set_search_terms() + concept.set_optional_terms() + except Exception as e: + log.error(concept) + raise e + # for query in self.node_to_element_queries: + # log.info(query) + # casting_config = query['casting_config'] + # tranql_source = query['tranql_source'] + # dug_element_type = query['output_dug_type'] + # extracted_dug_elements += crawler.expand_to_dug_element( + # concept=concept, + # casting_config=casting_config, + # dug_element_type=dug_element_type, + # tranql_source=tranql_source + # ) + concept.clean() + percent_complete = int((counter / total) * 100) + if percent_complete % 10 == 0: + log.info("%d%%", percent_complete) + log.info("Crawling %s done", data_set_name) + storage.write_object(obj=jsonpickle.encode(concepts, indent=2), + path=output_file) + log.info ("Concepts serialized to %s", output_file) + # storage.write_object(obj=jsonpickle.encode(extracted_dug_elements, + # indent=2), + # path=extracted_output_file) + # log.info("Extracted elements serialized to %s", extracted_output_file) + + def _index_concepts(self, concepts): + """Submit concepts and their KG answers to ElasticSearch. + + Uses the bulk API rather than dug's per-doc index_concept / + index_kg_answer: those issue an exists check plus a write per + document, so a 60k-concept dataset became ~150k sequential round + trips and ran for an hour. Bulk sends 1000 docs per request. + + Dropping the exists check is deliberate. index_concept skips + documents already present, which silently keeps stale copies; these + tasks always rebuild the indexes wholesale, so an id-keyed upsert is + both correct and what we want. + """ + from elasticsearch.helpers import bulk + + if self.index_obj is None: + self.index_obj: Index = self.factory.build_indexer_obj() + + log.info("Indexing %d concepts", len(concepts)) + + def actions(): + for concept_id, concept in concepts.items(): + yield {'_op_type': 'index', + '_index': self.concepts_index, + '_id': concept_id, + '_source': concept.get_searchable_dict()} + for kg_answer_id, kg_answer in concept.kg_answers.items(): + targets = (kg_answer.get_node_names(include_curie=False) + + kg_answer.get_node_synonyms( + include_curie=False)) + yield {'_op_type': 'index', + '_index': self.kg_index, + '_id': f"{concept_id}_{kg_answer_id}", + '_source': { + 'concept_id': concept_id, + 'search_targets': dedupe_and_sort(targets), + 'knowledge_graph': kg_answer.get_kg()}} + + indexed, errors = bulk(self.index_obj.es, actions(), + chunk_size=1000, raise_on_error=False, + request_timeout=120) + if errors: + log.error("%d document(s) failed to index; first: %s", + len(errors), errors[0]) + raise PipelineException( + f"Bulk indexing failed for {len(errors)} document(s)") + log.info("Done Indexing concepts: %d document(s)", indexed) + + def _validate_indexed_concepts(self, elements, concepts): + """ + Validates linked concepts are searchable + :param elements: Annotated dug elements + :param concepts: Crawled (expanded) concepts + :return: + """ + # 1 . Find concepts with KG <= 10% of all concepts, + # <= because we might have no results for some concepts from tranql + sample_concepts = {key: value for key, value + in concepts.items() if value.kg_answers} + if len(concepts) == 0: + log.info("No Concepts found.") + return + log.info("Found only %d Concepts with Knowledge graph out of %d. %d%%", + len(sample_concepts), len(concepts), + (len(sample_concepts) / len(concepts)) * 100) + # 2. pick variables that have concepts in the sample concepts set + sample_elements = {} + for element in elements: + if not isinstance(element, DugVariable): + continue + for concept in element.concepts: + # add elements that have kg + if concept in sample_concepts: + sample_elements[concept] = sample_elements.get( + concept, set()) + sample_elements[concept].add(element.id) + + # Time for some validation + for curie in concepts: + concept = concepts[curie] + if not concept.kg_answers: + continue + if curie not in sample_elements: + # Concept had no variable elements associated + continue + search_terms = [] + for key in concept.kg_answers: + kg_object = concept.kg_answers[key] + search_terms += kg_object.get_node_names() + search_terms += kg_object.get_node_synonyms() + # reduce(lambda x,y: x + y, [[node.get("name")] + # + node.get("synonyms", []) + # for node in concept.kg_answers[ + # "knowledge_graph"]["nodes"]], []) + # validation here is that for any of these nodes we should get back + # the variable. + # make unique + search_terms_cap = 10 + search_terms = list(set(search_terms))[:search_terms_cap] + log.debug("Using %d Search terms for concept %s", len(search_terms), + str(curie)) + for search_term in search_terms: + # avoids elastic failure due to some reserved characters + # 'search_phase_execution_exception', + # 'token_mgr_error: Lexical error ... + search_term = re.sub(r'-', ' ', search_term) + search_term = re.sub(r'[^a-zA-Z0-9/\.<>_:\ ]+', '', search_term) + + searched_element_ids = self._search_elements(curie, search_term) + + present = bool([x for x in sample_elements[curie] + if x in searched_element_ids]) + if not present: + log.error("Did not find expected variable %s " + "in search result.", + str(curie)) + log.error("Concept id : %s, Search term: %s", + str(concept.id), search_term) + import json + log.error("search_element_ids: %s", json.dumps(searched_element_ids)) + log.error("sample_elements: %s", json.dumps({k: list(v) for k,v in sample_elements.items()})) + raise PipelineException( + f"Validation error - Did not find {curie} for" + f" Concept id : {concept.id}, " + f"Search term: {search_term}") + + def clear_index(self, index_id): + "Delete the index specified by index_id from ES" + # lazy init: clearing can run as the first ES touch of a task + if self.search_obj is None: + self.search_obj: Search = self.factory.build_search_obj() + if self.index_obj is None: + self.index_obj: Index = self.factory.build_indexer_obj() + exists = self.event_loop.run_until_complete( + self.search_obj.es.indices.exists(index=index_id)) + if exists: + log.info("Deleting index %s", str(index_id)) + response = self.event_loop.run_until_complete( + self.search_obj.es.indices.delete(index=index_id)) + log.info("Cleared Elastic : %s", str(response)) + log.info("Re-initializing the indicies") + self.index_obj.init_indices() + + def clear_variables_index(self): + "Delete the variables index from ES" + self.clear_index(self.variables_index) + + def clear_kg_index(self): + "Delete the KG index from ES" + self.clear_index(self.kg_index) + + def clear_concepts_index(self): + "Delete the concepts index from ES" + self.clear_index(self.concepts_index) + + def clear_all_es_indexes(self, to_string=False, input_data_path=None, + output_data_path=None): + """Wipe every ES index ahead of a full rebuild from the files + remaining in lakefs. Callable as an Airflow task method.""" + for index_id in (self.variables_index, self.studies_index, + self.sections_index, self.concepts_index, + self.kg_index): + self.clear_index(index_id) + return self.log_stream.getvalue() if to_string else '' + + #### + # Methods above this are directly from what used to be + # dug_helpers.dug_utils.Dug. Methods below are consolidated from what used + # to be dug_helpers.dug_utils.DugUtil. These are intented to be the "top + # level" interface to Roger, which Airflow DAGs or other orchestrators can + # call directly. + + def _fetch_s3_file(self, filename, output_dir): + "Fetch a file from s3 to output_dir" + log.info("Fetching %s", filename) + output_name = filename.split('/')[-1] + output_path = output_dir / output_name + self.s3_utils.get( + str(filename), + str(output_path), + ) + if self.unzip_source: + log.info("Unzipping %s", str(output_path)) + with tarfile.open(str(output_path)) as tar: + tar.extractall(path=output_dir) + return output_path + + def _fetch_remote_file(self, filename, output_dir, current_version): + "Fetch a file from a location using FileFetcher" + log.info("Fetching %s", filename) + # fetch from stars + remote_host = self.config.annotation_base_data_uri + fetch = FileFetcher( + remote_host=remote_host, + remote_dir=current_version, + local_dir=output_dir) + output_path = fetch(filename) + if self.unzip_source: + log.info("Unzipping %s", str(output_path)) + with tarfile.open(str(output_path)) as tar: + tar.extractall(path=output_dir) + return output_path + + def get_versioned_files(self): + """ Fetches a dug input data files to input file directory + """ + meta_data = storage.read_relative_object("../../metadata.yaml") + output_dir: Path = storage.dug_input_files_path( + self.get_files_dir()) + data_store = self.config.dug_inputs.data_source + + # clear dir + storage.clear_dir(output_dir) + data_sets = self.config.dug_inputs.data_sets + log.info("dataset: %s", data_sets) + pulled_files = [] + for data_set in data_sets: + data_set_name, current_version = data_set.split(':') + for item in meta_data["dug_inputs"]["versions"]: + if (item["version"] == current_version and + item["name"] == data_set_name and + item["format"] == self.get_data_format()): + if data_store == "s3": + for filename in item["files"]["s3"]: + pulled_files.append( + self._fetch_s3_file(filename, output_dir)) + else: + for filename in item["files"]["stars"]: + pulled_files.append( + self.fetch_remote_file(filename, output_dir, + current_version)) + return [str(filename) for filename in pulled_files] + + @staticmethod + def input_file_filter(file_name): + """Default filter for """ + return file_name.endswith('.xml') + + def get_objects(self, input_data_path=None): + """Retrieve initial source objects for parsing + + This is a default method that will be overridden by subclasses + frequently, it is expected. + """ + if not input_data_path: + input_data_path = storage.dug_input_files_path( + self.get_files_dir()) + files = storage.get_files_recursive( + self.input_file_filter, + input_data_path) + return sorted([str(f) for f in files]) + + def annotate(self, to_string=False, files=None, input_data_path=None, + output_data_path=None): + "Annotate files with the appropriate parsers and crawlers" + log.debug("annotate called with files %s, input path %s, " + "output path $s", files, str(input_data_path), + str(output_data_path)) + if files is None: + files = self.get_objects(input_data_path=input_data_path) + self.annotate_files(parsable_files=files, + output_data_path=output_data_path) + output_log = self.log_stream.getvalue() if to_string else '' + return output_log + + def index_variables(self, to_string=False, element_object_files=None, + input_data_path=None, output_data_path=None): + """Index variables from element object files for pipeline + + if element_object_files is specified, only those files are + indexed. Otherwise, if the input_data_path is supplied, elements files + under that path are indexed. If neither is supplied, the expanded + concepts directory is searched for elements files that have been + updated with KG-derived optional terms from the crawl step. + """ + # self.clear_variables_index() + if element_object_files is None: + element_object_files = storage.dug_expanded_elements_objects( + input_data_path, format='txt') + for file_ in element_object_files: + self.index_elements(file_) + self.record_ingest_date(self.variables_index, self.studies_index, + self.sections_index) + output_log = self.log_stream.getvalue() if to_string else '' + return output_log + + def validate_indexed_variables(self, to_string=None, + element_object_files=None, + input_data_path=None, + output_data_path=None): + "Validate output from index variables task for pipeline" + if not element_object_files: + element_object_files = storage.dug_expanded_elements_objects( + input_data_path, format='txt') + for file_ in element_object_files: + log.info("Validating %s", str(file_)) + self.validate_indexed_element_file(file_) + output_log = self.log_stream.getvalue() if to_string else '' + return output_log + + def validate_indexed_concepts(self, config=None, to_string=None, + input_data_path=None, output_data_path=None): + """ + Entry for validate concepts + """ + get_data_set_name = lambda file: ( + os.path.split(os.path.dirname(file))[-1]) + expanded_concepts_files_dict = { + get_data_set_name(file): file for file in + storage.dug_expanded_concept_objects(data_path=input_data_path, + format='txt') + } + annotated_elements_files_dict = { + get_data_set_name(file): file for file in + storage.dug_expanded_elements_objects(data_path=input_data_path, + format='txt') + } + try: + assert (len(expanded_concepts_files_dict) == + len(annotated_elements_files_dict)) + except: + log.error("Files Annotated Elements files and " + "expanded concepts files, should be pairs") + if len(expanded_concepts_files_dict) > len(annotated_elements_files_dict): + log.error("Some Annotated Elements files " + "(from load_and_annotate task) are missing") + else: + log.error("Some Expanded Concepts files (from crawl task) are missing") + log.error(f"Annotated Datasets : {list(annotated_elements_files_dict.keys())}") + log.error(f"Expanded Concepts Datasets: {list(expanded_concepts_files_dict.keys())}") + exit(-1) + for data_set_name in annotated_elements_files_dict: + log.debug(f"Reading concepts and elements for dataset {data_set_name}") + elements_file_path = annotated_elements_files_dict[data_set_name] + concepts_file_path = expanded_concepts_files_dict[data_set_name] + dug_elements = jsonpickle.decode(storage.read_object(elements_file_path)) + dug_concepts = jsonpickle.decode(storage.read_object(concepts_file_path)) + log.debug(f"Read {len(dug_elements)} elements, and {len(dug_concepts)} Concepts") + log.info(f"Validating {data_set_name}") + self._validate_indexed_concepts(elements=dug_elements, concepts=dug_concepts) + output_log = self.log_stream.getvalue() if to_string else '' + return output_log + + def make_kg_tagged(self, to_string=False, elements_files=None, + input_data_path=None, output_data_path=None): + "Create tagged knowledge graphs from elements" + if not output_data_path: + output_data_path = storage.dug_kgx_path("") + storage.clear_dir(output_data_path) + log.info("Starting building KGX files") + + if not elements_files: + elements_files = storage.dug_elements_objects(input_data_path, format='txt') + log.info(f"found {len(elements_files)} files : {elements_files}") + for file_ in elements_files: + elements = jsonpickle.decode(storage.read_object(file_)) + if "topmed_" in file_: + kg = self.make_tagged_kg(elements) + else: + kg = self.convert_to_kgx_json(elements) + dug_base_file_name = file_.split(os.path.sep)[-2] + output_file_path = os.path.join(output_data_path, + dug_base_file_name + '_kgx.json') + storage.write_object(kg, output_file_path) + log.info("Wrote %d and %d edges, to %s", len(kg['nodes']), + len(kg['edges']), output_file_path) + output_log = self.log_stream.getvalue() if to_string else '' + return output_log + + def crawl_tranql(self, to_string=False, concept_files=None, + input_data_path=None, output_data_path=None): + "Perform the tranql crawl" + if not concept_files: + concept_files = storage.dug_concepts_objects( + input_data_path, format='txt') + + if output_data_path: + crawl_dir = os.path.join(output_data_path, 'crawl_output') + expanded_concepts_dir = os.path.join(output_data_path, + 'expanded_concepts') + else: + crawl_dir = storage.dug_crawl_path('crawl_output') + expanded_concepts_dir = storage.dug_expanded_concepts_path("") + log.info("Clearing crawl output dir %s", crawl_dir) + storage.clear_dir(crawl_dir) + + log.info("Clearing expanded concepts dir: %s", expanded_concepts_dir) + storage.clear_dir(expanded_concepts_dir) + + log.info("Crawling Dug Concepts, found %d file(s).", + len(concept_files)) + for file_ in concept_files: + objects = storage.read_object(file_) + objects = objects or {} + if not objects: + log.info(f'no concepts in {file_}') + data_set = jsonpickle.decode(objects) + original_variables_dataset_name = os.path.split( + os.path.dirname(file_))[-1] + self.crawl_concepts(concepts=data_set, + data_set_name=original_variables_dataset_name, + output_path= output_data_path) + + # After expanding concepts with KG answers, update the + # corresponding elements' optional_terms so that KG-derived + # search terms are present when elements are later indexed. + # This mirrors what Crawler.crawl() does after concept expansion. + # The updated elements are written to the expanded concepts + # directory (alongside expanded_concepts.txt) rather than + # mutating the annotate step's output. + annotation_elements_file = os.path.join( + os.path.dirname(file_), 'elements.txt') + expanded_elements_file_name = os.path.join( + original_variables_dataset_name, 'elements.txt') + if not output_data_path: + expanded_elements_file = ( + storage.dug_expanded_concepts_path( + expanded_elements_file_name)) + else: + expanded_elements_file = os.path.join( + output_data_path, expanded_elements_file_name) + if os.path.exists(annotation_elements_file): + log.info("Updating element optional terms from expanded " + "concepts for %s", original_variables_dataset_name) + elements = jsonpickle.decode( + storage.read_object(annotation_elements_file)) + for element in elements: + if isinstance(element, DugConcept): + continue + # Replace each element's concept references with + # the expanded versions that now carry kg_answers. + for concept_id in list(element.concepts.keys()): + if concept_id in data_set: + element.concepts[concept_id] = data_set[ + concept_id] + element.set_optional_terms() + storage.write_object( + jsonpickle.encode(elements, indent=2), + expanded_elements_file) + log.info("Updated elements serialized to %s", + expanded_elements_file) + else: + log.warning("Elements file not found at %s, skipping " + "optional terms update", + annotation_elements_file) + + output_log = self.log_stream.getvalue() if to_string else '' + return output_log + + def index_concepts(self, to_string=False, + input_data_path=None, output_data_path=None): + "Index concepts from expanded concept files" + # These are concepts that have knowledge graphs from tranql + # clear out concepts and kg indicies from previous runs + # self.clear_concepts_index() + # self.clear_kg_index() + expanded_concepts_files = storage.dug_expanded_concept_objects( + input_data_path, format="txt") + for file_ in expanded_concepts_files: + concepts = jsonpickle.decode(storage.read_object(file_)) + self._index_concepts(concepts=concepts) + + if self.config.indexing.node_to_element_queries: + log.info("*******************") + + extracted_elements_files = storage.dug_extracted_elements_objects( + data_path=input_data_path) + log.info(f"{extracted_elements_files}") + for file_ in extracted_elements_files: + log.info(f"reading file {file_}") + self.index_elements(file_) + self.record_ingest_date(self.variables_index, self.studies_index, + self.sections_index) + self.record_ingest_date(self.concepts_index, self.kg_index) + output_log = self.log_stream.getvalue() if to_string else '' + return output_log + +class DDM2Pipeline(DugPipeline): + """Base class for pipelines working on Dug Data Model v2""" + + @staticmethod + def input_file_filter(file_name): + """Retrieve initial source objects for parsing in .dug.json format + + This ideally removes the need for specialized get_objects methods in + DDM2 pipelines. + """ + return file_name.endswith('.dug.json'), diff --git a/src/roger/pipelines/bdc.py b/src/roger/pipelines/bdc.py new file mode 100644 index 00000000..bc30cf44 --- /dev/null +++ b/src/roger/pipelines/bdc.py @@ -0,0 +1,19 @@ +"Pipeline for BDC-dbGap data" + +from roger.pipelines import DugPipeline +from roger.core import storage + +class bdcPipeline(DugPipeline): + "Pipeline for BDC-dbGap data set" + pipeline_name = "bdc" + parser_name = "dbgap" + + def get_objects(self, input_data_path=None): + if not input_data_path: + input_data_path = storage.dug_dd_xml_path() + files = storage.get_files_recursive( + lambda file_name: ( + not file_name.startswith('._') + and file_name.endswith('.xml')), + input_data_path) + return sorted([str(f) for f in files]) diff --git a/src/roger/pipelines/bdc_pipelines.py b/src/roger/pipelines/bdc_pipelines.py new file mode 100644 index 00000000..d4c6436d --- /dev/null +++ b/src/roger/pipelines/bdc_pipelines.py @@ -0,0 +1,58 @@ +"Dug pipeline for dbGaP data set" + +from roger.pipelines import DugPipeline + +class BIOLINCCdbGaPPipeline(DugPipeline): + "Pipeline for the dbGaP data set" + pipeline_name = 'bdc-biolincc' + parser_name = 'biolincc' + + +class covid19dbGaPPipeline(DugPipeline): + "Pipeline for the dbGaP data set" + pipeline_name = 'bdc-covid19' + parser_name = 'covid19' + +class dirDbGaPPipeline(DugPipeline): + pipeline_name = "bdc-dir" + parser_name = "dir" + +class LungMapDbGaPPipeline(DugPipeline): + pipeline_name = "bdc-lungmap" + parser_name = "lungmap" + +class nsrrDbGaPPipeline(DugPipeline): + pipeline_name = "bdc-nsrr" + parser_name = "nsrr" + +class ParentDbGaPPipeline(DugPipeline): + pipeline_name = "bdc-parent" + parser_name = "parent" + +class PCGCDbGaPPipeline(DugPipeline): + pipeline_name = "pcgc-dbgap" + parser_name = "pcgc" + +class RecoverDbGaPPipeline(DugPipeline): + pipeline_name = "bdc-recover" + parser_name = "recover" + +class TopmedDBGaPPipeline(DugPipeline): + pipeline_name = "bdc-topmed" + parser_name = "topmeddbgap" + +class CureSCPipeline(DugPipeline): + pipeline_name = "bdc-curesc" + parser_name = "curesc" + +class HeartFailurePipeline(DugPipeline): + pipeline_name = "bdc-heartfailure" + parser_name = "heartfailure" + +class ImagingPipeline(DugPipeline): + pipeline_name = "bdc-imaging" + parser_name = "imaging" + +class RedsPipeline(DugPipeline): + pipeline_name = "bdc-reds" + parser_name = "reds" \ No newline at end of file diff --git a/src/roger/pipelines/crdc.py b/src/roger/pipelines/crdc.py new file mode 100644 index 00000000..2143cf7b --- /dev/null +++ b/src/roger/pipelines/crdc.py @@ -0,0 +1,19 @@ +"Pipeline for Cancer Commons data" + +from roger.pipelines import DugPipeline +from roger.core import storage + +class CRDCPipeline(DugPipeline): + "Pipeline for Cancer Commons data set" + pipeline_name = "crdc" + parser_name = "crdc" + + def get_objects(self, input_data_path=None): + if not input_data_path: + input_data_path = storage.dug_crdc_path() + files = storage.get_files_recursive( + lambda file_name: ( + not file_name.startswith('GapExchange_') + and file_name.endswith('.xml')), + input_data_path) + return sorted([str(f) for f in files]) diff --git a/src/roger/pipelines/ctn.py b/src/roger/pipelines/ctn.py new file mode 100644 index 00000000..25918062 --- /dev/null +++ b/src/roger/pipelines/ctn.py @@ -0,0 +1,10 @@ +"Pipeline for Clinical trials network data" + +from roger.pipelines import DugPipeline + +class CTNPipeline(DugPipeline): + "Pipeline for Clinical trials nework data set" + pipeline_name = "ctn" + parser_name = "ctn" + + diff --git a/src/roger/pipelines/db_gap.py b/src/roger/pipelines/db_gap.py new file mode 100644 index 00000000..7c1db504 --- /dev/null +++ b/src/roger/pipelines/db_gap.py @@ -0,0 +1,10 @@ +"Dug pipeline for dbGaP data set" + +from roger.pipelines import DugPipeline + +class dbGaPPipeline(DugPipeline): + "Pipeline for the dbGaP data set" + + pipeline_name = 'dbGaP' + parser_name = 'DbGaP' + files_dir = 'db_gap' diff --git a/src/roger/pipelines/heal_cdes.py b/src/roger/pipelines/heal_cdes.py new file mode 100644 index 00000000..275dff59 --- /dev/null +++ b/src/roger/pipelines/heal_cdes.py @@ -0,0 +1,17 @@ +"Pipeline to ingest HEAL in new data model" + +from roger.pipelines import DDM2Pipeline +from roger.core import storage + +class HealStudiesDDM2Pipeline(DDM2Pipeline): + "Pipeline for HEAL data using dug data model v2" + pipeline_name = "heal-cdes" + parser_name = "heal-ddm2" + + def get_objects(self, input_data_path=None): + if not input_data_path: + input_data_path = storage.dug_heal_study_path() + files = storage.get_files_recursive( + lambda file_name: file_name.endswith('.dug.json'), + input_data_path) + return sorted([str(f) for f in files]) diff --git a/src/roger/pipelines/heal_research_programs.py b/src/roger/pipelines/heal_research_programs.py new file mode 100644 index 00000000..bfec3f83 --- /dev/null +++ b/src/roger/pipelines/heal_research_programs.py @@ -0,0 +1,16 @@ +"Pipeline for Heal-studies data" + +from roger.pipelines import DugPipeline +from roger.core import storage + +class HealResearchProgramPipeline(DugPipeline): + "Pipeline for Heal-research-programs data set" + pipeline_name = "heal-mds-research-networks" + parser_name = "heal-research" + + def get_objects(self, input_data_path=None): + if not input_data_path: + input_data_path = storage.dug_heal_research_program_path() + files = storage.get_files_recursive(lambda file_name: file_name.endswith('.xml'), + input_data_path) + return sorted([str(f) for f in files]) \ No newline at end of file diff --git a/src/roger/pipelines/heal_studies.py b/src/roger/pipelines/heal_studies.py new file mode 100644 index 00000000..3f9c7f81 --- /dev/null +++ b/src/roger/pipelines/heal_studies.py @@ -0,0 +1,17 @@ +"Pipeline for Heal-studies data" + +from roger.pipelines import DDM2Pipeline +from roger.core import storage + +class HealStudiesPipeline(DDM2Pipeline): + "Pipeline for Heal-studies data set" + pipeline_name = "heal-mds-studies" + parser_name = "heal-ddm2" + + def get_objects(self, input_data_path=None): + if not input_data_path: + input_data_path = storage.dug_heal_study_path() + files = storage.get_files_recursive( + lambda file_name: file_name.endswith('.dug.json'), + input_data_path) + return sorted([str(f) for f in files]) diff --git a/src/roger/pipelines/kfdrc.py b/src/roger/pipelines/kfdrc.py new file mode 100644 index 00000000..bcb0b7ac --- /dev/null +++ b/src/roger/pipelines/kfdrc.py @@ -0,0 +1,19 @@ +"Pipeline for KDFRC data" + +from roger.pipelines import DugPipeline +from roger.core import storage + +class kfdrcPipeline(DugPipeline): + "Pipeline for KDFRC data set" + pipeline_name = "kfdrc" + parser_name = "kfdrc" + + def get_objects(self, input_data_path=None): + if not input_data_path: + input_data_path = storage.dug_kfdrc_path() + files = storage.get_files_recursive( + lambda file_name: ( + not file_name.startswith('GapExchange_') + and file_name.endswith('.xml')), + input_data_path) + return sorted([str(f) for f in files]) diff --git a/src/roger/pipelines/nida.py b/src/roger/pipelines/nida.py new file mode 100644 index 00000000..b2e841bd --- /dev/null +++ b/src/roger/pipelines/nida.py @@ -0,0 +1,18 @@ +"NIDA data set pipeline definition" + +from roger.pipelines import DugPipeline +from roger.core import storage + +class NIDAPipeline(DugPipeline): + "NIDA data pipeline" + + pipeline_name = 'nida' + parser_name = 'NIDA' + + def get_objects(self, input_data_path=None): + "Return list of NIDA source files" + if not input_data_path: + input_data_path = storage.dug_input_files_path( + self.get_files_dir()) + files = sorted(storage.get_files_recursive(lambda x: 'NIDA-' in x , input_data_path)) + return files diff --git a/src/roger/pipelines/picsure_test.py b/src/roger/pipelines/picsure_test.py new file mode 100644 index 00000000..21e04b64 --- /dev/null +++ b/src/roger/pipelines/picsure_test.py @@ -0,0 +1,27 @@ +from roger.pipelines import DugPipeline +from roger.core import storage +from roger.logger import logger + + +class PicSure(DugPipeline): + "Pipeline for BACPAC data set" + pipeline_name = "bdc-test6" #lakefs + parser_name = "dbgap" + files_dir = "anvil" + + def get_objects(self, input_data_path=None): + """Retrieve anvil objects + + This code is imported from roger.core.storage.dug_anvil_objects + """ + if not input_data_path: + input_data_path = storage.dug_input_files_path( + self.files_dir) + files = storage.get_files_recursive( + lambda file_name: ( + not file_name.startswith('GapExchange_') + and file_name.endswith('.xml')), + input_data_path) + logger.info("**********") + logger.info(files) + return sorted([str(f) for f in files]) diff --git a/src/roger/pipelines/radx.py b/src/roger/pipelines/radx.py new file mode 100644 index 00000000..7ffae159 --- /dev/null +++ b/src/roger/pipelines/radx.py @@ -0,0 +1,18 @@ +"Pipeline for BACPAC data" + +from roger.pipelines import DugPipeline +from roger.core import storage + + +class RadxPipeline(DugPipeline): + "Pipeline for Radx data set" + pipeline_name = "radx" + parser_name = "radx" + + def get_objects(self, input_data_path=None): + if not input_data_path: + input_data_path = storage.dug_kfdrc_path() + files = storage.get_files_recursive( + lambda file_name: file_name.endswith('.json'), + input_data_path) + return sorted([str(f) for f in files]) diff --git a/src/roger/pipelines/sparc.py b/src/roger/pipelines/sparc.py new file mode 100644 index 00000000..d1c9c950 --- /dev/null +++ b/src/roger/pipelines/sparc.py @@ -0,0 +1,17 @@ +"Pipeline for Sparc data" + +from roger.pipelines import DugPipeline +from roger.core import storage + +class SparcPipeline(DugPipeline): + "Pipeline for Sparc data set" + pipeline_name = "sparc" + parser_name = "SciCrunch" + + def get_objects(self, input_data_path=None): + if not input_data_path: + input_data_path = storage.dug_heal_study_path() + files = storage.get_files_recursive( + lambda x: True, input_data_path + ) + return sorted([str(f) for f in files]) diff --git a/src/roger/pipelines/topmed.py b/src/roger/pipelines/topmed.py new file mode 100644 index 00000000..90b3e515 --- /dev/null +++ b/src/roger/pipelines/topmed.py @@ -0,0 +1,41 @@ +"Pipeline for Topmed data" + +from roger.pipelines import DugPipeline +from roger.pipelines.base import log, os +import jsonpickle +from roger.core import storage +from roger.logger import logger +class TopmedPipeline(DugPipeline): + "Pipeline for Topmed data set" + pipeline_name = "topmed" + parser_name = "TOPMedTag" + + def get_objects(self, input_data_path=None): + if not input_data_path: + input_data_path = str(storage.dug_input_files_path('topmed')) + files =storage.get_files_recursive( + lambda file_name: file_name.endswith('.csv'), + input_data_path) + return sorted([str(x) for x in files]) + + def make_kg_tagged(self, to_string=False, elements_files=None, + input_data_path=None, output_data_path=None): + "Create tagged knowledge graphs from elements" + log.info("Override base.make_kg_tagged called") + if not output_data_path: + output_data_path = storage.dug_kgx_path("") + storage.clear_dir(output_data_path) + if not elements_files: + elements_files = storage.dug_elements_objects(input_data_path, format='txt') + for file_ in elements_files: + elements = jsonpickle.decode(storage.read_object(file_)) + kg = self.make_tagged_kg(elements) + dug_base_file_name = file_.split(os.path.sep)[-2] + output_file_path = os.path.join(output_data_path, + dug_base_file_name + '_kgx.json') + storage.write_object(kg, output_file_path) + log.info("Wrote %d and %d edges, to %s", len(kg['nodes']), + len(kg['edges']), output_file_path) + output_log = self.log_stream.getvalue() if to_string else '' + return output_log + diff --git a/src/roger/pvc.yaml b/src/roger/pvc.yaml new file mode 100644 index 00000000..691fed1b --- /dev/null +++ b/src/roger/pvc.yaml @@ -0,0 +1,10 @@ +apiVersion: v1 +kind: PersistentVolumeClaim +metadata: + name: search-data +spec: + accessModes: + - ReadWriteMany + resources: + requests: + storage: 10Mi diff --git a/src/roger/tasks.py b/src/roger/tasks.py new file mode 100755 index 00000000..37f60f66 --- /dev/null +++ b/src/roger/tasks.py @@ -0,0 +1,1147 @@ +# Tasks and methods related to Airflow implementations of Roger + +import os +import json +from datetime import datetime +from functools import partial +from typing import Union +from pathlib import Path +import glob +import shutil + +# Airflow 3.x - prefer provider imports and new public types +from airflow.providers.standard.operators.python import PythonOperator +from airflow.providers.standard.operators.empty import EmptyOperator +from airflow.exceptions import AirflowSkipException +from airflow.sdk import TaskGroup +from airflow.models import DAG +from airflow.models.taskinstance import TaskInstance +from airflow.providers.standard.operators.bash import BashOperator +from airflow.utils.context import Context # type: ignore +try: + # Task SDK Variable works on Airflow 3.x workers (no direct DB access) + from airflow.sdk import Variable +except ImportError: + from airflow.models import Variable + +from roger.config import config, RogerConfig +from roger.logger import get_logger +from roger.pipelines.base import DugPipeline +from avalon.mainoperations import put_files, LakeFsWrapper, get_files +from lakefs_sdk.configuration import Configuration +from lakefs_sdk.models.merge import Merge +from lakefs_sdk.models.path_list import PathList +from lakefs_sdk.models.commit_creation import CommitCreation +from lakefs_sdk.exceptions import NotFoundException + +logger = get_logger() + +default_args = { + 'owner': 'RENCI', + 'start_date': datetime(2025, 1, 1) +} + +# dotfile: storage.py readers glob '*.json' / '**/*.json' over pulled task +# outputs (e.g. kgx_objects) and glob skips dotfiles, so this never gets +# parsed as pipeline data +REMOVED_FILES_MANIFEST = ".removed_files.json" + + +def task_wrapper(python_callable, **kwargs): + """ + Overrides configuration with config from airflow. + :param python_callable: + :param kwargs: + :return: + """ + # get dag config provided + dag_run = kwargs.get('dag_run') + pass_conf = kwargs.get('pass_conf', True) + if config.lakefs_config.enabled: + # get input path + input_data_path = generate_dir_name_from_task_instance( + kwargs['ti'], + roger_config=config, + suffix='input' + ) + # get output path from task id run id dag id combo + output_data_path = generate_dir_name_from_task_instance( + kwargs['ti'], + roger_config=config, + suffix='output' + ) + else: + input_data_path, output_data_path = None, None + # cast it to a path object + func_args = { + 'input_data_path': input_data_path, + 'output_data_path': output_data_path, + 'to_string': kwargs.get('to_string') + } + logger.info(f"Task function args: {func_args}") + # overrides values + config.dag_run = dag_run + # splat, so this works for both callable shapes: roger.core functions + # (bulk_load, merge_nodes, ...) and execute_pipeline_method + if pass_conf: + return python_callable(**func_args, config=config) + return python_callable(**func_args) + + +def get_executor_config(data_path='/opt/airflow/share/data'): + """ Get an executor configuration. + :param annotations: Annotations to attach to the executor. + :returns: Returns a KubernetesExecutor if K8s configured, None otherwise. + """ + env_var_prefix = config.OS_VAR_PREFIX + # based on environment set on scheduler pod, make secrets for worker pod + secrets_map = [{ + "secret_name_ref": "ELASTIC_SEARCH_PASSWORD_SECRET", + "secret_key_ref": "ELASTIC_SEARCH_PASSWORD_SECRET_KEY", + "env_var_name": f"{env_var_prefix}ELASTIC__SEARCH_PASSWORD" + }, { + "secret_name_ref": "REDIS_PASSWORD_SECRET", + "secret_key_ref": "REDIS_PASSWORD_SECRET_KEY", + "env_var_name": f"{env_var_prefix}REDISGRAPH_PASSWORD" + }] + secrets = [] + for secret in secrets_map: + secret_name = os.environ.get(secret["secret_name_ref"], False) + secret_key_name = os.environ.get(secret["secret_key_ref"], False) + if secret_name and secret_key_name: + secrets.append({ + "name": secret["env_var_name"], + "valueFrom": { + "secretKeyRef": { + "name": secret_name, + "key": secret_key_name + } + }}) + + k8s_executor_config = { + "KubernetesExecutor": { + "envs": secrets, + } + } + return k8s_executor_config + + +def memory_override(limit: str, request: str = None) -> dict: + """executor_config bumping only this task's memory. + + Everything else (image, volumes, env, service account) is inherited from + the chart's worker pod template; this patches the 'base' container so one + heavy task does not force the default up for every task. Keep request + well under limit: the namespace quota counts requests.memory and + limits.memory separately. + """ + from kubernetes.client import models as k8s + return {"pod_override": k8s.V1Pod(spec=k8s.V1PodSpec(containers=[ + k8s.V1Container( + name="base", + resources=k8s.V1ResourceRequirements( + requests={"memory": request or "1Gi"}, + limits={"memory": limit}))]))} + + +def init_lakefs_client(config: RogerConfig) -> LakeFsWrapper: + configuration = Configuration() + configuration.username = config.lakefs_config.access_key_id + configuration.password = config.lakefs_config.secret_access_key + configuration.host = config.lakefs_config.host + the_lake = LakeFsWrapper(configuration=configuration) + return the_lake + + +def pagination_helper(page_fetcher, **kwargs): + """Helper function to iterate over paginated results""" + while True: + resp = page_fetcher(**kwargs) + yield from resp.results + if not resp.pagination.has_more: + break + kwargs['after'] = resp.pagination.next_offset + + +def incremental_state_key(dag_id: str, task_id: str, repo: str, + branch: str) -> str: + """Airflow Variable key holding the last source commit consumed by a + task. task_id is the group-qualified id, unique within a dag.""" + return f"roger_incr::{dag_id}::{task_id}::{repo}@{branch}" + + +def get_state_file_path(task_instance: TaskInstance) -> Union[str, None]: + path = generate_dir_name_from_task_instance( + task_instance, roger_config=config, suffix='state') + return str(path) if path else None + + +def read_state_file(task_instance: TaskInstance) -> dict: + path = get_state_file_path(task_instance) + if path and os.path.isfile(path): + with open(path) as f: + return json.load(f) + return {} + + +def write_state_file(task_instance: TaskInstance, state: dict): + path = get_state_file_path(task_instance) + if not path: + return + with open(path, 'w') as f: + json.dump(state, f) + + +def resolve_ref_tip(client: LakeFsWrapper, repo: str, ref: str) -> str: + """Resolve any ref (branch, tag or commit id) to a commit id. + + Tags are tried first. lakefs validates log_commits' ref argument as a + *branch id*, so a tag whose name is not a legal branch name -- the + dotted versions this pipeline uses, e.g. kgx.data_sets + 'baseline-graph:v7.0' -- is rejected with a 400 before the tag is ever + looked up. branches_api.get_branch has the same problem, which is why + neither API alone is enough. + """ + try: + return client._client.tags_api.get_tag(repository=repo, tag=ref).commit_id + except NotFoundException: + pass # not a tag; a branch or commit id then + results = client._client.refs_api.log_commits( + repository=repo, ref=ref, amount=1).results + return results[0].id if results else ref + + +def get_changed_files(client: LakeFsWrapper, repo: str, from_ref: str, + to_ref: str, prefixes=None) -> dict: + """Diff two refs, returning paths bucketed as added/changed/removed. + + Unlike avalon's get_changes this keeps 'removed' entries and filters by + path prefixes. Prefixes are normalized with a trailing '/' so that + 'a/task' does not match 'a/task_b/...'; '*' or empty means no filter. + """ + changes = {'added': [], 'changed': [], 'removed': []} + if not prefixes or '*' in prefixes: + norm = None + else: + norm = tuple(p if p.endswith('/') else p + '/' for p in prefixes) + for diff in pagination_helper(client._client.refs_api.diff_refs, + repository=repo, left_ref=from_ref, + right_ref=to_ref): + if norm and not diff.path.startswith(norm): + continue + if diff.type in changes: + changes[diff.type].append(diff.path) + return changes + + +def find_sibling_files(client: LakeFsWrapper, repo: str, ref: str, + downloaded, marker="GapExchange_") -> list: + """dbGaP data dicts need a sibling GapExchange_ file in the same + lakefs directory for study name/description. Incremental diffs only + carry the changed data dicts, so list each affected dir and pull any + marker file not already downloaded.""" + extra = set() + seen_dirs = set() + for path in downloaded: + d = path.rsplit('/', 1)[0] if '/' in path else '' + if d in seen_dirs: + continue + seen_dirs.add(d) + list_prefix = d + '/' if d else '' + for obj in pagination_helper( + client._client.objects_api.list_objects, + repository=repo, ref=ref, prefix=list_prefix, + delimiter='/', amount=1000): + name = obj.path.rsplit('/', 1)[-1] + if name.startswith(marker) and obj.path not in downloaded: + extra.add(obj.path) + return sorted(extra) + + +def removed_bases(removed: dict, dag_id: str, upstream_ids) -> set: + """Map removed source paths to the dataset base names their derived + outputs are keyed by. + + External-repo sources use annotate's convention (filename minus last + extension); roger-repo sources use the first path segment under the + upstream task dir (e.g. {dag}/{task}//elements.txt -> ). + """ + bases = set() + roger_repo = config.lakefs_config.repo + for src_repo, paths in removed.items(): + for p in paths: + if src_repo == roger_repo: + rel = p + for uid in upstream_ids: + prefix = f"{dag_id}/{uid}/" + if p.startswith(prefix): + rel = p[len(prefix):] + break + base = rel.split('/')[0] + if base == rel: + # top-level file, not a dataset dir + base = '.'.join(rel.split('.')[:-1]) + else: + base = '.'.join(os.path.basename(p).split('.')[:-1]) + if base and not base.startswith('.'): + bases.add(base) + return bases + + +def orphaned_output_paths(existing, remote_path: str, local_path: str) -> list: + """Remote objects under this task's prefix that the local output dir no + longer holds. + + The bulk-load CSV names encode data-dependent counters + ({type}.csv-{group}-{flush}), so every run writes new names and put_files + never overwrites the previous ones. Left alone they pile up and the + loader sees the same node id in files from several runs. + """ + local = set() + for root, _, files in os.walk(local_path): + for name in files: + local.add(os.path.relpath(os.path.join(root, name), local_path)) + orphans = [] + for f in existing: + rel = f[len(remote_path):] if f.startswith(remote_path) else f + if rel not in local: + orphans.append(f) + return orphans + + +def delete_objects(client: LakeFsWrapper, repo: str, branch: str, + paths: list, message: str): + "Delete paths on branch and commit. Chunked; lakefs caps the path list." + for i in range(0, len(paths), 1000): + client._client.objects_api.delete_objects( + repository=repo, branch=branch, + path_list=PathList(paths=paths[i:i + 1000])) + client._client.commits_api.commit( + repository=repo, branch=branch, + commit_creation=CommitCreation(message=message)) + + +def stale_output_paths(existing, remote_path: str, bases: set) -> list: + """Pick output objects derived from removed sources: anything under + / or the _kgx.json flat file make_kg_tagged emits.""" + stale = [] + for f in existing: + rel = f[len(remote_path):] if f.startswith(remote_path) else f + for b in bases: + if rel.startswith(b + '/') or rel == f"{b}_kgx.json": + stale.append(f) + break + return stale + + +def _get_last_consumed(key: str) -> Union[str, None]: + try: + return Variable.get(key, default=None) + except TypeError: + try: + return Variable.get(key, default_var=None) + except Exception: + return None + except Exception: + return None + + +def _advance_state_variables(state: dict): + for key, entry in state.get('entries', {}).items(): + Variable.set(key, entry['commit_id']) + logger.info("Recorded last ingested commit %s for %s", + entry['commit_id'], key) + + +def _merge_had_no_changes(exc): + """True when lakefs rejected a merge because there was nothing to apply. + + lakefs answers an empty merge with 400 `update branch : no changes`. + That is what a task produces when its output already matches the branch, + which is normal for deterministic work over unchanged input -- not a + failure, and it must not fail the task. + """ + return (getattr(exc, "status", None) == 400 + and "no changes" in str(exc)) + + +def avalon_commit_callback(context: Context, **kwargs): + client: LakeFsWrapper = init_lakefs_client(config=config) + state = read_state_file(context['ti']) + # now files have been processed, + # this part should + # get the out path of the task + local_path = str(generate_dir_name_from_task_instance( + context['ti'], + roger_config=config, + suffix='output')).rstrip('/') + '/' + task_id = context['ti'].task_id + dag_id = context['ti'].dag_id + run_id = context['ti'].run_id + # normalize run/dag/task ids for branch name + run_id_normalized = run_id.replace('-', '_').replace(':', '_').replace('+', '_').replace('.', '_') + dag_id_normalized = dag_id.replace('-', '_').replace(':', '_').replace('+', '_').replace('.', '_') + task_id_normalized = task_id.replace('-', '_').replace(':', '_').replace('+', '_').replace('.', '_') + temp_branch_name = f'{dag_id_normalized}_{task_id_normalized}_{run_id_normalized}' + remote_path = f'{dag_id}/{task_id}/' + + branch = config.lakefs_config.branch + repo = config.lakefs_config.repo + + # record source-data removals alongside outputs; rewritten every run so + # the committed manifest always reflects the latest diff window. Also + # guarantees put_files has at least one file on removal-only runs. + if state: + os.makedirs(local_path, exist_ok=True) + with open(local_path + REMOVED_FILES_MANIFEST, 'w') as f: + json.dump(state.get('removed', {}), f, indent=2) + + # real source commit(s) consumed by this task, recorded by + # setup_input_data; falls back to branch name for runs without state + # (e.g. manual repository_id overrides) + consumed = ",".join(sorted( + {e['commit_id'] for e in state.get('entries', {}).values()})) + + logger.info("Pushing local path %s to %s@%s in %s dir", + local_path, repo, temp_branch_name, remote_path) + put_files( + local_path=local_path, + remote_path=remote_path, + task_name=task_id, + task_args=[""], + pipeline_id=dag_id, + task_docker_image="docker-image", + s3storage=False, + lake_fs_client=client, + branch=temp_branch_name, + repo=repo, + commit_id=consumed or branch, + source_branch_name=branch + ) + + # tasks whose output filenames shift between runs must mirror the local + # dir, not accumulate; otherwise consumers read several runs at once + if kwargs.get('clear_output_prefix'): + existing = client.get_filelist(repository=repo, + branch=temp_branch_name, + remote_path=remote_path) + orphans = orphaned_output_paths(existing, remote_path, local_path) + if orphans: + logger.info("Deleting %d output(s) from previous runs under %s", + len(orphans), remote_path) + delete_objects(client, repo, temp_branch_name, orphans, + f"drop superseded {task_id} outputs") + + # drop derived outputs whose source files were deleted upstream; done on + # the temp branch so the removal merges atomically with this run's + # outputs and propagates to downstream tasks via their input diffs + if state.get('removed'): + bases = removed_bases(state['removed'], dag_id, + context['ti'].task.upstream_task_ids) + existing = client.get_filelist(repository=repo, + branch=temp_branch_name, + remote_path=remote_path) + stale = stale_output_paths(existing, remote_path, bases) + if stale: + logger.info("Deleting %d stale outputs for removed sources: %s", + len(stale), stale) + delete_objects(client, repo, temp_branch_name, stale, + f"remove stale {task_id} outputs " + f"for deleted sources") + + for diff in pagination_helper(client._client.refs_api.diff_refs, + repository=repo, left_ref=branch, + right_ref=temp_branch_name): + logger.info("Diff: " + str(diff)) + + try: + merge = Merge(**{"strategy": "source-wins"}) + client._client.refs_api.merge_into_branch(repository=repo, + source_ref=temp_branch_name, + destination_branch=branch, + merge=merge + ) + + logger.info(f"merged branch {temp_branch_name} into {branch}") + # only advance incremental state once outputs are safely merged; a + # failed merge leaves the Variables untouched so the next run + # re-processes the same commit window (idempotent) + _advance_state_variables(state) + except Exception as e: + if _merge_had_no_changes(e): + # lakefs 400s an empty merge. It means the output we just + # produced is byte-identical to what is already on the branch, + # which is a correct outcome, not a failure -- deterministic + # annotation over unchanged input lands here every time. State + # still advances: the input was consumed and the branch holds + # the right content. + logger.info("Nothing to merge from %s into %s; branch already " + "matches this output", temp_branch_name, branch) + _advance_state_variables(state) + else: + # never swallow anything else: the clean_up below deletes the + # local output, so a logged-and-ignored merge failure destroyed + # the work and still reported success. Raising from post_execute + # fails the task, which keeps the output (on_failure_callback + # passes keep_output=True) so the retry resumes from it. + logger.error(e) + raise + finally: + client._client.branches_api.delete_branch( + repository=repo, + branch=temp_branch_name + ) + + logger.info(f"deleted temp branch {temp_branch_name}") + logger.info(f"deleting local dir {local_path}") + + # cleanup local dirs, including any left by earlier tries whose outputs + # were hard-linked into this one and are now committed + clean_up(context, all_tries=True, **kwargs) + + +def record_state_callback(context: Context, **kwargs): + """Success callback for tasks with no lakefs output: advance the + incremental state Variables and clean local dirs.""" + state = read_state_file(context['ti']) + _advance_state_variables(state) + if state.get('removed'): + logger.warning("Upstream removals not propagated to indexes: %s", + state['removed']) + clean_up(context, all_tries=True, **kwargs) + + +def try_dir_pattern(task_instance: TaskInstance, suffix: str): + """Glob matching this task's dir for every try of the current dag run. + + generate_dir_name_from_task_instance stamps the try number into the path, + so a retry gets a fresh directory and cannot see what the previous try + produced. This is how we find the previous tries. + """ + root_data_dir = os.getenv("ROGER_DATA_DIR", "").rstrip('/') + return (f"{root_data_dir}/{task_instance.dag_id}_{task_instance.task_id}" + f"_{task_instance.run_id}_*_{suffix}") + + +def reuse_prior_try_outputs(task_instance: TaskInstance): + """Hard-link outputs from earlier tries of this task into the current one. + + Annotation is the expensive step -- tens of seconds per input file, weeks + for a large dbGaP study -- and its output is only committed to lakefs when + the whole task succeeds. So a task that died at file 40,000 of 61,597 used + to discard all 40,000 finished files, and the retry began again at zero. + + Making the finished work visible in the current try's output dir means the + pipeline's own skip check (DugPipeline.annotation_is_complete) passes over + it, and the eventual successful commit includes it. Hard links keep this + free in both time and disk; a copy is the fallback for filesystems that + refuse them (nothing in-cluster does, but the local dev path is a bind + mount). + + Returns the number of files made available. + """ + current = str(generate_dir_name_from_task_instance( + task_instance, roger_config=config, suffix='output')) + if not current: + return 0 + reused = 0 + for prior in sorted(glob.glob(try_dir_pattern(task_instance, 'output'))): + if os.path.abspath(prior) == os.path.abspath(current): + continue + for src in glob.glob(prior.rstrip('/') + '/**', recursive=True): + if not os.path.isfile(src): + continue + dest = os.path.join(current, + os.path.relpath(src, prior)) + if os.path.exists(dest): + continue + os.makedirs(os.path.dirname(dest), exist_ok=True) + try: + os.link(src, dest) + except OSError: + shutil.copy2(src, dest) + reused += 1 + if reused: + logger.info("Reused %d output file(s) from earlier tries of %s", + reused, task_instance.task_id) + return reused + + +def clean_up(context: Context, keep_output=False, all_tries=False, **kwargs): + """Remove this task's local input and output dirs. + + :param keep_output: leave the output dir in place. Used on failure, so a + retry can pick up the work already done (see + reuse_prior_try_outputs); the successful commit cleans it up. + :param all_tries: also remove the dirs left behind by earlier tries of + this task, whose contents have by then been hard-linked into this + try's output dir and committed. + """ + task_instance = context['ti'] + input_dir = str(generate_dir_name_from_task_instance( + task_instance, roger_config=config, suffix='input')).rstrip('/') + '/' + output_dir = str(generate_dir_name_from_task_instance( + task_instance, roger_config=config, suffix='output')).rstrip('/') + '/' + + dirs_to_clean = [input_dir] + if all_tries: + dirs_to_clean += glob.glob(try_dir_pattern(task_instance, 'input')) + if not keep_output: + dirs_to_clean.append(output_dir) + if all_tries: + dirs_to_clean += glob.glob(try_dir_pattern(task_instance, + 'output')) + else: + logger.info("Keeping %s so a retry can resume from it", output_dir) + + files_to_clean = [] + for a_dir in dict.fromkeys(dirs_to_clean): + files_to_clean += glob.glob(a_dir.rstrip('/') + '/**', recursive=True) + files_to_clean.append(a_dir) + for f in files_to_clean: + if os.path.exists(f): + shutil.rmtree(f, ignore_errors=True) + state_file = get_state_file_path(task_instance) + if state_file and os.path.isfile(state_file): + os.remove(state_file) + + +def generate_dir_name_from_task_instance(task_instance: TaskInstance, + roger_config: RogerConfig, suffix: str): + # if lakefs is not enabled just return none so methods default to using + # local dir structure. + if not roger_config.lakefs_config.enabled: + return None + root_data_dir = os.getenv("ROGER_DATA_DIR").rstrip('/') + task_id = task_instance.task_id + dag_id = task_instance.dag_id + run_id = task_instance.run_id + try_number = task_instance.try_number + return Path( + f"{root_data_dir}/{dag_id}_{task_id}_{run_id}_{try_number}_{suffix}") + + +def setup_input_data(context: Context, exec_conf): + logger.info(""" + - Figures out the task name and id, + - find its data dependencies + - clean up and create in and out dir + - put dependency data in input dir + - if for some reason data was not found raise an exception + """) + logger.info(">>> context") + logger.info(context) + + task_instance: TaskInstance = context['ti'] + input_dir = str(generate_dir_name_from_task_instance( + task_instance, roger_config=config, suffix="input")) + os.makedirs(input_dir, exist_ok=True) + + # a retry starts in a fresh try dir; carry forward whatever earlier tries + # of this same task already finished so it resumes instead of restarting + reuse_prior_try_outputs(task_instance) + + client = init_lakefs_client(config=config) + repos = exec_conf.get('repos', []) + dag_params = context.get("params", {}) + dag_id = task_instance.dag_id + task_id = task_instance.task_id + + if dag_params.get("repository_id"): + logger.info(">>> repository_id supplied. Overriding repo.") + repos = [{ + 'repo': dag_params.get("repository_id"), + 'branch': dag_params.get("branch_name"), + 'commitid_from': dag_params.get("commitid_from"), + 'commitid_to': dag_params.get("commitid_to") + }] + + if not repos or len(repos) == 0: + branch = config.lakefs_config.branch + repo = config.lakefs_config.repo + upstream_ids = task_instance.task.upstream_task_ids + repos = [{ + 'repo': repo, + 'branch': branch, + 'path': f'{dag_id}/{upstream_id}', + 'commitid_from': None, + 'commitid_to': None + } for upstream_id in upstream_ids] + + for r in repos: + if not r.get('path'): + r['path'] = '*' + if not os.path.exists(input_dir + f'/{r["repo"]}'): + os.mkdir(input_dir + f'/{r["repo"]}') + logger.info(f"repos : {repos}") + + logger.info(">>> start of downloading data") + if dag_params.get("repository_id"): + # manual override: honor the explicit repo/branch/commit range, + # bypassing incremental state entirely + for r in repos: + logger.info("downloading %s from %s@%s to %s", + r['path'], r['repo'], r['branch'], input_dir) + get_files( + local_path=input_dir + f'/{r["repo"]}', + remote_path=r['path'], + branch=r['branch'], + repo=r['repo'], + changes_only=r.get("commitid_from") is not None, + changes_from=r.get("commitid_from"), + changes_to=r.get("commitid_to"), + lake_fs_client=client + ) + logger.info(">>> end of downloading data") + return + + incremental = (bool(dag_params.get("incremental")) + and exec_conf.get('incremental_pull', True)) + # one tip resolution / state key / diff per source ref, even when a task + # pulls several upstream paths from the same repo+branch + groups = {} + for r in repos: + groups.setdefault((r['repo'], r['branch']), []).append(r['path']) + + state = {'entries': {}, 'removed': {}} + any_change = False + for (repo, branch), prefixes in groups.items(): + local_path = input_dir + f'/{repo}' + tip = resolve_ref_tip(client, repo, branch) + key = incremental_state_key(dag_id, task_id, repo, branch) + last = _get_last_consumed(key) if incremental else None + changes = None + if last and last == tip: + logger.info("No new commits on %s@%s since %s", repo, branch, tip) + elif last: + try: + changes = get_changed_files(client, repo, last, tip, prefixes) + except NotFoundException: + logger.warning( + "Commit %s not found on %s; falling back to full " + "download", last, repo) + last = None + if last and changes is not None: + to_download = changes['added'] + changes['changed'] + logger.info("Incremental %s@%s %s..%s: %d added, %d changed, " + "%d removed", repo, branch, last, tip, + len(changes['added']), len(changes['changed']), + len(changes['removed'])) + if to_download: + siblings = find_sibling_files(client, repo, tip, to_download) + if siblings: + logger.info("Adding %d sibling file(s): %s", + len(siblings), siblings) + to_download = to_download + siblings + client.download_files( + remote_files=to_download, + local_path=local_path, + repository=repo, + branch_or_commit_id=tip) + if changes['removed']: + state['removed'].setdefault( + repo, []).extend(changes['removed']) + if to_download or changes['removed']: + any_change = True + elif not last: + # first run, incremental disabled, or unreachable last commit: + # full download pinned to the resolved tip + for prefix in prefixes: + logger.info("downloading %s from %s@%s to %s", + prefix, repo, tip, input_dir) + get_files( + local_path=local_path, + remote_path=prefix, + branch=tip, + repo=repo, + changes_only=False, + lake_fs_client=client + ) + any_change = True + state['entries'][key] = { + 'repo': repo, 'branch': branch, 'commit_id': tip} + logger.info(">>> end of downloading data") + + if incremental and not any_change: + raise AirflowSkipException( + "No changes in source refs since last successful run") + write_state_file(task_instance, state) + + +def create_python_task(dag, name, a_callable, func_kwargs=None, + external_repos=None, pass_conf=True, + no_output_files=False, no_input_files=False, + incremental_pull=True, clear_output_prefix=False, + memory=None, resumable=False): + """ Create a python task. + :param func_kwargs: additional arguments for callable. + :param dag: dag to add task to. + :param name: The name of the task. + :param a_callable: The code to run in this task. + :param no_input_files: skip the lakefs input download entirely. + :param resumable: keep the output dir when the task fails, so a retry + can pick up where it left off. Only true for annotate, which is + the one task with a skip check (annotation_is_complete). For the + others the retained output is never reused and is pure disk cost: + three failed crawls held 47GB and filled the shared volume, which + is what made them fail in the first place. + :param incremental_pull: when False the task always downloads its full + inputs even if the dag runs with incremental=True (needed for tasks + that rebuild state from scratch, e.g. ES indexing after a wipe). + :param clear_output_prefix: mirror the local output dir into lakefs, + deleting objects from previous runs. Needed when output filenames + vary run to run (the bulk-load CSVs) and would otherwise accumulate. + :param memory: memory limit for this task's pod, e.g. '15Gi'. Omit to + take the chart's worker default. + """ + + if external_repos is None: + external_repos = {} + + # these are actual arguments passed down to the task function + op_kwargs = { + "python_callable": a_callable, + "to_string": True, + "pass_conf": pass_conf + } + if func_kwargs is None: + func_kwargs = {} + op_kwargs.update(func_kwargs) + + python_operator_args = { + "task_id": name, + "python_callable": task_wrapper, + # executor_config example left commented; fill if needed + "dag": dag, + } + if memory: + python_operator_args["executor_config"] = memory_override(memory) + + if config.lakefs_config.enabled: + pre_exec_conf = { + 'repos': [], + 'incremental_pull': incremental_pull + } + if external_repos: + pre_exec_conf['repos'] = [{ + 'repo': r['name'], + 'branch': r['branch'], + 'path': r.get('path', '*') + } for r in external_repos] + + if not no_input_files: + pre_exec = partial(setup_input_data, exec_conf=pre_exec_conf) + # pre_execute will be called with context -> partial keeps exec_conf fixed + python_operator_args['pre_execute'] = pre_exec + + # pass fixed kwargs into partials so resulting callback accepts (context,) + # keep the output dir on failure: the next try hard-links it in and + # skips the work already done (annotation is the expensive step) + python_operator_args['on_failure_callback'] = partial( + clean_up, keep_output=resumable, **op_kwargs) + # pre_execute creates the input dir before it can raise + # AirflowSkipException; clean it up on skip too + python_operator_args['on_skipped_callback'] = partial(clean_up, **op_kwargs) + # post_execute, not on_success_callback. Airflow runs post_execute + # inside _execute_task, before it records end_date and releases + # downstream; success callbacks run in finalize(), after. Committing + # there meant downstream tasks read the branch before the output + # landed -- BulkLoad loaded an edgeless graph 105s early, and + # make_kgx built kgx from annotations 4.8h stale. Airflow also + # swallows exceptions from state-change callbacks (it only logs + # them), so a failed upload or merge used to leave the task green; + # from post_execute it fails the task instead. + if not no_output_files: + commit = partial(avalon_commit_callback, + clear_output_prefix=clear_output_prefix, + **op_kwargs) + else: + commit = partial(record_state_callback, **op_kwargs) + # the hook is called as (context, result); our callbacks take context + python_operator_args['post_execute'] = ( + lambda context, result=None, _c=commit: _c(context)) + + python_operator_args["op_kwargs"] = op_kwargs + + return PythonOperator(**python_operator_args) + +def execute_pipeline_method(pipeline_class, configparam, method_name, + input_data_path=None, output_data_path=None, + to_string=False, **pipeline_kwargs): + """ + Lazy execution wrapper. + Initializes the heavy pipeline class and executes the method ONLY inside the K8s worker pod. + """ + logger.info(f"Initializing {pipeline_class.__name__} for method {method_name}") + + # 1. The class initialization happens safely here, ignored by the Scheduler + with pipeline_class(config=configparam, **pipeline_kwargs) as pipeline: + + # 2. Grab the requested method (e.g., pipeline.annotate) + method_to_call = getattr(pipeline, method_name) + + # 3. Run it with the Airflow context args + return method_to_call(input_data_path=input_data_path, + output_data_path=output_data_path, + to_string=to_string) + + + +def file_task_group_id(name): + return f"{name}_dataset_pipeline_task_group" + + +def create_pipeline_taskgroup( + dag, + pipeline_class: type, + configparam: RogerConfig, + **kwargs): + """Emit the file-based (lakefs-committed, incremental) task group for + the specified pipeline_class: annotate -> crawl, make_kgx. + + ES indexing/validation lives in create_es_taskgroup and runs after a + global index wipe, rebuilding elastic from the files left in lakefs.""" + name = pipeline_class.pipeline_name + input_dataset_version = pipeline_class.input_version + + with TaskGroup(group_id=file_task_group_id(name)) as tg: + + # --- 1. Annotate Task --- + annotate_callable = partial( + execute_pipeline_method, + pipeline_class=pipeline_class, + configparam=configparam, + method_name='annotate', + **kwargs + ) + annotate_task = create_python_task( + dag, + f"annotate_{name}_files", + annotate_callable, + external_repos=[{ + 'name': getattr(pipeline_class, 'pipeline_name'), + 'branch': input_dataset_version + }], + # annotate_workers threads each hold a whole parsed file; the + # chart default was sized for the serial annotator + memory=configparam.annotation.annotate_memory, + # the only task that can resume: annotation_is_complete skips + # input files whose output already exists + resumable=True, + pass_conf=False) + + # --- 2. Make KGX Task --- + make_kgx_callable = partial( + execute_pipeline_method, + pipeline_class=pipeline_class, + configparam=configparam, + method_name='make_kg_tagged', + **kwargs + ) + make_kgx_task = create_python_task( + dag, + f"make_kgx_{name}", + make_kgx_callable, + # holds every element of the dataset in memory to build the kgx; + # OOMKilled at the 2Gi chart default on bdc-recover + memory=configparam.annotation.annotate_memory, + pass_conf=False) + make_kgx_task.set_upstream(annotate_task) + + # --- 3. Crawl Task --- + crawl_callable = partial( + execute_pipeline_method, + pipeline_class=pipeline_class, + configparam=configparam, + method_name='crawl_tranql', + **kwargs + ) + crawl_task = create_python_task( + dag, + f"crawl_{name}", + crawl_callable, + # expands every concept through tranql, accumulating answers + memory=configparam.annotation.annotate_memory, + pass_conf=False) + crawl_task.set_upstream(annotate_task) + + # --- 4. Complete Task --- + # none_failed: a group skipped for "no new data" still completes + # green; genuine failures still propagate as upstream_failed + complete_task = EmptyOperator(task_id=f"complete_{name}", + trigger_rule="none_failed") + complete_task.set_upstream((make_kgx_task, crawl_task)) + + return tg + + +def create_es_wipe_task(dag, pipeline_class: type, configparam: RogerConfig, + **kwargs): + """Single task that wipes all ES indexes before the per-dataset rebuild + tasks repopulate them from lakefs. Any pipeline class works: index names + come from shared config.""" + wipe_callable = partial( + execute_pipeline_method, + pipeline_class=pipeline_class, + configparam=configparam, + method_name='clear_all_es_indexes', + **kwargs + ) + return create_python_task( + dag, + "wipe_es_indexes", + wipe_callable, + pass_conf=False, + no_output_files=True, + no_input_files=True) + + +def create_es_taskgroup( + dag, + pipeline_class: type, + configparam: RogerConfig, + **kwargs): + """Emit the elastic task group for a pipeline: full (non-incremental) + pulls of this run's file outputs from lakefs, indexed into the freshly + wiped indexes so ES always mirrors what lakefs holds.""" + name = pipeline_class.pipeline_name + repo = config.lakefs_config.repo + branch = config.lakefs_config.branch + file_group = file_task_group_id(name) + # Everything ES needs lives in the crawl output: expanded_concepts.txt + # AND an elements.txt whose optional_terms carry the KG-derived search + # terms (crawl_tranql rewrites it after expanding concepts). The + # annotate copy has empty optional_terms, so indexing that one makes + # validate_indexed_concepts unsatisfiable -- it searches by KG node + # name. storage.dug_expanded_elements_objects globs "**/elements.txt" + # and cannot tell the two apart, so only the crawl prefix is pulled. + crawl_path = f"{dag.dag_id}/{file_group}.crawl_{name}" + + def full_pull(*paths): + return { + 'external_repos': [ + {'name': repo, 'branch': branch, 'path': p} for p in paths], + 'pass_conf': False, + 'no_output_files': True, + 'incremental_pull': False, + } + + with TaskGroup(group_id=f"{name}_es_index_task_group") as tg: + + index_variables_task = create_python_task( + dag, + f"index_{name}_variables", + partial(execute_pipeline_method, + pipeline_class=pipeline_class, + configparam=configparam, + method_name='index_variables', + **kwargs), + **full_pull(crawl_path)) + + validate_index_variables_task = create_python_task( + dag, + f"validate_{name}_index_variables", + partial(execute_pipeline_method, + pipeline_class=pipeline_class, + configparam=configparam, + method_name='validate_indexed_variables', + **kwargs), + **full_pull(crawl_path)) + validate_index_variables_task.set_upstream(index_variables_task) + + index_concepts_task = create_python_task( + dag, + f"index_{name}_concepts", + partial(execute_pipeline_method, + pipeline_class=pipeline_class, + configparam=configparam, + method_name='index_concepts', + **kwargs), + **full_pull(crawl_path)) + + validate_index_concepts_task = create_python_task( + dag, + f"validate_{name}_index_concepts", + partial(execute_pipeline_method, + pipeline_class=pipeline_class, + configparam=configparam, + method_name='validate_indexed_concepts', + **kwargs), + **full_pull(crawl_path)) + validate_index_concepts_task.set_upstream(index_concepts_task) + # it asserts variables are findable by KG-derived terms, so it + # searches the variables index -- which must be fully populated + # first, not just the concepts index + validate_index_concepts_task.set_upstream(index_variables_task) + + return tg + + +# ponytail: prefixes string-coupled to create_pipeline_taskgroup's group/task +# naming; breaks only if annotate_and_index renames its taskgroup/tasks. +ANNOTATE_DAG_ID = "annotate_and_index" + + +def _annotate_index_source_path(name: str, task: str) -> str: + """Runtime-repo prefix where annotate_and_index committed a task's output: + {dag_id}/{group_id}.{task_id}/""" + group = f"{name}_dataset_pipeline_task_group" + return f"{ANNOTATE_DAG_ID}/{group}.{task}" + + +def create_index_only_taskgroup( + dag, + pipeline_class: type, + configparam: RogerConfig, + **kwargs): + """Re-index ES from annotate_and_index outputs already present in the + runtime repo (e.g. after merging dev->prod). Pulls annotate/crawl outputs + by explicit path from the configured runtime repo+branch; no annotate or + crawl is re-run.""" + name = pipeline_class.pipeline_name + repo = configparam.lakefs_config.repo + branch = configparam.lakefs_config.branch + + def index_task(task_name, method_name, source_tasks): + callable_ = partial( + execute_pipeline_method, + pipeline_class=pipeline_class, + configparam=configparam, + method_name=method_name, + **kwargs) + return create_python_task( + dag, task_name, callable_, + external_repos=[{ + 'name': repo, + 'branch': branch, + 'path': _annotate_index_source_path(name, src), + } for src in source_tasks], + pass_conf=False, + no_output_files=True) + + with TaskGroup(group_id=f"{name}_index_only_task_group") as tg: + # all four read crawl output; see create_es_taskgroup for why + crawl_src = f"crawl_{name}" + + index_variables_task = index_task( + f"index_{name}_variables", 'index_variables', [crawl_src]) + validate_variables_task = index_task( + f"validate_{name}_index_variables", + 'validate_indexed_variables', [crawl_src]) + validate_variables_task.set_upstream(index_variables_task) + + index_concepts_task = index_task( + f"index_{name}_concepts", 'index_concepts', [crawl_src]) + validate_concepts_task = index_task( + f"validate_{name}_index_concepts", + 'validate_indexed_concepts', [crawl_src]) + validate_concepts_task.set_upstream(index_concepts_task) + # searches the variables index too; see create_es_taskgroup + validate_concepts_task.set_upstream(index_variables_task) + + complete_task = EmptyOperator(task_id=f"complete_{name}", + trigger_rule="none_failed") + complete_task.set_upstream( + (validate_variables_task, validate_concepts_task)) + + return tg diff --git a/src/roger/utils/__init__.py b/src/roger/utils/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/src/roger/utils/batched_annotator.py b/src/roger/utils/batched_annotator.py new file mode 100644 index 00000000..8a57610b --- /dev/null +++ b/src/roger/utils/batched_annotator.py @@ -0,0 +1,178 @@ +"""Batch the per-identifier lookups dug makes while annotating an element. + +dug resolves identifiers one at a time (`AnnotateSapbert.__call__`): + + for entity, raw_identifiers in raw_identifiers_dict.items(): + for identifier in raw_identifiers: + norm_id = self.normalizer(identifier, http_session) + norm_id.synonyms = self.synonym_finder(norm_id.id, http_session) + +Both services accept a list and answer it in about the time they take to +answer one. Measured against the deployed translator-dev services, per curie: + + n=1 normalize 10.34 ms synonyms 25.82 ms + n=200 normalize 0.13 ms synonyms 0.13 ms + +An element with 150 identifiers spends ~5.4 s in that loop and ~50 ms +batched. The two phases stay serial with respect to each other, because +synonyms are looked up by the *normalized* curie and so normalization has to +finish first -- but that is 2 requests instead of 300. + +sapbert is deliberately left alone: the service takes a single `text` and +422s on any list form, and it is called once per classified entity rather +than once per identifier, so it is not what makes the loop expensive. +""" + +from urllib.parse import urlsplit, parse_qs + +from dug.core.annotators.utils.biolink_purl_util import BioLinkPURLerizer + +from roger.logger import get_logger + +log = get_logger() + +# 200 measured fine on both services; the curve is flat well before this, so +# this is about bounding request size, not about finding an optimum. +BATCH_SIZE = 200 + +# The node normalizer's GET query defaults. dug builds a GET url and only +# sets some of these, so the rest are whatever GET defaults to -- and GET +# and POST do not agree (drug_chemical_conflate is true on GET, false on +# POST). Sending the full set keeps the batched POST answering exactly what +# dug's GET answers. +GET_DEFAULTS = { + "conflate": True, + "drug_chemical_conflate": True, + "description": False, + "individual_types": False, + "include_taxa": True, +} + + +def _chunks(items, size): + for i in range(0, len(items), size): + yield items[i:i + size] + + +class BatchedAnnotator: + """Wraps dug's sapbert annotator, batching normalize and synonym lookups. + + Everything except the identifier loop is delegated to the wrapped + annotator, including response parsing -- `handle_response` on dug's own + normalizer and synonym finder is reused so the greenlist behaviour, the + biolink type coercion and the in-place mutation of DugIdentifier stay + exactly as they are today. + + If either batch call fails the whole element falls back to dug's serial + path. Classification and sapbert responses are cached by then, so the + fallback re-runs cheaply. + """ + + def __init__(self, annotator): + self._inner = annotator + + def __getattr__(self, name): + # urls, thresholds, bagel config: whatever we do not override is read + # straight off the wrapped annotator + return getattr(self._inner, name) + + def __call__(self, text, http_session): + inner = self._inner + classifiers = inner.text_classification(text, http_session) + raw = inner.annotate_classifiers(classifiers, http_session) + if not raw: + log.warning("Failed to annotate: %s", text) + return [] + + identifiers = [i for ids in raw.values() for i in ids] + try: + normalized = self._normalize_batch(identifiers, http_session) + except Exception as e: + log.warning("Batch normalize failed, falling back to dug's " + "per-identifier path: %s", e) + return inner(text, http_session) + + # Phase 1 -- normalize. dug's own handle_response does the parsing, so + # a curie missing from the batch response yields None exactly as a + # failed single lookup would. + kept = {} + for entity, raw_identifiers in raw.items(): + for identifier in raw_identifiers: + norm_id = inner.normalizer.handle_response(identifier, + normalized) + if norm_id is None: + log.warning("Failed to normalize: %s", identifier.id) + if identifier.id_type not in inner.ontology_greenlist: + continue + norm_id = identifier + norm_id.purl = BioLinkPURLerizer.get_curie_purl(norm_id.id) + kept.setdefault(entity, []).append(norm_id) + + # Phase 2 -- synonyms, keyed by the normalized curie, which is why + # this cannot be folded into the pass above. + wanted = [i.id for ids in kept.values() for i in ids] + try: + synonyms = self._synonyms_batch(wanted, http_session) + except Exception as e: + log.warning("Batch synonym lookup failed, falling back to dug's " + "per-identifier path: %s", e) + return inner(text, http_session) + for ids in kept.values(): + for identifier in ids: + identifier.synonyms = inner.synonym_finder.handle_response( + identifier.id, synonyms) + + if inner.bagel_enabled: + # matches dug: bagel runs per classified entity, including ones + # where nothing survived normalization + for entity in raw: + kept[entity] = inner.bagel(description_text=text, + entity=entity, + ids=kept.get(entity, []), + http_session=http_session) + + return [i for ids in kept.values() for i in ids] + + def _normalizer_endpoint(self): + """(url, flags) for the normalizer's POST form. + + dug stores the normalizer as a GET url ending in `curie=`. The same + service answers POST on the same path with {"curies": [...]}, and the + query flags move into the body. + + The flags have to be sent in full, because the service does not + default them the same way on both verbs -- notably + drug_chemical_conflate defaults to true on GET and false on POST. + Omitting it silently stopped conflating drug/chemical curies: + CHEBI:3759 came back as itself where dug's GET resolved it to + CHEBI:37941. So start from the GET defaults, which are what dug + gets today, and let the url override them. + """ + split = urlsplit(self._inner.normalizer.url) + query = parse_qs(split.query) + flags = dict(GET_DEFAULTS) + for name in flags: + if name in query: + flags[name] = str(query[name][0]).strip().lower() == "true" + url = f"{split.scheme}://{split.netloc}{split.path}" + return url, flags + + def _normalize_batch(self, identifiers, http_session): + curies = list(dict.fromkeys(i.id for i in identifiers)) + url, flags = self._normalizer_endpoint() + out = {} + for chunk in _chunks(curies, BATCH_SIZE): + response = http_session.post(url, json=dict(flags, curies=chunk)) + response.raise_for_status() + out.update(response.json() or {}) + return out + + def _synonyms_batch(self, curies, http_session): + url = self._inner.synonym_finder.url + out = {} + for chunk in _chunks(list(dict.fromkeys(curies)), BATCH_SIZE): + response = http_session.post(url, + json={"preferred_curies": chunk}) + response.raise_for_status() + out.update(response.json() or {}) + return out diff --git a/src/roger/utils/http_utils.py b/src/roger/utils/http_utils.py new file mode 100644 index 00000000..513c0451 --- /dev/null +++ b/src/roger/utils/http_utils.py @@ -0,0 +1,141 @@ +"Hardening for the outbound HTTP session used by the annotation pipeline" + +import socket + +from requests.adapters import HTTPAdapter +from urllib3.connection import HTTPConnection +from urllib3.util.retry import Retry + +from roger.logger import get_logger + +log = get_logger() + +# The annotation services are read-only lookups, so replaying a POST is safe. +RETRY_METHODS = frozenset(["GET", "POST"]) +RETRY_STATUSES = (429, 502, 503, 504) + +# Keepalive probing: start after 60s idle, then every 15s, give up after 4. +KEEPALIVE_IDLE = 60 +KEEPALIVE_INTERVAL = 15 +KEEPALIVE_COUNT = 4 + + +def keepalive_socket_options(): + """Socket options enabling TCP keepalive on top of urllib3's defaults. + + A socket with no keepalive has no timer of any kind, so a peer that + disappears without sending a FIN or RST is invisible to the client and a + blocking read waits forever. Keepalive gives the kernel its own way to + notice, independent of any application level timeout. + """ + options = list(HTTPConnection.default_socket_options) + options.append((socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1)) + # These three are Linux only. Dev machines are often macOS, where the + # names are absent or spelled differently, so probe rather than assume. + for name, value in (("TCP_KEEPIDLE", KEEPALIVE_IDLE), + ("TCP_KEEPINTVL", KEEPALIVE_INTERVAL), + ("TCP_KEEPCNT", KEEPALIVE_COUNT)): + option = getattr(socket, name, None) + if option is not None: + options.append((socket.IPPROTO_TCP, option, value)) + return options + + +class TimeoutRetryAdapter(HTTPAdapter): + """Adapter applying a default timeout and TCP keepalive to every request. + + requests has no session level timeout, so any call site that omits one + blocks indefinitely. Defaulting it here covers all of them, including the + call sites inside dug, none of which pass a timeout. + """ + + def __init__(self, *args, timeout=None, **kwargs): + self._timeout = timeout + super().__init__(*args, **kwargs) + + def init_poolmanager(self, *args, **kwargs): + kwargs.setdefault('socket_options', keepalive_socket_options()) + super().init_poolmanager(*args, **kwargs) + + def send(self, request, **kwargs): + if kwargs.get('timeout') is None: + kwargs['timeout'] = self._timeout + return super().send(request, **kwargs) + + +def harden_session(session, connect_timeout, read_timeout, retries, + backoff_factor): + """Mount a timeout/retry/keepalive adapter on `session` and return it. + + Annotation talks to several in-cluster services over long lived + keep-alive connections. When one of those connections breaks silently, + with the peer gone but no FIN or RST arriving, an unbounded read hangs + until a human notices, which has taken days. A read timeout turns that + into an error, and the retries absorb the transient case so that a single + blip does not fail a crawl that takes hours. + + Retries matter here specifically because dug's Crawler.annotate_elements + has no per-element error handling, so an exception that escapes discards + the progress of the whole run. + """ + retry = Retry( + total=retries, + connect=retries, + read=retries, + status=retries, + backoff_factor=backoff_factor, + status_forcelist=RETRY_STATUSES, + allowed_methods=RETRY_METHODS, + raise_on_status=False, + ) + adapter = TimeoutRetryAdapter( + max_retries=retry, + timeout=(connect_timeout, read_timeout), + ) + session.mount('http://', adapter) + session.mount('https://', adapter) + log.info( + "HTTP session hardened: connect_timeout=%ss read_timeout=%ss " + "retries=%s backoff_factor=%s", + connect_timeout, read_timeout, retries, backoff_factor) + return session + + +# The annotation services are all read-only lookups whose answer depends only +# on the request body, so a POST is cacheable in exactly the way a GET is. +CACHEABLE_METHODS = ('GET', 'HEAD', 'POST') + + +def enable_post_caching(session, expire_seconds=0): + """Let `session` cache POST responses, and report whether it can. + + dug builds the annotation session with requests_cache, but requests_cache + caches only GET and HEAD unless told otherwise. Three of the four + annotation calls -- token classification, sapbert, synonym lookup -- are + POSTs and so never hit the cache. Only node normalization is a GET + (DefaultNormalizer.make_request) and was already being cached. + + The cost of that is not marginal. dbGaP parsers emit the study element + into every one of a study's data-dict files, so annotating bdc-parent + re-ran the same 24 study descriptions 61,597 times. The study element + takes ~53s to annotate against ~1.4s for the variable the file actually + contributes, which is 96% of a 39-day run spent recomputing 24 answers. + + Setting expire_seconds also bounds the normalizer GETs, which dug cached + with no expiry at all. See AnnotationConfig for why that matters to a + redis shared with the graph. + + Returns True if caching was turned on. A plain requests.Session has no + `settings`, which is the no-lakefs/no-redis test path, and is left alone. + """ + settings = getattr(session, 'settings', None) + if settings is None: + log.warning("HTTP session is not a CachedSession; annotation " + "responses will not be cached") + return False + settings.allowable_methods = CACHEABLE_METHODS + if expire_seconds: + settings.expire_after = expire_seconds + log.info("HTTP response cache enabled for %s (expire_after=%s)", + ",".join(CACHEABLE_METHODS), settings.expire_after) + return True diff --git a/src/roger/utils/s3_utils.py b/src/roger/utils/s3_utils.py new file mode 100644 index 00000000..f0f7277b --- /dev/null +++ b/src/roger/utils/s3_utils.py @@ -0,0 +1,45 @@ +from contextlib import contextmanager + +import boto3 + +from roger.config import S3Config + + +class S3Utils: + + def __init__( + self, + s3_config: S3Config + ): + self.config = s3_config + + @contextmanager + def connect( + self, + ): + session = boto3.session.Session( + aws_access_key_id=self.config.access_key, + aws_secret_access_key=self.config.secret_key, + ) + + s3 = session.resource( + 's3', + endpoint_url=self.config.host, + ) + bucket = s3.Bucket(self.config.bucket) + yield bucket + + def get(self, remote_file_name: str, local_file_name: str): + with self.connect() as bucket: + bucket.download_file(remote_file_name, local_file_name) + + def put(self, local_file_name: str, remote_file_name: str): + with self.connect() as bucket: + bucket.upload_file(local_file_name, remote_file_name) + + def ls(self): + with self.connect() as bucket: + return [ + obj + for obj in bucket.objects.all() + ] diff --git a/tests/__init__.py b/tests/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/conftest.py b/tests/conftest.py new file mode 100644 index 00000000..bdc954cb --- /dev/null +++ b/tests/conftest.py @@ -0,0 +1 @@ +pythonpath = "dags" diff --git a/tests/integration/__init__.py b/tests/integration/__init__.py new file mode 100644 index 00000000..32e25657 --- /dev/null +++ b/tests/integration/__init__.py @@ -0,0 +1,3 @@ +from pathlib import Path + +TEST_DATA_DIR = (Path(__file__).parent / 'data').resolve() \ No newline at end of file diff --git a/tests/integration/conftest.py b/tests/integration/conftest.py new file mode 100644 index 00000000..edb68d21 --- /dev/null +++ b/tests/integration/conftest.py @@ -0,0 +1,86 @@ +import os +import glob +import json + +from roger.core.enums import SchemaType + +class BiolinkMock: + def __init__(self): + self.leafs = [ + 'chemical_substance', + 'molecular_activity', + 'gene', + 'biological_process', + 'disease', + 'phenotypic_feature' + ] + + def get_leaf_class(self, class_names): + for y in self.leafs: + if y in class_names: + return y + return class_names[0] + + def find_biolink_leaves(self, biolink_concepts): + return set([concept for concept in biolink_concepts + if concept in ['named_thing', 'some_other_type']]) + +category = None +predicates = None +file_content_assertions = {} +kgx_files = [] +merged_files = [] +merge_file_test_dir = '' +schema = { + SchemaType.PREDICATE: {}, + SchemaType.CATEGORY: {} +} + +def kgx_objects(): + return [os.path.join(*os.path.split(__file__)[:-1], 'data', file) + for file in kgx_files] + +def merged_objects(file_type, path=None): + """ A list of merged KGX objects. """ + if not path: + merged_pattern = merge_path(f"**/{file_type}.jsonl") + else: + merged_pattern = merge_path(f"**/{file_type}.jsonl", path=path) + # this thing should always return one edges or nodes file (based on file_type) + try: + return sorted(glob.glob(merged_pattern, recursive=True))[0] + except IndexError: + raise ValueError(f"Could not find merged KGX of type {file_type} " + f"in {merged_pattern}") + +def bulk_path(*args, **kwargs): + return os.path.join(*os.path.split(__file__)[:-1], 'data', 'bulk') + +def is_up_to_date(*args, **kwargs): + return False + +def schema_path(name, *args, **kwargs): + return name + +def read_schema(schema_type: SchemaType, *args, **kwargs): + return conftest.schema[schema_type] + +def read_object(path, *args, **kwargs): + import json + with open(path) as f: + return json.load(f) + +def write_object(dictionary, file_name): + print(dictionary, file_name) + print(file_content_assertions) + assert file_content_assertions[file_name] == dictionary + +def merge_path(file_name): + return os.path.join(*os.path.split(__file__)[:-1], 'data', 'merge', + merge_file_test_dir, file_name) + +def json_line_iter(jsonl_file_path): + f = open(file=jsonl_file_path, mode='r') + for line in f: + yield json.loads(line) + f.close() diff --git a/tests/integration/data/merge/conflicting_prop_types__edges__schema__kgx/edges.jsonl b/tests/integration/data/merge/conflicting_prop_types__edges__schema__kgx/edges.jsonl new file mode 100644 index 00000000..e0477a06 --- /dev/null +++ b/tests/integration/data/merge/conflicting_prop_types__edges__schema__kgx/edges.jsonl @@ -0,0 +1,2 @@ +{"id": "edge_1", "edge_label": "edge_type_1", "list_vs_str": [], "list_vs_int": [], "list_vs_bool": [], "list_vs_float": [], "str_vs_float": "", "str_vs_bool": "", "str_vs_int": "", "int_vs_bool": 0, "int_vs_float": 0, "float_vs_bool": 0, "predicate": "related_to"} +{"id": "edge_2", "edge_label": "edge_type_1", "list_vs_str": "", "list_vs_int": 0, "list_vs_bool": true, "list_vs_float": 0.0, "str_vs_float": 0.0, "str_vs_bool": false, "str_vs_int": 0, "int_vs_bool": true, "int_vs_float": 0.0, "float_vs_bool": true, "predicate": "related_to" } \ No newline at end of file diff --git a/tests/integration/data/merge/conflicting_prop_types__edges__schema__kgx/expected.json b/tests/integration/data/merge/conflicting_prop_types__edges__schema__kgx/expected.json new file mode 100644 index 00000000..49c8e2b3 --- /dev/null +++ b/tests/integration/data/merge/conflicting_prop_types__edges__schema__kgx/expected.json @@ -0,0 +1,20 @@ +{ + "predicate-schema.json": { + "related_to": { + "id": "str", + "edge_label": "str", + "list_vs_str": "list", + "list_vs_int": "list", + "list_vs_bool": "list", + "list_vs_float": "list", + "str_vs_float": "str", + "str_vs_bool": "str", + "str_vs_int": "str", + "int_vs_bool": "str", + "int_vs_float": "str", + "float_vs_bool": "str", + "predicate": "str" + } + }, + "category-schema.json": {} +} \ No newline at end of file diff --git a/tests/integration/data/merge/conflicting_prop_types__edges__schema__kgx/nodes.jsonl b/tests/integration/data/merge/conflicting_prop_types__edges__schema__kgx/nodes.jsonl new file mode 100644 index 00000000..e69de29b diff --git a/tests/integration/data/merge/conflicting_prop_types__nodes__schema__kgx/edges.jsonl b/tests/integration/data/merge/conflicting_prop_types__nodes__schema__kgx/edges.jsonl new file mode 100644 index 00000000..e69de29b diff --git a/tests/integration/data/merge/conflicting_prop_types__nodes__schema__kgx/expected.json b/tests/integration/data/merge/conflicting_prop_types__nodes__schema__kgx/expected.json new file mode 100644 index 00000000..7d7979b5 --- /dev/null +++ b/tests/integration/data/merge/conflicting_prop_types__nodes__schema__kgx/expected.json @@ -0,0 +1,19 @@ +{ + "category-schema.json": { + "named_thing": { + "id": "str", + "category": "list", + "list_vs_str": "list", + "list_vs_int": "list", + "list_vs_bool": "list", + "list_vs_float": "list", + "str_vs_float": "str", + "str_vs_bool": "str", + "str_vs_int": "str", + "int_vs_bool": "str", + "int_vs_float": "str", + "float_vs_bool": "str" + } + }, + "predicate-schema.json": {} +} \ No newline at end of file diff --git a/tests/integration/data/merge/conflicting_prop_types__nodes__schema__kgx/nodes.jsonl b/tests/integration/data/merge/conflicting_prop_types__nodes__schema__kgx/nodes.jsonl new file mode 100644 index 00000000..ed5aae5d --- /dev/null +++ b/tests/integration/data/merge/conflicting_prop_types__nodes__schema__kgx/nodes.jsonl @@ -0,0 +1,2 @@ +{"id": "node_1", "category": ["named_thing"], "list_vs_str": [], "list_vs_int": [], "list_vs_bool": [], "list_vs_float": [], "str_vs_float": "", "str_vs_bool": "", "str_vs_int": "", "int_vs_bool": 0, "int_vs_float": 0, "float_vs_bool": 0} +{"id": "node_1", "category": ["named_thing"], "list_vs_str": "", "list_vs_int": 0, "list_vs_bool": true, "list_vs_float": 0.0, "str_vs_float": 0.0, "str_vs_bool": false, "str_vs_int": 0, "int_vs_bool": true, "int_vs_float": 0.0, "float_vs_bool": true} \ No newline at end of file diff --git a/tests/integration/data/merge/non_conflicting_prop_types__schema__kgx/edges.jsonl b/tests/integration/data/merge/non_conflicting_prop_types__schema__kgx/edges.jsonl new file mode 100644 index 00000000..63ab7769 --- /dev/null +++ b/tests/integration/data/merge/non_conflicting_prop_types__schema__kgx/edges.jsonl @@ -0,0 +1,4 @@ +{"edge_label": "edge_label_1", "id": "x", "bool_attr": false, "int_attr": 0, "float_attr": 0.0 , "predicate": "edge_label_1"} +{"edge_label": "edge_label_1", "id": "x3", "str_attr": "str", "list_attr": [], "predicate": "edge_label_1"} +{"edge_label": "edge_label_2", "id": "x4", "str_attr": "str", "predicate": "edge_label_2"} +{"edge_label": "edge_label_2", "id": "x3", "bool_attr": true, "float_attr": 2.33, "int_attr": 3092, "str_att": "name", "predicate": "edge_label_2"} \ No newline at end of file diff --git a/tests/integration/data/merge/non_conflicting_prop_types__schema__kgx/expected.json b/tests/integration/data/merge/non_conflicting_prop_types__schema__kgx/expected.json new file mode 100644 index 00000000..6fcf9cd5 --- /dev/null +++ b/tests/integration/data/merge/non_conflicting_prop_types__schema__kgx/expected.json @@ -0,0 +1,43 @@ +{ + "category-schema.json": { + "named_thing": { + "str_attr": "str", + "list_attr": "list", + "bool_attr": "bool", + "int_attr": "int", + "float_attr": "float", + "id": "str", + "category": "list" + }, + "some_other_type": { + "id": "str", + "category": "list", + "attr_1": "str", + "attr_2": "list", + "attr_3": "bool", + "attr_4": "int" + } + }, + "predicate-schema.json": { + "edge_label_1": { + "id": "str", + "edge_label": "str", + "str_attr": "str", + "list_attr": "list", + "bool_attr": "bool", + "int_attr": "int", + "float_attr": "float", + "predicate": "str" + }, + "edge_label_2": { + "id": "str", + "str_attr": "str", + "edge_label": "str", + "bool_attr": "bool", + "float_attr": "float", + "int_attr": "int", + "str_att": "str", + "predicate": "str" + } + } +} \ No newline at end of file diff --git a/tests/integration/data/merge/non_conflicting_prop_types__schema__kgx/nodes.jsonl b/tests/integration/data/merge/non_conflicting_prop_types__schema__kgx/nodes.jsonl new file mode 100644 index 00000000..1670a2be --- /dev/null +++ b/tests/integration/data/merge/non_conflicting_prop_types__schema__kgx/nodes.jsonl @@ -0,0 +1,3 @@ +{"id": "ID1", "category": ["named_thing"], "list_attr": [], "bool_attr": false, "int_attr": 0} +{"id": "ID2", "category": ["named_thing"], "str_attr": "", "float_attr": 0.0} +{"id": "Id3", "category": ["some_other_type"], "attr_1": "", "attr_2": [], "attr_3": true, "attr_4": 1} \ No newline at end of file diff --git a/tests/integration/test_KGX_Model.py b/tests/integration/test_KGX_Model.py new file mode 100644 index 00000000..4fc5716f --- /dev/null +++ b/tests/integration/test_KGX_Model.py @@ -0,0 +1,42 @@ +import json +import pytest +from unittest.mock import patch + +from roger.models.kgx import KGXModel +from . import conftest + + +@pytest.fixture +def kgx_model(): + biolink = conftest.BiolinkMock() + kgx_model = KGXModel(biolink=biolink, config={}) + return kgx_model + +def setup_mock_and_run_create_schema(test_files_dir, kgx_model: KGXModel): + + with patch('roger.models.kgx.storage', conftest): + conftest.merge_file_test_dir = test_files_dir + with open(conftest.merge_path("expected.json")) as f: + expected = json.load(f) + conftest.file_content_assertions = expected + kgx_model.create_schema() + +def test_create_schema_plain(kgx_model: KGXModel): + file_name = 'non_conflicting_prop_types__schema__kgx' + setup_mock_and_run_create_schema(file_name, kgx_model=kgx_model) + +def test_create_schema_conflicting_nodes(kgx_model: KGXModel): + file_name = 'conflicting_prop_types__nodes__schema__kgx' + setup_mock_and_run_create_schema(file_name, kgx_model=kgx_model) + +def test_create_schema_conflicting_edges(kgx_model: KGXModel): + file_name = 'conflicting_prop_types__edges__schema__kgx' + setup_mock_and_run_create_schema(file_name, kgx_model=kgx_model) + +def test_merge(kgx_model: KGXModel): + with patch('roger.models.kgx.storage', conftest): + conftest.kgx_files = [ + 'data_1.merge.kgx.json', + 'data_2.merge.kgx.json' + ] + #TODO add tests for merge nodes diff --git a/tests/integration/test_bulk_loader.py b/tests/integration/test_bulk_loader.py new file mode 100644 index 00000000..ae99573f --- /dev/null +++ b/tests/integration/test_bulk_loader.py @@ -0,0 +1,111 @@ +import pytest +from unittest.mock import patch + +from roger.core import BulkLoad +from . import conftest + + +@pytest.fixture +def bulk_loader(): + biolink = conftest.BiolinkMock() + return BulkLoad(biolink=biolink, config={'separator': 30}) + + +def test_create_redis_schema(): + test_schema = { + 'concept': { + 'attribute0': 'list', + 'attribute1': 'str', + 'attribute2': 'int', + 'attribute3': 'bool' + } + } + redis_schema = BulkLoad.create_redis_schema_header(test_schema['concept'], is_relation=False) + assert 'attribute0:ARRAY' in redis_schema + assert 'attribute1:STRING' in redis_schema + assert 'attribute2:INT' in redis_schema + assert 'attribute3:BOOL' in redis_schema + + redis_schema = BulkLoad.create_redis_schema_header(test_schema['concept'], is_relation=True) + assert 'attribute0:ARRAY' in redis_schema + assert 'attribute1:STRING' in redis_schema + assert 'attribute2:INT' in redis_schema + assert 'attribute3:BOOL' in redis_schema + + # should add these columns to relationships + assert 'internal_start_id:START_ID' in redis_schema + assert 'internal_end_id:END_ID' in redis_schema + + +def test_group_by_set_attr(): + items = [ + { # we need to make sure that empty values are the only ones ignored + # not values that evaluate to false. + 'id': 0, + 'attr_1': '', + 'attr_2': 2, + 'attr_3': [], + 'attr_4': False, + 'attr_5': None + }, + { + 'id': 1, + 'attr_1': 'a', + 'attr_2': 'b', + 'attr_3': 'c', + 'attr_4': '' + } + ] + # first group is attr_2, attr_4, 'id' + group_1 = frozenset(['attr_2', 'attr_4', 'id']) + # second group is attr_1, attr_2, attr_3 , 'id' + group_2 = frozenset(['attr_1', 'attr_2', 'attr_3', 'id']) + grouping, invalid_keys = BulkLoad.group_items_by_attributes_set(objects=items, + processed_object_ids=set()) + assert group_1 in grouping + assert group_2 in grouping + + assert items[0] in grouping[group_1] + assert items[1] in grouping[group_2] + + +def test_write_bulk_nodes(bulk_loader: BulkLoad): + nodes_schema = { + "named_thing": { + "id": "str", + "str": "str", + "list_attr": "list", + "bool_attr": "bool", + "float_attr": "float", + "int_attr": "int" + } + } + node_objects = { + "named_thing": [ + { + "id": "ID:1", + "str": "name", + "list_attr": ["x"], + "bool_attr": False, + "float_attr": 0.1, + "int_attr": 0 + } + ] + } + with patch('roger.core.bulkload.storage', conftest): + bulk_path = conftest.bulk_path() + state = {} + bulk_loader.write_bulk(bulk_path=bulk_path, + obj_map=node_objects, + schema=nodes_schema, + state=state, + is_relation=False) + assert len(state['file_paths']) > 0 + # @TODO add assertions. + # with open(os.path.join(bulk_path,'named_thing_csv-0-1')) + + + + + + diff --git a/tests/integration/test_type_conversion_util.py b/tests/integration/test_type_conversion_util.py new file mode 100644 index 00000000..ab4e122f --- /dev/null +++ b/tests/integration/test_type_conversion_util.py @@ -0,0 +1,49 @@ +from roger.components.data_conversion_utils import TypeConversionUtil + + +def test_type_comparision(): + datatype_1 = list.__name__ + datatype_2 = str.__name__ + datatype_3 = bool.__name__ + datatype_4 = float.__name__ + datatype_5 = int.__name__ + # list should always come first + assert datatype_1 == TypeConversionUtil.compare_types(datatype_1, datatype_2) + assert datatype_1 == TypeConversionUtil.compare_types(datatype_1, datatype_3) + assert datatype_1 == TypeConversionUtil.compare_types(datatype_1, datatype_4) + assert datatype_1 == TypeConversionUtil.compare_types(datatype_1, datatype_5) + + # then string + assert datatype_2 == TypeConversionUtil.compare_types(datatype_2, datatype_3) + assert datatype_2 == TypeConversionUtil.compare_types(datatype_2, datatype_4) + assert datatype_2 == TypeConversionUtil.compare_types(datatype_2, datatype_5) + + # the rest should always be casted up to string + assert datatype_2 == TypeConversionUtil.compare_types(datatype_3, datatype_4) + assert datatype_2 == TypeConversionUtil.compare_types(datatype_4, datatype_5) + assert datatype_2 == TypeConversionUtil.compare_types(datatype_5, datatype_3) + + # should raise error when sent 'Unknown' data types + bogus_dt = "bogus" + try: + TypeConversionUtil.compare_types(bogus_dt, datatype_1) + except AssertionError as error: + exception_raised = True + assert exception_raised + try: + TypeConversionUtil.compare_types(datatype_1, bogus_dt) + except AssertionError as error: + exception_raised = True + assert exception_raised + + +def test_casting_values(): + castable = [ + ["True", bool.__name__, True], + [1 , bool.__name__, True], + [1.0, bool.__name__, True], + [[], bool.__name__, False] + ] + for items in castable: + assert items[-1] == TypeConversionUtil.cast(*items[:-1]) # cast (value, type) + diff --git a/tests/test_redis_query.cypher b/tests/test_redis_query.cypher new file mode 100644 index 00000000..509df1fa --- /dev/null +++ b/tests/test_redis_query.cypher @@ -0,0 +1,5 @@ +MATCH (c{id:'HP:0032316'}) return c + +MATCH (disease:`Disease` {`id`: 'MONDO:0004979'}) WITH disease MATCH (disease)-[e1_disease_phenotypic_feature]-(phenotypic_feature:`PhenotypicFeature` {}) +WITH disease AS disease, phenotypic_feature AS phenotypic_feature, collect(e1_disease_phenotypic_feature) AS e1_disease_phenotypic_feature +RETURN disease,phenotypic_feature,e1_disease_phenotypic_feature,labels(disease) AS type__disease,labels(phenotypic_feature) AS type__phenotypic_feature,[edge in e1_disease_phenotypic_feature | type(edge)] AS type__e1_disease_phenotypic_feature,[edge in e1_disease_phenotypic_feature | [startNode(edge).id, endNode(edge).id]] AS id_pairs__e1_disease_phenotypic_feature \ No newline at end of file diff --git a/tests/unit/__init__.py b/tests/unit/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/unit/test_annotate_files.py b/tests/unit/test_annotate_files.py new file mode 100644 index 00000000..7a683e7a --- /dev/null +++ b/tests/unit/test_annotate_files.py @@ -0,0 +1,143 @@ +"""Unit tests for DugPipeline.annotate_files: resume and file-level threading. + +Annotation is by far the most expensive step in roger. Measured on the live +bdc-parent run: ~55s per input file, 61,597 files, i.e. 39 days -- and 96% of +that was re-annotating the same 24 dbGaP study descriptions once per +data-dict file, because dug's parsers emit the study element into every file. + +These tests need dug/airflow importable (e.g. inside the roger image); they +skip cleanly elsewhere. +""" + +import os +import threading +from types import SimpleNamespace + +import pytest + +base = pytest.importorskip("roger.pipelines.base") + + +class StubPipeline(base.DugPipeline): + """A DugPipeline with the heavy __init__ replaced. + + __init__ builds a DugFactory, a BiolinkModel and tranql queries, none of + which annotate_files touches. + """ + + pipeline_name = "stub" + + def __init__(self, workers=1): + self.annotation_conf = SimpleNamespace(annotate_workers=workers) + self._thread_local = threading.local() + self.annotated = [] + self.threads = set() + self.lock = threading.Lock() + + def thread_annotation_context(self): + return ("session", "annotator") + + def get_parser(self): + return lambda path: [] + + def annotate_one_file(self, parse_file, parser, output_data_path, + index=0, total=0): + with self.lock: + self.annotated.append(parse_file) + self.threads.add(threading.current_thread().name) + elements, concepts = self.annotation_output_paths(parse_file, + output_data_path) + for path in (elements, concepts): + os.makedirs(os.path.dirname(path), exist_ok=True) + with open(path, "w") as handle: + handle.write("[]") + + +def source(tmp_path, name): + path = tmp_path / name + path.write_text("") + return str(path) + + +def test_annotates_every_file_when_nothing_is_done(tmp_path): + files = [source(tmp_path, f"f{n}.xml") for n in range(3)] + pipeline = StubPipeline() + pipeline.annotate_files(files, output_data_path=str(tmp_path / "out")) + assert sorted(pipeline.annotated) == sorted(files) + + +def test_skips_files_already_fully_annotated(tmp_path): + """The resume path: outputs hard-linked in from a previous try must not be + recomputed.""" + out = tmp_path / "out" + files = [source(tmp_path, f"f{n}.xml") for n in range(3)] + for name in ("f0", "f1"): + done = out / name + done.mkdir(parents=True) + (done / "elements.txt").write_text("[]") + (done / "concepts.txt").write_text("[]") + + pipeline = StubPipeline() + pipeline.annotate_files(files, output_data_path=str(out)) + assert pipeline.annotated == [files[2]] + + +def test_half_written_output_is_redone(tmp_path): + """elements.txt is written before concepts.txt, so a task killed between + the two left a directory that looks started but is unusable.""" + out = tmp_path / "out" + partial = out / "f0" + partial.mkdir(parents=True) + (partial / "elements.txt").write_text("[]") + + files = [source(tmp_path, "f0.xml")] + pipeline = StubPipeline() + pipeline.annotate_files(files, output_data_path=str(out)) + assert pipeline.annotated == files + + +def test_empty_output_file_is_redone(tmp_path): + "A zero-byte pickle is a failed write, not finished work." + out = tmp_path / "out" + done = out / "f0" + done.mkdir(parents=True) + (done / "elements.txt").write_text("") + (done / "concepts.txt").write_text("") + + pipeline = StubPipeline() + pipeline.annotate_files([source(tmp_path, "f0.xml")], + output_data_path=str(out)) + assert len(pipeline.annotated) == 1 + + +def test_files_are_annotated_concurrently(tmp_path): + """Annotation is almost entirely http wait, so threads scale it despite + the GIL. Files are independent: own parse, own crawler, own output dir.""" + files = [source(tmp_path, f"f{n}.xml") for n in range(8)] + pipeline = StubPipeline(workers=4) + pipeline.annotate_files(files, output_data_path=str(tmp_path / "out")) + assert sorted(pipeline.annotated) == sorted(files) + assert len(pipeline.threads) > 1, pipeline.threads + + +def test_worker_failure_is_not_swallowed(tmp_path): + """dug's annotate_elements has no per-element error handling, so a raised + exception means that file produced nothing; the task must fail.""" + + class Failing(StubPipeline): + def annotate_one_file(self, parse_file, *args, **kwargs): + if parse_file.endswith("f1.xml"): + raise RuntimeError("annotator exploded") + return super().annotate_one_file(parse_file, *args, **kwargs) + + files = [source(tmp_path, f"f{n}.xml") for n in range(4)] + with pytest.raises(RuntimeError, match="annotator exploded"): + Failing(workers=4).annotate_files( + files, output_data_path=str(tmp_path / "out")) + + +def test_worker_count_never_drops_below_one(tmp_path): + "min(workers, len(pending)) is 0 for an empty list; ThreadPoolExecutor " + "rejects max_workers=0." + StubPipeline(workers=4).annotate_files([], + output_data_path=str(tmp_path)) diff --git a/tests/unit/test_batched_annotator.py b/tests/unit/test_batched_annotator.py new file mode 100644 index 00000000..ca754649 --- /dev/null +++ b/tests/unit/test_batched_annotator.py @@ -0,0 +1,144 @@ +"""Batched identifier lookups must match dug's serial path. + +dug calls the normalizer and the name resolution service once per +identifier. Both take a list, so this collapses ~2 requests per identifier +into 2 per element. The risk is not speed, it is drifting from dug's +semantics -- greenlist handling, dropped identifiers, synonym keying -- so +these tests use dug's own handle_response parsers. +""" +from unittest import mock + +import pytest + +from dug.core.annotators._base import (DefaultNormalizer, DefaultSynonymFinder, + DugIdentifier) +from roger.utils.batched_annotator import BatchedAnnotator + +NORM_URL = ("http://norm.local:8080/get_normalized_nodes" + "?conflate=false&description=true&curie=") +SYN_URL = "http://names.local:2433/synonyms" + + +class FakeResponse: + def __init__(self, payload): + self._payload = payload + + def raise_for_status(self): + pass + + def json(self): + return self._payload + + +class FakeSession: + """Answers the two batch endpoints, recording every call.""" + + def __init__(self, normalized, synonyms, fail=None): + self.normalized, self.synonyms, self.fail = normalized, synonyms, fail + self.calls = [] + + def post(self, url, json=None, **kw): + self.calls.append((url, json)) + if self.fail and self.fail in url: + raise RuntimeError("service down") + if "normalized_nodes" in url: + return FakeResponse({c: self.normalized.get(c) + for c in json["curies"]}) + return FakeResponse({c: self.synonyms.get(c, {"names": []}) + for c in json["preferred_curies"]}) + + +def _norm_entry(curie, label, typ="biolink:Disease"): + return {"id": {"identifier": curie, "label": label}, + "equivalent_identifiers": [{"identifier": curie}], + "type": [typ]} + + +def make_annotator(greenlist=()): + """An inner annotator with dug's real parsers, faked network methods.""" + normalizer = object.__new__(DefaultNormalizer) + normalizer.url = NORM_URL + synonym_finder = object.__new__(DefaultSynonymFinder) + synonym_finder.url = SYN_URL + + inner = mock.MagicMock() + inner.normalizer = normalizer + inner.synonym_finder = synonym_finder + inner.ontology_greenlist = list(greenlist) + inner.bagel_enabled = False + inner.text_classification.return_value = [{"text": "x", "bl_type": "y"}] + inner.annotate_classifiers.return_value = { + "asthma": [DugIdentifier(id="MONDO:1", label="asthma"), + DugIdentifier(id="MONDO:2", label="wheeze")], + "aspirin": [DugIdentifier(id="CHEBI:1", label="aspirin")], + } + return inner + + +def test_one_request_per_service_not_per_identifier(): + inner = make_annotator() + session = FakeSession( + normalized={"MONDO:1": _norm_entry("MONDO:1", "Asthma"), + "MONDO:2": _norm_entry("MONDO:2", "Wheeze"), + "CHEBI:1": _norm_entry("CHEBI:1", "Aspirin")}, + synonyms={"MONDO:1": {"names": ["asthma", "bronchial asthma"]}}) + + out = BatchedAnnotator(inner)("some text", session) + + norm_calls = [c for c in session.calls if "normalized_nodes" in c[0]] + syn_calls = [c for c in session.calls if c[0] == SYN_URL] + assert len(norm_calls) == 1, "normalization should be a single request" + assert len(syn_calls) == 1, "synonym lookup should be a single request" + # 3 identifiers, one request each way + assert sorted(norm_calls[0][1]["curies"]) == ["CHEBI:1", "MONDO:1", "MONDO:2"] + body = norm_calls[0][1] + assert body["conflate"] is False # from the url + assert body["description"] is True # from the url + # Not in the url, and GET/POST disagree on the default: GET conflates + # drug/chemical curies, POST does not. Sending it explicitly is what keeps + # CHEBI:3759 resolving to CHEBI:37941 the way dug's GET does. + assert body["drug_chemical_conflate"] is True + assert body["include_taxa"] is True + assert body["individual_types"] is False + + assert [i.id for i in out] == ["MONDO:1", "MONDO:2", "CHEBI:1"] + assert [i.label for i in out] == ["Asthma", "Wheeze", "Aspirin"] + by_id = {i.id: i for i in out} + assert by_id["MONDO:1"].synonyms == ["asthma", "bronchial asthma"] + assert by_id["CHEBI:1"].synonyms == [] + + +def test_unnormalizable_identifier_is_dropped_unless_greenlisted(): + inner = make_annotator() + session = FakeSession(normalized={"MONDO:1": _norm_entry("MONDO:1", "Asthma")}, + synonyms={}) + out = BatchedAnnotator(inner)("t", session) + assert [i.id for i in out] == ["MONDO:1"], "unnormalized ids must be dropped" + + inner = make_annotator(greenlist=["CHEBI"]) + session = FakeSession(normalized={"MONDO:1": _norm_entry("MONDO:1", "Asthma")}, + synonyms={}) + out = BatchedAnnotator(inner)("t", session) + assert [i.id for i in out] == ["MONDO:1", "CHEBI:1"], "greenlist must survive" + + +@pytest.mark.parametrize("failing", ["normalized_nodes", "synonyms"]) +def test_falls_back_to_dug_when_a_batch_call_fails(failing): + inner = make_annotator() + inner.return_value = ["serial result"] + session = FakeSession( + normalized={"MONDO:1": _norm_entry("MONDO:1", "Asthma"), + "MONDO:2": _norm_entry("MONDO:2", "W"), + "CHEBI:1": _norm_entry("CHEBI:1", "A")}, + synonyms={}, fail=failing) + + out = BatchedAnnotator(inner)("some text", session) + + assert out == ["serial result"] + inner.assert_called_once_with("some text", session) + + +def test_delegates_unknown_attributes(): + inner = make_annotator() + inner.score_threshold = 0.8 + assert BatchedAnnotator(inner).score_threshold == 0.8 diff --git a/tests/unit/test_bulkload_edgeless.py b/tests/unit/test_bulkload_edgeless.py new file mode 100644 index 00000000..6fd62927 --- /dev/null +++ b/tests/unit/test_bulkload_edgeless.py @@ -0,0 +1,62 @@ +"""Bulk loading nodes with no edges must fail loudly. + +A build whose edge csvs had not been committed yet used to load nodes only +and report success, which is how the graph ended up with 3.9M nodes and zero +relationships. +""" +from unittest import mock +import pytest + +from roger.core.bulkload import BulkLoad +from roger.config import config + + +def _bulk(tmp_path, nodes=(), edges=()): + for kind, names in (('nodes', nodes), ('edges', edges)): + d = tmp_path / 'knowledge_graph_build' / f'CreateBulkLoad{kind.title()}' / kind + d.mkdir(parents=True, exist_ok=True) + for n in names: + (d / n).write_text('id\n1\n') + return tmp_path + + +def test_nodes_without_edges_raises(tmp_path): + _bulk(tmp_path, nodes=['biolink~Gene.csv-0-1']) + with pytest.raises(ValueError, match='edgeless'): + BulkLoad(config).insert(input_data_path=tmp_path) + + +def test_nodes_with_edges_passes_the_guard(tmp_path): + _bulk(tmp_path, nodes=['biolink~Gene.csv-0-1'], + edges=['biolink~treats.csv-0-1']) + # Gets past the guard and fails later trying to reach redis, which is + # enough: the guard is what this test is about. + with pytest.raises(Exception) as exc: + BulkLoad(config).insert(input_data_path=tmp_path) + assert 'edgeless' not in str(exc.value) + + +def test_index_labels_are_backticked(tmp_path, monkeypatch): + """Biolink labels contain a dot, so an unquoted one is a Cypher syntax + error. falkordb interpolates the label straight into the index pattern + and the loader only *prints* the failure, so this regresses silently.""" + from roger.core import bulkload as bulkload_mod + + _bulk(tmp_path, nodes=['biolink~Gene.csv-0-1'], + edges=['biolink~treats.csv-0-1']) + captured = {} + monkeypatch.setattr(bulkload_mod, 'bulk_insert', + lambda args, **kw: captured.setdefault('args', args)) + bulk = BulkLoad(config) + bulk.get_redisgraph = mock.MagicMock() + bulk.biolink = mock.MagicMock() + bulk.biolink.toolkit.get_ancestors.return_value = [] + bulk.insert(input_data_path=tmp_path) + + idx = [a for a in captured['args'] if a.startswith(('-i ', '-f '))] + assert idx, "no index arguments were built" + assert [a for a in idx if a.startswith('-f ')], "no fulltext index args" + for a in idx: + label = a.split(' ', 1)[1].rsplit(':', 1)[0] + assert label.startswith('`') and label.endswith('`'), \ + f"label not backticked, falkordb will reject it: {a!r}" diff --git a/tests/unit/test_commit_hook.py b/tests/unit/test_commit_hook.py new file mode 100644 index 00000000..3fb2d38d --- /dev/null +++ b/tests/unit/test_commit_hook.py @@ -0,0 +1,90 @@ +"""Output must be committed from post_execute, not on_success_callback. + +Airflow records end_date and releases downstream tasks before success +callbacks run, so committing there let downstream read the branch before the +output landed. It also only logs exceptions from state-change callbacks, so a +failed merge left the task green. +""" +from unittest import mock + +import pytest + + +def _make_task(**kw): + from roger import tasks + with mock.patch.object(tasks, 'PythonOperator', lambda **a: a): + return tasks.create_python_task( + mock.MagicMock(), "T", lambda **_: None, **kw) + + +@pytest.mark.parametrize("no_output_files", [False, True]) +def test_commit_runs_from_post_execute(no_output_files, monkeypatch): + from roger import tasks + monkeypatch.setattr(tasks.config, 'lakefs_config', + mock.MagicMock(enabled=True)) + args = _make_task(no_output_files=no_output_files) + assert 'post_execute' in args + assert 'on_success_callback' not in args + + called = [] + target = 'record_state_callback' if no_output_files else 'avalon_commit_callback' + monkeypatch.setattr(tasks, target, lambda ctx, **k: called.append(ctx)) + # rebuild so the partial closes over the patched function + args = _make_task(no_output_files=no_output_files) + args['post_execute']({'ti': 'x'}, None) + assert called == [{'ti': 'x'}] + + +def test_post_execute_accepts_the_result_arg(monkeypatch): + """Airflow calls the hook as (context, result); ours takes context.""" + from roger import tasks + monkeypatch.setattr(tasks.config, 'lakefs_config', + mock.MagicMock(enabled=True)) + monkeypatch.setattr(tasks, 'avalon_commit_callback', lambda ctx, **k: None) + hook = _make_task()['post_execute'] + hook({'ti': 'x'}, "some result") # must not raise TypeError + + +class _LakefsError(Exception): + """Shaped like lakefs_sdk's BadRequestException.""" + def __init__(self, status, body): + super().__init__(f"({status})\nHTTP response body: {body}") + self.status = status + + +def test_empty_merge_is_not_a_failure(): + """lakefs 400s a merge with nothing to apply. + + That happens whenever a task's output already matches the branch, which + is normal for deterministic work over unchanged input. Failing the task + there stalls the dag for no reason. + """ + from roger.tasks import _merge_had_no_changes + assert _merge_had_no_changes( + _LakefsError(400, '{"message":"update branch main: no changes"}')) + + +def test_other_merge_failures_still_raise(): + from roger.tasks import _merge_had_no_changes + assert not _merge_had_no_changes( + _LakefsError(400, '{"message":"branch not found"}')) + assert not _merge_had_no_changes( + _LakefsError(500, '{"message":"no changes"}')) # wrong status + assert not _merge_had_no_changes(RuntimeError("no changes")) # no status + + +def test_output_is_kept_on_failure_only_when_resumable(monkeypatch): + """keep_output is the resume mechanism for annotate, and only annotate. + + crawl and make_kgx redo everything on retry, so retaining their output + buys nothing and costs disk -- three failed crawls held 47GB and filled + the shared volume, which is what had made them fail. + """ + from roger import tasks + monkeypatch.setattr(tasks.config, 'lakefs_config', + mock.MagicMock(enabled=True)) + for resumable, expected in ((True, True), (False, False)): + args = _make_task(resumable=resumable) + assert args['on_failure_callback'].keywords['keep_output'] is expected + # default must be the safe one + assert _make_task()['on_failure_callback'].keywords['keep_output'] is False diff --git a/tests/unit/test_config.py b/tests/unit/test_config.py new file mode 100644 index 00000000..e1e36276 --- /dev/null +++ b/tests/unit/test_config.py @@ -0,0 +1,101 @@ +import os + +from roger.config import AnnotationConfig, RedisConfig, RogerConfig + + +def test_merge(): + dict_a = { + 'redis': { + 'host': 'redis', + 'port': 6379, + 'user': 'admin', + 'password': 'pass1' + } + } + dict_b = { + 'redis': { + 'port': 6389, + 'password': 'pass2' + }, + 'elasticsearch': { + 'host': 'elastic', + 'port': 9200 + } + } + + assert RogerConfig.merge_dicts(dict_a, dict_b) == { + 'redis': { + 'host': 'redis', + 'port': 6389, + 'user': 'admin', + 'password': 'pass2' + }, + 'elasticsearch': { + 'host': 'elastic', + 'port': 9200 + } + } + + +def test_get_overrides(): + prefix = "TEST_VALUES_" + assert RogerConfig.get_override_data(prefix) == {} + + os.environ[f"{prefix}REDIS_HOST"] = 'http://redis.svc' + os.environ[f"{prefix}REDIS_PORT"] = '6379' + os.environ[f"{prefix}REDIS_USER"] = 'redis-admin' + os.environ[f"{prefix}REDIS_PASSWORD"] = 'admin-pass' + os.environ[f"{prefix}ELASTIC__SEARCH_HOST"] = 'http://elastic.svc' + + actual = RogerConfig.get_override_data(prefix) + expected = { + 'redis': { + 'host': 'http://redis.svc', + 'port': '6379', + 'user': 'redis-admin', + 'password': 'admin-pass', + }, + 'elastic_search': { + 'host': 'http://elastic.svc', + } + } + assert actual == expected + + +def test_redis_conf(): + redis_conf = RedisConfig(**{}) + assert redis_conf.username == "" + assert redis_conf.password == "" + assert redis_conf.host == "redis" + assert redis_conf.graph == "test" + assert redis_conf.port == 6379 + + redis_conf = RedisConfig(**{"port": "6379"}) + assert redis_conf.port == 6379 + + + +def test_annotation_cache_flags_coerced_from_environment_strings(): + """Env vars arrive as strings, and 'false' is truthy -- which would leave + POST caching on for someone who explicitly turned it off.""" + conf = AnnotationConfig( + cache_post_requests='false', + http_cache_expire_seconds='3600', + annotate_workers='8', + ) + assert conf.cache_post_requests is False + assert conf.http_cache_expire_seconds == 3600 + assert conf.annotate_workers == 8 + + assert AnnotationConfig(cache_post_requests='true').cache_post_requests + # a zero or negative worker count would break ThreadPoolExecutor + assert AnnotationConfig(annotate_workers='0').annotate_workers == 1 + + +def test_annotation_cache_defaults_are_safe(): + """POST caching on, and a nonzero expiry so requests_cache writes redis + entries with SETEX -- keeping the annotation cache evictable while the + graph keys sharing that redis are not.""" + conf = AnnotationConfig() + assert conf.cache_post_requests is True + assert conf.http_cache_expire_seconds > 0 diff --git a/tests/unit/test_http_utils.py b/tests/unit/test_http_utils.py new file mode 100644 index 00000000..db3b8994 --- /dev/null +++ b/tests/unit/test_http_utils.py @@ -0,0 +1,196 @@ +import socket +import threading + +import pytest +from requests_cache import CachedSession + +from roger.config import AnnotationConfig +from roger.utils.http_utils import (CACHEABLE_METHODS, TimeoutRetryAdapter, + enable_post_caching, harden_session) + + +@pytest.fixture +def blackhole_server(): + """A server whose first connection goes silent, like the observed failure. + + It accepts the request and acknowledges it at the TCP layer, then never + responds and never closes. Later connections are answered normally, which + matches production: the backends stayed healthy while one pooled + connection was silently broken. + """ + srv = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + srv.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + srv.bind(('127.0.0.1', 0)) + srv.listen(8) + state = {'attempts': 0, 'held': []} + body = b'{"denotations":[]}' + + def serve(): + while True: + try: + conn, _ = srv.accept() + except OSError: + return + conn.recv(65535) + state['attempts'] += 1 + if state['attempts'] == 1: + state['held'].append(conn) + continue + conn.sendall( + b'HTTP/1.1 200 OK\r\nContent-Type: application/json\r\n' + b'Content-Length: %d\r\n\r\n%s' % (len(body), body)) + conn.close() + + threading.Thread(target=serve, daemon=True).start() + yield f"http://127.0.0.1:{srv.getsockname()[1]}/annotate/", state + for conn in state['held']: + conn.close() + srv.close() + + +@pytest.fixture +def counting_server(): + "A server that answers every POST and counts how many it received." + srv = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + srv.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + srv.bind(('127.0.0.1', 0)) + srv.listen(8) + state = {'requests': 0} + + def serve(): + while True: + try: + conn, _ = srv.accept() + except OSError: + return + try: + conn.recv(65535) + state['requests'] += 1 + body = b'{"denotations":[{"n":%d}]}' % state['requests'] + conn.sendall( + b'HTTP/1.1 200 OK\r\nContent-Type: application/json\r\n' + b'Content-Length: %d\r\n\r\n%s' % (len(body), body)) + except OSError: + pass + finally: + conn.close() + + thread = threading.Thread(target=serve, daemon=True) + thread.start() + yield f'http://127.0.0.1:{srv.getsockname()[1]}/annotate', state + srv.close() + + +def _session(tmp_path, read_timeout=0.3, retries=2): + return harden_session( + CachedSession(cache_name=str(tmp_path / 'cache')), + connect_timeout=0.5, + read_timeout=read_timeout, + retries=retries, + backoff_factor=0, + ) + + +def test_adapter_mounted_with_timeout(tmp_path): + session = _session(tmp_path) + for scheme in ('http://', 'https://'): + adapter = session.get_adapter(scheme) + assert isinstance(adapter, TimeoutRetryAdapter) + assert adapter._timeout == (0.5, 0.3) + + +def test_keepalive_enabled(tmp_path): + session = _session(tmp_path) + options = session.get_adapter( + 'http://').poolmanager.connection_pool_kw['socket_options'] + assert (socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1) in options + + +def test_silent_peer_does_not_hang_forever(tmp_path, blackhole_server): + """Without a timeout this call never returns; that caused multi-day stalls.""" + url, _ = blackhole_server + session = _session(tmp_path, retries=0) + with pytest.raises(Exception) as exc: + session.post(url, json={'text': 'x'}) # no timeout passed, as in dug + assert 'Timeout' in type(exc.value).__name__ or 'Connection' in type( + exc.value).__name__ + + +def test_retry_recovers_from_silent_peer(tmp_path, blackhole_server): + """A broken pooled connection should be retried on a fresh one, not fail.""" + url, state = blackhole_server + session = _session(tmp_path, retries=2) + response = session.post(url, json={'text': 'x'}) + assert response.status_code == 200 + assert state['attempts'] >= 2 + + +def test_timeouts_coerced_from_environment_strings(): + """Env vars arrive as strings; urllib3 rejects a string timeout outright.""" + conf = AnnotationConfig( + http_connect_timeout='5', + http_read_timeout='45.5', + http_retries='7', + http_retry_backoff='2', + ) + assert conf.http_connect_timeout == 5.0 + assert conf.http_read_timeout == 45.5 + assert conf.http_retries == 7 + assert conf.http_retry_backoff == 2.0 + assert isinstance(conf.http_retries, int) + + +# --- POST response caching ------------------------------------------------- +# Every annotation call is a POST and requests_cache caches only GET/HEAD by +# default, so dug's CachedSession never returned a cached annotation. dbGaP +# parsers emit the study element into every data-dict file of a study, so +# bdc-parent re-annotated the same 24 study descriptions 61,597 times at +# ~53s each, against ~1.4s for the variable each file actually contributes. + +def test_post_caching_off_by_default_in_requests_cache(tmp_path): + "Guard the premise: this is the upstream default the fix works around." + session = CachedSession(cache_name=str(tmp_path / 'c')) + assert 'POST' not in session.settings.allowable_methods + + +def test_enable_post_caching_allows_post(tmp_path): + session = CachedSession(cache_name=str(tmp_path / 'c')) + assert enable_post_caching(session) is True + assert set(CACHEABLE_METHODS) <= set(session.settings.allowable_methods) + + +def test_enable_post_caching_sets_expiry_only_when_asked(tmp_path): + never = CachedSession(cache_name=str(tmp_path / 'a')) + enable_post_caching(never, expire_seconds=0) + assert never.settings.expire_after in (None, -1) + + ttl = CachedSession(cache_name=str(tmp_path / 'b')) + enable_post_caching(ttl, expire_seconds=3600) + assert ttl.settings.expire_after == 3600 + + +def test_enable_post_caching_tolerates_plain_session(): + "The no-redis/no-cache path must warn, not raise." + import requests + assert enable_post_caching(requests.Session()) is False + + +def test_identical_annotation_post_is_served_from_cache(tmp_path, + counting_server): + """The whole point: the second identical study annotation must not hit + the network.""" + url, state = counting_server + session = harden_session( + CachedSession(cache_name=str(tmp_path / 'cache')), + connect_timeout=0.5, read_timeout=2, retries=0, backoff_factor=0) + enable_post_caching(session) + + payload = {'text': 'Framingham Heart Study, offspring cohort'} + first = session.post(url, json=payload) + second = session.post(url, json=payload) + other = session.post(url, json={'text': 'a different variable'}) + + assert first.json() == second.json() + assert getattr(second, 'from_cache', False) is True + # two network calls: the first payload and the different one + assert state['requests'] == 2, state diff --git a/tests/unit/test_snapshot_pause.py b/tests/unit/test_snapshot_pause.py new file mode 100644 index 00000000..dd738d5f --- /dev/null +++ b/tests/unit/test_snapshot_pause.py @@ -0,0 +1,46 @@ +"""RDB snapshots are paused for the bulk load and restored after. + +bgsave forks repeatedly under bulk-load write volume and copy-on-write on a +graph this size can OOMKill the pod. Restoring matters more than pausing: a +load that leaves snapshots off is a silent durability change. +""" +from unittest import mock + +import pytest + +from roger.config import config +from roger.core.bulkload import BulkLoad + + +def _loader(save=b'300 100000'): + bulk = BulkLoad(config) + client = mock.MagicMock() + client.config_get.return_value = {b'save': save} + bulk.get_redisgraph = mock.MagicMock(return_value=mock.MagicMock(r=client)) + return bulk, client + + +def test_pauses_then_restores(): + bulk, client = _loader() + with bulk.snapshots_paused(): + client.config_set.assert_called_once_with('save', '') + assert client.config_set.call_args_list[-1] == mock.call('save', '300 100000') + + +def test_restores_when_the_load_raises(): + bulk, client = _loader() + with pytest.raises(RuntimeError): + with bulk.snapshots_paused(): + raise RuntimeError("loader blew up") + assert client.config_set.call_args_list[-1] == mock.call('save', '300 100000') + + +def test_load_still_runs_if_config_is_denied(): + """A redis that refuses CONFIG must not block the load.""" + bulk, client = _loader() + client.config_get.side_effect = Exception("NOPERM") + ran = [] + with bulk.snapshots_paused(): + ran.append(True) + assert ran == [True] + client.config_set.assert_not_called() # nothing to restore diff --git a/tests/unit/test_tasks_incremental.py b/tests/unit/test_tasks_incremental.py new file mode 100644 index 00000000..7b6bbc3e --- /dev/null +++ b/tests/unit/test_tasks_incremental.py @@ -0,0 +1,634 @@ +"""Unit tests for incremental (delta) lakefs ingestion in roger.tasks. + +These tests need airflow/avalon/lakefs_sdk importable (e.g. inside the +Roger image); they skip cleanly elsewhere. +""" + +import os +from functools import partial +from types import SimpleNamespace +from unittest.mock import MagicMock + +import pytest + +tasks = pytest.importorskip("roger.tasks") + + +def make_ti(dag_id="annotate_and_index", + task_id="tg.annotate_topmed_files", + run_id="manual__1", try_number=1, upstream_ids=()): + return SimpleNamespace( + dag_id=dag_id, task_id=task_id, run_id=run_id, + try_number=try_number, + task=SimpleNamespace(upstream_task_ids=list(upstream_ids))) + + +def paged_diff(entries, page_size=2): + """Fake refs_api.diff_refs returning paginated results.""" + pages = ([entries[i:i + page_size] + for i in range(0, len(entries), page_size)] or [[]]) + + def fetch(**kwargs): + idx = int(kwargs.get('after') or 0) + return SimpleNamespace( + results=pages[idx], + pagination=SimpleNamespace(has_more=idx + 1 < len(pages), + next_offset=idx + 1)) + return fetch + + +def paged_objects(paths, page_size=2): + """Fake objects_api.list_objects returning paginated object listings.""" + pages = ([paths[i:i + page_size] + for i in range(0, len(paths), page_size)] or [[]]) + + def fetch(**kwargs): + idx = int(kwargs.get('after') or 0) + return SimpleNamespace( + results=[SimpleNamespace(path=p) for p in pages[idx]], + pagination=SimpleNamespace(has_more=idx + 1 < len(pages), + next_offset=idx + 1)) + return fetch + + +def _tag_not_found(repository, tag): + from lakefs_sdk.exceptions import NotFoundException + raise NotFoundException(status=404) + + +class FakeClient: + def __init__(self, diff_entries=None, tip="tipA", objects=None): + self.downloads = [] + self._client = SimpleNamespace( + refs_api=SimpleNamespace( + diff_refs=paged_diff(diff_entries or []), + log_commits=lambda repository, ref, amount: + SimpleNamespace(results=[SimpleNamespace(id=tip)])), + # resolve_ref_tip asks the tags API first; these fakes use + # branch names, so the tag lookup misses and it falls through + tags_api=SimpleNamespace(get_tag=_tag_not_found), + objects_api=SimpleNamespace( + list_objects=paged_objects(objects or []))) + + def download_files(self, remote_files, local_path, repository, + branch_or_commit_id): + self.downloads.append((tuple(remote_files), local_path, repository, + branch_or_commit_id)) + + +@pytest.fixture +def lakefs_env(monkeypatch, tmp_path): + monkeypatch.setenv("ROGER_DATA_DIR", str(tmp_path)) + monkeypatch.setattr(tasks.config.lakefs_config, "enabled", True) + monkeypatch.setattr(tasks.config.lakefs_config, "repo", "roger-out") + monkeypatch.setattr(tasks.config.lakefs_config, "branch", "main") + return tmp_path + + +def diff_entry(diff_type, path): + return SimpleNamespace(type=diff_type, path=path) + + +def test_incremental_state_key(): + key = tasks.incremental_state_key( + "annotate_and_index", "tg.annotate_topmed_files", "topmed", "v2.0") + assert key == ("roger_incr::annotate_and_index::" + "tg.annotate_topmed_files::topmed@v2.0") + other = tasks.incremental_state_key( + "annotate_and_index", "tg.annotate_topmed_files", "topmed", "v3.0") + assert key != other + + +def test_resolve_ref_tip(): + assert tasks.resolve_ref_tip(FakeClient(tip="abc123"), 'r', 'v1.0') \ + == "abc123" + empty = FakeClient() + empty._client.refs_api.log_commits = ( + lambda repository, ref, amount: SimpleNamespace(results=[])) + # empty history: fall back to the ref itself + assert tasks.resolve_ref_tip(empty, 'r', 'v1.0') == 'v1.0' + + +def test_get_changed_files_buckets_and_prefix(): + entries = [ + diff_entry('added', 'a/task/f1'), + diff_entry('changed', 'a/task/f2'), + diff_entry('removed', 'a/task/f3'), + diff_entry('added', 'a/task_b/f4'), + diff_entry('conflict', 'a/task/f5'), + ] + client = FakeClient(diff_entries=entries) + + # prefix without trailing slash must not match the 'a/task_b' sibling + changes = tasks.get_changed_files(client, 'repo', 'c1', 'c2', + prefixes=['a/task']) + assert changes == {'added': ['a/task/f1'], 'changed': ['a/task/f2'], + 'removed': ['a/task/f3']} + + # '*' or no prefixes means no filtering; conflict entries ignored + for prefixes in (['*'], None): + changes = tasks.get_changed_files(client, 'repo', 'c1', 'c2', + prefixes=prefixes) + assert changes['added'] == ['a/task/f1', 'a/task_b/f4'] + assert changes['removed'] == ['a/task/f3'] + + +def test_setup_input_data_skips_when_no_change(monkeypatch, lakefs_env): + client = FakeClient(tip="tipA") + monkeypatch.setattr(tasks, "init_lakefs_client", lambda config: client) + monkeypatch.setattr(tasks, "_get_last_consumed", lambda key: "tipA") + get_files = MagicMock() + monkeypatch.setattr(tasks, "get_files", get_files) + + context = {'ti': make_ti(), 'params': {'incremental': True}} + exec_conf = {'repos': [{'repo': 'topmed', 'branch': 'v2.0', 'path': '*'}]} + with pytest.raises(tasks.AirflowSkipException): + tasks.setup_input_data(context, exec_conf) + get_files.assert_not_called() + assert client.downloads == [] + + +def test_setup_input_data_first_run_full_download(monkeypatch, lakefs_env): + client = FakeClient(tip="tipA") + monkeypatch.setattr(tasks, "init_lakefs_client", lambda config: client) + monkeypatch.setattr(tasks, "_get_last_consumed", lambda key: None) + get_files = MagicMock() + monkeypatch.setattr(tasks, "get_files", get_files) + + ti = make_ti() + context = {'ti': ti, 'params': {'incremental': True}} + exec_conf = {'repos': [{'repo': 'topmed', 'branch': 'v2.0', 'path': '*'}]} + tasks.setup_input_data(context, exec_conf) + + get_files.assert_called_once() + call = get_files.call_args.kwargs + assert call['branch'] == "tipA" # pinned to resolved tip + assert call['changes_only'] is False + assert call['repo'] == 'topmed' + + state = tasks.read_state_file(ti) + key = tasks.incremental_state_key(ti.dag_id, ti.task_id, 'topmed', 'v2.0') + assert state['entries'][key]['commit_id'] == "tipA" + + +def test_setup_input_data_incremental_diff_download(monkeypatch, lakefs_env): + entries = [diff_entry('added', 'f1'), diff_entry('removed', 'f2')] + client = FakeClient(diff_entries=entries, tip="tipNew") + monkeypatch.setattr(tasks, "init_lakefs_client", lambda config: client) + monkeypatch.setattr(tasks, "_get_last_consumed", lambda key: "tipOld") + get_files = MagicMock() + monkeypatch.setattr(tasks, "get_files", get_files) + + ti = make_ti() + context = {'ti': ti, 'params': {'incremental': True}} + exec_conf = {'repos': [{'repo': 'topmed', 'branch': 'v2.0', 'path': '*'}]} + tasks.setup_input_data(context, exec_conf) + + get_files.assert_not_called() + assert len(client.downloads) == 1 + remote_files, _local, repo, ref = client.downloads[0] + assert remote_files == ('f1',) + assert repo == 'topmed' + assert ref == "tipNew" + + state = tasks.read_state_file(ti) + assert state['removed'] == {'topmed': ['f2']} + key = tasks.incremental_state_key(ti.dag_id, ti.task_id, 'topmed', 'v2.0') + assert state['entries'][key]['commit_id'] == "tipNew" + + +def test_find_sibling_files(): + # data dict changed in a study dir; GapExchange sibling lives alongside it + objects = ['s/phs1.v1/data_dict.xml', + 's/phs1.v1/GapExchange_phs1.v1.xml', + 's/phs1.v1/other.xml'] + client = FakeClient(objects=objects) + siblings = tasks.find_sibling_files( + client, 'repo', 'tip', ['s/phs1.v1/data_dict.xml']) + assert siblings == ['s/phs1.v1/GapExchange_phs1.v1.xml'] + + # already-downloaded marker is not duplicated + assert tasks.find_sibling_files( + client, 'repo', 'tip', + ['s/phs1.v1/GapExchange_phs1.v1.xml']) == [] + + +def test_setup_input_data_incremental_pulls_gap_exchange(monkeypatch, + lakefs_env): + entries = [diff_entry('changed', 's/phs1.v1/data_dict.xml')] + objects = ['s/phs1.v1/data_dict.xml', + 's/phs1.v1/GapExchange_phs1.v1.xml'] + client = FakeClient(diff_entries=entries, tip="tipNew", objects=objects) + monkeypatch.setattr(tasks, "init_lakefs_client", lambda config: client) + monkeypatch.setattr(tasks, "_get_last_consumed", lambda key: "tipOld") + monkeypatch.setattr(tasks, "get_files", MagicMock()) + + context = {'ti': make_ti(), 'params': {'incremental': True}} + exec_conf = {'repos': [{'repo': 'topmed', 'branch': 'v2.0', 'path': '*'}]} + tasks.setup_input_data(context, exec_conf) + + remote_files, _local, _repo, _ref = client.downloads[0] + assert set(remote_files) == {'s/phs1.v1/data_dict.xml', + 's/phs1.v1/GapExchange_phs1.v1.xml'} + + +def test_setup_input_data_incremental_pull_false(monkeypatch, lakefs_env): + # ES rebuild tasks: even with incremental=True and no upstream changes, + # incremental_pull=False must force a full pinned download (no skip) + client = FakeClient(tip="tipA") + monkeypatch.setattr(tasks, "init_lakefs_client", lambda config: client) + monkeypatch.setattr(tasks, "_get_last_consumed", lambda key: "tipA") + get_files = MagicMock() + monkeypatch.setattr(tasks, "get_files", get_files) + + context = {'ti': make_ti(), 'params': {'incremental': True}} + exec_conf = {'repos': [{'repo': 'roger-out', 'branch': 'main', + 'path': 'dag/task'}], + 'incremental_pull': False} + tasks.setup_input_data(context, exec_conf) + + get_files.assert_called_once() + call = get_files.call_args.kwargs + assert call['branch'] == "tipA" + assert call['changes_only'] is False + + +def test_removed_bases(lakefs_env): + removed = { + # external source repo: filename minus last extension + 'topmed': ['some/dir/study_one.xml', 'study_two.csv'], + # roger repo: first segment under the upstream task path + 'roger-out': [ + 'annotate_and_index/tg.annotate_x_files/study_three/elements.txt', + 'annotate_and_index/tg.annotate_x_files/.removed_files.json', + ], + } + bases = tasks.removed_bases( + removed, 'annotate_and_index', ['tg.annotate_x_files']) + assert bases == {'study_one', 'study_two', 'study_three'} + + +def test_stale_output_paths(): + remote = 'annotate_and_index/tg.crawl_x/' + existing = [ + remote + 'study_one/expanded_concepts.txt', + remote + 'study_one/elements.txt', + remote + 'study_one_kgx.json', + remote + 'study_one_extra/elements.txt', # prefix sibling: keep + remote + 'study_two/elements.txt', # not removed: keep + remote + '.removed_files.json', # manifest: keep + ] + stale = tasks.stale_output_paths(existing, remote, {'study_one'}) + assert stale == [ + remote + 'study_one/expanded_concepts.txt', + remote + 'study_one/elements.txt', + remote + 'study_one_kgx.json', + ] + + +def test_setup_input_data_manual_override_precedence(monkeypatch, lakefs_env): + client = FakeClient() + monkeypatch.setattr(tasks, "init_lakefs_client", lambda config: client) + last_consumed = MagicMock() + monkeypatch.setattr(tasks, "_get_last_consumed", last_consumed) + get_files = MagicMock() + monkeypatch.setattr(tasks, "get_files", get_files) + + ti = make_ti() + context = {'ti': ti, 'params': { + 'repository_id': 'ext-repo', 'branch_name': 'b', + 'commitid_from': 'c1', 'commitid_to': 'c2', 'incremental': True}} + tasks.setup_input_data(context, {'repos': []}) + + last_consumed.assert_not_called() + get_files.assert_called_once() + call = get_files.call_args.kwargs + assert call['repo'] == 'ext-repo' + assert call['branch'] == 'b' + assert call['changes_only'] is True + assert call['changes_from'] == 'c1' + assert call['changes_to'] == 'c2' + assert tasks.read_state_file(ti) == {} + + +def test_record_state_callback_sets_variables_and_cleans(monkeypatch, + lakefs_env): + ti = make_ti() + key = tasks.incremental_state_key(ti.dag_id, ti.task_id, 'topmed', 'v2.0') + tasks.write_state_file(ti, { + 'entries': {key: {'repo': 'topmed', 'branch': 'v2.0', + 'commit_id': 'tipA'}}, + 'removed': {}}) + state_file = tasks.get_state_file_path(ti) + assert os.path.isfile(state_file) + + recorded = [] + monkeypatch.setattr( + tasks, "Variable", + SimpleNamespace(set=lambda k, v: recorded.append((k, v)))) + + tasks.record_state_callback({'ti': ti}) + + assert recorded == [(key, 'tipA')] + assert not os.path.exists(state_file) + + +def test_state_file_noop_when_lakefs_disabled(monkeypatch): + monkeypatch.setattr(tasks.config.lakefs_config, "enabled", False) + ti = make_ti() + assert tasks.get_state_file_path(ti) is None + tasks.write_state_file(ti, {'entries': {}}) # must not raise + assert tasks.read_state_file(ti) == {} + + +def test_task_wrapper_calls_core_function_shape(monkeypatch): + """roger.core functions take explicit kwargs, not a task_kwargs dict.""" + monkeypatch.setattr(tasks.config.lakefs_config, "enabled", False) + seen = {} + + def core_like(to_string=False, config=None, input_data_path=None, + output_data_path=None): + seen.update(to_string=to_string, config=config, + input_data_path=input_data_path, + output_data_path=output_data_path) + return "ok" + + assert tasks.task_wrapper(core_like, dag_run=None, to_string=True) == "ok" + assert seen['to_string'] is True + assert seen['config'] is tasks.config + assert seen['input_data_path'] is None + + +def test_task_wrapper_calls_pipeline_method_shape(monkeypatch): + "The annotate/index path goes through execute_pipeline_method." + monkeypatch.setattr(tasks.config.lakefs_config, "enabled", False) + seen = {} + + class FakePipeline: + def __init__(self, config=None, **kwargs): + pass + + def __enter__(self): + return self + + def __exit__(self, *exc): + return False + + def annotate(self, to_string=False, input_data_path=None, + output_data_path=None): + seen.update(to_string=to_string, + input_data_path=input_data_path, + output_data_path=output_data_path) + return "annotated" + + callable_ = partial(tasks.execute_pipeline_method, + pipeline_class=FakePipeline, + configparam=tasks.config, + method_name='annotate') + out = tasks.task_wrapper(callable_, dag_run=None, to_string=True, + pass_conf=False) + assert out == "annotated" + assert seen['to_string'] is True + + +def test_orphaned_output_paths(tmp_path): + """Objects from previous runs, whose names no longer exist locally.""" + (tmp_path / "nodes").mkdir() + (tmp_path / "nodes" / "biolink~Gene.csv-0-4").write_text("x") + (tmp_path / "nodes" / "biolink~Gene.csv-1-4").write_text("x") + + remote = "knowledge_graph_build/CreateBulkLoadNodes/" + existing = [ + remote + "nodes/biolink~Gene.csv-0-4", # current run, keep + remote + "nodes/biolink~Gene.csv-1-4", # current run, keep + remote + "nodes/biolink~Gene.csv-0-3", # previous run, drop + remote + "nodes/biolink~AnatomicalEntity.csv-0-31", # older, drop + ] + orphans = tasks.orphaned_output_paths(existing, remote, str(tmp_path)) + assert orphans == [remote + "nodes/biolink~Gene.csv-0-3", + remote + "nodes/biolink~AnatomicalEntity.csv-0-31"] + + +def test_orphaned_output_paths_empty_when_in_sync(tmp_path): + (tmp_path / "a.csv").write_text("x") + remote = "dag/task/" + assert tasks.orphaned_output_paths([remote + "a.csv"], remote, + str(tmp_path)) == [] + + +def test_memory_override_patches_base_container(): + pytest.importorskip("kubernetes") + cfg = tasks.memory_override("15Gi") + pod = cfg["pod_override"] + container = pod.spec.containers[0] + assert container.name == "base" + assert container.resources.limits == {"memory": "15Gi"} + # request stays small so the namespace quota is not reserved wholesale + assert container.resources.requests == {"memory": "1Gi"} + + +def test_es_taskgroup_pulls_crawl_outputs_only(monkeypatch, lakefs_env): + """index_variables must read crawl's expanded elements.txt, not + annotate's: only the crawl copy carries KG-derived optional_terms, and + the storage glob ('**/elements.txt') cannot tell them apart.""" + recorded = {} + + def fake_create_python_task(dag, name, a_callable, **kw): + recorded[name] = [r['path'] for r in (kw.get('external_repos') or [])] + return MagicMock(name=name) + + monkeypatch.setattr(tasks, "create_python_task", fake_create_python_task) + monkeypatch.setattr(tasks, "TaskGroup", MagicMock()) + monkeypatch.setattr(tasks, "EmptyOperator", MagicMock()) + + class FakePipeline: + pipeline_name = "heal-mds-studies" + input_version = "main" + + dag = SimpleNamespace(dag_id="annotate_and_index") + tasks.create_es_taskgroup(dag, FakePipeline, tasks.config) + + assert set(recorded) == { + "index_heal-mds-studies_variables", + "validate_heal-mds-studies_index_variables", + "index_heal-mds-studies_concepts", + "validate_heal-mds-studies_index_concepts"}, recorded + for name, paths in recorded.items(): + assert paths, f"{name} pulls nothing" + assert all("crawl_heal-mds-studies" in p for p in paths), (name, paths) + assert not any("annotate_heal-mds-studies_files" in p + for p in paths), (name, paths) + + +def test_merge_searchable_docs_unions_lists_keeps_new_scalars(): + """The same CDE id appears in several element files with different + concepts and parents (BRTHDTC in adult- and pediatric-demographic); + indexing must union those, not overwrite.""" + base = pytest.importorskip("roger.pipelines.base") + + adult = {"id": "BRTHDTC", "name": "BRTHDTC", "description": "old", + "identifiers": ["UMLS:C0005615"], "parents": ["adult-demographic"], + "search_terms": ["Birth"], "optional_terms": [], "programs": [], + "tags": []} + pediatric = {"id": "BRTHDTC", "name": "BRTHDTC", "description": "new", + "identifiers": ["UMLS:C0011008"], + "parents": ["pediatric-demographic"], + "search_terms": ["Date of birth"], "optional_terms": [], + "programs": [], "tags": []} + + merged = base.merge_searchable_docs(adult, pediatric) + assert merged["identifiers"] == sorted(["UMLS:C0005615", "UMLS:C0011008"]) + assert merged["parents"] == sorted(["adult-demographic", + "pediatric-demographic"]) + assert set(merged["search_terms"]) == {"Birth", "Date of birth"} + # scalars come from the newer document + assert merged["description"] == "new" + + +def test_merge_searchable_docs_dedupes_tags(): + base = pytest.importorskip("roger.pipelines.base") + tag = {"category": "c", "value": "v"} + merged = base.merge_searchable_docs({"tags": [tag]}, {"tags": [dict(tag)]}) + assert merged["tags"] == [tag] + + +def test_validate_concepts_waits_for_index_variables(monkeypatch, + lakefs_env): + """validate_indexed_concepts searches the VARIABLES index, so it must + wait for index_variables -- not just index_concepts. Otherwise it runs + against a freshly wiped, half-populated index and finds nothing.""" + made = {} + + def fake_create_python_task(dag, name, a_callable, **kw): + t = MagicMock(name=name) + t.upstreams = [] + t.set_upstream = lambda other: t.upstreams.append(other) + made[name] = t + return t + + monkeypatch.setattr(tasks, "create_python_task", fake_create_python_task) + monkeypatch.setattr(tasks, "TaskGroup", MagicMock()) + monkeypatch.setattr(tasks, "EmptyOperator", MagicMock()) + + class FakePipeline: + pipeline_name = "heal-cdes" + input_version = "main" + + tasks.create_es_taskgroup(SimpleNamespace(dag_id="annotate_and_index"), + FakePipeline, tasks.config) + + validate = made["validate_heal-cdes_index_concepts"] + assert made["index_heal-cdes_concepts"] in validate.upstreams + assert made["index_heal-cdes_variables"] in validate.upstreams + + +def test_quiet_noisy_loggers_caps_http_chatter(): + """requests_cache logs 5 DEBUG lines per HTTP call; a single annotate task + produced a 2.9 GB log that way, which evicted the api-server (750Mi + ephemeral) when the log was opened.""" + import logging + rl = pytest.importorskip("roger.logger") + + logging.getLogger("requests_cache.policy.actions").setLevel(logging.DEBUG) + rl.quiet_noisy_loggers() + # child loggers inherit the cap from the configured parent + assert not logging.getLogger( + "requests_cache.policy.actions").isEnabledFor(logging.DEBUG) + assert not logging.getLogger("httpcore.http11").isEnabledFor(logging.DEBUG) + # roger's own logger is untouched + assert logging.getLogger("roger").isEnabledFor(logging.INFO) + + +def test_resolve_ref_tip_prefers_tags(monkeypatch): + """lakefs validates log_commits' ref as a *branch id*, so a dotted tag + like 'v7.0' (kgx.data_sets 'baseline-graph:v7.0') 400s before the tag is + looked up. Tags must be resolved through the tags API.""" + from lakefs_sdk.exceptions import NotFoundException + + calls = [] + + class Tags: + def get_tag(self, repository, tag): + calls.append(('get_tag', tag)) + if tag == 'v7.0': + return SimpleNamespace(commit_id='tagcommit') + raise NotFoundException(status=404) + + class Refs: + def log_commits(self, repository, ref, amount): + calls.append(('log_commits', ref)) + assert ref != 'v7.0', "dotted tag must not reach log_commits" + return SimpleNamespace( + results=[SimpleNamespace(id='branchcommit')]) + + client = SimpleNamespace(_client=SimpleNamespace(tags_api=Tags(), + refs_api=Refs())) + + assert tasks.resolve_ref_tip(client, 'baseline-graph', 'v7.0') == 'tagcommit' + assert tasks.resolve_ref_tip(client, 'baseline-graph', 'main') == 'branchcommit' + assert ('log_commits', 'main') in calls + + +# --- resuming an interrupted annotate -------------------------------------- +# Annotation costs tens of seconds per input file, so a bdc-parent run is +# weeks long. Its output only reaches lakefs when the whole task succeeds, so +# a task that died at file 40,000 of 61,597 used to discard all of it. + +def _write(path, text="x"): + os.makedirs(os.path.dirname(path), exist_ok=True) + with open(path, "w") as handle: + handle.write(text) + + +def test_reuse_prior_try_outputs_carries_finished_work_forward(lakefs_env): + prior = lakefs_env / ("annotate_and_index_tg.annotate_topmed_files" + "_manual__1_1_output") + _write(str(prior / "phs000007_dd" / "elements.txt")) + _write(str(prior / "phs000007_dd" / "concepts.txt")) + + ti = make_ti(try_number=2) + assert tasks.reuse_prior_try_outputs(ti) == 2 + + current = lakefs_env / ("annotate_and_index_tg.annotate_topmed_files" + "_manual__1_2_output") + assert (current / "phs000007_dd" / "elements.txt").is_file() + assert (current / "phs000007_dd" / "concepts.txt").is_file() + + +def test_reuse_prior_try_outputs_does_not_overwrite_current_try(lakefs_env): + prior = lakefs_env / ("annotate_and_index_tg.annotate_topmed_files" + "_manual__1_1_output") + current = lakefs_env / ("annotate_and_index_tg.annotate_topmed_files" + "_manual__1_2_output") + _write(str(prior / "d" / "elements.txt"), "old") + _write(str(current / "d" / "elements.txt"), "new") + + assert tasks.reuse_prior_try_outputs(make_ti(try_number=2)) == 0 + assert (current / "d" / "elements.txt").read_text() == "new" + + +def test_clean_up_keeps_output_on_failure(lakefs_env): + "The failure callback must not destroy what a retry can resume from." + ti = make_ti(try_number=1) + base = "annotate_and_index_tg.annotate_topmed_files_manual__1_1_" + _write(str(lakefs_env / (base + "input") / "src.xml")) + _write(str(lakefs_env / (base + "output") / "d" / "elements.txt")) + + tasks.clean_up({'ti': ti}, keep_output=True) + + assert not (lakefs_env / (base + "input")).exists() + assert (lakefs_env / (base + "output") / "d" / "elements.txt").is_file() + + +def test_clean_up_all_tries_clears_every_try_dir(lakefs_env): + "A successful commit has everything; leaving try dirs would leak the PVC." + for try_number in (1, 2): + base = ("annotate_and_index_tg.annotate_topmed_files_manual__1_" + f"{try_number}_") + _write(str(lakefs_env / (base + "input") / "src.xml")) + _write(str(lakefs_env / (base + "output") / "d" / "elements.txt")) + + tasks.clean_up({'ti': make_ti(try_number=2)}, all_tries=True) + + leftovers = list(lakefs_env.glob("annotate_and_index_*")) + assert leftovers == [], leftovers diff --git a/tranql-schema.yaml b/tranql-schema.yaml new file mode 100644 index 00000000..79d9d575 --- /dev/null +++ b/tranql-schema.yaml @@ -0,0 +1,12 @@ +schema: + redis: + doc: | + Roger is a knowledge graph built by aggregeting several kgx formatted knowledge graphs from several sources. + url: "redis:" + redis: true + redis_connection_params: + # Host here is the service name in the docker composed container. + host: redis + port: 6379 + # SET USERNAME and PASSWORD + # via REDIS_USERNAME , REDIS_PASSWORD Env vars (i.e capitialize service name) diff --git a/tranql_translate.py b/tranql_translate.py deleted file mode 100644 index e9fe67e6..00000000 --- a/tranql_translate.py +++ /dev/null @@ -1,71 +0,0 @@ -# -*- coding: utf-8 -*- -# - -""" -An Airflow workflow for the Roger Translator KGX data pipeline. -""" - -import os -import subprocess -from airflow.operators.bash_operator import BashOperator -from airflow.contrib.example_dags.libs.helper import print_stuff -from airflow.models import DAG -from airflow.operators.python_operator import PythonOperator -from airflow.utils.dates import days_ago -from roger.core import RogerUtil - -default_args = { - 'owner': 'RENCI', - 'start_date': days_ago(1) -} - -""" Build the workflow's tasks and DAG. """ -with DAG( - dag_id='tranql_translate', - default_args=default_args, - schedule_interval=None -) as dag: - - """ Configure use of KubernetesExecutor. """ - at_k8s=False - - def get_executor_config (annotations=None): - """ Get an executor configuration. - :param annotations: Annotations to attach to the executor. - :returns: Returns a KubernetesExecutor if K8s is configured and None otherwise. - """ - k8s_executor_config = { - "KubernetesExecutor": { - "annotations": annotations - } - } - return k8s_executor_config if at_k8s else None - - def create_python_task (name, a_callable): - """ Create a python task. - :param name: The name of the task. - :param a_callable: The code to run in this task. - """ - return PythonOperator( - task_id=name, - python_callable=a_callable, - op_kwargs={ 'to_string' : True }, - executor_config=get_executor_config (annotations={ - "task_name" : name - }) - ) - - """ Build the workflow tasks. """ - intro = BashOperator(task_id='Intro', bash_command='echo running tranql translator') - get_kgx = create_python_task ("GetSource", RogerUtil.get_kgx) - create_schema = create_python_task ("CreateSchema", RogerUtil.create_schema) - merge_nodes = create_python_task ("MergeNodes", RogerUtil.merge_nodes) - create_bulk_load = create_python_task ("CreateBulkLoad", RogerUtil.create_bulk_load) - bulk_load = create_python_task ("BulkLoad", RogerUtil.bulk_load) - validate = create_python_task ("Validate", RogerUtil.validate) - finish = BashOperator (task_id='Finish', bash_command='echo finish') - - """ Build the DAG. """ - intro >> get_kgx >> [ create_schema, merge_nodes ] >> create_bulk_load >> \ - bulk_load >> validate >> finish -