Servidor de embeddings com FastAPI, Transformers e FAISS: batching e quantização

Contexto e motivacao

Criar um servidor de embeddings é tarefa comum em projetos de re-ranking, busca semântica e RAG. A implementação simples — um endpoint que recebe texto, chama o modelo e retorna um vetor — funciona em protótipos, mas escala mal. Latência, custos de GPU/CPU e consultas concorrentes derrubam a solução quando a carga aumenta. Este artigo adota tom de code review: primeiro mostro o jeito errado (anti-pattern), depois demonstro a correção com padrões robustos.

Por que modelar corretamente

Sem batching e sem controle de concorrência você paga muito mais por GPU, perde throughput por overhead de chamadas, e enfrenta picos de latência. Além disso, indexação sem quantização exige memória linear ao tamanho do dataset. Quantização reduz memória e melhora custo sem sacrificar drasticamente a qualidade.

Requisitos deste artigo

Vou assumir experiência básica em Python e FastAPI. Vamos usar transformers (HuggingFace), FAISS e FastAPI. O foco é patterns, trade-offs e exemplos funcionais: endpoints corretos, batching, fila assíncrona e criação de índice FAISS com quantização (PQ/OPQ + IVF).

Como funciona por baixo dos panos

Para otimizar um servidor de embeddings precisamos entender três componentes essenciais: o modelo (compute), a orquestração do trabalho (batching e filas) e o indexador (FAISS). Cada camada tem custos e características próprias.

Modelo: inferencia e custo

Transformers geram embedding extraindo estados ocultos e agregando tokens (mean pooling, CLS, etc.). A inferência é tipicamente CPU-bound quando em FP32 e I/O-bound por chamadas a hardware. Grandes lotes (batches) amortizam o custo de preparação de tensores e chamadas para GPU, reduzindo latência média por item.

Orquestração: por que batching importa

Batching aumenta throughput ao reduzir overhead por item. Mas tamanhos de batch muito grandes aumentam latência tail e consumo de memória. É preciso um mecanismo que agrupe requisições próximas no tempo, limite o tempo máximo de espera e proteja o backend via circuit-breaker e limitação.

Implementacao passo a passo

Aqui mostro um anti-pattern rápido e depois a versão correta com fila assíncrona, worker de batching e uso de ThreadPool/ProcessPool para operações CPU-bound. Também incluo geração de index FAISS com quantização.

Anti-pattern: endpoint síncrono que chama o modelo por requisição

Contexto: um dev novo implementa o endpoint assim. Parece simples, mas é uma bomba de performance sob carga.

from fastapi import FastAPI, Request
from transformers import AutoTokenizer, AutoModel
import torch

app = FastAPI()

tokenizer = AutoTokenizer.from_pretrained('sentence-transformers/all-MiniLM-L6-v2')
model = AutoModel.from_pretrained('sentence-transformers/all-MiniLM-L6-v2')

@app.post('/embed')
async def embed(req: Request):
    data = await req.json()
    text = data['text']
    inputs = tokenizer(text, return_tensors='pt')
    with torch.no_grad():
        out = model(**inputs)
    # mean pooling simplificado
    emb = out.last_hidden_state.mean(dim=1).squeeze().cpu().tolist()
    return {'embedding': emb}

O que acontece: cada requisição faz tokenização e inferência isoladas, sem aproveitar batching. Em alta concorrência, a GPU não é usada eficientemente e a fila do servidor explode.

Correção: fila assíncrona e worker de batching

Contexto: vamos criar um componente que recebe requisições, empacota entradas em lotes e executa inferência em batch. A API responde com promessa (futuro) quando o batch for processado.

import asyncio
from concurrent.futures import ThreadPoolExecutor
from fastapi import FastAPI, Request
from transformers import AutoTokenizer, AutoModel
import torch

app = FastAPI()

tokenizer = AutoTokenizer.from_pretrained('sentence-transformers/all-MiniLM-L6-v2')
model = AutoModel.from_pretrained('sentence-transformers/all-MiniLM-L6-v2').eval()

# Parametros de batching
MAX_BATCH = 32
MAX_WAIT_MS = 20

queue = asyncio.Queue()
executor = ThreadPoolExecutor(max_workers=2)

async def batch_worker():
    while True:
        # coleta itens com timeout para formar um batch
        batch = []
        first = await queue.get()
        batch.append(first)
        t_start = asyncio.get_event_loop().time()
        while len(batch) < MAX_BATCH:
            timeout = MAX_WAIT_MS / 1000 - (asyncio.get_event_loop().time() - t_start)
            if timeout <= 0:
                break
            try:
                item = await asyncio.wait_for(queue.get(), timeout)
                batch.append(item)
            except asyncio.TimeoutError:
                break
        # Processa batch em thread pool (tokenizacao e inferencia sao CPU/GPU-bound)
        texts, futs = zip(*batch)
        loop = asyncio.get_event_loop()
        emb_batch = await loop.run_in_executor(executor, compute_embeddings, list(texts))
        for fut, emb in zip(futs, emb_batch):
            fut.set_result(emb)

def compute_embeddings(texts):
    # executa tokenizacao e inferencia em lote (rodando em thread para nao bloquear loop)
    with torch.no_grad():
        inputs = tokenizer(texts, padding=True, truncation=True, return_tensors='pt')
        out = model(**inputs)
        emb = out.last_hidden_state.mean(dim=1).cpu()
        return emb.tolist()

@app.on_event('startup')
async def startup():
    asyncio.create_task(batch_worker())

@app.post('/embed')
async def embed(req: Request):
    data = await req.json()
    text = data['text']
    loop = asyncio.get_event_loop()
    fut = loop.create_future()
    await queue.put((text, fut))
    emb = await fut
    return {'embedding': emb}

O que acontece: requisições são agrupadas em batches de até MAX_BATCH ou até MAX_WAIT_MS milissegundos. A função compute_embeddings roda em executor, permitindo que a event loop processe outras requisições. Latência média cai, throughput aumenta.

Detalhes de produção: limites e proteção

Precisamos adicionar timeouts, circuit breaker, limites de tamanho de payload e policies de retry. Também convém expor métricas (prometheus) para monitorar tamanho de batch e latência tail.

Patterns e boas praticas

Vou listar padrões e depois exemplificar a aplicação prática com FAISS e quantização. Mantenha a regra: medir antes de otimizar. Use load testing com cenários realistas.

Pattern: batching com deadline e fallback

Recomendo aplicar um deadline por requisição. Se exceder, você escolhe entre retornar erro, responder com fallback (embedding aproximado), ou priorizar requisição e enviar sem batching. Implementar esse fallback evita clientes travados e garante SLOs.

Pattern: separa inferencia e indexacao

Mantenha o servidor de embeddings separado do serviço que indexa ou consulta FAISS. Assim, escalamos CPU/GPU e indexador de forma independente. Utilize filas (Kafka/RabbitMQ) para ingestion em massa e workers que alimentam FAISS off-line ou em batches.

Boa prática: mensure throughput por GPU, latência 95/99, taxa de ocupação do batch. Não confie apenas em latência média.

Criando índice FAISS com quantização

Contexto: FAISS oferece várias estratégias: IndexFlat para baselines, IVFPQ para combinação de particionamento e Product Quantization, e OPQ para rotacionar vetores antes da quantização. Abaixo um pipeline para IVF + PQ com OPQ. Esse código cria e treina o índice com float32 arrays.

import numpy as np
import faiss

# dados de exemplo (num_vectors, dim)
xb = np.random.random((10000, 384)).astype('float32')

# parametros
d = xb.shape[1]
nlist = 256  # numero de centroides
m = 16       # numero de subquantizers (m*8 bits = 128 bits total)
nbits = 8

# cria OPQ matrix para rotacao (opcional, melhora densidade)
opq_matrix = faiss.OPQMatrix(d, m)

# cria index IVF + PQ: primeiro builda o quantizer (IndexFlatL2)
quantizer = faiss.IndexFlatL2(d)
index = faiss.IndexIVFPQ(quantizer, d, nlist, m, nbits)

# treina OPQ e index
opq_matrix.train(xb)
xb_opq = opq_matrix.apply_py(xb)
index.train(xb_opq)
index.add(xb_opq)

# salva index + opq
axis = faiss.IndexIDMap2(index)
faiss.write_index(opq_matrix, 'opq_matrix.bin')
faiss.write_index(index, 'index_ivfpq.bin')

O que acontece: OPQ melhora a compactacao ao rotacionar os vetores antes da PQ; IVFPQ combina clustering e PQ para busca rápida em grandes coleções. Treine com dados representativos (não com amostras triviais).

Consulta: busca aproximada e re-ranking

Contexto: após indexar com PQ, você consulta convertendo o vetor em OPQ space e usando search. Recomendo re-ranking com coseno exato nos top-k retornados.

def search_topk(query_vec, k=10):
    q = query_vec.astype('float32')
    # aplicar OPQ
    q_opq = opq_matrix.apply_py(q.reshape(1, -1))
    D, I = index.search(q_opq, k)
    # re-rank com calculo exato (se necessario) usando vetores originais
    return D[0], I[0]

O que acontece: search retorna ids aproximados rapidamente. Se for necessario, extraia vetores originais dos ids e compute uma similaridade exata nos top-k para garantir qualidade.

Armadilhas comuns e como evitar

Aqui mostro erros que aparecem em code reviews e como corrigi-los. Listarei armadilhas práticas e medidas para evitá-las.

1) Fazer tokenizacao por requisição (sem batching)

Problema: tokenizacao em massa em threads concorrentes consome CPU e memória, gerando contenção e GC. Solução: empacotar strings e tokenizar em lote dentro do worker, como demonstrei no pattern correto. Além disso, reutilize o tokenizer e evite re-carregar o modelo por request.

2) Usar async/await diretamente para tarefas CPU bound

Problema: colocar compute pesado diretamente no loop async bloqueia o event loop. Solução: delegue tarefas para ThreadPoolExecutor ou ProcessPoolExecutor. Para operações com PyTorch que fazem uso de GPU, rodar em thread é aceitável, pois a chamada ao CUDA libera o GIL e a thread permite que o loop siga. Para CPU puro (numpy, faiss CPU), use ProcessPool para evitar GIL.

3) Indexar em tempo real sem strategy

Problema: adicionar vetores um a um em índices grandes (IVF) frequentemente requer re-treinamento ou muda a performance. Solução: ingestion em batches e treinamento offline periódica. Se for necessário atualizar online, prefira estruturas que suportam add sem re-treino (IndexFlat) para dados pequenos, ou mantenha um buffer de novos vetores e faça merges regulares.

• Monitoramento: trackeie tamanho do batch, latência P50/P95/P99 e utilização de GPU/CPU.
• Configuração: experimente MAX_BATCH entre 8-64 dependendo da memória e do modelo.
• Quantização: PQ com 8 bits por subquantizer (nbits=8) costuma ser um bom trade-off.

Regra prática: se o custo de CPU/GPU por chamada for alto, agregar é grátis; se a latência máxima é crítica, prefira ajustar deadline e fallback.

Conclusao e checklist prático

Resumo direto: evite endpoints síncronos que chamam o modelo por requisição. Use batching com uma fila assíncrona, execute inferência em executor apropriado, separe indexador e serviço de inferência, e use FAISS com quantização para reduzir custo de memória. Meça antes e depois.

Checklist para deploy

1. Implementar fila e worker de batching com MAX_BATCH e MAX_WAIT_MS configuráveis.
2. Executar tokenizacao e inferencia em thread/process pool (dependendo do workload).
3. Separar streaming de indexacao: ingestion -> buffer -> treinamento do índice offline.
4. Testar IVFPQ/OPQ com dados reais; ajustar nlist, m e nbits conforme recall e memória.
5. Adicionar métricas (batch size, latência p95/p99, erros) e circuit-breaker.

Conclusao pratica

Implementar um servidor de embeddings eficiente é, em grande parte, engenharia de orquestração: batching, deadlines, execução adequada de tarefas CPU/GPU e indexação eficiente via FAISS com quantização. A diferença entre um protótipo e um serviço de produção está em detalhes: como você agrupa requisições, protege o loop, e organiza a indexação. As snippets e patterns acima formam um conjunto mínimo viável para mover do anti-pattern para uma base escalável.

Recomendo iniciar com medições simples, adotar batching com fila e depois aplicar quantização no índice. Isso reduz custos e melhora throughput sem sacrificar qualidade de forma brusca.