From 18d59e5879dcfab12e565a0de75b5c1ccafadeba Mon Sep 17 00:00:00 2001 From: Julian Miller Date: Tue, 23 Jun 2026 20:00:20 +0200 Subject: [PATCH 01/18] Use KvikIO for CAGRA-ACE --- .../all_cuda-129_arch-aarch64.yaml | 1 + .../all_cuda-129_arch-x86_64.yaml | 2 + .../all_cuda-133_arch-aarch64.yaml | 2 + .../all_cuda-133_arch-x86_64.yaml | 2 + .../bench_ann_cuda-129_arch-x86_64.yaml | 1 + .../bench_ann_cuda-133_arch-aarch64.yaml | 1 + .../bench_ann_cuda-133_arch-x86_64.yaml | 1 + .../environments/go_cuda-129_arch-x86_64.yaml | 1 + .../go_cuda-133_arch-aarch64.yaml | 1 + .../environments/go_cuda-133_arch-x86_64.yaml | 1 + .../rust_cuda-129_arch-x86_64.yaml | 1 + .../rust_cuda-133_arch-aarch64.yaml | 1 + .../rust_cuda-133_arch-x86_64.yaml | 1 + conda/recipes/libcuvs/recipe.yaml | 22 ++ cpp/CMakeLists.txt | 4 + cpp/cmake/thirdparty/get_kvikio.cmake | 40 +++ cpp/include/cuvs/util/file_io.hpp | 239 ++++++++++++--- .../neighbors/detail/cagra/cagra_build.cuh | 287 ++++++++++-------- cpp/src/neighbors/detail/hnsw.hpp | 87 +++--- cpp/src/util/file_io.cpp | 220 +++++++++++--- cpp/tests/CMakeLists.txt | 7 + cpp/tests/util/file_io_test.cpp | 218 +++++++++++++ dependencies.yaml | 44 +++ python/libcuvs/pyproject.toml | 2 + 24 files changed, 929 insertions(+), 257 deletions(-) create mode 100644 cpp/cmake/thirdparty/get_kvikio.cmake create mode 100644 cpp/tests/util/file_io_test.cpp diff --git a/conda/environments/all_cuda-129_arch-aarch64.yaml b/conda/environments/all_cuda-129_arch-aarch64.yaml index bf34a36008..39e300eade 100644 --- a/conda/environments/all_cuda-129_arch-aarch64.yaml +++ b/conda/environments/all_cuda-129_arch-aarch64.yaml @@ -27,6 +27,7 @@ dependencies: - libcurand-dev - libcusolver-dev - libcusparse-dev +- libkvikio==26.8.*,>=0.0.0a0 - libnvjitlink-dev - libopenblas<=0.3.30 - librmm==26.8.*,>=0.0.0a0 diff --git a/conda/environments/all_cuda-129_arch-x86_64.yaml b/conda/environments/all_cuda-129_arch-x86_64.yaml index 8c83164f58..f372dfdeb7 100644 --- a/conda/environments/all_cuda-129_arch-x86_64.yaml +++ b/conda/environments/all_cuda-129_arch-x86_64.yaml @@ -24,9 +24,11 @@ dependencies: - go - libclang==20.1.8 - libcublas-dev +- libcufile-dev - libcurand-dev - libcusolver-dev - libcusparse-dev +- libkvikio==26.8.*,>=0.0.0a0 - libnvjitlink-dev - librmm==26.8.*,>=0.0.0a0 - make diff --git a/conda/environments/all_cuda-133_arch-aarch64.yaml b/conda/environments/all_cuda-133_arch-aarch64.yaml index 2f10891439..1da8d32820 100644 --- a/conda/environments/all_cuda-133_arch-aarch64.yaml +++ b/conda/environments/all_cuda-133_arch-aarch64.yaml @@ -24,9 +24,11 @@ dependencies: - go - libclang==20.1.8 - libcublas-dev +- libcufile-dev - libcurand-dev - libcusolver-dev - libcusparse-dev +- libkvikio==26.8.*,>=0.0.0a0 - libnvjitlink-dev - libopenblas<=0.3.30 - librmm==26.8.*,>=0.0.0a0 diff --git a/conda/environments/all_cuda-133_arch-x86_64.yaml b/conda/environments/all_cuda-133_arch-x86_64.yaml index 0d98002ba2..77173ae228 100644 --- a/conda/environments/all_cuda-133_arch-x86_64.yaml +++ b/conda/environments/all_cuda-133_arch-x86_64.yaml @@ -24,9 +24,11 @@ dependencies: - go - libclang==20.1.8 - libcublas-dev +- libcufile-dev - libcurand-dev - libcusolver-dev - libcusparse-dev +- libkvikio==26.8.*,>=0.0.0a0 - libnvjitlink-dev - librmm==26.8.*,>=0.0.0a0 - make diff --git a/conda/environments/bench_ann_cuda-129_arch-x86_64.yaml b/conda/environments/bench_ann_cuda-129_arch-x86_64.yaml index d84bf87234..009f31cd5f 100644 --- a/conda/environments/bench_ann_cuda-129_arch-x86_64.yaml +++ b/conda/environments/bench_ann_cuda-129_arch-x86_64.yaml @@ -28,6 +28,7 @@ dependencies: - libboost-devel=1.87 - libclang==20.1.8 - libcublas-dev +- libcufile-dev - libcurand-dev - libcusolver-dev - libcusparse-dev diff --git a/conda/environments/bench_ann_cuda-133_arch-aarch64.yaml b/conda/environments/bench_ann_cuda-133_arch-aarch64.yaml index c874298d7d..31c5e64c21 100644 --- a/conda/environments/bench_ann_cuda-133_arch-aarch64.yaml +++ b/conda/environments/bench_ann_cuda-133_arch-aarch64.yaml @@ -26,6 +26,7 @@ dependencies: - h5py>=3.8.0 - libclang==20.1.8 - libcublas-dev +- libcufile-dev - libcurand-dev - libcusolver-dev - libcusparse-dev diff --git a/conda/environments/bench_ann_cuda-133_arch-x86_64.yaml b/conda/environments/bench_ann_cuda-133_arch-x86_64.yaml index 43b2dbce46..e2a344e6d6 100644 --- a/conda/environments/bench_ann_cuda-133_arch-x86_64.yaml +++ b/conda/environments/bench_ann_cuda-133_arch-x86_64.yaml @@ -28,6 +28,7 @@ dependencies: - libboost-devel=1.87 - libclang==20.1.8 - libcublas-dev +- libcufile-dev - libcurand-dev - libcusolver-dev - libcusparse-dev diff --git a/conda/environments/go_cuda-129_arch-x86_64.yaml b/conda/environments/go_cuda-129_arch-x86_64.yaml index e99615ab2c..eb9fcbf319 100644 --- a/conda/environments/go_cuda-129_arch-x86_64.yaml +++ b/conda/environments/go_cuda-129_arch-x86_64.yaml @@ -22,6 +22,7 @@ dependencies: - go - libclang==20.1.8 - libcublas-dev +- libcufile-dev - libcurand-dev - libcusolver-dev - libcusparse-dev diff --git a/conda/environments/go_cuda-133_arch-aarch64.yaml b/conda/environments/go_cuda-133_arch-aarch64.yaml index eeb5511537..05b00c1537 100644 --- a/conda/environments/go_cuda-133_arch-aarch64.yaml +++ b/conda/environments/go_cuda-133_arch-aarch64.yaml @@ -22,6 +22,7 @@ dependencies: - go - libclang==20.1.8 - libcublas-dev +- libcufile-dev - libcurand-dev - libcusolver-dev - libcusparse-dev diff --git a/conda/environments/go_cuda-133_arch-x86_64.yaml b/conda/environments/go_cuda-133_arch-x86_64.yaml index 5fd2153ffa..6b02c21a87 100644 --- a/conda/environments/go_cuda-133_arch-x86_64.yaml +++ b/conda/environments/go_cuda-133_arch-x86_64.yaml @@ -22,6 +22,7 @@ dependencies: - go - libclang==20.1.8 - libcublas-dev +- libcufile-dev - libcurand-dev - libcusolver-dev - libcusparse-dev diff --git a/conda/environments/rust_cuda-129_arch-x86_64.yaml b/conda/environments/rust_cuda-129_arch-x86_64.yaml index a9e9e2142c..13c2c77147 100644 --- a/conda/environments/rust_cuda-129_arch-x86_64.yaml +++ b/conda/environments/rust_cuda-129_arch-x86_64.yaml @@ -16,6 +16,7 @@ dependencies: - cxx-compiler - gcc_linux-64=14.* - libcublas-dev +- libcufile-dev - libcurand-dev - libcusolver-dev - libcusparse-dev diff --git a/conda/environments/rust_cuda-133_arch-aarch64.yaml b/conda/environments/rust_cuda-133_arch-aarch64.yaml index 3daa10d2be..0223fba1ca 100644 --- a/conda/environments/rust_cuda-133_arch-aarch64.yaml +++ b/conda/environments/rust_cuda-133_arch-aarch64.yaml @@ -16,6 +16,7 @@ dependencies: - cxx-compiler - gcc_linux-aarch64=14.* - libcublas-dev +- libcufile-dev - libcurand-dev - libcusolver-dev - libcusparse-dev diff --git a/conda/environments/rust_cuda-133_arch-x86_64.yaml b/conda/environments/rust_cuda-133_arch-x86_64.yaml index 9e0fe2d529..8aa9341870 100644 --- a/conda/environments/rust_cuda-133_arch-x86_64.yaml +++ b/conda/environments/rust_cuda-133_arch-x86_64.yaml @@ -16,6 +16,7 @@ dependencies: - cxx-compiler - gcc_linux-64=14.* - libcublas-dev +- libcufile-dev - libcurand-dev - libcusolver-dev - libcusparse-dev diff --git a/conda/recipes/libcuvs/recipe.yaml b/conda/recipes/libcuvs/recipe.yaml index aa7a37db44..3d645fa093 100644 --- a/conda/recipes/libcuvs/recipe.yaml +++ b/conda/recipes/libcuvs/recipe.yaml @@ -75,6 +75,12 @@ cache: - libnvjitlink-dev - librmm =${{ minor_version }} - libraft-headers =${{ minor_version }} + - libkvikio =${{ minor_version }} + # libcufile (GPUDirect Storage backend that kvikio compiles against): CUDA 13 on all arches, + # CUDA 12 on x86_64 only. + - if: cuda_major == "13" or linux64 + then: + - libcufile-dev - nccl ${{ nccl_version }} - cuda-version =${{ cuda_version }} - cuda-cudart-dev @@ -174,6 +180,10 @@ outputs: - ${{ pin_subpackage("libcuvs-headers", exact=True) }} - librmm =${{ minor_version }} - libraft-headers =${{ minor_version }} + - libkvikio =${{ minor_version }} + - if: cuda_major == "13" or linux64 + then: + - libcufile-dev - nccl ${{ nccl_version }} - cuda-version =${{ cuda_version }} - cuda-cudart-dev @@ -189,6 +199,10 @@ outputs: - ${{ pin_subpackage("libcuvs-headers", exact=True) }} - libraft-headers =${{ minor_version }} - librmm =${{ minor_version }} + - libkvikio =${{ minor_version }} + - if: cuda_major == "13" or linux64 + then: + - libcufile - cuda-nvrtc - nccl - libcublas @@ -235,6 +249,10 @@ outputs: - ${{ pin_subpackage("libcuvs-headers", exact=True) }} - librmm =${{ minor_version }} - libraft-headers =${{ minor_version }} + - libkvikio =${{ minor_version }} + - if: cuda_major == "13" or linux64 + then: + - libcufile-dev - nccl ${{ nccl_version }} - cuda-version =${{ cuda_version }} - cuda-cudart-dev @@ -250,6 +268,10 @@ outputs: - ${{ pin_subpackage("libcuvs-headers", exact=True) }} - libraft-headers =${{ minor_version }} - librmm =${{ minor_version }} + - libkvikio =${{ minor_version }} + - if: cuda_major == "13" or linux64 + then: + - libcufile - cuda-nvrtc - nccl - libcublas diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt index 89cbadfcfc..39f4139985 100644 --- a/cpp/CMakeLists.txt +++ b/cpp/CMakeLists.txt @@ -189,6 +189,7 @@ if(NOT BUILD_CPU_ONLY) include(${rapids-cmake-dir}/cpm/cccl.cmake) rapids_cpm_cccl() include(cmake/thirdparty/get_raft.cmake) + include(cmake/thirdparty/get_kvikio.cmake) include(cmake/thirdparty/get_cutlass.cmake) include(${rapids-cmake-dir}/cpm/cuco.cmake) rapids_cpm_cuco() @@ -1463,6 +1464,7 @@ if(NOT BUILD_CPU_ONLY) ${CUVS_CTK_MATH_DEPENDENCIES} $ $ + kvikio::kvikio ) target_include_directories( @@ -1541,6 +1543,7 @@ if(NOT BUILD_CPU_ONLY) $<$:CUDA::nvtx3> PRIVATE $ $ $ CUDA::nvJitLink CUDA::nvrtc + $ ) # ensure CUDA symbols aren't relocated to the middle of the debug build binaries @@ -1600,6 +1603,7 @@ SECTIONS $<$:CUDA::nvtx3> $ $ + $ ) endif() diff --git a/cpp/cmake/thirdparty/get_kvikio.cmake b/cpp/cmake/thirdparty/get_kvikio.cmake new file mode 100644 index 0000000000..9fb46a0dcf --- /dev/null +++ b/cpp/cmake/thirdparty/get_kvikio.cmake @@ -0,0 +1,40 @@ +# ============================================================================= +# cmake-format: off +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. +# SPDX-License-Identifier: Apache-2.0 +# cmake-format: on + +# Use RAPIDS_VERSION_MAJOR_MINOR from rapids_config.cmake +set(KVIKIO_VERSION "${RAPIDS_VERSION_MAJOR_MINOR}") +set(KVIKIO_FORK "rapidsai") +set(KVIKIO_PINNED_TAG "${rapids-cmake-checkout-tag}") + +function(find_and_configure_kvikio) + set(oneValueArgs VERSION FORK PINNED_TAG) + cmake_parse_arguments(PKG "${options}" "${oneValueArgs}" "${multiValueArgs}" ${ARGN}) + + # ---------------------------------------------------------------------------- + # KvikIO provides the GPUDirect Storage (cuFile) and POSIX file I/O backends used by the CAGRA + # ACE disk-mode build. It is a private, build-time dependency: it is linked into libcuvs via + # $ and therefore is NOT part of cuvs' exported interface + # (no export set / find_dependency is generated). At runtime the conda 'libkvikio' package + # provides the shared library. + # ---------------------------------------------------------------------------- + rapids_cpm_find( + kvikio ${PKG_VERSION} + GLOBAL_TARGETS kvikio::kvikio + CPM_ARGS + EXCLUDE_FROM_ALL TRUE + GIT_REPOSITORY https://github.com/${PKG_FORK}/kvikio.git + GIT_TAG ${PKG_PINNED_TAG} + GIT_SHALLOW TRUE + SOURCE_SUBDIR cpp + OPTIONS "KvikIO_BUILD_EXAMPLES OFF" "KvikIO_REMOTE_SUPPORT OFF" + ) +endfunction() + +# Change pinned tag here to test a commit in CI. +# To use a different KvikIO locally, set the CMake variable CPM_kvikio_SOURCE=/path/to/local/kvikio +find_and_configure_kvikio( + VERSION ${KVIKIO_VERSION}.00 FORK ${KVIKIO_FORK} PINNED_TAG ${KVIKIO_PINNED_TAG} +) diff --git a/cpp/include/cuvs/util/file_io.hpp b/cpp/include/cuvs/util/file_io.hpp index f02e9c2d7d..fc336a3935 100644 --- a/cpp/include/cuvs/util/file_io.hpp +++ b/cpp/include/cuvs/util/file_io.hpp @@ -11,6 +11,7 @@ #include #include #include +#include #include #include #include @@ -26,19 +27,41 @@ namespace CUVS_EXPORT cuvs { namespace util { + +/** + * @brief Alignment (in bytes) for the numpy data body in on-disk ACE artifacts. + * + * The numpy header is space-padded so that the array data begins on this boundary. This keeps the + * file numpy-compatible (readers use the stored HEADER_LEN) while letting kvikio's O_DIRECT / + * GPUDirect Storage path transfer the aligned interior directly. 4096 ensures support for most + * storage devices. + */ +inline constexpr size_t kNumpyDataAlignment = 4096; + +/** @brief Round @p x up to the next multiple of @p alignment (must be a power of two). */ +inline constexpr size_t numpy_align_up(size_t x, size_t alignment = kNumpyDataAlignment) noexcept +{ + return (x + alignment - 1) & ~(alignment - 1); +} + /** * @brief Streambuf that reads from a POSIX file descriptor */ class fd_streambuf : public std::streambuf { - int fd_; + int fd_ = -1; std::unique_ptr buffer_; - size_t buffer_size_; + size_t buffer_size_ = 0; protected: int_type underflow() override { - if (gptr() < egptr()) return traits_type::to_int_type(*gptr()); - ssize_t n = ::read(fd_, buffer_.get(), buffer_size_); + if (gptr() < egptr()) { return traits_type::to_int_type(*gptr()); } + if (fd_ == -1 || !buffer_) { return traits_type::eof(); } + + ssize_t n = 0; + do { + n = ::read(fd_, buffer_.get(), buffer_size_); + } while (n < 0 && errno == EINTR); if (n <= 0) return traits_type::eof(); setg(buffer_.get(), buffer_.get(), buffer_.get() + n); return traits_type::to_int_type(*gptr()); @@ -48,18 +71,43 @@ class fd_streambuf : public std::streambuf { explicit fd_streambuf(int fd, size_t buffer_size = 8192) : fd_(fd), buffer_(new char[buffer_size]), buffer_size_(buffer_size) { + RAFT_EXPECTS(buffer_size > 0, "fd_streambuf buffer size must be greater than zero"); setg(buffer_.get(), buffer_.get(), buffer_.get()); } - ~fd_streambuf() + ~fd_streambuf() override { close(); } + + fd_streambuf(const fd_streambuf&) = delete; + fd_streambuf& operator=(const fd_streambuf&) = delete; + + fd_streambuf(fd_streambuf&& other) noexcept + : fd_(std::exchange(other.fd_, -1)), + buffer_(std::move(other.buffer_)), + buffer_size_(std::exchange(other.buffer_size_, 0)) { - if (fd_ != -1) ::close(fd_); + setg(buffer_.get(), buffer_.get(), buffer_.get()); } - fd_streambuf(const fd_streambuf&) = delete; - fd_streambuf& operator=(const fd_streambuf&) = delete; - fd_streambuf(fd_streambuf&&) noexcept = default; - fd_streambuf& operator=(fd_streambuf&&) noexcept = default; + fd_streambuf& operator=(fd_streambuf&& other) noexcept + { + if (this != &other) { + close(); + fd_ = std::exchange(other.fd_, -1); + buffer_ = std::move(other.buffer_); + buffer_size_ = std::exchange(other.buffer_size_, 0); + setg(buffer_.get(), buffer_.get(), buffer_.get()); + } + return *this; + } + + private: + void close() noexcept + { + if (fd_ != -1) { + ::close(fd_); + fd_ = -1; + } + } }; /** @@ -91,8 +139,9 @@ class fd_istream : public std::istream { /** * @brief RAII wrapper for POSIX file descriptors * - * Manages file descriptor lifecycle with automatic cleanup. - * Non-copyable, move-only. + * Manages file descriptor lifecycle with automatic cleanup. Used to own the lifetime of disk-backed + * ACE artifacts and to parse their numpy headers; the bulk data transfers go through kvikio (see + * ::read_large_file / ::write_large_file). Non-copyable, move-only. */ class file_descriptor { public: @@ -173,19 +222,25 @@ class file_descriptor { /** * @brief Create a numpy file with pre-allocated space and write the header. * - * Opens a file, writes a numpy header for the given shape/dtype, and pre-allocates - * space for the data. This is useful for memory-mapped or streaming writes. + * Opens a file, writes a numpy header for the given shape/dtype, and pre-allocates space for the + * data. The numpy header is space-padded so that the data body begins on a ::kNumpyDataAlignment + * boundary, which keeps the file numpy-compatible (readers use the stored HEADER_LEN) while letting + * kvikio transfer the aligned interior via O_DIRECT / GPUDirect Storage. The data region is rounded + * up to a block boundary so that an O_DIRECT write of the trailing block stays within the + * allocated file. The returned descriptor owns the file lifetime; bulk data is written separately + * through kvikio (::write_large_file). * * @tparam T Data type for the numpy array * @param path File path to create * @param shape Shape of the numpy array (e.g., {rows, cols} for 2D) - * @return Pair of (file_descriptor, header_size) + * @return Pair of (file_descriptor, header_size) where header_size is the data offset in bytes */ template std::pair create_numpy_file(const std::string& path, const std::vector& shape) { - // Open file + // Open the file for the header write + preallocation. Bulk data is written via kvikio (which + // opens its own descriptor, including an O_DIRECT one when supported). file_descriptor fd(path, O_CREAT | O_RDWR | O_TRUNC, 0644); // Build header @@ -196,42 +251,86 @@ std::pair create_numpy_file(const std::string& path, std::stringstream ss; raft::detail::numpy_serializer::write_header(ss, header); std::string header_str = ss.str(); - size_t header_size = header_str.size(); + + RAFT_EXPECTS((kNumpyDataAlignment & (kNumpyDataAlignment - 1)) == 0, + "kNumpyDataAlignment must be a power of two"); + + // Re-pad the numpy v1.0 header so the data body starts on a block boundary. + // Layout: [6 bytes magic][2 bytes version][2 bytes HEADER_LEN][dict padded to HEADER_LEN]. + RAFT_EXPECTS(header_str.size() >= 10 && static_cast(header_str[6]) == 0x01, + "Expected a numpy v1.0 header to align the data body"); + std::string dict = header_str.substr(10); + while (!dict.empty() && (dict.back() == '\n' || dict.back() == ' ')) { + dict.pop_back(); + } + const size_t min_total = 10 + dict.size() + 1; // +1 for the terminating newline + const size_t aligned_total = numpy_align_up(min_total); + const size_t header_len = aligned_total - 10; + RAFT_EXPECTS(header_len <= 0xFFFF, "Aligned numpy v1.0 header is too large"); + + std::string padded_dict = dict; + padded_dict.append(header_len - dict.size() - 1, ' '); + padded_dict.push_back('\n'); + + std::string aligned_header = header_str.substr(0, 8); // magic + version + aligned_header.push_back(static_cast(header_len & 0xFF)); + aligned_header.push_back(static_cast((header_len >> 8) & 0xFF)); + aligned_header.append(padded_dict); + header_str = std::move(aligned_header); + + const size_t header_size = header_str.size(); // Calculate data size from shape + auto checked_mul = [](size_t lhs, size_t rhs) { + RAFT_EXPECTS(rhs == 0 || lhs <= std::numeric_limits::max() / rhs, + "Numpy file size calculation overflowed"); + return lhs * rhs; + }; size_t data_bytes = sizeof(T); for (auto dim : shape) { - data_bytes *= dim; - } - - // Pre-allocate file space - if (posix_fallocate(fd.get(), 0, header_size + data_bytes) != 0) { - RAFT_FAIL("Failed to pre-allocate space for file: %s", path.c_str()); + data_bytes = checked_mul(data_bytes, dim); } - // Seek to beginning and write header - if (lseek(fd.get(), 0, SEEK_SET) == -1) { - RAFT_FAIL("Failed to seek to beginning of file: %s", path.c_str()); + // Pre-allocate file space. The data region is rounded up to a block boundary so read-modify-write + // of the trailing block during an O_DIRECT write stays within the allocated file. + const size_t padded_data_bytes = numpy_align_up(data_bytes); + RAFT_EXPECTS(header_size <= std::numeric_limits::max() - padded_data_bytes, + "Numpy file size calculation overflowed"); + const size_t alloc_bytes = header_size + padded_data_bytes; + const int fallocate_status = posix_fallocate(fd.get(), 0, alloc_bytes); + if (fallocate_status != 0) { + RAFT_FAIL( + "Failed to pre-allocate space for file %s: %s", path.c_str(), strerror(fallocate_status)); } - ssize_t written = write(fd.get(), header_str.data(), header_str.size()); - if (written < 0 || static_cast(written) != header_str.size()) { - RAFT_FAIL("Failed to write numpy header to file: %s", path.c_str()); + // Write the small numpy header (one-time, buffered). + const char* hp = header_str.data(); + size_t hremaining = header_size; + off_t hoff = 0; + while (hremaining > 0) { + const size_t chunk = std::min(hremaining, static_cast(SSIZE_MAX)); + const ssize_t w = ::pwrite(fd.get(), hp, chunk, hoff); + if (w < 0 && errno == EINTR) { continue; } + RAFT_EXPECTS(w > 0, "Failed to write numpy header to %s: %s", path.c_str(), strerror(errno)); + hp += w; + hoff += w; + hremaining -= static_cast(w); } return {std::move(fd), header_size}; } /** - * @brief Read large file in chunks using pread + * @brief Read a region of a file into @p dest_ptr through kvikio. * - * Reads a file in chunks to avoid issues with very large reads. - * Uses pread for thread-safe, offset-based reading. + * Routed through kvikio::FileHandle, so it uses GPUDirect Storage (cuFile) when @p dest_ptr is + * device memory on a GDS-capable system, and the POSIX + threadpool backend (with O_DIRECT when + * available) otherwise. @p dest_ptr may be host or device memory. * - * @param fd File descriptor to read from - * @param dest_ptr Destination buffer - * @param total_bytes Total bytes to read - * @param file_offset Starting offset in file + * @param fd File descriptor identifying the file (its path is used to open the handle). + * @param dest_ptr Destination buffer (host or device). + * @param total_bytes Total bytes to read. + * @param file_offset Starting offset in file. */ void read_large_file(const file_descriptor& fd, void* dest_ptr, @@ -239,35 +338,79 @@ void read_large_file(const file_descriptor& fd, const uint64_t file_offset); /** - * @brief Write large file in chunks using pwrite + * @brief Write @p data_ptr to a region of a (pre-created) file through kvikio. * - * Writes data to a file in chunks to avoid issues with very large writes. - * Uses pwrite for thread-safe, offset-based writing. + * Counterpart of ::read_large_file. The file must already exist (e.g. created by + * ::create_numpy_file); the handle is opened in read+write mode so the existing header and + * preallocation are preserved. @p data_ptr may be host or device memory. * - * @param fd File descriptor to write to - * @param data_ptr Source data buffer - * @param total_bytes Total bytes to write - * @param file_offset Starting offset in file + * @param fd File descriptor identifying the file (its path is used to open the handle). + * @param data_ptr Source data buffer (host or device). + * @param total_bytes Total bytes to write. + * @param file_offset Starting offset in file. */ void write_large_file(const file_descriptor& fd, const void* data_ptr, const size_t total_bytes, const uint64_t file_offset); +/** + * @brief Sequential std::ostream backed by kvikio. + * + * A std::ostream whose bytes are staged into a large buffer and written to disk through kvikio, + * which bypasses the page cache via O_DIRECT when supported (and falls back to buffered POSIX + * writes otherwise). Because it is a std::ostream it can be passed anywhere an ostream is expected + * (e.g. the hnswlib serializer), routing that output through kvikio. Non-copyable, non-movable. + */ +class kvikio_ofstream : public std::ostream { + public: + /** + * @brief Open @p path for writing (created/truncated). + * + * @param path Output file path. + * @param buffer_size Staging-buffer capacity in bytes; full buffers are written at aligned + * offsets. + */ + explicit kvikio_ofstream(const std::string& path, size_t buffer_size = (size_t(32) << 20)); + + ~kvikio_ofstream() override; + + kvikio_ofstream(const kvikio_ofstream&) = delete; + kvikio_ofstream& operator=(const kvikio_ofstream&) = delete; + kvikio_ofstream(kvikio_ofstream&&) = delete; + kvikio_ofstream& operator=(kvikio_ofstream&&) = delete; + + /** @brief Flush any remaining staged bytes and close the file. */ + void close(); + + /** @brief Total number of logical bytes written so far. */ + [[nodiscard]] size_t bytes_written() const noexcept; + + private: + class sbuf; + std::unique_ptr buf_; +}; + /** * @brief Buffered output stream wrapper * - * Wraps an std::ostream with a buffer to improve write performance by - * reducing the number of system calls. Automatically flushes on destruction. - * Non-copyable, non-movable. + * Wraps an std::ostream with a buffer to improve write performance by reducing the number of + * system calls. Automatically flushes on destruction. Used by the hnswlib serializer. Non-copyable, + * non-movable. */ class buffered_ofstream { public: - buffered_ofstream(std::ostream* os, size_t buffer_size) : os_(os), buffer_(buffer_size), pos_(0) + buffered_ofstream(std::ostream* os, size_t buffer_size) : buffer_(buffer_size), os_(os), pos_(0) { } - ~buffered_ofstream() noexcept { flush(); } + ~buffered_ofstream() noexcept + { + try { + flush(); + } catch (...) { + } + } buffered_ofstream(const buffered_ofstream& res) = delete; auto operator=(const buffered_ofstream& other) -> buffered_ofstream& = delete; diff --git a/cpp/src/neighbors/detail/cagra/cagra_build.cuh b/cpp/src/neighbors/detail/cagra/cagra_build.cuh index e119d7e6ca..36404811ea 100644 --- a/cpp/src/neighbors/detail/cagra/cagra_build.cuh +++ b/cpp/src/neighbors/detail/cagra/cagra_build.cuh @@ -28,6 +28,8 @@ #include #include +#include + // TODO: This shouldn't be calling spatial/knn APIs #include "../ann_utils.cuh" @@ -35,6 +37,7 @@ #include #include +#include #include #include #include @@ -674,127 +677,109 @@ void ace_reorder_and_store_dataset( // ACE: Load partition dataset and augmented dataset from disk template void ace_load_partition_dataset_from_disk( - raft::resources const& res, - const std::string& build_dir, size_t partition_id, size_t dataset_dim, raft::host_matrix_view partition_histogram, raft::host_vector_view core_partition_offsets, raft::host_vector_view augmented_partition_offsets, - raft::host_matrix_view sub_dataset) + const cuvs::util::file_descriptor& reordered_fd, + const cuvs::util::file_descriptor& augmented_fd, + size_t reordered_header_size, + size_t augmented_header_size, + T* sub_dataset_ptr) { - size_t n_partitions = partition_histogram.extent(0); - RAFT_LOG_DEBUG("ACE: Loading partition %lu dataset from disk", partition_id); - size_t core_size = partition_histogram(partition_id, 0); - size_t augmented_size = partition_histogram(partition_id, 1); - size_t total_partition_size = core_size + augmented_size; + const std::string reordered_path = reordered_fd.get_path(); + const std::string augmented_path = augmented_fd.get_path(); + RAFT_EXPECTS(sub_dataset_ptr != nullptr, "ACE: sub-dataset destination must not be null"); + RAFT_EXPECTS(reordered_fd.is_valid() && !reordered_path.empty(), + "ACE: reordered dataset file descriptor is not valid"); + RAFT_EXPECTS(augmented_fd.is_valid() && !augmented_path.empty(), + "ACE: augmented dataset file descriptor is not valid"); + RAFT_EXPECTS(partition_id < static_cast(partition_histogram.extent(0)), + "ACE: partition id is out of range"); + RAFT_EXPECTS(partition_id < static_cast(core_partition_offsets.extent(0)), + "ACE: core partition offset is out of range"); + RAFT_EXPECTS(partition_id < static_cast(augmented_partition_offsets.extent(0)), + "ACE: augmented partition offset is out of range"); + + size_t core_size = partition_histogram(partition_id, 0); + size_t augmented_size = partition_histogram(partition_id, 1); RAFT_LOG_DEBUG("ACE: Partition %lu: %lu core + %lu augmented = %lu total vectors", partition_id, core_size, augmented_size, - total_partition_size); - - RAFT_EXPECTS(static_cast(sub_dataset.extent(0)) == total_partition_size, - "sub_dataset rows (%lu) must match total partition size (%lu)", - sub_dataset.extent(0), - total_partition_size); - RAFT_EXPECTS(static_cast(sub_dataset.extent(1)) == dataset_dim, - "sub_dataset columns (%lu) must match dataset dimensions (%lu)", - sub_dataset.extent(1), - dataset_dim); + core_size + augmented_size); const size_t vector_size = dataset_dim * sizeof(T); - const std::string reordered_dataset_path = build_dir + "/reordered_dataset.npy"; - const std::string augmented_dataset_path = build_dir + "/augmented_dataset.npy"; - - if (!std::filesystem::exists(reordered_dataset_path)) { - RAFT_FAIL("ACE: Required file does not exist: %s", reordered_dataset_path.c_str()); - } - if (!std::filesystem::exists(augmented_dataset_path)) { - RAFT_FAIL("ACE: Required file does not exist: %s", augmented_dataset_path.c_str()); - } - - size_t core_header_size = 0; - size_t augmented_header_size = 0; - size_t core_file_offset = 0; - size_t augmented_file_offset = 0; - { - std::ifstream is(reordered_dataset_path, std::ios::in | std::ios::binary); - if (!is) { RAFT_FAIL("Cannot open file %s", reordered_dataset_path.c_str()); } - auto start_pos = is.tellg(); - raft::detail::numpy_serializer::read_header(is); - core_header_size = static_cast(is.tellg() - start_pos); - } - { - std::ifstream is(augmented_dataset_path, std::ios::in | std::ios::binary); - if (!is) { RAFT_FAIL("Cannot open file %s", augmented_dataset_path.c_str()); } - auto start_pos = is.tellg(); - raft::detail::numpy_serializer::read_header(is); - augmented_header_size = static_cast(is.tellg() - start_pos); - } - - for (size_t p = 0; p < partition_id; p++) { - core_file_offset += partition_histogram(p, 0); - augmented_file_offset += partition_histogram(p, 1); - } - - core_file_offset *= vector_size; - augmented_file_offset *= vector_size; - - core_file_offset += core_header_size; - augmented_file_offset += augmented_header_size; + const size_t core_file_offset = + reordered_header_size + static_cast(core_partition_offsets(partition_id)) * vector_size; + const size_t augmented_file_offset = + augmented_header_size + + static_cast(augmented_partition_offsets(partition_id)) * vector_size; RAFT_LOG_DEBUG("ACE: Core file offset: %lu bytes, Augmented file offset: %lu bytes", core_file_offset, augmented_file_offset); - // Read core and augmented data in parallel - std::exception_ptr core_exception = nullptr; - std::exception_ptr augmented_exception = nullptr; + // Read core and augmented vectors through kvikio. @p sub_dataset_ptr may be host or device + // memory; kvikio uses GPUDirect Storage (cuFile) for device destinations on a GDS-capable system + // and the POSIX + threadpool backend (with O_DIRECT when available) otherwise. The two file + // regions are independent, so issue both reads before waiting for either one. + const size_t core_bytes = core_size * vector_size; + const size_t augmented_bytes = augmented_size * vector_size; + T* augmented_dest = sub_dataset_ptr + (core_size * dataset_dim); + + auto expect_complete_read = [partition_id](const char* name, size_t expected, size_t actual) { + RAFT_EXPECTS(actual == expected, + "ACE: Short %s read for partition %lu: expected %zu bytes, got %zu", + name, + partition_id, + expected, + actual); + }; -#pragma omp parallel sections - { -#pragma omp section - { - try { - if (core_size > 0) { - RAFT_LOG_DEBUG( - "ACE: Reading %lu core vectors from offset %lu", core_size, core_file_offset); - cuvs::util::file_descriptor reordered_fd(reordered_dataset_path, O_RDONLY); - const size_t core_bytes = core_size * vector_size; - cuvs::util::read_large_file( - reordered_fd, sub_dataset.data_handle(), core_bytes, core_file_offset); - } - } catch (...) { - core_exception = std::current_exception(); - } + if (core_bytes > 0 && augmented_bytes > 0) { + RAFT_LOG_DEBUG("ACE: Reading %lu core vectors from offset %lu", core_size, core_file_offset); + RAFT_LOG_DEBUG( + "ACE: Reading %lu augmented vectors from offset %lu", augmented_size, augmented_file_offset); + kvikio::FileHandle reordered_handle(reordered_path, "r"); + kvikio::FileHandle augmented_handle(augmented_path, "r"); + auto core_future = reordered_handle.pread(sub_dataset_ptr, core_bytes, core_file_offset); + auto augmented_future = + augmented_handle.pread(augmented_dest, augmented_bytes, augmented_file_offset); + std::exception_ptr read_exception = nullptr; + size_t core_read = 0; + size_t augmented_read = 0; + try { + core_read = core_future.get(); + } catch (...) { + read_exception = std::current_exception(); } -#pragma omp section - { - try { - if (augmented_size > 0) { - RAFT_LOG_DEBUG("ACE: Reading %lu augmented vectors from offset %lu", - augmented_size, - augmented_file_offset); - cuvs::util::file_descriptor augmented_fd(augmented_dataset_path, O_RDONLY); - const size_t augmented_bytes = augmented_size * vector_size; - T* augmented_dest = sub_dataset.data_handle() + (core_size * dataset_dim); - cuvs::util::read_large_file( - augmented_fd, augmented_dest, augmented_bytes, augmented_file_offset); - } - } catch (...) { - augmented_exception = std::current_exception(); - } + try { + augmented_read = augmented_future.get(); + } catch (...) { + if (!read_exception) { read_exception = std::current_exception(); } } + if (read_exception) { std::rethrow_exception(read_exception); } + expect_complete_read("core", core_bytes, core_read); + expect_complete_read("augmented", augmented_bytes, augmented_read); + } else if (core_bytes > 0) { + RAFT_LOG_DEBUG("ACE: Reading %lu core vectors from offset %lu", core_size, core_file_offset); + kvikio::FileHandle reordered_handle(reordered_path, "r"); + auto core_future = reordered_handle.pread(sub_dataset_ptr, core_bytes, core_file_offset); + expect_complete_read("core", core_bytes, core_future.get()); + } else if (augmented_bytes > 0) { + RAFT_LOG_DEBUG( + "ACE: Reading %lu augmented vectors from offset %lu", augmented_size, augmented_file_offset); + kvikio::FileHandle augmented_handle(augmented_path, "r"); + auto augmented_future = + augmented_handle.pread(augmented_dest, augmented_bytes, augmented_file_offset); + expect_complete_read("augmented", augmented_bytes, augmented_future.get()); } - - // Check for exceptions from parallel sections - if (core_exception) { std::rethrow_exception(core_exception); } - if (augmented_exception) { std::rethrow_exception(augmented_exception); } } // Memory requirements for ACE operation @@ -1229,7 +1214,9 @@ index build_ace(raft::resources const& res, // Mark for cleanup if we fail after creating the directory cleanup_on_failure = true; - // Create numpy files with pre-allocated space + // Create numpy files with pre-allocated space. All bulk I/O on these artifacts goes through + // kvikio, which bypasses the page cache via O_DIRECT (and uses GPUDirect Storage for device + // transfers) when supported, falling back to buffered POSIX I/O otherwise. std::tie(reordered_fd, reordered_header_size) = cuvs::util::create_numpy_file( build_dir + "/reordered_dataset.npy", {dataset_size, dataset_dim}); @@ -1351,20 +1338,78 @@ index build_ace(raft::resources const& res, core_sub_dataset_size, augmented_sub_dataset_size); - auto sub_dataset = raft::make_host_matrix(sub_dataset_size, dataset_dim); + // Index parameters depend only on partition sizes, not on where the dataset lives. + cuvs::neighbors::cagra::index_params sub_index_params; + sub_index_params = cuvs::neighbors::cagra::index_params::from_hnsw_params( + raft::make_extents(sub_dataset_size, dataset_dim), + graph_degree / 2, + ef_construction, + cuvs::neighbors::cagra::hnsw_heuristic_type::SAME_GRAPH_FOOTPRINT, + params.metric); + sub_index_params.attach_dataset_on_build = false; + sub_index_params.guarantee_connectivity = params.guarantee_connectivity; - if (use_disk_mode) { - // Load partition dataset from disk files - ace_load_partition_dataset_from_disk(res, - build_dir, - partition_id, - dataset_dim, - partition_histogram.view(), - core_partition_offsets.view(), - augmented_partition_offsets.view(), - sub_dataset.view()); - } else { - // Gather partition dataset from memory + // In disk mode, try to read this partition's vectors straight into device memory so CAGRA + // builds from a device view. This enables GPUDirect Storage (NVMe->GPU DMA via kvikio/cuFile) + // on a GDS-capable system and avoids a host staging copy. Only attempt it when the vectors + // comfortably fit free device memory (leaving headroom for the CAGRA build working set); + // otherwise fall back to a host read + CAGRA's internal host->device streaming. + auto read_end = start; + bool used_device_read = false; + const size_t sub_ds_bytes = sub_dataset_size * dataset_dim * sizeof(T); + auto sub_index = [&]() { + if (use_disk_mode) { + const size_t free_gpu_bytes = + ace_params.max_gpu_memory_gb > 0 + ? static_cast(ace_params.max_gpu_memory_gb * (1ULL << 30)) + : rmm::available_device_memory().first; + if (sub_ds_bytes < static_cast(0.4 * free_gpu_bytes)) { + try { + auto sub_dataset_dev = + raft::make_device_matrix(res, sub_dataset_size, dataset_dim); + ace_load_partition_dataset_from_disk(partition_id, + dataset_dim, + partition_histogram.view(), + core_partition_offsets.view(), + augmented_partition_offsets.view(), + reordered_fd, + augmented_fd, + reordered_header_size, + augmented_header_size, + sub_dataset_dev.data_handle()); + used_device_read = true; + read_end = std::chrono::high_resolution_clock::now(); + return cuvs::neighbors::cagra::build( + res, sub_index_params, raft::make_const_mdspan(sub_dataset_dev.view())); + } catch (const std::bad_alloc&) { + RAFT_LOG_WARN( + "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&) { + RAFT_LOG_WARN( + "ACE: device allocation failed for partition %lu direct read; falling back to a " + "host read", + partition_id); + } + } + auto sub_dataset = raft::make_host_matrix(sub_dataset_size, dataset_dim); + ace_load_partition_dataset_from_disk(partition_id, + dataset_dim, + partition_histogram.view(), + core_partition_offsets.view(), + augmented_partition_offsets.view(), + reordered_fd, + augmented_fd, + reordered_header_size, + augmented_header_size, + sub_dataset.data_handle()); + read_end = std::chrono::high_resolution_clock::now(); + return cuvs::neighbors::cagra::build( + res, sub_index_params, raft::make_const_mdspan(sub_dataset.view())); + } + // In-memory mode: gather this partition's vectors from the host dataset. + auto sub_dataset = raft::make_host_matrix(sub_dataset_size, dataset_dim); ace_gather_partition_dataset(core_sub_dataset_size, augmented_sub_dataset_size, dataset_dim, @@ -1375,25 +1420,17 @@ index build_ace(raft::resources const& res, core_partition_offsets.view(), augmented_partition_offsets.view(), sub_dataset.view()); + read_end = std::chrono::high_resolution_clock::now(); + return cuvs::neighbors::cagra::build( + res, sub_index_params, raft::make_const_mdspan(sub_dataset.view())); + }(); + if (used_device_read) { + RAFT_LOG_DEBUG("ACE: partition %lu read directly into device memory (GDS path)", + partition_id); } - auto read_end = std::chrono::high_resolution_clock::now(); auto read_elapsed = std::chrono::duration_cast(read_end - start).count(); - // Create index for this partition - cuvs::neighbors::cagra::index_params sub_index_params; - sub_index_params = cuvs::neighbors::cagra::index_params::from_hnsw_params( - raft::make_extents(sub_dataset_size, dataset_dim), - graph_degree / 2, - ef_construction, - cuvs::neighbors::cagra::hnsw_heuristic_type::SAME_GRAPH_FOOTPRINT, - params.metric); - sub_index_params.attach_dataset_on_build = false; - sub_index_params.guarantee_connectivity = params.guarantee_connectivity; - - auto sub_index = cuvs::neighbors::cagra::build( - res, sub_index_params, raft::make_const_mdspan(sub_dataset.view())); - auto optimize_end = std::chrono::high_resolution_clock::now(); auto optimize_elapsed = std::chrono::duration_cast(optimize_end - read_end).count(); diff --git a/cpp/src/neighbors/detail/hnsw.hpp b/cpp/src/neighbors/detail/hnsw.hpp index d2ceee7bd2..789a1d09a6 100644 --- a/cpp/src/neighbors/detail/hnsw.hpp +++ b/cpp/src/neighbors/detail/hnsw.hpp @@ -14,6 +14,8 @@ #include #include +#include + #include #include #include @@ -722,9 +724,12 @@ void serialize_to_hnswlib_from_disk(raft::resources const& res, RAFT_EXPECTS(!dataset_path.empty(), "Unable to get path from dataset file descriptor"); RAFT_EXPECTS(!mapping_path.empty(), "Unable to get path from mapping file descriptor"); - int graph_fd = graph_fd_opt->get(); - int dataset_fd = dataset_fd_opt->get(); - int label_fd = mapping_fd_opt->get(); + // Open kvikio handles for the disk-backed artifacts. Reads here target host buffers (the hnswlib + // layout is assembled on the CPU), so kvikio uses its POSIX + threadpool backend, with O_DIRECT + // when available; it handles any alignment internally. + kvikio::FileHandle graph_kv(graph_path, "r"); + kvikio::FileHandle dataset_kv(dataset_path, "r"); + kvikio::FileHandle label_kv(mapping_path, "r"); // Read headers from files to get dimensions size_t graph_header_size = 0; @@ -813,44 +818,47 @@ void serialize_to_hnswlib_from_disk(raft::resources const& res, raft::host_matrix_view graph_buf, raft::host_matrix_view dataset_buf, raft::host_vector_view label_buf) { - const size_t graph_bytes = rows_to_read * graph_degree_int * sizeof(IdxT); - const size_t dataset_bytes = rows_to_read * dim * sizeof(T); - const size_t label_bytes = rows_to_read * sizeof(uint32_t); + RAFT_EXPECTS(start_row >= 0 && rows_to_read >= 0, + "Batch start row and row count must be non-negative"); + const size_t row = static_cast(start_row); + const size_t rows = static_cast(rows_to_read); + const size_t graph_degree = static_cast(graph_degree_int); + const size_t dim_size = static_cast(dim); + + const size_t graph_bytes = rows * graph_degree * sizeof(IdxT); + const size_t dataset_bytes = rows * dim_size * sizeof(T); + const size_t label_bytes = rows * sizeof(uint32_t); - const off_t graph_offset = graph_header_size + start_row * graph_degree_int * sizeof(IdxT); - const off_t dataset_offset = dataset_header_size + start_row * dim * sizeof(T); - const off_t label_offset = label_header_size + start_row * sizeof(uint32_t); + const size_t graph_offset = graph_header_size + row * graph_degree * sizeof(IdxT); + const size_t dataset_offset = dataset_header_size + row * dim_size * sizeof(T); + const size_t label_offset = label_header_size + row * sizeof(uint32_t); RAFT_LOG_DEBUG("Reading batch: row=%ld, rows=%ld", start_row, rows_to_read); -#pragma omp parallel sections num_threads(3) - { -#pragma omp section - { - ssize_t bytes_read = pread(graph_fd, graph_buf.data_handle(), graph_bytes, graph_offset); - RAFT_EXPECTS(bytes_read == static_cast(graph_bytes), - "Failed to read graph data: expected %zu, got %zd", - graph_bytes, - bytes_read); - } -#pragma omp section - { - ssize_t bytes_read = - pread(dataset_fd, dataset_buf.data_handle(), dataset_bytes, dataset_offset); - RAFT_EXPECTS(bytes_read == static_cast(dataset_bytes), - "Failed to read dataset data: expected %zu, got %zd", - dataset_bytes, - bytes_read); - } -#pragma omp section - { - ssize_t bytes_read = pread(label_fd, label_buf.data_handle(), label_bytes, label_offset); - RAFT_EXPECTS(bytes_read == static_cast(label_bytes), - "Failed to read label data: expected %zu, got %zd", - label_bytes, - bytes_read); - } - } + // Issue the three reads concurrently through kvikio (its threadpool parallelizes each), then + // wait for all to complete. + auto graph_future = graph_kv.pread(graph_buf.data_handle(), graph_bytes, graph_offset); + auto dataset_future = + dataset_kv.pread(dataset_buf.data_handle(), dataset_bytes, dataset_offset); + auto label_future = label_kv.pread(label_buf.data_handle(), label_bytes, label_offset); + const size_t graph_read = graph_future.get(); + const size_t dataset_read = dataset_future.get(); + const size_t label_read = label_future.get(); + RAFT_EXPECTS(graph_read == graph_bytes, + "Short graph read at row %ld: expected %zu, got %zu", + start_row, + graph_bytes, + graph_read); + RAFT_EXPECTS(dataset_read == dataset_bytes, + "Short dataset read at row %ld: expected %zu, got %zu", + start_row, + dataset_bytes, + dataset_read); + RAFT_EXPECTS(label_read == label_bytes, + "Short label read at row %ld: expected %zu, got %zu", + start_row, + label_bytes, + label_read); }; serialize_to_hnswlib_batched( @@ -1258,7 +1266,10 @@ std::unique_ptr> from_cagra( std::string index_filename = (std::filesystem::path(index_directory) / "hnsw_index.bin").string(); - std::ofstream of(index_filename, std::ios::out | std::ios::binary); + // Route the disk-backed hnswlib output through kvikio (bypassing the page cache via O_DIRECT + // when supported) so all ACE disk I/O goes through kvikio. kvikio_ofstream is a std::ostream + // backed by a kvikio writer, so it plugs into the shared serializer below. + cuvs::util::kvikio_ofstream of(index_filename); RAFT_EXPECTS(of, "Cannot open file %s", index_filename.c_str()); diff --git a/cpp/src/util/file_io.cpp b/cpp/src/util/file_io.cpp index d924527e72..459094e481 100644 --- a/cpp/src/util/file_io.cpp +++ b/cpp/src/util/file_io.cpp @@ -1,16 +1,17 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ #include +#include + #include +#include #include -#include - -#include -#include +#include +#include namespace cuvs::util { @@ -22,28 +23,19 @@ void read_large_file(const file_descriptor& fd, RAFT_EXPECTS(total_bytes > 0, "Total bytes must be greater than 0"); 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"); - const size_t read_chunk_size = std::min(1024 * 1024 * 1024, SSIZE_MAX); - size_t bytes_remaining = total_bytes; - size_t offset = 0; - - while (bytes_remaining > 0) { - const size_t chunk_size = std::min(read_chunk_size, bytes_remaining); - const uint64_t file_pos = file_offset + offset; - const ssize_t bytes_read = - pread(fd.get(), reinterpret_cast(dest_ptr) + offset, chunk_size, file_pos); - - RAFT_EXPECTS( - bytes_read != -1, "Failed to read from file at offset %lu: %s", file_pos, strerror(errno)); - RAFT_EXPECTS(bytes_read == static_cast(chunk_size), - "Incomplete read from file. Expected %zu bytes, got %zd at offset %lu", - chunk_size, - bytes_read, - file_pos); - - bytes_remaining -= chunk_size; - offset += chunk_size; - } + // kvikio selects GPUDirect Storage (cuFile) for device destinations on a GDS-capable system, and + // the POSIX + threadpool backend (with O_DIRECT when available) otherwise. The destination may be + // host or device memory; kvikio detects which. + kvikio::FileHandle handle(path, "r"); + const size_t bytes_read = handle.pread(dest_ptr, total_bytes, file_offset).get(); + RAFT_EXPECTS(bytes_read == total_bytes, + "Incomplete read from %s: expected %zu bytes, got %zu", + path.c_str(), + total_bytes, + bytes_read); } void write_large_file(const file_descriptor& fd, @@ -54,28 +46,164 @@ void write_large_file(const file_descriptor& fd, RAFT_EXPECTS(total_bytes > 0, "Total bytes must be greater than 0"); RAFT_EXPECTS(data_ptr != nullptr, "Data 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"); + + // Open in read+write mode ("r+") so the existing numpy header and preallocation are preserved + // (kvikio's "w" mode would truncate). The source may be host or device memory. + kvikio::FileHandle handle(path, "r+"); + const size_t bytes_written = handle.pwrite(data_ptr, total_bytes, file_offset).get(); + RAFT_EXPECTS(bytes_written == total_bytes, + "Incomplete write to %s: expected %zu bytes, wrote %zu", + path.c_str(), + total_bytes, + bytes_written); +} + +// std::streambuf that stages output into a large buffer and writes full buffers to disk through +// kvikio at an increasing file offset. The trailing partial buffer is written on sync()/close(). +class kvikio_ofstream::sbuf : public std::streambuf { + public: + sbuf(const std::string& path, size_t cap) + : handle_(path, "w"), buffer_(std::max(cap, kNumpyDataAlignment)) + { + RAFT_EXPECTS(buffer_.size() <= static_cast(std::numeric_limits::max()), + "kvikio_ofstream buffer size must fit in std::streambuf::pbump"); + setp(buffer_.data(), buffer_.data() + buffer_.size()); + } + + ~sbuf() override + { + try { + close(); + } catch (...) { + // Swallow during destruction. + } + } + + void close() + { + if (closed_) { return; } + flush_buffer(); + handle_.close(); + closed_ = true; + } + + // Total logical bytes accepted = already-written + currently-staged. + [[nodiscard]] size_t bytes_written() const noexcept + { + return offset_ + static_cast(pptr() - pbase()); + } + + protected: + int_type overflow(int_type ch) override + { + RAFT_EXPECTS(!closed_, "kvikio_ofstream: write attempted after close"); + flush_buffer(); + if (!traits_type::eq_int_type(ch, traits_type::eof())) { + *pptr() = traits_type::to_char_type(ch); + pbump(1); + } + return traits_type::not_eof(ch); + } + + std::streamsize xsputn(const char* input, std::streamsize count) override + { + RAFT_EXPECTS(!closed_, "kvikio_ofstream: write attempted after close"); + if (count <= 0) { return 0; } + + const auto requested = count; + auto* current = input; + size_t remaining = static_cast(count); + + while (remaining > 0) { + // If the caller hands us a large contiguous chunk, flush pending staged bytes and pass the + // chunk straight to kvikio. This avoids std::streambuf's byte-at-a-time fallback path. + if (remaining >= buffer_.size()) { + flush_buffer(); + write_direct(current, remaining); + return requested; + } + + size_t available = static_cast(epptr() - pptr()); + if (available == 0) { + flush_buffer(); + available = static_cast(epptr() - pptr()); + } + + const size_t n = std::min(remaining, available); + std::memcpy(pptr(), current, n); + pbump(static_cast(n)); + current += n; + remaining -= n; + } - const size_t write_chunk_size = std::min(1024 * 1024 * 1024, SSIZE_MAX); - size_t bytes_remaining = total_bytes; - size_t offset = 0; - - while (bytes_remaining > 0) { - const size_t chunk_size = std::min(write_chunk_size, bytes_remaining); - const uint64_t file_pos = file_offset + offset; - const ssize_t chunk_written = - pwrite(fd.get(), reinterpret_cast(data_ptr) + offset, chunk_size, file_pos); - - RAFT_EXPECTS( - chunk_written != -1, "Failed to write to file at offset %lu: %s", file_pos, strerror(errno)); - RAFT_EXPECTS(chunk_written == static_cast(chunk_size), - "Incomplete write to file. Expected %zu bytes, wrote %zd at offset %lu", - chunk_size, - chunk_written, - file_pos); - - bytes_remaining -= chunk_size; - offset += chunk_size; + return requested; + } + + int sync() override + { + try { + flush_buffer(); + return 0; + } catch (...) { + return -1; + } + } + + // Support tellp(): report the current output position. + pos_type seekoff(off_type off, std::ios_base::seekdir dir, std::ios_base::openmode which) override + { + if ((which & std::ios_base::out) && dir == std::ios_base::cur && off == 0) { + return pos_type(static_cast(bytes_written())); + } + return pos_type(off_type(-1)); + } + + private: + void flush_buffer() + { + const size_t n = static_cast(pptr() - pbase()); + if (n > 0) { + write_direct(pbase(), n); + setp(buffer_.data(), buffer_.data() + buffer_.size()); + } + } + + void write_direct(const char* data, size_t n) + { + RAFT_EXPECTS(!closed_, "kvikio_ofstream: write attempted after close"); + const size_t w = handle_.pwrite(data, n, offset_).get(); + RAFT_EXPECTS(w == n, "kvikio_ofstream: short write (expected %zu, wrote %zu)", n, w); + offset_ += n; + } + + kvikio::FileHandle handle_; + std::vector buffer_; + size_t offset_ = 0; + bool closed_ = false; +}; + +kvikio_ofstream::kvikio_ofstream(const std::string& path, size_t buffer_size) + : std::ostream(nullptr), buf_(std::make_unique(path, buffer_size)) +{ + rdbuf(buf_.get()); +} + +kvikio_ofstream::~kvikio_ofstream() +{ + try { + if (buf_) { buf_->close(); } + } catch (...) { + // Swallow during destruction. } } +void kvikio_ofstream::close() +{ + if (buf_) { buf_->close(); } +} + +size_t kvikio_ofstream::bytes_written() const noexcept { return buf_ ? buf_->bytes_written() : 0; } + } // namespace cuvs::util diff --git a/cpp/tests/CMakeLists.txt b/cpp/tests/CMakeLists.txt index 13a07b10b5..ec2dbddf4a 100644 --- a/cpp/tests/CMakeLists.txt +++ b/cpp/tests/CMakeLists.txt @@ -393,6 +393,13 @@ ConfigureTest( PERCENT 100 ) +ConfigureTest( + NAME UTIL_TEST + PATH util/file_io_test.cpp + GPUS 1 + PERCENT 100 +) + # ################################################################################################## # Install tests #################################################################################### # ################################################################################################## diff --git a/cpp/tests/util/file_io_test.cpp b/cpp/tests/util/file_io_test.cpp new file mode 100644 index 0000000000..79809e79c5 --- /dev/null +++ b/cpp/tests/util/file_io_test.cpp @@ -0,0 +1,218 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include + +#include +#include +#include +#include +#include + +#include + +#include +#include +#include +#include +#include +#include +#include + +namespace cuvs::util { + +namespace { + +// Deterministic byte pattern so reads can be validated independently of the writer. +std::vector make_pattern(size_t n, uint8_t seed) +{ + std::vector v(n); + for (size_t i = 0; i < n; i++) { + v[i] = static_cast((i * 131u + seed * 7u + 17u) & 0xFFu); + } + return v; +} + +// Create a unique writable scratch directory. Prefer the current working directory (typically the +// on-disk build tree, where O_DIRECT is usually supported) and fall back to the system temp dir. +class scratch_dir { + public: + scratch_dir() + { + const std::string name = ".cuvs_file_io_test_" + std::to_string(::getpid()); + std::error_code ec; + std::filesystem::path base = std::filesystem::current_path(ec); + if (ec) { base = std::filesystem::temp_directory_path(); } + path_ = base / name; + std::filesystem::create_directories(path_, ec); + if (ec || !std::filesystem::is_directory(path_)) { + path_ = std::filesystem::temp_directory_path() / name; + std::filesystem::create_directories(path_); + } + } + + ~scratch_dir() + { + std::error_code ec; + std::filesystem::remove_all(path_, ec); + } + + [[nodiscard]] std::string file(const std::string& name) const { return (path_ / name).string(); } + [[nodiscard]] std::string dir() const { return path_.string(); } + + private: + std::filesystem::path path_; +}; + +std::vector read_whole_file(const std::string& path) +{ + std::ifstream is(path, std::ios::in | std::ios::binary); + EXPECT_TRUE(is.good()) << "cannot open " << path; + return std::vector((std::istreambuf_iterator(is)), std::istreambuf_iterator()); +} + +} // namespace + +// create_numpy_file must produce a numpy-compatible header whose data body begins on a block +// boundary (so kvikio's O_DIRECT / GDS interior is aligned), and the file must be readable back +// through the numpy deserializer. +TEST(FileIO, CreateNumpyFileAlignedHeader) +{ + scratch_dir scratch; + const std::string path = scratch.file("aligned.npy"); + const size_t rows = 1234; + const size_t cols = 17; + auto [fd, header_size] = create_numpy_file(path, {rows, cols}); + EXPECT_TRUE(fd.is_valid()); + EXPECT_EQ(header_size % kNumpyDataAlignment, 0u) + << "numpy data body must start on a " << kNumpyDataAlignment << "-byte boundary"; + + // The numpy deserializer must recover the shape from the (re-padded) header. + auto stream = fd.make_istream(); + auto header = raft::detail::numpy_serializer::read_header(stream); + ASSERT_EQ(header.shape.size(), 2u); + EXPECT_EQ(header.shape[0], rows); + EXPECT_EQ(header.shape[1], cols); +} + +// write_large_file followed by read_large_file with host buffers must round-trip the data exactly. +TEST(FileIO, HostReadWriteRoundTrip) +{ + scratch_dir scratch; + const std::string path = scratch.file("host.npy"); + const size_t rows = 4096; + const size_t cols = 33; + const size_t n = rows * cols; + + auto [fd, header_size] = create_numpy_file(path, {rows, cols}); + + std::vector src(n); + for (size_t i = 0; i < n; i++) { + src[i] = static_cast(i % 1000) * 0.5f; + } + write_large_file(fd, src.data(), n * sizeof(float), header_size); + + std::vector dst(n, -1.0f); + read_large_file(fd, dst.data(), n * sizeof(float), header_size); + + EXPECT_EQ(src, dst); +} + +// read_large_file must also fill device memory (kvikio uses GPUDirect Storage when available, and +// stages through a host bounce buffer in compatibility mode). The data must match after copying +// back to the host. +TEST(FileIO, DeviceReadRoundTrip) +{ + raft::resources res; + scratch_dir scratch; + const std::string path = scratch.file("device.npy"); + const size_t rows = 5000; + const size_t cols = 24; + const size_t n = rows * cols; + + auto [fd, header_size] = create_numpy_file(path, {rows, cols}); + + std::vector src(n); + for (size_t i = 0; i < n; i++) { + src[i] = static_cast((i * 7) % 4096); + } + write_large_file(fd, src.data(), n * sizeof(float), header_size); + + auto dev = raft::make_device_vector(res, n); + read_large_file(fd, dev.data_handle(), n * sizeof(float), header_size); + + std::vector dst(n, -1.0f); + raft::copy(dst.data(), dev.data_handle(), n, raft::resource::get_cuda_stream(res)); + raft::resource::sync_stream(res); + + EXPECT_EQ(src, dst); +} + +// kvikio_ofstream must reproduce exactly the bytes handed to write(), across many small writes that +// span multiple internal buffer flushes, and report the correct logical size. +TEST(FileIO, KvikioOfstreamRoundTrip) +{ + scratch_dir scratch; + const std::string path = scratch.file("stream.bin"); + std::vector data = make_pattern(5'000'003, 42); + const size_t cap = size_t(1) << 20; // 1 MiB staging buffer -> several flushes + + { + kvikio_ofstream os(path, cap); + size_t pos = 0; + while (pos < data.size()) { + const size_t chunk = std::min(7777, data.size() - pos); + os.write(data.data() + pos, chunk); + pos += chunk; + } + os.flush(); + EXPECT_EQ(os.bytes_written(), data.size()); + EXPECT_EQ(os.tellp(), std::streampos(static_cast(data.size()))); + os.close(); + } + + const std::vector got = read_whole_file(path); + ASSERT_EQ(got.size(), data.size()); + EXPECT_EQ(got, data); +} + +// A large single write must bypass the small staging buffer correctly and still round-trip. +TEST(FileIO, KvikioOfstreamLargeSingleWrite) +{ + scratch_dir scratch; + const std::string path = scratch.file("stream_large.bin"); + std::vector data = make_pattern((size_t(8) << 20) + 123, 7); + + { + kvikio_ofstream os(path, size_t(1) << 20); + os.write(data.data(), data.size()); + EXPECT_EQ(os.bytes_written(), data.size()); + EXPECT_EQ(os.tellp(), std::streampos(static_cast(data.size()))); + os.close(); + } + + const std::vector got = read_whole_file(path); + ASSERT_EQ(got.size(), data.size()); + EXPECT_EQ(got, data); +} + +// Defensive checks on the bulk helpers. +TEST(FileIO, InvalidArguments) +{ + scratch_dir scratch; + auto [fd, header_size] = create_numpy_file(scratch.file("args.npy"), {16, 4}); + char buf[8] = {0}; + EXPECT_THROW(read_large_file(fd, buf, 0, header_size), raft::exception); + EXPECT_THROW(write_large_file(fd, buf, 0, header_size), raft::exception); + EXPECT_THROW(read_large_file(fd, nullptr, 8, header_size), raft::exception); + + EXPECT_NO_THROW(read_large_file(fd, buf, sizeof(buf), header_size)); + const int dup_fd = ::dup(fd.get()); + ASSERT_NE(dup_fd, -1); + file_descriptor pathless_fd(dup_fd); + EXPECT_THROW(read_large_file(pathless_fd, buf, sizeof(buf), header_size), raft::exception); +} + +} // namespace cuvs::util diff --git a/dependencies.yaml b/dependencies.yaml index 744e4d9227..820813c39c 100644 --- a/dependencies.yaml +++ b/dependencies.yaml @@ -17,6 +17,7 @@ files: - depends_on_cuda_python - depends_on_cupy - depends_on_librmm + - depends_on_libkvikio - depends_on_pylibraft - depends_on_nccl - docs @@ -140,6 +141,7 @@ files: includes: - depends_on_libraft - depends_on_librmm + - depends_on_libkvikio - depends_on_nccl - rapids_build py_run_libcuvs: @@ -151,6 +153,7 @@ files: - cuda_wheels - depends_on_libraft - depends_on_librmm + - depends_on_libkvikio - depends_on_nccl py_build_cuvs: output: pyproject @@ -400,6 +403,24 @@ dependencies: - libcusolver-dev - libcusparse-dev - libnvjitlink-dev + specific: + # libcufile provides the GPUDirect Storage (cuFile) backend that kvikio compiles against. + # It is packaged for CUDA 12 on x86_64 only, and for CUDA 13 on all arches. + - output_types: [conda] + matrices: + - matrix: + cuda: "12.*" + arch: x86_64 + packages: + - libcufile-dev + - matrix: + cuda: "12.*" + arch: aarch64 + packages: + - matrix: + cuda: "13.*" + packages: + - libcufile-dev cuda_wheels: specific: # cuVS needs 'nvJitLink>={whatever-cuvs-was-built-against}' at runtime, and mixing @@ -686,6 +707,29 @@ dependencies: packages: - librmm-cu13==26.8.*,>=0.0.0a0 - {matrix: null, packages: [*librmm_unsuffixed]} + depends_on_libkvikio: + common: + - output_types: conda + packages: + - &libkvikio_unsuffixed libkvikio==26.8.*,>=0.0.0a0 + - output_types: requirements + packages: + # pip recognizes the index as a global option for the requirements.txt file + - --extra-index-url=https://pypi.anaconda.org/rapidsai-wheels-nightly/simple + specific: + - output_types: [requirements, pyproject] + matrices: + - matrix: + cuda: "12.*" + cuda_suffixed: "true" + packages: + - libkvikio-cu12==26.8.*,>=0.0.0a0 + - matrix: + cuda: "13.*" + cuda_suffixed: "true" + packages: + - libkvikio-cu13==26.8.*,>=0.0.0a0 + - {matrix: null, packages: [*libkvikio_unsuffixed]} depends_on_pylibraft: common: - output_types: conda diff --git a/python/libcuvs/pyproject.toml b/python/libcuvs/pyproject.toml index 5025daa66d..d41e39976f 100644 --- a/python/libcuvs/pyproject.toml +++ b/python/libcuvs/pyproject.toml @@ -20,6 +20,7 @@ license = "Apache-2.0" requires-python = ">=3.11" dependencies = [ "cuda-toolkit[cublas,curand,cusolver,cusparse,nvrtc]==13.*", + "libkvikio==26.8.*,>=0.0.0a0", "libraft==26.8.*,>=0.0.0a0", "librmm==26.8.*,>=0.0.0a0", "nvidia-nvjitlink>=13.0,<14", @@ -81,6 +82,7 @@ regex = "(?P.*)" build-backend = "scikit_build_core.build" requires = [ "cmake>=4.0", + "libkvikio==26.8.*,>=0.0.0a0", "libraft==26.8.*,>=0.0.0a0", "librmm==26.8.*,>=0.0.0a0", "ninja", From c6df26012592f58cccb2b76d7c7a36b3326b24f4 Mon Sep 17 00:00:00 2001 From: Julian Miller Date: Wed, 24 Jun 2026 10:30:24 +0200 Subject: [PATCH 02/18] Address review feedback --- ci/release/update-version.sh | 1 + .../all_cuda-129_arch-x86_64.yaml | 1 - .../all_cuda-133_arch-aarch64.yaml | 1 - .../all_cuda-133_arch-x86_64.yaml | 1 - .../bench_ann_cuda-129_arch-x86_64.yaml | 1 - .../bench_ann_cuda-133_arch-aarch64.yaml | 1 - .../bench_ann_cuda-133_arch-x86_64.yaml | 1 - .../environments/go_cuda-129_arch-x86_64.yaml | 1 - .../go_cuda-133_arch-aarch64.yaml | 1 - .../environments/go_cuda-133_arch-x86_64.yaml | 1 - .../rust_cuda-129_arch-x86_64.yaml | 1 - .../rust_cuda-133_arch-aarch64.yaml | 1 - .../rust_cuda-133_arch-x86_64.yaml | 1 - conda/recipes/libcuvs/recipe.yaml | 17 ------- cpp/CMakeLists.txt | 7 ++- cpp/src/neighbors/detail/hnsw.hpp | 21 ++++++-- cpp/src/util/file_io.cpp | 49 +++++++++++++++++++ cpp/tests/util/file_io_test.cpp | 23 +++++++++ dependencies.yaml | 18 ------- python/libcuvs/libcuvs/load.py | 4 +- 20 files changed, 98 insertions(+), 54 deletions(-) diff --git a/ci/release/update-version.sh b/ci/release/update-version.sh index f36c14f771..de1802adb4 100755 --- a/ci/release/update-version.sh +++ b/ci/release/update-version.sh @@ -110,6 +110,7 @@ DEPENDENCIES=( cuvs-bench libcuvs libcuvs-tests + libkvikio libraft librmm pylibraft diff --git a/conda/environments/all_cuda-129_arch-x86_64.yaml b/conda/environments/all_cuda-129_arch-x86_64.yaml index f372dfdeb7..9e8d26acf8 100644 --- a/conda/environments/all_cuda-129_arch-x86_64.yaml +++ b/conda/environments/all_cuda-129_arch-x86_64.yaml @@ -24,7 +24,6 @@ dependencies: - go - libclang==20.1.8 - libcublas-dev -- libcufile-dev - libcurand-dev - libcusolver-dev - libcusparse-dev diff --git a/conda/environments/all_cuda-133_arch-aarch64.yaml b/conda/environments/all_cuda-133_arch-aarch64.yaml index 1da8d32820..3befeac713 100644 --- a/conda/environments/all_cuda-133_arch-aarch64.yaml +++ b/conda/environments/all_cuda-133_arch-aarch64.yaml @@ -24,7 +24,6 @@ dependencies: - go - libclang==20.1.8 - libcublas-dev -- libcufile-dev - libcurand-dev - libcusolver-dev - libcusparse-dev diff --git a/conda/environments/all_cuda-133_arch-x86_64.yaml b/conda/environments/all_cuda-133_arch-x86_64.yaml index 77173ae228..63508294d5 100644 --- a/conda/environments/all_cuda-133_arch-x86_64.yaml +++ b/conda/environments/all_cuda-133_arch-x86_64.yaml @@ -24,7 +24,6 @@ dependencies: - go - libclang==20.1.8 - libcublas-dev -- libcufile-dev - libcurand-dev - libcusolver-dev - libcusparse-dev diff --git a/conda/environments/bench_ann_cuda-129_arch-x86_64.yaml b/conda/environments/bench_ann_cuda-129_arch-x86_64.yaml index 009f31cd5f..d84bf87234 100644 --- a/conda/environments/bench_ann_cuda-129_arch-x86_64.yaml +++ b/conda/environments/bench_ann_cuda-129_arch-x86_64.yaml @@ -28,7 +28,6 @@ dependencies: - libboost-devel=1.87 - libclang==20.1.8 - libcublas-dev -- libcufile-dev - libcurand-dev - libcusolver-dev - libcusparse-dev diff --git a/conda/environments/bench_ann_cuda-133_arch-aarch64.yaml b/conda/environments/bench_ann_cuda-133_arch-aarch64.yaml index 31c5e64c21..c874298d7d 100644 --- a/conda/environments/bench_ann_cuda-133_arch-aarch64.yaml +++ b/conda/environments/bench_ann_cuda-133_arch-aarch64.yaml @@ -26,7 +26,6 @@ dependencies: - h5py>=3.8.0 - libclang==20.1.8 - libcublas-dev -- libcufile-dev - libcurand-dev - libcusolver-dev - libcusparse-dev diff --git a/conda/environments/bench_ann_cuda-133_arch-x86_64.yaml b/conda/environments/bench_ann_cuda-133_arch-x86_64.yaml index e2a344e6d6..43b2dbce46 100644 --- a/conda/environments/bench_ann_cuda-133_arch-x86_64.yaml +++ b/conda/environments/bench_ann_cuda-133_arch-x86_64.yaml @@ -28,7 +28,6 @@ dependencies: - libboost-devel=1.87 - libclang==20.1.8 - libcublas-dev -- libcufile-dev - libcurand-dev - libcusolver-dev - libcusparse-dev diff --git a/conda/environments/go_cuda-129_arch-x86_64.yaml b/conda/environments/go_cuda-129_arch-x86_64.yaml index eb9fcbf319..e99615ab2c 100644 --- a/conda/environments/go_cuda-129_arch-x86_64.yaml +++ b/conda/environments/go_cuda-129_arch-x86_64.yaml @@ -22,7 +22,6 @@ dependencies: - go - libclang==20.1.8 - libcublas-dev -- libcufile-dev - libcurand-dev - libcusolver-dev - libcusparse-dev diff --git a/conda/environments/go_cuda-133_arch-aarch64.yaml b/conda/environments/go_cuda-133_arch-aarch64.yaml index 05b00c1537..eeb5511537 100644 --- a/conda/environments/go_cuda-133_arch-aarch64.yaml +++ b/conda/environments/go_cuda-133_arch-aarch64.yaml @@ -22,7 +22,6 @@ dependencies: - go - libclang==20.1.8 - libcublas-dev -- libcufile-dev - libcurand-dev - libcusolver-dev - libcusparse-dev diff --git a/conda/environments/go_cuda-133_arch-x86_64.yaml b/conda/environments/go_cuda-133_arch-x86_64.yaml index 6b02c21a87..5fd2153ffa 100644 --- a/conda/environments/go_cuda-133_arch-x86_64.yaml +++ b/conda/environments/go_cuda-133_arch-x86_64.yaml @@ -22,7 +22,6 @@ dependencies: - go - libclang==20.1.8 - libcublas-dev -- libcufile-dev - libcurand-dev - libcusolver-dev - libcusparse-dev diff --git a/conda/environments/rust_cuda-129_arch-x86_64.yaml b/conda/environments/rust_cuda-129_arch-x86_64.yaml index 13c2c77147..a9e9e2142c 100644 --- a/conda/environments/rust_cuda-129_arch-x86_64.yaml +++ b/conda/environments/rust_cuda-129_arch-x86_64.yaml @@ -16,7 +16,6 @@ dependencies: - cxx-compiler - gcc_linux-64=14.* - libcublas-dev -- libcufile-dev - libcurand-dev - libcusolver-dev - libcusparse-dev diff --git a/conda/environments/rust_cuda-133_arch-aarch64.yaml b/conda/environments/rust_cuda-133_arch-aarch64.yaml index 0223fba1ca..3daa10d2be 100644 --- a/conda/environments/rust_cuda-133_arch-aarch64.yaml +++ b/conda/environments/rust_cuda-133_arch-aarch64.yaml @@ -16,7 +16,6 @@ dependencies: - cxx-compiler - gcc_linux-aarch64=14.* - libcublas-dev -- libcufile-dev - libcurand-dev - libcusolver-dev - libcusparse-dev diff --git a/conda/environments/rust_cuda-133_arch-x86_64.yaml b/conda/environments/rust_cuda-133_arch-x86_64.yaml index 8aa9341870..9e0fe2d529 100644 --- a/conda/environments/rust_cuda-133_arch-x86_64.yaml +++ b/conda/environments/rust_cuda-133_arch-x86_64.yaml @@ -16,7 +16,6 @@ dependencies: - cxx-compiler - gcc_linux-64=14.* - libcublas-dev -- libcufile-dev - libcurand-dev - libcusolver-dev - libcusparse-dev diff --git a/conda/recipes/libcuvs/recipe.yaml b/conda/recipes/libcuvs/recipe.yaml index 3d645fa093..7d0542984e 100644 --- a/conda/recipes/libcuvs/recipe.yaml +++ b/conda/recipes/libcuvs/recipe.yaml @@ -76,11 +76,6 @@ cache: - librmm =${{ minor_version }} - libraft-headers =${{ minor_version }} - libkvikio =${{ minor_version }} - # libcufile (GPUDirect Storage backend that kvikio compiles against): CUDA 13 on all arches, - # CUDA 12 on x86_64 only. - - if: cuda_major == "13" or linux64 - then: - - libcufile-dev - nccl ${{ nccl_version }} - cuda-version =${{ cuda_version }} - cuda-cudart-dev @@ -181,9 +176,6 @@ outputs: - librmm =${{ minor_version }} - libraft-headers =${{ minor_version }} - libkvikio =${{ minor_version }} - - if: cuda_major == "13" or linux64 - then: - - libcufile-dev - nccl ${{ nccl_version }} - cuda-version =${{ cuda_version }} - cuda-cudart-dev @@ -200,9 +192,6 @@ outputs: - libraft-headers =${{ minor_version }} - librmm =${{ minor_version }} - libkvikio =${{ minor_version }} - - if: cuda_major == "13" or linux64 - then: - - libcufile - cuda-nvrtc - nccl - libcublas @@ -250,9 +239,6 @@ outputs: - librmm =${{ minor_version }} - libraft-headers =${{ minor_version }} - libkvikio =${{ minor_version }} - - if: cuda_major == "13" or linux64 - then: - - libcufile-dev - nccl ${{ nccl_version }} - cuda-version =${{ cuda_version }} - cuda-cudart-dev @@ -269,9 +255,6 @@ outputs: - libraft-headers =${{ minor_version }} - librmm =${{ minor_version }} - libkvikio =${{ minor_version }} - - if: cuda_major == "13" or linux64 - then: - - libcufile - cuda-nvrtc - nccl - libcublas diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt index 39f4139985..ecbf377fcc 100644 --- a/cpp/CMakeLists.txt +++ b/cpp/CMakeLists.txt @@ -1541,8 +1541,11 @@ if(NOT BUILD_CPU_ONLY) $> $> $<$:CUDA::nvtx3> - PRIVATE $ $ - $ CUDA::nvJitLink CUDA::nvrtc + PRIVATE $ + $ + $ + CUDA::nvJitLink + CUDA::nvrtc $ ) diff --git a/cpp/src/neighbors/detail/hnsw.hpp b/cpp/src/neighbors/detail/hnsw.hpp index 789a1d09a6..8e55d56a4d 100644 --- a/cpp/src/neighbors/detail/hnsw.hpp +++ b/cpp/src/neighbors/detail/hnsw.hpp @@ -29,6 +29,7 @@ #include #include +#include #include #include #include @@ -840,10 +841,22 @@ void serialize_to_hnswlib_from_disk(raft::resources const& res, auto graph_future = graph_kv.pread(graph_buf.data_handle(), graph_bytes, graph_offset); auto dataset_future = dataset_kv.pread(dataset_buf.data_handle(), dataset_bytes, dataset_offset); - auto label_future = label_kv.pread(label_buf.data_handle(), label_bytes, label_offset); - const size_t graph_read = graph_future.get(); - const size_t dataset_read = dataset_future.get(); - const size_t label_read = label_future.get(); + auto label_future = label_kv.pread(label_buf.data_handle(), label_bytes, label_offset); + + // Drain all three futures before propagating any failure. + std::exception_ptr read_error; + auto drain = [&read_error](auto& fut) -> size_t { + try { + return fut.get(); + } catch (...) { + if (!read_error) { read_error = std::current_exception(); } + return 0; + } + }; + const size_t graph_read = drain(graph_future); + const size_t dataset_read = drain(dataset_future); + const size_t label_read = drain(label_future); + if (read_error) { std::rethrow_exception(read_error); } RAFT_EXPECTS(graph_read == graph_bytes, "Short graph read at row %ld: expected %zu, got %zu", start_row, diff --git a/cpp/src/util/file_io.cpp b/cpp/src/util/file_io.cpp index 459094e481..55432fd6ed 100644 --- a/cpp/src/util/file_io.cpp +++ b/cpp/src/util/file_io.cpp @@ -11,9 +11,56 @@ #include #include #include +#include #include namespace cuvs::util { +namespace { + +struct file_identity { + dev_t device; + ino_t inode; +}; + +file_identity get_file_identity(int fd, const char* description) +{ + struct stat status{}; + RAFT_EXPECTS(::fstat(fd, &status) == 0, + "Failed to stat %s file descriptor: %s", + description, + std::strerror(errno)); + return {status.st_dev, status.st_ino}; +} + +void expect_matching_file_identity(file_identity expected, + int actual_fd, + const char* actual_description, + const std::string& path) +{ + RAFT_EXPECTS(actual_fd >= 0, "kvikio did not open %s file descriptor", actual_description); + + const auto actual = get_file_identity(actual_fd, actual_description); + RAFT_EXPECTS(actual.device == expected.device && actual.inode == expected.inode, + "File path changed while opening %s for kvikio I/O", + path.c_str()); +} + +void validate_kvikio_handle_matches_fd(const file_descriptor& fd, + const kvikio::FileHandle& handle, + const std::string& path) +{ + const auto expected = get_file_identity(fd.get(), "source"); + + const int buffered_fd = handle.fd(false); + expect_matching_file_identity(expected, buffered_fd, "kvikio buffered", path); + + const int direct_fd = handle.fd(true); + if (direct_fd >= 0 && direct_fd != buffered_fd) { + expect_matching_file_identity(expected, direct_fd, "kvikio direct", path); + } +} + +} // namespace void read_large_file(const file_descriptor& fd, void* dest_ptr, @@ -30,6 +77,7 @@ void read_large_file(const file_descriptor& fd, // the POSIX + threadpool backend (with O_DIRECT when available) otherwise. The destination may be // host or device memory; kvikio detects which. kvikio::FileHandle handle(path, "r"); + validate_kvikio_handle_matches_fd(fd, handle, path); const size_t bytes_read = handle.pread(dest_ptr, total_bytes, file_offset).get(); RAFT_EXPECTS(bytes_read == total_bytes, "Incomplete read from %s: expected %zu bytes, got %zu", @@ -52,6 +100,7 @@ void write_large_file(const file_descriptor& fd, // Open in read+write mode ("r+") so the existing numpy header and preallocation are preserved // (kvikio's "w" mode would truncate). The source may be host or device memory. kvikio::FileHandle handle(path, "r+"); + validate_kvikio_handle_matches_fd(fd, handle, path); const size_t bytes_written = handle.pwrite(data_ptr, total_bytes, file_offset).get(); RAFT_EXPECTS(bytes_written == total_bytes, "Incomplete write to %s: expected %zu bytes, wrote %zu", diff --git a/cpp/tests/util/file_io_test.cpp b/cpp/tests/util/file_io_test.cpp index 79809e79c5..fd84ce9321 100644 --- a/cpp/tests/util/file_io_test.cpp +++ b/cpp/tests/util/file_io_test.cpp @@ -213,6 +213,29 @@ TEST(FileIO, InvalidArguments) ASSERT_NE(dup_fd, -1); file_descriptor pathless_fd(dup_fd); EXPECT_THROW(read_large_file(pathless_fd, buf, sizeof(buf), header_size), raft::exception); + EXPECT_THROW(write_large_file(pathless_fd, buf, sizeof(buf), header_size), raft::exception); +} + +TEST(FileIO, RejectsReplacedPathForBulkIO) +{ + scratch_dir scratch; + const std::string path = scratch.file("replace.npy"); + auto [fd, header_size] = create_numpy_file(path, {16, 4}); + char buf[8] = {0}; + + std::error_code ec; + ASSERT_TRUE(std::filesystem::remove(path, ec)); + ASSERT_FALSE(ec) << ec.message(); + + std::ofstream replacement(path, std::ios::out | std::ios::binary); + ASSERT_TRUE(replacement.good()) << "cannot create replacement " << path; + std::vector replacement_data(header_size + sizeof(buf), 0); + replacement.write(replacement_data.data(), replacement_data.size()); + replacement.close(); + ASSERT_TRUE(replacement.good()) << "cannot write replacement " << path; + + EXPECT_THROW(read_large_file(fd, buf, sizeof(buf), header_size), raft::exception); + EXPECT_THROW(write_large_file(fd, buf, sizeof(buf), header_size), raft::exception); } } // namespace cuvs::util diff --git a/dependencies.yaml b/dependencies.yaml index 820813c39c..a159b0ded2 100644 --- a/dependencies.yaml +++ b/dependencies.yaml @@ -403,24 +403,6 @@ dependencies: - libcusolver-dev - libcusparse-dev - libnvjitlink-dev - specific: - # libcufile provides the GPUDirect Storage (cuFile) backend that kvikio compiles against. - # It is packaged for CUDA 12 on x86_64 only, and for CUDA 13 on all arches. - - output_types: [conda] - matrices: - - matrix: - cuda: "12.*" - arch: x86_64 - packages: - - libcufile-dev - - matrix: - cuda: "12.*" - arch: aarch64 - packages: - - matrix: - cuda: "13.*" - packages: - - libcufile-dev cuda_wheels: specific: # cuVS needs 'nvJitLink>={whatever-cuvs-was-built-against}' at runtime, and mixing diff --git a/python/libcuvs/libcuvs/load.py b/python/libcuvs/libcuvs/load.py index 3e66761a2a..db33c758a6 100644 --- a/python/libcuvs/libcuvs/load.py +++ b/python/libcuvs/libcuvs/load.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. # SPDX-License-Identifier: Apache-2.0 # @@ -37,9 +37,11 @@ def load_library(): try: # These libraries must be loaded before libcuvs because libcuvs # references their symbols + import libkvikio import libraft import librmm + libkvikio.load_library() librmm.load_library() libraft.load_library() except ModuleNotFoundError: From 7496f4febb6cda2156c1fa76779400b29d546146 Mon Sep 17 00:00:00 2001 From: Julian Miller Date: Wed, 24 Jun 2026 11:57:52 +0200 Subject: [PATCH 03/18] Fix missing C kvikio private linkage --- c/CMakeLists.txt | 1 + 1 file changed, 1 insertion(+) diff --git a/c/CMakeLists.txt b/c/CMakeLists.txt index be4bc7a051..4c9cda6747 100644 --- a/c/CMakeLists.txt +++ b/c/CMakeLists.txt @@ -171,6 +171,7 @@ target_link_libraries( $,cuvs::cuvs_static,cuvs::cuvs> # $ # enforce we shouldn't use raft symbols $ + $> ) # ################################################################################################## From 66f1fe1b09c94a72eeb850a5b9f247aa115cd010 Mon Sep 17 00:00:00 2001 From: Julian Miller Date: Wed, 24 Jun 2026 13:00:24 +0200 Subject: [PATCH 04/18] Fix C kvikio include and wheel build --- c/CMakeLists.txt | 4 ++++ ci/build_wheel.sh | 1 + 2 files changed, 5 insertions(+) diff --git a/c/CMakeLists.txt b/c/CMakeLists.txt index 4c9cda6747..74cab676d2 100644 --- a/c/CMakeLists.txt +++ b/c/CMakeLists.txt @@ -71,6 +71,10 @@ if(BUILD_CAGRA_HNSWLIB) include(../cpp/cmake/thirdparty/get_hnswlib.cmake) endif() +if(CUVSC_STATIC_CUVS_LIBRARY AND NOT TARGET kvikio::kvikio) + include(../cpp/cmake/thirdparty/get_kvikio.cmake) +endif() + if(BUILD_MG_ALGOS AND CUVSC_STATIC_CUVS_LIBRARY) rapids_find_generate_module( NCCL diff --git a/ci/build_wheel.sh b/ci/build_wheel.sh index 6adee1fe34..ca3903fd1f 100755 --- a/ci/build_wheel.sh +++ b/ci/build_wheel.sh @@ -40,6 +40,7 @@ EXCLUDE_ARGS=( --exclude "libcurand.so.*" --exclude "libcusolver.so.*" --exclude "libcusparse.so.*" + --exclude "libkvikio.so*" --exclude "libnccl.so.*" --exclude "libnvJitLink.so.*" --exclude "libnvrtc.so.*" From 46b5cf2ce762ba20537515001c2baa8760b9ecfc Mon Sep 17 00:00:00 2001 From: Julian Miller Date: Wed, 24 Jun 2026 15:28:00 +0200 Subject: [PATCH 05/18] Fix C kvikio include order --- c/CMakeLists.txt | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/c/CMakeLists.txt b/c/CMakeLists.txt index 74cab676d2..45b4c29d28 100644 --- a/c/CMakeLists.txt +++ b/c/CMakeLists.txt @@ -71,10 +71,6 @@ if(BUILD_CAGRA_HNSWLIB) include(../cpp/cmake/thirdparty/get_hnswlib.cmake) endif() -if(CUVSC_STATIC_CUVS_LIBRARY AND NOT TARGET kvikio::kvikio) - include(../cpp/cmake/thirdparty/get_kvikio.cmake) -endif() - if(BUILD_MG_ALGOS AND CUVSC_STATIC_CUVS_LIBRARY) rapids_find_generate_module( NCCL @@ -219,6 +215,10 @@ if(BUILD_TESTS) enable_language(CUDA) find_package(CUDAToolkit REQUIRED) + if(CUVSC_STATIC_CUVS_LIBRARY AND NOT TARGET kvikio::kvikio) + include(../cpp/cmake/thirdparty/get_kvikio.cmake) + endif() + enable_testing() add_subdirectory(tests) endif() From f9ee2cb97cd6b9e87500f5c6be6c743bbf373b81 Mon Sep 17 00:00:00 2001 From: Julian Miller Date: Wed, 24 Jun 2026 18:49:00 +0200 Subject: [PATCH 06/18] Add C kvikio runtime library --- c/CMakeLists.txt | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/c/CMakeLists.txt b/c/CMakeLists.txt index 45b4c29d28..44e3cef0f9 100644 --- a/c/CMakeLists.txt +++ b/c/CMakeLists.txt @@ -191,6 +191,19 @@ if(PROJECT_IS_TOP_LEVEL) COMPONENT c_api EXPORT cuvs-c-exports ) + if(CUVSC_STATIC_CUVS_LIBRARY AND TARGET kvikio::kvikio) + get_target_property(_cuvs_kvikio_aliased_target kvikio::kvikio ALIASED_TARGET) + if(_cuvs_kvikio_aliased_target) + set(_cuvs_kvikio_runtime_target "${_cuvs_kvikio_aliased_target}") + else() + set(_cuvs_kvikio_runtime_target kvikio::kvikio) + endif() + install( + FILES "$" + DESTINATION ${lib_dir} + COMPONENT c_api + ) + endif() install( DIRECTORY include/cuvs COMPONENT c_api From 515ec5590ccf8cd8e6b3754e3fe2efee00c8633f Mon Sep 17 00:00:00 2001 From: Julian Miller Date: Wed, 1 Jul 2026 11:32:31 +0200 Subject: [PATCH 07/18] Try GDS or direct I/O paths if available --- c/CMakeLists.txt | 2 +- ci/build_wheel.sh | 2 +- ci/release/update-version.sh | 2 +- conda/recipes/libcuvs/recipe.yaml | 2 +- cpp/CMakeLists.txt | 2 +- cpp/cmake/thirdparty/get_kvikio.cmake | 4 +- cpp/include/cuvs/util/file_io.hpp | 2 +- .../neighbors/detail/cagra/cagra_build.cuh | 15 ++-- cpp/src/neighbors/detail/hnsw.hpp | 2 +- cpp/src/util/file_io.cpp | 8 +- cpp/src/util/kvikio_io.hpp | 78 +++++++++++++++++++ cpp/tests/CMakeLists.txt | 2 +- cpp/tests/util/file_io_test.cpp | 2 +- python/libcuvs/libcuvs/load.py | 2 +- python/libcuvs/pyproject.toml | 2 +- 15 files changed, 107 insertions(+), 20 deletions(-) create mode 100644 cpp/src/util/kvikio_io.hpp diff --git a/c/CMakeLists.txt b/c/CMakeLists.txt index 44e3cef0f9..069a8c6b66 100644 --- a/c/CMakeLists.txt +++ b/c/CMakeLists.txt @@ -1,6 +1,6 @@ # ============================================================================= # cmake-format: off -# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # cmake-format: on # ============================================================================= diff --git a/ci/build_wheel.sh b/ci/build_wheel.sh index ca3903fd1f..87206e9616 100755 --- a/ci/build_wheel.sh +++ b/ci/build_wheel.sh @@ -1,5 +1,5 @@ #!/bin/bash -# SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 set -euo pipefail diff --git a/ci/release/update-version.sh b/ci/release/update-version.sh index de1802adb4..358bc23f85 100755 --- a/ci/release/update-version.sh +++ b/ci/release/update-version.sh @@ -1,5 +1,5 @@ #!/bin/bash -# SPDX-FileCopyrightText: Copyright (c) 2020-2026, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2020-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 ######################## # CUVS Version Updater # diff --git a/conda/recipes/libcuvs/recipe.yaml b/conda/recipes/libcuvs/recipe.yaml index 7d0542984e..7909dc01f6 100644 --- a/conda/recipes/libcuvs/recipe.yaml +++ b/conda/recipes/libcuvs/recipe.yaml @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 schema_version: 1 diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt index ecbf377fcc..0a3bd143fa 100644 --- a/cpp/CMakeLists.txt +++ b/cpp/CMakeLists.txt @@ -1,6 +1,6 @@ # ============================================================================= # cmake-format: off -# SPDX-FileCopyrightText: Copyright (c) 2020-2026, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2020-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # cmake-format: on cmake_minimum_required(VERSION 4.0 FATAL_ERROR) diff --git a/cpp/cmake/thirdparty/get_kvikio.cmake b/cpp/cmake/thirdparty/get_kvikio.cmake index 9fb46a0dcf..8898bb4c15 100644 --- a/cpp/cmake/thirdparty/get_kvikio.cmake +++ b/cpp/cmake/thirdparty/get_kvikio.cmake @@ -1,6 +1,6 @@ # ============================================================================= # cmake-format: off -# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # cmake-format: on @@ -29,6 +29,8 @@ function(find_and_configure_kvikio) GIT_TAG ${PKG_PINNED_TAG} GIT_SHALLOW TRUE SOURCE_SUBDIR cpp + # Force KvikIO checkout to get public headers. + PATCH_COMMAND "" OPTIONS "KvikIO_BUILD_EXAMPLES OFF" "KvikIO_REMOTE_SUPPORT OFF" ) endfunction() diff --git a/cpp/include/cuvs/util/file_io.hpp b/cpp/include/cuvs/util/file_io.hpp index 338cdc2679..6653f0be12 100644 --- a/cpp/include/cuvs/util/file_io.hpp +++ b/cpp/include/cuvs/util/file_io.hpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ #pragma once diff --git a/cpp/src/neighbors/detail/cagra/cagra_build.cuh b/cpp/src/neighbors/detail/cagra/cagra_build.cuh index 36404811ea..6f9b60cb72 100644 --- a/cpp/src/neighbors/detail/cagra/cagra_build.cuh +++ b/cpp/src/neighbors/detail/cagra/cagra_build.cuh @@ -1,11 +1,12 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ #pragma once #include "../../../core/nvtx.hpp" #include "../../../preprocessing/quantize/vpq_build-ext.cuh" +#include "../../../util/kvikio_io.hpp" #include "graph_core.cuh" #include @@ -746,8 +747,10 @@ void ace_load_partition_dataset_from_disk( RAFT_LOG_DEBUG("ACE: Reading %lu core vectors from offset %lu", core_size, core_file_offset); RAFT_LOG_DEBUG( "ACE: Reading %lu augmented vectors from offset %lu", augmented_size, augmented_file_offset); - kvikio::FileHandle reordered_handle(reordered_path, "r"); - kvikio::FileHandle augmented_handle(augmented_path, "r"); + auto reordered_handle = + cuvs::util::detail::open_kvikio_file_for_ace_io(reordered_path, "r", sub_dataset_ptr); + auto augmented_handle = + cuvs::util::detail::open_kvikio_file_for_ace_io(augmented_path, "r", augmented_dest); auto core_future = reordered_handle.pread(sub_dataset_ptr, core_bytes, core_file_offset); auto augmented_future = augmented_handle.pread(augmented_dest, augmented_bytes, augmented_file_offset); @@ -769,13 +772,15 @@ void ace_load_partition_dataset_from_disk( expect_complete_read("augmented", augmented_bytes, augmented_read); } else if (core_bytes > 0) { RAFT_LOG_DEBUG("ACE: Reading %lu core vectors from offset %lu", core_size, core_file_offset); - kvikio::FileHandle reordered_handle(reordered_path, "r"); + auto reordered_handle = + cuvs::util::detail::open_kvikio_file_for_ace_io(reordered_path, "r", sub_dataset_ptr); auto core_future = reordered_handle.pread(sub_dataset_ptr, core_bytes, core_file_offset); expect_complete_read("core", core_bytes, core_future.get()); } else if (augmented_bytes > 0) { RAFT_LOG_DEBUG( "ACE: Reading %lu augmented vectors from offset %lu", augmented_size, augmented_file_offset); - kvikio::FileHandle augmented_handle(augmented_path, "r"); + auto augmented_handle = + cuvs::util::detail::open_kvikio_file_for_ace_io(augmented_path, "r", augmented_dest); auto augmented_future = augmented_handle.pread(augmented_dest, augmented_bytes, augmented_file_offset); expect_complete_read("augmented", augmented_bytes, augmented_future.get()); diff --git a/cpp/src/neighbors/detail/hnsw.hpp b/cpp/src/neighbors/detail/hnsw.hpp index 730be54262..654d312702 100644 --- a/cpp/src/neighbors/detail/hnsw.hpp +++ b/cpp/src/neighbors/detail/hnsw.hpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/src/util/file_io.cpp b/cpp/src/util/file_io.cpp index 55432fd6ed..fbb9417332 100644 --- a/cpp/src/util/file_io.cpp +++ b/cpp/src/util/file_io.cpp @@ -1,10 +1,12 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ #include +#include "kvikio_io.hpp" + #include #include @@ -76,7 +78,7 @@ void read_large_file(const file_descriptor& fd, // kvikio selects GPUDirect Storage (cuFile) for device destinations on a GDS-capable system, and // the POSIX + threadpool backend (with O_DIRECT when available) otherwise. The destination may be // host or device memory; kvikio detects which. - kvikio::FileHandle handle(path, "r"); + auto handle = detail::open_kvikio_file_for_ace_io(path, "r", dest_ptr); validate_kvikio_handle_matches_fd(fd, handle, path); const size_t bytes_read = handle.pread(dest_ptr, total_bytes, file_offset).get(); RAFT_EXPECTS(bytes_read == total_bytes, @@ -99,7 +101,7 @@ void write_large_file(const file_descriptor& fd, // Open in read+write mode ("r+") so the existing numpy header and preallocation are preserved // (kvikio's "w" mode would truncate). The source may be host or device memory. - kvikio::FileHandle handle(path, "r+"); + auto handle = detail::open_kvikio_file_for_ace_io(path, "r+", data_ptr); validate_kvikio_handle_matches_fd(fd, handle, path); const size_t bytes_written = handle.pwrite(data_ptr, total_bytes, file_offset).get(); RAFT_EXPECTS(bytes_written == total_bytes, diff --git a/cpp/src/util/kvikio_io.hpp b/cpp/src/util/kvikio_io.hpp new file mode 100644 index 0000000000..6b4a279ded --- /dev/null +++ b/cpp/src/util/kvikio_io.hpp @@ -0,0 +1,78 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +#pragma once + +#if __has_include() +#include +#define CUVS_KVIKIO_HAS_COMPAT_MODE_HEADER 1 +#else +#define CUVS_KVIKIO_HAS_COMPAT_MODE_HEADER 0 +#endif + +#include + +#if __has_include() +#include +#define CUVS_KVIKIO_HAS_UTILS_HEADER 1 +#else +#define CUVS_KVIKIO_HAS_UTILS_HEADER 0 +#endif + +#include +#include + +namespace cuvs::util::detail { + +#if CUVS_KVIKIO_HAS_COMPAT_MODE_HEADER +template +auto open_kvikio_file_compat_off(const std::string& path, const std::string& flags, int) + -> decltype(Handle(path, flags, Handle::m644, kvikio::CompatMode::OFF)) +{ + return Handle(path, flags, Handle::m644, kvikio::CompatMode::OFF); +} +#endif + +template +auto open_kvikio_file_compat_off(const std::string& path, const std::string& flags, long) + -> decltype(Handle(path, flags, Handle::m644, false)) +{ + return Handle(path, flags, Handle::m644, false); +} + +inline kvikio::FileHandle open_kvikio_file_compat_off(const std::string& path, + const std::string& flags, + ...) +{ + return kvikio::FileHandle(path, flags); +} + +inline bool is_kvikio_device_memory(const void* buffer) +{ +#if CUVS_KVIKIO_HAS_UTILS_HEADER + return buffer != nullptr && !kvikio::is_host_memory(buffer); +#else + return false; +#endif +} + +inline kvikio::FileHandle open_kvikio_file_for_ace_io(const std::string& path, + const std::string& flags, + const void* buffer) +{ + if (!is_kvikio_device_memory(buffer)) { return kvikio::FileHandle(path, flags); } + + // Prefer GDS for device transfers, but retain KvikIO's automatic POSIX fallback when the + // current system cannot open the file with compatibility mode disabled. + try { + return open_kvikio_file_compat_off(path, flags, 0); + } catch (const std::exception&) { + return kvikio::FileHandle(path, flags); + } +} + +} // namespace cuvs::util::detail + +#undef CUVS_KVIKIO_HAS_COMPAT_MODE_HEADER +#undef CUVS_KVIKIO_HAS_UTILS_HEADER diff --git a/cpp/tests/CMakeLists.txt b/cpp/tests/CMakeLists.txt index ec2dbddf4a..db8ea693cb 100644 --- a/cpp/tests/CMakeLists.txt +++ b/cpp/tests/CMakeLists.txt @@ -1,6 +1,6 @@ # ============================================================================= # cmake-format: off -# SPDX-FileCopyrightText: Copyright (c) 2021-2026, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2021-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # cmake-format: on # ============================================================================= diff --git a/cpp/tests/util/file_io_test.cpp b/cpp/tests/util/file_io_test.cpp index fd84ce9321..baa684690f 100644 --- a/cpp/tests/util/file_io_test.cpp +++ b/cpp/tests/util/file_io_test.cpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/python/libcuvs/libcuvs/load.py b/python/libcuvs/libcuvs/load.py index db33c758a6..43f9d76d4f 100644 --- a/python/libcuvs/libcuvs/load.py +++ b/python/libcuvs/libcuvs/load.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # diff --git a/python/libcuvs/pyproject.toml b/python/libcuvs/pyproject.toml index d41e39976f..2ccb29d26a 100644 --- a/python/libcuvs/pyproject.toml +++ b/python/libcuvs/pyproject.toml @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 [build-system] From f7d726b4cb252c05b64fca6f1764655d3be540c4 Mon Sep 17 00:00:00 2001 From: Julian Miller Date: Wed, 1 Jul 2026 11:48:29 +0200 Subject: [PATCH 08/18] Add one-time warning if GDS is unavailable --- cpp/src/util/kvikio_io.hpp | 13 ++++++++++++- 1 file changed, 12 insertions(+), 1 deletion(-) diff --git a/cpp/src/util/kvikio_io.hpp b/cpp/src/util/kvikio_io.hpp index 6b4a279ded..4ffbf4d41f 100644 --- a/cpp/src/util/kvikio_io.hpp +++ b/cpp/src/util/kvikio_io.hpp @@ -4,6 +4,8 @@ */ #pragma once +#include + #if __has_include() #include #define CUVS_KVIKIO_HAS_COMPAT_MODE_HEADER 1 @@ -21,6 +23,7 @@ #endif #include +#include #include namespace cuvs::util::detail { @@ -67,7 +70,15 @@ inline kvikio::FileHandle open_kvikio_file_for_ace_io(const std::string& path, // current system cannot open the file with compatibility mode disabled. try { return open_kvikio_file_compat_off(path, flags, 0); - } catch (const std::exception&) { + } catch (const std::exception& e) { + static std::once_flag warning_once; + std::call_once(warning_once, [&] { + RAFT_LOG_WARN( + "ACE: GDS is unavailable for %s (%s); falling back to KvikIO's default I/O mode. " + "Further GDS fallback warnings are suppressed.", + path.c_str(), + e.what()); + }); return kvikio::FileHandle(path, flags); } } From c43b3f02e4c12fd4a73983eddf4554cc8d58f203 Mon Sep 17 00:00:00 2001 From: Julian Miller Date: Thu, 2 Jul 2026 15:28:18 +0200 Subject: [PATCH 09/18] Improve ACE GDS usage - Fix wrongly-typed throughput calculation. - Keep sub_search_graph on GPU to use GDS. --- .../neighbors/detail/cagra/cagra_build.cuh | 145 ++++++++++++------ 1 file changed, 99 insertions(+), 46 deletions(-) diff --git a/cpp/src/neighbors/detail/cagra/cagra_build.cuh b/cpp/src/neighbors/detail/cagra/cagra_build.cuh index 6f9b60cb72..289580f219 100644 --- a/cpp/src/neighbors/detail/cagra/cagra_build.cuh +++ b/cpp/src/neighbors/detail/cagra/cagra_build.cuh @@ -18,6 +18,7 @@ #include #include #include +#include #include #include @@ -429,35 +430,84 @@ void ace_adjust_sub_graph_ids( } } -// ACE: Adjust ids in sub search graph in place for disk version -template +// ACE: Adjust IDs into global reordered IDs on the device for the disk version. +template +__global__ void ace_adjust_sub_graph_ids_disk_kernel(const GraphIdxT* sub_search_graph, + IdxT* adjusted_search_graph, + size_t graph_edges, + size_t core_sub_dataset_size, + IdxT core_partition_offset, + const IdxT* augmented_reordered_ids) +{ + const size_t edge_id = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (edge_id >= graph_edges) { return; } + + const GraphIdxT neighbor = sub_search_graph[edge_id]; + if (static_cast(neighbor) < core_sub_dataset_size) { + adjusted_search_graph[edge_id] = static_cast(neighbor) + core_partition_offset; + } else { + adjusted_search_graph[edge_id] = + augmented_reordered_ids[static_cast(neighbor) - core_sub_dataset_size]; + } +} + +template void ace_adjust_sub_graph_ids_disk( + raft::resources const& res, size_t core_sub_dataset_size, size_t augmented_sub_dataset_size, size_t graph_degree, size_t partition_id, - raft::host_matrix_view sub_search_graph, + raft::device_matrix_view sub_search_graph, + raft::device_matrix_view adjusted_search_graph, raft::host_vector_view core_partition_offsets, raft::host_vector_view augmented_partition_offsets, raft::host_vector_view augmented_backward_mapping, raft::host_vector_view core_forward_mapping) { + RAFT_EXPECTS(static_cast(sub_search_graph.extent(0)) >= core_sub_dataset_size, + "ACE: source graph has fewer rows than the core partition"); + RAFT_EXPECTS(static_cast(sub_search_graph.extent(1)) == graph_degree, + "ACE: source graph degree does not match the requested graph degree"); + RAFT_EXPECTS(static_cast(adjusted_search_graph.extent(0)) == core_sub_dataset_size && + static_cast(adjusted_search_graph.extent(1)) == graph_degree, + "ACE: adjusted graph shape does not match the core partition"); + + const size_t augmented_offset = augmented_partition_offsets(partition_id); + RAFT_EXPECTS(augmented_offset + augmented_sub_dataset_size <= + static_cast(augmented_backward_mapping.extent(0)), + "ACE: augmented partition exceeds the backward mapping"); + + // Resolve the global mapping once per augmented vector. This compact table is small enough to + // copy to the GPU even when the full dataset mappings are too large to fit in device memory. + auto augmented_reordered_ids = raft::make_host_vector(augmented_sub_dataset_size); #pragma omp parallel for - for (size_t i = 0; i < core_sub_dataset_size; i++) { - for (size_t k = 0; k < graph_degree; k++) { - size_t j = sub_search_graph(i, k); - if (j < core_sub_dataset_size) { - // core partition neighbor: local → core reordered - sub_search_graph(i, k) = j + core_partition_offsets(partition_id); - } else { - // Augmented partition neighbor: local → augmented reordered→ original → core reordered - size_t j_augmented = j - core_sub_dataset_size; - size_t j_original = - augmented_backward_mapping(j_augmented + augmented_partition_offsets(partition_id)); - sub_search_graph(i, k) = core_forward_mapping(j_original); - } - } + for (size_t i = 0; i < augmented_sub_dataset_size; i++) { + const size_t original_id = augmented_backward_mapping(augmented_offset + i); + augmented_reordered_ids(i) = core_forward_mapping(original_id); } + + auto augmented_reordered_ids_dev = + raft::make_device_vector(res, augmented_sub_dataset_size); + raft::copy(res, augmented_reordered_ids_dev.view(), augmented_reordered_ids.view()); + + const size_t graph_edges = core_sub_dataset_size * graph_degree; + if (graph_edges == 0) { return; } + + constexpr uint32_t block_size = 256; + const dim3 grid_size(raft::ceildiv(graph_edges, static_cast(block_size))); + ace_adjust_sub_graph_ids_disk_kernel<<>>( + sub_search_graph.data_handle(), + adjusted_search_graph.data_handle(), + graph_edges, + core_sub_dataset_size, + core_partition_offsets(partition_id), + augmented_reordered_ids_dev.data_handle()); + RAFT_CUDA_TRY(cudaPeekAtLastError()); + raft::resource::sync_stream(res); } // ACE: Reorder dataset based on partition assignments and store to disk @@ -1440,28 +1490,44 @@ index build_ace(raft::resources const& res, auto optimize_elapsed = std::chrono::duration_cast(optimize_end - read_end).count(); - // Copy graph edges for core members of this partition - auto sub_search_graph = - raft::make_host_matrix(core_sub_dataset_size, graph_degree); - cudaStream_t stream = raft::resource::get_cuda_stream(res); - raft::copy( - res, - raft::make_host_vector_view(sub_search_graph.data_handle(), sub_search_graph.size()), - raft::make_device_vector_view(sub_index.graph().data_handle(), sub_search_graph.size())); - raft::resource::sync_stream(res, stream); - + auto adjust_end = optimize_end; + auto write_elapsed = 0L; + const size_t graph_bytes = core_sub_dataset_size * graph_degree * sizeof(IdxT); if (use_disk_mode) { - // Adjust IDs in sub_search_graph in place for disk storage - ace_adjust_sub_graph_ids_disk(core_sub_dataset_size, + // Keep the graph on the GPU, adjust IDs there, and write it directly through GDS. + auto adjusted_search_graph = + raft::make_device_matrix(res, core_sub_dataset_size, graph_degree); + ace_adjust_sub_graph_ids_disk(res, + core_sub_dataset_size, augmented_sub_dataset_size, graph_degree, partition_id, - sub_search_graph.view(), + sub_index.graph(), + adjusted_search_graph.view(), core_partition_offsets.view(), augmented_partition_offsets.view(), augmented_backward_mapping.view(), core_forward_mapping.view()); + adjust_end = std::chrono::high_resolution_clock::now(); + + const size_t graph_offset = + static_cast(core_partition_offsets(partition_id)) * graph_degree * sizeof(IdxT) + + graph_header_size; + cuvs::util::write_large_file( + graph_fd, adjusted_search_graph.data_handle(), graph_bytes, graph_offset); + auto write_end = std::chrono::high_resolution_clock::now(); + write_elapsed = + std::chrono::duration_cast(write_end - adjust_end).count(); } else { + // The in-memory result is host-resident, so copy the core graph back before remapping IDs. + auto sub_search_graph = + raft::make_host_matrix(core_sub_dataset_size, graph_degree); + raft::copy( + res, + raft::make_host_vector_view(sub_search_graph.data_handle(), sub_search_graph.size()), + raft::make_device_vector_view(sub_index.graph().data_handle(), sub_search_graph.size())); + raft::resource::sync_stream(res); + // Adjust IDs in sub_search_graph and save to search_graph ace_adjust_sub_graph_ids(core_sub_dataset_size, augmented_sub_dataset_size, @@ -1473,33 +1539,20 @@ index build_ace(raft::resources const& res, augmented_partition_offsets.view(), core_backward_mapping.view(), augmented_backward_mapping.view()); + adjust_end = std::chrono::high_resolution_clock::now(); } - auto adjust_end = std::chrono::high_resolution_clock::now(); auto adjust_elapsed = std::chrono::duration_cast(adjust_end - optimize_end).count(); - if (use_disk_mode) { - const size_t graph_offset = - static_cast(core_partition_offsets(partition_id)) * graph_degree * sizeof(IdxT) + - graph_header_size; - const size_t graph_bytes = core_sub_dataset_size * graph_degree * sizeof(IdxT); - cuvs::util::write_large_file( - graph_fd, sub_search_graph.data_handle(), graph_bytes, graph_offset); - } - - auto end = std::chrono::high_resolution_clock::now(); - auto write_elapsed = - std::chrono::duration_cast(end - adjust_end).count(); + auto end = std::chrono::high_resolution_clock::now(); auto elapsed_ms = std::chrono::duration_cast(end - start).count(); double read_throughput = read_elapsed > 0 ? to_mib(sub_dataset_size * dataset_dim * sizeof(T)) / (read_elapsed / 1000.0) : 0.0; double write_throughput = - write_elapsed > 0 - ? to_mib(core_sub_dataset_size * dataset_dim * sizeof(T)) / (write_elapsed / 1000.0) - : 0.0; + write_elapsed > 0 ? to_mib(graph_bytes) / (write_elapsed / 1000.0) : 0.0; RAFT_LOG_INFO( "ACE: Partition %4lu (%8lu + %8lu) completed in %6ld ms: read %6ld ms (%7.1f MiB/s), " "optimize %6ld ms, adjust %6ld ms, write %6ld ms (%7.1f MiB/s)", From 84458dfb1bec10b3af08126273e8a5d7baf4cb0b Mon Sep 17 00:00:00 2001 From: Julian Miller Date: Fri, 3 Jul 2026 09:56:54 +0200 Subject: [PATCH 10/18] Fix kvikio in C standalone CI tests --- ci/build_standalone_c.sh | 9 ++++++++- cpp/CMakeLists.txt | 18 ++++++++++++++++++ 2 files changed, 26 insertions(+), 1 deletion(-) diff --git a/ci/build_standalone_c.sh b/ci/build_standalone_c.sh index 33123cd79c..fbc473119d 100755 --- a/ci/build_standalone_c.sh +++ b/ci/build_standalone_c.sh @@ -1,5 +1,5 @@ #!/bin/bash -# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 set -euo pipefail @@ -86,6 +86,13 @@ sccache --stop-server >/dev/null 2>&1 || true rapids-logger "Begin c install" cmake --install c/build --prefix c/build/install +# libcuvs_c contains libcuvs_static, whose ACE implementation uses the private KvikIO shared +# library. Bundle that runtime library without adding KvikIO headers or CMake metadata to the +# standalone C artifact. +cmake --install cpp/build \ + --prefix c/build/install \ + --component cuvs_standalone_runtime + # need to install the tests if [ "${BUILD_C_LIB_TESTS}" != "OFF" ]; then cmake --install c/build --prefix c/build/install --component testing diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt index 0a3bd143fa..87c41fc306 100644 --- a/cpp/CMakeLists.txt +++ b/cpp/CMakeLists.txt @@ -1647,6 +1647,24 @@ target_compile_definitions(cuvs::cuvs INTERFACE $<$:NVTX_ENAB include(GNUInstallDirs) include(CPack) + # The standalone C artifact links libcuvs_static into libcuvs_c, so private shared dependencies + # of the static archive must be bundled explicitly. Keep this component out of normal cuVS + # installs, where package managers provide the KvikIO runtime dependency. + if(TARGET kvikio::kvikio) + get_target_property(_cuvs_kvikio_aliased_target kvikio::kvikio ALIASED_TARGET) + if(_cuvs_kvikio_aliased_target) + set(_cuvs_standalone_kvikio_target "${_cuvs_kvikio_aliased_target}") + else() + set(_cuvs_standalone_kvikio_target kvikio::kvikio) + endif() + install( + FILES "$" + DESTINATION ${lib_dir} + COMPONENT cuvs_standalone_runtime + EXCLUDE_FROM_ALL + ) + endif() + set(target_names cuvs cuvs_static cuvs_cpp_headers cuvs_c) set(component_names cuvs_shared cuvs_static cuvs_cpp_headers c_api) set(export_names cuvs-shared-exports cuvs-static-exports cuvs-cpp-headers-exports cuvs-c-exports) From 669bd2c3656dd54567bb9484e5b866069f29cb5f Mon Sep 17 00:00:00 2001 From: Julian Miller Date: Fri, 3 Jul 2026 17:19:07 +0200 Subject: [PATCH 11/18] Add CAGRA and hnswlib serialization support --- cpp/include/cuvs/util/file_io.hpp | 20 +- .../detail/cagra/cagra_serialize.cuh | 398 +++++++++++++----- .../neighbors/detail/dataset_serialize.hpp | 83 +++- cpp/src/util/file_io.cpp | 54 ++- cpp/src/util/kvikio_io.hpp | 35 +- cpp/src/util/kvikio_serialize.hpp | 62 +++ cpp/tests/util/file_io_test.cpp | 34 ++ python/libcuvs/libcuvs/load.py | 38 +- 8 files changed, 587 insertions(+), 137 deletions(-) create mode 100644 cpp/src/util/kvikio_serialize.hpp diff --git a/cpp/include/cuvs/util/file_io.hpp b/cpp/include/cuvs/util/file_io.hpp index 6653f0be12..1b00293221 100644 --- a/cpp/include/cuvs/util/file_io.hpp +++ b/cpp/include/cuvs/util/file_io.hpp @@ -358,10 +358,11 @@ void write_large_file(const file_descriptor& fd, /** * @brief Sequential std::ostream backed by kvikio. * - * A std::ostream whose bytes are staged into a large buffer and written to disk through kvikio, + * Ordinary stream output is staged into a large host buffer and written to disk through kvikio, * which bypasses the page cache via O_DIRECT when supported (and falls back to buffered POSIX - * writes otherwise). Because it is a std::ostream it can be passed anywhere an ostream is expected - * (e.g. the hnswlib serializer), routing that output through kvikio. Non-copyable, non-movable. + * writes otherwise). Device buffers passed to write_device() use GPUDirect Storage when available. + * Because this is a std::ostream it can be passed anywhere an ostream is expected (e.g. the + * hnswlib serializer), routing that output through kvikio. Non-copyable, non-movable. */ class kvikio_ofstream : public std::ostream { public: @@ -384,6 +385,19 @@ class kvikio_ofstream : public std::ostream { /** @brief Flush any remaining staged bytes and close the file. */ void close(); + /** + * @brief Write a device-memory region at the current stream position. + * + * Pending host output is flushed first, then @p data is written through a compatibility-off + * KvikIO handle so GPUDirect Storage is used when available. The operation is synchronous and + * advances the same logical position used by ordinary std::ostream writes. + * + * @param data Device pointer to the source data. + * @param size Number of bytes to write. + * @return This stream. + */ + kvikio_ofstream& write_device(const void* data, size_t size); + /** @brief Total number of logical bytes written so far. */ [[nodiscard]] size_t bytes_written() const noexcept; diff --git a/cpp/src/neighbors/detail/cagra/cagra_serialize.cuh b/cpp/src/neighbors/detail/cagra/cagra_serialize.cuh index f106b82500..fd1270a006 100644 --- a/cpp/src/neighbors/detail/cagra/cagra_serialize.cuh +++ b/cpp/src/neighbors/detail/cagra/cagra_serialize.cuh @@ -1,11 +1,12 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ #pragma once #include +#include #include #include #include @@ -17,19 +18,38 @@ #include #include "../../../core/nvtx.hpp" +#include "../../../util/kvikio_serialize.hpp" #include "../../../util/serialize_validation.hpp" #include "../dataset_serialize.hpp" +#include +#include +#include +#include #include #include #include #include -#include namespace cuvs::neighbors::cagra::detail { constexpr int serialization_version = 5; +template +void serialize_index_mdspan(raft::resources const& res, std::ostream& os, const MdspanT& mdspan) +{ + raft::serialize_mdspan(res, os, mdspan); +} + +template +void serialize_index_mdspan( + raft::resources const& res, + cuvs::util::kvikio_ofstream& os, + const raft::device_mdspan& mdspan) +{ + cuvs::util::detail::serialize_device_mdspan(res, os, mdspan); +} + /** * Save the index to file. * @@ -40,11 +60,11 @@ constexpr int serialization_version = 5; * @param[in] index_ CAGRA index * */ -template -void serialize(raft::resources const& res, - std::ostream& os, - const index& index_, - bool include_dataset) +template +void serialize_impl(raft::resources const& res, + OutputStream& os, + const index& index_, + bool include_dataset) { raft::common::nvtx::range fun_scope("cagra::serialize"); @@ -66,7 +86,7 @@ void serialize(raft::resources const& res, raft::serialize_scalar(res, os, index_.graph_degree()); raft::serialize_scalar(res, os, index_.metric()); - raft::serialize_mdspan(res, os, index_.graph()); + serialize_index_mdspan(res, os, index_.graph()); include_dataset &= (index_.data().n_rows() > 0); bool has_source_indices = index_.source_indices().has_value(); @@ -80,7 +100,16 @@ void serialize(raft::resources const& res, RAFT_LOG_DEBUG("Saving CAGRA index WITHOUT dataset"); } - if (has_source_indices) { raft::serialize_mdspan(res, os, index_.source_indices().value()); } + if (has_source_indices) { serialize_index_mdspan(res, os, index_.source_indices().value()); } +} + +template +void serialize(raft::resources const& res, + std::ostream& os, + const index& index_, + bool include_dataset) +{ + serialize_impl(res, os, index_, include_dataset); } template @@ -93,30 +122,20 @@ void serialize(raft::resources const& res, "Cannot serialize a disk-backed CAGRA index. Convert it with " "cuvs::neighbors::hnsw::from_cagra() and load it into memory via " "cuvs::neighbors::hnsw::deserialize() before serialization."); - std::ofstream of(filename, std::ios::out | std::ios::binary); + cuvs::util::kvikio_ofstream of(filename); if (!of) { RAFT_FAIL("Cannot open file %s", filename.c_str()); } - detail::serialize(res, of, index_, include_dataset); + serialize_impl(res, of, index_, include_dataset); of.close(); if (!of) { RAFT_FAIL("Error writing output %s", filename.c_str()); } } template -void serialize_to_hnswlib( - raft::resources const& res, - std::ostream& os, - const cuvs::neighbors::cagra::index& index_, - std::optional> dataset) +void write_hnswlib_header(std::ostream& os, + const cuvs::neighbors::cagra::index& index_, + int dim) { - // static_assert(std::is_same_v or std::is_same_v, - // "An hnswlib index can only be trained with int32 or uint32 IdxT"); - int dim = (dataset) ? dataset->extent(1) : index_.dim(); - raft::common::nvtx::range fun_scope("cagra::serialize"); - RAFT_LOG_DEBUG("Saving CAGRA index to hnswlib format, size %zu, dim %u", - static_cast(index_.size()), - dim); - // offset_level_0 std::size_t offset_level_0 = 0; os.write(reinterpret_cast(&offset_level_0), sizeof(std::size_t)); @@ -159,86 +178,263 @@ void serialize_to_hnswlib( // efConstruction, can be anything std::size_t efConstruction = 500; os.write(reinterpret_cast(&efConstruction), sizeof(std::size_t)); +} + +inline void log_hnswlib_progress(size_t completed_rows, + size_t total_rows, + size_t bytes_written, + const std::chrono::system_clock::time_point& start_clock, + size_t& next_report_offset) +{ + if (completed_rows < next_report_offset || completed_rows == 0) { return; } + + const auto elapsed = std::chrono::duration_cast( + std::chrono::system_clock::now() - start_clock) + .count() * + 1e-6; + if (elapsed <= 0) { return; } + + constexpr double gib = double(size_t{1} << 30); + const double throughput = bytes_written / gib / elapsed; + const double rows_throughput = completed_rows / elapsed; + const double eta = (total_rows - completed_rows) / rows_throughput; + RAFT_LOG_DEBUG( + "# Writing rows %12zu / %12zu (%3.2f %%), %3.2f GiB/sec, ETA %d:%3.1f, written %3.2f GiB\r", + completed_rows, + total_rows, + completed_rows / static_cast(total_rows) * 100, + throughput, + int(eta / 60), + std::fmod(eta, 60.0), + bytes_written / gib); + + const size_t report_interval = std::max(1, total_rows / 10); + next_report_offset += report_interval; +} - // Remove padding before saving the dataset - raft::host_matrix host_dataset = raft::make_host_matrix(0, 0); - raft::host_matrix_view host_dataset_view; +template +void write_hnswlib_rows_host( + raft::resources const& res, + std::ostream& os, + const cuvs::neighbors::cagra::index& index_, + std::optional> dataset) +{ + const size_t n_rows = static_cast(index_.size()); + const int64_t dim = dataset ? dataset->extent(1) : index_.dim(); + const int64_t graph_degree = index_.graph_degree(); + const size_t row_size = + sizeof(int) + graph_degree * sizeof(IdxT) + dim * sizeof(T) + sizeof(size_t); + if (n_rows == 0) { return; } + + const size_t batch_rows = std::min( + n_rows, std::max(1, cuvs::util::detail::kDeviceSerializationBatchBytes / row_size)); + auto graph_buffer = + raft::make_host_matrix(batch_rows, graph_degree); + auto dataset_buffer = raft::make_host_matrix(batch_rows, dim); + + const auto graph = index_.graph(); + RAFT_EXPECTS(static_cast(graph.extent(0)) == n_rows, + "CAGRA graph rows (%zu) do not match index size (%zu)", + static_cast(graph.extent(0)), + n_rows); + + const T* device_dataset = nullptr; + int64_t dataset_stride = 0; if (dataset) { - host_dataset_view = *dataset; + RAFT_EXPECTS(static_cast(dataset->extent(0)) == n_rows, + "CAGRA dataset rows (%zu) do not match index size (%zu)", + static_cast(dataset->extent(0)), + n_rows); } else { - auto dataset = index_.dataset(); - RAFT_EXPECTS(dataset.size() > 0, + const auto dataset_view = index_.dataset(); + RAFT_EXPECTS(dataset_view.size() > 0, "Invalid CAGRA dataset of size 0 during serialization, shape %zux%zu", - static_cast(dataset.extent(0)), - static_cast(dataset.extent(1))); - host_dataset = raft::make_host_matrix(dataset.extent(0), dataset.extent(1)); - raft::copy_matrix(host_dataset.data_handle(), - host_dataset.extent(1), - dataset.data_handle(), - dataset.stride(0), - host_dataset.extent(1), - dataset.extent(0), - raft::resource::get_cuda_stream(res)); - raft::resource::sync_stream(res); - host_dataset_view = raft::make_const_mdspan(host_dataset.view()); + static_cast(dataset_view.extent(0)), + static_cast(dataset_view.extent(1))); + RAFT_EXPECTS(static_cast(dataset_view.extent(0)) == n_rows, + "CAGRA dataset rows (%zu) do not match index size (%zu)", + static_cast(dataset_view.extent(0)), + n_rows); + device_dataset = dataset_view.data_handle(); + dataset_stride = dataset_view.stride(0); } - auto graph = index_.graph(); - auto host_graph = - raft::make_host_matrix(graph.extent(0), graph.extent(1)); - raft::copy(res, host_graph.view(), graph); - raft::resource::sync_stream(res); - - size_t d_report_offset = index_.size() / 10; // Report progress in 10% steps. - size_t next_report_offset = d_report_offset; - const auto start_clock = std::chrono::system_clock::now(); - // Write one dataset and graph row at a time - RAFT_EXPECTS(host_graph.stride(1) == 1, "serialize_to_hnswlib expects row_major graph"); - RAFT_EXPECTS(host_dataset_view.stride(1) == 1, "serialize_to_hnswlib expects row_major dataset"); - - size_t bytes_written = 0; - float GiB = 1 << 30; - for (std::size_t i = 0; i < index_.size(); i++) { - auto graph_degree = static_cast(index_.graph_degree()); - os.write(reinterpret_cast(&graph_degree), sizeof(int)); - - IdxT* graph_row = &host_graph(i, 0); - os.write(reinterpret_cast(graph_row), sizeof(IdxT) * index_.graph_degree()); - - const T* data_row = &host_dataset_view(i, 0); - os.write(reinterpret_cast(data_row), sizeof(T) * dim); - os.write(reinterpret_cast(&i), sizeof(std::size_t)); - - bytes_written += - dim * sizeof(T) + index_.graph_degree() * sizeof(IdxT) + sizeof(int) + sizeof(size_t); - const auto end_clock = std::chrono::system_clock::now(); - if (!os.good()) { RAFT_FAIL("Error writing HNSW file, row %zu", i); } - if (i > next_report_offset) { - next_report_offset += d_report_offset; - const auto time = - std::chrono::duration_cast(end_clock - start_clock).count() * - 1e-6; - float throughput = bytes_written / GiB / time; - float rows_throughput = i / time; - float ETA = (index_.size() - i) / rows_throughput; - RAFT_LOG_DEBUG( - "# Writing rows %12lu / %12lu (%3.2f %%), %3.2f GiB/sec, ETA %d:%3.1f, written %3.2f GiB\r", - i, - index_.size(), - i / static_cast(index_.size()) * 100, - throughput, - int(ETA / 60), - std::fmod(ETA, 60.0f), - bytes_written / GiB); + + const auto stream = raft::resource::get_cuda_stream(res); + size_t bytes_written = 0; + size_t next_report = std::max(1, n_rows / 10); + const auto start_clock = std::chrono::system_clock::now(); + const auto graph_degree_int = static_cast(graph_degree); + + for (size_t first_row = 0; first_row < n_rows; first_row += batch_rows) { + const size_t rows = std::min(batch_rows, n_rows - first_row); + raft::copy_matrix(graph_buffer.data_handle(), + graph_degree, + graph.data_handle() + first_row * graph_degree, + graph_degree, + graph_degree, + rows, + stream); + if (!dataset) { + raft::copy_matrix(dataset_buffer.data_handle(), + dim, + device_dataset + first_row * dataset_stride, + dataset_stride, + dim, + rows, + stream); } + raft::resource::sync_stream(res); + + for (size_t batch_row = 0; batch_row < rows; ++batch_row) { + const size_t row = first_row + batch_row; + os.write(reinterpret_cast(&graph_degree_int), sizeof(int)); + os.write(reinterpret_cast(&graph_buffer(batch_row, 0)), + graph_degree * sizeof(IdxT)); + + const T* data_row = + dataset ? dataset->data_handle() + row * dataset->stride(0) : &dataset_buffer(batch_row, 0); + os.write(reinterpret_cast(data_row), dim * sizeof(T)); + os.write(reinterpret_cast(&row), sizeof(size_t)); + } + + RAFT_EXPECTS(os.good(), "Error writing HNSW file at row %zu", first_row + rows - 1); + bytes_written += rows * row_size; + log_hnswlib_progress(first_row + rows, n_rows, bytes_written, start_clock, next_report); + } +} + +template +__device__ void write_unaligned_value(uint8_t* output, ValueT value) +{ + const auto* bytes = reinterpret_cast(&value); +#pragma unroll + for (size_t i = 0; i < sizeof(ValueT); ++i) { + output[i] = bytes[i]; + } +} + +template +__global__ void pack_hnswlib_rows(uint8_t* output, + size_t row_size, + const uint32_t* graph, + const T* dataset, + size_t first_row, + size_t rows, + uint32_t graph_degree, + uint32_t dim, + int64_t dataset_stride) +{ + const size_t warps_per_block = blockDim.x / warpSize; + const size_t warp = threadIdx.x / warpSize; + const size_t lane = threadIdx.x % warpSize; + const size_t batch_row = blockIdx.x * warps_per_block + warp; + if (batch_row >= rows) { return; } + + const size_t row = first_row + batch_row; + auto* row_output = output + batch_row * row_size; + if (lane == 0) { + write_unaligned_value(row_output, static_cast(graph_degree)); + const size_t label_offset = sizeof(int) + graph_degree * sizeof(IdxT) + dim * sizeof(T); + write_unaligned_value(row_output + label_offset, row); + } + + const size_t graph_offset = sizeof(int); + for (size_t col = lane; col < graph_degree; col += warpSize) { + write_unaligned_value(row_output + graph_offset + col * sizeof(IdxT), + static_cast(graph[row * static_cast(graph_degree) + col])); + } + + const size_t dataset_offset = graph_offset + graph_degree * sizeof(IdxT); + for (size_t col = lane; col < dim; col += warpSize) { + write_unaligned_value(row_output + dataset_offset + col * sizeof(T), + dataset[row * static_cast(dataset_stride) + col]); + } +} + +template +void write_hnswlib_rows_device(raft::resources const& res, + cuvs::util::kvikio_ofstream& os, + const cuvs::neighbors::cagra::index& index_) +{ + const size_t n_rows = static_cast(index_.size()); + const uint32_t dim = index_.dim(); + const uint32_t graph_degree = index_.graph_degree(); + const size_t row_size = + sizeof(int) + graph_degree * sizeof(IdxT) + dim * sizeof(T) + sizeof(size_t); + if (n_rows == 0) { return; } + + const auto graph = index_.graph(); + const auto dataset = index_.dataset(); + RAFT_EXPECTS(dataset.size() > 0, + "Invalid CAGRA dataset of size 0 during serialization, shape %zux%zu", + static_cast(dataset.extent(0)), + static_cast(dataset.extent(1))); + RAFT_EXPECTS(static_cast(graph.extent(0)) == n_rows && + static_cast(dataset.extent(0)) == n_rows, + "CAGRA graph and dataset rows must match index size"); + + const size_t batch_rows = std::min( + n_rows, std::max(1, cuvs::util::detail::kDeviceSerializationBatchBytes / row_size)); + auto output = raft::make_device_vector(res, batch_rows * row_size); + const auto stream = raft::resource::get_cuda_stream(res); + + size_t bytes_written = 0; + size_t next_report = std::max(1, n_rows / 10); + const auto start_clock = std::chrono::system_clock::now(); + constexpr int block_size = 256; + constexpr int warps_per_block = block_size / 32; + + for (size_t first_row = 0; first_row < n_rows; first_row += batch_rows) { + const size_t rows = std::min(batch_rows, n_rows - first_row); + const size_t blocks = (rows + warps_per_block - 1) / warps_per_block; + pack_hnswlib_rows + <<(blocks), block_size, 0, stream>>>(output.data_handle(), + row_size, + graph.data_handle(), + dataset.data_handle(), + first_row, + rows, + graph_degree, + dim, + dataset.stride(0)); + RAFT_CUDA_TRY(cudaPeekAtLastError()); + raft::resource::sync_stream(res); + + const size_t bytes = rows * row_size; + os.write_device(output.data_handle(), bytes); + bytes_written += bytes; + log_hnswlib_progress(first_row + rows, n_rows, bytes_written, start_clock, next_report); } +} - for (std::size_t i = 0; i < index_.size(); i++) { - // zeroes - auto zero = 0; - os.write(reinterpret_cast(&zero), sizeof(int)); +inline void write_hnswlib_empty_levels(std::ostream& os, size_t n_rows) +{ + constexpr size_t chunk_size = 16 * 1024; + const std::array zeros{}; + for (size_t first_row = 0; first_row < n_rows; first_row += chunk_size) { + const size_t rows = std::min(chunk_size, n_rows - first_row); + os.write(reinterpret_cast(zeros.data()), rows * sizeof(int)); } } +template +void serialize_to_hnswlib( + raft::resources const& res, + std::ostream& os, + const cuvs::neighbors::cagra::index& index_, + std::optional> dataset) +{ + const int dim = dataset ? dataset->extent(1) : index_.dim(); + raft::common::nvtx::range fun_scope("cagra::serialize"); + RAFT_LOG_DEBUG("Saving CAGRA index to hnswlib format, size %zu, dim %d", + static_cast(index_.size()), + dim); + + write_hnswlib_header(os, index_, dim); + write_hnswlib_rows_host(res, os, index_, dataset); + write_hnswlib_empty_levels(os, static_cast(index_.size())); +} + template void serialize_to_hnswlib( raft::resources const& res, @@ -246,10 +442,22 @@ void serialize_to_hnswlib( const cuvs::neighbors::cagra::index& index_, std::optional> dataset) { - std::ofstream of(filename, std::ios::out | std::ios::binary); + const int dim = dataset ? dataset->extent(1) : index_.dim(); + raft::common::nvtx::range fun_scope("cagra::serialize"); + RAFT_LOG_DEBUG("Saving CAGRA index to hnswlib format, size %zu, dim %d", + static_cast(index_.size()), + dim); + + cuvs::util::kvikio_ofstream of(filename); if (!of) { RAFT_FAIL("Cannot open file %s", filename.c_str()); } - detail::serialize_to_hnswlib(res, of, index_, dataset); + write_hnswlib_header(of, index_, dim); + if (dataset) { + write_hnswlib_rows_host(res, of, index_, dataset); + } else { + write_hnswlib_rows_device(res, of, index_); + } + write_hnswlib_empty_levels(of, static_cast(index_.size())); of.close(); if (!of) { RAFT_FAIL("Error writing output %s", filename.c_str()); } diff --git a/cpp/src/neighbors/detail/dataset_serialize.hpp b/cpp/src/neighbors/detail/dataset_serialize.hpp index 00032ae9d2..184dce9966 100644 --- a/cpp/src/neighbors/detail/dataset_serialize.hpp +++ b/cpp/src/neighbors/detail/dataset_serialize.hpp @@ -1,11 +1,13 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ #pragma once #include +#include +#include #include #include #include @@ -15,9 +17,13 @@ #include +#include +#include #include #include +#include "../../util/kvikio_serialize.hpp" + namespace cuvs::neighbors::detail { using dataset_instance_tag = uint32_t; @@ -56,6 +62,45 @@ void serialize(const raft::resources& res, raft::serialize_mdspan(res, os, dst.view()); } +template +void serialize(const raft::resources& res, + cuvs::util::kvikio_ofstream& os, + const strided_dataset& dataset) +{ + const auto n_rows = dataset.n_rows(); + const auto dim = dataset.dim(); + const auto stride = dataset.stride(); + raft::serialize_scalar(res, os, n_rows); + raft::serialize_scalar(res, os, dim); + raft::serialize_scalar(res, os, stride); + + cuvs::util::detail::write_numpy_header( + os, {static_cast(n_rows), static_cast(dim)}); + if (n_rows == 0 || dim == 0) { return; } + + const auto src = dataset.view(); + const size_t row_len = static_cast(dim) * sizeof(DataT); + if (stride == dim) { + raft::resource::sync_stream(res); + os.write_device(src.data_handle(), static_cast(n_rows) * row_len); + return; + } + + const auto batch_rows = static_cast(std::min( + static_cast(n_rows), + std::max(1, cuvs::util::detail::kDeviceSerializationBatchBytes / row_len))); + auto packed = raft::make_device_matrix(res, batch_rows, dim); + const auto stream = raft::resource::get_cuda_stream(res); + + for (IdxT first_row = 0; first_row < n_rows; first_row += batch_rows) { + const auto rows = std::min(batch_rows, n_rows - first_row); + raft::copy_matrix( + packed.data_handle(), dim, src.data_handle() + first_row * stride, stride, dim, rows, stream); + raft::resource::sync_stream(res); + os.write_device(packed.data_handle(), static_cast(rows) * row_len); + } +} + template void serialize(const raft::resources& res, std::ostream& os, @@ -72,8 +117,26 @@ void serialize(const raft::resources& res, raft::serialize_mdspan(res, os, make_const_mdspan(dataset.data.view())); } -template -void serialize(const raft::resources& res, std::ostream& os, const dataset& dataset) +template +void serialize(const raft::resources& res, + cuvs::util::kvikio_ofstream& os, + const vpq_dataset& dataset) +{ + raft::serialize_scalar(res, os, dataset.n_rows()); + raft::serialize_scalar(res, os, dataset.dim()); + raft::serialize_scalar(res, os, dataset.vq_n_centers()); + raft::serialize_scalar(res, os, dataset.pq_n_centers()); + raft::serialize_scalar(res, os, dataset.pq_len()); + raft::serialize_scalar(res, os, dataset.encoded_row_length()); + cuvs::util::detail::serialize_device_mdspan( + res, os, make_const_mdspan(dataset.vq_code_book.view())); + cuvs::util::detail::serialize_device_mdspan( + res, os, make_const_mdspan(dataset.pq_code_book.view())); + cuvs::util::detail::serialize_device_mdspan(res, os, make_const_mdspan(dataset.data.view())); +} + +template +void serialize_dataset(const raft::resources& res, OutputStream& os, const dataset& dataset) { if (auto x = dynamic_cast*>(&dataset); x != nullptr) { raft::serialize_scalar(res, os, kSerializeEmptyDataset); @@ -112,6 +175,20 @@ void serialize(const raft::resources& res, std::ostream& os, const dataset RAFT_FAIL("unsupported dataset type."); } +template +void serialize(const raft::resources& res, std::ostream& os, const dataset& dataset) +{ + serialize_dataset(res, os, dataset); +} + +template +void serialize(const raft::resources& res, + cuvs::util::kvikio_ofstream& os, + const dataset& dataset) +{ + serialize_dataset(res, os, dataset); +} + template auto deserialize_empty(raft::resources const& res, std::istream& is) -> std::unique_ptr> diff --git a/cpp/src/util/file_io.cpp b/cpp/src/util/file_io.cpp index fbb9417332..3eb364fc1a 100644 --- a/cpp/src/util/file_io.cpp +++ b/cpp/src/util/file_io.cpp @@ -62,6 +62,21 @@ void validate_kvikio_handle_matches_fd(const file_descriptor& fd, } } +void validate_kvikio_handles_match(const kvikio::FileHandle& expected, + const kvikio::FileHandle& actual, + const std::string& path) +{ + const auto expected_identity = get_file_identity(expected.fd(false), "stream"); + + const int buffered_fd = actual.fd(false); + expect_matching_file_identity(expected_identity, buffered_fd, "kvikio buffered", path); + + const int direct_fd = actual.fd(true); + if (direct_fd >= 0 && direct_fd != buffered_fd) { + expect_matching_file_identity(expected_identity, direct_fd, "kvikio direct", path); + } +} + } // namespace void read_large_file(const file_descriptor& fd, @@ -116,7 +131,7 @@ void write_large_file(const file_descriptor& fd, class kvikio_ofstream::sbuf : public std::streambuf { public: sbuf(const std::string& path, size_t cap) - : handle_(path, "w"), buffer_(std::max(cap, kNumpyDataAlignment)) + : path_(path), handle_(path, "w"), buffer_(std::max(cap, kNumpyDataAlignment)) { RAFT_EXPECTS(buffer_.size() <= static_cast(std::numeric_limits::max()), "kvikio_ofstream buffer size must fit in std::streambuf::pbump"); @@ -136,10 +151,31 @@ class kvikio_ofstream::sbuf : public std::streambuf { { if (closed_) { return; } flush_buffer(); + if (device_handle_) { device_handle_->close(); } handle_.close(); closed_ = true; } + void write_device(const void* data, size_t n) + { + RAFT_EXPECTS(!closed_, "kvikio_ofstream: write attempted after close"); + RAFT_EXPECTS(data != nullptr || n == 0, "kvikio_ofstream: device pointer must not be null"); + if (n == 0) { return; } + + flush_buffer(); + if (!detail::is_kvikio_device_memory(data)) { + write_at_current_offset(handle_, data, n); + return; + } + + if (!device_handle_) { + device_handle_ = + std::make_unique(detail::open_kvikio_file_for_device_io(path_, "r+")); + validate_kvikio_handles_match(handle_, *device_handle_, path_); + } + write_at_current_offset(*device_handle_, data, n); + } + // Total logical bytes accepted = already-written + currently-staged. [[nodiscard]] size_t bytes_written() const noexcept { @@ -172,7 +208,7 @@ class kvikio_ofstream::sbuf : public std::streambuf { // chunk straight to kvikio. This avoids std::streambuf's byte-at-a-time fallback path. if (remaining >= buffer_.size()) { flush_buffer(); - write_direct(current, remaining); + write_at_current_offset(handle_, current, remaining); return requested; } @@ -216,20 +252,22 @@ class kvikio_ofstream::sbuf : public std::streambuf { { const size_t n = static_cast(pptr() - pbase()); if (n > 0) { - write_direct(pbase(), n); + write_at_current_offset(handle_, pbase(), n); setp(buffer_.data(), buffer_.data() + buffer_.size()); } } - void write_direct(const char* data, size_t n) + void write_at_current_offset(kvikio::FileHandle& handle, const void* data, size_t n) { RAFT_EXPECTS(!closed_, "kvikio_ofstream: write attempted after close"); - const size_t w = handle_.pwrite(data, n, offset_).get(); + const size_t w = handle.pwrite(data, n, offset_).get(); RAFT_EXPECTS(w == n, "kvikio_ofstream: short write (expected %zu, wrote %zu)", n, w); offset_ += n; } + std::string path_; kvikio::FileHandle handle_; + std::unique_ptr device_handle_; std::vector buffer_; size_t offset_ = 0; bool closed_ = false; @@ -255,6 +293,12 @@ void kvikio_ofstream::close() if (buf_) { buf_->close(); } } +kvikio_ofstream& kvikio_ofstream::write_device(const void* data, size_t size) +{ + buf_->write_device(data, size); + return *this; +} + size_t kvikio_ofstream::bytes_written() const noexcept { return buf_ ? buf_->bytes_written() : 0; } } // namespace cuvs::util diff --git a/cpp/src/util/kvikio_io.hpp b/cpp/src/util/kvikio_io.hpp index 4ffbf4d41f..a964d37cae 100644 --- a/cpp/src/util/kvikio_io.hpp +++ b/cpp/src/util/kvikio_io.hpp @@ -51,21 +51,9 @@ inline kvikio::FileHandle open_kvikio_file_compat_off(const std::string& path, return kvikio::FileHandle(path, flags); } -inline bool is_kvikio_device_memory(const void* buffer) +inline kvikio::FileHandle open_kvikio_file_for_device_io(const std::string& path, + const std::string& flags) { -#if CUVS_KVIKIO_HAS_UTILS_HEADER - return buffer != nullptr && !kvikio::is_host_memory(buffer); -#else - return false; -#endif -} - -inline kvikio::FileHandle open_kvikio_file_for_ace_io(const std::string& path, - const std::string& flags, - const void* buffer) -{ - if (!is_kvikio_device_memory(buffer)) { return kvikio::FileHandle(path, flags); } - // Prefer GDS for device transfers, but retain KvikIO's automatic POSIX fallback when the // current system cannot open the file with compatibility mode disabled. try { @@ -74,7 +62,7 @@ inline kvikio::FileHandle open_kvikio_file_for_ace_io(const std::string& path, static std::once_flag warning_once; std::call_once(warning_once, [&] { RAFT_LOG_WARN( - "ACE: GDS is unavailable for %s (%s); falling back to KvikIO's default I/O mode. " + "GDS is unavailable for %s (%s); falling back to KvikIO's default I/O mode. " "Further GDS fallback warnings are suppressed.", path.c_str(), e.what()); @@ -83,6 +71,23 @@ inline kvikio::FileHandle open_kvikio_file_for_ace_io(const std::string& path, } } +inline bool is_kvikio_device_memory(const void* buffer) +{ +#if CUVS_KVIKIO_HAS_UTILS_HEADER + return buffer != nullptr && !kvikio::is_host_memory(buffer); +#else + return false; +#endif +} + +inline kvikio::FileHandle open_kvikio_file_for_ace_io(const std::string& path, + const std::string& flags, + const void* buffer) +{ + if (!is_kvikio_device_memory(buffer)) { return kvikio::FileHandle(path, flags); } + return open_kvikio_file_for_device_io(path, flags); +} + } // namespace cuvs::util::detail #undef CUVS_KVIKIO_HAS_COMPAT_MODE_HEADER diff --git a/cpp/src/util/kvikio_serialize.hpp b/cpp/src/util/kvikio_serialize.hpp new file mode 100644 index 0000000000..fcddfc8e40 --- /dev/null +++ b/cpp/src/util/kvikio_serialize.hpp @@ -0,0 +1,62 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +#pragma once + +#include + +#include +#include +#include +#include + +#include +#include +#include +#include + +namespace cuvs::util::detail { + +inline constexpr size_t kDeviceSerializationBatchBytes = size_t{64} << 20; + +template +raft::numpy_serializer::header_t get_numpy_header(const std::vector& shape, + bool fortran_order = false) +{ + const auto dtype = raft::numpy_serializer::get_numpy_dtype(); + return raft::numpy_serializer::header_t{dtype, fortran_order, shape}; +} + +template +void write_numpy_header(std::ostream& os, const std::vector& shape) +{ + raft::numpy_serializer::write_header(os, get_numpy_header(shape)); +} + +template +void serialize_device_mdspan( + const raft::resources& res, + kvikio_ofstream& os, + const raft::device_mdspan& obj) +{ + static_assert(std::is_same_v || + std::is_same_v, + "The serializer only supports row-major and column-major layouts"); + + using obj_t = raft::device_mdspan; + std::vector shape; + shape.reserve(obj.rank()); + for (typename obj_t::rank_type i = 0; i < obj.rank(); ++i) { + shape.push_back(static_cast(obj.extent(i))); + } + + const bool fortran_order = std::is_same_v; + raft::numpy_serializer::write_header(os, get_numpy_header(shape, fortran_order)); + + raft::resource::sync_stream(res); + os.write_device(obj.data_handle(), obj.size() * sizeof(ElementType)); + RAFT_EXPECTS(os.good(), "Error writing content of device mdspan"); +} + +} // namespace cuvs::util::detail diff --git a/cpp/tests/util/file_io_test.cpp b/cpp/tests/util/file_io_test.cpp index baa684690f..5f24b918fe 100644 --- a/cpp/tests/util/file_io_test.cpp +++ b/cpp/tests/util/file_io_test.cpp @@ -198,6 +198,40 @@ TEST(FileIO, KvikioOfstreamLargeSingleWrite) EXPECT_EQ(got, data); } +// Host stream output and direct device output must share one ordered logical file position. +TEST(FileIO, KvikioOfstreamMixedHostAndDeviceWrites) +{ + raft::resources res; + scratch_dir scratch; + const std::string path = scratch.file("stream_device.bin"); + const std::vector prefix{'h', 'e', 'a', 'd'}; + const std::vector payload = make_pattern((size_t{3} << 20) + 17, 19); + const std::vector suffix{'t', 'a', 'i', 'l'}; + + auto device_payload = raft::make_device_vector(res, payload.size()); + raft::copy(device_payload.data_handle(), + payload.data(), + payload.size(), + raft::resource::get_cuda_stream(res)); + raft::resource::sync_stream(res); + + { + kvikio_ofstream os(path, size_t{1} << 20); + os.write(prefix.data(), prefix.size()); + os.write_device(device_payload.data_handle(), payload.size()); + os.write(suffix.data(), suffix.size()); + EXPECT_EQ(os.bytes_written(), prefix.size() + payload.size() + suffix.size()); + os.close(); + } + + std::vector expected; + expected.reserve(prefix.size() + payload.size() + suffix.size()); + expected.insert(expected.end(), prefix.begin(), prefix.end()); + expected.insert(expected.end(), payload.begin(), payload.end()); + expected.insert(expected.end(), suffix.begin(), suffix.end()); + EXPECT_EQ(read_whole_file(path), expected); +} + // Defensive checks on the bulk helpers. TEST(FileIO, InvalidArguments) { diff --git a/python/libcuvs/libcuvs/load.py b/python/libcuvs/libcuvs/load.py index 43f9d76d4f..16b16abfee 100644 --- a/python/libcuvs/libcuvs/load.py +++ b/python/libcuvs/libcuvs/load.py @@ -3,6 +3,7 @@ # import ctypes +import importlib import os # Loading with RTLD_LOCAL adds the library itself to the loader's @@ -32,24 +33,29 @@ def _load_wheel_installation(soname: str): return None -def load_library(): - """Dynamically load libcuvs.so and its dependencies""" - try: - # These libraries must be loaded before libcuvs because libcuvs - # references their symbols - import libkvikio - import libraft - import librmm +def _load_dependency(module_name: str): + """Import ``module_name`` and call its ``load_library()``. - libkvikio.load_library() - librmm.load_library() - libraft.load_library() + A missing module is ignored: these runtime dependencies might be + satisfied by conda packages (which do not have any Python modules) + instead of wheels. In that situation, assume that the libraries are in + a place where the loader can find them. Each dependency is loaded + independently so that a missing module does not prevent the others + from loading. + """ + try: + module = importlib.import_module(module_name) except ModuleNotFoundError: - # These runtime dependencies might be satisfied by conda packages - # (which do not have any Python modules) instead of wheels. In that - # situation, assume that the libraries are in a place where the loader - # can find them. - pass + return + module.load_library() + + +def load_library(): + """Dynamically load libcuvs.so and its dependencies""" + # These libraries must be loaded before libcuvs because libcuvs + # references their symbols. + for module_name in ("libkvikio", "librmm", "libraft"): + _load_dependency(module_name) prefer_system_installation = ( os.getenv("RAPIDS_LIBCUVS_PREFER_SYSTEM_LIBRARY", "false").lower() From 18c951918b29af533ceffc4c739de3ed312b0487 Mon Sep 17 00:00:00 2001 From: Julian Miller Date: Tue, 7 Jul 2026 08:44:33 +0200 Subject: [PATCH 12/18] Enhance kvikio_ofstream doc and testing --- cpp/include/cuvs/util/file_io.hpp | 6 ++-- cpp/tests/util/file_io_test.cpp | 50 +++++++++++++++++++++++++++++++ 2 files changed, 54 insertions(+), 2 deletions(-) diff --git a/cpp/include/cuvs/util/file_io.hpp b/cpp/include/cuvs/util/file_io.hpp index 1b00293221..910f39c424 100644 --- a/cpp/include/cuvs/util/file_io.hpp +++ b/cpp/include/cuvs/util/file_io.hpp @@ -361,8 +361,10 @@ void write_large_file(const file_descriptor& fd, * Ordinary stream output is staged into a large host buffer and written to disk through kvikio, * which bypasses the page cache via O_DIRECT when supported (and falls back to buffered POSIX * writes otherwise). Device buffers passed to write_device() use GPUDirect Storage when available. - * Because this is a std::ostream it can be passed anywhere an ostream is expected (e.g. the - * hnswlib serializer), routing that output through kvikio. Non-copyable, non-movable. + * This can be passed to APIs accepting a std::ostream& for sequential output (e.g. the hnswlib + * serializer). It supports querying the current output position, but not random-access seeking. + * std::fstream-specific APIs such as is_open() are not part of its interface. Non-copyable, + * non-movable. */ class kvikio_ofstream : public std::ostream { public: diff --git a/cpp/tests/util/file_io_test.cpp b/cpp/tests/util/file_io_test.cpp index 5f24b918fe..bf374a7c06 100644 --- a/cpp/tests/util/file_io_test.cpp +++ b/cpp/tests/util/file_io_test.cpp @@ -18,6 +18,7 @@ #include #include #include +#include #include #include @@ -178,6 +179,55 @@ TEST(FileIO, KvikioOfstreamRoundTrip) EXPECT_EQ(got, data); } +// Verify the std::ostream substitution used by serializers, including formatted and unformatted +// sequential output and current-position queries. +TEST(FileIO, KvikioOfstreamSequentialOstreamInterface) +{ + static_assert(std::is_base_of_v); + static_assert(std::is_convertible_v); + static_assert(!std::is_constructible_v); + + scratch_dir scratch; + const std::string path = scratch.file("stream_interface.bin"); + + { + kvikio_ofstream file(path, 4096); + std::ostream& os = file; + os << "value=" << 42 << '\n'; + os.write("tail", 4); + os.seekp(0, std::ios_base::cur); + EXPECT_TRUE(os.good()); + EXPECT_EQ(os.tellp(), std::streampos(13)); + os.flush(); + file.close(); + } + + const std::vector expected{ + 'v', 'a', 'l', 'u', 'e', '=', '4', '2', '\n', 't', 'a', 'i', 'l'}; + EXPECT_EQ(read_whole_file(path), expected); +} + +TEST(FileIO, KvikioOfstreamRejectsRandomAccessSeeking) +{ + scratch_dir scratch; + const std::string path = scratch.file("stream_seek.bin"); + + { + kvikio_ofstream os(path, 4096); + os << "before"; + os.seekp(0, std::ios_base::beg); + EXPECT_TRUE(os.fail()); + + os.clear(); + EXPECT_EQ(os.tellp(), std::streampos(6)); + os << "after"; + os.close(); + } + + const std::vector expected{'b', 'e', 'f', 'o', 'r', 'e', 'a', 'f', 't', 'e', 'r'}; + EXPECT_EQ(read_whole_file(path), expected); +} + // A large single write must bypass the small staging buffer correctly and still round-trip. TEST(FileIO, KvikioOfstreamLargeSingleWrite) { From 4321d3227608fc755c89906381bbd0bae2a19c92 Mon Sep 17 00:00:00 2001 From: Julian Miller Date: Tue, 7 Jul 2026 08:49:29 +0200 Subject: [PATCH 13/18] Fix cmake-format --- cpp/CMakeLists.txt | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt index 87c41fc306..77d7bfa34e 100644 --- a/cpp/CMakeLists.txt +++ b/cpp/CMakeLists.txt @@ -1647,9 +1647,9 @@ target_compile_definitions(cuvs::cuvs INTERFACE $<$:NVTX_ENAB include(GNUInstallDirs) include(CPack) - # The standalone C artifact links libcuvs_static into libcuvs_c, so private shared dependencies - # of the static archive must be bundled explicitly. Keep this component out of normal cuVS - # installs, where package managers provide the KvikIO runtime dependency. + # The standalone C artifact links libcuvs_static into libcuvs_c, so private shared dependencies of + # the static archive must be bundled explicitly. Keep this component out of normal cuVS installs, + # where package managers provide the KvikIO runtime dependency. if(TARGET kvikio::kvikio) get_target_property(_cuvs_kvikio_aliased_target kvikio::kvikio ALIASED_TARGET) if(_cuvs_kvikio_aliased_target) From fa82bd3b19bf1f02afb3a8adec519e528a7318a4 Mon Sep 17 00:00:00 2001 From: Julian Miller Date: Tue, 7 Jul 2026 11:17:52 +0200 Subject: [PATCH 14/18] Fix ACE graph degree for small datasets --- cpp/src/neighbors/detail/cagra/cagra_build.cuh | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/cpp/src/neighbors/detail/cagra/cagra_build.cuh b/cpp/src/neighbors/detail/cagra/cagra_build.cuh index 756d2ce83d..cb6d0b6a92 100644 --- a/cpp/src/neighbors/detail/cagra/cagra_build.cuh +++ b/cpp/src/neighbors/detail/cagra/cagra_build.cuh @@ -1396,12 +1396,16 @@ index build_ace(raft::resources const& res, // Index parameters depend only on partition sizes, not on where the dataset lives. cuvs::neighbors::cagra::index_params sub_index_params; - sub_index_params = cuvs::neighbors::cagra::index_params::from_hnsw_params( + const auto sub_hnsw_m = static_cast((graph_degree + 1) / 2); + sub_index_params = cuvs::neighbors::cagra::index_params::from_hnsw_params( raft::make_extents(sub_dataset_size, dataset_dim), - graph_degree / 2, + sub_hnsw_m, ef_construction, cuvs::neighbors::cagra::hnsw_heuristic_type::SAME_GRAPH_FOOTPRINT, params.metric); + // SAME_GRAPH_FOOTPRINT produces 2 * M edges. Preserve an odd graph degree after CAGRA's + // small-dataset validation reduces it (for example, 64 -> 3 for a four-row dataset). + sub_index_params.graph_degree = graph_degree; sub_index_params.attach_dataset_on_build = false; sub_index_params.guarantee_connectivity = params.guarantee_connectivity; From 1021d49d46470b07107cd788126490f54570d9f4 Mon Sep 17 00:00:00 2001 From: Julian Miller Date: Wed, 8 Jul 2026 11:40:40 +0200 Subject: [PATCH 15/18] Deserialize with kvikio --- ci/build_standalone_c.sh | 7 +- cpp/include/cuvs/neighbors/brute_force.hpp | 9 +- cpp/include/cuvs/neighbors/cagra.hpp | 11 +++ cpp/include/cuvs/util/file_io.hpp | 60 +++++++++++ cpp/src/neighbors/brute_force.cu | 18 +++- cpp/src/neighbors/brute_force_serialize.cu | 46 ++++----- .../detail/cagra/cagra_serialize.cuh | 99 +++++++++++++------ .../neighbors/detail/dataset_serialize.hpp | 98 ++++++++++++++---- .../neighbors/ivf_flat/ivf_flat_serialize.cuh | 34 ++++--- cpp/src/neighbors/ivf_list.cuh | 26 ++++- cpp/src/neighbors/ivf_pq/ivf_pq_serialize.cuh | 40 ++++---- cpp/src/util/file_io.cpp | 52 ++++++++++ cpp/src/util/kvikio_serialize.hpp | 63 ++++++++++++ cpp/tests/util/file_io_test.cpp | 43 ++++++++ 14 files changed, 493 insertions(+), 113 deletions(-) diff --git a/ci/build_standalone_c.sh b/ci/build_standalone_c.sh index fbc473119d..e4818b30e6 100755 --- a/ci/build_standalone_c.sh +++ b/ci/build_standalone_c.sh @@ -89,9 +89,12 @@ cmake --install c/build --prefix c/build/install # libcuvs_c contains libcuvs_static, whose ACE implementation uses the private KvikIO shared # library. Bundle that runtime library without adding KvikIO headers or CMake metadata to the # standalone C artifact. -cmake --install cpp/build \ +if ! cmake --install cpp/build \ --prefix c/build/install \ - --component cuvs_standalone_runtime + --component cuvs_standalone_runtime; then + echo "Error: failed to install component 'cuvs_standalone_runtime' from cpp/build" >&2 + exit 1 +fi # need to install the tests if [ "${BUILD_C_LIB_TESTS}" != "OFF" ]; then diff --git a/cpp/include/cuvs/neighbors/brute_force.hpp b/cpp/include/cuvs/neighbors/brute_force.hpp index 635447d0ed..d30ea21320 100644 --- a/cpp/include/cuvs/neighbors/brute_force.hpp +++ b/cpp/include/cuvs/neighbors/brute_force.hpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -69,6 +69,13 @@ struct CUVS_EXPORT index : cuvs::neighbors::index { cuvs::distance::DistanceType metric, DistT metric_arg = 0.0); + /** Construct a brute force index and take ownership of a row-major device dataset. */ + index(raft::resources const& res, + raft::device_matrix&& dataset, + std::optional>&& norms, + cuvs::distance::DistanceType metric, + DistT metric_arg = 0.0); + /** Construct a brute force index from dataset * * Constructs a brute force index from a dataset. This lets us precompute norms for diff --git a/cpp/include/cuvs/neighbors/cagra.hpp b/cpp/include/cuvs/neighbors/cagra.hpp index d1937cba27..0de957583e 100644 --- a/cpp/include/cuvs/neighbors/cagra.hpp +++ b/cpp/include/cuvs/neighbors/cagra.hpp @@ -697,6 +697,17 @@ struct CUVS_EXPORT index : cuvs::neighbors::index { graph_view_ = knn_graph; } + /** + * Replace the graph with a new graph and take ownership of its allocation. + * + * The passed device matrix is moved into the index, which will manage its lifetime. + */ + void update_graph(raft::device_matrix&& knn_graph) + { + graph_ = std::move(knn_graph); + graph_view_ = graph_.view(); + } + /** * Replace the graph with a new graph. * diff --git a/cpp/include/cuvs/util/file_io.hpp b/cpp/include/cuvs/util/file_io.hpp index 910f39c424..3cad81bf21 100644 --- a/cpp/include/cuvs/util/file_io.hpp +++ b/cpp/include/cuvs/util/file_io.hpp @@ -68,6 +68,37 @@ class fd_streambuf : public std::streambuf { return traits_type::to_int_type(*gptr()); } + pos_type seekoff(off_type off, std::ios_base::seekdir dir, std::ios_base::openmode which) override + { + if (fd_ == -1 || (which & std::ios_base::in) == 0) { return pos_type(off_type(-1)); } + + int whence = SEEK_SET; + if (dir == std::ios_base::cur) { + whence = SEEK_CUR; + // read() advances the descriptor past the buffered bytes. Account for unread bytes so + // tellg()/seekg() operate on the stream's logical position. + off -= static_cast(egptr() - gptr()); + } else if (dir == std::ios_base::end) { + whence = SEEK_END; + } else if (dir != std::ios_base::beg) { + return pos_type(off_type(-1)); + } + + off_t position = 0; + do { + position = ::lseek(fd_, static_cast(off), whence); + } while (position == static_cast(-1) && errno == EINTR); + if (position == static_cast(-1)) { return pos_type(off_type(-1)); } + + setg(buffer_.get(), buffer_.get(), buffer_.get()); + return pos_type(position); + } + + pos_type seekpos(pos_type pos, std::ios_base::openmode which) override + { + return seekoff(static_cast(pos), std::ios_base::beg, which); + } + public: explicit fd_streambuf(int fd, size_t buffer_size = 8192) : fd_(fd), buffer_(new char[buffer_size]), buffer_size_(buffer_size) @@ -220,6 +251,35 @@ class file_descriptor { std::string path_; }; +/** + * @brief Sequential file reader supporting mixed stream and direct-to-device reads. + * + * Small metadata can be consumed through stream(), while read_device() transfers the next bytes + * through KvikIO into device memory. Both operations advance one logical file position. KvikIO + * uses GPUDirect Storage when available and falls back to its compatible I/O path otherwise. + * Non-copyable, non-movable. + */ +class kvikio_file_reader { + public: + explicit kvikio_file_reader(const std::string& path); + ~kvikio_file_reader(); + + kvikio_file_reader(const kvikio_file_reader&) = delete; + kvikio_file_reader& operator=(const kvikio_file_reader&) = delete; + kvikio_file_reader(kvikio_file_reader&&) = delete; + kvikio_file_reader& operator=(kvikio_file_reader&&) = delete; + + /** @brief Stream used to consume metadata at the current logical file position. */ + [[nodiscard]] std::istream& stream(); + + /** @brief Read the next @p size bytes directly into @p data and advance the file position. */ + void read_device(void* data, size_t size); + + private: + class impl; + std::unique_ptr impl_; +}; + /** * @brief Create a numpy file with pre-allocated space and write the header. * diff --git a/cpp/src/neighbors/brute_force.cu b/cpp/src/neighbors/brute_force.cu index 2f9000acf7..d499df5796 100644 --- a/cpp/src/neighbors/brute_force.cu +++ b/cpp/src/neighbors/brute_force.cu @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -46,6 +46,22 @@ index::index(raft::resources const& res, raft::resource::sync_stream(res); } +template +index::index(raft::resources const&, + raft::device_matrix&& dataset, + std::optional>&& norms, + cuvs::distance::DistanceType metric, + DistT metric_arg) + : cuvs::neighbors::index(), + metric_(metric), + dataset_(std::move(dataset)), + norms_(std::move(norms)), + metric_arg_(metric_arg) +{ + dataset_view_ = raft::make_const_mdspan(dataset_.view()); + if (norms_) { norms_view_ = raft::make_const_mdspan(norms_->view()); } +} + template index::index(raft::resources const& res, raft::device_matrix_view dataset, diff --git a/cpp/src/neighbors/brute_force_serialize.cu b/cpp/src/neighbors/brute_force_serialize.cu index 1b7595ee11..06a18ab4d4 100644 --- a/cpp/src/neighbors/brute_force_serialize.cu +++ b/cpp/src/neighbors/brute_force_serialize.cu @@ -1,8 +1,9 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ +#include "../util/kvikio_serialize.hpp" #include "../util/serialize_validation.hpp" #include @@ -80,9 +81,11 @@ void serialize(raft::resources const& handle, serialize(handle, os, index, include_dataset); } -template -auto deserialize(raft::resources const& handle, std::istream& is) +template +auto deserialize(raft::resources const& handle, Input& input) { + auto& is = cuvs::util::detail::input_stream(input); + char dtype_string[4]; RAFT_EXPECTS(is.read(dtype_string, 4), "brute_force::deserialize: failed to read dtype prefix"); RAFT_EXPECTS(cuvs::util::validate_serialized_dtype(dtype_string, sizeof(dtype_string)), @@ -107,7 +110,7 @@ auto deserialize(raft::resources const& handle, std::istream& is) static_cast(metric)); auto metric_arg = raft::deserialize_scalar(handle, is); - auto dataset_storage = raft::make_host_matrix(std::int64_t{}, std::int64_t{}); + auto dataset_storage = raft::make_device_matrix(handle, std::int64_t{}, std::int64_t{}); auto include_dataset = raft::deserialize_scalar(handle, is); if (include_dataset) { RAFT_EXPECTS(cuvs::util::is_mul_no_overflow( @@ -117,27 +120,18 @@ auto deserialize(raft::resources const& handle, std::istream& is) static_cast(rows), static_cast(dim), sizeof(T)); - dataset_storage = raft::make_host_matrix(rows, dim); - raft::deserialize_mdspan(handle, is, dataset_storage.view()); + dataset_storage = raft::make_device_matrix(handle, rows, dim); + cuvs::util::detail::deserialize_mdspan(handle, input, dataset_storage.view()); } - auto has_norms = raft::deserialize_scalar(handle, is); - auto norms_storage = has_norms ? std::optional{raft::make_host_vector(rows)} - : std::optional>{}; - // TODO(wphicks): Use mdbuffer here when available - auto norms_storage_dev = + auto has_norms = raft::deserialize_scalar(handle, is); + auto norms_storage = has_norms ? std::optional{raft::make_device_vector(handle, rows)} : std::optional>{}; - if (has_norms) { - raft::deserialize_mdspan(handle, is, norms_storage->view()); - raft::copy(handle, norms_storage_dev->view(), norms_storage->view()); - } + if (has_norms) { cuvs::util::detail::deserialize_mdspan(handle, input, norms_storage->view()); } - auto result = index(handle, - raft::make_const_mdspan(dataset_storage.view()), - std::move(norms_storage_dev), - metric, - metric_arg); + auto result = index( + handle, std::move(dataset_storage), std::move(norms_storage), metric, metric_arg); raft::resource::sync_stream(handle); return result; @@ -147,20 +141,16 @@ void deserialize(raft::resources const& handle, const std::string& filename, cuvs::neighbors::brute_force::index* index) { - auto is = std::ifstream{filename, std::ios::in | std::ios::binary}; - RAFT_EXPECTS(is, "Cannot open file %s", filename.c_str()); - - *index = deserialize(handle, is); + cuvs::util::kvikio_file_reader reader(filename); + *index = deserialize(handle, reader); } void deserialize(raft::resources const& handle, const std::string& filename, cuvs::neighbors::brute_force::index* index) { - auto is = std::ifstream{filename, std::ios::in | std::ios::binary}; - RAFT_EXPECTS(is, "Cannot open file %s", filename.c_str()); - - *index = deserialize(handle, is); + cuvs::util::kvikio_file_reader reader(filename); + *index = deserialize(handle, reader); } void deserialize(raft::resources const& handle, diff --git a/cpp/src/neighbors/detail/cagra/cagra_serialize.cuh b/cpp/src/neighbors/detail/cagra/cagra_serialize.cuh index fd1270a006..d04447b902 100644 --- a/cpp/src/neighbors/detail/cagra/cagra_serialize.cuh +++ b/cpp/src/neighbors/detail/cagra/cagra_serialize.cuh @@ -473,10 +473,61 @@ void serialize_to_hnswlib( * */ template -void deserialize(raft::resources const& res, std::istream& is, index* index_) +void deserialize_graph(raft::resources const& res, + std::istream& is, + IdxT n_rows, + uint32_t graph_degree, + index* index_) +{ + using graph_index_type = typename index::graph_index_type; + auto graph = raft::make_host_matrix(n_rows, graph_degree); + deserialize_mdspan(res, is, graph.view()); + index_->update_graph(res, raft::make_const_mdspan(graph.view())); +} + +template +void deserialize_graph(raft::resources const& res, + cuvs::util::kvikio_file_reader& reader, + IdxT n_rows, + uint32_t graph_degree, + index* index_) +{ + using graph_index_type = typename index::graph_index_type; + auto graph = raft::make_device_matrix(res, n_rows, graph_degree); + cuvs::util::detail::deserialize_device_mdspan(res, reader, graph.view()); + index_->update_graph(std::move(graph)); +} + +template +void deserialize_source_indices(raft::resources const& res, + std::istream& is, + IdxT n_rows, + index* index_) +{ + auto source_indices = raft::make_host_vector(n_rows); + deserialize_mdspan(res, is, source_indices.view()); + index_->update_source_indices(res, raft::make_const_mdspan(source_indices.view())); + raft::resource::sync_stream(res); +} + +template +void deserialize_source_indices(raft::resources const& res, + cuvs::util::kvikio_file_reader& reader, + IdxT n_rows, + index* index_) +{ + auto source_indices = raft::make_device_vector(res, n_rows); + cuvs::util::detail::deserialize_device_mdspan(res, reader, source_indices.view()); + index_->update_source_indices(std::move(source_indices)); +} + +template +void deserialize_impl(raft::resources const& res, Input& input, index* index_) { raft::common::nvtx::range fun_scope("cagra::deserialize"); + auto& is = cuvs::util::detail::input_stream(input); + char dtype_string[4]; RAFT_EXPECTS(is.read(dtype_string, 4), "cagra::deserialize: failed to read dtype prefix"); RAFT_EXPECTS(cuvs::util::validate_serialized_dtype(dtype_string, sizeof(dtype_string)), @@ -498,46 +549,38 @@ void deserialize(raft::resources const& res, std::istream& is, index* i "cagra::deserialize: graph_degree=%u exceeds maximum %u", graph_degree, cuvs::util::kMaxGraphDegree); - RAFT_EXPECTS( - cuvs::util::is_mul_no_overflow( - static_cast(n_rows), static_cast(graph_degree), sizeof(IdxT)), - "cagra::deserialize: integer overflow in n_rows*graph_degree*sizeof(IdxT) " - "(n_rows=%lld, graph_degree=%u, sizeof(IdxT)=%zu)", - static_cast(n_rows), - graph_degree, - sizeof(IdxT)); - - auto graph = raft::make_host_matrix(n_rows, graph_degree); - deserialize_mdspan(res, is, graph.view()); + RAFT_EXPECTS(cuvs::util::is_mul_no_overflow(static_cast(n_rows), + static_cast(graph_degree), + sizeof(typename index::graph_index_type)), + "cagra::deserialize: integer overflow in graph allocation " + "(n_rows=%lld, graph_degree=%u, element_size=%zu)", + static_cast(n_rows), + graph_degree, + sizeof(typename index::graph_index_type)); *index_ = index(res, metric); - index_->update_graph(res, raft::make_const_mdspan(graph.view())); + deserialize_graph(res, input, n_rows, graph_degree, index_); auto content_map = raft::deserialize_scalar(res, is); bool has_dataset = content_map & 0x1u; if (has_dataset) { - index_->update_dataset(res, cuvs::neighbors::detail::deserialize_dataset(res, is)); + index_->update_dataset(res, cuvs::neighbors::detail::deserialize_dataset(res, input)); } bool has_source_indices = content_map & 0x2u; - if (has_source_indices) { - auto source_indices = raft::make_host_vector(n_rows); - deserialize_mdspan(res, is, source_indices.view()); - index_->update_source_indices(res, raft::make_const_mdspan(source_indices.view())); - raft::resource::sync_stream( - res); // Don't let the vector out of the scope before the copy is finished - } + if (has_source_indices) { deserialize_source_indices(res, input, n_rows, index_); } } template -void deserialize(raft::resources const& res, const std::string& filename, index* index_) +void deserialize(raft::resources const& res, std::istream& is, index* index_) { - std::ifstream is(filename, std::ios::in | std::ios::binary); - - if (!is) { RAFT_FAIL("Cannot open file %s", filename.c_str()); } - - detail::deserialize(res, is, index_); + deserialize_impl(res, is, index_); +} - is.close(); +template +void deserialize(raft::resources const& res, const std::string& filename, index* index_) +{ + cuvs::util::kvikio_file_reader reader(filename); + deserialize_impl(res, reader, index_); } } // namespace cuvs::neighbors::cagra::detail diff --git a/cpp/src/neighbors/detail/dataset_serialize.hpp b/cpp/src/neighbors/detail/dataset_serialize.hpp index 184dce9966..64e6263633 100644 --- a/cpp/src/neighbors/detail/dataset_serialize.hpp +++ b/cpp/src/neighbors/detail/dataset_serialize.hpp @@ -21,6 +21,7 @@ #include #include #include +#include #include "../../util/kvikio_serialize.hpp" @@ -189,10 +190,11 @@ void serialize(const raft::resources& res, serialize_dataset(res, os, dataset); } -template -auto deserialize_empty(raft::resources const& res, std::istream& is) +template +auto deserialize_empty(raft::resources const& res, Input& input) -> std::unique_ptr> { + auto& is = cuvs::util::detail::input_stream(input); auto suggested_dim = raft::deserialize_scalar(res, is); return std::make_unique>(suggested_dim); } @@ -209,10 +211,72 @@ auto deserialize_strided(raft::resources const& res, std::istream& is) return make_strided_dataset(res, std::move(host_array), stride); } -template -auto deserialize_vpq(raft::resources const& res, std::istream& is) +template +auto deserialize_strided(raft::resources const& res, cuvs::util::kvikio_file_reader& reader) + -> std::unique_ptr> +{ + auto& is = reader.stream(); + const auto rows = raft::deserialize_scalar(res, is); + const auto dim = raft::deserialize_scalar(res, is); + const auto stride = raft::deserialize_scalar(res, is); + if constexpr (std::is_signed_v) { + RAFT_EXPECTS(rows >= 0, "Failed to deserialize dataset: row count must not be negative."); + } + RAFT_EXPECTS( + dim <= stride, "Failed to deserialize dataset: dimension %u exceeds stride %u.", dim, stride); + + auto data = raft::make_device_matrix(res, rows, static_cast(stride)); + cuvs::util::detail::validate_numpy_header( + reader, {static_cast(rows), static_cast(dim)}); + + const auto stream = raft::resource::get_cuda_stream(res); + if (data.size() > 0) { + RAFT_CUDA_TRY(cudaMemsetAsync(data.data_handle(), 0, data.size() * sizeof(DataT), stream)); + raft::resource::sync_stream(res); + } + + if (rows > 0 && dim > 0) { + const size_t row_bytes = static_cast(dim) * sizeof(DataT); + if (stride == dim) { + reader.read_device(data.data_handle(), static_cast(rows) * row_bytes); + } else { + const auto batch_rows = static_cast(std::min( + static_cast(rows), + std::max(1, cuvs::util::detail::kDeviceSerializationBatchBytes / row_bytes))); + auto packed = raft::make_device_matrix(res, batch_rows, static_cast(dim)); + raft::resource::sync_stream(res); + + for (IdxT first_row = 0; first_row < rows; first_row += batch_rows) { + const auto batch_size = std::min(batch_rows, rows - first_row); + reader.read_device(packed.data_handle(), static_cast(batch_size) * row_bytes); + raft::copy_matrix(data.data_handle() + first_row * stride, + stride, + packed.data_handle(), + dim, + dim, + batch_size, + stream); + raft::resource::sync_stream(res); + } + } + } + + const auto logical_extents = raft::make_extents(rows, static_cast(dim)); + const auto layout = raft::make_strided_layout( + logical_extents, cuda::std::array{static_cast(stride), 1}); + using storage_type = decltype(data); + using owning_dataset_type = owning_dataset; + return std::make_unique(std::move(data), layout); +} + +template +auto deserialize_vpq(raft::resources const& res, Input& input) -> std::unique_ptr> { + auto& is = cuvs::util::detail::input_stream(input); auto n_rows = raft::deserialize_scalar(res, is); auto dim = raft::deserialize_scalar(res, is); auto vq_n_centers = raft::deserialize_scalar(res, is); @@ -227,28 +291,28 @@ auto deserialize_vpq(raft::resources const& res, std::istream& is) auto data = raft::make_device_matrix(res, n_rows, encoded_row_length); - raft::deserialize_mdspan(res, is, vq_code_book.view()); - raft::deserialize_mdspan(res, is, pq_code_book.view()); - raft::deserialize_mdspan(res, is, data.view()); + cuvs::util::detail::deserialize_mdspan(res, input, vq_code_book.view()); + cuvs::util::detail::deserialize_mdspan(res, input, pq_code_book.view()); + cuvs::util::detail::deserialize_mdspan(res, input, data.view()); return std::make_unique>( std::move(vq_code_book), std::move(pq_code_book), std::move(data)); } -template -auto deserialize_dataset(raft::resources const& res, std::istream& is) - -> std::unique_ptr> +template +auto deserialize_dataset(raft::resources const& res, Input& input) -> std::unique_ptr> { + auto& is = cuvs::util::detail::input_stream(input); const auto tag = raft::deserialize_scalar(res, is); switch (tag) { - case kSerializeEmptyDataset: return deserialize_empty(res, is); + case kSerializeEmptyDataset: return deserialize_empty(res, input); case kSerializeStridedDataset: { const auto dtype = raft::deserialize_scalar(res, is); switch (dtype) { - case CUDA_R_32F: return deserialize_strided(res, is); - case CUDA_R_16F: return deserialize_strided(res, is); - case CUDA_R_8I: return deserialize_strided(res, is); - case CUDA_R_8U: return deserialize_strided(res, is); + case CUDA_R_32F: return deserialize_strided(res, input); + case CUDA_R_16F: return deserialize_strided(res, input); + case CUDA_R_8I: return deserialize_strided(res, input); + case CUDA_R_8U: return deserialize_strided(res, input); default: RAFT_FAIL("Failed to deserialize dataset: unsupported strided dataset element type %d.", static_cast(dtype)); @@ -257,8 +321,8 @@ auto deserialize_dataset(raft::resources const& res, std::istream& is) case kSerializeVPQDataset: { const auto dtype = raft::deserialize_scalar(res, is); switch (dtype) { - case CUDA_R_32F: return deserialize_vpq(res, is); - case CUDA_R_16F: return deserialize_vpq(res, is); + case CUDA_R_32F: return deserialize_vpq(res, input); + case CUDA_R_16F: return deserialize_vpq(res, input); default: RAFT_FAIL("Failed to deserialize dataset: unsupported VPQ dtype %d.", static_cast(dtype)); diff --git a/cpp/src/neighbors/ivf_flat/ivf_flat_serialize.cuh b/cpp/src/neighbors/ivf_flat/ivf_flat_serialize.cuh index ef1898c17f..b05b8dc33f 100644 --- a/cpp/src/neighbors/ivf_flat/ivf_flat_serialize.cuh +++ b/cpp/src/neighbors/ivf_flat/ivf_flat_serialize.cuh @@ -1,10 +1,11 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ #pragma once +#include "../../util/kvikio_serialize.hpp" #include "../../util/serialize_validation.hpp" #include "../ivf_common.cuh" #include "../ivf_list.cuh" @@ -104,9 +105,11 @@ void serialize(raft::resources const& handle, * @param[in] index_ IVF-Flat index * */ -template -auto deserialize(raft::resources const& handle, std::istream& is) -> index +template +auto deserialize_impl(raft::resources const& handle, Input& input) -> index { + auto& is = cuvs::util::detail::input_stream(input); + char dtype_string[4]; RAFT_EXPECTS(is.read(dtype_string, 4), "ivf_flat::deserialize: failed to read dtype prefix"); RAFT_EXPECTS(cuvs::util::validate_serialized_dtype(dtype_string, sizeof(dtype_string)), @@ -139,7 +142,7 @@ auto deserialize(raft::resources const& handle, std::istream& is) -> index index_ = index(handle, metric, n_lists, adaptive_centers, cma, dim); - deserialize_mdspan(handle, is, index_.centers()); + cuvs::util::detail::deserialize_mdspan(handle, input, index_.centers()); bool has_norms = raft::deserialize_scalar(handle, is); if (has_norms) { index_.allocate_center_norms(handle); @@ -147,15 +150,15 @@ auto deserialize(raft::resources const& handle, std::istream& is) -> index list_device_spec{index_.dim(), cma}; list_spec list_store_spec{index_.dim(), true}; for (uint32_t label = 0; label < index_.n_lists(); label++) { - ivf::deserialize_list(handle, is, index_.lists()[label], list_store_spec, list_device_spec); + ivf::deserialize_list(handle, input, index_.lists()[label], list_store_spec, list_device_spec); } raft::resource::sync_stream(handle); @@ -165,17 +168,16 @@ auto deserialize(raft::resources const& handle, std::istream& is) -> index -auto deserialize(raft::resources const& handle, const std::string& filename) -> index +auto deserialize(raft::resources const& handle, std::istream& is) -> index { - std::ifstream is(filename, std::ios::in | std::ios::binary); - - if (!is) { RAFT_FAIL("Cannot open file %s", filename.c_str()); } - - auto index = detail::deserialize(handle, is); - - is.close(); + return deserialize_impl(handle, is); +} - return index; +template +auto deserialize(raft::resources const& handle, const std::string& filename) -> index +{ + cuvs::util::kvikio_file_reader reader(filename); + return deserialize_impl(handle, reader); } } // namespace cuvs::neighbors::ivf_flat::detail diff --git a/cpp/src/neighbors/ivf_list.cuh b/cpp/src/neighbors/ivf_list.cuh index 24691463ff..7b9ded0343 100644 --- a/cpp/src/neighbors/ivf_list.cuh +++ b/cpp/src/neighbors/ivf_list.cuh @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -22,6 +22,7 @@ #include #include +#include "../util/kvikio_serialize.hpp" #include "ivf_common.cuh" #include @@ -174,4 +175,27 @@ enable_if_valid_list_t deserialize_list(const raft::resources& handle, // Make sure the data is copied from host to device before the host arrays get out of the scope. raft::resource::sync_stream(handle); } + +template +enable_if_valid_list_t deserialize_list(const raft::resources& handle, + cuvs::util::kvikio_file_reader& reader, + std::shared_ptr& ld, + const typename ListT::spec_type& store_spec, + const typename ListT::spec_type& device_spec) +{ + using size_type = typename ListT::size_type; + auto& is = reader.stream(); + const auto size = raft::deserialize_scalar(handle, is); + if (size == 0) { return ld.reset(); } + + std::make_shared(handle, device_spec, size).swap(ld); + + const auto data_extents = store_spec.make_list_extents(size); + auto data_view = + raft::make_mdspan( + ld->data.data_handle(), data_extents); + auto indices_view = raft::make_device_vector_view(ld->indices.data_handle(), size); + cuvs::util::detail::deserialize_device_mdspan(handle, reader, data_view); + cuvs::util::detail::deserialize_device_mdspan(handle, reader, indices_view); +} } // namespace cuvs::neighbors::ivf diff --git a/cpp/src/neighbors/ivf_pq/ivf_pq_serialize.cuh b/cpp/src/neighbors/ivf_pq/ivf_pq_serialize.cuh index f20f2bf81e..c3481ac697 100644 --- a/cpp/src/neighbors/ivf_pq/ivf_pq_serialize.cuh +++ b/cpp/src/neighbors/ivf_pq/ivf_pq_serialize.cuh @@ -1,10 +1,11 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2022-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ #pragma once +#include "../../util/kvikio_serialize.hpp" #include "../../util/serialize_validation.hpp" #include "../ivf_common.cuh" #include "../ivf_list.cuh" @@ -120,9 +121,11 @@ void serialize(raft::resources const& handle_, * @param[in] is input stream * */ -template -auto deserialize(raft::resources const& handle_, std::istream& is) -> index +template +auto deserialize_impl(raft::resources const& handle_, Input& input) -> index { + auto& is = cuvs::util::detail::input_stream(input); + auto ver = raft::deserialize_scalar(handle_, is); if (ver != kSerializationVersion) { RAFT_FAIL("serialization version mismatch %d vs. %d", ver, kSerializationVersion); @@ -178,17 +181,17 @@ auto deserialize(raft::resources const& handle_, std::istream& is) -> indexpq_centers()); - raft::deserialize_mdspan(handle_, is, impl->centers()); - raft::deserialize_mdspan(handle_, is, impl->centers_rot()); - raft::deserialize_mdspan(handle_, is, impl->rotation_matrix()); - raft::deserialize_mdspan(handle_, is, impl->list_sizes()); + cuvs::util::detail::deserialize_mdspan(handle_, input, impl->pq_centers()); + cuvs::util::detail::deserialize_mdspan(handle_, input, impl->centers()); + cuvs::util::detail::deserialize_mdspan(handle_, input, impl->centers_rot()); + cuvs::util::detail::deserialize_mdspan(handle_, input, impl->rotation_matrix()); + cuvs::util::detail::deserialize_mdspan(handle_, input, impl->list_sizes()); if (codes_layout == list_layout::FLAT) { auto list_device_spec = list_spec_flat{pq_bits, pq_dim, cma}; auto list_store_spec = list_spec_flat{pq_bits, pq_dim, true}; for (auto& list_data_base_ptr : impl->lists()) { std::shared_ptr> typed_list; - ivf::deserialize_list(handle_, is, typed_list, list_store_spec, list_device_spec); + ivf::deserialize_list(handle_, input, typed_list, list_store_spec, list_device_spec); list_data_base_ptr = typed_list; } } else { @@ -196,7 +199,7 @@ auto deserialize(raft::resources const& handle_, std::istream& is) -> index{pq_bits, pq_dim, true}; for (auto& list_data_base_ptr : impl->lists()) { std::shared_ptr> typed_list; - ivf::deserialize_list(handle_, is, typed_list, list_store_spec, list_device_spec); + ivf::deserialize_list(handle_, input, typed_list, list_store_spec, list_device_spec); list_data_base_ptr = typed_list; } } @@ -210,6 +213,12 @@ auto deserialize(raft::resources const& handle_, std::istream& is) -> index +auto deserialize(raft::resources const& handle_, std::istream& is) -> index +{ + return deserialize_impl(handle_, is); +} + /** * Load index from file. * @@ -222,15 +231,8 @@ auto deserialize(raft::resources const& handle_, std::istream& is) -> index auto deserialize(raft::resources const& handle_, const std::string& filename) -> index { - std::ifstream infile(filename, std::ios::in | std::ios::binary); - - if (!infile) { RAFT_FAIL("Cannot open file %s", filename.c_str()); } - - auto index = detail::deserialize(handle_, infile); - - infile.close(); - - return index; + cuvs::util::kvikio_file_reader reader(filename); + return deserialize_impl(handle_, reader); } } // namespace cuvs::neighbors::ivf_pq::detail diff --git a/cpp/src/util/file_io.cpp b/cpp/src/util/file_io.cpp index 3eb364fc1a..e410536280 100644 --- a/cpp/src/util/file_io.cpp +++ b/cpp/src/util/file_io.cpp @@ -126,6 +126,58 @@ void write_large_file(const file_descriptor& fd, bytes_written); } +class kvikio_file_reader::impl { + public: + explicit impl(const std::string& path) + : fd_(path, O_RDONLY), + stream_(fd_.make_istream()), + handle_(detail::open_kvikio_file_for_device_io(path, "r")) + { + validate_kvikio_handle_matches_fd(fd_, handle_, path); + } + + std::istream& stream() { return stream_; } + + void read_device(void* data, size_t size) + { + RAFT_EXPECTS(data != nullptr || size == 0, "kvikio_file_reader: destination must not be null"); + if (size == 0) { return; } + RAFT_EXPECTS(size <= static_cast(std::numeric_limits::max()), + "kvikio_file_reader: read size exceeds stream offset range"); + + const auto position = stream_.tellg(); + RAFT_EXPECTS(position != std::istream::pos_type(-1), + "kvikio_file_reader: failed to determine the current file position"); + const auto offset = static_cast(position); + RAFT_EXPECTS(offset >= 0, "kvikio_file_reader: invalid negative file position"); + + const size_t bytes_read = handle_.pread(data, size, static_cast(offset)).get(); + RAFT_EXPECTS(bytes_read == size, + "kvikio_file_reader: short read (expected %zu, read %zu)", + size, + bytes_read); + + stream_.seekg(static_cast(size), std::ios_base::cur); + RAFT_EXPECTS(stream_.good(), "kvikio_file_reader: failed to advance the file position"); + } + + private: + file_descriptor fd_; + fd_istream stream_; + kvikio::FileHandle handle_; +}; + +kvikio_file_reader::kvikio_file_reader(const std::string& path) + : impl_(std::make_unique(path)) +{ +} + +kvikio_file_reader::~kvikio_file_reader() = default; + +std::istream& kvikio_file_reader::stream() { return impl_->stream(); } + +void kvikio_file_reader::read_device(void* data, size_t size) { impl_->read_device(data, size); } + // std::streambuf that stages output into a large buffer and writes full buffers to disk through // kvikio at an increasing file offset. The trailing partial buffer is written on sync()/close(). class kvikio_ofstream::sbuf : public std::streambuf { diff --git a/cpp/src/util/kvikio_serialize.hpp b/cpp/src/util/kvikio_serialize.hpp index fcddfc8e40..1b27ee6670 100644 --- a/cpp/src/util/kvikio_serialize.hpp +++ b/cpp/src/util/kvikio_serialize.hpp @@ -12,7 +12,9 @@ #include #include +#include #include +#include #include #include @@ -20,6 +22,10 @@ namespace cuvs::util::detail { inline constexpr size_t kDeviceSerializationBatchBytes = size_t{64} << 20; +inline std::istream& input_stream(std::istream& is) { return is; } + +inline std::istream& input_stream(kvikio_file_reader& reader) { return reader.stream(); } + template raft::numpy_serializer::header_t get_numpy_header(const std::vector& shape, bool fortran_order = false) @@ -34,6 +40,21 @@ void write_numpy_header(std::ostream& os, const std::vector& shape) raft::numpy_serializer::write_header(os, get_numpy_header(shape)); } +template +void validate_numpy_header(kvikio_file_reader& reader, + const std::vector& shape, + bool fortran_order = false) +{ + const auto expected = get_numpy_header(shape, fortran_order); + const auto actual = raft::numpy_serializer::read_header(reader.stream()); + RAFT_EXPECTS(actual.dtype.to_string() == expected.dtype.to_string(), + "Unexpected numpy dtype: expected %s, found %s", + expected.dtype.to_string().c_str(), + actual.dtype.to_string().c_str()); + RAFT_EXPECTS(actual.fortran_order == expected.fortran_order, "Unexpected numpy memory order"); + RAFT_EXPECTS(actual.shape == expected.shape, "Unexpected numpy array shape"); +} + template void serialize_device_mdspan( const raft::resources& res, @@ -59,4 +80,46 @@ void serialize_device_mdspan( RAFT_EXPECTS(os.good(), "Error writing content of device mdspan"); } +template +void deserialize_device_mdspan( + const raft::resources& res, + kvikio_file_reader& reader, + raft::device_mdspan obj) +{ + static_assert(std::is_same_v || + std::is_same_v, + "The serializer only supports row-major and column-major layouts"); + + using obj_t = raft::device_mdspan; + std::vector shape; + shape.reserve(obj.rank()); + for (typename obj_t::rank_type i = 0; i < obj.rank(); ++i) { + shape.push_back(static_cast(obj.extent(i))); + } + + const bool fortran_order = std::is_same_v; + validate_numpy_header(reader, shape, fortran_order); + + // KvikIO operates outside the RAFT stream. Ensure stream-ordered allocation and initialization + // of the destination have completed before handing its pointer to GDS. + raft::resource::sync_stream(res); + reader.read_device(obj.data_handle(), obj.size() * sizeof(ElementType)); +} + +template +void deserialize_mdspan(const raft::resources& res, + std::istream& is, + raft::device_mdspan obj) +{ + raft::deserialize_mdspan(res, is, obj); +} + +template +void deserialize_mdspan(const raft::resources& res, + kvikio_file_reader& reader, + raft::device_mdspan obj) +{ + deserialize_device_mdspan(res, reader, obj); +} + } // namespace cuvs::util::detail diff --git a/cpp/tests/util/file_io_test.cpp b/cpp/tests/util/file_io_test.cpp index bf374a7c06..00598f7345 100644 --- a/cpp/tests/util/file_io_test.cpp +++ b/cpp/tests/util/file_io_test.cpp @@ -282,6 +282,49 @@ TEST(FileIO, KvikioOfstreamMixedHostAndDeviceWrites) EXPECT_EQ(read_whole_file(path), expected); } +// Buffered metadata reads and direct device reads must advance one shared logical position. The +// short prefix deliberately leaves unread bytes in fd_streambuf's read-ahead buffer. +TEST(FileIO, KvikioFileReaderMixedStreamAndDeviceReads) +{ + raft::resources res; + scratch_dir scratch; + const std::string path = scratch.file("reader_device.bin"); + const std::vector prefix{'m', 'e', 't', 'a', 'd', 'a', 't', 'a'}; + const std::vector payload = make_pattern((size_t{3} << 20) + 17, 31); + const std::vector suffix{'t', 'r', 'a', 'i', 'l', 'e', 'r'}; + + { + std::ofstream os(path, std::ios::out | std::ios::binary); + ASSERT_TRUE(os.good()); + os.write(prefix.data(), prefix.size()); + os.write(payload.data(), payload.size()); + os.write(suffix.data(), suffix.size()); + } + + auto device_payload = raft::make_device_vector(res, payload.size()); + raft::resource::sync_stream(res); + kvikio_file_reader reader(path); + + std::vector prefix_read(prefix.size()); + reader.stream().read(prefix_read.data(), prefix_read.size()); + ASSERT_TRUE(reader.stream().good()); + reader.read_device(device_payload.data_handle(), payload.size()); + std::vector suffix_read(suffix.size()); + reader.stream().read(suffix_read.data(), suffix_read.size()); + ASSERT_TRUE(reader.stream().good()); + + std::vector payload_read(payload.size()); + raft::copy(payload_read.data(), + device_payload.data_handle(), + payload_read.size(), + raft::resource::get_cuda_stream(res)); + raft::resource::sync_stream(res); + + EXPECT_EQ(prefix_read, prefix); + EXPECT_EQ(payload_read, payload); + EXPECT_EQ(suffix_read, suffix); +} + // Defensive checks on the bulk helpers. TEST(FileIO, InvalidArguments) { From ef5ecd1925fbc7b97139661b2cf582238670b82d Mon Sep 17 00:00:00 2001 From: Julian Miller Date: Wed, 8 Jul 2026 15:07:54 +0200 Subject: [PATCH 16/18] Support cgroup mem limits --- cpp/include/cuvs/util/host_memory.hpp | 35 ++- .../neighbors/detail/cagra/cagra_build.cuh | 23 +- cpp/src/neighbors/detail/hnsw.hpp | 15 +- cpp/src/util/host_memory.cpp | 285 +++++++++++++++++- cpp/src/util/host_memory.hpp | 30 ++ cpp/tests/CMakeLists.txt | 2 +- cpp/tests/util/host_memory_test.cpp | 158 ++++++++++ 7 files changed, 524 insertions(+), 24 deletions(-) create mode 100644 cpp/src/util/host_memory.hpp create mode 100644 cpp/tests/util/host_memory_test.cpp diff --git a/cpp/include/cuvs/util/host_memory.hpp b/cpp/include/cuvs/util/host_memory.hpp index 06e253c1db..501004ad45 100644 --- a/cpp/include/cuvs/util/host_memory.hpp +++ b/cpp/include/cuvs/util/host_memory.hpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ #pragma once @@ -7,18 +7,39 @@ #include #include -#include -#include + +#include +#include namespace CUVS_EXPORT cuvs { namespace util { +/** @brief Snapshot of host and cgroup memory available to the current process. */ +struct host_memory_info { + /** Host-wide MemAvailable from /proc/meminfo. */ + size_t system_available; + /** Effective memory available after applying any cgroup limit. */ + size_t available; + /** Hard limit of the most constrained cgroup ancestor, when finite. */ + std::optional cgroup_limit; + /** Current usage charged to the most constrained cgroup ancestor. */ + std::optional cgroup_current; +}; + +/** + * @brief Get host memory available to the current process. + * + * Combines host-wide MemAvailable with cgroup v1 or v2 memory headroom. When a finite cgroup limit + * is found, the effective value is `min(MemAvailable, limit - current_usage)`. Visible ancestor + * cgroups are considered because a parent can impose a tighter limit than the process's leaf + * cgroup. + */ +host_memory_info get_host_memory_info(); + /** - * @brief Get available host memory from /proc/meminfo + * @brief Get effective available host memory. * - * Queries the system for available memory by reading /proc/meminfo. - * This is useful for determining how much host memory can be used - * for buffering or temporary storage. + * Convenience wrapper around get_host_memory_info(). * * @return Available memory in bytes */ diff --git a/cpp/src/neighbors/detail/cagra/cagra_build.cuh b/cpp/src/neighbors/detail/cagra/cagra_build.cuh index cb6d0b6a92..4943ee31d8 100644 --- a/cpp/src/neighbors/detail/cagra/cagra_build.cuh +++ b/cpp/src/neighbors/detail/cagra/cagra_build.cuh @@ -874,13 +874,30 @@ bool ace_check_use_disk_mode(raft::resources const& res, bool guarantee_connectivity, ace_memory_requirements& mem) { + const auto host_memory = cuvs::util::get_host_memory_info(); + RAFT_EXPECTS(host_memory.available > 0, + "ACE: No host memory is available within the current system or cgroup limit"); + if (host_memory.cgroup_limit.has_value() && host_memory.cgroup_current.has_value()) { + const size_t cgroup_headroom = *host_memory.cgroup_current < *host_memory.cgroup_limit + ? *host_memory.cgroup_limit - *host_memory.cgroup_current + : 0; + RAFT_LOG_INFO( + "ACE: Cgroup host memory limit: %.2f GiB, current: %.2f GiB, headroom: %.2f GiB; " + "system available: %.2f GiB, effective available: %.2f GiB", + to_gib(*host_memory.cgroup_limit), + to_gib(*host_memory.cgroup_current), + to_gib(cgroup_headroom), + to_gib(host_memory.system_available), + to_gib(host_memory.available)); + } + // Use overridden memory limits if provided (> 0), otherwise query actual system memory if (max_host_memory_gb.has_value() && max_host_memory_gb.value() > 0) { - auto actual_available_host_memory = cuvs::util::get_free_host_memory(); + auto actual_available_host_memory = host_memory.available; auto configured_host_memory = static_cast(max_host_memory_gb.value() * (1ULL << 30)); if (actual_available_host_memory < configured_host_memory) { RAFT_LOG_WARN( - "ACE: Actual host memory (%zu GiB) is less than configured limit (%zu GiB). " + "ACE: Actual host memory (%.2f GiB) is less than configured limit (%.2f GiB). " "Using actual host memory.", to_gib(actual_available_host_memory), to_gib(configured_host_memory)); @@ -891,7 +908,7 @@ bool ace_check_use_disk_mode(raft::resources const& res, mem.available_host_memory = configured_host_memory; } } else { - mem.available_host_memory = cuvs::util::get_free_host_memory(); + mem.available_host_memory = host_memory.available; } size_t sub_partition_size = static_cast(imbalance_factor * vector_expansion_factor * diff --git a/cpp/src/neighbors/detail/hnsw.hpp b/cpp/src/neighbors/detail/hnsw.hpp index 654d312702..35bb6370fa 100644 --- a/cpp/src/neighbors/detail/hnsw.hpp +++ b/cpp/src/neighbors/detail/hnsw.hpp @@ -1240,8 +1240,19 @@ inline std::pair get_available_memory( { size_t available_host_memory = cuvs::util::get_free_host_memory(); if (max_host_memory_gb.has_value() && max_host_memory_gb.value() > 0) { - available_host_memory = static_cast(max_host_memory_gb.value() * (1ULL << 30)); - RAFT_LOG_INFO("ACE: Using overridden host memory limit: %.2f GiB", max_host_memory_gb.value()); + const auto configured_host_memory = + static_cast(max_host_memory_gb.value() * (1ULL << 30)); + if (available_host_memory < configured_host_memory) { + RAFT_LOG_WARN( + "ACE: Actual host memory (%.2f GiB) is less than configured limit (%.2f GiB). Using " + "actual host memory.", + static_cast(available_host_memory) / (1ULL << 30), + max_host_memory_gb.value()); + } else { + available_host_memory = configured_host_memory; + RAFT_LOG_INFO("ACE: Using overridden host memory limit: %.2f GiB", + max_host_memory_gb.value()); + } } // Note: We use total device memory rather than free memory because RMM pools // and other allocators may report artificially low free memory. The assumption diff --git a/cpp/src/util/host_memory.cpp b/cpp/src/util/host_memory.cpp index 10630749eb..7edef19f48 100644 --- a/cpp/src/util/host_memory.cpp +++ b/cpp/src/util/host_memory.cpp @@ -1,29 +1,292 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ #include +#include "host_memory.hpp" + +#include +#include +#include +#include +#include #include +#include +#include +#include #include +#include +#include +#include +#include namespace cuvs::util { +namespace { -size_t get_free_host_memory() +struct cgroup_membership { + detail::cgroup_version version; + std::filesystem::path path; +}; + +struct cgroup_mount { + detail::cgroup_version version; + std::filesystem::path root; + std::filesystem::path mount_point; +}; + +std::optional parse_size(std::string_view text) +{ + while (!text.empty() && std::isspace(static_cast(text.front()))) { + text.remove_prefix(1); + } + while (!text.empty() && std::isspace(static_cast(text.back()))) { + text.remove_suffix(1); + } + if (text.empty()) { return std::nullopt; } + + size_t value = 0; + const auto [end, error] = std::from_chars(text.data(), text.data() + text.size(), value); + if (error != std::errc{} || end != text.data() + text.size()) { return std::nullopt; } + return value; +} + +std::optional read_first_token(const std::filesystem::path& path) +{ + std::ifstream input(path); + std::string token; + if (!(input >> token)) { return std::nullopt; } + return token; +} + +bool contains_controller(std::string_view controllers, std::string_view expected) +{ + while (!controllers.empty()) { + const auto separator = controllers.find(','); + const auto controller = controllers.substr(0, separator); + if (controller == expected) { return true; } + if (separator == std::string_view::npos) { break; } + controllers.remove_prefix(separator + 1); + } + return false; +} + +std::vector split_whitespace(std::string_view text) +{ + std::istringstream stream{std::string{text}}; + std::vector fields; + for (std::string field; stream >> field;) { + fields.push_back(std::move(field)); + } + return fields; +} + +std::string unescape_mount_path(std::string_view path) +{ + std::string result; + result.reserve(path.size()); + for (size_t i = 0; i < path.size(); ++i) { + if (path[i] == '\\' && i + 3 < path.size() && path[i + 1] >= '0' && path[i + 1] <= '7' && + path[i + 2] >= '0' && path[i + 2] <= '7' && path[i + 3] >= '0' && path[i + 3] <= '7') { + const auto value = + static_cast((path[i + 1] - '0') * 64 + (path[i + 2] - '0') * 8 + (path[i + 3] - '0')); + result.push_back(value); + i += 3; + } else { + result.push_back(path[i]); + } + } + return result; +} + +std::vector parse_cgroup_memberships(const std::filesystem::path& path) +{ + std::ifstream input(path); + std::vector memberships; + for (std::string line; std::getline(input, line);) { + const auto first_separator = line.find(':'); + const auto second_separator = first_separator == std::string::npos + ? std::string::npos + : line.find(':', first_separator + 1); + if (second_separator == std::string::npos) { continue; } + + const std::string_view hierarchy{line.data(), first_separator}; + const std::string_view controllers{line.data() + first_separator + 1, + second_separator - first_separator - 1}; + const std::string_view cgroup_path{line.data() + second_separator + 1, + line.size() - second_separator - 1}; + if (cgroup_path.empty() || cgroup_path.front() != '/') { continue; } + + if (hierarchy == "0" && controllers.empty()) { + memberships.push_back({detail::cgroup_version::v2, std::filesystem::path{cgroup_path}}); + } else if (contains_controller(controllers, "memory")) { + memberships.push_back({detail::cgroup_version::v1, std::filesystem::path{cgroup_path}}); + } + } + return memberships; +} + +std::vector parse_cgroup_mounts(const std::filesystem::path& path) +{ + std::ifstream input(path); + std::vector mounts; + for (std::string line; std::getline(input, line);) { + const auto separator = line.find(" - "); + if (separator == std::string::npos) { continue; } + + const auto left = split_whitespace(std::string_view{line}.substr(0, separator)); + const auto right = split_whitespace(std::string_view{line}.substr(separator + 3)); + if (left.size() < 6 || right.size() < 3) { continue; } + + std::optional version; + if (right[0] == "cgroup2") { + version = detail::cgroup_version::v2; + } else if (right[0] == "cgroup" && (contains_controller(left[5], "memory") || + contains_controller(right[2], "memory"))) { + version = detail::cgroup_version::v1; + } + if (!version.has_value()) { continue; } + + mounts.push_back({*version, + std::filesystem::path{unescape_mount_path(left[3])}.lexically_normal(), + std::filesystem::path{unescape_mount_path(left[4])}.lexically_normal()}); + } + return mounts; +} + +bool is_path_prefix(const std::filesystem::path& prefix, const std::filesystem::path& path) +{ + return std::mismatch(prefix.begin(), prefix.end(), path.begin(), path.end()).first == + prefix.end(); +} + +std::optional resolve_cgroup_path(const cgroup_membership& membership, + const cgroup_mount& mount) +{ + if (membership.version != mount.version) { return std::nullopt; } + + const auto member_path = membership.path.lexically_normal(); + if (is_path_prefix(mount.root, member_path)) { + const auto relative = member_path.lexically_relative(mount.root); + if (relative.empty()) { return std::nullopt; } + return relative == "." ? mount.mount_point : (mount.mount_point / relative).lexically_normal(); + } + + // In a cgroup namespace the membership path is relative to the namespace root, while mountinfo + // can still expose the original cgroup as the mount root. + return (mount.mount_point / member_path.relative_path()).lexically_normal(); +} + +bool is_v1_unlimited(size_t limit) +{ + // The v1 controller represents unlimited with a page-aligned value close to LONG_MAX. + return limit >= static_cast(std::numeric_limits::max() / 2); +} + +std::optional inspect_cgroup_hierarchy( + detail::cgroup_version version, + std::filesystem::path current_path, + const std::filesystem::path& mount_point) +{ + std::optional result; + if (!is_path_prefix(mount_point, current_path)) { return result; } + + while (true) { + const auto limit_file = + current_path / + (version == detail::cgroup_version::v2 ? "memory.max" : "memory.limit_in_bytes"); + const auto current_file = + current_path / + (version == detail::cgroup_version::v2 ? "memory.current" : "memory.usage_in_bytes"); + const auto limit_token = read_first_token(limit_file); + const auto current_token = read_first_token(current_file); + + if (limit_token.has_value() && *limit_token != "max" && current_token.has_value()) { + const auto limit = parse_size(*limit_token); + const auto current = parse_size(*current_token); + if (limit.has_value() && current.has_value() && + !(version == detail::cgroup_version::v1 && is_v1_unlimited(*limit))) { + const size_t available = *current < *limit ? *limit - *current : 0; + if (!result.has_value() || available < result->available) { + result = + detail::cgroup_memory_info{*limit, *current, available, version, current_path.string()}; + } + } + } + + if (current_path == mount_point) { break; } + const auto parent = current_path.parent_path(); + if (parent == current_path || !is_path_prefix(mount_point, parent)) { break; } + current_path = parent; + } + return result; +} + +size_t get_system_available_memory() { - size_t available_memory = 0; std::ifstream meminfo("/proc/meminfo"); - std::string line; - while (std::getline(meminfo, line)) { - if (line.find("MemAvailable:") != std::string::npos) { - available_memory = std::stoull(line.substr(line.find(":") + 1)); + for (std::string line; std::getline(meminfo, line);) { + constexpr std::string_view key = "MemAvailable:"; + if (line.rfind(key, 0) != 0) { continue; } + + std::istringstream value_stream{line.substr(key.size())}; + size_t kilobytes = 0; + std::string unit; + if (value_stream >> kilobytes >> unit && unit == "kB" && + kilobytes <= std::numeric_limits::max() / 1024) { + return kilobytes * 1024; + } + } + RAFT_FAIL("Failed to get available memory from /proc/meminfo"); +} + +} // namespace + +namespace detail { + +std::optional get_cgroup_memory_info( + const std::filesystem::path& proc_self_cgroup, const std::filesystem::path& proc_self_mountinfo) +{ + const auto memberships = parse_cgroup_memberships(proc_self_cgroup); + const auto mounts = parse_cgroup_mounts(proc_self_mountinfo); + std::optional result; + + for (const auto& membership : memberships) { + for (const auto& mount : mounts) { + const auto resolved_path = resolve_cgroup_path(membership, mount); + if (!resolved_path.has_value()) { continue; } + const auto candidate = + inspect_cgroup_hierarchy(membership.version, *resolved_path, mount.mount_point); + if (candidate.has_value() && + (!result.has_value() || candidate->available < result->available)) { + result = candidate; + } } } - available_memory *= 1024; - meminfo.close(); - RAFT_EXPECTS(available_memory > 0, "Failed to get available memory from /proc/meminfo"); - return available_memory; + return result; +} + +} // namespace detail + +host_memory_info get_host_memory_info() +{ + const size_t system_available = get_system_available_memory(); + host_memory_info result{system_available, system_available, std::nullopt, std::nullopt}; + if (const auto cgroup = detail::get_cgroup_memory_info(); cgroup.has_value()) { + result.cgroup_limit = cgroup->limit; + result.cgroup_current = cgroup->current; + result.available = std::min(system_available, cgroup->available); + } + return result; +} + +size_t get_free_host_memory() +{ + const auto memory = get_host_memory_info(); + RAFT_EXPECTS(memory.available > 0, + "No host memory is available within the current system or cgroup limit"); + return memory.available; } } // namespace cuvs::util diff --git a/cpp/src/util/host_memory.hpp b/cpp/src/util/host_memory.hpp new file mode 100644 index 0000000000..0b084465f6 --- /dev/null +++ b/cpp/src/util/host_memory.hpp @@ -0,0 +1,30 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +#pragma once + +#include + +#include +#include +#include +#include + +namespace cuvs::util::detail { + +enum class cgroup_version { v1, v2 }; + +struct cgroup_memory_info { + size_t limit; + size_t current; + size_t available; + cgroup_version version; + std::string path; +}; + +CUVS_EXPORT std::optional get_cgroup_memory_info( + const std::filesystem::path& proc_self_cgroup = "/proc/self/cgroup", + const std::filesystem::path& proc_self_mountinfo = "/proc/self/mountinfo"); + +} // namespace cuvs::util::detail diff --git a/cpp/tests/CMakeLists.txt b/cpp/tests/CMakeLists.txt index db8ea693cb..20a7377a67 100644 --- a/cpp/tests/CMakeLists.txt +++ b/cpp/tests/CMakeLists.txt @@ -395,7 +395,7 @@ ConfigureTest( ConfigureTest( NAME UTIL_TEST - PATH util/file_io_test.cpp + PATH util/file_io_test.cpp util/host_memory_test.cpp GPUS 1 PERCENT 100 ) diff --git a/cpp/tests/util/host_memory_test.cpp b/cpp/tests/util/host_memory_test.cpp new file mode 100644 index 0000000000..e7edb5039d --- /dev/null +++ b/cpp/tests/util/host_memory_test.cpp @@ -0,0 +1,158 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "../../src/util/host_memory.hpp" + +#include + +#include +#include +#include +#include +#include +#include + +namespace cuvs::util::detail { +namespace { + +class cgroup_fixture { + public: + cgroup_fixture() + : root_(std::filesystem::temp_directory_path() / + ("cuvs_cgroup_test_" + std::to_string(::getpid()) + "_" + + std::to_string(reinterpret_cast(this)))) + { + std::filesystem::create_directories(root_); + } + + ~cgroup_fixture() + { + std::error_code error; + std::filesystem::remove_all(root_, error); + } + + [[nodiscard]] std::filesystem::path path(const std::string& name) const { return root_ / name; } + + void write(const std::filesystem::path& path, const std::string& content) + { + std::filesystem::create_directories(path.parent_path()); + std::ofstream output(path); + ASSERT_TRUE(output.good()); + output << content; + ASSERT_TRUE(output.good()); + } + + private: + std::filesystem::path root_; +}; + +TEST(HostMemory, CgroupV2UsesTightestAncestorHeadroom) +{ + cgroup_fixture fixture; + const auto mount = fixture.path("cgroup2"); + const auto leaf = mount / "tenant/job"; + + fixture.write(fixture.path("self.cgroup"), "0::/tenant/job\n"); + fixture.write( + fixture.path("self.mountinfo"), + "29 23 0:26 / " + mount.string() + " rw,nosuid,nodev,noexec,relatime - cgroup2 cgroup rw\n"); + fixture.write(leaf / "memory.max", "900\n"); + fixture.write(leaf / "memory.current", "200\n"); + fixture.write(mount / "tenant/memory.max", "1000\n"); + fixture.write(mount / "tenant/memory.current", "400\n"); + fixture.write(mount / "memory.max", "max\n"); + fixture.write(mount / "memory.current", "500\n"); + + const auto result = + get_cgroup_memory_info(fixture.path("self.cgroup"), fixture.path("self.mountinfo")); + ASSERT_TRUE(result.has_value()); + EXPECT_EQ(result->version, cgroup_version::v2); + EXPECT_EQ(result->limit, 1000); + EXPECT_EQ(result->current, 400); + EXPECT_EQ(result->available, 600); + EXPECT_EQ(result->path, (mount / "tenant").string()); +} + +TEST(HostMemory, CgroupV2ResolvesNamespacedRoot) +{ + cgroup_fixture fixture; + const auto mount = fixture.path("cgroup2"); + const auto leaf = mount / "job"; + + fixture.write(fixture.path("self.cgroup"), "0::/job\n"); + fixture.write(fixture.path("self.mountinfo"), + "29 23 0:26 /docker/container " + mount.string() + + " rw,nosuid,nodev,noexec,relatime - cgroup2 cgroup rw\n"); + fixture.write(leaf / "memory.max", "2048\n"); + fixture.write(leaf / "memory.current", "512\n"); + + const auto result = + get_cgroup_memory_info(fixture.path("self.cgroup"), fixture.path("self.mountinfo")); + ASSERT_TRUE(result.has_value()); + EXPECT_EQ(result->version, cgroup_version::v2); + EXPECT_EQ(result->available, 1536); + EXPECT_EQ(result->path, leaf.string()); +} + +TEST(HostMemory, CgroupV1UsesMemoryControllerAndSaturatesHeadroom) +{ + cgroup_fixture fixture; + const auto mount = fixture.path("memory"); + const auto leaf = mount / "container"; + + fixture.write(fixture.path("self.cgroup"), + "7:cpu:/docker/container\n8:memory:/docker/container\n"); + fixture.write(fixture.path("self.mountinfo"), + "31 23 0:28 /docker " + mount.string() + + " rw,nosuid,nodev,noexec,relatime - cgroup cgroup rw,memory\n"); + fixture.write(leaf / "memory.limit_in_bytes", "4096\n"); + fixture.write(leaf / "memory.usage_in_bytes", "8192\n"); + + const auto result = + get_cgroup_memory_info(fixture.path("self.cgroup"), fixture.path("self.mountinfo")); + ASSERT_TRUE(result.has_value()); + EXPECT_EQ(result->version, cgroup_version::v1); + EXPECT_EQ(result->limit, 4096); + EXPECT_EQ(result->current, 8192); + EXPECT_EQ(result->available, 0); +} + +TEST(HostMemory, UnlimitedAndMalformedLimitsAreIgnored) +{ + cgroup_fixture fixture; + const auto mount = fixture.path("cgroup2"); + const auto leaf = mount / "job"; + + fixture.write(fixture.path("self.cgroup"), "0::/job\n"); + fixture.write( + fixture.path("self.mountinfo"), + "29 23 0:26 / " + mount.string() + " rw,nosuid,nodev,noexec,relatime - cgroup2 cgroup rw\n"); + fixture.write(leaf / "memory.max", "not-a-limit\n"); + fixture.write(leaf / "memory.current", "100\n"); + fixture.write(mount / "memory.max", "max\n"); + fixture.write(mount / "memory.current", "100\n"); + + EXPECT_FALSE(get_cgroup_memory_info(fixture.path("self.cgroup"), fixture.path("self.mountinfo")) + .has_value()); +} + +TEST(HostMemory, CgroupV1UnlimitedSentinelIsIgnored) +{ + cgroup_fixture fixture; + const auto mount = fixture.path("memory"); + + fixture.write(fixture.path("self.cgroup"), "8:memory:/container\n"); + fixture.write(fixture.path("self.mountinfo"), + "31 23 0:28 / " + mount.string() + + " rw,nosuid,nodev,noexec,relatime - cgroup cgroup rw,memory\n"); + fixture.write(mount / "container/memory.limit_in_bytes", "9223372036854771712\n"); + fixture.write(mount / "container/memory.usage_in_bytes", "1024\n"); + + EXPECT_FALSE(get_cgroup_memory_info(fixture.path("self.cgroup"), fixture.path("self.mountinfo")) + .has_value()); +} + +} // namespace +} // namespace cuvs::util::detail From c1bf4f5d21bd1ff6d90921ad53932f5a41c09c2e Mon Sep 17 00:00:00 2001 From: Julian Miller Date: Wed, 8 Jul 2026 15:48:23 +0200 Subject: [PATCH 17/18] Fix data conversion warning in read_device --- cpp/src/util/file_io.cpp | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/cpp/src/util/file_io.cpp b/cpp/src/util/file_io.cpp index e410536280..25a2b55744 100644 --- a/cpp/src/util/file_io.cpp +++ b/cpp/src/util/file_io.cpp @@ -153,9 +153,9 @@ class kvikio_file_reader::impl { const size_t bytes_read = handle_.pread(data, size, static_cast(offset)).get(); RAFT_EXPECTS(bytes_read == size, - "kvikio_file_reader: short read (expected %zu, read %zu)", - size, - bytes_read); + "kvikio_file_reader: short read (expected %llu, read %llu)", + static_cast(size), + static_cast(bytes_read)); stream_.seekg(static_cast(size), std::ios_base::cur); RAFT_EXPECTS(stream_.good(), "kvikio_file_reader: failed to advance the file position"); From 446c8a9f1ce46ee847dad3b2dc15be37bbf1f944 Mon Sep 17 00:00:00 2001 From: Julian Miller Date: Wed, 8 Jul 2026 19:50:00 +0200 Subject: [PATCH 18/18] Account for cgroup reclaimable mem --- cpp/include/cuvs/util/host_memory.hpp | 10 ++- .../neighbors/detail/cagra/cagra_build.cuh | 15 ++-- cpp/src/util/host_memory.cpp | 79 +++++++++++++++++-- cpp/src/util/host_memory.hpp | 2 + cpp/tests/util/host_memory_test.cpp | 46 +++++++++-- 5 files changed, 132 insertions(+), 20 deletions(-) diff --git a/cpp/include/cuvs/util/host_memory.hpp b/cpp/include/cuvs/util/host_memory.hpp index 501004ad45..626d1cde60 100644 --- a/cpp/include/cuvs/util/host_memory.hpp +++ b/cpp/include/cuvs/util/host_memory.hpp @@ -24,13 +24,19 @@ struct host_memory_info { std::optional cgroup_limit; /** Current usage charged to the most constrained cgroup ancestor. */ std::optional cgroup_current; + /** Clean file cache treated as reclaimable for capacity planning. */ + std::optional cgroup_reclaimable_file; + /** Current cgroup usage after excluding reclaimable file cache. */ + std::optional cgroup_working_set; }; /** * @brief Get host memory available to the current process. * - * Combines host-wide MemAvailable with cgroup v1 or v2 memory headroom. When a finite cgroup limit - * is found, the effective value is `min(MemAvailable, limit - current_usage)`. Visible ancestor + * Combines host-wide MemAvailable with cgroup v1 or v2 memory headroom. For cgroup v2, clean file + * cache from memory.stat excludes shmem, dirty, writeback, and unevictable pages from the total + * file cache. Cgroup v1 uses hierarchical inactive file cache when available. When a finite cgroup + * limit is found, the effective value is `min(MemAvailable, limit - working_set)`. Visible ancestor * cgroups are considered because a parent can impose a tighter limit than the process's leaf * cgroup. */ diff --git a/cpp/src/neighbors/detail/cagra/cagra_build.cuh b/cpp/src/neighbors/detail/cagra/cagra_build.cuh index 4943ee31d8..938753bd1d 100644 --- a/cpp/src/neighbors/detail/cagra/cagra_build.cuh +++ b/cpp/src/neighbors/detail/cagra/cagra_build.cuh @@ -877,15 +877,20 @@ bool ace_check_use_disk_mode(raft::resources const& res, const auto host_memory = cuvs::util::get_host_memory_info(); RAFT_EXPECTS(host_memory.available > 0, "ACE: No host memory is available within the current system or cgroup limit"); - if (host_memory.cgroup_limit.has_value() && host_memory.cgroup_current.has_value()) { - const size_t cgroup_headroom = *host_memory.cgroup_current < *host_memory.cgroup_limit - ? *host_memory.cgroup_limit - *host_memory.cgroup_current + if (host_memory.cgroup_limit.has_value() && host_memory.cgroup_current.has_value() && + host_memory.cgroup_reclaimable_file.has_value() && + host_memory.cgroup_working_set.has_value()) { + const size_t cgroup_headroom = *host_memory.cgroup_working_set < *host_memory.cgroup_limit + ? *host_memory.cgroup_limit - *host_memory.cgroup_working_set : 0; RAFT_LOG_INFO( - "ACE: Cgroup host memory limit: %.2f GiB, current: %.2f GiB, headroom: %.2f GiB; " - "system available: %.2f GiB, effective available: %.2f GiB", + "ACE: Cgroup host memory limit: %.2f GiB, current: %.2f GiB, reclaimable file cache: %.2f " + "GiB, working set: %.2f GiB, headroom: %.2f GiB; system available: %.2f GiB, effective " + "available: %.2f GiB", to_gib(*host_memory.cgroup_limit), to_gib(*host_memory.cgroup_current), + to_gib(*host_memory.cgroup_reclaimable_file), + to_gib(*host_memory.cgroup_working_set), to_gib(cgroup_headroom), to_gib(host_memory.system_available), to_gib(host_memory.available)); diff --git a/cpp/src/util/host_memory.cpp b/cpp/src/util/host_memory.cpp index 7edef19f48..2d1e7499b6 100644 --- a/cpp/src/util/host_memory.cpp +++ b/cpp/src/util/host_memory.cpp @@ -184,6 +184,60 @@ bool is_v1_unlimited(size_t limit) return limit >= static_cast(std::numeric_limits::max() / 2); } +size_t subtract_saturating(size_t value, size_t amount) +{ + return amount < value ? value - amount : 0; +} + +size_t get_reclaimable_file_memory(const std::filesystem::path& path, + detail::cgroup_version version) +{ + std::ifstream input(path / "memory.stat"); + std::optional file; + size_t shmem = 0; + size_t file_dirty = 0; + size_t file_writeback = 0; + size_t unevictable = 0; + std::optional inactive_file; + std::optional total_inactive_file; + for (std::string key, value; input >> key >> value;) { + const auto parsed_value = parse_size(value); + if (!parsed_value.has_value()) { continue; } + if (key == "file") { + file = *parsed_value; + } else if (key == "shmem") { + shmem = *parsed_value; + } else if (key == "file_dirty") { + file_dirty = *parsed_value; + } else if (key == "file_writeback") { + file_writeback = *parsed_value; + } else if (key == "unevictable") { + unevictable = *parsed_value; + } else if (key == "inactive_file") { + inactive_file = *parsed_value; + } else if (key == "total_inactive_file") { + total_inactive_file = *parsed_value; + } + } + + // v1's total_ counter includes descendants and therefore matches hierarchical usage accounting. + if (version == detail::cgroup_version::v1 && total_inactive_file.has_value()) { + return *total_inactive_file; + } + + if (version == detail::cgroup_version::v2 && file.has_value()) { + size_t clean_file = *file; + clean_file = subtract_saturating(clean_file, shmem); + clean_file = subtract_saturating(clean_file, file_dirty); + clean_file = subtract_saturating(clean_file, file_writeback); + clean_file = subtract_saturating(clean_file, unevictable); + return clean_file; + } + + // Older or restricted cgroup interfaces may not expose the full v2 breakdown. + return inactive_file.value_or(0); +} + std::optional inspect_cgroup_hierarchy( detail::cgroup_version version, std::filesystem::path current_path, @@ -207,10 +261,18 @@ std::optional inspect_cgroup_hierarchy( const auto current = parse_size(*current_token); if (limit.has_value() && current.has_value() && !(version == detail::cgroup_version::v1 && is_v1_unlimited(*limit))) { - const size_t available = *current < *limit ? *limit - *current : 0; + const size_t reclaimable_file = + std::min(*current, get_reclaimable_file_memory(current_path, version)); + const size_t working_set = *current - reclaimable_file; + const size_t available = working_set < *limit ? *limit - working_set : 0; if (!result.has_value() || available < result->available) { - result = - detail::cgroup_memory_info{*limit, *current, available, version, current_path.string()}; + result = detail::cgroup_memory_info{*limit, + *current, + reclaimable_file, + working_set, + available, + version, + current_path.string()}; } } } @@ -272,11 +334,14 @@ std::optional get_cgroup_memory_info( host_memory_info get_host_memory_info() { const size_t system_available = get_system_available_memory(); - host_memory_info result{system_available, system_available, std::nullopt, std::nullopt}; + host_memory_info result{ + system_available, system_available, std::nullopt, std::nullopt, std::nullopt, std::nullopt}; if (const auto cgroup = detail::get_cgroup_memory_info(); cgroup.has_value()) { - result.cgroup_limit = cgroup->limit; - result.cgroup_current = cgroup->current; - result.available = std::min(system_available, cgroup->available); + result.cgroup_limit = cgroup->limit; + result.cgroup_current = cgroup->current; + result.cgroup_reclaimable_file = cgroup->reclaimable_file; + result.cgroup_working_set = cgroup->working_set; + result.available = std::min(system_available, cgroup->available); } return result; } diff --git a/cpp/src/util/host_memory.hpp b/cpp/src/util/host_memory.hpp index 0b084465f6..353be1d387 100644 --- a/cpp/src/util/host_memory.hpp +++ b/cpp/src/util/host_memory.hpp @@ -18,6 +18,8 @@ enum class cgroup_version { v1, v2 }; struct cgroup_memory_info { size_t limit; size_t current; + size_t reclaimable_file; + size_t working_set; size_t available; cgroup_version version; std::string path; diff --git a/cpp/tests/util/host_memory_test.cpp b/cpp/tests/util/host_memory_test.cpp index e7edb5039d..09bd9e84b9 100644 --- a/cpp/tests/util/host_memory_test.cpp +++ b/cpp/tests/util/host_memory_test.cpp @@ -59,9 +59,13 @@ TEST(HostMemory, CgroupV2UsesTightestAncestorHeadroom) fixture.path("self.mountinfo"), "29 23 0:26 / " + mount.string() + " rw,nosuid,nodev,noexec,relatime - cgroup2 cgroup rw\n"); fixture.write(leaf / "memory.max", "900\n"); - fixture.write(leaf / "memory.current", "200\n"); + fixture.write(leaf / "memory.current", "800\n"); + fixture.write(leaf / "memory.stat", "file 700\ninactive_file 20\nactive_file 680\n"); fixture.write(mount / "tenant/memory.max", "1000\n"); - fixture.write(mount / "tenant/memory.current", "400\n"); + fixture.write(mount / "tenant/memory.current", "900\n"); + fixture.write(mount / "tenant/memory.stat", + "file 250\nshmem 20\nfile_dirty 10\nfile_writeback 5\nunevictable 15\n" + "inactive_file 50\nactive_file 200\n"); fixture.write(mount / "memory.max", "max\n"); fixture.write(mount / "memory.current", "500\n"); @@ -70,8 +74,10 @@ TEST(HostMemory, CgroupV2UsesTightestAncestorHeadroom) ASSERT_TRUE(result.has_value()); EXPECT_EQ(result->version, cgroup_version::v2); EXPECT_EQ(result->limit, 1000); - EXPECT_EQ(result->current, 400); - EXPECT_EQ(result->available, 600); + EXPECT_EQ(result->current, 900); + EXPECT_EQ(result->reclaimable_file, 200); + EXPECT_EQ(result->working_set, 700); + EXPECT_EQ(result->available, 300); EXPECT_EQ(result->path, (mount / "tenant").string()); } @@ -92,11 +98,36 @@ TEST(HostMemory, CgroupV2ResolvesNamespacedRoot) get_cgroup_memory_info(fixture.path("self.cgroup"), fixture.path("self.mountinfo")); ASSERT_TRUE(result.has_value()); EXPECT_EQ(result->version, cgroup_version::v2); + EXPECT_EQ(result->reclaimable_file, 0); + EXPECT_EQ(result->working_set, 512); EXPECT_EQ(result->available, 1536); EXPECT_EQ(result->path, leaf.string()); } -TEST(HostMemory, CgroupV1UsesMemoryControllerAndSaturatesHeadroom) +TEST(HostMemory, CgroupV2ProtectedFileMemoryDoesNotUnderflow) +{ + cgroup_fixture fixture; + const auto mount = fixture.path("cgroup2"); + const auto leaf = mount / "job"; + + fixture.write(fixture.path("self.cgroup"), "0::/job\n"); + fixture.write( + fixture.path("self.mountinfo"), + "29 23 0:26 / " + mount.string() + " rw,nosuid,nodev,noexec,relatime - cgroup2 cgroup rw\n"); + fixture.write(leaf / "memory.max", "100\n"); + fixture.write(leaf / "memory.current", "120\n"); + fixture.write(leaf / "memory.stat", + "file 100\nshmem 80\nfile_dirty 30\nfile_writeback 10\nunevictable 5\n"); + + const auto result = + get_cgroup_memory_info(fixture.path("self.cgroup"), fixture.path("self.mountinfo")); + ASSERT_TRUE(result.has_value()); + EXPECT_EQ(result->reclaimable_file, 0); + EXPECT_EQ(result->working_set, 120); + EXPECT_EQ(result->available, 0); +} + +TEST(HostMemory, CgroupV1UsesHierarchicalInactiveFile) { cgroup_fixture fixture; const auto mount = fixture.path("memory"); @@ -109,6 +140,7 @@ TEST(HostMemory, CgroupV1UsesMemoryControllerAndSaturatesHeadroom) " rw,nosuid,nodev,noexec,relatime - cgroup cgroup rw,memory\n"); fixture.write(leaf / "memory.limit_in_bytes", "4096\n"); fixture.write(leaf / "memory.usage_in_bytes", "8192\n"); + fixture.write(leaf / "memory.stat", "inactive_file 128\ntotal_inactive_file 5000\n"); const auto result = get_cgroup_memory_info(fixture.path("self.cgroup"), fixture.path("self.mountinfo")); @@ -116,7 +148,9 @@ TEST(HostMemory, CgroupV1UsesMemoryControllerAndSaturatesHeadroom) EXPECT_EQ(result->version, cgroup_version::v1); EXPECT_EQ(result->limit, 4096); EXPECT_EQ(result->current, 8192); - EXPECT_EQ(result->available, 0); + EXPECT_EQ(result->reclaimable_file, 5000); + EXPECT_EQ(result->working_set, 3192); + EXPECT_EQ(result->available, 904); } TEST(HostMemory, UnlimitedAndMalformedLimitsAreIgnored)