214 lines
7.8 KiB
Python
214 lines
7.8 KiB
Python
#!/usr/bin/env python3
|
|
"""Multi-turn session benchmark for OpenAI-compatible APIs.
|
|
|
|
Runs concurrent multi-turn chat sessions and reports per-turn,
|
|
per-session, and aggregate throughput metrics.
|
|
"""
|
|
import argparse
|
|
import asyncio
|
|
import json
|
|
import sys
|
|
import time
|
|
|
|
from openai import AsyncOpenAI
|
|
|
|
INITIAL_PROMPTS = [
|
|
"Write a Python function that implements a lock-free concurrent hash map using compare-and-swap operations. Include proper memory ordering.",
|
|
"Explain the mathematical foundations of diffusion models in machine learning. Start from the forward process and derive the reverse process.",
|
|
"Design a distributed consensus protocol for a system with Byzantine fault tolerance. Describe the phases and prove the safety properties.",
|
|
"Implement a B+ tree in Rust with support for range queries, bulk loading, and concurrent access using optimistic locking.",
|
|
"Analyze the computational complexity of the Aho-Corasick algorithm and compare it to naive multi-pattern matching. Provide the proof.",
|
|
"Write a CUDA kernel for flash attention with causal masking that handles variable sequence lengths within a batch.",
|
|
"Derive the optimal batch size for gradient descent given a fixed compute budget, following the scaling laws from Kaplan et al.",
|
|
"Design an LSM-tree based key-value store with write-ahead logging, compaction strategies, and bloom filters for read optimization.",
|
|
]
|
|
|
|
FOLLOW_UP_PROMPTS = [
|
|
"Can you explain the most complex part of that in more detail?",
|
|
"What are the main failure modes and how would you handle them?",
|
|
"Now optimize that for a production environment with 10x the scale.",
|
|
"Write comprehensive tests for the core logic you described.",
|
|
"What are the tradeoffs compared to the most common alternative approach?",
|
|
]
|
|
|
|
|
|
async def run_session(
|
|
client: AsyncOpenAI,
|
|
session_id: int,
|
|
model: str,
|
|
turns_per_session: int,
|
|
max_tokens: int,
|
|
temperature: float,
|
|
timeout_seconds: float,
|
|
) -> dict:
|
|
messages = []
|
|
turn_results = []
|
|
session_start = time.monotonic()
|
|
deadline = session_start + timeout_seconds
|
|
|
|
initial_prompt = INITIAL_PROMPTS[session_id % len(INITIAL_PROMPTS)]
|
|
|
|
for turn_idx in range(turns_per_session):
|
|
if turn_idx == 0:
|
|
user_content = initial_prompt
|
|
else:
|
|
user_content = FOLLOW_UP_PROMPTS[(turn_idx - 1) % len(FOLLOW_UP_PROMPTS)]
|
|
|
|
messages.append({"role": "user", "content": user_content})
|
|
|
|
remaining = deadline - time.monotonic()
|
|
if remaining <= 0:
|
|
break
|
|
|
|
turn_start = time.monotonic()
|
|
try:
|
|
response = await asyncio.wait_for(
|
|
client.chat.completions.create(
|
|
model=model,
|
|
messages=messages,
|
|
max_tokens=max_tokens,
|
|
temperature=temperature,
|
|
),
|
|
timeout=remaining,
|
|
)
|
|
except (asyncio.TimeoutError, Exception) as exc:
|
|
turn_results.append({
|
|
"turn": turn_idx + 1,
|
|
"error": f"{type(exc).__name__}: {exc}",
|
|
})
|
|
break
|
|
|
|
turn_wall = time.monotonic() - turn_start
|
|
usage = response.usage
|
|
completion_tokens = usage.completion_tokens if usage else 0
|
|
prompt_tokens = usage.prompt_tokens if usage else 0
|
|
tok_per_sec = completion_tokens / turn_wall if turn_wall > 0 else 0.0
|
|
|
|
turn_results.append({
|
|
"turn": turn_idx + 1,
|
|
"prompt_tokens": prompt_tokens,
|
|
"completion_tokens": completion_tokens,
|
|
"wall_seconds": round(turn_wall, 3),
|
|
"tok_per_sec": round(tok_per_sec, 1),
|
|
})
|
|
|
|
assistant_content = response.choices[0].message.content or ""
|
|
messages.append({"role": "assistant", "content": assistant_content})
|
|
|
|
total_completion = sum(
|
|
t.get("completion_tokens", 0) for t in turn_results
|
|
)
|
|
total_wall = time.monotonic() - session_start
|
|
turns_completed = sum(1 for t in turn_results if "error" not in t)
|
|
avg_tok_per_sec = total_completion / total_wall if total_wall > 0 else 0.0
|
|
|
|
return {
|
|
"session_id": session_id,
|
|
"turns": turn_results,
|
|
"total_completion_tokens": total_completion,
|
|
"total_wall_seconds": round(total_wall, 3),
|
|
"avg_tok_per_sec": round(avg_tok_per_sec, 1),
|
|
"turns_completed": turns_completed,
|
|
}
|
|
|
|
|
|
async def run_benchmark(args: argparse.Namespace) -> dict:
|
|
client = AsyncOpenAI(base_url=args.base_url, api_key="unused")
|
|
|
|
tasks = [
|
|
run_session(
|
|
client=client,
|
|
session_id=i,
|
|
model=args.model,
|
|
turns_per_session=args.turns_per_session,
|
|
max_tokens=args.max_tokens,
|
|
temperature=args.temperature,
|
|
timeout_seconds=args.timeout_seconds,
|
|
)
|
|
for i in range(args.sessions)
|
|
]
|
|
|
|
wall_start = time.monotonic()
|
|
session_results = await asyncio.gather(*tasks)
|
|
wall_total = time.monotonic() - wall_start
|
|
|
|
total_completion = sum(s["total_completion_tokens"] for s in session_results)
|
|
turns_completed = sum(s["turns_completed"] for s in session_results)
|
|
sessions_completed = sum(
|
|
1 for s in session_results if s["turns_completed"] == args.turns_per_session
|
|
)
|
|
per_session_rates = [
|
|
s["avg_tok_per_sec"]
|
|
for s in session_results
|
|
if s["turns_completed"] > 0
|
|
]
|
|
mean_per_session = (
|
|
sum(per_session_rates) / len(per_session_rates)
|
|
if per_session_rates
|
|
else 0.0
|
|
)
|
|
|
|
return {
|
|
"config": {
|
|
"sessions": args.sessions,
|
|
"turns_per_session": args.turns_per_session,
|
|
"max_tokens": args.max_tokens,
|
|
"temperature": args.temperature,
|
|
"model": args.model,
|
|
},
|
|
"sessions": session_results,
|
|
"aggregate": {
|
|
"total_completion_tokens": total_completion,
|
|
"total_wall_seconds": round(wall_total, 3),
|
|
"aggregate_tok_per_sec": round(
|
|
total_completion / wall_total if wall_total > 0 else 0.0, 1
|
|
),
|
|
"mean_per_session_tok_per_sec": round(mean_per_session, 1),
|
|
"sessions_completed": sessions_completed,
|
|
"turns_completed": turns_completed,
|
|
},
|
|
}
|
|
|
|
|
|
def main() -> int:
|
|
parser = argparse.ArgumentParser(
|
|
description="Multi-turn session benchmark for OpenAI-compatible APIs"
|
|
)
|
|
parser.add_argument("--base-url", default="http://127.0.0.1:8262/v1")
|
|
parser.add_argument("--model", default="kimi-k2.6-amd-dflash")
|
|
parser.add_argument("--sessions", type=int, default=4)
|
|
parser.add_argument("--turns-per-session", type=int, default=5)
|
|
parser.add_argument("--max-tokens", type=int, default=512)
|
|
parser.add_argument("--temperature", type=float, default=0)
|
|
parser.add_argument("--output-json", type=str, default=None)
|
|
parser.add_argument("--timeout-seconds", type=float, default=3600)
|
|
args = parser.parse_args()
|
|
|
|
result = asyncio.run(run_benchmark(args))
|
|
|
|
output = json.dumps(result, indent=2)
|
|
print(output)
|
|
|
|
if args.output_json:
|
|
with open(args.output_json, "w") as f:
|
|
f.write(output)
|
|
f.write("\n")
|
|
print(f"\nResults written to {args.output_json}", file=sys.stderr)
|
|
|
|
agg = result["aggregate"]
|
|
print(
|
|
f"\n--- Summary ---\n"
|
|
f"Sessions: {agg['sessions_completed']}/{args.sessions} completed\n"
|
|
f"Turns: {agg['turns_completed']}/{args.sessions * args.turns_per_session}\n"
|
|
f"Aggregate throughput: {agg['aggregate_tok_per_sec']} tok/s\n"
|
|
f"Mean per-session: {agg['mean_per_session_tok_per_sec']} tok/s\n"
|
|
f"Wall time: {agg['total_wall_seconds']}s",
|
|
file=sys.stderr,
|
|
)
|
|
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|