Suy luận trên máy chủ web

Máy chủ web là một hệ thống chờ các yêu cầu và xử lý chúng khi chúng đến. Điều này có nghĩa là bạn có thể sử dụng Pipeline như một bộ máy suy luận trên máy chủ web, vì bạn có thể dùng một trình lặp (tương tự cách bạn lặp qua một tập dữ liệu) để xử lý từng yêu cầu đến.

Việc thiết kế một web server với Pipeline là điều đặc biệt vì về cơ bản chúng khác nhau. Web server được đa hợp (đa luồng, bất đồng bộ, v.v.) để xử lý nhiều yêu cầu đồng thời. Ngược lại, Pipeline và mô hình bên dưới nó không được thiết kế cho tính song song vì chúng tiêu tốn nhiều bộ nhớ. Tốt nhất là cấp cho một Pipeline tất cả các tài nguyên khả dụng khi chúng đang chạy hoặc cho một tác vụ tính toán nặng.

Hướng dẫn này cho thấy cách khắc phục sự khác biệt này bằng cách sử dụng một máy chủ web để xử lý tải nhẹ hơn của việc nhận và gửi yêu cầu, và có một luồng duy nhất để xử lý tải nặng hơn của việc chạy Pipeline.

Tạo một máy chủ

Starlette là một framework nhẹ để xây dựng máy chủ web. Bạn có thể sử dụng bất kỳ framework nào khác, nhưng có thể sẽ phải thực hiện một số thay đổi đối với mã bên dưới.

Trước khi bắt đầu, hãy đảm bảo rằng Starlette và uvicorn đã được cài đặt.

!pip install starlette uvicorn

Bây giờ bạn có thể tạo một web server đơn giản trong tệp server.py. Điểm mấu chốt là chỉ tải mô hình một lần để ngăn các bản sao không cần thiết của nó làm tiêu tốn bộ nhớ.

Tạo một pipeline để điền vào token bị che, [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))

Khởi động máy chủ bằng lệnh sau.

uvicorn server:app

Truy vấn máy chủ bằng một yêu cầu POST.

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

Điều này nên trả về kết quả bên dưới.

[{'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.'}]

Xếp hàng các yêu cầu

Cơ chế xếp hàng của máy chủ có thể được sử dụng cho một số ứng dụng thú vị như dynamic batching. Dynamic batching sẽ tích lũy một vài yêu cầu trước khi xử lý chúng bằng Pipeline.

Ví dụ bên dưới được viết bằng mã giả để dễ đọc hơn là tối ưu hiệu suất, cụ thể, bạn sẽ nhận thấy rằng:

Không có giới hạn về kích thước lô (batch size).
Timeout được đặt lại ở mỗi lần lấy hàng đợi, do đó bạn có thể phải chờ lâu hơn nhiều so với giá trị timeout trước khi xử lý một yêu cầu. Điều này cũng sẽ làm trì hoãn yêu cầu suy luận đầu tiên bằng đúng khoảng thời gian đó. Máy chủ web luôn chờ 1ms ngay cả khi hàng đợi trống, điều này không hiệu quả vì khoảng thời gian đó có thể được dùng để bắt đầu suy luận. Tuy nhiên, điều này có thể hợp lý nếu việc nhóm (batching) là thiết yếu cho trường hợp sử dụng của bạn.

Tốt hơn là nên có một hạn chót 1ms duy nhất, thay vì đặt lại nó ở mỗi lần fetch, như minh họa bên dưới.

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)

Kiểm tra lỗi

Có rất nhiều điều có thể xảy ra sai trong môi trường production. Bạn có thể gặp lỗi out-of-memory, hết dung lượng, không tải được model, cấu hình model không chính xác, truy vấn không chính xác, và còn rất nhiều vấn đề khác.

Việc thêm các câu lệnh try...except sẽ hữu ích để trả về các lỗi này cho người dùng nhằm phục vụ việc gỡ lỗi. Lưu ý rằng điều này có thể là một rủi ro bảo mật nếu bạn không nên tiết lộ một số thông tin nhất định.

Ngắt mạch

Hãy cố gắng trả về lỗi 503 hoặc 504 khi máy chủ quá tải thay vì bắt người dùng chờ vô hạn.

Việc triển khai các loại lỗi này tương đối đơn giản vì chỉ có một hàng đợi duy nhất. Hãy xem xét kích thước hàng đợi để xác định thời điểm bắt đầu trả về lỗi trước khi máy chủ của bạn gặp sự cố do quá tải.

Chặn luồng chính

PyTorch không nhận thức được tính bất đồng bộ, do đó việc tính toán sẽ chặn luồng chính chạy.

Vì lý do này, tốt hơn là nên chạy PyTorch trên một thread hoặc process riêng biệt. Khi quá trình suy luận của một yêu cầu duy nhất đặc biệt dài (hơn 1 giây), điều này càng quan trọng hơn vì nó có nghĩa là mọi truy vấn trong suốt quá trình suy luận đều phải chờ 1 giây trước khi thậm chí nhận được một lỗi.

Đóng gói động

Dynamic batching có thể rất hiệu quả khi được sử dụng trong đúng ngữ cảnh, nhưng không cần thiết khi bạn chỉ truyền 1 request mỗi lần (xem batch inference để biết thêm chi tiết).

Update on GitHub