fecee94af1
* feat: finance bench * docs: results * chore: ignroe data README * feat: fix financebench * feat: laion, also required idmaps support * style: format * style: format * fix: resolve ruff linting errors - Remove unused variables in benchmark scripts - Rename unused loop variables to follow convention * feat: enron email bench * experiments for running DiskANN & BM25 on Arch 4090 * style: format * chore(ci): remove paru-bin submodule and config to fix checkout --recurse-submodules * docs: data * docs: data updated * fix: as package * fix(ci): only run pre-commit * chore: use http url of astchunk; use group for some dev deps * fix(ci): should checkout modules as well since `uv sync` checks * fix(ci): run with lint only * fix: find links to install wheels available * CI: force local wheels in uv install step * CI: install local wheels via file paths * CI: pick wheels matching current Python tag * CI: handle python tag mismatches for local wheels * CI: use matrix python venv and set macOS deployment target * CI: revert install step to match main * CI: use uv group install with local wheel selection * CI: rely on setup-uv for Python and tighten group install * CI: install build deps with uv python interpreter * CI: use temporary uv venv for build deps * CI: add build venv scripts path for wheel repair
473 lines
19 KiB
Python
473 lines
19 KiB
Python
import contextlib
|
|
import logging
|
|
import os
|
|
import struct
|
|
import sys
|
|
from pathlib import Path
|
|
from typing import Any, Literal, Optional
|
|
|
|
import numpy as np
|
|
import psutil
|
|
from leann.interface import (
|
|
LeannBackendBuilderInterface,
|
|
LeannBackendFactoryInterface,
|
|
LeannBackendSearcherInterface,
|
|
)
|
|
from leann.registry import register_backend
|
|
from leann.searcher_base import BaseSearcher
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
@contextlib.contextmanager
|
|
def suppress_cpp_output_if_needed():
|
|
"""Suppress C++ stdout/stderr based on LEANN_LOG_LEVEL"""
|
|
# In CI we avoid fiddling with low-level file descriptors to prevent aborts
|
|
if os.getenv("CI") == "true":
|
|
yield
|
|
return
|
|
|
|
log_level = os.getenv("LEANN_LOG_LEVEL", "WARNING").upper()
|
|
|
|
# Only suppress if log level is WARNING or higher (ERROR, CRITICAL)
|
|
should_suppress = log_level in ["WARNING", "ERROR", "CRITICAL"]
|
|
|
|
if not should_suppress:
|
|
# Don't suppress, just yield
|
|
yield
|
|
return
|
|
|
|
# Save original file descriptors
|
|
stdout_fd = sys.stdout.fileno()
|
|
stderr_fd = sys.stderr.fileno()
|
|
|
|
# Save original stdout/stderr
|
|
stdout_dup = os.dup(stdout_fd)
|
|
stderr_dup = os.dup(stderr_fd)
|
|
|
|
try:
|
|
# Redirect to /dev/null
|
|
devnull = os.open(os.devnull, os.O_WRONLY)
|
|
os.dup2(devnull, stdout_fd)
|
|
os.dup2(devnull, stderr_fd)
|
|
os.close(devnull)
|
|
|
|
yield
|
|
|
|
finally:
|
|
# Restore original file descriptors
|
|
os.dup2(stdout_dup, stdout_fd)
|
|
os.dup2(stderr_dup, stderr_fd)
|
|
os.close(stdout_dup)
|
|
os.close(stderr_dup)
|
|
|
|
|
|
def _get_diskann_metrics():
|
|
from . import _diskannpy as diskannpy # type: ignore
|
|
|
|
return {
|
|
"mips": diskannpy.Metric.INNER_PRODUCT,
|
|
"l2": diskannpy.Metric.L2,
|
|
"cosine": diskannpy.Metric.COSINE,
|
|
}
|
|
|
|
|
|
@contextlib.contextmanager
|
|
def chdir(path):
|
|
original_dir = os.getcwd()
|
|
os.chdir(path)
|
|
try:
|
|
yield
|
|
finally:
|
|
os.chdir(original_dir)
|
|
|
|
|
|
def _write_vectors_to_bin(data: np.ndarray, file_path: Path):
|
|
num_vectors, dim = data.shape
|
|
with open(file_path, "wb") as f:
|
|
f.write(struct.pack("I", num_vectors))
|
|
f.write(struct.pack("I", dim))
|
|
f.write(data.tobytes())
|
|
|
|
|
|
def _calculate_smart_memory_config(data: np.ndarray) -> tuple[float, float]:
|
|
"""
|
|
Calculate smart memory configuration for DiskANN based on data size and system specs.
|
|
|
|
Args:
|
|
data: The embedding data array
|
|
|
|
Returns:
|
|
tuple: (search_memory_maximum, build_memory_maximum) in GB
|
|
"""
|
|
num_vectors, dim = data.shape
|
|
|
|
# Calculate embedding storage size
|
|
embedding_size_bytes = num_vectors * dim * 4 # float32 = 4 bytes
|
|
embedding_size_gb = embedding_size_bytes / (1024**3)
|
|
|
|
# search_memory_maximum: 1/10 of embedding size for optimal PQ compression
|
|
# This controls Product Quantization size - smaller means more compression
|
|
search_memory_gb = max(0.1, embedding_size_gb / 10) # At least 100MB
|
|
|
|
# build_memory_maximum: Based on available system RAM for sharding control
|
|
# This controls how much memory DiskANN uses during index construction
|
|
available_memory_gb = psutil.virtual_memory().available / (1024**3)
|
|
total_memory_gb = psutil.virtual_memory().total / (1024**3)
|
|
|
|
# Use 50% of available memory, but at least 2GB and at most 75% of total
|
|
build_memory_gb = max(2.0, min(available_memory_gb * 0.5, total_memory_gb * 0.75))
|
|
|
|
logger.info(
|
|
f"Smart memory config - Data: {embedding_size_gb:.2f}GB, "
|
|
f"Search mem: {search_memory_gb:.2f}GB (PQ control), "
|
|
f"Build mem: {build_memory_gb:.2f}GB (sharding control)"
|
|
)
|
|
|
|
return search_memory_gb, build_memory_gb
|
|
|
|
|
|
@register_backend("diskann")
|
|
class DiskannBackend(LeannBackendFactoryInterface):
|
|
@staticmethod
|
|
def builder(**kwargs) -> LeannBackendBuilderInterface:
|
|
return DiskannBuilder(**kwargs)
|
|
|
|
@staticmethod
|
|
def searcher(index_path: str, **kwargs) -> LeannBackendSearcherInterface:
|
|
return DiskannSearcher(index_path, **kwargs)
|
|
|
|
|
|
class DiskannBuilder(LeannBackendBuilderInterface):
|
|
def __init__(self, **kwargs):
|
|
self.build_params = kwargs
|
|
|
|
def _safe_cleanup_after_partition(self, index_dir: Path, index_prefix: str):
|
|
"""
|
|
Safely cleanup files after partition.
|
|
In partition mode, C++ doesn't read _disk.index content,
|
|
so we can delete it if all derived files exist.
|
|
"""
|
|
disk_index_file = index_dir / f"{index_prefix}_disk.index"
|
|
beam_search_file = index_dir / f"{index_prefix}_disk_beam_search.index"
|
|
|
|
# Required files that C++ partition mode needs
|
|
# Note: C++ generates these with _disk.index suffix
|
|
disk_suffix = "_disk.index"
|
|
required_files = [
|
|
f"{index_prefix}{disk_suffix}_medoids.bin", # Critical: assert fails if missing
|
|
# Note: _centroids.bin is not created in single-shot build - C++ handles this automatically
|
|
f"{index_prefix}_pq_pivots.bin", # PQ table
|
|
f"{index_prefix}_pq_compressed.bin", # PQ compressed vectors
|
|
]
|
|
|
|
# Check if all required files exist
|
|
missing_files = []
|
|
for filename in required_files:
|
|
file_path = index_dir / filename
|
|
if not file_path.exists():
|
|
missing_files.append(filename)
|
|
|
|
if missing_files:
|
|
logger.warning(
|
|
f"Cannot safely delete _disk.index - missing required files: {missing_files}"
|
|
)
|
|
logger.info("Keeping all original files for safety")
|
|
return
|
|
|
|
# Calculate space savings
|
|
space_saved = 0
|
|
files_to_delete = []
|
|
|
|
if disk_index_file.exists():
|
|
space_saved += disk_index_file.stat().st_size
|
|
files_to_delete.append(disk_index_file)
|
|
|
|
if beam_search_file.exists():
|
|
space_saved += beam_search_file.stat().st_size
|
|
files_to_delete.append(beam_search_file)
|
|
|
|
# Safe to delete!
|
|
for file_to_delete in files_to_delete:
|
|
try:
|
|
os.remove(file_to_delete)
|
|
logger.info(f"✅ Safely deleted: {file_to_delete.name}")
|
|
except Exception as e:
|
|
logger.warning(f"Failed to delete {file_to_delete.name}: {e}")
|
|
|
|
if space_saved > 0:
|
|
space_saved_mb = space_saved / (1024 * 1024)
|
|
logger.info(f"💾 Space saved: {space_saved_mb:.1f} MB")
|
|
|
|
# Show what files are kept
|
|
logger.info("📁 Kept essential files for partition mode:")
|
|
for filename in required_files:
|
|
file_path = index_dir / filename
|
|
if file_path.exists():
|
|
size_mb = file_path.stat().st_size / (1024 * 1024)
|
|
logger.info(f" - {filename} ({size_mb:.1f} MB)")
|
|
|
|
def build(self, data: np.ndarray, ids: list[str], index_path: str, **kwargs):
|
|
path = Path(index_path)
|
|
index_dir = path.parent
|
|
index_prefix = path.stem
|
|
index_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
if data.dtype != np.float32:
|
|
logger.warning(f"Converting data to float32, shape: {data.shape}")
|
|
data = data.astype(np.float32)
|
|
|
|
data_filename = f"{index_prefix}_data.bin"
|
|
_write_vectors_to_bin(data, index_dir / data_filename)
|
|
|
|
build_kwargs = {**self.build_params, **kwargs}
|
|
|
|
# Extract is_recompute from nested backend_kwargs if needed
|
|
is_recompute = build_kwargs.get("is_recompute", False)
|
|
if not is_recompute and "backend_kwargs" in build_kwargs:
|
|
is_recompute = build_kwargs["backend_kwargs"].get("is_recompute", False)
|
|
|
|
# Flatten all backend_kwargs parameters to top level for compatibility
|
|
if "backend_kwargs" in build_kwargs:
|
|
nested_params = build_kwargs.pop("backend_kwargs")
|
|
build_kwargs.update(nested_params)
|
|
|
|
metric_enum = _get_diskann_metrics().get(
|
|
build_kwargs.get("distance_metric", "mips").lower()
|
|
)
|
|
if metric_enum is None:
|
|
raise ValueError(
|
|
f"Unsupported distance_metric '{build_kwargs.get('distance_metric', 'unknown')}'."
|
|
)
|
|
|
|
# Calculate smart memory configuration if not explicitly provided
|
|
if (
|
|
"search_memory_maximum" not in build_kwargs
|
|
or "build_memory_maximum" not in build_kwargs
|
|
):
|
|
smart_search_mem, smart_build_mem = _calculate_smart_memory_config(data)
|
|
else:
|
|
smart_search_mem = build_kwargs.get("search_memory_maximum", 4.0)
|
|
smart_build_mem = build_kwargs.get("build_memory_maximum", 8.0)
|
|
|
|
try:
|
|
from . import _diskannpy as diskannpy # type: ignore
|
|
|
|
with chdir(index_dir):
|
|
diskannpy.build_disk_float_index(
|
|
metric_enum,
|
|
data_filename,
|
|
index_prefix,
|
|
build_kwargs.get("complexity", 64),
|
|
build_kwargs.get("graph_degree", 32),
|
|
build_kwargs.get("search_memory_maximum", smart_search_mem),
|
|
build_kwargs.get("build_memory_maximum", smart_build_mem),
|
|
build_kwargs.get("num_threads", 8),
|
|
build_kwargs.get("pq_disk_bytes", 0),
|
|
"",
|
|
)
|
|
|
|
# Auto-partition if is_recompute is enabled
|
|
if build_kwargs.get("is_recompute", False):
|
|
logger.info("is_recompute=True, starting automatic graph partitioning...")
|
|
from .graph_partition import partition_graph
|
|
|
|
# Partition the index using absolute paths
|
|
# Convert to absolute paths to avoid issues with working directory changes
|
|
absolute_index_dir = Path(index_dir).resolve()
|
|
absolute_index_prefix_path = str(absolute_index_dir / index_prefix)
|
|
disk_graph_path, partition_bin_path = partition_graph(
|
|
index_prefix_path=absolute_index_prefix_path,
|
|
output_dir=str(absolute_index_dir),
|
|
partition_prefix=index_prefix,
|
|
)
|
|
|
|
# Safe cleanup: In partition mode, C++ doesn't read _disk.index content
|
|
# but still needs the derived files (_medoids.bin, _centroids.bin, etc.)
|
|
self._safe_cleanup_after_partition(index_dir, index_prefix)
|
|
|
|
logger.info("✅ Graph partitioning completed successfully!")
|
|
logger.info(f" - Disk graph: {disk_graph_path}")
|
|
logger.info(f" - Partition file: {partition_bin_path}")
|
|
|
|
finally:
|
|
temp_data_file = index_dir / data_filename
|
|
if temp_data_file.exists():
|
|
os.remove(temp_data_file)
|
|
logger.debug(f"Cleaned up temporary data file: {temp_data_file}")
|
|
|
|
|
|
class DiskannSearcher(BaseSearcher):
|
|
def __init__(self, index_path: str, **kwargs):
|
|
super().__init__(
|
|
index_path,
|
|
backend_module_name="leann_backend_diskann.diskann_embedding_server",
|
|
**kwargs,
|
|
)
|
|
|
|
# Initialize DiskANN index with suppressed C++ output based on log level
|
|
with suppress_cpp_output_if_needed():
|
|
from . import _diskannpy as diskannpy # type: ignore
|
|
|
|
distance_metric = kwargs.get("distance_metric", "mips").lower()
|
|
metric_enum = _get_diskann_metrics().get(distance_metric)
|
|
if metric_enum is None:
|
|
raise ValueError(f"Unsupported distance_metric '{distance_metric}'.")
|
|
|
|
self.num_threads = kwargs.get("num_threads", 8)
|
|
|
|
# For DiskANN, we need to reinitialize the index when zmq_port changes
|
|
# Store the initialization parameters for later use
|
|
# Note: C++ load method expects the BASE path (without _disk.index suffix)
|
|
# C++ internally constructs: index_prefix + "_disk.index"
|
|
index_name = self.index_path.stem # "simple_test.leann" -> "simple_test"
|
|
diskann_index_prefix = str(self.index_dir / index_name) # /path/to/simple_test
|
|
full_index_prefix = diskann_index_prefix # /path/to/simple_test (base path)
|
|
|
|
# Auto-detect partition files and set partition_prefix
|
|
partition_graph_file = self.index_dir / f"{index_name}_disk_graph.index"
|
|
partition_bin_file = self.index_dir / f"{index_name}_partition.bin"
|
|
|
|
partition_prefix = ""
|
|
if partition_graph_file.exists() and partition_bin_file.exists():
|
|
# C++ expects full path prefix, not just filename
|
|
partition_prefix = str(self.index_dir / index_name) # /path/to/simple_test
|
|
logger.info(
|
|
f"✅ Detected partition files, using partition_prefix='{partition_prefix}'"
|
|
)
|
|
else:
|
|
logger.debug("No partition files detected, using standard index files")
|
|
|
|
self._init_params = {
|
|
"metric_enum": metric_enum,
|
|
"full_index_prefix": full_index_prefix,
|
|
"num_threads": self.num_threads,
|
|
"num_nodes_to_cache": kwargs.get("num_nodes_to_cache", 0),
|
|
# 1 -> initialize cache using sample_data; 2 -> ready cache without init; others disable cache
|
|
"cache_mechanism": kwargs.get("cache_mechanism", 1),
|
|
"pq_prefix": "",
|
|
"partition_prefix": partition_prefix,
|
|
}
|
|
|
|
# Log partition configuration for debugging
|
|
if partition_prefix:
|
|
logger.info(
|
|
f"✅ Detected partition files, using partition_prefix='{partition_prefix}'"
|
|
)
|
|
self._diskannpy = diskannpy
|
|
self._current_zmq_port = None
|
|
self._index = None
|
|
logger.debug("DiskANN searcher initialized (index will be loaded on first search)")
|
|
|
|
def _ensure_index_loaded(self, zmq_port: int):
|
|
"""Ensure the index is loaded with the correct zmq_port."""
|
|
if self._index is None or self._current_zmq_port != zmq_port:
|
|
# Need to (re)load the index with the correct zmq_port
|
|
with suppress_cpp_output_if_needed():
|
|
if self._index is not None:
|
|
logger.debug(f"Reloading DiskANN index with new zmq_port: {zmq_port}")
|
|
else:
|
|
logger.debug(f"Loading DiskANN index with zmq_port: {zmq_port}")
|
|
|
|
self._index = self._diskannpy.StaticDiskFloatIndex(
|
|
self._init_params["metric_enum"],
|
|
self._init_params["full_index_prefix"],
|
|
self._init_params["num_threads"],
|
|
self._init_params["num_nodes_to_cache"],
|
|
self._init_params["cache_mechanism"],
|
|
zmq_port,
|
|
self._init_params["pq_prefix"],
|
|
self._init_params["partition_prefix"],
|
|
)
|
|
self._current_zmq_port = zmq_port
|
|
|
|
def search(
|
|
self,
|
|
query: np.ndarray,
|
|
top_k: int,
|
|
complexity: int = 64,
|
|
beam_width: int = 1,
|
|
prune_ratio: float = 0.0,
|
|
recompute_embeddings: bool = False,
|
|
pruning_strategy: Literal["global", "local", "proportional"] = "global",
|
|
zmq_port: Optional[int] = None,
|
|
batch_recompute: bool = False,
|
|
dedup_node_dis: bool = False,
|
|
**kwargs,
|
|
) -> dict[str, Any]:
|
|
"""
|
|
Search for nearest neighbors using DiskANN index.
|
|
|
|
Args:
|
|
query: Query vectors (B, D) where B is batch size, D is dimension
|
|
top_k: Number of nearest neighbors to return
|
|
complexity: Search complexity/candidate list size, higher = more accurate but slower
|
|
beam_width: Number of parallel IO requests per iteration
|
|
prune_ratio: Ratio of neighbors to prune via approximate distance (0.0-1.0)
|
|
recompute_embeddings: Whether to fetch fresh embeddings from server
|
|
pruning_strategy: PQ candidate selection strategy:
|
|
- "global": Use global pruning strategy (default)
|
|
- "local": Use local pruning strategy
|
|
- "proportional": Not supported in DiskANN, falls back to global
|
|
zmq_port: ZMQ port for embedding server communication. Must be provided if recompute_embeddings is True.
|
|
batch_recompute: Whether to batch neighbor recomputation (DiskANN-specific)
|
|
dedup_node_dis: Whether to cache and reuse distance computations (DiskANN-specific)
|
|
**kwargs: Additional DiskANN-specific parameters (for legacy compatibility)
|
|
|
|
Returns:
|
|
Dict with 'labels' (list of lists) and 'distances' (ndarray)
|
|
"""
|
|
# Handle zmq_port compatibility: Ensure index is loaded with correct port
|
|
if recompute_embeddings:
|
|
if zmq_port is None:
|
|
raise ValueError("zmq_port must be provided if recompute_embeddings is True")
|
|
self._ensure_index_loaded(zmq_port)
|
|
else:
|
|
# If not recomputing, we still need an index, use a default port
|
|
if self._index is None:
|
|
self._ensure_index_loaded(6666) # Default port when not recomputing
|
|
|
|
# DiskANN doesn't support "proportional" strategy
|
|
if pruning_strategy == "proportional":
|
|
raise NotImplementedError(
|
|
"DiskANN backend does not support 'proportional' pruning strategy. Use 'global' or 'local' instead."
|
|
)
|
|
|
|
if query.dtype != np.float32:
|
|
query = query.astype(np.float32)
|
|
|
|
# Map pruning_strategy to DiskANN's global_pruning parameter
|
|
if pruning_strategy == "local":
|
|
use_global_pruning = False
|
|
else: # "global"
|
|
use_global_pruning = True
|
|
|
|
# Strategy:
|
|
# - Traversal always uses PQ distances
|
|
# - If recompute_embeddings=True, do a single final rerank via deferred fetch
|
|
# (fetch embeddings for the final candidate set only)
|
|
# - Do not recompute neighbor distances along the path
|
|
use_deferred_fetch = True if recompute_embeddings else False
|
|
recompute_neighors = False # Expected typo. For backward compatibility.
|
|
|
|
with suppress_cpp_output_if_needed():
|
|
labels, distances = self._index.batch_search(
|
|
query,
|
|
query.shape[0],
|
|
top_k,
|
|
complexity,
|
|
beam_width,
|
|
self.num_threads,
|
|
use_deferred_fetch,
|
|
kwargs.get("skip_search_reorder", False),
|
|
recompute_neighors,
|
|
dedup_node_dis,
|
|
prune_ratio,
|
|
batch_recompute,
|
|
use_global_pruning,
|
|
)
|
|
|
|
string_labels = [[str(int_label) for int_label in batch_labels] for batch_labels in labels]
|
|
|
|
return {"labels": string_labels, "distances": distances}
|