Faiss 分布式磁盘索引流水线：从 500M 向量分布式 k-means 到 1T 级 OnDisk IVF 索引的完整构建

本系列文章聚焦 Faiss 仓库 benchs/distributed_ondisk 中面向 1T 级向量数据集的"分布式训练 + 磁盘存储 + 分布式检索"完整流水线。读完后，你将掌握：如何用 Python 层 RPC 把 k-means 的分配阶段拆分到多台机器、如何构造"随机旋转 + HNSW 量化器 + 6-bit 标量量化"的可训练空索引、如何用"垂直切片建索引、水平切片落盘"的两步法拼装磁盘 IVF 索引，以及如何用 VStackInvertedLists 组合检索并进一步把倒排列表分配到多个检索服务器上。

整体思路：为什么需要"磁盘 + 分布式"

当数据量达到 1B（Deep1B 演示集）乃至 1T 向量时，任何单机内存索引都不可行。这条流水线的核心策略可以概括为四点：

1. 训练数据可以远大于内存：向量文件以 memory-mapped（np.memmap/bvecs_mmap/fvecs_mmap）方式读取，文件通过分布式文件系统对全部机器可见；
2. k-means 的分配阶段可以水平拆分：把训练矩阵切成若干段分给不同机器/GPU 做分配，主节点汇总结果更新质心，这只需要在 Python 层重实现 k-means 主循环，无需改动 Faiss C++ 核心；
3. 索引构建可以完全并行：先把全量向量按行切成"垂直切片"（vslice），每个切片独立建一个完整 IVF 索引；再按倒排列表维度切成"水平切片"（hslice），把每个倒排列表的子集写入连续的 OnDiskInvertedLists 文件；
4. 检索也可以水平拆分：查询机只做粗量化（coarse quantization），把要访问的倒排列表按负载均衡地分给多个检索服务器，各自读自己的盘、返回结果后在主节点合并。

整个目录只有 8 个文件，每个文件对应流水线的一个阶段：

文件	流水线角色
distributed_kmeans.py	分布式 k-means 的 client/server 实现与本地测试
run_on_cluster.bash	面向 Slurm 的全部运行入口（8 个 todo）
make_trained_index.py	由质心文件构造"空的可训练索引"
make_index_vslice.py	为数据的一个行区间建立独立索引
merge_to_ondisk.py	把 200 个 vslice 的倒排列表合并成 1 个 hslice
combined_index.py	查询侧：把 50 个 hslice 的倒排列表 vstack 进空索引
search_server.py	分布式检索的 server 与"按列表拆分"客户端
distributed_query_demo.py	分布式检索演示与 recall/耗时评测

所有脚本均为 Python 3，依赖一个包含 faiss 与 scipy（稀疏矩阵，用于 Python k-means 中心更新）的安装。集群调度命令默认写死为 Slurm，README 明确指出"换其他调度器应该相当直接"——实际上只有 run_on_cluster.bash 一个文件依赖 Slurm 的 srun/sbatch 与环境变量，其余脚本都是普通 Python。

阶段一：分布式 k-means

目标：把 500M 训练向量聚成 10M 质心（Deep1B 生产运行实际是 50M 向量聚 1M 质心）。分布式方式非常简单：把训练向量按机器切分，各机器做"分配到最近质心"这一步（这是计算大头），主节点合成各机器的分配结果并更新质心。

三个核心构件

distributed_kmeans.py 依赖 faiss.contrib.clustering 中的抽象（见 contrib/clustering.py）：

• DatasetAssign：对训练矩阵的封装，只暴露 k-means 主循环需要的三件事——count()、dim()、assign_to(centroids, weights) 返回 (I, D, sum_per_centroid)（即分配 ID、距离、各质心的样本和）。它既可以是纯 CPU 的本地矩阵，也可以是 GPU 版本 DatasetAssignGPU，还可以是远端对象的 RPC 句柄——这正是分布式的支点；
• DatasetAssignDispatch：distributed_kmeans.py#L26-L72 中定义的"再分发器"，把多个 DatasetAssign 组合成一个：
  • assign_to 用 imap（可选 ThreadPool 并行）把每个分片自己的 assign_to 并发执行，再把各分片的 I、D 用 np.hstack 拼接、sum_per_centroid 逐元素相加；
  • get_subset 用 np.searchsorted 判断每个全局下标落在哪个分片上，再向对应分片取子集。 构造时 self.cs = np.cumsum([0] + self.sizes) 记录了各分片的全局偏移，这就是"本地切 5 段"与"远程 5 台机器"在客户端看来完全同构的原因；
• AssignServer：distributed_kmeans.py#L75-L83 继承 faiss.contrib.rpc.Server，通过 __getattr__ 把自身属性转发给内部 assign 对象。效果是：客户端拿到的是一个"透明代理"，对它调用 assign_to(...)、count()、dim() 都会被序列化成 RPC 请求发到远端。DatasetAssignDispatch([client, client, ...], True) 中每个元素就是这样的 RPC client。

主循环本身（k-means 迭代、中心更新、不平衡处理）由 contrib/clustering.py 里的 Python 版 kmeans() 完成，README 说明它"紧跟着 Faiss C++ 实现重做，效率不会有显著差距"。

命令行参数

distributed_kmeans.py 的 argparse 定义了四组参数，使用时需要完整理解：

参数组	参数	说明
general	--test	执行本地自检（逗号分隔的数字，见下）
	--k	质心数
	--seed	随机种子，默认 1234
	--niter	迭代次数，默认 20
	--gpu	使用哪块 GPU：-2 不用 GPU，-1 用全部
I/O	--indata	输入数据文件，支持 fvecs / bvecs / npy，以 mmap 方式读入
	--i0 / --i1	参与聚类的向量区间 [i0, i1)，这是数据切分的基本手段
	--out	质心输出文件（.npy）
	--store_each_iteration	每轮迭代存一次质心检查点
server	--server	以 server 模式运行（只做分配，不跑主循环）
	--port	监听端口，默认 12345
	--when_ready	就绪后把 host:port 写入该文件（供上层探测）
	--ipv4	强制 IPv4
client	--client	以 client 模式运行（连接远端 server 后跑主循环）
	--servers	空格分隔的 host:port 列表

注意 main() 里的分支逻辑：给 --indata 时按扩展名 mmap 读入（.bvecs → bvecs_mmap、.fvecs → fvecs_mmap、.npy → np.load(..., mmap_mode="r")），再按 --i0/--i1 截取区间；给 --client 时则对 --servers 逐条 rpc.Client(host, port, v6=not args.ipv4) 建连。二者互斥，缺了会直接 raise AssertionError。

本地四级自检

--test 模式（distributed_kmeans.py#L86-L134 的 do_test）把"同一份数据的 k-means 结果"用四种方式各跑一遍，逐级打开抽象层，用来验证每一层封装的正确性。若本地有 BigANN 学习集文件则使用它，否则退回 SyntheticDataset(128, 100000, 0, 0) 合成数据，并对前 10 万向量做了"前一半全部替换成同一向量"的故意偏斜处理，用来压力测试切分代码。

# 0：Faiss C++ Kmeans 参考实现
python distributed_kmeans.py --test 0
# 1：Python 主循环 + 本地 DatasetAssign
python distributed_kmeans.py --test 1
# 2：DatasetAssignDispatch，本地切 5 段
python distributed_kmeans.py --test 2
# 3：同样用 dispatch，但每段放一块 GPU（ThreadPool 并行）
python distributed_kmeans.py --test 3

对应源码可以看清每一级打开的能力：test 0 直接 faiss.Kmeans(d, 1000, niter=20, verbose=True).train(...)；test 2 是 DatasetAssignDispatch([DatasetAssign(xx[20000*i:20000*(i+1)]) for i in range(5)], False)；test 3 则把 10 万向量按 GPU 数均分给 DatasetAssignGPU 并以 in_parallel=True 建池。四个测试的输出应一致（README 中给出了参考输出 gist 的样子）。

单机模拟分布式：test_kmeans_2

run_on_cluster.bash 用三个 todo 演示"真·分布式"的最小闭环，全部在一台机器上完成：

# 非分布基线：CPU 单机
bash run_on_cluster.bash test_kmeans_0
# 用满本机全部 GPU
bash run_on_cluster.bash test_kmeans_1
# 每块 GPU 起一个 server，client 通过 RPC 连接
bash run_on_cluster.bash test_kmeans_2

脚本头部设好测试参数 nvec=1000000, k=4000（真实运行时注释里改为 nvec=50000000, k=1000000）。test_kmeans_2 的做法值得细看：

1. ngpu=$(echo /dev/nvidia? | wc -w) 数 GPU，对每块 GPU gpu 计算区间 i0 = nvec*gpu/ngpu、i1 = nvec*(gpu+1)/ngpu、端口 baseport+gpu（baseport=12012）；
2. 后台拉起 python distributed_kmeans.py --indata $traindata --i0 $i0 --i1 $i1 --server --gpu $gpu --port $port --ipv4，并把 localhost:$port 累积进 hostports；
3. sleep 5s 保证 server 就绪后，执行 python distributed_kmeans.py --client --servers "$hostports" --k $k --ipv4 跑主循环；
4. 用 trap 'kill -HUP 0' 0 在脚本退出时回收后台进程。

这就是"单机多进程模拟集群"的完整样板：每个 server 只持有数据的一个区间的 mmap，客户端把远端 client 句柄直接塞进 DatasetAssignDispatch，主循环无感知。

Slurm 上真正跨机运行

slurm_distributed_kmeans 这个 todo（run_on_cluster.bash#L78-L86）向 Slurm 要 5 台机器、每台 4 GPU、40 核、100G 内存、48 小时：

srun -n5 \
     --time=48:00:00 \
     --cpus-per-task=40 --gres=gpu:4 --mem=100G \
     --partition=priority \
     -l bash $(realpath $0) slurm_within_kmeans_server

被 srun 在所有节点上执行的 slurm_within_kmeans_server 分支里，各机器通过 Slurm 环境变量确定自己的身份与分工：

• nserv=$SLURM_NPROCS：server 总数；rank=$SLURM_PROCID：本机编号；
• 按 rank 均分数据区间、把端口设为 baseport + rank；
• 非 0 号机：只跑 distributed_kmeans.py --server（阻塞在 RPC 服务上）；
• 0 号机：先把自己那份区间也起成后台 server，再解析 SLURM_TASKS_PER_NODE 与 SLURM_JOB_NODELIST（脚本里专门写了 parse_tasks_per_node() 处理 N*xhost 这种压缩写法，并用 scontrol show hostnames 展开节点列表），拼出全部 host:port 列表，sleep 20s 后跑 client 主循环；client 结束时 scancel $SLURM_JOBID 主动杀掉整个作业。

Deep1B 生产聚类的实际运行

deep1b_clustering 是生产命令（run_on_cluster.bash#L167-L176）：同样走 slurm_within_kmeans_server，但把 nserv 设为 20 台机器，并在末尾透传 --out $workdir/1M_centroids.npy。README 记录的真实运行是 50M 向量聚 1M 质心，输出日志形如：

  Iteration 19 (898.92 s, search 875.71 s): objective=1.33601e+07 imbalance=1.303 nsplit=0
 0: writing centroids to /checkpoint/matthijs/ondisk_distributed/1M_centroids.npy

README 对这次运行的分析很实在：总时长 899 秒中 876 秒是计算，但数据传输开销不可忽略——原因是 Python 实现本身以及 RPC 协议没有针对 broadcast/gather 做优化（不像 MPI），因此"传输 + 质心计算"阶段占比偏高。作者同时强调这是一个"能在大多数集群上跑起来"的简单实现。

阶段二：构造空的可训练索引

质心文件得到后，还不能直接建库——需要一个"知道怎么搜、但还没有数据"的索引，供后续切片脚本反复加载填充。make_trained_index.py 只做三件事，全部代码不到 60 行：

1. 随机旋转。faiss.RandomRotationMatrix(d, d) 以种子 1234 初始化，对质心做 rrot.apply_py(centroids)。目的（README 原文）是"把各维度的尺度拉平"，让后续标量量化在每一维上表现均衡；旋转可以放在聚类之后做，因为"簇也会被同样旋转"，不破坏聚类结果。
2. 用 HNSW 包装质心做量化器。faiss.IndexHNSWFlat(d, 32)，并设置 efSearch=1024、efConstruction=200，把 100 万质心 add 进去。查询时用它加速 CPU 上的粗量化（"向量归属哪个簇"）；
3. 训练 6-bit 标量量化器。最终索引是：

index = faiss.IndexPreTransform(
    rrot,
    faiss.IndexIVFScalarQuantizer(
        quantizer, d, ncent, faiss.ScalarQuantizer.QT_6bit
    ),
)

即 IndexPreTransform(RandomRotationMatrix, IndexIVFScalarQuantizer(quantizer=IndexHNSWFlat, QT_6bit))。index.train(xt) 用 256k 向量完成 SQ 码本训练（fvecs_mmap 读入 learn.fvecs 的前 256000 行后 np.ascontiguousarray 拷贝），最后 faiss.write_index(index, workdir + "trained.faissindex") 落盘。

这套"旋转 + HNSW 粗量化器 + 6bit SQ"的组件选择与后面 combined_index.py 的检索配置互相呼应：加载时会把 quantizer.hnsw.efSearch 再设成 1024（combined_index.py#L60-L61），保证粗量化精度；6-bit SQ 则把每条倒排码长压到 d/2 字节，是"磁盘 IO 换压缩率"的关键。

阶段三：垂直切片建索引（vslice）

所谓 vslice（vertical slice），是按行切开的那部分数据：Deep1B 的 10 亿向量被切成 200 个各含 5M 向量的切片，每个切片单独构建成一个完整、独立的 IVF 索引文件。由于切片之间没有任何依赖，这一步"可以以暴力并行方式跑，对顺序没有任何约束"——run_on_cluster.bash#L178-L210 里就是循环 200 次，每次 sbatch 提交一个单节点作业：

# 每个 vslice 作业的核心命令（脚本自动生成）
srun python -u make_index_vslice.py \
             --inputindex $workdir/trained.faissindex \
             --input $deep1bdir/base.fvecs \
             --nt 40 \
             --i0 $i0 --i1 $i1 \
             -o $workdir/vslices/slice$i.faissindex

资源规格为 -n1 --cpus-per-task=40 --mem=200G，日志写到 vslices/slice$i.log。

make_index_vslice.py 的实现里有三个工程细节值得注意：

• 限速读取：produce_batches() 以 --bs（默认 2^18 条）为批从 mmap 文件取数；rate_limited_iter()（make_index_vslice.py#L40-L58）用单线程 ThreadPool 提前预取下一批，让"读盘"与"索引写入"重叠，避免加库速度被数据库/磁盘读取拖住。README 特别指出：对真实数据集，数据会从 DBMS 读取，此时"读取与建索引并行做是值得的，因为读取非常慢"；
• 显式 ID：index.add_with_ids(x, ids) 直接用全局行号作为向量 ID，保证切片合并后 ID 语义一致；
• 瘦身量化器：默认构建完会把 index_ivf.quantizer.reset() 掉（--keepquantizer 可关闭此行为），因为 200 个 vslice 文件若各自保留一份 HNSW 质心数据，空间浪费巨大；量化器只需要保留"训练好的参数"，数据留一份即可。脚本结束前还会打印 invlists.imbalance_factor() 与 print_stats()，用于核对负载分布。

阶段四：水平切片合并落盘（hslice）

200 个 vslice 各有自己完整的 1M 个倒排列表，但每条列表里只有 1/200 的数据。下一步按列表维度合并：把倒排列表切成 50 个水平切片（hslice），第 i 个 hslice 包含全部 1M 条列表的第 [i0, i1) 条。命令：

bash run_on_cluster.bash make_index_hslices

对应 run_on_cluster.bash#L212-L246：nlist=1000000, nslice=50，每个作业的核心命令是

srun python -u merge_to_ondisk.py \
             --input $workdir/vslices/slice{0..199}.faissindex \
             --nt 20 \
             --l0 $i0 --l1 $i1 \
             --output $workdir/hslices/slice$i.faissindex \
             --outputIL $workdir/hslices/slice$i.invlists

注意 --input 一次传入全部 200 个 vslice 文件，但每个作业只处理自己负责的那段列表区间（--l0/--l1）。merge_to_ondisk.py 的实现要点：

• mmap 读取输入：faiss.read_index(fname, faiss.IO_FLAG_MMAP | faiss.IO_FLAG_READ_ONLY) 配 20 线程 ThreadPool 并行加载 200 个 vslice；
• 摘出并裁剪倒排列表：faiss.downcast_InvertedLists(index_ivf.invlists) 取出列表后 il.this.own() 接管所有权、index_ivf.invlists = None 释放索引本体；非全量区间时先 il.crop_invlists(l0, l1) 裁出自己的列表段（merge_to_ondisk.py#L54-L64）；
• 写 OnDisk 列表：faiss.OnDiskInvertedLists(il0.nlist, il0.code_size, outputIL) 以第一个 hslice 的 nlist/code_size 为准，il.merge_from(ils.data(), ils.size(), True) 把 200 份（已裁剪的）倒排表数据合并进这一个连续磁盘文件；
• 回装索引：把合并后的 OnDiskInvertedLists 换进第一个加载的索引（replace_invlists 语义：设 invlists = il、own_invlists = False、修正 nlist/ntotal），只序列化索引骨架（不含列表数据，数据已在 .invlists 文件里）。

产物是 50 对 hslices/slice$i.faissindex + hslices/slice$i.invlists，这就是最终检索时真正从盘上读的数据。

阶段五：查询索引——CombinedIndex 组装

检索侧的关键是 combined_index.py 中的 CombinedIndex 类。它的 __init__（combined_index.py#L20-L61）做四步：

1. 依次 faiss.read_index(fname) 读入全部 50 个 hslice 索引（CombinedIndexDeep1B 子类把文件名生成为 hslices/slice%d.faissindex，range(50)），并用 faiss.extract_index_ivf(index).invlists 取出各自的倒排表句柄，收集进 faiss.InvertedListsPtrVector；
2. self.big_il = faiss.VStackInvertedLists(ilv.size(), ilv.data())——这是组装的精髓：VStackInvertedLists 把多份倒排表按列表号拼成一个逻辑上完整、物理上分布在 50 个文件的倒排表，不需要再次重写磁盘数据；
3. 读入空索引 trained.faissindex（mmap/只读），计算 ntotal = self.big_il.compute_ntotal()，然后 index_ivf.replace_invlists(self.big_il, False) 把倒排表换进索引、置 index_ivf.parallel_mode = 1（README 注释："看起来一直开着比较合理"，即按查询并行而非按列表并行）；
4. 若提供了 masked_index_fname，还会额外套一层 MaskedInvertedLists，用另一个索引的倒排表做"掩码/覆盖"，本演示中为 None。

CombinedIndex 还刻意把字段与方法暴露成 Python 方法（set_nprobe、set_parallel_mode、set_omp_num_threads、set_prefetch_nthread 等），README 说明这是"为了能被 RPC 调用"——这些方法名与 search_server.py 里客户端远程调用的名字一一对应。

直接运行即可做 recall/耗时评测（__main__ 块会加载 Deep1B 的查询集与 groundtruth，对 nprobe = 1, 10, 100, 1000 各测一次）：

python combined_index.py

README 记录的真实输出：

reading /checkpoint/matthijs/ondisk_distributed//hslices/slice49.faissindex
loading empty index /checkpoint/matthijs/ondisk_distributed/trained.faissindex
replace invlists
loaded index of size  1000000000
nprobe=1 1-recall@1=0.2904 t=12.35s
nnprobe=10 1-recall@1=0.6499 t=17.67s
nprobe=100 1-recall@1=0.8673 t=29.23s
nprobe=1000 1-recall@1=0.9132 t=129.58s

结论直白：磁盘检索比内存检索慢得多（nprobe=1 也要 12 秒，nprobe=1000 要 130 秒），瓶颈是 IO。这也就引出了最后一步。

阶段六：分布式检索（search server）

思路：把查询也按"倒排列表"水平拆分到多台检索服务器上，每台只读自己负责的列表段，最终只把有效结果（k 个）回传给主节点，从而大幅降低主机的磁盘带宽需求。

启动 server（search_server.py）：MyServer 与前面的 AssignServer 同款——rpc.Server 子类用 __getattr__ 把 CombinedIndexDeep1B 的方法透出去；server 侧先完整加载一遍 CombinedIndexDeep1B() 再 rpc.run_server(...) 监听端口。集群入口：

bash run_on_cluster.bash run_search_servers   # srun -n3，64 核、100G 内存，端口 12012

server 参数支持 --port（默认 12012）、--when_ready_dir（就绪后按 --rank 写 host:port 到目录下的文件，便于客户端发现）、--ipv4、--rank。

客户端：search_server.py#L96-L170 的 SplitPerListIndex 是"本地粗量化 + 远端细检"的编排器，其 search(x, k) 的调用链是：

1. self.index.transform_and_assign(x)：在本地索引上做预处理与粗量化——CombinedIndex.transform_and_assign（combined_index.py#L73-L85）先应用 IndexPreTransform 的第一级变换（随机旋转），再用 quantizer.search(xq, index_ivf.nprobe) 得到 list_nos 与 coarse_dis；
2. balance_lists(list_nos)：按各倒排列表的 list_size 取权重，用 distribute_weights()（贪心：按权重从大到小依次塞进当前最轻的 bin）把待访问列表均衡地分给 ni 个远端子索引；
3. 对每个子索引，把不属于自己的列表置为 -1（list_nos_i[assign != i] = -1），远程调用 ivf_search_preassigned(xqo, list_nos_i, coarse_dis, k)——CombinedIndex 里对应的方法直接调 C++ 层 index_ivf.search_preassigned（用 faiss.swig_ptr 传裸指针）；
4. 主节点用 faiss.ResultHeap 收集各远端返回的 (D, I)，finalize() 得到最终 top-k；range_search 同理，走 ivf_range_search_preassigned 并把各远端的 (lims, D, I) 按查询拼接。

演示脚本 distributed_query_demo.py 的用法：命令行传入各 server 主机名，本地加载一份 CombinedIndexDeep1B（用于粗量化），对每台 rpc.Client(host, 12012, v6=False) 建连，先 get_ntotal() 自检连通性，再 search_server.SplitPerListIndex(ci, clients) 包装，最后统一设置 parallel_mode=1、prefetch_nthread=0、omp_num_threads=64，对 nprobe = 1, 10, 100, 1000 评测 recall 与耗时。README 说明：输出中的 "MiB" 数字表示该次检索实际从盘上读入的数据量；在 1B 这个规模上分布式检索收益还不明显，但"在 10 倍以上的数据集上，差距会显著起来"。

适用边界与实现取舍

• 数据集与路径：全部脚本默认写死 Deep1B 的路径（如 /datasets01_101/simsearch/041218/deep1b/、/checkpoint/matthijs/ondisk_distributed/），复用到自己的数据前需要改脚本顶部变量与 run_on_cluster.bash 的 traindata/workdir，并把 fvecs/bvecs/npy 文件放到各节点可见的共享文件系统上；
• 调度系统：仅 run_on_cluster.bash 绑定 Slurm（srun/sbatch/SLURM_NPROCS/SLURM_PROCID/SLURM_JOB_NODELIST/scontrol/scancel），换成其他调度器时只需重写这一个文件里"如何分配机器并在每台机器上起同一命令"的部分；
• 通信开销：如 Deep1B 聚类那次运行所示，RPC 协议的逐向量传参与 Python 层中心更新是主要非计算开销，README 明确承认这相对 MPI 式 broadcast/gather 是低效的，但换来的是"在大多数集群上直接可跑"的简单性；
• 与生产系统的差异：README 结尾指出，这套代码包含构建"可扩展到 1T 向量索引"的核心组件，但相比实际生产使用的索引做了若干简化（如单机 CombinedIndex 的 vstack 组装在 1T 规模下需要更细粒度的调度，此处从略）。

整条流水线的价值在于：它把"1T 级向量索引"拆成了六个彼此解耦、每一步都可以独立验证的阶段——分布式 k-means（有四级本地自检）、可训练空索引、200 路并行 vslice、50 路并行 hslice、mmap vstack 组装、按列表拆分的分布式检索，全部建立在 Faiss 现成的 InvertedLists 抽象（crop_invlists、merge_from、OnDiskInvertedLists、VStackInvertedLists、search_preassigned）之上，几乎没有自定义的 C++ 代码，因此对希望在自己的集群上复刻这条流水线的工程师来说，是一份可以直接逐文件对照源码执行的参考实现。