ai-infra-course

模块 7 配套代码 — 推理服务平台

文件说明

文件	内容	对应 PPT
ai_gateway.py	Flask 简易 AI 网关 — Random/Consistent Hash LB + Token Bucket + 健康检查 + 故障转移	第 42 页 [动手]
mock_vllm.py	vLLM Mock 后端 — 无需 GPU，本地测试网关用	—

环境要求

• Python ≥ 3.10
• Flask + requests: pip install flask requests
• vLLM ≥ 0.6.0 (真实后端，需要 GPU；测试网关功能可用 mock_vllm.py，无需 GPU)

运行方法

# 方式 1: 使用 mock 后端 (推荐，无需 GPU)
python3 mock_vllm.py --port 8001 &
python3 mock_vllm.py --port 8002 &
python3 ai_gateway.py

# 方式 2: 使用真实 vLLM 后端
vllm serve Qwen/Qwen2.5-0.5B-Instruct --port 8001 &
vllm serve Qwen/Qwen2.5-0.5B-Instruct --port 8002 &
python3 ai_gateway.py

测试

测试 1: 基本路由 (观察两个后端的负载分配)

for i in $(seq 1 10); do
  curl -s http://localhost:8080/v1/chat/completions \
    -H "Content-Type: application/json" \
    -d '{"model":"test","messages":[{"role":"user","content":"say hi"}],"max_tokens":10}' &
done

测试 2: 限流 (观察 429 响应)

for i in $(seq 1 15); do
  curl -s -w "\nHTTP %{http_code}\n" http://localhost:8080/v1/chat/completions \
    -H "Content-Type: application/json" \
    -H "Authorization: Bearer test-key" \
    -d '{"model":"test","messages":[{"role":"user","content":"hi"}],"max_tokens":5}'
done

测试 3: 故障转移 (kill 8001 后端，观察自动切换)

kill %1   # 停掉 8001 后端
curl -s http://localhost:8080/health  # 查看后端状态
curl -s http://localhost:8080/v1/chat/completions \
    -H "Content-Type: application/json" \
    -d '{"model":"test","messages":[{"role":"user","content":"hello"}],"max_tokens":10}'

负载均衡策略

通过环境变量 LB_STRATEGY 切换:

# 加权随机 (默认) — 无 Cache 亲和性
LB_STRATEGY=random python3 ai_gateway.py

# 一致性哈希 — 相同 Session-ID → 固定后端 (KV Cache 热命中)
LB_STRATEGY=consistent-hash python3 ai_gateway.py

验证一致性哈希效果:

# 同一 session 的所有请求打到同一个后端
for i in $(seq 1 5); do
  curl -s http://localhost:8080/v1/chat/completions \
    -H "Content-Type: application/json" \
    -H "Authorization: Bearer test-key" \
    -H "X-Session-ID: user-alice" \
    -d '{"model":"test","messages":[{"role":"user","content":"hi"}]}'
done

代码结构

ai_gateway.py   (~140 行)
├── BACKENDS         # 后端地址 + 权重配置
├── TokenBucket      # Token Bucket 限流器 (线程安全)
├── health_check()   # 主动健康检查 (GET /v1/models)
├── select_backend() # 加权随机选择健康后端
└── Flask 路由:
    ├── /v1/chat/completions  # 认证 → 限流 → 路由 → 转发 (流式)
    └── /health               # 网关状态 + 后端健康