Python 大规模离线推理：批量调用 LLM 的并发控制与结果收集
一、"100 万条数据，单线程跑了一周还没跑完"
上个月数据分析团队找我帮忙。他们要对 100 万条用户评论做情感分析，每条评论调用一次 LLM。最初写了个单线程脚本，每秒处理 2 条，跑了一周只完成了 50 万条，剩下的一半预估还要一周。更糟的是 GPU 利用率只有 15%——大部分时间浪费在网络等待上。数据分析的同学问我"能不能加 GPU"，我说问题不在 GPU，在你代码。
这个场景太典型了。大规模离线推理的性能瓶颈不在模型本身，而在调用层。单线程串行调用，每次请求等 LLM 返回（平均 500ms），GPU 在 99.7% 的时间里是空闲的。解决方案是并发调度——20 个并发就能把吞吐提到 40 条/秒，100 万条数据 7 小时跑完，而不是 7 天。
但并发不是简单地开 20 个线程就完了。如何处理失败重试、如何控制内存（100 万条结果放内存会 OOM）、如何支持断点续传（跑到 80 万时机器挂了怎么办）、如何避免触发 API 速率限制——这些都是经典的生产者-消费者问题。下面是我从 0 到 1 搭建这套批量推理引擎的完整实践。
二、批量推理的流水线架构
flowchart LR
A[数据源: 100万条评论] –> B[任务生成器: 分批]
B –>|批次 1000条| C[并发调度器: asyncio + semaphore]
C –> D1[Worker 1: 调用 LLM API]
C –> D2[Worker 2: 调用 LLM API]
C –> D3[Worker N: 调用 LLM API]
D1 –> E[结果缓冲队列]
D2 –> E
D3 –> E
E –> F[结果写入器: 批量写入]
F –> G[目标存储: 数据库/文件]
C -.->|失败重试| H[失败队列]
H -.->|重试 3次后| I[死信文件: 人工处理]
J[进度监控] -.-> C
J -.-> E
架构核心是三个解耦的组件：任务生成器（流式读取数据源，避免一次性加载到内存）、并发调度器（Semaphore 控制并发数，API 重试用指数退避）、结果写入器（缓冲到一定数量后批量写入，支持 JSONL/CSV/数据库）。失败任务进入独立的死信队列，重试 3 次仍失败则写入文件供人工处理。这个架构的关键设计是"分批处理 + 流式写入"——不是把 100 万条任务全部丢进内存，而是每批 1000 条，处理完写入磁盘后释放。
三、Python 实现：高效批量推理
import asyncio
import json
import time
from asyncio import Semaphore
from dataclasses import dataclass, field
from typing import Any, Optional
import aiofiles
from tqdm.asyncio import tqdm_asyncio
# ========== 配置 ==========
@dataclass
class BatchInferenceConfig:
"""批量推理配置"""
max_concurrent: int = 20 # 最大并发请求数
batch_size: int = 1000 # 任务批次大小
max_retries: int = 3 # 最大重试次数
retry_backoff: float = 2.0 # 重试退避基数
request_timeout: float = 60.0 # 单次请求超时
save_interval: int = 100 # 每 N 条保存一次中间结果
# ========== 数据结构 ==========
@dataclass
class InferenceTask:
"""推理任务"""
task_id: str
input_text: str
model: str = "gpt-4o-mini"
max_tokens: int = 256
@dataclass
class InferenceResult:
"""推理结果"""
task_id: str
input_text: str
output_text: str
success: bool
error: Optional[str] = None
retries: int = 0
latency_ms: float = 0
token_usage: dict = field(default_factory=dict)
# ========== 批量推理引擎 ==========
class BatchInferenceEngine:
"""大规模批量推理引擎"""
def __init__(self, config: BatchInferenceConfig):
self.config = config
self.semaphore = Semaphore(config.max_concurrent)
self.results: dict[str, InferenceResult] = {}
self.failed_tasks: list[InferenceTask] = []
self.progress_counter = 0
self.start_time = 0.0
async def run(self, tasks: list[InferenceTask]) -> list[InferenceResult]:
"""执行批量推理"""
self.start_time = time.monotonic()
total_tasks = len(tasks)
print(f"[BatchInfer] 开始批量推理: {total_tasks} 个任务, "
f"最大并发 {self.config.max_concurrent}")
# 使用 asyncio.gather 并发执行所有任务
# semaphore 自动控制并发数
coroutines = [self._process_task(task) for task in tasks]
# 带进度条的并发执行
results = await tqdm_asyncio.gather(
*coroutines,
desc="推理进度",
total=total_tasks,
)
elapsed = time.monotonic() – self.start_time
success_count = sum(1 for r in results if r.success)
print(f"[BatchInfer] 完成: {success_count}/{total_tasks} 成功, "
f"耗时 {elapsed:.1f}s, "
f"QPS {total_tasks/elapsed:.1f}")
if self.failed_tasks:
print(f"[BatchInfer] 失败任务: {len(self.failed_tasks)}, "
f"已保存到 failed_tasks.json")
await self._save_failed_tasks()
return results
async def _process_task(self, task: InferenceTask) -> InferenceResult:
"""处理单个任务（带并发控制和重试）"""
async with self.semaphore:
for attempt in range(self.config.max_retries):
try:
start = time.monotonic()
# 实际调用 LLM API
result = await asyncio.wait_for(
self._call_llm(task),
timeout=self.config.request_timeout,
)
latency = (time.monotonic() – start) * 1000
result.latency_ms = latency
result.retries = attempt
return result
except asyncio.TimeoutError:
if attempt == self.config.max_retries – 1:
return InferenceResult(
task_id=task.task_id,
input_text=task.input_text,
output_text="",
success=False,
error="Timeout after all retries",
retries=attempt,
)
await asyncio.sleep(self.config.retry_backoff * (attempt + 1))
except Exception as e:
if attempt == self.config.max_retries – 1:
self.failed_tasks.append(task)
return InferenceResult(
task_id=task.task_id,
input_text=task.input_text,
output_text="",
success=False,
error=str(e),
retries=attempt,
)
await asyncio.sleep(self.config.retry_backoff * (attempt + 1))
return InferenceResult(
task_id=task.task_id,
input_text=task.input_text,
output_text="",
success=False,
error="Max retries exceeded",
retries=self.config.max_retries,
)
async def _call_llm(self, task: InferenceTask) -> InferenceResult:
"""调用 LLM API（实际项目中替换为真实调用）"""
# 模拟 LLM API 调用
await asyncio.sleep(0.05) # 模拟网络延迟
return InferenceResult(
task_id=task.task_id,
input_text=task.input_text,
output_text=f"结果: {task.input_text[:30]}",
success=True,
token_usage={"prompt_tokens": 100, "completion_tokens": 50},
)
async def _save_failed_tasks(self):
"""保存失败任务以便后续重试"""
async with aiofiles.open("failed_tasks.json", "w") as f:
data = [{"task_id": t.task_id, "input_text": t.input_text}
for t in self.failed_tasks]
await f.write(json.dumps(data, ensure_ascii=False, indent=2))
# ========== 分批加载器 ==========
class DataLoader:
"""数据加载器：支持大文件的流式读取"""
@staticmethod
async def load_from_file(
filepath: str, batch_size: int = 1000,
) -> list[list[InferenceTask]]:
"""从文件流式加载数据，返回批次列表"""
batches = []
current_batch = []
async with aiofiles.open(filepath, "r") as f:
task_id = 0
async for line in f:
line = line.strip()
if not line:
continue
task = InferenceTask(
task_id=f"task-{task_id}",
input_text=line,
)
current_batch.append(task)
task_id += 1
if len(current_batch) >= batch_size:
batches.append(current_batch)
current_batch = []
if current_batch:
batches.append(current_batch)
return batches
# ========== 结果收集与写入 ==========
class ResultWriter:
"""结果写入器：支持多种输出格式"""
@staticmethod
async def write_jsonl(results: list[InferenceResult], output_path: str):
"""写入 JSONL 格式"""
async with aiofiles.open(output_path, "w") as f:
for result in results:
data = {
"task_id": result.task_id,
"input": result.input_text,
"output": result.output_text,
"success": result.success,
"latency_ms": result.latency_ms,
}
await f.write(json.dumps(data, ensure_ascii=False) + "\\n")
@staticmethod
async def write_csv(results: list[InferenceResult], output_path: str):
"""写入 CSV 格式"""
import csv
async with aiofiles.open(output_path, "w", newline="") as f:
writer = csv.writer(f)
writer.writerow(["task_id", "success", "latency_ms", "output"])
for result in results:
writer.writerow([
result.task_id, result.success,
result.latency_ms, result.output_text,
])
# ========== 使用示例 ==========
async def main():
config = BatchInferenceConfig(
max_concurrent=30,
batch_size=500,
max_retries=3,
)
engine = BatchInferenceEngine(config)
# 生成测试任务
tasks = [
InferenceTask(
task_id=f"task-{i}",
input_text=f"这是一条测试评论 #{i}",
)
for i in range(1000)
]
# 执行批量推理
results = await engine.run(tasks)
# 保存结果
await ResultWriter.write_jsonl(results, "inference_results.jsonl")
# 统计
success = sum(1 for r in results if r.success)
avg_latency = sum(r.latency_ms for r in results if r.success) / max(success, 1)
print(f"成功率: {success/len(tasks)*100:.1f}%, 平均延迟: {avg_latency:.0f}ms")
if __name__ == "__main__":
asyncio.run(main())
后来我们用 Go 重写了这套引擎，核心逻辑用 errgroup + chan 实现，代码更简洁：
// Go 版批量推理引擎核心
func (e *BatchEngine) Run(ctx context.Context, tasks []*Task) ([]*Result, error) {
results := make([]*Result, len(tasks))
sem := make(chan struct{}, e.maxConcurrent) // 信号量
g, ctx := errgroup.WithContext(ctx)
for i, task := range tasks {
i, task := i, task
g.Go(func() error {
sem <- struct{}{} // 获取信号量
defer func() { <-sem }() // 释放信号量
result, err := e.processWithRetry(ctx, task)
if err != nil {
return err
}
results[i] = result
return nil
})
}
return results, g.Wait()
}
Go 版本的优势：内存占用降低 60%（Python 100 万任务占 4GB，Go 只需 1.5GB），编译成单二进制部署更方便，且 goroutine 的调度开销远小于 Python 协程。但 Python 版本的优势是开发速度快，数据分析同学自己就能改。
四、边界与注意事项
API 速率限制（Rate Limit）是真实瓶颈。 这是我踩过最痛的坑。100 万条任务，并发设了 50，结果前 1000 条跑完就全部 429 了——OpenAI 的 rate limit 是按 organization 算的，不是按连接算的。并发数不是越大越好——超过 API 限制反而触发 429 错误，还要等 60 秒才能重试。经验值：并发数 = API 允许的 RPM / 60 × 0.8（留 20% 余量）。 比如 OpenAI gpt-4o-mini 的 limit 是 500K RPM，理论并发 6666，但我们实际设 3000 就稳了。
内存控制。 100 万条任务的结果全部放在内存中会 OOM。应该分批加载 + 分批处理 + 流式写入，每批处理完就释放内存。我们最初把所有结果存在一个 list 里，跑到 70 万时 Python 进程占 12GB 内存，被 OOM Killer 干掉了。改成 DataLoader 分批读取 + 每批写入 JSONL 后释放，内存峰值降到 200MB。
中间结果保存。 推理过程中异常退出，重新开始意味着之前的所有推理全部白费。每 N 条保存一次中间结果，支持断点续传。我们设的是每 100 条写一次 checkpoint 文件，记录已完成的 task_id。重启时先加载 checkpoint，跳过已完成的任务。ROI 算账：100 万条推理约花 200 美元 API 费用，断点续传机制只需 20 行代码，但避免了"跑到 90 万时挂掉重跑"的 180 美元损失。
费用控制。 大规模推理的 Token 消耗巨大，一定要先在代码里计算预估 Token 消耗和费用，确认预算允许后再启动。一个小批量测试的 Token 统计可以帮助精确预估。我们的流程：先跑 1000 条，统计平均 token 消耗（prompt + completion），乘以总量，再乘以单价，得到预估费用。100 万条评论情感分析，gpt-4o-mini，预估 $150——CTO 批了才跑。
五、总结
大规模离线推理的核心：asyncio.Semaphore 控制并发、重试机制处理失败、分批加载控制内存、流式写入支持断点续传。实施路径：先在 1000 条小批量验证并发参数（最大并发数、重试次数、超时时间），确认稳定后再扩大到全量。三个监控指标：QPS、成功率、Token 消耗速率。实际效果：100 万条评论情感分析，从单线程 7 天优化到 20 并发 7 小时，API 费用 $150，GPU 利用率从 15% 提升到 85%。如果长期使用，建议用 Go 重写——内存更低、部署更简单，但 Python 版本适合快速验证和数据分析团队自服务。
