[Benchmark] Add agentic rollout simulator and offline explorer (#40034)
This commit is contained in:
Byron Hsu
@@ -0,0 +1,433 @@
|
||||
# /// script
|
||||
# requires-python = ">=3.11"
|
||||
# dependencies = ["prometheus-client"]
|
||||
# ///
|
||||
"""Build an offline Rollout Simulator and Explorer HTML from a recording."""
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import math
|
||||
import statistics
|
||||
from pathlib import Path
|
||||
|
||||
from metrics import (
|
||||
counter_rate,
|
||||
)
|
||||
from metrics import gauge as metric_gauge
|
||||
from metrics import (
|
||||
hit_percentages,
|
||||
metric_values,
|
||||
occupancy,
|
||||
values,
|
||||
)
|
||||
|
||||
|
||||
def read_rows(path):
|
||||
if not path.exists():
|
||||
return []
|
||||
rows = []
|
||||
for line in path.read_text().splitlines():
|
||||
if line.strip():
|
||||
try:
|
||||
rows.append(json.loads(line))
|
||||
except json.JSONDecodeError:
|
||||
# Interrupted writes remain visible as recording errors.
|
||||
rows.append({"error": "Incomplete JSONL record"})
|
||||
return rows
|
||||
|
||||
|
||||
def panels_for(reqs, scrapes, origin, duration):
|
||||
series = []
|
||||
parsed = []
|
||||
for row in sorted(scrapes, key=lambda r: r["timestamp"]):
|
||||
try:
|
||||
samples = metric_values(row.get("text", "")) if not row.get("error") else {}
|
||||
except ValueError:
|
||||
samples = {}
|
||||
parsed.append((row["timestamp"] - origin, samples))
|
||||
|
||||
def gauge(name, mult=1):
|
||||
return [
|
||||
[t, None if (v := metric_gauge(ss, name)) is None else v * mult]
|
||||
for t, ss in parsed
|
||||
]
|
||||
|
||||
def rate(name):
|
||||
return [
|
||||
[t1, counter_rate(s0, s1, "sglang:" + name, t1 - t0)]
|
||||
for (t0, s0), (t1, s1) in zip(parsed, parsed[1:])
|
||||
]
|
||||
|
||||
cache_hits = [[], [], []]
|
||||
for (t0, s0), (t1, s1) in zip(parsed, parsed[1:]):
|
||||
for points, value in zip(
|
||||
cache_hits, hit_percentages(s0, s1) if t1 > t0 else [None] * 3
|
||||
):
|
||||
points.append([t1, value])
|
||||
l1_occupied, l1_capacity, l2_occupied, l2_capacity = [], [], [], []
|
||||
l1_usage, l2_usage = [], []
|
||||
for t, samples in parsed:
|
||||
active, evictable, capacity, usage1 = occupancy(
|
||||
samples, ("kv_used_tokens", "kv_evictable_tokens", "max_total_num_tokens")
|
||||
)
|
||||
host, host_capacity, usage2 = occupancy(
|
||||
samples, ("hicache_host_used_tokens", "hicache_host_total_tokens")
|
||||
)
|
||||
l1_occupied.append([t, active + evictable if active is not None else None])
|
||||
l1_capacity.append([t, metric_gauge(samples, "max_total_num_tokens")])
|
||||
l2_occupied.append([t, metric_gauge(samples, "hicache_host_used_tokens")])
|
||||
l2_capacity.append([t, metric_gauge(samples, "hicache_host_total_tokens")])
|
||||
l1_usage.append([t, usage1])
|
||||
l2_usage.append([t, usage2])
|
||||
|
||||
bins = [(i, min(i + 2, duration)) for i in range(0, math.ceil(duration), 2)]
|
||||
|
||||
def distribution(observations):
|
||||
buckets = [[] for _ in bins]
|
||||
for elapsed, value in observations:
|
||||
if 0 <= elapsed < duration:
|
||||
buckets[int(elapsed // 2)].append(value)
|
||||
groups = [
|
||||
((a + b) / 2, sorted(values)) for (a, b), values in zip(bins, buckets)
|
||||
]
|
||||
return [
|
||||
{
|
||||
"name": name,
|
||||
"points": [
|
||||
[
|
||||
t,
|
||||
(
|
||||
(
|
||||
statistics.mean(v)
|
||||
if q is None
|
||||
else v[math.ceil(q * len(v)) - 1]
|
||||
)
|
||||
if v
|
||||
else None
|
||||
),
|
||||
]
|
||||
for t, v in groups
|
||||
],
|
||||
}
|
||||
for name, q in (("avg", None), ("p50", 0.5), ("p90", 0.9), ("p99", 0.99))
|
||||
]
|
||||
|
||||
def dist(key, scale=1):
|
||||
return distribution(
|
||||
[
|
||||
(
|
||||
r.get(
|
||||
"completed_at",
|
||||
r.get(
|
||||
"failed_at",
|
||||
r.get("first_token_at", r.get("submitted_at", origin)),
|
||||
),
|
||||
)
|
||||
- origin,
|
||||
r[key] * scale,
|
||||
)
|
||||
for r in reqs
|
||||
if r.get(key) is not None
|
||||
]
|
||||
)
|
||||
|
||||
def add(title, unit, source, items):
|
||||
series.append({"title": title, "unit": unit, "source": source, "series": items})
|
||||
|
||||
def one(name, points, **style):
|
||||
return {"name": name, "points": points, **style}
|
||||
|
||||
add(
|
||||
"QPS",
|
||||
"req/s",
|
||||
"Successful client completions / 2-second window",
|
||||
[
|
||||
one(
|
||||
"requests",
|
||||
[
|
||||
[
|
||||
(a + b) / 2,
|
||||
sum(
|
||||
a <= r.get("completed_at", float("inf")) - origin < b
|
||||
and not r.get("error")
|
||||
for r in reqs
|
||||
)
|
||||
/ (b - a),
|
||||
]
|
||||
for a, b in bins
|
||||
],
|
||||
)
|
||||
],
|
||||
)
|
||||
add(
|
||||
"Running And Queued Requests",
|
||||
"requests",
|
||||
"Server num_running_reqs / num_queue_reqs gauges",
|
||||
[
|
||||
one("running", gauge("num_running_reqs")),
|
||||
one("queued", gauge("num_queue_reqs")),
|
||||
],
|
||||
)
|
||||
add(
|
||||
"Input Length",
|
||||
"tokens",
|
||||
"Client context, grouped by completion/failure time",
|
||||
dist("context_tokens"),
|
||||
)
|
||||
reqs = [
|
||||
{**r, "output_tokens": (r.get("meta_info") or {}).get("completion_tokens")}
|
||||
for r in reqs
|
||||
]
|
||||
add(
|
||||
"Output Length",
|
||||
"tokens",
|
||||
"Observed output token counts; includes partial requests",
|
||||
dist("output_tokens"),
|
||||
)
|
||||
add(
|
||||
"TTFT",
|
||||
"ms",
|
||||
"Client submit to first output; includes partial requests with measured TTFT",
|
||||
dist("ttft_s", 1000),
|
||||
)
|
||||
# Report the distribution of observed per-token stream intervals, not GPU timing.
|
||||
itl = []
|
||||
for r in reqs:
|
||||
for (t0, n0), (t1, n1) in zip(r.get("events", []), r.get("events", [])[1:]):
|
||||
if n1 > n0:
|
||||
itl.append(
|
||||
(r["submitted_at"] + t1 - origin, (t1 - t0) / (n1 - n0) * 1000)
|
||||
)
|
||||
add(
|
||||
"ITL",
|
||||
"ms",
|
||||
"Client stream intervals / received token increments; not kernel timing",
|
||||
distribution(itl),
|
||||
)
|
||||
add(
|
||||
"Prompt / Input Throughput",
|
||||
"tokens/s",
|
||||
"Server counter deltas / scrape interval; cached sums all tiers",
|
||||
[
|
||||
one("prompt", rate("prompt_tokens_total")),
|
||||
one("cached", rate("cached_tokens_total")),
|
||||
],
|
||||
)
|
||||
add(
|
||||
"Output Throughput",
|
||||
"tokens/s",
|
||||
"Server generation counter deltas / scrape interval",
|
||||
[one("total", rate("generation_tokens_total"))],
|
||||
)
|
||||
add(
|
||||
"Cache Hit Rate",
|
||||
"%",
|
||||
"Counter deltas from prefill_effective_tokens_total; each tier / all prompt tokens. Total includes storage hits.",
|
||||
[
|
||||
one("L1 GPU", cache_hits[0], color=0),
|
||||
one("L2 CPU", cache_hits[1], color=1),
|
||||
one("Total", cache_hits[2], color=2, dashed=True),
|
||||
],
|
||||
)
|
||||
add(
|
||||
"E2E Request Latency",
|
||||
"ms",
|
||||
"Client submit to completed stream",
|
||||
dist("latency_s", 1000),
|
||||
)
|
||||
add(
|
||||
"KV Usage",
|
||||
"%",
|
||||
"L1: (active + evictable) / GPU capacity. L2: host used / host capacity.",
|
||||
[one("L1 GPU", l1_usage, color=0), one("L2 CPU", l2_usage, color=1)],
|
||||
)
|
||||
add(
|
||||
"KV Tokens",
|
||||
"tokens",
|
||||
"L1 occupied = active + evictable GPU tokens. L2 occupied = host used tokens. Dashed lines show capacity.",
|
||||
[
|
||||
one("L1 occupied", l1_occupied, color=0),
|
||||
one("L1 capacity", l1_capacity, color=0, dashed=True),
|
||||
one("L2 occupied", l2_occupied, color=1),
|
||||
one("L2 capacity", l2_capacity, color=1, dashed=True),
|
||||
],
|
||||
)
|
||||
add(
|
||||
"Mamba Usage",
|
||||
"%",
|
||||
"Maximum server mamba_usage across reported ranks",
|
||||
[
|
||||
one(
|
||||
"usage",
|
||||
[
|
||||
[
|
||||
t,
|
||||
(
|
||||
None
|
||||
if not (v := list(values(ss, "mamba_usage").values()))
|
||||
else max(v) * 100
|
||||
),
|
||||
]
|
||||
for t, ss in parsed
|
||||
],
|
||||
)
|
||||
],
|
||||
)
|
||||
add(
|
||||
"Mamba Tokens",
|
||||
"tokens",
|
||||
"Server Mamba token gauges",
|
||||
[
|
||||
one("used", gauge("mamba_used_tokens")),
|
||||
one("available", gauge("mamba_available_tokens")),
|
||||
one("evictable", gauge("mamba_evictable_tokens")),
|
||||
],
|
||||
)
|
||||
add(
|
||||
"Spec Acceptance Length",
|
||||
"tokens",
|
||||
"Unweighted mean spec_accept_length across reported ranks",
|
||||
[
|
||||
one(
|
||||
"accepted",
|
||||
[
|
||||
[
|
||||
t,
|
||||
(
|
||||
statistics.mean(v)
|
||||
if (v := list(values(ss, "spec_accept_length").values()))
|
||||
else None
|
||||
),
|
||||
]
|
||||
for t, ss in parsed
|
||||
],
|
||||
)
|
||||
],
|
||||
)
|
||||
return series
|
||||
|
||||
|
||||
def load_run(directory):
|
||||
manifest = json.loads((directory / "manifest.json").read_text())
|
||||
requests = read_rows(directory / "requests.jsonl")
|
||||
scrapes = read_rows(directory / "metrics.jsonl")
|
||||
sessions = read_rows(directory / "sessions.jsonl")
|
||||
origin = manifest["started_at"]
|
||||
stamps = [origin] + [
|
||||
r[k]
|
||||
for r in requests + scrapes + sessions
|
||||
for k in ("completed_at", "failed_at", "submitted_at", "timestamp")
|
||||
if isinstance(r.get(k), (int, float))
|
||||
]
|
||||
stamps.extend(
|
||||
r["submitted_at"] + event[0]
|
||||
for r in requests
|
||||
if "submitted_at" in r
|
||||
for event in r.get("events", [])
|
||||
)
|
||||
duration = max(0.001, manifest.get("finished_at", max(stamps)) - origin)
|
||||
info = manifest.get("server_info", {})
|
||||
args = manifest.get("arguments", {})
|
||||
server = info.get("server_args", info)
|
||||
rows = []
|
||||
for r in requests:
|
||||
meta = r.get("meta_info") or {}
|
||||
rank = meta.get("dp_rank")
|
||||
if rank is None:
|
||||
rank = r.get("rank")
|
||||
if rank is None and info.get("dp_size") == 1:
|
||||
rank = 0
|
||||
phases, unavailable = [], []
|
||||
for kind, first, last in (
|
||||
("Tool call", "tool_started_at", "tool_completed_at"),
|
||||
("Wait", "client_wait_started_at", "submitted_at"),
|
||||
(
|
||||
"Sampling",
|
||||
"submitted_at",
|
||||
"completed_at" if "completed_at" in r else "failed_at",
|
||||
),
|
||||
):
|
||||
if kind == "Tool call" and r.get("turn") == 0:
|
||||
continue
|
||||
if (
|
||||
r.get(first) is not None
|
||||
and r.get(last) is not None
|
||||
and r[last] >= r[first]
|
||||
):
|
||||
phases.append(
|
||||
{"type": kind, "start": r[first] - origin, "end": r[last] - origin}
|
||||
)
|
||||
else:
|
||||
unavailable.append(kind)
|
||||
rows.append(
|
||||
{
|
||||
"conversation": r.get("conversation", "unknown"),
|
||||
"turn": r.get("turn", "?"),
|
||||
"worker": rank,
|
||||
"phases": phases,
|
||||
"unavailable": unavailable,
|
||||
"error": r.get("error"),
|
||||
"ttft": r.get("ttft_s"),
|
||||
"latency": r.get("latency_s"),
|
||||
"wait": r.get("client_wait_s"),
|
||||
"context": r.get("context_tokens"),
|
||||
"output": meta.get("completion_tokens"),
|
||||
"hits": meta.get("cached_tokens_details") or {},
|
||||
}
|
||||
)
|
||||
exporters = {}
|
||||
for scrape in scrapes:
|
||||
if "timestamp" in scrape:
|
||||
exporters.setdefault(scrape.get("url", "unknown"), []).append(scrape)
|
||||
if not exporters:
|
||||
exporters["No metrics recorded"] = []
|
||||
summary = {
|
||||
"model": server.get("model_path", args.get("tokenizer", "Unknown model")),
|
||||
"duration": duration,
|
||||
"requests": len(requests),
|
||||
"conversations": len({r["conversation"] for r in rows}),
|
||||
"turns": args.get("turns", "?"),
|
||||
"concurrency": args.get("concurrency", "?"),
|
||||
"page_size": server.get("page_size", "?"),
|
||||
"status": (
|
||||
manifest.get("status", "incomplete")
|
||||
if manifest.get("finished_at")
|
||||
else "incomplete"
|
||||
),
|
||||
"errors": sum(bool(r.get("error")) for r in requests + sessions + scrapes),
|
||||
"error": manifest.get("error"),
|
||||
}
|
||||
return {
|
||||
"summary": summary,
|
||||
"rows": rows,
|
||||
"exporters": {
|
||||
url: panels_for(requests, records, origin, duration)
|
||||
for url, records in exporters.items()
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def build(directory, output):
|
||||
data = load_run(directory)
|
||||
assets = Path(__file__).parent
|
||||
encoded = json.dumps(data, allow_nan=False).replace("<", r"\u003c")
|
||||
html = (assets / "viewer.html").read_text()
|
||||
html = html.replace("__CSS__", (assets / "viewer.css").read_text())
|
||||
html = html.replace("__JS__", (assets / "viewer.js").read_text())
|
||||
html = html.replace("__REAL_DATA__", encoded)
|
||||
output.write_text(html)
|
||||
return data
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = argparse.ArgumentParser(description=__doc__)
|
||||
parser.add_argument("directory", type=Path, help="Recorded run directory")
|
||||
parser.add_argument(
|
||||
"--output",
|
||||
type=Path,
|
||||
default=Path("rollout.html"),
|
||||
help="Self-contained HTML output",
|
||||
)
|
||||
args = parser.parse_args()
|
||||
build(args.directory, args.output)
|
||||
print(args.output)
|
||||
Reference in New Issue
Block a user