ARTICLE DETAIL

资讯详情

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

FastAPI+GPU推理服务并发控制:从显存溢出到动态批处理实战

FastAPI+GPU推理服务并发控制:从显存溢出到动态批处理实战 1. 从一次显存溢出事故说起为什么并发控制是GPU推理服务的生死线去年帮一个团队排查线上问题他们的场景很典型用 FastAPI 包了一个视觉推理接口模型是 YOLO 系列的检测网络单卡 24G 显存平时 QPS 不高的时候跑得好好的。结果做了一次活动推广请求量从每秒几个涨到几十个服务在 40 秒内直接崩了日志里清一色是CUDA out of memory紧接着进程被系统 OOM Killer 干掉连带着健康检查都挂了。这个事故的核心不是模型太大也不是显卡太差而是并发控制完全缺位。FastAPI 默认是异步框架一个请求进来就交给事件循环处理如果推理代码是同步阻塞的绝大多数 PyTorch 推理都是FastAPI 会把它丢到线程池里跑。线程池默认能开 40 个线程意味着理论上同时可以有 40 个推理任务抢同一块 GPU 的显存。每个任务加载输入张量、中间激活、输出结果哪怕单个只占 500MB40 个叠起来就是 20G显存瞬间见底。所以这篇内容想聊的不是FastAPI 怎么调用 GPU这种入门话题而是当请求量上来之后怎么让 GPU 推理服务稳住不崩。涉及的核心点包括并发模型怎么选、信号量怎么放、批处理怎么做、显存怎么监控、超时怎么兜底。适合已经能把模型跑起来、但一上量就出问题的后端和算法同学也适合正在做推理服务架构选型的人参考。我下面讲的方案都是基于单机单卡或多卡的常见部署形态用 FastAPI PyTorch 这套最普遍的组合来举例思路换成 vLLM、TensorRT 或者其他推理引擎同样适用因为并发控制的本质是资源调度跟具体框架关系不大。2. 并发失控的根因拆解FastAPI 的异步模型和 GPU 的独占性天然冲突2.1 FastAPI 的线程池机制到底怎么运作很多人对 FastAPI 有个误解觉得它异步就等于高并发不占资源。实际上 FastAPI 基于 Starlette底层是 asyncio 事件循环。当你定义一个async def的路由时它跑在事件循环里当你定义的是普通def路由时Starlette 会把它扔进anyio的线程池执行默认容量是 40 个线程这个数字来自anyio的默认配置可以通过run_in_threadpool的 limiter 调整。问题就出在这里。PyTorch 的推理调用是同步阻塞的你写model(input)这一行CPU 会一直等到 GPU 算完才返回。如果你把它写在async def里整个事件循环会被卡住其他请求全部排队吞吐量惨不忍睹如果你写在普通def里40 个线程可以同时发起推理显存直接爆炸。这就是那个团队踩的坑——他们用的是普通def以为线程池会帮忙限流其实线程池限的是线程数不是显存。2.2 GPU 显存为什么经不起并发叠加要理解显存溢出得先搞清楚一次推理到底占了多少显存。以一张 24G 的卡跑一个中等规模的检测模型为例显存占用大致分四块占用类型典型大小是否随并发增长模型权重200MB ~ 2GB否加载一次常驻CUDA 上下文与 cuDNN 句柄300MB ~ 800MB否进程级输入张量与中间激活100MB ~ 1GB/请求是线性增长输出结果与缓存几十 MB/请求是线性增长关键在第三、四行。模型权重是共享的但每个请求的输入张量和前向传播的中间激活是独立的。批大小batch size为 1 时可能只占 300MB但同时跑 20 个请求就相当于隐式地把 batch 拉到了 20激活内存按比例膨胀。更麻烦的是 PyTorch 的缓存分配器不会立刻把释放的显存还给系统torch.cuda.memory_allocated()降下来了torch.cuda.memory_reserved()还占着碎片一多新的请求就分配不到连续显存。2.3 三种典型崩溃场景对照我把实际遇到过的崩溃场景整理成一张表方便你对号入座场景现象根因突发流量打崩请求一多立刻 OOM恢复后正常无并发上限线程池全开慢请求拖垮个别请求推理慢显存迟迟不释放无超时机制长任务堆积内存碎片累积跑几小时后逐渐 OOM重启就好缓存分配器碎片未定期清理这三种场景的解法不一样但底层都指向同一件事你必须显式地控制同一时刻有多少个推理任务在跑。这个数字不是拍脑袋定的而是根据显存预算反推出来的。3. 并发控制方案选型信号量、队列还是批处理3.1 方案一asyncio.Semaphore 做请求级限流最直接的做法是在推理入口加一个信号量。信号量的值就是允许同时进行的推理任务数比如设成 4那么第 5 个请求就得等前面某个完成才能进。import asyncio from fastapi import FastAPI import torch app FastAPI() # 根据显存预算设定并发上限 INFER_SEMAPHORE asyncio.Semaphore(4) app.post(/infer) async def infer(payload: dict): async with INFER_SEMAPHORE: # 把同步推理丢到线程池避免阻塞事件循环 result await asyncio.to_thread(run_model, payload) return result def run_model(payload): with torch.inference_mode(): tensor preprocess(payload).cuda() out model(tensor) return postprocess(out)这里有几个细节值得说。第一信号量必须定义在模块级别不能写在路由函数里否则每个请求都会新建一个信号量等于没限。第二asyncio.to_thread是 Python 3.9 的写法低版本用loop.run_in_executor。第三torch.inference_mode()比torch.no_grad()更省显存因为它连版本计数都不记录实测能省 5% 到 10% 的激活内存。信号量方案的优点是简单、可控缺点是并发数固定吞吐量上不去。如果单个请求推理耗时 100ms并发 4 的话理论 QPS 只有 40而且请求之间是串行等待的GPU 利用率可能只有 30% 到 50%。3.2 方案二请求队列 单 worker 串行推理另一种极端是彻底串行所有请求进队列一个 worker 挨个处理。这样显存占用最稳定永远不会 OOM但延迟会随队列长度线性增长。import asyncio from collections import deque queue deque() queue_event asyncio.Event() async def worker(): while True: if not queue: queue_event.clear() await queue_event.wait() continue future, payload queue.popleft() try: result await asyncio.to_thread(run_model, payload) future.set_result(result) except Exception as e: future.set_exception(e) app.on_event(startup) async def start_worker(): asyncio.create_task(worker()) app.post(/infer) async def infer(payload: dict): loop asyncio.get_event_loop() future loop.create_future() queue.append((future, payload)) queue_event.set() return await future这个方案适合推理耗时短、请求量不大、但对稳定性要求极高的场景。它的致命伤是 GPU 利用率低因为推理和预处理、后处理是串行的GPU 在 CPU 干活的时候是空闲的。3.3 方案三动态批处理把并发转成批大小真正能兼顾吞吐和显存的做法是动态批处理dynamic batching。思路是不限制并发请求数但把短时间内到达的请求攒成一个 batch一次性送进 GPU。这样显存占用由 batch size 决定而不是由并发数决定。import asyncio import torch BATCH_WINDOW 0.02 # 20ms 攒批窗口 MAX_BATCH 8 pending [] batch_lock asyncio.Lock() async def batch_worker(): while True: await asyncio.sleep(BATCH_WINDOW) async with batch_lock: if not pending: continue batch pending[:MAX_BATCH] del pending[:MAX_BATCH] futures [f for f, _ in batch] payloads [p for _, p in batch] try: results await asyncio.to_thread(run_batch, payloads) for f, r in zip(futures, results): f.set_result(r) except Exception as e: for f in futures: f.set_exception(e) app.post(/infer) async def infer(payload: dict): loop asyncio.get_event_loop() future loop.create_future() async with batch_lock: pending.append((future, payload)) return await future动态批处理的关键参数是攒批窗口和最大批大小。窗口太短攒不到请求太长延迟高最大批大小要根据显存反推。以 24G 卡跑检测模型为例batch1 占 1.2Gbatch8 占 3.5Gbatch16 占 6G那么设 MAX_BATCH8 是比较稳的留足余量给碎片和其他开销。3.4 三种方案怎么选方案吞吐延迟显存稳定性适用场景信号量限流中低好请求量中等延迟敏感串行队列低高极好请求量小稳定性优先动态批处理高中好请求量大吞吐优先我的建议是先用信号量把服务稳住再逐步上动态批处理。信号量是保命措施动态批处理是性能优化两者可以叠加使用——信号量控制总并发批处理在并发内部做聚合。4. 显存预算怎么算从模型大小反推并发上限4.1 显存占用的四块账要设定合理的并发数得先把显存账算清楚。我用一个实际项目的数据举例模型是 ResNet50 分类网络输入 224x224FP16 推理显卡是 24G 的消费级卡项目实测占用说明模型权重FP1650MB25M 参数 x 2 字节CUDA 上下文约 400MB进程启动后固定cuDNN 工作区约 200MB首次推理后固定单请求激活约 180MBbatch1 前向传播单请求输入输出约 20MB张量 结果固定开销约 650MB单请求约 200MB。24G 卡留 2G 给系统和其他进程可用 22G那么理论上并发上限是 (22000 - 650) / 200 ≈ 106。但这是理论值实际要打对折甚至更多因为显存碎片会让可用连续块变小PyTorch 缓存分配器会预留比实际需要更多的显存突发流量下请求不是均匀到达的峰值可能翻倍所以实际设并发 20 到 30 比较稳妥。这个数字不是一次定死的要压测后调整。4.2 用代码实测显存占用别靠猜直接写个脚本测import torch def measure_memory(model, input_shape, batch_size, dtypetorch.float16): torch.cuda.empty_cache() torch.cuda.reset_peak_memory_stats() base torch.cuda.memory_allocated() dummy torch.randn(batch_size, *input_shape, dtypedtype).cuda() with torch.inference_mode(): _ model(dummy) peak torch.cuda.max_memory_allocated() print(fbatch{batch_size}, base{base/1e6:.1f}MB, fpeak{peak/1e6:.1f}MB, delta{(peak-base)/1e6:.1f}MB) del dummy torch.cuda.empty_cache() for bs in [1, 2, 4, 8, 16]: measure_memory(model, (3, 224, 224), bs)跑一遍就能看到显存随 batch 的增长曲线。如果增长是线性的说明激活内存是主导如果增长很快超过线性说明有中间变量没释放得检查代码里有没有在循环里累积张量。4.3 留多少余量才安全我的经验是至少留 20% 的显存余量。原因有三个一是 PyTorch 的缓存分配器在碎片化时可能申请不到连续块即使总空闲显存够二是推理过程中会有临时的显存峰值比如某些算子的 workspace三是如果服务还跑着其他东西比如日志、监控会分走一部分。具体到配置上可以用环境变量控制# 限制 PyTorch 可用的显存比例防止吃满 export PYTORCH_CUDA_ALLOC_CONFmax_split_size_mb:128max_split_size_mb这个参数控制缓存分配器切分内存块的最大粒度设小一点能减少碎片但会增加分配次数。128MB 是个比较平衡的值实测能明显降低长时间运行后的 OOM 概率。5. 完整实操搭一个带并发控制的 FastAPI 推理服务5.1 项目目录结构先把目录搭清楚避免所有代码堆在一个文件里infer-service/ ├── app/ │ ├── __init__.py │ ├── main.py # FastAPI 入口 │ ├── config.py # 配置项 │ ├── model.py # 模型加载与推理 │ ├── scheduler.py # 并发调度 │ └── schemas.py # 请求响应模型 ├── tests/ │ └── test_infer.py ├── requirements.txt └── run.sh这个结构的好处是模型加载、调度逻辑、路由分离改并发策略不用动业务代码。5.2 模型加载全局单例启动时预热模型必须在服务启动时加载一次不能每个请求都加载。用 FastAPI 的 lifespan 机制# app/model.py import torch import torch.nn as nn class InferModel: _instance None def __init__(self): self.device torch.device(cuda:0) self.model self._build_model() self.model.eval() self.model.to(self.device) self._warmup() classmethod def get(cls): if cls._instance is None: cls._instance cls() return cls._instance def _build_model(self): # 换成你的模型 from torchvision.models import resnet50 return resnet50(weightsNone) def _warmup(self): # 预热触发 cuDNN 算法选择避免首个请求慢 dummy torch.randn(1, 3, 224, 224).to(self.device) with torch.inference_mode(): for _ in range(3): _ self.model(dummy) torch.cuda.synchronize()预热这一步很多人会忽略但它很关键。第一次推理时 cuDNN 会做算法基准测试可能耗时几秒如果不预热第一个真实请求会超时。预热 3 次基本能让后续推理稳定在正常耗时。5.3 调度器信号量 动态批处理组合# app/scheduler.py import asyncio import torch from app.model import InferModel class InferScheduler: def __init__(self, max_concurrent4, max_batch8, window0.02): self.sem asyncio.Semaphore(max_concurrent) self.max_batch max_batch self.window window self.pending [] self.lock asyncio.Lock() self.model InferModel.get() async def submit(self, payload): loop asyncio.get_event_loop() future loop.create_future() async with self.lock: self.pending.append((future, payload)) return await future async def run(self): while True: await asyncio.sleep(self.window) async with self.lock: if not self.pending: continue batch self.pending[:self.max_batch] del self.pending[:self.max_batch] async with self.sem: await self._process(batch) async def _process(self, batch): futures [f for f, _ in batch] payloads [p for _, p in batch] try: results await asyncio.to_thread(self._infer, payloads) for f, r in zip(futures, results): if not f.done(): f.set_result(r) except Exception as e: for f in futures: if not f.done(): f.set_exception(e) def _infer(self, payloads): tensors torch.stack([self._preprocess(p) for p in payloads]) tensors tensors.to(self.model.device) with torch.inference_mode(): outputs self.model(tensors) return [self._postprocess(o) for o in outputs] def _preprocess(self, payload): # 你的预处理逻辑 return torch.randn(3, 224, 224) def _postprocess(self, output): return {class: int(output.argmax().item())}这里信号量和批处理是叠加的信号量控制同时进行的批次数批处理控制每批的大小。max_concurrent4、max_batch8意味着峰值显存是 4 批 x 8 样本的激活比 32 个独立请求省得多。5.4 路由与超时兜底# app/main.py import asyncio from contextlib import asynccontextmanager from fastapi import FastAPI, HTTPException from app.scheduler import InferScheduler from app.schemas import InferRequest, InferResponse scheduler None asynccontextmanager async def lifespan(app: FastAPI): global scheduler scheduler InferScheduler(max_concurrent4, max_batch8) task asyncio.create_task(scheduler.run()) yield task.cancel() app FastAPI(lifespanlifespan) app.post(/infer, response_modelInferResponse) async def infer(req: InferRequest): try: result await asyncio.wait_for( scheduler.submit(req.dict()), timeout5.0 ) return result except asyncio.TimeoutError: raise HTTPException(status_code504, detailinference timeout) except torch.cuda.OutOfMemoryError: torch.cuda.empty_cache() raise HTTPException(status_code503, detailgpu busy, retry later)超时兜底是必须的。没有超时一个卡住的请求会一直占着信号量把整个服务拖死。5 秒是个经验值根据你的模型耗时调整一般设成 P99 耗时的 2 到 3 倍。5.5 启动脚本与压测# run.sh #!/bin/bash export PYTORCH_CUDA_ALLOC_CONFmax_split_size_mb:128 uvicorn app.main:app --host 0.0.0.0 --port 8000 --workers 1注意--workers 1。多 worker 会各自加载一份模型显存翻倍除非你有多张卡做数据并行否则单卡就单 worker。压测用wrk或locustwrk -t4 -c50 -d30s -s post.lua http://localhost:8000/infer观察三个指标QPS、P99 延迟、nvidia-smi的显存占用。如果显存稳定在 70% 以下说明并发上限设得合理如果逼近 90%赶紧调小。6. 踩坑实录那些文档里不会写的经验6.1 常见问题速查表问题排查方向解决启动就 OOM模型太大或有多份副本检查 worker 数确认单例跑一段时间 OOM显存碎片设 max_split_size_mb定期 empty_cache首个请求超时未预热启动时跑几次 dummy 推理并发上不去信号量太小或线程池满调大信号量检查 anyio 线程池日志丢失uvicorn 缓冲加 --log-level info用 logging 配置请求延迟抖动大批处理窗口不合理调小 window或按负载动态调整6.2 显存碎片这个隐形杀手碎片问题最阴险因为它不是立刻爆而是跑几小时后突然爆。原因是 PyTorch 的缓存分配器会把释放的显存块留着复用但如果请求的显存大小不一大块被切碎新的请求要连续大块时就分配不到。解法有两个。一是设max_split_size_mb限制切分粒度二是定期调用torch.cuda.empty_cache()但别频繁调因为它会强制同步影响性能。我的做法是在每次 OOM 异常捕获后调一次平时不调。try: outputs self.model(tensors) except torch.cuda.OutOfMemoryError: torch.cuda.empty_cache() raise6.3 别在 async 函数里直接跑推理这个坑我见过太多次。有人写app.post(/infer) async def infer(req): result model(req.tensor) # 错误阻塞事件循环 return result这样写事件循环被阻塞所有请求串行QPS 低到离谱但显存不会爆——因为根本没并发。正确的做法是await asyncio.to_thread(...)让推理跑在线程池里事件循环继续接收请求。6.4 监控显存别等崩了才知道加一个健康检查接口暴露显存状态app.get(/health) async def health(): allocated torch.cuda.memory_allocated() / 1e9 reserved torch.cuda.memory_reserved() / 1e9 total torch.cuda.get_device_properties(0).total_memory / 1e9 return { allocated_gb: round(allocated, 2), reserved_gb: round(reserved, 2), total_gb: round(total, 2), usage_pct: round(reserved / total * 100, 1) }配合 Prometheus 抓取设个告警阈值比如 reserved 超过 85% 就报警。这样能在崩之前介入而不是等用户投诉。6.5 多卡场景的额外注意点如果有多张卡别用DataParallel它会把 batch 拆到各卡但主卡显存占用更高容易主卡先爆。用DistributedDataParallel或者干脆手动分卡每个 worker 绑定一张卡前面挂个负载均衡。import os os.environ[CUDA_VISIBLE_DEVICES] 0 # 每个进程指定一张卡多进程启动时每个进程设不同的CUDA_VISIBLE_DEVICES这样每个进程只看到一张卡显存互不干扰。7. 性能调优的进阶思路7.1 用 FP16 或 INT8 把显存砍一半如果显存实在紧张量化是最直接的手段。FP16 推理能把模型权重和激活都砍一半精度损失通常小于 1%。INT8 更激进但需要校准适合对精度不敏感的场景。model.half() # 转 FP16 dummy dummy.half()注意输入张量也要转 FP16否则会报类型不匹配。实测 ResNet50 从 FP32 转 FP16显存从 1.2G 降到 0.7G推理速度还快了 30%。7.2 动态调整并发数固定并发数在流量波动时不够灵活。可以根据显存使用率动态调整信号量的值def adjust_concurrency(current_pct): if current_pct 85: return max(1, current_concurrency - 1) elif current_pct 50: return min(max_concurrent, current_concurrency 1) return current_concurrency这个逻辑放在一个后台任务里每 10 秒跑一次。注意调整要平滑别频繁抖动。7.3 请求优先级与降级高优先级请求比如付费用户可以走单独的信号量低优先级的在显存紧张时直接返回 503。这样保证核心业务不被拖垮。HIGH_PRIORITY_SEM asyncio.Semaphore(2) LOW_PRIORITY_SEM asyncio.Semaphore(2) app.post(/infer) async def infer(req): sem HIGH_PRIORITY_SEM if req.priority high else LOW_PRIORITY_SEM async with sem: ...这套组合拳打下来服务在流量翻 10 倍的情况下也能稳住最多是低优先级请求被拒绝核心链路不受影响。8. 我个人的几条实操心得第一并发上限宁小勿大。设小了顶多吞吐低一点设大了直接崩恢复成本高得多。我一般先设一个保守值压测后逐步往上调每次加 2观察 10 分钟。第二批处理窗口别超过 50ms。超过这个数用户能感知到延迟。20ms 是个甜点值既能攒到请求又不影响体验。第三一定要有降级路径。显存爆了不能直接 500要返回 503 让客户端重试同时触发告警。我见过太多服务 OOM 后直接进程退出连日志都来不及写。第四压测要模拟真实流量分布。用固定 QPS 压测和用泊松分布压测结果差很多。真实流量是突发的压测时要用wrk的--latency配合随机延迟才能暴露并发控制的问题。第五监控比调优重要。你不可能预判所有流量模式但只要有监控就能在问题发生时快速定位。显存使用率、请求队列长度、推理耗时 P99这三个指标必须盯着。这套方案我在几个项目里都跑过单卡 24G 跑检测模型稳定支撑 50 到 80 QPS显存长期维持在 60% 到 70%没再出过 OOM。核心就一句话把并发当成一种需要显式管理的资源而不是让框架替你决定。
返回列表