跳到主要内容

数据截至 (上游 commit dc85934f318c)

03 — 可插拔后端

本章讲什么: LEANN 怎么做到"装一个包就多一个后端",后端要遵守什么契约,以及内置的三个后端各自适合什么场景。

1. 插件机制:约定优于配置​

1.1 它要解决的小问题​

向量索引算法太多了,而且各有各的原生依赖(FAISS、DiskANN、CUDA)。核心包不能硬依赖它们,否则装个 LEANN 要编译半天。

1.2 思路​

用分发包名当约定。 启动时扫一遍已安装的所有 distribution,凡是名字以 leann-backend- 开头的就 import 一下;import 的副作用会触发装饰器注册。

# 示意,非源码:自动发现的两步
for dist in importlib.metadata.distributions():
    if dist.metadata["name"].startswith("leann-backend-"):
        importlib.import_module(name.replace("-", "_"))   # import 即注册

真实实现在 registry.py:61-80;注册装饰器是 register_backend(registry.py:50-58),往模块级字典 BACKEND_REGISTRY 里塞(registry.py:18)。import 失败静默吞掉(registry.py:77-79)——没装 CUDA 的机器上 GPU 后端 import 不了是正常的。

触发点在包入口:from .registry import ... ; autodiscover_backends()(packages/leann-core/src/leann/__init__.py)。

1.3 后端使用者这一侧​

LeannBuilder 和 LeannSearcher 都只从字典里取工厂,取不到就报错(api.py:419-422、:1191-1193)。查询可用后端用 get_registered_backends()(api.py:43-45)。

2. 契约:三个抽象接口​

interface.py 里定了三个 ABC:

接口要实现什么位置
LeannBackendFactoryInterface两个静态工厂:builder() / searcher(index_path)interface.py:96
LeannBackendBuilderInterfacebuild(data, ids, index_path, **kw)interface.py:7
LeannBackendSearcherInterfacesearch(...)、_ensure_server_running(...)、compute_query_embedding(...)interface.py:23

后两个方法各后端不必重写——BaseSearcher 已经实现好了通用逻辑:加载 meta、校验维度、建 EmbeddingServerManager、query 向量计算与降级(searcher_base.py:12-102)。后端实际只需要写 __init__ 和 search。

search 的返回字段约定写在 docstring 里:{"labels": [...], "distances": [...]}(interface.py:69-70)。但labels 里装的是字符串 passage ID、不是后端内部的整数标签这条类型约定,接口本身没写——它只体现在实现里:HNSW 拿 FAISS 返回的整数标签查 _id_map 换成字符串,查不到(或压根没有 id map)就把整数直接转成字符串兜底(hnsw_backend.py:274-287)。也就是说,标签到 passage ID 的映射由各后端自己完成。

2.1 一个最小后端长什么样​

# 示意,非源码:注册一个新后端所需的全部骨架
@register_backend("myann")
class MyBackend(LeannBackendFactoryInterface):
    @staticmethod
    def builder(**kw):
        return MyBuilder(**kw)

    @staticmethod
    def searcher(index_path, **kw):
        return MySearcher(index_path, **kw)

把它放进一个叫 leann-backend-myann 的包里,装上就能用 backend_name="myann"。核心一行代码都不用改。

3. 三个内置后端的取舍​

3.1 对比表​

维度HNSW(默认)IVFDiskANN
底层FAISS IndexHNSWFlat(作者 fork)FAISS IndexIVFFlatDiskANN(作者 fork)
省存储的手段CSR + 丢弃向量段 + 全程重算无(向量照存)PQ 压缩 + 图分区 + 延迟取回重排
增量新增仅 non-compact 支持追加支持不支持
增量删除不支持支持不支持
适合追求最小体积的静态语料经常改动的代码库 / 笔记超内存规模数据集
注册名hnswivfdiskann
代码hnsw_backend.py:38ivf_backend.py:85diskann_backend.py:131

CLI 的 --backend-name 三选一,默认 hnsw(cli.py:303-309);MCP 的 build 工具则默认 ivf,因为它主要服务"边写代码边更新索引"的场景(mcp.py:119-124)。

3.2 HNSW:最省的那个​

HNSW(Hierarchical Navigable Small World,分层可导航小世界图) 是主流的图式近邻索引:建一个多层图,上层稀疏用来快速跳转,底层稠密用来精搜。

LEANN 在它上面加的东西已在 01 讲透:CSR 重写 + 丢向量段。搜索侧的细节在 02。

两个开关实际可达的组合只有三种:

is_compactis_recompute结果
TrueTrue默认。CSR + 无向量,必须重算
FalseTrue原布局 + 无向量(只调 prune_hnsw_embeddings_inplace)
FalseFalse普通 FAISS 索引,可追加

分支在 hnsw_backend.py:102-105。

第四种组合 is_compact=True + is_recompute=False(即 convert_to_csr 的 prune_embeddings=False 分支,CSR + 保留向量)在 HNSW 路径上不可达:构造 HNSWBuilder 时就会检测到这个组合,强制把 is_compact 改成 False 并告警(hnsw_backend.py:58-64);更上游的 LeannBuilder 构造函数也做了同样的归一化,好让写进 meta 的标志位一致(api.py:407-417)。所以调用方即使显式传了这一对,落到 build() 时也已经变成了表里最后一行。

3.3 IVF:为增量而生​

IVF(Inverted File,倒排文件) 先把向量空间聚成 nlist 个簇,查询时只探其中 nprobe 个簇。

LEANN 选它的理由只有一个:能删。 建索引时显式设置 DirectMap.Hashtable:

ivf = faiss.IndexIVFFlat(quantizer, dim, nlist, metric)
ivf.train(data)
ivf.set_direct_map_type(faiss.DirectMap.Hashtable)   # 关键这一行
ivf.add_with_ids(data, np.arange(n, dtype=np.int64))

这段是真实源码的骨架(ivf_backend.py:126-130)。模块 docstring 明说:用 IndexIVFFlat + DirectMap.Hashtable 而不是 IndexIDMap2,就是为了让 remove_ids 正常工作(ivf_backend.py:1-6)。

配套的两个模块级函数:

函数干什么位置
add_vectors追加向量,分配递增的整数 IDivf_backend.py:200
remove_ids按 passage ID 删除ivf_backend.py:238

两者共用一份 JSON 的 ID 映射文件 <prefix>.ivf_id_map.json,里面存 id_to_passage / passage_to_id / next_id(ivf_backend.py:55-82)。删除时不回收 next_id,新增永远拿新号(ivf_backend.py:275-276 的注释)。

remove_ids 还做了一致性检查:FAISS 实际删掉的条数和预期不符就打警告(ivf_backend.py:285-290)。

注意 IVF 后端不省存储——IndexIVFFlat 老老实实存着全部向量。它换来的是可变性。搜索时 nprobe 默认取 min(complexity, nlist)(ivf_backend.py:175)。

3.4 DiskANN:大数据集那个​

DiskANN 是面向"索引比内存大"的磁盘图索引,靠 PQ(Product Quantization,乘积量化:把高维向量切段后各段独立聚类,用簇心编号代替原值) 压缩出一份能全放内存的近似向量做遍历,再回磁盘取精确数据重排。

LEANN 的搜索策略写得很直白(diskann_backend.py:536-541 的注释):

  • 遍历始终用 PQ 距离;
  • recompute_embeddings=True 时,只在最后对候选集做一次 deferred fetch 重排;
  • 不沿途重算邻居距离(recompute_neighors = False,:542,变量名的拼写错误是为了向后兼容而保留的)。

这和 HNSW 的"沿途全程重算"是两种截然不同的精度—延迟取舍。

三个 DiskANN 特有的工程细节:

① 建索引跑在子进程里。 原生崩溃不该静默杀掉 CLI,所以用 mp.Process + 队列回传状态,退出码非 0 或队列为空都当失败处理(diskann_backend.py:271-337)。建完还校验期望的 _disk.index 文件确实存在(:339-346)。

② 内存参数自动推算。 search_memory_maximum 取向量总大小的 1/10(控制 PQ 压缩率),build_memory_maximum 取可用内存的 50%(控制分片)(_calculate_smart_memory_config,:94-128)。

③ is_recompute=True 时自动做图分区。 调 partition_graph,内部是两个 C++ 可执行文件:partitioner 做 LDG 分区、index_relayout 按分区重排布局(graph_partition.py:174-226)。分区完成后,_disk.index 本体可以安全删掉——分区模式下 C++ 不读它,只要 _medoids.bin / _pq_pivots.bin / _pq_compressed.bin 还在(diskann_backend.py:146-209)。

搜索侧会自动探测分区文件是否存在来决定要不要传 partition_prefix(diskann_backend.py:405-415)。索引是懒加载的:第一次 search 时才按当时的 zmq 端口加载,端口变了要重载(_ensure_index_loaded,:453-473)。

DiskANN 不支持 proportional 剪枝策略,传了直接抛 NotImplementedError(diskann_backend.py:521-525)。

3.5 还有两个 GPU 后端​

packages/leann-backend-flashlib(Triton/CuteDSL 上的精确 k-NN)和 packages/leann-backend-flashlib-ivf(GPU 版 IVF-Flat)在根 pyproject.toml 里各占一个可选 extra:flashlib(pyproject.toml:69-72)和 flashlib-ivf(:74-77),两者都要求 CUDA GPU。本文不展开。

顺带说清工作区的挂载口径,免得误会:[tool.uv.sources] 里以 editable 路径挂进来的是 leann-core 和四个后端(diskann / hnsw / flashlib / flashlib-ivf)加 astchunk,唯独没有 leann-backend-ivf(pyproject.toml:97-103)——尽管 packages/leann-backend-ivf/ 这个目录是存在的。

4. 索引发现:leann list 怎么找到索引​

和后端注册无关但同属 registry.py 的一块功能:扫盘找索引。

iter_index_meta_files 遍历目录树找 *.leann.meta.json,默认深度上限 5 层,并且永远跳过 .git / .venv / venv / node_modules / __pycache__ / Library 这 6 个目录(registry.py:19-22 的 DEFAULT_INDEX_SCAN_DEPTH 与 INDEX_SCAN_SKIP_DIRS、:25-47 的遍历)。

register_project_directory 把有索引的项目目录登记进 ~/.leann/projects.json(registry.py:83-134)。它有个小优化:先看有没有 .leann/indexes 目录,有就直接登记,免掉那次目录扫描(registry.py:104-106)。

5. 本章代码地图​

主题文件符号
注册装饰器packages/leann-core/src/leann/registry.py:50register_backend
自动发现packages/leann-core/src/leann/registry.py:61autodiscover_backends
注册表packages/leann-core/src/leann/registry.py:18BACKEND_REGISTRY
索引扫描packages/leann-core/src/leann/registry.py:25iter_index_meta_files
扫描跳过目录packages/leann-core/src/leann/registry.py:20INDEX_SCAN_SKIP_DIRS
三个接口packages/leann-core/src/leann/interface.py:7LeannBackend*Interface
搜索器公共实现packages/leann-core/src/leann/searcher_base.py:12BaseSearcher
HNSW 工厂packages/leann-backend-hnsw/leann_backend_hnsw/hnsw_backend.py:38HNSWBackend
HNSW 开关归一化packages/leann-backend-hnsw/leann_backend_hnsw/hnsw_backend.py:49HNSWBuilder.__init__
整数标签 → passage IDpackages/leann-backend-hnsw/leann_backend_hnsw/hnsw_backend.py:274HNSWSearcher.search 的 map_label
IVF 工厂与建图packages/leann-backend-ivf/leann_backend_ivf/ivf_backend.py:85IVFBackend / IVFBuilder.build
IVF 增删packages/leann-backend-ivf/leann_backend_ivf/ivf_backend.py:200add_vectors / remove_ids
IVF ID 映射packages/leann-backend-ivf/leann_backend_ivf/ivf_backend.py:55_load_id_map / _save_id_map
DiskANN 建索引packages/leann-backend-diskann/leann_backend_diskann/diskann_backend.py:211DiskannBuilder.build
DiskANN 搜索packages/leann-backend-diskann/leann_backend_diskann/diskann_backend.py:475DiskannSearcher.search
内存参数推算packages/leann-backend-diskann/leann_backend_diskann/diskann_backend.py:94_calculate_smart_memory_config
图分区packages/leann-backend-diskann/leann_backend_diskann/graph_partition.py:89GraphPartitioner.partition_graph