AI infrastructure

You’ll recognise it when

• The problem is a piece of a model server: a batcher, a streaming client, a limiter, a cache, a sampler, an eval dashboard.
• Rules are exact and numeric: “at least”, “at most”, “within 60,000 ms”, “ties go to the lower index”.
• Costs are counted in tokens, not in requests or entries.
• It comes in levels, and a later level changes an earlier rule (“the token budget replaces the count rule”).
• Time arrives as an argument (now, t_ms), and nothing runs in the background.

It’s often confused with machine learning itself. You never train anything here. These are careful systems problems, closest to simulation, rolling counters and caches.

The idea

Think of a shuttle bus at an airport. It leaves when it’s full, or when the first passenger has waited ten minutes, whichever comes first. The driver doesn’t watch a clock all day: every time someone walks up, the driver first checks whether the bus should already have left, sends it if so, and only then lets the new passenger on.

Almost every AI-infra round is that driver. The class has no clock of its own. Each call brings the time, and the first thing every method does is catch up: apply whatever became due by now (timeouts, expiries, promotions). Then it applies the call’s own rule. Two more habits run through all of them: measure capacity in tokens, because that’s what fills a GPU, and make every choice deterministic: seeded randomness and explicit tie-breaks, so tests can check exact answers.

How it works

Build a request batcher. Requests wait in waiting. A batch leaves when it reaches max_batch requests, when the tokens waiting reach max_tokens, or when the oldest request has waited max_wait. A request that would push the tokens over the budget doesn’t join: the current batch leaves first.

1. Catch up first. submit(now, ...) and tick(now) begin by checking the oldest request: if it arrived at t0 and t0 + max_wait <= now, everything waiting leaves as one batch. Doing this before anything else means a request arriving exactly at a deadline can’t sneak into a batch that should already have gone.
2. Protect the budget. If something is waiting and waiting_tokens + tokens > max_tokens, cut the current batch before adding the new request.
3. Join. Append the request to waiting and add its tokens to waiting_tokens.
4. Full? If waiting now holds max_batch requests, or waiting_tokens >= max_tokens, cut. A single request bigger than the whole budget therefore leaves alone, right away.
5. Return what left during this call, in order. One call can send two batches: a timeout batch and then a full one.

With max_batch = 3, max_tokens = 100, max_wait = 50:

Why it’s correct: after every call, the waiting requests are under every limit (fewer than max_batch, fewer than max_tokens tokens, and the oldest not yet due). Each call restores that, in a fixed order: time first, then budget, then the new request.

class Batcher:
    """Groups requests into batches by count, by token budget and by waiting time.
    It has no clock of its own: time moves only when a method is called."""

    def __init__(self, max_batch, max_tokens, max_wait):
        self.max_batch = max_batch
        self.max_tokens = max_tokens
        self.max_wait = max_wait
        self.waiting = []             # (arrival time, request id), oldest first
        self.waiting_tokens = 0

    def _cut(self):
        batch = [req for _, req in self.waiting]
        self.waiting, self.waiting_tokens = [], 0
        return batch

    def _catch_up(self, now):
        # Every rule that depends on time runs here, at the top of every call.
        if self.waiting and self.waiting[0][0] + self.max_wait <= now:
            return [self._cut()]
        return []

    def submit(self, now, req, tokens):
        out = self._catch_up(now)
        if self.waiting and self.waiting_tokens + tokens > self.max_tokens:
            out.append(self._cut())   # the new request would overflow: send what's there
        self.waiting.append((now, req))
        self.waiting_tokens += tokens
        if len(self.waiting) == self.max_batch or self.waiting_tokens >= self.max_tokens:
            out.append(self._cut())   # full by count or by tokens
        return out

    def tick(self, now):
        return self._catch_up(now)

    def flush(self):
        return [self._cut()] if self.waiting else []

#include <string>
#include <utility>
#include <vector>
using namespace std;

using Batch = vector<string>;

// Groups requests into batches by count, by token budget and by waiting time.
// It has no clock of its own: time moves only when a method is called.
class Batcher {
    int max_batch;
    long long max_tokens, max_wait;
    vector<pair<long long, string>> waiting;   // (arrival time, request id), oldest first
    long long waiting_tokens = 0;

    Batch cut() {
        Batch batch;
        for (auto& [t, req] : waiting) batch.push_back(req);
        waiting.clear();
        waiting_tokens = 0;
        return batch;
    }

    // Every rule that depends on time runs here, at the top of every call.
    vector<Batch> catch_up(long long now) {
        if (!waiting.empty() && waiting.front().first + max_wait <= now) return {cut()};
        return {};
    }

public:
    Batcher(int max_batch, long long max_tokens, long long max_wait)
        : max_batch(max_batch), max_tokens(max_tokens), max_wait(max_wait) {}

    vector<Batch> submit(long long now, const string& req, long long tokens) {
        vector<Batch> out = catch_up(now);
        if (!waiting.empty() && waiting_tokens + tokens > max_tokens)
            out.push_back(cut());              // the new request would overflow: send what's there
        waiting.push_back({now, req});
        waiting_tokens += tokens;
        if ((int)waiting.size() == max_batch || waiting_tokens >= max_tokens)
            out.push_back(cut());              // full by count or by tokens
        return out;
    }

    vector<Batch> tick(long long now) { return catch_up(now); }

    vector<Batch> flush() {
        if (waiting.empty()) return {};
        return {cut()};
    }
};

import java.util.*;

// Groups requests into batches by count, by token budget and by waiting time.
// It has no clock of its own: time moves only when a method is called.
class Batcher {
    private record Request(long arrival, String id) {}

    private final int maxBatch;
    private final long maxTokens, maxWait;
    private List<Request> waiting = new ArrayList<>();   // oldest first
    private long waitingTokens = 0;

    Batcher(int maxBatch, long maxTokens, long maxWait) {
        this.maxBatch = maxBatch;
        this.maxTokens = maxTokens;
        this.maxWait = maxWait;
    }

    private List<String> cut() {
        List<String> batch = new ArrayList<>();
        for (Request r : waiting) batch.add(r.id());
        waiting = new ArrayList<>();
        waitingTokens = 0;
        return batch;
    }

    // Every rule that depends on time runs here, at the top of every call.
    private List<List<String>> catchUp(long now) {
        List<List<String>> out = new ArrayList<>();
        if (!waiting.isEmpty() && waiting.get(0).arrival() + maxWait <= now) out.add(cut());
        return out;
    }

    List<List<String>> submit(long now, String req, long tokens) {
        List<List<String>> out = catchUp(now);
        if (!waiting.isEmpty() && waitingTokens + tokens > maxTokens)
            out.add(cut());                    // the new request would overflow: send what's there
        waiting.add(new Request(now, req));
        waitingTokens += tokens;
        if (waiting.size() == maxBatch || waitingTokens >= maxTokens)
            out.add(cut());                    // full by count or by tokens
        return out;
    }

    List<List<String>> tick(long now) { return catchUp(now); }

    List<List<String>> flush() {
        List<List<String>> out = new ArrayList<>();
        if (!waiting.isEmpty()) out.add(cut());
        return out;
    }
}

The clock never runs by itself: if nobody calls tick, a due batch leaves at the next call. Real servers call tick from a timer. In an interview, say that out loud; the tests drive time with explicit calls.

The other building blocks

• Streaming parsers. A streamed response arrives in chunks cut anywhere, even in the middle of a line. Keep the unfinished tail in a buffer and act only on complete lines. This short version copies the tail again on every chunk; when one line can be huge, keep a list of pieces and join them once a newline arrives, so each chunk costs O(its length).

  class LineBuffer:
      def __init__(self):
          self.partial = ""
  
      def feed(self, chunk):
          *lines, self.partial = (self.partial + chunk).split("\n")
          return lines            # only complete lines; the tail waits
  
  b = LineBuffer()
  b.feed("data: he")              # []
  b.feed("llo\ndata: wor")        # ['data: hello']

• Tokens-per-minute limits. Two sliding windows per key, one over requests and one over tokens, and a request passes only if it fits both. “When can I retry?” is either now or the moment one of the recorded requests turns 60 seconds old. See rolling counters.

• Prompt caches. Prompts are token lists, and prompts that start the same share work. Store them in a trie of tokens, and when memory runs out evict the least recently used leaf first, counting capacity in tokens.

• Sampling. Turn logits into probabilities with a softmax, subtracting the largest logit first so exp can’t overflow, and dividing by the temperature. Apply the filters in a fixed order (temperature, then top-k, then top-p), and draw from a seeded generator like random.Random(seed) so every run is reproducible.

• Eval harnesses. Results stream in and retries replace earlier results, so key them by (case, run). Decide whether a pass rate averages runs or cases; levels love to change that.

• Napkin math. Capacity estimates: requests per second at peak, storage over a retention period, GPUs needed with headroom. Use integers, keep units straight, and round once, at the end.

Why it’s O(1) amortized

Each request is appended once and leaves in exactly one batch, so the work per request is O(1), plus O(batch size) to build the batch it leaves in, which is paid once per request overall. The catch-up step looks only at the oldest request, O(1), because a timeout sends everything that’s waiting. waiting_tokens is kept as a running sum, so the budget check never adds up the queue. Space is O(requests waiting), at most max_batch.

Common mistakes

self.waiting.append(req); out = self._catch_up(now)   # ✗ joins a late batch
out = self._catch_up(now); self.waiting.append(req)   # ✓ time first

if len(self.waiting) >= limit: ...             # ✗ counts requests
if self.waiting_tokens + tokens > budget: ...  # ✓ counts what fills the GPU

for line in chunk.split("\n"): handle(line)     # ✗ the last piece may be half a line
*lines, self.partial = (self.partial + chunk).split("\n")   # ✓ keep the tail

u = random.random()                 # ✗ shared, unseeded state
self.rng = random.Random(seed)      # ✓ one generator per sampler, seeded

Variations

• Priority lanes. Interactive traffic goes first, bulk jobs wait, and a bulk request that has waited too long gets promoted. Promotion is one more thing to catch up on at the top of every call.
• Splitting by budget. Instead of flushing before an overflow, cut the longest run from the front of the queue that fits the budget, and repeat while enough is waiting.
• Resuming streams. After a reconnect the server may repeat events. Remember the last event id you processed and drop anything not newer.
• Pinned cache entries. System prompts every request shares shouldn’t be evicted: count pins per node and skip pinned nodes when choosing a victim.
• Error bars. A pass rate from 10 runs means less than one from 10,000. A Wilson score interval gives a range that stays sensible at 0% and 100%.

Climb the ladder

Our AI infrastructure problems in ladder order. Each one is a multi-level OA.

1. Napkin math: peak load, storage and GPU counts, rounding once at the end.
2. Streaming response parser: buffers, complete lines, and resuming after a reconnect.
3. Dynamic request batcher: count, timeout, token budget and priority lanes.
4. Eval harness: latest result per run, flaky cases, per-case averages, confidence intervals.
5. Tokens-per-minute limiter: two windows per key, per-key limits, and retry times.
6. Next-token sampler: greedy, temperature, top-k and top-p with a seeded generator.
7. Prompt prefix cache: a token trie with statistics, LRU eviction of leaves, and pins.