Skip to content

Run SubstructLibrary queries concurrently within a GPU-memory budget - #359

Merged
scal444 merged 4 commits into
NVIDIA-BioNeMo:mainfrom
scal444:substructlib-3b-concurrent-queries
Oct 8, 2026
Merged

scal444 merged 4 commits into
NVIDIA-BioNeMo:mainfrom
scal444:substructlib-3b-concurrent-queries

Conversation

@scal444

@scal444 scal444 commented Oct 7, 2026

Copy link
Copy Markdown
Collaborator

Instead of serializing every call, let queries run concurrently. Each finalize() works out how many queries fit below 85% of each GPU's memory, counting the batches it is about to replace as freed, and gives every GPU that many query slots, each a search workspace plus its match flags. A query takes one slot per GPU, and queries with recursive SMARTS also take one of a smaller number of slots for their scratch memory. maxConcurrentQueries() reports the limit.

Queries share the library lock; additions and finalize() take it exclusively, passing a gate first so queries that keep arriving cannot starve them. If a finalize() fails after resizing the slots, the previous slots are rebuilt.

The memory estimates live next to the search they describe, which also now resolves the default executors per runner in one place.

Instead of serializing every call, let queries run concurrently. Each
finalize() works out how many queries fit below 85% of each GPU's memory,
counting the batches it is about to replace as freed, and gives every GPU
that many query slots, each a search workspace plus its match flags. A
query takes one slot per GPU, and queries with recursive SMARTS also take
one of a smaller number of slots for their scratch memory.
maxConcurrentQueries() reports the limit.

Queries share the library lock; additions and finalize() take it
exclusively, passing a gate first so queries that keep arriving cannot
starve them. If a finalize() fails after resizing the slots, the previous
slots are rebuilt.

The memory estimates live next to the search they describe, which also
now resolves the default executors per runner in one place.
@greptile-apps

greptile-apps Bot commented Oct 7, 2026 •

Copy link
Copy Markdown
Contributor

RetriggerConfidence Score: 5/5

[High risk] Adds concurrent query execution with GPU memory pooling.

The PR appears safe to merge, with no outstanding findings from this re-review.

Summary

The PR lets SubstructLibrary queries run concurrently within a GPU-memory budget.

  • Adds per-GPU query slots and a separate limit for recursive queries.
  • Keeps additions and finalize() exclusive and restores slots after a failed rebuild.
  • Since the previous review, disables the test that incorrectly requires multiple slots on every GPU.
  • The earlier reader-start and sequential-reuse concerns are addressed by readersReady and the corrected test description.

Reviews (4) · Last reviewed commit: "Disable memory-dependent query concurren..." · Reviewed by Greptile

Comment thread tests/test_substruct_library.cu
Comment thread tests/test_substruct_library.cu Outdated
@scal444
scal444 requested a review from evasnow1992 October 7, 2026 21:19
@scal444

scal444 commented Oct 7, 2026

Copy link
Copy Markdown
Collaborator Author

PR 4 for #350

Comment thread src/substruct/substruct_library.cpp Outdated

@evasnow1992 evasnow1992 left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

One clarification question. Changes look good to me.

Comment thread tests/test_substruct_library.cu
@scal444
scal444 merged commit e662fd2 into NVIDIA-BioNeMo:main Oct 8, 2026
15 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants