自建AI中转并发限流脚本防止GPU满载宕机故障
自建AI中转并发限流脚本编写指南，防止GPU满载宕机
自建AI中转服务时，如果上游API请求瞬间涌入，GPU很容易被“挤爆”，轻则推理变慢，重则整机宕机。
今天咱就自己动手写一个并发限流脚本，放在中转层前面，把住关口，保护GPU安全。
本文适合没用过限流代码的新手，跟着步骤走就行。
这个脚本能解决什么问题？
AI中转服务本质是转发请求到后端的GPU推理节点。
如果不做控制，哪怕后端只有一块GPU，前端也可能瞬间来几十个并发请求。
GPU显存和算力有限，超载就会OOM甚至死机。并发限流脚本就像一个交通警察，每个时间窗口只放行指定数量的请求，剩下的排队或返回429错误。
这样GPU始终工作在安全水位。
动手前需要准备的3样东西
- 一台中转服务器（Linux系统，CentOS 7+ 或 Ubuntu 18.04+）
- Python 3.7以上（用来跑限流脚本）
- Flask或FastAPI（用作中转API接口）
如果服务器上还没有Python环境，先运行以下命令安装依赖：
# Debian/Ubuntu
sudo apt update && sudo apt install python3 python3-pip -y
# CentOS/RHEL
sudo yum install python3 python3-pip -y然后安装Flask和redis（推荐用Redis实现分布式限流，如果没有Redis，也可以用内存队列，但多实例时必须用Redis）：
pip3 install flask redis分步实现并发限流脚本
第一步：设计限流逻辑
我们采用令牌桶算法：每秒向桶里放N个令牌，每个请求必须取走一个令牌，令牌不够就等待或拒绝。
这里简化成固定窗口计数，按并发数控制（适合小块GPU）。
第二步：编写核心限流代码
创建一个文件 rate_limiter.py，内容如下：
import redis
import time
class ConcurrencyLimiter:
def __init__(self, redis_host='localhost', redis_port=6379, max_concurrent=4):
self.r = redis.Redis(host=redis_host, port=redis_port, decode_responses=True)
self.max_concurrent = max_concurrent
self.key = 'gpu:concurrent' # 计数器键
def try_acquire(self) -> bool:
"""尝试获取一个并发槽位，成功返回True，失败返回False"""
current = self.r.incr(self.key)
if current <= self.max_concurrent:
# 设置过期时间，防止死锁（以秒为单位）
self.r.expire(self.key, 10) # 10秒后自动归零
return True
else:
# 超额了，回退减1
self.r.decr(self.key)
return False
def release(self):
"""请求完成后释放一个槽位"""
self.r.decr(self.key)第三步：集成到中转API中
假如你有一个Flask中转接口，加入限流装饰器（或中间件）。
app.py：
from flask import Flask, request, jsonify
from rate_limiter import ConcurrencyLimiter
import time
app = Flask(__name__)
limiter = ConcurrencyLimiter(max_concurrent=4) # 设置最大并发4
@app.route('/v1/chat/completions', methods=['POST'])
def proxy():
if not limiter.try_acquire():
return jsonify({'error': 'Too Many Requests', 'message': '服务器繁忙，请稍后再试'}), 429
try:
# 这里替换为你的实际转发逻辑（调用后端GPU推理）
# 模拟耗时
time.sleep(2)
result = {'status': 'ok', 'data': '模拟推理结果'}
return jsonify(result)
finally:
limiter.release()
if __name__ == '__main__':
app.run(host='0.0.0.0', port=8000)重要说明：try_acquire 和 release 务必成对出现，推荐用 try...finally 保证释放。
避坑指南——新手最容易犯的3个错误
- 忘记释放槽位：如果请求异常退出而没有执行
release()，槽位会一直占用，导致后续所有请求都被限死。必须在finally块里释放，或者用上下文管理器自动释放。 - 并发数设置过大：先从小往大调。比如单块A100显存80GB，推理大模型可以先设并发2，观察GPU利用率不超95%。切勿照搬别人的数值。
- Redis未启动：如果使用Redis做计数，记得检查Redis服务是否运行。启动命令：
sudo systemctl start redis。
怎么验证脚本生效了？
用 ab 或 wrk 模拟并发请求，观察限制效果。
# 安装ab工具
sudo apt install apache2-utils -y # Ubuntu
sudo yum install httpd-tools -y # CentOS
# 发送20个并发请求（-c 20），总共100个请求（-n 100）
ab -n 100 -c 20 http://你的服务器IP:8000/v1/chat/completions正常情况你会看到部分请求返回429（Too Many Requests），成功请求的数量不会超过并发限制（这里是4）。
同时用 nvidia-smi 监控GPU利用率，满载不会持续100%。
如果一切正常，你的自建AI中转就有了“防爆”能力。
后续可以结合 prometheus + grafana 做可视化监控，效果更直观。
常见问题解答
Q：没有Redis怎么办？
可以用Python内置的 threading.Semaphore 或 asyncio.Semaphore 做单机限流。但多进程或多服务器必须用Redis。
Q：并发数设为4还是8？
没有标准答案。先在低压力下测试，观察GPU显存占用和利用率，逐步上调直到找到安全阈值。建议留20%余量。
Q：限流后用户等待超时怎么办？
可以配合队列（如Celery）或异步任务，把请求放入队列顺序处理，避免直接拒绝。但队列长度也要限制，防止内存泄漏。
如果你正在处理自建AI中转的GPU过载问题，建议先按本文步骤完整执行，再根据自己的环境调整并发数和超时策略；
遇到异常时优先回看避坑和高频问题部分。
（注意：以上代码中的Redis地址和端口需根据实际环境修改，并发数请自行测试确定，不要直接复制为生产设置。
）
---
延伸阅读：如果你还想了解如何为限流脚本加上API防爬虫策略，或者配置Nginx反向代理时的限流方案，可以继续查看本站相关教程。
