diff --git a/CMakeLists.txt b/CMakeLists.txt index 956ee5599..a58a334b6 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -122,14 +122,18 @@ else() endif() message(STATUS "RABITQ_ARCH_FLAG: ${RABITQ_ARCH_FLAG}") -# DiskAnn support (Linux x86_64 only; libaio loaded at runtime via dlopen) -if(CMAKE_SYSTEM_NAME STREQUAL "Linux" AND CMAKE_SYSTEM_PROCESSOR MATCHES "x86_64|i686|i386" AND NOT ANDROID AND NOT IOS) +# DiskAnn support: +# - Linux (x86_64, i686, i386, aarch64, arm64), using libaio when available +# and synchronous pread otherwise +# - macOS (x86_64, ARM64/Apple Silicon), using POSIX AIO with aio_suspend +if((CMAKE_SYSTEM_NAME STREQUAL "Linux" AND CMAKE_SYSTEM_PROCESSOR MATCHES "x86_64|i686|i386|aarch64|arm64" AND NOT ANDROID AND NOT IOS) + OR (CMAKE_SYSTEM_NAME STREQUAL "Darwin" AND NOT IOS)) set(DISKANN_SUPPORTED ON) add_definitions(-DDISKANN_SUPPORTED=1) else() set(DISKANN_SUPPORTED OFF) add_definitions(-DDISKANN_SUPPORTED=0) - message(STATUS "DiskAnn support disabled - only supported on Linux x86_64") + message(STATUS "DiskAnn support disabled - supported on Linux x86/x86_64/ARM64 and macOS") endif() message(STATUS "DISKANN_SUPPORTED: ${DISKANN_SUPPORTED}") diff --git a/examples/c/diskann_example.c b/examples/c/diskann_example.c index 5ef80562f..facabbeb1 100644 --- a/examples/c/diskann_example.c +++ b/examples/c/diskann_example.c @@ -21,8 +21,10 @@ * a Vamana graph structure combined with product quantization (PQ) to * achieve high recall with efficient disk I/O. * - * NOTE: DiskANN requires Linux x86_64 with libaio. On other platforms the - * example will compile but the runtime plugin will fail to load. + * NOTE: DiskANN supports Linux x86/x86_64/ARM64 (libaio when available, + * synchronous pread fallback otherwise) and macOS (POSIX AIO with + * aio_suspend()). + * Schema validation rejects DiskANN on other platforms. * * Workflow demonstrated: * 1. Create collection schema with DiskANN-indexed vector field @@ -265,9 +267,6 @@ int main(void) { &results, &result_count); if (error != ZVEC_OK) { handle_error(error, "executing DiskANN query"); - printf( - " (This is expected on non-Linux platforms — DiskANN requires " - "libaio)\n"); } else { printf(" Query returned %zu results:\n", result_count); for (size_t r = 0; r < result_count && r < 5; r++) { diff --git a/python/tests/detail/fixture_helper.py b/python/tests/detail/fixture_helper.py index f5e6c08e4..2aed7bf66 100644 --- a/python/tests/detail/fixture_helper.py +++ b/python/tests/detail/fixture_helper.py @@ -2,11 +2,10 @@ import logging import platform -DISKANN_SUPPORTED = platform.system() == "Linux" and platform.machine() in ( - "x86_64", - "AMD64", - "i686", - "i386", +_MACHINE = platform.machine().lower() +DISKANN_SUPPORTED = platform.system() == "Darwin" or ( + platform.system() == "Linux" + and _MACHINE in ("x86_64", "amd64", "i686", "i386", "aarch64", "arm64") ) from typing import Any, Generator @@ -27,7 +26,9 @@ def _ensure_diskann_runtime_or_reason() -> str | None: _DISKANN_PRELOAD_DONE = True if not DISKANN_SUPPORTED: - _DISKANN_PRELOAD_REASON = "DiskAnn only supported on Linux x86_64" + _DISKANN_PRELOAD_REASON = ( + "DiskAnn is supported on Linux x86/x86_64/ARM64 and macOS" + ) return _DISKANN_PRELOAD_REASON _DISKANN_PRELOAD_REASON = None diff --git a/python/tests/test_collection_diskann.py b/python/tests/test_collection_diskann.py index 8daf8c770..b948bb27f 100644 --- a/python/tests/test_collection_diskann.py +++ b/python/tests/test_collection_diskann.py @@ -15,16 +15,13 @@ Mirrors ``test_collection_hnsw_rabitq.py`` but targets the DiskAnn index. -Two platform-level prerequisites are enforced at module import time: +Platform behavior: -1. DiskAnn is currently built only for Linux x86_64 — other platforms are - skipped wholesale. -2. libaio is loaded eagerly (via dlopen) inside DiskAnnBuilder::init() / - DiskAnnStreamer::init(). If libaio is missing, DiskAnn falls back to - synchronous pread() — the tests still run but with degraded performance. +1. Linux x86/x86_64/ARM64 probes libaio at runtime and falls back to + synchronous pread() when libaio is unavailable. +2. macOS uses the system POSIX AIO implementation with aio_suspend() waits. -If either prerequisite fails the whole module is skipped so the rest of the -test-suite is not affected. +The module is skipped on platforms where DiskAnn is not built. """ from __future__ import annotations @@ -38,9 +35,14 @@ # --------------------------------------------------------------------------- # # Platform gating (must happen BEFORE we touch zvec). # --------------------------------------------------------------------------- # +_machine = platform.machine().lower() +_diskann_supported = sys.platform == "darwin" or ( + sys.platform == "linux" + and _machine in ("x86_64", "amd64", "i686", "i386", "aarch64", "arm64") +) pytestmark = pytest.mark.skipif( - not (sys.platform == "linux" and platform.machine() in ("x86_64", "AMD64")), - reason="DiskAnn plugin is only supported on Linux x86_64", + not _diskann_supported, + reason="DiskAnn is supported on Linux x86/x86_64/ARM64 and macOS", ) import zvec # noqa: E402 diff --git a/python/tests/test_typing.py b/python/tests/test_typing.py index d566d5efc..ab821f66c 100644 --- a/python/tests/test_typing.py +++ b/python/tests/test_typing.py @@ -34,6 +34,7 @@ (DataType.FLOAT, "FLOAT"), (IndexType.HNSW, "HNSW"), (IOBackendType.PREAD, "PREAD"), + (IOBackendType.POSIX_AIO, "POSIX_AIO"), (MetricType.COSINE, "COSINE"), (QuantizeType.INT8, "INT8"), (StatusCode.OK, "OK"), @@ -49,6 +50,7 @@ def test_enum_names(member, name): (DataType.FLOAT, 8), (IndexType.HNSW, 1), (IOBackendType.PREAD, 0), + (IOBackendType.POSIX_AIO, 2), (MetricType.COSINE, 3), (QuantizeType.INT8, 2), (StatusCode.OK, 0), @@ -112,7 +114,7 @@ def test_index_type_has_member(member): assert member in IndexType.__members__ -@pytest.mark.parametrize("member", ["PREAD", "LIBAIO"]) +@pytest.mark.parametrize("member", ["PREAD", "LIBAIO", "POSIX_AIO"]) def test_io_backend_type_has_member(member): assert member in IOBackendType.__members__ diff --git a/python/zvec/__init__.pyi b/python/zvec/__init__.pyi index 09e828dae..e31b68da1 100644 --- a/python/zvec/__init__.pyi +++ b/python/zvec/__init__.pyi @@ -52,6 +52,7 @@ from .zvec import create_and_open, init, open def io_backend_type() -> IOBackendType: """Returns the current I/O backend type for DiskAnn async disk reads as an IOBackendType enum (zvec.typing.IOBackendType). + On macOS this is IOBackendType.POSIX_AIO. On Linux this is IOBackendType.LIBAIO if libaio is available, IOBackendType.PREAD otherwise.""" def io_backend_description() -> str: diff --git a/python/zvec/typing/__init__.pyi b/python/zvec/typing/__init__.pyi index fab03f4b9..0e969a7a7 100644 --- a/python/zvec/typing/__init__.pyi +++ b/python/zvec/typing/__init__.pyi @@ -130,6 +130,7 @@ class IOBackendType: - PREAD: Synchronous pread() — no async I/O. - LIBAIO: libaio loaded at runtime via dlopen(). + - POSIX_AIO: macOS POSIX AIO with aio_suspend() completion waits. Examples: >>> from zvec.typing import IOBackendType @@ -142,13 +143,16 @@ class IOBackendType: PREAD LIBAIO + + POSIX_AIO """ LIBAIO: typing.ClassVar[IOBackendType] # value = + POSIX_AIO: typing.ClassVar[IOBackendType] # value = PREAD: typing.ClassVar[IOBackendType] # value = __members__: typing.ClassVar[ dict[str, IOBackendType] - ] # value = {'PREAD': , 'LIBAIO': } + ] # value = {'PREAD': , 'LIBAIO': , 'POSIX_AIO': } def __eq__(self, other: typing.Any) -> bool: ... def __getstate__(self) -> int: ... diff --git a/src/ailego/io/io_backend_def.h b/src/ailego/io/io_backend_def.h index d795c3046..4155066e5 100644 --- a/src/ailego/io/io_backend_def.h +++ b/src/ailego/io/io_backend_def.h @@ -32,6 +32,10 @@ #include #include +#if defined(__APPLE__) +#include +#endif + namespace zvec { namespace ailego { @@ -40,6 +44,8 @@ inline const char *IOBackendTypeName(IOBackendType type) { switch (type) { case IOBackendType::kLibAio: return "libaio"; + case IOBackendType::kPosixAio: + return "posix_aio"; case IOBackendType::kPread: return "pread"; } @@ -47,24 +53,31 @@ inline const char *IOBackendTypeName(IOBackendType type) { } // Returns a human-readable description for the given backend type. -// When the backend is kPread, includes installation guidance for libaio. +// On Linux, the kPread description includes installation guidance for libaio. inline const char *IOBackendDescription(IOBackendType type) { switch (type) { case IOBackendType::kLibAio: return "libaio async I/O backend loaded at runtime via dlopen()."; + case IOBackendType::kPosixAio: + return "macOS POSIX AIO backend with aio_suspend() completion waits."; case IOBackendType::kPread: +#if defined(__linux) || defined(__linux__) return "No async I/O backend available. Install libaio (e.g. " "'apt-get install libaio1', or 'libaio1t64' on Ubuntu 24.04+) " "and retry. DiskAnn will fall back to synchronous pread() \u2014 " "performance will be degraded."; +#else + return "Synchronous pread() backend; no async I/O backend is " + "available."; +#endif } return "Unknown I/O backend."; } // Singleton that loads and queries an I/O backend on demand. // -// available() (no arg) tries the best backend with priority (libaio > pread) -// and returns the loaded backend type. +// available() (no arg) returns the platform backend. On Linux it tries libaio +// before falling back to pread; on macOS POSIX AIO is provided by the system. // available(IOBackendType) tries a specific backend. // Use type() / name() to query the loaded backend without triggering a load. class IOBackend { @@ -74,8 +87,8 @@ class IOBackend { return instance; } - // Try to load the best available backend (libaio > pread). - // Returns the loaded backend type. + // Return the selected platform backend. Linux probes libaio before falling + // back to pread; macOS uses the system POSIX AIO backend. // Idempotent — if already loaded, returns immediately. IOBackendType available() { if (type_ != IOBackendType::kPread) { @@ -84,13 +97,19 @@ class IOBackend { return available(IOBackendType::kLibAio); } - // Try to load the requested backend. Returns the loaded backend type - // (may differ from requested if the load failed — falls back to kPread). + // Try to load the requested backend. Returns the selected backend type, + // which may differ from the request when it is unavailable on this platform. // Idempotent — if the same backend is already loaded, returns immediately. IOBackendType available(IOBackendType requested) { if (type_ == requested && type_ != IOBackendType::kPread) { return type_; } +#if defined(__APPLE__) && TARGET_OS_OSX + // POSIX AIO is provided by Darwin; no user-installed runtime dependency + // needs to be probed. + type_ = IOBackendType::kPosixAio; + return type_; +#endif #if defined(__linux) || defined(__linux__) if (requested == IOBackendType::kLibAio) { if (LibAioLoader::Instance().load() && @@ -112,6 +131,10 @@ class IOBackend { return available() == IOBackendType::kLibAio; } + bool is_posix_aio() { + return available() == IOBackendType::kPosixAio; + } + // Returns the loaded backend type. IOBackendType type() const { return type_; @@ -130,7 +153,13 @@ class IOBackend { private: IOBackend() = default; - IOBackendType type_{IOBackendType::kPread}; + IOBackendType type_{ +#if defined(__APPLE__) && TARGET_OS_OSX + IOBackendType::kPosixAio +#else + IOBackendType::kPread +#endif + }; }; } // namespace ailego diff --git a/src/binding/python/model/common/python_config.cc b/src/binding/python/model/common/python_config.cc index d6ad1f48b..3f9f073fc 100644 --- a/src/binding/python/model/common/python_config.cc +++ b/src/binding/python/model/common/python_config.cc @@ -228,7 +228,8 @@ void ZVecPyConfig::Initialize(pybind11::module_ &m) { }, "Returns the current I/O backend type for DiskAnn async disk reads " "as an IOBackendType enum (zvec.typing.IOBackendType). " - "IOBackendType.LIBAIO if libaio is available, " + "On macOS this is IOBackendType.POSIX_AIO; on Linux this is " + "IOBackendType.LIBAIO when libaio is available and " "IOBackendType.PREAD otherwise."); // Returns a human-readable description of the I/O backend, including @@ -242,4 +243,4 @@ void ZVecPyConfig::Initialize(pybind11::module_ &m) { } -} // namespace zvec \ No newline at end of file +} // namespace zvec diff --git a/src/binding/python/typing/python_type.cc b/src/binding/python/typing/python_type.cc index d20cac755..58a260b72 100644 --- a/src/binding/python/typing/python_type.cc +++ b/src/binding/python/typing/python_type.cc @@ -146,6 +146,7 @@ Enumeration of supported I/O backend types for DiskAnn async disk reads. - PREAD: Synchronous pread() \u2014 no async I/O. - LIBAIO: libaio loaded at runtime via dlopen(). +- POSIX_AIO: macOS POSIX AIO with aio_suspend() completion waits. Examples: >>> from zvec.typing import IOBackendType @@ -153,7 +154,8 @@ Enumeration of supported I/O backend types for DiskAnn async disk reads. IOBackendType.LIBAIO )pbdoc") .value("PREAD", ailego::IOBackendType::kPread) - .value("LIBAIO", ailego::IOBackendType::kLibAio); + .value("LIBAIO", ailego::IOBackendType::kLibAio) + .value("POSIX_AIO", ailego::IOBackendType::kPosixAio); } void ZVecPyTyping::bind_status(py::module_ &m) { @@ -245,4 +247,4 @@ Construct a status with the given code and optional message. }); } -} // namespace zvec \ No newline at end of file +} // namespace zvec diff --git a/src/core/CMakeLists.txt b/src/core/CMakeLists.txt index 1441b1d98..f29596ce3 100644 --- a/src/core/CMakeLists.txt +++ b/src/core/CMakeLists.txt @@ -74,7 +74,8 @@ if(NOT DISKANN_SUPPORTED) endif() set(ZVEC_CORE_LIBS zvec_ailego zvec_turbo sparsehash magic_enum rabitqlib) -# The DiskAnn runtime loader uses dlopen/dlsym, so link libdl on Linux. +# The DiskAnn runtime loader (LibAioLoader) uses dlopen/dlsym, so link libdl +# on Linux. On macOS, POSIX AIO is provided by the system. if(CMAKE_SYSTEM_NAME STREQUAL "Linux") list(APPEND ZVEC_CORE_LIBS ${CMAKE_DL_LIBS}) endif() diff --git a/src/core/algorithm/CMakeLists.txt b/src/core/algorithm/CMakeLists.txt index f874eba62..5d5654a9a 100644 --- a/src/core/algorithm/CMakeLists.txt +++ b/src/core/algorithm/CMakeLists.txt @@ -17,7 +17,7 @@ else() # Empty stub library for unsupported platforms file(WRITE ${CMAKE_CURRENT_BINARY_DIR}/diskann_stub.cc "// Stub implementation for unsupported platforms\n" - "// DiskAnn only supports Linux x86_64\n" + "// DiskAnn supports Linux x86/x86_64/ARM64 and macOS\n" "namespace zvec { namespace core { /* empty namespace for compatibility */ } }\n" ) diff --git a/src/core/algorithm/diskann/CMakeLists.txt b/src/core/algorithm/diskann/CMakeLists.txt index 5f1b7bf7f..bb1746c63 100644 --- a/src/core/algorithm/diskann/CMakeLists.txt +++ b/src/core/algorithm/diskann/CMakeLists.txt @@ -14,7 +14,7 @@ file(GLOB_RECURSE ALL_SRCS *.cc *.c) # ${CMAKE_DL_LIBS} for dlopen/dlsym/dlclose. set(CORE_KNN_DISKANN_LIBS core_framework core_knn_cluster) -if(CMAKE_SYSTEM_NAME STREQUAL "Linux" AND CMAKE_SYSTEM_PROCESSOR MATCHES "x86_64|i686|i386") +if(CMAKE_SYSTEM_NAME STREQUAL "Linux" AND CMAKE_SYSTEM_PROCESSOR MATCHES "x86_64|i686|i386|aarch64|arm64") list(APPEND CORE_KNN_DISKANN_LIBS ${CMAKE_DL_LIBS}) endif() diff --git a/src/core/algorithm/diskann/diskann_builder.cc b/src/core/algorithm/diskann/diskann_builder.cc index b3597f1c5..4ab837b16 100644 --- a/src/core/algorithm/diskann/diskann_builder.cc +++ b/src/core/algorithm/diskann/diskann_builder.cc @@ -17,6 +17,7 @@ #include #include #include +#include #include #include #include @@ -32,7 +33,22 @@ namespace core { int DiskAnnBuilder::init(const IndexMeta &meta, const ailego::Params ¶ms) { LOG_INFO("Begin DiskAnnBuilder::init"); - log_diskann_io_backend(); +#if defined(__linux__) || defined(__linux) + // Eagerly probe the I/O backend at init time so the user gets immediate + // feedback about whether async I/O is available. IOBackend auto-probes on + // first Instance() access. + auto &backend = ailego::IOBackend::Instance(); + if (backend.available() == ailego::IOBackendType::kPread) { + LOG_WARN( + "DiskAnn: no async I/O backend available. Install libaio (e.g. " + "'apt-get install libaio1', or 'libaio1t64' on Ubuntu 24.04+) and " + "retry. DiskAnn will fall back to synchronous pread() — performance " + "will be degraded."); + } else { + LOG_INFO("DiskAnn: I/O backend '%s' loaded — async I/O enabled.", + backend.name()); + } +#endif params.get(PARAM_DISKANN_BUILDER_MAX_DEGREE, &max_degree_); params.get(PARAM_DISKANN_BUILDER_LIST_SIZE, &list_size_); @@ -271,24 +287,25 @@ int DiskAnnBuilder::build_internal(IndexThreads::Pointer threads) { { std::unique_lock lk(mutex_); - while (finished.load() < entity_.doc_cnt()) { + while (finished.load() < entity_.doc_cnt() && + !error_.load(std::memory_order_acquire)) { cond_.wait_until(lk, std::chrono::system_clock::now() + std::chrono::seconds(check_interval_secs_)); - if (error_.load(std::memory_order_acquire)) { - LOG_ERROR("Failed to build index while waiting finish"); - return errcode_; - } LOG_INFO("Built cnt %zu, finished percent %.3f%%", (size_t)finished.load(), finished.load() * 100.0f / entity_.doc_cnt()); } } + // A worker may fail before processing its entire shard. Always wait for the + // remaining workers before returning so they cannot access this builder, + // its mutex, or the stack-owned finished counter after their lifetimes end. + task_group->wait_finish(); + if (error_.load(std::memory_order_acquire)) { LOG_ERROR("Failed to build index while waiting finish"); return errcode_; } - task_group->wait_finish(); return 0; } @@ -308,24 +325,24 @@ int DiskAnnBuilder::prune_internal(IndexThreads::Pointer threads) { { std::unique_lock lk(mutex_); - while (finished.load() < entity_.doc_cnt()) { + while (finished.load() < entity_.doc_cnt() && + !error_.load(std::memory_order_acquire)) { cond_.wait_until(lk, std::chrono::system_clock::now() + std::chrono::seconds(check_interval_secs_)); - if (error_.load(std::memory_order_acquire)) { - LOG_ERROR("Failed to prune index while waiting finish"); - return errcode_; - } LOG_INFO("Prune cnt %zu, finished percent %.3f%%", (size_t)finished.load(), finished.load() * 100.0f / entity_.doc_cnt()); } } + // See build_internal(): error paths must not let worker closures outlive the + // builder or the stack-owned completion counter. + task_group->wait_finish(); + if (error_.load(std::memory_order_acquire)) { LOG_ERROR("Failed to prune index while waiting finish"); return errcode_; } - task_group->wait_finish(); return 0; } diff --git a/src/core/algorithm/diskann/diskann_builder_entity.cc b/src/core/algorithm/diskann/diskann_builder_entity.cc index 4381293a5..992372a9a 100644 --- a/src/core/algorithm/diskann/diskann_builder_entity.cc +++ b/src/core/algorithm/diskann/diskann_builder_entity.cc @@ -53,7 +53,7 @@ int DiskAnnBuilderEntity::add_vector(diskann_key_t key, const void *vec) { keys_buffer_.append(reinterpret_cast(&key), sizeof(key)); uint32_t neighbor_cnt = 0; - std::vector neighbor{max_build_degree_, 0}; + std::vector neighbor(max_build_degree_, 0); neighbors_buffer_.append(reinterpret_cast(&neighbor_cnt), sizeof(uint32_t)); diff --git a/src/core/algorithm/diskann/diskann_context.h b/src/core/algorithm/diskann/diskann_context.h index dd824ff23..1f4bc75ac 100644 --- a/src/core/algorithm/diskann/diskann_context.h +++ b/src/core/algorithm/diskann/diskann_context.h @@ -349,7 +349,13 @@ class DiskAnnContext : public IndexContext, uint32_t group_num_{0}; std::map group_topk_heaps_{}; - IOContext io_ctx_{0}; + IOContext io_ctx_{ +#if defined(__APPLE__) && TARGET_OS_OSX + -1 +#else + 0 +#endif + }; SearchStats query_stats_; float *pq_table_dist_buffer_{nullptr}; diff --git a/src/core/algorithm/diskann/diskann_entity.h b/src/core/algorithm/diskann/diskann_entity.h index af302290d..cdaeaed33 100644 --- a/src/core/algorithm/diskann/diskann_entity.h +++ b/src/core/algorithm/diskann/diskann_entity.h @@ -72,17 +72,6 @@ struct DiskAnnMetaHeader { clear(); } - DiskAnnMetaHeader(const DiskAnnMetaHeader &header) { - memcpy(this, &header, sizeof(header)); - } - - DiskAnnMetaHeader &operator=(const DiskAnnMetaHeader &header) { - if (this != &header) { - memcpy(this, &header, sizeof(header)); - } - return *this; - } - void reset() { doc_cnt = 0U; } @@ -104,15 +93,6 @@ struct DiskAnnPqMeta { clear(); } - DiskAnnPqMeta(const DiskAnnPqMeta &meta) { - memcpy(this, &meta, sizeof(meta)); - } - - DiskAnnPqMeta &operator=(const DiskAnnPqMeta &meta) { - memcpy(this, &meta, sizeof(meta)); - return *this; - } - void clear() { memset(this, 0, sizeof(DiskAnnPqMeta)); } diff --git a/src/core/algorithm/diskann/diskann_file_reader.cc b/src/core/algorithm/diskann/diskann_file_reader.cc index 85be8c334..813f645fb 100644 --- a/src/core/algorithm/diskann/diskann_file_reader.cc +++ b/src/core/algorithm/diskann/diskann_file_reader.cc @@ -13,14 +13,15 @@ // limitations under the License. #include "diskann_file_reader.h" -#include +#include +#include #include -#include #include -#include #include -#include #include +#if defined(__APPLE__) && TARGET_OS_OSX +#include +#endif #define MAX_EVENTS 1024 @@ -35,33 +36,34 @@ typedef struct iocb iocb_t; // regardless of which entry point (setup_io_ctx or register_thread) // triggers it first. static std::once_flag g_io_backend_log_once; +#elif defined(__APPLE__) && TARGET_OS_OSX +static std::once_flag g_io_backend_log_once; #endif -void log_diskann_io_backend() { -#if (defined(__linux) || defined(__linux__)) - auto &backend = ailego::IOBackend::Instance(); - if (backend.is_pread()) { - LOG_WARN( - "DiskAnn: no async I/O backend available. Install libaio (e.g. " - "'apt-get install libaio1', or 'libaio1t64' on Ubuntu 24.04+) and " - "retry. DiskAnn will fall back to synchronous pread() — performance " - "will be degraded."); - } else { - LOG_INFO("DiskAnn: I/O backend '%s' loaded — async I/O enabled.", - backend.name()); - } -#endif -} - int setup_io_ctx(IOContext &ctx) { #if (defined(__linux) || defined(__linux__)) - std::call_once(g_io_backend_log_once, log_diskann_io_backend); - if (ailego::IOBackend::Instance().is_pread()) { + auto &backend = ailego::IOBackend::Instance(); + std::call_once(g_io_backend_log_once, [&backend] { + if (backend.available() != ailego::IOBackendType::kPread) { + LOG_INFO("DiskAnn I/O backend: %s (async I/O enabled)", backend.name()); + } else { + LOG_WARN( + "DiskAnn I/O backend: synchronous pread (no async I/O available)"); + } + }); + if (backend.available() == ailego::IOBackendType::kPread) { return 0; } int ret = LibAioLoader::Instance().io_setup(MAX_EVENTS, &ctx); - return ret; +#elif defined(__APPLE__) && TARGET_OS_OSX + std::call_once(g_io_backend_log_once, [] { + LOG_INFO("DiskAnn I/O backend: macOS POSIX AIO with aio_suspend"); + }); + // POSIX AIO does not require a persistent per-context kernel object on + // macOS. Keep a placeholder value for the cross-platform IOContext API. + ctx = 0; + return 0; #else return 0; #endif @@ -69,18 +71,25 @@ int setup_io_ctx(IOContext &ctx) { int destroy_io_ctx(IOContext &ctx) { #if (defined(__linux) || defined(__linux__)) - if (ailego::IOBackend::Instance().is_pread()) { + if (ailego::IOBackend::Instance().available() == + ailego::IOBackendType::kPread) { return 0; } int ret = LibAioLoader::Instance().io_destroy(ctx); - return ret; +#elif defined(__APPLE__) && TARGET_OS_OSX + ctx = 0; + return 0; #else return 0; #endif } -static int execute_io_pread(int fd, std::vector &read_reqs) { +static int execute_io_pread(int fd, std::vector &read_reqs, + ailego::IOBackendType *used_backend = nullptr) { + if (used_backend != nullptr) { + *used_backend = ailego::IOBackendType::kPread; + } for (auto &req : read_reqs) { ssize_t bytes_read = ::pread(fd, req.buf, req.len, req.offset); if (bytes_read < 0) { @@ -98,11 +107,149 @@ static int execute_io_pread(int fd, std::vector &read_reqs) { return 0; } -int execute_io(IOContext ctx, int fd, std::vector &read_reqs, - uint64_t n_retries = 0) { +#if defined(__APPLE__) && TARGET_OS_OSX +// POSIX AIO on Darwin accepts a maximum of AIO_LISTIO_MAX requests in the +// interfaces used to wait for a batch. Keeping the same bound here also avoids +// exhausting the process-wide AIO request limit when several search threads +// run concurrently. +static constexpr size_t kMacAioBatchSize = AIO_LISTIO_MAX; + +static bool validate_aio_result(int aio_err, int64_t result, + const AlignedRead &req) { + if (aio_err == 0 && result == static_cast(req.len)) { + return true; + } + + LOG_WARN( + "macOS aio request failed; aio_error=%d, %s, result=%ld, " + "expected=%lu, offset=%lu", + aio_err, aio_err == 0 ? "success" : ::strerror(aio_err), (long)result, + (unsigned long)req.len, (unsigned long)req.offset); + return false; +} + +// Wait for and reap every submitted request. Requests and their destination +// buffers must remain alive until aio_error() no longer reports EINPROGRESS, +// and each completed request must be reaped exactly once with aio_return(). +static bool drain_aio_requests(std::vector &cbs, + const std::vector &read_reqs, + size_t req_begin, size_t submitted_count) { + size_t completed_count = 0; + bool all_ok = true; + std::vector completed(submitted_count, 0); + + while (completed_count < submitted_count) { + bool made_progress = false; + std::vector pending; + pending.reserve(submitted_count - completed_count); + + for (size_t i = 0; i < submitted_count; ++i) { + if (completed[i] != 0) { + continue; + } + + int aio_err = ::aio_error(&cbs[i]); + if (aio_err == EINPROGRESS) { + pending.push_back(&cbs[i]); + continue; + } + + ssize_t result = ::aio_return(&cbs[i]); + all_ok = validate_aio_result(aio_err, result, read_reqs[req_begin + i]) && + all_ok; + completed[i] = 1; + ++completed_count; + made_progress = true; + } + + if (completed_count == submitted_count || made_progress) { + continue; + } + + int ret; + do { + ret = ::aio_suspend(pending.data(), static_cast(pending.size()), + nullptr); + } while (ret == -1 && errno == EINTR); + if (ret == -1) { + LOG_ERROR("aio_suspend failed while draining requests; errno=%d, %s", + errno, ::strerror(errno)); + struct timespec pause = {0, 1000000}; // 1 ms + ::nanosleep(&pause, nullptr); + all_ok = false; + } + } + + return all_ok; +} + +// Submit a batch through POSIX aio_read() and wait for completion with the +// portable aio_suspend() API supported by all target macOS SDK versions. +static int execute_io_aio_suspend(int fd, std::vector &read_reqs, + uint64_t n_retries, + ailego::IOBackendType *used_backend) { + for (size_t req_begin = 0; req_begin < read_reqs.size(); + req_begin += kMacAioBatchSize) { + size_t n_ops = std::min(kMacAioBatchSize, read_reqs.size() - req_begin); + std::vector cbs(n_ops); + size_t submitted_count = 0; + bool submission_ok = true; + + for (size_t i = 0; i < n_ops; ++i) { + std::memset(&cbs[i], 0, sizeof(cbs[i])); + cbs[i].aio_fildes = fd; + cbs[i].aio_offset = static_cast(read_reqs[req_begin + i].offset); + cbs[i].aio_buf = read_reqs[req_begin + i].buf; + cbs[i].aio_nbytes = read_reqs[req_begin + i].len; + cbs[i].aio_reqprio = 0; + cbs[i].aio_lio_opcode = LIO_READ; + cbs[i].aio_sigevent.sigev_notify = SIGEV_NONE; + + uint64_t tries = 0; + while (::aio_read(&cbs[i]) == -1) { + int submit_errno = errno; + if ((submit_errno == EINTR || submit_errno == EAGAIN) && + tries < n_retries) { + ++tries; + continue; + } + LOG_WARN( + "aio_read submission failed; errno=%d, %s, offset=%lu, " + "len=%lu; falling back to pread after draining submitted I/O", + submit_errno, ::strerror(submit_errno), + (unsigned long)read_reqs[req_begin + i].offset, + (unsigned long)read_reqs[req_begin + i].len); + submission_ok = false; + break; + } + if (!submission_ok) { + break; + } + ++submitted_count; + } + + bool all_ok = + drain_aio_requests(cbs, read_reqs, req_begin, submitted_count); + + if (!submission_ok || !all_ok) { + return execute_io_pread(fd, read_reqs, used_backend); + } + } + + if (used_backend != nullptr) { + *used_backend = ailego::IOBackendType::kPosixAio; + } + return 0; +} +#endif // __APPLE__ + +int execute_io(IOContext &ctx, int fd, std::vector &read_reqs, + uint64_t n_retries = 0, + ailego::IOBackendType *used_backend = nullptr) { #if (defined(__linux) || defined(__linux__)) - if (ailego::IOBackend::Instance().is_pread()) { - return execute_io_pread(fd, read_reqs); + if (ailego::IOBackend::Instance().available() == + ailego::IOBackendType::kPread) { + return execute_io_pread(fd, read_reqs, used_backend); } uint64_t iters = DiskAnnUtil::div_round_up(read_reqs.size(), MAX_EVENTS); @@ -139,7 +286,7 @@ int execute_io(IOContext ctx, int fd, std::vector &read_reqs, "io_submit failed; returned: %d, expected=%lu. falling back to " "pread", ret, n_ops); - return execute_io_pread(fd, read_reqs); + return execute_io_pread(fd, read_reqs, used_backend); } // Phase 2: io_getevents with retry (never re-submits). @@ -158,7 +305,7 @@ int execute_io(IOContext ctx, int fd, std::vector &read_reqs, "io_getevents failed; returned: %d, expected=%lu, errno=%d, %s, " "falling back to pread", ret, n_ops, errno, ::strerror(-ret)); - return execute_io_pread(fd, read_reqs); + return execute_io_pread(fd, read_reqs, used_backend); } // Phase 3: verify each completed read (res must equal requested length). @@ -173,13 +320,21 @@ int execute_io(IOContext ctx, int fd, std::vector &read_reqs, } } if (!all_ok) { - return execute_io_pread(fd, read_reqs); + return execute_io_pread(fd, read_reqs, used_backend); } } + if (used_backend != nullptr) { + *used_backend = ailego::IOBackendType::kLibAio; + } return 0; +#elif defined(__APPLE__) && TARGET_OS_OSX + (void)ctx; + return execute_io_aio_suspend(fd, read_reqs, n_retries, used_backend); #else - return execute_io_pread(fd, read_reqs); + (void)ctx; + (void)n_retries; + return execute_io_pread(fd, read_reqs, used_backend); #endif } @@ -216,14 +371,21 @@ void LinuxAlignedFileReader::register_thread() { std::unique_lock lk(ctx_mut); if (ctx_map.find(thread_id) != ctx_map.end()) { LOG_ERROR("multiple calls to register_thread from the same thread"); - return; } IOContext ctx = nullptr; - std::call_once(g_io_backend_log_once, log_diskann_io_backend); - if (ailego::IOBackend::Instance().is_pread()) { + auto &backend = ailego::IOBackend::Instance(); + std::call_once(g_io_backend_log_once, [&backend] { + if (backend.available() != ailego::IOBackendType::kPread) { + LOG_INFO("DiskAnn I/O backend: %s (async I/O enabled)", backend.name()); + } else { + LOG_WARN( + "DiskAnn I/O backend: synchronous pread (no async I/O available)"); + } + }); + if (backend.available() == ailego::IOBackendType::kPread) { lk.unlock(); return; } @@ -238,10 +400,21 @@ void LinuxAlignedFileReader::register_thread() { } } else { LOG_INFO("allocating ctx: %lu", (uint64_t)ctx); - ctx_map[thread_id] = ctx; } + lk.unlock(); +#elif defined(__APPLE__) && TARGET_OS_OSX + auto thread_id = std::this_thread::get_id(); + std::unique_lock lk(ctx_mut); + if (ctx_map.find(thread_id) != ctx_map.end()) { + LOG_ERROR("multiple calls to register_thread from the same thread"); + return; + } + std::call_once(g_io_backend_log_once, [] { + LOG_INFO("DiskAnn I/O backend: macOS POSIX AIO with aio_suspend"); + }); + ctx_map[thread_id] = 0; lk.unlock(); #endif } @@ -268,6 +441,20 @@ void LinuxAlignedFileReader::deregister_thread() { LibAioLoader::Instance().io_destroy(ctx); } LOG_INFO("returned ctx from thread"); +#elif defined(__APPLE__) && TARGET_OS_OSX + auto thread_id = std::this_thread::get_id(); + + { + std::lock_guard lk(ctx_mut); + auto it = ctx_map.find(thread_id); + if (it == ctx_map.end()) { + LOG_ERROR("deregister_thread: thread not registered"); + return; + } + ctx_map.erase(it); + } + + LOG_INFO("deregistered POSIX AIO thread"); #endif } @@ -283,6 +470,9 @@ void LinuxAlignedFileReader::deregister_all_threads() { } } ctx_map.clear(); +#elif defined(__APPLE__) && TARGET_OS_OSX + std::unique_lock lk(ctx_mut); + ctx_map.clear(); #endif } @@ -312,6 +502,32 @@ void LinuxAlignedFileReader::open(const std::string &fname) { ::strerror(errno)); } +#if defined(__APPLE__) && TARGET_OS_OSX + // macOS has no O_DIRECT. F_NOCACHE is its closest per-file equivalent: it + // asks the kernel to minimize caching for I/O through this descriptor. This + // is advisory rather than a guarantee that every read reaches the device. + // Disable read-ahead as well because DiskAnn performs random reads. + // + // Do not mmap the entire index and call msync(MS_INVALIDATE) here. That does + // not provide a reliable global cache eviction guarantee and makes open time + // and virtual-address usage scale with the size of the index. + if (this->file_desc != -1) { + if (::fcntl(this->file_desc, F_NOCACHE, 1) == -1) { + LOG_WARN( + "fcntl(F_NOCACHE) failed for %s (errno=%d: %s); reads will use " + "the page cache", + fname.c_str(), errno, ::strerror(errno)); + } else { + LOG_INFO("DiskAnn macOS: F_NOCACHE enabled for %s", fname.c_str()); + } + + if (::fcntl(this->file_desc, F_RDAHEAD, 0) == -1) { + LOG_WARN("fcntl(F_RDAHEAD, 0) failed for %s (errno=%d: %s)", + fname.c_str(), errno, ::strerror(errno)); + } + } +#endif + LOG_INFO("Opened file : %s", fname.c_str()); } @@ -323,7 +539,8 @@ void LinuxAlignedFileReader::close() { } int LinuxAlignedFileReader::read(std::vector &read_reqs, - IOContext &ctx, bool async) { + IOContext &ctx, bool async, + ailego::IOBackendType *used_backend) { if (async == true) { LOG_WARN("Async currently not supported"); } @@ -333,11 +550,10 @@ int LinuxAlignedFileReader::read(std::vector &read_reqs, return IndexError_Runtime; } - int ret = execute_io(ctx, this->file_desc, read_reqs); + int ret = execute_io(ctx, this->file_desc, read_reqs, 0, used_backend); return ret; } - } // namespace core } // namespace zvec diff --git a/src/core/algorithm/diskann/diskann_file_reader.h b/src/core/algorithm/diskann/diskann_file_reader.h index a1cb7c91a..16d089ef5 100644 --- a/src/core/algorithm/diskann/diskann_file_reader.h +++ b/src/core/algorithm/diskann/diskann_file_reader.h @@ -15,23 +15,34 @@ #define MAX_IO_DEPTH 128 -#include +#if defined(__APPLE__) +#include +#endif #if (defined(__linux) || defined(__linux__)) #include // dlopen-based libaio wrapper #endif -#include -#include +#include +#include +#include +#include #include +#include #include #include "diskann_util.h" namespace zvec { namespace core { +// On Linux, IOContext is the libaio io_context_t. +// On macOS, IOContext is an unused int placeholder because POSIX AIO does not +// require a persistent per-context resource. +// On other platforms, IOContext is a uint32_t placeholder. #if (defined(__linux) || defined(__linux__)) typedef io_context_t IOContext; +#elif defined(__APPLE__) && TARGET_OS_OSX +typedef int IOContext; #else typedef uint32_t IOContext; #endif @@ -39,10 +50,6 @@ typedef uint32_t IOContext; int setup_io_ctx(IOContext &ctx); int destroy_io_ctx(IOContext &ctx); -// Log the current DiskAnn I/O backend status (async vs. synchronous pread). -// Probes the backend on first call. No-op on non-Linux platforms. -void log_diskann_io_backend(); - struct AlignedRead { uint64_t offset; uint64_t len; @@ -52,9 +59,12 @@ struct AlignedRead { AlignedRead(uint64_t offset, uint64_t len, void *buf) : offset(offset), len(len), buf(buf) { +#if defined(__linux__) || defined(__linux) + // O_DIRECT requires 512-byte alignment on Linux. ailego_assert(static_cast(offset) % 512 == 0); ailego_assert(static_cast(len) % 512 == 0); ailego_assert(reinterpret_cast(buf) % 512 == 0); +#endif } }; @@ -75,10 +85,17 @@ class AlignedFileReader { virtual void open(const std::string &fname) = 0; virtual void close() = 0; + // used_backend, when provided, receives the backend that completed this + // call, including a synchronous fallback from an async backend. virtual int read(std::vector &read_reqs, IOContext &ctx, - bool async = false) = 0; + bool async = false, + ailego::IOBackendType *used_backend = nullptr) = 0; }; +// Reader implementation used on all supported platforms. +// On Linux (x86_64 and ARM64) it uses libaio for asynchronous batch I/O. +// On macOS (including ARM/Apple Silicon) it submits POSIX AIO requests and +// waits for completion with aio_suspend(). class LinuxAlignedFileReader : public AlignedFileReader { private: int file_desc; @@ -100,7 +117,7 @@ class LinuxAlignedFileReader : public AlignedFileReader { void close(); int read(std::vector &read_reqs, IOContext &ctx, - bool async = false); + bool async = false, ailego::IOBackendType *used_backend = nullptr); }; } // namespace core diff --git a/src/core/algorithm/diskann/diskann_indexer.cc b/src/core/algorithm/diskann/diskann_indexer.cc index 32e7ff67c..936779050 100644 --- a/src/core/algorithm/diskann/diskann_indexer.cc +++ b/src/core/algorithm/diskann/diskann_indexer.cc @@ -1086,8 +1086,14 @@ int DiskAnnIndexer::cached_beam_search_by_group(DiskAnnContext *ctx) { io_timer.reset(); - reader_->read(frontier_read_reqs, io_ctx); // synchronous IO linux + int read_ret = reader_->read(frontier_read_reqs, io_ctx); stats.io_us += io_timer.micro_seconds(); + if (read_ret != 0) { + LOG_ERROR("cached_beam_search_by_group: reader_->read failed, ret=%d", + read_ret); + ctx->set_error(true); + return IndexError_Runtime; + } } for (auto &cached_neighbor : cached_neighbors) { diff --git a/src/core/algorithm/diskann/diskann_indexer.h b/src/core/algorithm/diskann/diskann_indexer.h index c372d288f..abd8039d6 100644 --- a/src/core/algorithm/diskann/diskann_indexer.h +++ b/src/core/algorithm/diskann/diskann_indexer.h @@ -91,7 +91,13 @@ class DiskAnnIndexer { PQTable::Pointer pq_table_; - IOContext init_ctx_{0}; + IOContext init_ctx_{ +#if defined(__APPLE__) && TARGET_OS_OSX + -1 +#else + 0 +#endif + }; std::vector neighbor_cache_buffer_; void *coord_cache_buf_{nullptr}; diff --git a/src/core/algorithm/diskann/diskann_pq_trainer.cc b/src/core/algorithm/diskann/diskann_pq_trainer.cc index 73ca01656..c84744cb8 100644 --- a/src/core/algorithm/diskann/diskann_pq_trainer.cc +++ b/src/core/algorithm/diskann/diskann_pq_trainer.cc @@ -153,7 +153,7 @@ int DiskAnnPqTrainer::convert_pivot_data( cluster * dim + chunk_offsets[chunk]; const T *feature_ptr = reinterpret_cast(centroids[idx].feature()); - for (size_t d = 0; d <= chunk_dims[chunk]; ++d) { + for (size_t d = 0; d < chunk_dims[chunk]; ++d) { pivot_data_ptr[d] = feature_ptr[d]; } } diff --git a/src/core/algorithm/diskann/diskann_searcher.cc b/src/core/algorithm/diskann/diskann_searcher.cc index a34c546e5..13d5d1b3a 100644 --- a/src/core/algorithm/diskann/diskann_searcher.cc +++ b/src/core/algorithm/diskann/diskann_searcher.cc @@ -13,6 +13,7 @@ // limitations under the License. #include "diskann_searcher.h" +#include #include "diskann_context.h" #include "diskann_indexer.h" #include "diskann_params.h" @@ -25,7 +26,22 @@ DiskAnnSearcher::DiskAnnSearcher() {} DiskAnnSearcher::~DiskAnnSearcher() {} int DiskAnnSearcher::init(const ailego::Params &search_params) { - log_diskann_io_backend(); +#if defined(__linux__) || defined(__linux) + // Eagerly probe the I/O backend at init time so the user gets immediate + // feedback about whether async I/O is available. IOBackend auto-probes on + // first Instance() access. + auto &backend = ailego::IOBackend::Instance(); + if (backend.available() == ailego::IOBackendType::kPread) { + LOG_WARN( + "DiskAnn: no async I/O backend available. Install libaio (e.g. " + "'apt-get install libaio1', or 'libaio1t64' on Ubuntu 24.04+) and " + "retry. DiskAnn will fall back to synchronous pread() — performance " + "will be degraded."); + } else { + LOG_INFO("DiskAnn: I/O backend '%s' loaded — async I/O enabled.", + backend.name()); + } +#endif search_params.get(PARAM_DISKANN_SEARCHER_LIST_SIZE, &list_size_); search_params.get(PARAM_DISKANN_SEARCHER_CACHE_NODE_NUM, &cache_nodes_num_); @@ -161,10 +177,10 @@ int DiskAnnSearcher::search_impl(const void *query, const IndexQueryMeta &qmeta, for (uint32_t i = 0; i < count; i++) { ctx->reset_query(query); - diskann_indexer_->knn_search(ctx); - - if (ailego_unlikely(ctx->error())) { - return IndexError_Runtime; + int ret = diskann_indexer_->knn_search(ctx); + if (ailego_unlikely(ret != 0 || ctx->error())) { + LOG_ERROR("DiskAnn knn search failed, ret=%d", ret); + return ret != 0 ? ret : IndexError_Runtime; } ctx->topk_to_result(i); diff --git a/src/core/algorithm/diskann/diskann_searcher_entity.cc b/src/core/algorithm/diskann/diskann_searcher_entity.cc index c9e49deba..3e353a216 100644 --- a/src/core/algorithm/diskann/diskann_searcher_entity.cc +++ b/src/core/algorithm/diskann/diskann_searcher_entity.cc @@ -398,8 +398,8 @@ const void *DiskAnnSearcherEntity::get_vector(diskann_id_t id) const { const void *vec; if (ailego_unlikely(vector_segment_->read(total_offset, &vec, read_size) != read_size)) { - LOG_ERROR("Read vector from segment failed, id: %u, offset: %lu", id, - total_offset); + LOG_ERROR("Read vector from segment failed, id: %u, offset: %llu", id, + (unsigned long long)total_offset); return nullptr; } diff --git a/src/core/algorithm/diskann/diskann_streamer.cc b/src/core/algorithm/diskann/diskann_streamer.cc index 82e97dcd6..ca312e4b5 100644 --- a/src/core/algorithm/diskann/diskann_streamer.cc +++ b/src/core/algorithm/diskann/diskann_streamer.cc @@ -13,6 +13,7 @@ // limitations under the License. #include "diskann_streamer.h" +#include #include "diskann_context.h" #include "diskann_index_provider.h" #include "diskann_indexer.h" @@ -29,7 +30,22 @@ int DiskAnnStreamer::init(const IndexMeta &meta, const ailego::Params &search_params) { meta_ = meta; - log_diskann_io_backend(); +#if defined(__linux__) || defined(__linux) + // Eagerly probe the I/O backend at init time so the user gets immediate + // feedback about whether async I/O is available. IOBackend auto-probes on + // first Instance() access. + auto &backend = ailego::IOBackend::Instance(); + if (backend.available() == ailego::IOBackendType::kPread) { + LOG_WARN( + "DiskAnn: no async I/O backend available. Install libaio (e.g. " + "'apt-get install libaio1', or 'libaio1t64' on Ubuntu 24.04+) and " + "retry. DiskAnn will fall back to synchronous pread() — performance " + "will be degraded."); + } else { + LOG_INFO("DiskAnn: I/O backend '%s' loaded — async I/O enabled.", + backend.name()); + } +#endif search_params.get(PARAM_DISKANN_SEARCHER_LIST_SIZE, &list_size_); search_params.get(PARAM_DISKANN_SEARCHER_CACHE_NODE_NUM, &cache_nodes_num_); @@ -160,7 +176,11 @@ int DiskAnnStreamer::search_impl(const void *query, const IndexQueryMeta &qmeta, for (uint32_t i = 0; i < count; i++) { ctx->reset_query(query); - diskann_indexer_->knn_search(ctx); + int ret = diskann_indexer_->knn_search(ctx); + if (ailego_unlikely(ret != 0 || ctx->error())) { + LOG_ERROR("DiskAnn knn search failed, ret=%d", ret); + return ret != 0 ? ret : IndexError_Runtime; + } ctx->topk_to_result(i); @@ -338,6 +358,7 @@ IndexSearcher::Context::Pointer DiskAnnStreamer::create_context() const { } ctx->set_list_size(list_size_); + ctx->set_magic(magic_); return Context::Pointer(ctx); } diff --git a/src/core/algorithm/diskann/diskann_util.h b/src/core/algorithm/diskann/diskann_util.h index a02130bf0..3e2ecd652 100644 --- a/src/core/algorithm/diskann/diskann_util.h +++ b/src/core/algorithm/diskann/diskann_util.h @@ -35,7 +35,13 @@ class DiskAnnUtil { } static inline void alloc_aligned(void **ptr, size_t size, size_t align) { - *ptr = ::aligned_alloc(align, size); + // C11 aligned_alloc() requires size to be an integral multiple of + // alignment. This is true on Linux (glibc relaxes the requirement) + // but NOT on macOS, where aligned_alloc(32, 16) returns NULL. + // Use posix_memalign() which is portable and has no such restriction. + if (::posix_memalign(ptr, align, size) != 0) { + *ptr = nullptr; + } } static inline void free_aligned(void *ptr) { diff --git a/src/core/interface/indexes/diskann_index.cc b/src/core/interface/indexes/diskann_index.cc index 5af446733..f1553b6ff 100644 --- a/src/core/interface/indexes/diskann_index.cc +++ b/src/core/interface/indexes/diskann_index.cc @@ -271,4 +271,4 @@ int DiskAnnIndex::Merge(const std::vector &indexes, return 0; } -} // namespace zvec::core_interface \ No newline at end of file +} // namespace zvec::core_interface diff --git a/src/db/index/common/schema.cc b/src/db/index/common/schema.cc index 532696803..875520f63 100644 --- a/src/db/index/common/schema.cc +++ b/src/db/index/common/schema.cc @@ -198,19 +198,19 @@ Status FieldSchema::validate() const { } if (index_params_->type() == IndexType::DISKANN) { - // DiskAnn requires Linux x86_64/i686/i386. The CMake variable + // DiskAnn supports Linux x86/x86_64/ARM64 and macOS. The CMake variable // DISKANN_SUPPORTED (defined in the top-level CMakeLists.txt) is the // single source of truth for platform eligibility — it is also used by // index_factory.cc to conditionally compile the DiskAnn index // registration. Using the same macro here ensures that schema // validation and index registration agree on supported platforms. // - // libaio is loaded eagerly (via dlopen) inside DiskAnnBuilder::init() - // and DiskAnnStreamer::init(); if libaio is missing, DiskAnn falls - // back to synchronous pread() with degraded performance. + // Linux probes libaio at runtime and falls back to synchronous pread() + // when it is unavailable. macOS uses system POSIX AIO with + // aio_suspend() completion waits. #if !DISKANN_SUPPORTED return Status::NotSupported( - "DiskAnn is not supported on this platform (Linux x86_64 only)"); + "DiskAnn is supported only on Linux x86/x86_64/ARM64 and macOS"); #endif } diff --git a/src/include/zvec/ailego/io/io_backend.h b/src/include/zvec/ailego/io/io_backend.h index 008b01514..c177c424a 100644 --- a/src/include/zvec/ailego/io/io_backend.h +++ b/src/include/zvec/ailego/io/io_backend.h @@ -30,12 +30,14 @@ namespace ailego { // Supported I/O backend types. enum class IOBackendType { - kPread, // Synchronous pread() — no async I/O - kLibAio, // libaio loaded at runtime via dlopen() + kPread = 0, // Synchronous pread() — no async I/O + kLibAio = 1, // Linux libaio loaded at runtime via dlopen() + kPosixAio = 2 // macOS POSIX AIO with aio_suspend() completion waits }; // Returns the currently active I/O backend type. -// Triggers backend initialization on first call (libaio > pread). +// On Linux, triggers backend initialization on first call (libaio > pread). +// On macOS, returns the system POSIX AIO backend. IOBackendType current_io_backend_type(); // Returns a human-readable description of the currently active I/O backend. diff --git a/src/include/zvec/c_api.h b/src/include/zvec/c_api.h index 2d048689b..9b8d4a410 100644 --- a/src/include/zvec/c_api.h +++ b/src/include/zvec/c_api.h @@ -789,14 +789,16 @@ typedef uint32_t zvec_io_backend_type_t; 0 /**< Synchronous pread() \u2014 no async I/O */ #define ZVEC_IO_BACKEND_TYPE_LIBAIO \ 1 /**< libaio loaded at runtime via dlopen() */ +#define ZVEC_IO_BACKEND_TYPE_POSIX_AIO \ + 2 /**< macOS POSIX AIO with aio_suspend() completion waits */ /** * @brief Get the current I/O backend type for DiskAnn async disk reads. * - * Pure introspection \u2014 no side effects, no install hints. + * On Linux, the first call may probe for libaio. No configuration is changed + * and no installation hints are emitted. * - * @return zvec_io_backend_type_t The loaded backend type - * (ZVEC_IO_BACKEND_TYPE_LIBAIO or ZVEC_IO_BACKEND_TYPE_PREAD). + * @return zvec_io_backend_type_t The selected backend type. */ ZVEC_EXPORT zvec_io_backend_type_t ZVEC_CALL zvec_get_io_backend_type(void); @@ -805,7 +807,7 @@ ZVEC_EXPORT zvec_io_backend_type_t ZVEC_CALL zvec_get_io_backend_type(void); * * @param type The backend type code. * @return Thread-local string valid until the next call on this thread; - * "libaio", "pread", or "unknown". + * "libaio", "posix_aio", "pread", or "unknown". */ ZVEC_EXPORT const char *ZVEC_CALL zvec_get_io_backend_type_name(zvec_io_backend_type_t type); @@ -813,7 +815,8 @@ zvec_get_io_backend_type_name(zvec_io_backend_type_t type); /** * @brief Get a human-readable description of the current I/O backend. * - * When only pread is available, includes installation guidance for libaio. + * On Linux, the pread description includes installation guidance for libaio. + * Other platforms return a generic synchronous-backend description. * * @return Thread-local string valid until the next call on this thread. */ diff --git a/src/include/zvec/db/collection.h b/src/include/zvec/db/collection.h index 83a289c52..581a9449b 100644 --- a/src/include/zvec/db/collection.h +++ b/src/include/zvec/db/collection.h @@ -119,4 +119,4 @@ class Collection { const std::string &column_name) const = 0; }; -} // namespace zvec \ No newline at end of file +} // namespace zvec diff --git a/tests/ailego/io/io_backend_test.cc b/tests/ailego/io/io_backend_test.cc new file mode 100644 index 000000000..800696bde --- /dev/null +++ b/tests/ailego/io/io_backend_test.cc @@ -0,0 +1,52 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include +#include +#include +#include + +#if defined(__APPLE__) +#include +#endif + +namespace zvec { +namespace ailego { + +TEST(IOBackendTest, StableBackendValues) { + EXPECT_EQ(static_cast(IOBackendType::kPread), 0U); + EXPECT_EQ(static_cast(IOBackendType::kLibAio), 1U); + EXPECT_EQ(static_cast(IOBackendType::kPosixAio), 2U); +} + +TEST(IOBackendTest, BackendNames) { + EXPECT_STREQ(IOBackendTypeName(IOBackendType::kPread), "pread"); + EXPECT_STREQ(IOBackendTypeName(IOBackendType::kLibAio), "libaio"); + EXPECT_STREQ(IOBackendTypeName(IOBackendType::kPosixAio), "posix_aio"); +} + +#if defined(__APPLE__) && TARGET_OS_OSX +TEST(IOBackendTest, MacOSUsesPosixAio) { + EXPECT_EQ(current_io_backend_type(), IOBackendType::kPosixAio); + EXPECT_NE(current_io_backend_description().find("aio_suspend"), + std::string::npos); +} +#elif defined(__APPLE__) +TEST(IOBackendTest, NonMacOSAppleUsesPread) { + EXPECT_EQ(current_io_backend_type(), IOBackendType::kPread); +} +#endif + +} // namespace ailego +} // namespace zvec diff --git a/tests/c/c_api_test.c b/tests/c/c_api_test.c index 7365dcf1a..f63a47482 100644 --- a/tests/c/c_api_test.c +++ b/tests/c/c_api_test.c @@ -20,6 +20,10 @@ #include #include +#if defined(__APPLE__) +#include +#endif + // Platform-specific headers #ifdef _WIN32 #include @@ -130,6 +134,30 @@ void test_version_functions(void) { TEST_END(); } +void test_io_backend_functions(void) { + TEST_START(); + + zvec_io_backend_type_t type = zvec_get_io_backend_type(); + const char *name = zvec_get_io_backend_type_name(type); + const char *description = zvec_get_io_backend_description(); + + TEST_ASSERT(name != NULL); + TEST_ASSERT(description != NULL); +#if defined(__APPLE__) && TARGET_OS_OSX + TEST_ASSERT(type == ZVEC_IO_BACKEND_TYPE_POSIX_AIO); + TEST_ASSERT(strcmp(name, "posix_aio") == 0); +#elif defined(__linux__) || defined(__linux) + TEST_ASSERT(type == ZVEC_IO_BACKEND_TYPE_LIBAIO || + type == ZVEC_IO_BACKEND_TYPE_PREAD); +#else + TEST_ASSERT(type == ZVEC_IO_BACKEND_TYPE_PREAD); +#endif + TEST_ASSERT(strcmp(zvec_get_io_backend_type_name(UINT32_MAX), "unknown") == + 0); + + TEST_END(); +} + void test_error_handling_functions(void) { TEST_START(); @@ -6375,6 +6403,7 @@ int main(void) { printf("Cleanup completed.\n\n"); test_version_functions(); + test_io_backend_functions(); test_error_handling_functions(); test_zvec_config(); test_zvec_initialize(); diff --git a/tests/core/algorithm/diskann/diskann_file_reader_test.cc b/tests/core/algorithm/diskann/diskann_file_reader_test.cc new file mode 100644 index 000000000..2eb2d76c7 --- /dev/null +++ b/tests/core/algorithm/diskann/diskann_file_reader_test.cc @@ -0,0 +1,150 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "diskann_file_reader.h" +#include +#include +#include +#include +#include +#include +#include + +using namespace zvec::core; + +namespace { + +class TemporaryFile { + public: + TemporaryFile() : fd_(::mkstemp(path_)) {} + + ~TemporaryFile() { + if (fd_ >= 0) { + ::close(fd_); + } + ::unlink(path_); + } + + int fd() const { + return fd_; + } + + const char *path() const { + return path_; + } + + void close() { + if (fd_ >= 0) { + ::close(fd_); + fd_ = -1; + } + } + + private: + char path_[64] = "DiskAnnFileReaderTest.XXXXXX"; + int fd_; +}; + +} // namespace + +TEST(DiskAnnFileReaderTest, BatchAlignedReads) { + constexpr size_t kPageSize = 4096; + constexpr size_t kPageCount = 32; + + TemporaryFile file; + ASSERT_GE(file.fd(), 0); + + std::vector source(kPageSize * kPageCount); + for (size_t page = 0; page < kPageCount; ++page) { + std::memset(source.data() + page * kPageSize, static_cast(page + 1), + kPageSize); + } + + size_t written = 0; + while (written < source.size()) { + ssize_t ret = ::pwrite(file.fd(), source.data() + written, + source.size() - written, written); + ASSERT_GT(ret, 0); + written += static_cast(ret); + } + ASSERT_EQ(::fsync(file.fd()), 0); + file.close(); + + void *raw_buffer = nullptr; + ASSERT_EQ(::posix_memalign(&raw_buffer, kPageSize, source.size()), 0); + std::unique_ptr output(raw_buffer, &std::free); + + std::vector requests; + requests.reserve(kPageCount); + for (size_t i = 0; i < kPageCount; ++i) { + // Use a non-sequential page order and more than AIO_LISTIO_MAX requests + // so the macOS backend has to submit multiple batches. + size_t source_page = (i * 7) % kPageCount; + requests.emplace_back(source_page * kPageSize, kPageSize, + static_cast(output.get()) + i * kPageSize); + } + + LinuxAlignedFileReader reader; + reader.open(file.path()); + IOContext ctx{ +#if defined(__APPLE__) && TARGET_OS_OSX + 0 +#else + nullptr +#endif + }; + ASSERT_EQ(setup_io_ctx(ctx), 0); + zvec::ailego::IOBackendType used_backend = + zvec::ailego::IOBackendType::kPread; + ASSERT_EQ(reader.read(requests, ctx, false, &used_backend), 0); +#if defined(__APPLE__) && TARGET_OS_OSX + EXPECT_EQ(used_backend, zvec::ailego::IOBackendType::kPosixAio); +#endif + + for (size_t i = 0; i < kPageCount; ++i) { + size_t source_page = (i * 7) % kPageCount; + const auto *page = + static_cast(output.get()) + i * kPageSize; + for (size_t byte = 0; byte < kPageSize; ++byte) { + ASSERT_EQ(page[byte], static_cast(source_page + 1)); + } + } + + EXPECT_EQ(destroy_io_ctx(ctx), 0); + reader.close(); +} + +#if defined(__APPLE__) && TARGET_OS_OSX +TEST(DiskAnnFileReaderTest, ShortAioReadFallsBackToPread) { + TemporaryFile file; + ASSERT_GE(file.fd(), 0); + + constexpr char kSource[] = "DiskAnn pread fallback"; + ASSERT_EQ(::pwrite(file.fd(), kSource, sizeof(kSource), 0), + static_cast(sizeof(kSource))); + file.close(); + + char output[sizeof(kSource) * 2] = {}; + std::vector requests{{0, sizeof(output), output}}; + LinuxAlignedFileReader reader; + reader.open(file.path()); + IOContext ctx = 0; + zvec::ailego::IOBackendType used_backend = + zvec::ailego::IOBackendType::kPosixAio; + + EXPECT_NE(reader.read(requests, ctx, false, &used_backend), 0); + EXPECT_EQ(used_backend, zvec::ailego::IOBackendType::kPread); + reader.close(); +} +#endif diff --git a/tests/core/algorithm/diskann/diskann_searcher_test.cc b/tests/core/algorithm/diskann/diskann_searcher_test.cc index ff432886f..8804e4407 100644 --- a/tests/core/algorithm/diskann/diskann_searcher_test.cc +++ b/tests/core/algorithm/diskann/diskann_searcher_test.cc @@ -16,6 +16,7 @@ #include #include #include +#include #include #include #include @@ -220,6 +221,32 @@ TEST_F(DiskAnnSearcherTest, TestGeneral) { EXPECT_GT(recall, 0.90f); EXPECT_GT(topk1Recall, 0.80f); EXPECT_GT(cost, 2.0f); + + // A context created by the streamer must carry the streamer's magic so it + // can be reused instead of being recreated on every search. + IndexStreamer::Pointer streamer = + IndexFactory::CreateStreamer("DiskAnnStreamer"); + ASSERT_NE(streamer, nullptr); + ASSERT_EQ(0, streamer->init(*_index_meta_ptr, search_params)); + + auto streamer_storage = IndexFactory::CreateStorage("FileReadStorage"); + ASSERT_EQ(0, streamer_storage->open(path, false)); + ASSERT_EQ(0, streamer->open(streamer_storage)); + + auto streamer_ctx = streamer->create_context(); + ASSERT_NE(streamer_ctx, nullptr); + streamer_ctx->set_topk(topk); + auto *original_ctx = streamer_ctx.get(); + + ASSERT_EQ(0, streamer->search_impl(vec.data(), qmeta, streamer_ctx)); + EXPECT_EQ(original_ctx, streamer_ctx.get()); + ASSERT_EQ(0, streamer->search_impl(vec.data(), qmeta, streamer_ctx)); + EXPECT_EQ(original_ctx, streamer_ctx.get()); + + // I/O failures from the indexer must be propagated by the streamer instead + // of being converted into a successful search with incomplete results. + ASSERT_EQ(0, ::truncate(path.c_str(), 0)); + EXPECT_NE(0, streamer->search_impl(vec.data(), qmeta, streamer_ctx)); } TEST_F(DiskAnnSearcherTest, TestNodeCache) {