Python 大规模离线推理批量调用 LLM 的并发控制与结果收集一、100 万条数据单线程跑了一周还没跑完上个月数据分析团队找我帮忙。他们要对 100 万条用户评论做情感分析每条评论调用一次 LLM。最初写了个单线程脚本每秒处理 2 条跑了一周只完成了 50 万条剩下的一半预估还要一周。更糟的是 GPU 利用率只有 15%——大部分时间浪费在网络等待上。数据分析的同学问我能不能加 GPU我说问题不在 GPU在你代码。这个场景太典型了。大规模离线推理的性能瓶颈不在模型本身而在调用层。单线程串行调用每次请求等 LLM 返回平均 500msGPU 在 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_factorydict) # 批量推理引擎 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推理进度, totaltotal_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, fQPS {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), timeoutself.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_idtask.task_id, input_texttask.input_text, output_text, successFalse, errorTimeout after all retries, retriesattempt, ) 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_idtask.task_id, input_texttask.input_text, output_text, successFalse, errorstr(e), retriesattempt, ) await asyncio.sleep(self.config.retry_backoff * (attempt 1)) return InferenceResult( task_idtask.task_id, input_texttask.input_text, output_text, successFalse, errorMax retries exceeded, retriesself.config.max_retries, ) async def _call_llm(self, task: InferenceTask) - InferenceResult: 调用 LLM API实际项目中替换为真实调用 # 模拟 LLM API 调用 await asyncio.sleep(0.05) # 模拟网络延迟 return InferenceResult( task_idtask.task_id, input_texttask.input_text, output_textf结果: {task.input_text[:30]}, successTrue, 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_asciiFalse, indent2)) # 分批加载器 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_idftask-{task_id}, input_textline, ) 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_asciiFalse) \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_concurrent30, batch_size500, max_retries3, ) engine BatchInferenceEngine(config) # 生成测试任务 tasks [ InferenceTask( task_idftask-{i}, input_textf这是一条测试评论 #{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 重写了这套引擎核心逻辑用errgroupchan实现代码更简洁// 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 万任务占 4GBGo 只需 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 费用 $150GPU 利用率从 15% 提升到 85%。如果长期使用建议用 Go 重写——内存更低、部署更简单但 Python 版本适合快速验证和数据分析团队自服务。
