Problem 115 · LLM Systems

Async LLM Client: Token Bucket, Semaphore, Retries with Jitter

HardAI engineering

01Problem

Write an asyncio `LLMClient(transport, tpm, max_concurrency, max_retries, base_delay, max_delay, rng, clock)`. `complete(prompt, max_tokens)` estimates tokens (`len(prompt) // 4 + 1 + max_tokens`), waits on a tokens-per-minute bucket, limits in-flight calls with a semaphore, and calls `await transport(prompt, max_tokens) -> (text, tokens_used)`. On `RetryableError` it refunds the estimate and retries with full-jitter exponential backoff (or the error's `retry_after`), up to `max_retries`; other exceptions propagate immediately. After success it reconciles the bucket with the real usage. Use an injected fake clock so tests take milliseconds.

02Key insight

Provider limits are on tokens per minute as well as requests, so concurrency alone does not prevent 429s; you need a token bucket sized in tokens and a semaphore for in-flight calls. The bucket's wait loop sits behind a lock so a large request is not starved by a stream of small ones. Full jitter (uniform between 0 and the exponential cap) decorrelates retries from many workers, avoiding the thundering herd that fixed backoff creates; `retry_after` from the server wins when present. Reconciling estimated against actual tokens keeps the bucket honest, and dependency-injected time makes all of it deterministic to test.

03Solution

solution.pypython85 lines

import asyncio
import random
from typing import Awaitable, Callable

Transport = Callable[[str, int], Awaitable[tuple[str, int]]]


class RetryableError(Exception):
    def __init__(self, msg: str, retry_after: float | None = None) -> None:
        super().__init__(msg)
        self.retry_after = retry_after


class FakeClock:
    def __init__(self) -> None:
        self.t = 0.0

    def now(self) -> float:
        return self.t

    async def sleep(self, seconds: float) -> None:
        target = self.t + max(0.0, seconds)
        await asyncio.sleep(0)
        self.t = max(self.t, target)


class TokenBucket:
    def __init__(self, tokens_per_minute: int, clock: FakeClock) -> None:
        self.capacity = float(tokens_per_minute)
        self.rate = tokens_per_minute / 60.0
        self.tokens = self.capacity
        self.clock = clock
        self.updated = clock.now()
        self._lock = asyncio.Lock()

    def _refill(self) -> None:
        t = self.clock.now()
        self.tokens = min(self.capacity, self.tokens + (t - self.updated) * self.rate)
        self.updated = t

    async def acquire(self, n: int) -> None:
        if n > self.capacity:
            raise ValueError("request larger than the per-minute budget")
        async with self._lock:
            while True:
                self._refill()
                if self.tokens >= n - 1e-9:
                    self.tokens -= n
                    return
                await self.clock.sleep((n - self.tokens) / self.rate)

    def adjust(self, delta: float) -> None:
        self._refill()
        self.tokens = min(self.capacity, self.tokens + delta)


class LLMClient:
    def __init__(self, transport: Transport, tpm: int, max_concurrency: int, max_retries: int = 4,
                 base_delay: float = 0.5, max_delay: float = 20.0, rng: random.Random | None = None,
                 clock: FakeClock | None = None) -> None:
        self.transport = transport
        self.clock = clock or FakeClock()
        self.bucket = TokenBucket(tpm, self.clock)
        self.sem = asyncio.Semaphore(max_concurrency)
        self.max_retries, self.base_delay, self.max_delay = max_retries, base_delay, max_delay
        self.rng = rng or random.Random(0)

    async def complete(self, prompt: str, max_tokens: int) -> str:
        estimate = len(prompt) // 4 + 1 + max_tokens
        async with self.sem:
            for attempt in range(self.max_retries + 1):
                await self.bucket.acquire(estimate)
                try:
                    text, used = await self.transport(prompt, max_tokens)
                except RetryableError as e:
                    self.bucket.adjust(estimate)
                    if attempt == self.max_retries:
                        raise
                    cap = min(self.max_delay, self.base_delay * 2 ** attempt)
                    delay = e.retry_after if e.retry_after is not None else self.rng.uniform(0, cap)
                    await self.clock.sleep(delay)
                    continue
                self.bucket.adjust(estimate - used)
                return text
        raise AssertionError("unreachable")

Try it first.

TimeO(1) per call plus waiting; bucket refill is O(1)

SpaceO(max_concurrency) in-flight requests plus O(waiters)

tests.py · 10 passing

tests.pypython63 lines

async def scenario_concurrency():
    clock = FakeClock()
    state = {"in": 0, "max": 0}
    async def transport(prompt, max_tokens):
        state["in"] += 1
        state["max"] = max(state["max"], state["in"])
        await clock.sleep(1.0)
        state["in"] -= 1
        return prompt.upper(), 10
    client = LLMClient(transport, tpm=1_000_000, max_concurrency=3, clock=clock)
    out = await asyncio.gather(*(client.complete(f"p{i}", 5) for i in range(10)))
    assert out == [f"P{i}" for i in range(10)]
    assert state["max"] == 3
asyncio.run(scenario_concurrency())

async def scenario_rate_limit():
    clock = FakeClock()
    async def transport(prompt, max_tokens):
        return "ok", len(prompt) // 4 + 1 + max_tokens
    client = LLMClient(transport, tpm=1000, max_concurrency=10, clock=clock)
    await asyncio.gather(*(client.complete("x" * 40, 189) for _ in range(10)))
    assert 60.0 - 1e-6 <= clock.now() < 61.0, clock.now()
asyncio.run(scenario_rate_limit())

async def scenario_retries():
    clock = FakeClock()
    calls = {"n": 0}
    async def flaky(prompt, max_tokens):
        calls["n"] += 1
        if calls["n"] == 1:
            raise RetryableError("429", retry_after=5.0)
        if calls["n"] == 2:
            raise RetryableError("503")
        return "done", 3
    client = LLMClient(flaky, tpm=10_000, max_concurrency=2, rng=random.Random(1), clock=clock)
    assert await client.complete("hi", 10) == "done"
    assert calls["n"] == 3 and 5.0 <= clock.now() <= 6.0
    async def broken(prompt, max_tokens):
        calls["n"] += 1
        raise ValueError("bad request")
    calls["n"] = 0
    client.transport = broken
    try:
        await client.complete("hi", 10)
        assert False
    except ValueError:
        assert calls["n"] == 1
    async def down(prompt, max_tokens):
        calls["n"] += 1
        raise RetryableError("503")
    calls["n"] = 0
    client2 = LLMClient(down, tpm=10_000, max_concurrency=1, max_retries=3, clock=FakeClock())
    try:
        await client2.complete("hi", 10)
        assert False
    except RetryableError:
        assert calls["n"] == 4
    try:
        await client2.complete("x", 20_000)
        assert False
    except ValueError:
        pass
asyncio.run(scenario_retries())

04Staff-level follow-up

This client sleeps between retries while holding its semaphore slot; discuss releasing it so retries do not block fresh work, and add a circuit breaker so a provider outage fails fast instead of queueing. Across many processes the bucket must be shared (Redis with a Lua script, or the provider's rate-limit headers as the source of truth). Finally, make retries safe: a timeout after the provider started generating still bills you, so attach idempotency keys where supported and cap total retry spend per request.