Build a tiny model server you can read end to end

"""Keep the copy of the model, the chip and the 400 invented chats inside ../widget.html equal to the ones in tiny_server.py.

    python3 embed_widget.py          rewrite the data block in widget.html
    python3 embed_widget.py --check  exit 1 if widget.html differs (test_widget.py runs the same check)
"""
import json
import re
import sys
from pathlib import Path

import show
import tiny_server as ts

WIDGET = Path(__file__).resolve().parent.parent / "widget.html"
CONFIG_BLOCK = re.compile(r'(<script type="application/json" id="ts-config">)(.*?)(</script>)', re.S)


def config():
    """What the page needs to run the same toy server: the model, the chip, the drawn chats and the settings it offers."""
    b = ts.tail.batching
    return {
        "params": b.PARAMS, "bytesPerWeight": b.FP16_BYTES, "kvBytes": b.KV_BYTES_PER_TOKEN, "peak": b.A100_PEAK_FLOPS, "bandwidth": b.A100_BANDWIDTH,
        "pageTokens": ts.PAGE_TOKENS, "minQueue": ts.MIN_QUEUE, "columns": show.COLUMNS, "gaps": ts.WORKLOAD["gaps"], "prompts": ts.WORKLOAD["prompts"], "replies": ts.WORKLOAD["replies"],
        "rates": show.RATES, "pageStops": show.PAGE_STOPS, "limitStops": show.LIMIT_STOPS, "rerolls": show.REROLLS, "opening": show.OPENING,
    }


def embed(html):
    """html with the data block rewritten from tiny_server.py."""
    settings = json.dumps(config(), ensure_ascii=False, separators=(", ", ": "))
    return CONFIG_BLOCK.sub(lambda m: m.group(1) + "\n" + settings + "\n" + m.group(3), html, count=1)


if __name__ == "__main__":
    current = WIDGET.read_text(encoding="utf-8")
    updated = embed(current)
    if "--check" in sys.argv:
        sys.exit(0 if updated == current else 1)
    WIDGET.write_text(updated, encoding="utf-8")
    print("widget.html data block now matches tiny_server.py")

"""Draw the four figures of lesson 31 from tiny_server.py, so that no number in a figure is typed by hand.

    python3 make_figs.py          rewrite ../figs/*.svg
    python3 make_figs.py --check  exit 1 if a figure on disk differs (test_tiny_server.py runs the same check)

All four figures are 600 units wide. Text is at least 22 units (20 for axis ticks) in ink or grey on white or a light fill; orange and green are used on marks
only, and every mark that colour tells apart is also named in words.
"""
import sys
from pathlib import Path
from xml.sax.saxutils import escape

import tiny_server as ts
from show import WALL_MOVES, WALLS, ms_text, seconds_text

FIGS = Path(__file__).resolve().parent.parent / "figs"
INK, GREY, ORANGE, GREEN, FILL, PANEL = "#0a1626", "#44566b", "#b36200", "#00855F", "#d2dae4", "#ebf0f6"
DASH = "10 8"
WINDOW = 3.0                                                                   # seconds of the heavy run that figure 2 shows


def num(x: float) -> str:
    return f"{x:.2f}".rstrip("0").rstrip(".")


# The SVG namespace is a name that every standalone .svg file must carry, not an address this program visits. It is written in two pieces so that a plain
# search of code/ for the start of a web address finds nothing (test_tiny_server.py runs that search).
NAMESPACE = "ht" + "tp://www.w3.org/2000/svg"


def frame(width: int, height: int, title: str, desc: str, body: list[str], defs: str = "") -> str:
    head = (f'<svg xmlns="{NAMESPACE}" viewBox="0 0 {width} {height}" width="{width}" height="{height}" role="img" '
            f'font-family="system-ui, sans-serif" font-size="22" fill="{GREY}">')
    return "\n".join([head, f"  <title>{escape(title)}</title>", f"  <desc>{escape(desc)}</desc>", *([f"  {defs}"] if defs else []),
                      '  <rect width="100%" height="100%" fill="#ffffff"/>', *("  " + line for line in body), "</svg>", ""])


def text(x: float, y: float, content: str, anchor: str = "start", color: str = INK, bold: bool = False, size: int = 22) -> str:
    extra = (f' text-anchor="{anchor}"' if anchor != "start" else "") + (' font-weight="bold"' if bold else "") + (f' font-size="{size}"' if size != 22 else "")
    return f'<text x="{num(x)}" y="{num(y)}"{extra} fill="{color}">{escape(content)}</text>'


ARROW = (f'<defs><marker id="head" viewBox="0 0 10 10" refX="9" refY="5" markerWidth="7" markerHeight="7" orient="auto-start-reverse">'
         f'<path d="M0 0L10 5L0 10z" fill="{INK}"/></marker></defs>')


def arrow(x1: float, y1: float, x2: float, y2: float, both: bool = False) -> str:
    start = ' marker-start="url(#head)"' if both else ""
    return f'<path d="M{num(x1)} {num(y1)}L{num(x2)} {num(y2)}" stroke="{INK}" stroke-width="3"{start} marker-end="url(#head)"/>'


def box(x: float, y: float, w: float, h: float, badge: int, lines: list[str]) -> list[str]:
    out = [f'<rect x="{x}" y="{y}" width="{w}" height="{h}" fill="{PANEL}" stroke="{INK}" stroke-width="2"/>',
           f'<circle cx="{x + 26}" cy="{y + 28}" r="16" fill="{INK}"/>', text(x + 26, y + 36, str(badge), "middle", "#ffffff", True)]
    out.append(text(x + 52, y + 36, lines[0], bold=True))
    out.extend(text(x + 16, y + 36 + 30 * k, line, color=GREY) for k, line in enumerate(lines[1:], 1))
    return out


# --- figure 1: the server as a kitchen -----------------------------------------------------------------------------------------

def fig_kitchen() -> str:
    body = box(20, 16, 560, 82, 6, [f"load generator: {ts.HEAVY} chats a second", "the dining room sends the orders"])
    body.append(arrow(135, 98, 135, 142))
    body += box(20, 146, 230, 136, 1, ["Queue", "the order rail", "chats wait here", "queue"])
    body += box(350, 146, 230, 136, 3, ["Batcher", "the cook", f"{ts.LIMIT} chats at most", "serve()"])
    body += box(350, 346, 230, 136, 4, ["Cache", "the shelf", f"{ts.PAGES} pages", "Pages"])
    body += box(20, 346, 230, 136, 5, ["Metrics", "the clipboard", "p50, p99, queue", "report()"])
    body.append(arrow(250, 214, 350, 214))
    body.append(text(300, 196, "2 admit", "middle"))
    body.append(arrow(465, 282, 465, 346, both=True))
    body.append(text(449, 322, "pages", "end"))
    body.append(arrow(372, 282, 250, 360))
    desc = (f"A map of the server as a kitchen, six numbered parts. The load generator (6), the dining room, sends {ts.HEAVY} chats a second into the queue, the order rail (1). The batcher, "
            f"the cook (3), admits a chat when a place and pages are free (2) and steps up to {ts.LIMIT} chats at once. The cache, the shelf (4), holds {ts.PAGES} pages of "
            f"{ts.PAGE_TOKENS} tokens, taken and returned by the batcher. The metrics, the clipboard (5), record each step.")
    return frame(600, 498, "A tiny model server as a kitchen: queue, batcher, cache and metrics", desc, body, ARROW)


# --- figure 2: the cache fills ------------------------------------------------------------------------------------------------

def fig_cache_fills() -> str:
    run = ts.serve(ts.WORKLOAD, ts.HEAVY)
    steps = [s for s in run.steps if s.start <= WINDOW]
    full = next(s.start for s in run.steps if "cache" in s.blocked)
    x0, x1 = 84, 580
    px = lambda t: x0 + t / WINDOW * (x1 - x0)
    top, base = 84, 270
    py = lambda pages: base - pages / ts.PAGES * (base - top)
    qtop, qbase, qmax = 356, 506, 50
    qy = lambda chats: qbase - chats / qmax * (qbase - qtop)
    body = [text(20, 40, "cache pages promised", color=GREY)]
    for tick in (0, ts.PAGES // 2, ts.PAGES):
        body.append(f'<path d="M{x0} {num(py(tick))}H{x1}" stroke="{FILL}" stroke-width="2"/>')
        body.append(text(x0 - 8, py(tick) + 7, str(tick), "end", GREY, size=20))
    body.append(f'<path d="M{x0} {num(py(ts.PAGES))}H{x1}" stroke="{ORANGE}" stroke-width="4" stroke-dasharray="{DASH}"/>')
    body.append(text(x0 + 8, py(ts.PAGES) - 12, f"{ts.PAGES} pages: the cache is full", bold=True))
    points = " L".join(f"{num(px(s.start))} {num(py(s.held + s.owed))}" for s in steps)
    body.append(f'<path d="M{points}" fill="none" stroke="{INK}" stroke-width="3"/>')
    body.append(f'<path d="M{num(px(full))} {base}V{top}" stroke="{GREEN}" stroke-width="4"/>')
    body.append(text(px(full) + 12, base - 30, f"cache full at {seconds_text(full)} s", bold=True))
    body.append(text(20, 326, "chats waiting in the queue", color=GREY))
    for tick in (0, qmax // 2, qmax):
        body.append(f'<path d="M{x0} {num(qy(tick))}H{x1}" stroke="{FILL}" stroke-width="2"/>')
        body.append(text(x0 - 8, qy(tick) + 7, str(tick), "end", GREY, size=20))
    points = " L".join(f"{num(px(s.start))} {num(qy(s.queue))}" for s in steps)
    body.append(f'<path d="M{points}" fill="none" stroke="{INK}" stroke-width="3"/>')
    body.append(f'<path d="M{num(px(full))} {qbase}V{qtop}" stroke="{GREEN}" stroke-width="4"/>')
    last = steps[-1]
    body.append(text(x1 - 8, qy(last.queue) - 14, f"{last.queue} waiting at {WINDOW:g} s", "end", bold=True))
    for base_y in (base, qbase):
        body.append(f'<path d="M{x0} {base_y}H{x1}" stroke="{GREY}" stroke-width="2"/>')
    for tick in range(int(WINDOW) + 1):
        body.append(text(px(tick), qbase + 30, str(tick), "middle", GREY, size=20))
    body.append(text((x0 + x1) / 2, qbase + 62, "seconds since the first chat", "middle", GREY))
    desc = (f"The first {WINDOW:g} seconds of the heavy run, {ts.HEAVY} chats a second. Top: cache pages promised (held plus owed) climb to the {ts.PAGES}-page limit and stay at it from "
            f"{seconds_text(full)} s, marked by a green line. Bottom: the queue is empty until then and has {last.queue} chats waiting at {WINDOW:g} s.")
    return frame(600, 576, f"The cache is full at {seconds_text(full)} s and the queue grows", desc, body)


# --- figure 3: the first word against the words between ------------------------------------------------------------------------

def fig_two_waits() -> str:
    light, heavy = ts.load_test(ts.LIGHT), ts.load_test(ts.HEAVY)
    x0 = 20
    panels = (("p99 wait for the first word", "ttft_p99", 3.5, 480, lambda v: f"{seconds_text(v)} s"),
              ("p99 time between words", "tbt_p99", 0.035, 480, lambda v: f"{ms_text(v)} ms"))
    body, y, facts = [], 36, []
    for title, key, top_value, longest, show in panels:
        body.append(text(20, y, title, bold=True))
        for k, (name, m, color) in enumerate(((f"{ts.LIGHT} chats a second", light, GREEN), (f"{ts.HEAVY} chats a second", heavy, ORANGE))):
            top = y + 40 + k * 84
            width = max(m[key] / top_value * longest, 3)
            body.append(text(20, top - 8, name, color=GREY))
            body.append(f'<rect x="{x0}" y="{top}" width="{num(width)}" height="34" fill="{PANEL}" stroke="{color}" stroke-width="4"/>')
            body.append(text(x0 + width + 12, top + 25, show(m[key]), bold=True))
        facts.append(f"{title}: {show(light[key])} at {ts.LIGHT} chats a second, {show(heavy[key])} at {ts.HEAVY}")
        y += 214
    desc = ("Two bar charts, each on its own scale, for the same toy server at a light and a heavy rate. " + "; ".join(facts) +
            f". At {ts.PAGES} pages and {ts.LIMIT} chats, the wait for the first word grows about {round(heavy['ttft_p99'] / light['ttft_p99'], -1):.0f} times; the time between words, every gap of every chat pooled, grows by about half.")
    return frame(600, y - 28, "At the heavy rate the first word waits far longer and the gaps between words grow by half", desc, body)


# --- figure 4: the wall moves ------------------------------------------------------------------------------------------------

def fig_walls() -> str:
    rows = [(pages, limit, ts.load_test(ts.HEAVY, pages, limit)) for pages, limit in WALL_MOVES]
    longest, top_value, x0 = 290, 3.5, 20
    body, y, facts = [], 36, []
    for pages, limit, m in rows:
        wall = WALLS[ts.bottleneck(m)]
        body.append(text(20, y, f"{pages} pages, up to {limit} chats", bold=True))
        width = max(m["ttft_p99"] / top_value * longest, 3)
        body.append(f'<rect x="{x0}" y="{y + 16}" width="{num(width)}" height="60" fill="{PANEL}" stroke="{INK}" stroke-width="2"/>')
        body.append(text(x0 + width + 14, y + 44, f"{seconds_text(m['ttft_p99'])} s", bold=True))
        body.append(text(x0 + width + 14, y + 72, f"wall: {wall}"))
        facts.append(f"{pages} pages and {limit} chats: {seconds_text(m['ttft_p99'])} s, wall {wall}")
        y += 118
    desc = (f"The p99 wait for the first word at {ts.HEAVY} chats a second, for four settings of cache pages and batch limit, on one scale. " + "; ".join(facts) +
            ". More cache moves the wall from the cache to the batch limit; more cache again changes nothing.")
    return frame(600, y - 8, "Changing one setting moves the wall", desc, body)


FIGURES = {
    "kitchen.svg": fig_kitchen,
    "cache-fills.svg": fig_cache_fills,
    "two-waits.svg": fig_two_waits,
    "wall-moves.svg": fig_walls,
}


def main() -> int:
    stale = []
    for name, draw in FIGURES.items():
        path, svg = FIGS / name, draw()
        if "--check" in sys.argv:
            if not path.exists() or path.read_text(encoding="utf-8") != svg:
                stale.append(name)
        else:
            FIGS.mkdir(exist_ok=True)
            path.write_text(svg, encoding="utf-8")
    if "--check" in sys.argv:
        print("figures differ from tiny_server.py: " + ", ".join(stale) if stale else "figures match tiny_server.py")
        return 1 if stale else 0
    print("wrote " + ", ".join(FIGURES))
    return 0


if __name__ == "__main__":
    sys.exit(main())

"""Print every number the lesson quotes. Run: python3 show_numbers.py"""
import statistics

import show
import tiny_server as ts

W = ts.WORKLOAD
ARRIVE = ts.intended_times(W["gaps"], ts.HEAVY)


def runs() -> dict:
    """The settings the lesson walks through, each as (run, metrics)."""
    out = {}
    for name, rate, pages, limit in (("light", ts.LIGHT, ts.PAGES, ts.LIMIT), ("heavy", ts.HEAVY, ts.PAGES, ts.LIMIT), ("480 pages", ts.HEAVY, 480, ts.LIMIT),
                                     ("960 pages", ts.HEAVY, 960, ts.LIMIT), ("480 pages, limit 24", ts.HEAVY, 480, 24)):
        run = ts.serve(W, rate, pages, limit)
        out[name] = (run, ts.report(W, rate, run))
    return out


def heavy_facts(run) -> dict:
    """What the heavy run does while the cache is full: the queue at set times, and how many chats waited."""
    first_full = next(k for k, s in enumerate(run.steps) if "cache" in s.blocked)
    at = lambda t: next(s for s in reversed(run.steps) if s.start <= t).queue
    waits = [f - a for f, a in zip(run.first, ARRIVE)]
    return {"full_step": first_full + 1, "full": run.steps[first_full], "queue_at": {t: at(t) for t in (0.5, 1, 2, 3, 5, 10)},
            "over_1s": sum(1 for w in waits if w > 1), "longest": max(waits), "arrive_end": ARRIVE[-1], "end": run.steps[-1].start + run.steps[-1].seconds,
            "steps": len(run.steps), "full_steps_first_3s": sum(1 for s in run.steps if s.start < 3 and "cache" in s.blocked), "steps_first_3s": sum(1 for s in run.steps if s.start < 3)}


def main() -> None:
    pages = [ts.pages_for(p + r) for p, r in zip(W["prompts"], W["replies"])]
    print(f"workload: {ts.N_CHATS} chats from lesson 27's seed; prompts {min(W['prompts'])} to {max(W['prompts'])} tokens; replies {min(W['replies'])} to {max(W['replies'])}; "
          f"whole chats {min(p + r for p, r in zip(W['prompts'], W['replies']))} to {max(p + r for p, r in zip(W['prompts'], W['replies']))} tokens, mean {statistics.fmean(p + r for p, r in zip(W['prompts'], W['replies'])):.1f}")
    print(f"pages per chat: mean {statistics.fmean(pages):.2f}, most {max(pages)}; the cache: {ts.PAGES} pages = {ts.PAGES * 2} MiB (a page is 16 tokens x 128 KiB = 2 MiB); a real card holds 25,720")
    print(f"step time with 1 chat and a 34-token prompt: {ts.step_seconds([], 34) * 1000:.2f} ms; with 16 chats of 220 tokens: {ts.step_seconds([220] * 16, 0) * 1000:.2f} ms")
    print(f"the 400 chats arrive over {ARRIVE[-1]:.1f} s at {ts.HEAVY} a second and over {ts.intended_times(W['gaps'], ts.LIGHT)[-1]:.1f} s at {ts.LIGHT}")
    print()
    for name, (run, m) in runs().items():
        print(f"{name:20} wall {ts.bottleneck(m):6} first word p50 {m['ttft_p50']:.3f} p99 {m['ttft_p99']:.3f} s | between words p50 {m['tbt_p50'] * 1000:.1f} p99 {m['tbt_p99'] * 1000:.1f} ms | "
              f"queue {m['queue_max']} pages {m['pages_max']} | full at {m['cache_full_at']} | cache steps {m['cache_steps']} batch steps {m['batch_steps']} | served {m['served']:.1f} a second | ends {run.steps[-1].start + run.steps[-1].seconds:.2f} s, {len(run.steps)} steps")
    print()
    print("the two kinds of TBT (ms): every gap pooled (the toy) / one average per chat (p50, p99 of 400), and the steps behind the gaps")
    for name, (run, m) in runs().items():
        per_chat = [(d - f) / (r - 1) for d, f, r in zip(run.done, run.first, W["replies"])]
        alone, reading = [s.seconds for s in run.steps if not s.joined], [s.seconds for s in run.steps if s.joined]
        print(f"{name:20} per gap {m['tbt_p50'] * 1000:.1f} / {m['tbt_p99'] * 1000:.1f} | per chat {ts.percentile(per_chat, 50) * 1000:.1f} / {ts.percentile(per_chat, 99) * 1000:.1f} | "
              f"steps with no new chat {len(alone)}, mean {statistics.fmean(alone) * 1000:.2f}, longest {max(alone) * 1000:.2f} | steps that read prompts {len(reading)}, "
              f"mean {statistics.fmean(reading) * 1000:.2f}, longest {max(reading) * 1000:.2f}")
    print()
    print("heavy run:", heavy_facts(runs()["heavy"][0]))
    print()
    print("widget shuffles at 40 chats a second (p99 first word, wall):")
    for pages_, limit in [(ts.PAGES, ts.LIMIT), *show.WALL_MOVES[1:]]:
        ds = [show.describe(ts.HEAVY, pages_, limit, k) for k in range(len(show.REROLLS))]
        print(f"  {pages_} pages, limit {limit}: p99 {min(d['metrics']['ttft_p99'] for d in ds):.2f} to {max(d['metrics']['ttft_p99'] for d in ds):.2f} s; walls", [d["lines"][2] for d in ds])


if __name__ == "__main__":
    main()

"""The settings, words and chart numbers of the widget: lesson 31 (build-a-tiny-model-server).

`describe` runs the tiny server at one setting and returns everything the widget shows. The widget (../widget.html) repeats `serve`, `report` and
`describe` in JavaScript, and test_widget.py runs both and compares them. Standard library only.
"""
from tiny_server import HEAVY, LIMIT, N_CHATS, PAGES, WORKLOAD, bottleneck, report, serve, tail

RATES, PAGE_STOPS, LIMIT_STOPS = [10, 20, 30, 40], [120, 240, 480, 960], [8, 16, 24, 48]     # the three sliders' stops
OPENING = {"rate": HEAVY, "pages": PAGES, "limit": LIMIT}                                       # the widget opens on the lesson's heavy run
REROLLS = tail.draw_rerolls(count=6, n=N_CHATS)                                                 # "Shuffle the chats": same chats, another order
COLUMNS = 48                                                                                    # slices of the run the chart shows
WALL_MOVES = [(240, 16), (480, 16), (960, 16), (480, 24)]                                       # the settings the lesson walks through, at the heavy rate
WALLS = {"none": "none", "cache": "the cache", "batch": "the batch limit", "both": "both"}

seconds_text = tail.seconds_text


def ms_text(seconds: float) -> str:
    return f"{seconds * 1000:.1f}"


def chats_text(n: int) -> str:
    return f"{n} chat" if n == 1 else f"{n} chats"


def columns(run, pages: int) -> dict:
    """The chart: for each slice of the run, the most pages promised (held plus owed) and the longest queue, and the slice where the cache first filled."""
    last = run.steps[-1]
    end = last.start + last.seconds
    promised, queue, full = [0] * COLUMNS, [0] * COLUMNS, -1
    for s in run.steps:
        c = min(COLUMNS - 1, int(s.start * COLUMNS / end))
        promised[c] = max(promised[c], s.held + s.owed)
        queue[c] = max(queue[c], s.queue)
        if full < 0 and "cache" in s.blocked:
            full = c
    return {"end": end, "promised": promised, "queue": queue, "full": full}


def headline(m: dict, rate: int, limit: int) -> str:
    slow = f"1\u00a0chat in 100 waits more than {seconds_text(m['ttft_p99'])}\u00a0s for its first word."
    wall = bottleneck(m)
    if wall == "none":
        return f"At {rate}\u00a0chats a second no wall binds. {slow}"
    if wall == "cache":
        return f"At {rate}\u00a0chats a second the cache fills at {seconds_text(m['cache_full_at'])}\u00a0s and new chats wait. {slow}"
    if wall == "batch":
        return f"At {rate}\u00a0chats a second new chats mostly wait for a place in the batch, which holds {limit}. {slow}"
    return f"At {rate}\u00a0chats a second new chats wait for cache pages and for a place in the batch. {slow}"


def key_table(m: dict) -> list:
    return [["First word", seconds_text(m["ttft_p50"]) + " s", seconds_text(m["ttft_p99"]) + " s"],
            ["Between words", ms_text(m["tbt_p50"]) + " ms", ms_text(m["tbt_p99"]) + " ms"]]


def key_lines(m: dict, pages: int) -> list:
    return [f"Longest queue: {chats_text(m['queue_max'])}", f"Cache pages used, most: {m['pages_max']} of {pages}", f"Wall: {WALLS[bottleneck(m)]}"]


def describe(rate: int, pages: int, limit: int, reroll: int = 0) -> dict:
    a, b = REROLLS[reroll]
    workload = tail.shuffled(WORKLOAD, a, b)
    run = serve(workload, rate, pages, limit)
    m = report(workload, rate, run)
    return {"head": headline(m, rate, limit), "table": key_table(m), "lines": key_lines(m, pages), "chart": columns(run, pages),
            "metrics": {k: m[k] for k in ("ttft_p50", "ttft_p99", "tbt_p50", "tbt_p99", "queue_max", "pages_max", "cache_full_at", "served", "cache_steps", "batch_steps")}}

"""A tiny model server you can read end to end: lesson 31 (build-a-tiny-model-server).

Three parts around a toy model: a cache of fixed-size pages (lesson 22), a batcher that steps every admitted chat together (lesson 24)
and metrics that report the waits (lesson 27). The "model" is a toy: a chat asks for a fixed number of tokens and each step gives every
admitted chat its next one, at the step time of lesson 24 (A100 80GB, Llama 3.1 8B). All of it runs in this one process on a simulated
clock, with no network code: the "load generator" is a list of send times, so nothing here can be pointed at a real service.
The widget (../widget.html) repeats `serve` and `report` in JavaScript; test_widget.py compares them. Standard library only.
"""
from __future__ import annotations

import importlib.util
import sys
from collections import deque, namedtuple
from pathlib import Path

# --- lesson 27's workload, percentiles and (through it) lesson 24's step time, loaded by path -----------------------------------

def _load_lesson_27():
    path = Path(__file__).resolve().parent.parent.parent / "why-average-speed-is-misleading" / "code" / "tail_sim.py"
    spec = importlib.util.spec_from_file_location("tail_sim_lesson27", path)
    module = importlib.util.module_from_spec(spec)
    sys.modules[spec.name] = module
    keep, sys.dont_write_bytecode = sys.dont_write_bytecode, True
    try:
        spec.loader.exec_module(module)
    finally:
        sys.dont_write_bytecode = keep
    return module


tail = _load_lesson_27()
step_seconds, percentile, intended_times = tail.step_seconds, tail.percentile, tail.intended_times

PAGE_TOKENS = 16                        # lesson 22 and lesson 28: tokens per cache page, 2 MiB of cache at 128 KiB a token
N_CHATS = 400                           # the first chats of lesson 27's seeded workload
WORKLOAD = {key: values[:N_CHATS] for key, values in tail.WORKLOAD.items()}
PAGES, LIMIT = 240, 16                  # the toy's cache (480 MiB; a real card holds 25,720 pages) and the most chats stepped together
LIGHT, HEAVY = 10, 40                   # chats sent per second


def pages_for(tokens: int) -> int:
    return -(-tokens // PAGE_TOKENS)


# --- the cache ----------------------------------------------------------------------------------------------------------------

class Pages:
    """Lesson 22's free list and page tables, as counts only: no keys and values are stored."""

    def __init__(self, total: int):
        self.total, self.free, self.tables = total, deque(range(total)), {}

    def grow(self, chat: int, tokens: int) -> None:
        """A chat holding `tokens` tokens needs enough pages for them: it takes one more whenever its last page is full."""
        table = self.tables.setdefault(chat, [])
        while len(table) * PAGE_TOKENS < tokens:
            table.append(self.free.popleft())

    def release(self, chat: int) -> None:
        self.free.extend(self.tables.pop(chat))

    def used(self) -> int:
        return self.total - len(self.free)


# --- the batcher ------------------------------------------------------------------------------------------------------------

# One row per step: when it started, how long it took, chats that got a token, of those how many just joined, chats still waiting,
# pages in use, pages promised but not yet taken, and what kept the next chat out ("cache", "batch", "cache and batch" or "" for nothing).
Step = namedtuple("Step", "start seconds stepped joined queue held owed blocked")
Run = namedtuple("Run", "first done steps")      # when each chat's first and last token came out, and every step


def serve(workload: dict, rate: float, pages: int = PAGES, limit: int = LIMIT) -> Run:
    prompts, replies, arrive = workload["prompts"], workload["replies"], intended_times(workload["gaps"], rate)
    n = len(prompts)
    assert pages >= pages_for(max(p + r for p, r in zip(prompts, replies))), "one chat needs more pages than the cache has"
    cache, queue, running, made = Pages(pages), deque(), [], [0] * n
    first, done, steps = [0.0] * n, [0.0] * n, []
    promised, now, sent, finished = 0, 0.0, 0, 0
    while finished < n:
        if not queue and not running and arrive[sent] > now:
            now = arrive[sent]                                          # nothing to do: wait for the next chat
        while sent < n and arrive[sent] <= now:
            queue.append(sent)                                          # 1. a request arrives and waits in the queue
            sent += 1
        joining, blocked = [], ""
        while queue:                                                    # 2. admit it when a place and all its pages are free
            need = pages_for(prompts[queue[0]] + replies[queue[0]])
            no_place, no_pages = len(running) + len(joining) == limit, promised + need > pages
            blocked = " and ".join(wall for wall, hit in (("cache", no_pages), ("batch", no_place)) if hit)
            if blocked:
                break
            promised += need
            joining.append(queue.popleft())
        seconds = step_seconds([prompts[i] + made[i] for i in running], sum(prompts[i] for i in joining))      # 3. one step, all together
        start, now = now, now + seconds
        running += joining
        for i in running:
            made[i] += 1
            cache.grow(i, prompts[i] + made[i])                         # 4. a page more when the last one is full
        for i in joining:
            first[i] = now
        stepped, ended = len(running), [i for i in running if made[i] == replies[i]]
        for i in ended:
            done[i] = now
            cache.release(i)                                            # 4. all its pages back when the chat ends
            promised -= pages_for(prompts[i] + replies[i])
            running.remove(i)
        finished += len(ended)
        steps.append(Step(start, seconds, stepped, len(joining), len(queue), cache.used(), promised - cache.used(), blocked))
    return Run(first, done, steps)


# --- the metrics ------------------------------------------------------------------------------------------------------------

def report(workload: dict, rate: float, run: Run) -> dict:
    """5. What the manager's clipboard shows: time to first token and time between tokens (every gap of every chat, pooled; p50, p99), the longest queue, the most pages used."""
    arrive = intended_times(workload["gaps"], rate)
    ttft = [f - a for f, a in zip(run.first, arrive)]
    tbt = [s.seconds for s in run.steps for _ in range(s.stepped - s.joined)]       # a running chat gets a word every step: its gap is the step's length
    cache_full = [s.start for s in run.steps if "cache" in s.blocked]
    return {"ttft_p50": percentile(ttft, 50), "ttft_p99": percentile(ttft, 99), "tbt_p50": percentile(tbt, 50), "tbt_p99": percentile(tbt, 99),
            "queue_max": max(s.queue for s in run.steps), "pages_max": max(s.held for s in run.steps),
            "cache_full_at": cache_full[0] if cache_full else None, "served": len(ttft) / max(run.done),
            "cache_steps": len(cache_full), "batch_steps": sum(1 for s in run.steps if "batch" in s.blocked)}


def load_test(rate: float, pages: int = PAGES, limit: int = LIMIT, workload: dict = WORKLOAD) -> dict:
    """6. The in-process load generator: chats arrive at random times that average `rate` a second, and the server reports."""
    return report(workload, rate, serve(workload, rate, pages, limit))


MIN_QUEUE = 4                           # a wall is named only if this many chats (1 in 100 of the 400) waited at once


def bottleneck(metrics: dict) -> str:
    """Which wall held the next chat out in more steps: "cache" (memory binds), "batch" (the batch limit binds), "both" (neither more than twice
    as often as the other) or "none" (fewer than MIN_QUEUE chats ever waited at once)."""
    c, b = metrics["cache_steps"], metrics["batch_steps"]
    return "none" if metrics["queue_max"] < MIN_QUEUE else "both" if 2 * min(c, b) >= max(c, b) else "cache" if c > b else "batch"

"""Follow chats through the tiny server, step by step, then read its metrics. Run it from this folder: python3 trace.py

It uses the lesson's own functions and adds no rules of its own: serve() runs the server and report() reads the clipboard.
"""
from tiny_server import HEAVY, LIGHT, LIMIT, PAGES, WORKLOAD, bottleneck, intended_times, load_test, pages_for, serve

MORE_PAGES = 480  # the cache of the second heavy run

arrive = intended_times(WORKLOAD["gaps"], HEAVY)
run = serve(WORKLOAD, HEAVY)
print(f"heavy load: {HEAVY} chats a second, a cache of {PAGES} pages, at most {LIMIT} chats per step")
print("\nthe first chats arrive")
for i in range(4):
    prompt, reply = WORKLOAD["prompts"][i], WORKLOAD["replies"][i]
    print(f"  chat {i}: sent at {arrive[i]:.3f} s, prompt {prompt} tokens, reply {reply} tokens, needs {pages_for(prompt + reply)} pages")


def show(k):
    s = run.steps[k]
    wait = f"  next chat waits for: {s.blocked}" if s.blocked else ""
    print(f"  {k + 1:>4}  {s.start:.3f} s  {s.seconds * 1000:>5.1f} ms  {s.stepped:>5}  {s.joined:>3}  {s.queue:>5}  {s.held:>3} + {s.owed:<3}  {PAGES - s.held - s.owed:>3}{wait}")


header = "  step  start     took     chats  new  queue  pages held + owed  left"
print(f"\nthe first steps\n{header}")
for k in range(5):
    show(k)
full = next(k for k, s in enumerate(run.steps) if "cache" in s.blocked)
print(f"\nthe cache fills at step {full + 1}\n{header}")
for k in range(full - 1, full + 3):
    show(k)

print(f"\nthe clipboard: the same {len(WORKLOAD['gaps'])} chats, three runs")
print("  run                  rate  first word p50 / p99   between words p50 / p99   queue  pages  wall")
for name, rate, pages in (("light", LIGHT, PAGES), ("heavy", HEAVY, PAGES), (f"heavy, {MORE_PAGES} pages", HEAVY, MORE_PAGES)):
    m = load_test(rate, pages)
    print(f"  {name:<19} {rate:>4}  {m['ttft_p50']:>8.3f} s {m['ttft_p99']:>7.3f} s  "
          f"{m['tbt_p50'] * 1000:>10.1f} ms {m['tbt_p99'] * 1000:>6.1f} ms  {m['queue_max']:>5}  {m['pages_max']:>5}  {bottleneck(m)}")

"""Guards for tiny_server.py: the cache against lesson 22's, the batcher's promises and counts, the server against lesson 24's, the metrics on a case worked by hand,
the story the lesson tells, and that every number the lesson, its trace and its figures print is the one the code gives."""
import importlib.util
import io
import itertools
import math
import re
import statistics
import subprocess
import sys
import xml.etree.ElementTree as ET
from contextlib import redirect_stdout
from functools import lru_cache
from pathlib import Path

import numpy as np
import pytest

import make_figs
import show
import tiny_server as ts

HERE = Path(__file__).resolve().parent
LESSON = HERE.parent
LESSONS = LESSON.parent
W = ts.WORKLOAD
N = ts.N_CHATS
SOURCE = (HERE / "tiny_server.py").read_text(encoding="utf-8")
secs, ms = show.seconds_text, show.ms_text
SHORT, LONG = 0.0078, 0.0085             # a step with no new chat takes between these many seconds, in all of the widget's runs


@lru_cache(maxsize=None)
def run(rate, pages=ts.PAGES, limit=ts.LIMIT):
    return ts.serve(W, rate, pages, limit)


@lru_cache(maxsize=None)
def metrics(rate, pages=ts.PAGES, limit=ts.LIMIT):
    return ts.report(W, rate, run(rate, pages, limit))


def load_by_path(name, relative):
    spec = importlib.util.spec_from_file_location(name, LESSONS / relative)
    module = importlib.util.module_from_spec(spec)
    sys.modules[name] = module
    keep, sys.dont_write_bytecode = sys.dont_write_bytecode, True
    try:
        spec.loader.exec_module(module)
    finally:
        sys.dont_write_bytecode = keep
    return module


def test_pages_follow_lesson_22_block_manager():
    """Guards the cache drifting from the pool lesson 22 explained: on one seeded script of chats that grow a token at a time and end, the page tables and the
    free list (its order too) of `Pages` equal those of lesson 22's `BlockManager` after every single step."""
    paged = load_by_path("paged_kv_lesson22", "paged-kv-cache/code/paged_kv.py")
    rng = np.random.default_rng(31)
    total = 40
    manager, pages = paged.BlockManager(total, ts.PAGE_TOKENS, 1), ts.Pages(total)
    lengths = [int(x) for x in rng.integers(1, 90, size=14)]            # tokens each chat will hold when it ends
    live, promised, nxt, tokens, checked = {}, 0, 0, {}, 0
    while nxt < len(lengths) or live:
        while nxt < len(lengths) and promised + ts.pages_for(lengths[nxt]) <= total:
            manager.add_request(nxt)
            live[nxt], tokens[nxt] = lengths[nxt], 0
            promised += ts.pages_for(lengths[nxt])
            nxt += 1
        for chat in list(live):
            manager.append(chat, np.zeros(1), np.zeros(1))
            tokens[chat] += 1
            pages.grow(chat, tokens[chat])
            if tokens[chat] == live[chat]:
                manager.free(chat)
                pages.release(chat)
                promised -= ts.pages_for(live.pop(chat))
            started = {c: t for c, t in manager.tables.items() if t}       # lesson 22 opens an empty table at admission, `Pages` at the first token
            assert pages.tables == started and list(pages.free) == list(manager.free_list), (chat, tokens[chat])
            checked += 1
    assert checked == sum(lengths) and pages.used() == 0


SETTINGS = [(10, ts.PAGES, ts.LIMIT), (40, ts.PAGES, ts.LIMIT), (40, 480, 24), (40, 52, 48), (40, 60, 3), (25, 120, 8)]


def test_every_chat_is_served_once_and_no_step_over_admits():
    """Guards a request admitted with no free page, pages leaked after a chat ends, a step serving a chat twice or skipping one, and a step with more chats than the
    limit. The cache of 52 pages only just fits the longest chat, so any over-admission would run the free list dry."""
    arrive = {rate: ts.intended_times(W["gaps"], rate) for rate, _, _ in SETTINGS}
    for rate, pages, limit in SETTINGS:
        r, steps = run(rate, pages, limit), run(rate, pages, limit).steps
        where = (rate, pages, limit)
        assert all(0 <= s.held and 0 <= s.owed and s.held + s.owed <= pages and s.stepped <= limit for s in steps), where
        assert sum(s.stepped for s in steps) == sum(W["replies"]) and sum(s.joined for s in steps) == N, where
        assert steps[-1].held == 0 and steps[-1].owed == 0, where
        assert all(a < f <= d for a, f, d in zip(arrive[rate], r.first, r.done)), where


def test_toy_gives_the_times_of_lesson_24_continuous_batching():
    """Guards the toy drifting from the scheduler lesson 24 explained: with a cache that never fills and 8 places, every chat's first word and last word must equal
    that lesson's `simulate`."""
    b = ts.tail.batching
    for rate in (10, 40):
        arrive = ts.intended_times(W["gaps"], rate)
        reference = b.simulate([b.Request(t, p, r) for t, p, r in zip(arrive, W["prompts"], W["replies"])], 8, "continuous")
        mine = run(rate, 5000, 8)
        assert mine.first == reference.first_token and mine.done == reference.finished and len(mine.steps) == len(reference.steps), rate


def test_a_case_worked_by_hand():
    """Guards the metrics and the wall labels: two chats of 10 prompt tokens and 2 and 3 reply tokens, sent together, with room for one chat at a time (batch), one
    chat's pages (cache) or both. The times are sums of lesson 24's step times, the gaps between words are the three single gaps (not each chat's average), and the
    label follows the steps held out once enough chats queued."""
    work = {"gaps": [0.0, 0.0], "prompts": [10, 10], "replies": [2, 3]}
    read, step11, step12 = ts.step_seconds([], 10), ts.step_seconds([11], 0), ts.step_seconds([12], 0)
    first = [read, read + step11 + read]
    done = [read + step11, first[1] + step11 + step12]
    for pages, limit, wall, held_out in [(100, 1, "batch", (0, 2)), (1, 8, "cache", (2, 0)), (1, 1, "both", (2, 2))]:
        r = ts.serve(work, 1, pages, limit)
        m = ts.report(work, 1, r)
        assert r.first == pytest.approx(first, abs=1e-12) and r.done == pytest.approx(done, abs=1e-12), wall
        assert (m["cache_steps"], m["batch_steps"]) == held_out and ts.bottleneck({**m, "queue_max": ts.MIN_QUEUE}) == wall and ts.bottleneck(m) == "none", wall
        assert m["ttft_p99"] == pytest.approx(first[1]) and m["ttft_p50"] == pytest.approx(first[0])
        assert (m["tbt_p50"], m["tbt_p99"]) == pytest.approx((step11, step12)) and m["queue_max"] == 1 and m["pages_max"] == 1       # gaps: step11, step11, step12
    wall_of = lambda c, b, q: ts.bottleneck({"cache_steps": c, "batch_steps": b, "queue_max": q})
    assert [wall_of(c, b, ts.MIN_QUEUE) for c, b in [(10, 5), (10, 4), (4, 10), (5, 10), (1, 0)]] == ["both", "cache", "batch", "both", "cache"]
    assert [wall_of(10, 0, q) for q in (0, 1, ts.MIN_QUEUE - 1, ts.MIN_QUEUE)] == ["none", "none", "none", "cache"] and ts.MIN_QUEUE == 4     # the floor: fewer than 4 chats queued names no wall
    assert [ts.pages_for(t) for t in (1, 16, 17, 150)] == [1, 1, 2, 10]


def test_a_cache_smaller_than_the_longest_chat_is_refused():
    """Guards a server that would wait forever: a chat that can never fit must stop the run at the start, not hang it."""
    longest = max(p + r for p, r in zip(W["prompts"], W["replies"]))
    assert ts.pages_for(longest) == 52
    with pytest.raises(AssertionError, match="needs more pages"):
        ts.serve(W, 40, 51, 16)


def test_the_story_reproduces():
    """Guards the lesson's claims: no wall at the light rate, the cache at the heavy one, the batch limit once the cache doubles, nothing more from 960 pages, both walls at
    a limit of 24, a first-word wait that grows about 80 times where the gap between words grows by about half, and the same walls under every one of the widget's 6 orders.
    The sweep of all 384 widget runs guards the reason given for the gaps: a step with no new chat takes about 8 ms, a step that reads prompts takes longer on average."""
    light, heavy, more, most, high = metrics(10), metrics(40), metrics(40, 480), metrics(40, 960), metrics(40, 480, 24)
    assert [ts.bottleneck(m) for m in (light, heavy, more, most, high)] == ["none", "cache", "batch", "batch", "both"]
    assert more == most and light["queue_max"] == 0 and light["cache_full_at"] is None and more["cache_full_at"] is None
    assert 0.38 < heavy["cache_full_at"] < 0.40 and heavy["queue_max"] > more["queue_max"] > high["queue_max"]
    assert heavy["ttft_p99"] > 50 * light["ttft_p99"] and 1.4 < heavy["tbt_p99"] / light["tbt_p99"] < 1.6
    assert heavy["ttft_p99"] > more["ttft_p99"] > high["ttft_p99"] and heavy["tbt_p99"] < more["tbt_p99"] < high["tbt_p99"]
    assert heavy["served"] < ts.HEAVY and more["served"] > heavy["served"]
    walls = {setting: set() for setting in show.WALL_MOVES}
    for k in range(len(show.REROLLS)):
        for pages, limit in show.WALL_MOVES:
            walls[(pages, limit)].add(show.describe(ts.HEAVY, pages, limit, k)["lines"][2])
    assert [walls[s] for s in show.WALL_MOVES[:3]] == [{"Wall: the cache"}, {"Wall: the batch limit"}, {"Wall: the batch limit"}]
    assert walls[(480, 24)] == {"Wall: both", "Wall: the cache"}
    for order in show.REROLLS:
        shuffled = ts.tail.shuffled(W, *order)
        for rate, pages, limit in itertools.product(show.RATES, show.PAGE_STOPS, show.LIMIT_STOPS):
            steps = ts.serve(shuffled, rate, pages, limit).steps
            alone, reading = [s.seconds for s in steps if not s.joined], [s.seconds for s in steps if s.joined]
            assert SHORT < min(alone) and max(alone) < LONG < statistics.fmean(reading) / 1.3, (order, rate, pages, limit)


def reading():
    prose = (LESSON / "lesson.md").read_text(encoding="utf-8")
    return prose, {name: (LESSON / "figs" / name).read_text(encoding="utf-8") for name in make_figs.FIGURES}


def per_chat_p99(r):
    """The other kind of TBT, one average per chat, as the lesson's Go deeper quotes it (p99 over the 400 chats)."""
    return ts.percentile([(d - f) / (n - 1) for d, f, n in zip(r.done, r.first, W["replies"])], 99)


def test_every_number_the_lesson_and_its_figures_quote_comes_from_the_code():
    """Guards prose, trace and figure numbers drifting from the code: each fragment is built from a value computed here and must appear in the lesson or the figure."""
    prose, figures = reading()
    light, heavy, more, most, high = metrics(10), metrics(40), metrics(40, 480), metrics(40, 960), metrics(40, 480, 24)
    arrive = ts.intended_times(W["gaps"], ts.HEAVY)
    heavy_run = run(40)
    full = next(k for k, s in enumerate(heavy_run.steps) if "cache" in s.blocked)
    s35, s4, s1 = heavy_run.steps[full], heavy_run.steps[3], heavy_run.steps[0]
    at3 = next(s for s in reversed(heavy_run.steps) if s.start <= 3.0).queue
    lines = len(SOURCE.splitlines())
    whole = [p + r for p, r in zip(W["prompts"], W["replies"])]
    ratio = heavy["ttft_p99"] / light["ttft_p99"]
    next_chat = sum(s.joined for s in heavy_run.steps[:full])
    shuffled = [{show.describe(40, 480, 24, k)["lines"][2] for k in range(len(show.REROLLS))}]
    mib = ts.PAGES * ts.PAGE_TOKENS * ts.tail.batching.KV_BYTES_PER_TOKEN // 2 ** 20
    table = [f"| {p} | {l} | {w} | {secs(m['ttft_p99'])} | {ms(m['tbt_p99'])} |" for (p, l), w, m in
             zip(show.WALL_MOVES, ("cache", "batch", "batch", "both"), (heavy, more, most, high))]
    rerolled = [(show.describe(40, 480, 16, k)["metrics"]["ttft_p99"], show.describe(40, 240, 16, k)["metrics"]["ttft_p99"] / show.describe(10, 240, 16, k)["metrics"]["ttft_p99"])
                for k in range(len(show.REROLLS))]
    low120 = ts.load_test(10, 120)
    few = ts.load_test(10, 240, 8)
    per_gap = f"{ms(light['tbt_p99'])} and {ms(heavy['tbt_p99'])} ms"
    per_chat = f"{ms(per_chat_p99(run(10)))} and {ms(per_chat_p99(heavy_run))} ms"
    sweep = len(show.REROLLS) * len(show.RATES) * len(show.PAGE_STOPS) * len(show.LIMIT_STOPS)
    gaps = sum(W["replies"]) - N
    wanted = {"lesson.md": [
        f"one {lines}-line Python program", f"`tiny_server.py` is {lines} lines", f"1 chat in 100 (p99) waits more than {secs(light['ttft_p99'])} s for its first word",
        f"the cache fills {secs(heavy['cache_full_at'])} s in", f"that wait is {secs(heavy['ttft_p99'])} s, about 80 times longer", f"the p99 gap between two words grows only from {ms(light['tbt_p99'])} to {ms(heavy['tbt_p99'])} ms",
        f"cuts the wait to {secs(more['ttft_p99'])} s", f"{ts.HEAVY} chats a second, {ts.PAGES} cache pages, {ts.LIMIT} chats at most per step",
        f"the cache fills at {secs(heavy['cache_full_at'])} s and 1 chat in 100 waits more than {secs(heavy['ttft_p99'])} s", f"the wait is {secs(light['ttft_p99'])} s",
        f"{secs(more['ttft_p99'])} s at 480 pages becomes {secs(min(r[0] for r in rerolled))} to {secs(max(r[0] for r in rerolled))} s, and about 80 times becomes "
        f"{min(r[1] for r in rerolled):.0f} to {max(r[1] for r in rerolled):.0f}",
        f"takes {ms(ts.step_seconds([220] * 16, 0))} ms where 1 chat of 34 tokens takes {ms(ts.step_seconds([34], 0))} ms",
        f"arrive within {arrive[3] * 1000:.0f} ms", f"chat 0 needs {ts.pages_for(whole[0])} (44 tokens at 16 a page), chat 2 needs {ts.pages_for(whole[2])}",
        f"Trace step 1 takes {ms(s1.seconds)} ms with one chat. Trace step 4 takes {ms(s4.seconds)} ms", f"{W['prompts'][2]} and {W['prompts'][3]} tokens",
        f"At trace step {full + 1}, {s35.held} pages are held and {s35.owed} more are owed", f"Only {ts.PAGES - s35.held - s35.owed} of {ts.PAGES} are left and the next chat in line needs "
        f"{ts.pages_for(whole[next_chat])}", f"with {s35.stepped} of the {ts.LIMIT} places used", f"reach the ceiling at {secs(heavy['cache_full_at'])} s",
        f"the queue reaches {at3} chats by 3 s", "about 80 times as long", f"the p99 gap between words grows by about half, from {ms(light['tbt_p99'])} to {ms(heavy['tbt_p99'])} ms",
        "A step with no new chat takes about 8 ms, so the slow gaps are the steps that read new prompts", "Every running chat waits through it",
        f"They peak at {more['pages_max']} pages in use", f"the p99 wait is {secs(more['ttft_p99'])} s", f"960 pages gives the same {secs(most['ttft_p99'])} s",
        f"cuts the wait to {secs(high['ttft_p99'])} s", f"at a small price: the p99 gap between words rises from {ms(heavy['tbt_p99'])} to {ms(more['tbt_p99'])} to {ms(high['tbt_p99'])} ms",
        f"{heavy['pages_max']} of {ts.PAGES} pages held", f"{more['pages_max']} of 480 held", f"Even the median chat waited {secs(heavy['ttft_p50'])} s",
        f"the p99 gap between words by about half. Averaging each chat's gaps would show {ms(per_chat_p99(run(10)))} to {ms(per_chat_p99(heavy_run))} ms and hide the slow steps",
        f"a limit of 24 cut the wait to {secs(high['ttft_p99'])} s while the p99 gap between words rose to {ms(high['tbt_p99'])} ms",
        f"prompts of {min(W['prompts'])} to {max(W['prompts'])} tokens, replies of {min(W['replies'])} to {max(W['replies'])}, and {sum(whole) / N:.0f} tokens on average",
        f"their {sum(whole):,} tokens", f"The toy's {ts.PAGES} pages are {mib} MiB", f"the {math.ceil(99 * N / 100)}th shortest and {N - math.ceil(99 * N / 100)} chats are slower",
        f"TBT pools every gap between two words of every chat, {gaps:,} of them", f"chat 0 is promised {ts.pages_for(whole[0])} pages when it is admitted, chat 2 is promised {ts.pages_for(whole[2])}",
        f"at least {ts.MIN_QUEUE} chats (`MIN_QUEUE`, 1 in 100 of the {N}) ever waited at once", "where one chat ever waited, would blame the batch limit",
        f"{low120['queue_max']} chats queue and the p99 wait is {secs(low120['ttft_p99'])} s", "At 120 pages the same rate still names the cache",
        f"held a chat out in {heavy['cache_steps']:,} steps and the batch limit in {heavy['batch_steps']}", f"the limit {more['batch_steps']} steps",
        f"the counts are {high['cache_steps']} and {high['batch_steps']}: both", f"{len(shuffled[0])}", "reads \"both\" in 4 of 6 orders and \"the cache\" in the others",
        f"finishes {heavy['served']:.1f} chats a second on average, fewer than the {ts.HEAVY} sent", f"| {ts.PAGES} | {ts.LIMIT} | cache | {secs(heavy['ttft_p99'])} | {ms(heavy['tbt_p99'])} |",
        "| Pages | Limit | Wall | p99 first word (s) | p99 gap (ms) |",
        f"at 10 and 40 chats a second the p99 is {per_gap}", f"Its p99 is {per_chat}", "records both: a histogram with one sample per streamed output", "chunked prefill, described on the vLLM Optimization page below",
        f"a sweep of the widget's {sweep} runs checks that a step with no new chat takes {SHORT * 1000:.1f} to {LONG * 1000:.1f} ms",
        "`vllm:inter_token_latency_seconds`, which records one sample per streamed output; the per-chat average above lines up with `vllm:request_time_per_output_token_seconds`", "large enough to use the GPU well and small enough to avoid much waste in most workloads",
        "Lesson 22 and [lesson 28](../what-model-server-settings-change/) call a page a block", "The toy knows each reply's length, so it can promise all of its pages; a real server cannot, and fills later (Go deeper)",
        "The dining room sends the orders", "The cook takes an order off the rail only if two things hold: the shelf has room for it, and a plate is free",
        "configuration/optimization/", "read 2026-10-20)", "Figure 3. At 240 pages and 16 chats, the wait for the first word grows about 80 times; the gap between words grows by about half",
        *table, f"{ts.pages_for(150)} pages", f"{most['pages_max']} of the 960 pages", "At 480 pages and a limit of 24"] +
        ["Peak use is %d pages of 960" % most["pages_max"], f"to {secs(high['ttft_p99'])} s, and raises the p99 gap between words from {ms(more['tbt_p99'])} to {ms(high['tbt_p99'])} ms",
         "the p99 gap between words only by about half", "Only steps that read new prompts are slow, and every running chat waits through them"],
        "kitchen.svg": [f"load generator: {ts.HEAVY} chats a second", "the dining room sends the orders", "the order rail", f"{ts.LIMIT} chats at most", f"{ts.PAGES} pages", "p50, p99, queue"],
        "cache-fills.svg": [f"cache full at {secs(heavy['cache_full_at'])} s", f"{at3} waiting at 3 s", f"{ts.PAGES} pages: the cache is full", "cache pages promised"],
        "two-waits.svg": [f"{secs(light['ttft_p99'])} s", f"{secs(heavy['ttft_p99'])} s", f"{ms(light['tbt_p99'])} ms", f"{ms(heavy['tbt_p99'])} ms", "about 80 times", "grows by about half"],
        "wall-moves.svg": [f"{p} pages, up to {l} chats" for p, l in show.WALL_MOVES] + [f"{secs(m['ttft_p99'])} s" for m in (heavy, more, high)] +
                          [f"wall: {w}" for w in ("the cache", "the batch limit", "both")],
    }
    assert 75 <= ratio <= 85 and 1.4 < heavy["tbt_p99"] / light["tbt_p99"] < 1.6 and lines < 150 and "25,720" in (LESSONS / "what-model-server-settings-change" / "lesson.md").read_text(encoding="utf-8")
    assert ts.MIN_QUEUE == N // 100 and few["queue_max"] == 1 and ts.bottleneck(few) == "none" and ts.bottleneck(low120) == "cache" and 1.0 < low120["ttft_p99"] < 1.3
    assert (round(min(r[0] for r in rerolled), 2), round(max(r[0] for r in rerolled), 2)) == (0.39, 1.22) and sweep == 384
    texts = {"lesson.md": prose, **figures}
    missing = [f"{name}: {fragment!r}" for name, fragments in wanted.items() for fragment in fragments if fragment not in texts[name]]
    assert not missing, "not found: " + "; ".join(missing)


def test_the_figures_are_the_ones_the_code_draws_and_follow_the_house_style():
    """Guards a figure going stale or off-palette: the files on disk equal what make_figs draws, the lesson shows each one, and each parses, has a title and a description,
    a white ground, no script, no web address, only the house colours and no text under 20 units."""
    prose, figures = reading()
    assert set(figures) == set(re.findall(r"\]\(figs/([\w-]+\.svg)\)", prose)) == set(make_figs.FIGURES)
    allowed = {make_figs.INK, make_figs.GREY, make_figs.ORANGE, make_figs.GREEN, make_figs.FILL, make_figs.PANEL, "#ffffff"}
    web_address = re.compile("ht" + r"tps?://")
    for name, svg in figures.items():
        assert svg == make_figs.FIGURES[name](), name
        root = ET.fromstring(svg)
        assert root.find("{%s}title" % make_figs.NAMESPACE) is not None and root.find("{%s}desc" % make_figs.NAMESPACE) is not None and 'fill="#ffffff"' in svg, name
        assert "<script" not in svg and not web_address.search(svg.replace(make_figs.NAMESPACE, "")) and root.get("width") == "600", name
        assert {c.lower() for c in re.findall(r"#[0-9a-fA-F]{6}\b", svg)} <= {c.lower() for c in allowed}, name
        assert all(int(size) >= 20 for size in re.findall(r'font-size="(\d+)"', svg)), name
    for path in HERE.glob("*.py"):
        assert not web_address.search(path.read_text(encoding="utf-8")), path.name


def table_rows(prose):
    return re.findall(r"^\| \d\. [^|]+\| `(\w+\.py):(\d+)`<br>`(.+?)` \|", prose, re.M)


def test_code_blocks_are_slices_of_the_real_file_and_the_table_points_at_its_lines():
    """Guards a `file:line` in the code table pointing at the wrong line, and a code block in Go deeper that no longer matches the file."""
    prose, _ = reading()
    blocks = re.findall(r"```python\n(.*?)\n```", prose, re.S)
    assert len(blocks) == 5 and all(block in SOURCE for block in blocks)
    rows = table_rows(prose)
    assert [int(n) for _, n, _ in rows] == [85, 90, 96, 57, 120, 131]
    for name, number, code in rows:
        line = (HERE / name).read_text(encoding="utf-8").splitlines()[int(number) - 1]
        assert re.sub(r"\s+# .*$", "", line).strip() == code, (name, number)
    assert "self.free.extend(self.tables.pop(chat))" in SOURCE.splitlines()[59] and "Line 60 gives them all back" in prose


def run_trace(source):
    out = io.StringIO()
    with redirect_stdout(out):
        exec(compile(source, "trace.py", "exec"), {"__name__": "__main__"})
    return out.getvalue()


def test_trace_is_the_block_in_the_lesson_and_the_your_turn_answer_is_right():
    """Guards the pasted run drifting from `python3 trace.py` (the real command, from the code folder) and the Your turn answer: changing line 7 to 960 pages must change
    only the label of the last row, and the details block must be that block."""
    prose, _ = reading()
    real = subprocess.run([sys.executable, "-B", "trace.py"], cwd=HERE, capture_output=True, text=True, check=True).stdout
    pasted = re.search(r"python3 trace\.py\n```\n\nPhone readers: slide the box sideways to see every column\.\n\n```text\n(.*?)\n```", prose, re.S).group(1)
    assert real == pasted + "\n"
    trace = (HERE / "trace.py").read_text(encoding="utf-8")
    assert trace.splitlines()[6] == "MORE_PAGES = 480  # the cache of the second heavy run" and "change line 7 to `MORE_PAGES = 960`" in prose
    changed = run_trace(trace.replace("MORE_PAGES = 480", "MORE_PAGES = 960", 1))
    start = real.index("the clipboard")
    assert changed[:start] == real[:start] and changed[start:] == real[start:].replace("heavy, 480 pages", "heavy, 960 pages")
    answer = re.search(r"<summary>What you should see</summary>.*?```text\n(.*?)\n```", prose, re.S).group(1)
    assert answer + "\n" == changed[start:] and "343" in answer