Từng mẩu một: dữ liệu chảy qua Python như thế nào

Generator dừng giữa chừng thế nào, async generator có hai kiểu tạm dừng ra sao, vì sao finally chạy muộn, backpressure đi từ yield xuống socket, asyncio.Queue, StreamingResponse của FastAPI, và một chatbot trả lời từng token, kèm thí nghiệm chạy thật.

Bài Python chạy đồng thời thế nào tự tay "lái" một coroutine bằng send() và viết một event loop 40 dòng. Bài Bên trong CPython xem hàm, frame và bytecode. Bài này ghép hai thứ đó lại quanh một từ khóa: yield. Nó là thứ đứng sau generator, coroutine, async generator, và cả cách FastAPI gửi một file 1 GB hay từng token của chatbot về trình duyệt mà không cần giữ hết trong RAM.

Lộ trình: vì sao cần stream, yield làm gì bên trong CPython, async generator có hai kiểu tạm dừng ra sao, chỗ nguy hiểm nhất là dọn dẹp, backpressure đi từ yield xuống tận socket, asyncio.Queue, rồi lên StreamingResponse của FastAPI, ghép tất cả lại thành một chatbot trả lời từng token, và đo đạc.

Mọi con số trong bài đều đo trên máy thật: Linux (WSL2), 12 nhân, Python 3.12 và 3.14, FastAPI 0.143, Starlette 1.8, uvicorn 0.54; ví dụ chatbot dùng transformers 5.20 và PyTorch 2.14 chạy bằng CPU.

Trả về một triệu dòng

Một endpoint xuất CSV một triệu dòng (khoảng 45 MB). Cách thứ nhất dựng cả file trong bộ nhớ rồi trả về. Cách thứ hai trả từng mẻ 1.000 dòng ngay khi có:

from fastapi import FastAPI
from fastapi.responses import PlainTextResponse, StreamingResponse

app = FastAPI()


def row(i: int) -> str:
    return f"{i},user{i},user{i}@example.com,{i * 7 % 1000}\n"


@app.get("/csv/list")
async def csv_list(n: int = 1_000_000):
    return PlainTextResponse("".join([row(i) for i in range(n)]))  # the whole file in memory first


async def rows(n: int, batch: int):
    buf = []
    for i in range(n):
        buf.append(row(i))
        if len(buf) == batch:
            yield "".join(buf)  # hand one chunk out, then pause here
            buf.clear()
    if buf:
        yield "".join(buf)


@app.get("/csv/stream")
async def csv_stream(n: int = 1_000_000, batch: int = 1000):
    return StreamingResponse(rows(n, batch), media_type="text/csv")

Mỗi lần đo chạy một server mới, rồi đọc bộ nhớ cao nhất của process (VmHWM trong /proc/<pid>/status):

Cách trả Byte đầu tiên tới sau Xong sau Bộ nhớ cao nhất của server
Dựng cả file rồi trả 371–1.142 ms 0,40–1,17 s 192 MiB
Stream từng mẻ 1.000 dòng 9–11 ms 0,32–0,36 s 51 MiB

Server lúc rảnh dùng 50 MiB. Bản dựng cả file tốn thêm khoảng 140 MiB cho một request: một list một triệu chuỗi, cộng chuỗi ghép, cộng bản đã mã hóa thành byte. Mười người tải cùng lúc là hơn 1 GB. Bản stream gần như không tốn thêm gì, người dùng nhận byte đầu tiên sau 10 ms, và tổng thời gian còn ngắn hơn vì việc tạo dữ liệu và việc gửi đi diễn ra xen kẽ nhau.

Tất cả nhờ một từ khóa: yield.

yield làm gì bên trong CPython

Một hàm có yield ở bất kỳ đâu trong thân thì gọi nó không chạy dòng nào cả; nó trả về một generator:

import inspect


def numbers():
    total = 0
    for n in range(3):
        total += n
        sent = yield n
        print("  got back:", sent)
    return total


g = numbers()
print(inspect.getgeneratorstate(g))
print("first:", next(g))
print(inspect.getgeneratorstate(g), g.gi_frame.f_locals)
print("send:", g.send("hello"))
print("locals now:", g.gi_frame.f_locals)
try:
    g.send(None); g.send(None)
except StopIteration as stop:
    print("finished, return value:", stop.value)
GEN_CREATED
first: 0
GEN_SUSPENDED {'total': 0, 'n': 0}
  got back: hello
send: 1
locals now: {'total': 1, 'n': 1, 'sent': 'hello'}
  got back: None
  got back: None
finished, return value: 3

Bài trước có nói: mỗi lần gọi hàm, CPython tạo một frame chứa biến cục bộ và vị trí lệnh đang chạy, và frame thường biến mất khi hàm trả về. Generator khác ở đúng một điểm: frame của nó nằm bên trong object generator, trên heap, nên sống qua nhiều lần chạy. Bytecode cho thấy rõ:

  6           RETURN_GENERATOR       # create the generator, move the frame into it, return it
              POP_TOP
      L1:     RESUME                   0
  ...
 10           LOAD_FAST_BORROW         1 (n)
              YIELD_VALUE              0   # hand n out, remember this spot, leave the frame
              RESUME                   5   # the next send() continues here
              STORE_FAST               2 (sent)

next(g) hay g.send(x) chạy frame tới YIELD_VALUE kế tiếp. Lệnh đó lưu vị trí, trả giá trị ra ngoài và rời frame mà không hủy nó. Lần sau, giá trị được send vào trở thành kết quả của biểu thức yield. Khi hàm return, CPython ném StopIteration mang giá trị trả về; đó là cách vòng for biết đã hết. Một generator chỉ tốn 232 byte trên 3.14 (224 trên 3.12), nên giữ hàng nghìn cái đang dừng cũng không đáng kể.

Ngoài send, generator còn hai cách nữa để người bên ngoài can thiệp:

  • g.throw(exc) ném một exception vào đúng chỗ yield đang dừng.
  • g.close() ném GeneratorExit vào chỗ đó, để các khối finally và with bên trong chạy.
def lines():
    try:
        yield "a"
        yield "b"
    finally:
        print("  finally: cleanup")


g = lines()
next(g)
g = None  # nobody holds it any more
print("--- after the drop")
  finally: cleanup
--- after the drop

Không ai gọi close(), nhưng finally vẫn chạy ngay: khi bộ đếm tham chiếu về 0, CPython tự đóng generator. Hãy nhớ chi tiết này, vì với async generator thì nó không còn đúng.

Generator, coroutine, async generator: cùng một cỗ máy

Bài đầu của series đã lái một coroutine bằng send(None). Đó không phải trùng hợp: coroutine là generator được đổi tên. Ba loại hàm chỉ khác nhau ở một cờ trong code object:

Viết Cờ trong co_flags Gọi ra
def có yield CO_GENERATOR generator: dùng với for
async def CO_COROUTINE coroutine: dùng với await
async def có yield CO_ASYNC_GENERATOR async generator: dùng với async for

Cả ba dùng chung RETURN_GENERATOR, frame trên heap, YIELD_VALUE và RESUME. Bytecode của await x cũng chỉ là một vòng yield quen thuộc:

GET_AWAITABLE                 # get the awaitable's iterator
LOAD_CONST       None
L2: SEND        (to L5)       # pass the value down; if the inner one finished, jump out
L3: YIELD_VALUE               # the inner one paused: pause us too, pass its value up
L4: RESUME
    JUMP_BACKWARD_NO_INTERRUPT (to L2)
L5: END_SEND                  # the result of the await

Đây chính là yield from của generator. Nó giải thích câu "await chuyền yêu cầu ra ngoài qua từng tầng" ở bài 1: mỗi tầng await là một vòng SEND/YIELD_VALUE, chuyền giá trị lên tầng trên cho tới event loop.

Async generator: hai kiểu tạm dừng

Async generator có một điều khó: trong cùng một hàm có hai kiểu tạm dừng, đi về hai hướng khác nhau.

  • await: chờ I/O. Yêu cầu phải đi ra tới event loop.
  • yield: có một giá trị. Nó phải đi tới vòng async for đang đọc.

Cách tốt nhất để thấy hai đường đó là tự tay lái một async generator, không cần asyncio, giống bài 1:

import types


@types.coroutine
def wait_for(what):
    """Pause and hand `what` to whoever drives us (in asyncio: the event loop)."""
    reply = yield what
    return reply


async def rows():
    for page in (1, 2):
        batch = await wait_for(f"SELECT ... PAGE {page}")  # pause kind 1: goes out to the driver
        for row in batch:
            yield row  # pause kind 2: goes to whoever runs `async for`


agen = rows()
step = agen.__anext__()  # an awaitable: "work out the next item"
print("anext() gives:", type(step).__name__)
print("driver got request:", step.send(None))  # runs until the await inside rows()
try:
    step.send(["an", "binh"])  # answer the query: rows() runs on until its yield
except StopIteration as item:
    print("item:", item.value)

step = agen.__anext__()
try:
    step.send(None)  # the next row is already in memory: no request, straight to the yield
except StopIteration as item:
    print("item:", item.value)

step = agen.__anext__()
print("driver got request:", step.send(None))
try:
    step.send([])  # empty page 2: rows() ends
except StopAsyncIteration:
    print("StopAsyncIteration: the async for loop ends")
anext() gives: async_generator_asend
driver got request: SELECT ... PAGE 1
item: an
item: binh
driver got request: SELECT ... PAGE 2
StopAsyncIteration: the async for loop ends

async for row in rows() làm đúng những bước trên: mỗi vòng gọi __anext__() để lấy một awaitable, rồi await nó. Bên trong awaitable đó:

  • rows() gặp await: giá trị đi xuyên qua, lên tới event loop, y như một coroutine thường. Vòng async for không hề thấy.
  • rows() gặp yield: awaitable kết thúc bằng StopIteration(row), nên với vòng async for, await cho ra row.

Vậy làm sao CPython phân biệt được, khi cả hai đều là YIELD_VALUE? Xem bytecode của dòng yield row:

 16           LOAD_FAST_BORROW         2 (row)
              CALL_INTRINSIC_1         4 (INTRINSIC_ASYNC_GEN_WRAP)
              YIELD_VALUE              0

Giá trị của yield được bọc trong một object đặc biệt trước khi đi ra. Object async_generator_asend nhìn thấy lớp bọc thì biết đó là một phần tử cho async for; không có lớp bọc thì đó là việc của event loop, chuyền tiếp lên trên.

Async generator có asend(), athrow() và aclose(), là bản await được của send, throw, close. Hai thứ của generator thường bị cấm ở đây: yield from báo SyntaxError: 'yield from' inside async function, còn return có giá trị báo SyntaxError: 'return' with value in async generator.

Phần nguy hiểm: dọn dẹp

Async generator hay giữ tài nguyên: một kết nối database, một file, một lock. Vậy khi bên đọc dừng giữa chừng (break, exception, client ngắt kết nối), khối finally chạy lúc nào? Thử ba cách viết, trong asyncio:

import asyncio
from contextlib import aclosing


async def rows(name):
    print(f"  [{name}] borrow a DB connection")
    try:
        for i in range(1_000):
            await asyncio.sleep(0)  # stands in for fetching the next row
            yield i
    finally:
        print(f"  [{name}] finally: give the connection back (task: {asyncio.current_task().get_name()})")


async def first_rows_held():
    stream = rows("held")  # a variable keeps the generator alive
    async for row in stream:
        if row == 2:
            break
    print("  handler: done with the rows, now doing slow work...")
    await asyncio.sleep(0.5)
    print("  handler: returns")


async def first_rows_inline():
    async for row in rows("inline"):  # no variable: its refcount drops to 0 at break
        if row == 2:
            break
    print("  handler: after the loop")
    await asyncio.sleep(0)
    print("  handler: after one more await")


async def first_rows_aclosing():
    async with aclosing(rows("aclosing")) as stream:
        async for row in stream:
            if row == 2:
                break
    print("  handler: after the loop")
first_rows_held
  [held] borrow a DB connection
  handler: done with the rows, now doing slow work...
  handler: returns
  [held] finally: give the connection back (task: Task-2)
first_rows_inline
  [inline] borrow a DB connection
  handler: after the loop
  handler: after one more await
  [inline] finally: give the connection back (task: Task-3)
first_rows_aclosing
  [aclosing] borrow a DB connection
  [aclosing] finally: give the connection back (task: first_rows_aclosing)
  handler: after the loop

Ba điều đáng chú ý:

  • held: kết nối bị giữ suốt 0,5 giây handler làm việc khác, chỉ được trả khi handler kết thúc và biến stream biến mất. Với một pool 10 kết nối và 10 request như vậy cùng lúc, request thứ 11 phải đứng chờ.
  • inline: dù không còn ai giữ generator, finally vẫn không chạy ngay như với generator thường. Nó chạy sau, khi handler đã đi tiếp.
  • Cả hai trường hợp, finally chạy trong một task khác (Task-2, Task-3), không phải task của handler.

Lý do: finally của async generator có thể cần await, ví dụ await conn.close(), mà lúc bộ đếm tham chiếu về 0 thì không có event loop nào đang chạy generator đó. CPython không thể await từ bên trong bộ dọn rác. Nên nó cho event loop đăng ký hai hook qua sys.set_asyncgen_hooks():

# asyncio/base_events.py (3.14)
def _asyncgen_firstiter_hook(self, agen):
    ...
    self._asyncgens.add(agen)  # remember every async generator that started

def _asyncgen_finalizer_hook(self, agen):
    self._asyncgens.discard(agen)
    if not self.is_closed():
        self.call_soon_threadsafe(self.create_task, agen.aclose())  # close it later, in a new task

Khi một async generator đang dở dang bị thu hồi, asyncio lên lịch aclose() trong một task mới. Còn khi asyncio.run() kết thúc, nó gọi loop.shutdown_asyncgens() để đóng những cái còn sót. Hệ quả thực tế:

  • Tài nguyên được trả muộn, vào lúc không đoán trước được.
  • Nếu finally dùng thứ gì đó đã bị dọn trước nó (một session database đã đóng cùng request), lỗi xuất hiện ở một task lạ, rất khó lần ra.

Cách sửa là contextlib.aclosing (từ Python 3.10). Nó gọi await agen.aclose() ngay khi ra khỏi khối async with, trong chính task của bạn, như trường hợp thứ ba. Quy tắc: async generator nào giữ tài nguyên thì đọc nó trong async with aclosing(...).

Backpressure: ai bắt ai chờ

Generator chỉ chạy khi có người hỏi phần tử kế tiếp. Nhờ vậy, stream có sẵn một cơ chế tự hãm, gọi là backpressure: bên đọc chậm thì bên tạo dữ liệu tự chậm theo. Nhưng nó chỉ hoạt động nếu mọi tầng ở giữa đều chịu chờ.

asyncio Streams và drain()

Tầng thấp nhất trong asyncio mà bạn dùng trực tiếp được là Streams: asyncio.open_connection() và asyncio.start_server() trả về một cặp StreamReader/StreamWriter. Bài 1 dùng chúng để viết client. Điểm cần biết: writer.write() không bao giờ chờ. Nó chỉ chép dữ liệu vào bộ đệm của transport rồi trả về ngay. Muốn chờ, phải gọi await writer.drain().

Thử một server gửi 200 MiB cho một client chưa đọc gì trong giây đầu:

CHUNK = b"x" * 64 * 1024


async def serve(reader, writer, use_drain):
    for i in range(3_200):  # 200 MiB
        writer.write(i.to_bytes(4) + CHUNK[4:])  # never waits: it only appends to the buffer
        if use_drain:
            await writer.drain()  # waits while the buffer is above the high-water mark
    await writer.drain()
    writer.close()


async def slow_client(port):
    reader, writer = await asyncio.open_connection("127.0.0.1", port)
    await asyncio.sleep(1)  # the client doesn't read for a second
    while await reader.read(1 << 20):
        pass
Dữ liệu nằm chờ trong bộ đệm (cao nhất) Bộ nhớ cao nhất
Không drain() 197,5 MiB 218 MiB
Có drain() 0,1 MiB 21 MiB

Không có drain(), server dồn cả 200 MiB vào RAM trong lúc client chưa đọc. Có drain(), bộ đệm không bao giờ vượt quá vài chục KiB. writer.transport.get_write_buffer_limits() trả về (16384, 65536): khi bộ đệm vượt 64 KiB, transport gọi pause_writing() và drain() bắt đầu chờ; khi nó xuống dưới 16 KiB, resume_writing() cho chạy tiếp.

Từ yield xuống socket

uvicorn làm đúng như vậy. Mỗi lần ứng dụng gọi send() để gửi một mẩu body, nó kiểm tra cờ này trước:

# uvicorn/protocols/http/httptools_impl.py
async def send(self, message):
    if self.flow.write_paused and not self.disconnected:
        await self.flow.drain()  # the transport said its buffer is over 64 KiB: wait
    ...

Ghép lại, đây là chuỗi chờ của một StreamingResponse:

  1. Client đọc chậm, nên bộ đệm nhận của nó trong kernel đầy.
  2. Bộ đệm gửi của server trong kernel cũng đầy.
  3. Bộ đệm của transport vượt 64 KiB, uvicorn đặt cờ write_paused.
  4. await send(...) của Starlette đứng chờ trong drain().
  5. Vòng async for không gọi __anext__() nữa.
  6. Generator của bạn đứng yên tại yield, không tạo thêm dữ liệu.

Đo thật: một endpoint stream 625 MiB theo mẩu 64 KiB, và một client chỉ đọc header rồi ngừng đọc:

client not reading for 0.5 s: generator produced 9.75 MiB of 625 MiB
client not reading for 1 s: generator produced 9.75 MiB of 625 MiB
client not reading for 2 s: generator produced 9.75 MiB of 625 MiB

Generator dừng ở 9,75 MiB và đứng yên. Con số đó xấp xỉ tổng các bộ đệm trên đường đi: trên máy này kernel cho bộ đệm gửi tối đa 4 MiB (tcp_wmem) và bộ đệm nhận tối đa 6 MiB (tcp_rmem), cộng 64 KiB của uvicorn.

asyncio.Queue: khi bên tạo và bên đọc là hai task

Generator gắn bên tạo và bên đọc vào cùng một nhịp: một phần tử được hỏi, một phần tử được tạo. Nhiều khi không như vậy:

  • Dữ liệu được đẩy tới, qua callback hoặc từ một thread, chứ không phải được kéo ra.
  • Nhiều nguồn cùng đổ về một chỗ, hoặc một nguồn chia cho nhiều worker.
  • Bên tạo chạy nền, độc lập với từng request.

Lúc đó, hai bên là hai task, nối với nhau bằng asyncio.Queue.

Queue không giới hạn là một cái bẫy

Một producer đẩy một triệu sự kiện, một consumer mất 1 ms cho mỗi sự kiện:

async def producer(queue, stats):
    for i in range(1_000_000):
        await queue.put(f"event {i}")  # with maxsize, waits here while the queue is full
        stats["produced"] += 1


async def consumer(queue):
    while True:
        await queue.get()
        await asyncio.sleep(0.001)  # 1 ms per event: a slow client, a slow API


queue = asyncio.Queue(maxsize=0)  # 0 = unbounded; then the same with maxsize=100
Python 3.12 Sau 1 giây: đã tạo Đang nằm trong queue Bộ nhớ
maxsize=0 (không giới hạn) 1.000.000 999.561 90 MiB
maxsize=100 933 100 31 MiB

Queue không giới hạn gây ra hai vấn đề:

  • Bộ nhớ phình theo khoảng cách giữa hai bên.
  • Event loop bị chặn. Producer tạo xong cả triệu sự kiện trong 0,68 s trên 3.12 (0,34 s trên 3.14), và suốt thời gian đó consumer không được chạy lần nào. await queue.put() trên một queue không đầy trả về ngay mà không nhường loop. await chỉ nhường loop khi nó thực sự phải chờ.

Với maxsize=100, put() phải chờ khi queue đầy, nên producer tự chạy theo nhịp của consumer. Gần như mọi queue nối hai bên có tốc độ khác nhau đều nên có maxsize.

Biến dữ liệu "đẩy" thành async generator "kéo"

Một trường hợp rất hay gặp: một SDK đồng bộ gọi callback cho từng token, từ thread riêng của nó, trong khi bạn cần một async generator để đưa vào StreamingResponse:

import asyncio
import threading
import time


def sdk_stream(on_token, on_done):
    """A blocking SDK that pushes tokens to a callback from its own thread."""
    for token in "Xin chào , mình là chatbot".split():
        time.sleep(0.1)
        on_token(token)
    on_done()


async def tokens():
    """Turn the push-style callbacks into a pull-style async generator."""
    loop = asyncio.get_running_loop()
    queue = asyncio.Queue(maxsize=16)

    def on_token(token):  # runs in the SDK's thread: never touch the queue directly from there
        asyncio.run_coroutine_threadsafe(queue.put(token), loop).result()  # blocks the thread while full

    def on_done():
        loop.call_soon_threadsafe(queue.shutdown)  # 3.13+: wakes the reader with QueueShutDown

    threading.Thread(target=sdk_stream, args=(on_token, on_done), daemon=True).start()
    while True:
        try:
            yield await queue.get()
        except asyncio.QueueShutDown:
            return


async def main():
    start = time.perf_counter()
    async for token in tokens():
        print(f"  {time.perf_counter() - start:.1f} s  {token}")


asyncio.run(main())
  0.1 s  Xin
  0.2 s  chào
  0.3 s  ,
  0.4 s  mình
  0.5 s  là
  0.6 s  chatbot

Ba chi tiết quan trọng:

  • asyncio.Queue không an toàn khi dùng từ thread khác. Từ thread của SDK, phải nhờ loop làm hộ, qua run_coroutine_threadsafe hoặc call_soon_threadsafe.
  • .result() làm thread của SDK đứng chờ khi queue đầy. Backpressure đi xuyên qua ranh giới thread.
  • queue.shutdown() (Python 3.13+) là cách báo "hết" gọn gàng. Mọi get() đang chờ hoặc gọi sau đó ném QueueShutDown, thay cho việc tự chế một giá trị đánh dấu kết thúc.

Nhiều worker, một queue

shutdown() cũng giải quyết bài toán cũ của mô hình nhiều worker: làm sao cho các worker đang chờ việc biết là đã xong để thoát.

async def worker(name, queue):
    try:
        while True:
            job = await queue.get()
            await asyncio.sleep(0.01)
            print(f"  {name} did {job}")
            queue.task_done()
    except asyncio.QueueShutDown:
        print(f"  {name}: queue shut down, stopping")


async def main():
    queue = asyncio.Queue()
    async with asyncio.TaskGroup() as tg:
        for n in range(3):
            tg.create_task(worker(f"worker-{n}", queue))
        for job in range(5):
            queue.put_nowait(f"job {job}")
        await queue.join()  # every job got task_done()
        queue.shutdown()  # wake idle workers so the TaskGroup can finish
  worker-0 did job 0
  worker-1 did job 1
  worker-2 did job 2
  worker-0 did job 3
  worker-1 did job 4
  worker-2: queue shut down, stopping
  worker-0: queue shut down, stopping
  worker-1: queue shut down, stopping

Lên lại FastAPI: StreamingResponse từ bên trong

Đây là phần lõi của StreamingResponse trong Starlette 1.8, rút gọn:

class StreamingResponse(Response):
    def __init__(self, content, ...):
        if isinstance(content, AsyncIterable):
            self.body_iterator = content
        else:
            self.body_iterator = iterate_in_threadpool(content)  # a plain generator: each next() in a thread

    async def stream_response(self, send):
        await send({"type": "http.response.start", "status": self.status_code, "headers": self.raw_headers})
        async for chunk in self.body_iterator:
            await send({"type": "http.response.body", "body": chunk, "more_body": True})
        await send({"type": "http.response.body", "body": b"", "more_body": False})

    async def __call__(self, scope, receive, send):
        # uvicorn speaks ASGI 2.3: run the stream and a disconnect listener side by side
        async with create_collapsing_task_group() as task_group:
            async def wrap(func):
                await func()
                task_group.cancel_scope.cancel()  # whichever finishes first cancels the other

            task_group.start_soon(wrap, partial(self.stream_response, send))
            await wrap(partial(self.listen_for_disconnect, receive))

Nó chỉ là một vòng async for quanh generator của bạn, mỗi phần tử một lần await send(...) với more_body: True. uvicorn gửi mỗi phần tử thành một mẩu của chunked encoding: một dòng ghi độ dài (dạng hex), dữ liệu, rồi \r\n. Ba hệ quả đáng biết như sau.

Generator def thường: mỗi phần tử một lần đổi thread

Nếu content là generator thường, Starlette bọc nó trong iterate_in_threadpool, nên mỗi next() chạy trong thread pool của AnyIO, giống endpoint def ở bài 1. Không sai, vì generator đồng bộ có thể chặn, nhưng mỗi lần đổi thread tốn tiền. Đo ở mục cuối.

Client ngắt kết nối: lệnh hủy chỉ chen vào được ở await

uvicorn khai báo ASGI phiên bản 2.3, nên Starlette chạy song song hai việc: stream body, và nghe receive(). Khi client đóng kết nối, receive() trả về http.disconnect, và task group hủy phần stream. Nhưng hủy trong asyncio nghĩa là ném CancelledError vào chỗ task đang chờ. Task không chờ ở đâu thì không có chỗ để ném.

Thí nghiệm: endpoint 625 MiB ở trên, client ngừng đọc 2 giây rồi cắt kết nối. Hai phiên bản generator: một cái chỉ tạo dữ liệu trong bộ nhớ, một cái thêm await asyncio.sleep(0) mỗi mẩu:

async def big_chunks(pause: bool):
    try:
        for i in range(10_000):
            if pause:
                await asyncio.sleep(0)  # a real await: the loop gets control back here
            yield i.to_bytes(4) + b"x" * (64 * 1024 - 4)
    finally:
        ...  # record how far we got
Generator Đã tạo khi client đi Đã tạo tổng cộng CancelledError trong generator
Không có await thật 9,75 MiB 625 MiB Không
Có await asyncio.sleep(0) 9,44 MiB 9,44 MiB Có, finally chạy ngay

Phiên bản không có await tạo nốt hơn 600 MiB sau khi client đã đi, không gửi cho ai cả. Lý do: khi kết nối đã đóng, send() của uvicorn trả về ngay (if self.disconnected: return) mà không chờ gì, generator cũng không chờ gì, nên vòng async for chạy một mạch tới hết và lệnh hủy không có chỗ nào để chen vào. Suốt khoảng đó, event loop cũng không phục vụ được request nào khác.

Generator thật lấy dữ liệu từ database hay API thường có await thật nên ít gặp chuyện này. Nhưng một generator đi qua dữ liệu đã có sẵn trong bộ nhớ, hay chỉ tính toán, thì gặp đúng như trên. Một await asyncio.sleep(0) mỗi mẩu là đủ để nó dừng được.

Ví dụ trọn vẹn: chatbot trả lời từng token

Mọi mảnh ở trên ghép lại trong trường hợp quen thuộc nhất: một chatbot trả lời từng chữ. Ví dụ này chạy một mô hình ngôn ngữ thật, Qwen2.5-0.5B-Instruct, ngay trên CPU bằng thư viện transformers. Đây là mô hình rất nhỏ nên câu trả lời không hay; điều cần xem ở đây là thời gian và đường đi của từng token, không phải nội dung.

Vì sao phải stream

Mô hình sinh từng token một. Câu hỏi "Viết một đoạn ngắn giới thiệu ngôn ngữ lập trình Python cho người mới bắt đầu" cho ra 180 token:

Thời gian
Mẩu chữ đầu tiên tới trình duyệt 0,69–0,71 s
Cả câu trả lời xong 10,34–10,51 s

Không stream, người dùng nhìn màn hình trống hơn 10 giây. Có stream, chữ bắt đầu hiện sau 0,7 giây rồi chạy dần, mỗi mẩu cách nhau khoảng 0,05–0,1 s:

   0.69 s  'Tất '
   0.79 s  'nhiên, '
   0.84 s  'tôi '
   0.90 s  'sẽ '
   0.95 s  'giúp '
   1.00 s  'bạn '

Các API như OpenAI hay Anthropic sinh nhanh hơn một mô hình chạy trên CPU laptop nhiều, nhưng tỉ lệ vẫn vậy: một câu trả lời dài vẫn mất nhiều giây, còn token đầu tiên đến gần như ngay lập tức.

Mô hình chạy trong thread, token đi qua queue

model.generate() là một hàm đồng bộ chạy vài giây tới vài chục giây, nên nó không được chạy trên event loop. Nó chạy trong thread riêng, và gửi chữ ra qua một "streamer": một object mà generate() gọi mỗi khi có thêm chữ. Đây đúng là trường hợp "đẩy" ở mục asyncio.Queue: thread của mô hình đẩy chữ vào queue, còn async generator kéo ra và gửi đi dạng SSE.

Server-Sent Events là định dạng stream đơn giản nhất mà trình duyệt hiểu sẵn: Content-Type: text/event-stream, mỗi sự kiện là các dòng data: ... kết thúc bằng một dòng trống.

import asyncio
import json
import threading

import torch
from fastapi import FastAPI
from fastapi.responses import StreamingResponse
from transformers import AutoModelForCausalLM, AutoTokenizer, StoppingCriteria, TextStreamer

MODEL = "Qwen/Qwen2.5-0.5B-Instruct"
tokenizer = AutoTokenizer.from_pretrained(MODEL)
model = AutoModelForCausalLM.from_pretrained(MODEL)
app = FastAPI()


class QueueStreamer(TextStreamer):
    """Gets each new piece of text from the model's thread and puts it on an asyncio.Queue."""

    def __init__(self, queue, loop):
        super().__init__(tokenizer, skip_prompt=True, skip_special_tokens=True)
        self.queue, self.loop = queue, loop

    def on_finalized_text(self, text, stream_end=False):
        try:
            if text:
                # waits (in the model's thread) while the queue is full: backpressure on the model
                asyncio.run_coroutine_threadsafe(self.queue.put(text), self.loop).result()
            if stream_end:
                self.loop.call_soon_threadsafe(self.queue.shutdown)
        except asyncio.QueueShutDown:
            pass  # the reader is gone; StopWhen below ends generation at the next token


class StopWhen(StoppingCriteria):
    """Checked by generate() after every token: stop as soon as the event is set."""

    def __init__(self, event):
        self.event = event

    def __call__(self, input_ids, scores, **kwargs):
        return torch.full((input_ids.shape[0],), self.event.is_set(), dtype=torch.bool)


async def answer(question: str):
    loop = asyncio.get_running_loop()
    queue = asyncio.Queue(maxsize=16)
    stop = threading.Event()
    inputs = tokenizer.apply_chat_template(
        [{"role": "user", "content": question}], add_generation_prompt=True, return_tensors="pt", return_dict=True
    )

    def generate():
        model.generate(
            **inputs, max_new_tokens=180, do_sample=False,
            streamer=QueueStreamer(queue, loop), stopping_criteria=[StopWhen(stop)],
        )

    threading.Thread(target=generate, daemon=True).start()  # the model blocks: run it in its own thread
    try:
        while True:
            try:
                text = await queue.get()  # a real await: cancellation lands here if the user leaves
            except asyncio.QueueShutDown:
                yield "data: [DONE]\n\n"
                return
            yield f"data: {json.dumps(text, ensure_ascii=False)}\n\n"
    finally:
        stop.set()  # the user left (or we're done): stop generating
        queue.shutdown(immediate=True)  # and unblock the model's thread if it's waiting on a full queue


@app.get("/chat")
async def chat(q: str):
    return StreamingResponse(answer(q), media_type="text/event-stream")

Mỗi mảnh trong đoạn code này đều đã gặp ở trên:

  • Thread của mô hình không đụng thẳng vào queue: nó nhờ event loop qua run_coroutine_threadsafe và call_soon_threadsafe.
  • maxsize=16 và .result(): nếu người dùng đọc chậm, thread của mô hình đứng chờ. Đó là backpressure đi tới tận mô hình.
  • await queue.get() là một await thật, nên khi người dùng rời đi, lệnh hủy của Starlette có chỗ để chen vào.
  • queue.shutdown() báo hết khi mô hình xong. shutdown(immediate=True) trong finally gỡ thread của mô hình ra nếu nó đang chờ một queue đầy.

Phía trình duyệt

<p id="answer"></p>
<script>
const answer = document.querySelector("#answer");
const events = new EventSource("/chat?q=" + encodeURIComponent("Viết một đoạn ngắn giới thiệu Python."));
events.onmessage = (event) => {
  if (event.data === "[DONE]") return events.close();  // otherwise EventSource reconnects and asks again
  answer.textContent += JSON.parse(event.data);
};
</script>

Sự kiện [DONE] và events.close() không phải để trang trí. Khi server đóng stream, EventSource coi đó là mất kết nối và tự kết nối lại. Thử trong Chromium với một stream 3 chữ: không gọi close(), trong 10 giây trang gửi lại request 3–4 lần, mỗi lần khoảng 3 giây một lần, và câu trả lời hiện ra 4 lần nối đuôi nhau. Với một chatbot thật, mỗi lần như vậy là một lần chạy mô hình.

Khi người dùng rời đi

Người dùng đóng tab sau khi nhận 5 mẩu chữ. So sánh đoạn code ở trên với một bản không có finally:

Token khi người dùng đi 2 giây sau Kết thúc
Có finally: stop.set() + shutdown(immediate=True) 8 9 Dừng ở 9 token
Không có finally 8 25 Kẹt mãi ở 25 token, thread không bao giờ kết thúc

Có finally, lệnh hủy chen vào ở await queue.get(), finally bật cờ dừng, và generate() dừng ngay ở token kế tiếp. Không có finally, mô hình tiếp tục sinh chữ cho không ai đọc, cho tới khi queue 16 chỗ đầy. Lúc đó thread của mô hình đứng chờ trong put() mãi mãi: một thread bị rò rỉ, giữ bộ nhớ và các tensor của nó, với mỗi người dùng rời trang giữa chừng. Nếu queue không giới hạn, nó sẽ không kẹt mà chạy hết 180 token, tốn CPU (hay GPU) cho một câu trả lời không ai xem.

Với API của OpenAI hay Anthropic, cấu trúc giống hệt, chỉ khác là SDK đã có sẵn async iterator, nên không cần thread và queue. Điều cần giữ là finally (hoặc async with) để đóng stream tới API ngay khi người dùng rời đi, thay vì đọc tiếp tới hết một câu trả lời không ai xem.

Qua proxy

Giữa người dùng và uvicorn thường có một reverse proxy, và proxy có thể gom dữ liệu lại trước khi gửi. Để đo cho sạch, thí nghiệm này thay mô hình bằng một generator giả, mỗi token cách nhau đúng 0,1 s, 19 token, client có gửi Accept-Encoding: gzip:

Đường đi Token đầu tiên tới sau Token về
Thẳng vào uvicorn 0,1 s Từng cái một
Caddy 2.11, encode gzip 0,1 s Từng cái một, đã nén
nginx 1.31, gzip on 1,91 s Cả câu cùng lúc, trong một lần nhận
nginx 1.31, gzip on, app gửi header X-Accel-Buffering: no 0,1 s Từng cái một, đã nén

Với nginx bật gzip, người dùng nhìn màn hình trống gần 2 giây rồi thấy cả câu hiện ra cùng lúc, đúng cảm giác "chatbot bị đơ", trong khi server vẫn stream đúng. Header X-Accel-Buffering: no bảo nginx không gom cho riêng response đó. Caddy (website này chạy sau Caddy với encode gzip) gửi từng token ngay. GZipMiddleware của Starlette 1.8 cũng an toàn: nó bỏ qua text/event-stream, và với stream khác thì flush sau mỗi mẩu.

Bài học chung: khi stream "không chạy", hãy đo ở từng chặng (thẳng vào app, qua proxy, qua CDN) trước khi sửa code.

Đo đạc

Một lần yield tốn bao nhiêu? Lặp 5 triệu phần tử, lấy lần nhanh nhất trong 5 lần:

Cách lặp Python 3.12 Python 3.14
Tạo list rồi for 32,0 ns/phần tử 30,7 ns/phần tử
for qua generator 22,6 ns/phần tử 14,2 ns/phần tử
async for qua async generator 99,6 ns/phần tử 86,7 ns/phần tử

Bản thân yield rất rẻ. Async generator đắt hơn khoảng 4–6 lần vì mỗi phần tử đi qua một awaitable asend và một StopIteration, nhưng vẫn dưới 0,1 µs.

Còn trong StreamingResponse thì sao? Cùng CSV ở đầu bài, đổi kích thước mẩu và loại generator:

Generator Mỗi mẩu Số dòng Thời gian
async def 1.000 dòng 1.000.000 0,32–0,36 s
def thường 1.000 dòng 1.000.000 0,40–0,45 s
async def 1 dòng 1.000.000 3,08–3,17 s
async def 1 dòng 100.000 0,32–0,34 s
def thường 1 dòng 100.000 9,71–9,76 s
  • Chi phí nằm ở mỗi mẩu, không ở mỗi yield. Mỗi mẩu là một lần send() qua Starlette, uvicorn và transport, khoảng 3 µs. Yield từng dòng chậm hơn khoảng 10 lần so với gom 1.000 dòng. Nó còn tốn thêm 6 MB trên đường truyền, vì mỗi mẩu chunked encoding có phần đầu và đuôi riêng.
  • Generator def thường với mẩu nhỏ là tệ nhất: khoảng 97 µs mỗi dòng, vì mỗi dòng là một lần đổi sang thread pool và quay về. Với một triệu dòng, ước tính hơn một phút rưỡi. Gom mẩu lại thì gần bằng async def.

Quy tắc: gom thành mẩu cỡ vài chục KiB (trừ khi cần gửi ngay từng phần như token chatbot), và dùng async def nếu nguồn dữ liệu cho phép.

Tóm lại

Khái niệm Thực chất
Generator Một frame sống trên heap; YIELD_VALUE lưu vị trí và rời frame, send() chạy tiếp
Coroutine Cùng cỗ máy đó; await là một vòng SEND/YIELD_VALUE chuyền giá trị lên tới event loop
Async generator Hai kiểu tạm dừng: await đi ra event loop, yield (được bọc lại) đi tới async for
Dọn dẹp Generator thường đóng ngay khi không ai giữ; async generator được đóng sau, trong task khác. Dùng aclosing
Backpressure Bên đọc chậm thì drain() chờ, send() chờ, generator đứng ở yield
asyncio.Queue Nối hai task; luôn đặt maxsize; shutdown() để báo hết (3.13+)
StreamingResponse Một vòng async for quanh generator; mỗi phần tử một lần send()
Hủy khi client đi Chỉ chen vào được ở await thực sự chờ
Chatbot Mô hình trong thread → queue có giới hạn → async generator → SSE; finally dừng mô hình khi người dùng đi

Mọi thứ trong bài quy về hai câu hỏi đã gặp ở bài 1: dữ liệu đang nằm ở đâu trong lúc chờ, và ai đang chờ ai? Stream tốt là stream mà ở mọi tầng, câu trả lời đều là "chỉ một mẩu nhỏ" và "bên nhanh chờ bên chậm".