#!/usr/bin/env python3
"""Concurrency load generator for an OpenAI-compatible chat completions endpoint.

Purpose: drive a served model at one or more concurrency levels and report the numbers
    Part 9 cares about - output token throughput, request throughput, goodput against a
    stated objective, time to first token and time per output token - with percentiles
    rather than means, plus the prompt and cached token counts the server reports.
Platform: all (spark, strix, mac, nvidia). Pure Python standard library, 3.9 or later:
    no pip install.
Minimum memory: 16 GB on the machine running the server. The generator itself is tiny
    and may run on a second machine on the same network.
Assumes: a server already listening and answering POST /v1/chat/completions with
    streaming, reachable over plain HTTP (this is a localhost or LAN lab tool; it does
    not speak TLS). The model name must be the one the server reports at /v1/models.
    Results are appended to the lab notebook as one JSON line per concurrency level.

Usage:
    python3 load-test.py --base-url http://127.0.0.1:8080/v1 --model local-chat \
        --concurrency 1,5,10,20 --requests 40 --min-rounds 4 --max-tokens 128 \
        --prompt-set mixed --label llama-baseline --labbook labbook.md

    python3 load-test.py --base-url http://127.0.0.1:8000/v1 --model local-chat \
        --concurrency 1,10 --requests 40 --max-tokens 64 \
        --prompt-set shared-prefix --label vllm-prefix-shared --labbook labbook.md

    An API key, if the server needs one, is read from an environment variable named by
    --api-key-env. No key is ever written to this file or to the notebook.

Method, stated once:
    * Closed loop. Each level starts `concurrency` workers that each send a request, wait
      for it to finish and send the next, so that number of requests is in flight until
      the queue drains. Requests per level = max(--requests, --min-rounds x concurrency).
    * Before the first level, --warmup requests (default 1) with a short prompt that shares
      nothing with the prompt sets are sent and discarded.
    * Time to first token (TTFT) is measured to the first streamed delta carrying any text:
      answer text ("content") or reasoning text ("reasoning_content" or "reasoning"), so a
      model that thinks first is timed the same way on every engine.
    * Output tokens come from the server's usage block (requested with
      stream_options.include_usage); where a server sends none, one text delta is counted
      as one token and the notebook line says "deltas" instead of "usage".
    * Time per output token (TPOT) for one request = (end - TTFT) / (output tokens - 1),
      the same definition vllm bench serve uses. Percentiles are nearest-rank.
    * Output tokens per second = all output tokens at the level / wall-clock of the level.
    * Goodput = requests per second that completed with TTFT <= --slo-ttft and
      TPOT <= --slo-tpot.
    * Mean in flight = sum of request durations / wall-clock: close to `concurrency` when
      the loop held the load, lower when the client or the drain at the end did not.
    * Cached tokens: usage.prompt_tokens_details.cached_tokens (vLLM with
      --enable-prompt-tokens-details, llama-server, mlx_lm.server), or llama-server's
      timings.cache_n where that is the only count sent.
"""

from __future__ import annotations

import argparse
import asyncio
import json
import math
import os
import secrets
import statistics
import sys
import time
from collections import Counter
from urllib.parse import urlsplit

LAB_ID = "part-09/lab-serve-twenty-concurrent-users"
TEXT_FIELDS = ("content", "reasoning_content", "reasoning")

# --------------------------------------------------------------------------- prompts

TOOLS = [
    ("read_metrics", "Read the Prometheus metrics endpoint of one inference server and return the series whose names match a pattern.", {"server": "string: the server name from the inventory", "pattern": "string: a regular expression matched against metric names"}),
    ("tail_log", "Return the last lines of a server's log, optionally filtered by level.", {"server": "string", "lines": "integer: how many lines, at most 500", "level": "string: one of debug, info, warning, error"}),
    ("restart_server", "Stop and start one inference server with its recorded configuration. Refuses while requests are running unless force is true.", {"server": "string", "force": "boolean"}),
    ("set_option", "Change one recorded serving option for the next restart and return the previous value.", {"server": "string", "option": "string: for example max-num-seqs or ctx-size", "value": "string"}),
    ("memory_report", "Report total, used and free accelerator memory, and the part of it held by the key-value cache.", {"server": "string"}),
    ("list_models", "List the models a server reports at /v1/models, with their context length.", {"server": "string"}),
    ("run_load_test", "Start a load test against a server and return the path of the results once it finishes.", {"server": "string", "concurrency": "array of integers", "requests": "integer", "max_tokens": "integer"}),
    ("compare_runs", "Compare two recorded load-test runs level by level and return the differences in throughput and latency.", {"first_label": "string", "second_label": "string"}),
    ("disk_usage", "Report the disk space used by the model directory and by each model file in it.", {"path": "string"}),
    ("checksum_file", "Compute the SHA-256 of a model file and compare it with the checksum recorded beside it.", {"path": "string"}),
    ("gpu_temperature", "Return the accelerator temperature, power draw and any throttling reasons reported by the driver.", {"device": "integer"}),
    ("open_ticket", "Record an operational issue with a title, a severity and the evidence gathered so far.", {"title": "string", "severity": "string: one of low, medium, high", "evidence": "string"}),
    ("slot_status", "Return each parallel slot of a llama-server instance with its context size, whether it is processing and how many tokens it holds.", {"server": "string"}),
    ("queue_depth", "Return how many requests are running and how many are waiting on a server, sampled once per second for a number of seconds.", {"server": "string", "seconds": "integer: at most 60"}),
    ("prefix_cache_stats", "Return prefix cache queries, hits and the hit rate since the server started.", {"server": "string"}),
    ("download_model", "Download a model file from a named repository into the model directory and verify its checksum.", {"repository": "string", "file": "string"}),
    ("delete_model_file", "Delete one model file after confirming that no running server has it loaded.", {"path": "string", "confirm": "boolean"}),
    ("network_throughput", "Measure throughput between this machine and another on the local network for a number of seconds.", {"peer": "string: a host name from the inventory", "seconds": "integer"}),
    ("schedule_benchmark", "Schedule a benchmark to run at a quiet time and record it in the lab notebook when it finishes.", {"when": "string: an ISO 8601 time", "benchmark": "string"}),
    ("summarise_notebook", "Return the notebook lines recorded for one lab, grouped by label, as a table.", {"lab": "string", "label": "string"}),
]

RULES = [
    "Answer in at most four sentences unless the operator asks for a table.",
    "Quote numbers only from tool results, never from memory, and name the tool that produced them.",
    "Prefer the smallest change that tests one hypothesis, and say what result would refute it.",
    "Never restart a server that has requests running without asking the operator first.",
    "When memory is short, reduce concurrency or context before reducing model precision.",
    "Report latency as a median and a ninetieth percentile, never as a mean alone.",
]

# About 1,365 tokens under the Qwen3 tokeniser: a system-prompt-and-tool-schema preamble
# of the kind an agent sends on every step.
SHARED_PREAMBLE = (
    "You are the operations assistant for a small local inference service: one open-weight "
    "model served over an OpenAI-compatible interface on one machine. You can call the tools "
    "described below by answering with a JSON object naming the tool and its arguments.\n\n"
    "Rules:\n" + "\n".join(f"{i + 1}. {rule}" for i, rule in enumerate(RULES))
    + "\n\nTools:\n"
    + json.dumps([{"name": n, "description": d, "parameters": p} for n, d, p in TOOLS], indent=2)
)

SHARED_QUESTIONS = [
    "What happens to the key-value cache when I double the context length?",
    "Why is prompt processing so much faster than generation?",
    "Which measurement should I quote to a user who says the service feels slow?",
    "What does it mean when the number of waiting requests stays above zero?",
    "Why does a second concurrent request cost less than twice as much?",
    "How do I decide between more concurrent slots and a longer context?",
    "What is the first thing to check when memory runs out at startup?",
    "Why does the mean latency mislead me about a serving system?",
]

MIXED_PROMPTS = [
    "Explain in three sentences why decode speed is limited by memory bandwidth.",
    "Write a short shell function that reports how much disk a directory uses.",
    "Summarise the difference between tensor parallel and pipeline parallel.",
    "Give three reasons a language model server might refuse to start.",
    "What is a KV cache, and why does it grow with the conversation?",
    "Write a Python function that returns the median of a list, without imports.",
    "List four things that belong in a benchmark result besides the number itself.",
    "Explain quantisation to somebody who knows what a floating point number is.",
    "Describe what continuous batching does, in plain language.",
    "Name three failure modes of retrieval-augmented generation.",
    "Write a regular expression that matches an ISO date, and explain it.",
    "Why might a smaller model that fits beat a larger one that does not?",
    "Give a checklist for exposing a local API endpoint to a home network.",
    "What is the difference between a context window and a context length?",
    "Explain why two runs with the same seed can still produce different text.",
    "Write a short bash loop that retries a command three times with a pause.",
    "What does a high standard deviation in a benchmark run usually indicate?",
    "Describe the trade-off between batch size and per-user latency.",
    "Explain what a chat template does and what goes wrong without one.",
    "Give three signs that a machine is thermally limited during a benchmark.",
    "What is speculative decoding, and when does it not help?",
    "Explain the difference between total and active parameters.",
    "Write a JSON object describing a machine with a chip, memory and an operating system.",
    "Why is goodput a more useful headline number than raw throughput?",
]

WARMUP_PROMPT = "Reply with one word: ready"


def build_prompts(prompt_set: str, count: int) -> list[str]:
    """Returns `count` prompts, cycling the chosen set.

    mixed: 24 short, different questions, each behind a random tag, so that a prompt
        repeated later in the run (or in an earlier run) is not a cache hit.
    shared-prefix: the preamble, then one of eight questions; everything before the
        question is byte-identical across requests, so a prefix cache can reuse it.
    unique-prefix: the same text with a random tag as its first line, different for every
        request and every run, so nothing after the chat template's opening tokens can be
        reused. It is the control for shared-prefix: same length, no reusable prefix.
    """
    if prompt_set == "shared-prefix":
        pool = [f"{SHARED_PREAMBLE}\n\nQuestion: {q}" for q in SHARED_QUESTIONS]
        return [pool[i % len(pool)] for i in range(count)]
    if prompt_set == "unique-prefix":
        return [
            f"Request {secrets.token_hex(8)}.\n{SHARED_PREAMBLE}\n\n"
            f"Question: {SHARED_QUESTIONS[i % len(SHARED_QUESTIONS)]}"
            for i in range(count)
        ]
    if prompt_set == "mixed":
        return [
            f"Request {secrets.token_hex(8)}. {MIXED_PROMPTS[i % len(MIXED_PROMPTS)]}"
            for i in range(count)
        ]
    raise ValueError(f"unknown prompt set: {prompt_set}")


# ------------------------------------------------------------------- minimal HTTP/1.1


class RequestFailed(Exception):
    """Raised when a request did not complete with a usable stream."""


async def _read_headers(reader: asyncio.StreamReader) -> tuple[int, bool]:
    """Reads the status line and headers. Returns (status, chunked)."""
    status_line = await reader.readline()
    if not status_line:
        raise RequestFailed("server closed the connection without a response")
    parts = status_line.decode("latin-1", "replace").split()
    if len(parts) < 2 or not parts[1].isdigit():
        raise RequestFailed(f"unparseable status line: {status_line!r}")
    status = int(parts[1])
    chunked = False
    while True:
        line = await reader.readline()
        if line in (b"\r\n", b"\n", b""):
            break
        lowered = line.lower()
        if lowered.startswith(b"transfer-encoding:") and b"chunked" in lowered:
            chunked = True
    return status, chunked


async def _body_lines(reader: asyncio.StreamReader, chunked: bool):
    """Yields body lines, decoding chunked transfer encoding where it is used."""
    buf = b""
    if chunked:
        while True:
            size_line = await reader.readline()
            if not size_line:
                break
            head = size_line.strip().split(b";")[0]
            if not head:
                continue
            try:
                size = int(head, 16)
            except ValueError:
                break
            if size == 0:
                await reader.readline()
                break
            buf += await reader.readexactly(size)
            await reader.readexactly(2)
            while b"\n" in buf:
                line, buf = buf.split(b"\n", 1)
                yield line.rstrip(b"\r")
    else:
        while True:
            chunk = await reader.read(65536)
            if not chunk:
                break
            buf += chunk
            while b"\n" in buf:
                line, buf = buf.split(b"\n", 1)
                yield line.rstrip(b"\r")
    if buf.strip():
        yield buf.rstrip(b"\r")


def _cached_from(event: dict):
    """The cached prompt token count in a streamed event, if the server sent one."""
    usage = event.get("usage") or {}
    details = usage.get("prompt_tokens_details") or {}
    if isinstance(details, dict) and details.get("cached_tokens") is not None:
        return int(details["cached_tokens"])
    timings = event.get("timings") or {}
    if timings.get("cache_n") is not None:
        return int(timings["cache_n"])
    return None


async def stream_completion(cfg: dict, prompt: str, max_tokens: int) -> dict:
    """Sends one streaming chat completion and times it.

    Returns a dict with ok, ttft, total, tokens, token_source, prompt_tokens, cached,
    finish and (on failure) error.
    """
    body = json.dumps(
        {
            "model": cfg["model"],
            "messages": [{"role": "user", "content": prompt}],
            "max_tokens": max_tokens,
            "temperature": cfg["temperature"],
            "stream": True,
            "stream_options": {"include_usage": True},
        }
    ).encode("utf-8")

    headers = [
        f"POST {cfg['path']} HTTP/1.1",
        f"Host: {cfg['host']}:{cfg['port']}",
        "Content-Type: application/json",
        "Accept: text/event-stream",
        f"Content-Length: {len(body)}",
        "Connection: close",
    ]
    if cfg["api_key"]:
        headers.append(f"Authorization: Bearer {cfg['api_key']}")

    started = time.perf_counter()
    ttft = None
    deltas = 0
    usage_tokens = None
    prompt_tokens = None
    cached = None
    finish = None
    writer = None
    try:
        reader, writer = await asyncio.wait_for(
            asyncio.open_connection(cfg["host"], cfg["port"]), timeout=cfg["timeout"]
        )
        writer.write(("\r\n".join(headers) + "\r\n\r\n").encode("latin-1") + body)
        await writer.drain()

        status, chunked = await asyncio.wait_for(_read_headers(reader), timeout=cfg["timeout"])
        if status != 200:
            detail = b""
            try:
                detail = await asyncio.wait_for(reader.read(300), timeout=5)
            except (OSError, asyncio.TimeoutError):
                pass
            raise RequestFailed(f"HTTP {status} {detail.decode('utf-8', 'replace').strip()}")

        async for raw in _body_lines(reader, chunked):
            if not raw.startswith(b"data:"):
                continue
            payload = raw[5:].strip()
            if payload in (b"[DONE]", b""):
                continue
            try:
                event = json.loads(payload)
            except json.JSONDecodeError:
                continue
            if event.get("error"):
                raise RequestFailed(f"server error in stream: {str(event['error'])[:200]}")
            for choice in event.get("choices") or []:
                delta = choice.get("delta") or {}
                if any(delta.get(field) for field in TEXT_FIELDS) or delta.get("tool_calls"):
                    deltas += 1
                    if ttft is None:
                        ttft = time.perf_counter() - started
                if choice.get("finish_reason"):
                    finish = choice["finish_reason"]
            usage = event.get("usage")
            if usage:
                if usage.get("completion_tokens") is not None:
                    usage_tokens = int(usage["completion_tokens"])
                if usage.get("prompt_tokens") is not None:
                    prompt_tokens = int(usage["prompt_tokens"])
            found = _cached_from(event)
            if found is not None:
                cached = found
    except (OSError, asyncio.TimeoutError, asyncio.IncompleteReadError, RequestFailed) as exc:
        return {"ok": False, "error": f"{type(exc).__name__}: {exc}"}
    finally:
        if writer is not None:
            writer.close()
            try:
                await writer.wait_closed()
            except OSError:
                pass

    total = time.perf_counter() - started
    if deltas == 0 or ttft is None:
        return {"ok": False, "error": "stream produced no content deltas"}
    tokens = usage_tokens if usage_tokens else deltas
    return {
        "ok": True,
        "ttft": ttft,
        "total": total,
        "tokens": tokens,
        "token_source": "usage" if usage_tokens else "deltas",
        "prompt_tokens": prompt_tokens,
        "cached": cached,
        "finish": finish or "none",
    }


# ------------------------------------------------------------------------ measurement


def percentile(values: list[float], p: float) -> float:
    """Nearest-rank percentile: the smallest value with at least p% of values at or below it."""
    if not values:
        return 0.0
    ordered = sorted(values)
    rank = max(1, min(len(ordered), math.ceil(p / 100.0 * len(ordered))))
    return ordered[rank - 1]


def _summary(values: list[float], digits: int) -> dict:
    return {
        "p50": round(percentile(values, 50), digits),
        "p90": round(percentile(values, 90), digits),
        "p99": round(percentile(values, 99), digits),
        "mean": round(statistics.fmean(values), digits) if values else 0.0,
    }


def _token_source(sources: Counter):
    """"usage" or "deltas" when every request agrees, the counts when they do not, "none" if empty."""
    if not sources:
        return "none"
    if len(sources) == 1:
        return next(iter(sources))
    return dict(sources)


def _mean_or_none(values: list) -> float | None:
    present = [v for v in values if v is not None]
    return round(statistics.fmean(present), 1) if present else None


async def warm_up(cfg: dict, count: int) -> None:
    """Sends `count` short requests one at a time and discards their timings."""
    for _ in range(count):
        await stream_completion(cfg, WARMUP_PROMPT, 8)


async def run_level(cfg: dict, concurrency: int, prompts: list[str]) -> dict:
    """Runs every prompt through `concurrency` workers and aggregates the results."""
    queue: asyncio.Queue = asyncio.Queue()
    for prompt in prompts:
        queue.put_nowait(prompt)
    results: list[dict] = []

    async def worker() -> None:
        while True:
            try:
                prompt = queue.get_nowait()
            except asyncio.QueueEmpty:
                return
            results.append(await stream_completion(cfg, prompt, cfg["max_tokens"]))
            queue.task_done()

    started = time.perf_counter()
    await asyncio.gather(*[worker() for _ in range(concurrency)])
    wall = time.perf_counter() - started

    ok = [r for r in results if r["ok"]]
    failed = [r for r in results if not r["ok"]]
    ttfts = [r["ttft"] for r in ok]
    tpots = {id(r): (r["total"] - r["ttft"]) / (r["tokens"] - 1) for r in ok if r["tokens"] > 1}
    out_tokens = sum(r["tokens"] for r in ok)
    good = [
        r for r in ok
        if r["ttft"] <= cfg["slo_ttft"] and tpots.get(id(r), 0.0) <= cfg["slo_tpot"]
    ]
    sources = Counter(r["token_source"] for r in ok)

    return {
        "concurrency": concurrency,
        "requests": len(results),
        "completed": len(ok),
        "failed": len(failed),
        "wall_s": round(wall, 3),
        "output_tokens": out_tokens,
        "output_tokens_per_s": round(out_tokens / wall, 2) if wall > 0 else 0.0,
        "requests_per_s": round(len(ok) / wall, 3) if wall > 0 else 0.0,
        "goodput_requests_per_s": round(len(good) / wall, 3) if wall > 0 else 0.0,
        "within_slo": len(good),
        "mean_in_flight": round(sum(r["total"] for r in ok) / wall, 2) if wall > 0 else 0.0,
        "ttft_s": _summary(ttfts, 4),
        "tpot_s": _summary(list(tpots.values()), 5),
        "prompt_tokens_mean": _mean_or_none([r["prompt_tokens"] for r in ok]),
        "cached_tokens_mean": _mean_or_none([r["cached"] for r in ok]),
        "token_counts": _token_source(sources),
        "finish_reasons": dict(Counter(r["finish"] for r in ok)),
        "first_error": failed[0]["error"] if failed else None,
    }


def print_level(row: dict) -> None:
    """Prints one concurrency level as two aligned lines."""
    print(
        f"  c={row['concurrency']:>3}  "
        f"ok {row['completed']:>3}/{row['requests']:<3} "
        f"fail {row['failed']:<3} "
        f"wall {row['wall_s']:>7.2f}s  "
        f"out/s {row['output_tokens_per_s']:>8.2f}  "
        f"req/s {row['requests_per_s']:>6.3f}  "
        f"TTFT p50 {row['ttft_s']['p50']:.3f} p90 {row['ttft_s']['p90']:.3f}  "
        f"TPOT p50 {row['tpot_s']['p50']:.4f} p90 {row['tpot_s']['p90']:.4f}"
    )
    cached = row["cached_tokens_mean"]
    prompt = row["prompt_tokens_mean"]
    print(
        f"         goodput {row['goodput_requests_per_s']:.3f} req/s "
        f"({row['within_slo']}/{row['completed']} within SLO)  "
        f"in flight {row['mean_in_flight']:.1f}  "
        f"prompt {prompt if prompt is not None else '?'} tok, "
        f"cached {cached if cached is not None else '?'}  "
        f"tokens from {row['token_counts']}  finish {row['finish_reasons']}"
    )
    if row["first_error"]:
        print(f"         first error: {row['first_error']}")


def append_labbook(path: str, record: dict) -> None:
    """Appends one JSON line, in the course notebook format."""
    with open(path, "a", encoding="utf-8") as handle:
        handle.write(json.dumps(record, sort_keys=True) + "\n")


def parse_args(argv: list[str]) -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    parser.add_argument("--base-url", default="http://127.0.0.1:8000/v1",
                        help="OpenAI-compatible base URL, ending in /v1")
    parser.add_argument("--model", required=True, help="model name the server reports")
    parser.add_argument("--concurrency", default="1,5,10,20",
                        help="comma-separated concurrency levels to run in order")
    parser.add_argument("--requests", type=int, default=40,
                        help="requests sent at each concurrency level (at least)")
    parser.add_argument("--min-rounds", type=int, default=0,
                        help="send at least this many requests per worker at each level")
    parser.add_argument("--max-tokens", type=int, default=128)
    parser.add_argument("--temperature", type=float, default=0.0)
    parser.add_argument("--timeout", type=float, default=300.0)
    parser.add_argument("--prompt-set", default="mixed",
                        choices=["mixed", "shared-prefix", "unique-prefix"])
    parser.add_argument("--warmup", type=int, default=1,
                        help="discarded short requests sent before the first level")
    parser.add_argument("--slo-ttft", type=float, default=2.0,
                        help="goodput objective for time to first token, in seconds")
    parser.add_argument("--slo-tpot", type=float, default=0.1,
                        help="goodput objective for time per output token, in seconds")
    parser.add_argument("--label", default="run", help="tag written into the notebook line")
    parser.add_argument("--engine", default="unknown",
                        help="engine name recorded in the notebook line, e.g. vllm")
    parser.add_argument("--engine-version", default="unknown",
                        help="engine version recorded in the notebook line")
    parser.add_argument("--host-desc", default="unknown",
                        help="machine description recorded in the notebook line")
    parser.add_argument("--quant", default="unknown",
                        help="weight format recorded in the notebook line, e.g. Q4_K_M or AWQ")
    parser.add_argument("--context-length", type=int, default=0,
                        help="context length per sequence the server was started with")
    parser.add_argument("--server-slots", type=int, default=0,
                        help="sequences the server may run at once (slots or max-num-seqs)")
    parser.add_argument("--api-key-env", default="LOADTEST_API_KEY",
                        help="environment variable holding the API key, if the server needs one")
    parser.add_argument("--labbook", default="labbook.md",
                        help="file to append one JSON line per concurrency level to")
    return parser.parse_args(argv)


def main(argv: list[str]) -> int:
    args = parse_args(argv)

    parts = urlsplit(args.base_url)
    if parts.scheme != "http":
        print("This tool speaks plain HTTP only; point it at a localhost or LAN endpoint.",
              file=sys.stderr)
        return 2
    if not parts.hostname:
        print(f"Could not read a host from --base-url {args.base_url!r}.", file=sys.stderr)
        return 2
    if args.requests < 1 or args.max_tokens < 2:
        print("--requests must be at least 1 and --max-tokens at least 2.", file=sys.stderr)
        return 2

    cfg = {
        "host": parts.hostname,
        "port": parts.port or 80,
        "path": (parts.path.rstrip("/") or "/v1") + "/chat/completions",
        "model": args.model,
        "max_tokens": args.max_tokens,
        "temperature": args.temperature,
        "timeout": args.timeout,
        "slo_ttft": args.slo_ttft,
        "slo_tpot": args.slo_tpot,
        "api_key": os.environ.get(args.api_key_env, ""),
    }

    try:
        levels = [int(x) for x in args.concurrency.split(",") if x.strip()]
    except ValueError:
        print(f"--concurrency must be a comma-separated list of integers, got {args.concurrency!r}",
              file=sys.stderr)
        return 2
    if not levels or min(levels) < 1:
        print("--concurrency listed no levels, or a level below 1.", file=sys.stderr)
        return 2

    print(f"==> {args.label}: {args.model} at {args.base_url}")
    print(f"    prompt set {args.prompt_set}, at least {args.requests} request(s) per level"
          f"{f' and {args.min_rounds} per worker' if args.min_rounds else ''}, "
          f"max_tokens {args.max_tokens}, temperature {args.temperature}")
    print(f"    SLO for goodput: TTFT <= {args.slo_ttft} s and TPOT <= {args.slo_tpot} s; "
          f"{args.warmup} warm-up request(s)")

    if args.warmup > 0:
        asyncio.run(warm_up(cfg, args.warmup))

    exit_code = 0
    for level in levels:
        count = max(args.requests, args.min_rounds * level)
        prompts = build_prompts(args.prompt_set, count)
        row = asyncio.run(run_level(cfg, level, prompts))
        print_level(row)
        if row["completed"] == 0:
            exit_code = 1
        record = {
            "lab": LAB_ID,
            "label": args.label,
            "engine": args.engine,
            "engine_version": args.engine_version,
            "host": args.host_desc,
            "quant": args.quant,
            "context_length": args.context_length or None,
            "server_slots": args.server_slots or None,
            "model": args.model,
            "base_url": args.base_url,
            "prompt_set": args.prompt_set,
            "max_tokens": args.max_tokens,
            "temperature": args.temperature,
            "slo_ttft_s": args.slo_ttft,
            "slo_tpot_s": args.slo_tpot,
            "recorded_at": time.strftime("%Y-%m-%dT%H:%M:%S%z"),
        }
        record.update(row)
        append_labbook(args.labbook, record)

    print(f"    appended {len(levels)} line(s) to {args.labbook}")
    if exit_code:
        print("    at least one level completed no requests; see the error above.",
              file=sys.stderr)
    return exit_code


if __name__ == "__main__":
    sys.exit(main(sys.argv[1:]))