Benchmark streaming inference under load

"""Replay generic workload scenarios with assumed coefficients; use benchmark.py for actual HTTP measurements."""

from __future__ import annotations

from dataclasses import asdict, dataclass
import json
import math
import random


@dataclass(frozen=True)
class WorkloadRequest:
    request_id: str
    arrival_ms: float
    shared_prefix_tokens: int
    unique_prompt_tokens: int
    output_tokens: int


@dataclass(frozen=True)
class EngineProfile:
    name: str
    prefill_tokens_per_ms: float
    prefix_cache_hit_efficiency: float
    base_tpot_ms: float
    batch_tpot_slope_ms: float
    max_batch_slots: int


@dataclass(frozen=True)
class BenchmarkSummary:
    engine: str
    workload: str
    total_requests: int
    duration_s: float
    ttft_p50_ms: float
    ttft_p99_ms: float
    tpot_p50_ms: float
    tpot_p99_ms: float
    raw_output_tokens_per_s: float
    request_throughput_rps: float
    slo_goodput_rps: float
    slo_attainment_ratio: float


SCENARIO_PROFILES: dict[str, EngineProfile] = {
    "baseline": EngineProfile(
        name="baseline",
        prefill_tokens_per_ms=26.0,
        prefix_cache_hit_efficiency=0.45,
        base_tpot_ms=17.5,
        batch_tpot_slope_ms=0.62,
        max_batch_slots=16,
    ),
    "prefix_reuse": EngineProfile(
        name="prefix_reuse",
        prefill_tokens_per_ms=27.0,
        prefix_cache_hit_efficiency=0.92,
        base_tpot_ms=16.8,
        batch_tpot_slope_ms=0.58,
        max_batch_slots=16,
    ),
    "faster_decode": EngineProfile(
        name="faster_decode",
        prefill_tokens_per_ms=34.0,
        prefix_cache_hit_efficiency=0.50,
        base_tpot_ms=12.4,
        batch_tpot_slope_ms=0.44,
        max_batch_slots=16,
    ),
}


def build_agent_trace(
    num_requests: int = 80, seed: int = 42, mean_rps: float = 8.5
) -> list[WorkloadRequest]:
    rng = random.Random(seed)
    trace: list[WorkloadRequest] = []
    current_ms = 0.0
    for idx in range(num_requests):
        # Bursty gamma-like arrivals: alternating bursts (short inter-arrival) and lulls.
        burst_factor = 0.32 if (idx % 10) in (2, 3, 4, 5, 6) else 1.55
        inter_arrival_ms = rng.expovariate(mean_rps / 1000.0) * burst_factor
        current_ms += inter_arrival_ms
        shared_prefix = 3072 if idx % 4 != 0 else 2048
        unique_prompt = 256 + (idx % 5) * 64
        output_len = 96 + (idx % 4) * 32
        trace.append(
            WorkloadRequest(
                request_id=f"req-{idx:03d}",
                arrival_ms=round(current_ms, 2),
                shared_prefix_tokens=shared_prefix,
                unique_prompt_tokens=unique_prompt,
                output_tokens=output_len,
            )
        )
    return trace


def _percentile(sorted_vals: list[float], q: float) -> float:
    if not sorted_vals:
        return 0.0
    pos = min(len(sorted_vals) - 1, max(0, int(math.ceil(q * len(sorted_vals)) - 1)))
    return sorted_vals[pos]


def evaluate_engine_benchmark(
    trace: list[WorkloadRequest],
    profile: EngineProfile,
    workload_name: str = "agent-shared-prefix",
    ttft_slo_ms: float = 500.0,
    tpot_slo_ms: float = 35.0,
    closed_loop: bool = False,
) -> BenchmarkSummary:
    if profile.max_batch_slots < 1 or profile.prefill_tokens_per_ms <= 0:
        raise ValueError("Capacity and prefill rate must be positive")
    if not trace:
        raise ValueError("Workload trace must not be empty")

    prefill_free_at_ms = 0.0
    active_completions_ms: list[float] = []
    ttft_samples: list[float] = []
    tpot_samples: list[float] = []
    slo_passed = 0
    total_output_tokens = 0
    first_arrival_ms = trace[0].arrival_ms
    last_completion_ms = 0.0

    for req in trace:
        effective_arrival_ms = last_completion_ms if closed_loop else req.arrival_ms
        # Enforce admission capacity instead of merely capping the reported batch.
        occupied = sorted(t for t in active_completions_ms if t > effective_arrival_ms)
        admitted_at_ms = effective_arrival_ms
        if len(occupied) >= profile.max_batch_slots:
            admitted_at_ms = occupied[len(occupied) - profile.max_batch_slots]
        active_completions_ms = [t for t in active_completions_ms if t > admitted_at_ms]
        active_batch = min(profile.max_batch_slots, len(active_completions_ms) + 1)

        queue_wait_ms = max(admitted_at_ms, prefill_free_at_ms) - effective_arrival_ms
        uncached_shared = req.shared_prefix_tokens * (
            1.0 - profile.prefix_cache_hit_efficiency
        )
        effective_prefill_tokens = uncached_shared + req.unique_prompt_tokens
        prefill_ms = 3.0 + (effective_prefill_tokens / profile.prefill_tokens_per_ms)

        ttft_ms = queue_wait_ms + prefill_ms
        prefill_free_at_ms = effective_arrival_ms + ttft_ms

        tpot_ms = profile.base_tpot_ms + profile.batch_tpot_slope_ms * active_batch
        decode_ms = (req.output_tokens - 1) * tpot_ms
        completion_ms = effective_arrival_ms + ttft_ms + decode_ms

        active_completions_ms.append(completion_ms)
        last_completion_ms = max(last_completion_ms, completion_ms)
        ttft_samples.append(round(ttft_ms, 2))
        tpot_samples.append(round(tpot_ms, 2))
        total_output_tokens += req.output_tokens

        if ttft_ms <= ttft_slo_ms and tpot_ms <= tpot_slo_ms:
            slo_passed += 1

    ttft_sorted = sorted(ttft_samples)
    tpot_sorted = sorted(tpot_samples)
    duration_s = max(
        0.001,
        (last_completion_ms - (0.0 if closed_loop else first_arrival_ms)) / 1000.0,
    )

    return BenchmarkSummary(
        engine=profile.name,
        workload=workload_name,
        total_requests=len(trace),
        duration_s=round(duration_s, 3),
        ttft_p50_ms=round(_percentile(ttft_sorted, 0.50), 2),
        ttft_p99_ms=round(_percentile(ttft_sorted, 0.99), 2),
        tpot_p50_ms=round(_percentile(tpot_sorted, 0.50), 2),
        tpot_p99_ms=round(_percentile(tpot_sorted, 0.99), 2),
        raw_output_tokens_per_s=round(total_output_tokens / duration_s, 2),
        request_throughput_rps=round(len(trace) / duration_s, 3),
        slo_goodput_rps=round(slo_passed / duration_s, 3),
        slo_attainment_ratio=round(slo_passed / len(trace), 4),
    )


def main() -> None:
    trace = build_agent_trace(num_requests=80, seed=42, mean_rps=8.5)
    summaries = {
        name: asdict(evaluate_engine_benchmark(trace, profile))
        for name, profile in SCENARIO_PROFILES.items()
    }

    overbatched_trace = build_agent_trace(num_requests=80, seed=42, mean_rps=15.0)
    overbatched_profile = EngineProfile(
        name="overbatched_unbounded",
        prefill_tokens_per_ms=42.0,
        prefix_cache_hit_efficiency=0.50,
        base_tpot_ms=22.0,
        batch_tpot_slope_ms=0.45,
        max_batch_slots=48,
    )
    overbatched = evaluate_engine_benchmark(
        overbatched_trace,
        overbatched_profile,
        workload_name="overbatched-tail-violation",
    )
    closed_loop_baseline = evaluate_engine_benchmark(
        trace,
        SCENARIO_PROFILES["baseline"],
        workload_name="closed-loop-illusion",
        closed_loop=True,
    )

    report = {
        "experiment": "synthetic-arrival-and-goodput-scenarios",
        "execution_mode": "CPU simulation; arbitrary scenario coefficients, no serving engine executed",
        "scenario_assumptions": {k: asdict(v) for k, v in SCENARIO_PROFILES.items()},
        "slo_targets": {"ttft_slo_ms": 500.0, "tpot_slo_ms": 35.0},
        "trace_requests": len(trace),
        "scenario_comparisons": summaries,
        "diagnostics": {
            "overbatched_unbounded": asdict(overbatched),
            "closed_loop_baseline": asdict(closed_loop_baseline),
        },
    }
    print(
        json.dumps(
            {
                "execution_mode": "CPU algorithm with controlled inputs; any latency/capacity coefficients are assumptions",
                **report,
            },
            indent=2,
        )
    )


if __name__ == "__main__":
    main()

"""Measure an actual SSE endpoint. A frame can contain multiple tokens: do not call frame gaps TPOT."""

import argparse
import concurrent.futures
import json
import math
import time
import urllib.error
import urllib.request


def request_one(url, prompt, limit):
    start = time.perf_counter()
    arrivals = []
    pieces = []
    saw_done = False
    req = urllib.request.Request(
        url,
        data=json.dumps({"prompt": prompt, "max_new_tokens": limit}).encode(),
        headers={"Content-Type": "application/json"},
    )
    try:
        with urllib.request.urlopen(req, timeout=120) as response:
            data = []
            for raw in response:
                line = raw.decode("utf-8").rstrip("\r\n")
                if line.startswith("data:"):
                    data.append(line[5:].lstrip(" "))
                elif not line and data:
                    event = "\n".join(data)
                    data = []
                    if event == "[DONE]":
                        saw_done = True
                        break
                    payload = json.loads(event)
                    if not isinstance(payload,dict): raise ValueError("Expected an event object")
                    piece = payload.get("token", "")
                    if not isinstance(piece,str):raise ValueError("Expected text content")
                    if piece:
                        arrivals.append(time.perf_counter() - start)
                        pieces.append(piece)
        if not saw_done:raise ValueError("Stream ended without a completion marker")
        end = time.perf_counter()
        return {
            "status": "completed",
            "ttft_ms": round(arrivals[0] * 1000, 3) if arrivals else None,
            "event_gaps_ms": [
                round((b - a) * 1000, 3) for a, b in zip(arrivals, arrivals[1:])
            ],
            "e2e_ms": round((end - start) * 1000, 3),
            "text": "".join(pieces),
            "content_events": len(pieces),
        }
    except (urllib.error.URLError, TimeoutError, ValueError) as e:
        return {
            "status": "failed",
            "error": str(e),
            "e2e_ms": round((time.perf_counter() - start) * 1000, 3),
        }


def run(url, prompts, limit=24, request_rate=1.0, max_in_flight=4, ttft_slo_ms=500):
    if not prompts or request_rate <= 0 or max_in_flight < 1:
        raise ValueError("Workload, rate and capacity must be positive")
    started = time.perf_counter()
    futures = []
    results = []
    local_shed = 0
    with concurrent.futures.ThreadPoolExecutor(max_workers=max_in_flight) as executor:
        for i, prompt in enumerate(prompts):
            delay = started + i / request_rate - time.perf_counter()
            if delay > 0:
                time.sleep(delay)
            # Do not silently queue inside the load generator: disclose local saturation.
            if sum(not f.done() for f in futures) >= max_in_flight:
                local_shed += 1
                continue
            futures.append(executor.submit(request_one, url, prompt, limit))
        results = [f.result() for f in futures]
    duration = time.perf_counter() - started
    passed = [
        r
        for r in results
        if r["status"] == "completed"
        and r["ttft_ms"] is not None
        and r["ttft_ms"] <= ttft_slo_ms
    ]
    return {
        "execution_mode": "live HTTP requests",
        "offered_requests": len(prompts),
        "client_shed_requests": local_shed,
        "duration_s": round(duration, 3),
        "ttft_slo_ms": ttft_slo_ms,
        "ttft_goodput_rps": round(len(passed) / duration, 4),
        "requests": results,
        "limits": "TTFT-only goodput; SSE events are not necessarily model tokens. Client saturation and failed responses are disclosed.",
    }


def main():
    p = argparse.ArgumentParser(description=__doc__)
    p.add_argument("--url", default="http://127.0.0.1:8766/generate")
    p.add_argument("--requests", type=int, default=6)
    p.add_argument("--rate", type=float, default=1)
    p.add_argument("--output", required=True)
    a = p.parse_args()
    report = run(
        a.url,
        ["Explain how a cache reduces repeated work."] * a.requests,
        request_rate=a.rate,
    )
    with open(a.output, "w") as f:
        json.dump(report, f, indent=2)
        f.write("\n")


if __name__ == "__main__":
    main()

"""Checks the scheduler model, not any serving engine's performance."""

import unittest
from lab import (
    SCENARIO_PROFILES,
    EngineProfile,
    WorkloadRequest,
    build_agent_trace,
    evaluate_engine_benchmark,
)


class BenchmarkTests(unittest.TestCase):
    def test_closed_loop_benchmark_hides_queueing(self):
        # A single closed-loop client waits for the WHOLE prior response.
        profile = EngineProfile("fixed", 1, 0, 10, 0, 1)
        trace = [WorkloadRequest("a", 0, 0, 7, 3), WorkloadRequest("b", 0, 0, 7, 3)]
        result = evaluate_engine_benchmark(trace, profile, closed_loop=True)
        self.assertEqual(result.duration_s, 0.06)  # 2 * (10 prefill + 20 decode)
        self.assertEqual(result.ttft_p99_ms, 10)
        opened = evaluate_engine_benchmark(trace, profile)
        self.assertEqual(opened.ttft_p99_ms, 40)  # second arrival waits for admission

    def test_throughput_without_slo_masks_tail_violations(self):
        trace = build_agent_trace()
        result = evaluate_engine_benchmark(
            trace, SCENARIO_PROFILES["baseline"], tpot_slo_ms=1
        )
        self.assertGreater(result.raw_output_tokens_per_s, 0)
        self.assertEqual(result.slo_goodput_rps, 0)

    def test_admission_capacity_changes_queue_wait(self):
        trace = [WorkloadRequest(str(i), 0, 0, 7, 3) for i in range(3)]
        one = evaluate_engine_benchmark(trace, EngineProfile("one", 1, 0, 10, 0, 1))
        three = evaluate_engine_benchmark(trace, EngineProfile("three", 1, 0, 10, 0, 3))
        self.assertEqual(one.ttft_p99_ms, 70)
        self.assertEqual(three.ttft_p99_ms, 30)

    def test_invalid_capacity_is_rejected(self):
        with self.assertRaises(ValueError):
            evaluate_engine_benchmark(
                build_agent_trace(), EngineProfile("invalid", 1, 0, 10, 0, 0)
            )


class TransportCompletionTests(unittest.TestCase):
    def test_eof_without_done_is_failed(self):
        import io
        from unittest.mock import patch
        from benchmark import request_one

        with patch(
            "urllib.request.urlopen",
            return_value=io.BytesIO(b'data: {"token":"partial"}\n\n'),
        ):
            result = request_one("http://localhost/generate", "prompt", 24)
        self.assertEqual(result["status"], "failed")
        self.assertIn("completion marker", result["error"])

    def test_complete_stream_requires_done_marker(self):
        import io
        from unittest.mock import patch
        from benchmark import request_one

        with patch(
            "urllib.request.urlopen",
            return_value=io.BytesIO(b'data: {"token":"complete"}\n\ndata: [DONE]\n\n'),
        ):
            result = request_one("http://localhost/generate", "prompt", 24)
        self.assertEqual(result["status"], "completed")
        self.assertEqual(result["text"], "complete")


if __name__ == "__main__":
    unittest.main()


class GuardBoundaryTests(unittest.TestCase):
    def test_rejects_invalid_boundary_inputs_with_specific_reason(self):
        import lab
        cases = [('import benchmark\nbenchmark.run("unused",[])', 'Workload, rate'), ('evaluate_engine_benchmark([], EngineProfile("test",1,0,1,0,1))', 'Workload trace')]
        for code, reason in cases:
            with self.subTest(code=code):
                with self.assertRaisesRegex((ValueError, MemoryError, KeyError, RuntimeError), reason):
                    exec(code, vars(lab).copy())