FastAPI GPU推理并发控制实战:避免请求一多就显存溢出

31 阅读4分钟

推理部署在单请求测试时响应正常,一接入多个用户就出现显存溢出、超时甚至服务退出,是GPU服务器租用场景中的常见问题。GPU并不适合让所有请求毫无限制地同时执行。本文使用FastAPI、Semaphore和单进程模型服务,搭建一套基础并发控制方案。

一、问题背景

每个请求都会创建输入张量、中间激活和输出缓存。并发提高后,这些对象会同时占用显存。大模型训练更关注参数和梯度,在线服务还要考虑上下文长度、生成Token数与请求叠加。

很多开发者把Web服务Worker数量直接调大,但每个进程可能各加载一份模型,反而迅速耗尽显存。GPU算力平台提供的是计算资源,应用仍需设置入口队列、并发上限和超时策略。

二、环境准备

准备Linux、Python、CUDA、PyTorch、FastAPI和Uvicorn:

pip install fastapi uvicorn
nvidia-smi

Python 3.12 + CUDA 12.8开发镜像,包含Ubuntu 24.04、JupyterLab和SSH,并支持继续安装PyTorch等框架。当前GPU云服务器和镜像资源可通过官网查看。

三、实操步骤

步骤1:模型只加载一次

import torch
from fastapi import FastAPI

app = FastAPI()
device = "cuda"
model = load_model().to(device).eval()

不要在每个请求中重新加载权重。模型应在进程启动时初始化,并在启动完成后再对外提供服务。

步骤2:添加并发闸门

import asyncio
from pydantic import BaseModel
from starlette.concurrency import run_in_threadpool

gate = asyncio.Semaphore(1)

class RequestBody(BaseModel):
    text: str

def infer(text):
    with torch.inference_mode():
        return model_generate(model, text)

@app.post("/generate")
async def generate(req: RequestBody):
    async with gate:
        result = await run_in_threadpool(infer, req.text)
        return {"result": result}

并发上限应从1开始压测,再根据峰值显存和延迟逐步提高。Semaphore限制的是同时进入模型的请求数,等待中的请求不会立即占用推理显存。

步骤3:限制输入与输出

在接口层检查文本长度,并为生成任务设置最大Token数。无限长输入不仅增加单次显存,还会让后续请求长时间排队。多模态服务还应限制图片数量、尺寸和文件类型。

if len(req.text) > 8000:
    raise HTTPException(status_code=400, detail="input too long")

阈值应依据模型和业务压测确定。

步骤4:使用单Worker启动

uvicorn app:app --host 0.0.0.0 --port 8000 --workers 1

单GPU服务优先从一个Worker开始。如果使用多个Worker,每个进程通常会各自加载模型。需要扩容时,应先评估多实例、多GPU或专用推理框架,而不是只增加Web进程。

步骤5:增加健康检查

@app.get("/health")
def health():
    return {
        "ready": model is not None,
        "cuda": torch.cuda.is_available()
    }

健康检查验证服务就绪,不应触发完整生成。还应记录请求耗时、排队时间、错误码和峰值显存,为后续容量判断提供数据。

步骤6:执行阶梯压测

依次使用1、2、4等并发级别测试,记录吞吐、首Token延迟、完整响应时间与显存峰值。每次只改变一个变量。若准备从测试转向正式AI应用部署,应使用接近真实业务长度的输入,而不是只测短句。

四、常见问题

1. 设置Semaphore后吞吐没有提高

Semaphore主要用于保护资源。若要提升吞吐,可在显存允许时增加并发,或使用支持动态批处理的推理引擎。

2. 服务启动就占满显存

先检查模型精度、权重大小和是否有多个Worker,再确认GPU上是否残留其他进程。

3. 请求排队时间过长

应设置超时和队列长度,并在入口快速拒绝超出能力的请求,避免所有用户无限等待。

4. 深度学习模型可以开多个Worker吗

可以,但要计算每个进程的模型副本和运行显存,不能按CPU Web服务的经验直接设置。

五、总结

稳定的推理部署需要控制并发、限制输入、保持单次模型加载,并通过阶梯压测确定容量。润云智算提供按需GPU资源和开发镜像,可支持模型验证、大模型训练及推理服务测试;实际并发能力仍应由模型、上下文、显存和业务数据共同决定。

FAQ

Q1:Semaphore设置为多少合适?

从1开始,根据峰值显存、延迟和吞吐逐级测试后确定。

Q2:为什么不能直接增加Uvicorn Worker?

多个进程通常会重复加载模型,造成显存成倍增加。

Q3:排队请求会占用GPU显存吗?

尚未进入模型计算的请求通常只占用服务端内存,不会创建完整推理张量。

Q4:AI算力平台能保证固定并发吗?

不能只凭GPU型号保证。还需结合模型版本、输入长度、生成长度和服务框架压测。