63 lines
2.5 KiB
Python
63 lines
2.5 KiB
Python
"""Time to first token and decode speed of an OpenAI-compatible server, as JSON lines.
|
|
|
|
usage (prain installed): python3 scripts/bench.py http://127.0.0.1:8000 [--tokens 128] [--runs 3]
|
|
Each run sends a fresh prompt (a random prefix defeats the prefix cache) and decodes exactly
|
|
--tokens tokens (ignore_eos), so numbers compare across runs and configs. The last line holds
|
|
the medians. --context FILE... puts the files' text (up to --context-chars) before every
|
|
prompt and the random tag after it: run 1 prefills the whole context, later runs find it in
|
|
the prefix cache, like an agent's system prompt and tools.
|
|
"""
|
|
|
|
import argparse
|
|
import json
|
|
import statistics
|
|
import time
|
|
import uuid
|
|
|
|
from prain.chat import chat_body, first_model, post, stream_text
|
|
|
|
PROMPTS = [
|
|
"Explain how a CPU cache works, step by step, for a beginner.",
|
|
"Write a short story about a lighthouse keeper who finds a message in a bottle.",
|
|
"Schreibe eine kurze Anleitung, wie man Brot backt.",
|
|
]
|
|
|
|
|
|
def run(base_url, model, prompt, tokens, context=""):
|
|
text = f"{context}[{uuid.uuid4().hex[:8]}] {prompt}"
|
|
body = chat_body(model, [{"role": "user", "content": text}], max_tokens=tokens,
|
|
temperature=0, ignore_eos=True)
|
|
start = time.monotonic()
|
|
stamps = []
|
|
with post(base_url, "/v1/chat/completions", body, timeout=3600) as response:
|
|
for _ in stream_text(response):
|
|
stamps.append(time.monotonic())
|
|
return {"ttft_s": round(stamps[0] - start, 3),
|
|
"decode_tok_s": round((tokens - 1) / max(stamps[-1] - stamps[0], 1e-9), 3)}
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("base_url")
|
|
parser.add_argument("--tokens", type=int, default=128)
|
|
parser.add_argument("--runs", type=int, default=3)
|
|
parser.add_argument("--context", nargs="*", default=[], help="files to put before each prompt")
|
|
parser.add_argument("--context-chars", type=int, default=48000)
|
|
args = parser.parse_args()
|
|
model = first_model(args.base_url)
|
|
context = ""
|
|
for path in args.context:
|
|
with open(path) as f:
|
|
context += f.read()
|
|
context = context[: args.context_chars] + "\n\n" if context else ""
|
|
results = []
|
|
for i in range(args.runs):
|
|
results.append(run(args.base_url, model, PROMPTS[i % len(PROMPTS)], args.tokens, context))
|
|
print(json.dumps(results[-1]), flush=True)
|
|
print(json.dumps({key: statistics.median(r[key] for r in results) for key in results[0]}
|
|
| {"median_of": len(results), "tokens": args.tokens}))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|