444 lines
15 KiB
Python
444 lines
15 KiB
Python
# /// script
|
|
# requires-python = ">=3.11"
|
|
# dependencies = ["aiohttp", "transformers"]
|
|
# ///
|
|
"""Synthetic multi-turn HTTP workload. See README.md for examples."""
|
|
|
|
import argparse
|
|
import asyncio
|
|
import hashlib
|
|
import importlib.metadata
|
|
import json
|
|
import random
|
|
import sys
|
|
import time
|
|
import uuid
|
|
from pathlib import Path
|
|
|
|
import aiohttp
|
|
|
|
|
|
def write_json(path, value):
|
|
path.write_text(json.dumps(value, indent=2, sort_keys=True) + "\n")
|
|
|
|
|
|
def record(file, value):
|
|
file.write(json.dumps(value, sort_keys=True) + "\n")
|
|
file.flush()
|
|
|
|
|
|
def synthetic_tokens(tokenizer, seed, conversation, turn, length):
|
|
text = (
|
|
f"Session {seed}/{conversation}. Tool result {turn}: "
|
|
"The local search returned a synthetic measurement. "
|
|
"Read the result and continue the investigation. "
|
|
)
|
|
ids = tokenizer.encode(text, add_special_tokens=False)
|
|
if not ids:
|
|
raise ValueError("Tokenizer produced no input tokens")
|
|
return (ids * ((length + len(ids) - 1) // len(ids)))[:length]
|
|
|
|
|
|
async def sse_events(content):
|
|
# aiohttp iterates complete lines, even when TCP splits a UTF-8 character.
|
|
data = []
|
|
async for raw in content:
|
|
line = raw.decode("utf-8").rstrip("\r\n")
|
|
if not line:
|
|
if data:
|
|
yield "\n".join(data)
|
|
data.clear()
|
|
elif line.startswith("data:"):
|
|
data.append(line[5:].lstrip(" "))
|
|
if data:
|
|
raise ValueError("Truncated SSE frame")
|
|
|
|
|
|
async def check_response(response):
|
|
if response.status >= 400:
|
|
detail = (await response.text())[:2000]
|
|
raise ValueError(f"HTTP {response.status} {response.url}: {detail}")
|
|
|
|
|
|
async def json_request(http, method, url, **kwargs):
|
|
async with http.request(method, url, **kwargs) as response:
|
|
await check_response(response)
|
|
return await response.json()
|
|
|
|
|
|
async def generate(http, url, payload, row, incremental, dp_size=None):
|
|
start = time.perf_counter()
|
|
row.update(submitted_at=time.time(), events=[])
|
|
count, ids, meta, done = 0, [], {}, False
|
|
async with http.post(url + "/generate", json=payload) as response:
|
|
await check_response(response)
|
|
async for event in sse_events(response.content):
|
|
if event == "[DONE]":
|
|
done = True
|
|
break
|
|
message = json.loads(event)
|
|
if "error" in message:
|
|
raise ValueError(str(message["error"]))
|
|
meta = message["meta_info"]
|
|
current = meta["completion_tokens"]
|
|
if current < count:
|
|
raise ValueError("Completion count moved backwards")
|
|
if current > count:
|
|
row["events"].append([time.perf_counter() - start, current])
|
|
if not count:
|
|
row["first_token_at"] = time.time()
|
|
row["ttft_s"] = row["events"][-1][0]
|
|
count = current
|
|
if "output_ids" in message:
|
|
if incremental:
|
|
ids.extend(message["output_ids"])
|
|
else:
|
|
ids = message["output_ids"]
|
|
row["meta_info"] = meta
|
|
row.update(completed_at=time.time(), latency_s=time.perf_counter() - start)
|
|
if (
|
|
not done
|
|
or not count
|
|
or (meta.get("finish_reason") or {}).get("type") != "length"
|
|
):
|
|
raise ValueError(
|
|
f"Incomplete or aborted generation: {meta.get('finish_reason')}"
|
|
)
|
|
expected = payload["sampling_params"]["max_new_tokens"]
|
|
if count != expected or len(ids) != expected:
|
|
raise ValueError(
|
|
f"Expected {expected} tokens, got count={count}, IDs={len(ids)}"
|
|
)
|
|
if meta["prompt_tokens"] != row["context_tokens"]:
|
|
raise ValueError(f"Wrong context length: {meta['prompt_tokens']}")
|
|
if (
|
|
row["rank"] is not None
|
|
and meta.get("dp_rank") != row["rank"]
|
|
and not (dp_size == 1 and row["rank"] == 0 and meta.get("dp_rank") is None)
|
|
):
|
|
raise ValueError(f"Wrong DP rank: {meta.get('dp_rank')}")
|
|
events = row["events"]
|
|
row["avg_token_time_s"] = (
|
|
(events[-1][0] - events[0][0]) / (count - events[0][1])
|
|
if count > events[0][1]
|
|
else None
|
|
)
|
|
row["output_sha256"] = hashlib.sha256(json.dumps(ids).encode()).hexdigest()
|
|
return meta["id"], ids
|
|
|
|
|
|
async def scrape_metrics(http, url, file, stopped):
|
|
while not stopped.is_set():
|
|
started = time.monotonic()
|
|
row = {"timestamp": time.time(), "url": url}
|
|
try:
|
|
async with http.get(url) as response:
|
|
response.raise_for_status()
|
|
row["text"] = await response.text()
|
|
except (aiohttp.ClientError, OSError, asyncio.TimeoutError) as exc:
|
|
row["error"] = str(exc)
|
|
row["duration_s"] = time.monotonic() - started
|
|
record(file, row)
|
|
try:
|
|
await asyncio.wait_for(stopped.wait(), max(0.01, 1 - row["duration_s"]))
|
|
except asyncio.TimeoutError:
|
|
pass
|
|
|
|
|
|
def max_context_tokens(args):
|
|
return (
|
|
args.initial_tokens
|
|
+ (args.turns - 1) * (args.tool_tokens + args.output_tokens)
|
|
+ args.output_tokens
|
|
)
|
|
|
|
|
|
async def conversation(
|
|
http, tokenizer, args, info, slot, semaphore, requests, sessions
|
|
):
|
|
rank = slot % info["dp_size"] if not args.disable_dp_sticky_routing else None
|
|
rng = random.Random(f"{args.seed}/{slot}")
|
|
await asyncio.sleep(args.start_spread * slot / args.conversations)
|
|
sid, rid, history = None, None, []
|
|
identity = {"conversation": slot, "rank": rank}
|
|
try:
|
|
if args.mode != "full-history":
|
|
sid = uuid.uuid4().hex
|
|
opened = await json_request(
|
|
http,
|
|
"POST",
|
|
args.base_url + "/open_session",
|
|
json={
|
|
"session_id": sid,
|
|
"streaming": args.mode == "streaming",
|
|
# Required by older servers, but unused by session execution.
|
|
"capacity_of_str_len": 1000,
|
|
},
|
|
)
|
|
if opened != sid:
|
|
raise ValueError(f"Unexpected open response: {opened}")
|
|
record(
|
|
sessions,
|
|
{**identity, "event": "open", "id": sid, "timestamp": time.time()},
|
|
)
|
|
for turn in range(args.turns):
|
|
delay = rng.uniform(*args.tool_delay) if turn else 0
|
|
tool_started_at = time.time()
|
|
delay_start = time.perf_counter()
|
|
await asyncio.sleep(delay)
|
|
actual_delay = time.perf_counter() - delay_start
|
|
tool_completed_at = time.time()
|
|
length = args.tool_tokens if turn else args.initial_tokens
|
|
chunk = synthetic_tokens(tokenizer, args.seed, slot, turn, length)
|
|
history.extend(chunk)
|
|
row = {
|
|
**identity,
|
|
"turn": turn,
|
|
"context_tokens": len(history),
|
|
"tool_delay_s": delay,
|
|
"tool_started_at": tool_started_at if turn else None,
|
|
"tool_completed_at": tool_completed_at if turn else None,
|
|
"actual_tool_delay_s": actual_delay,
|
|
"input_tokens": length,
|
|
}
|
|
payload = {
|
|
"input_ids": history if args.mode == "full-history" else chunk,
|
|
"stream": True,
|
|
"sampling_params": {
|
|
"temperature": 0,
|
|
"ignore_eos": True,
|
|
"max_new_tokens": args.output_tokens,
|
|
},
|
|
}
|
|
if sid is not None:
|
|
payload["session_params"] = {"id": sid, "rid": rid}
|
|
if rank is not None:
|
|
payload["routed_dp_rank"] = rank
|
|
row["client_wait_started_at"] = time.time()
|
|
waiting = time.perf_counter()
|
|
try:
|
|
async with semaphore:
|
|
row["client_wait_s"] = time.perf_counter() - waiting
|
|
rid, output = await generate(
|
|
http,
|
|
args.base_url,
|
|
payload,
|
|
row,
|
|
info.get("incremental_streaming_output", False),
|
|
dp_size=info.get("dp_size"),
|
|
)
|
|
history.extend(output)
|
|
except BaseException as exc:
|
|
row["failed_at"] = time.time()
|
|
row["error"] = f"{type(exc).__name__}: {exc}"
|
|
raise
|
|
finally:
|
|
record(requests, row)
|
|
finally:
|
|
if sid is not None:
|
|
event = {
|
|
**identity,
|
|
"event": "close",
|
|
"id": sid,
|
|
"timestamp": time.time(),
|
|
}
|
|
try:
|
|
async with http.post(
|
|
args.base_url + "/close_session",
|
|
json={"session_id": sid},
|
|
timeout=aiohttp.ClientTimeout(total=10),
|
|
) as response:
|
|
response.raise_for_status()
|
|
except (aiohttp.ClientError, OSError, asyncio.TimeoutError) as exc:
|
|
event["error"] = str(exc)
|
|
raise
|
|
finally:
|
|
record(sessions, event)
|
|
|
|
|
|
async def run(args, tokenizer):
|
|
args.output_dir.mkdir(parents=True, exist_ok=False)
|
|
manifest = {
|
|
"arguments": {
|
|
k: str(v) if isinstance(v, Path) else v for k, v in vars(args).items()
|
|
},
|
|
"started_at": time.time(),
|
|
"status": "running",
|
|
"python": sys.version,
|
|
"script_sha256": hashlib.sha256(Path(__file__).read_bytes()).hexdigest(),
|
|
"dependencies": {
|
|
n: importlib.metadata.version(n) for n in ("aiohttp", "transformers")
|
|
},
|
|
}
|
|
stopped = asyncio.Event()
|
|
manifest_path = args.output_dir / "manifest.json"
|
|
write_json(manifest_path, manifest)
|
|
try:
|
|
async with (
|
|
aiohttp.ClientSession(
|
|
timeout=aiohttp.ClientTimeout(total=args.timeout),
|
|
connector=aiohttp.TCPConnector(limit=0),
|
|
) as http,
|
|
aiohttp.ClientSession(
|
|
timeout=aiohttp.ClientTimeout(total=10)
|
|
) as metrics_http,
|
|
):
|
|
info = await json_request(http, "GET", args.base_url + "/server_info")
|
|
manifest["server_info"] = info
|
|
context = max_context_tokens(args)
|
|
if info.get("context_length") and context > info["context_length"]:
|
|
raise ValueError(
|
|
f"Workload requires {context} tokens; server context is too small"
|
|
)
|
|
if not args.disable_dp_sticky_routing and info.get("dp_size", 0) < 1:
|
|
raise ValueError("DP sticky routing requires server_info.dp_size")
|
|
write_json(manifest_path, manifest)
|
|
semaphore = asyncio.Semaphore(args.concurrency)
|
|
with (
|
|
(args.output_dir / "requests.jsonl").open("w") as requests,
|
|
(args.output_dir / "sessions.jsonl").open("w") as sessions,
|
|
(args.output_dir / "metrics.jsonl").open("w") as metrics,
|
|
):
|
|
scrapers = [
|
|
asyncio.create_task(
|
|
scrape_metrics(metrics_http, url, metrics, stopped)
|
|
)
|
|
for url in (args.metrics_url or [args.base_url + "/metrics"])
|
|
]
|
|
try:
|
|
async with asyncio.TaskGroup() as group:
|
|
for slot in range(args.conversations):
|
|
group.create_task(
|
|
conversation(
|
|
http,
|
|
tokenizer,
|
|
args,
|
|
info,
|
|
slot,
|
|
semaphore,
|
|
requests,
|
|
sessions,
|
|
)
|
|
)
|
|
finally:
|
|
stopped.set()
|
|
await asyncio.gather(*scrapers)
|
|
manifest["status"] = "completed"
|
|
except BaseException as exc:
|
|
manifest.update(status="failed", error=repr(exc))
|
|
raise
|
|
finally:
|
|
manifest["finished_at"] = time.time()
|
|
write_json(manifest_path, manifest)
|
|
print(f"{manifest['status']}: {args.output_dir}", flush=True)
|
|
|
|
|
|
def parse_args(argv=None):
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
|
parser.add_argument(
|
|
"--base-url", default="http://127.0.0.1:30000", help="SGLang server URL"
|
|
)
|
|
parser.add_argument(
|
|
"--tokenizer",
|
|
required=True,
|
|
help="Model ID or tokenizer path matching the server",
|
|
)
|
|
parser.add_argument(
|
|
"--trust-remote-code", action="store_true", help="Allow custom tokenizer code"
|
|
)
|
|
parser.add_argument(
|
|
"--mode",
|
|
choices=("full-history", "ordinary", "streaming"),
|
|
default="ordinary",
|
|
help="Send full history or use ordinary/streaming sessions",
|
|
)
|
|
parser.add_argument(
|
|
"--conversations",
|
|
type=int,
|
|
default=8,
|
|
help="Number of independent conversations",
|
|
)
|
|
parser.add_argument(
|
|
"--concurrency",
|
|
type=int,
|
|
default=8,
|
|
help="Maximum in-flight generation requests",
|
|
)
|
|
parser.add_argument(
|
|
"--turns",
|
|
type=int,
|
|
default=8,
|
|
help="Generation turns per conversation",
|
|
)
|
|
parser.add_argument(
|
|
"--initial-tokens",
|
|
type=int,
|
|
default=2048,
|
|
help="Tokens in each initial prompt",
|
|
)
|
|
parser.add_argument(
|
|
"--tool-tokens",
|
|
type=int,
|
|
default=256,
|
|
help="Synthetic tool-result tokens appended each turn",
|
|
)
|
|
parser.add_argument(
|
|
"--output-tokens",
|
|
type=int,
|
|
default=128,
|
|
help="Exact number of tokens generated per turn",
|
|
)
|
|
parser.add_argument(
|
|
"--tool-delay",
|
|
type=float,
|
|
nargs=2,
|
|
default=[1, 3],
|
|
metavar=("MIN", "MAX"),
|
|
help="Range of simulated tool waits between turns, in seconds",
|
|
)
|
|
parser.add_argument(
|
|
"--start-spread",
|
|
type=float,
|
|
default=5,
|
|
help="Seconds over which conversation starts are staggered",
|
|
)
|
|
parser.add_argument(
|
|
"--seed", type=int, default=1, help="Seed for synthetic inputs and tool delays"
|
|
)
|
|
parser.add_argument(
|
|
"--disable-dp-sticky-routing",
|
|
action="store_true",
|
|
help="Disable default routing of conversation i to data-parallel rank i %% dp_size",
|
|
)
|
|
parser.add_argument(
|
|
"--timeout",
|
|
type=float,
|
|
default=600,
|
|
help="HTTP request timeout in seconds",
|
|
)
|
|
parser.add_argument(
|
|
"--metrics-url",
|
|
action="append",
|
|
help="Metrics endpoint; repeat for distinct exporters (default: BASE_URL/metrics)",
|
|
)
|
|
parser.add_argument(
|
|
"--output-dir",
|
|
type=Path,
|
|
required=True,
|
|
help="New directory for request records, metrics, and run settings",
|
|
)
|
|
args = parser.parse_args(argv)
|
|
args.base_url = args.base_url.rstrip("/")
|
|
return args
|
|
|
|
|
|
if __name__ == "__main__":
|
|
from transformers import AutoTokenizer
|
|
|
|
args = parse_args()
|
|
tokenizer = AutoTokenizer.from_pretrained(
|
|
args.tokenizer,
|
|
trust_remote_code=args.trust_remote_code,
|
|
)
|
|
asyncio.run(run(args, tokenizer))
|