refactor(diskann): decouple from libaio via dlopen (#532)

Co-authored-by: Zefeng Yin <yinzefeng.yzf@alibaba-inc.com>
This commit is contained in:
rayx 2026-07-16 17:58:25 +08:00 committed by GitHub
parent 23538ab876
commit ec8a78ee08
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
49 changed files with 1108 additions and 809 deletions

View File

@ -60,14 +60,6 @@ jobs:
sudo apt-get install -y clang libomp-dev sudo apt-get install -y clang libomp-dev
shell: bash shell: bash
- name: Install AIO
if: runner.os == 'Linux' && runner.arch == 'X64'
run: |
sudo apt-get update
sudo apt-get install -y --no-install-recommends \
libaio-dev
shell: bash
- name: Print CPU info - name: Print CPU info
if: runner.os == 'Linux' if: runner.os == 'Linux'
run: lscpu run: lscpu
@ -183,3 +175,41 @@ jobs:
./c_api_index_example ./c_api_index_example
./c_api_optimized_example ./c_api_optimized_example
shell: bash shell: bash
# ------------------------------------------------------------------ #
# DiskAnn libaio round: install libaio and re-run tests
# ------------------------------------------------------------------ #
- name: Install libaio runtime
if: matrix.platform == 'linux-x64'
run: |
sudo apt-get update -y
# libaio1t64 on Ubuntu 24.04+ (t64 transition), libaio1 on older
sudo apt-get install -y libaio1t64 || sudo apt-get install -y libaio1
echo "=== libaio status (should be present) ==="
dpkg -l | grep -i libaio || true
ldconfig -p | grep libaio || true
shell: bash
- name: Run DiskAnn C++ Tests (w/ libaio)
if: matrix.platform == 'linux-x64'
run: |
cd "$GITHUB_WORKSPACE/build"
ctest -R diskann --output-on-failure --parallel $NPROC
shell: bash
- name: Run DiskAnn Python Tests (w/ libaio)
if: matrix.platform == 'linux-x64'
run: |
cd "$GITHUB_WORKSPACE"
python -m pytest python/tests/test_collection_diskann.py -v
shell: bash
# Verify installing libaio does not affect existing non-DiskAnn tests.
- name: Run HNSW Tests (w/ libaio)
if: matrix.platform == 'linux-x64'
run: |
cd "$GITHUB_WORKSPACE/build"
ctest -R hnsw --output-on-failure --parallel $NPROC
cd "$GITHUB_WORKSPACE"
python -m pytest python/tests/test_hnsw_contiguous_memory.py -v
shell: bash

View File

@ -34,13 +34,6 @@ jobs:
cache: 'pip' cache: 'pip'
cache-dependency-path: 'pyproject.toml' cache-dependency-path: 'pyproject.toml'
- name: Install system dependencies
run: |
sudo apt-get update
sudo apt-get install -y --no-install-recommends \
libaio-dev
shell: bash
- name: Set up environment variables - name: Set up environment variables
run: | run: |
NPROC=$(nproc 2>/dev/null || echo 2) NPROC=$(nproc 2>/dev/null || echo 2)

View File

@ -47,7 +47,7 @@ jobs:
if: steps.changed_files.outputs.any_changed == 'true' if: steps.changed_files.outputs.any_changed == 'true'
run: | run: |
sudo apt-get update sudo apt-get update
sudo apt-get install -y clang-tidy=1:18.0-59~exp2 cmake ninja-build libomp-dev libaio-dev sudo apt-get install -y clang-tidy=1:18.0-59~exp2 cmake ninja-build libomp-dev
- name: Setup ccache - name: Setup ccache
if: steps.changed_files.outputs.any_changed == 'true' if: steps.changed_files.outputs.any_changed == 'true'

View File

@ -48,8 +48,8 @@ jobs:
- name: Install system dependencies - name: Install system dependencies
run: | run: |
sudo apt-get update sudo apt-get update
sudo apt-get install -y --no-install-recommends \ sudo apt-get install -y --no-install-recommends lcov
lcov libaio-dev sudo apt-get install -y libaio1t64 || sudo apt-get install -y libaio1
shell: bash shell: bash
- name: Install dependencies - name: Install dependencies

View File

@ -122,7 +122,7 @@ else()
endif() endif()
message(STATUS "RABITQ_ARCH_FLAG: ${RABITQ_ARCH_FLAG}") message(STATUS "RABITQ_ARCH_FLAG: ${RABITQ_ARCH_FLAG}")
# DiskAnn support (Linux x86_64 only, requires libaio) # 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) if(CMAKE_SYSTEM_NAME STREQUAL "Linux" AND CMAKE_SYSTEM_PROCESSOR MATCHES "x86_64|i686|i386" AND NOT ANDROID AND NOT IOS)
set(DISKANN_SUPPORTED ON) set(DISKANN_SUPPORTED ON)
add_definitions(-DDISKANN_SUPPORTED=1) add_definitions(-DDISKANN_SUPPORTED=1)
@ -180,24 +180,9 @@ if(BUILD_PYTHON_BINDINGS)
install(TARGETS _zvec LIBRARY DESTINATION ${ZVEC_PY_INSTALL_DIR}/zvec install(TARGETS _zvec LIBRARY DESTINATION ${ZVEC_PY_INSTALL_DIR}/zvec
COMPONENT python) COMPONENT python)
# DiskAnn ships as a runtime-loaded shared module # DiskAnn is now statically linked into _zvec.so via --whole-archive
# (libzvec_diskann_plugin.so) that is brought online implicitly the # (see src/binding/python/CMakeLists.txt -> core_knn_diskann_static),
# first time a DiskAnn index is created users never call any load # so no separate runtime .so needs to be installed in the wheel.
# function. The Python extension resolves the module next to _zvec.so
# (see the $ORIGIN rpath in src/binding/python/CMakeLists.txt); the
# module must therefore be installed alongside _zvec.so, i.e. inside the
# zvec package directory as well.
#
# Gate on DISKANN_SUPPORTED, not on the target's existence: on unsupported
# platforms (e.g. macOS / ARM64) the core_knn_diskann target is still
# defined, but built from an empty stub (src/core/algorithm/CMakeLists.txt)
# with zero exported symbols and a runtime load path compiled out
# (#if DISKANN_SUPPORTED). Shipping that stub is pure dead weight, so it is
# only packaged where DiskAnn is real currently Linux x86_64 with libaio.
if(DISKANN_SUPPORTED)
install(TARGETS core_knn_diskann LIBRARY DESTINATION ${ZVEC_PY_INSTALL_DIR}/zvec
COMPONENT python)
endif()
# Bundle cppjieba's dictionary files so the `jieba` FTS tokenizer works # Bundle cppjieba's dictionary files so the `jieba` FTS tokenizer works
# out of the box. python/zvec/__init__.py resolves this directory via # out of the box. python/zvec/__init__.py resolves this directory via
# importlib.resources and registers it with set_default_jieba_dict_dir(). # importlib.resources and registers it with set_default_jieba_dict_dir().

View File

@ -17,7 +17,6 @@ get_filename_component(ZVEC_ROOT_DIR "${CMAKE_CURRENT_LIST_DIR}/../.." ABSOLUTE)
set(ZVEC_INCLUDE_DIR ${ZVEC_ROOT_DIR}/src/include) set(ZVEC_INCLUDE_DIR ${ZVEC_ROOT_DIR}/src/include)
set(ZVEC_LIB_DIR ${ZVEC_ROOT_DIR}/${HOST_BUILD_DIR}/lib) set(ZVEC_LIB_DIR ${ZVEC_ROOT_DIR}/${HOST_BUILD_DIR}/lib)
# Add include and library search paths
include_directories(${ZVEC_INCLUDE_DIR}) include_directories(${ZVEC_INCLUDE_DIR})
set(ZVEC_LIB_SEARCH_DIRS ${ZVEC_LIB_DIR}) set(ZVEC_LIB_SEARCH_DIRS ${ZVEC_LIB_DIR})

View File

@ -181,11 +181,6 @@ test-command = "cd {project} && pytest python/tests -v --tb=short"
build-verbosity = 1 build-verbosity = 1
[tool.cibuildwheel.linux] [tool.cibuildwheel.linux]
# libaio is required by the C++ backend; install it inside the manylinux
# container (manylinux_2_28 is AlmaLinux 8 based, so use dnf/libaio-devel).
# libaio-devel lives in BaseOS; disable EPEL so a flaky EPEL mirror cannot
# break metadata refresh.
before-all = "dnf install -y --disablerepo=epel libaio-devel"
archs = ["auto"] archs = ["auto"]
environment = { CMAKE_GENERATOR = "Unix Makefiles", CMAKE_BUILD_PARALLEL_LEVEL = "16" } environment = { CMAKE_GENERATOR = "Unix Makefiles", CMAKE_BUILD_PARALLEL_LEVEL = "16" }
manylinux-x86_64-image = "manylinux_2_28" manylinux-x86_64-image = "manylinux_2_28"

View File

@ -13,19 +13,14 @@ from typing import Any, Generator
from zvec.typing import DataType, StatusCode, MetricType, QuantizeType from zvec.typing import DataType, StatusCode, MetricType, QuantizeType
import zvec import zvec
# Cache the DiskAnn plugin preload status so we pay the load cost once per
# test session. The plugin normally auto-loads on first DiskAnn use, but we
# preload it explicitly here so a missing libaio / misplaced plugin .so
# surfaces as a clear pytest skip instead of a confusing
# "Create vector column indexer failed" deep inside the collection code path.
_DISKANN_PRELOAD_REASON: str | None = None _DISKANN_PRELOAD_REASON: str | None = None
_DISKANN_PRELOAD_DONE: bool = False _DISKANN_PRELOAD_DONE: bool = False
def _ensure_diskann_runtime_or_reason() -> str | None: def _ensure_diskann_runtime_or_reason() -> str | None:
"""Preload the DiskAnn plugin and return None on success or a human-readable """Check whether DiskAnn is available on this platform and return None
skip reason on failure. Idempotent across calls.""" on success or a human-readable skip reason on failure. Idempotent across
calls."""
global _DISKANN_PRELOAD_DONE, _DISKANN_PRELOAD_REASON global _DISKANN_PRELOAD_DONE, _DISKANN_PRELOAD_REASON
if _DISKANN_PRELOAD_DONE: if _DISKANN_PRELOAD_DONE:
return _DISKANN_PRELOAD_REASON return _DISKANN_PRELOAD_REASON
@ -35,22 +30,6 @@ def _ensure_diskann_runtime_or_reason() -> str | None:
_DISKANN_PRELOAD_REASON = "DiskAnn only supported on Linux x86_64" _DISKANN_PRELOAD_REASON = "DiskAnn only supported on Linux x86_64"
return _DISKANN_PRELOAD_REASON return _DISKANN_PRELOAD_REASON
if not zvec.is_libaio_available():
_DISKANN_PRELOAD_REASON = (
"libaio is not available on this host; DiskAnn cannot run. "
"Install libaio1 (or libaio1t64 on Ubuntu 24.04+) and retry."
)
return _DISKANN_PRELOAD_REASON
status = zvec.load_diskann_plugin()
if status != zvec.DISKANN_PLUGIN_OK:
_DISKANN_PRELOAD_REASON = (
f"Failed to load DiskAnn plugin (status={status}); "
"check that libzvec_diskann_plugin.so is installed alongside "
"_zvec.so in the Python site-packages directory."
)
return _DISKANN_PRELOAD_REASON
_DISKANN_PRELOAD_REASON = None _DISKANN_PRELOAD_REASON = None
return None return None
@ -148,8 +127,7 @@ def full_schema_new(request) -> CollectionSchema:
else: else:
nullable, has_index, vector_index = True, False, HnswIndexParam() nullable, has_index, vector_index = True, False, HnswIndexParam()
# Skip DiskAnn tests on unsupported platforms or when the runtime cannot # Skip DiskAnn tests on unsupported platforms.
# be brought up (missing libaio, plugin .so not installed, etc.).
from zvec.model.param import DiskAnnIndexParam from zvec.model.param import DiskAnnIndexParam
if isinstance(vector_index, DiskAnnIndexParam): if isinstance(vector_index, DiskAnnIndexParam):

View File

@ -13,17 +13,15 @@
# limitations under the License. # limitations under the License.
"""End-to-end collection tests for the DiskAnn index. """End-to-end collection tests for the DiskAnn index.
Mirrors ``test_collection_hnsw_rabitq.py`` but targets the DiskAnn plugin. Mirrors ``test_collection_hnsw_rabitq.py`` but targets the DiskAnn index.
Two platform-level prerequisites are enforced at module import time: Two platform-level prerequisites are enforced at module import time:
1. DiskAnn is currently built only for Linux x86_64 other platforms are 1. DiskAnn is currently built only for Linux x86_64 other platforms are
skipped wholesale. skipped wholesale.
2. The DiskAnn backend lives in a *runtime-loaded* plugin 2. libaio is loaded eagerly (via dlopen) inside DiskAnnBuilder::init() /
(``libzvec_diskann_plugin.so``). It must be loaded with ``RTLD_GLOBAL | DiskAnnStreamer::init(). If libaio is missing, DiskAnn falls back to
RTLD_NOW`` BEFORE ``import zvec`` so that the plugin's ``IndexFactory`` synchronous pread() the tests still run but with degraded performance.
singleton is unified with the one inside ``_zvec.so``. After ``import
zvec`` we must also call ``zvec.load_diskann_plugin()`` exactly once.
If either prerequisite fails the whole module is skipped so the rest of the If either prerequisite fails the whole module is skipped so the rest of the
test-suite is not affected. test-suite is not affected.
@ -32,7 +30,6 @@ test-suite is not affected.
from __future__ import annotations from __future__ import annotations
import math import math
import os
import platform import platform
import sys import sys
@ -46,13 +43,6 @@ pytestmark = pytest.mark.skipif(
reason="DiskAnn plugin is only supported on Linux x86_64", reason="DiskAnn plugin is only supported on Linux x86_64",
) )
# Promote all symbols in subsequently-loaded DSOs to the global namespace and
# resolve relocations eagerly. This is REQUIRED so the DiskAnn plugin can see
# the ``IndexFactory`` singleton that lives in ``_zvec.so`` and vice versa.
# See: DiskAnn RTLD_GLOBAL + RTLD_NOW Requirement.
if sys.platform == "linux":
sys.setdlopenflags(sys.getdlopenflags() | os.RTLD_GLOBAL | os.RTLD_NOW)
import zvec # noqa: E402 import zvec # noqa: E402
from zvec import ( # noqa: E402 from zvec import ( # noqa: E402

View File

@ -17,6 +17,7 @@ import pytest
from zvec import ( from zvec import (
DataType, DataType,
IndexType, IndexType,
IOBackendType,
MetricType, MetricType,
QuantizeType, QuantizeType,
Status, Status,
@ -32,6 +33,7 @@ from zvec import (
[ [
(DataType.FLOAT, "FLOAT"), (DataType.FLOAT, "FLOAT"),
(IndexType.HNSW, "HNSW"), (IndexType.HNSW, "HNSW"),
(IOBackendType.PREAD, "PREAD"),
(MetricType.COSINE, "COSINE"), (MetricType.COSINE, "COSINE"),
(QuantizeType.INT8, "INT8"), (QuantizeType.INT8, "INT8"),
(StatusCode.OK, "OK"), (StatusCode.OK, "OK"),
@ -46,6 +48,7 @@ def test_enum_names(member, name):
[ [
(DataType.FLOAT, 8), (DataType.FLOAT, 8),
(IndexType.HNSW, 1), (IndexType.HNSW, 1),
(IOBackendType.PREAD, 0),
(MetricType.COSINE, 3), (MetricType.COSINE, 3),
(QuantizeType.INT8, 2), (QuantizeType.INT8, 2),
(StatusCode.OK, 0), (StatusCode.OK, 0),
@ -91,11 +94,29 @@ def test_data_type_has_member(member):
assert member in DataType.__members__ assert member in DataType.__members__
@pytest.mark.parametrize("member", ["HNSW", "IVF", "FLAT", "INVERT"]) @pytest.mark.parametrize(
"member",
[
"UNDEFINED",
"HNSW",
"IVF",
"FLAT",
"HNSW_RABITQ",
"DISKANN",
"VAMANA",
"INVERT",
"FTS",
],
)
def test_index_type_has_member(member): def test_index_type_has_member(member):
assert member in IndexType.__members__ assert member in IndexType.__members__
@pytest.mark.parametrize("member", ["PREAD", "LIBAIO"])
def test_io_backend_type_has_member(member):
assert member in IOBackendType.__members__
@pytest.mark.parametrize("member", ["FP16", "INT8", "INT4", "UNDEFINED"]) @pytest.mark.parametrize("member", ["FP16", "INT8", "INT4", "UNDEFINED"])
def test_quantize_type_has_member(member): def test_quantize_type_has_member(member):
assert member in QuantizeType.__members__ assert member in QuantizeType.__members__

View File

@ -39,6 +39,8 @@ try:
from zvec._zvec import ( from zvec._zvec import (
get_default_jieba_dict_dir, get_default_jieba_dict_dir,
io_backend_description,
io_backend_type,
set_default_jieba_dict_dir, set_default_jieba_dict_dir,
) )
@ -52,21 +54,6 @@ except Exception:
# Public API — grouped by category # Public API — grouped by category
# ============================== # ==============================
# —— DiskAnn runtime plugin ——
# Re-export the plugin management entry points defined by the C++ extension.
# DiskAnn normally auto-loads on first use; these APIs let tests and
# diagnostic tools preload the plugin and get a clear error if libaio is
# missing or the plugin shared object cannot be located.
from zvec._zvec import (
DISKANN_PLUGIN_DLOPEN_FAILED,
DISKANN_PLUGIN_LIBAIO_MISSING,
DISKANN_PLUGIN_OK,
DISKANN_PLUGIN_UNSUPPORTED_PLATFORM,
is_diskann_plugin_loaded,
is_libaio_available,
load_diskann_plugin,
)
from . import model as model from . import model as model
# —— Extensions —— # —— Extensions ——
@ -131,6 +118,7 @@ from .tool import require_module
from .typing import ( from .typing import (
DataType, DataType,
IndexType, IndexType,
IOBackendType,
MetricType, MetricType,
QuantizeType, QuantizeType,
Status, Status,
@ -151,6 +139,8 @@ __all__ = [
"open", "open",
"set_default_jieba_dict_dir", "set_default_jieba_dict_dir",
"get_default_jieba_dict_dir", "get_default_jieba_dict_dir",
"io_backend_type",
"io_backend_description",
# Core classes # Core classes
"Collection", "Collection",
"Doc", "Doc",
@ -204,6 +194,7 @@ __all__ = [
"QwenReRanker", "QwenReRanker",
# Typing # Typing
"DataType", "DataType",
"IOBackendType",
"MetricType", "MetricType",
"QuantizeType", "QuantizeType",
"IndexType", "IndexType",
@ -213,14 +204,6 @@ __all__ = [
"StatusCode", "StatusCode",
# Tools # Tools
"require_module", "require_module",
# DiskAnn plugin
"load_diskann_plugin",
"is_diskann_plugin_loaded",
"is_libaio_available",
"DISKANN_PLUGIN_OK",
"DISKANN_PLUGIN_UNSUPPORTED_PLATFORM",
"DISKANN_PLUGIN_LIBAIO_MISSING",
"DISKANN_PLUGIN_DLOPEN_FAILED",
] ]
# ============================== # ==============================

View File

@ -40,6 +40,7 @@ from .tool import require_module
from .typing import ( from .typing import (
DataType, DataType,
IndexType, IndexType,
IOBackendType,
MetricType, MetricType,
QuantizeType, QuantizeType,
Status, Status,
@ -48,6 +49,22 @@ from .typing import (
from .typing.enum import LogLevel, LogType from .typing.enum import LogLevel, LogType
from .zvec import create_and_open, init, open 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).
IOBackendType.LIBAIO if libaio is available, IOBackendType.PREAD otherwise."""
def io_backend_description() -> str:
"""Returns a human-readable description of the current I/O backend.
When only pread is available, includes instructions for installing
libaio to enable async I/O."""
def set_default_jieba_dict_dir(dir: str) -> None:
"""Register the process-wide default jieba dict directory."""
def get_default_jieba_dict_dir() -> str:
"""Read the currently registered default jieba dict directory."""
__all__: list = [ __all__: list = [
"AddColumnOption", "AddColumnOption",
"AlterColumnOption", "AlterColumnOption",
@ -71,6 +88,7 @@ __all__: list = [
"HnswQueryParam", "HnswQueryParam",
"HnswRabitqIndexParam", "HnswRabitqIndexParam",
"HnswRabitqQueryParam", "HnswRabitqQueryParam",
"IOBackendType",
"IVFIndexParam", "IVFIndexParam",
"IVFQueryParam", "IVFQueryParam",
"IndexOption", "IndexOption",
@ -93,9 +111,13 @@ __all__: list = [
"VectorSchema", "VectorSchema",
"WeightedReRanker", "WeightedReRanker",
"create_and_open", "create_and_open",
"get_default_jieba_dict_dir",
"init", "init",
"io_backend_description",
"io_backend_type",
"open", "open",
"require_module", "require_module",
"set_default_jieba_dict_dir",
] ]
class _Collection: class _Collection:

View File

@ -13,6 +13,8 @@ __all__: list[str] = [
"AddColumnOption", "AddColumnOption",
"AlterColumnOption", "AlterColumnOption",
"CollectionOption", "CollectionOption",
"DiskAnnIndexParam",
"DiskAnnQueryParam",
"FlatIndexParam", "FlatIndexParam",
"FtsIndexParam", "FtsIndexParam",
"FtsQueryParam", "FtsQueryParam",
@ -133,6 +135,124 @@ class CollectionOption:
@property @property
def read_only(self) -> bool: ... def read_only(self) -> bool: ...
class DiskAnnIndexParam(VectorIndexParam):
"""
Parameters for configuring a DiskAnn index.
DiskAnn stores compressed vector in memory and high-definition vector on disk. At query time,
only compressed vector will be loaded into memory. By this way, search memory at runtime is diminished.
Attributes:
metric_type (MetricType): Distance metric used for similarity computation.
Default is ``MetricType.IP`` (inner product).
max_degree (int): Maximum out-degree of each node in the Vamana graph.
Larger values improve recall at the cost of build time and index size.
Clamped to the range [1, 100]. Default is 100.
list_size (int): Candidate list size used during graph construction.
Larger values improve graph quality and recall at the cost of build time.
Clamped to the range [10, 100]. Default is 50.
pq_chunk_num (int): Number of PQ chunks used for product-quantizing the
in-memory compressed vectors. ``0`` means auto-pick based on dimension.
Clamped to the range [1, 1024]. Default is 0.
quantize_type (QuantizeType): Optional quantization type for vector
compression (e.g., FP16, INT8). Default is ``QuantizeType.UNDEFINED``.
Examples:
>>> from zvec.typing import MetricType, QuantizeType
>>> params = DiskAnnIndexParam(
... metric_type=MetricType.COSINE,
... max_degree=100,
... list_size=50,
... pq_chunk_num=8,
... quantize_type=QuantizeType.FP16
... )
>>> print(params.max_degree)
100
"""
def __getstate__(self) -> tuple: ...
def __init__(
self,
metric_type: zvec._zvec.typing.MetricType = ...,
max_degree: typing.SupportsInt = 100,
list_size: typing.SupportsInt = 50,
pq_chunk_num: typing.SupportsInt = 0,
quantize_type: zvec._zvec.typing.QuantizeType = ...,
quantizer_param: QuantizerParam = ...,
) -> None:
"""
Constructs a DiskAnnIndexParam instance.
Args:
metric_type (MetricType, optional): Distance metric. Defaults to MetricType.IP.
max_degree (int, optional): Maximum out-degree of each node in the Vamana
graph. Clamped to [1, 100]. Defaults to 100.
list_size (int, optional): Candidate list size used during graph
construction. Clamped to [10, 100]. Defaults to 50.
pq_chunk_num (int, optional): Number of PQ chunks for product
quantization. ``0`` means auto-pick based on dimension.
Clamped to [1, 1024]. Defaults to 0.
quantize_type (QuantizeType, optional): Vector quantization type.
Defaults to QuantizeType.UNDEFINED.
quantizer_param (QuantizerParam, optional): Quantizer configuration.
Defaults to QuantizerParam().
"""
def __repr__(self) -> str: ...
def __setstate__(self, arg0: tuple) -> None: ...
def to_dict(self) -> dict:
"""
Convert to dictionary with all fields
"""
@property
def max_degree(self) -> int:
"""int: Maximum out-degree of each node in the Vamana graph."""
@property
def list_size(self) -> int:
"""int: Candidate list size used during graph construction."""
@property
def pq_chunk_num(self) -> int:
"""int: Number of PQ chunks for product quantization."""
class DiskAnnQueryParam(QueryParam):
"""
Query parameters for DiskAnn index.
Attributes:
type (IndexType): Always ``IndexType.DISKANN``.
list_size (int): Beam-search candidate list size used at query time.
Higher values improve recall but increase latency. Default is 300.
Examples:
>>> params = DiskAnnQueryParam(list_size=20)
>>> print(params.list_size)
20
"""
def __getstate__(self) -> tuple: ...
def __init__(self, list_size: typing.SupportsInt = 300) -> None:
"""
Constructs a DiskAnnQueryParam instance.
Args:
list_size (int, optional): Beam-search candidate list size during
graph search. Higher values improve recall at the cost of latency.
Defaults to 300.
"""
def __repr__(self) -> str: ...
def __setstate__(self, arg0: tuple) -> None: ...
@property
def list_size(self) -> int:
"""int: Beam-search candidate list size during DiskAnn query."""
class FlatIndexParam(VectorIndexParam): class FlatIndexParam(VectorIndexParam):
""" """

View File

@ -16,6 +16,7 @@ from __future__ import annotations
from zvec._zvec.typing import ( from zvec._zvec.typing import (
DataType, DataType,
IndexType, IndexType,
IOBackendType,
MetricType, MetricType,
QuantizeType, QuantizeType,
Status, Status,
@ -24,6 +25,7 @@ from zvec._zvec.typing import (
__all__ = [ __all__ = [
"DataType", "DataType",
"IOBackendType",
"IndexType", "IndexType",
"MetricType", "MetricType",
"QuantizeType", "QuantizeType",

View File

@ -8,6 +8,7 @@ import typing
__all__: list[str] = [ __all__: list[str] = [
"DataType", "DataType",
"IOBackendType",
"IndexType", "IndexType",
"MetricType", "MetricType",
"QuantizeType", "QuantizeType",
@ -122,6 +123,48 @@ class DataType:
@property @property
def value(self) -> int: ... def value(self) -> int: ...
class IOBackendType:
"""
Enumeration of supported I/O backend types for DiskAnn async disk reads.
- PREAD: Synchronous pread() no async I/O.
- LIBAIO: libaio loaded at runtime via dlopen().
Examples:
>>> from zvec.typing import IOBackendType
>>> print(IOBackendType.LIBAIO)
IOBackendType.LIBAIO
Members:
PREAD
LIBAIO
"""
LIBAIO: typing.ClassVar[IOBackendType] # value = <IOBackendType.LIBAIO: 1>
PREAD: typing.ClassVar[IOBackendType] # value = <IOBackendType.PREAD: 0>
__members__: typing.ClassVar[
dict[str, IOBackendType]
] # value = {'PREAD': <IOBackendType.PREAD: 0>, 'LIBAIO': <IOBackendType.LIBAIO: 1>}
def __eq__(self, other: typing.Any) -> bool: ...
def __getstate__(self) -> int: ...
def __hash__(self) -> int: ...
def __index__(self) -> int: ...
def __init__(self, value: typing.SupportsInt) -> None: ...
def __int__(self) -> int: ...
def __ne__(self, other: typing.Any) -> bool: ...
def __repr__(self) -> str: ...
def __setstate__(self, state: typing.SupportsInt) -> None: ...
def __str__(self) -> str: ...
@property
def name(self) -> str: ...
@property
def value(self) -> int: ...
class IndexType: class IndexType:
""" """
@ -143,17 +186,29 @@ class IndexType:
FLAT FLAT
HNSW_RABITQ
DISKANN
VAMANA
INVERT INVERT
FTS
""" """
FLAT: typing.ClassVar[IndexType] # value = <IndexType.FLAT: 4> DISKANN: typing.ClassVar[IndexType] # value = <IndexType.DISKANN: 5>
FLAT: typing.ClassVar[IndexType] # value = <IndexType.FLAT: 3>
FTS: typing.ClassVar[IndexType] # value = <IndexType.FTS: 11>
HNSW: typing.ClassVar[IndexType] # value = <IndexType.HNSW: 1> HNSW: typing.ClassVar[IndexType] # value = <IndexType.HNSW: 1>
HNSW_RABITQ: typing.ClassVar[IndexType] # value = <IndexType.HNSW_RABITQ: 4>
INVERT: typing.ClassVar[IndexType] # value = <IndexType.INVERT: 10> INVERT: typing.ClassVar[IndexType] # value = <IndexType.INVERT: 10>
IVF: typing.ClassVar[IndexType] # value = <IndexType.IVF: 3> IVF: typing.ClassVar[IndexType] # value = <IndexType.IVF: 2>
UNDEFINED: typing.ClassVar[IndexType] # value = <IndexType.UNDEFINED: 0> UNDEFINED: typing.ClassVar[IndexType] # value = <IndexType.UNDEFINED: 0>
VAMANA: typing.ClassVar[IndexType] # value = <IndexType.VAMANA: 6>
__members__: typing.ClassVar[ __members__: typing.ClassVar[
dict[str, IndexType] dict[str, IndexType]
] # value = {'UNDEFINED': <IndexType.UNDEFINED: 0>, 'HNSW': <IndexType.HNSW: 1>, 'IVF': <IndexType.IVF: 3>, 'FLAT': <IndexType.FLAT: 4>, 'INVERT': <IndexType.INVERT: 10>} ] # value = {'UNDEFINED': <IndexType.UNDEFINED: 0>, 'HNSW': <IndexType.HNSW: 1>, 'IVF': <IndexType.IVF: 2>, 'FLAT': <IndexType.FLAT: 3>, 'HNSW_RABITQ': <IndexType.HNSW_RABITQ: 4>, 'DISKANN': <IndexType.DISKANN: 5>, 'VAMANA': <IndexType.VAMANA: 6>, 'INVERT': <IndexType.INVERT: 10>, 'FTS': <IndexType.FTS: 11>}
def __eq__(self, other: typing.Any) -> bool: ... def __eq__(self, other: typing.Any) -> bool: ...
def __getstate__(self) -> int: ... def __getstate__(self) -> int: ...
@ -236,15 +291,18 @@ class QuantizeType:
INT8 INT8
INT4 INT4
RABITQ
""" """
FP16: typing.ClassVar[QuantizeType] # value = <QuantizeType.FP16: 1> FP16: typing.ClassVar[QuantizeType] # value = <QuantizeType.FP16: 1>
INT4: typing.ClassVar[QuantizeType] # value = <QuantizeType.INT4: 3> INT4: typing.ClassVar[QuantizeType] # value = <QuantizeType.INT4: 3>
INT8: typing.ClassVar[QuantizeType] # value = <QuantizeType.INT8: 2> INT8: typing.ClassVar[QuantizeType] # value = <QuantizeType.INT8: 2>
RABITQ: typing.ClassVar[QuantizeType] # value = <QuantizeType.RABITQ: 4>
UNDEFINED: typing.ClassVar[QuantizeType] # value = <QuantizeType.UNDEFINED: 0> UNDEFINED: typing.ClassVar[QuantizeType] # value = <QuantizeType.UNDEFINED: 0>
__members__: typing.ClassVar[ __members__: typing.ClassVar[
dict[str, QuantizeType] dict[str, QuantizeType]
] # value = {'UNDEFINED': <QuantizeType.UNDEFINED: 0>, 'FP16': <QuantizeType.FP16: 1>, 'INT8': <QuantizeType.INT8: 2>, 'INT4': <QuantizeType.INT4: 3>} ] # value = {'UNDEFINED': <QuantizeType.UNDEFINED: 0>, 'FP16': <QuantizeType.FP16: 1>, 'INT8': <QuantizeType.INT8: 2>, 'INT4': <QuantizeType.INT4: 3>, 'RABITQ': <QuantizeType.RABITQ: 4>}
def __eq__(self, other: typing.Any) -> bool: ... def __eq__(self, other: typing.Any) -> bool: ...
def __getstate__(self) -> int: ... def __getstate__(self) -> int: ...

View File

@ -0,0 +1,39 @@
// 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.
// Implementation of the convenience helpers declared in the public header
// zvec/ailego/io/io_backend.h.
//
// This translation unit is the single place that bridges the dependency-free
// public header to the internal IOBackend singleton, so that public headers
// can expose current_io_backend_type() / current_io_backend_description()
// without pulling in libaio_loader or io_backend_def.h.
#include <ailego/io/io_backend_def.h>
#include <zvec/ailego/io/io_backend.h>
namespace zvec {
namespace ailego {
IOBackendType current_io_backend_type() {
return IOBackend::Instance().available();
}
std::string current_io_backend_description() {
auto type = IOBackend::Instance().available();
return IOBackendDescription(type);
}
} // namespace ailego
} // namespace zvec

View File

@ -0,0 +1,137 @@
// 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.
// Abstract I/O backend selector.
//
// Wraps the low-level loaders (LibAioLoader for libaio) and provides a uniform
// way to initialize, query, and report the active I/O backend. The actual I/O
// operations are still performed by the underlying loaders; this class is
// responsible only for backend initialization and reporting.
//
// When no async backend is available, the caller should fall back to
// synchronous pread().
//
// Usage:
// auto& backend = ailego::IOBackend::Instance();
// if (!backend.is_pread()) { ... }
// LOG_INFO("I/O backend: %s", backend.name());
#pragma once
#include <ailego/io/libaio_loader.h>
#include <zvec/ailego/io/io_backend.h>
namespace zvec {
namespace ailego {
// Returns a human-readable name for the given backend type.
inline const char *IOBackendTypeName(IOBackendType type) {
switch (type) {
case IOBackendType::kLibAio:
return "libaio";
case IOBackendType::kPread:
return "pread";
}
return "unknown";
}
// Returns a human-readable description for the given backend type.
// When the backend is kPread, 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::kPread:
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.";
}
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(IOBackendType) tries a specific backend.
// Use type() / name() to query the loaded backend without triggering a load.
class IOBackend {
public:
static IOBackend &Instance() {
static IOBackend instance;
return instance;
}
// Try to load the best available backend (libaio > pread).
// Returns the loaded backend type.
// Idempotent — if already loaded, returns immediately.
IOBackendType available() {
if (type_ != IOBackendType::kPread) {
return type_;
}
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).
// Idempotent — if the same backend is already loaded, returns immediately.
IOBackendType available(IOBackendType requested) {
if (type_ == requested && type_ != IOBackendType::kPread) {
return type_;
}
#if defined(__linux) || defined(__linux__)
if (requested == IOBackendType::kLibAio) {
if (LibAioLoader::Instance().load() &&
LibAioLoader::Instance().is_available()) {
type_ = IOBackendType::kLibAio;
return type_;
}
}
#endif
type_ = IOBackendType::kPread;
return type_;
}
bool is_pread() {
return available() == IOBackendType::kPread;
}
bool is_libaio() {
return available() == IOBackendType::kLibAio;
}
// Returns the loaded backend type.
IOBackendType type() const {
return type_;
}
// Human-readable name for the selected backend.
const char *name() const {
return IOBackendTypeName(type_);
}
// Human-readable description for the selected backend.
const char *description() const {
return IOBackendDescription(type_);
}
private:
IOBackend() = default;
IOBackendType type_{IOBackendType::kPread};
};
} // namespace ailego
} // namespace zvec

190
src/ailego/io/libaio_def.h Normal file
View File

@ -0,0 +1,190 @@
// 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.
// Private replacement for <libaio.h>.
//
// This header declares *only* the types, constants, and inline helpers that
// zvec needs from libaio. By doing so the project is completely decoupled
// from the system libaio-dev header: there is no `#include <libaio.h>` anywhere
// in the source tree, which means libaio-dev does not need to be installed at
// build time and the code is portable to cross-compilation environments that
// lack the header.
//
// The struct layouts (struct iocb, struct io_event, ...) are part of the Linux
// kernel ABI. They are copied verbatim from the upstream <libaio.h>, including
// the PADDED macros that handle architecture-specific padding. The inline
// helper io_prep_pread() is likewise copied — it only manipulates struct fields
// and does not call into the library.
#pragma once
#include <time.h> // struct timespec (used by io_getevents signature)
#include <cstring> // memset() — used by io_prep_pread() inline helper
#if defined(__linux) || defined(__linux__)
struct sockaddr;
struct iovec;
// ---------------------------------------------------------------------------
// Type and struct definitions copied from <libaio.h>
// ---------------------------------------------------------------------------
typedef struct io_context *io_context_t;
typedef enum io_iocb_cmd {
IO_CMD_PREAD = 0,
IO_CMD_PWRITE = 1,
IO_CMD_FSYNC = 2,
IO_CMD_FDSYNC = 3,
IO_CMD_POLL = 5,
IO_CMD_NOOP = 6,
IO_CMD_PREADV = 7,
IO_CMD_PWRITEV = 8,
} io_iocb_cmd_t;
// PADDED macros — copied verbatim from <libaio.h> to guarantee ABI-compatible
// struct layout on every supported architecture.
/* little endian, 32 bits */
#if defined(__i386__) || (defined(__x86_64__) && defined(__ILP32__)) || \
(defined(__arm__) && !defined(__ARMEB__)) || \
(defined(__sh__) && defined(__LITTLE_ENDIAN__)) || defined(__bfin__) || \
(defined(__MIPSEL__) && !defined(__mips64)) || defined(__cris__) || \
defined(__loongarch32) || (defined(__riscv) && __riscv_xlen == 32) || \
(defined(__GNUC__) && defined(__BYTE_ORDER__) && \
__BYTE_ORDER__ == __ORDER_LITTLE_ENDIAN__ && __SIZEOF_LONG__ == 4)
#define AIO_PADDED(x, y) \
x; \
unsigned y
#define AIO_PADDEDptr(x, y) \
x; \
unsigned y
#define AIO_PADDEDul(x, y) \
unsigned long x; \
unsigned y
/* little endian, 64 bits */
#elif defined(__ia64__) || defined(__x86_64__) || defined(__alpha__) || \
(defined(__mips64) && defined(__MIPSEL__)) || \
(defined(__aarch64__) && defined(__AARCH64EL__)) || \
defined(__loongarch64) || (defined(__riscv) && __riscv_xlen == 64) || \
(defined(__GNUC__) && defined(__BYTE_ORDER__) && \
__BYTE_ORDER__ == __ORDER_LITTLE_ENDIAN__ && __SIZEOF_LONG__ == 8)
#define AIO_PADDED(x, y) x, y
#define AIO_PADDEDptr(x, y) x
#define AIO_PADDEDul(x, y) unsigned long x
/* big endian, 64 bits */
#elif defined(__powerpc64__) || defined(__s390x__) || \
(defined(__hppa__) && defined(__arch64__)) || \
(defined(__sparc__) && defined(__arch64__)) || \
(defined(__mips64) && defined(__MIPSEB__)) || \
(defined(__aarch64__) && defined(__AARCH64EB__)) || \
(defined(__GNUC__) && defined(__BYTE_ORDER__) && \
__BYTE_ORDER__ == __ORDER_BIG_ENDIAN__ && __SIZEOF_LONG__ == 8)
#define AIO_PADDED(x, y) \
unsigned y; \
x
#define AIO_PADDEDptr(x, y) x
#define AIO_PADDEDul(x, y) unsigned long x
/* big endian, 32 bits */
#elif defined(__PPC__) || defined(__s390__) || \
(defined(__arm__) && defined(__ARMEB__)) || \
(defined(__sh__) && defined(__BIG_ENDIAN__)) || defined(__sparc__) || \
defined(__MIPSEB__) || defined(__m68k__) || defined(__hppa__) || \
defined(__frv__) || defined(__avr32__) || \
(defined(__GNUC__) && defined(__BYTE_ORDER__) && \
__BYTE_ORDER__ == __ORDER_BIG_ENDIAN__ && __SIZEOF_LONG__ == 4)
#define AIO_PADDED(x, y) \
unsigned y; \
x
#define AIO_PADDEDptr(x, y) \
unsigned y; \
x
#define AIO_PADDEDul(x, y) \
unsigned y; \
unsigned long x
#else
#error endianness?
#endif
struct io_iocb_poll {
AIO_PADDED(int events, __pad1);
};
struct io_iocb_sockaddr {
AIO_PADDEDptr(struct sockaddr *addr, __pad1);
AIO_PADDEDul(len, __pad2);
};
struct io_iocb_common {
AIO_PADDEDptr(void *buf, __pad1);
AIO_PADDEDul(nbytes, __pad2);
long long offset;
long long __pad3;
unsigned flags;
unsigned resfd;
};
struct io_iocb_vector {
AIO_PADDEDptr(const struct iovec *vec, __pad1);
AIO_PADDEDul(nr, __pad2);
long long offset;
};
struct iocb {
AIO_PADDEDptr(void *data, __pad1);
AIO_PADDED(unsigned key, aio_rw_flags);
short aio_lio_opcode;
short aio_reqprio;
int aio_fildes;
union {
struct io_iocb_common c;
struct io_iocb_vector v;
struct io_iocb_poll poll;
struct io_iocb_sockaddr saddr;
} u;
};
struct io_event {
AIO_PADDEDptr(void *data, __pad1);
AIO_PADDEDptr(struct iocb *obj, __pad2);
AIO_PADDEDul(res, __pad3);
AIO_PADDEDul(res2, __pad4);
};
#undef AIO_PADDED
#undef AIO_PADDEDptr
#undef AIO_PADDEDul
// Inline helper — copied from <libaio.h>. Only manipulates struct fields.
static inline void io_prep_pread(struct iocb *iocb, int fd, void *buf,
size_t count, long long offset) {
memset(iocb, 0, sizeof(*iocb));
iocb->aio_fildes = fd;
iocb->aio_lio_opcode = IO_CMD_PREAD;
iocb->aio_reqprio = 0;
iocb->u.c.buf = buf;
iocb->u.c.nbytes = count;
iocb->u.c.offset = offset;
}
// ---------------------------------------------------------------------------
// End: type and struct definitions from <libaio.h>
// ---------------------------------------------------------------------------
#endif // __linux__

View File

@ -0,0 +1,144 @@
// 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.
// dlopen-based wrapper for libaio. Instead of linking against -laio at build
// time, the DiskAnn plugin loads libaio at runtime via dlopen()/dlsym(). This
// removes libaio as a hard build- and load-time dependency: the plugin .so no
// longer carries a NEEDED entry for libaio.so.1, so it can be dlopen'ed on
// hosts that don't have libaio installed. Callers that actually exercise the
// async-I/O path will fall back to synchronous pread() when libaio is absent.
//
// All ABI-stable type and struct definitions (struct iocb, struct io_event,
// io_context_t, PADDED macros, io_prep_pread(), ...) now live in libaio_def.h —
// a private header that replaces <libaio.h> so the project has zero build-time
// dependency on libaio-dev.
#pragma once
#if defined(__linux) || defined(__linux__)
#include <dlfcn.h>
#include <atomic>
#include <cstring>
#include <mutex>
#include <ailego/io/libaio_def.h> // ABI-stable struct definitions (replaces <libaio.h>)
// Function-pointer typedefs for the four libaio syscalls used by DiskAnn.
typedef int (*aio_setup_fn)(int maxevents, io_context_t *ctxp);
typedef int (*aio_destroy_fn)(io_context_t ctx);
typedef int (*aio_submit_fn)(io_context_t ctx, long nr, struct iocb *ios[]);
typedef int (*aio_getevents_fn)(io_context_t ctx, long min_nr, long nr,
struct io_event *events,
struct timespec *timeout);
// Runtime loader for libaio. Thread-safe singleton that dlopen()'s libaio
// once and caches the function pointers. If libaio is not present on the
// host, all pointers remain nullptr and callers should fall back to
// synchronous I/O (pread).
//
// Usage:
// if (LibAioLoader::Instance().load()) {
// LibAioLoader::Instance().io_setup(...);
// }
class LibAioLoader {
public:
static LibAioLoader &Instance() {
static LibAioLoader instance;
return instance;
}
// Load (or confirm already loaded) libaio. Returns true on success.
// Thread-safe and idempotent.
bool load() {
if (available_.load(std::memory_order_acquire)) {
return true;
}
std::call_once(once_, [this] { this->try_load(); });
return available_.load(std::memory_order_relaxed);
}
bool is_available() const {
return available_.load(std::memory_order_acquire);
}
// Function pointers — nullptr until load() succeeds.
aio_setup_fn io_setup;
aio_destroy_fn io_destroy;
aio_submit_fn io_submit;
aio_getevents_fn io_getevents;
private:
LibAioLoader()
: io_setup(nullptr),
io_destroy(nullptr),
io_submit(nullptr),
io_getevents(nullptr) {}
~LibAioLoader() {
if (handle_ != nullptr) {
dlclose(handle_);
}
}
LibAioLoader(const LibAioLoader &) = delete;
LibAioLoader &operator=(const LibAioLoader &) = delete;
void try_load() {
// On Ubuntu 24.04 the libaio package was renamed with the t64 suffix
// (64-bit time_t transition), so probe both spellings.
static constexpr const char *kSonames[] = {
"libaio.so.1",
"libaio.so.1t64",
};
for (const char *soname : kSonames) {
void *h = dlopen(soname, RTLD_LAZY);
if (h == nullptr) {
continue;
}
io_setup = reinterpret_cast<aio_setup_fn>(dlsym(h, "io_setup"));
io_destroy = reinterpret_cast<aio_destroy_fn>(dlsym(h, "io_destroy"));
io_submit = reinterpret_cast<aio_submit_fn>(dlsym(h, "io_submit"));
// io_getevents may be redirected to io_getevents_time64 on 32-bit
// platforms compiled with _TIME_BITS=64.
io_getevents =
reinterpret_cast<aio_getevents_fn>(dlsym(h, "io_getevents"));
if (io_getevents == nullptr) {
io_getevents =
reinterpret_cast<aio_getevents_fn>(dlsym(h, "io_getevents_time64"));
}
if (io_setup && io_destroy && io_submit && io_getevents) {
handle_ = h;
available_.store(true, std::memory_order_release);
return;
}
// Some symbols missing — try the next soname.
dlclose(h);
io_setup = nullptr;
io_destroy = nullptr;
io_submit = nullptr;
io_getevents = nullptr;
}
}
std::once_flag once_;
std::atomic<bool> available_{false};
void *handle_{nullptr};
};
#endif // __linux__

View File

@ -29,6 +29,8 @@
#include <string> #include <string>
#include <unordered_map> #include <unordered_map>
#include <vector> #include <vector>
#include <ailego/io/io_backend_def.h>
#include <zvec/ailego/io/io_backend.h>
#include <zvec/db/collection.h> #include <zvec/db/collection.h>
#include <zvec/db/config.h> #include <zvec/db/config.h>
#include <zvec/db/doc.h> #include <zvec/db/doc.h>
@ -759,6 +761,25 @@ const char *zvec_get_default_jieba_dict_dir(void) {
return cached.c_str(); return cached.c_str();
} }
// =============================================================================
// I/O Backend Introspection
// =============================================================================
zvec_io_backend_type_t zvec_get_io_backend_type(void) {
auto type = zvec::ailego::current_io_backend_type();
return static_cast<zvec_io_backend_type_t>(static_cast<uint32_t>(type));
}
const char *zvec_get_io_backend_type_name(zvec_io_backend_type_t type) {
auto cpp_type = static_cast<zvec::ailego::IOBackendType>(type);
return zvec::ailego::IOBackendTypeName(cpp_type);
}
const char *zvec_get_io_backend_description(void) {
auto type = zvec::ailego::current_io_backend_type();
return zvec::ailego::IOBackendDescription(type);
}
// ============================================================================= // =============================================================================
// Error handling interface implementation // Error handling interface implementation
// ============================================================================= // =============================================================================
@ -2800,12 +2821,16 @@ const char *zvec_index_type_to_string(zvec_index_type_t index_type) {
return "IVF"; return "IVF";
case ZVEC_INDEX_TYPE_FLAT: case ZVEC_INDEX_TYPE_FLAT:
return "FLAT"; return "FLAT";
case ZVEC_INDEX_TYPE_HNSW_RABITQ:
return "HNSW_RABITQ";
case ZVEC_INDEX_TYPE_DISKANN:
return "DISKANN";
case ZVEC_INDEX_TYPE_VAMANA:
return "VAMANA";
case ZVEC_INDEX_TYPE_INVERT: case ZVEC_INDEX_TYPE_INVERT:
return "INVERT"; return "INVERT";
case ZVEC_INDEX_TYPE_FTS: case ZVEC_INDEX_TYPE_FTS:
return "FTS"; return "FTS";
case ZVEC_INDEX_TYPE_DISKANN:
return "DiskANN";
default: default:
return "UNKNOWN_INDEX_TYPE"; return "UNKNOWN_INDEX_TYPE";
} }

View File

@ -19,21 +19,6 @@ set(SRC_LISTS
pybind11_add_module(_zvec ${SRC_LISTS}) pybind11_add_module(_zvec ${SRC_LISTS})
# pybind11_add_module() defaults to CXX_VISIBILITY_PRESET=hidden +
# VISIBILITY_INLINES_HIDDEN=ON, which hides the compiler-generated helper
# symbols attached to inline functions (guard variables for static locals,
# vtables, typeinfo ...). The DiskAnn runtime plugin has its own copy of
# Factory<IndexStreamer>::Instance()'s guard variable; if _zvec.so's copy
# is hidden, the two guards are separate and the factory constructor runs
# twice during plugin load, wiping out the registrations that happened
# during _zvec.so import. We switch to default visibility here and rely on
# the version script (exports.map / exports.mac) to keep the dynamic
# symbol table small by exporting only zvec::* and PyInit_*.
set_target_properties(_zvec PROPERTIES
CXX_VISIBILITY_PRESET default
C_VISIBILITY_PRESET default
VISIBILITY_INLINES_HIDDEN OFF)
# Ensure any change to the linker version script (exports.map) triggers a # Ensure any change to the linker version script (exports.map) triggers a
# re-link of _zvec.so. target_link_options() alone is a command-line flag # re-link of _zvec.so. target_link_options() alone is a command-line flag
# and does not register the script file as a build dependency, so stale # and does not register the script file as a build dependency, so stale
@ -57,11 +42,11 @@ if (CMAKE_SYSTEM_NAME STREQUAL "Linux")
$<TARGET_FILE:core_knn_ivf_static> $<TARGET_FILE:core_knn_ivf_static>
$<TARGET_FILE:core_knn_vamana_static> $<TARGET_FILE:core_knn_vamana_static>
$<TARGET_FILE:core_knn_cluster_static> $<TARGET_FILE:core_knn_cluster_static>
$<TARGET_FILE:core_knn_diskann_static>
$<TARGET_FILE:core_mix_reducer_static> $<TARGET_FILE:core_mix_reducer_static>
$<TARGET_FILE:core_metric_static> $<TARGET_FILE:core_metric_static>
$<TARGET_FILE:core_utility_static> $<TARGET_FILE:core_utility_static>
$<TARGET_FILE:core_quantizer_static> $<TARGET_FILE:core_quantizer_static>
$<TARGET_FILE:core_plugin>
-Wl,--no-whole-archive -Wl,--no-whole-archive
zvec zvec
${CMAKE_DL_LIBS} ${CMAKE_DL_LIBS}
@ -69,21 +54,11 @@ if (CMAKE_SYSTEM_NAME STREQUAL "Linux")
target_link_options(_zvec PRIVATE target_link_options(_zvec PRIVATE
"LINKER:--version-script=${CMAKE_CURRENT_SOURCE_DIR}/exports.map" "LINKER:--version-script=${CMAKE_CURRENT_SOURCE_DIR}/exports.map"
) )
# DiskAnn is x86-only (it depends on libaio, unavailable on ARM64), so the # DiskAnn is now statically linked into _zvec.so via --whole-archive, so
# runtime plugin (libzvec_diskann_plugin.so) is only produced on non-ARM # no separate runtime .so is shipped in the wheel. libaio is still loaded
# builds. It is shipped as a runtime-loaded shared module and brought up # at runtime via dlopen() (see libaio_loader.h); if it is missing, DiskAnn
# implicitly the first time a DiskAnn index is created users never need to # fails cleanly with an actionable error while other index types
# call any load function. If libaio is missing at runtime the auto-load # (HNSW/IVF/Flat/Vamana) remain fully functional.
# fails cleanly and the error is surfaced only when DiskAnn is actually
# used; other index types (HNSW/IVF/Flat/Vamana) remain fully functional.
# The .so must therefore be discoverable next to the extension module,
# hence the $ORIGIN rpath below.
if (NOT CMAKE_SYSTEM_PROCESSOR MATCHES "aarch64|arm64|arm")
set_target_properties(_zvec PROPERTIES
BUILD_RPATH "$ORIGIN"
INSTALL_RPATH "$ORIGIN"
)
endif()
elseif (APPLE) elseif (APPLE)
target_link_libraries(_zvec PRIVATE target_link_libraries(_zvec PRIVATE
-Wl,-force_load,$<TARGET_FILE:core_knn_flat_static> -Wl,-force_load,$<TARGET_FILE:core_knn_flat_static>
@ -94,11 +69,11 @@ elseif (APPLE)
-Wl,-force_load,$<TARGET_FILE:core_knn_ivf_static> -Wl,-force_load,$<TARGET_FILE:core_knn_ivf_static>
-Wl,-force_load,$<TARGET_FILE:core_knn_vamana_static> -Wl,-force_load,$<TARGET_FILE:core_knn_vamana_static>
-Wl,-force_load,$<TARGET_FILE:core_knn_cluster_static> -Wl,-force_load,$<TARGET_FILE:core_knn_cluster_static>
-Wl,-force_load,$<TARGET_FILE:core_knn_diskann_static>
-Wl,-force_load,$<TARGET_FILE:core_mix_reducer_static> -Wl,-force_load,$<TARGET_FILE:core_mix_reducer_static>
-Wl,-force_load,$<TARGET_FILE:core_metric_static> -Wl,-force_load,$<TARGET_FILE:core_metric_static>
-Wl,-force_load,$<TARGET_FILE:core_utility_static> -Wl,-force_load,$<TARGET_FILE:core_utility_static>
-Wl,-force_load,$<TARGET_FILE:core_quantizer_static> -Wl,-force_load,$<TARGET_FILE:core_quantizer_static>
-Wl,-force_load,$<TARGET_FILE:core_plugin>
zvec zvec
) )
target_link_libraries(_zvec PRIVATE target_link_libraries(_zvec PRIVATE
@ -113,11 +88,11 @@ elseif (MSVC)
core_knn_ivf_static core_knn_ivf_static
core_knn_vamana_static core_knn_vamana_static
core_knn_cluster_static core_knn_cluster_static
core_knn_diskann_static
core_mix_reducer_static core_mix_reducer_static
core_metric_static core_metric_static
core_utility_static core_utility_static
core_quantizer_static core_quantizer_static
core_plugin
) )
target_link_libraries(_zvec PRIVATE target_link_libraries(_zvec PRIVATE
${_zvec_whole_archive_libs} ${_zvec_whole_archive_libs}

View File

@ -13,7 +13,6 @@
// limitations under the License. // limitations under the License.
#include <pybind11/pybind11.h> #include <pybind11/pybind11.h>
#include <zvec/plugin/diskann_plugin.h>
#include "python_collection.h" #include "python_collection.h"
#include "python_config.h" #include "python_config.h"
#include "python_doc.h" #include "python_doc.h"
@ -24,39 +23,6 @@
namespace zvec { namespace zvec {
namespace {
// Expose DiskAnn plugin management to Python. The DiskAnn runtime normally
// auto-loads on first use, but tests (and diagnostic tooling) need a way to
// force a load up-front and get actionable errors when libaio is missing or
// the plugin shared object cannot be located.
void InitializeDiskAnnPluginBindings(pybind11::module_ &m) {
m.def(
"load_diskann_plugin",
[](const std::string &path) { return ::zvec::LoadDiskAnnPlugin(path); },
pybind11::arg("path") = std::string(),
"Load the DiskAnn runtime plugin. Returns 0 on success or a negative "
"DiskAnnPluginStatus code on failure (unsupported platform, libaio "
"missing, or dlopen failure).");
m.def("is_diskann_plugin_loaded", &::zvec::IsDiskAnnPluginLoaded,
"Return True if the DiskAnn runtime plugin is currently loaded.");
m.def("is_libaio_available", &::zvec::IsLibAioAvailable,
"Return True if libaio is resolvable on this host (required by the "
"DiskAnn runtime).");
// Status constants so callers can compare against well-known codes without
// hard-coding integers.
m.attr("DISKANN_PLUGIN_OK") = static_cast<int>(::zvec::kDiskAnnPluginOk);
m.attr("DISKANN_PLUGIN_UNSUPPORTED_PLATFORM") =
static_cast<int>(::zvec::kDiskAnnPluginUnsupportedPlatform);
m.attr("DISKANN_PLUGIN_LIBAIO_MISSING") =
static_cast<int>(::zvec::kDiskAnnPluginLibAioMissing);
m.attr("DISKANN_PLUGIN_DLOPEN_FAILED") =
static_cast<int>(::zvec::kDiskAnnPluginDlopenFailed);
}
} // namespace
PYBIND11_MODULE(_zvec, m) { PYBIND11_MODULE(_zvec, m) {
m.doc() = "Zvec core module"; m.doc() = "Zvec core module";
@ -67,6 +33,5 @@ PYBIND11_MODULE(_zvec, m) {
ZVecPyConfig::Initialize(m); ZVecPyConfig::Initialize(m);
ZVecPyDoc::Initialize(m); ZVecPyDoc::Initialize(m);
ZVecPyCollection::Initialize(m); ZVecPyCollection::Initialize(m);
InitializeDiskAnnPluginBindings(m);
} }
} // namespace zvec } // namespace zvec

View File

@ -1,34 +1,6 @@
{ {
global: global:
# Python module entry point(s). PyInit_*; # export only python related functions
PyInit_*;
# Expose the full zvec C++ namespace so the DiskAnn runtime plugin
# (libzvec_diskann_plugin.so), which is dlopen()ed with
# RTLD_NOW | RTLD_GLOBAL after the interpreter has loaded _zvec.so,
# can resolve its undefined references against this module. Without
# this, the plugin fails to load with errors like
# undefined symbol: _ZN4zvec6ailego6Logger10LEVEL_INFOE
# because the default version script hides every internal symbol.
extern "C++" {
"zvec::*";
zvec::*;
# Also export the compiler-generated helper symbols that live
# alongside symbols in the zvec namespace (guard variables for
# static locals, vtables, typeinfo, VTT, construction vtables,
# thunks). Without these, the DiskAnn plugin and _zvec.so each
# get their own copy of e.g. the guard for Factory<T>::Instance's
# static local, which causes the factory constructor to run
# twice - wiping out registrations done during _zvec.so load.
"guard variable for zvec::*";
"vtable for zvec::*";
"VTT for zvec::*";
"typeinfo for zvec::*";
"typeinfo name for zvec::*";
"construction vtable for zvec::*";
"non-virtual thunk to zvec::*";
"virtual thunk to zvec::*";
};
local: local:
*; *;
}; };

View File

@ -10,9 +10,12 @@
// distributed under the License is distributed on an "AS IS" BASIS, // distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and // See the License for the specific language governing permissions and
// limitations under the License.#pragma once // limitations under the License.
#pragma once
#include <pybind11/pybind11.h> #include <pybind11/pybind11.h>
#include <zvec/ailego/io/io_backend.h>
#include <zvec/db/status.h> #include <zvec/db/status.h>
#include <zvec/db/type.h> #include <zvec/db/type.h>
@ -32,6 +35,7 @@ class ZVecPyTyping {
static void bind_index_types(py::module_ &m); static void bind_index_types(py::module_ &m);
static void bind_metric_types(py::module_ &m); static void bind_metric_types(py::module_ &m);
static void bind_quantize_types(py::module_ &m); static void bind_quantize_types(py::module_ &m);
static void bind_io_backend_types(py::module_ &m);
static void bind_status(py::module_ &m); static void bind_status(py::module_ &m);
}; };

View File

@ -14,6 +14,7 @@
#include "python_config.h" #include "python_config.h"
#include <pybind11/stl.h> #include <pybind11/stl.h>
#include <zvec/ailego/io/io_backend.h>
namespace zvec { namespace zvec {
@ -217,6 +218,27 @@ void ZVecPyConfig::Initialize(pybind11::module_ &m) {
"get_default_jieba_dict_dir", "get_default_jieba_dict_dir",
[]() -> std::string { return GlobalConfig::Instance().jieba_dict_dir(); }, []() -> std::string { return GlobalConfig::Instance().jieba_dict_dir(); },
"Read the currently registered default jieba dict directory."); "Read the currently registered default jieba dict directory.");
// Returns the current I/O backend type for DiskAnn async disk reads.
// Pure introspection \u2014 no side effects, no install hints.
m.def(
"io_backend_type",
[]() -> ailego::IOBackendType {
return ailego::current_io_backend_type();
},
"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, "
"IOBackendType.PREAD otherwise.");
// Returns a human-readable description of the I/O backend, including
// installation guidance for libaio when only pread is available.
m.def(
"io_backend_description",
[]() -> std::string { return ailego::current_io_backend_description(); },
"Returns a human-readable description of the current I/O backend. "
"When only pread is available, includes instructions for installing "
"libaio to enable async I/O.");
} }

View File

@ -24,6 +24,7 @@ void ZVecPyTyping::Initialize(pybind11::module_ &parent) {
bind_index_types(m); bind_index_types(m);
bind_metric_types(m); bind_metric_types(m);
bind_quantize_types(m); bind_quantize_types(m);
bind_io_backend_types(m);
bind_status(m); bind_status(m);
} }
@ -96,11 +97,13 @@ Examples:
)pbdoc") )pbdoc")
.value("UNDEFINED", IndexType::UNDEFINED) .value("UNDEFINED", IndexType::UNDEFINED)
.value("HNSW", IndexType::HNSW) .value("HNSW", IndexType::HNSW)
.value("HNSW_RABITQ", IndexType::HNSW_RABITQ)
.value("IVF", IndexType::IVF) .value("IVF", IndexType::IVF)
.value("FLAT", IndexType::FLAT) .value("FLAT", IndexType::FLAT)
.value("HNSW_RABITQ", IndexType::HNSW_RABITQ)
.value("DISKANN", IndexType::DISKANN)
.value("VAMANA", IndexType::VAMANA) .value("VAMANA", IndexType::VAMANA)
.value("INVERT", IndexType::INVERT); .value("INVERT", IndexType::INVERT)
.value("FTS", IndexType::FTS);
} }
void ZVecPyTyping::bind_metric_types(pybind11::module_ &m) { void ZVecPyTyping::bind_metric_types(pybind11::module_ &m) {
@ -137,6 +140,22 @@ Examples:
.value("RABITQ", QuantizeType::RABITQ); .value("RABITQ", QuantizeType::RABITQ);
} }
void ZVecPyTyping::bind_io_backend_types(py::module_ &m) {
py::enum_<ailego::IOBackendType>(m, "IOBackendType", R"pbdoc(
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().
Examples:
>>> from zvec.typing import IOBackendType
>>> print(IOBackendType.LIBAIO)
IOBackendType.LIBAIO
)pbdoc")
.value("PREAD", ailego::IOBackendType::kPread)
.value("LIBAIO", ailego::IOBackendType::kLibAio);
}
void ZVecPyTyping::bind_status(py::module_ &m) { void ZVecPyTyping::bind_status(py::module_ &m) {
// bind status code // bind status code
py::enum_<StatusCode>(m, "StatusCode", R"pbdoc( py::enum_<StatusCode>(m, "StatusCode", R"pbdoc(

View File

@ -49,7 +49,6 @@ cc_directory(quantizer)
cc_directory(utility) cc_directory(utility)
cc_directory(interface) cc_directory(interface)
cc_directory(mixed_reducer) cc_directory(mixed_reducer)
cc_directory(plugin)
git_version(GIT_SRCS_VER ${CMAKE_CURRENT_SOURCE_DIR}) git_version(GIT_SRCS_VER ${CMAKE_CURRENT_SOURCE_DIR})
file(GLOB_RECURSE ALL_CORE_SRCS *.cc *.c *.h) file(GLOB_RECURSE ALL_CORE_SRCS *.cc *.c *.h)
@ -63,17 +62,17 @@ endif()
# Always exclude algorithm/diskann implementation files from zvec_core. # Always exclude algorithm/diskann implementation files from zvec_core.
# The DiskAnn algorithm is provided by the separate core_knn_diskann library # The DiskAnn algorithm is provided by the separate core_knn_diskann library
# (real on Linux x86_64, stub on other platforms). Including them here causes # (STATIC+SHARED, real on Linux x86_64, stub on other platforms). The static
# duplicate symbols and missing -laio when test binaries link both zvec_core # variant is whole-archived into _zvec.so for the Python wheel; the shared
# (via zvec) and core_knn_diskann. # variant is used by C++ tools and tests. Including the sources here would
# cause duplicate symbols.
list(FILTER ALL_CORE_SRCS EXCLUDE REGEX ".*/algorithm/diskann/.*") list(FILTER ALL_CORE_SRCS EXCLUDE REGEX ".*/algorithm/diskann/.*")
if(NOT DISKANN_SUPPORTED) if(NOT DISKANN_SUPPORTED)
list(FILTER ALL_CORE_SRCS EXCLUDE REGEX ".*/interface/indexes/diskann_index\\.cc") list(FILTER ALL_CORE_SRCS EXCLUDE REGEX ".*/interface/indexes/diskann_index\\.cc")
endif() endif()
set(ZVEC_CORE_LIBS zvec_ailego zvec_turbo sparsehash magic_enum rabitqlib) set(ZVEC_CORE_LIBS zvec_ailego zvec_turbo sparsehash magic_enum rabitqlib)
# The plugin loader uses dlopen/dlsym, so link libdl on Linux. # The DiskAnn runtime loader uses dlopen/dlsym, so link libdl on Linux.
if(CMAKE_SYSTEM_NAME STREQUAL "Linux") if(CMAKE_SYSTEM_NAME STREQUAL "Linux")
list(APPEND ZVEC_CORE_LIBS ${CMAKE_DL_LIBS}) list(APPEND ZVEC_CORE_LIBS ${CMAKE_DL_LIBS})
endif() endif()

View File

@ -23,7 +23,12 @@ else()
if(MSVC) if(MSVC)
# MSVC: STATIC-only stub to avoid creating an empty DLL with no exports # MSVC: STATIC-only stub to avoid creating an empty DLL with no exports
# (MSVC linker fails to produce an import library when there are zero exports) # (MSVC linker fails to produce an import library when there are zero exports).
# cc_library with STATIC-only creates target "core_knn_diskann" but NOT the
# "core_knn_diskann_static" variant (that only happens with STATIC+SHARED).
# The Python binding references core_knn_diskann_static on all platforms, so
# create an ALIAS so the same target name works under MSVC as well.
# ($<TARGET_FILE:> supports ALIAS targets since CMake 3.18.)
cc_library( cc_library(
NAME core_knn_diskann NAME core_knn_diskann
STATIC STRICT ALWAYS_LINK STATIC STRICT ALWAYS_LINK
@ -32,6 +37,7 @@ else()
INCS . ${PROJECT_ROOT_DIR}/src ${PROJECT_ROOT_DIR}/src/core ${PROJECT_ROOT_DIR}/src/core/algorithm INCS . ${PROJECT_ROOT_DIR}/src ${PROJECT_ROOT_DIR}/src/core ${PROJECT_ROOT_DIR}/src/core/algorithm
VERSION "${PROXIMA_ZVEC_VERSION}" VERSION "${PROXIMA_ZVEC_VERSION}"
) )
add_library(core_knn_diskann_static ALIAS core_knn_diskann)
else() else()
cc_library( cc_library(
NAME core_knn_diskann NAME core_knn_diskann

View File

@ -1,57 +1,34 @@
include(${PROJECT_ROOT_DIR}/cmake/bazel.cmake) include(${PROJECT_ROOT_DIR}/cmake/bazel.cmake)
include(${PROJECT_ROOT_DIR}/cmake/option.cmake)
file(GLOB_RECURSE ALL_SRCS *.cc *.c) file(GLOB_RECURSE ALL_SRCS *.cc *.c)
# The DiskAnn plugin is loaded at runtime via zvec::LoadDiskAnnPlugin() with # DiskAnn is compiled as a STATIC+SHARED library (matching the pattern used by
# RTLD_GLOBAL, so its undefined references to internal zvec symbols # core_knn_hnsw, core_knn_cluster, etc.). The static variant
# (core_framework, core_knn_cluster, zvec_ailego, ...) are resolved at load # (core_knn_diskann_static) is whole-archived into _zvec.so for the Python
# time against the hosting binary (_zvec.so for the Python extension, the # wheel, so no separate runtime .so is needed. The shared variant is used by
# test executable for gtest, or libzvec_core for tools). # C++ tools and tests that link against it directly.
# #
# As a consequence the plugin .so must NOT carry NEEDED entries for those # libaio is loaded at runtime via dlopen()/dlsym() (see libaio_loader.h), so
# libs: otherwise the dynamic loader would try to resolve them from disk # we do NOT link against -laio. The only system-level link dependency is
# before invoking dlopen's RTLD_GLOBAL symbol sharing, and the Python wheel # ${CMAKE_DL_LIBS} for dlopen/dlsym/dlclose.
# (which only ships _zvec.so + the plugin) would fail to load. set(CORE_KNN_DISKANN_LIBS core_framework core_knn_cluster)
#
# We therefore link the plugin only against its system-level dependency
# (libaio) and rely on the global include dirs plus -Wl,--unresolved-symbols
# =ignore-all to let the linker build a shared library with unresolved
# references to internal APIs.
set(CORE_KNN_DISKANN_LIBS "")
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")
list(APPEND CORE_KNN_DISKANN_LIBS aio) list(APPEND CORE_KNN_DISKANN_LIBS ${CMAKE_DL_LIBS})
endif()
if(NOT APPLE)
set(CORE_KNN_DISKANN_LDFLAGS
"-Wl,--exclude-libs,libparquet.a:libarrow.a:libarrow_bundled_dependencies.a")
endif() endif()
cc_library( cc_library(
NAME core_knn_diskann NAME core_knn_diskann
SHARED STRICT STATIC SHARED STRICT ALWAYS_LINK
SRCS *.cc SRCS *.cc
LIBS ${CORE_KNN_DISKANN_LIBS} LIBS ${CORE_KNN_DISKANN_LIBS}
INCS . ${PROJECT_ROOT_DIR}/src/core ${PROJECT_ROOT_DIR}/src/core/algorithm INCS . ${PROJECT_ROOT_DIR}/src/core ${PROJECT_ROOT_DIR}/src/core/algorithm
LDFLAGS "${CORE_KNN_DISKANN_LDFLAGS}"
VERSION "${PROXIMA_ZVEC_VERSION}" VERSION "${PROXIMA_ZVEC_VERSION}"
)
# Internal zvec libs are referenced only for header availability; the actual
# symbol resolution happens at dlopen time. Expose their public include dirs
# without adding them as link dependencies.
foreach(_dep zvec_ailego core_framework core_knn_cluster)
if(TARGET ${_dep})
target_include_directories(core_knn_diskann PRIVATE
$<TARGET_PROPERTY:${_dep},INTERFACE_INCLUDE_DIRECTORIES>)
add_dependencies(core_knn_diskann ${_dep})
endif()
endforeach()
# Allow the plugin to have unresolved symbols that will be satisfied at
# dlopen(RTLD_NOW | RTLD_GLOBAL) time by the hosting binary.
if(CMAKE_SYSTEM_NAME STREQUAL "Linux")
target_link_options(core_knn_diskann PRIVATE
"LINKER:--unresolved-symbols=ignore-all")
endif()
# Rename the artifact to libzvec_diskann_plugin.so so it is clearly identified
# as an optional plugin that users load at runtime via zvec::LoadDiskAnnPlugin().
set_target_properties(core_knn_diskann PROPERTIES
OUTPUT_NAME zvec_diskann_plugin
) )

View File

@ -32,6 +32,8 @@ namespace core {
int DiskAnnBuilder::init(const IndexMeta &meta, const ailego::Params &params) { int DiskAnnBuilder::init(const IndexMeta &meta, const ailego::Params &params) {
LOG_INFO("Begin DiskAnnBuilder::init"); LOG_INFO("Begin DiskAnnBuilder::init");
log_diskann_io_backend();
params.get(PARAM_DISKANN_BUILDER_MAX_DEGREE, &max_degree_); params.get(PARAM_DISKANN_BUILDER_MAX_DEGREE, &max_degree_);
params.get(PARAM_DISKANN_BUILDER_LIST_SIZE, &list_size_); params.get(PARAM_DISKANN_BUILDER_LIST_SIZE, &list_size_);
params.get(PARAM_DISKANN_BUILDER_THREAD_COUNT, &build_thread_count_); params.get(PARAM_DISKANN_BUILDER_THREAD_COUNT, &build_thread_count_);

View File

@ -18,6 +18,8 @@
#include <cstdio> #include <cstdio>
#include <cstring> #include <cstring>
#include <iostream> #include <iostream>
#include <ailego/io/io_backend_def.h>
#include <zvec/ailego/io/io_backend.h>
#include <zvec/core/framework/index_logger.h> #include <zvec/core/framework/index_logger.h>
#define MAX_EVENTS 1024 #define MAX_EVENTS 1024
@ -28,11 +30,36 @@ namespace core {
#if (defined(__linux) || defined(__linux__)) #if (defined(__linux) || defined(__linux__))
typedef struct io_event io_event_t; typedef struct io_event io_event_t;
typedef struct iocb iocb_t; typedef struct iocb iocb_t;
// Ensures the I/O backend selection is logged exactly once per process,
// regardless of which entry point (setup_io_ctx or register_thread)
// triggers it first.
static std::once_flag g_io_backend_log_once;
#endif #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) { int setup_io_ctx(IOContext &ctx) {
#if (defined(__linux) || defined(__linux__)) #if (defined(__linux) || defined(__linux__))
int ret = io_setup(MAX_EVENTS, &ctx); std::call_once(g_io_backend_log_once, log_diskann_io_backend);
if (ailego::IOBackend::Instance().is_pread()) {
return 0;
}
int ret = LibAioLoader::Instance().io_setup(MAX_EVENTS, &ctx);
return ret; return ret;
#else #else
@ -42,7 +69,10 @@ int setup_io_ctx(IOContext &ctx) {
int destroy_io_ctx(IOContext &ctx) { int destroy_io_ctx(IOContext &ctx) {
#if (defined(__linux) || defined(__linux__)) #if (defined(__linux) || defined(__linux__))
int ret = io_destroy(ctx); if (ailego::IOBackend::Instance().is_pread()) {
return 0;
}
int ret = LibAioLoader::Instance().io_destroy(ctx);
return ret; return ret;
#else #else
@ -71,6 +101,9 @@ static int execute_io_pread(int fd, std::vector<AlignedRead> &read_reqs) {
int execute_io(IOContext ctx, int fd, std::vector<AlignedRead> &read_reqs, int execute_io(IOContext ctx, int fd, std::vector<AlignedRead> &read_reqs,
uint64_t n_retries = 0) { uint64_t n_retries = 0) {
#if (defined(__linux) || defined(__linux__)) #if (defined(__linux) || defined(__linux__))
if (ailego::IOBackend::Instance().is_pread()) {
return execute_io_pread(fd, read_reqs);
}
uint64_t iters = DiskAnnUtil::div_round_up(read_reqs.size(), MAX_EVENTS); uint64_t iters = DiskAnnUtil::div_round_up(read_reqs.size(), MAX_EVENTS);
for (uint64_t iter = 0; iter < iters; iter++) { for (uint64_t iter = 0; iter < iters; iter++) {
@ -93,7 +126,8 @@ int execute_io(IOContext ctx, int fd, std::vector<AlignedRead> &read_reqs,
size_t n_tries = 0; size_t n_tries = 0;
// Phase 1: io_submit with retry. // Phase 1: io_submit with retry.
while (true) { while (true) {
int ret = io_submit(ctx, (int64_t)n_ops, cbs.data()); int ret =
LibAioLoader::Instance().io_submit(ctx, (int64_t)n_ops, cbs.data());
if (ret == (int)n_ops) { if (ret == (int)n_ops) {
break; break;
} }
@ -111,8 +145,8 @@ int execute_io(IOContext ctx, int fd, std::vector<AlignedRead> &read_reqs,
// Phase 2: io_getevents with retry (never re-submits). // Phase 2: io_getevents with retry (never re-submits).
n_tries = 0; n_tries = 0;
while (true) { while (true) {
int ret = io_getevents(ctx, (int64_t)n_ops, (int64_t)n_ops, evts.data(), int ret = LibAioLoader::Instance().io_getevents(
nullptr); ctx, (int64_t)n_ops, (int64_t)n_ops, evts.data(), nullptr);
if (ret == (int)n_ops) { if (ret == (int)n_ops) {
break; break;
} }
@ -188,16 +222,19 @@ void LinuxAlignedFileReader::register_thread() {
IOContext ctx = nullptr; IOContext ctx = nullptr;
int ret = io_setup(MAX_EVENTS, &ctx); std::call_once(g_io_backend_log_once, log_diskann_io_backend);
if (ret != 0) { if (ailego::IOBackend::Instance().is_pread()) {
lk.unlock(); lk.unlock();
return;
}
int ret = LibAioLoader::Instance().io_setup(MAX_EVENTS, &ctx);
if (ret != 0) {
if (ret == -EAGAIN) { if (ret == -EAGAIN) {
LOG_ERROR( LOG_ERROR(
"io_setup failed with EAGAIN: Consider increasing " "io_setup failed with EAGAIN: Consider increasing "
"/proc/sys/fs/aio-max-nr"); "/proc/sys/fs/aio-max-nr");
} else { } else {
LOG_ERROR("io_setup failed; returned: %d, %s", ret, ::strerror(-ret)); LOG_ERROR("io_setup failed; returned: %d, %s", ret, ::strerror(-ret));
;
} }
} else { } else {
LOG_INFO("allocating ctx: %lu", (uint64_t)ctx); LOG_INFO("allocating ctx: %lu", (uint64_t)ctx);
@ -226,7 +263,10 @@ void LinuxAlignedFileReader::deregister_thread() {
} }
// io_destroy is a syscall; keep it outside the lock to avoid blocking others // io_destroy is a syscall; keep it outside the lock to avoid blocking others
io_destroy(ctx); if (ailego::IOBackend::Instance().available() !=
ailego::IOBackendType::kPread) {
LibAioLoader::Instance().io_destroy(ctx);
}
LOG_INFO("returned ctx from thread"); LOG_INFO("returned ctx from thread");
#endif #endif
} }
@ -234,9 +274,13 @@ void LinuxAlignedFileReader::deregister_thread() {
void LinuxAlignedFileReader::deregister_all_threads() { void LinuxAlignedFileReader::deregister_all_threads() {
#if (defined(__linux) || defined(__linux__)) #if (defined(__linux) || defined(__linux__))
std::unique_lock<std::mutex> lk(ctx_mut); std::unique_lock<std::mutex> lk(ctx_mut);
bool aio_available = ailego::IOBackend::Instance().available() !=
ailego::IOBackendType::kPread;
for (auto x = ctx_map.begin(); x != ctx_map.end(); x++) { for (auto x = ctx_map.begin(); x != ctx_map.end(); x++) {
IOContext ctx = x->second; IOContext ctx = x->second;
io_destroy(ctx); if (aio_available) {
LibAioLoader::Instance().io_destroy(ctx);
}
} }
ctx_map.clear(); ctx_map.clear();
#endif #endif

View File

@ -18,7 +18,7 @@
#include <fcntl.h> #include <fcntl.h>
#if (defined(__linux) || defined(__linux__)) #if (defined(__linux) || defined(__linux__))
#include <libaio.h> #include <ailego/io/libaio_loader.h> // dlopen-based libaio wrapper
#endif #endif
#include <unistd.h> #include <unistd.h>
@ -39,6 +39,10 @@ typedef uint32_t IOContext;
int setup_io_ctx(IOContext &ctx); int setup_io_ctx(IOContext &ctx);
int destroy_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 { struct AlignedRead {
uint64_t offset; uint64_t offset;
uint64_t len; uint64_t len;

View File

@ -25,6 +25,8 @@ DiskAnnSearcher::DiskAnnSearcher() {}
DiskAnnSearcher::~DiskAnnSearcher() {} DiskAnnSearcher::~DiskAnnSearcher() {}
int DiskAnnSearcher::init(const ailego::Params &search_params) { int DiskAnnSearcher::init(const ailego::Params &search_params) {
log_diskann_io_backend();
search_params.get(PARAM_DISKANN_SEARCHER_LIST_SIZE, &list_size_); search_params.get(PARAM_DISKANN_SEARCHER_LIST_SIZE, &list_size_);
search_params.get(PARAM_DISKANN_SEARCHER_CACHE_NODE_NUM, &cache_nodes_num_); search_params.get(PARAM_DISKANN_SEARCHER_CACHE_NODE_NUM, &cache_nodes_num_);
return 0; return 0;

View File

@ -28,6 +28,9 @@ DiskAnnStreamer::~DiskAnnStreamer() {}
int DiskAnnStreamer::init(const IndexMeta &meta, int DiskAnnStreamer::init(const IndexMeta &meta,
const ailego::Params &search_params) { const ailego::Params &search_params) {
meta_ = meta; meta_ = meta;
log_diskann_io_backend();
search_params.get(PARAM_DISKANN_SEARCHER_LIST_SIZE, &list_size_); search_params.get(PARAM_DISKANN_SEARCHER_LIST_SIZE, &list_size_);
search_params.get(PARAM_DISKANN_SEARCHER_CACHE_NODE_NUM, &cache_nodes_num_); search_params.get(PARAM_DISKANN_SEARCHER_CACHE_NODE_NUM, &cache_nodes_num_);
return 0; return 0;

View File

@ -1,10 +1,15 @@
include(${PROJECT_ROOT_DIR}/cmake/bazel.cmake) include(${PROJECT_ROOT_DIR}/cmake/bazel.cmake)
include(${PROJECT_ROOT_DIR}/cmake/option.cmake) include(${PROJECT_ROOT_DIR}/cmake/option.cmake)
set(CORE_INTERFACE_LIBS zvec_ailego core_framework sparsehash magic_enum rabitqlib)
if(CMAKE_SYSTEM_NAME STREQUAL "Linux")
list(APPEND CORE_INTERFACE_LIBS ${CMAKE_DL_LIBS})
endif()
cc_library( cc_library(
NAME core_interface STATIC STRICT ALWAYS_LINK NAME core_interface STATIC STRICT ALWAYS_LINK
SRCS *.cc indexes/*.cc SRCS *.cc indexes/*.cc
INCS . ${PROJECT_ROOT_DIR}/src/ ${PROJECT_ROOT_DIR}/src/core INCS . ${PROJECT_ROOT_DIR}/src/ ${PROJECT_ROOT_DIR}/src/core
LIBS zvec_ailego core_framework core_plugin sparsehash magic_enum rabitqlib LIBS ${CORE_INTERFACE_LIBS}
VERSION "${PROXIMA_ZVEC_VERSION}" VERSION "${PROXIMA_ZVEC_VERSION}"
) )

View File

@ -16,61 +16,12 @@
#include <mutex> #include <mutex>
#include <string> #include <string>
#include <zvec/core/interface/index.h> #include <zvec/core/interface/index.h>
#include <zvec/plugin/diskann_plugin.h>
#include "algorithm/diskann/diskann_params.h" #include "algorithm/diskann/diskann_params.h"
#include "holder_builder.h" #include "holder_builder.h"
namespace zvec::core_interface { namespace zvec::core_interface {
namespace {
// Implicitly bring the DiskAnn runtime online on first use. This keeps the
// DiskAnn index an ordinary public API (users just instantiate a
// DiskAnnIndexParam) while still letting the rest of the library — HNSW,
// IVF, Flat, Vamana — run on hosts that happen to lack libaio. On such
// hosts only DiskAnn fails, with a clear, actionable error message, and
// every other index type stays fully functional.
int EnsureDiskAnnRuntimeReady() {
static std::once_flag once;
static int cached_result = 0;
std::call_once(once, []() {
const int status = ::zvec::LoadDiskAnnPlugin();
if (status == kDiskAnnPluginOk) {
cached_result = 0;
return;
}
switch (status) {
case kDiskAnnPluginLibAioMissing:
LOG_ERROR(
"DiskAnn requires libaio at runtime, but it was not found on this "
"host. Install it (e.g. 'apt-get install libaio1' on "
"Debian/Ubuntu, "
"or 'libaio1t64' on Ubuntu 24.04+) and retry.");
break;
case kDiskAnnPluginUnsupportedPlatform:
LOG_ERROR("DiskAnn is only supported on Linux x86_64.");
break;
case kDiskAnnPluginDlopenFailed:
default:
LOG_ERROR("Failed to initialize the DiskAnn runtime (status=%d).",
status);
break;
}
cached_result = core::IndexError_Runtime;
});
return cached_result;
}
} // namespace
int DiskAnnIndex::CreateAndInitStreamer(const BaseIndexParam &param) { int DiskAnnIndex::CreateAndInitStreamer(const BaseIndexParam &param) {
// Fail fast and cleanly if the DiskAnn runtime cannot be brought up on
// this host (most commonly: libaio is missing). The rest of zvec keeps
// running; only DiskAnn is unusable.
if (int rc = EnsureDiskAnnRuntimeReady(); rc != 0) {
return rc;
}
if (is_sparse_) { if (is_sparse_) {
LOG_ERROR("Failed to create streamer. Sparse is not Supported."); LOG_ERROR("Failed to create streamer. Sparse is not Supported.");
return core::IndexError_Unsupported; return core::IndexError_Unsupported;

View File

@ -1,29 +0,0 @@
##
## Copyright (C) The Software Authors. All rights reserved.
##
## \file CMakeLists.txt
## \brief Build script for the zvec internal plugin-loading glue library.
## Lives inside the main zvec_core artifact and provides the
## implicit DiskAnn runtime bring-up used by DiskAnnIndex on first
## use. These APIs are NOT part of the public user surface: users
## simply instantiate ``DiskAnnIndexParam`` / ``DiskAnnIndex`` and
## the runtime (``libzvec_diskann_plugin.so``) is loaded behind
## the scenes.
##
include(${PROJECT_ROOT_DIR}/cmake/bazel.cmake)
include(${PROJECT_ROOT_DIR}/cmake/option.cmake)
set(CORE_PLUGIN_LIBS zvec_ailego core_framework)
if(CMAKE_SYSTEM_NAME STREQUAL "Linux")
list(APPEND CORE_PLUGIN_LIBS ${CMAKE_DL_LIBS})
endif()
cc_library(
NAME core_plugin
STATIC STRICT ALWAYS_LINK
SRCS *.cc
LIBS ${CORE_PLUGIN_LIBS}
INCS . ${PROJECT_ROOT_DIR}/src/core
VERSION "${PROXIMA_ZVEC_VERSION}"
)

View File

@ -1,322 +0,0 @@
// 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 <atomic>
#include <mutex>
#include <string>
#include <vector>
#include <zvec/core/framework/index_logger.h>
#include <zvec/plugin/diskann_plugin.h>
#if defined(__linux__) || defined(__linux) || defined(__APPLE__)
#include <dlfcn.h>
#include <unistd.h>
#endif
#if defined(__linux__) || defined(__linux)
#include <limits.h>
#endif
namespace zvec {
namespace {
#if defined(__linux__) || defined(__linux)
constexpr const char *kPluginFileName = "libzvec_diskann_plugin.so";
// Candidate soname list. On Ubuntu 24.04 the libaio package was renamed with
// the t64 suffix (64-bit time_t transition), so we probe both spellings.
constexpr const char *kLibAioSoNames[] = {
"libaio.so.1",
"libaio.so.1t64",
};
constexpr bool kPlatformSupportsDiskAnnPlugin = true;
#elif defined(__APPLE__)
[[maybe_unused]] constexpr const char *kPluginFileName =
"libzvec_diskann_plugin.dylib";
constexpr bool kPlatformSupportsDiskAnnPlugin = false;
#else
[[maybe_unused]] constexpr const char *kPluginFileName =
"zvec_diskann_plugin.dll";
constexpr bool kPlatformSupportsDiskAnnPlugin = false;
#endif
// Global plugin handle. Nullptr means "not loaded".
std::atomic<void *> g_plugin_handle{nullptr};
std::mutex g_plugin_mutex;
#if defined(__linux__) || defined(__linux)
// Resolve the directory containing the currently running executable, so we
// can look for the plugin next to it regardless of the working directory.
std::string GetExecutableDir() {
char buf[PATH_MAX];
ssize_t n = ::readlink("/proc/self/exe", buf, sizeof(buf) - 1);
if (n <= 0) {
return {};
}
buf[n] = '\0';
std::string path(buf);
auto slash = path.find_last_of('/');
if (slash == std::string::npos) {
return {};
}
return path.substr(0, slash);
}
// Resolve the directory containing the shared object that hosts this
// function. For Python wheels this is the directory of
// ``_zvec.cpython-*.so``; for regular C++ binaries it is the directory of
// ``libzvec_core.so``. In either case the DiskAnn plugin is shipped
// alongside, so this is the most reliable lookup location.
//
// NOTE: we pass the address of an *exported* function (LoadDiskAnnPlugin)
// rather than one in this anonymous namespace, because dladdr() on a symbol
// with internal linkage can report the main executable instead of the
// hosting shared object when the translation unit is whole-archived into
// another .so.
std::string ResolveHostingSoDir() {
::Dl_info info{};
if (::dladdr(reinterpret_cast<void *>(&::zvec::LoadDiskAnnPlugin), &info) ==
0 ||
info.dli_fname == nullptr) {
return {};
}
std::string path(info.dli_fname);
auto slash = path.find_last_of('/');
if (slash == std::string::npos) {
return {};
}
return path.substr(0, slash);
}
// Full path of the shared object that hosts LoadDiskAnnPlugin, or empty
// string on failure.
std::string ResolveHostingSoPath() {
::Dl_info info{};
if (::dladdr(reinterpret_cast<void *>(&::zvec::LoadDiskAnnPlugin), &info) ==
0 ||
info.dli_fname == nullptr) {
return {};
}
return std::string(info.dli_fname);
}
// Promote the hosting shared object (e.g. the Python extension module
// ``_zvec.cpython-*.so``) to the global symbol scope. Python loads C
// extensions with RTLD_LOCAL by default, which means their C++ symbols are
// invisible to subsequently dlopen(RTLD_GLOBAL)ed libraries. Without this
// promotion, the DiskAnn plugin's undefined references to zvec:: symbols
// cannot be resolved against the already-loaded host module and dlopen
// fails with messages like:
//
// undefined symbol: _ZN4zvec6ailego6Logger10LEVEL_INFOE
//
// Using RTLD_NOLOAD re-opens the existing image without loading it again,
// while RTLD_GLOBAL merges its symbols into the global scope. This is a
// no-op for hosts that were already loaded with RTLD_GLOBAL.
void PromoteHostingSoToGlobal() {
const std::string host = ResolveHostingSoPath();
if (host.empty()) {
return;
}
// When LoadDiskAnnPlugin is statically linked into the main executable,
// dladdr resolves to the executable itself. The main exe's symbols are
// already in the global scope by definition, so skip promotion.
const std::string exe_path = GetExecutableDir();
if (!exe_path.empty()) {
char exe_buf[PATH_MAX];
ssize_t n = ::readlink("/proc/self/exe", exe_buf, sizeof(exe_buf) - 1);
if (n > 0) {
exe_buf[n] = '\0';
// Compare resolved real paths to handle relative vs absolute.
char host_real[PATH_MAX];
char exe_real[PATH_MAX];
if (::realpath(host.c_str(), host_real) != nullptr &&
::realpath(exe_buf, exe_real) != nullptr &&
std::string(host_real) == std::string(exe_real)) {
// Host IS the main executable; symbols are already global.
return;
}
}
}
void *h = ::dlopen(host.c_str(), RTLD_NOW | RTLD_GLOBAL | RTLD_NOLOAD);
if (h == nullptr) {
const char *err = ::dlerror();
LOG_WARN("Could not promote host '%s' to RTLD_GLOBAL: %s", host.c_str(),
err ? err : "unknown");
return;
}
// We purposely keep the handle refcount incremented for the life of the
// process: there is no safe point at which to dlclose() it, and doing so
// would only decrement the count anyway (the image remains mapped because
// Python still holds its own reference).
(void)h;
}
// Build the list of candidate paths for the plugin.
std::vector<std::string> BuildCandidatePaths(const std::string &explicit_path) {
std::vector<std::string> candidates;
if (!explicit_path.empty()) {
candidates.push_back(explicit_path);
return candidates;
}
// Helper that pushes both ``<dir>/<plugin>`` and ``<dir>/../lib/<plugin>``
// to handle the conventional CMake build layout where executables live in
// ``bin/`` while shared objects (including this plugin) live in ``lib/``.
auto push_dir_candidates = [&candidates](const std::string &dir) {
if (dir.empty()) {
return;
}
candidates.push_back(dir + "/" + kPluginFileName);
candidates.push_back(dir + "/../lib/" + kPluginFileName);
};
// 1. Directory of the library that hosts LoadDiskAnnPlugin (e.g. the
// Python extension module or libzvec_core). This works for Python,
// C++ embedding, and most packaging layouts.
const std::string own_dir = ResolveHostingSoDir();
push_dir_candidates(own_dir);
// 2. Directory of the running executable. Useful for self-contained C++
// tools that drop the plugin next to their binary, as well as for the
// standard CMake bin/lib split (handled by ``../lib/`` above).
const std::string exe_dir = GetExecutableDir();
if (!exe_dir.empty() && exe_dir != own_dir) {
push_dir_candidates(exe_dir);
}
// 3. Fallback: rely on the dynamic linker's default search path
// (RPATH / LD_LIBRARY_PATH / /etc/ld.so.conf).
candidates.emplace_back(kPluginFileName);
return candidates;
}
#endif // linux
} // namespace
bool IsLibAioAvailable() {
#if defined(__linux__) || defined(__linux)
const char *kRequiredSymbols[] = {"io_setup", "io_submit", "io_getevents",
"io_destroy"};
for (const char *soname : kLibAioSoNames) {
// RTLD_LAZY keeps the cost low; we only need to know whether the library
// is resolvable and exposes the symbols DiskAnn actually calls.
void *handle = ::dlopen(soname, RTLD_LAZY);
if (handle == nullptr) {
continue;
}
bool ok = true;
for (const char *sym : kRequiredSymbols) {
if (::dlsym(handle, sym) == nullptr) {
ok = false;
break;
}
}
::dlclose(handle);
if (ok) {
return true;
}
}
return false;
#else
return false;
#endif
}
bool IsDiskAnnPluginLoaded() {
return g_plugin_handle.load(std::memory_order_acquire) != nullptr;
}
int LoadDiskAnnPlugin(const std::string &path) {
if (!kPlatformSupportsDiskAnnPlugin) {
LOG_ERROR(
"DiskAnn plugin is not supported on this platform; it is only "
"available on Linux x86_64 with libaio.");
return kDiskAnnPluginUnsupportedPlatform;
}
#if defined(__linux__) || defined(__linux)
// Fast path: already loaded.
if (g_plugin_handle.load(std::memory_order_acquire) != nullptr) {
return kDiskAnnPluginOk;
}
std::lock_guard<std::mutex> lock(g_plugin_mutex);
if (g_plugin_handle.load(std::memory_order_relaxed) != nullptr) {
return kDiskAnnPluginOk;
}
if (!IsLibAioAvailable()) {
LOG_ERROR(
"libaio is not available on this host; the DiskAnn runtime cannot be "
"activated. Install libaio1 (e.g. 'apt-get install libaio1', or "
"'libaio1t64' on Ubuntu 24.04+) and retry. This does not affect "
"other index types (HNSW, IVF, Flat, Vamana).");
return kDiskAnnPluginLibAioMissing;
}
const std::vector<std::string> candidates = BuildCandidatePaths(path);
// Ensure the hosting module's C++ symbols (zvec::*) are visible to the
// plugin at dlopen time. See PromoteHostingSoToGlobal() for the rationale.
PromoteHostingSoToGlobal();
void *handle = nullptr;
std::string last_error;
for (const std::string &candidate : candidates) {
// RTLD_GLOBAL so the plugin's factory registrations (which live in the
// plugin's own static-init code) can reference symbols from the main
// library, and any callers that later dlsym against the process can see
// the plugin's symbols.
handle = ::dlopen(candidate.c_str(), RTLD_NOW | RTLD_GLOBAL);
if (handle != nullptr) {
LOG_INFO("Loaded DiskAnn plugin from: %s", candidate.c_str());
break;
}
const char *err = ::dlerror();
last_error = err ? err : "unknown dlopen error";
LOG_DEBUG("dlopen(%s) failed: %s", candidate.c_str(), last_error.c_str());
}
if (handle == nullptr) {
LOG_ERROR("Failed to load DiskAnn plugin; last error: %s",
last_error.c_str());
return kDiskAnnPluginDlopenFailed;
}
g_plugin_handle.store(handle, std::memory_order_release);
return kDiskAnnPluginOk;
#else
(void)path;
return kDiskAnnPluginUnsupportedPlatform;
#endif
}
bool UnloadDiskAnnPlugin() {
#if defined(__linux__) || defined(__linux)
std::lock_guard<std::mutex> lock(g_plugin_mutex);
void *handle = g_plugin_handle.exchange(nullptr, std::memory_order_acq_rel);
if (handle == nullptr) {
return false;
}
if (::dlclose(handle) != 0) {
const char *err = ::dlerror();
LOG_WARN("dlclose for DiskAnn plugin returned non-zero: %s",
err ? err : "unknown");
}
return true;
#else
return false;
#endif
}
} // namespace zvec

View File

@ -20,7 +20,6 @@
#include <zvec/db/schema.h> #include <zvec/db/schema.h>
#include <zvec/db/status.h> #include <zvec/db/status.h>
#include <zvec/db/type.h> #include <zvec/db/type.h>
#include <zvec/plugin/diskann_plugin.h>
#include "ailego/internal/cpu_features.h" #include "ailego/internal/cpu_features.h"
#include "db/common/constants.h" #include "db/common/constants.h"
#include "db/common/typedef.h" #include "db/common/typedef.h"
@ -154,7 +153,8 @@ Status FieldSchema::validate() const {
support_dense_vector_index.end()) { support_dense_vector_index.end()) {
return Status::InvalidArgument( return Status::InvalidArgument(
"schema validate failed: dense_vector's index_params only " "schema validate failed: dense_vector's index_params only "
"support FLAT|HNSW|IVF index, but field[", "support FLAT|HNSW|HNSW_RABITQ|IVF|DISKANN|VAMANA index, but "
"field[",
name_, "]'s index_type is ", name_, "]'s index_type is ",
IndexTypeCodeBook::AsString(index_params_->type())); IndexTypeCodeBook::AsString(index_params_->type()));
} }
@ -198,30 +198,20 @@ Status FieldSchema::validate() const {
} }
if (index_params_->type() == IndexType::DISKANN) { if (index_params_->type() == IndexType::DISKANN) {
// Probe the DiskAnn runtime eagerly at creation time so unsupported // DiskAnn requires Linux x86_64/i686/i386. The CMake variable
// platforms (non Linux x86_64), missing libaio, or a missing plugin // DISKANN_SUPPORTED (defined in the top-level CMakeLists.txt) is the
// .so fail fast with a clear message instead of surfacing later during // single source of truth for platform eligibility — it is also used by
// optimize(). This reuses the same gate DiskAnnIndex applies on first // index_factory.cc to conditionally compile the DiskAnn index
// use (zvec::LoadDiskAnnPlugin, wrapped by EnsureDiskAnnRuntimeReady). // registration. Using the same macro here ensures that schema
// All validate() call sites are creation-time only, so triggering the // validation and index registration agree on supported platforms.
// plugin load here is safe (and idempotent/cached). //
const int rc = ::zvec::LoadDiskAnnPlugin(); // libaio is loaded eagerly (via dlopen) inside DiskAnnBuilder::init()
switch (rc) { // and DiskAnnStreamer::init(); if libaio is missing, DiskAnn falls
case kDiskAnnPluginOk: // back to synchronous pread() with degraded performance.
break; #if !DISKANN_SUPPORTED
case kDiskAnnPluginUnsupportedPlatform: return Status::NotSupported(
return Status::NotSupported( "DiskAnn is not supported on this platform (Linux x86_64 only)");
"DiskAnn is not supported on this platform (Linux x86_64 " #endif
"only)");
case kDiskAnnPluginLibAioMissing:
return Status::NotSupported(
"DiskAnn requires libaio at runtime, but it was not found on "
"this host. Install it (e.g. 'apt-get install libaio1', or "
"'libaio1t64' on Ubuntu 24.04+) and retry.");
default:
return Status::NotSupported(
"DiskAnn runtime could not be initialized on this host");
}
} }

View File

@ -0,0 +1,46 @@
// 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.
// I/O backend type enum.
//
// This is the public, dependency-free part of the I/O backend abstraction.
// It defines the IOBackendType enum and the convenience helpers
// current_io_backend_type() / current_io_backend_description() so that
// public headers can reference IOBackendType without pulling in the
// internal IOBackend singleton or libaio_loader.
#pragma once
#include <string>
#include <zvec/ailego/internal/platform.h>
namespace zvec {
namespace ailego {
// Supported I/O backend types.
enum class IOBackendType {
kPread, // Synchronous pread() — no async I/O
kLibAio, // libaio loaded at runtime via dlopen()
};
// Returns the currently active I/O backend type.
// Triggers backend initialization on first call (libaio > pread).
IOBackendType current_io_backend_type();
// Returns a human-readable description of the currently active I/O backend.
// When only pread is available, includes installation guidance for libaio.
std::string current_io_backend_description();
} // namespace ailego
} // namespace zvec

View File

@ -120,15 +120,18 @@ class Factory {
//! a guard variable (`_ZGVZN...E7factory`) alongside the object, and both //! a guard variable (`_ZGVZN...E7factory`) alongside the object, and both
//! must be unified across DSOs for the singleton to be shared. //! must be unified across DSOs for the singleton to be shared.
//! //!
//! In the Python extension build the _zvec.so version script exports the //! Historically (when DiskAnn was a runtime-loaded shared plugin) the
//! storage (`zvec::*` matches its demangled name) while the guard variable, //! _zvec.so version script exported the storage (`zvec::*` matches its
//! whose demangled form is `guard variable for zvec::...`, ends up hidden //! demangled name) while the guard variable, whose demangled form is
//! (compilers emit the guard in a COMDAT group whose visibility is not //! `guard variable for zvec::...`, ended up hidden (compilers emit the
//! upgraded by our version script). When libzvec_diskann_plugin.so is //! guard in a COMDAT group whose visibility is not upgraded by our version
//! loaded, _zvec.so and the plugin then share the `factory` storage but //! script). When the plugin was loaded, _zvec.so and the plugin shared the
//! each have their own guard; the plugin's still-zero guard triggers a //! `factory` storage but each had its own guard; the plugin's still-zero
//! second run of the Factory constructor on the shared storage, wiping //! guard triggered a second run of the Factory constructor on the shared
//! all registrations performed during _zvec.so import (e.g. FlatStreamer). //! storage, wiping out registrations performed during _zvec.so import.
//! DiskAnn is now statically linked into _zvec.so, so this DSO-splitting
//! issue no longer applies, but the atomic-pointer pattern is retained as
//! a defensive measure.
//! //!
//! A constant-initialized static std::atomic<T*> has NO guard variable //! A constant-initialized static std::atomic<T*> has NO guard variable
//! (its zero init is compile-time), so we use a leaked heap singleton with //! (its zero init is compile-time), so we use a leaked heap singleton with

View File

@ -774,6 +774,51 @@ ZVEC_EXPORT void ZVEC_CALL zvec_set_default_jieba_dict_dir(const char *dir);
*/ */
ZVEC_EXPORT const char *ZVEC_CALL zvec_get_default_jieba_dict_dir(void); ZVEC_EXPORT const char *ZVEC_CALL zvec_get_default_jieba_dict_dir(void);
// =============================================================================
// I/O Backend Introspection
// =============================================================================
/**
* @brief I/O backend type codes for DiskAnn async disk reads.
*
* Defined as uint32_t constants for consistent binary representation
* across C and C++ boundaries.
*/
typedef uint32_t zvec_io_backend_type_t;
#define ZVEC_IO_BACKEND_TYPE_PREAD \
0 /**< Synchronous pread() \u2014 no async I/O */
#define ZVEC_IO_BACKEND_TYPE_LIBAIO \
1 /**< libaio loaded at runtime via dlopen() */
/**
* @brief Get the current I/O backend type for DiskAnn async disk reads.
*
* Pure introspection \u2014 no side effects, no install hints.
*
* @return zvec_io_backend_type_t The loaded backend type
* (ZVEC_IO_BACKEND_TYPE_LIBAIO or ZVEC_IO_BACKEND_TYPE_PREAD).
*/
ZVEC_EXPORT zvec_io_backend_type_t ZVEC_CALL zvec_get_io_backend_type(void);
/**
* @brief Get a human-readable name for the given I/O backend type.
*
* @param type The backend type code.
* @return Thread-local string valid until the next call on this thread;
* "libaio", "pread", or "unknown".
*/
ZVEC_EXPORT const char *ZVEC_CALL
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.
*
* @return Thread-local string valid until the next call on this thread.
*/
ZVEC_EXPORT const char *ZVEC_CALL zvec_get_io_backend_description(void);
// ============================================================================= // =============================================================================
// Data Type Enumerations // Data Type Enumerations
// ============================================================================= // =============================================================================
@ -826,6 +871,7 @@ typedef uint32_t zvec_index_type_t;
#define ZVEC_INDEX_TYPE_HNSW 1 #define ZVEC_INDEX_TYPE_HNSW 1
#define ZVEC_INDEX_TYPE_IVF 2 #define ZVEC_INDEX_TYPE_IVF 2
#define ZVEC_INDEX_TYPE_FLAT 3 #define ZVEC_INDEX_TYPE_FLAT 3
#define ZVEC_INDEX_TYPE_HNSW_RABITQ 4
#define ZVEC_INDEX_TYPE_DISKANN 5 #define ZVEC_INDEX_TYPE_DISKANN 5
#define ZVEC_INDEX_TYPE_VAMANA 6 #define ZVEC_INDEX_TYPE_VAMANA 6
#define ZVEC_INDEX_TYPE_INVERT 10 #define ZVEC_INDEX_TYPE_INVERT 10

View File

@ -1,84 +0,0 @@
// 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.
#pragma once
#include <string>
// NOTE: The APIs declared in this header are INTERNAL to zvec. They are
// invoked implicitly by ``DiskAnnIndex`` on first use so that DiskAnn works
// out of the box, without users ever calling a ``load_diskann_plugin()`` /
// ``is_libaio_available()`` entry point. External callers should not depend
// on these symbols; they may change or be removed in future releases. On
// hosts missing libaio the bring-up fails cleanly and only DiskAnn becomes
// unavailable — other index types (HNSW / IVF / Flat / Vamana) keep working.
#if defined(_WIN32) || defined(__CYGWIN__)
#ifdef ZVEC_BUILD_SHARED
#define ZVEC_PLUGIN_EXPORT __declspec(dllexport)
#elif defined(ZVEC_USE_SHARED)
#define ZVEC_PLUGIN_EXPORT __declspec(dllimport)
#else
#define ZVEC_PLUGIN_EXPORT
#endif
#else
#if __GNUC__ >= 4
#define ZVEC_PLUGIN_EXPORT __attribute__((visibility("default")))
#else
#define ZVEC_PLUGIN_EXPORT
#endif
#endif
namespace zvec {
// Return codes for LoadDiskAnnPlugin().
enum DiskAnnPluginStatus {
kDiskAnnPluginOk = 0,
kDiskAnnPluginUnsupportedPlatform = -1,
kDiskAnnPluginLibAioMissing = -2,
kDiskAnnPluginDlopenFailed = -3,
};
// Returns true if libaio is present on the host and the minimum set of symbols
// required by the DiskAnn runtime (io_setup / io_submit / io_getevents /
// io_destroy) can be resolved at runtime.
//
// Internal probe used by ``LoadDiskAnnPlugin`` before attempting dlopen. Not
// part of the user-facing API.
ZVEC_PLUGIN_EXPORT bool IsLibAioAvailable();
// Load the DiskAnn runtime shared library (libzvec_diskann_plugin.so) via
// dlopen(). Invoked implicitly by ``DiskAnnIndex::CreateAndInitStreamer`` on
// first use; callers should not invoke it directly. The call is idempotent
// and returns ``kDiskAnnPluginOk`` when the runtime is already active.
//
// Search order when ``path`` is empty:
// 1. next to the currently running executable (``/proc/self/exe`` on Linux);
// 2. next to the hosting shared object (e.g. ``_zvec.cpython-*.so``);
// 3. the platform default dynamic-linker search path (RPATH,
// ``LD_LIBRARY_PATH``, ``/etc/ld.so.conf``, …).
ZVEC_PLUGIN_EXPORT int LoadDiskAnnPlugin(const std::string &path = "");
// Returns true if the DiskAnn runtime is currently loaded in this process.
// Internal diagnostic; not a user-facing API.
ZVEC_PLUGIN_EXPORT bool IsDiskAnnPluginLoaded();
// Unload the DiskAnn runtime. Internal only — unloading a library that has
// registered itself into global factory singletons is inherently racy;
// callers must guarantee that no DiskAnn objects are still alive and no
// background threads are executing DiskAnn code before calling this. Returns
// true when a live handle was released.
ZVEC_PLUGIN_EXPORT bool UnloadDiskAnnPlugin();
} // namespace zvec

View File

@ -6221,7 +6221,7 @@ void test_diskann_index_params_functions(void) {
// to_string should report DiskANN // to_string should report DiskANN
const char *type_str = zvec_index_type_to_string(ZVEC_INDEX_TYPE_DISKANN); const char *type_str = zvec_index_type_to_string(ZVEC_INDEX_TYPE_DISKANN);
TEST_ASSERT(type_str != NULL && strcmp(type_str, "DiskANN") == 0); TEST_ASSERT(type_str != NULL && strcmp(type_str, "DISKANN") == 0);
zvec_index_params_destroy(params); zvec_index_params_destroy(params);
TEST_END(); TEST_END();

View File

@ -7,7 +7,7 @@ foreach(CC_SRCS ${ALL_TEST_SRCS})
cc_gtest( cc_gtest(
NAME ${CC_TARGET} NAME ${CC_TARGET}
STRICT STRICT
LIBS zvec_ailego core_framework core_utility core_metric core_quantizer core_knn_cluster core_plugin core_knn_diskann LIBS zvec_ailego core_framework core_utility core_metric core_quantizer core_knn_cluster core_knn_diskann
SRCS ${CC_SRCS} SRCS ${CC_SRCS}
INCS . ${PROJECT_ROOT_DIR}/src/core ${PROJECT_ROOT_DIR}/src/core/algorithm/diskann INCS . ${PROJECT_ROOT_DIR}/src/core ${PROJECT_ROOT_DIR}/src/core/algorithm/diskann
) )

View File

@ -143,18 +143,6 @@ TEST_F(DiskAnnBuilderTest, SmallDatasetBuildTime) {
<< " ms — likely a lost-wakeup regression in progress loops."; << " ms — likely a lost-wakeup regression in progress loops.";
} }
// DiskAnn is now exposed implicitly: no caller ever invokes a
// ``LoadDiskAnnPlugin`` / ``IsLibAioAvailable`` API (those were removed from
// the public surface together with ``zvec.load_diskann_plugin()`` in Python).
// The only contract this test validates is the UX guarantee: once the DiskAnn
// module has been linked into the hosting binary (here, directly into the
// test via the ``core_knn_diskann`` target), its factory entries are
// registered automatically and the global ``IndexFactory`` can hand out a
// ``DiskAnnBuilder`` without any explicit setup step. On hosts missing
// libaio, DiskAnn would fail at the index-creation layer with a clear error
// while other index types (HNSW/IVF/Flat/Vamana) remain unaffected; that
// runtime branch lives in ``DiskAnnIndex::CreateAndInitStreamer`` and is
// covered by the higher-level interface tests.
TEST_F(DiskAnnBuilderTest, TestImplicitFactoryRegistration) { TEST_F(DiskAnnBuilderTest, TestImplicitFactoryRegistration) {
IndexBuilder::Pointer builder = IndexFactory::CreateBuilder("DiskAnnBuilder"); IndexBuilder::Pointer builder = IndexFactory::CreateBuilder("DiskAnnBuilder");
ASSERT_NE(builder, nullptr) ASSERT_NE(builder, nullptr)

View File

@ -8,7 +8,7 @@ foreach(CC_SRCS ${ALL_TEST_SRCS})
NAME ${CC_TARGET} NAME ${CC_TARGET}
STRICT STRICT
LIBS zvec_ailego core_framework core_metric core_interface core_knn_flat core_utility core_quantizer sparsehash core_knn_hnsw core_mix_reducer LIBS zvec_ailego core_framework core_metric core_interface core_knn_flat core_utility core_quantizer sparsehash core_knn_hnsw core_mix_reducer
core_knn_flat_sparse core_knn_hnsw_sparse core_knn_ivf core_knn_hnsw_rabitq core_knn_vamana core_knn_flat_sparse core_knn_hnsw_sparse core_knn_ivf core_knn_hnsw_rabitq core_knn_vamana core_knn_diskann
SRCS ${CC_SRCS} SRCS ${CC_SRCS}
INCS . ${PROJECT_ROOT_DIR}/src/core ${PROJECT_ROOT_DIR}/src/core/algorithm INCS . ${PROJECT_ROOT_DIR}/src/core ${PROJECT_ROOT_DIR}/src/core/algorithm
) )

View File

@ -885,7 +885,7 @@ int main(int argc, char *argv[]) {
: "debug"; : "debug";
transform(log_level.begin(), log_level.end(), log_level.begin(), ::tolower); transform(log_level.begin(), log_level.end(), log_level.begin(), ::tolower);
if (LOG_LEVEL.find(log_level) != LOG_LEVEL.end()) { if (LOG_LEVEL.find(log_level) != LOG_LEVEL.end()) {
IndexLoggerBroker::SetLevel(LOG_LEVEL[log_level]); zvec::ailego::LoggerBroker::SetLevel(LOG_LEVEL[log_level]);
} }
// Calculate Bench // Calculate Bench

View File

@ -1882,7 +1882,7 @@ int main(int argc, char *argv[]) {
: "debug"; : "debug";
transform(log_level.begin(), log_level.end(), log_level.begin(), ::tolower); transform(log_level.begin(), log_level.end(), log_level.begin(), ::tolower);
if (LOG_LEVEL.find(log_level) != LOG_LEVEL.end()) { if (LOG_LEVEL.find(log_level) != LOG_LEVEL.end()) {
IndexLoggerBroker::SetLevel(LOG_LEVEL[log_level]); zvec::ailego::LoggerBroker::SetLevel(LOG_LEVEL[log_level]);
} }
// Calculate Recall // Calculate Recall