[Bug]: /v1/embeddings: a multi-input request rejected with 429 still encodes the inputs it had queued #22133

System Info

CPU only, found by reading tensorrt_llm/serve/encode_batcher.py and OpenAIServer.openai_embedding on main (e3c25b0) and reproduced with EncodeBatcher directly. No GPU is needed to show it.

Who can help?

@xwang233

Information

• The official example scripts
• My own modified scripts

Tasks

• An officially supported task in the examples folder (such as GLUE/SQuAD, ...)
• My own task or dataset (give details below)

Reproduction

openai_embedding sends the inputs of one /v1/embeddings request to the batcher with asyncio.gather(*[batcher.submit(t) for t in token_ids_list]). Each submit enqueues on its own, so when the queue fills partway through the request, the inputs enqueued before the full one stay in the queue. The request is answered with 429, but the worker still runs encode() on those inputs, and until it does they hold queue slots, so other requests get 429 too.

The handler already validates lengths up front for this reason ("otherwise gather() cancels the awaiters but their already-queued inputs still run encode() (wasted GPU work)"), but the queue-full case still enqueues part of the request. The same holds after a mid-request failure or a cancelled caller: inputs still waiting in the queue are encoded even though nobody awaits them.

import asyncio

from tensorrt_llm.serve.encode_batcher import EncodeBatcher, QueueFullError


async def main():
    calls = []

    def encode_fn(batch):
        calls.append([list(t) for t in batch])
        return [sum(t) for t in batch]

    b = EncodeBatcher(encode_fn, max_batch_size=8, max_queue_delay=0.05,
                      max_queue_size=2)
    await b.start()
    # what openai_embedding does for a request with three inputs
    try:
        await asyncio.gather(*[b.submit(t) for t in ([1], [2], [3])])
    except QueueFullError as e:
        print("request got 429:", e)
    await asyncio.sleep(0.2)
    print("encode_fn calls after the 429:", calls)
    await b.shutdown()


asyncio.run(main())

request got 429: Embedding request queue is full; retry later.
encode_fn calls after the 429: [[[1], [2]]]

Expected behavior

A request rejected with 429 (or 400) leaves nothing in the queue, and inputs of a request that already ended (an input failed, or the caller was cancelled) are not encoded.

actual behavior

The inputs enqueued before the queue filled run through encode() after the client got 429, and take queue slots from other requests until then. Under load, where the queue is near full, this turns one rejected request into GPU work nobody reads plus extra 429s for other clients.

additional notes

A fix: add EncodeBatcher.submit_many() that validates the whole group, checks that the queue has room for all of it, and only then enqueues (no await between the check and the puts, so it is atomic on the event loop); when the group ends early it cancels its inputs still queued, and the worker drops requests whose future is already done instead of encoding them. openai_embedding calls submit_many(). I have a patch with CPU unit tests in tests/unittest/llmapi/apps/test_encode_batcher.py and will open a PR.

Before submitting a new issue...

• Make sure you already searched for relevant issues, and checked the documentation and examples for answers to frequently asked questions.