自建AI中转并发限流脚本防止GPU满载宕机

知识分享

2026-07-22 09:00

7 阅读

当你自己搭建了一个AI中转服务（比如基于one-api或new-api），后端挂着几块GPU卡跑模型推理时，最怕的就是前端突然涌入大量请求直接把GPU打满，导致服务崩溃甚至主机宕机。
本文就从零开始，带你写一个简单实用的并发限流脚本，放在中转层前面拦住超出的请求，让GPU永远跑在安全水位线上。

前置准备：你要有什么？

在动手之前，确认你的服务器已经装了这些：

中转服务：比如one-api，已经搭建好并能正常转发请求到后端推理机。
Python 3.6+：我们用来写限流脚本，系统自带的或者自己装的都行。
Nginx：如果你用Nginx做反向代理，限流脚本也可以作为中间件插入。但本文用更简单的方式——独立脚本配合supervisor守护。
GPU服务器：安装了nvidia-smi，用于后续验证。

搞清楚一个概念：并发限流不是降低模型精度，而是控制同一时刻能通过的请求数量。
超出部分要么排队，要么直接返回429状态码。
我们这里采用排队机制，避免客户端频繁重试造成雪崩。

编写限流脚本：核心步骤

我们使用Python的asyncio和aiohttp库来写一个异步代理脚本，自带令牌桶限流。
新建文件ai_proxy_limit.py，粘贴以下内容：

import asyncio
import aiohttp
from collections import deque
import time

# 配置区域
UPSTREAM_URL = "http://你的中转服务IP:端口"  # 你的中转服务地址
MAX_CONCURRENT = 4                         # 最大并发数，根据GPU显存/算力调整
TASK_QUEUE = deque()                       # 等待队列
current_concurrent = 0                     # 当前正在处理的请求数

async def handle_request(session, request_data):
    global current_concurrent
    # 模拟转发并接收响应
    async with session.post(UPSTREAM_URL, json=request_data) as resp:
        return await resp.text()

async def worker(session):
    global current_concurrent
    while True:
        if TASK_QUEUE and current_concurrent < MAX_CONCURRENT:
            request_data = TASK_QUEUE.popleft()
            current_concurrent += 1
            try:
                result = await handle_request(session, request_data)
                # 这里可以返回结果给客户端，本文为了简化省略websocket/回调
                print(f"完成请求，剩余并发：{current_concurrent-1}")
            finally:
                current_concurrent -= 1
        await asyncio.sleep(0.01)

async def main():
    # 模拟接收外部请求（实际场景可能用HTTP端口或队列消费）
    connector = aiohttp.TCPConnector(limit=100)
    async with aiohttp.ClientSession(connector=connector) as session:
        asyncio.create_task(worker(session))
        # 这里持续等待新请求加入队列... 实际使用时需要监听端口
        await asyncio.Event().wait()

if __name__ == "__main__":
    asyncio.run(main())

说明：这个脚本是一个基础框架，MAX_CONCURRENT就是你要设置的并发上限。
使用时需要改写接收外部请求的部分（比如使用aiohttp.web开一个HTTP端口）。
如果你已经用Nginx做反向代理，可以结合limit_req_zone做更简单的限流，但灵活性不如脚本。

小提示：如何让脚本常驻运行

把脚本保存后，用supervisor守护。
配置示例（/etc/supervisor/conf.d/ai_limit.conf）：

[program:ai_limit]
command=python3 /root/ai_proxy_limit.py
autostart=true
autorestart=true
stderr_logfile=/var/log/ai_limit.err.log
stdout_logfile=/var/log/ai_limit.out.log

然后运行：

supervisorctl reread
supervisorctl update
supervisorctl start ai_limit

避坑指南：这些细节一定要注意

并发数不是越大越好：先测出单次推理的显存占用，再推算安全并发。例如单次占用6GB，12GB显存最多开2并发，留一点余量。
限流后客户端状态码：如果改用直接拒绝的方式（非排队），一定要返回429 Too Many Requests，并在响应头里加Retry-After，让客户端主动等待，否则它会立即重试，反而加重压力。
注意Python GIL：上面用了异步，不影响I/O，但如果推理本身是CPU密集型（某些场景），建议用多进程方案。
测试时用ab或wrk压测：先用小并发验证队列是否正常，再用大并发看GPU利用率是否被控制在范围之内。
日志记录：给限流脚本加上日志输出，方便追查哪一次请求被限流、哪次失败。

效果验证：怎么知道限流有用？

启动限流脚本和你的中转服务。
用压测工具模拟请求：

# 安装ab
sudo apt install apache2-utils -y
# 假设中转服务入口是http://localhost:8080，发送100个请求，并发20
ab -n 100 -c 20 http://localhost:8080/v1/chat/completions

在另一窗口实时观察GPU使用率：

watch -n 1 nvidia-smi --query-gpu=utilization.gpu,memory.used --format=csv

理想结果：GPU利用率稳定在安全阈值（比如80%以下），压测过程中没有出现OOM或服务中断。同时压测工具会显示部分请求响应时间增加（排队导致）但无失败。

常见问题解答

Q：限流脚本会不会变成新的瓶颈？
A：会的，所以限流脚本本身要高效。Python异步框架足以应付每秒几百的请求量，单个脚本性能没问题。如果请求量更大，可以在Nginx层用limit_req双保险。

Q：我不想写Python，Nginx自带限流够用吗？
A：如果你只是简单地控制QPS（每秒请求数），Nginx的limit_req_zone完全够用。但如果你需要按GPU剩余显存动态调整并发，还是得用脚本来配合nvidia-smi数据。

Q：脚本里怎么把响应返回给客户端？
A：上面为了展示核心逻辑做了简化。实际你应该用aiohttp.web接收HTTP请求，转发后把响应写回客户端。网上有很多现成的异步代理例子，可以搜索“python async proxy openai”参考。

掌握了这几个步骤，你的自建AI中转服务就能在高并发下稳定运行。
遇到异常时优先检查MAX_CONCURRENT设置是否合理，以及中转服务本身有没有报错逻辑。

OneAPI多租户权限分离杜绝下游数据互通泄露

Ollama离线模型包无网络住宅主机部署教程