[Benchmark] Add agentic rollout simulator and offline explorer (#40034)
This commit is contained in:
Byron Hsu
@@ -0,0 +1,173 @@
|
||||
"""Explorer checks use small recordings, without GPUs or external services."""
|
||||
|
||||
import json
|
||||
import sys
|
||||
import tempfile
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
|
||||
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
|
||||
import explore
|
||||
from metrics import counter_rate, hit_percentages, metric_values, occupancy
|
||||
|
||||
|
||||
def counters(device=0, host=0, storage=0, uncached=0, rank="0"):
|
||||
return metric_values(
|
||||
"\n".join(
|
||||
f'sglang:prefill_effective_tokens_total{{mode="{mode}",dp_rank="{rank}"}} {value}'
|
||||
for mode, value in zip(
|
||||
("device_hit", "host_hit", "storage_hit", "input"),
|
||||
(device, host, storage, uncached),
|
||||
)
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
class MetricTests(unittest.TestCase):
|
||||
def test_occupancy_includes_evictable_and_matches_labels(self):
|
||||
samples = metric_values("""
|
||||
sglang:kv_used_tokens{rank="0"} 20
|
||||
sglang:kv_evictable_tokens{rank="0"} 30
|
||||
sglang:max_total_num_tokens{rank="0"} 100
|
||||
sglang:hicache_host_used_tokens{rank="0"} 80
|
||||
sglang:hicache_host_total_tokens{rank="0"} 200
|
||||
""")
|
||||
names = ("kv_used_tokens", "kv_evictable_tokens", "max_total_num_tokens")
|
||||
self.assertEqual(occupancy(samples, names), [20, 30, 100, 50])
|
||||
self.assertEqual(
|
||||
occupancy(
|
||||
samples, ("hicache_host_used_tokens", "hicache_host_total_tokens")
|
||||
),
|
||||
[80, 200, 40],
|
||||
)
|
||||
samples.update(metric_values('sglang:kv_used_tokens{rank="1"} 1'))
|
||||
self.assertEqual(occupancy(samples, names), [None] * 4)
|
||||
|
||||
def test_hits_include_storage_and_uncached_denominator(self):
|
||||
self.assertEqual(
|
||||
hit_percentages(counters(), counters(40, 20, 10, 30)), [40, 20, 70]
|
||||
)
|
||||
self.assertEqual(
|
||||
hit_percentages(counters(), counters(80, 0, 0, 20)), [80, 0, 80]
|
||||
)
|
||||
self.assertEqual(hit_percentages(counters(), counters()), [None] * 3)
|
||||
missing = counters(80, 0, 0, 20)
|
||||
missing.pop(next(k for k in missing if dict(k[1])["mode"] == "host_hit"))
|
||||
self.assertEqual(hit_percentages(counters(), missing), [None] * 3)
|
||||
|
||||
def test_per_label_reset_cannot_be_hidden_by_another_worker(self):
|
||||
before = counters(20) | counters(20, rank="1")
|
||||
after = counters(19) | counters(100, rank="1")
|
||||
self.assertEqual(hit_percentages(before, after), [None] * 3)
|
||||
self.assertIsNone(
|
||||
counter_rate(before, after, "sglang:prefill_effective_tokens_total", 1)
|
||||
)
|
||||
self.assertIsNone(
|
||||
counter_rate({}, after, "sglang:prefill_effective_tokens_total", 1)
|
||||
)
|
||||
|
||||
def test_scrape_error_breaks_rates(self):
|
||||
records = [
|
||||
{"timestamp": 0, "text": "sglang:generation_tokens_total 1"},
|
||||
{"timestamp": 1, "error": "timeout"},
|
||||
{"timestamp": 2, "text": "sglang:generation_tokens_total 9"},
|
||||
]
|
||||
panels = explore.panels_for([], records, 0, 3)
|
||||
rate = next(p for p in panels if p["title"] == "Output Throughput")
|
||||
self.assertEqual(rate["series"][0]["points"], [[1, None], [2, None]])
|
||||
|
||||
def test_nonfinite_missing_and_zero_capacity(self):
|
||||
self.assertEqual(metric_values("sglang:kv_used_tokens NaN"), {})
|
||||
self.assertEqual(occupancy({}, ("a", "b")), [None] * 3)
|
||||
self.assertEqual(
|
||||
occupancy(metric_values("sglang:a 1\nsglang:b 0"), ("a", "b")), [1, 0, None]
|
||||
)
|
||||
|
||||
|
||||
class RecordingTests(unittest.TestCase):
|
||||
def test_old_failed_recording_and_offline_assets(self):
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
path = Path(tmp)
|
||||
(path / "manifest.json").write_text(
|
||||
json.dumps(
|
||||
{
|
||||
"started_at": 10,
|
||||
"status": "failed",
|
||||
"finished_at": 15,
|
||||
"server_info": {"dp_size": 1},
|
||||
"arguments": {
|
||||
"tokenizer": "</script><script>alert(1)</script>"
|
||||
},
|
||||
}
|
||||
)
|
||||
)
|
||||
(path / "requests.jsonl").write_text(
|
||||
json.dumps(
|
||||
{
|
||||
"conversation": 7,
|
||||
"turn": 2,
|
||||
"submitted_at": 11,
|
||||
"failed_at": 14,
|
||||
"ttft_s": 0.5,
|
||||
"first_token_at": 11.5,
|
||||
"events": [[0.5, 1], [1, 2]],
|
||||
"error": "disconnected",
|
||||
"meta_info": {"completion_tokens": 2},
|
||||
}
|
||||
)
|
||||
+ "\n"
|
||||
)
|
||||
(path / "metrics.jsonl").write_text(
|
||||
"\n".join(
|
||||
json.dumps(
|
||||
{
|
||||
"timestamp": 12,
|
||||
"url": url,
|
||||
"text": f"sglang:num_running_reqs {n}",
|
||||
}
|
||||
)
|
||||
for url, n in (("one", 2), ("two", 3))
|
||||
)
|
||||
)
|
||||
output = path / "view.html"
|
||||
data = explore.build(path, output)
|
||||
row = data["rows"][0]
|
||||
self.assertEqual(row["worker"], 0)
|
||||
self.assertEqual(row["unavailable"], ["Tool call", "Wait"])
|
||||
self.assertEqual(
|
||||
row["phases"], [{"type": "Sampling", "start": 1, "end": 4}]
|
||||
)
|
||||
self.assertEqual(data["summary"]["errors"], 1)
|
||||
self.assertEqual(set(data["exporters"]), {"one", "two"})
|
||||
for url, expected in (("one", 2), ("two", 3)):
|
||||
panels = data["exporters"][url]
|
||||
self.assertEqual(len(panels), 15)
|
||||
running = next(
|
||||
p for p in panels if p["title"] == "Running And Queued Requests"
|
||||
)
|
||||
self.assertEqual(running["series"][0]["points"], [[2, expected]])
|
||||
ttft = next(p for p in panels if p["title"] == "TTFT")
|
||||
self.assertTrue(any(v == 500 for _, v in ttft["series"][0]["points"]))
|
||||
html = output.read_text()
|
||||
self.assertNotIn("</script><script>alert", html)
|
||||
self.assertNotIn("<script src", html)
|
||||
self.assertNotIn("<link", html)
|
||||
self.assertNotIn("fetch(", html)
|
||||
self.assertNotIn("__REAL_DATA__", html)
|
||||
self.assertIn("Timeline", html)
|
||||
self.assertIn("Engine Metrics", html)
|
||||
|
||||
def test_interrupted_empty_recording(self):
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
path = Path(tmp)
|
||||
(path / "manifest.json").write_text('{"started_at":1,"status":"running"}')
|
||||
(path / "requests.jsonl").write_text('{"conversation":')
|
||||
data = explore.load_run(path)
|
||||
self.assertEqual(data["summary"]["status"], "incomplete")
|
||||
self.assertEqual(data["summary"]["errors"], 1)
|
||||
self.assertEqual(data["rows"][0]["phases"], [])
|
||||
self.assertIsNone(data["rows"][0]["worker"])
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,341 @@
|
||||
"""CPU protocol tests using a local HTTP server; no model downloads."""
|
||||
|
||||
import builtins
|
||||
import json
|
||||
import sys
|
||||
import tempfile
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
|
||||
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
|
||||
import plot as plots
|
||||
import simulate as bench
|
||||
from aiohttp import web
|
||||
|
||||
|
||||
class Tokenizer:
|
||||
def encode(self, text, **kwargs):
|
||||
return list(text.encode())
|
||||
|
||||
|
||||
class ProtocolTests(unittest.IsolatedAsyncioTestCase):
|
||||
async def asyncSetUp(self):
|
||||
self.incremental = False
|
||||
self.dp_size = 2
|
||||
self.omit_rank = False
|
||||
self.failure = None
|
||||
self.sessions = {}
|
||||
self.closed = []
|
||||
self.payloads = []
|
||||
|
||||
async def info(request):
|
||||
return web.json_response(
|
||||
{
|
||||
"dp_size": self.dp_size,
|
||||
"context_length": 10000,
|
||||
"incremental_streaming_output": self.incremental,
|
||||
}
|
||||
)
|
||||
|
||||
async def open_session(request):
|
||||
payload = await request.json()
|
||||
if self.failure == "disabled":
|
||||
raise web.HTTPBadRequest(text="Streaming sessions are disabled")
|
||||
self.sessions[payload["session_id"]] = ([], None, payload["streaming"])
|
||||
return web.json_response(payload["session_id"])
|
||||
|
||||
async def close_session(request):
|
||||
self.closed.append((await request.json())["session_id"])
|
||||
return web.json_response(None)
|
||||
|
||||
async def generate(request):
|
||||
payload = await request.json()
|
||||
self.payloads.append(payload)
|
||||
ids = payload["input_ids"]
|
||||
sid = payload.get("session_params", {}).get("id")
|
||||
if sid:
|
||||
history, previous, streaming = self.sessions[sid]
|
||||
self.assertEqual(payload["session_params"]["rid"], previous)
|
||||
ids = history + ids
|
||||
rid = str(len(self.payloads))
|
||||
count = payload["sampling_params"]["max_new_tokens"]
|
||||
output = list(range(1000, 1000 + count))
|
||||
if sid:
|
||||
self.sessions[sid] = (ids + output, rid, streaming)
|
||||
response = web.StreamResponse(headers={"Content-Type": "text/event-stream"})
|
||||
await response.prepare(request)
|
||||
for n in (0, 1, count):
|
||||
meta = {
|
||||
"id": rid,
|
||||
"completion_tokens": n,
|
||||
"prompt_tokens": len(ids),
|
||||
"dp_rank": payload.get("routed_dp_rank", 0),
|
||||
"cached_tokens": 0,
|
||||
"finish_reason": {"type": "length"} if n == count else None,
|
||||
}
|
||||
if self.omit_rank:
|
||||
meta.pop("dp_rank")
|
||||
if self.failure == "context":
|
||||
meta["prompt_tokens"] += 1
|
||||
if self.failure == "rank":
|
||||
meta["dp_rank"] += 1
|
||||
if self.failure == "abort" and n == count:
|
||||
meta["finish_reason"] = {"type": "abort"}
|
||||
tokens = output[:n]
|
||||
if self.incremental and n == count:
|
||||
tokens = output[1:]
|
||||
frame = (
|
||||
"data: "
|
||||
+ json.dumps(
|
||||
{"meta_info": meta, "output_ids": tokens, "text": "é"},
|
||||
ensure_ascii=False,
|
||||
)
|
||||
+ "\r\n\r\n"
|
||||
).encode()
|
||||
for byte in frame:
|
||||
await response.write(bytes([byte]))
|
||||
if self.failure != "truncated":
|
||||
await response.write(b"data: [DONE]\n\n")
|
||||
return response
|
||||
|
||||
async def metrics(request):
|
||||
return web.Response(text='sglang:num_running_reqs{dp_rank="0"} 1\n')
|
||||
|
||||
app = web.Application()
|
||||
for method, path, handler in [
|
||||
("GET", "/server_info", info),
|
||||
("POST", "/open_session", open_session),
|
||||
("POST", "/close_session", close_session),
|
||||
("POST", "/generate", generate),
|
||||
("GET", "/metrics", metrics),
|
||||
]:
|
||||
app.router.add_route(method, path, handler)
|
||||
self.runner = web.AppRunner(app)
|
||||
await self.runner.setup()
|
||||
site = web.TCPSite(self.runner, "127.0.0.1", 0)
|
||||
await site.start()
|
||||
self.url = f"http://127.0.0.1:{self.runner.addresses[0][1]}"
|
||||
|
||||
async def asyncTearDown(self):
|
||||
await self.runner.cleanup()
|
||||
|
||||
def args(self, path, mode, extra=()):
|
||||
return bench.parse_args(
|
||||
[
|
||||
"--base-url",
|
||||
self.url,
|
||||
"--tokenizer",
|
||||
"fake",
|
||||
"--output-dir",
|
||||
str(path),
|
||||
"--mode",
|
||||
mode,
|
||||
"--conversations",
|
||||
"2",
|
||||
"--concurrency",
|
||||
"1",
|
||||
"--turns",
|
||||
"3",
|
||||
"--initial-tokens",
|
||||
"8",
|
||||
"--tool-tokens",
|
||||
"4",
|
||||
"--output-tokens",
|
||||
"2",
|
||||
"--tool-delay",
|
||||
"0",
|
||||
"0",
|
||||
"--start-spread",
|
||||
"0",
|
||||
]
|
||||
+ list(extra)
|
||||
)
|
||||
|
||||
async def test_all_modes_and_stream_formats_preserve_history(self):
|
||||
for incremental in (False, True):
|
||||
for mode in ("full-history", "ordinary", "streaming"):
|
||||
with (
|
||||
self.subTest(incremental=incremental, mode=mode),
|
||||
tempfile.TemporaryDirectory() as tmp,
|
||||
):
|
||||
self.incremental = incremental
|
||||
self.sessions.clear()
|
||||
self.closed.clear()
|
||||
self.payloads.clear()
|
||||
path = Path(tmp) / "run"
|
||||
await bench.run(self.args(path, mode), Tokenizer())
|
||||
rows = list(plots.read_rows(path / "requests.jsonl"))
|
||||
self.assertEqual(len(rows), 6)
|
||||
self.assertTrue(
|
||||
all(
|
||||
r["meta_info"]["dp_rank"] == r["conversation"] % 2
|
||||
for r in rows
|
||||
)
|
||||
)
|
||||
self.assertTrue(all("routed_dp_rank" in p for p in self.payloads))
|
||||
self.assertEqual({r["conversation"] for r in rows}, {0, 1})
|
||||
for conversation in (0, 1):
|
||||
self.assertEqual(
|
||||
[
|
||||
r["turn"]
|
||||
for r in rows
|
||||
if r["conversation"] == conversation
|
||||
],
|
||||
[0, 1, 2],
|
||||
)
|
||||
self.assertTrue(
|
||||
all(r["context_tokens"] == 8 + r["turn"] * 6 for r in rows)
|
||||
)
|
||||
self.assertTrue(
|
||||
all(r["meta_info"]["completion_tokens"] == 2 for r in rows)
|
||||
)
|
||||
self.assertEqual(
|
||||
len(self.closed), 0 if mode == "full-history" else 2
|
||||
)
|
||||
self.assertTrue(
|
||||
all(
|
||||
s[2] == (mode == "streaming")
|
||||
for s in self.sessions.values()
|
||||
)
|
||||
)
|
||||
if mode == "full-history":
|
||||
self.assertTrue(
|
||||
any(
|
||||
p["input_ids"][8:10] == [1000, 1001]
|
||||
for p in self.payloads
|
||||
)
|
||||
)
|
||||
plots.analyze(path, 30)
|
||||
self.assertEqual(
|
||||
json.loads((path / "summary.json").read_text())["status"],
|
||||
"completed",
|
||||
)
|
||||
|
||||
async def test_omitted_rank_requires_verified_single_worker(self):
|
||||
self.omit_rank = True
|
||||
for size in (1, 2):
|
||||
self.dp_size = size
|
||||
with self.subTest(dp_size=size), tempfile.TemporaryDirectory() as tmp:
|
||||
path = Path(tmp) / "run"
|
||||
if size == 1:
|
||||
await bench.run(self.args(path, "ordinary"), Tokenizer())
|
||||
rows = list(plots.read_rows(path / "requests.jsonl"))
|
||||
for row in rows:
|
||||
self.assertLessEqual(
|
||||
row["client_wait_started_at"], row["submitted_at"]
|
||||
)
|
||||
if row["turn"]:
|
||||
self.assertLessEqual(
|
||||
row["tool_started_at"], row["tool_completed_at"]
|
||||
)
|
||||
self.assertLessEqual(
|
||||
row["tool_completed_at"], row["client_wait_started_at"]
|
||||
)
|
||||
else:
|
||||
with self.assertRaises(builtins.ExceptionGroup):
|
||||
await bench.run(self.args(path, "ordinary"), Tokenizer())
|
||||
|
||||
async def test_disable_dp_sticky_routing(self):
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
args = self.args(
|
||||
Path(tmp) / "run", "ordinary", ["--disable-dp-sticky-routing"]
|
||||
)
|
||||
await bench.run(args, Tokenizer())
|
||||
self.assertEqual(len(self.payloads), 6)
|
||||
self.assertTrue(all("routed_dp_rank" not in p for p in self.payloads))
|
||||
|
||||
async def test_invalid_stream_fails_and_closes_owned_sessions(self):
|
||||
for failure in ("abort", "context", "rank", "truncated"):
|
||||
with self.subTest(failure=failure), tempfile.TemporaryDirectory() as tmp:
|
||||
self.failure = failure
|
||||
self.sessions.clear()
|
||||
self.closed.clear()
|
||||
path = Path(tmp) / "run"
|
||||
with self.assertRaises(builtins.ExceptionGroup):
|
||||
await bench.run(self.args(path, "ordinary"), Tokenizer())
|
||||
self.assertEqual(
|
||||
json.loads((path / "manifest.json").read_text())["status"], "failed"
|
||||
)
|
||||
self.assertEqual(set(self.closed), set(self.sessions))
|
||||
self.assertTrue(
|
||||
any("error" in r for r in plots.read_rows(path / "requests.jsonl"))
|
||||
)
|
||||
|
||||
async def test_server_rejection_is_preserved(self):
|
||||
self.failure = "disabled"
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
path = Path(tmp) / "run"
|
||||
with self.assertRaises(builtins.ExceptionGroup):
|
||||
await bench.run(self.args(path, "streaming"), Tokenizer())
|
||||
manifest = json.loads((path / "manifest.json").read_text())
|
||||
self.assertIn("Streaming sessions are disabled", manifest["error"])
|
||||
self.assertIn("HTTP 400", manifest["error"])
|
||||
|
||||
async def test_timeout_is_failure(self):
|
||||
# A timeout while entering the request must not become a successful sample.
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
with (
|
||||
tempfile.TemporaryDirectory() as tmp,
|
||||
patch.object(bench, "generate", new=AsyncMock(side_effect=TimeoutError)),
|
||||
):
|
||||
path = Path(tmp) / "run"
|
||||
with self.assertRaises(builtins.ExceptionGroup):
|
||||
await bench.run(self.args(path, "ordinary"), Tokenizer())
|
||||
self.assertEqual(set(self.closed), set(self.sessions))
|
||||
|
||||
|
||||
class MeasurementTests(unittest.TestCase):
|
||||
def test_inputs_are_repeatable_and_distinct(self):
|
||||
a = bench.synthetic_tokens(Tokenizer(), 1, 0, 0, 100)
|
||||
self.assertEqual(a, bench.synthetic_tokens(Tokenizer(), 1, 0, 0, 100))
|
||||
self.assertNotEqual(a, bench.synthetic_tokens(Tokenizer(), 1, 1, 0, 100))
|
||||
self.assertEqual(len(a), 100)
|
||||
|
||||
def test_token_weighted_cache_and_missing_data(self):
|
||||
rows = [
|
||||
{"meta_info": {"prompt_tokens": 100, "cached_tokens": 100}},
|
||||
{"meta_info": {"prompt_tokens": 900, "cached_tokens": 0}},
|
||||
]
|
||||
self.assertEqual(plots.cache_hit(rows, "total"), 0.1)
|
||||
self.assertIsNone(plots.cache_hit(rows, "host"))
|
||||
self.assertIsNone(plots.cache_hit([], "total"))
|
||||
|
||||
def test_throughput_uses_arrivals_and_partial_window_duration(self):
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
root = Path(tmp)
|
||||
bench.write_json(
|
||||
root / "manifest.json",
|
||||
{
|
||||
"status": "completed",
|
||||
"started_at": 0,
|
||||
"finished_at": 1.5,
|
||||
},
|
||||
)
|
||||
row = {
|
||||
"turn": 0,
|
||||
"context_tokens": 8,
|
||||
"submitted_at": 0,
|
||||
"first_token_at": 0.2,
|
||||
"completed_at": 1.4,
|
||||
"events": [[0.2, 1], [1.2, 3]],
|
||||
"ttft_s": 0.2,
|
||||
"avg_token_time_s": 0.5,
|
||||
"meta_info": {"prompt_tokens": 8, "completion_tokens": 3},
|
||||
}
|
||||
with (root / "requests.jsonl").open("w") as file:
|
||||
bench.record(file, row)
|
||||
_, windows, _ = plots.analyze(root, 1)
|
||||
self.assertEqual([w["output_tokens_s"] for w in windows], [1, 4])
|
||||
self.assertIsNone(windows[1]["ttft_p95_s"])
|
||||
|
||||
def test_counter_labels_resets_and_missing_series(self):
|
||||
before = plots.metric_values('x_total{rank="0"} 5\nx_total{rank="1"} 8\n')
|
||||
after = plots.metric_values('x_total{rank="0"} 7\nx_total{rank="1"} 12\n')
|
||||
self.assertEqual(plots.counter_rate(before, after, "x_total", 2), 3)
|
||||
self.assertIsNone(plots.counter_rate(after, before, "x_total", 2))
|
||||
self.assertIsNone(plots.counter_rate({}, after, "x_total", 2))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Reference in New Issue
Block a user