Embedding批量处理，大量文档向量化任务多进程加速脚本

当你有成千上万篇文档需要转成Embedding向量时，单进程逐条调用接口往往耗时几小时甚至更久。
本文提供一个可直接运行的Python多进程加速方案，讲解数据分片、并发参数、断点续跑和错误重试等关键逻辑，适合有批量文档向量化需求但脚本编写经验不足的读者参考。
实际测试中，在机器核数充足且接口允许并发的情况下，多进程通常能把整体耗时缩短到原来的三分之一以下。

为什么Embedding批量处理要做多进程加速

Embedding向量化是典型的I/O密集型任务，每调用一次接口就要等待网络返回。
如果只有单个进程逐条处理，CPU大部分时间都在空转，磁盘和网络带宽也利用不起来。
多进程的核心思路是让多个进程同时向Embedding接口发送请求，把等待时间重叠起来。

这里要特别注意的是，进程数不是越多越好。
接口服务商通常有并发上限，超出后会返回限流错误。
所以脚本里应该让进程数可配置，避免把接口打爆。

写脚本前要确认的3个条件

1. 机器核数：多进程不能无限开，建议按CPU核心数或接口并发上限来设置进程数，例如 os.cpu_count() 减一。
2. Embedding接口调用方式：无论是OpenAI兼容接口还是本地模型，都要确认并发限制和超时时间，避免被限流或报错。
3. 原始数据存储格式：建议先把所有文档转成统一的JSON或CSV，每行包含 doc_id 和 text，方便后续分片和断点续跑。

多进程加速脚本的核心代码拆解

整体思路是把文档列表切成多个分片，每个进程独立处理一个分片，最后汇总结果。
以下使用 concurrent.futures.ProcessPoolExecutor，比手写 multiprocessing.Pool 更容易管理返回值。

import json
import time
import os
from concurrent.futures import ProcessPoolExecutor, as_completed

def embed_texts(texts):
    # 这里替换成你自己的Embedding接口调用
    # 返回[(doc_id, vector), ...]
    results = []
    for doc_id, text in texts:
        vector = call_embedding_api(text)  # 你的实际调用
        results.append((doc_id, vector))
    return results

def parse_chunk(chunk):
    # 每个分片独立执行，避免共享变量冲突
    return embed_texts(chunk)

def main():
    with open("docs.json", "r", encoding="utf-8") as f:
        docs = json.load(f)  # 形如[{"doc_id": "1", "text": "..."}]
    # 按进程数分片
    num_workers = max(2, os.cpu_count() - 1)
    chunks = [docs[i::num_workers] for i in range(num_workers)]
    all_results = []
    with ProcessPoolExecutor(max_workers=num_workers) as executor:
        futures = [executor.submit(parse_chunk, c) for c in chunks]
        for future in as_completed(futures):
            all_results.extend(future.result())
    # 保存结果
    with open("embeddings.jsonl", "w", encoding="utf-8") as f:
        for doc_id, vec in all_results:
            f.write(json.dumps({"doc_id": doc_id, "vector": vec}) + "\n")

if __name__ == "__main__":
    main()

注意：如果你的Embedding接口客户端不是线程安全的，请在每个进程内部初始化，不要全局复用同一个连接。

运行后如何验证加速效果

跑完一遍之后，不要只看“好像变快了”，建议用下面两步做确认：

1. 记录单进程和多进程的耗时对比：同一个数据集分别跑两次，或看日志中的起止时间。如果多进程耗时反而更长，优先检查是否被接口限流。
2. 校验向量结果数量：统计生成的 embeddings.jsonl 行数，必须和输入文档数一致。还可以随机抽几条文本，重新生成向量比对相似度。

多进程Embedding最容易踩的坑

• 接口限流导致大量失败：不要一上来就把进程数调到CPU核数。先开2个进程测试，如果没有 429 或 Too Many Requests 再逐步增加。
• 数据分片不均：如果每个文档长度差异大，简单切分会导致某些进程运行很久。可以先按文本长度排序再分片，尽量让每个分片的总字符数接近。
• 异常导致整批中断：建议在 embed_texts 内部增加重试机制，比如单条失败后等待1秒重试3次，仍失败则写入 failed_ids.txt，方便后续补跑。
• 进程间复制大对象：如果传入的 chunk 是包含大量重复字段的DataFrame，序列化开销会抵消加速。最好只传 doc_id 和 text 两个字段的简单列表。

如果你正在处理大量文档向量化任务，建议先按本文步骤写一个小的数据分片测试脚本，确认接口并发限制后，再正式跑全量数据。
遇到异常时优先查看失败日志和重试记录，不要盲目调大进程数。