Problem 127 · Practical builds

Rate-Limited LLM Request Scheduler

HardPractical multi-stage build

01Problem

Build the scheduler that sits in front of an LLM provider. Requests carry a tenant, model, estimated tokens and a priority lane. Dispatch must respect a provider rate limit, per-tenant token quotas, priority with starvation protection, and track cost per tenant.

Stages — finish each before reading the next

1Rate limiter

submit(request) queues work; dispatch(now) returns the requests that may start, using a token bucket (rate per second, burst capacity). Time is passed in, never read from the wall clock.

2Per-tenant quotas

Each tenant has a token quota per fixed window. Reserve est_tokens at dispatch; skip (do not block on) requests whose tenant is over quota; reject at submit any request that could never fit. complete(id, input, output) reconciles the reservation with actual usage.

3Priority lanes

Lane 0 (interactive) is served before lane 1 (batch). A lower lane passed over starvation_limit times in a row while non-empty gets the next slot.

4Cost accounting

complete() prices actual input/output tokens per model in integer micro-USD and adds it to a per-tenant ledger, returning the request's cost.

02Key insight

Separate the three gates: the bucket decides whether anything may start, quotas decide which tenant may, and lane order decides who goes first. Reserve the estimate at dispatch and reconcile at completion, or one tenant can overrun by issuing many requests before any finish. Keep money in integers.

03Solution

solution.pypython98 lines

from collections import deque
from dataclasses import dataclass


@dataclass(frozen=True)
class Request:
    id: str
    tenant: str
    model: str
    est_tokens: int   # prompt + max_output, reserved against the tenant quota
    lane: int = 1     # 0 = interactive (highest priority)


class Scheduler:
    """Token-bucket rate limit + per-tenant token quotas + priority lanes + cost ledger."""

    def __init__(self, rate_per_s: float, burst: int, quotas: dict[str, int],
                 window_s: float, prices: dict[str, tuple[int, int]],
                 lanes: int = 2, starvation_limit: int = 3) -> None:
        self.rate, self.burst = rate_per_s, burst
        self.bucket, self.last = float(burst), 0.0
        self.quotas, self.window_s = quotas, window_s
        self.prices = prices  # model -> (micro-USD per input token, per output token)
        self.queues: list[deque[Request]] = [deque() for _ in range(lanes)]
        self.passed_over = [0] * lanes
        self.starvation_limit = starvation_limit
        self.window: dict[str, int] = {}
        self.used: dict[str, int] = {}
        self.inflight: dict[str, tuple[Request, int]] = {}
        self.cost: dict[str, int] = {}

    def submit(self, req: Request) -> None:
        if req.tenant not in self.quotas:
            raise KeyError(f"unknown tenant {req.tenant}")
        if req.model not in self.prices:
            raise KeyError(f"unpriced model {req.model}")
        if req.est_tokens > self.quotas[req.tenant]:
            raise ValueError("request can never fit the tenant quota")
        if not 0 <= req.lane < len(self.queues):
            raise ValueError("bad lane")
        self.queues[req.lane].append(req)

    def _refill(self, now: float) -> None:
        if now < self.last:
            raise ValueError("clock went backwards")
        self.bucket = min(self.burst, self.bucket + (now - self.last) * self.rate)
        self.last = now

    def _used(self, tenant: str, now: float) -> int:
        w = int(now // self.window_s)
        if self.window.get(tenant) != w:  # fixed window rolled over
            self.window[tenant], self.used[tenant] = w, 0
        return self.used[tenant]

    def _take_eligible(self, lane: int, now: float):
        q = self.queues[lane]
        for i, r in enumerate(q):  # skip over-quota tenants: no head-of-line blocking
            if self._used(r.tenant, now) + r.est_tokens <= self.quotas[r.tenant]:
                del q[i]
                return r
        return None

    def _lane_order(self) -> list[int]:
        n = len(self.queues)
        starving = [l for l in range(n) if self.queues[l] and self.passed_over[l] >= self.starvation_limit]
        return starving + [l for l in range(n) if l not in starving]

    def dispatch(self, now: float) -> list[Request]:
        self._refill(now)
        out: list[Request] = []
        while self.bucket >= 1:
            picked = None
            for lane in self._lane_order():
                r = self._take_eligible(lane, now)
                if r is not None:
                    picked = (lane, r)
                    break
            if picked is None:
                break
            lane, r = picked
            for lower in range(lane + 1, len(self.queues)):
                if self.queues[lower]:
                    self.passed_over[lower] += 1
            self.passed_over[lane] = 0
            self.bucket -= 1
            self.used[r.tenant] += r.est_tokens
            self.inflight[r.id] = (r, self.window[r.tenant])
            out.append(r)
        return out

    def complete(self, req_id: str, input_tokens: int, output_tokens: int) -> int:
        r, w = self.inflight.pop(req_id)
        if self.window.get(r.tenant) == w:  # reconcile reservation vs actual usage
            self.used[r.tenant] = max(0, self.used[r.tenant] + input_tokens + output_tokens - r.est_tokens)
        p_in, p_out = self.prices[r.model]
        micro_usd = input_tokens * p_in + output_tokens * p_out
        self.cost[r.tenant] = self.cost.get(r.tenant, 0) + micro_usd
        return micro_usd

Try it first.

Timedispatch O(k * q) for k dispatched and q queued in scanned lanes

SpaceO(q + tenants + in-flight)

tests.py · 12 passing

tests.pypython39 lines

prices = {"small": (1, 4), "large": (3, 15)}
ids = lambda rs: [r.id for r in rs]
# stage 1: rate limit (1 req/s, burst 2)
s = Scheduler(1.0, 2, {"a": 10_000}, 60, prices)
for i in range(3):
    s.submit(Request(f"r{i}", "a", "small", 10))
assert ids(s.dispatch(0.0)) == ["r0", "r1"]
assert s.dispatch(0.5) == []
assert ids(s.dispatch(1.0)) == ["r2"]
# stage 2: quotas; over-quota tenant does not block others in the same lane
s = Scheduler(100.0, 100, {"a": 100, "b": 100}, 60, prices)
s.submit(Request("a1", "a", "small", 60)); s.submit(Request("a2", "a", "small", 60))
s.submit(Request("b1", "b", "small", 60))
assert ids(s.dispatch(0.0)) == ["a1", "b1"]
s.complete("a1", 10, 10)            # actual 20 < reserved 60: refund 40
assert ids(s.dispatch(0.1)) == ["a2"]
s.submit(Request("a3", "a", "small", 60))
assert s.dispatch(1.0) == []        # a has 80 used in this window
assert ids(s.dispatch(60.0)) == ["a3"]  # new window
try:
    s.submit(Request("big", "a", "small", 101)); assert False
except ValueError:
    pass
# stage 3: priority lanes with starvation guard (limit 2)
s = Scheduler(1.0, 1, {"a": 10**6}, 60, prices, lanes=2, starvation_limit=2)
for i in range(4):
    s.submit(Request(f"hi{i}", "a", "small", 1, lane=0))
s.submit(Request("lo0", "a", "small", 1, lane=1))
order = [ids(s.dispatch(float(t)))[0] for t in range(5)]
assert order == ["hi0", "hi1", "lo0", "hi2", "hi3"]
# stage 4: cost accounting in integer micro-USD
s = Scheduler(10.0, 10, {"a": 10**6, "b": 10**6}, 60, prices)
s.submit(Request("x", "a", "large", 2000)); s.submit(Request("y", "b", "small", 2000))
s.submit(Request("z", "a", "small", 2000))
s.dispatch(0.0)
assert s.complete("x", 1000, 200) == 6000
assert s.complete("y", 500, 500) == 2500
assert s.complete("z", 100, 0) == 100
assert s.cost == {"a": 6100, "b": 2500}

04Staff-level follow-up

How does this scale past one process? Move the bucket and quota counters to Redis with atomic Lua scripts (or a central admission service), honour provider 429 Retry-After by shrinking the bucket rate, and replace FIFO scanning with per-tenant sub-queues and weighted fair queuing so a noisy tenant can't fill a lane.