数据截至 (上游 commit dc85934f318c)
02 — 一次查询的端到端
本章讲什么:
searcher.search("...")这一行背后到底发生了什么。重点是那个反直觉的部分——C++ 图遍历过程中反过来调用 Python。
1. 它要解决的小问题
01 讲完索引里没有向量。那问题就来了:图遍历要算"query 和候选节点有多像",没有向量怎么算?
答案:现算。而且只算遍历真正碰到的那些节点——这就是 "selective recomputation(选择性重算)"里"选择性"三个字的含义。
2. 全链路图
从上到下是时间顺序。虚线框里是本章的核心:一次查询里这个回环会跑很多轮。
query 文本
│
▼
① 确保嵌入服务在跑(能复用就复用)──▶ 拿到 zmq_port
│
▼
② 把 query 算成向量(也走这个服务)
│
▼
③ 后端在图上遍历(C++)
┌ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ┐
│ 需要某批节点的距离 │
│ │ ZMQ REQ [[ids],[query向量]] │
│ ▼ │
│ Python:查原文 → 算 embedding → 算距离│
│ │ ZMQ REP [[距离...]] │
│ ▼ │
│ C++ 继续遍历 │
└ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ┘
│
▼
④ top-k 整数标签 ─▶ ids.txt ─▶ passage_id ─▶ 偏移表 ─▶ 原文
│
▼
list[SearchResult]
3. 阶段①:把嵌入服务拉起来
3.1 谁负责
LeannSearcher.search 在 effective_recompute 为真时调 self.backend_impl._ensure_server_running(...),拿回真正可用的端口(api.py:1347-1355)。默认请求端口 5557(api.py:1249),但实际端口可能被让开(见 3.3)。
_ensure_server_running 定义在公共基类里(searcher_base.py:62),它做两件预处理:
- 从 kwargs 或 meta 里定出
distance_metric(searcher_base.py:73-77)。 - 把所有 prompt 模板从 provider options 里过滤掉——因为模板已经在
compute_query_embedding里加过了,服务端再加一次就重复了(searcher_base.py:79-86)。
3.2 服务进程怎么起
EmbeddingServerManager._build_server_command 拼出的是一条 python -m leann_backend_hnsw.hnsw_embedding_server --zmq-port ... --model-name ... --passages-file <meta.json> 命令(embedding_server_manager.py:451-480)。注意传的是 meta.json 路径,不是 passages 路径——服务端自己用 PassageManager 解析(hnsw_embedding_server.py:115-124),而且明确拒绝非 .meta.json 的输入。
provider options 通过环境变量 LEANN_EMBEDDING_OPTIONS 传进子进程(embedding_server_manager.py:510-512,服务端解析在 hnsw_embedding_server.py:48-55)。
3.3 daemon 复用:靠一个配置指纹
默认 use_daemon=True、TTL 900 秒(api.py:1160-1161)。复用逻辑是这样的:
配置指纹 = sha256({
backend_module_name,
model_name, embedding_mode, distance_metric,
passages_file 绝对路径,
provider_options,
passages 签名(meta 与各 passage/idx 文件的 mtime_ns + size)
})
│
▼
~/.leann/servers/<指纹>.json ← 记录 {pid, port, ttl, 指纹}
│
▼
pid 还活着 且 端口还通 ? ──否──▶ 删记录,起新进程
│是
▼
直接复用这个端口
指纹构造在 _build_config_signature(embedding_server_manager.py:349-369),其中 passages 签名由 _build_passages_signature 生成——它会把 meta 里列出的每个 passage / idx 文件的 mtime_ns 和 size 都收进去(embedding_server_manager.py:157-207)。语料一变,指纹就变,旧 daemon 自动不被复用。 这是防"拿旧原文算向量"的关键一道闸。
注册表读写在 _write_registry_record(:797)和 _adopt_registered_server(:820);采纳前会校验 pid 存活和端口可连(:842-846)。
3.4 检查—启动这段临界区上了两层锁
_registry_lock 同时用:
- 线程锁(模块级、按注册表 key 分):因为 POSIX 的
fcntl.flock是进程粒度,同进程内多线程并不互斥。 - 文件锁:POSIX 用
fcntl.flock,Windows 用msvcrt.locking。
理由写在 docstring 里(embedding_server_manager.py:728-740)。还配了陈旧锁回收:锁信息文件超过 600 秒或持有者 pid 已死就删掉(:768-781、_LOCK_STALE_SECONDS 在 :26)。
3.5 端口选择与 warmup
端口从请求值开始往上试 100 个,找第一个能 bind 的(_get_available_port,embedding_server_manager.py:101-111)。所以"请求 5557、实际 5561"是正常现象,增量更新路径会为此专门打日志(api.py:1104-1110)。
LeannSearcher 默认 enable_warmup=True,构造完立刻用 "__LEANN_WARMUP__" 跑一次嵌入把模型加载起来,失败只警告不抛(api.py:1227-1228、:1230-1238)。服务端也有对称的 warmup(hnsw_embedding_server.py:90-101)。
4. 阶段②:query 向量
compute_query_embedding(searcher_base.py:104)的顺序很讲究:
- 先套 prompt 模板(如果有),保证走服务和走本地回退两条路结果一致(
searcher_base.py:123-126)。 - 尝试走服务:
_compute_embedding_via_server发一个纯字符串列表过去,30 秒超时(searcher_base.py:164-194)。 - 服务失败就降级到本地直接算,只打印 警告(
searcher_base.py:149-151)。
模板本身的取值有一条优先级链,在 LeannSearcher.search 里:调用时传的 provider_options["prompt_template"] > meta 里的 query_prompt_template > meta 里的 prompt_template > 无(api.py:1362-1373)。
5. 阶段③:遍历 + 回调(核心)
5.1 Python 侧交给 C++ 什么
HNSWSearcher.search(hnsw_backend.py:174)把参数打包成 faiss.SearchParametersHNSW:
| 参数 | 来源 | 作用 |
|---|---|---|
efSearch | complexity(默认 64) | 候选列表大小,越大越准越慢 |
beam_size | beam_width(默认 1) | 并行搜索路径数 |
pq_pruning_ratio | prune_ratio | 用近似距离先砍掉多少邻居 |
local_prune / send_neigh_times_ratio | pruning_strategy | 三选一:global / local / proportional |
batch_size | batch_size | 邻居批处理大小,0 = 关 |
zmq_port | 上一步拿到的端口 | 回调打哪儿 |
映射代码在 hnsw_backend.py:226-257;端口还额外通过 self._index.set_zmq_port(zmq_port) 设了一次(:218-219)。
有一处针对 OpenAI 嵌入的特判:cosine 度量 + 模型名含 text-embedding / openai 时,关掉 check_relative_distance,防止分数区间太窄导致过早终止(hnsw_backend.py:232-240)。
索引本身是 mmap 打开的,加载时把 is_compact / is_recompute 通过 HNSWIndexConfig 告诉 C++(hnsw_backend.py:156-162)。
5.2 服务端认得 4 种消息
路由在 hnsw_embedding_server.py:371-396,按消息形状判别:
| 消息形状 | 含义 | 处理函数 |
|---|---|---|
["__QUERY_MODEL__"] | 问一下当前模型名 | 就地回 [model_name](:373-378) |
| 全是字符串的列表 | 直接给文本要向量 | _handle_text_embedding(:195) |
[[ids], [query向量]] | 要这批节点到 query 的距离 | _handle_distance_request(:211) |
| 其他 | 按 ID 要向量(回退路径) | _handle_embedding_by_id(:273) |
第三种是主路径。它的妙处在于:距离在 Python 侧就算完了,回传的是一个 float 列表而不是一大堆向量,省带宽。
5.3 距离怎么算
# 示意,非源码:两种度量的分支
if distance_metric == "l2":
d = np.sum(np.square(embeddings - query.reshape(1, -1)), axis=1) # 越小越近
else:
d = -np.dot(embeddings, query) # 内积取负,越小越近
真实实现在 hnsw_embedding_server.py:257-262。内积取负是为了让两种度量对 C++ 侧呈现同一个语义:越小越近。
5.4 节点 ID 怎么变成原文
C++ 给的是整数标签,要先经 ids.txt 映射成 passage ID,再经偏移表取原文:
整数标签 ──ids.txt 第 n 行──▶ passage_id ──偏移表──▶ JSONL 某一行 ──▶ text
映射函数是 _map_node_id(hnsw_embedding_server.py:151-159),ids.txt 的定位逻辑要连剥两次后缀(.meta.json 和 .leann)才拼出文件名(:134-141)。找不到 ids.txt 就退化成直接把整数当字符串用(:147)。
5.5 容错:哨兵距离
查不到 passage、文本为空、模型抛异常——这些情况都不让整次搜索崩掉,而是给那个位置填 1e9(hnsw_embedding_server.py:243-244、:266-267)。C++ 侧看到一个巨大的距离,自然就不会选它。
更深一层的兜底是 _build_safe_fallback(:180-193):连 msgpack 都解不开时,按上一次请求的类型和长度造一个形 状合法的空回复。这是必须的——ZMQ 的 REQ/REP 是严格轮换的,漏发一次回复,对面就永久卡住。
5.6 daemon 的自杀开关
主线程每 0.1 秒检查一次:若开了 daemon 模式且距上次请求超过 TTL,就优雅关闭(hnsw_embedding_server.py:472-480)。last_activity 在每次成功 recv 时更新(:353)。
6. 阶段④:回填原文
LeannSearcher.search 拿到 {"labels": [[...]], "distances": [[...]]} 后逐条查原文、组装 SearchResult;某个 ID 查不到就打红色错误日志跳过,不中断(api.py:1437-1471)。
还有一个前置的贴心处理:top_k 超过总文档数时自动下调并告警(api.py:1310-1317)。
7. 关键细节与坑
7.1 recompute 开关有两处,别用错
recompute_embeddings 既能在构造 LeannSearcher 时给(api.py:1159),也能在 search() 时给——但后者已废弃,传了会打警告(api.py:1336-1341)。正确做法是在构造时配。
7.2 剪枝索引上关掉 recompute 会直接报错
RuntimeError: Recompute is required for pruned/compact HNSW index.
判定在 hnsw_backend.py:210-214。原因很直白:文件里没有向量。错误信息里给的两条出路——加 --recompute,或者用 --no-recompute --no-compact 重建——是准确的。
7.3 可选的 query 日志
设了环境变量 LEANN_QUERY_LOG=<path>,每次搜索会追加一行 JSON:query、top_k、结果 ID 和分数,有向量时连向量一起(api.py:1221-1224、_log_query 在 :1490)。用于离线 benchmark 回放。
7.4 资源清理
LeannSearcher 同时实现了 cleanup()、__enter__/__exit__ 和 __del__(api.py:1640-1669)。docstring 特别点名 Windows:原生后端(如 DiskANN 的 mmap)不释放文件句柄的话,索引目录删不掉。
8. 本章代码地图
| 主题 | 文件 | 符号 |
|---|---|---|
| 搜索主入口 | packages/leann-core/src/leann/api.py:1240 | LeannSearcher.search |
| 预热 | packages/leann-core/src/leann/api.py:1230 | LeannSearcher.warmup |
| query 日志 | packages/leann-core/src/leann/api.py:1490 | _log_query |
| 起服务(公共) | packages/leann-core/src/leann/searcher_base.py:62 | BaseSearcher._ensure_server_running |
| query 向量 + 降级 | packages/leann-core/src/leann/searcher_base.py:104 | compute_query_embedding |
| ZMQ 客户端 | packages/leann-core/src/leann/searcher_base.py:164 | _compute_embedding_via_server |
| 服务进程管理 | packages/leann-core/src/leann/embedding_server_manager.py:242 | EmbeddingServerManager.start_server |
| 配置指纹 | packages/leann-core/src/leann/embedding_server_manager.py:349 | _build_config_signature |
| 语料签名 | packages/leann-core/src/leann/embedding_server_manager.py:157 | _build_passages_signature |
| daemon 采纳 | packages/leann-core/src/leann/embedding_server_manager.py:820 | _adopt_registered_server |
| 双层锁 | packages/leann-core/src/leann/embedding_server_manager.py:728 | _registry_lock |
| HNSW 搜索参数映射 | packages/leann-backend-hnsw/leann_backend_hnsw/hnsw_backend.py:174 | HNSWSearcher.search |
| ZMQ 服务主体 | packages/leann-backend-hnsw/leann_backend_hnsw/hnsw_embedding_server.py:58 | create_hnsw_embedding_server |
| 距离请求处理 | packages/leann-backend-hnsw/leann_backend_hnsw/hnsw_embedding_server.py:211 | _handle_distance_request |
| 按 ID 取向量 | packages/leann-backend-hnsw/leann_backend_hnsw/hnsw_embedding_server.py:273 | _handle_embedding_by_id |
| 安全回退回复 | packages/leann-backend-hnsw/leann_backend_hnsw/hnsw_embedding_server.py:180 | _build_safe_fallback |