Serving a Transformer with FastAPI and Gradio: latency under load

import importlib.util
import subprocess
import sys

TFS = "git+https://github.com/xmarva/transformers-from-scratch@main"
if importlib.util.find_spec("tfs") is None:
    subprocess.check_call([sys.executable, "-m", "pip", "install", "-q", TFS])

from tfs import env
from tfs.viz import table, token_chips, use_course_style

env.quiet()  # no progress bars or environment warnings in the outputs
SEED = 42
env.seed_everything(SEED)
COLORS = use_course_style()
WORK_DIR = env.work_dir()

import asyncio
import math
import os
import random
import shutil
import time
from pathlib import Path

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import torch
from huggingface_hub import hf_hub_download

for package in ["sacrebleu", "httpx"]:  # Kaggle has them, Colab may not
    if importlib.util.find_spec(package) is None:
        subprocess.check_call([sys.executable, "-m", "pip", "install", "-q", package])
import httpx
import sacrebleu

print("CPU cores Python sees:", os.cpu_count())
print("threads PyTorch uses:", torch.get_num_threads())

from transformers import MarianMTModel, MarianTokenizer

OPUS = "Helsinki-NLP/opus-mt-en-ru"
opus_tokenizer = MarianTokenizer.from_pretrained(OPUS)
opus = MarianMTModel.from_pretrained(OPUS).eval()

DATA_REVISION = "540be5b93bf77cd0846b02d55f2558a5596887a9"
test_path = hf_hub_download("xmarva/tfs-tatoeba-en-ru", "data/test.parquet",
                            repo_type="dataset", revision=DATA_REVISION)
test = pd.read_parquet(test_path)
probe = list(test.sample(30, random_state=SEED).en)


@torch.inference_mode()
def opus_translate(sentence):
    batch = opus_tokenizer([sentence], return_tensors="pt")
    output = opus.generate(**batch, num_beams=4, max_new_tokens=128)
    return opus_tokenizer.decode(output[0], skip_special_tokens=True)


def latencies(translate, sentences):
    times = []
    for sentence in sentences:
        start = time.perf_counter()
        translate(sentence)
        times.append(time.perf_counter() - start)
    return np.array(times)


opus_times = latencies(opus_translate, probe)
n_opus = sum(parameter.numel() for parameter in opus.parameters())
print(f"opus-mt: {n_opus / 1e6:.0f}M parameters, "
      f"median {np.median(opus_times) * 1000:.0f} ms per sentence")

REPO = WORK_DIR / "course"
shutil.rmtree(REPO, ignore_errors=True)
URL = "https://github.com/xmarva/transformers-from-scratch"
QUIET = ["-c", "advice.detachedHead=false"]
subprocess.run(["git", *QUIET, "clone", "-q", "--depth", "1", "--branch", "perevod-v0.1.0",
                "--filter=blob:none", "--sparse", URL, str(REPO)], check=True)
subprocess.run(["git", "sparse-checkout", "set", "perevod"], cwd=REPO, check=True)
os.chdir(REPO)


def write(path, text):
    Path(path).parent.mkdir(parents=True, exist_ok=True)
    Path(path).write_text(text.lstrip("\n"), encoding="utf-8")


!{sys.executable} -m pip install -q ./perevod
importlib.invalidate_caches()

import perevod
from perevod import Translator

translator = Translator.from_pretrained(device="cpu")
n_ours = sum(parameter.numel() for parameter in translator.model.parameters())
ours_times = latencies(translator, probe)
print(f"perevod {perevod.__version__}: {n_ours / 1e6:.1f}M parameters, "
      f"median {np.median(ours_times) * 1000:.0f} ms")

from perevod import decoding

model, tokenizer = translator.model, translator.tokenizer
LONG = ("The European Central Bank raised interest rates by a quarter of a percentage "
        "point on Thursday, citing persistent inflation in the services sector.")
src = torch.tensor([decoding.source_ids(tokenizer, LONG, model.config.max_len)])
pad_id, bos, eos = decoding.special_ids(tokenizer)

with torch.inference_mode():
    start = time.perf_counter()
    memory = model.encode(src)
    encode_ms = (time.perf_counter() - start) * 1000
    written = torch.full((1, 1), bos)
    plain_steps = []
    for _ in range(100):  # past the end of the translation too: only the clock matters here
        start = time.perf_counter()
        next_id = model.decode(written, memory, src)[:, -1].argmax(dim=-1)
        plain_steps.append((time.perf_counter() - start) * 1000)
        written = torch.cat([written, next_id[:, None]], dim=1)
print(f"encoder: {encode_ms:.1f} ms for {src.shape[1]} tokens")
print(f"decoder step 1: {plain_steps[0]:.1f} ms, step 100: {plain_steps[99]:.1f} ms")

%%writefile perevod/src/perevod/decoding.py
# Greedy decoding and beam search of lesson 3, batched as in lesson 4, with lesson 5's KV
# cache: every step feeds the decoder only the newest token. A source longer than the
# model's positions is cut to fit, instead of failing inside the model.

import torch
import torch.nn.functional as F

from .model import apply_rope


def special_ids(tokenizer):
    pad_id = tokenizer.token_to_id("<pad>")
    return pad_id, tokenizer.token_to_id("<bos>"), tokenizer.token_to_id("<eos>")


def source_ids(tokenizer, sentence, limit):
    return tokenizer.encode(sentence).ids[:limit]


def pad(sequences, pad_id, device):
    longest = max(len(ids) for ids in sequences)
    batch = torch.full((len(sequences), longest), pad_id, device=device)
    for row, ids in enumerate(sequences):
        batch[row, : len(ids)] = torch.tensor(ids, device=device)
    return batch


def batches(sources, batch_size):
    # sentences in batches of similar length, so that little of a batch is padding
    order = sorted(range(len(sources)), key=lambda i: len(sources[i]))
    return [order[start : start + batch_size] for start in range(0, len(order), batch_size)]


def length_penalty(length, alpha=0.6):
    # Wu et al. (2016): ((5 + |y|) / 6) ** alpha
    return ((5 + length) / 6) ** alpha

%%writefile -a perevod/src/perevod/decoding.py


def attend(attention, x, cache, mask, angles):
    # one attention sublayer for the newest token; the cache holds everything before it
    q = attention.split_heads(attention.w_q(x))
    if angles is not None:  # RoPE: self-attention rotates queries and keys by position
        q = apply_rope(q, angles)
    if "fixed" not in cache:  # self-attention: add the new token's key and value
        k = attention.split_heads(attention.w_k(x))
        v = attention.split_heads(attention.w_v(x))
        if angles is not None:
            k = apply_rope(k, angles)
        if "k" in cache:
            k = torch.cat([cache["k"], k], dim=2)
            v = torch.cat([cache["v"], v], dim=2)
        cache["k"], cache["v"] = k, v
    out = F.scaled_dot_product_attention(q, cache["k"], cache["v"], attn_mask=mask)
    batch, _, length, _ = out.shape
    return attention.w_o(out.transpose(1, 2).reshape(batch, length, -1))

%%writefile -a perevod/src/perevod/decoding.py


def new_caches(model, memory):
    # per decoder layer: an empty self-attention cache, and the cross-attention's keys and
    # values, which depend only on the source and are computed once
    caches = []
    for layer in model.decoder:
        attention = layer.cross_attention
        k = attention.split_heads(attention.w_k(memory))
        v = attention.split_heads(attention.w_v(memory))
        caches.append({"self": {}, "cross": {"fixed": True, "k": k, "v": v}})
    return caches


def decode_step(model, token, memory_mask, caches, position):
    # the decoder of model.py, in the same order, for one new token; dropout is off anyway
    x = model.embed(token)
    angles = model.angles[position : position + 1]
    for layer, cache in zip(model.decoder, caches, strict=True):
        x = x + attend(layer.self_attention, layer.norm1(x), cache["self"], None, angles)
        x = x + attend(layer.cross_attention, layer.norm2(x), cache["cross"], memory_mask,
                       None)
        x = x + layer.feed_forward(layer.norm3(x))
    return (model.decoder_norm(x) @ model.embedding.weight.T)[:, -1]

%%writefile -a perevod/src/perevod/decoding.py


def reorder(caches, rows):
    # the cached keys and values follow their hypotheses: row r of the new caches is row
    # rows[r] of the old ones
    for cache in caches:
        for part in cache.values():
            for key in ("k", "v"):
                if key in part:
                    part[key] = part[key][rows]

%%writefile -a perevod/src/perevod/decoding.py


@torch.inference_mode()
def greedy(model, tokenizer, sentences, max_tokens=128, batch_size=256):
    model.eval()
    device = next(model.parameters()).device
    pad_id, bos, eos = special_ids(tokenizer)
    sources = []
    for sentence in sentences:
        sources.append(source_ids(tokenizer, sentence, model.config.max_len))
    max_tokens = min(max_tokens, model.config.max_len - 1)
    translations = [None] * len(sentences)
    for indices in batches(sources, batch_size):
        src = pad([sources[i] for i in indices], pad_id, device)
        memory_mask = (src != pad_id)[:, None, None, :]
        caches = new_caches(model, model.encode(src))
        token = torch.full((len(indices), 1), bos, device=device)
        written = []
        done = torch.zeros(len(indices), dtype=torch.bool, device=device)
        for position in range(max_tokens):
            scores = decode_step(model, token, memory_mask, caches, position)
            next_ids = scores.argmax(dim=-1)
            next_ids[done] = pad_id  # a finished sentence only collects padding
            written.append(next_ids)
            done |= next_ids == eos
            token = next_ids[:, None]
            if done.all():
                break
        written = torch.stack(written, dim=1)
        for row, i in enumerate(indices):
            ids = []
            for token_id in written[row].tolist():
                if token_id in (eos, pad_id):
                    break
                ids.append(token_id)
            translations[i] = tokenizer.decode(ids)
    return translations

%%writefile -a perevod/src/perevod/decoding.py


@torch.inference_mode()
def beam_search(model, tokenizer, sentences, beam_size=4, alpha=0.6, max_tokens=128,
                batch_size=64):
    model.eval()
    device = next(model.parameters()).device
    pad_id, bos, eos = special_ids(tokenizer)
    sources = []
    for sentence in sentences:
        sources.append(source_ids(tokenizer, sentence, model.config.max_len))
    max_tokens = min(max_tokens, model.config.max_len - 1)
    translations = [None] * len(sentences)
    for indices in batches(sources, batch_size):
        src = pad([sources[i] for i in indices], pad_id, device)
        memory_mask = (src != pad_id)[:, None, None, :]
        caches = new_caches(model, model.encode(src))
        # per sentence: (tokens so far, total log-probability, its row in the caches)
        beams = []
        for k in range(len(indices)):
            beams.append([([bos], 0.0, k)])
        finished = [[] for _ in indices]
        active = [True for _ in indices]

%%writefile -a perevod/src/perevod/decoding.py

        for position in range(max_tokens):
            rows, owners, last = [], [], []
            for k in range(len(indices)):
                if active[k]:
                    for tokens, _, row in beams[k]:
                        rows.append(row)
                        owners.append(k)
                        last.append(tokens[-1])
            if not rows:
                break
            reorder(caches, torch.tensor(rows, device=device))
            token = torch.tensor(last, device=device)[:, None]
            owner = torch.tensor(owners, device=device)
            scores = decode_step(model, token, memory_mask[owner], caches, position)
            top = scores.float().log_softmax(dim=-1).topk(beam_size)
            values, ids = top.values.tolist(), top.indices.tolist()

            row = 0
            for k in range(len(indices)):
                if not active[k]:
                    continue
                candidates = []
                for tokens, total, _ in beams[k]:
                    for log_p, token_id in zip(values[row], ids[row], strict=True):
                        candidates.append((tokens + [token_id], total + log_p, row))
                    row += 1
                candidates.sort(key=lambda candidate: candidate[1], reverse=True)
                beams[k] = []
                for tokens, total, parent in candidates:
                    if tokens[-1] == eos:
                        finished[k].append((tokens, total))
                    else:
                        beams[k].append((tokens, total, parent))
                    if len(beams[k]) == beam_size:
                        break
                if len(finished[k]) >= beam_size:
                    active[k] = False

%%writefile -a perevod/src/perevod/decoding.py

        for k, i in enumerate(indices):

            def normalized(candidate):
                tokens, total = candidate[0], candidate[1]
                return total / length_penalty(len(tokens) - 1, alpha)  # without <bos>

            best = max(finished[k] or beams[k], key=normalized)
            translations[i] = tokenizer.decode(best[0][1:])  # decode drops <eos>
    return translations

def load_module(name, path):
    spec = importlib.util.spec_from_file_location(name, path)
    module = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(module)
    return module


old = decoding  # perevod 0.1.0, as installed
new = load_module("perevod.decoding_new", "perevod/src/perevod/decoding.py")
print("old:", old.__file__)
print("new:", new.__file__)

from sacrebleu.utils import get_source_file

tatoeba = test.sample(500, random_state=SEED)
news_file = Path(get_source_file("wmt20", "en-ru"))
news = random.Random(SEED).sample(news_file.read_text(encoding="utf-8").splitlines(), 30)

rows = []
for name, sentences in [("Tatoeba", list(tatoeba.en)), ("WMT20 news", news)]:
    start = time.perf_counter()
    before = old.beam_search(model, tokenizer, sentences)
    middle = time.perf_counter()
    after = new.beam_search(model, tokenizer, sentences)
    end = time.perf_counter()
    assert after == before, name
    rows.append({"test set": name, "sentences": len(sentences),
                 "without the cache, s": middle - start, "with it, s": end - middle,
                 "faster": (middle - start) / (end - middle)})
speed = pd.DataFrame(rows)
display(table(speed))

with torch.inference_mode():
    memory_mask = (src != pad_id)[:, None, None, :]
    caches = new.new_caches(model, model.encode(src))
    token = torch.full((1, 1), bos)
    cached_steps = []
    for position in range(100):
        start = time.perf_counter()
        scores = new.decode_step(model, token, memory_mask, caches, position)
        next_id = scores.argmax(dim=-1)
        cached_steps.append((time.perf_counter() - start) * 1000)
        token = next_id[:, None]

fig, ax = plt.subplots(figsize=(8, 3.2))
ax.plot(range(1, 101), plain_steps, color=COLORS["dec"], label="recompute everything")
ax.plot(range(1, 101), cached_steps, color=COLORS["enc"], label="with the KV cache")
ax.set_xlabel("decoding step")
ax.set_ylabel("ms per step")
ax.set_ylim(bottom=0)
ax.set_title("One decoding step, on a CPU")
ax.legend(frameon=False)
plt.show()

from torch.ao.quantization import quantize_dynamic

engines = torch.backends.quantized.supported_engines
if torch.backends.quantized.engine == "none" and "qnnpack" in engines:  # ARM, e.g. a Mac
    torch.backends.quantized.engine = "qnnpack"
quantized = quantize_dynamic(model, {torch.nn.Linear}, dtype=torch.qint8)
sentences = list(tatoeba.en)
start = time.perf_counter()
fp32 = new.beam_search(model, tokenizer, sentences)
fp32_seconds = time.perf_counter() - start
start = time.perf_counter()
int8 = new.beam_search(quantized, tokenizer, sentences)
int8_seconds = time.perf_counter() - start
changed = sum(a != b for a, b in zip(fp32, int8))
fp32_chrf = sacrebleu.corpus_chrf(fp32, [list(tatoeba.ru)]).score
int8_chrf = sacrebleu.corpus_chrf(int8, [list(tatoeba.ru)]).score
print(f"float32: {fp32_seconds:.1f} s, chrF {fp32_chrf:.1f}")
print(f"int8:    {int8_seconds:.1f} s, chrF {int8_chrf:.1f}, {changed} of {len(sentences)} "
      "translations changed")

sentences = list(test.sample(256, random_state=SEED + 1).en)
rows = []
for batch_size in [1, 2, 4, 8, 16, 32, 64]:
    start = time.perf_counter()
    new.beam_search(model, tokenizer, sentences, batch_size=batch_size)
    seconds = time.perf_counter() - start
    rows.append({"batch": batch_size, "sentences per second": len(sentences) / seconds,
                 "ms per batch": 1000 * seconds / math.ceil(len(sentences) / batch_size)})
by_batch = pd.DataFrame(rows)
display(table(by_batch, digits=1))

threads_before = torch.get_num_threads()
rows = []
for threads in [1, 2, 4]:
    torch.set_num_threads(threads)
    start = time.perf_counter()
    new.beam_search(model, tokenizer, sentences[:128], batch_size=16)
    rows.append({"threads": threads, "seconds": time.perf_counter() - start})
torch.set_num_threads(threads_before)
display(table(pd.DataFrame(rows)))

%%writefile perevod/src/perevod/server.py
import argparse
import asyncio
import time
from contextlib import asynccontextmanager

from fastapi import FastAPI
from pydantic import BaseModel, Field

from . import __version__
from .translator import MODEL, Translator


class TranslateRequest(BaseModel):
    # the model reads at most 256 tokens, about a thousand characters of English
    text: str = Field(max_length=1000, examples=["Tom is afraid of the dark."])


class TranslateResponse(BaseModel):
    translation: str

%%writefile -a perevod/src/perevod/server.py


class Batcher:
    """Translates requests that arrive close together as one batch.

    A batch closes when it holds max_batch requests or max_wait seconds after its first one.
    """

    def __init__(self, translator, max_batch=16, max_wait=0.005):
        self.translator = translator
        self.max_batch, self.max_wait = max_batch, max_wait
        self.queue = asyncio.Queue()
        self.batch_sizes = []

    async def translate(self, text):
        answer = asyncio.get_running_loop().create_future()
        await self.queue.put((text, answer))
        return await answer  # until run() gives the future its result

%%writefile -a perevod/src/perevod/server.py

    async def collect(self):
        batch = [await self.queue.get()]  # waits for the first request, however long
        deadline = time.monotonic() + self.max_wait
        while len(batch) < self.max_batch:
            remaining = deadline - time.monotonic()
            if remaining <= 0:
                break
            try:
                batch.append(await asyncio.wait_for(self.queue.get(), remaining))
            except asyncio.TimeoutError:
                break
        return batch

%%writefile -a perevod/src/perevod/server.py

    async def run(self):
        while True:
            batch = await self.collect()
            self.batch_sizes.append(len(batch))
            texts = [text for text, _ in batch]
            try:  # in a thread, so that the event loop keeps serving requests meanwhile
                translations = await asyncio.to_thread(self.translator.translate, texts)
            except Exception as error:  # noqa: BLE001, every request of the batch gets it
                for _, answer in batch:
                    if not answer.done():
                        answer.set_exception(error)
                continue
            for (_, answer), translation in zip(batch, translations, strict=True):
                if not answer.done():  # done already if its caller stopped waiting
                    answer.set_result(translation)

%%writefile -a perevod/src/perevod/server.py


def create_app(model=MODEL, revision=None, max_batch=16, max_wait=0.005, demo=False):
    @asynccontextmanager
    async def lifespan(app):
        # before yield: when the server starts; after it: when it stops
        translator = Translator.from_pretrained(model, revision)
        app.state.batcher = Batcher(translator, max_batch, max_wait)
        worker = asyncio.create_task(app.state.batcher.run())  # runs next to the handlers
        yield
        worker.cancel()

    app = FastAPI(title="perevod", version=__version__, lifespan=lifespan)

    @app.post("/translate")
    async def translate(request: TranslateRequest) -> TranslateResponse:
        translation = await app.state.batcher.translate(request.text)
        return TranslateResponse(translation=translation)

    @app.get("/health")
    async def health():
        return {"status": "ok", "version": __version__}

    if demo:
        import gradio as gr  # only the demo extra installs it

        from .demo import build_demo

        app = gr.mount_gradio_app(app, build_demo(app), path="/")
    return app

%%writefile -a perevod/src/perevod/server.py


def main(argv=None):
    parser = argparse.ArgumentParser(prog="perevod-server",
                                     description="Serve perevod over HTTP.")
    parser.add_argument("--host", default="127.0.0.1")
    parser.add_argument("--port", type=int, default=8000)
    parser.add_argument("--max-batch", type=int, default=16, help="1 turns batching off")
    parser.add_argument("--max-wait-ms", type=float, default=5.0)
    parser.add_argument("--demo", action="store_true", help="serve a Gradio page at /")
    args = parser.parse_args(argv)

    import uvicorn  # only the serve extra installs it; the tests don't need it

    app = create_app(max_batch=args.max_batch, max_wait=args.max_wait_ms / 1000,
                     demo=args.demo)
    uvicorn.run(app, host=args.host, port=args.port)


if __name__ == "__main__":
    main()

write("perevod/pyproject.toml", """
[project]
name = "perevod"
dynamic = ["version"]  # read from src/perevod/__init__.py, see [tool.hatch.version]
description = "English-to-Russian translation with a small Transformer built from scratch"
readme = "README.md"
license = "MIT"
requires-python = ">=3.10"
dependencies = [
    "torch>=2.1",
    "tokenizers>=0.15",
    "huggingface_hub>=0.20",
    "safetensors>=0.4",
]

[project.optional-dependencies]
serve = ["fastapi>=0.110", "uvicorn>=0.29"]
demo = ["fastapi>=0.110", "uvicorn>=0.29", "gradio>=4.40"]
test = ["pytest>=8", "ruff>=0.6", "sacrebleu>=2.4", "pandas>=2.0", "pyarrow>=14",
        "fastapi>=0.110", "httpx>=0.27"]

[project.scripts]
perevod = "perevod.cli:main"
perevod-server = "perevod.server:main"

[project.urls]
Source = "https://github.com/xmarva/transformers-from-scratch/tree/main/perevod"
Model = "https://huggingface.co/xmarva/perevod-en-ru"

[build-system]
requires = ["hatchling"]
build-backend = "hatchling.build"

[tool.hatch.version]
path = "src/perevod/__init__.py"  # the one place the version is written

[tool.pytest.ini_options]
markers = ["slow: downloads the published model and translates real sentences"]

[tool.ruff]
line-length = 92  # what a notebook cell shows: lesson 7 writes these files in its cells
""")

%%writefile perevod/src/perevod/__init__.py
__version__ = "0.2.0"


def __getattr__(name):
    # `import perevod` stays light, so that --help and --version answer at once;
    # PyTorch is loaded only when someone asks for the translator
    if name == "Translator":
        from .translator import Translator

        return Translator
    raise AttributeError(f"module 'perevod' has no attribute {name!r}")

!{sys.executable} -m pip install -q -e "./perevod[serve]"

def start_server(port, max_batch, max_wait_ms=5, demo=False):
    command = [sys.executable, "-m", "perevod.server", "--port", str(port),
               "--max-batch", str(max_batch), "--max-wait-ms", str(max_wait_ms)]
    if demo:
        command.append("--demo")
    log = Path(f"server-{port}.log")
    with log.open("w") as output:
        process = subprocess.Popen(command, stdout=output, stderr=subprocess.STDOUT)
    for _ in range(180):
        if process.poll() is not None:  # it exited: a port in use, a missing package...
            break
        try:
            if httpx.get(f"http://127.0.0.1:{port}/health").status_code == 200:
                return process
        except httpx.HTTPError:
            pass
        time.sleep(1)
    process.kill()
    raise RuntimeError("the server did not start:\n" + log.read_text()[-2000:])


server = start_server(8001, max_batch=1)
URL = "http://127.0.0.1:8001/translate"
answer = httpx.post(URL, json={"text": "Nobody understands me."})
print(answer.status_code, answer.json())
refused = httpx.post(URL, json={"text": "word " * 2000})
print(refused.status_code, refused.json()["detail"][0]["msg"])

spec = httpx.get("http://127.0.0.1:8001/openapi.json").json()
for path, operations in spec["paths"].items():
    for method, operation in operations.items():
        print(method.upper(), path, "-", operation["summary"])

def arrival_times(rate, seconds, seed):
    rng = random.Random(seed)
    times, moment = [], rng.expovariate(rate)
    while moment < seconds:
        times.append(moment)
        moment += rng.expovariate(rate)  # exponential gaps: a Poisson process
    return times


async def load_test(url, rate, seconds, sentences, seed=0):
    # no connection limit: with httpx's default of 100, requests would queue in the
    # client and the test would blame the server; no timeout: a slow answer is a
    # measurement, not an error
    limits = httpx.Limits(max_connections=None, max_keepalive_connections=None)
    async with httpx.AsyncClient(timeout=None, limits=limits) as client:
        start = time.perf_counter()

        async def one(k, at):
            await asyncio.sleep(max(0.0, start + at - time.perf_counter()))
            response = await client.post(url, json={"text": sentences[k % len(sentences)]})
            response.raise_for_status()
            return time.perf_counter() - (start + at)  # from when it was due, not sent

        tasks = []
        for k, at in enumerate(arrival_times(rate, seconds, seed)):
            tasks.append(one(k, at))
        return np.array(await asyncio.gather(*tasks))

traffic = list(test.sample(2_000, random_state=SEED + 2).en)
service_times, answers = [], []
with httpx.Client(timeout=None) as client:
    for sentence in traffic[:500]:
        start = time.perf_counter()
        response = client.post(URL, json={"text": sentence})
        service_times.append(time.perf_counter() - start)
        answers.append(response.json()["translation"])
service_times = np.array(service_times)
service = service_times.mean()
mu = 1 / service
median, p99 = np.percentile(service_times, [50, 99]) * 1000
print(f"service time: mean {service * 1000:.0f} ms, median {median:.0f} ms, "
      f"99th percentile {p99:.0f} ms")
print(f"one request at a time, at most {mu:.1f} requests per second")
slowest = service_times.argmax()
print(f"\nthe slowest, {service_times[slowest] * 1000:.0f} ms:")
print(traffic[slowest])
print(answers[slowest])

LOADS = [0.2, 0.4, 0.6, 0.7, 0.8, 0.9, 0.95]
one_at_a_time = {}
for load in LOADS:
    rate = load * mu
    one_at_a_time[rate] = await load_test(URL, rate, 30, traffic, seed=int(100 * load))
    p50, p90 = np.percentile(one_at_a_time[rate], [50, 90]) * 1000
    print(f"{rate:5.1f} requests/s, {load:.0%} of capacity: "
          f"p50 {p50:6.0f} ms, p90 {p90:6.0f} ms")

again = []
with httpx.Client(timeout=None) as client:
    for sentence in traffic[:100]:
        start = time.perf_counter()
        client.post(URL, json={"text": sentence}).raise_for_status()
        again.append(time.perf_counter() - start)
before, after = 1000 * service_times[:100].mean(), 1000 * np.mean(again)
print(f"the same 100 sentences: {before:.0f} ms before the load tests, "
      f"{after:.0f} ms after")

c_s = service_times.std() / service_times.mean()
print(f"coefficient of variation of the service time: {c_s:.2f}")
rows = []
for rate, times in one_at_a_time.items():
    rho = rate / mu
    kingman = service + rho / (1 - rho) * (1 + c_s**2) / 2 * service
    rows.append({"load, %": round(100 * rho), "measured mean, ms": 1000 * times.mean(),
                 "M/M/1, ms": 1000 * service / (1 - rho), "Kingman, ms": 1000 * kingman})
theory = pd.DataFrame(rows)
display(table(theory, digits=0))

def simulated_mean(arrivals, service_times):
    done, spent = 0.0, []
    for arrived, service_time in zip(arrivals, service_times, strict=True):
        done = max(arrived, done) + service_time  # on arrival, or when the last one ends
        spent.append(done - arrived)
    return np.mean(spent)


def simulate(arrivals):
    return simulated_mean(arrivals, rng.choice(service_times, len(arrivals)))


for load in [0.8, 0.9, 0.95]:
    print(f"relaxation time at {load:.0%} load: {service / (1 - load**0.5) ** 2:.0f} s")
rng = np.random.default_rng(SEED)
rows = []
for load, (rate, times) in zip(LOADS, one_at_a_time.items(), strict=True):
    own = arrival_times(rate, 30, seed=int(100 * load))  # the arrivals this test had
    sent = np.arange(len(own)) % len(service_times)  # request k sent sentence k
    replay = simulated_mean(own, service_times[sent])
    others = [simulate(arrival_times(rate, 30, seed=1_000 + k)) for k in range(200)]
    low, high = np.percentile(others, [10, 90])
    rows.append({"load, %": round(100 * load), "measured mean, ms": 1000 * times.mean(),
                 "replay, ms": 1000 * replay,
                 "8 in 10 tests: from": 1000 * low, "to": 1000 * high})
simulated = pd.DataFrame(rows)
display(table(simulated, digits=0))

rhos = np.linspace(0.05, 0.97, 100)
fig, ax = plt.subplots(figsize=(8, 3.8))
measured = [1000 * times.mean() for times in one_at_a_time.values()]
ax.plot(rhos, 1000 * (service + rhos / (1 - rhos) * (1 + c_s**2) / 2 * service),
        color=COLORS["att"], label="Kingman's approximation")
ax.plot(rhos, 1000 * service / (1 - rhos), color=COLORS["muted"], linestyle=(0, (3, 3)),
        label="M/M/1")
ax.vlines(LOADS, simulated["8 in 10 tests: from"], simulated["to"],
          color=COLORS["dec"], alpha=0.3, linewidth=7,
          label="8 in 10 simulated 30-second tests")
ax.plot(LOADS, simulated["replay, ms"], "x", color=COLORS["dec"], markersize=8,
        label="replay of the test")
ax.plot(LOADS, measured, "o", color=COLORS["dec"], label="measured mean")
ax.set_xlabel("utilization ρ = λ / μ")
ax.set_ylabel("mean latency, ms")
ax.set_ylim(0, 1.15 * max(max(measured), simulated["to"].max()))
ax.set_title("Latency against load, one request at a time")
ax.legend(frameon=False)
plt.show()

server.terminate()
server.wait()
server = start_server(8002, max_batch=16, max_wait_ms=5)
URL = "http://127.0.0.1:8002/translate"
batched = {}
for load in LOADS + [1.5, 2, 2.5, 3]:
    rate = load * mu
    batched[rate] = await load_test(URL, rate, 30, traffic, seed=int(100 * load))
    p50, p90 = np.percentile(batched[rate], [50, 90]) * 1000
    print(f"{rate:5.1f} requests/s, {load:.0%} of the capacity one at a time: "
          f"p50 {p50:6.0f} ms, p90 {p90:6.0f} ms")

fig, ax = plt.subplots(figsize=(8, 3.8))
for runs, colour, name in [(one_at_a_time, COLORS["dec"], "one at a time"),
                           (batched, COLORS["enc"], "dynamic batching, up to 16")]:
    rates = list(runs)
    p50 = [1000 * np.percentile(runs[rate], 50) for rate in rates]
    p90 = [1000 * np.percentile(runs[rate], 90) for rate in rates]
    ax.plot(rates, p50, "o-", color=colour, label=f"{name}: median")
    ax.plot(rates, p90, "o--", color=colour, alpha=0.6, label=f"{name}: 90th percentile")
ax.axvline(mu, color=COLORS["muted"], linewidth=1, linestyle=(0, (3, 3)))
ax.text(mu, ax.get_ylim()[1] * 0.92, " capacity one at a time", color=COLORS["muted"],
        fontsize=9)
ax.set_yscale("log")
ax.set_xlabel("requests per second")
ax.set_ylabel("latency, ms (log scale)")
ax.set_title("The same server, without and with dynamic batching")
ax.legend(frameon=False, fontsize=8)
plt.show()

%%writefile perevod/src/perevod/demo.py
import gradio as gr

EXAMPLES = ["My brother lives in a small town by the sea.",
            "Every morning I drink tea and read the news.",
            "The cat is sleeping on the windowsill.", "Tom is afraid of the dark."]
ABOUT = """# perevod
English to Russian with a small Transformer built from scratch in the course
[Transformers: from scratch to production][course]. It was trained on short everyday
sentences; its [model card][card] says where it fails.

[course]: https://github.com/xmarva/transformers-from-scratch
[card]: https://huggingface.co/xmarva/perevod-en-ru"""


def build_demo(app):
    async def translate(text):
        return await app.state.batcher.translate(text)  # through the same queue as the API

    with gr.Blocks(title="perevod") as demo:
        gr.Markdown(ABOUT)
        source = gr.Textbox(label="English", lines=3)
        target = gr.Textbox(label="Russian", lines=3)
        button = gr.Button("Translate", variant="primary")
        # no limit: Gradio would otherwise send its visitors one at a time
        button.click(translate, source, target, concurrency_limit=None)
        gr.Examples(EXAMPLES, source)
    return demo

server.terminate()
server.wait()
!{sys.executable} -m pip install -q -e "./perevod[demo]"
server = start_server(8003, max_batch=16, demo=True)
page = httpx.get("http://127.0.0.1:8003/")
answer = httpx.post("http://127.0.0.1:8003/translate", json={"text": "Do you dream?"})
server.terminate()
server.wait()
print("the page:", page.status_code, "| the API:", answer.json())

%%writefile perevod/tests/test_server.py
import asyncio

import httpx
import pytest
from fastapi.testclient import TestClient

from perevod.server import create_app


@pytest.fixture
def client(tiny_dir):
    app = create_app(str(tiny_dir), max_wait=0.05)
    with TestClient(app) as client:  # entering it runs the lifespan: model and batcher
        yield client


def test_the_health_check(client):
    assert client.get("/health").json()["status"] == "ok"


def test_the_api_translates_like_the_package(client, translator, pairs):
    english = next(iter(pairs))
    response = client.post("/translate", json={"text": english})
    assert response.json() == {"translation": translator.translate(english)}


def test_a_text_too_long_is_refused(client):
    assert client.post("/translate", json={"text": "tom " * 2000}).status_code == 422


def test_requests_in_one_batch_get_their_own_translations(tiny_dir, pairs):
    app = create_app(str(tiny_dir), max_wait=0.2)

    async def ask_all():
        async with app.router.lifespan_context(app):  # the model and the batcher
            transport = httpx.ASGITransport(app=app)
            base_url = "http://perevod"
            async with httpx.AsyncClient(transport=transport, base_url=base_url) as client:
                requests = []
                for english in pairs:
                    requests.append(client.post("/translate", json={"text": english}))
                return await asyncio.gather(*requests)

    responses = asyncio.run(ask_all())
    assert app.state.batcher.batch_sizes == [len(pairs)]  # one batch for all of them
    for response, russian in zip(responses, pairs.values(), strict=True):
        assert response.json()["translation"] == russian

%%writefile perevod/tests/test_translator.py
import pytest


def test_an_empty_input_has_an_empty_translation(translator):
    assert translator.translate(["", "   ", "tom reads"])[:2] == ["", ""]


def test_a_string_gives_a_string_and_a_list_gives_a_list(translator):
    assert isinstance(translator.translate("tom reads"), str)
    assert isinstance(translator.translate(["tom reads"]), list)


def test_spaces_and_line_breaks_change_nothing(translator):
    expected = translator.translate("tom reads a book")
    assert translator.translate("tom  reads\na book") == expected


def test_a_long_input_is_cut_with_a_warning(translator):
    long_text = " ".join(["tom reads a book ."] * 10)  # 50 tokens, the model has 32 places
    with pytest.warns(UserWarning, match="cut"):
        translation = translator.translate(long_text)
    assert isinstance(translation, str)


def test_the_translator_switches_dropout_off(translator):
    assert not translator.model.training