diff --git a/c/CMakeLists.txt b/c/CMakeLists.txt index 25e6455887..eab7560dfc 100644 --- a/c/CMakeLists.txt +++ b/c/CMakeLists.txt @@ -173,6 +173,7 @@ target_link_libraries( $,cuvs::cuvs_static,cuvs::cuvs> # $ # enforce we shouldn't use raft symbols $ + $> ) # ################################################################################################## @@ -192,6 +193,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 @@ -216,6 +230,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() diff --git a/ci/build_standalone_c.sh b/ci/build_standalone_c.sh index 94cffe1d64..2b8e0863f9 100755 --- a/ci/build_standalone_c.sh +++ b/ci/build_standalone_c.sh @@ -86,6 +86,16 @@ 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. +if ! cmake --install cpp/build \ + --prefix c/build/install \ + --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 cmake --install c/build --prefix c/build/install --component testing diff --git a/ci/build_wheel.sh b/ci/build_wheel.sh index b8e09827bd..a59b04248c 100755 --- a/ci/build_wheel.sh +++ b/ci/build_wheel.sh @@ -41,6 +41,7 @@ EXCLUDE_ARGS=( --exclude "libcurand.so.*" --exclude "libcusolver.so.*" --exclude "libcusparse.so.*" + --exclude "libkvikio.so*" --exclude "libnccl.so.*" --exclude "libnvJitLink.so.*" --exclude "libnvrtc.so.*" diff --git a/ci/release/update-version.sh b/ci/release/update-version.sh index 03cbb95f92..3a5d51baa0 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-aarch64.yaml b/conda/environments/all_cuda-129_arch-aarch64.yaml index 4c152e43da..0bf1117ba1 100644 --- a/conda/environments/all_cuda-129_arch-aarch64.yaml +++ b/conda/environments/all_cuda-129_arch-aarch64.yaml @@ -26,6 +26,7 @@ dependencies: - libcurand-dev - libcusolver-dev - libcusparse-dev +- libkvikio==26.10.*,>=0.0.0a0 - libnvjitlink-dev - libopenblas<=0.3.30 - librmm==26.10.*,>=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 37f0d5ab08..98d80d81a6 100644 --- a/conda/environments/all_cuda-129_arch-x86_64.yaml +++ b/conda/environments/all_cuda-129_arch-x86_64.yaml @@ -26,6 +26,7 @@ dependencies: - libcurand-dev - libcusolver-dev - libcusparse-dev +- libkvikio==26.10.*,>=0.0.0a0 - libnvjitlink-dev - librmm==26.10.*,>=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 545dcf5326..d95c6854bd 100644 --- a/conda/environments/all_cuda-133_arch-aarch64.yaml +++ b/conda/environments/all_cuda-133_arch-aarch64.yaml @@ -26,6 +26,7 @@ dependencies: - libcurand-dev - libcusolver-dev - libcusparse-dev +- libkvikio==26.10.*,>=0.0.0a0 - libnvjitlink-dev - libopenblas<=0.3.30 - librmm==26.10.*,>=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 7b06a9c20e..69d8f84605 100644 --- a/conda/environments/all_cuda-133_arch-x86_64.yaml +++ b/conda/environments/all_cuda-133_arch-x86_64.yaml @@ -26,6 +26,7 @@ dependencies: - libcurand-dev - libcusolver-dev - libcusparse-dev +- libkvikio==26.10.*,>=0.0.0a0 - libnvjitlink-dev - librmm==26.10.*,>=0.0.0a0 - make diff --git a/conda/recipes/libcuvs/recipe.yaml b/conda/recipes/libcuvs/recipe.yaml index 29aebfd06a..93a69ea1c0 100644 --- a/conda/recipes/libcuvs/recipe.yaml +++ b/conda/recipes/libcuvs/recipe.yaml @@ -75,6 +75,7 @@ cache: - libnvjitlink-dev - librmm =${{ minor_version }} - libraft-headers =${{ minor_version }} + - libkvikio =${{ minor_version }} - nccl ${{ nccl_version }} - cuda-version =${{ cuda_version }} - cuda-cudart-dev @@ -174,6 +175,7 @@ outputs: - ${{ pin_subpackage("libcuvs-headers", exact=True) }} - librmm =${{ minor_version }} - libraft-headers =${{ minor_version }} + - libkvikio =${{ minor_version }} - nccl ${{ nccl_version }} - cuda-version =${{ cuda_version }} - cuda-cudart-dev @@ -189,6 +191,7 @@ outputs: - ${{ pin_subpackage("libcuvs-headers", exact=True) }} - libraft-headers =${{ minor_version }} - librmm =${{ minor_version }} + - libkvikio =${{ minor_version }} - cuda-nvrtc - nccl - libcublas @@ -236,6 +239,7 @@ outputs: - ${{ pin_subpackage("libcuvs-headers", exact=True) }} - librmm =${{ minor_version }} - libraft-headers =${{ minor_version }} + - libkvikio =${{ minor_version }} - nccl ${{ nccl_version }} - cuda-version =${{ cuda_version }} - cuda-cudart-dev @@ -251,6 +255,7 @@ outputs: - ${{ pin_subpackage("libcuvs-headers", exact=True) }} - libraft-headers =${{ minor_version }} - librmm =${{ minor_version }} + - libkvikio =${{ minor_version }} - cuda-nvrtc - nccl - libcublas diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt index 4d51f17836..69ede8404e 100644 --- a/cpp/CMakeLists.txt +++ b/cpp/CMakeLists.txt @@ -187,6 +187,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() @@ -1494,6 +1495,7 @@ if(NOT BUILD_CPU_ONLY) ${CUVS_CTK_MATH_DEPENDENCIES} $ $ + kvikio::kvikio ) target_include_directories( @@ -1569,8 +1571,13 @@ if(NOT BUILD_CPU_ONLY) $> $> $<$:CUDA::nvtx3> - PRIVATE $ $ - $ CUDA::nvJitLink CUDA::nvrtc rtcx::rtcx + PRIVATE $ + $ + $ + CUDA::nvJitLink + CUDA::nvrtc + rtcx::rtcx + $ ) # ensure CUDA symbols aren't relocated to the middle of the debug build binaries @@ -1632,6 +1639,7 @@ SECTIONS $<$:CUDA::nvtx3> $ $ + kvikio::kvikio ) endif() @@ -1672,6 +1680,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) diff --git a/cpp/cmake/thirdparty/get_kvikio.cmake b/cpp/cmake/thirdparty/get_kvikio.cmake new file mode 100644 index 0000000000..130421ecd2 --- /dev/null +++ b/cpp/cmake/thirdparty/get_kvikio.cmake @@ -0,0 +1,43 @@ +# ============================================================================= +# cmake-format: off +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# 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. Shared cuVS builds consume it privately, while installed static cuVS + # targets retain it as a link-only dependency. At runtime the conda 'libkvikio' package provides + # both the shared library and its CMake package configuration. + # ---------------------------------------------------------------------------- + rapids_cpm_find( + kvikio ${PKG_VERSION} + GLOBAL_TARGETS kvikio::kvikio + BUILD_EXPORT_SET cuvs-static-exports + INSTALL_EXPORT_SET cuvs-static-exports + 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 + # Force KvikIO checkout to get public headers. + PATCH_COMMAND "" + 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/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/common.hpp b/cpp/include/cuvs/neighbors/common.hpp index 3be6cf1fa1..1b943afe30 100644 --- a/cpp/include/cuvs/neighbors/common.hpp +++ b/cpp/include/cuvs/neighbors/common.hpp @@ -1768,6 +1768,13 @@ enable_if_valid_list_t serialize_list( const typename ListT::spec_type& store_spec, std::optional size_override = std::nullopt); +/** + * Deserialize a list from an arbitrary input stream. + * + * This compatibility path stages list data through host memory because a std::istream does not + * expose a portable file path or descriptor. Index filename overloads use KvikIO and transfer list + * payloads directly to device memory when GDS is available. + */ template enable_if_valid_list_t deserialize_list(const raft::resources& handle, std::istream& is, diff --git a/cpp/include/cuvs/util/file_io.hpp b/cpp/include/cuvs/util/file_io.hpp index a7d67ec2c0..6ea85c49f9 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 @@ -12,6 +12,7 @@ #include #include #include +#include #include #include #include @@ -27,40 +28,132 @@ 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()); } + 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) { + 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; + + /** + * @brief Move-construct from @p other. + * + * Transfers ownership of the file descriptor and buffer state. The moved-from object is left in + * a valid empty state. + */ + fd_streambuf(fd_streambuf&& other) noexcept + : std::streambuf(std::move(other)), + fd_(std::exchange(other.fd_, -1)), + buffer_(std::move(other.buffer_)), + buffer_size_(std::exchange(other.buffer_size_, 0)) { - if (fd_ != -1) ::close(fd_); + other.setg(nullptr, nullptr, nullptr); + } + + /** + * @brief Move-assign from @p other. + * + * Transfers ownership of the file descriptor and buffer state. The moved-from object is left in + * a valid empty state. + */ + fd_streambuf& operator=(fd_streambuf&& other) noexcept + { + if (this != &other) { + close(); + std::streambuf::operator=(std::move(other)); + fd_ = std::exchange(other.fd_, -1); + buffer_ = std::move(other.buffer_); + buffer_size_ = std::exchange(other.buffer_size_, 0); + other.setg(nullptr, nullptr, nullptr); + } + return *this; } - 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; + private: + void close() noexcept + { + if (fd_ != -1) { + ::close(fd_); + fd_ = -1; + } + } }; /** @@ -92,8 +185,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: @@ -171,22 +265,57 @@ 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. * - * 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 @@ -197,42 +326,87 @@ std::pair create_numpy_file(const std::string& path, std::stringstream ss; raft::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; + data_bytes = checked_mul(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()); + // 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)); } - // 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()); - } - - 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. * - * Reads a file in chunks to avoid issues with very large reads. - * Uses pread for thread-safe, offset-based reading. + * Path-backed descriptors are routed through kvikio::FileHandle, so they use GPUDirect Storage + * (cuFile) when @p dest_ptr is device memory on a GDS-capable system, and KvikIO's compatible I/O + * backend otherwise. Descriptors constructed from a raw file descriptor without a path retain the + * POSIX pread fallback and require host 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 and, when available, its path. + * @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, @@ -240,35 +414,96 @@ 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. * - * 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. Path-backed descriptors use KvikIO and may transfer host or device + * memory. Descriptors without a path retain the POSIX pwrite fallback and require host 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 and, when available, its path. + * @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. + * + * 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. + * 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: + /** + * @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 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; + + 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/include/cuvs/util/host_memory.hpp b/cpp/include/cuvs/util/host_memory.hpp index 06e253c1db..626d1cde60 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,45 @@ #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; + /** 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. 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. + */ +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/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..f88609f498 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 @@ -18,9 +19,9 @@ namespace cuvs::neighbors::brute_force { int constexpr serialization_version = 0; -template +template void serialize(raft::resources const& handle, - std::ostream& os, + Output& os, const index& index, bool include_dataset = true) { @@ -37,10 +38,10 @@ void serialize(raft::resources const& handle, raft::serialize_scalar(handle, os, index.metric()); raft::serialize_scalar(handle, os, index.metric_arg()); raft::serialize_scalar(handle, os, include_dataset); - if (include_dataset) { raft::serialize_mdspan(handle, os, index.dataset()); } + if (include_dataset) { cuvs::util::detail::serialize_mdspan(handle, os, index.dataset()); } auto has_norms = index.has_norms(); raft::serialize_scalar(handle, os, has_norms); - if (has_norms) { raft::serialize_mdspan(handle, os, index.norms()); } + if (has_norms) { cuvs::util::detail::serialize_mdspan(handle, os, index.norms()); } raft::resource::sync_stream(handle); } @@ -49,9 +50,11 @@ void serialize(raft::resources const& handle, const index& index, bool include_dataset) { - auto os = std::ofstream{filename, std::ios::out | std::ios::binary}; + cuvs::util::kvikio_ofstream os(filename); RAFT_EXPECTS(os, "Cannot open file %s", filename.c_str()); serialize(handle, os, index, include_dataset); + os.close(); + RAFT_EXPECTS(os, "Error writing output %s", filename.c_str()); } void serialize(raft::resources const& handle, @@ -59,9 +62,11 @@ void serialize(raft::resources const& handle, const index& index, bool include_dataset) { - auto os = std::ofstream{filename, std::ios::out | std::ios::binary}; + cuvs::util::kvikio_ofstream os(filename); RAFT_EXPECTS(os, "Cannot open file %s", filename.c_str()); serialize(handle, os, index, include_dataset); + os.close(); + RAFT_EXPECTS(os, "Error writing output %s", filename.c_str()); } void serialize(raft::resources const& handle, @@ -80,9 +85,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 +114,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 +124,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 +145,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_build.cuh b/cpp/src/neighbors/detail/cagra/cagra_build.cuh index 8705926a41..c7743fd2d3 100644 --- a/cpp/src/neighbors/detail/cagra/cagra_build.cuh +++ b/cpp/src/neighbors/detail/cagra/cagra_build.cuh @@ -5,6 +5,7 @@ #pragma once #include "../../../core/nvtx.hpp" +#include "../../../util/kvikio_io.hpp" #include "../../ivf_pq/ivf_pq_fp16_overflow.cuh" #include "graph_core.cuh" #include @@ -20,6 +21,7 @@ #include #include #include +#include #include #include @@ -31,6 +33,8 @@ #include #include +#include + // TODO: This shouldn't be calling spatial/knn APIs #include "../ann_utils.cuh" @@ -38,12 +42,14 @@ #include #include +#include #include #include #include #include #include #include +#include #include #include @@ -434,35 +440,82 @@ 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"); + + 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 @@ -688,127 +741,111 @@ 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::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::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; + 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); -#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(); - } + 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); + }; + + 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); + 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); + 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); + auto reordered_handle = + cuvs::util::detail::open_kvikio_file_for_ace_io(reordered_path, "r", sub_dataset_ptr); + expect_complete_read( + "core", + core_bytes, + reordered_handle.pread(sub_dataset_ptr, core_bytes, core_file_offset).get()); + } else if (augmented_bytes > 0) { + RAFT_LOG_DEBUG( + "ACE: Reading %lu augmented vectors from offset %lu", augmented_size, augmented_file_offset); + auto augmented_handle = + cuvs::util::detail::open_kvikio_file_for_ace_io(augmented_path, "r", augmented_dest); + expect_complete_read( + "augmented", + augmented_bytes, + augmented_handle.pread(augmented_dest, augmented_bytes, augmented_file_offset).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 @@ -847,13 +884,35 @@ 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() && + 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, 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)); + } + // 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)); @@ -864,7 +923,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 * @@ -1385,20 +1444,91 @@ auto build_ace(raft::resources const& res, const index_params& params, DatasetVi core_sub_dataset_size, augmented_sub_dataset_size); - auto sub_dataset = raft::make_host_matrix(sub_dataset_size, dataset_dim); + // Create index for this partition + auto sub_index_params = cuvs::neighbors::cagra::index_params::from_dataset( + raft::make_extents(sub_dataset_size, dataset_dim), + graph_degree, + params.metric, + ef_construction / 16); + 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 + auto read_end = start; + bool used_device_read = false; + const size_t sub_ds_bytes = sub_dataset_size * dataset_dim * sizeof(T); + using device_sub_index_t = cuvs::neighbors::cagra::device_padded_index; + using host_sub_index_t = cuvs::neighbors::cagra::host_standard_index; + using sub_index_t = std::variant; + auto sub_index = [&]() -> sub_index_t { + if (use_disk_mode) { + const size_t current_free_gpu = rmm::available_device_memory().first; + const size_t configured_gpu = + ace_params.max_gpu_memory_gb > 0 + ? static_cast(ace_params.max_gpu_memory_gb * (1ULL << 30)) + : current_free_gpu; + const size_t free_gpu_bytes = std::min(current_free_gpu, configured_gpu); + if (sub_ds_bytes < static_cast(0.4 * free_gpu_bytes)) { + try { + auto sub_dataset_tight = + raft::make_device_matrix(res, sub_dataset_size, dataset_dim); + raft::resource::sync_stream(res); + 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_tight.data_handle()); + read_end = std::chrono::high_resolution_clock::now(); + auto sub_dataset_view = raft::make_const_mdspan(sub_dataset_tight.view()); + std::unique_ptr> + sub_dataset_padded; + auto sub_dataset_dev = [&]() { + if (cuvs::neighbors::matrix_row_width_matches_cagra_required(sub_dataset_view)) { + return cuvs::neighbors::make_device_padded_dataset_view(res, sub_dataset_view); + } + sub_dataset_padded = + cuvs::neighbors::make_device_padded_dataset(res, sub_dataset_view); + return sub_dataset_padded->as_dataset_view(); + }(); + auto direct_index = + ::cuvs::neighbors::cagra::detail::build_from_device_matrix( + res, sub_index_params, sub_dataset_dev); + used_device_read = true; + return sub_index_t{std::in_place_type, std::move(direct_index)}; + } catch (const std::bad_alloc& e) { + RAFT_LOG_WARN( + "ACE: partition %lu did not fit in device memory for a direct (GDS) read: %s; " + "falling back to a host read", + partition_id, + e.what()); + } + } + + 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(); + auto sub_dataset_view = cuvs::neighbors::make_host_standard_dataset_view( + raft::make_const_mdspan(sub_dataset.view())); + auto host_index = + ::cuvs::neighbors::cagra::build(res, sub_index_params, sub_dataset_view); + static_assert(std::is_same_v); + return sub_index_t{std::in_place_type, std::move(host_index)}; + } + + 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, @@ -1409,54 +1539,60 @@ auto build_ace(raft::resources const& res, const index_params& params, DatasetVi core_partition_offsets.view(), augmented_partition_offsets.view(), sub_dataset.view()); + read_end = std::chrono::high_resolution_clock::now(); + auto sub_dataset_view = cuvs::neighbors::make_host_standard_dataset_view( + raft::make_const_mdspan(sub_dataset.view())); + auto host_index = ::cuvs::neighbors::cagra::build(res, sub_index_params, sub_dataset_view); + static_assert(std::is_same_v); + return sub_index_t{std::in_place_type, std::move(host_index)}; + }(); + auto sub_graph = std::visit([](auto const& index) { return index.graph(); }, sub_index); + 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 - auto sub_index_params = cuvs::neighbors::cagra::index_params::from_dataset( - raft::make_extents(sub_dataset_size, dataset_dim), - graph_degree, - params.metric, - ef_construction / 16); - sub_index_params.attach_dataset_on_build = false; - sub_index_params.guarantee_connectivity = params.guarantee_connectivity; - - // Keep the partition host-resident so IVF-PQ and NN-descent can consume it in batches. - // Iterative CAGRA uploads and pads the partition inside build_from_host_matrix. - auto sub_dataset_view = cuvs::neighbors::make_host_standard_dataset_view( - raft::make_const_mdspan(sub_dataset.view())); - auto sub_index = ::cuvs::neighbors::cagra::build(res, sub_index_params, sub_dataset_view); - static_assert( - std::is_same_v>); - auto optimize_end = std::chrono::high_resolution_clock::now(); 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, + 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_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 { + 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_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, @@ -1468,33 +1604,20 @@ auto build_ace(raft::resources const& res, const index_params& params, DatasetVi 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)", diff --git a/cpp/src/neighbors/detail/cagra/cagra_merge.cuh b/cpp/src/neighbors/detail/cagra/cagra_merge.cuh index 6b701dde8a..235bda11e1 100644 --- a/cpp/src/neighbors/detail/cagra/cagra_merge.cuh +++ b/cpp/src/neighbors/detail/cagra/cagra_merge.cuh @@ -6,7 +6,6 @@ #include -#include "cagra_build.cuh" #include "cagra_merge_scaffold.cuh" #include "graph_core.cuh" @@ -150,6 +149,15 @@ cuvs::neighbors::cagra::index merge_rebuild( } }; + auto build_merged_index = [&] { + auto build_params = params; + build_params.attach_dataset_on_build = false; + auto index = ::cuvs::neighbors::cagra::build(handle, build_params, merged_dataset); + index.update_device_dataset_same_layout(handle, merged_dataset); + RAFT_LOG_DEBUG("cagra merge: using device memory for merged dataset"); + return index; + }; + cudaStream_t stream = raft::resource::get_cuda_stream(handle); if (bitset_filtered) { @@ -183,11 +191,7 @@ cuvs::neighbors::cagra::index merge_rebuild( raft::matrix::copy_rows( handle, raft::make_const_mdspan(staging.view()), output_view, indices_view); - auto index = ::cuvs::neighbors::cagra::detail::build_from_device_matrix( - handle, params, merged_dataset); - index.update_device_dataset_same_layout(handle, merged_dataset); - RAFT_LOG_DEBUG("cagra merge: using device memory for merged dataset"); - return index; + return build_merged_index(); } RAFT_CUDA_TRY(cudaMemsetAsync( @@ -196,11 +200,7 @@ cuvs::neighbors::cagra::index merge_rebuild( static_cast(final_rows) * static_cast(stride) * sizeof(T), stream)); merge_dataset(output_view.data_handle(), static_cast(stride)); - auto index = ::cuvs::neighbors::cagra::detail::build_from_device_matrix( - handle, params, merged_dataset); - index.update_device_dataset_same_layout(handle, merged_dataset); - RAFT_LOG_DEBUG("cagra merge: using device memory for merged dataset"); - return index; + return build_merged_index(); } struct fastener_preflight_result { diff --git a/cpp/src/neighbors/detail/cagra/cagra_serialize.cuh b/cpp/src/neighbors/detail/cagra/cagra_serialize.cuh index f2e0c4f07b..be74555d33 100644 --- a/cpp/src/neighbors/detail/cagra/cagra_serialize.cuh +++ b/cpp/src/neighbors/detail/cagra/cagra_serialize.cuh @@ -6,6 +6,7 @@ #pragma once #include +#include #include #include #include @@ -17,10 +18,14 @@ #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 @@ -72,6 +77,12 @@ constexpr bool is_valid_serialized_dataset_kind(std::uint32_t raw) return raw <= static_cast(kind::host_standard); } +template +void serialize_index_mdspan(raft::resources const& res, std::ostream& os, MdspanT const& mdspan) +{ + cuvs::util::detail::serialize_mdspan(res, os, mdspan); +} + /** * Save the index to file. * @@ -113,7 +124,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()); bool has_source_indices = index_.source_indices().has_value(); uint32_t content_map = 0x1u * include_dataset + 0x2u * has_source_indices; @@ -134,7 +145,7 @@ 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 @@ -147,7 +158,7 @@ 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); @@ -157,21 +168,8 @@ void serialize(raft::resources const& res, } template -void serialize_to_hnswlib( - raft::resources const& res, - std::ostream& os, - CagraIndexT const& index_, - std::optional> dataset) +void write_hnswlib_header(std::ostream& os, CagraIndexT const& index_, int dim) { - static_assert(is_cagra_hnsw_serialize_index_v, - "serialize_to_hnswlib requires a dense device or host padded CAGRA index"); - - 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)); @@ -214,110 +212,273 @@ 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, + std::chrono::system_clock::time_point const& start_clock, + size_t& next_report_offset) +{ + if (completed_rows < next_report_offset || completed_rows == 0) { return; } + + auto const 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); + double const throughput = bytes_written / gib / elapsed; + double const rows_throughput = completed_rows / elapsed; + double const 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); + + next_report_offset += std::max(1, total_rows / 10); +} + +template +void write_hnswlib_rows_host( + raft::resources const& res, + std::ostream& os, + CagraIndexT const& index_, + std::optional> dataset) +{ + static_assert(is_cagra_hnsw_serialize_index_v, + "serialize_to_hnswlib requires a dense device or host padded CAGRA index"); - // Remove padding before saving the dataset - raft::host_matrix host_dataset = raft::make_host_matrix(0, 0); - raft::host_matrix_view host_dataset_view; + auto const n_rows = static_cast(index_.size()); + auto const dim = static_cast(dataset ? dataset->extent(1) : index_.dim()); + auto const graph_degree = static_cast(index_.graph_degree()); + auto const row_size = + sizeof(int) + graph_degree * sizeof(IdxT) + dim * sizeof(T) + sizeof(size_t); + if (n_rows == 0) { return; } + + auto const 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); + + auto const 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); + + T const* dataset_data = nullptr; + int64_t dataset_stride = 0; + bool dataset_is_device = false; if (dataset) { - host_dataset_view = *dataset; - } else if constexpr (is_device_cagra_hnsw_serialize_index_v) { - auto dataset_view = index_.dataset(); + dataset_data = dataset->data_handle(); + dataset_stride = dataset->stride(0); + 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 const dataset_view = index_.dataset(); RAFT_EXPECTS(dataset_view.n_rows() > 0, "Invalid CAGRA dataset of size 0 during serialization, shape %zux%zu", static_cast(dataset_view.n_rows()), static_cast(dataset_view.dim())); - host_dataset = raft::make_host_matrix(dataset_view.n_rows(), dataset_view.dim()); - raft::copy_matrix(host_dataset.data_handle(), - host_dataset.extent(1), - dataset_view.view().data_handle(), - dataset_view.stride(), - host_dataset.extent(1), - dataset_view.n_rows(), - raft::resource::get_cuda_stream(res)); - raft::resource::sync_stream(res); - host_dataset_view = raft::make_const_mdspan(host_dataset.view()); - } else if constexpr (is_host_cagra_hnsw_serialize_index_v) { - auto dataset_view = index_.dataset(); - RAFT_EXPECTS(dataset_view.n_rows() > 0, - "Invalid CAGRA dataset of size 0 during serialization, shape %zux%zu", + RAFT_EXPECTS(static_cast(dataset_view.n_rows()) == n_rows, + "CAGRA dataset rows (%zu) do not match index size (%zu)", static_cast(dataset_view.n_rows()), - static_cast(dataset_view.dim())); - auto const n_rows = static_cast(dataset_view.n_rows()); - auto const logical_dim = static_cast(dataset_view.dim()); - auto const stride = static_cast(dataset_view.stride()); - auto const* src = dataset_view.view().data_handle(); - if (stride == logical_dim) { - host_dataset_view = - raft::make_host_matrix_view(src, n_rows, logical_dim); - } else { - // Padded host layout: compact the rows so the writer below sees contiguous vectors. - host_dataset = raft::make_host_matrix(n_rows, logical_dim); - for (int64_t i = 0; i < n_rows; i++) { - std::copy_n(src + i * stride, logical_dim, &host_dataset(i, 0)); - } - host_dataset_view = raft::make_const_mdspan(host_dataset.view()); - } - } else { - static_assert(is_cagra_hnsw_serialize_index_v, - "serialize_to_hnswlib: unsupported CagraIndexT"); + n_rows); + dataset_data = dataset_view.view().data_handle(); + dataset_stride = dataset_view.stride(); + dataset_is_device = is_device_cagra_hnsw_serialize_index_v; } - 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); + + auto const stream = raft::resource::get_cuda_stream(res); + auto const graph_degree_int = static_cast(graph_degree); + size_t bytes_written = 0; + size_t next_report = std::max(1, n_rows / 10); + auto const start_clock = std::chrono::system_clock::now(); + + for (size_t first_row = 0; first_row < n_rows; first_row += batch_rows) { + auto const rows = std::min(batch_rows, n_rows - first_row); + raft::copy_matrix(graph_buffer.data_handle(), + graph_degree, + graph.data_handle() + first_row * graph.stride(0), + graph.stride(0), + graph_degree, + rows, + stream); + if (dataset_is_device) { + raft::copy_matrix(dataset_buffer.data_handle(), + dim, + dataset_data + 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) { + auto const 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)); + T const* data_row = + dataset_is_device ? &dataset_buffer(batch_row, 0) : dataset_data + row * dataset_stride; + 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) +{ + auto const* 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, + GraphT const* graph, + T const* dataset, + size_t first_row, + size_t rows, + uint32_t graph_degree, + uint32_t dim, + int64_t dataset_stride) +{ + auto const warps_per_block = blockDim.x / warpSize; + auto const warp = threadIdx.x / warpSize; + auto const lane = threadIdx.x % warpSize; + auto const batch_row = blockIdx.x * warps_per_block + warp; + if (batch_row >= rows) { return; } + + auto const 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)); + auto const label_offset = sizeof(int) + graph_degree * sizeof(IdxT) + dim * sizeof(T); + write_unaligned_value(row_output + label_offset, row); + } + + auto const 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])); + } + + auto const 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, + CagraIndexT const& index_) +{ + static_assert(is_device_cagra_hnsw_serialize_index_v, + "write_hnswlib_rows_device requires a device-backed CAGRA index"); + + auto const n_rows = static_cast(index_.size()); + auto const dim = index_.dim(); + auto const graph_degree = index_.graph_degree(); + auto const row_size = + sizeof(int) + graph_degree * sizeof(IdxT) + dim * sizeof(T) + sizeof(size_t); + if (n_rows == 0) { return; } + + auto const graph = index_.graph(); + auto const dataset = index_.dataset(); + RAFT_EXPECTS(dataset.n_rows() > 0, + "Invalid CAGRA dataset of size 0 during serialization, shape %zux%zu", + static_cast(dataset.n_rows()), + static_cast(dataset.dim())); + RAFT_EXPECTS(static_cast(graph.extent(0)) == n_rows && + static_cast(dataset.n_rows()) == n_rows, + "CAGRA graph and dataset rows must match index size"); + + auto const 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); + auto const stream = raft::resource::get_cuda_stream(res); + + size_t bytes_written = 0; + size_t next_report = std::max(1, n_rows / 10); + auto const 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) { + auto const rows = std::min(batch_rows, n_rows - first_row); + auto const 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.view().data_handle(), + first_row, + rows, + graph_degree, + dim, + dataset.stride()); + RAFT_CUDA_TRY(cudaPeekAtLastError()); + raft::resource::sync_stream(res); + + auto const 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; + std::array const zeros{}; + for (size_t first_row = 0; first_row < n_rows; first_row += chunk_size) { + auto const 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, + CagraIndexT const& index_, + std::optional> dataset) +{ + static_assert(is_cagra_hnsw_serialize_index_v, + "serialize_to_hnswlib requires a dense device or host padded CAGRA index"); + + int const 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, @@ -325,10 +486,29 @@ void serialize_to_hnswlib( CagraIndexT const& index_, std::optional> dataset) { - std::ofstream of(filename, std::ios::out | std::ios::binary); + static_assert(is_cagra_hnsw_serialize_index_v, + "serialize_to_hnswlib requires a dense device or host padded CAGRA index"); + + int const 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 constexpr (is_device_cagra_hnsw_serialize_index_v) { + if (dataset) { + write_hnswlib_rows_host(res, of, index_, dataset); + } else { + write_hnswlib_rows_device(res, of, index_); + } + } else { + write_hnswlib_rows_host(res, of, index_, dataset); + } + write_hnswlib_empty_levels(of, static_cast(index_.size())); of.close(); if (!of) { RAFT_FAIL("Error writing output %s", filename.c_str()); } @@ -343,14 +523,15 @@ void serialize_to_hnswlib( * @param[in] index_ CAGRA index * */ -template -void deserialize( +template +void deserialize_impl( raft::resources const& res, - std::istream& is, + Input& input, cuvs::neighbors::cagra::index* index_, std::unique_ptr>* out_dataset = nullptr) { 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"); @@ -379,72 +560,107 @@ void deserialize( "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()); - - auto content_map = raft::deserialize_scalar(res, is); - bool has_dataset = content_map & 0x1u; - using kind = cuvs::neighbors::cagra::serialized_dataset_kind; - RAFT_EXPECTS(has_dataset == (dataset_kind != kind::none), - "cagra::deserialize: dataset kind and content map disagree"); - - using owner_t = cuvs::neighbors::owning_dataset_for_view_t; - std::unique_ptr dataset_owner{}; - if (has_dataset) { - if (out_dataset == nullptr) { - cuvs::neighbors::detail::skip_dense_dataset(res, is); + using index_t = cuvs::neighbors::cagra::index; + using graph_type_t = typename index_t::graph_index_type; + RAFT_EXPECTS(cuvs::util::is_mul_no_overflow(static_cast(n_rows), + static_cast(graph_degree), + sizeof(graph_type_t)), + "cagra::deserialize: integer overflow in graph allocation " + "(n_rows=%lld, graph_degree=%u, element_size=%zu)", + static_cast(n_rows), + graph_degree, + sizeof(graph_type_t)); + + auto finish_deserialize = [&](auto graph) { + auto content_map = raft::deserialize_scalar(res, is); + bool has_dataset = content_map & 0x1u; + using kind = cuvs::neighbors::cagra::serialized_dataset_kind; + RAFT_EXPECTS(has_dataset == (dataset_kind != kind::none), + "cagra::deserialize: dataset kind and content map disagree"); + + using owner_t = cuvs::neighbors::owning_dataset_for_view_t; + std::unique_ptr dataset_owner{}; + if (has_dataset) { + if (out_dataset == nullptr) { + cuvs::neighbors::detail::skip_dense_dataset(res, is); + } else { + auto const expected_kind = serialized_dataset_kind_for_view(); + RAFT_EXPECTS( + dataset_kind == expected_kind, + "cagra::deserialize: serialized dataset kind %u does not match requested kind %u", + dataset_kind_raw, + static_cast(expected_kind)); + if constexpr (cuvs::neighbors::is_device_padded_dataset_view_v) { + dataset_owner = + cuvs::neighbors::detail::deserialize_padded_dataset(res, input); + } else if constexpr (cuvs::neighbors::is_device_standard_dataset_view_v) { + dataset_owner = + cuvs::neighbors::detail::deserialize_standard_dataset(res, input); + } else if constexpr (cuvs::neighbors::is_host_padded_dataset_view_v) { + dataset_owner = + cuvs::neighbors::detail::deserialize_host_padded_dataset(res, input); + } else if constexpr (cuvs::neighbors::is_host_standard_dataset_view_v) { + dataset_owner = + cuvs::neighbors::detail::deserialize_host_standard_dataset(res, input); + } else { + static_assert(sizeof(DatasetViewT) == 0, + "deserialize: dataset deserialization is not implemented for this view"); + } + } + } + + if (dataset_owner) { + *index_ = index_t( + res, metric, dataset_owner->as_dataset_view(), raft::make_const_mdspan(graph.view())); } else { - auto const expected_kind = serialized_dataset_kind_for_view(); - RAFT_EXPECTS( - dataset_kind == expected_kind, - "cagra::deserialize: serialized dataset kind %u does not match requested kind %u", - dataset_kind_raw, - static_cast(expected_kind)); - if constexpr (cuvs::neighbors::is_device_padded_dataset_view_v) { - dataset_owner = cuvs::neighbors::detail::deserialize_padded_dataset(res, is); - } else if constexpr (cuvs::neighbors::is_device_standard_dataset_view_v) { - dataset_owner = cuvs::neighbors::detail::deserialize_standard_dataset(res, is); - } else if constexpr (cuvs::neighbors::is_host_padded_dataset_view_v) { - dataset_owner = - cuvs::neighbors::detail::deserialize_host_padded_dataset(res, is); - } else if constexpr (cuvs::neighbors::is_host_standard_dataset_view_v) { - dataset_owner = - cuvs::neighbors::detail::deserialize_host_standard_dataset(res, is); + *index_ = index_t(res, metric); + if constexpr (raft::is_device_mdspan_v) { + index_->update_graph(res, std::move(graph)); } else { - static_assert(sizeof(DatasetViewT) == 0, - "deserialize: dataset deserialization is not implemented for this view"); + index_->update_graph(res, raft::make_const_mdspan(graph.view())); } } - } - if (dataset_owner) { - *index_ = cuvs::neighbors::cagra::index( - res, metric, dataset_owner->as_dataset_view(), raft::make_const_mdspan(graph.view())); + if constexpr (raft::is_device_mdspan_v) { + if (dataset_owner) { index_->update_graph(res, std::move(graph)); } + if (content_map & 0x2u) { + auto source_indices = raft::make_device_vector(res, n_rows); + cuvs::util::detail::deserialize_mdspan(res, input, source_indices.view()); + index_->update_source_indices(std::move(source_indices)); + } + } else { + std::optional> source_indices; + if (content_map & 0x2u) { + source_indices.emplace(raft::make_host_vector(n_rows)); + raft::deserialize_mdspan(res, is, source_indices->view()); + index_->update_source_indices(res, raft::make_const_mdspan(source_indices->view())); + } + // Graph and source-index updates can enqueue copies from host staging. Keep both staging + // buffers alive through this synchronization. + raft::resource::sync_stream(res); + } + if (dataset_owner) { *out_dataset = std::move(dataset_owner); } + }; + + if constexpr (std::is_same_v, cuvs::util::kvikio_file_reader>) { + auto graph = raft::make_device_matrix(res, n_rows, graph_degree); + cuvs::util::detail::deserialize_mdspan(res, input, graph.view()); + finish_deserialize(std::move(graph)); } else { - *index_ = cuvs::neighbors::cagra::index(res, metric); - index_->update_graph(res, raft::make_const_mdspan(graph.view())); + auto graph = raft::make_host_matrix(n_rows, graph_degree); + raft::deserialize_mdspan(res, is, graph.view()); + finish_deserialize(std::move(graph)); } +} - bool has_source_indices = content_map & 0x2u; - std::optional> source_indices; - if (has_source_indices) { - source_indices.emplace(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())); - } - // Graph and source-index updates can enqueue copies from host staging. Keep both staging buffers - // alive through this single synchronization. - raft::resource::sync_stream(res); - if (dataset_owner) { *out_dataset = std::move(dataset_owner); } +template +void deserialize( + raft::resources const& res, + std::istream& is, + cuvs::neighbors::cagra::index* index_, + std::unique_ptr>* out_dataset = nullptr) +{ + deserialize_impl(res, is, index_, out_dataset); } template @@ -454,12 +670,7 @@ void deserialize( cuvs::neighbors::cagra::index* index_, std::unique_ptr>* out_dataset = nullptr) { - 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_, out_dataset); - - is.close(); + cuvs::util::kvikio_file_reader reader(filename); + deserialize_impl(res, reader, index_, out_dataset); } } // namespace cuvs::neighbors::cagra::detail diff --git a/cpp/src/neighbors/detail/dataset_serialize.hpp b/cpp/src/neighbors/detail/dataset_serialize.hpp index 05c71e0213..12889d4f4e 100644 --- a/cpp/src/neighbors/detail/dataset_serialize.hpp +++ b/cpp/src/neighbors/detail/dataset_serialize.hpp @@ -5,6 +5,8 @@ #pragma once #include +#include +#include #include #include #include @@ -24,6 +26,8 @@ #include #include +#include "../../util/kvikio_serialize.hpp" + namespace cuvs::neighbors::detail { using dataset_instance_tag = uint32_t; @@ -97,6 +101,36 @@ void serialize(const raft::resources& res, std::ostream& os, ViewT const& datase if (elements == 0) { return; } if constexpr (cuvs::neighbors::is_device_dataset_view_v) { + if (auto* kvikio_stream = dynamic_cast(&os); + kvikio_stream != nullptr) { + auto const row_bytes = static_cast(dim) * sizeof(DataT); + if (stride == dim) { + raft::resource::sync_stream(res); + kvikio_stream->write_device(src.data_handle(), elements * sizeof(DataT)); + return; + } + + auto const batch_rows = static_cast(std::min( + static_cast(n_rows), + std::max(1, cuvs::util::detail::kDeviceSerializationBatchBytes / row_bytes))); + auto packed = raft::make_device_matrix(res, batch_rows, dim); + auto const stream = raft::resource::get_cuda_stream(res); + for (IdxT first_row = 0; first_row < n_rows; first_row += batch_rows) { + auto const 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); + kvikio_stream->write_device(packed.data_handle(), + static_cast(rows) * row_bytes); + } + return; + } + auto staging = raft::make_host_matrix(n_rows, dim); raft::copy_matrix(staging.data_handle(), dim, @@ -261,6 +295,53 @@ auto deserialize_device_dense(raft::resources const& res, std::istream& is) return std::make_unique(std::move(storage), metadata.dim); } +template +auto deserialize_device_dense(raft::resources const& res, cuvs::util::kvikio_file_reader& reader) + -> std::unique_ptr +{ + auto const metadata = deserialize_dense_payload_metadata(res, reader.stream()); + auto storage = raft::make_device_matrix(res, metadata.n_rows, metadata.stride); + auto const stream = raft::resource::get_cuda_stream(res); + + if (metadata.stride > metadata.dim && metadata.n_rows > 0) { + RAFT_CUDA_TRY(cudaMemset2DAsync(storage.data_handle() + metadata.dim, + metadata.stride * sizeof(DataT), + 0, + (metadata.stride - metadata.dim) * sizeof(DataT), + metadata.n_rows, + stream)); + } + raft::resource::sync_stream(res); + + if (metadata.elements == 0) { + return std::make_unique(std::move(storage), metadata.dim); + } + + auto const row_bytes = static_cast(metadata.dim) * sizeof(DataT); + if (metadata.stride == metadata.dim) { + reader.read_device(storage.data_handle(), metadata.elements * sizeof(DataT)); + } else { + auto const batch_rows = static_cast(std::min( + static_cast(metadata.n_rows), + std::max(1, cuvs::util::detail::kDeviceSerializationBatchBytes / row_bytes))); + auto packed = raft::make_device_matrix(res, batch_rows, metadata.dim); + raft::resource::sync_stream(res); + for (IdxT first_row = 0; first_row < metadata.n_rows; first_row += batch_rows) { + auto const rows = std::min(batch_rows, metadata.n_rows - first_row); + reader.read_device(packed.data_handle(), static_cast(rows) * row_bytes); + raft::copy_matrix(storage.data_handle() + first_row * metadata.stride, + metadata.stride, + packed.data_handle(), + metadata.dim, + metadata.dim, + rows, + stream); + raft::resource::sync_stream(res); + } + } + return std::make_unique(std::move(storage), metadata.dim); +} + template auto deserialize_host_dense(raft::resources const& res, std::istream& is) -> std::unique_ptr @@ -305,10 +386,11 @@ auto deserialize_vpq(raft::resources const& res, std::istream& is) std::move(vq_code_book), std::move(pq_code_book), std::move(data)); } -template -auto deserialize_dense_dataset(raft::resources const& res, std::istream& is) +template +auto deserialize_dense_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); RAFT_EXPECTS(tag == kSerializeStridedDataset, "deserialize_dataset: expected strided tag, got %u", @@ -323,9 +405,9 @@ auto deserialize_dense_dataset(raft::resources const& res, std::istream& is) static_cast(dtype), static_cast(expected_dtype)); if constexpr (std::is_same_v>) { - return deserialize_device_dense(res, is); + return deserialize_device_dense(res, input); } else if constexpr (std::is_same_v>) { - return deserialize_device_dense(res, is); + return deserialize_device_dense(res, input); } else if constexpr (std::is_same_v>) { return deserialize_host_dense(res, is); } else if constexpr (std::is_same_v>) { @@ -359,32 +441,32 @@ void skip_dense_dataset(raft::resources const& res, std::istream& is) // dense owner. When a new dataset kind is supported, add a matching overload of // deserialize_dataset here rather than extending this one — overload dispatch replaces the old // type-erased variant routing. -template -auto deserialize_padded_dataset(raft::resources const& res, std::istream& is) +template +auto deserialize_padded_dataset(raft::resources const& res, Input& input) -> std::unique_ptr> { - return deserialize_dense_dataset>(res, is); + return deserialize_dense_dataset>(res, input); } -template -auto deserialize_standard_dataset(raft::resources const& res, std::istream& is) +template +auto deserialize_standard_dataset(raft::resources const& res, Input& input) -> std::unique_ptr> { - return deserialize_dense_dataset>(res, is); + return deserialize_dense_dataset>(res, input); } -template -auto deserialize_host_padded_dataset(raft::resources const& res, std::istream& is) +template +auto deserialize_host_padded_dataset(raft::resources const& res, Input& input) -> std::unique_ptr> { - return deserialize_dense_dataset>(res, is); + return deserialize_dense_dataset>(res, input); } -template -auto deserialize_host_standard_dataset(raft::resources const& res, std::istream& is) +template +auto deserialize_host_standard_dataset(raft::resources const& res, Input& input) -> std::unique_ptr> { - return deserialize_dense_dataset>(res, is); + return deserialize_dense_dataset>(res, input); } } // namespace cuvs::neighbors::detail diff --git a/cpp/src/neighbors/detail/hnsw.hpp b/cpp/src/neighbors/detail/hnsw.hpp index 8328fac5d5..be6753ec8b 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 @@ -27,6 +29,7 @@ #include #include +#include #include #include #include @@ -745,9 +748,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; @@ -836,44 +842,59 @@ 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); + // 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); + + // 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; } -#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); - } - } + }; + 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, + 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( @@ -1253,8 +1274,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 @@ -1299,7 +1331,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()); @@ -1399,7 +1434,7 @@ 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); + cuvs::util::kvikio_ofstream of(index_filename); RAFT_EXPECTS(of, "Cannot open file %s", index_filename.c_str()); serialize_to_hnswlib_from_inmem(res, of, params, cagra_index, dataset); diff --git a/cpp/src/neighbors/detail/vamana/vamana_serialize.cuh b/cpp/src/neighbors/detail/vamana/vamana_serialize.cuh index a81b2252b5..fcc3eec9fd 100644 --- a/cpp/src/neighbors/detail/vamana/vamana_serialize.cuh +++ b/cpp/src/neighbors/detail/vamana/vamana_serialize.cuh @@ -8,6 +8,7 @@ #include "vamana_structs.cuh" #include +#include #include #include #include @@ -18,28 +19,61 @@ #include #include +#include "../../../util/kvikio_serialize.hpp" #include "../dataset_serialize.hpp" +#include #include #include -#include +#include #include + namespace cuvs::neighbors::vamana::detail { -// write matrix containing dataset to file -template -void to_file(const std::string& dataset_base_file, raft::host_matrix& dataset) +template +void serialize_dataset_view(raft::resources const& res, + DatasetView dataset, + const std::string& dataset_base_file) { - std::ofstream dataset_of(dataset_base_file, std::ios::out | std::ios::binary); + auto const n_rows = dataset.extent(0); + auto const dim = dataset.extent(1); + RAFT_EXPECTS(n_rows <= std::numeric_limits::max() && dim <= std::numeric_limits::max(), + "Vamana dataset shape (%lld, %lld) exceeds the on-disk format limits", + static_cast(n_rows), + static_cast(dim)); + + cuvs::util::kvikio_ofstream dataset_of(dataset_base_file); if (!dataset_of) { RAFT_FAIL("Cannot open file %s", dataset_base_file.c_str()); } - size_t dataset_file_offset = 0; - int size = static_cast(dataset.extent(0)); - int dim = static_cast(dataset.extent(1)); - dataset_of.seekp(dataset_file_offset, dataset_of.beg); + int const size = static_cast(n_rows); + int const cols = static_cast(dim); dataset_of.write((char*)&size, sizeof(int)); - dataset_of.write((char*)&dim, sizeof(int)); - for (int i = 0; i < size; i++) { - dataset_of.write((char*)(dataset.data_handle() + i * dataset.extent(1)), dim * sizeof(T)); + dataset_of.write((char*)&cols, sizeof(int)); + + if (n_rows > 0 && dim > 0) { + auto const stride = dataset.stride(0); + auto const row_bytes = static_cast(dim) * sizeof(T); + if (stride == dim) { + raft::resource::sync_stream(res); + dataset_of.write_device(dataset.data_handle(), static_cast(n_rows) * row_bytes); + } else { + auto const batch_rows = static_cast(std::min( + static_cast(n_rows), + std::max(1, cuvs::util::detail::kDeviceSerializationBatchBytes / row_bytes))); + auto packed = raft::make_device_matrix(res, batch_rows, dim); + auto const stream = raft::resource::get_cuda_stream(res); + for (int64_t first_row = 0; first_row < n_rows; first_row += batch_rows) { + auto const rows = std::min(batch_rows, n_rows - first_row); + raft::copy_matrix(packed.data_handle(), + dim, + dataset.data_handle() + first_row * stride, + stride, + dim, + rows, + stream); + raft::resource::sync_stream(res); + dataset_of.write_device(packed.data_handle(), static_cast(rows) * row_bytes); + } + } } dataset_of.close(); if (!dataset_of) { RAFT_FAIL("Error writing output %s", dataset_base_file.c_str()); } @@ -61,22 +95,8 @@ void serialize_dataset(raft::resources const& res, const std::string& dataset_base_file) { if (dataset == nullptr) { return; } - // try allocating a buffer for the dataset on host try { - auto nrows = dataset->n_rows(); - auto dim = dataset->dim(); - auto stride = dataset->stride(); - auto d_data = dataset->view(); - auto h_dataset = raft::make_host_matrix(nrows, dim); - raft::copy_matrix(h_dataset.data_handle(), - dim, - d_data.data_handle(), - stride, - dim, - nrows, - raft::resource::get_cuda_stream(res)); - raft::resource::sync_stream(res); - to_file(dataset_base_file, h_dataset); + serialize_dataset_view(res, dataset->view(), dataset_base_file); } catch (std::bad_alloc& e) { RAFT_LOG_INFO("Failed to serialize dataset"); } catch (raft::logic_error& e) { @@ -88,12 +108,8 @@ void serialize_dataset(raft::resources const& res, raft::device_matrix_view dataset, const std::string& dataset_base_file) { - // try allocating a buffer for the dataset on host try { - auto h_dataset = raft::make_host_matrix(dataset.extent(0), dataset.extent(1)); - raft::copy(res, h_dataset.view(), dataset); - raft::resource::sync_stream(res); - to_file(dataset_base_file, h_dataset); + serialize_dataset_view(res, dataset, dataset_base_file); } catch (std::bad_alloc& e) { RAFT_LOG_INFO("Failed to serialize dataset"); } catch (raft::logic_error& e) { @@ -119,7 +135,7 @@ void serialize_sector_aligned( const HostMatT& h_graph, const cuvs::neighbors::device_padded_dataset_view& dataset, const uint64_t medoid, - std::ofstream& output_writer) + std::ostream& output_writer) { if constexpr (!std::is_same_v) { RAFT_FAIL("serialization is only implemented for uint32_t graph"); @@ -190,17 +206,18 @@ void serialize_sector_aligned( output_meta.push_back(static_cast(append_reorder_data)); output_meta.push_back(disk_index_file_size); - // zero out first sector of output file - output_writer.seekp(0, output_writer.beg); - output_writer.write(sector_buf.get(), sector_len); - // write metadata to first sector - output_writer.seekp(0, output_writer.beg); + // Build and write the complete metadata sector so the KvikIO stream remains sequential. + memset(sector_buf.get(), 0, sector_len); const int metadata_size = static_cast(output_meta.size()); const int metadata_ndims = 1; - output_writer.write((char*)&metadata_size, sizeof(int)); - output_writer.write((char*)&metadata_ndims, sizeof(int)); - output_writer.write((char*)output_meta.data(), sizeof(uint64_t) * output_meta.size()); - output_writer.seekp(sector_len, output_writer.beg); + size_t metadata_offset = 0; + memcpy(sector_buf.get() + metadata_offset, &metadata_size, sizeof(int)); + metadata_offset += sizeof(int); + memcpy(sector_buf.get() + metadata_offset, &metadata_ndims, sizeof(int)); + metadata_offset += sizeof(int); + memcpy( + sector_buf.get() + metadata_offset, output_meta.data(), sizeof(uint64_t) * output_meta.size()); + output_writer.write(sector_buf.get(), sector_len); if (nnodes_per_sector > 0) { uint64_t cur_node_id = 0; @@ -316,7 +333,7 @@ void serialize(raft::resources const& res, if (sector_aligned) { // Write graph to disk index file with file name suffix according to DiskANN build_disk_index const std::string index_file_name = file_name + "_disk.index"; - std::ofstream index_of(index_file_name, std::ios::out | std::ios::binary); + cuvs::util::kvikio_ofstream index_of(index_file_name); RAFT_EXPECTS(index_of, "Cannot open file %s", index_file_name.c_str()); serialize_sector_aligned( @@ -329,26 +346,14 @@ void serialize(raft::resources const& res, return; } - // Write graph to first index file (format from MSFT DiskANN OSS) - std::ofstream index_of(file_name, std::ios::out | std::ios::binary); - RAFT_EXPECTS(index_of, "Cannot open file %s", file_name.c_str()); - - size_t file_offset = 0; - index_of.seekp(file_offset, index_of.beg); - uint32_t max_degree = 0; - size_t index_size = 24; // Starting metadata - uint32_t start = static_cast(index_.medoid()); - size_t num_frozen_points = 0; - uint32_t max_observed_degree = 0; - - index_of.write((char*)&index_size, sizeof(uint64_t)); - index_of.write((char*)&max_observed_degree, sizeof(uint32_t)); - index_of.write((char*)&start, sizeof(uint32_t)); - index_of.write((char*)&num_frozen_points, sizeof(size_t)); - - size_t total_edges = 0; - size_t num_sparse = 0; - size_t num_single = 0; + // Compute graph metadata before opening the sequential KvikIO stream. + uint32_t max_degree = 0; + size_t index_size = 24; // Starting metadata + uint32_t start = static_cast(index_.medoid()); + size_t num_frozen_points = 0; + size_t total_edges = 0; + size_t num_sparse = 0; + size_t num_single = 0; for (uint32_t i = 0; i < h_graph.extent(0); i++) { uint32_t node_edges = 0; @@ -360,18 +365,28 @@ void serialize(raft::resources const& res, if (node_edges < 2) num_single++; total_edges += node_edges; + max_degree = node_edges > max_degree ? (uint32_t)node_edges : max_degree; + index_size += (size_t)(sizeof(uint32_t) * (node_edges + 1)); + } + + cuvs::util::kvikio_ofstream index_of(file_name); + RAFT_EXPECTS(index_of, "Cannot open file %s", file_name.c_str()); + index_of.write((char*)&index_size, sizeof(uint64_t)); + index_of.write((char*)&max_degree, sizeof(uint32_t)); + index_of.write((char*)&start, sizeof(uint32_t)); + index_of.write((char*)&num_frozen_points, sizeof(size_t)); + + for (uint32_t i = 0; i < h_graph.extent(0); i++) { + uint32_t node_edges = 0; + for (; node_edges < h_graph.extent(1); node_edges++) { + if (h_graph(i, node_edges) == raft::upper_bound()) { break; } + } index_of.write((char*)&node_edges, sizeof(uint32_t)); if constexpr (!std::is_same_v) { RAFT_FAIL("serialization is only implemented for uint32_t graph"); } index_of.write((char*)&h_graph(i, 0), node_edges * sizeof(uint32_t)); - - max_degree = node_edges > max_degree ? (uint32_t)node_edges : max_degree; - index_size += (size_t)(sizeof(uint32_t) * (node_edges + 1)); } - index_of.seekp(file_offset, index_of.beg); - index_of.write((char*)&index_size, sizeof(uint64_t)); - index_of.write((char*)&max_degree, sizeof(uint32_t)); RAFT_LOG_DEBUG( "Wrote file out, index size:%lu, max_degree:%u, num_sparse:%ld, num_single:%ld, total " diff --git a/cpp/src/neighbors/iface/iface.hpp b/cpp/src/neighbors/iface/iface.hpp index 5d5ee76406..172fc81e23 100644 --- a/cpp/src/neighbors/iface/iface.hpp +++ b/cpp/src/neighbors/iface/iface.hpp @@ -278,23 +278,20 @@ void deserialize(const raft::resources& handle, { std::lock_guard lock(*interface.mutex_); - std::ifstream is(filename, std::ios::in | std::ios::binary); - if (!is) { RAFT_FAIL("Cannot open file %s", filename.c_str()); } - if constexpr (std::is_same>::value) { ivf_flat::index idx(handle); - ivf_flat::deserialize(handle, is, &idx); + ivf_flat::deserialize(handle, filename, &idx); resource::sync_stream(handle); interface.index_.emplace(std::move(idx)); } else if constexpr (std::is_same>::value) { ivf_pq::index idx(handle); - ivf_pq::deserialize(handle, is, &idx); + ivf_pq::deserialize(handle, filename, &idx); resource::sync_stream(handle); interface.index_.emplace(std::move(idx)); } else if constexpr (std::is_same>::value) { cagra::device_padded_index idx(handle); std::unique_ptr> out_dataset; - cagra::deserialize(handle, is, &idx, &out_dataset); + cagra::deserialize(handle, filename, &idx, &out_dataset); interface.cagra_owned_padded_dataset_.reset(); interface.cagra_owned_standard_dataset_.reset(); if (out_dataset) { interface.cagra_owned_padded_dataset_ = std::move(out_dataset); } @@ -303,15 +300,13 @@ void deserialize(const raft::resources& handle, } else if constexpr (std::is_same>::value) { cagra::device_standard_index idx(handle); std::unique_ptr> out_dataset; - cagra::deserialize(handle, is, &idx, &out_dataset); + cagra::deserialize(handle, filename, &idx, &out_dataset); interface.cagra_owned_padded_dataset_.reset(); interface.cagra_owned_standard_dataset_.reset(); if (out_dataset) { interface.cagra_owned_standard_dataset_ = std::move(out_dataset); } resource::sync_stream(handle); interface.index_.emplace(std::move(idx)); } - - is.close(); } }; // namespace cuvs::neighbors diff --git a/cpp/src/neighbors/ivf_flat/ivf_flat_serialize.cuh b/cpp/src/neighbors/ivf_flat/ivf_flat_serialize.cuh index 70b789ddc7..7f9388031a 100644 --- a/cpp/src/neighbors/ivf_flat/ivf_flat_serialize.cuh +++ b/cpp/src/neighbors/ivf_flat/ivf_flat_serialize.cuh @@ -5,6 +5,7 @@ #pragma once +#include "../../util/kvikio_serialize.hpp" #include "../../util/serialize_validation.hpp" #include "../ivf_common.cuh" #include "../ivf_list.cuh" @@ -38,8 +39,8 @@ constexpr int serialization_version = 5; * @param[in] index_ IVF-Flat index * */ -template -void serialize(raft::resources const& handle, std::ostream& os, const index& index_) +template +void serialize(raft::resources const& handle, Output& os, const index& index_) { RAFT_LOG_DEBUG( "Saving IVF-Flat index, size %zu, dim %u", static_cast(index_.size()), index_.dim()); @@ -55,11 +56,11 @@ void serialize(raft::resources const& handle, std::ostream& os, const index(index_.list_sizes().extent(0)); raft::copy(handle, sizes_host.view(), index_.list_sizes()); raft::resource::sync_stream(handle); - serialize_mdspan(handle, os, sizes_host.view()); + cuvs::util::detail::serialize_mdspan(handle, os, sizes_host.view()); list_spec list_store_spec{index_.dim(), true}; for (uint32_t label = 0; label < index_.n_lists(); label++) { @@ -81,7 +82,7 @@ void serialize(raft::resources const& handle, const std::string& filename, const index& index_) { - 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(handle, of, index_); @@ -99,9 +100,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)), @@ -134,7 +137,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); @@ -142,15 +145,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); @@ -160,17 +163,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 8ff0465e47..4117f07a53 100644 --- a/cpp/src/neighbors/ivf_list.cuh +++ b/cpp/src/neighbors/ivf_list.cuh @@ -22,6 +22,7 @@ #include #include +#include "../util/kvikio_serialize.hpp" #include "ivf_common.cuh" #include @@ -111,6 +112,23 @@ enable_if_valid_list_t serialize_list(const raft::resources& handle, const typename ListT::spec_type& store_spec, std::optional size_override) { + if (auto* kvikio_stream = dynamic_cast(&os); + kvikio_stream != nullptr) { + using size_type = typename ListT::size_type; + const auto size = size_override.value_or(ld.size.load()); + raft::serialize_scalar(handle, *kvikio_stream, size); + if (size == 0) { return; } + + const auto data_extents = store_spec.make_list_extents(size); + const auto data_view = + raft::make_mdspan( + ld.data.data_handle(), data_extents); + const auto indices_view = raft::make_device_vector_view(ld.indices.data_handle(), size); + cuvs::util::detail::serialize_device_mdspan(handle, *kvikio_stream, data_view); + cuvs::util::detail::serialize_device_mdspan(handle, *kvikio_stream, indices_view); + return; + } + using size_type = typename ListT::size_type; auto size = size_override.value_or(ld.size.load()); raft::serialize_scalar(handle, os, size); @@ -153,6 +171,9 @@ enable_if_valid_list_t deserialize_list(const raft::resources& handle, const typename ListT::spec_type& store_spec, const typename ListT::spec_type& device_spec) { + // Public stream deserializers accept arbitrary streams, which need not be backed by a file that + // KvikIO can reopen. Keep this host-staged path for those callers; filename overloads use the + // KvikIO reader overload below. using size_type = typename ListT::size_type; auto size = raft::deserialize_scalar(handle, is); if (size == 0) { return ld.reset(); } @@ -174,4 +195,29 @@ 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) +{ + // Path-backed deserialization reaches this overload and reads list payloads directly into device + // memory through GDS when available (with KvikIO's compatible I/O fallback otherwise). + 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..a238137813 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" @@ -37,8 +38,8 @@ constexpr int kSerializationVersion = 4; * @param[in] index IVF-PQ index * */ -template -void serialize(raft::resources const& handle_, std::ostream& os, const index& index) +template +void serialize(raft::resources const& handle_, Output& os, const index& index) { RAFT_LOG_DEBUG("Size %zu, dim %d, pq_dim %d, pq_bits %d", static_cast(index.size()), @@ -58,16 +59,16 @@ void serialize(raft::resources const& handle_, std::ostream& os, const index(index.list_sizes().extents()); raft::copy(handle_, sizes_host.view(), index.list_sizes()); raft::resource::sync_stream(handle_); - raft::serialize_mdspan(handle_, os, sizes_host.view()); + cuvs::util::detail::serialize_mdspan(handle_, os, sizes_host.view()); // NOTE: We use static_cast here because serialize_list requires the concrete list type // to access the spec_type for determining the serialized data layout. if (index.codes_layout() == list_layout::FLAT) { @@ -101,7 +102,7 @@ void serialize(raft::resources const& handle_, const std::string& filename, const index& index) { - 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(handle_, of, index); @@ -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/neighbors/ivf_rabitq/gpu_index/initializer_gpu.cu b/cpp/src/neighbors/ivf_rabitq/gpu_index/initializer_gpu.cu index 6755bfc3de..7d05173a69 100644 --- a/cpp/src/neighbors/ivf_rabitq/gpu_index/initializer_gpu.cu +++ b/cpp/src/neighbors/ivf_rabitq/gpu_index/initializer_gpu.cu @@ -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 */ @@ -48,18 +48,12 @@ void FlatInitializerGPU::LoadCentroids(std::ifstream& input, const char* filenam raft::resource::sync_stream(handle_); } -// Saves the centroids to a file. -// The centroids are first copied from device memory to a temporary host buffer, -// and then the entire block of data is written to the output stream. -void FlatInitializerGPU::SaveCentroids(std::ofstream& output, const char* filename) const +// Saves the centroids directly from device memory. +void FlatInitializerGPU::SaveCentroids(cuvs::util::kvikio_ofstream& output, + const char* filename) const { - // Allocate temporary host buffer for centroids. - auto host_buf = raft::make_host_vector(K * D); - // Copy centroids from device to host. - raft::copy(host_buf.data_handle(), centroids_.data_handle(), data_elements(), stream_); raft::resource::sync_stream(handle_); - // Write the raw data to the file. - output.write(reinterpret_cast(host_buf.data_handle()), data_bytes()); + output.write_device(centroids_.data_handle(), data_bytes()); RAFT_EXPECTS(static_cast(output), "Failed to write centroids to file: %s", filename); } diff --git a/cpp/src/neighbors/ivf_rabitq/gpu_index/initializer_gpu.cuh b/cpp/src/neighbors/ivf_rabitq/gpu_index/initializer_gpu.cuh index 0640148814..d2bbbcf345 100644 --- a/cpp/src/neighbors/ivf_rabitq/gpu_index/initializer_gpu.cuh +++ b/cpp/src/neighbors/ivf_rabitq/gpu_index/initializer_gpu.cuh @@ -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 */ @@ -11,6 +11,8 @@ #include "../defines.hpp" +#include + #include #include #include @@ -73,7 +75,7 @@ class InitializerGPU { * @param save * @param filename */ - virtual void SaveCentroids(std::ofstream& output, const char* filename) const = 0; + virtual void SaveCentroids(cuvs::util::kvikio_ofstream& output, const char* filename) const = 0; protected: size_t D; // Dimension @@ -93,7 +95,7 @@ class FlatInitializerGPU : public InitializerGPU { void LoadCentroids(std::ifstream& input, const char* filename) override; - void SaveCentroids(std::ofstream& output, const char* filename) const override; + void SaveCentroids(cuvs::util::kvikio_ofstream& output, const char* filename) const override; private: // D, K are inherited from parent diff --git a/cpp/src/neighbors/ivf_rabitq/gpu_index/ivf_gpu.cu b/cpp/src/neighbors/ivf_rabitq/gpu_index/ivf_gpu.cu index 49070acbcf..230f7ffd03 100644 --- a/cpp/src/neighbors/ivf_rabitq/gpu_index/ivf_gpu.cu +++ b/cpp/src/neighbors/ivf_rabitq/gpu_index/ivf_gpu.cu @@ -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 */ @@ -12,6 +12,7 @@ #include "../utils/reductions.cuh" #include "ivf_gpu.cuh" #include "searcher_gpu.cuh" +#include #include #include @@ -344,8 +345,8 @@ void IVFGPU::save(const char* filename) const { RAFT_EXPECTS(num_centroids > 0, "IVF index has not been constructed"); - std::ofstream output(filename, std::ios::binary); - RAFT_EXPECTS(output.is_open(), "failed to open file for saving: %s", filename); + cuvs::util::kvikio_ofstream output(filename); + RAFT_EXPECTS(output, "failed to open file for saving: %s", filename); auto write_exact = [&](const void* ptr, size_t n_bytes) { output.write(reinterpret_cast(ptr), n_bytes); @@ -386,37 +387,14 @@ void IVFGPU::save(const char* filename) const // for batch data size_t short_factors_size = GetShortDataFactorBytesBatch(); - // Allocate temporary host buffers. - auto h_short_data_buf = raft::make_host_vector(short_data_size); - auto h_long_code_buf = raft::make_host_vector(long_code_size); - auto h_ex_factor_buf = - raft::make_host_vector(ex_factor_size / sizeof(ExFactor)); - auto h_ids_buf = raft::make_host_vector(ids_size / sizeof(PID)); - auto h_short_factors_batch_buf = raft::make_host_vector(short_factors_size); - - // Copy device data to host (SoA layout). - raft::copy(h_short_data_buf.data_handle(), - reinterpret_cast(short_data_.data_handle()), - short_data_size, - stream_); - raft::copy(h_short_factors_batch_buf.data_handle(), - reinterpret_cast(short_factors_batch_.data_handle()), - short_factors_size, - stream_); - raft::copy(h_long_code_buf.data_handle(), long_code_.data_handle(), long_code_size, stream_); - raft::copy(h_ex_factor_buf.data_handle(), - ex_factor_.data_handle(), - ex_factor_size / sizeof(ExFactor), - stream_); - raft::copy(h_ids_buf.data_handle(), ids_.data_handle(), ids_size / sizeof(PID), stream_); raft::resource::sync_stream(handle_); // Write raw arrays to file. - write_exact(h_short_data_buf.data_handle(), short_data_size); - write_exact(h_short_factors_batch_buf.data_handle(), short_factors_size); - write_exact(h_long_code_buf.data_handle(), long_code_size); - write_exact(h_ex_factor_buf.data_handle(), ex_factor_size); - write_exact(h_ids_buf.data_handle(), ids_size); + output.write_device(short_data_.data_handle(), short_data_size); + output.write_device(short_factors_batch_.data_handle(), short_factors_size); + output.write_device(long_code_.data_handle(), long_code_size); + output.write_device(ex_factor_.data_handle(), ex_factor_size); + output.write_device(ids_.data_handle(), ids_size); output.close(); } diff --git a/cpp/src/neighbors/ivf_rabitq/gpu_index/rotator_gpu.cu b/cpp/src/neighbors/ivf_rabitq/gpu_index/rotator_gpu.cu index 59a3e575d9..e095c35205 100644 --- a/cpp/src/neighbors/ivf_rabitq/gpu_index/rotator_gpu.cu +++ b/cpp/src/neighbors/ivf_rabitq/gpu_index/rotator_gpu.cu @@ -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 */ @@ -44,14 +44,10 @@ void RotatorGPU::load(raft::resources const& handle, std::ifstream& input) raft::resource::sync_stream(handle); } -void RotatorGPU::save(raft::resources const& handle, std::ofstream& output) const +void RotatorGPU::save(raft::resources const& handle, cuvs::util::kvikio_ofstream& output) const { - auto stream = raft::resource::get_cuda_stream(handle); - auto host_buf = raft::make_host_vector(D * D); - raft::copy(host_buf.data_handle(), rotation_matrix_.data_handle(), D * D, stream); raft::resource::sync_stream(handle); - output.write(reinterpret_cast(host_buf.data_handle()), - static_cast(sizeof(float) * D * D)); + output.write_device(rotation_matrix_.data_handle(), sizeof(float) * D * D); RAFT_EXPECTS(static_cast(output), "failed to write rotator matrix"); } diff --git a/cpp/src/neighbors/ivf_rabitq/gpu_index/rotator_gpu.cuh b/cpp/src/neighbors/ivf_rabitq/gpu_index/rotator_gpu.cuh index 8db4bb464c..e945d91c91 100644 --- a/cpp/src/neighbors/ivf_rabitq/gpu_index/rotator_gpu.cuh +++ b/cpp/src/neighbors/ivf_rabitq/gpu_index/rotator_gpu.cuh @@ -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 */ @@ -11,6 +11,8 @@ #include "../defines.hpp" +#include + #include #include #include @@ -59,7 +61,7 @@ class RotatorGPU { * The function copies the rotation matrix from device memory, transposes it from column-major to * row-major, and writes it to the file. */ - void save(raft::resources const& handle, std::ofstream& output) const; + void save(raft::resources const& handle, cuvs::util::kvikio_ofstream& output) const; // Rotate matrix A and store the result in RAND_A. // A and RAND_A are device pointers representing matrices of size N x D. diff --git a/cpp/src/neighbors/ivf_sq/ivf_sq_serialize.cuh b/cpp/src/neighbors/ivf_sq/ivf_sq_serialize.cuh index 22d89ba483..992e9d818c 100644 --- a/cpp/src/neighbors/ivf_sq/ivf_sq_serialize.cuh +++ b/cpp/src/neighbors/ivf_sq/ivf_sq_serialize.cuh @@ -5,6 +5,7 @@ #pragma once +#include "../../util/kvikio_serialize.hpp" #include "../../util/serialize_validation.hpp" #include "../ivf_common.cuh" #include "../ivf_list.cuh" @@ -22,8 +23,8 @@ namespace cuvs::neighbors::ivf_sq::detail { constexpr int serialization_version = 2; -template -void serialize(raft::resources const& handle, std::ostream& os, const index& index_) +template +void serialize(raft::resources const& handle, Output& os, const index& index_) { RAFT_LOG_DEBUG( "Saving IVF-SQ index, size %zu, dim %u", static_cast(index_.size()), index_.dim()); @@ -38,19 +39,19 @@ void serialize(raft::resources const& handle, std::ostream& os, const index(index_.list_sizes().extent(0)); raft::copy(sizes_host.data_handle(), @@ -58,7 +59,7 @@ void serialize(raft::resources const& handle, std::ostream& os, const index list_store_spec{index_.dim(), true}; for (uint32_t label = 0; label < index_.n_lists(); label++) { @@ -72,16 +73,18 @@ void serialize(raft::resources const& handle, const std::string& filename, const index& index_) { - 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(handle, of, index_); of.close(); if (!of) { RAFT_FAIL("Error writing output %s", filename.c_str()); } } -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_sq::deserialize: failed to read dtype prefix"); RAFT_EXPECTS(cuvs::util::validate_serialized_dtype(dtype_string, sizeof(dtype_string)), @@ -116,7 +119,7 @@ auto deserialize(raft::resources const& handle, std::istream& is) -> index index_ = index(handle, metric, n_lists, dim, cma); - 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) { @@ -125,19 +128,19 @@ 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); @@ -146,14 +149,17 @@ 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); +} + template auto deserialize(raft::resources const& handle, const std::string& filename) -> 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 index; + cuvs::util::kvikio_file_reader reader(filename); + return deserialize_impl(handle, reader); } } // namespace cuvs::neighbors::ivf_sq::detail diff --git a/cpp/src/neighbors/mg/snmg.cuh b/cpp/src/neighbors/mg/snmg.cuh index ae7814acb8..01e771fbf9 100644 --- a/cpp/src/neighbors/mg/snmg.cuh +++ b/cpp/src/neighbors/mg/snmg.cuh @@ -24,6 +24,7 @@ #include #include #include +#include #include @@ -782,7 +783,7 @@ void serialize(const raft::resources& clique, const mg_index& index, const std::string& filename) { - 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()); } std::string dtype_string = raft::numpy_serializer::get_numpy_dtype().to_string(); diff --git a/cpp/src/neighbors/scann/detail/scann_serialize.cuh b/cpp/src/neighbors/scann/detail/scann_serialize.cuh index c3b4fdf357..bdb8e1be5f 100644 --- a/cpp/src/neighbors/scann/detail/scann_serialize.cuh +++ b/cpp/src/neighbors/scann/detail/scann_serialize.cuh @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -14,28 +14,17 @@ #include #include +#include "../../../util/kvikio_serialize.hpp" + #include #include #include -#include #include namespace cuvs::neighbors::experimental::scann::detail { constexpr int kSerializationVersion = 1; -namespace { -// Helper for opening an ofstream -std::ofstream open_file(std::string file_name) -{ - std::ofstream file_of(file_name, std::ios::out | std::ios::binary); - - if (!file_of) { RAFT_FAIL("Cannot open file %s", file_name.c_str()); } - - return file_of; -} -} // namespace - // Helper for serializing device/host matrix to a given file template , raft::row_major, Accessor> mat) { - auto mat_of = open_file(file_path); + cuvs::util::kvikio_ofstream mat_of(file_path.string()); + RAFT_EXPECTS(mat_of, "Cannot open file %s", file_path.string().c_str()); - raft::serialize_mdspan(res, mat_of, mat); + cuvs::util::detail::serialize_mdspan(res, mat_of, mat); mat_of.close(); + RAFT_EXPECTS(mat_of, "Error writing output %s", file_path.string().c_str()); } // ScaNN stores primary and soar cluster assignments interleaved in a single vector. @@ -81,11 +72,13 @@ void save_labels(raft::resources const& res, return soar_label; }); - auto labels_of = open_file(labels_path); + cuvs::util::kvikio_ofstream labels_of(labels_path.string()); + RAFT_EXPECTS(labels_of, "Cannot open file %s", labels_path.string().c_str()); - raft::serialize_mdspan(res, labels_of, combined_labels.view()); + cuvs::util::detail::serialize_device_mdspan(res, labels_of, combined_labels.view()); labels_of.close(); + RAFT_EXPECTS(labels_of, "Error writing output %s", labels_path.string().c_str()); } /** @@ -111,13 +104,16 @@ void serialize(raft::resources const& res, // Metadata std::filesystem::path scann_path(scann_assets_dir); - auto metadata_of = open_file(scann_path / "cuvs_metadata.bin"); + const auto metadata_path = scann_path / "cuvs_metadata.bin"; + cuvs::util::kvikio_ofstream metadata_of(metadata_path.string()); + RAFT_EXPECTS(metadata_of, "Cannot open file %s", metadata_path.string().c_str()); raft::serialize_scalar(res, metadata_of, kSerializationVersion); raft::serialize_scalar(res, metadata_of, index_.dim()); raft::serialize_scalar(res, metadata_of, index_.pq_dim()); metadata_of.close(); + RAFT_EXPECTS(metadata_of, "Error writing output %s", metadata_path.string().c_str()); // kmeans cluster centers serialize_matrix(res, scann_path / "centers.npy", index_.centers()); diff --git a/cpp/src/util/file_io.cpp b/cpp/src/util/file_io.cpp index d924527e72..cb7dba3c23 100644 --- a/cpp/src/util/file_io.cpp +++ b/cpp/src/util/file_io.cpp @@ -1,18 +1,151 @@ /* - * 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 "kvikio_io.hpp" + +#include + #include +#include #include -#include - -#include -#include +#include +#include +#include namespace cuvs::util { +namespace { + +constexpr size_t kPosixTransferSize = size_t{1} << 30; + +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); + } +} + +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); + } +} + +off_t checked_posix_offset(uint64_t file_offset, size_t buffer_offset) +{ + RAFT_EXPECTS(buffer_offset <= std::numeric_limits::max() - file_offset, + "File offset overflow"); + const uint64_t position = file_offset + buffer_offset; + RAFT_EXPECTS(position <= static_cast(std::numeric_limits::max()), + "File offset exceeds the POSIX offset range"); + return static_cast(position); +} + +void read_large_file_posix(const file_descriptor& fd, + void* dest_ptr, + size_t total_bytes, + uint64_t file_offset) +{ + auto* destination = static_cast(dest_ptr); + size_t bytes_read = 0; + + while (bytes_read < total_bytes) { + const size_t chunk_size = std::min(kPosixTransferSize, total_bytes - bytes_read); + const off_t position = checked_posix_offset(file_offset, bytes_read); + ssize_t result = 0; + do { + result = ::pread(fd.get(), destination + bytes_read, chunk_size, position); + } while (result < 0 && errno == EINTR); + + RAFT_EXPECTS(result >= 0, + "Failed to read from file descriptor at offset %llu: %s", + static_cast(position), + std::strerror(errno)); + RAFT_EXPECTS(result > 0, + "Incomplete read from file descriptor: expected %zu bytes, got %zu", + total_bytes, + bytes_read); + bytes_read += static_cast(result); + } +} + +void write_large_file_posix(const file_descriptor& fd, + const void* data_ptr, + size_t total_bytes, + uint64_t file_offset) +{ + const auto* source = static_cast(data_ptr); + size_t bytes_written = 0; + + while (bytes_written < total_bytes) { + const size_t chunk_size = std::min(kPosixTransferSize, total_bytes - bytes_written); + const off_t position = checked_posix_offset(file_offset, bytes_written); + ssize_t result = 0; + do { + result = ::pwrite(fd.get(), source + bytes_written, chunk_size, position); + } while (result < 0 && errno == EINTR); + + RAFT_EXPECTS(result >= 0, + "Failed to write to file descriptor at offset %llu: %s", + static_cast(position), + std::strerror(errno)); + RAFT_EXPECTS(result > 0, + "Incomplete write to file descriptor: expected %zu bytes, wrote %zu", + total_bytes, + bytes_written); + bytes_written += static_cast(result); + } +} + +} // namespace void read_large_file(const file_descriptor& fd, void* dest_ptr, @@ -22,28 +155,25 @@ 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 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; + const std::string path = fd.get_path(); + if (path.empty()) { + RAFT_EXPECTS(!detail::is_kvikio_device_memory(dest_ptr), + "Pathless POSIX I/O requires host memory"); + read_large_file_posix(fd, dest_ptr, total_bytes, file_offset); + return; } + + // 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. + 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, + "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 +184,251 @@ 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(); + if (path.empty()) { + RAFT_EXPECTS(!detail::is_kvikio_device_memory(data_ptr), + "Pathless POSIX I/O requires host memory"); + write_large_file_posix(fd, data_ptr, total_bytes, file_offset); + return; + } + + // 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. + 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, + "Incomplete write to %s: expected %zu bytes, wrote %zu", + path.c_str(), + total_bytes, + 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_; } - const size_t write_chunk_size = std::min(1024 * 1024 * 1024, SSIZE_MAX); - size_t bytes_remaining = total_bytes; - size_t offset = 0; + 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"); - 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); + 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"); - 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); + const size_t bytes_read = handle_.pread(data, size, static_cast(offset)).get(); + RAFT_EXPECTS(bytes_read == size, + "kvikio_file_reader: short read (expected %llu, read %llu)", + static_cast(size), + static_cast(bytes_read)); - bytes_remaining -= chunk_size; - offset += chunk_size; + 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 { + public: + sbuf(const std::string& path, size_t cap) + : 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"); + setp(buffer_.data(), buffer_.data() + buffer_.size()); + } + + ~sbuf() override + { + try { + close(); + } catch (...) { + // Swallow during destruction. + } + } + + void close() + { + 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 + { + 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_at_current_offset(handle_, 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; + } + + 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_at_current_offset(handle_, pbase(), n); + setp(buffer_.data(), buffer_.data() + buffer_.size()); + } + } + + 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(); + 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; +}; + +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(); } +} + +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/host_memory.cpp b/cpp/src/util/host_memory.cpp index 10630749eb..2d1e7499b6 100644 --- a/cpp/src/util/host_memory.cpp +++ b/cpp/src/util/host_memory.cpp @@ -1,29 +1,357 @@ /* - * 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); +} + +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, + 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 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, + reclaimable_file, + working_set, + 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; } } - available_memory *= 1024; - meminfo.close(); - RAFT_EXPECTS(available_memory > 0, "Failed to get available memory from /proc/meminfo"); - return available_memory; + 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; + } + } + } + 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, 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.cgroup_reclaimable_file = cgroup->reclaimable_file; + result.cgroup_working_set = cgroup->working_set; + 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..353be1d387 --- /dev/null +++ b/cpp/src/util/host_memory.hpp @@ -0,0 +1,32 @@ +/* + * 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 reclaimable_file; + size_t working_set; + 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/src/util/kvikio_io.hpp b/cpp/src/util/kvikio_io.hpp new file mode 100644 index 0000000000..a964d37cae --- /dev/null +++ b/cpp/src/util/kvikio_io.hpp @@ -0,0 +1,94 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +#pragma once + +#include + +#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 +#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 kvikio::FileHandle open_kvikio_file_for_device_io(const std::string& path, + const std::string& 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& e) { + static std::once_flag warning_once; + std::call_once(warning_once, [&] { + RAFT_LOG_WARN( + "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); + } +} + +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 +#undef CUVS_KVIKIO_HAS_UTILS_HEADER diff --git a/cpp/src/util/kvikio_serialize.hpp b/cpp/src/util/kvikio_serialize.hpp new file mode 100644 index 0000000000..5161e6102d --- /dev/null +++ b/cpp/src/util/kvikio_serialize.hpp @@ -0,0 +1,145 @@ +/* + * 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 +#include +#include + +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) +{ + 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 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, + 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"); +} + +template +void serialize_mdspan(const raft::resources& res, std::ostream& os, const MdspanT& obj) +{ + if constexpr (raft::is_device_mdspan_v) { + if (auto* kvikio_stream = dynamic_cast(&os); kvikio_stream != nullptr) { + return serialize_device_mdspan(res, *kvikio_stream, obj); + } + } + raft::serialize_mdspan(res, os, obj); +} + +template +void serialize_mdspan( + const raft::resources& res, + kvikio_ofstream& os, + const raft::device_mdspan& obj) +{ + serialize_device_mdspan(res, os, obj); +} + +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/CMakeLists.txt b/cpp/tests/CMakeLists.txt index b8f571e4a9..744bc6a7a2 100644 --- a/cpp/tests/CMakeLists.txt +++ b/cpp/tests/CMakeLists.txt @@ -435,6 +435,13 @@ ConfigureTest( PERCENT 100 ) +ConfigureTest( + NAME UTIL_TEST + PATH util/file_io_test.cpp util/host_memory_test.cpp + GPUS 1 + PERCENT 100 +) + # ################################################################################################## # Install tests #################################################################################### # ################################################################################################## diff --git a/cpp/tests/neighbors/ann_ivf_flat/test_float_int64_t.cu b/cpp/tests/neighbors/ann_ivf_flat/test_float_int64_t.cu index d4499ed6d9..cf24a64d5d 100644 --- a/cpp/tests/neighbors/ann_ivf_flat/test_float_int64_t.cu +++ b/cpp/tests/neighbors/ann_ivf_flat/test_float_int64_t.cu @@ -11,7 +11,9 @@ #include #include +#include #include +#include namespace cuvs::neighbors::ivf_flat { @@ -25,6 +27,27 @@ TEST_P(AnnIVFFlatTestF_float, AnnIVFFlat) INSTANTIATE_TEST_CASE_P(AnnIVFFlatTest, AnnIVFFlatTestF_float, ::testing::ValuesIn(inputs)); +TEST(IVFFlatSerialization, StreamRoundTrip) +{ + raft::resources handle; + auto dataset = raft::make_host_matrix(32, 4); + std::iota(dataset.data_handle(), dataset.data_handle() + dataset.size(), 0.0f); + + index_params params; + params.n_lists = 4; + const auto index = build(handle, params, raft::make_const_mdspan(dataset.view())); + + std::stringstream stream(std::ios::in | std::ios::out | std::ios::binary); + serialize(handle, stream, index); + stream.seekg(0); + + ivf_flat::index loaded(handle); + deserialize(handle, stream, &loaded); + EXPECT_EQ(loaded.size(), index.size()); + EXPECT_EQ(loaded.dim(), index.dim()); + EXPECT_EQ(loaded.n_lists(), index.n_lists()); +} + TEST(AnnIVFFlatTest, RepeatedExtendCopyPreservesSharedListWithinCapacity) { raft::resources handle; diff --git a/cpp/tests/util/file_io_test.cpp b/cpp/tests/util/file_io_test.cpp new file mode 100644 index 0000000000..770b2b0271 --- /dev/null +++ b/cpp/tests/util/file_io_test.cpp @@ -0,0 +1,501 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include + +#include "../../src/util/kvikio_serialize.hpp" + +#include +#include +#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()); +} + +void expect_stream_contents(std::istream& stream, const std::vector& expected, size_t offset) +{ + ASSERT_LE(offset, expected.size()); + std::vector actual(expected.size() - offset); + stream.read(actual.data(), static_cast(actual.size())); + ASSERT_EQ(stream.gcount(), static_cast(actual.size())); + EXPECT_TRUE(std::equal(actual.begin(), actual.end(), expected.begin() + offset)); +} + +} // 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); +} + +// Descriptors constructed from raw file descriptors do not have a path that KvikIO can reopen. +// Preserve the public bulk-I/O API for those callers through positioned POSIX I/O. +TEST(FileIO, PathlessDescriptorReadWriteRoundTrip) +{ + scratch_dir scratch; + file_descriptor original(scratch.file("pathless.bin"), O_CREAT | O_RDWR | O_TRUNC, 0644); + const int dup_fd = ::dup(original.get()); + ASSERT_NE(dup_fd, -1); + file_descriptor pathless_fd(dup_fd); + ASSERT_TRUE(pathless_fd.get_path().empty()); + + const std::vector src = make_pattern(16391, 13); + const uint64_t offset = 37; + write_large_file(pathless_fd, src.data(), src.size(), offset); + + std::vector dst(src.size()); + read_large_file(pathless_fd, dst.data(), dst.size(), offset); + EXPECT_EQ(dst, src); +} + +// Pathless descriptors cannot reopen the file through KvikIO, so device pointers must be rejected +// before the POSIX pread/pwrite fallback runs. +TEST(FileIO, PathlessDescriptorRejectsDeviceMemory) +{ + raft::resources res; + scratch_dir scratch; + file_descriptor original(scratch.file("pathless_device.bin"), O_CREAT | O_RDWR | O_TRUNC, 0644); + const int dup_fd = ::dup(original.get()); + ASSERT_NE(dup_fd, -1); + file_descriptor pathless_fd(dup_fd); + ASSERT_TRUE(pathless_fd.get_path().empty()); + + auto device = raft::make_device_vector(res, 128); + auto expect_host_memory_error = [](auto&& op) { + try { + op(); + FAIL() << "expected pathless POSIX I/O to reject device memory"; + } catch (const raft::exception& e) { + EXPECT_NE(std::string(e.what()).find("host memory"), std::string::npos) << e.what(); + } + }; + expect_host_memory_error([&] { write_large_file(pathless_fd, device.data_handle(), 128, 0); }); + expect_host_memory_error([&] { read_large_file(pathless_fd, device.data_handle(), 128, 0); }); +} + +// Moving a stream after a short read must retain bytes already buffered beyond the logical stream +// position. The underlying descriptor has advanced to the end of that read-ahead buffer. +TEST(FileIO, FdIstreamMovePreservesUnreadBuffer) +{ + scratch_dir scratch; + const std::string path = scratch.file("move_stream.bin"); + const std::vector data = make_pattern(32771, 29); + { + std::ofstream output(path, std::ios::out | std::ios::binary); + ASSERT_TRUE(output.good()); + output.write(data.data(), static_cast(data.size())); + } + + constexpr size_t prefix_size = 37; + { + file_descriptor fd(path, O_RDONLY); + auto source = fd.make_istream(); + std::vector prefix(prefix_size); + source.read(prefix.data(), static_cast(prefix.size())); + ASSERT_EQ(source.gcount(), static_cast(prefix.size())); + EXPECT_TRUE(std::equal(prefix.begin(), prefix.end(), data.begin())); + + fd_istream destination(std::move(source)); + expect_stream_contents(destination, data, prefix_size); + } + + { + file_descriptor fd(path, O_RDONLY); + auto source = fd.make_istream(); + std::vector prefix(prefix_size); + source.read(prefix.data(), static_cast(prefix.size())); + ASSERT_EQ(source.gcount(), static_cast(prefix.size())); + + file_descriptor destination_fd(path, O_RDONLY); + auto destination = destination_fd.make_istream(); + destination = std::move(source); + expect_stream_contents(destination, data, prefix_size); + } + + { + file_descriptor fd(path, O_RDONLY); + auto source = fd.make_istream(); + std::vector all(data.size()); + source.read(all.data(), static_cast(all.size())); + ASSERT_EQ(source.gcount(), static_cast(all.size())); + + fd_istream destination(std::move(source)); + expect_stream_contents(destination, data, data.size()); + } +} + +// 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); +} + +// 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, KvikioDeviceMdspanSerializationThroughOstream) +{ + raft::resources res; + scratch_dir scratch; + const std::string path = scratch.file("device_mdspan.npy"); + const int64_t size = 4096; + + auto source_host = raft::make_host_vector(size); + for (int64_t i = 0; i < size; ++i) { + source_host(i) = static_cast(i * 17 + 3); + } + auto source_device = raft::make_device_vector(res, size); + raft::copy(res, source_device.view(), source_host.view()); + + { + kvikio_ofstream file(path); + std::ostream& os = file; + detail::serialize_mdspan(res, os, raft::make_const_mdspan(source_device.view())); + file.close(); + } + + auto result = raft::make_host_vector(size); + std::ifstream input(path, std::ios::in | std::ios::binary); + raft::deserialize_mdspan(res, input, result.view()); + EXPECT_TRUE(std::equal(source_host.data_handle(), + source_host.data_handle() + source_host.size(), + result.data_handle())); +} + +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) +{ + 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); +} + +// 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); +} + +// 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) +{ + 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)); +} + +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/cpp/tests/util/host_memory_test.cpp b/cpp/tests/util/host_memory_test.cpp new file mode 100644 index 0000000000..09bd9e84b9 --- /dev/null +++ b/cpp/tests/util/host_memory_test.cpp @@ -0,0 +1,192 @@ +/* + * 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", "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", "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"); + + 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, 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()); +} + +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->reclaimable_file, 0); + EXPECT_EQ(result->working_set, 512); + EXPECT_EQ(result->available, 1536); + EXPECT_EQ(result->path, leaf.string()); +} + +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"); + 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"); + 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")); + 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->reclaimable_file, 5000); + EXPECT_EQ(result->working_set, 3192); + EXPECT_EQ(result->available, 904); +} + +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 diff --git a/dependencies.yaml b/dependencies.yaml index a58e00cb58..6978603d41 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 @@ -753,6 +756,29 @@ dependencies: packages: - librmm-cu13==26.10.*,>=0.0.0a0 - {matrix: null, packages: [*librmm_unsuffixed]} + depends_on_libkvikio: + common: + - output_types: conda + packages: + - &libkvikio_unsuffixed libkvikio==26.10.*,>=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.10.*,>=0.0.0a0 + - matrix: + cuda: "13.*" + cuda_suffixed: "true" + packages: + - libkvikio-cu13==26.10.*,>=0.0.0a0 + - {matrix: null, packages: [*libkvikio_unsuffixed]} depends_on_pylibraft: common: - output_types: conda diff --git a/python/libcuvs/libcuvs/load.py b/python/libcuvs/libcuvs/load.py index e4dfa66527..b5ac3571f9 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,35 @@ def _load_wheel_installation(soname: str): return None -def load_library(): - """Dynamically load libcuvs.so and its dependencies""" +def _load_dependency(module_name: str): + """Load an optional wheel-provided native dependency.""" try: - # These libraries must be loaded before libcuvs because libcuvs - # references their symbols - import libraft - import librmm - from cuda.pathfinder import load_nvidia_dynamic_lib + module = importlib.import_module(module_name) + except ModuleNotFoundError: + # Conda packages provide the native libraries without Python loader + # modules, in which case the system loader is expected to find them. + return + module.load_library() + - librmm.load_library() - libraft.load_library() - load_nvidia_dynamic_lib("nvrtc") +def _load_nvrtc(): + """Load NVRTC from a CUDA wheel when cuda-pathfinder is available.""" + try: + from cuda.pathfinder import load_nvidia_dynamic_lib 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 + load_nvidia_dynamic_lib("nvrtc") + + +def load_library(): + """Dynamically load libcuvs.so and its dependencies""" + # These libraries must be loaded before libcuvs because libcuvs + # references their symbols. Load them independently so a dependency + # supplied by conda does not prevent wheel-provided dependencies from + # being preloaded. + for module_name in ("libkvikio", "librmm", "libraft"): + _load_dependency(module_name) + _load_nvrtc() prefer_system_installation = ( os.getenv("RAPIDS_LIBCUVS_PREFER_SYSTEM_LIBRARY", "false").lower() diff --git a/python/libcuvs/pyproject.toml b/python/libcuvs/pyproject.toml index 6e8789349c..f9b3b3ee23 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.10.*,>=0.0.0a0", "libraft==26.10.*,>=0.0.0a0", "librmm==26.10.*,>=0.0.0a0", "nvidia-nvjitlink>=13.3,<14", @@ -81,6 +82,7 @@ regex = "(?P.*)" build-backend = "scikit_build_core.build" requires = [ "cmake>=4.0", + "libkvikio==26.10.*,>=0.0.0a0", "libraft==26.10.*,>=0.0.0a0", "librmm==26.10.*,>=0.0.0a0", "ninja",