实战踩坑:Python构建LLM流式代理网关的内存泄漏与反压优化

2次阅读
没有评论

共计 5157 个字符,预计需要花费 13 分钟才能阅读完成。

最近在给团队私有化部署的大模型推理集群做接入层统一网关,原本以为拿 FastAPI 写个 StreamingResponse 转发下游 vLLM 或 Ollama 的 SSE(Server-Sent Events)只是半天的工作量。结果灰度上线当天,压测并发刚顶到 200,网关节点的物理内存直接一路狂飙,最后触发 Linux OOM Killer 强杀了 worker 进程。

顺藤摸瓜排查后发现,问题不是出在模型推理算力上,而是隐藏在 Python 异步 IO 与 HTTP 流式转发机制的夹缝里:当下游客户端因网络弱或前端渲染阻塞导致消费速率变慢时,服务端如果缺乏有效的背压(Backpressure)与超时断开机制,FastAPI 就会在内存中疯狂积压未消费的 chunk 缓冲区。今天把这套生产级代理网关的架构演进、核心代码和排障避坑细节完整盘一遍。

一、为什么简单的流式转发会演变成内存黑洞?

很多开源脚本和教程里给出的流式代理代码通常长这样:

# 常见但有严重隐患的流式转发写法
@app.post("/v1/chat/completions")
async def chat_proxy(request: Request):
    client = httpx.AsyncClient()
    req_data = await request.json()
    
    async def event_generator():
        async with client.stream("POST", UPSTREAM_URL, json=req_data) as resp:
            async for chunk in resp.aiter_bytes():
                yield chunk

    return StreamingResponse(event_generator(), media_type="text/event-stream")

这段代码在本地单并发测试时毫无破绽,但一旦放进生产环境,至少有三个致命弱点:

  1. httpx 客户端生命周期失控:每次请求实例化一个 AsyncClient,连接池根本没复用,高并发下 TCP 连接处于 TIME_WAIT 状态,句柄直接耗尽。
  2. 客户端断连感知迟钝(Zombie Stream):当用户在网页端刷新、关闭标签页或点了“停止生成”,由于未捕获 request.is_disconnected(),生成器依然在傻傻向模型推理后端拉取 token,下游继续烧显卡算力,上游代理网关继续缓存垃圾字节。
  3. 缓冲区无背压机制:下游客户端接收慢,网关就用极其昂贵的内存去兜底存储这些未消费的 token 片段,直至把堆内存撑爆。

二、高并发流式代理架构核心设计

要构建一个稳固的 Python 流式代理网关,核心思路是:全局长连接池 + 请求断连感知环 + 异步队列背压限流 + 异常兜底终止。

我们在工程中选用了 httpx.AsyncClient 配备 Limits 参数作为连接池,同时使用 asyncio.Queue(maxsize=N) 构建有限缓冲区。如果消费速度跟不上生产速度,生产者会自动挂起(await queue.put() 阻塞),从而将背压自然传导至 upstream 推理后端。

2.1 生产级流式代理核心网关代码

import asyncio
import logging
from typing import AsyncGenerator
from contextlib import asynccontextmanager
from fastapi import FastAPI, Request, HTTPException
from fastapi.responses import StreamingResponse
import httpx

logger = logging.getLogger("gateway")

# 全局共享 HTTP 客户端,配置合理的高并发连接池与超时时间
http_client: httpx.AsyncClient = None

@asynccontextmanager
async def lifespan(app: FastAPI):
    global http_client
    limits = httpx.Limits(max_keepalive_connections=500, max_connections=1000)
    timeout = httpx.Timeout(connect=5.0, read=60.0, write=5.0, pool=5.0)
    http_client = httpx.AsyncClient(limits=limits, timeout=timeout)
    yield
    await http_client.aclose()

app = FastAPI(lifespan=lifespan)

UPSTREAM_LLM_URL = "http://127.0.0.1:8000/v1/chat/completions"

async def stream_processor(
    request: Request,
    payload: dict,
    max_buffer_size: int = 10
) -> AsyncGenerator[bytes, None]:
    """流式数据处理器:集成主动断连检测与背压控制"""
    queue: asyncio.Queue = asyncio.Queue(maxsize=max_buffer_size)
    stop_event = asyncio.Event()

    async def fetch_upstream():
        try:
            headers = {"Authorization": request.headers.get("Authorization", "")}
            async with http_client.stream("POST", UPSTREAM_LLM_URL, json=payload, headers=headers) as resp:
                if resp.status_code != 200:
                    error_detail = await resp.aread()
                    logger.error(f"上游模型节点返回异常: {resp.status_code}, {error_detail.decode('utf-8', errors='ignore')}")
                    await queue.put(f"data: {{\"error\": \"Upstream returned {resp.status_code}\"}}\n\n".encode('utf-8'))
                    return

                async for chunk in resp.aiter_raw():
                    if stop_event.is_set():
                        break
                    # 若客户端主动断连,提前释放上游算力
                    if await request.is_disconnected():
                        logger.warning("客户端已断开连接,主动终止上游推理流")
                        stop_event.set()
                        break
                    
                    # 关键点:队列满时阻塞,阻断无限制的内存缓冲
                    await queue.put(chunk)
        except asyncio.CancelledError:
            pass
        except Exception as e:
            logger.exception(f"上游读取异常: {e}")
        finally:
            # 无论何种异常退出,写入 None 作为消费终止符
            await queue.put(None)

    # 启动后台生产者任务
    fetch_task = asyncio.create_task(fetch_upstream())

    try:
        while True:
            # 如果请求已中断,直接退出循环
            if await request.is_disconnected():
                stop_event.set()
                break

            try:
                # 设置单次 token 等待超时,防止上游假死导致连接长期挂起
                chunk = await asyncio.wait_for(queue.get(), timeout=30.0)
            except asyncio.TimeoutError:
                logger.error("读取上游流式数据超时(30s 无新数据)")
                yield b"data: {\"error\": \"Gateway stream timeout\"}\n\n"
                break

            if chunk is None:
                break
            
            yield chunk
            queue.task_done()
    finally:
        stop_event.set()
        fetch_task.cancel()
        # 等待后台任务彻底清理
        try:
            await fetch_task
        except asyncio.CancelledError:
            pass

@app.post("/v1/chat/completions")
async def chat_proxy(request: Request):
    try:
        body = await request.json()
    except Exception:
        raise HTTPException(status_code=400, detail="Invalid JSON body")

    # 如果是非流式请求,走普通转发;这里重点针对流式
    if not body.get("stream", False):
        resp = await http_client.post(UPSTREAM_LLM_URL, json=body)
        return resp.json()

    return StreamingResponse(stream_processor(request, body),
        media_type="text/event-stream",
        headers={
            "Cache-Control": "no-cache",
            "Connection": "keep-alive",
            "X-Accel-Buffering": "no"  # 禁用反向代理(如 Nginx)的缓冲机制
        }
    )

三、实战压测与排障避坑细节

1. 为什么一定要加 X-Accel-Buffering: no?

如果你的架构是 客户端 -> Nginx -> Python 网关 -> vLLM,你可能会发现不管 Python 代码怎么写,前端都必须等待 4KB 或 8KB 数据攒满后才能一次性收到一坨文字,完全没有打字机效果。这是因为 Nginx 默认启用了 proxy buffering。加上这个 HTTP 头后,Nginx 会立刻将收到的 chunk 吐给客户端,避免在中间层再次发生积压。

2. 压测对比:原生直转发 vs 背压队列控制

我在本地 4 台 8 核 16G 节点上,使用 Locust 模拟了 300 个并发慢速客户端(故意限制客户端接收速率在 100 字节 / 秒,模拟弱网移动端用户)进行压力测试,对比网关节点的表现:

网关实现方案 并发峰值内存 模型算力浪费率 (中断后继续生成) 单 Worker 吞吐 (QPS) 稳定性状态
未做背压与断连感知(基础版) 14.2 GB (直接触发 OOM) 约 42% 41 Worker 频繁挂掉重启
背压队列 + 断连清理(优化版) 850 MB (平稳恒定) < 2% 118 平稳运行 6 小时无抖动

可以看到,内存消耗从压测几分钟就爆内存,直接压制到了 1GB 以内。核心就在于 queue.maxsize = 10 严格锁死了单个长连接在内存里滞留的 payload 数量,上游推理引擎因为写不进 TCP 缓冲区,自身的生成也会随之受限甚至在客户端断连后直接被 kill,把算力和内存双向节省了下来。

3. Gunicorn 部署时的工作模式选型踩坑

在生产环境部署时,我们用 Gunicorn 守护 FastAPI。很多人随手抄配置用了 sync 或 gevent,这对于纯 IO 绑定的 SSE 网关来说是大忌。

# 生产推荐配置启动命令
gunicorn main:app \
  --workers 4 \
  --worker-class uvicorn.workers.UvicornWorker \
  --bind 0.0.0.0:8080 \
  --backlog 2048 \
  --keep-alive 65 \
  --timeout 120

一定要明确使用 uvicorn.workers.UvicornWorker。由于大模型生成一个几千 token 的回答耗时常常在 15~45 秒左右,属于典型的 极低并发密度、极高持续时长 的连接。Worker 数量不需要设太多(通常 CPU 核心数 * 1 ~ 2),重点是让单 Worker 依靠 uvloop 事件循环高效承载数千个持续不断流的 TCP 会话。

总结与避坑心得

构建大模型流式代理,看似只是简单的 HTTP 文本拼接,但本质上是一场精细的 异步生产 - 消费速率平衡战。我们在写后端代理层时,务必保持两个肌肉记忆:

  1. 永远假设客户端随时会拔网线:及时监听 request.is_disconnected() 并联动终止上游任务,保护宝贵的 GPU 推理资源。
  2. 永远给所有的异步队列设置 maxsize:任何没有设定容量上限的 asyncio.Queue 或内存缓存,在不可控的公网环境与慢速客户端面前,都是一颗随时会被引爆的 OOM 炸弹。
正文完
 0
评论(没有评论)