532 lines
23 KiB
Python
532 lines
23 KiB
Python
import numpy as np
|
|
import os
|
|
import json
|
|
import struct
|
|
from pathlib import Path
|
|
from typing import Dict, Any
|
|
import contextlib
|
|
import threading
|
|
import time
|
|
import atexit
|
|
import socket
|
|
import subprocess
|
|
import sys
|
|
|
|
from .convert_to_csr import convert_hnsw_graph_to_csr
|
|
|
|
from leann.registry import register_backend
|
|
from leann.interface import (
|
|
LeannBackendFactoryInterface,
|
|
LeannBackendBuilderInterface,
|
|
LeannBackendSearcherInterface
|
|
)
|
|
|
|
def get_metric_map():
|
|
from . import faiss
|
|
return {
|
|
"mips": faiss.METRIC_INNER_PRODUCT,
|
|
"l2": faiss.METRIC_L2,
|
|
"cosine": faiss.METRIC_INNER_PRODUCT,
|
|
}
|
|
|
|
def _check_port(port: int) -> bool:
|
|
"""Check if a port is in use"""
|
|
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
|
|
return s.connect_ex(('localhost', port)) == 0
|
|
|
|
class HNSWEmbeddingServerManager:
|
|
"""
|
|
HNSW-specific embedding server manager that handles the lifecycle of the embedding server process.
|
|
Mirrors the DiskANN EmbeddingServerManager architecture.
|
|
"""
|
|
def __init__(self):
|
|
self.server_process = None
|
|
self.server_port = None
|
|
atexit.register(self.stop_server)
|
|
|
|
def start_server(self, port=5556, model_name="sentence-transformers/all-mpnet-base-v2", passages_file=None):
|
|
"""
|
|
Start the HNSW embedding server process.
|
|
|
|
Args:
|
|
port: ZMQ port for the server
|
|
model_name: Name of the embedding model to use
|
|
passages_file: Optional path to passages JSON file
|
|
"""
|
|
if self.server_process and self.server_process.poll() is None:
|
|
print(f"INFO: Reusing existing HNSW server process for this session (PID {self.server_process.pid})")
|
|
return True
|
|
|
|
# Check if port is already in use
|
|
if _check_port(port):
|
|
print(f"WARNING: Port {port} is already in use. Assuming an external HNSW server is running and connecting to it.")
|
|
return True
|
|
|
|
print(f"INFO: Starting session-level HNSW embedding server as a background process...")
|
|
|
|
try:
|
|
command = [
|
|
sys.executable,
|
|
"-m", "leann_backend_hnsw.hnsw_embedding_server",
|
|
"--zmq-port", str(port),
|
|
"--model-name", model_name
|
|
]
|
|
|
|
if passages_file:
|
|
command.extend(["--passages-file", str(passages_file)])
|
|
|
|
project_root = Path(__file__).parent.parent.parent.parent
|
|
print(f"INFO: Running HNSW command from project root: {project_root}")
|
|
|
|
self.server_process = subprocess.Popen(
|
|
command,
|
|
cwd=project_root,
|
|
stdout=subprocess.PIPE,
|
|
stderr=subprocess.PIPE,
|
|
text=True,
|
|
encoding='utf-8'
|
|
)
|
|
self.server_port = port
|
|
print(f"INFO: HNSW server process started with PID: {self.server_process.pid}")
|
|
|
|
max_wait, wait_interval = 30, 0.5
|
|
for _ in range(int(max_wait / wait_interval)):
|
|
if _check_port(port):
|
|
print(f"✅ HNSW embedding server is up and ready for this session.")
|
|
log_thread = threading.Thread(target=self._log_monitor, daemon=True)
|
|
log_thread.start()
|
|
return True
|
|
if self.server_process.poll() is not None:
|
|
print("❌ ERROR: HNSW server process terminated unexpectedly during startup.")
|
|
self._log_monitor()
|
|
return False
|
|
time.sleep(wait_interval)
|
|
|
|
print(f"❌ ERROR: HNSW server process failed to start listening within {max_wait} seconds.")
|
|
self.stop_server()
|
|
return False
|
|
|
|
except Exception as e:
|
|
print(f"❌ ERROR: Failed to start HNSW embedding server process: {e}")
|
|
return False
|
|
|
|
def _log_monitor(self):
|
|
"""Monitor server logs"""
|
|
if not self.server_process:
|
|
return
|
|
try:
|
|
if self.server_process.stdout:
|
|
for line in iter(self.server_process.stdout.readline, ''):
|
|
print(f"[HNSWEmbeddingServer LOG]: {line.strip()}")
|
|
self.server_process.stdout.close()
|
|
if self.server_process.stderr:
|
|
for line in iter(self.server_process.stderr.readline, ''):
|
|
print(f"[HNSWEmbeddingServer ERROR]: {line.strip()}")
|
|
self.server_process.stderr.close()
|
|
except Exception as e:
|
|
print(f"HNSW Log monitor error: {e}")
|
|
|
|
def stop_server(self):
|
|
"""Stop the HNSW embedding server process"""
|
|
if self.server_process and self.server_process.poll() is None:
|
|
print(f"INFO: Terminating HNSW session server process (PID: {self.server_process.pid})...")
|
|
self.server_process.terminate()
|
|
try:
|
|
self.server_process.wait(timeout=5)
|
|
print("INFO: HNSW server process terminated.")
|
|
except subprocess.TimeoutExpired:
|
|
print("WARNING: HNSW server process did not terminate gracefully, killing it.")
|
|
self.server_process.kill()
|
|
self.server_process = None
|
|
|
|
@register_backend("hnsw")
|
|
class HNSWBackend(LeannBackendFactoryInterface):
|
|
@staticmethod
|
|
def builder(**kwargs) -> LeannBackendBuilderInterface:
|
|
return HNSWBuilder(**kwargs)
|
|
|
|
@staticmethod
|
|
def searcher(index_path: str, **kwargs) -> LeannBackendSearcherInterface:
|
|
path = Path(index_path)
|
|
meta_path = path.parent / f"{path.name}.meta.json"
|
|
if not meta_path.exists():
|
|
raise FileNotFoundError(f"Leann metadata file not found at {meta_path}. Cannot infer vector dimension for searcher.")
|
|
|
|
with open(meta_path, 'r') as f:
|
|
meta = json.load(f)
|
|
|
|
try:
|
|
from sentence_transformers import SentenceTransformer
|
|
model = SentenceTransformer(meta.get("embedding_model"))
|
|
dimensions = model.get_sentence_embedding_dimension()
|
|
kwargs['dimensions'] = dimensions
|
|
except ImportError:
|
|
raise ImportError("sentence-transformers is required to infer embedding dimensions. Please install it.")
|
|
except Exception as e:
|
|
raise RuntimeError(f"Could not load SentenceTransformer model to get dimension: {e}")
|
|
|
|
return HNSWSearcher(index_path, **kwargs)
|
|
|
|
class HNSWBuilder(LeannBackendBuilderInterface):
|
|
def __init__(self, **kwargs):
|
|
self.build_params = kwargs.copy()
|
|
|
|
# --- Configuration defaults with standardized names ---
|
|
# Apply defaults and write them back to the build_params dict
|
|
# so they can be saved in the metadata file by LeannBuilder.
|
|
self.is_compact = self.build_params.setdefault("is_compact", True)
|
|
self.is_recompute = self.build_params.setdefault("is_recompute", True) # Default: prune embeddings
|
|
|
|
# --- Additional Options ---
|
|
self.is_skip_neighbors = self.build_params.setdefault("is_skip_neighbors", False)
|
|
self.disk_cache_ratio = self.build_params.setdefault("disk_cache_ratio", 0.0)
|
|
self.external_storage_path = self.build_params.get("external_storage_path", None)
|
|
|
|
# --- Standard HNSW parameters ---
|
|
self.M = self.build_params.setdefault("M", 32)
|
|
self.efConstruction = self.build_params.setdefault("efConstruction", 200)
|
|
self.distance_metric = self.build_params.setdefault("distance_metric", "mips")
|
|
|
|
if self.is_skip_neighbors and not self.is_compact:
|
|
raise ValueError("is_skip_neighbors can only be used with is_compact=True")
|
|
|
|
if self.is_recompute and not self.is_compact:
|
|
raise ValueError("is_recompute requires is_compact=True for efficiency")
|
|
|
|
def build(self, data: np.ndarray, index_path: str, **kwargs):
|
|
"""Build HNSW index using FAISS"""
|
|
from . import faiss
|
|
|
|
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:
|
|
data = data.astype(np.float32)
|
|
if not data.flags['C_CONTIGUOUS']:
|
|
data = np.ascontiguousarray(data)
|
|
|
|
metric_str = self.distance_metric.lower()
|
|
metric_enum = get_metric_map().get(metric_str)
|
|
print('metric_enum', metric_enum,' metric_str', metric_str)
|
|
if metric_enum is None:
|
|
raise ValueError(f"Unsupported distance_metric '{metric_str}'.")
|
|
|
|
M = self.M
|
|
efConstruction = self.efConstruction
|
|
dim = data.shape[1]
|
|
|
|
print(f"INFO: Building HNSW index for {data.shape[0]} vectors with metric {metric_enum}...")
|
|
|
|
try:
|
|
if metric_enum == faiss.METRIC_INNER_PRODUCT:
|
|
index = faiss.IndexHNSWFlat(dim, M, metric_enum)
|
|
else: # L2
|
|
index = faiss.IndexHNSWFlat(dim, M, metric_enum)
|
|
|
|
index.hnsw.efConstruction = efConstruction
|
|
|
|
if metric_str == "cosine":
|
|
faiss.normalize_L2(data)
|
|
|
|
print('starting to add vectors to index')
|
|
index.add(data.shape[0], faiss.swig_ptr(data))
|
|
print('vectors added to index')
|
|
|
|
index_file = index_dir / f"{index_prefix}.index"
|
|
faiss.write_index(index, str(index_file))
|
|
|
|
print(f"✅ HNSW index built successfully at '{index_file}'")
|
|
|
|
if self.is_compact:
|
|
self._convert_to_csr(index_file)
|
|
|
|
# Generate passages file for recompute mode
|
|
if self.is_recompute:
|
|
self._generate_passages_file(index_dir, index_prefix, **kwargs)
|
|
|
|
except Exception as e:
|
|
print(f"💥 ERROR: HNSW index build failed. Exception: {e}")
|
|
raise
|
|
|
|
def _convert_to_csr(self, index_file: Path):
|
|
"""Convert built index to CSR format"""
|
|
try:
|
|
mode_str = "CSR-pruned" if self.is_recompute else "CSR-standard"
|
|
print(f"INFO: Converting HNSW index to {mode_str} format...")
|
|
|
|
csr_temp_file = index_file.with_suffix(".csr.tmp")
|
|
|
|
success = convert_hnsw_graph_to_csr(
|
|
str(index_file),
|
|
str(csr_temp_file),
|
|
prune_embeddings=self.is_recompute
|
|
)
|
|
|
|
if success:
|
|
print("✅ CSR conversion successful.")
|
|
import shutil
|
|
shutil.move(str(csr_temp_file), str(index_file))
|
|
print(f"INFO: Replaced original index with {mode_str} version at '{index_file}'")
|
|
else:
|
|
# Clean up and fail fast
|
|
if csr_temp_file.exists():
|
|
os.remove(csr_temp_file)
|
|
raise RuntimeError("CSR conversion failed - cannot proceed with compact format")
|
|
|
|
except Exception as e:
|
|
print(f"💥 ERROR: CSR conversion failed. Exception: {e}")
|
|
raise
|
|
|
|
def _generate_passages_file(self, index_dir: Path, index_prefix: str, **kwargs):
|
|
"""Generate passages file for recompute mode"""
|
|
try:
|
|
chunks = kwargs.get('chunks', [])
|
|
if not chunks:
|
|
print("INFO: No chunks data provided, skipping passages file generation")
|
|
return
|
|
|
|
# Generate node_id to text mapping
|
|
passages_data = {}
|
|
for node_id, chunk in enumerate(chunks):
|
|
passages_data[str(node_id)] = chunk["text"]
|
|
|
|
# Save passages file
|
|
passages_file = index_dir / f"{index_prefix}.passages.json"
|
|
with open(passages_file, 'w', encoding='utf-8') as f:
|
|
json.dump(passages_data, f, ensure_ascii=False, indent=2)
|
|
|
|
print(f"✅ Generated passages file for recompute mode at '{passages_file}' ({len(passages_data)} passages)")
|
|
|
|
except Exception as e:
|
|
print(f"💥 ERROR: Failed to generate passages file. Exception: {e}")
|
|
# Don't raise - this is not critical for index building
|
|
pass
|
|
|
|
class HNSWSearcher(LeannBackendSearcherInterface):
|
|
def _get_index_storage_status(self, index_file: Path) -> tuple[bool, bool]:
|
|
"""
|
|
Robustly determines the index's storage status by parsing the file.
|
|
|
|
Returns:
|
|
A tuple (is_compact, is_pruned).
|
|
"""
|
|
if not index_file.exists():
|
|
return False, False
|
|
|
|
with open(index_file, 'rb') as f:
|
|
try:
|
|
def read_struct(fmt):
|
|
size = struct.calcsize(fmt)
|
|
data = f.read(size)
|
|
if len(data) != size:
|
|
raise EOFError(f"File ended unexpectedly reading struct fmt '{fmt}'.")
|
|
return struct.unpack(fmt, data)[0]
|
|
|
|
def skip_vector(element_size):
|
|
count = read_struct('<Q')
|
|
f.seek(count * element_size, 1)
|
|
|
|
# 1. Read up to the compact flag
|
|
read_struct('<I'); read_struct('<i'); read_struct('<q');
|
|
read_struct('<q'); read_struct('<q'); read_struct('<?')
|
|
metric_type = read_struct('<i')
|
|
if metric_type > 1: read_struct('<f')
|
|
skip_vector(8); skip_vector(4); skip_vector(4)
|
|
|
|
# 2. Check if there's a compact flag byte
|
|
# Try to read the compact flag, but handle both old and new formats
|
|
pos_before_compact = f.tell()
|
|
try:
|
|
is_compact = read_struct('<?')
|
|
print(f"INFO: Detected is_compact flag as: {is_compact}")
|
|
except (EOFError, struct.error):
|
|
# Old format without compact flag - assume non-compact
|
|
f.seek(pos_before_compact)
|
|
is_compact = False
|
|
print(f"INFO: No compact flag found, assuming is_compact=False")
|
|
|
|
# 3. Read storage FourCC to determine if pruned
|
|
is_pruned = False
|
|
try:
|
|
if is_compact:
|
|
# For compact, we need to skip pointers and scalars to get to the storage FourCC
|
|
skip_vector(8) # level_ptr
|
|
skip_vector(8) # node_offsets
|
|
read_struct('<i'); read_struct('<i'); read_struct('<i');
|
|
read_struct('<i'); read_struct('<i')
|
|
storage_fourcc = read_struct('<I')
|
|
else:
|
|
# For non-compact, we need to read the flag probe, then skip offsets and neighbors
|
|
pos_before_probe = f.tell()
|
|
flag_byte = f.read(1)
|
|
if not (flag_byte and flag_byte == b'\x00'):
|
|
f.seek(pos_before_probe)
|
|
skip_vector(8); skip_vector(4) # offsets, neighbors
|
|
read_struct('<i'); read_struct('<i'); read_struct('<i');
|
|
read_struct('<i'); read_struct('<i')
|
|
# Now we are at the storage. The entire rest is storage blob.
|
|
storage_fourcc = struct.unpack('<I', f.read(4))[0]
|
|
|
|
NULL_INDEX_FOURCC = int.from_bytes(b'null', 'little')
|
|
if storage_fourcc == NULL_INDEX_FOURCC:
|
|
is_pruned = True
|
|
except (EOFError, struct.error):
|
|
# Cannot determine pruning status, assume not pruned
|
|
pass
|
|
|
|
print(f"INFO: Detected is_pruned as: {is_pruned}")
|
|
return is_compact, is_pruned
|
|
|
|
except (EOFError, struct.error) as e:
|
|
print(f"WARNING: Could not parse index file to detect format: {e}. Assuming standard, not pruned.")
|
|
return False, False
|
|
|
|
def __init__(self, index_path: str, **kwargs):
|
|
from . import faiss
|
|
path = Path(index_path)
|
|
index_dir = path.parent
|
|
index_prefix = path.stem
|
|
|
|
# Store configuration and paths for later use
|
|
self.config = kwargs.copy()
|
|
self.config["index_path"] = index_path
|
|
self.index_dir = index_dir
|
|
self.index_prefix = index_prefix
|
|
|
|
metric_str = self.config.get("distance_metric", "mips").lower()
|
|
metric_enum = get_metric_map().get(metric_str)
|
|
if metric_enum is None:
|
|
raise ValueError(f"Unsupported distance_metric '{metric_str}'.")
|
|
|
|
dimensions = self.config.get("dimensions")
|
|
if not dimensions:
|
|
raise ValueError("Vector dimension not provided to HNSWSearcher.")
|
|
|
|
index_file = index_dir / f"{index_prefix}.index"
|
|
if not index_file.exists():
|
|
raise FileNotFoundError(f"HNSW index file not found at {index_file}")
|
|
|
|
self.is_compact, self.is_pruned = self._get_index_storage_status(index_file)
|
|
|
|
# Validate configuration constraints
|
|
if not self.is_compact and self.config.get("is_skip_neighbors", False):
|
|
raise ValueError("is_skip_neighbors can only be used with is_compact=True")
|
|
|
|
if self.config.get("is_recompute", False) and self.config.get("external_storage_path"):
|
|
raise ValueError("Cannot use both is_recompute and external_storage_path simultaneously")
|
|
|
|
hnsw_config = faiss.HNSWIndexConfig()
|
|
hnsw_config.is_compact = self.is_compact
|
|
|
|
# Apply additional configuration options with strict validation
|
|
hnsw_config.is_skip_neighbors = self.config.get("is_skip_neighbors", False)
|
|
# If index is pruned, force recompute mode regardless of user setting
|
|
hnsw_config.is_recompute = self.is_pruned or self.config.get("is_recompute", False)
|
|
hnsw_config.disk_cache_ratio = self.config.get("disk_cache_ratio", 0.0)
|
|
hnsw_config.external_storage_path = self.config.get("external_storage_path")
|
|
hnsw_config.zmq_port = self.config.get("zmq_port", 5557)
|
|
|
|
# CRITICAL ASSERTION: If index is pruned, recompute MUST be enabled
|
|
if self.is_pruned and not hnsw_config.is_recompute:
|
|
raise RuntimeError("Index is pruned (embeddings removed) but recompute is disabled. This is impossible - recompute must be enabled for pruned indices.")
|
|
|
|
print(f"INFO: Loading index with is_compact={self.is_compact}, is_pruned={self.is_pruned}")
|
|
print(f"INFO: Config - skip_neighbors={hnsw_config.is_skip_neighbors}, recompute={hnsw_config.is_recompute}")
|
|
|
|
self._index = faiss.read_index(str(index_file), faiss.IO_FLAG_MMAP, hnsw_config)
|
|
|
|
if self.is_compact:
|
|
print("✅ Compact CSR format HNSW index loaded successfully.")
|
|
else:
|
|
print("✅ Standard HNSW index loaded successfully.")
|
|
|
|
self.metric_str = metric_str
|
|
self.embedding_server_manager = HNSWEmbeddingServerManager()
|
|
|
|
def _get_index_file(self, index_dir: Path, index_prefix: str) -> Path:
|
|
"""Get the appropriate index file path based on format"""
|
|
# We always use the same filename now, format is detected internally
|
|
return index_dir / f"{index_prefix}.index"
|
|
|
|
def search(self, query: np.ndarray, top_k: int, **kwargs) -> Dict[str, Any]:
|
|
"""Search using HNSW index with optional recompute functionality"""
|
|
from . import faiss
|
|
# Merge config with search-time kwargs
|
|
search_config = self.config.copy()
|
|
search_config.update(kwargs)
|
|
|
|
ef = search_config.get("ef", 200) # Size of the dynamic candidate list for search
|
|
|
|
# Recompute parameters
|
|
zmq_port = search_config.get("zmq_port", 5557)
|
|
embedding_model = search_config.get("embedding_model", "sentence-transformers/all-mpnet-base-v2")
|
|
passages_file = search_config.get("passages_file", None)
|
|
|
|
# For recompute mode, try to find the passages file automatically
|
|
if self.is_pruned and not passages_file:
|
|
potential_passages_file = self.index_dir / f"{self.index_prefix}.passages.json"
|
|
print(f"DEBUG: Checking for passages file at: {potential_passages_file}")
|
|
if potential_passages_file.exists():
|
|
passages_file = str(potential_passages_file)
|
|
print(f"INFO: Found passages file for recompute mode: {passages_file}")
|
|
else:
|
|
print(f"WARNING: No passages file found for recompute mode at {potential_passages_file}")
|
|
|
|
# If index is pruned (embeddings removed), we MUST start embedding server for recompute
|
|
if self.is_pruned:
|
|
print(f"INFO: Index is pruned - starting embedding server for recompute")
|
|
|
|
# CRITICAL: Check passages file exists - fail fast if not
|
|
if not passages_file:
|
|
raise RuntimeError(f"FATAL: Index is pruned but no passages file found. Cannot proceed with recompute mode.")
|
|
|
|
# Check if server is already running first
|
|
if _check_port(zmq_port):
|
|
print(f"INFO: Embedding server already running on port {zmq_port}")
|
|
else:
|
|
if not self.embedding_server_manager.start_server(zmq_port, embedding_model, passages_file):
|
|
raise RuntimeError(f"Failed to start HNSW embedding server on port {zmq_port}")
|
|
|
|
# Give server extra time to fully initialize
|
|
print(f"INFO: Waiting for embedding server to fully initialize...")
|
|
time.sleep(3)
|
|
|
|
# Final verification
|
|
if not _check_port(zmq_port):
|
|
raise RuntimeError(f"Embedding server failed to start listening on port {zmq_port}")
|
|
else:
|
|
print(f"INFO: Index has embeddings stored - no recompute needed")
|
|
|
|
if query.dtype != np.float32:
|
|
query = query.astype(np.float32)
|
|
if query.ndim == 1:
|
|
query = np.expand_dims(query, axis=0)
|
|
|
|
# Normalize query if using cosine similarity
|
|
if self.metric_str == "cosine":
|
|
faiss.normalize_L2(query)
|
|
|
|
try:
|
|
# Set search parameter
|
|
self._index.hnsw.efSearch = ef
|
|
|
|
# Prepare output arrays for the older FAISS SWIG API
|
|
batch_size = query.shape[0]
|
|
distances = np.empty((batch_size, top_k), dtype=np.float32)
|
|
labels = np.empty((batch_size, top_k), dtype=np.int64)
|
|
|
|
# Use standard FAISS search - recompute is handled internally by FAISS
|
|
self._index.search(query.shape[0], faiss.swig_ptr(query), top_k, faiss.swig_ptr(distances), faiss.swig_ptr(labels))
|
|
|
|
return {"labels": labels, "distances": distances}
|
|
|
|
except Exception as e:
|
|
print(f"💥 ERROR: HNSW search failed. Exception: {e}")
|
|
raise
|
|
|
|
def __del__(self):
|
|
if hasattr(self, 'embedding_server_manager'):
|
|
self.embedding_server_manager.stop_server() |