From 003025d7cafdde978dd9ac8f9a160f327009b2a7 Mon Sep 17 00:00:00 2001 From: ScissorJack-ever <1533201944@qq.com> Date: Wed, 22 Jul 2026 08:02:05 +0000 Subject: [PATCH] feat(diskann): support cross-platform updates --- CMakeLists.txt | 3 +- cmake/AlayaDependencies.cmake | 18 +- cmake/AlayaDiskANN.cmake | 47 ++++ conanfile.py | 5 +- include/index/graph/diskann/disk_page_io.hpp | 223 +++++++++++++++--- include/index/graph/diskann/diskann_index.hpp | 19 +- .../graph/laser/utils/iocp_file_reader.hpp | 2 +- pyproject.toml | 2 +- python/tests/diskann/test_update.py | 34 ++- tests/diskann/CMakeLists.txt | 10 + .../diskann/test_diskann_portable_update.cpp | 81 +++++++ tests/diskann/test_diskann_update_e2e.cpp | 57 +++-- 12 files changed, 424 insertions(+), 77 deletions(-) create mode 100644 cmake/AlayaDiskANN.cmake create mode 100644 tests/diskann/test_diskann_portable_update.cpp diff --git a/CMakeLists.txt b/CMakeLists.txt index 42a29c8f..f6f7f436 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -41,6 +41,7 @@ alaya_preflight() include(cmake/AlayaToolchain.cmake) include(cmake/AlayaFlags.cmake) include(cmake/AlayaDependencies.cmake) +include(cmake/AlayaDiskANN.cmake) include(cmake/AlayaLaser.cmake) include(cmake/Codecov.cmake) include(cmake/PrintSummary.cmake) @@ -52,7 +53,7 @@ add_library(AlayaLite INTERFACE) target_include_directories( AlayaLite INTERFACE $ $ ) -target_link_libraries(AlayaLite INTERFACE ${THIRD_PARTY_LIBS}) +target_link_libraries(AlayaLite INTERFACE ${THIRD_PARTY_LIBS} alaya_diskann) target_compile_features(AlayaLite INTERFACE cxx_std_20) if(BUILD_TESTING) diff --git a/cmake/AlayaDependencies.cmake b/cmake/AlayaDependencies.cmake index ab1e1f7b..552a482d 100644 --- a/cmake/AlayaDependencies.cmake +++ b/cmake/AlayaDependencies.cmake @@ -6,8 +6,8 @@ # # Packages resolve through the Conan dependency provider registered in AlayaConan.cmake: the first find_package below # triggers `conan install` transparently and the generated config packages land in ${CMAKE_BINARY_DIR}/conan. OpenMP -# (system / Homebrew) and libaio (system, handled in AlayaLaser.cmake) are not in the Conan graph and fall through the -# provider to CMake's builtin lookup. THIRD_PARTY_LIBS stays a plain list because test helper code and the AlayaLite +# (system / Homebrew) and libaio (system, handled by the disk-I/O targets) are not in the Conan graph and fall through +# the provider to CMake's builtin lookup. THIRD_PARTY_LIBS stays a plain list because test helper code and the AlayaLite # INTERFACE target both consume it. include_guard(GLOBAL) @@ -28,6 +28,7 @@ find_package(spdlog REQUIRED) find_package(Eigen3 REQUIRED) find_package(OpenMP QUIET) find_package(RocksDB REQUIRED) +find_package(libcoro REQUIRED) if(BUILD_PYTHON) # Reuse the interpreter alaya_preflight() resolved (FindPython / Python_* family) instead of letting pybind11 run its @@ -36,14 +37,19 @@ if(BUILD_PYTHON) find_package(pybind11 REQUIRED) endif() -set(THIRD_PARTY_LIBS spdlog::spdlog_header_only concurrentqueue::concurrentqueue Eigen3::Eigen RocksDB::rocksdb) +set(THIRD_PARTY_LIBS + spdlog::spdlog_header_only + concurrentqueue::concurrentqueue + Eigen3::Eigen + RocksDB::rocksdb + libcoro::libcoro +) if(TARGET OpenMP::OpenMP_CXX) list(APPEND THIRD_PARTY_LIBS OpenMP::OpenMP_CXX) endif() -if(UNIX AND NOT APPLE) - find_package(libcoro REQUIRED) +if(CMAKE_SYSTEM_NAME STREQUAL "Linux") find_package(liburing REQUIRED) - list(APPEND THIRD_PARTY_LIBS libcoro::libcoro liburing::liburing) + list(APPEND THIRD_PARTY_LIBS liburing::liburing) endif() diff --git a/cmake/AlayaDiskANN.cmake b/cmake/AlayaDiskANN.cmake new file mode 100644 index 00000000..239e7197 --- /dev/null +++ b/cmake/AlayaDiskANN.cmake @@ -0,0 +1,47 @@ +# SPDX-FileCopyrightText: 2025 AlayaDB.AI +# +# SPDX-License-Identifier: AGPL-3.0-only + +# DiskANN's portable reader backend and platform runtime dependencies. + +include_guard(GLOBAL) + +set(_alaya_diskann_backend_libs libcoro::libcoro) + +if(CMAKE_SYSTEM_NAME STREQUAL "Windows") + set(_alaya_diskann_backend_definition ALAYA_LASER_USE_IOCP=1) + set(_alaya_diskann_backend_message "IOCP + Win32 positioned writes") +elseif(CMAKE_SYSTEM_NAME STREQUAL "Darwin" OR ALAYA_LASER_USE_THREADPOOL) + set(_alaya_diskann_backend_definition ALAYA_LASER_USE_THREADPOOL=1) + set(_alaya_diskann_backend_message "thread-pool reads + POSIX positioned writes") +elseif(CMAKE_SYSTEM_NAME STREQUAL "Linux") + find_library(AIO_LIBRARY aio) + find_path(AIO_INCLUDE_DIR libaio.h) + if(NOT AIO_LIBRARY OR NOT AIO_INCLUDE_DIR) + message(FATAL_ERROR "DiskANN requires libaio for the default Linux reader. Install libaio-dev/libaio-devel, " + "or configure with -DALAYA_LASER_USE_THREADPOOL=ON." + ) + endif() + if(NOT TARGET AIO::aio) + add_library(AIO::aio UNKNOWN IMPORTED) + set_target_properties( + AIO::aio PROPERTIES IMPORTED_LOCATION "${AIO_LIBRARY}" INTERFACE_INCLUDE_DIRECTORIES "${AIO_INCLUDE_DIR}" + ) + endif() + list(APPEND _alaya_diskann_backend_libs AIO::aio liburing::liburing) + set(_alaya_diskann_backend_definition ALAYA_LASER_USE_LIBAIO=1) + set(_alaya_diskann_backend_message "libaio reads + O_DIRECT positioned writes") +else() + set(_alaya_diskann_backend_definition ALAYA_LASER_USE_THREADPOOL=1) + set(_alaya_diskann_backend_message "portable thread-pool reads + POSIX positioned writes") +endif() + +message(STATUS "DiskANN I/O backend: ${_alaya_diskann_backend_message}") + +add_library(alaya_diskann INTERFACE) +target_include_directories( + alaya_diskann INTERFACE $ $ +) +target_link_libraries(alaya_diskann INTERFACE ${_alaya_diskann_backend_libs}) +target_compile_definitions(alaya_diskann INTERFACE ${_alaya_diskann_backend_definition}) +target_compile_features(alaya_diskann INTERFACE cxx_std_20) diff --git a/conanfile.py b/conanfile.py index 7cfc915d..16a604ae 100644 --- a/conanfile.py +++ b/conanfile.py @@ -42,9 +42,10 @@ def requirements(self): self.requires("lz4/1.9.4") self.requires("zstd/1.5.6") self.requires("rocksdb/10.5.1") - # io_uring and coroutine scheduler support are Linux-only in CMake. + # DiskANN uses libcoro's portable task/thread-pool primitives on every + # platform; only its io_uring accelerator is Linux-specific. + self.requires("libcoro/0.14.1") if self.settings.os == "Linux": - self.requires("libcoro/0.14.1") self.requires("liburing/2.13") # OpenMP support diff --git a/include/index/graph/diskann/disk_page_io.hpp b/include/index/graph/diskann/disk_page_io.hpp index 40507fad..7d3589d2 100644 --- a/include/index/graph/diskann/disk_page_io.hpp +++ b/include/index/graph/diskann/disk_page_io.hpp @@ -7,22 +7,21 @@ * @brief Sector-aligned read-modify-write of node records in diskann.index. * * `DiskPageIO` (design D6) is the single place that writes the disk index during - * in-place updates. It owns its own `O_DIRECT | O_RDWR` file descriptor (separate - * from the search reader's read-only descriptor); in serial-update mode (a global - * mutex serialises search and update) the two descriptors never race, and because - * both use O_DIRECT a completed pwrite is visible to a subsequent pread. + * in-place updates. It owns a read-write file handle separate from the search + * reader. Linux uses O_DIRECT, macOS uses positioned POSIX I/O, and Windows uses + * overlapped positioned I/O so concurrent page operations never share a mutable + * file offset. * * Each operation read-modify-writes one sector-aligned page so co-resident nodes - * survive updates. Appends extend the file with ftruncate. Linux O_DIRECT is - * required for the private syscall path; non-Linux update calls throw loudly. + * survive updates. Appends extend the file with the platform's native resize API. * * Concurrency: state is sharded by page offset (mutex + LRU page cache + RMW * scratch per shard), so parallel reconnect workers touching different pages * proceed independently — the analog of Yi's per-buffer locks; a single global * mutex here was measured to flatline update throughput regardless of worker * count. Pages map to exactly one shard, which preserves per-page RMW atomicity. - * pread/pwrite are positional and thread-safe on one fd; ftruncate extension is - * serialized by a dedicated file mutex (lock order: shard -> file, always). + * Platform reads and writes are positional and thread-safe on one handle; file + * extension is serialized by a dedicated mutex (lock order: shard -> file). */ #pragma once @@ -47,7 +46,12 @@ #include "coro/task.hpp" #include "coro/thread_pool.hpp" #include "coro/when_all.hpp" -#if defined(__linux__) +#if defined(_WIN32) + #ifndef NOMINMAX + #define NOMINMAX + #endif + #include +#else #include #include #include @@ -109,7 +113,7 @@ class DiskPageIO { return node_from_page(shard.page_buf, id); } - /// Attach a UringReactor: page misses in the awaitable paths become + /// Attach a Linux UringReactor: page misses in the awaitable paths become /// suspending io_uring reads instead of pool-thread-blocking preads. /// nullptr (default) keeps every path on blocking pread. void set_reactor(UringReactor *reactor) { @@ -129,7 +133,7 @@ class DiskPageIO { } /// Awaitable single-node read. With a reactor the miss suspends on io_uring; - /// otherwise the blocking O_DIRECT pread executes on @p pool. + /// otherwise the platform's blocking positioned read executes on @p pool. coro::task read_node_async(uint32_t id, coro::thread_pool &pool) { #if defined(__linux__) if (reactor_ != nullptr) { @@ -230,7 +234,7 @@ class DiskPageIO { } /// Read multiple node records using coroutine tasks over private page buffers. - /// Shared cache state is protected briefly; O_DIRECT pread runs outside the mutex + /// Shared cache state is protected briefly; positioned I/O runs outside the mutex /// and is guarded by a page epoch before cached state is reused. std::vector read_nodes_async(const std::vector &ids, uint32_t threads) { return read_nodes_async(ids.data(), static_cast(ids.size()), threads); @@ -485,6 +489,7 @@ class DiskPageIO { write_page_to_disk(shard, page_off, page); }); } + sync_file(); } [[nodiscard]] uint64_t file_size() const { return file_size_.load(std::memory_order_acquire); } @@ -523,7 +528,7 @@ class DiskPageIO { DiskPageCache cache; std::unordered_map versions; char *page_buf = nullptr; ///< RMW scratch, guarded by mutex - char *flush_buf = nullptr; ///< aligned bounce buffer for O_DIRECT cache flushes + char *flush_buf = nullptr; ///< aligned bounce buffer for direct/unbuffered cache flushes }; Shard &shard_for(uint64_t page_off) { @@ -879,7 +884,7 @@ class DiskPageIO { { std::lock_guard file_lock(file_mutex_); if (page_end > file_size_.load(std::memory_order_acquire)) { - extend_to(page_end); // ftruncate; OS zero-fills the new region + extend_to(page_end); // the platform resize API zero-fills the new region std::memset(shard.page_buf, 0, geom_.page_size); return; } @@ -887,7 +892,7 @@ class DiskPageIO { read_page_locked(shard, id); // another thread extended past us meanwhile } - // ---- platform-gated syscalls (Linux O_DIRECT) ---- + // ---- platform file operations ---- void open_rw(const std::string &path); void read_page_locked(Shard &shard, uint32_t id); void read_page_locked_off(Shard &shard, uint64_t page_off); @@ -895,29 +900,39 @@ class DiskPageIO { void write_page_locked(Shard &shard, uint32_t id); void write_page_to_disk(Shard &shard, uint64_t page_off, const char *page); void extend_to(uint64_t new_size); + void sync_file(); void close_fd(); DiskLayoutGeometry geom_; +#if defined(_WIN32) + HANDLE file_handle_ = INVALID_HANDLE_VALUE; +#else int fd_ = -1; +#endif std::atomic file_size_{0}; uint32_t num_shards_ = kNumShards; bool cache_enabled_ = false; std::vector> shards_; std::unordered_map> vec_cache_; mutable std::mutex vec_mutex_; - std::mutex file_mutex_; ///< serializes ftruncate extension (after shard lock) + std::mutex file_mutex_; ///< serializes file extension (after shard lock) #if defined(__linux__) UringReactor *reactor_ = nullptr; ///< not owned; nullptr = blocking pread paths #endif uint32_t update_fallback_threads_ = 8; ///< blocking fallback of read_neighbors_batch_async }; -#if defined(__linux__) +#if !defined(_WIN32) inline void DiskPageIO::open_rw(const std::string &path) { - fd_ = ::open(path.c_str(), O_DIRECT | O_RDWR); // NOLINT(hicpp-vararg) + #if defined(__linux__) + constexpr int kOpenFlags = O_DIRECT | O_RDWR; + #else + constexpr int kOpenFlags = O_RDWR; + #endif + fd_ = ::open(path.c_str(), kOpenFlags); // NOLINT(hicpp-vararg) if (fd_ < 0) { - throw std::runtime_error("DiskPageIO::open_rw: cannot open (O_DIRECT|O_RDWR) " + path); + throw std::runtime_error("DiskPageIO::open_rw: cannot open read-write " + path); } struct stat st{}; if (::fstat(fd_, &st) != 0) { @@ -992,6 +1007,12 @@ inline void DiskPageIO::extend_to(uint64_t new_size) { file_size_.store(new_size, std::memory_order_release); } +inline void DiskPageIO::sync_file() { + if (::fsync(fd_) != 0) { + throw std::runtime_error("DiskPageIO::sync_file: fsync failed"); + } +} + inline void DiskPageIO::close_fd() { if (fd_ >= 0) { ::close(fd_); @@ -999,31 +1020,161 @@ inline void DiskPageIO::close_fd() { } } -#else // !__linux__ : in-place updates require Linux O_DIRECT — fail loudly when used. +#else -inline void DiskPageIO::open_rw(const std::string &) { - throw std::runtime_error("DiskPageIO: in-place DiskANN updates require Linux (O_DIRECT)"); +inline void DiskPageIO::open_rw(const std::string &path) { + const int wide_len = ::MultiByteToWideChar(CP_UTF8, 0, path.c_str(), -1, nullptr, 0); + if (wide_len <= 0) { + throw std::runtime_error("DiskPageIO::open_rw: invalid UTF-8 path"); + } + std::wstring wide_path(static_cast(wide_len), L'\0'); + ::MultiByteToWideChar(CP_UTF8, 0, path.c_str(), -1, wide_path.data(), wide_len); + if (!wide_path.empty() && wide_path.back() == L'\0') { + wide_path.pop_back(); + } + + file_handle_ = ::CreateFileW(wide_path.c_str(), + GENERIC_READ | GENERIC_WRITE, + FILE_SHARE_READ | FILE_SHARE_WRITE, + nullptr, + OPEN_EXISTING, + FILE_FLAG_NO_BUFFERING | FILE_FLAG_OVERLAPPED, + nullptr); + if (file_handle_ == INVALID_HANDLE_VALUE) { + throw std::runtime_error("DiskPageIO::open_rw: CreateFileW failed: Win32 error " + + std::to_string(::GetLastError())); + } + LARGE_INTEGER size{}; + if (!::GetFileSizeEx(file_handle_, &size)) { + const DWORD error = ::GetLastError(); + ::CloseHandle(file_handle_); + file_handle_ = INVALID_HANDLE_VALUE; + throw std::runtime_error("DiskPageIO::open_rw: GetFileSizeEx failed: Win32 error " + + std::to_string(error)); + } + file_size_.store(static_cast(size.QuadPart), std::memory_order_release); +} + +inline void DiskPageIO::read_page_locked(Shard &shard, uint32_t id) { + read_page_locked_off(shard, geom_.get_page_offset(id)); } -inline void DiskPageIO::read_page_locked(Shard &, uint32_t) { - throw std::runtime_error("DiskPageIO: unsupported platform (needs Linux O_DIRECT)"); + +inline void DiskPageIO::read_page_locked_off(Shard &shard, uint64_t page_off) { + if (shard.cache.read(page_off, shard.page_buf, geom_.page_size)) { + return; + } + read_page_from_disk(page_off, shard.page_buf); + shard.cache.write(page_off, + shard.page_buf, + geom_.page_size, + false, + [this, &shard](uint64_t off, const char *page) { + write_page_to_disk(shard, off, page); + }); +} + +inline void DiskPageIO::read_page_from_disk(uint64_t page_off, char *page) const { + OVERLAPPED overlapped{}; + overlapped.Offset = static_cast(page_off & 0xffffffffULL); + overlapped.OffsetHigh = static_cast(page_off >> 32U); + overlapped.hEvent = ::CreateEventW(nullptr, TRUE, FALSE, nullptr); + if (overlapped.hEvent == nullptr) { + throw std::runtime_error("DiskPageIO::read_page: CreateEventW failed: Win32 error " + + std::to_string(::GetLastError())); + } + DWORD transferred = 0; + BOOL ok = ::ReadFile(file_handle_, + page, + static_cast(geom_.page_size), + &transferred, + &overlapped); + if (!ok && ::GetLastError() == ERROR_IO_PENDING) { + ok = ::GetOverlappedResult(file_handle_, &overlapped, &transferred, TRUE); + } + const DWORD error = ok ? ERROR_SUCCESS : ::GetLastError(); + ::CloseHandle(overlapped.hEvent); + if (!ok || transferred != geom_.page_size) { + throw std::runtime_error("DiskPageIO::read_page: short/failed ReadFile at " + + std::to_string(page_off) + " (got " + std::to_string(transferred) + + ", Win32 error " + std::to_string(error) + ")"); + } } -inline void DiskPageIO::read_page_locked_off(Shard &, uint64_t) { - throw std::runtime_error("DiskPageIO: unsupported platform (needs Linux O_DIRECT)"); + +inline void DiskPageIO::write_page_locked(Shard &shard, uint32_t id) { + const uint64_t off = geom_.get_page_offset(id); + if (shard.cache.enabled()) { + shard.cache.write(off, + shard.page_buf, + geom_.page_size, + true, + [this, &shard](uint64_t page_off, const char *page) { + write_page_to_disk(shard, page_off, page); + }); + ++shard.versions[off]; + return; + } + write_page_to_disk(shard, off, shard.page_buf); + ++shard.versions[off]; } -inline void DiskPageIO::read_page_from_disk(uint64_t, char *) const { - throw std::runtime_error("DiskPageIO: unsupported platform (needs Linux O_DIRECT)"); + +inline void DiskPageIO::write_page_to_disk(Shard &shard, uint64_t page_off, const char *page) { + const char *write_buf = page; + if (page != shard.page_buf) { + std::memcpy(shard.flush_buf, page, geom_.page_size); + write_buf = shard.flush_buf; + } + OVERLAPPED overlapped{}; + overlapped.Offset = static_cast(page_off & 0xffffffffULL); + overlapped.OffsetHigh = static_cast(page_off >> 32U); + overlapped.hEvent = ::CreateEventW(nullptr, TRUE, FALSE, nullptr); + if (overlapped.hEvent == nullptr) { + throw std::runtime_error("DiskPageIO::write_page: CreateEventW failed: Win32 error " + + std::to_string(::GetLastError())); + } + DWORD transferred = 0; + BOOL ok = ::WriteFile(file_handle_, + write_buf, + static_cast(geom_.page_size), + &transferred, + &overlapped); + if (!ok && ::GetLastError() == ERROR_IO_PENDING) { + ok = ::GetOverlappedResult(file_handle_, &overlapped, &transferred, TRUE); + } + const DWORD error = ok ? ERROR_SUCCESS : ::GetLastError(); + ::CloseHandle(overlapped.hEvent); + if (!ok || transferred != geom_.page_size) { + throw std::runtime_error("DiskPageIO::write_page: short/failed WriteFile at " + + std::to_string(page_off) + " (got " + std::to_string(transferred) + + ", Win32 error " + std::to_string(error) + ")"); + } } -inline void DiskPageIO::write_page_locked(Shard &, uint32_t) { - throw std::runtime_error("DiskPageIO: unsupported platform (needs Linux O_DIRECT)"); + +inline void DiskPageIO::extend_to(uint64_t new_size) { + FILE_END_OF_FILE_INFO end_info{}; + end_info.EndOfFile.QuadPart = static_cast(new_size); + if (!::SetFileInformationByHandle(file_handle_, FileEndOfFileInfo, &end_info, sizeof(end_info))) { + throw std::runtime_error( + "DiskPageIO::extend_to: SetFileInformationByHandle failed: " + "Win32 error " + + std::to_string(::GetLastError())); + } + file_size_.store(new_size, std::memory_order_release); } -inline void DiskPageIO::write_page_to_disk(Shard &, uint64_t, const char *) { - throw std::runtime_error("DiskPageIO: unsupported platform (needs Linux O_DIRECT)"); + +inline void DiskPageIO::sync_file() { + if (!::FlushFileBuffers(file_handle_)) { + throw std::runtime_error("DiskPageIO::sync_file: FlushFileBuffers failed: Win32 error " + + std::to_string(::GetLastError())); + } } -inline void DiskPageIO::extend_to(uint64_t) { - throw std::runtime_error("DiskPageIO: unsupported platform (needs Linux O_DIRECT)"); + +inline void DiskPageIO::close_fd() { + if (file_handle_ != INVALID_HANDLE_VALUE) { + ::CloseHandle(file_handle_); + file_handle_ = INVALID_HANDLE_VALUE; + } } -inline void DiskPageIO::close_fd() {} -#endif // __linux__ +#endif } // namespace alaya::diskann diff --git a/include/index/graph/diskann/diskann_index.hpp b/include/index/graph/diskann/diskann_index.hpp index ff177a77..c4cde0f1 100644 --- a/include/index/graph/diskann/diskann_index.hpp +++ b/include/index/graph/diskann/diskann_index.hpp @@ -325,6 +325,14 @@ class DiskANNIndex { } read_ids(path(index_dir, "ids.bin"), max_slot_id_); + const std::string slots_path = path(index_dir, "slots.bin"); + if (std::filesystem::exists(slots_path)) { + slot_alloc_.load(slots_path); + max_slot_id_ = std::max(max_slot_id_, slot_alloc_.next_fresh_id()); + } else { + slot_alloc_.reset(static_cast(max_slot_id_)); + } + rebuild_label_lookup_unlocked(); cache_.load(path(index_dir, "cache_ids.bin"), path(index_dir, "cache_nodes.bin")); cache_.configure_geometry(dim_, max_degree_); if (has_pq_) { @@ -1087,7 +1095,7 @@ class DiskANNIndex { // zero set bits means the snapshot would be all-live — skip the copy. The // bitmap's capacity never shrinks, so after the first update round every // query would otherwise pay a full-capacity copy even with no tombstones. - if (updatable_ && slot_alloc_.tombstone().count() > 0) { + if (slot_alloc_.tombstone().count() > 0) { slot_alloc_.tombstone().snapshot_into(snapshot.tombstone); snapshot.has_tombstone = true; } @@ -1163,15 +1171,6 @@ class DiskANNIndex { } #endif - const std::string slots_path = path(index_dir, "slots.bin"); - if (std::filesystem::exists(slots_path)) { - slot_alloc_.load(slots_path); // restore free list + next id + tombstones - max_slot_id_ = std::max(max_slot_id_, slot_alloc_.next_fresh_id()); - resize_thread_data_slot_capacity(); - } else { - slot_alloc_.reset(static_cast(max_slot_id_)); - } - rebuild_label_lookup_unlocked(); updatable_ = true; } diff --git a/include/index/graph/laser/utils/iocp_file_reader.hpp b/include/index/graph/laser/utils/iocp_file_reader.hpp index 917dbd31..5cf5fd09 100644 --- a/include/index/graph/laser/utils/iocp_file_reader.hpp +++ b/include/index/graph/laser/utils/iocp_file_reader.hpp @@ -271,7 +271,7 @@ class IOCPFileReader : public AlignedFileReader { hFile_ = ::CreateFileW(wname.c_str(), GENERIC_READ, - FILE_SHARE_READ, + FILE_SHARE_READ | FILE_SHARE_WRITE, nullptr, OPEN_EXISTING, FILE_FLAG_NO_BUFFERING | FILE_FLAG_OVERLAPPED, diff --git a/pyproject.toml b/pyproject.toml index 11f7f7fa..9b188432 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -142,7 +142,7 @@ skip = ["*-musllinux*", "*-manylinux_i686"] # package cache warm across the per-python-version builds and across CI runs. container-engine = "docker; create_args: -v ~/.conan2:/root/.conan2" test-requires = ["pytest", "pytest-asyncio", "pyyaml"] -test-command = "python -c \"import platform, sys, alayalite; from alayalite import laser, vamana; is_linux_x86 = sys.platform.startswith('linux') and platform.machine() in {'x86_64', 'AMD64'}; is_windows_x86 = sys.platform == 'win32' and platform.machine() in {'x86_64', 'AMD64'}; laser_expected = is_linux_x86 or is_windows_x86; assert not (laser_expected and laser.RawIndex is None), 'LASER should be available on this lane'; simd = laser.selected_simd() if is_linux_x86 else 'not-tested'; print(alayalite.__version__); print(f'laser_simd={simd}', flush=True); assert (not is_linux_x86) or simd in {'avx512', 'avx2'}\"" +test-command = "python -c \"import platform, sys, alayalite; from alayalite import laser, vamana; is_linux_x86 = sys.platform.startswith('linux') and platform.machine() in {'x86_64', 'AMD64'}; is_windows_x86 = sys.platform == 'win32' and platform.machine() in {'x86_64', 'AMD64'}; laser_expected = is_linux_x86 or is_windows_x86; assert not (laser_expected and laser.RawIndex is None), 'LASER should be available on this lane'; simd = laser.selected_simd() if is_linux_x86 else 'not-tested'; print(alayalite.__version__); print(f'laser_simd={simd}', flush=True); assert (not is_linux_x86) or simd in {'avx512', 'avx2'}\" && python -m pytest \"{project}/python/tests/diskann/test_update.py\"" [tool.cibuildwheel.linux] manylinux-x86_64-image = "manylinux_2_28" diff --git a/python/tests/diskann/test_update.py b/python/tests/diskann/test_update.py index 4014b30b..4f4c936c 100644 --- a/python/tests/diskann/test_update.py +++ b/python/tests/diskann/test_update.py @@ -1,13 +1,9 @@ """Tests for external-ID based DiskANN update operations.""" -import sys - import numpy as np import pytest from alayalite import diskann -pytestmark = pytest.mark.skipif(sys.platform != "linux", reason="DiskANN updates require Linux") - def _build_index(path): rng = np.random.default_rng(42) @@ -86,3 +82,33 @@ def test_update_arrays_require_exact_dtype_and_layout(tmp_path): index.insert(np.ones(16, dtype=np.float64), 8000) with pytest.raises(TypeError, match="uint64"): index.batch_remove(np.array([1001], dtype=np.int64)) + + +def test_default_open_uses_portable_update_backend(tmp_path): + path = tmp_path / "index" + built_index, _ = _build_index(path) + del built_index + + index = diskann.Index.open(str(path)) + assert index.updatable + vector = np.random.default_rng(10).random(16, dtype=np.float32) + index.insert(vector, 9000) + assert index.contains(9000) + index.remove(9000) + index.flush() + + +def test_read_only_reopen_restores_external_id_mapping_and_deletions(tmp_path): + path = tmp_path / "index" + index, vectors = _build_index(path) + index.remove(1001) + index.flush() + del index + + params = diskann.LoadParams() + params.updatable = False + reopened = diskann.Index.open(str(path), params) + assert reopened.contains(1000) + assert not reopened.contains(1001) + labels, _ = reopened.search(vectors[1], 10) + assert 1001 not in labels diff --git a/tests/diskann/CMakeLists.txt b/tests/diskann/CMakeLists.txt index 7b096418..7cec0e2c 100644 --- a/tests/diskann/CMakeLists.txt +++ b/tests/diskann/CMakeLists.txt @@ -33,6 +33,11 @@ alaya_cc_target( GTEST SRCS test_diskann_update_trace.cpp ) +alaya_cc_target( + test_diskann_portable_update + GTEST + SRCS test_diskann_portable_update.cpp +) alaya_add_test( NAME test_diskann_layout @@ -59,6 +64,11 @@ alaya_add_test( TARGET test_diskann_update_trace LABELS diskann ) +alaya_add_test( + NAME test_diskann_portable_update + TARGET test_diskann_portable_update + LABELS diskann +) if(ALAYA_ENABLE_LASER) alaya_cc_target( diff --git a/tests/diskann/test_diskann_portable_update.cpp b/tests/diskann/test_diskann_portable_update.cpp new file mode 100644 index 00000000..dffd8b4e --- /dev/null +++ b/tests/diskann/test_diskann_portable_update.cpp @@ -0,0 +1,81 @@ +// SPDX-FileCopyrightText: 2025 AlayaDB.AI +// +// SPDX-License-Identifier: AGPL-3.0-only + +#include "index/graph/diskann/diskann_index.hpp" + +#include + +#include +#include +#include +#include +#include +#include + +namespace { + +using alaya::diskann::DiskANNBuildParams; +using alaya::diskann::DiskANNIndex; +using alaya::diskann::DiskANNLoadParams; +using alaya::diskann::DiskANNUpdateIO; + +class DiskANNPortableUpdateTest : public ::testing::Test { + protected: + void SetUp() override { + static std::atomic counter{0}; + dir_ = std::filesystem::temp_directory_path() / + ("diskann_portable_update_" + std::to_string(counter.fetch_add(1))); + vectors_.resize(kCount * kDim); + std::mt19937 rng(42); + std::uniform_real_distribution distribution(-1.0f, 1.0f); + for (float &value : vectors_) { + value = distribution(rng); + } + labels_.resize(kCount); + std::iota(labels_.begin(), labels_.end(), uint64_t{1000}); + + DiskANNBuildParams build_params; + build_params.R = 16; + build_params.L = 32; + DiskANNIndex::build(dir_.string(), vectors_.data(), labels_.data(), kCount, kDim, build_params); + } + + void TearDown() override { + std::error_code error; + std::filesystem::remove_all(dir_, error); + } + + static constexpr uint64_t kCount = 96; + static constexpr uint64_t kDim = 16; + std::filesystem::path dir_; + std::vector vectors_; + std::vector labels_; +}; + +TEST_F(DiskANNPortableUpdateTest, ExternalIdUpdatesPersistAcrossWritableAndReadOnlyReloads) { + DiskANNLoadParams update_params; + update_params.updatable = true; + update_params.update_io = DiskANNUpdateIO::kBlocking; + update_params.num_threads = 2; + update_params.update_insert_threads = 2; + update_params.update_reconnect_threads = 2; + + { + DiskANNIndex index; + index.load(dir_.string(), update_params); + const std::vector inserted(vectors_.begin(), vectors_.begin() + kDim); + index.insert(inserted.data(), 9000); + EXPECT_TRUE(index.contains_label(9000)); + index.remove_by_label(1001); + EXPECT_FALSE(index.contains_label(1001)); + index.flush(); + } + + DiskANNIndex read_only; + read_only.load(dir_.string()); + EXPECT_TRUE(read_only.contains_label(9000)); + EXPECT_FALSE(read_only.contains_label(1001)); +} + +} // namespace diff --git a/tests/diskann/test_diskann_update_e2e.cpp b/tests/diskann/test_diskann_update_e2e.cpp index 2f16b95e..11b3873c 100644 --- a/tests/diskann/test_diskann_update_e2e.cpp +++ b/tests/diskann/test_diskann_update_e2e.cpp @@ -88,11 +88,13 @@ class UpdateE2ETest : public ::testing::Test { lp.updatable = true; // ALAYA_DISKANN_UPDATE_IO=blocking|uring|auto lets CI exercise both update // I/O backends with the same suite (default auto = uring when available). - if (const char *mode = std::getenv("ALAYA_DISKANN_UPDATE_IO")) { - if (std::string_view(mode) == "blocking") { - lp.update_io = alaya::diskann::DiskANNUpdateIO::kBlocking; - } else if (std::string_view(mode) == "uring") { - lp.update_io = alaya::diskann::DiskANNUpdateIO::kUring; + if (lp.update_io == alaya::diskann::DiskANNUpdateIO::kAuto) { + if (const char *mode = std::getenv("ALAYA_DISKANN_UPDATE_IO")) { + if (std::string_view(mode) == "blocking") { + lp.update_io = alaya::diskann::DiskANNUpdateIO::kBlocking; + } else if (std::string_view(mode) == "uring") { + lp.update_io = alaya::diskann::DiskANNUpdateIO::kUring; + } } } idx_ = std::make_unique(); @@ -504,8 +506,34 @@ TEST_F(UpdateE2ETest, PQSearchSkipsDeletedLabels) { } } +TEST_F(UpdateE2ETest, PipelinedSearchRequiresUring) { + DiskANNLoadParams load_params; + load_params.update_io = alaya::diskann::DiskANNUpdateIO::kBlocking; + build_and_load(/*n=*/100, /*dim=*/16, /*r=*/16, load_params); + + const auto queries = make_vectors(/*n=*/2, /*dim=*/16, /*seed=*/7); + std::vector labels(4); + std::vector distances(4); + DiskANNSearchParams search_params; + search_params.rerank = false; + EXPECT_THROW(idx_->search_pipelined(queries.data(), + /*n_queries=*/2, + /*top_k=*/2, + labels.data(), + distances.data(), + /*num_threads=*/1, + /*pipeline=*/2, + search_params), + std::runtime_error); +} + TEST_F(UpdateE2ETest, PipelinedSearchMatchesBatchSearch) { - build_and_load(/*n=*/600, /*dim=*/32, /*r=*/32, {}, /*pq_n_chunks=*/8); + if (!alaya::UringReactor::is_available()) { + GTEST_SKIP() << "io_uring not available on this kernel"; + } + DiskANNLoadParams load_params; + load_params.update_io = alaya::diskann::DiskANNUpdateIO::kUring; + build_and_load(/*n=*/600, /*dim=*/32, /*r=*/32, load_params, /*pq_n_chunks=*/8); ASSERT_TRUE(idx_->has_pq()); // Post-update state: labels beyond the base range plus live tombstones. @@ -538,19 +566,16 @@ TEST_F(UpdateE2ETest, PipelinedSearchMatchesBatchSearch) { std::vector ref_l(kNq * kK); std::vector ref_d(kNq * kK); - idx_->batch_search( - queries.data(), kNq, kK, ref_l.data(), ref_d.data(), /*num_threads=*/2, ref_sp); + idx_->batch_search(queries.data(), + kNq, + kK, + ref_l.data(), + ref_d.data(), + /*num_threads=*/2, + ref_sp); std::vector pipe_l(kNq * kK, 0); std::vector pipe_d(kNq * kK, 0.0F); - const char *mode = std::getenv("ALAYA_DISKANN_UPDATE_IO"); - if (mode != nullptr && std::string_view(mode) == "blocking") { - // No reactor in blocking mode: the pipelined path must refuse loudly. - EXPECT_THROW(idx_->search_pipelined( - queries.data(), kNq, kK, pipe_l.data(), pipe_d.data(), 2, 8, sp), - std::runtime_error); - return; - } idx_->search_pipelined(queries.data(), kNq, kK,