FastAPI 构建生产级 LLM 流式代理:解决连接泄漏、假死与断连背压

2次阅读
没有评论

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

上个月我们将内部的 AI Agent 网关迁移至自研的 FastAPI 服务,主要承载前端对 vLLM 与云端大模型的流式推理(Server-Sent Events, SSE)调用。上线刚满一周,Prometheus 就拉响了报警:网关 Pod 内存持续飙升,TIME_WAIT 与 CLOSE_WAIT 句柄数破万,最终导致整个代理层假死,所有流式请求全线阻塞。

表面上看,流式转发不过是 StreamingResponse(upstream_stream()) 几行代码的事。但一进真实高并发场景,客户端中途关掉网页、弱网抖动、上游模型吞吐抖动,都会让缺乏保护的异步协程直接沦为“僵尸任务”。这篇记录整理了排查现场及最终重构出来的生产级流式代理实现方案。

一、现场排查:谁动了我的 Socket 句柄?

出现故障时,我通过 lsof -i -a -p <pid> 查看 FastAPI 进程,发现与上游推理集群建立了数千条长连接,但当时并发在线用户其实只有几十个。顺着 Trace 链深入,最终定位出三个致命设计缺陷:

1. 客户端断开连接,后台还在为爱发电

前端用户点下“停止生成”,或者直接关掉浏览器标签页,浏览器会向网关发送一个 TCP RST 或 FIN 包。然而,默认的 StreamingResponse 如果直接包一个 httpx.stream() 生成器,FastAPI 并不会自动打断内部对大模型的读取!上游 vLLM 依然在竭尽全力地生成 Token,网关也在继续读取、试图推送,直到整个几千 Token 的推理完成,连接才被迫关闭。不仅白白烧光算力,连接也在内存里悬挂了数分钟。

2. 每次请求新建 Client,连接池瞬间打爆

部分初学者喜欢在每个请求路由里写 async with httpx.AsyncClient() as client:。在高并发 SSE 场景下,频繁创建与销毁 TLS 握手会导致巨额延迟,同时在高频创建销毁连接的交界处,操作系统堆积了大量 TIME_WAIT 状态的端口,直接触发 Linux 的 local port reservation 瓶颈。

3. Nginx 缓冲截流摧毁首字延迟 (TTFT)

网关前挂了反向代理 Nginx,默认开启了 proxy_buffering。上游模型吐出的一个个 chunk 被 Nginx 硬生生攒满 4KB 才往浏览器发。前端看起来就像卡死了一样,等 10 秒钟突然“嘭”地弹出一大段文本,体验彻底归零。

二、架构重构:安全健壮的流式转发代码

针对上述痛点,我们重新设计了请求生命周期管理:全局单例连接池、基于 request.is_disconnected() 的主动背压打断,以及优雅异常退出机制。

import asyncio
import logging
from contextlib import asynccontextmanager
from typing import AsyncGenerator

import httpx
from fastapi import FastAPI, Request, HTTPException
from fastapi.responses import StreamingResponse

logger = logging.getLogger("proxy")

# 全局复用高性能连接池
class UpstreamPool:
    client: httpx.AsyncClient | None = None

    @classmethod
    def get_client(cls) -> httpx.AsyncClient:
        if cls.client is None or cls.client.is_closed:
            # 针对大模型 SSE 场景调整 keepalive 与 limits
            cls.client = httpx.AsyncClient(
                limits=httpx.Limits(
                    max_keepalive_connections=200,
                    max_connections=1000,
                    keepalive_expiry=60.0
                ),
                # 流式响应严禁设置过短的 read timeout,但 connect 必须快速失败
                timeout=httpx.Timeout(connect=5.0, read=120.0, write=5.0, pool=10.0),
                http2=False  # 根据上游支持情况决定,内网多数 HTTP/1.1 更稳健
            )
        return cls.client

@asynccontextmanager
async def lifespan(app: FastAPI):
    # 初始化连接池
    UpstreamPool.get_client()
    yield
    # 服务优雅下线,排空连接
    if UpstreamPool.client:
        await UpstreamPool.client.aclose()

app = FastAPI(lifespan=lifespan)

async def stream_generator(request: Request, payload: dict) -> AsyncGenerator[bytes, None]:
    client = UpstreamPool.get_client()
    upstream_url = "http://internal-vllm-node:8000/v1/chat/completions"

    try:
        async with client.stream(
            "POST",
            upstream_url,
            json=payload,
            headers={"Authorization": "Bearer internal-token"}
        ) as response:
            
            if response.status_code != 200:
                error_detail = await response.aread()
                logger.error(f"Upstream error: {response.status_code} - {error_detail.decode()}")
                yield f"data: {{\"error\": \"Upstream returned {response.status_code}\"}}\n\n".encode()
                return

            # 按行迭代 SSE 数据流
            async for line in response.aiter_lines():
                # 核心防空转逻辑:每接收一个 chunk,先探测下游客户端是否存活
                if await request.is_disconnected():
                    logger.warning("Downstream client disconnected. Terminating upstream stream.")
                    break

                if line:
                    yield f"{line}\r\n\r\n".encode("utf-8")
                    
    except httpx.ReadTimeout:
        logger.error("Timeout reading from upstream LLM.")
        yield b"data: {\"error\": \"Upstream read timeout\"}\n\n"
    except asyncio.CancelledError:
        # 当客户端强制中断且框架取消协程时的兜底拦截
        logger.info("Streaming task was cancelled by client abort.")
        raise
    except Exception as e:
        logger.exception(f"Unexpected streaming error: {str(e)}")
        yield b"data: {\"error\": \"Internal gateway error\"}\n\n"

@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):
        raise HTTPException(status_code=400, detail="Only stream=True supported via this route")

    return StreamingResponse(stream_generator(request, body),
        media_type="text/event-stream",
        headers={
            # 必须透传的标准 SSE 头部
            "Cache-Control": "no-cache",
            "Connection": "keep-alive",
            "Content-Type": "text/event-stream",
            # 告诉 Nginx 不要对该响应做内存缓冲
            "X-Accel-Buffering": "no",
        }
    )

三、边沿反向代理与容器底层调优

光把 Python 逻辑改对还不够,只要架构链路上存在中间节点,就必须针对流式特化配置。

1. 搞定 Nginx 的首字延迟与无响应断开

在公网出口 Nginx 的配置中,需要关闭缓冲并延长超时时间。特别是 proxy_read_timeout,模型生成长文本时推理节点两段输出之间可能存在几十秒的空档(思考模型如 DeepSeek-R1 更明显),默认的 60 秒极易断连:

location /v1/chat/completions {
    proxy_pass http://fastapi_backend;
    
    # 关闭响应缓冲,保障 TTFT 为 0 延迟
    proxy_buffering off;
    proxy_cache off;
    
    # 建立持久化长连接通道
    proxy_http_version 1.1;
    proxy_set_header Connection "";
    proxy_set_header Host $host;

    # 思考类大模型需要放宽读取超时
    proxy_read_timeout 600s;
    proxy_send_timeout 600s;

    # 启用 TCP 探测,快速释放异常死链接
    tcp_nodelay on;
}

2. Uvicorn Worker 选型与事件循环

在部署容器镜像时,必须替换掉 Python 自带的 asyncio 默认事件循环,直接上 uvloop。对于 IO 密集的代理层,单 Worker 承载万级并发不是问题。我测试得出的最佳实践是配合 Gunicorn 挂载 UvicornWorker:

# Dockerfile 中的启动指令
exec gunicorn app:app \
    --workers 4 \
    --worker-class uvicorn.workers.UvicornWorker \
    --bind 0.0.0.0:8080 \
    --timeout 120 \
    --keep-alive 75

四、压测与验证:模拟暴力断连

重构完成后,我使用 Locust 编写了一个压测脚本:模拟 500 个并发用户发起请求,但设置每个连接在接收到第 5 个 SSE chunk 时 强行掐断 TCP 链路(模拟用户点击“取消”或页面切换)。

指标 旧实现(裸 StreamingResponse) 新方案(连接池 + 断连主动探测)
Pod 内存占用 (5min 压测) 持续上涨至 1.8GB 触发 OOM 平稳维持在 120MB ~ 150MB
TIME_WAIT 句柄数量 9,400+ 个 < 150 个
上游 Token 浪费率 94%(模型仍生成完整响应) < 2%(断连后在 1 个 chunk 内中止)
首字响应 (TTFT) 波动极大(部分被 Nginx 截留) 稳定在 220ms 以内

总结与避坑心得

在 Python 异步生态中处理 SSE 长流式代理,必须将 生命周期控制权 放在第一位。以下是上生产前需要逐项核对的检查清单:

  • 上游客户端严禁临时实例化:确保 httpx.AsyncClient 维持全局连接池,否则在超高频并发下 TCP 握手和端口耗尽会教你做人。
  • 生成器每帧检查存活:异步生成器内部务必每轮迭代调用 await request.is_disconnected(),发现下游挂断立即 break,及时止损后端 GPU 资源。
  • 网关透传特定头部:X-Accel-Buffering: no 是配合 Nginx 的杀手锏,能省掉单独为该路由定制 Nginx 配置的大半精力。
  • Timeout 分级设定:connect 必须严格限制在 3~5 秒以便快速重试,而 read 针对大模型应至少给到 120 秒甚至更长,避免遇到生成长文本时偶发中断。
正文完
 0
评论(没有评论)