hillst/bionemo noodles #458

skothenhill-nv · 2024-11-20T15:50:22Z

NvFaidx

Adds a memmapped faidx reader with python bindings.

Context

PyFaidx or any buffered read based index is not process safe, and therefore does not play nice with pytorch dataloaders.

Due to the order of operations, the underlying file handle is shared between processes, when seek() is called to perform random lookups, this can cause unexpected behavior in the forked processes.
Ref: mdshw5/pyfaidx#211

For a good solution we need three things:
1) Safe index creation, in multi-process or multi-node scenarios, this should be restricted to a single node where all workers block until it is complete (not implemented above)
2) Index object instantion must be fast.
3) Read-only use of the index object must be both thread safe and process safe with python.

Memmap backing

Using a backend of Memmaps provide the following benefits (Credit: @malcolmgreaves)

The kernel implements this mechanism by using page faults
Each read in a mmap'd file results in a page fault: there's nothing in memory to read!
The kernel handles this page fault by going to the disk, reading the file in the specified offset + index, then returning to the user process with what it just read, preventing penalties from context switching.

These three usecases are ideal for pytorch DataLoaders!

Perf Benchmarks

nvfaidx query faster by 5.2x
nvfaidx instantiation faster by 16x

Metric	pyfaidx (s)	nvfaidx (s)
Query/s	1.668e-05	3.239e-06
Instantiation Time	0.01001	0.0006313

TODO:

Dockerfile needs to be reworked to play nice with cargo.

Usage:

Import NvFaidx, pass in your fasta file path, and go fast.

from bionemo.noodles import NvFaidx
index = NvFaidx('sub-packages/bionemo-noodles/tests/bionemo/noodles/data/sample.fasta')
assert index['chr1'][0] == 'A'
assert index['chr1'][1:4] == 'CTG'
assert index['chr1'][1:10000] == 'CTGACTGACTG'
assert index['chr1'][0:-1] == 'ACTGACTGACT'
assert index['chr1'][:-1] == 'ACTGACTGACT'
assert index['chr1'][-1:] == 'G'

…kages

Initial commit! - adds the PyO3 wrapper around noodles-fasta - adds a python class that mimics the dict-like behavior of pyfaidx - adds a long test (best used with hg38 https://hgdownload.soe.ucsc.edu/goldenPath/hg38/bigZips/), or use whatever big reference. Build: ``` cd sub-packages/bionemo-noodles/noodles_fasta_wrapper maturin develop ``` This should install it into your local python environment, then you can run the associated tests! TODO: - flesh out equality tests with pyfaidx - profile performance - get the build system working correctly

…ramework into �mg/bionemo_noodles_pyO3

on process safety.

Dockerfile

Annotated dockerfile with places we should change expose the rust reader as a part of __init__.py

malcolmgreaves

First parital review.

Dockerfile

sub-packages/bionemo-noodles/src/bionemo/noodles/nvfaidx.py

sub-packages/bionemo-noodles/src/bionemo/noodles/__init__.py

sub-packages/bionemo-noodles/src/bionemo/noodles/nvfaidx.py

sub-packages/bionemo-noodles/src/bionemo/noodles/test_nvfaidx.py

… method to create faidx objects.

Co-authored-by: Malcolm Greaves <malcolmgreaves@users.noreply.github.com> Signed-off-by: Steven Kothen-Hill <148821680+skothenhill-nv@users.noreply.github.com>

…to hillst/bionemo_noodles

Dockerfile

sub-packages/bionemo-noodles/src/bionemo/noodles/nvfaidx.py

Dockerfile

sub-packages/bionemo-noodles/src/bionemo/noodles/test_nvfaidx.py

sub-packages/bionemo-noodles/src/bionemo/noodles/nvfaidx.py

sub-packages/bionemo-noodles/src/bionemo/noodles/test_nvfaidx.py

Dockerfile

edawson

Some changes requested, but overall in good shape. Recommend refactoring PyRecord to a more descriptive name, adding some small tests, and cleaning up a little branchiness. I also noticed a lot of errors get pushed down to the rust class - is this desired?

sub-packages/bionemo-fw/pyproject.toml

sub-packages/bionemo-noodles/VERSION

sub-packages/bionemo-noodles/rust/src/lib.rs

sub-packages/bionemo-noodles/src/bionemo/noodles/nvfaidx.py

sub-packages/bionemo-noodles/tests/bionemo/noodles/test_nvfaidx.py

skothenhill-nv · 2024-11-27T22:59:54Z

/build-ci

skothenhill-nv · 2024-11-27T23:53:22Z

/build-ci

skothenhill-nv · 2024-11-28T04:17:47Z

/build-ci

skothenhill-nv · 2024-11-28T05:07:00Z

/build-ci

skothenhill-nv · 2024-11-28T05:07:41Z

/build-ci

malcolmgreaves and others added 10 commits November 14, 2024 17:22

inital creation: create-bionemo-project -p bionemo-noodles -l sub-pac…

9431205

…kages

wip maturin-based build

08d6c1b

TODO add "bionemo-noodles" back to this & toplevel & tach.toml !!!

2250617

Adds rust implementation of NvFaidx.

fa28690

add nvfaidx class and associated test

c71954b

fix bionemo-fw deps & tach.toml, incorporate bionemo-noodles

10e72be

wip, cannot build on osx

5c726e4

Merge branch 'mg/bionemo_noodles_pyO3' of github.com:NVIDIA/bionemo-f…

5d9011e

…ramework into �mg/bionemo_noodles_pyO3

Adds tests and sets the build system up for this library. Still working

444fc75

on process safety.

skothenhill-nv requested review from jstjohn, malcolmgreaves, ohadmo, pstjohn, trvachov and dorotat-nv as code owners November 20, 2024 15:50

pstjohn reviewed Nov 20, 2024

View reviewed changes

Dockerfile Outdated Show resolved Hide resolved

skothenhill-nv added 4 commits November 20, 2024 18:46

checkpointing

c667e51

Started implementing mmap and associated functions.

3ebc5bc

Annotated dockerfile with places we should change expose the rust reader as a part of __init__.py

All tests are passing.

be4e680

cleanup python tests

6651c9a

skothenhill-nv changed the base branch from mg/bionemo_noodles_pyO3 to main November 22, 2024 23:24

skothenhill-nv requested review from DejunL, farhadrgh and guoqing-zhou as code owners November 22, 2024 23:24

some refacotring and additional tests.

cb8825c

malcolmgreaves reviewed Nov 25, 2024

View reviewed changes

skothenhill-nv and others added 3 commits November 25, 2024 16:42

Updates interfaces to handle explicit/implicit faidx usage, exposes a…

3b56a7f

… method to create faidx objects.

Update sub-packages/bionemo-noodles/src/bionemo/noodles/nvfaidx.py

62ad067

Co-authored-by: Malcolm Greaves <malcolmgreaves@users.noreply.github.com> Signed-off-by: Steven Kothen-Hill <148821680+skothenhill-nv@users.noreply.github.com>

add mmap test

9384b27

skothenhill-nv added 5 commits November 27, 2024 12:59

bringing back old comment

da78afb

revert missing lines

7bf3238

add missing period

b755a0d

Merge branch 'main' of https://github.com/NVIDIA/bionemo-framework in…

d50d12e

…to hillst/bionemo_noodles

cleaning up whitespace

c20131b

trvachov reviewed Nov 27, 2024

View reviewed changes

Dockerfile Show resolved Hide resolved

fix init file and invalid import

b1485aa

pstjohn reviewed Nov 27, 2024

View reviewed changes

Dockerfile Show resolved Hide resolved

sub-packages/bionemo-noodles/src/bionemo/noodles/nvfaidx.py Outdated Show resolved Hide resolved

wrap docstrings in nvfaidx

f6e802a

skothenhill-nv commented Nov 27, 2024

View reviewed changes

formatting

27bdd85

jstjohn approved these changes Nov 27, 2024

View reviewed changes

skothenhill-nv added 2 commits November 27, 2024 13:43

formatting

5677f54

add codeowners entry for noodles

d2d8c64

farhadrgh approved these changes Nov 27, 2024

View reviewed changes

pstjohn approved these changes Nov 27, 2024

View reviewed changes

final dockerfile change, removes build artifacts

bb7ea81

edawson reviewed Nov 27, 2024

View reviewed changes

address erics comments

cb33526

skothenhill-nv enabled auto-merge (squash) November 27, 2024 23:04

Merge branch 'main' into hillst/bionemo_noodles

2b3a11b

Merge branch 'main' into hillst/bionemo_noodles

e9e0fde

Merge branch 'main' into hillst/bionemo_noodles

7e449a6

skothenhill-nv merged commit 47ca4a5 into main Nov 28, 2024
4 checks passed

skothenhill-nv deleted the hillst/bionemo_noodles branch November 28, 2024 05:56

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

hillst/bionemo noodles #458

hillst/bionemo noodles #458

skothenhill-nv commented Nov 20, 2024 •

edited

Loading

malcolmgreaves left a comment

edawson left a comment

skothenhill-nv commented Nov 27, 2024

skothenhill-nv commented Nov 27, 2024

skothenhill-nv commented Nov 28, 2024

skothenhill-nv commented Nov 28, 2024

skothenhill-nv commented Nov 28, 2024

hillst/bionemo noodles #458

hillst/bionemo noodles #458

Conversation

skothenhill-nv commented Nov 20, 2024 • edited Loading

NvFaidx

Context

Memmap backing

Perf Benchmarks

TODO:

Usage:

malcolmgreaves left a comment

Choose a reason for hiding this comment

edawson left a comment

Choose a reason for hiding this comment

skothenhill-nv commented Nov 27, 2024

skothenhill-nv commented Nov 27, 2024

skothenhill-nv commented Nov 28, 2024

skothenhill-nv commented Nov 28, 2024

skothenhill-nv commented Nov 28, 2024

skothenhill-nv commented Nov 20, 2024 •

edited

Loading