LLM Load Testing and SLOs: TTFT, TPOT, Open-Loop Traffic, Goodput and Capacity

"""Load tests for Kestrel Legal's model server. You write the functions marked TODO, one step at a time."""
import asyncio
import json
import math
import random
import time

import httpx
import numpy as np

URL = "http://127.0.0.1:8300/v1/chat/completions"
SLO = {"ttft_p95": 1.0, "tpot_p95": 0.050, "error_rate": 0.01}   # seconds, seconds, fraction


def prompt_of(tokens):
    """A prompt the server counts as `tokens` tokens (it counts one per word)."""
    return "word " * tokens


def load_traffic(path="traffic.jsonl"):
    """A day of production requests: [{"prompt_tokens": ..., "output_tokens": ...}, ...]."""
    with open(path) as f:
        return [json.loads(line) for line in f if line.strip()]


# ---------- Step 1: one request, three numbers ----------

async def stream_one(client, url, prompt_tokens, output_tokens):
    """Send one streaming chat request and time it. Returns {"status", "ttft", "tpot", "e2e", "output_tokens"}:
    ttft = seconds from sending to the first content token, e2e = seconds to the end of the stream,
    tpot = (e2e - ttft) / (output_tokens - 1), 0.0 for a single token. A request that fails has its HTTP status
    (0 if the connection failed) and None for the three times."""
    # TODO (Step 1): build the body (model, messages, max_tokens, stream=True); start = time.perf_counter();
    # async with client.stream("POST", url, json=body) as r: return a failure record if r.status_code != 200;
    # read r.aiter_lines(), json.loads each "data:" line up to [DONE], count chunks whose delta has content
    # and time the first. Catch httpx.HTTPError as status 0.
    raise NotImplementedError("Step 1: write stream_one()")


# ---------- Step 2: an open-loop load generator ----------

async def run_load(url, schedule):
    """Send every request of schedule, a list of (at_seconds, prompt_tokens, output_tokens), at its time,
    whether or not earlier requests have finished. Returns stream_one()'s records in schedule order, each with
    "at" added."""
    raise NotImplementedError("run_load() arrives in Step 2")


# ---------- Step 3: traffic that looks like production ----------

def poisson_schedule(rate, duration, traffic, seed=0):
    """Arrivals of a Poisson process with `rate` requests per second for `duration` seconds, each with the lengths
    of a random row of traffic: a list of (at, prompt_tokens, output_tokens), times ascending and below duration."""
    raise NotImplementedError("poisson_schedule() arrives in Step 3")


# ---------- Step 4: the SLO report ----------

def summarize(records, duration):
    """{"requests", "error_rate", "ttft_p50", "ttft_p95", "tpot_p50", "tpot_p95", "e2e_p95", "throughput_rps",
    "tokens_per_s"}. Latency percentiles are over successful requests (np.percentile), inf if none succeeded;
    throughput counts successful requests and their output tokens per second of duration."""
    raise NotImplementedError("summarize() arrives in Step 4")


def meets_slo(summary, slo=SLO):
    """{target: True/False} for every target in slo: the summary's value must be at or below the target."""
    raise NotImplementedError("meets_slo() arrives in Step 4")


def goodput(records, duration, slo=SLO):
    """Requests per second of duration that succeeded with their own ttft and tpot within the SLO's p95 targets."""
    raise NotImplementedError("goodput() arrives in Step 4")


# ---------- Step 5: capacity and replicas ----------

def find_capacity(measure, lo, hi, tol=0.5, slo=SLO):
    """The highest request rate that meets every SLO target, to within tol. measure(rate) runs a load test and
    returns its summarize(). Assumes hi fails; returns None if lo already fails."""
    raise NotImplementedError("find_capacity() arrives in Step 5")


def replicas_for(peak_rate, capacity, headroom=0.75):
    """Replicas needed to serve peak_rate when each may run at only headroom x its measured capacity."""
    raise NotImplementedError("replicas_for() arrives in Step 5")