Add HNSW Layered Index Support - #2148
Conversation
|
Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually. Contributors can view more details about this message here. |
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughThis PR adds GPU_LAYERED_ON_DISK support for HNSW, including public API changes, layered artifact serialization and deserialization, benchmark save handling, ACE build updates, tests, and a new example. ChangesLayered HNSW disk-backed flow
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes Suggested labels
Suggested reviewers
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
cpp/src/neighbors/detail/hnsw.hpp (2)
2396-2405: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick winHIGH: Validate
enterpoint_nodebefore storing it in hnswlib state.Have you considered range-checking
metadata.enterpoint_nodeagainst[0, n_rows)during deserialize? A malformed layered artifact can currently survive load with an invalid entrypoint and then walk hnswlib with an out-of-bounds internal id on the first search.Suggested fix
RAFT_EXPECTS(metadata.n_rows > 0, "Layered HNSW artifact must contain at least one row"); RAFT_EXPECTS(metadata.dim > 0, "Layered HNSW artifact must contain at least one dimension"); + RAFT_EXPECTS(metadata.enterpoint_node >= 0 && + static_cast<size_t>(metadata.enterpoint_node) < metadata.n_rows, + "Layered HNSW artifact enterpoint_node (%d) is outside [0, %zu)", + metadata.enterpoint_node, + metadata.n_rows); RAFT_EXPECTS(static_cast<size_t>(dim) == metadata.dim, "Layered HNSW artifact dim (%zu) does not match requested dim (%d)", metadata.dim, dim);As per coding guidelines,
Input validation must check for negative or invalid dimensions, null pointers, and invalid parameter combinations before GPU operations.Also applies to: 2499-2504
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@cpp/src/neighbors/detail/hnsw.hpp` around lines 2396 - 2405, Add a deserialize-time validation for metadata.enterpoint_node in the layered HNSW load path before it is written into hnswlib state. In the same checks near the existing RAFT_EXPECTS guards in the layered HNSW deserialization logic, verify that enterpoint_node is within [0, metadata.n_rows) and reject malformed artifacts early, using the relevant deserialize/load helper and the hnswlib state initialization code that stores the entrypoint.Source: Coding guidelines
727-735: 🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick winHIGH: Reject
.npydatasets whose dtype or memory order does not match the layered artifact.Have you considered carrying the
.npyheader’s dtype/layout throughopen_npy_file()and validating it here? Right now only the shape is checked, so a same-shaped but wrong-typed or Fortran-ordered dataset will be streamed as row-majorTbytes and silently corrupt the reconstructed index. As per coding guidelines,Data layout (row-major vs column-major) must be explicitly verified and handled in memory accessandData format parameters (row-major vs column-major, memory layout) must be explicit in function signatures and documentation; ambiguous data layout assumptions should be clarified or eliminated.Also applies to: 2428-2459
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@cpp/src/neighbors/detail/hnsw.hpp` around lines 727 - 735, open_npy_file() currently returns only the shape, so layered_artifact::load_from_npz() can accept same-shaped but wrong-typed or Fortran-ordered .npy data and corrupt the index. Update open_npy_file() to carry the numpy header’s dtype and layout/stride metadata alongside shape, then validate those fields in the layered artifact loading path before streaming bytes into T. Use the existing symbols open_npy_file(), npy_file, and layered_artifact::load_from_npz() to enforce that only the expected row-major dtype/layout is accepted, and reject mismatches early.Source: Coding guidelines
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Outside diff comments:
In `@cpp/src/neighbors/detail/hnsw.hpp`:
- Around line 2396-2405: Add a deserialize-time validation for
metadata.enterpoint_node in the layered HNSW load path before it is written into
hnswlib state. In the same checks near the existing RAFT_EXPECTS guards in the
layered HNSW deserialization logic, verify that enterpoint_node is within [0,
metadata.n_rows) and reject malformed artifacts early, using the relevant
deserialize/load helper and the hnswlib state initialization code that stores
the entrypoint.
- Around line 727-735: open_npy_file() currently returns only the shape, so
layered_artifact::load_from_npz() can accept same-shaped but wrong-typed or
Fortran-ordered .npy data and corrupt the index. Update open_npy_file() to carry
the numpy header’s dtype and layout/stride metadata alongside shape, then
validate those fields in the layered artifact loading path before streaming
bytes into T. Use the existing symbols open_npy_file(), npy_file, and
layered_artifact::load_from_npz() to enforce that only the expected row-major
dtype/layout is accepted, and reject mismatches early.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: b15a3474-efb8-43fe-8b06-6a6283890e30
📒 Files selected for processing (10)
cpp/bench/ann/src/cuvs/cuvs_cagra_hnswlib.cucpp/bench/ann/src/cuvs/cuvs_cagra_hnswlib_wrapper.hcpp/src/neighbors/detail/cagra/cagra_build.cuhcpp/src/neighbors/detail/hnsw.hppcpp/tests/neighbors/ann_hnsw_ace.cuhcpp/tests/neighbors/ann_hnsw_ace/test_float_uint32_t.cucpp/tests/neighbors/ann_hnsw_ace/test_half_uint32_t.cucpp/tests/neighbors/ann_hnsw_ace/test_int8_t_uint32_t.cucpp/tests/neighbors/ann_hnsw_ace/test_uint8_t_uint32_t.cuexamples/cpp/CMakeLists.txt
🚧 Files skipped from review as they are similar to previous changes (6)
- cpp/tests/neighbors/ann_hnsw_ace/test_float_uint32_t.cu
- examples/cpp/CMakeLists.txt
- cpp/bench/ann/src/cuvs/cuvs_cagra_hnswlib.cu
- cpp/tests/neighbors/ann_hnsw_ace.cuh
- cpp/bench/ann/src/cuvs/cuvs_cagra_hnswlib_wrapper.h
- cpp/src/neighbors/detail/cagra/cagra_build.cuh
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@cpp/include/cuvs/util/numpy_dtype.hpp`:
- Around line 24-35: The public API in numpy_dtype.hpp is missing required
Doxygen coverage, so add complete Doxygen comments for numpy_dtype_string,
make_numpy_header_from_dtype, and make_numpy_header_string, including parameter
descriptions, return values, and any side effects/contract details; place the
docs directly above the function declarations and ensure the public header
exposes the serialized-header behavior clearly for downstream users.
- Around line 24-35: These new public helpers expose stateful STL types in the
cuVS API, so either move numpy_dtype_string and make_numpy_header_from_dtype out
of the public header or change their interfaces to use stateless inputs instead
of std::string and std::vector<size_t>. Update the API surface in
cpp/include/cuvs/util/numpy_dtype.hpp and the related helpers around the same
area to rely on POD-style parameters, raft::resources, pointers, or
mdspan-compatible inputs so the public-header contract is preserved.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 80b27030-35d3-4c2c-9a98-ec15b4faa45a
📒 Files selected for processing (10)
cpp/include/cuvs/util/file_io.hppcpp/include/cuvs/util/numpy_dtype.hppcpp/src/neighbors/brute_force_serialize.cucpp/src/neighbors/detail/cagra/cagra_serialize.cuhcpp/src/neighbors/detail/hnsw.hppcpp/src/neighbors/ivf_flat/ivf_flat_serialize.cuhcpp/src/neighbors/ivf_sq/ivf_sq_serialize.cuhcpp/src/neighbors/mg/snmg.cuhcpp/src/util/serialize_validation.hppexamples/cpp/CMakeLists.txt
🚧 Files skipped from review as they are similar to previous changes (2)
- examples/cpp/CMakeLists.txt
- cpp/src/neighbors/detail/hnsw.hpp
tfeher
added
feature request
tfeher
left a comment
tfeher
left a comment
There was a problem hiding this comment.
Thanks Julian for this PR! This is an important feature that will reduce data transfer times while reading ACE results. I am still reviewing, but sharing the first few comments.
I think we need more clarity in how we name this option, and when it is used. Even if we build in memory (without ACE), it would be useful to save only the graph (and hierarchy) to file.
tfeher
left a comment
tfeher
left a comment
There was a problem hiding this comment.
Thanks Julian, the implementation looks good now, but I think we can still improve parameter naming to make the API easier to understand for new users.
Keep the GPU-built hierarchy separate from vectors so a search node can reconstruct an hnswlib index from a local dataset.
c8ea429 to
b579ce2
Compare
Thanks, @tfeher. I like these simplifications. I think I've addressed all comments. Please double-check and approve/merge this PR if you're happy with it. Please note that I had to rebase and force-push since the 26.12 tag was reachable before due to the merged main branch. This failed the |
5 participants
The CAGRA graph built by the disk-backed ACE algorithm partitions the dataset. Thus, the CAGRA graph uses the reordered index space. Building a HNSW index using
hnsw::from_cagrauses the reordered dataset and CAGRA graph. Downstream consumers building an HNSW index would therefore require the reordered dataset, which is typically large when requiring the disk-backed ACE algorithm. Thus, building only the layers of the HNSW index without the dataset and moving this to the search node can minimize the network transfers for downstream consumers if they have the original dataset locally available. Thehnsw::deserializestep then takes the layered index and combines it with the local dataset to form a hnswlib compatible search index.Artifact Layout
Layered HNSW Serialization
The layered serializer creates
hnsw_index.cuvsfrom the disk-backed ACE graph..cuvsfile and write the fixed header and metadata.levelssequentially.dataset_mapping.npysequentially intoreordered_to_original.cagra_graph.npysource-sequentially in ACE reordered row order.base_nodes[row] = reordered_to_original[ace_reordered_row]base_links[row]upper_nodesandupper_linkswith node IDs and neighbor IDs converted back to original IDs.This keeps remapping, link padding, and upper-layer KNN work on the build node.
Deserialization
The search node reads:
hnsw_index.cuvsindex_params.dataset_pathThe loader:
levelssequentially.levels[original_id]base_nodesandbase_linkssequentially.get_linklist0(base_node_id).upper_nodesandupper_linkssequentially by layer.get_linklist(node_id, level).The search node does no graph remapping, no level generation, no link padding, and no KNN work.
Disk Access Patterns
Build node:
reordered_dataset.npyandaugmented_dataset.npy.cagra_graph.npywhen creating the final layered artifact.hnsw_index.cuvs.Search node:
hnsw_index.cuvs.Runtime Requirements
Only
hnsw_index.cuvsis copied to the search node. ACE temporary files remain build-node-only.The search node must have the original dataset in original row order and must provide that path through
index_params.dataset_path.Misc
Unifies the logging format of the ACE algorithm.