ARTICLE DETAIL

资讯详情

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

FastAPI GPU推理并发控制:从显存溢出到稳定服务的实战指南

FastAPI GPU推理并发控制:从显存溢出到稳定服务的实战指南 1. 从一次线上事故说起为什么并发控制是GPU推理服务的生死线去年冬天我帮一个做图像识别的团队排查线上问题。他们的服务用FastAPI包了一个YOLO模型部署在一台单卡A10上平时跑得好好的结果某天市场部做活动流量瞬间翻了五倍服务直接挂了。日志里赫然写着CUDA out of memory显存被瞬间打满连带着把同卡上的其他服务也拖死了。事后复盘问题根本不在模型本身而在于他们压根没做并发控制——每个请求进来就无脑往GPU上塞显存不炸才怪。这个场景太典型了。FastAPI本身是异步框架天生适合做高并发接口但GPU是个独占型资源它的显存容量是固定的计算单元也是有限的。你不可能像CPU那样开几十个线程同时跑GPU的并发模型和CPU完全不是一回事。所以FastAPI GPU推理这个组合核心矛盾就在于Web层的异步高并发和GPU层的串行独占之间需要一个中间层来做调度和限流。这个中间层做得好不好直接决定了你的服务是稳定运行还是三天两头崩。这篇文章就是围绕这个矛盾展开的。我会从架构设计、并发控制策略、显存管理、实操代码、问题排查几个维度把“FastAPI GPU推理并发控制”这件事讲透。适合谁看如果你正在用FastAPI部署深度学习模型或者准备把本地的推理脚本包装成线上服务又或者你已经上线了但经常遇到显存溢出、请求超时、服务崩溃的问题那这篇内容就是为你写的。我会尽量用大白话把原理讲清楚同时给出可以直接抄的代码和配置让你看完就能动手改自己的项目。2. 整体架构设计FastAPI和GPU之间到底该怎么配合2.1 为什么不能直接在路由函数里跑推理很多人的第一版代码是这样的写一个/predict接口在函数里直接加载模型、处理输入、调用model()、返回结果。本地测试没问题一上线就出事儿。原因很简单FastAPI默认用Uvicorn跑Uvicorn是异步服务器它的事件循环是单线程的。你在路由函数里写同步的GPU推理代码整个事件循环就被阻塞了其他请求全部排队等着。这时候如果并发量上来请求队列越堆越长每个请求都占着内存不释放最后要么超时要么OOM。更糟糕的是如果你用了async def定义路由但在里面调用了同步的GPU推理函数FastAPI会把它丢到线程池里跑。线程池默认是40个线程意味着最多有40个请求同时往GPU上塞。一张24G显存的卡跑一个YOLOv8-large大概占4G40个并发就是160G显存直接爆炸。所以绝对不能把GPU推理裸奔在路由函数里必须加一层控制。2.2 三种主流方案对比信号量、队列、独立推理进程我试过三种方案各有优劣下面用表格对比一下。方案实现方式优点缺点适用场景信号量控制用asyncio.Semaphore限制同时推理的请求数实现简单改动小请求排队时占用内存无法控制显存粒度小模型、显存充足、并发量不大请求队列用asyncio.Queue做生产者消费者单独协程消费解耦彻底可做优先级调度需要处理超时和队列满的情况中大型模型、需要精细控制独立推理进程把模型放在单独进程通过IPC通信隔离性好崩溃不影响Web层通信开销大部署复杂多模型、多卡、生产级服务我个人的选择是请求队列 单推理协程的方案。理由是这样既能控制并发数队列消费者只有一个天然串行又能通过队列长度做背压还能在消费者里做批处理优化。信号量方案虽然简单但它只是限制了同时进入推理的请求数排队中的请求依然占着内存而且没法做批处理。独立进程方案太重了除非你有多个模型要跑否则没必要。2.3 核心设计一个队列、一个消费者、一个显存监控我的架构是这样的FastAPI启动时创建一个asyncio.Queue同时启动一个后台协程作为消费者。每个请求进来后把输入数据和一个Future对象打包丢进队列然后await future等待结果。消费者协程从队列里取任务调用GPU推理把结果设置到Future里。这样所有GPU调用都在同一个协程里串行执行显存占用是可控的。同时我会加一个显存监控协程定期用pynvml查显存使用率。如果超过阈值比如90%就暂停从队列取任务等显存降下来再继续。这个机制能有效防止显存泄漏导致的累积溢出。另外队列本身要设一个最大长度比如100满了之后新请求直接返回503避免无限堆积。提示队列消费者协程要用asyncio.create_task启动并且要处理好异常不能让消费者挂掉。消费者挂了整个服务就废了。3. 核心细节解析显存管理、批处理与超时控制3.1 显存到底被谁吃了模型权重、激活值、缓存要控制显存先得知道显存花在哪了。以PyTorch为例显存占用主要分三块模型权重、前向传播的激活值、CUDA上下文和缓存。模型权重是固定的加载后就占着比如ResNet50大概100MYOLOv8x大概260M。激活值跟batch size和输入尺寸有关batch越大、输入越大激活值越多。CUDA上下文大概占300-500M这是固定开销。很多人忽略的是PyTorch的缓存分配器。PyTorch为了加速显存分配会缓存已经释放的显存块不立即还给系统。所以你用nvidia-smi看显存占用很高但实际上模型可能只用了很少。这时候如果你用torch.cuda.empty_cache()可以释放缓存但频繁调用会影响性能。我的建议是在服务启动时预热一次让缓存稳定下来之后不要频繁清缓存。还有一个坑是多线程下的显存分配。如果你用了线程池每个线程都可能触发显存分配导致缓存碎片化。所以一定要确保GPU推理在单线程/单协程里执行。3.2 批处理提升吞吐的关键但别贪心批处理是提升GPU利用率最有效的手段。单张图片推理可能只要5ms但GPU利用率只有10%因为大部分时间在等数据传输。如果把32张图片拼成一个batch推理时间可能只要20ms吞吐量提升8倍。但batch不是越大越好因为显存占用是线性增长的。你需要找到一个平衡点。我的做法是在队列消费者里做动态批处理。消费者不是取一个任务就推理一次而是等一小段时间比如10ms把队列里攒的任务凑成一个batch。如果队列里任务多batch就大任务少batch就小。这样既能提升吞吐又不会因为固定大batch导致显存溢出。具体实现可以用asyncio.wait_for加超时或者用queue.qsize()判断。计算最大batch size的公式大概是max_batch (可用显存 - 模型权重 - CUDA上下文) / 单样本激活值。单样本激活值需要实测你可以用不同batch跑一下看显存增长曲线。一般来说留20%的显存余量比较安全。3.3 超时控制别让一个慢请求拖死整个队列队列方案有个隐患如果某个请求推理特别慢比如输入图片特别大它会占着消费者后面的请求全部等着。所以必须加超时。我的做法是给每个任务设一个超时时间比如5秒。消费者取到任务后用asyncio.wait_for包住推理调用超时就抛异常把Future设置成超时错误。同时在请求端也要设超时如果await future超过一定时间比如10秒直接返回504不等了。另外队列本身也要设最大等待时间。如果队列里任务的平均等待时间超过阈值说明系统过载了应该直接拒绝新请求。这个可以通过监控队列长度和消费者处理速度来动态调整。注意超时时间要根据你的模型和业务来定。实时性要求高的场景超时要短离线批处理场景可以长一点。但无论如何不能无限等。4. 实操过程从零搭建一个带并发控制的FastAPI推理服务4.1 项目目录结构与环境准备先看目录结构我习惯这样组织project/ ├── app/ │ ├── __init__.py │ ├── main.py # FastAPI入口 │ ├── config.py # 配置 │ ├── inference.py # 推理核心逻辑 │ ├── queue_manager.py # 队列管理 │ └── monitor.py # 显存监控 ├── models/ │ └── yolov8n.pt # 模型文件 ├── requirements.txt └── run.sh环境准备很简单装这几个包就行pip install fastapi uvicorn torch torchvision pynvml numpy opencv-python如果你用YOLO再加个ultralytics。注意PyTorch要装GPU版别装成CPU版了。装完之后用torch.cuda.is_available()验证一下。4.2 核心代码队列管理器与推理消费者先写队列管理器。这个类负责创建队列、启动消费者、提交任务。import asyncio import logging from typing import Any, Optional logger logging.getLogger(__name__) class InferenceQueue: def __init__(self, max_size: int 100, num_workers: int 1): self.queue asyncio.Queue(maxsizemax_size) self.num_workers num_workers self.workers [] self._running False async def start(self, inference_fn): self._running True for i in range(self.num_workers): task asyncio.create_task(self._worker(inference_fn, i)) self.workers.append(task) async def _worker(self, inference_fn, worker_id: int): while self._running: try: item await self.queue.get() if item is None: break data, future item try: result await inference_fn(data) if not future.done(): future.set_result(result) except Exception as e: logger.exception(fWorker {worker_id} inference error) if not future.done(): future.set_exception(e) finally: self.queue.task_done() except asyncio.CancelledError: break except Exception: logger.exception(fWorker {worker_id} unexpected error) async def submit(self, data: Any, timeout: float 10.0) - Any: loop asyncio.get_event_loop() future loop.create_future() try: self.queue.put_nowait((data, future)) except asyncio.QueueFull: raise RuntimeError(Queue is full, please retry later) try: return await asyncio.wait_for(future, timeouttimeout) except asyncio.TimeoutError: raise RuntimeError(Inference timeout) async def stop(self): self._running False for _ in self.workers: await self.queue.put(None) await asyncio.gather(*self.workers, return_exceptionsTrue)这个队列管理器有几个关键点maxsize控制队列长度满了直接抛异常submit方法带超时防止无限等待worker里用try/except包住推理保证一个任务出错不影响后续任务。4.3 推理函数加载模型、预处理、批处理推理函数是核心。我这里用YOLO举例但思路通用。import torch import numpy as np from PIL import Image class InferenceEngine: def __init__(self, model_path: str, device: str cuda): self.device device self.model self._load_model(model_path) self.model.eval() # 预热 self._warmup() def _load_model(self, path): from ultralytics import YOLO model YOLO(path) model.to(self.device) return model def _warmup(self): dummy np.random.randint(0, 255, (640, 640, 3), dtypenp.uint8) for _ in range(3): self.model(dummy, verboseFalse) torch.cuda.synchronize() async def infer(self, image: np.ndarray): # 这里用run_in_executor把同步推理放到线程池避免阻塞事件循环 loop asyncio.get_event_loop() result await loop.run_in_executor(None, self._infer_sync, image) return result def _infer_sync(self, image: np.ndarray): with torch.no_grad(): results self.model(image, verboseFalse) return results[0].boxes.data.cpu().numpy()注意这里用了run_in_executor因为YOLO的推理是同步的直接调用会阻塞事件循环。虽然我们只有一个worker但放到线程池里可以让事件循环继续处理其他请求比如健康检查。不过要小心如果线程池里同时有多个推理任务还是会并发。所以我的做法是队列worker只有一个它串行地调用infer而infer内部用线程池执行同步推理。这样实际上GPU调用还是串行的但事件循环不被阻塞。4.4 显存监控与动态限流显存监控用pynvml实现import pynvml import asyncio class GPUMonitor: def __init__(self, device_id: int 0, threshold: float 0.9): pynvml.nvmlInit() self.handle pynvml.nvmlDeviceGetHandleByIndex(device_id) self.threshold threshold self._paused False def get_memory_info(self): info pynvml.nvmlDeviceGetMemoryInfo(self.handle) return info.used / info.total async def monitor_loop(self, queue: InferenceQueue): while True: usage self.get_memory_info() if usage self.threshold and not self._paused: self._paused True logger.warning(fGPU memory usage {usage:.2%}, pausing queue) elif usage self.threshold - 0.1 and self._paused: self._paused False logger.info(fGPU memory usage {usage:.2%}, resuming queue) await asyncio.sleep(1)然后在worker里检查_paused状态如果暂停了就等一会儿再取任务。这个机制能有效防止显存缓慢增长导致的溢出。4.5 FastAPI入口整合所有组件最后是FastAPI的入口from fastapi import FastAPI, UploadFile, HTTPException from contextlib import asynccontextmanager import numpy as np import cv2 queue_manager None engine None monitor None asynccontextmanager async def lifespan(app: FastAPI): global queue_manager, engine, monitor engine InferenceEngine(models/yolov8n.pt) queue_manager InferenceQueue(max_size100, num_workers1) await queue_manager.start(engine.infer) monitor GPUMonitor(threshold0.9) asyncio.create_task(monitor.monitor_loop(queue_manager)) yield await queue_manager.stop() app FastAPI(lifespanlifespan) app.post(/predict) async def predict(file: UploadFile): contents await file.read() nparr np.frombuffer(contents, np.uint8) image cv2.imdecode(nparr, cv2.IMREAD_COLOR) if image is None: raise HTTPException(status_code400, detailInvalid image) try: result await queue_manager.submit(image, timeout10.0) except RuntimeError as e: raise HTTPException(status_code503, detailstr(e)) return {boxes: result.tolist()}这样整个服务就搭起来了。启动命令uvicorn app.main:app --host 0.0.0.0 --port 8000 --workers 1注意--workers一定要设成1因为多个worker会创建多个进程每个进程都会加载模型显存直接翻倍。如果你要多worker得用共享内存或者独立推理服务的方式。5. 常见问题与排查技巧实录5.1 显存溢出排查速查表现象可能原因排查方法解决方案启动就OOM模型太大或显存被占用nvidia-smi看显存换小模型或清显存运行一段时间后OOM显存泄漏监控显存增长曲线检查是否有未释放的tensor并发高时OOM并发数太多看队列长度和worker数限制并发加队列批处理时OOMbatch太大逐步增大batch测试减小batch或动态batch随机OOM缓存碎片用torch.cuda.memory_summary()定期empty_cache5.2 请求超时和队列堆积怎么破队列堆积通常是因为消费者处理速度跟不上请求速度。这时候你要看两个指标队列长度和单请求处理时间。如果队列长度持续增长说明系统过载了。解决方案有几个一是加机器多卡并行二是优化模型用更小的模型或量化三是做限流在入口层用令牌桶或漏桶算法限制QPS。超时问题往往是某个请求特别慢导致的。你可以在推理函数里加日志记录每个请求的处理时间。如果发现某些请求特别慢可能是输入尺寸太大可以在预处理阶段做resize。5.3 那些年我踩过的坑第一个坑是在异步函数里直接调用同步GPU代码。我一开始写了个async def predict里面直接model(image)结果事件循环被阻塞健康检查都超时了。后来改成run_in_executor才解决。第二个坑是忘了设队列最大长度。有次压测请求疯狂进来队列无限增长内存直接爆了。后来加了maxsize100满了就返回503服务反而稳定了。第三个坑是多worker导致显存翻倍。我一开始用--workers 4结果四个进程各加载一个模型显存直接满了。后来改成单worker用队列做并发控制效果更好。第四个坑是PyTorch缓存导致显存虚高。nvidia-smi显示显存用了20G但实际上模型只用了5G。后来用torch.cuda.memory_allocated()看实际用量才发现是缓存。这个不影响运行但会误导排查。提示排查显存问题时一定要区分reserved和allocated。allocated是实际用的reserved是缓存。用torch.cuda.memory_summary()可以看详细分布。5.4 性能优化的几个小技巧第一个技巧是用半精度。model.half()可以把显存占用减半推理速度也能提升。但要注意有些模型对精度敏感可能会掉点。可以先在测试集上验证一下。第二个技巧是用TensorRT或ONNX Runtime。这两个推理引擎比原生PyTorch快很多显存占用也更低。但转换过程有点麻烦适合对性能要求高的场景。第三个技巧是输入预处理放到CPU。图片解码、resize这些操作在CPU上做不要占用GPU。GPU只负责推理。第四个技巧是用固定尺寸输入。动态尺寸会导致每次推理都重新分配显存容易产生碎片。固定成640x640或者416x416显存分配更稳定。6. 进阶扩展从单卡到多卡从单模型到多模型6.1 多卡并行的两种思路如果你有多张卡有两种方案。一种是数据并行每个卡跑一个模型副本请求分发到不同卡上。这种适合模型小、并发高的场景。另一种是模型并行把大模型拆到多张卡上适合单卡放不下的大模型。数据并行实现简单用torch.nn.DataParallel或者手动分发就行。模型并行复杂得多一般用accelerate或deepspeed。我的建议是如果单卡能放下模型优先用数据并行。每张卡一个队列请求按轮询或最少连接分发。这样线性扩展实现也简单。6.2 多模型服务的路由策略如果你的服务要跑多个模型比如一个检测模型、一个分类模型可以用一个队列但不同worker处理不同模型或者每个模型一个队列。我倾向于每个模型一个队列因为不同模型的显存占用和推理时间不一样分开控制更灵活。路由层根据请求的模型名分发到对应队列。6.3 监控与告警别等崩了才发现生产环境一定要加监控。我一般监控这几个指标GPU利用率、显存使用率、队列长度、请求延迟、错误率。用Prometheus Grafana做可视化设置告警阈值。比如显存超过90%持续30秒就告警队列长度超过50就告警。这样能在问题恶化前介入。另外日志要打全。每个请求的ID、输入尺寸、推理时间、显存变化都记下来。出问题的时候这些日志就是排查的依据。6.4 容器化部署的注意事项如果用Docker部署要注意几点。一是基础镜像要带CUDA比如nvidia/cuda:12.1-runtime。二是启动时要加--gpus all。三是显存限制不能用--memory要用NVIDIA_VISIBLE_DEVICES或者CUDA_VISIBLE_DEVICES。四是容器内nvidia-smi可能不可用要装pynvml来监控。还有一个坑是容器内的共享内存。PyTorch的DataLoader默认用共享内存如果/dev/shm太小会报错。启动时加--shm-size8g。我个人在实际操作中的体会是并发控制这件事核心不是技术多复杂而是要对GPU的脾气有敬畏心。它不像CPU那样可以随便压榨你得顺着它的性子来。队列、限流、监控这三板斧下去基本能解决90%的显存溢出问题。剩下的10%靠的是对模型和业务的深入理解知道什么时候该批处理什么时候该降级什么时候该拒绝请求。这些经验都是在一次次线上事故里攒出来的。
返回列表