Add KvikIO to Improve I/O Throughput - #2257
Conversation
|
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:
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: 📒 Files selected for processing (1)
🚧 Files skipped from review as they are similar to previous changes (1)
📝 WalkthroughSummary by CodeRabbit
WalkthroughIntegrates KvikIO into cuVS build, runtime, file I/O, neighbor disk workflows, tests, and packaging. Large-file I/O now uses kvikio handles, CAGRA and HNSW disk paths use kvikio-based reads and writes, and dependency declarations are updated across CMake, conda, Python, and release tooling. ChangesKvikIO GPU Direct Storage Integration
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.
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/src/neighbors/detail/hnsw.hpp`:
- Around line 840-846: The code calls get() on graph_future, dataset_future, and
label_future sequentially, but if graph_future.get() throws an exception, the
subsequent calls to dataset_future.get() and label_future.get() are never
executed, leaving those async operations unresolved in the background. To fix
this, ensure all three futures are awaited before any exception can propagate by
either using try-catch blocks to drain all futures before re-throwing, or by
restructuring the code to call all three get() methods within a context that
guarantees all are executed (such as a scoped guard or utility function that
ensures all futures are consumed).
In `@cpp/src/util/file_io.cpp`:
- Around line 26-38: The bulk I/O operations are reopening files by path using
kvikio::FileHandle instead of preserving the already-open file descriptor, which
creates a TOCTOU race condition where the path could be replaced between open
and I/O. Fix this by either using the existing file descriptor directly instead
of reopening via path, or by verifying that the inode and device have not
changed before proceeding with the I/O operation. This change breaks the API
contract for descriptor-only usage (where get_path() is unavailable), so add
proper deprecation warnings at the affected call sites (around line 26-38 and
line 49-61 in file_io.cpp) and update the migration guide to document the
required changes for users relying on descriptor-only semantics.
🪄 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: 300f4901-414b-474c-924a-8a674e657be5
📒 Files selected for processing (24)
conda/environments/all_cuda-129_arch-aarch64.yamlconda/environments/all_cuda-129_arch-x86_64.yamlconda/environments/all_cuda-133_arch-aarch64.yamlconda/environments/all_cuda-133_arch-x86_64.yamlconda/environments/bench_ann_cuda-129_arch-x86_64.yamlconda/environments/bench_ann_cuda-133_arch-aarch64.yamlconda/environments/bench_ann_cuda-133_arch-x86_64.yamlconda/environments/go_cuda-129_arch-x86_64.yamlconda/environments/go_cuda-133_arch-aarch64.yamlconda/environments/go_cuda-133_arch-x86_64.yamlconda/environments/rust_cuda-129_arch-x86_64.yamlconda/environments/rust_cuda-133_arch-aarch64.yamlconda/environments/rust_cuda-133_arch-x86_64.yamlconda/recipes/libcuvs/recipe.yamlcpp/CMakeLists.txtcpp/cmake/thirdparty/get_kvikio.cmakecpp/include/cuvs/util/file_io.hppcpp/src/neighbors/detail/cagra/cagra_build.cuhcpp/src/neighbors/detail/hnsw.hppcpp/src/util/file_io.cppcpp/tests/CMakeLists.txtcpp/tests/util/file_io_test.cppdependencies.yamlpython/libcuvs/pyproject.toml
jameslamb
left a comment
jameslamb
left a comment
There was a problem hiding this comment.
I left a few small initial comments from a packaging-codeowners perspective. In general, not opposed to having kvikio be a dependency of cuvs if cuVS maintainers are ok with it.
I'd be happy to review again once cuVS maintainers have more thoroughly gone through this and given their opinions.
cjnolet
added
improvement
# Conflicts: # cpp/src/neighbors/detail/cagra/cagra_build.cuh
# Conflicts: # cpp/CMakeLists.txt # cpp/include/cuvs/neighbors/cagra.hpp
dantegd
left a comment
dantegd
left a comment
There was a problem hiding this comment.
The integration looks well designed, the pr looks great overall.
| $<$<BOOL:${CUVS_NVTX}>:CUDA::nvtx3> | ||
| $<COMPILE_ONLY:nvidia::cutlass::cutlass> | ||
| $<COMPILE_ONLY:cuco::cuco> | ||
| $<BUILD_LOCAL_INTERFACE:kvikio::kvikio> |
There was a problem hiding this comment.
Should we preserve KvikIO in the installed static target’s transitive link interface?
cuvs_static now contains references to compiled kvikio::* symbols, but $<BUILD_LOCAL_INTERFACE:...> evaluates to empty when the target is consumed from another build system. Installing libkvikio supplies the DSO, but it does not add it to a downstream executable’s final link line, so consumers of the installed cuvs::cuvs_static target can get unresolved symbols when an affected archive member is selected, right?
There was a problem hiding this comment.
Good catch, thanks. cuvs_static now retains KvikIO as a link-only dependency, and KvikIO is registered in the static build/install export sets. The shared target continues to keep it private.
| RAFT_EXPECTS(dest_ptr != nullptr, "Destination pointer must not be nullptr"); | ||
| RAFT_EXPECTS(fd.is_valid(), "File descriptor must be valid"); | ||
| const std::string path = fd.get_path(); | ||
| RAFT_EXPECTS(!path.empty(), "File descriptor must have an associated path for kvikio I/O"); |
There was a problem hiding this comment.
The inode/device validation looks good for the path replacement race, but should we also retain the previous POSIX pread path for pathless descriptors, and likewise retain pwrite in the write helper?
file_descriptor(int) and these bulk helpers are installed public APIs if I'm not mistaken, and descriptor only callers worked before this PR and the new test at file_io_test.cpp now checks that both operations throw. That breaks raw external FDs and memfd-style descriptors in a PR labeled non-breaking, or we can make it breaking of course. What do you think?
There was a problem hiding this comment.
Yes, that might be helpful in the future. I've added a POSIX pread/pwrite path for pathless descriptors. I've also extended the tests.
| buffer_(std::move(other.buffer_)), | ||
| buffer_size_(std::exchange(other.buffer_size_, 0)) | ||
| { | ||
| setg(buffer_.get(), buffer_.get(), buffer_.get()); |
There was a problem hiding this comment.
I think we need to preserve the unread portion of the buffer when moving the stream. After a buffered read, fd_ has already advanced past the full 8 KiB buffer, so resetting the get area here causes any unread bytes to be skipped. It looks like move assignment has the same issue. Would it make sense to add a test that reads a prefix, moves the stream, and then verifies the remaining contents?
There was a problem hiding this comment.
Agreed. Move construction and assignment now transfer the underlying std::streambuf get area, preserving unread buffered bytes. Both operations have regression coverage.
| "ACE: partition %lu did not fit in device memory for a direct (GDS) read; " | ||
| "falling back to a host read", | ||
| partition_id); | ||
| } catch (const raft::logic_error&) { |
There was a problem hiding this comment.
Could we narrow this catch, or at least log the original exception?
There was a problem hiding this comment.
The fallback now catches only std::bad_alloc, which includes RMM allocation failures, and logs the original exception.
6 participants
Motivation:
raft::sample_rowsin each partition. ACE uses large files leading to high page-cache pressure and memory fragmentation. Direct file I/O skips the page cache and resolves the potentially failing pinned memory allocations.Changes:
.npydata regions for direct I/O and GDS-friendly transfers.std::istream/std::ostreamAPIs are unchanged.Package impact:
libkvikioforcondaandlibkvikio-cu12/libkvikio-cu13for wheels).cuda-pathfinderandcuFile, totaling approximately 3.7-4.2 MiB of additional downloads.condasolve adds 14 packages totaling approximately 16.1 MiB compressed, primarily ICU (12.1 MiB) through KvikIO’slibcurldependency chain.rapids-loggerpackages are reused.