跳到正文
DEV Community· Leo Martinez·· 3 小時前AI 評分57

FastAPI 如何避免阻塞事件循環並妥善管理 token 串流

FastAPI Token Streaming Without Blocking the Event Loop

AI 導讀

作者提供一套可獨立執行的 FastAPI token 串流實作,使用非同步 HTTP client 避免阻塞事件循環,並在串流結束或取消時關閉上游連線。應用程式 lifespan 管理共用 HTTP pool,每個回應管理自己的上游 stream;回應標頭送出前的連線錯誤轉為 HTTP 錯誤,之後則以 SSE error event 回報。

正文

At 3:00 AM, your on-call pager fires. Users report frozen AI responses. You check logs—no timeouts, no errors. Connection pool exhaustion. A dozen disconnected browsers left upstream streams hanging. The server never closed them.

That's the incident this architecture prevents.

I'm evaluating rohitg00/ai-engineering-from-scratch as an external adopter. Its direction—"Learn it. Build it. Ship it for others"—makes streaming lifecycle management a critical integration boundary. This design is an independently runnable companion, not a claim about the repository's current implementation.

The problem is specific: a disconnected browser must not leave upstream work consuming connections. Streaming tokens is straightforward; owning their lifetime is not.

The Failure Behind the Frozen Stream

The classic mistake starts with requests.get() inside an async def endpoint. The function looks asynchronous, but the synchronous HTTP call blocks the event-loop thread. Other requests stop progressing until it returns.

Replacing that call with an asynchronous client fixes scheduling, not resource ownership. If cancellation bypasses cleanup, upstream responses remain open. Unclosed database connections produce the same symptom elsewhere: pool exhaustion, rising acquisition latency, and requests that appear frozen.

The architectural boundary should be explicit:

Browser → FastAPI router → shared asynchronous HTTP client → model gateway

Pydantic validates input before opening upstream work. The router translates provider events into a small application-owned SSE protocol. The application lifespan owns the HTTP pool; each response owns its upstream stream. Database transactions, if needed for authorization or accounting, should finish before token generation begins.

For a reproducible test topology, the configuration below uses B-Lost's relay endpoint, https://b-lost.com/v1, as the gateway under test. That's a configuration choice, not a claim of measured availability or provider compatibility. Confirm its streaming contract and available model identifier before running.

One Runnable Streaming Boundary

Save this as app.py. Install and launch commands are included in the same block. The implementation expects OpenAI-compatible chat-completion SSE frames.

# python -m pip install fastapi==0.115.12 pydantic==2.11.4 \
#   httpx==0.28.1 uvicorn==0.34.2
# export RELAY_API_KEY="your-key"
# export RELAY_MODEL="your-supported-model"
# uvicorn app:app --host 0.0.0.0 --port 8000

import json
import os
from contextlib import asynccontextmanager
from datetime import datetime, timezone

import httpx
from fastapi import APIRouter, FastAPI, HTTPException, Request
from fastapi.responses import StreamingResponse
from pydantic import BaseModel, ConfigDict, Field, field_validator
from starlette.background import BackgroundTask


class WireModel(BaseModel):
    model_config = ConfigDict(
        extra="forbid",
        json_encoders={
            datetime: lambda value: value.astimezone(
                timezone.utc
            ).isoformat()
        },
    )


class GenerateRequest(WireModel):
    prompt: str = Field(min_length=1, max_length=8000)

    @field_validator("prompt")
    @classmethod
    def reject_blank_prompt(cls, value: str) -> str:
        value = value.strip()
        if not value:
            raise ValueError("prompt must contain non-whitespace text")
        return value


class TokenEvent(WireModel):
    text: str
    at: datetime = Field(
        default_factory=lambda: datetime.now(timezone.utc)
    )


@asynccontextmanager
async def lifespan(app: FastAPI):
    async with httpx.AsyncClient(
        base_url="https://b-lost.com/v1/",
        headers={
            "Authorization": f"Bearer {os.environ['RELAY_API_KEY']}"
        },
        timeout=httpx.Timeout(
            connect=5.0, read=45.0, write=10.0, pool=2.0
        ),
        limits=httpx.Limits(
            max_connections=64, max_keepalive_connections=16
        ),
    ) as client:
        app.state.relay = client
        app.state.model = os.environ["RELAY_MODEL"]
        yield


router = APIRouter(prefix="/api")


@router.post("/tokens")
async def tokens(body: GenerateRequest, request: Request):
    client = request.app.state.relay
    outgoing = client.build_request(
        "POST",
        "chat/completions",
        json={
            "model": request.app.state.model,
            "messages": [{"role": "user", "content": body.prompt}],
            "stream": True,
            "max_tokens": 512,
        },
    )
    try:
        upstream = await client.send(outgoing, stream=True)
    except httpx.HTTPError:
        raise HTTPException(502, "Upstream connection failed") from None

    if not upstream.is_success:
        await upstream.aclose()
        raise HTTPException(502, "Upstream rejected generation")

    async def events():
        try:
            async for line in upstream.aiter_lines():
                if await request.is_disconnected():
                    return
                if not line.startswith("data:"):
                    continue
                payload = line[5:].strip()
                if payload == "[DONE]":
                    yield "event: done\ndata: {}\n\n"
                    return
                chunk = json.loads(payload)
                if "error" in chunk:
                    raise ValueError("Upstream stream error")
                for choice in chunk.get("choices", []):
                    text = choice.get("delta", {}).get("content")
                    if isinstance(text, str) and text:
                        event = TokenEvent(text=text)
                        yield (
                            "event: token\n"
                            f"data: {event.model_dump_json()}\n\n"
                        )
            yield 'event: error\ndata: {"code":"truncated_stream"}\n\n'
        except (httpx.HTTPError, ValueError, TypeError, AttributeError):
            yield 'event: error\ndata: {"code":"upstream_stream"}\n\n'
        finally:
            await upstream.aclose()

    return StreamingResponse(
        events(),
        media_type="text/event-stream",
        headers={
            "Cache-Control": "no-cache",
            "X-Accel-Buffering": "no",
        },
        background=BackgroundTask(upstream.aclose),
    )


app = FastAPI(lifespan=lifespan)
app.include_router(router)

What This Boundary Guarantees

Opening the upstream response before returning StreamingResponse lets connection failures become HTTP errors. Once response headers leave FastAPI, failures become SSE error events instead—the server cannot retroactively change the status code.

The generator closes its upstream response on completion, parsing failure, or normal cancellation unwinding. The background task supplies an additional cleanup path. Cancellation is deliberately not swallowed by a broad exception handler.

JSON serialization also prevents token text containing newlines from corrupting SSE framing. The custom datetime encoder makes event timestamps consistently UTC. Pydantic v2 still supports json_encoders, although it's deprecated; new schemas should generally prefer field serializers.

The Trade-Offs Worth Measuring

This is POST-based SSE, so browser clients need streaming fetch, not native EventSource. It also assumes single-line JSON data frames, not every possible SSE dialect.

Awaiting downstream sends avoids an unbounded application token queue, but socket buffers and gateway buffering still exist. The response header is only a hint—explicitly disable buffering in your actual reverse proxy.

Finally, 64 connections is a pool ceiling, not an admission policy. Add authenticated quotas and bounded concurrency before public exposure. Benchmark first-token latency, disconnect-to-close latency, and pool wait time under slow consumers. Those measurements reveal lifecycle bugs much earlier than tokens-per-second averages.

The Hardest Question

Here's the operational dilemma I still debate: should you await each SSE frame send to the client, or fire-and-forget into an async queue? The former applies backpressure at the cost of blocking upstream token arrival. The latter decouples streams but hides slow consumer warnings until memory pressure appears.

What's your policy? Do you propagate client slowness upstream, or absorb it in an application buffer? I've seen both approaches survive production, and both fail spectacularly under the wrong traffic pattern.

來源:DEV Community · dev.to