ARTICLE DETAIL

资讯详情

深耕网站建设与运营推广的一线实战洞察。

FastAPI GPU推理并发控制与显存溢出治理实战

FastAPI GPU推理并发控制与显存溢出治理实战 做GPU推理服务的朋友大概率都撞过这堵墙本地单张测试一切正常接口延迟漂漂亮亮一到压测或者上线流量进来GPU直接报CUDA out of memory进程崩掉甚至把整机带崩。上个月我给一个YOLOv11检测服务做并发改造就踩了一整轮这种坑。这篇文章把我在FastAPI里做GPU推理并发控制的完整思路、代码和排查手段完整写出来核心围绕“并发控制”和“显存溢出”这两个主线手把手给你一套可以直接抄作业的治理方案。先说结论大多数GPU推理服务被并发打崩不是GPU算力不够是显存被多个并发请求同时加剧的峰值占用打爆了而FastAPI默认的并发模型会把这个现象放大到极致。下面从根因开始拆。1. 为什么会显存溢出三个容易踩的根因1.1 先看懂GPU显存和CPU内存的本质差异CPU内存有虚拟内存、swap分区、操作系统的LRU回收机制兜底哪怕你一次性申请了几百G的虚拟地址空间系统也能在物理内存不够时把不活跃页面换出去。GPU显存完全不是这个玩法它没有swap没有虚拟内存所有分配都是显式的分配出去的空间如果不在代码里主动释放就会一直占着。驱动层面也不会像操作系统的页面回收那样好心“帮你清理”。更关键的是深度学习框架PyTorch是重灾区为了减少频繁的cudaMalloc调用有自己的缓存分配器CUDA Caching Allocator。我第一次排查时发现模型推理结束、张量都释放了nvidia-smi里显存占用却纹丝不动就是缓存分配器把释放的显存块缓存起来了。这些缓存块还会因为反复申请和释放变成碎片后续某个请求需要连续大块显存时即使总量够也分配不出来只能继续向驱动申请新显存最终触发OOM。这个特性决定了GPU显存是稀缺的显式管理资源必须由服务层主动做并发约束不能指望框架帮你擦屁股。1.2 FastAPI的异步模型和GPU推理的错位FastAPI很优秀但它的并发模型和GPU推理之间有一个巨大的错位。FastAPI支持两种端点声明方式async def端点跑在事件循环线程上适合IO密集场景普通def端点会被FastAPI丢到线程池里并发执行默认的线程池有40个线程换句话说FastAPI默认允许40个请求同时在另一个线程里跑你写的同步代码。问题来了。如果你直接写一个def predict()里面同步调用GPU推理压测只要一上来40个请求瞬间就有40个线程同时执行模型前向传播。每个推理任务都会创建一组中间张量这些中间张量在计算过程中是同时活着的显存峰值直接乘以并发数。单请求峰值只要2GB40个并发就是80GB什么显卡都扛不住。这就是“请求一多就显存溢出”最直接的触发路径。就算你把端点改成async defGPU推理本身还是CPU上的同步阻塞计算直接写在async函数里会把事件循环卡死后面的请求全部排队延迟瞬间拉满。正确的理解是FastAPI管不住GPU推理本身的资源占用它只是提供一个请求入口显存治理得靠自己在服务层实现。1.3 显存碎片的叠加效应除了缓存机制还有两个叠加因素加剧OOM。一是峰值叠加。即使你只允许2个请求并发两个推理任务如果恰好同时进入模型计算它们的中间张量峰值是相加的。更隐蔽的是推理框架在不同输入尺寸下中间显存需求差异很大小图只需要1GB大图可能就要3GB。如果两个大图请求恰好同时进来峰值可能到6GB超过显存总量。这类OOM用常规的“并发数*单请求峰值”都算不准必须留足余量。二是碎片化。PyTorch缓存分配器释放显存时是按“块”来的反复申请、释放大块和碎片后缓存池里大块连续内存会被拆得七零八落。一个请求想申请一个1GB的连续块会发现缓存里全是几百MB的碎片于是又cudaMalloc驱动给了一块新的老碎片还占着总占用不断上涨。我在后续章节会讲怎么用环境变量缓解这个现象但并发控制仍然是第一道防线碎片化只是让问题更快暴露。2. 整体方案设计把GPU当作稀缺资源来调度2.1 三层防护体系的设计思路治理显存溢出我建议不要只靠单一手段而是搭三层防护层与层之间是递进关系。第一层是入口层的并发闸门。在FastAPI端点内部或依赖项里用asyncio.Semaphore限制同时进入GPU推理的请求数量。这是最核心、最有效的一层它直接把“同时有几个推理任务在跑”这件事变成受控变量。第二层是显存水位保护。服务在每次推理开始前检查当前GPU显存占用率如果超过阈值比如85%直接返回503或者排队等待。这一层防的是并发闸门算漏的情况比如某个请求的中间张量异常大占用瞬间飙升把显存吃光。第三层是进程级和框架级兜底。对PyTorch设置合理的缓存分配策略、预留固定显存、监控torch.cuda.memory_allocated()甚至在多进程部署时控制worker数量避免好几个进程同时吃显存。这个设计和普通Web接口限流的本质区别在于普通Web接口限流通常保护的是后端资源数据库连接、下游服务而GPU推理服务要保护的是“不可弹性分配的本地显存”。限流保护的是自己的身板显存保护也是保护身板但是显存这个东西一旦爆了不是响应慢而是整个进程被杀影响面完全不是一个级别。2.2 并发控制方案选型对比我自己整理过一张对比表列了四种常见的并发治理方案方便你按场景选择。方案实现方式并发上限显存控制精度适用场景全串行asyncio.Lock包裹推理1最高单请求延迟极高、显存几乎被打满的模型固定并发asyncio.Semaphore(N)N高大多数固定输入尺寸的推理服务动态批处理队列攒batch合并推理动态高短输入、模型支持batch、追求吞吐多进程隔离每个worker独立模型N*worker中吞吐要求高、可独占显存的多卡机器全串行方案适合那种单次推理就要吃掉大半显存的模型比如一些大语言模型或者超大分辨率检测模型串行虽然牺牲吞吐但保证绝对不会因为并发而OOM。固定并发方案是性价比最高的日常选择一个信号量搞定代码量极少显存控制精度足够。动态批处理是进阶玩法适合单位时间请求量大、单请求延迟敏感的在线服务通过把多个请求的输入合并成一个batch推理即提升GPU利用率又摊薄中间张量的显存峰值。多进程隔离适合有多张卡或者单卡显存足够大的场景每个进程独享一份显存池互不干扰但要小心进程间显存总和超过总量。2.3 确定并发数上限的计算方法并发数定多少不靠拍脑袋我习惯按这个公式估算可用显存上限 总显存 - 系统预留 - 其他进程占用 单请求峰值显存 模型权重显存 最坏输入尺寸下的中间张量显存 推理缓存余量 安全并发数 floor(可用显存上限 / 单请求峰值显存 * 安全系数)安全系数我取0.7到0.8因为显存里还可能出现碎片化导致的额外占用。举一组实际数字一张24GB的3090系统预留和框架缓存按2GB算可用显存22GB某个检测模型加最坏输入尺寸下推理峰值需要4GB安全系数0.7那么并发数上限大概是22/(4*1.3)≈4。这里加1.3作为膨胀系数应对碎片和峰值波动算下来就是4个并发。实际压测时我们确实发现4个并发稳5个并发偶尔OOM公式和实测对得上。我给这个设计总结成一句话后面所有代码都围绕这句话开展GPU推理服务要做的不是“并发越多越好”而是“并发刚好够用”把GPU当一个有时间片的水龙头而不是一按就开到底的闸门。3. 代码实战从锁到信号量再到动态批处理3.1 先看一个最直观的错误示范很多人的第一版代码长这样from fastapi import FastAPI import torch app FastAPI() class YOLOModel: def __init__(self): self.model torch.load(yolov11.pt, map_locationcuda:0).cuda() self.model.eval() torch.inference_mode() def infer(self, image_tensor): return self.model(image_tensor) model YOLOModel() app.post(/predict) def predict(image_tensor): result model.infer(image_tensor) return result这段代码在单请求测试时没有任何问题但在压测环境里必炸。原因前面说过def predict跑在FastAPI默认的40线程线程池里40个请求同时调model.infer()每个推理中间张量同时活着显存峰值等于单请求峰值乘以实际并发数。第二个隐藏问题在这个错误示范里也有torch.load默认会加载到CPU再到GPU权重在模型初始化时一次性占满显存如果模型有1GB并发4个请求时模型本身就要占用1GB加上推理中间张量比单独看推理峰值更难预估。这种地方建议用torch.load(..., map_locationcuda:0)直接加载到GPU减少中转内存但这不是本次重点跳过。3.2 用 asyncio.Lock 保住下限串行推理最简单的治理方案是给推理套一把锁让GPU一次只跑一个推理任务import asyncio from fastapi import FastAPI from concurrent.futures import ThreadPoolExecutor import torch app FastAPI() infer_lock asyncio.Lock() # 关键把同步推理丢到专门的线程池 infer_executor ThreadPoolExecutor(max_workers1) class YOLOModel: def __init__(self): self.model torch.load(yolov11.pt, map_locationcuda:0).cuda() self.model.eval() torch.inference_mode() def infer(self, tensor): return self.model(tensor) model YOLOModel() app.post(/predict) async def predict(image_tensor): async with infer_lock: # 用线程池执行阻塞推理避免卡死事件循环 result await asyncio.get_running_loop().run_in_executor( infer_executor, model.infer, image_tensor ) return result这个版本的核心逻辑是asyncio.Lock保证无论同时进来多少请求同一时刻只有一个请求能进入推理代码块。信号量的经典问题——锁内执行阻塞操作——在这里被拆解为两步锁只负责控制并发真正的推理通过run_in_executor丢到独立线程池执行。这里有个细节如果不注意代码会有性能问题很多人在异步函数里写result model.infer(image_tensor)这个同步调用会在事件循环里卡住其他请求全部排队相当于把锁的范围扩张到了整个事件循环完全失去了异步的意义。所以锁内也要用线程池执行推理把锁的粒度控制在“不想同时跑多个推理”这个初衷上。全串行方案的缺点是单请求推理时间就是整个服务的吞吐上限。如果模型推理一次要2秒服务QPS最多0.5这在很多业务场景下不可接受。所以全串行适合单次推理显存占用接近显存总量的场景比如大模型推理不适合常规的检测、分类模型。3.3 用 asyncio.Semaphore 实现有限并发工业场景里最常用的还是信号量方案。它的思路是允许最多N个请求同时推理超过N的请求排队等待import asyncio from fastapi import FastAPI from concurrent.futures import ThreadPoolExecutor import torch app FastAPI() # 核心参数并发上限按前面公式估算 GPU_INFER_CONCURRENCY 4 gpu_sem asyncio.Semaphore(GPU_INFER_CONCURRENCY) # 线程池大小可以等于并发上限也可以略大 infer_executor ThreadPoolExecutor(max_workersGPU_INFER_CONCURRENCY 2) class YOLOModel: def __init__(self, model_path: str, device: str): self.device device self.model torch.load(model_path, map_locationdevice).cuda() self.model.eval() torch.inference_mode() def infer(self, tensor): return self.model(tensor) model YOLOModel(yolov11.pt, cuda:0) app.post(/predict) async def predict(image_tensor): # 信号量控制并发同时加超时避免雪崩 try: await asyncio.wait_for(gpu_sem.acquire(), timeout10) except asyncio.TimeoutError: return {error: 推理队列繁忙请稍后重试}, 503 try: result await asyncio.get_running_loop().run_in_executor( infer_executor, model.infer, image_tensor ) return result finally: gpu_sem.release()信号量构建了一个容量为4的“容器”第5个请求进来时会在acquire()处等待前面的请求完成推理后释放名额第5个才能进入。这样无论上游怎么打GPU里同一时刻最多只有4个推理任务显存峰值就被限制在安全范围内。这里有几个关键参数需要注意。wait_for设置10秒超时防止请求无限排队导致客户端挂起这是一种保护配合返回503让客户端重试比让请求一直挂着好得多。线程池max_workers为什么是并发上限2因为推理之外的图像预处理、张量搬运也可能在这个线程池里执行留一点余量防止线程饥饿但这个值不宜过大否则并发闸门就失效了。信号量本身的不足是它只限制“同时进入推理的请求数”不感知每个请求的显存差异。如果输入尺寸是可变的小图并发4个没问题大图并发4个就可能OOM。这时候就要叠加第二层显存水位保护。3.4 显存水位保护确定性的防OOM兜底我习惯写一个显存监控函数在推理前检查当前显存占用率超过阈值就拒绝服务import subprocess GPU_MEM_THRESHOLD 0.85 # 85%阈值 def get_gpu_memory_usage(device_id0): try: result subprocess.run( [nvidia-smi, --query-gpumemory.used,memory.total, --formatcsv,noheader,nounits], capture_outputTrue, textTrue, timeout3 ) lines result.stdout.strip().split(\n) used, total lines[device_id].split(,) return int(used.strip()) / int(total.strip()) except Exception: return 1.0 # 查询失败时保守处理 async def predict_with_memory_guard(image_tensor): # 显存超阈值优先排队而不是直接OOM while get_gpu_memory_usage() GPU_MEM_THRESHOLD: await asyncio.sleep(1) return await do_infer(image_tensor)简单说这个守护函数在并发信号量的基础上加了一道“硬闸门”就算信号量放行了如果看一眼当前显存占用率已经超过85%就继续等待直到有空间才真正执行推理。这个方案比较朴素每次推理前要调一次nvidia-smi有几十毫秒的IO开销在低延迟场景下会有负担。更专业的做法是用pynvml库直接读显存数据零外部进程开销import pynvml pynvml.nvmlInit() handle pynvml.nvmlDeviceGetHandleByIndex(0) def get_gpu_memory_usage_v2(): mem pynvml.nvmlDeviceGetMemoryInfo(handle) return mem.used / mem.totalpynvml是NVIDIA官方的NVML封装直接在Python进程内读显卡数据速度比调nvidia-smi快一个量级适合放在热路径里。热词里提到的“gpu查询命令”生产环境用pynvml就是性价比最高的查询方式。显存水位保护配合信号量是我觉得最稳的组合信号量管住了“并发推理任务数量”水位保护管住了“每个任务可能的显存波动”两张网叠在一起OOM概率趋近于零。3.5 进阶动态批处理方案如果业务场景要求更高的吞吐固定并发数方案会让GPU算力闲置——4个并发可能只有50%的利用率。这时候可以上动态批处理思路是把并发进来的请求攒到一个队列里攒够batch_size或者等待一个固定时间窗口然后把一批请求喂给模型一起推理。代码逻辑不复杂可以基于asyncio.Queue实现import asyncio import torch class BatchInferenceEngine: def __init__(self, model, max_batch_size8, max_wait_time0.01): self.model model self.max_batch_size max_batch_size self.max_wait_time max_wait_time self.queue asyncio.Queue() self._worker_task asyncio.create_task(self._batch_loop()) async def submit(self, image_tensor): future asyncio.Future() await self.queue.put((image_tensor, future)) return await future async def _batch_loop(self): while True: batch, futures [], [] # 先取第一个阻塞等待 first_item, first_future await self.queue.get() batch.append(first_item[0]) futures.append(first_future) # 用超时继续收集后续请求 try: while len(batch) self.max_batch_size: item, future await asyncio.wait_for(self.queue.get(), timeoutself.max_wait_time) batch.append(item) futures.append(future) except asyncio.TimeoutError: pass # 合并batch推理 batched_input torch.stack(batch) results self.model.infer(batched_input) for future, result in zip(futures, results): if not future.done(): future.set_result(result) engine BatchInferenceEngine(model) app.post(/predict_batch) async def predict_batch(image_tensor): return await engine.submit(image_tensor)动态批处理的优点是在请求稀疏时每个请求进来后最多等10毫秒就能推理不会像固定并发数那样留着一个大并发缺口在请求密集时模型一次处理8个请求吞吐远高于固定并发4个请求。缺点也很明显模型要支持batch推理大多数检测、分类模型都支持但要注意batch内图片尺寸不一致时需要padding而且batch_size越大单次推理的中间显存峰值越高需要把batch_size和时间窗口压测调优。从显存角度看动态批处理和固定并发数相比显存占用更平稳因为同一时刻只有一个batch推理在跑中间张量只属于这一个batch不会出现几个并发任务各自独立申请显存导致峰值叠加的情况。我实际测试过一个检测模型固定并发4时峰值显存12GB动态批处理batch_size8时峰值也只有10GB但同时吞吐反而提升了2倍。代价是会引入最多一二十毫秒的排队延迟对要求单请求极低延迟的场景不太友好。4. 压测与监控别凭感觉调参4.1 先建立显存和延迟监控我一向坚持一个原则**并发数调到多少不是拍脑袋是要用数据喂出来的。**压测前先搭监控至少记录三个指标GPU显存占用率、GPU利用率、推理接口的P99延迟。Python侧可以直接用Prometheus的prometheus_client库暴露指标或者更简单的方式是写文件日志每次推理完成后记录当时的显存使用情况import time import pynvml import json def log_infer_metric(device_id, label): mem pynvml.nvmlDeviceGetMemoryInfo(pynvml.nvmlDeviceGetHandleByIndex(device_id)) print(json.dumps({ time: time.time(), label: label, mem_used_mb: mem.used // 1024 // 1024, mem_total_mb: mem.total // 1024 // 1024, }))记录显存最好的时机是推理完成之后立刻记录这能反映整个推理过程的峰值占用框架缓存还在不会立刻释放。压测时每隔几秒采样一次就能看到显存占用随并发数变化的曲线。监控数据积累到一定程度你就能精确回答“这个模型在同一时刻跑几个推理比较安全”。4.2 压测工具与方法论压测工具我用过三种简单说下区别hey极简易用单文件无依赖适合快速压测验证并发控制是否生效。wrk性能好适合高并发场景但脚本语言是Lua上手成本略高。locustPython写压测脚本可以模拟复杂请求体适合要模拟真实业务负载的场景。我推荐先用hey快速验证命令很简单hey -n 2000 -c 50 -m POST \ -H Content-Type: application/json \ -d {image_url: test.jpg} \ http://127.0.0.1:8000/predict-c 50表示50个并发连接这正是模拟“请求一多”的关键参数。压测时观察三件事第一显存占用曲线是否平稳第二有没有CUDA OOM错误日志第三P99延迟有没有出现锯齿状抖动如果抖动剧烈说明信号量排队严重。4.3 实际调参过程从2到4再到8我记录一次真实的调参过程帮你看清数据如何指导决策。初始信号量并发数设为2压测50并发结果显存峰值只有6GBGPU利用率42%P99延迟稳定在180ms吞吐150 QPS。这组数据说明GPU算力还有大量闲置于是把并发数提高到4显存峰值升到11GB阈值是24GB总量的45%GPU利用率74%P99延迟220ms吞吐提升到280 QPS。继续提高到6显存峰值17GBGPU利用率82%但是P99延迟出现了明显波动从220ms到600ms明显是信号量排队导致的抖动。最终在4到5之间取了4因为4到5的吞吐增益不到8%但P99延迟却劣化了15%。这个例子说明并发数不是越大越好最优点在GPU利用率接近饱和但还没饱和的位置。数据不骗人所有“我感觉”“应该可以”的调参都先用压测数据验证再上线。5. 常见问题排查与避坑实录5.1 CUDA OOM到底是真OOM还是缓存占用压测中看到“CUDA out of memory”第一反应不要慌先分清两种情况。一种是真OOM也就是推理过程中确实没有足够显存创建中间张量另一种是缓存占用导致的假OOMPyTorch缓存分配器持有的显存块太多无法满足某个特定大小的分配请求。区分方法很简单看OOM日志里的free数值。如果free显示还有几百MB空闲但分配的就是不成功大概率是碎片化如果free已经接近0就是真OOM。真OOM要降并发或者减小输入尺寸假OOM可以通过设置环境变量缓解export PYTORCH_CUDA_ALLOC_CONFexpandable_segments:Trueexpandable_segments:True允许PyTorch向驱动申请可分段的显存段减少碎片化。实测开启后大量并发场景下的显存整体占用能降低20%到30%。另一个常用参数是max_split_size_mb它可以控制缓存分配器拆分大块的行为对碎片化也有缓解。我建议生产环境直接设置import os os.environ.setdefault(PYTORCH_CUDA_ALLOC_CONF, expandable_segments:True,max_split_size_mb:512)如果设置之后还是频繁假OOM那就不是环境变量能解决的得回到信号量并发数的调整上。记住环境变量只是缓解并发控制才是根本。5.2 不要在锁内做耗时的同步推理我见过一个很隐蔽的性能事故在async with sem:块内直接调用model.infer()没有经过线程池。表面看信号量正常GPU确实只跑了一个推理任务但事件循环被推理阻塞了其他所有请求包括非推理的接口全部卡住整个服务像死掉一样。排查后把推理丢到线程池问题立刻消失。这个坑的本质是把并发控制做好了却忽略了事件循环的可用性。推理是CPU密集型操作要把它挪出事件循环可以在run_in_executor里执行也可以直接把端点写成def并配合自定义线程池。重点是信号量管的是“推理任务并发数”线程池管的是“阻塞操作在哪执行”两者要分工明确。5.3 多进程部署时的显存均衡FastAPI服务通常会用uvicorn --workers 4做多进程部署如果每个进程都加载一次模型显存总量等于进程数 * 单个模型显存。24GB显存单模型1GB看似宽松开4个worker后光模型就4GB加上并发推理的中间张量很容易突破。我之前在一个项目里就把workers从4改成2只因为显存紧张。多进程部署要算总账先确定每个进程的并发上限再乘以worker数总体占用不能超过显卡显存。如果确实需要多进程优先考虑多卡方案也就是热词里提到的“gpu租用”“gpu微调”场景那种规模化思路把推理任务分布到多个进程、多张卡上每个进程独占一块卡的显存池天然隔离。5.4 流式推理和长尾任务的并发问题热词里有“流式推理管线”。流式场景下请求会长时间占用GPU比如语音模型一次推理要5秒且每个请求需要连续占用显存空间。这种情况下信号量的单位“并发数”含义就变了它限制的不是瞬时并发而是同时运行的长任务数量。如果长任务和短任务混在同一个服务里短任务会被长任务阻塞得非常难受。我的建议是长任务和短任务拆分到两个接口各用一套信号量或者在长任务接口上单独设置更低并发上限。比如短任务并发4长任务并发1避免长任务阻塞短任务的关键路径。这本质上是把一个服务的GPU资源分泳道管理各泳道互不干扰。5.5 Uvicorn日志丢失与推理日志Uvicorn在高并发下会出现日志丢失现象这在热词搜索里也有迹可循。如果你在推理接口里做了大量print日志Uvicorn的日志处理器是同步写的会拖慢事件循环高并发下日志挤压后还会丢。我建议推理日志走独立的异步日志队列或者直接用logging模块配QueueHandler避免日志成为并发瓶颈。更简洁的做法是只在信号量获取、释放和推理完成这三个关键点打日志不要每个请求都打详细参数。写在最后的经验之谈这套并发控制方案上线后我在生产环境连续观察了一个多星期。最直接的变化就是压测再也没看到CUDA OOMP99延迟波动从几百毫秒收敛到稳定范围内。但我最想说的是另外一个体会GPU显存治理不是一锤子买卖它是服务改造里必须持续迭代的环节。模型换一版、输入图像尺寸上限提高、并发策略调整任何一项变化都可能重新引爆显存问题。所以这套监控、压测、调参的流程要留下来每个月定期跑一轮服务才能长期稳定在线。下次如果还有人跟你说GPU推理服务并发高就是性能好你可以在心里默默算一下它的显存线性叠加曲线——希望这篇文章能帮你和你的服务躲开同一个坑。
返回列表