FastAPI: inferência de modelos em streaming com asyncio, chunking e SSE

1) Problema real que resolve

Contexto técnico

Aplicações que fazem inferência com modelos grandes (LLMs, transformadores, modelos de voz/texto) precisam devolver respostas com latência baixa e experiência progressiva. Usuários esperam ver tokens incrementais ou progresso antes do resultado final — especialmente em interfaces conversacionais, aplicativos de assistente e pipelines de ETL que fazem streaming de saídas.

Quando a inferência é feita num endpoint HTTP tradicional que retorna um JSON só no final, o cliente fica bloqueado até o processamento terminar. Isso aumenta latência percebida, degrada UX e dificulta streaming de grandes volumes de texto.

Impacto em produção

Implementações ingênuas causam consumo alto de memória (buffering completo do output), picos de latência e dificuldade para cancelar tarefas quando o cliente desiste. Em cenários com muitos usuários simultâneos, a escalabilidade é prejudicada: threads/loops ocupados, filas longas e spikes em custo de infraestrutura.

Uma solução baseada em streaming com SSE (Server-Sent Events) e asyncio permite transmitir tokens progressivamente, aplicar chunking para reduzir overhead de mensagens e controlar concorrência para evitar sobrecarregar o modelo.

2) Abordagem ingênua vs abordagem otimizada

Abordagem ingênua

A abordagem ingênua executa a inferência de forma síncrona ou bloqueante, coleta todo o resultado em memória e só então retorna um JSON com o texto completo. É simples e funciona para protótipos, mas tem vários problemas:

• Latência percebida alta: usuário só vê saída depois de tudo pronto.
• Uso de memória: buffers crescentes para respostas grandes.
• Sem cancelamento fácil: difícil parar a inferência se o cliente desconectar.

Exemplo típico: endpoint que chama model.generate(...) e retorna texto. Funciona, mas não escala nem dá boa UX.

Abordagem otimizada

A abordagem otimizada usa um pipeline baseado em asyncio: o modelo produz tokens incrementalmente (token stream), cada token entra numa fila (asyncio.Queue), o endpoint consome a fila e faz chunking — agrupando tokens antes de enviar pelo SSE — para reduzir mensagens pequenas e overhead. Além disso, usamos um semáforo (asyncio.Semaphore) para limitar concorrência de inferência, timeouts e checagem de desconexão para cancelar trabalho.

Vantagens:

• Baixa latência perceptível ao usuário.
• Menos overhead de mensagens via chunking.
• Controle de concorrência e cancelamento em caso de desconexão.

3) Implementação com código real

Simulador de modelo (async token generator)

Antes de integrar um modelo real, crie um simulador assíncrono que gera tokens com delay. Isso facilita testes e demonstra o padrão sem dependências externas.

import asyncio

async def fake_model_stream(prompt: str):
    """Simula um gerador de tokens do modelo.

    Cada yield representa um token gerado incrementalmente.
    """
    words = (prompt + "").split()
    for i, w in enumerate(words):
        await asyncio.sleep(0.05)  # latência simulada por token
        yield w + (" " if i < len(words) - 1 else "")

O código acima gera tokens com pause para simular latência do modelo. Em produção, substitua por streaming do SDK do modelo (ex: generate_stream em libs que suportam callback async).

Servidor FastAPI: SSE + asyncio.Queue + chunking

Este trecho é o núcleo: cria um endpoint de streaming que inicia uma tarefa de inferência, coloca tokens numa fila e envia eventos SSE agrupados por chunk_size.

from fastapi import FastAPI, Request, HTTPException
from starlette.responses import StreamingResponse
import asyncio
import json

app = FastAPI()

# Limita quantas inferências rodarão simultaneamente
INFER_SEMAPHORE = asyncio.Semaphore(4)

async def stream_generator(prompt: str, request: Request, chunk_size: int = 4):
    """Gera dados formato SSE com chunking.

    - Lança uma tarefa de inferência que coloca tokens em uma queue.
    - Consome a queue, agrupa tokens em chunks e envia como eventos SSE.
    - Verifica request.is_disconnected() para cancelar.
    """
    queue = asyncio.Queue()
    stop_event = asyncio.Event()

    async def producer():
        # Protege o modelo com um semáforo
        async with INFER_SEMAPHORE:
            try:
                async for token in fake_model_stream(prompt):
                    await queue.put(token)
                    if stop_event.is_set():
                        break
            finally:
                await queue.put(None)  # sentinel

    prod_task = asyncio.create_task(producer())

    buffer = []
    sent_any = False
    while True:
        # se o cliente desconectar, sinaliza produtor e cancela
        if await request.is_disconnected():
            stop_event.set()
            prod_task.cancel()
            break

        try:
            token = await asyncio.wait_for(queue.get(), timeout=5.0)
        except asyncio.TimeoutError:
            # envia keep-alive ocasionalmente para manter conexão
            yield "data: {}\n\n".format(json.dumps({"keepalive": True}))
            continue

        if token is None:
            # fim do stream
            if buffer:
                data = ''.join(buffer)
                yield "data: {}\n\n".format(json.dumps({"text": data}))
            break

        buffer.append(token)
        # envia quando buffer atingir chunk_size
        if len(buffer) >= chunk_size:
            data = ''.join(buffer)
            yield "data: {}\n\n".format(json.dumps({"text": data}))
            buffer = []
            sent_any = True

    # cleanup
    if not prod_task.done():
        prod_task.cancel()
    return

@app.post('/stream')
async def stream_endpoint(request: Request):
    body = await request.json()
    prompt = body.get('prompt', '')
    if not prompt:
        raise HTTPException(status_code=400, detail='prompt required')

    headers = { 'Content-Type': 'text/event-stream', 'Cache-Control': 'no-cache' }
    return StreamingResponse(stream_generator(prompt, request, chunk_size=8), headers=headers)

O código acima provê um endpoint POST /stream que aceita JSON {"prompt": "..."} e responde com SSE. Cada evento tem campo data: com JSON que contém {"text": "..."}.

Cliente simples em Python que consome SSE

Um consumidor simples usando httpx em modo streaming lê linhas e reconstrói eventos SSE. Útil para testes ou scripts automação.

import httpx

def read_sse_stream(url: str, prompt: str):
    with httpx.stream('POST', url, json={'prompt': prompt}, timeout=None) as r:
        for line in r.iter_lines():
            if not line:
                continue
            s = line.decode('utf-8') if isinstance(line, bytes) else line
            if s.startswith('data: '):
                payload = s[len('data: '):]
                try:
                    obj = json.loads(payload)
                    if 'text' in obj:
                        print(obj['text'], end='', flush=True)
                except Exception:
                    pass

# Uso: read_sse_stream('http://localhost:8000/stream', 'hello world this is a test')

O cliente lê linhas, filtra as que começam com 'data: ' e decodifica payload JSON. Em produção, prefira bibliotecas SSE ou web frameworks no frontend que já tratam eventos.

4) Testes e validacao

Testes unitários e integração (pytest + httpx AsyncClient)

Automatize validação do streaming com testes assíncronos: verifique que o stream produz eventos parciais e final.

import pytest
import asyncio
from fastapi.testclient import TestClient

client = TestClient(app)

def test_stream_basic():
    resp = client.post('/stream', json={'prompt': 'a b c d e f g h'})
    assert resp.status_code == 200
    # o corpo é streaming; TestClient consome tudo automaticamente
    text = resp.text
    assert '"text"' in text

Para testes assíncronos mais realistas (sem bloquear), use httpx.AsyncClient e verifique fragmentos conforme chegam. Mocking do model também é recomendado para controlar latências.

Benchmark e validação de performance

Para validar throughput e latência, rode benchs que simulem concorrência usando asyncio.gather para disparar várias requisições simultâneas. Medir tempo até o primeiro evento (TTFP) é essencial para UX.

Exemplo rápido (load test mínimo):

import time
import asyncio
import httpx

async def single_call(client, url):
    start = time.perf_counter()
    async with client.stream('POST', url, json={'prompt': 'x ' * 100}, timeout=None) as r:
        # time until first chunk
        async for line in r.aiter_lines():
            if line and line.startswith('data: '):
                return time.perf_counter() - start
    return None

async def run_load(n):
    async with httpx.AsyncClient() as client:
        tasks = [single_call(client, 'http://localhost:8000/stream') for _ in range(n)]
        return await asyncio.gather(*tasks)

# asyncio.run(run_load(10))

Esse benchmark simples retorna tempo até o primeiro chunk para N clientes paralelos. Ajuste INFER_SEMAPHORE para equilibrar latência e throughput.

5) Consideracoes de producao

Escalabilidade e recursos

Em produção, considere:

• Semáforos para limitar inferências simultâneas por instância (evita OOM e saturação GPU/CPU).
• Pool de workers ou filas externas (RabbitMQ, Redis streams, Celery) quando inferências são pesadas e podem ser enfileiradas.
• Autoscaling baseado em métricas de latência e fila – prefira HPA configurado por custom metrics (GPU utilization, queue length).

O padrão apresentado é ótimo para latência e UX, mas em cargas muito altas combine com arquitetura assíncrona de filas para smoothing.

Observabilidade, timeouts e retries

Instrumente com métricas e logs estruturados:

• TTFP (time to first partial output)
• Tempo total por inferência
• Taxa de cancelamento (cliente desconectou)

Implemente timeouts para produtores e consumidores. Em deploys com GPU, monitore memória e introduza circuit-breakers para evitar que falhas em modelos degradem serviço.

Erros comuns, armadilhas e dicas práticas

Listo os problemas que vejo com frequência:

1. Não verificar desconexão: se você não checar request.is_disconnected(), continuará produzindo tokens inutilmente.
2. Chunk muito pequeno: envia muitos eventos pequenos, aumentando overhead e latência de rede.
3. Chunk muito grande: perde a vantagem de streaming; o usuário espera muito.
4. Sem controle de concorrência: modelos saturam memória/CPU/GPU em picos de tráfego.
5. Não limpar tasks: esqueça de cancelar tarefas produtoras e vazamentos acontecem.

Use chunk_size adaptativo: inicie com um valor conservador (4-16 tokens) e ajuste com base no tamanho médio do token e latência de rede. Combine com keep-alive SSE para conexões longas.

Dica prática: implemente um mecanismo de backpressure simples — se a queue ficar maior que X, reduza a taxa de aceitação de novas requisições ou faça enfileiramento externo.

Conclusao prática

Resumo das decisões

Streaming SSE + asyncio.Queue + chunking oferece uma solução balanceada entre UX e eficiência. O padrão:

• Cria uma tarefa produtora que injeta tokens em uma fila.
• Endpoint consome a fila, agrupa tokens e emite SSE com JSON.
• Semáforos controlam concorrência; request.is_disconnected() permite cancelamento.

Esse modelo reduz latência percebida, controla uso de recursos e melhora experiência em interfaces interativas.

Próximos passos recomendados

Para evoluir em produção, considere:

• Substituir simulador por SDK do modelo com streaming nativo.
• Adotar filas externas para smoothing de picos.
• Implementar métricas TTFP e alertas baseados em percentis.

Com essas práticas você terá um serviço mais resiliente, previsível e com boa UX para inferência em tempo real.

Recap rápido

Implementar inferência em streaming com FastAPI exige atenção a concorrência, chunking e cancelamento. Use asyncio primitives (Queue, Semaphore, Event), StreamingResponse com SSE e valide com testes de integração e benchmarks. Evite buffering completo e garanta limpeza de tasks para prevenir leaks.

Pronto para testar: copie os exemplos, substitua fake_model_stream pelo seu gerador real e ajuste semáforo/chunk_size conforme capacidade do host. Em seguida, rode benchmarks e monitore TTFP para calibrar parâmetros.