1. 从一次显存溢出事故说起FastAPI GPU推理并发控制到底在解决什么问题去年帮一个团队排查线上问题场景很典型他们用FastAPI包了一个视觉推理服务模型是YOLO系列的检测网络部署在一张24G显存的卡上。压测的时候单请求跑得好好的延迟稳定在80毫秒左右结果一上生产QPS刚过8服务就开始报CUDA out of memory进程直接挂掉重启之后又挂循环往复。运维那边看到的现象是服务隔几分钟就重启一次开发这边看到的是显存莫名其妙就满了两边扯了好几天。这个问题的本质其实不是模型太大也不是卡太小而是并发控制没做。FastAPI默认是异步框架请求进来之后如果直接丢给GPU去跑多个请求会同时抢占显存每个请求都要分配输入张量、中间激活、输出张量几个请求叠在一起24G瞬间就见底了。你要知道推理阶段的显存占用不是线性的它跟batch size、序列长度、模型结构都有关系一个请求可能占2G两个请求可能占5G三个请求可能就直接爆了。所以这篇东西想聊的就是FastAPI GPU推理场景下怎么做并发控制才能既不浪费算力又不把显存撑爆。适合谁看如果你正在用FastAPI做模型服务化或者准备把本地的推理脚本包装成HTTP接口对外提供又或者你已经上线了但被OOM搞得焦头烂额那这篇内容应该能帮你少走一些弯路。我会从整体设计思路讲起然后拆解几个核心的并发控制手段再给出一套可以直接抄的实操方案最后把我自己踩过的坑和排查技巧整理出来。需要提前说明的是并发控制这件事没有银弹不同的模型、不同的硬件、不同的业务场景最优解是不一样的。我下面讲的是通用思路加上一些经过验证的具体做法你拿去之后需要根据自己的实际情况调整参数。2. 整体设计思路为什么不能简单粗暴地加锁2.1 先搞清楚瓶颈到底在哪很多人一遇到OOM第一反应是加个锁吧用threading.Lock或者asyncio.Lock把推理部分包起来保证同一时刻只有一个请求在用GPU。这个做法确实能解决OOM但代价是吞吐量直接掉到1GPU利用率可能只有百分之十几等于花了大价钱买的卡在摸鱼。所以在动手之前得先分清楚你的瓶颈是什么。GPU推理的瓶颈通常有三个显存容量这是最硬的约束显存不够就是不够没有任何办法绕过只能减少同时驻留的请求数。计算单元利用率GPU的SM单元是有限的如果batch size太小计算单元跑不满算力就浪费了。数据传输带宽CPU到GPU的数据拷贝、GPU到CPU的结果回传这些也是有开销的尤其是输入数据大的时候。大部分OOM问题核心矛盾是显存容量和并发请求数之间的冲突。你要做的是找到一个平衡点在显存能承受的范围内尽可能让GPU忙起来。2.2 三种常见的并发控制策略对比我整理了一下实际项目中用过的几种策略各有各的适用场景策略核心做法优点缺点适用场景全局串行锁所有推理请求排队一次只跑一个实现简单绝对不会OOM吞吐量极低GPU利用率差请求量很小、对延迟不敏感的场景信号量限流用Semaphore限制同时执行的请求数实现简单能控制并发度并发数需要手动调不够灵活请求量中等、模型显存占用稳定的场景动态批处理把多个请求攒成一个batch一起推理GPU利用率高吞吐量大实现复杂延迟会增加请求量大、对吞吐量要求高的场景我个人的建议是先从信号量限流入手把服务跑稳然后再根据实际压测数据决定要不要上动态批处理。因为动态批处理的复杂度真的不低涉及到请求队列管理、超时控制、batch组装、结果拆分分发等一系列问题如果请求量没到那个级别投入产出比不划算。2.3 FastAPI的异步模型对并发控制的影响这里有个容易被忽略的点FastAPI是基于asyncio的它的请求处理是异步的。如果你在async def的路由函数里直接调用同步的GPU推理代码会阻塞整个事件循环导致其他请求连进来都进不来。这个行为跟Flask那种多线程模型完全不一样。所以正确的做法是把GPU推理放到线程池或者进程池里执行用run_in_executor或者asyncio.to_thread来调用这样事件循环不会被阻塞FastAPI才能正常处理其他请求的接入和排队。但这里又有一个坑如果你用线程池多个线程同时调用CUDAPyTorch的CUDA上下文是线程安全的但显存分配是全局的多个线程同时分配显存还是可能OOM。所以线程池的并发度必须和信号量的并发度配合起来用不能各管各的。3. 核心细节解析信号量、线程池与显存估算3.1 信号量限流的正确用法信号量Semaphore是控制并发数最直接的工具。在FastAPI里你可以用asyncio.Semaphore来限制同时执行的推理请求数。基本用法是这样的import asyncio from fastapi import FastAPI app FastAPI() semaphore asyncio.Semaphore(2) # 最多同时2个推理请求 app.post(/predict) async def predict(data: dict): async with semaphore: result await asyncio.to_thread(run_inference, data) return result这段代码看起来很简单但有几个细节需要注意。第一Semaphore的数值怎么定。这个不能拍脑袋得根据你的模型显存占用和显卡总显存来算。假设你的模型加载后占8G每个请求的输入输出和中间激活占3G显卡是24G那么理论上最多能同时跑(24 - 8) / 3 5个请求。但实际中要留一些余量给CUDA上下文、显存碎片、其他进程等所以建议取理论值的60%到70%也就是3个左右。第二Semaphore要放在全局还是每个请求创建。必须是全局的放在模块级别。如果你在路由函数里创建Semaphore那每个请求都会有自己的Semaphore完全起不到限流作用。第三acquire和release要配对。用async with是最安全的写法即使推理过程中抛异常Semaphore也会被正确释放。如果你手动调用acquire()和release()一定要用try/finally包起来否则一旦异常Semaphore就永远不会释放后续请求全部卡死。3.2 线程池的配置与陷阱前面说了GPU推理要放到线程池里跑避免阻塞事件循环。asyncio.to_thread底层用的是默认的ThreadPoolExecutor默认最大线程数是min(32, os.cpu_count() 4)。这个数值对于GPU推理场景来说通常太大了因为你的瓶颈在GPU不在CPU线程开多了反而会增加显存竞争。我一般会显式创建一个ThreadPoolExecutor把max_workers设置成和Semaphore一样的值from concurrent.futures import ThreadPoolExecutor import asyncio executor ThreadPoolExecutor(max_workers2) app.post(/predict) async def predict(data: dict): async with semaphore: loop asyncio.get_event_loop() result await loop.run_in_executor(executor, run_inference, data) return result这里有个坑ThreadPoolExecutor的max_workers和Semaphore的数值如果不一致会出现线程等待或者显存竞争。如果线程数大于Semaphore数多出来的线程会卡在Semaphore上浪费资源如果线程数小于Semaphore数Semaphore还没满但线程已经用完了请求会排队等线程实际并发度上不去。所以这两个值最好保持一致。还有一个更隐蔽的坑PyTorch的CUDA操作在某些情况下不是完全线程安全的。虽然PyTorch官方说CUDA上下文是线程安全的但在多线程环境下如果多个线程同时做显存分配和释放可能会遇到显存碎片问题。表现就是明明总显存够但就是分配不出来。解决办法是尽量让每个线程的显存分配模式一致或者用torch.cuda.empty_cache()定期清理但这个操作本身也有开销不能频繁调用。3.3 显存占用的估算方法要合理设置并发数你得知道每个请求大概占多少显存。最准确的方法当然是实测但在上线之前你可以用下面的方法做个粗略估算。对于推理阶段显存占用主要来自这几块模型参数这部分是固定的加载后就占着。可以用torch.cuda.memory_allocated()查看。输入张量跟输入数据的shape有关。比如一张640x640x3的图片float32的话就是640 * 640 * 3 * 4 4.9MB不大。中间激活这部分是大头跟模型结构有关。对于CNN来说中间激活的显存占用可能比模型参数还大。输出张量通常不大但如果输出是分割掩码或者高分辨率特征图也会占不少。CUDA上下文和显存碎片这部分通常占1G到2G不可忽略。实测的方法很简单加载模型跑一个请求记录torch.cuda.max_memory_allocated()然后再跑一个并发请求再记录两个值的差就是每个请求的增量显存占用。用这个增量去算最大并发数比较靠谱。注意torch.cuda.max_memory_allocated()返回的是峰值显存不是当前显存。如果你想看当前显存用torch.cuda.memory_allocated()。另外PyTorch的显存缓存机制会导致memory_allocated和nvidia-smi看到的数值不一致这是正常的以nvidia-smi为准。4. 实操过程从零搭建一个带并发控制的FastAPI推理服务4.1 项目目录结构先说一下目录结构这个是我自己常用的组织方式比较清晰fastapi-gpu-inference/ ├── app/ │ ├── __init__.py │ ├── main.py # FastAPI入口 │ ├── config.py # 配置参数 │ ├── model.py # 模型加载和推理逻辑 │ ├── schemas.py # 请求/响应数据结构 │ └── concurrency.py # 并发控制相关 ├── requirements.txt └── run.sh把并发控制单独抽一个模块出来是为了方便调整参数和复用。模型加载和推理逻辑也单独放这样换模型的时候不用动主逻辑。4.2 模型加载与预热模型加载这块有个关键点一定要在服务启动的时候加载不要在请求里加载。我见过有人在路由函数里torch.load()模型每个请求加载一次那显存不爆才怪。# app/model.py import torch import torchvision.models as models class InferenceModel: def __init__(self, devicecuda): self.device device self.model None def load(self): self.model models.resnet50(pretrainedTrue) self.model.to(self.device) self.model.eval() # 预热跑一次空推理让CUDA上下文和显存分配器初始化 dummy torch.randn(1, 3, 224, 224).to(self.device) with torch.no_grad(): self.model(dummy) torch.cuda.synchronize() torch.no_grad() def predict(self, input_tensor): input_tensor input_tensor.to(self.device) output self.model(input_tensor) return output.cpu()预热这一步很多人会省略但它其实很重要。第一次推理的时候CUDA会做很多初始化工作包括编译kernel、分配显存池等耗时可能是正常推理的几十倍。如果不预热第一个请求的延迟会非常高而且显存分配模式也不稳定。预热之后显存分配器已经建立好了缓存后续请求的显存分配会快很多也更稳定。4.3 并发控制模块的实现这是核心部分我把信号量、线程池、显存监控整合在一起# app/concurrency.py import asyncio from concurrent.futures import ThreadPoolExecutor import torch class GPUConcurrencyController: def __init__(self, max_concurrent2, max_workers2): self.semaphore asyncio.Semaphore(max_concurrent) self.executor ThreadPoolExecutor(max_workersmax_workers) self.max_concurrent max_concurrent async def run(self, func, *args, **kwargs): async with self.semaphore: loop asyncio.get_event_loop() return await loop.run_in_executor( self.executor, lambda: func(*args, **kwargs) ) def get_gpu_memory_info(self): if not torch.cuda.is_available(): return {} return { allocated: torch.cuda.memory_allocated() / 1024**3, reserved: torch.cuda.memory_reserved() / 1024**3, max_allocated: torch.cuda.max_memory_allocated() / 1024**3, }这个控制器把信号量和线程池绑在一起调用的时候只需要await controller.run(inference_func, data)它会自动处理限流和线程调度。4.4 主服务的组装# app/main.py from fastapi import FastAPI, HTTPException from app.model import InferenceModel from app.concurrency import GPUConcurrencyController import torch app FastAPI() model InferenceModel() controller GPUConcurrencyController(max_concurrent2, max_workers2) app.on_event(startup) async def startup(): model.load() app.post(/predict) async def predict(data: dict): try: input_tensor preprocess(data) output await controller.run(model.predict, input_tensor) return {result: postprocess(output)} except torch.cuda.OutOfMemoryError: raise HTTPException(status_code503, detailGPU memory exhausted, please retry later) except Exception as e: raise HTTPException(status_code500, detailstr(e)) app.get(/health) async def health(): return { status: ok, gpu_memory: controller.get_gpu_memory_info() }这里加了两个东西一个是OOM的异常捕获返回503让客户端重试另一个是健康检查接口可以实时看到显存占用情况方便排查问题。4.5 参数调优的实操过程参数调优不能靠猜得用数据说话。我的做法是写一个简单的压测脚本逐步增加并发数观察显存和延迟的变化# 压测脚本示例 import asyncio import aiohttp import time async def send_request(session, url, data): start time.time() async with session.post(url, jsondata) as resp: result await resp.json() return time.time() - start async def benchmark(concurrency): url http://localhost:8000/predict data {input: test} async with aiohttp.ClientSession() as session: tasks [send_request(session, url, data) for _ in range(concurrency)] latencies await asyncio.gather(*tasks) return latencies # 逐步增加并发 for c in [1, 2, 3, 4, 5]: latencies asyncio.run(benchmark(c)) print(fConcurrency {c}: avg{sum(latencies)/len(latencies):.3f}s, max{max(latencies):.3f}s)跑这个脚本的时候同时用nvidia-smi -l 1监控显存变化。你会看到随着并发数增加显存占用逐步上升延迟也会上升。当显存接近上限或者延迟急剧增加的时候就说明到了瓶颈当前的并发数就是上限。我实测下来对于ResNet50这种级别的模型24G卡上并发数设3到4比较合适再高就会开始出现OOM或者延迟抖动。但你的模型可能不一样一定要自己测。5. 常见问题与排查技巧实录5.1 OOM了但显存看起来还有空余这是最让人头疼的情况nvidia-smi显示还有好几个G的空闲显存但程序就是报OOM。原因通常是显存碎片。PyTorch的显存分配器会缓存已经分配的显存块如果请求的显存大小和缓存块不匹配即使总空闲显存够也分配不出来。解决办法有几个一是设置PYTORCH_CUDA_ALLOC_CONFexpandable_segments:True让分配器支持可扩展的显存段减少碎片二是定期调用torch.cuda.empty_cache()但这会释放缓存导致后续分配变慢不能频繁用三是尽量让每个请求的显存分配模式一致比如固定输入尺寸避免动态shape导致的碎片。5.2 服务跑一段时间后越来越慢这个问题的原因可能是显存泄漏。虽然PyTorch的Tensor在Python GC的时候会自动释放显存但如果你的代码里有循环引用或者把Tensor存到了全局变量里GC就回收不了。排查方法是定期打印torch.cuda.memory_allocated()看它是不是在持续增长。如果是就用gc.collect()手动触发垃圾回收看能不能降下来。如果降不下来那就是有地方持有Tensor的引用没释放。另外一个可能的原因是CUDA上下文切换开销。如果你的服务同时跑了多个模型或者频繁地在CPU和GPU之间拷贝数据上下文切换的开销会累积。解决办法是尽量把数据预处理放在GPU上做减少拷贝次数。5.3 并发数设了但实际不生效有时候你设了Semaphore为3但实际观察到的并发数只有1或者2。这种情况通常是线程池的max_workers设小了或者事件循环被阻塞了。检查方法很简单在推理函数里打印当前线程ID和时间戳看是不是真的有多个线程在同时执行。如果没有那就是线程池的问题。如果事件循环被阻塞了其他请求连Semaphore都进不来表现就是请求排队时间很长但GPU利用率不高。5.4 常见问题速查表现象可能原因排查方法解决方案CUDA OOM但显存有空余显存碎片查看memory_reserved和memory_allocated的差值设置expandable_segments固定输入尺寸服务越跑越慢显存泄漏监控memory_allocated趋势检查全局变量引用手动gc并发数不生效线程池配置错误打印线程ID调整max_workers与Semaphore一致第一个请求特别慢未预热观察首次推理耗时启动时跑一次空推理延迟抖动大显存竞争监控显存和GPU利用率降低并发数或改用动态批处理5.5 几个我踩过的坑第一个坑是在async函数里直接调用同步的CUDA操作。这个会导致事件循环阻塞表现是服务看起来在跑但所有请求都卡住。解决办法就是用run_in_executor或者to_thread包一层。第二个坑是Semaphore的release没放在finally里。有一次推理过程中抛了异常Semaphore没释放后续所有请求都卡在acquire上服务直接假死。后来改成async with就没这个问题了。第三个坑是模型加载了两次。因为用了app.on_event(startup)但在某些部署方式下这个事件会触发多次导致模型被加载多次显存直接翻倍。解决办法是在加载前检查模型是否已经加载或者用单例模式。第四个坑是没有设置CUDA_VISIBLE_DEVICES。如果机器上有多个GPU不设置这个环境变量的话PyTorch默认会用第0张卡可能跟其他服务冲突。建议在启动脚本里显式指定。提示如果你的服务部署在容器里记得给容器设置--gpus all或者--gpus device0否则容器里看不到GPU。另外容器里的显存限制和宿主机可能不一致需要额外注意。6. 进阶方向从限流到动态批处理信号量限流能解决OOM问题但它的吞吐量上限是固定的。如果你的请求量很大而且对吞吐量有要求那就需要考虑动态批处理了。动态批处理的核心思想是把短时间内到达的多个请求攒成一个batch一起送给GPU推理这样GPU的计算单元利用率会高很多。比如单请求推理需要80毫秒batch size为4的时候可能只需要120毫秒平均每个请求30毫秒吞吐量提升了将近3倍。实现动态批处理需要几个组件一个请求队列一个批处理调度器一个结果分发机制。请求进来之后先入队调度器根据当前队列长度和等待时间决定是否组batch推理完成后把结果按请求ID分发回去。这块复杂度不低我后面会单独写一篇来讲。另外一个方向是多模型共享GPU。如果你的服务需要同时跑多个模型可以用显存池化的方式让多个模型共享同一块显存按需分配。这个需要更底层的CUDA编程或者用NVIDIA的MPSMulti-Process Service门槛比较高一般场景用不上。我个人在实际操作中的体会是先把限流做稳再考虑优化吞吐。很多团队一上来就想搞动态批处理结果限流都没做好服务天天OOM反而耽误了上线时间。限流是保底批处理是锦上添花顺序不能反。
