Inferência em servidor web

Um servidor web é um sistema que aguarda solicitações e as atende conforme chegam. Isso significa que você pode usar o Pipeline como um mecanismo de inferência em um servidor web, pois é possível usar um iterador (de forma semelhante a como você iteraria sobre um conjunto de dados) para processar cada solicitação recebida.

Projetar um servidor web com Pipeline é único, pois eles são fundamentalmente diferentes. Servidores web são multiplexados (multithreaded, async, etc.) para lidar com múltiplas solicitações simultaneamente. Pipeline e o modelo subjacente, por outro lado, não foram projetados para paralelismo, pois consomem muita memória. O ideal é fornecer a um Pipeline todos os recursos disponíveis quando ele estiver em execução ou para uma tarefa intensiva em computação.

Este guia mostra como contornar essa diferença usando um servidor web para lidar com a carga mais leve de receber e enviar solicitações, e uma única thread para lidar com a carga mais pesada de executar Pipeline.

Criar um servidor

Starlette é um framework leve para a construção de servidores web. Você pode usar qualquer outro framework que desejar, mas talvez precise fazer algumas alterações no código abaixo.

Antes de começar, verifique se o Starlette e o uvicorn estão instalados.

!pip install starlette uvicorn

Agora você pode criar um servidor web simples em um arquivo server.py. O ponto-chave é carregar o modelo apenas uma vez para evitar que cópias desnecessárias dele consumam memória.

Crie um pipeline para preencher o token mascarado, [MASK].

from starlette.applications import Starlette
from starlette.responses import JSONResponse
from starlette.routing import Route
from transformers import pipeline
import asyncio

async def homepage(request):
    payload = await request.body()
    string = payload.decode("utf-8")
    response_q = asyncio.Queue()
    await request.app.model_queue.put((string, response_q))
    output = await response_q.get()
    return JSONResponse(output)

async def server_loop(q):
    pipe = pipeline(task="fill-mask",model="google-bert/bert-base-uncased")
    while True:
        (string, response_q) = await q.get()
        out = pipe(string)
        await response_q.put(out)

app = Starlette(
    routes=[
        Route("/", homepage, methods=["POST"]),
    ],
)

@app.on_event("startup")
async def startup_event():
    q = asyncio.Queue()
    app.model_queue = q
    asyncio.create_task(server_loop(q))

Inicie o servidor com o seguinte comando.

uvicorn server:app

Consulte o servidor com uma requisição POST.

curl -X POST -d "Paris is the [MASK] of France." http://localhost:8000/

Isso deve retornar a saída abaixo.

[{'score': 0.9969332218170166,
  'token': 3007,
  'token_str': 'capital',
  'sequence': 'paris is the capital of france.'},
 {'score': 0.0005914849461987615,
  'token': 2540,
  'token_str': 'heart',
  'sequence': 'paris is the heart of france.'},
 {'score': 0.00043787318281829357,
  'token': 2415,
  'token_str': 'center',
  'sequence': 'paris is the center of france.'},
 {'score': 0.0003378340043127537,
  'token': 2803,
  'token_str': 'centre',
  'sequence': 'paris is the centre of france.'},
 {'score': 0.00026995912776328623,
  'token': 2103,
  'token_str': 'city',
  'sequence': 'paris is the city of france.'}]

Fila de solicitações

O mecanismo de fila do servidor pode ser usado em algumas aplicações interessantes, como o loteamento dinâmico. O loteamento dinâmico acumula várias solicitações antes de processá-las com Pipeline.

O exemplo abaixo está escrito em pseudocódigo para facilitar a leitura, em vez de priorizar o desempenho. Em particular, você notará que:

1. Não há limite de tamanho de lote.

2. O timeout é reiniciado a cada busca na fila, portanto, você pode acabar esperando muito mais tempo que o valor de timeout antes de processar uma solicitação. Isso também atrasaria a primeira solicitação de inferência por esse mesmo período. O servidor web sempre espera 1 ms, mesmo que a fila esteja vazia, o que é ineficiente, pois esse tempo poderia ser usado para iniciar a inferência. No entanto, isso pode fazer sentido se o batching for essencial para o seu caso de uso.

   Seria melhor ter um único prazo de 1 ms, em vez de redefini-lo a cada busca, conforme mostrado abaixo.

async def server_loop(q):
    pipe = pipeline(task="fill-mask", model="google-bert/bert-base-uncased")
    while True:
        (string, rq) = await q.get()
        strings = []
        queues = []
        strings.append(string)
        queues.append(rq)
        while True:
            try:
                (string, rq) = await asyncio.wait_for(q.get(), timeout=1)
            except asyncio.exceptions.TimeoutError:
                break
            strings.append(string)
            queues.append(rq)
        outs = pipe(strings, batch_size=len(strings))
        for rq, out in zip(queues, outs):
            await rq.put(out)

Verificação de erros

Muitas coisas podem dar errado em produção. Você pode esgotar a memória, esgotar o espaço em disco, falhar ao carregar um modelo, ter uma configuração de modelo incorreta, ter uma consulta incorreta, e assim por diante.

Adicionar declarações try...except é útil para retornar esses erros ao usuário para depuração. Tenha em mente que isso pode representar um risco de segurança se você não deve revelar certas informações.

Circuit breaking

Tente retornar um erro 503 ou 504 quando o servidor estiver sobrecarregado, em vez de forçar o usuário a aguardar indefinidamente.

É relativamente simples implementar esses tipos de erro, pois se trata de uma única fila. Verifique o tamanho da fila para determinar quando começar a retornar erros antes de o servidor falhar sob carga.

Bloquear a thread principal

O PyTorch não é assíncrono, portanto a computação bloqueará a execução da thread principal.

Por esse motivo, é melhor executar o PyTorch em uma thread ou processo separado. Quando a inferência de uma única solicitação é especialmente longa (mais de 1s), isso se torna ainda mais importante, pois significa que cada consulta durante a inferência deve esperar 1s antes de receber um erro.

Loteamento dinâmico

O loteamento dinâmico pode ser muito eficaz quando usado no contexto adequado, mas não é necessário quando você envia apenas 1 solicitação por vez (veja inferência em lote para mais detalhes).