Batching Requests to a Model Server

How dynamic batching trades a little latency for much higher throughput, with a small asyncio implementation.

Published
Reading time
3 min

A model running on a GPU is far more efficient when it processes many inputs at once. A forward pass on a batch of thirty-two often takes barely longer than a pass on a batch of one. A web server, on the other hand, receives requests one at a time. Dynamic batching bridges the two: we hold incoming requests for a few milliseconds, group them, run one forward pass, and hand each caller its own result.

The problem

Without batching, each request runs its own forward pass. Under light load this is fine. Under heavy load the requests queue up behind each other, the GPU spends most of its time on small, inefficient kernels, and latency climbs even though the hardware is far from busy.

How dynamic batching works

The server keeps a queue of pending requests and a single worker that drains it. The worker waits until one of two things happens, then runs a batch.

Maximum batch size

If the queue reaches a fixed size, the worker runs immediately. This bounds memory use and keeps each forward pass within what the model was tested with.

Maximum wait

If the oldest request has waited longer than a fixed time, the worker runs whatever it has. This bounds the extra latency that batching adds when traffic is light. A few milliseconds is a common starting point.

A small implementation

The sketch below uses asyncio. Each caller puts its input and a future on the queue, then awaits the future. The worker collects up to max_size items or waits up to max_wait seconds, whichever comes first.

import asyncio

class Batcher:
    def __init__(self, predict, max_size=32, max_wait=0.005):
        self.predict = predict          # takes a list, returns a list
        self.max_size = max_size
        self.max_wait = max_wait
        self.queue = asyncio.Queue()

    async def submit(self, item):
        future = asyncio.get_running_loop().create_future()
        await self.queue.put((item, future))
        return await future

    async def run(self):
        while True:
            item, future = await self.queue.get()
            batch = [(item, future)]
            deadline = asyncio.get_running_loop().time() + self.max_wait
            while len(batch) < self.max_size:
                timeout = deadline - asyncio.get_running_loop().time()
                if timeout <= 0:
                    break
                try:
                    batch.append(await asyncio.wait_for(self.queue.get(), timeout))
                except asyncio.TimeoutError:
                    break
            outputs = self.predict([i for i, _ in batch])
            for (_, f), out in zip(batch, outputs):
                f.set_result(out)

In a real server predict would run in a thread or a separate process so it does not block the event loop, and errors would be passed to every future in the batch. The structure stays the same.

Measuring the trade-off

Batching is a trade: each request may wait up to max_wait longer, and in exchange the server handles many more requests per second. The only way to choose the two settings is to measure. A load test that sweeps the request rate and records the 50th and 99th percentile latency for a few values of max_wait shows quickly where the gains stop.

A good way to start is to set max_size to the largest batch the model handles comfortably in memory, and max_wait to a small fraction of the latency budget. If the service promises responses within one hundred milliseconds and the forward pass takes thirty, waiting five or ten milliseconds to fill a batch is affordable. Waiting fifty is not.

Throughput usually improves steeply up to a batch size of eight or sixteen and then flattens. Tail latency is the number to watch; a setting that helps the median can still hurt the slowest one percent.

When not to batch

Batching helps when the model is large relative to a single input and the hardware is parallel. For small models on a CPU, or for inputs whose sizes vary so much that padding wastes most of the batch, the gain can be small or negative. Sorting or bucketing requests by input length before batching recovers some of that waste, at the cost of more code.

It is also worth remembering that batching changes failure behaviour. One malformed input can fail a whole batch, so the worker should catch errors per item where possible, or retry the rest of the batch without the bad input. A batcher that is fast in the common case but turns one bad request into thirty failed ones is not an improvement.