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

Co-authored-by: Zefeng Yin <yinzefeng.yzf@alibaba-inc.com>

Author: richyreachy

.github/workflows/03-macos-linux-build.yml

@@ -60,14 +60,6 @@ jobs:
           sudo apt-get install -y clang libomp-dev
         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
         if: runner.os == 'Linux'
         run: lscpu
@@ -183,3 +175,41 @@ jobs:
           ./c_api_index_example
           ./c_api_optimized_example
         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

.github/workflows/08-cmake-subproject-integration.yml

@@ -34,13 +34,6 @@ jobs:
           cache: 'pip'
           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
         run: |
           NPROC=$(nproc 2>/dev/null || echo 2)

.github/workflows/clang_tidy.yml

@@ -47,7 +47,7 @@ jobs:
         if: steps.changed_files.outputs.any_changed == 'true'
         run: |
           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
         if: steps.changed_files.outputs.any_changed == 'true'

.github/workflows/nightly_coverage.yml

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

CMakeLists.txt

@@ -122,7 +122,7 @@ else()
 endif()
 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)
   set(DISKANN_SUPPORTED ON)
   add_definitions(-DDISKANN_SUPPORTED=1)
@@ -180,24 +180,9 @@ if(BUILD_PYTHON_BINDINGS)
     install(TARGETS _zvec LIBRARY DESTINATION ${ZVEC_PY_INSTALL_DIR}/zvec
             COMPONENT python)
 
-    # DiskAnn ships as a runtime-loaded shared module
-    # (libzvec_diskann_plugin.so) that is brought online implicitly the
-    # first time a DiskAnn index is created — users never call any load
-    # 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()
+    # DiskAnn is now statically linked into _zvec.so via --whole-archive
+    # (see src/binding/python/CMakeLists.txt -> core_knn_diskann_static),
+    # so no separate runtime .so needs to be installed in the wheel.
     # Bundle cppjieba's dictionary files so the `jieba` FTS tokenizer works
     # out of the box. python/zvec/__init__.py resolves this directory via
     # importlib.resources and registers it with set_default_jieba_dict_dir().

examples/c++/CMakeLists.txt

@@ -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_LIB_DIR ${ZVEC_ROOT_DIR}/${HOST_BUILD_DIR}/lib)
 
-# Add include and library search paths
 include_directories(${ZVEC_INCLUDE_DIR})
 set(ZVEC_LIB_SEARCH_DIRS ${ZVEC_LIB_DIR})
 

pyproject.toml

@@ -181,11 +181,6 @@ test-command = "cd {project} && pytest python/tests -v --tb=short"
 build-verbosity = 1
 
 [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"]
 environment = { CMAKE_GENERATOR = "Unix Makefiles", CMAKE_BUILD_PARALLEL_LEVEL = "16" }
 manylinux-x86_64-image = "manylinux_2_28"

python/tests/detail/fixture_helper.py

@@ -13,19 +13,14 @@
 from zvec.typing import DataType, StatusCode, MetricType, QuantizeType
 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_DONE: bool = False
 
 
 def _ensure_diskann_runtime_or_reason() -> str | None:
-    """Preload the DiskAnn plugin and return None on success or a human-readable
-    skip reason on failure. Idempotent across calls."""
+    """Check whether DiskAnn is available on this platform and return None
+    on success or a human-readable skip reason on failure. Idempotent across
+    calls."""
     global _DISKANN_PRELOAD_DONE, _DISKANN_PRELOAD_REASON
     if _DISKANN_PRELOAD_DONE:
         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"
         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
     return None
 
@@ -148,8 +127,7 @@ def full_schema_new(request) -> CollectionSchema:
     else:
         nullable, has_index, vector_index = True, False, HnswIndexParam()
 
-    # Skip DiskAnn tests on unsupported platforms or when the runtime cannot
-    # be brought up (missing libaio, plugin .so not installed, etc.).
+    # Skip DiskAnn tests on unsupported platforms.
     from zvec.model.param import DiskAnnIndexParam
 
     if isinstance(vector_index, DiskAnnIndexParam):

python/tests/test_collection_diskann.py

@@ -13,17 +13,15 @@
 # limitations under the License.
 """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:
 
 1. DiskAnn is currently built only for Linux x86_64 — other platforms are
    skipped wholesale.
-2. The DiskAnn backend lives in a *runtime-loaded* plugin
-   (``libzvec_diskann_plugin.so``). It must be loaded with ``RTLD_GLOBAL |
-   RTLD_NOW`` BEFORE ``import zvec`` so that the plugin's ``IndexFactory``
-   singleton is unified with the one inside ``_zvec.so``. After ``import
-   zvec`` we must also call ``zvec.load_diskann_plugin()`` exactly once.
+2. libaio is loaded eagerly (via dlopen) inside DiskAnnBuilder::init() /
+   DiskAnnStreamer::init(). If libaio is missing, DiskAnn falls back to
+   synchronous pread() — the tests still run but with degraded performance.
 
 If either prerequisite fails the whole module is skipped so the rest of the
 test-suite is not affected.
@@ -32,7 +30,6 @@
 from __future__ import annotations
 
 import math
-import os
 import platform
 import sys
 
@@ -46,13 +43,6 @@
     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
 
 from zvec import (  # noqa: E402

python/tests/test_typing.py

@@ -17,6 +17,7 @@
 from zvec import (
     DataType,
     IndexType,
+    IOBackendType,
     MetricType,
     QuantizeType,
     Status,
@@ -32,6 +33,7 @@
     [
         (DataType.FLOAT, "FLOAT"),
         (IndexType.HNSW, "HNSW"),
+        (IOBackendType.PREAD, "PREAD"),
         (MetricType.COSINE, "COSINE"),
         (QuantizeType.INT8, "INT8"),
         (StatusCode.OK, "OK"),
@@ -46,6 +48,7 @@ def test_enum_names(member, name):
     [
         (DataType.FLOAT, 8),
         (IndexType.HNSW, 1),
+        (IOBackendType.PREAD, 0),
         (MetricType.COSINE, 3),
         (QuantizeType.INT8, 2),
         (StatusCode.OK, 0),
@@ -91,11 +94,29 @@ def test_data_type_has_member(member):
     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):
     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"])
 def test_quantize_type_has_member(member):
     assert member in QuantizeType.__members__