Python 多进程协程混合架构GPU 推理与异步 I/O 的协作调度方案一、深度引言与场景痛点大家好我是赵咕咕。我之前遇到一个典型的资源错配问题——我们的 Embedding 服务需要同时做两件事用 GPU 跑模型推理把文本转成向量然后把向量写入远端的 Qdrant 向量数据库。这个服务部署在一台 8 核 32G 的机器上挂了一张 A10 GPU。理论上性能应该不错但实际上——GPU 利用率只有 40%而我们的 P99 延迟却在 800ms 以上。为什么 GPU 这么闲还这么慢因为 GPU 推理 10ms 就完了但这之后网络 I/O 要等 200ms。这 200ms 里 GPU 完全空转CPU 也只能干瞪眼——这就是计算和 I/O 的节奏不匹配问题。今天就来聊聊在 Python 里怎么把多进程、协程、GPU 推理三者优雅地协调起来让 GPU 一直在跑、网络一直在传、CPU 一直在调度谁都别闲着。二、底层机制与原理深度剖析2.1 问题拆解先分析一下典型的 Embedding 服务的数据流flowchart TD subgraph HTTP Layer A[HTTP 请求到达] end subgraph CPU Layer A -- B[协程接收请求] B -- C[文本预处理] B2[协程处理响应] -- H[HTTP 响应] end subgraph GPU Layer C -- D[GPU 推理/Embedding] D -- E[向量结果] end subgraph I/O Layer E -- F[写入向量数据库] F -- G[写入缓存/日志] end G -- B2关键矛盾GPU 推理步骤 D是计算密集型受 GIL 限制不能在协程里直接跑。但步骤 C、F、G 都是 I/O 密集型用协程效率最高。这就是为什么需要多进程协程的混搭。2.2 架构设计生产者-消费者管道flowchart LR subgraph 主进程 (asyncio) S[FastAPI Server] -- Q1[asyncio.Queue / 请求队列] end subgraph GPU Worker 进程 1 Q1 -.-|进程间通信| W1_IN[共享内存队列] W1_IN -- W1[GPU 推理] W1 -- W1_OUT[结果队列] end subgraph GPU Worker 进程 2 Q1 -.-|进程间通信| W2_IN[共享内存队列] W2_IN -- W2[GPU 推理] W2 -- W2_OUT[结果队列] end subgraph 主进程 (asyncio) W1_OUT -.-|结果回传| Q2[结果处理] Q2 -- I1[写入向量库] I1 -- R[返回响应] end style S fill:#e1f5fe style W1 fill:#fff3e0 style W2 fill:#fff3e0 style I1 fill:#e8f5e9核心设计思路主进程跑 asyncio负责 HTTP 接收、请求分发、结果收集、I/O 操作GPU Worker 进程每个进程独占一段 GPU 显存只做推理不碰 I/O进程间通信用torch.multiprocessing的共享内存队列避免序列化开销协程池主进程中用asyncio.Semaphore控制并发度防止 GPU Worker 被淹没2.3 GPU 推理的批处理优化单条推理 GPU 利用率低的核心原因是没有批处理。一条文本的 embedding 计算只用到了 GPU 的一小部分计算单元剩下的都在摸鱼。解决方案是动态攒批主进程收到请求后不立即发给 GPU而是攒够一批或达到等待上限再一次性发送。sequenceDiagram participant M as 主进程 participant B as 攒批缓冲区 participant G as GPU Worker M-B: 请求1到达 M-B: 请求2到达 Note over B: 等待攒批窗口 5ms M-B: 请求3到达 B-G: 发送批次 [req1, req2, req3] G-G: GPU 批量推理 G-B: 返回结果 [vec1, vec2, vec3] B-M: 分发各请求结果三、生产级代码实现import asyncio import multiprocessing as mp from dataclasses import dataclass from typing import Any import numpy as np import torch from torch.multiprocessing import Queue as TorchQueue # ── 1. GPU Worker 进程 ──────────────────────────── dataclass class BatchRequest: request_id: str texts: list[str] result_queue: TorchQueue | None None # 回传通道 class GPUWorker(mp.Process): 独占一个 GPU 的推理 Worker 进程 def __init__( self, gpu_id: int, task_queue: TorchQueue, batch_size: int 32, max_wait_ms: float 10.0, ): super().__init__(daemonTrue) self._gpu_id gpu_id self._task_queue task_queue self._batch_size batch_size self._max_wait_ms max_wait_ms def run(self): Worker 进程主循环 torch.cuda.set_device(self._gpu_id) # 在这里加载模型每个 Worker 一份模型副本 from transformers import AutoTokenizer, AutoModel model AutoModel.from_pretrained( BAAI/bge-large-zh-v1.5 ).cuda() model.eval() tokenizer AutoTokenizer.from_pretrained( BAAI/bge-large-zh-v1.5 ) while True: batch self._collect_batch() if not batch: continue try: # 收集所有文本 all_texts [] request_ids [] result_queues [] for req in batch: all_texts.extend(req.texts) request_ids.append(req.request_id) result_queues.append(req.result_queue) # GPU 批量推理 with torch.no_grad(): inputs tokenizer( all_texts, paddingTrue, truncationTrue, max_length512, return_tensorspt, ).to(fcuda:{self._gpu_id}) outputs model(**inputs) embeddings outputs.last_hidden_state[:, 0, :].cpu().numpy() # 拆分结果并回传 idx 0 for rid, rq in zip(request_ids, result_queues): num_texts len(batch[request_ids.index(rid)].texts) vecs embeddings[idx : idx num_texts].tolist() idx num_texts rq.put({request_id: rid, embeddings: vecs, error: None}) except Exception as exc: for rid, rq in zip(request_ids, result_queues): rq.put({request_id: rid, error: str(exc)}) def _collect_batch(self) - list[BatchRequest]: 攒批等待 batch_size 个请求或达到 max_wait_ms batch [] deadline asyncio.get_event_loop().time() if False else 0 import time deadline time.monotonic() self._max_wait_ms / 1000 while len(batch) self._batch_size: remaining deadline - time.monotonic() if remaining 0: break try: req self._task_queue.get(timeoutremaining) batch.append(req) except Exception: break # 超时返回已有的批次 return batch # ── 2. 主进程asyncio 调度器 ───────────────────── class EmbeddingService: 多进程 GPU asyncio I/O 混合服务 def __init__( self, num_workers: int 2, batch_size: int 32, max_concurrent: int 100, ): self._num_workers num_workers self._batch_size batch_size # 进程间通信队列 self._task_queue TorchQueue(maxsize1000) self._workers: list[GPUWorker] [] # 协程并发控制 self._semaphore asyncio.Semaphore(max_concurrent) # 分派策略Round-robin self._next_worker 0 async def start(self): 启动所有 GPU Worker 进程 for gpu_id in range(self._num_workers): worker GPUWorker( gpu_idgpu_id, task_queueself._task_queue, batch_sizeself._batch_size, ) worker.start() self._workers.append(worker) async def embed(self, texts: list[str]) - list[list[float]]: 异步 embedding 入口 async with self._semaphore: result_queue TorchQueue(maxsize1) request_id str(id(texts)) # 简化版 # 提交任务到 GPU Worker self._task_queue.put(BatchRequest( request_idrequest_id, textstexts, result_queueresult_queue, )) # 等待结果用 asyncio.to_thread 避免阻塞事件循环 result await asyncio.to_thread(result_queue.get) if result.get(error): raise RuntimeError(fGPU 推理失败: {result[error]}) return result[embeddings] async def process_batch(self, requests: list[dict]) - list[dict]: 并行处理多个请求 tasks [] for req in requests: task self.embed(req[texts]) tasks.append(task) results await asyncio.gather(*tasks, return_exceptionsTrue) response [] for req, res in zip(requests, results): if isinstance(res, Exception): response.append({id: req.get(id), error: str(res)}) else: response.append({id: req.get(id), embeddings: res}) return response def shutdown(self): for worker in self._workers: worker.terminate() worker.join(timeout5) # ── 3. 共享内存优化 ─────────────────────────────── class SharedMemoryBuffer: 对大文本传输使用共享内存避免序列化开销 def __init__(self, max_size_mb: int 1024): import multiprocessing.shared_memory as shm self._max_size max_size_mb * 1024 * 1024 self._buffer shm.SharedMemory( createTrue, sizeself._max_size ) def write(self, data: bytes) - int: 写入数据返回偏移量 if len(data) self._max_size: raise ValueError(f数据大小 {len(data)} 超过缓冲区限制) self._buffer.buf[:len(data)] data return len(data) def read(self, offset: int, size: int) - bytes: return bytes(self._buffer.buf[offset:offset size]) # ── 4. 完整 FastAPI 集成 ───────────────────────── # from fastapi import FastAPI # # app FastAPI() # service EmbeddingService(num_workers2, batch_size32) # # app.on_event(startup) # async def startup(): # await service.start() # # app.post(/embed) # async def embed_endpoint(texts: list[str]): # embeddings await service.embed(texts) # # 这里可以异步写入向量数据库 # asyncio.create_task(write_to_qdrant(texts, embeddings)) # return {embeddings: embeddings}四、边界分析与架构权衡GPU Worker 数量的选择不是越多越好。每个 GPU Worker 都需要一块 GPU 显存来加载模型。一张 A1024G跑 bge-large 大约能支持 2-3 个 Worker。Worker 太多反而会因为显存争抢导致频繁的显存换页。攒批窗口的权衡攒批窗口设得太长比如 50msP50 延迟会明显升高因为请求在等待攒批。设得太短比如 2ms批处理效果差GPU 利用率上不去。经验值是 5-10ms能覆盖大部分请求间隔又不至于让用户感知到延迟。进程间通信的开销TorchQueue底层用的是 pickle 序列化对于大数组会有不小的开销。如果你的输入文本是长文本 1000 字建议用共享内存传文本数据队列里只传指针偏移量 大小。CPU 和 GPU 的亲和性不同 GPU 在不同 NUMA 节点上Worker 进程应该绑定到对应 GPU 最近的 CPU 核心。用os.sched_setaffinity和taskset做亲和性绑定可以降低 5-10% 的推理延迟。错误隔离GPU Worker 进程崩溃不应该影响主进程的服务可用性。需要在主进程这边加心跳检测和自动重启。如果有 Worker 连续失败 N 次应该降级——剩余的请求由其他 Worker 处理同时触发告警。五、总结Python 多进程 asyncio 协程的混合架构本质上是一种计算和 I/O 的解耦架构。GPU 专心做推理协程专心做网络 I/O中间用队列和共享内存做管道连接。核心要点主进程跑 asyncio 处理 I/OGPU Worker 进程跑 torch 做推理动态攒批是提高 GPU 利用率的关键手段用torch.multiprocessing.Queue做进程间通信巨大数据走共享内存Worker 要做错误隔离和自动重启攒批窗口5-10ms和 Worker 数量需要根据实际负载调优这种架构不仅适用于 Embedding 服务所有GPU 计算 网络 I/O的场景——LLM 推理、图像处理、语音识别——都可以用这个套路。把 GPU 喂饱把网络跑满你的机器才算物尽其用。下一篇预告RAG 多路召回融合向量、关键词、图召回怎么统一排序。