chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,196 @@
|
||||
# Multi-Turn LLM Benchmark
|
||||
|
||||
A benchmark tool for OpenAI-compatible LLM inference servers that supports
|
||||
multi-turn conversations with configurable prefix cache hit rates, input/output
|
||||
sequence lengths, and cross-session prefix sharing.
|
||||
|
||||
## Entry Point
|
||||
|
||||
```
|
||||
python -m ray.llm._internal.serve.benchmark.cli [OPTIONS]
|
||||
```
|
||||
|
||||
## Modes
|
||||
|
||||
| Command | Mode | Description |
|
||||
|---------|------|-------------|
|
||||
| `... -s` | Smoke | Single request health check |
|
||||
| `... --concurrency 8 ...` | Direct (concurrency) | Closed-loop concurrency benchmark |
|
||||
| `... --request-rate 10 ...` | Direct (rate) | Constant-QPS benchmark |
|
||||
| `... -i` | Interactive server | Long-running server with UNIX socket control |
|
||||
| `... -i --client` | Interactive client | Connect to server; REPL or `--cmd` one-shot |
|
||||
|
||||
## Quick Examples
|
||||
|
||||
### Smoke test
|
||||
|
||||
```bash
|
||||
python -m ray.llm._internal.serve.benchmark.cli -s \
|
||||
-u http://localhost:8000 -m my-model
|
||||
```
|
||||
|
||||
### Concurrency benchmark
|
||||
|
||||
```bash
|
||||
python -m ray.llm._internal.serve.benchmark.cli \
|
||||
-u http://localhost:8000 -m meta-llama/Llama-3-8B-Instruct \
|
||||
--concurrency 8 --num-sessions 200 \
|
||||
--isl 2000 --osl 200 --hit-rate 0.85 --num-turns 5 \
|
||||
--think-time 1.0 --save-result results.json
|
||||
```
|
||||
|
||||
### Rate benchmark
|
||||
|
||||
```bash
|
||||
python -m ray.llm._internal.serve.benchmark.cli \
|
||||
-u http://localhost:8000 -m meta-llama/Llama-3-8B-Instruct \
|
||||
--request-rate 10 --duration 120 \
|
||||
--isl 2000 --osl 200 --hit-rate 0.85 --num-turns 5 \
|
||||
--warm-up 10 --save-result results.json
|
||||
```
|
||||
|
||||
### Interactive server
|
||||
|
||||
```bash
|
||||
python -m ray.llm._internal.serve.benchmark.cli -i \
|
||||
-u http://localhost:8000 -m meta-llama/Llama-3-8B-Instruct \
|
||||
--isl 2000 --osl 200 --hit-rate 0.85 --num-turns 5
|
||||
```
|
||||
|
||||
### Interactive client (REPL)
|
||||
|
||||
```bash
|
||||
python -m ray.llm._internal.serve.benchmark.cli -i --client
|
||||
```
|
||||
|
||||
### Interactive client (one-shot)
|
||||
|
||||
```bash
|
||||
python -m ray.llm._internal.serve.benchmark.cli -i --client --cmd "rate 10"
|
||||
python -m ray.llm._internal.serve.benchmark.cli -i --client --cmd "status"
|
||||
```
|
||||
|
||||
## Workload Parameters
|
||||
|
||||
All workload parameters use **simple mode**: you specify user-facing values
|
||||
and the tool derives internal parameters (per-turn user tokens `u` and
|
||||
system prompt tokens `s`) automatically.
|
||||
|
||||
| Parameter | Flag | Description |
|
||||
|-----------|------|-------------|
|
||||
| ISL | `--isl` | Average input sequence length (tokens) across all turns |
|
||||
| OSL | `--osl` | Output tokens per turn |
|
||||
| Hit rate | `--hit-rate` | Target prefix cache hit rate [0, 1] |
|
||||
| Shared system prompt ratio | `--shared-system-prompt-ratio` | Fraction of system prompt shared across sessions (default: 0.0) |
|
||||
| Num turns | `--num-turns` | Number of turns per conversation session |
|
||||
| Think time | `--think-time` | Simulated user think-time between turns in seconds (default: 0) |
|
||||
| First chunk threshold | `--first-chunk-threshold` | Number of SSE content chunks before recording first-chunk latency (default: 16) |
|
||||
|
||||
The solver derives `user_tokens` (new user tokens per turn) and `sys_tokens`
|
||||
(total system prompt tokens) from these inputs. The `print_summary()` output
|
||||
shows the resolved per-turn token breakdown including cached vs. new tokens at
|
||||
each turn.
|
||||
|
||||
## Tokenizer
|
||||
|
||||
By default, `--tokenizer` is `None`, which causes the tool to use the
|
||||
`--model` value as the HuggingFace tokenizer name. This works when `--model`
|
||||
is a valid HuggingFace model ID (e.g., `meta-llama/Llama-3-8B-Instruct`).
|
||||
|
||||
Provide `--tokenizer` explicitly when:
|
||||
- The `--model` value is an alias or deployment name that is not a valid
|
||||
HuggingFace repo (e.g., `--model my-deployment --tokenizer meta-llama/Llama-3-8B-Instruct`).
|
||||
- You want to use a local tokenizer path.
|
||||
|
||||
## Warm-Up Strategies
|
||||
|
||||
### Concurrency mode
|
||||
|
||||
Warm-up is **automatic** using entropy-based detection. The tool monitors the
|
||||
distribution of active turns across concurrent sessions. Once the Shannon
|
||||
entropy of the turn distribution reaches 50% of its theoretical maximum, the
|
||||
pool is considered at steady state and measurement begins. All requests
|
||||
dispatched before that point are discarded.
|
||||
|
||||
### Rate mode
|
||||
|
||||
Warm-up is **time-based** via the `--warm-up` flag (in seconds). All requests
|
||||
whose dispatch time falls within the warm-up window are excluded from reported
|
||||
metrics. Set this to allow the server's KV cache to fill and stabilize.
|
||||
|
||||
### Interactive mode
|
||||
|
||||
Warm-up is **manual**. The operator starts traffic with `rate <qps>`, waits
|
||||
for the system to stabilize, then explicitly starts a measurement window with
|
||||
`start` or `measure <n>`.
|
||||
|
||||
## Interactive Commands
|
||||
|
||||
| Command | Description |
|
||||
|---------|-------------|
|
||||
| `help` | Show available commands |
|
||||
| `rate <qps>` | Set target request rate (0 to pause) |
|
||||
| `start` | Start open-ended measurement window |
|
||||
| `measure <n>` | Start measurement capturing next `n` completed requests |
|
||||
| `stop` | Stop measurement and print summary |
|
||||
| `status` | Show current state: QPS, inflight, completed, measured |
|
||||
| `workload [k=v ...]` | Show or update workload parameters (e.g., `workload isl=3000 osl=300`) |
|
||||
| `save [path]` | Save last measurement window to JSON |
|
||||
| `save-dir <path>` | Set default directory for saved results |
|
||||
| `quit` | Stop the benchmark server |
|
||||
|
||||
## JSON Result Schema
|
||||
|
||||
Results saved with `--save-result` (direct mode) contain these top-level keys:
|
||||
|
||||
| Key | Description |
|
||||
|-----|-------------|
|
||||
| `config` | Run configuration (concurrency/rate, model, etc.) |
|
||||
| `spec` | Resolved workload spec with per-turn token breakdown |
|
||||
| `first_chunk_threshold` | Number of chunks before recording first-chunk latency |
|
||||
| `benchmark` | Run metadata: total requests, duration, warm-up info |
|
||||
| `stats` | Aggregate latency statistics (avg, P50, P90, P99 for TTFT, FC, TPOT, latency) |
|
||||
| `per_turn` | Per-turn breakdown of count, avg ISL, and latency percentiles |
|
||||
| `raw_metrics` | Array of per-request metrics (session_id, turn, all latency fields, token counts) |
|
||||
|
||||
Interactive mode saves with `save` produce a similar structure with a `window`
|
||||
summary instead of `benchmark`/`stats`/`per_turn`.
|
||||
|
||||
## Typical Workflow
|
||||
|
||||
1. **Smoke test** to verify connectivity:
|
||||
```bash
|
||||
python -m ray.llm._internal.serve.benchmark.cli -s -u http://localhost:8000 -m my-model
|
||||
```
|
||||
|
||||
2. **Direct benchmark** for a fixed workload:
|
||||
```bash
|
||||
python -m ray.llm._internal.serve.benchmark.cli \
|
||||
--concurrency 8 --num-sessions 200 \
|
||||
--isl 2000 --osl 200 --hit-rate 0.85 --num-turns 5 \
|
||||
-u http://localhost:8000 -m meta-llama/Llama-3-8B-Instruct \
|
||||
--save-result concurrency_8.json
|
||||
```
|
||||
|
||||
3. **Interactive mode** for exploratory testing:
|
||||
```bash
|
||||
# Terminal 1: start server
|
||||
python -m ray.llm._internal.serve.benchmark.cli -i \
|
||||
--isl 2000 --osl 200 --hit-rate 0.85 --num-turns 5 \
|
||||
-u http://localhost:8000 -m meta-llama/Llama-3-8B-Instruct
|
||||
|
||||
# Terminal 2: control
|
||||
python -m ray.llm._internal.serve.benchmark.cli -i --client
|
||||
benchctl> rate 5
|
||||
benchctl> measure 500
|
||||
benchctl> status
|
||||
benchctl> save results_qps5.json
|
||||
benchctl> rate 10
|
||||
benchctl> measure 500
|
||||
benchctl> save results_qps10.json
|
||||
benchctl> quit
|
||||
```
|
||||
|
||||
4. **Sweep** over multiple configurations: write an external script that loops
|
||||
over the CLI with different parameters. The tool does not include built-in
|
||||
sweep orchestration.
|
||||
@@ -0,0 +1,7 @@
|
||||
"""CLI entry point for the multi-turn OpenAI-compatible HTTP benchmark.
|
||||
|
||||
Example: python -m ray.llm._internal.serve.benchmark --help
|
||||
"""
|
||||
from ray.llm._internal.serve.benchmark.cli import main
|
||||
|
||||
main()
|
||||
@@ -0,0 +1,233 @@
|
||||
"""CLI entry point for the multi-turn OpenAI-compatible HTTP benchmark."""
|
||||
|
||||
import argparse
|
||||
import sys
|
||||
|
||||
|
||||
def build_parser() -> argparse.ArgumentParser:
|
||||
parser = argparse.ArgumentParser(
|
||||
prog="python -m ray.llm._internal.serve.benchmark",
|
||||
description="Multi-turn OpenAI-compatible HTTP benchmark",
|
||||
)
|
||||
|
||||
## Mode flags ##
|
||||
mode = parser.add_argument_group("mode")
|
||||
mode.add_argument(
|
||||
"-s",
|
||||
"--smoke",
|
||||
action="store_true",
|
||||
help="Smoke test (single request, exit)",
|
||||
)
|
||||
mode.add_argument(
|
||||
"-i",
|
||||
"--interactive",
|
||||
action="store_true",
|
||||
help="Interactive mode (server by default)",
|
||||
)
|
||||
mode.add_argument(
|
||||
"--client",
|
||||
action="store_true",
|
||||
help="Interactive client mode (used with -i)",
|
||||
)
|
||||
|
||||
## Server / API ##
|
||||
server = parser.add_argument_group("server/API")
|
||||
server.add_argument(
|
||||
"-u",
|
||||
"--base-url",
|
||||
default="http://127.0.0.1:8000",
|
||||
help="Base URL of the OpenAI-compatible API (default: %(default)s)",
|
||||
)
|
||||
server.add_argument(
|
||||
"-m",
|
||||
"--model",
|
||||
default=None,
|
||||
help="Model name to send in requests (required except for -i --client)",
|
||||
)
|
||||
server.add_argument(
|
||||
"--tokenizer",
|
||||
default=None,
|
||||
help="HuggingFace tokenizer name/path (default: same as --model)",
|
||||
)
|
||||
server.add_argument(
|
||||
"--api-key",
|
||||
default=None,
|
||||
help="API key for Authorization header (default: None)",
|
||||
)
|
||||
|
||||
## Workload ##
|
||||
workload = parser.add_argument_group("workload")
|
||||
workload.add_argument(
|
||||
"--isl",
|
||||
type=int,
|
||||
default=1000,
|
||||
help="Average input sequence length (default: %(default)s)",
|
||||
)
|
||||
workload.add_argument(
|
||||
"--hit-rate",
|
||||
type=float,
|
||||
default=0.5,
|
||||
help="Prefix cache hit rate [0, 1] (default: %(default)s)",
|
||||
)
|
||||
workload.add_argument(
|
||||
"--num-turns",
|
||||
type=int,
|
||||
default=1,
|
||||
help="Number of turns per session (default: %(default)s)",
|
||||
)
|
||||
workload.add_argument(
|
||||
"--osl",
|
||||
type=int,
|
||||
default=100,
|
||||
help="Output tokens per turn (default: %(default)s)",
|
||||
)
|
||||
workload.add_argument(
|
||||
"--shared-system-prompt-ratio",
|
||||
dest="shared_system_prompt_ratio",
|
||||
type=float,
|
||||
default=1.0,
|
||||
help="Fraction of the system prompt shared across all sessions "
|
||||
"(1.0 = identical, 0.0 = all unique) (default: %(default)s)",
|
||||
)
|
||||
workload.add_argument(
|
||||
"--think-time",
|
||||
type=float,
|
||||
default=0.0,
|
||||
help="Simulated user think-time between turns in seconds (default: %(default)s)",
|
||||
)
|
||||
workload.add_argument(
|
||||
"-fc",
|
||||
"--first-chunk-threshold",
|
||||
type=int,
|
||||
default=16,
|
||||
help="Number of content chunks before recording first-chunk latency (default: %(default)s)",
|
||||
)
|
||||
|
||||
## Traffic ##
|
||||
traffic = parser.add_argument_group("traffic")
|
||||
traffic.add_argument(
|
||||
"--concurrency",
|
||||
type=int,
|
||||
default=None,
|
||||
help="Number of concurrent sessions",
|
||||
)
|
||||
traffic.add_argument(
|
||||
"--request-rate",
|
||||
type=float,
|
||||
default=None,
|
||||
help="Request rate (requests per second)",
|
||||
)
|
||||
traffic.add_argument(
|
||||
"--duration",
|
||||
type=float,
|
||||
default=None,
|
||||
help="Duration in seconds",
|
||||
)
|
||||
traffic.add_argument(
|
||||
"--num-sessions",
|
||||
type=int,
|
||||
default=None,
|
||||
help="Total number of sessions to run",
|
||||
)
|
||||
traffic.add_argument(
|
||||
"--warm-up",
|
||||
type=float,
|
||||
default=0,
|
||||
help="Warm-up period in seconds (default: %(default)s)",
|
||||
)
|
||||
traffic.add_argument(
|
||||
"--warmup-jitter-max",
|
||||
type=float,
|
||||
default=10.0,
|
||||
help="Max random delay (seconds) between turns during entropy warm-up "
|
||||
"in concurrency mode. Jitter desynchronizes sessions so the benchmark "
|
||||
"reaches steady-state faster (default: %(default)s)",
|
||||
)
|
||||
traffic.add_argument(
|
||||
"--ramp-interval",
|
||||
type=float,
|
||||
default=-1,
|
||||
help="Seconds between launching successive sessions at benchmark start. "
|
||||
"Use this to avoid a thundering-herd of simultaneous first requests. "
|
||||
"-1 = auto-derive from request rate or concurrency (default: %(default)s)",
|
||||
)
|
||||
|
||||
## Interactive-only ##
|
||||
interactive = parser.add_argument_group("interactive-only")
|
||||
interactive.add_argument(
|
||||
"--status-interval",
|
||||
type=int,
|
||||
default=5,
|
||||
help="Status reporting interval in seconds (default: %(default)s)",
|
||||
)
|
||||
interactive.add_argument(
|
||||
"--cmd",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Command to send in interactive client mode",
|
||||
)
|
||||
interactive.add_argument(
|
||||
"--log-failures",
|
||||
action="store_true",
|
||||
help="Log individual request failures",
|
||||
)
|
||||
interactive.add_argument(
|
||||
"--seed",
|
||||
type=int,
|
||||
default=None,
|
||||
help="Random seed for reproducibility",
|
||||
)
|
||||
interactive.add_argument(
|
||||
"--save-result",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Filename to save results",
|
||||
)
|
||||
interactive.add_argument(
|
||||
"--save-dir",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Directory to save results",
|
||||
)
|
||||
interactive.add_argument(
|
||||
"--num-workers",
|
||||
type=int,
|
||||
default=1,
|
||||
help="Number of process-pool workers for conversation generation (default: %(default)s)",
|
||||
)
|
||||
|
||||
return parser
|
||||
|
||||
|
||||
def main() -> None:
|
||||
parser = build_parser()
|
||||
args = parser.parse_args()
|
||||
|
||||
if args.interactive and args.client:
|
||||
from ray.llm._internal.serve.benchmark.interactive import run_interactive_client
|
||||
|
||||
sys.exit(run_interactive_client(args))
|
||||
|
||||
# All other modes require --model
|
||||
if not args.model:
|
||||
parser.error("--model is required (except for -i --client mode)")
|
||||
|
||||
if args.smoke:
|
||||
from ray.llm._internal.serve.benchmark.runners import run_smoke
|
||||
|
||||
sys.exit(run_smoke(args))
|
||||
elif args.interactive:
|
||||
from ray.llm._internal.serve.benchmark.interactive import run_interactive_server
|
||||
|
||||
sys.exit(run_interactive_server(args))
|
||||
elif args.concurrency or args.request_rate:
|
||||
from ray.llm._internal.serve.benchmark.runners import run_direct
|
||||
|
||||
sys.exit(run_direct(args))
|
||||
else:
|
||||
parser.print_help()
|
||||
sys.exit(1)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,122 @@
|
||||
"""HTTP client for OpenAI-compatible chat completion endpoints."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import time
|
||||
from typing import Optional
|
||||
|
||||
import aiohttp
|
||||
|
||||
from ray.llm._internal.serve.benchmark.models import TurnResult
|
||||
|
||||
|
||||
async def send_chat_completion(
|
||||
session: aiohttp.ClientSession,
|
||||
base_url: str,
|
||||
model: str,
|
||||
messages: list[dict[str, str]],
|
||||
session_id: str = "",
|
||||
max_tokens: int = 256,
|
||||
first_chunk_threshold: int = 16,
|
||||
timeout_sec: int = 300,
|
||||
api_key: Optional[str] = None,
|
||||
) -> TurnResult:
|
||||
"""Send a streaming chat completion request and collect metrics."""
|
||||
url = f"{base_url}/v1/chat/completions"
|
||||
payload = {
|
||||
"model": model,
|
||||
"messages": messages,
|
||||
"max_tokens": max_tokens,
|
||||
"stream": True,
|
||||
"stream_options": {"include_usage": True},
|
||||
"temperature": 0.0,
|
||||
}
|
||||
|
||||
headers: dict[str, str] = {
|
||||
"Content-Type": "application/json",
|
||||
}
|
||||
if api_key:
|
||||
headers["Authorization"] = f"Bearer {api_key}"
|
||||
if session_id:
|
||||
headers["X-Session-Id"] = session_id
|
||||
|
||||
timeout = aiohttp.ClientTimeout(total=timeout_sec)
|
||||
|
||||
start_ns = time.perf_counter_ns()
|
||||
ttft_ns: Optional[int] = None
|
||||
fc_ns: Optional[int] = None
|
||||
content_chunk_count = 0
|
||||
chunk_times: list[int] = []
|
||||
generated_text = ""
|
||||
input_tokens = 0
|
||||
output_tokens = 0
|
||||
prev_ts = start_ns
|
||||
|
||||
async with session.post(
|
||||
url, json=payload, headers=headers, timeout=timeout
|
||||
) as resp:
|
||||
if resp.status != 200:
|
||||
body = await resp.text()
|
||||
raise RuntimeError(f"HTTP {resp.status}: {body[:500]}")
|
||||
|
||||
async for raw_line in resp.content:
|
||||
line = raw_line.strip()
|
||||
if not line:
|
||||
continue
|
||||
text = line.decode("utf-8", errors="replace")
|
||||
if not text.startswith("data: "):
|
||||
continue
|
||||
data_str = text[6:]
|
||||
if data_str == "[DONE]":
|
||||
continue
|
||||
|
||||
try:
|
||||
data = json.loads(data_str)
|
||||
except json.JSONDecodeError:
|
||||
continue
|
||||
|
||||
usage = data.get("usage")
|
||||
if usage:
|
||||
input_tokens = usage.get("prompt_tokens", input_tokens)
|
||||
output_tokens = usage.get("completion_tokens", output_tokens)
|
||||
|
||||
choices = data.get("choices", [])
|
||||
if not choices:
|
||||
continue
|
||||
|
||||
delta = choices[0].get("delta", {})
|
||||
content = delta.get("content") or delta.get("reasoning")
|
||||
if content:
|
||||
now_ns = time.perf_counter_ns()
|
||||
content_chunk_count += 1
|
||||
if ttft_ns is None:
|
||||
ttft_ns = now_ns - start_ns
|
||||
else:
|
||||
chunk_times.append(now_ns - prev_ts)
|
||||
if fc_ns is None and content_chunk_count >= first_chunk_threshold:
|
||||
fc_ns = now_ns - start_ns
|
||||
prev_ts = now_ns
|
||||
generated_text += content
|
||||
|
||||
end_ns = time.perf_counter_ns()
|
||||
latency_ns = end_ns - start_ns
|
||||
|
||||
if ttft_ns is None:
|
||||
ttft_ns = latency_ns
|
||||
if fc_ns is None:
|
||||
fc_ns = latency_ns
|
||||
|
||||
itl_ms_list = [t / 1e6 for t in chunk_times]
|
||||
itl_ms = sum(itl_ms_list) / len(itl_ms_list) if itl_ms_list else 0.0
|
||||
|
||||
return TurnResult(
|
||||
ttft_ms=ttft_ns / 1e6,
|
||||
fc_ms=fc_ns / 1e6,
|
||||
itl_ms=itl_ms,
|
||||
e2e_latency_ms=latency_ns / 1e6,
|
||||
input_tokens=input_tokens,
|
||||
output_tokens=output_tokens,
|
||||
generated_text=generated_text,
|
||||
itl_ms_list=itl_ms_list,
|
||||
)
|
||||
@@ -0,0 +1,824 @@
|
||||
"""Interactive server and client for the multi-turn benchmark.
|
||||
|
||||
The interactive server runs a long-lived benchmark loop whose QPS, workload
|
||||
parameters, and measurement windows are controlled at runtime via a UNIX
|
||||
domain socket. The interactive client connects to that socket (either as an
|
||||
interactive REPL or for one-shot commands).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import asyncio
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import random
|
||||
import time
|
||||
from dataclasses import dataclass, field
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
from typing import Optional
|
||||
|
||||
import aiohttp
|
||||
import numpy as np
|
||||
|
||||
from ray.llm._internal.serve.benchmark.metrics import (
|
||||
serialize_raw_metrics,
|
||||
summarize_metrics,
|
||||
)
|
||||
from ray.llm._internal.serve.benchmark.models import TurnMetric, WorkloadSpec
|
||||
from ray.llm._internal.serve.benchmark.text_gen import (
|
||||
Conversation,
|
||||
TextGenerator,
|
||||
conversation_factory,
|
||||
)
|
||||
from ray.llm._internal.serve.benchmark.turn import execute_single_turn
|
||||
|
||||
try:
|
||||
from prompt_toolkit import PromptSession
|
||||
from prompt_toolkit.history import FileHistory
|
||||
except ImportError:
|
||||
PromptSession = None # type: ignore[assignment,misc]
|
||||
FileHistory = None # type: ignore[assignment,misc]
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Control socket path
|
||||
# ---------------------------------------------------------------------------
|
||||
_DEFAULT_CONTROL_SOCKET = "/tmp/interactive_rate_bench.sock"
|
||||
|
||||
|
||||
def _control_socket_path() -> str:
|
||||
return os.environ.get("RAY_BENCH_CONTROL_SOCKET", _DEFAULT_CONTROL_SOCKET)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Process-pool worker helpers (module-level so they are picklable)
|
||||
# ---------------------------------------------------------------------------
|
||||
_worker_tokenizer = None
|
||||
_worker_text_gen: Optional[TextGenerator] = None
|
||||
|
||||
|
||||
def _pool_initializer(tokenizer_name: str, base_seed: int) -> None:
|
||||
"""Called once per worker process to load the tokenizer and seed RNG."""
|
||||
global _worker_tokenizer, _worker_text_gen
|
||||
from transformers import AutoTokenizer
|
||||
|
||||
_worker_tokenizer = AutoTokenizer.from_pretrained(
|
||||
tokenizer_name, trust_remote_code=True
|
||||
)
|
||||
_worker_text_gen = TextGenerator(_worker_tokenizer)
|
||||
proc_seed = (base_seed + os.getpid()) % (2**32)
|
||||
random.seed(proc_seed)
|
||||
np.random.seed(proc_seed)
|
||||
|
||||
|
||||
def _create_conv_in_worker(
|
||||
session_idx: int,
|
||||
spec: WorkloadSpec,
|
||||
shared_system_text: str,
|
||||
) -> Conversation:
|
||||
"""Create a Conversation inside a worker process."""
|
||||
return conversation_factory(session_idx, spec, shared_system_text, _worker_text_gen)
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Interactive-mode runtime state & helpers
|
||||
# ============================================================================
|
||||
|
||||
|
||||
@dataclass
|
||||
class RuntimeState:
|
||||
current_qps: float = 0.0
|
||||
total_completed: int = 0
|
||||
total_failed: int = 0
|
||||
inflight: int = 0
|
||||
measurement_active: bool = False
|
||||
measurement_start_ns: Optional[int] = None
|
||||
measurement_metrics: list[TurnMetric] = field(default_factory=list)
|
||||
measurement_target_requests: Optional[int] = None
|
||||
last_window_metrics: list[TurnMetric] = field(default_factory=list)
|
||||
last_window_elapsed_s: float = 0.0
|
||||
last_notice: Optional[str] = None
|
||||
save_dir: Optional[str] = None
|
||||
|
||||
|
||||
def _save_window_result(
|
||||
path: str,
|
||||
args: argparse.Namespace,
|
||||
spec: WorkloadSpec,
|
||||
metrics: list[TurnMetric],
|
||||
elapsed_s: float,
|
||||
runtime_qps: float = 0.0,
|
||||
) -> None:
|
||||
payload = {
|
||||
"mode": "interactive_rate",
|
||||
"saved_at_epoch_s": time.time(),
|
||||
"config": {
|
||||
"base_url": args.base_url,
|
||||
"model": args.model,
|
||||
"tokenizer": getattr(args, "tokenizer", None) or args.model,
|
||||
"first_chunk_threshold": args.first_chunk_threshold,
|
||||
"num_turns": args.num_turns,
|
||||
"osl": args.osl,
|
||||
"shared_system_prompt_ratio": args.shared_system_prompt_ratio,
|
||||
"isl": args.isl,
|
||||
"hit_rate": args.hit_rate,
|
||||
"runtime_qps": runtime_qps,
|
||||
},
|
||||
"spec": spec.summary(),
|
||||
"window": summarize_metrics(metrics, elapsed_s),
|
||||
"raw_metrics": serialize_raw_metrics(metrics),
|
||||
}
|
||||
p = Path(path)
|
||||
p.parent.mkdir(parents=True, exist_ok=True)
|
||||
with p.open("w") as f:
|
||||
json.dump(payload, f, indent=2)
|
||||
print(f"Saved measurement window to {path}")
|
||||
|
||||
|
||||
def _build_spec(
|
||||
args: argparse.Namespace, overrides: Optional[dict] = None
|
||||
) -> WorkloadSpec:
|
||||
"""Build and resolve a WorkloadSpec from args, optionally merging overrides."""
|
||||
kw = dict(
|
||||
num_sessions=1,
|
||||
duration_s=1.0,
|
||||
num_turns=args.num_turns,
|
||||
osl=args.osl,
|
||||
think_time=0.0,
|
||||
concurrency=None,
|
||||
request_rate=1.0,
|
||||
ramp_interval=0.0,
|
||||
shared_system_prompt_ratio=args.shared_system_prompt_ratio,
|
||||
isl=args.isl,
|
||||
hit_rate=args.hit_rate,
|
||||
)
|
||||
if overrides:
|
||||
kw.update(overrides)
|
||||
spec = WorkloadSpec(**kw)
|
||||
spec.resolve()
|
||||
return spec
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Command handler (extracted for testability)
|
||||
# ============================================================================
|
||||
|
||||
|
||||
class CommandHandler:
|
||||
"""Handles interactive benchmark commands.
|
||||
|
||||
Extracted from the ``run_interactive`` closure so that command parsing,
|
||||
state mutation, and response formatting can be unit-tested without
|
||||
starting a real server or HTTP session.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
runtime: RuntimeState,
|
||||
workload: dict,
|
||||
args: argparse.Namespace,
|
||||
text_gen: Optional[TextGenerator] = None,
|
||||
rate_changed: Optional[asyncio.Event] = None,
|
||||
workload_changed: Optional[asyncio.Event] = None,
|
||||
stop_event: Optional[asyncio.Event] = None,
|
||||
):
|
||||
self.runtime = runtime
|
||||
self.workload = workload
|
||||
self.args = args
|
||||
self.text_gen = text_gen
|
||||
self.rate_changed = rate_changed or asyncio.Event()
|
||||
self.workload_changed = workload_changed or asyncio.Event()
|
||||
self.stop_event = stop_event or asyncio.Event()
|
||||
|
||||
def resolve_save_path(self, raw: Optional[str]) -> str:
|
||||
if raw:
|
||||
expanded = str(Path(raw).expanduser())
|
||||
if "/" in expanded or expanded.startswith("."):
|
||||
return expanded
|
||||
return str(Path(self.runtime.save_dir) / expanded)
|
||||
|
||||
if self.args.save_result:
|
||||
return str(Path(self.args.save_result).expanduser())
|
||||
|
||||
ts = datetime.now().strftime("%Y%m%d_%H%M%S")
|
||||
qps_label = f"{self.runtime.current_qps:.2f}".replace(".", "p")
|
||||
return str(
|
||||
Path(self.runtime.save_dir)
|
||||
/ f"interactive_measure_qps{qps_label}_{ts}.json"
|
||||
)
|
||||
|
||||
async def handle(self, cmd: str) -> str: # noqa: C901
|
||||
"""Process a single command string and return the response."""
|
||||
cmd = cmd.strip()
|
||||
if not cmd:
|
||||
return "empty command"
|
||||
|
||||
parts = cmd.split()
|
||||
op = parts[0].lower()
|
||||
|
||||
if op == "help":
|
||||
return (
|
||||
"Commands: help, rate <qps>, start, measure <n>, stop, "
|
||||
"status, save [path|name], save-dir <path>, quit\n"
|
||||
"Workload: workload [isl=N] [osl=N] [hit-rate=F] "
|
||||
"[sharing=F] [num-turns=N]\n"
|
||||
" e.g. workload isl=2000 osl=200 hit-rate=0.5\n"
|
||||
" All params optional; unspecified ones keep their current values.\n"
|
||||
" workload (no args) prints current workload spec."
|
||||
)
|
||||
if op == "rate":
|
||||
if len(parts) != 2:
|
||||
return "Usage: rate <qps>"
|
||||
try:
|
||||
new_qps = float(parts[1])
|
||||
if new_qps < 0:
|
||||
raise ValueError()
|
||||
except ValueError:
|
||||
return "QPS must be a non-negative number."
|
||||
self.runtime.current_qps = new_qps
|
||||
self.rate_changed.set()
|
||||
return f"Set target qps={new_qps:.3f}"
|
||||
if op == "start":
|
||||
self.runtime.measurement_active = True
|
||||
self.runtime.measurement_start_ns = time.perf_counter_ns()
|
||||
self.runtime.measurement_metrics = []
|
||||
self.runtime.measurement_target_requests = None
|
||||
self.runtime.last_notice = None
|
||||
return "Measurement started."
|
||||
if op == "measure":
|
||||
if len(parts) != 2:
|
||||
return "Usage: measure <num_requests>"
|
||||
try:
|
||||
tgt = int(parts[1])
|
||||
if tgt <= 0:
|
||||
raise ValueError()
|
||||
except ValueError:
|
||||
return "measure requires a positive integer."
|
||||
self.runtime.measurement_active = True
|
||||
self.runtime.measurement_start_ns = time.perf_counter_ns()
|
||||
self.runtime.measurement_metrics = []
|
||||
self.runtime.measurement_target_requests = tgt
|
||||
self.runtime.last_notice = None
|
||||
return f"Measurement started: capturing next {tgt} completed requests."
|
||||
if op == "stop":
|
||||
if not self.runtime.measurement_active:
|
||||
return "Measurement is not active."
|
||||
self.runtime.measurement_active = False
|
||||
end_ns = time.perf_counter_ns()
|
||||
start_ns = self.runtime.measurement_start_ns or end_ns
|
||||
self.runtime.last_window_elapsed_s = (end_ns - start_ns) / 1e9
|
||||
self.runtime.last_window_metrics = list(self.runtime.measurement_metrics)
|
||||
self.runtime.measurement_target_requests = None
|
||||
summary = summarize_metrics(
|
||||
list(self.runtime.last_window_metrics),
|
||||
self.runtime.last_window_elapsed_s,
|
||||
)
|
||||
return f"Measurement stopped.\n{json.dumps(summary, indent=2)}"
|
||||
if op == "status":
|
||||
cur = self.workload["spec"]
|
||||
status = (
|
||||
f"qps={self.runtime.current_qps:.2f} "
|
||||
f"inflight={self.runtime.inflight} "
|
||||
f"completed={self.runtime.total_completed} "
|
||||
f"failed={self.runtime.total_failed} "
|
||||
f"measured={len(self.runtime.measurement_metrics)} "
|
||||
f"active={self.runtime.measurement_active} "
|
||||
f"target={self.runtime.measurement_target_requests} "
|
||||
f"save_dir={self.runtime.save_dir}\n"
|
||||
f"workload: isl={cur.isl} osl={cur.osl} hit-rate={cur.hit_rate} "
|
||||
f"sharing={cur.shared_system_prompt_ratio} num-turns={cur.num_turns}"
|
||||
)
|
||||
if self.runtime.last_notice:
|
||||
status += f"\n{self.runtime.last_notice}"
|
||||
self.runtime.last_notice = None
|
||||
return status
|
||||
if op == "save-dir":
|
||||
if len(parts) != 2:
|
||||
return "Usage: save-dir <path>"
|
||||
new_dir = str(Path(parts[1]).expanduser())
|
||||
self.runtime.save_dir = new_dir
|
||||
return f"Set save_dir={self.runtime.save_dir}"
|
||||
if op == "save":
|
||||
if len(parts) > 2:
|
||||
return "Usage: save [path.json|name.json]"
|
||||
|
||||
if (
|
||||
self.runtime.measurement_active
|
||||
and self.runtime.measurement_start_ns is not None
|
||||
):
|
||||
el = (time.perf_counter_ns() - self.runtime.measurement_start_ns) / 1e9
|
||||
mlist = list(self.runtime.measurement_metrics)
|
||||
else:
|
||||
el = self.runtime.last_window_elapsed_s
|
||||
mlist = list(self.runtime.last_window_metrics)
|
||||
if not mlist:
|
||||
return "No measured window data to save."
|
||||
save_path = self.resolve_save_path(parts[1] if len(parts) == 2 else None)
|
||||
_save_window_result(
|
||||
save_path,
|
||||
self.args,
|
||||
self.workload["spec"],
|
||||
mlist,
|
||||
el,
|
||||
runtime_qps=self.runtime.current_qps,
|
||||
)
|
||||
return f"Saved measurement window to {save_path}"
|
||||
if op == "workload":
|
||||
cur = self.workload["spec"]
|
||||
if len(parts) == 1:
|
||||
return (
|
||||
f"isl={cur.isl} osl={cur.osl} hit-rate={cur.hit_rate} "
|
||||
f"sharing={cur.shared_system_prompt_ratio} num-turns={cur.num_turns}"
|
||||
)
|
||||
_param_aliases = {
|
||||
"isl": "isl",
|
||||
"osl": "osl",
|
||||
"hit-rate": "hit_rate",
|
||||
"hitrate": "hit_rate",
|
||||
"hit_rate": "hit_rate",
|
||||
"sharing": "shared_system_prompt_ratio",
|
||||
"shared-system-prompt-ratio": "shared_system_prompt_ratio",
|
||||
"shared_system_prompt_ratio": "shared_system_prompt_ratio",
|
||||
"num-turns": "num_turns",
|
||||
"num_turns": "num_turns",
|
||||
}
|
||||
overrides: dict = {}
|
||||
errors: list[str] = []
|
||||
for token in parts[1:]:
|
||||
if "=" not in token:
|
||||
errors.append(f"bad token {token!r} (expected key=value)")
|
||||
continue
|
||||
k, _, v = token.partition("=")
|
||||
mapped = _param_aliases.get(k.lower())
|
||||
if mapped is None:
|
||||
errors.append(f"unknown param {k!r}")
|
||||
continue
|
||||
try:
|
||||
overrides[mapped] = (
|
||||
int(v) if mapped in ("isl", "osl", "num_turns") else float(v)
|
||||
)
|
||||
except ValueError:
|
||||
errors.append(f"invalid value for {k}: {v!r}")
|
||||
if errors:
|
||||
return "Error: " + "; ".join(errors)
|
||||
merged = dict(
|
||||
isl=cur.isl,
|
||||
osl=cur.osl,
|
||||
hit_rate=cur.hit_rate,
|
||||
shared_system_prompt_ratio=cur.shared_system_prompt_ratio,
|
||||
num_turns=cur.num_turns,
|
||||
)
|
||||
merged.update(overrides)
|
||||
try:
|
||||
new_spec = _build_spec(self.args, merged)
|
||||
except Exception as e:
|
||||
return f"Invalid workload spec: {e}"
|
||||
if self.text_gen is not None:
|
||||
new_sst = self.text_gen.generate(new_spec.shared_s)
|
||||
else:
|
||||
new_sst = ""
|
||||
self.workload["spec"] = new_spec
|
||||
self.workload["shared_system_text"] = new_sst
|
||||
self.workload_changed.set()
|
||||
new_spec.print_summary()
|
||||
return (
|
||||
f"Workload updated: isl={new_spec.isl} osl={new_spec.osl} "
|
||||
f"hit-rate={new_spec.hit_rate} "
|
||||
f"sharing={new_spec.shared_system_prompt_ratio} "
|
||||
f"num-turns={new_spec.num_turns}"
|
||||
)
|
||||
if op in ("quit", "exit"):
|
||||
self.stop_event.set()
|
||||
return "Stopping benchmark..."
|
||||
return f"Unknown command: {op}"
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Interactive server
|
||||
# ============================================================================
|
||||
|
||||
|
||||
async def run_interactive(args: argparse.Namespace) -> None:
|
||||
spec = _build_spec(args)
|
||||
spec.print_summary()
|
||||
print("Interactive mode: starts idle. Use 'rate <qps>' to begin sending traffic.")
|
||||
|
||||
from concurrent.futures import ProcessPoolExecutor
|
||||
|
||||
from transformers import AutoTokenizer
|
||||
|
||||
tokenizer_name: str = args.tokenizer if args.tokenizer else args.model
|
||||
|
||||
if args.seed is None:
|
||||
args.seed = random.randint(0, 2**31 - 1)
|
||||
random.seed(args.seed)
|
||||
np.random.seed(args.seed % (2**32))
|
||||
print(f"Seed: {args.seed}")
|
||||
|
||||
print(f"Loading tokenizer: {tokenizer_name}")
|
||||
tokenizer = AutoTokenizer.from_pretrained(tokenizer_name, trust_remote_code=True)
|
||||
text_gen = TextGenerator(tokenizer)
|
||||
|
||||
shared_system_text = text_gen.generate(spec.shared_s)
|
||||
bench_start_ns = time.perf_counter_ns()
|
||||
|
||||
workload: dict = {"spec": spec, "shared_system_text": shared_system_text}
|
||||
workload_changed = asyncio.Event()
|
||||
|
||||
default_save_dir = args.save_dir
|
||||
if default_save_dir is None and args.save_result:
|
||||
default_save_dir = str(Path(args.save_result).parent)
|
||||
if default_save_dir is None:
|
||||
default_save_dir = os.getcwd()
|
||||
|
||||
runtime = RuntimeState(
|
||||
current_qps=0.0,
|
||||
save_dir=str(Path(default_save_dir).expanduser()),
|
||||
)
|
||||
stop_event = asyncio.Event()
|
||||
rate_changed = asyncio.Event()
|
||||
ready_queue: asyncio.Queue[tuple[Conversation, int]] = asyncio.Queue()
|
||||
next_session_idx = 0
|
||||
running_tasks: set[asyncio.Task] = set()
|
||||
|
||||
num_workers = args.num_workers
|
||||
print(f"Starting process pool with {num_workers} workers")
|
||||
cpu_pool = ProcessPoolExecutor(
|
||||
max_workers=num_workers,
|
||||
initializer=_pool_initializer,
|
||||
initargs=(tokenizer_name, args.seed),
|
||||
)
|
||||
loop = asyncio.get_running_loop()
|
||||
|
||||
def _next_session_idx() -> int:
|
||||
nonlocal next_session_idx
|
||||
idx = next_session_idx
|
||||
next_session_idx += 1
|
||||
return idx
|
||||
|
||||
async def next_conv_async() -> Conversation:
|
||||
idx = _next_session_idx()
|
||||
s = workload["spec"]
|
||||
sst = workload["shared_system_text"]
|
||||
return await loop.run_in_executor(
|
||||
cpu_pool,
|
||||
_create_conv_in_worker,
|
||||
idx,
|
||||
s,
|
||||
sst,
|
||||
)
|
||||
|
||||
async def prefill_queue() -> None:
|
||||
while not stop_event.is_set():
|
||||
if workload_changed.is_set():
|
||||
workload_changed.clear()
|
||||
drained = 0
|
||||
while not ready_queue.empty():
|
||||
try:
|
||||
ready_queue.get_nowait()
|
||||
drained += 1
|
||||
except asyncio.QueueEmpty:
|
||||
break
|
||||
if drained:
|
||||
print(
|
||||
f"[workload] drained {drained} stale conversations from queue.",
|
||||
flush=True,
|
||||
)
|
||||
|
||||
qps = runtime.current_qps
|
||||
if qps <= 0:
|
||||
await asyncio.sleep(0.2)
|
||||
continue
|
||||
s = workload["spec"]
|
||||
sst = workload["shared_system_text"]
|
||||
target = max(8, int(qps * 2))
|
||||
current = ready_queue.qsize()
|
||||
if current < target:
|
||||
batch_size = min(target - current, num_workers * 2)
|
||||
idxs = [_next_session_idx() for _ in range(batch_size)]
|
||||
futs = [
|
||||
loop.run_in_executor(
|
||||
cpu_pool,
|
||||
_create_conv_in_worker,
|
||||
idx,
|
||||
s,
|
||||
sst,
|
||||
)
|
||||
for idx in idxs
|
||||
]
|
||||
for fut in asyncio.as_completed(futs):
|
||||
try:
|
||||
conv = await fut
|
||||
await ready_queue.put((conv, 0))
|
||||
except Exception as e:
|
||||
logger.warning("Failed to create conversation in worker: %s", e)
|
||||
await asyncio.sleep(0.02)
|
||||
|
||||
async def execute_turn(
|
||||
conv: Conversation, turn_idx: int, http_session: aiohttp.ClientSession
|
||||
) -> None:
|
||||
cur_spec = workload["spec"]
|
||||
runtime.inflight += 1
|
||||
try:
|
||||
outcome = await execute_single_turn(
|
||||
http_session=http_session,
|
||||
conv=conv,
|
||||
turn_idx=turn_idx,
|
||||
base_url=args.base_url,
|
||||
model=args.model,
|
||||
max_tokens=cur_spec.osl,
|
||||
bench_start_ns=bench_start_ns,
|
||||
first_chunk_threshold=args.first_chunk_threshold,
|
||||
api_key=getattr(args, "api_key", None),
|
||||
)
|
||||
metric = outcome.metric
|
||||
auto_complete_summary: Optional[str] = None
|
||||
runtime.total_completed += 1
|
||||
if runtime.measurement_active:
|
||||
target = runtime.measurement_target_requests
|
||||
if target is None:
|
||||
runtime.measurement_metrics.append(metric)
|
||||
elif len(runtime.measurement_metrics) < target:
|
||||
runtime.measurement_metrics.append(metric)
|
||||
|
||||
if target is not None and len(runtime.measurement_metrics) >= target:
|
||||
runtime.measurement_active = False
|
||||
end_ns = time.perf_counter_ns()
|
||||
start_ns = runtime.measurement_start_ns or end_ns
|
||||
runtime.last_window_elapsed_s = (end_ns - start_ns) / 1e9
|
||||
runtime.last_window_metrics = list(
|
||||
runtime.measurement_metrics[:target]
|
||||
)
|
||||
runtime.measurement_target_requests = None
|
||||
summary = summarize_metrics(
|
||||
runtime.last_window_metrics,
|
||||
runtime.last_window_elapsed_s,
|
||||
)
|
||||
auto_complete_summary = json.dumps(summary, indent=2)
|
||||
runtime.last_notice = (
|
||||
f"measurement auto-complete ({target} req):\n"
|
||||
f"{auto_complete_summary}"
|
||||
)
|
||||
|
||||
if auto_complete_summary is not None:
|
||||
print("Measurement auto-complete:")
|
||||
print(auto_complete_summary)
|
||||
|
||||
next_turn = turn_idx + 1
|
||||
if not stop_event.is_set():
|
||||
if next_turn < cur_spec.num_turns:
|
||||
await ready_queue.put((conv, next_turn))
|
||||
else:
|
||||
conv = await next_conv_async()
|
||||
await ready_queue.put((conv, 0))
|
||||
except Exception as e:
|
||||
if args.log_failures:
|
||||
print(
|
||||
f"[request-failed] session={conv.session_id} turn={turn_idx}: {e}"
|
||||
)
|
||||
runtime.total_failed += 1
|
||||
if not stop_event.is_set():
|
||||
conv = await next_conv_async()
|
||||
await ready_queue.put((conv, 0))
|
||||
finally:
|
||||
runtime.inflight -= 1
|
||||
|
||||
async def pacer(http_session: aiohttp.ClientSession) -> None:
|
||||
next_dispatch = time.perf_counter()
|
||||
while not stop_event.is_set():
|
||||
qps = runtime.current_qps
|
||||
|
||||
if qps <= 0:
|
||||
await asyncio.sleep(0.1)
|
||||
next_dispatch = time.perf_counter() + 0.05
|
||||
continue
|
||||
|
||||
if rate_changed.is_set():
|
||||
rate_changed.clear()
|
||||
next_dispatch = time.perf_counter() + (1.0 / qps)
|
||||
|
||||
now = time.perf_counter()
|
||||
if next_dispatch < now - 1.0:
|
||||
next_dispatch = now
|
||||
|
||||
wait = next_dispatch - now
|
||||
if wait > 0:
|
||||
await asyncio.sleep(wait)
|
||||
if stop_event.is_set():
|
||||
break
|
||||
|
||||
try:
|
||||
conv, turn_idx = ready_queue.get_nowait()
|
||||
except asyncio.QueueEmpty:
|
||||
next_dispatch += 1.0 / qps
|
||||
continue
|
||||
|
||||
t = asyncio.create_task(execute_turn(conv, turn_idx, http_session))
|
||||
running_tasks.add(t)
|
||||
t.add_done_callback(running_tasks.discard)
|
||||
next_dispatch += 1.0 / qps
|
||||
|
||||
async def reporter() -> None:
|
||||
if args.status_interval <= 0:
|
||||
return
|
||||
while not stop_event.is_set():
|
||||
await asyncio.sleep(args.status_interval)
|
||||
print(
|
||||
"status: "
|
||||
f"qps={runtime.current_qps:.2f} "
|
||||
f"inflight={runtime.inflight} "
|
||||
f"completed={runtime.total_completed} "
|
||||
f"failed={runtime.total_failed} "
|
||||
f"measured={len(runtime.measurement_metrics)} "
|
||||
f"active={runtime.measurement_active}",
|
||||
flush=True,
|
||||
)
|
||||
|
||||
cmd_handler = CommandHandler(
|
||||
runtime=runtime,
|
||||
workload=workload,
|
||||
args=args,
|
||||
text_gen=text_gen,
|
||||
rate_changed=rate_changed,
|
||||
workload_changed=workload_changed,
|
||||
stop_event=stop_event,
|
||||
)
|
||||
handle_command = cmd_handler.handle
|
||||
|
||||
async def stdin_command_loop() -> None:
|
||||
print(
|
||||
"Interactive commands: help | rate <qps> | start | measure <n> | "
|
||||
"stop | status | workload [k=v ...] | save [path|name] | "
|
||||
"save-dir <path> | quit"
|
||||
)
|
||||
while not stop_event.is_set():
|
||||
raw = await asyncio.to_thread(input, "bench> ")
|
||||
resp = await handle_command(raw)
|
||||
if resp:
|
||||
print(resp)
|
||||
|
||||
async def socket_command_handler(
|
||||
reader: asyncio.StreamReader, writer: asyncio.StreamWriter
|
||||
) -> None:
|
||||
try:
|
||||
data = await reader.read(4096)
|
||||
cmd = data.decode("utf-8", errors="replace").strip()
|
||||
resp = await handle_command(cmd)
|
||||
writer.write((resp + "\n").encode("utf-8"))
|
||||
await writer.drain()
|
||||
finally:
|
||||
writer.close()
|
||||
await writer.wait_closed()
|
||||
|
||||
# Seed the ready queue
|
||||
seed_count = 2
|
||||
seed_idxs = [_next_session_idx() for _ in range(seed_count)]
|
||||
seed_futs = [
|
||||
loop.run_in_executor(
|
||||
cpu_pool,
|
||||
_create_conv_in_worker,
|
||||
idx,
|
||||
spec,
|
||||
shared_system_text,
|
||||
)
|
||||
for idx in seed_idxs
|
||||
]
|
||||
for conv in await asyncio.gather(*seed_futs):
|
||||
await ready_queue.put((conv, 0))
|
||||
|
||||
control_socket = _control_socket_path()
|
||||
if os.path.exists(control_socket):
|
||||
os.unlink(control_socket)
|
||||
socket_server = await asyncio.start_unix_server(
|
||||
socket_command_handler,
|
||||
path=control_socket,
|
||||
)
|
||||
print(f"Control socket listening at: {control_socket}")
|
||||
|
||||
stdin_control = getattr(args, "stdin_control", False)
|
||||
if not stdin_control:
|
||||
print("Use a second terminal with --client to send commands.")
|
||||
|
||||
connector = aiohttp.TCPConnector(limit=0)
|
||||
async with aiohttp.ClientSession(connector=connector) as http_session:
|
||||
background_tasks = [
|
||||
asyncio.create_task(pacer(http_session)),
|
||||
asyncio.create_task(reporter()),
|
||||
asyncio.create_task(socket_server.serve_forever()),
|
||||
asyncio.create_task(prefill_queue()),
|
||||
]
|
||||
stdin_task = (
|
||||
asyncio.create_task(stdin_command_loop()) if stdin_control else None
|
||||
)
|
||||
|
||||
if stdin_task is not None:
|
||||
await stdin_task
|
||||
else:
|
||||
await stop_event.wait()
|
||||
|
||||
stop_event.set()
|
||||
socket_server.close()
|
||||
await socket_server.wait_closed()
|
||||
await asyncio.gather(*background_tasks, return_exceptions=True)
|
||||
|
||||
if running_tasks:
|
||||
print(f"Waiting for {len(running_tasks)} in-flight request task(s)...")
|
||||
await asyncio.gather(*list(running_tasks), return_exceptions=True)
|
||||
|
||||
cpu_pool.shutdown(wait=False)
|
||||
if os.path.exists(control_socket):
|
||||
os.unlink(control_socket)
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Interactive client
|
||||
# ============================================================================
|
||||
|
||||
|
||||
async def _send_command_once(control_socket: str, cmd: str) -> str:
|
||||
reader, writer = await asyncio.open_unix_connection(control_socket)
|
||||
writer.write(cmd.encode("utf-8"))
|
||||
await writer.drain()
|
||||
if writer.can_write_eof():
|
||||
writer.write_eof()
|
||||
data = await reader.read()
|
||||
writer.close()
|
||||
await writer.wait_closed()
|
||||
return data.decode("utf-8", errors="replace").strip()
|
||||
|
||||
|
||||
async def run_client(args: argparse.Namespace) -> None:
|
||||
control_socket = _control_socket_path()
|
||||
print(f"Connected to control socket: {control_socket}")
|
||||
print(
|
||||
"Type commands: help, rate <qps>, start, measure <n>, stop, status, "
|
||||
"save [path|name], save-dir <path>, quit"
|
||||
)
|
||||
session = None
|
||||
if PromptSession is not None and FileHistory is not None:
|
||||
history_path = str(Path("~/.interactive_rate_bench_history").expanduser())
|
||||
session = PromptSession(history=FileHistory(history_path))
|
||||
else:
|
||||
print(
|
||||
"prompt_toolkit not installed; using basic input(). "
|
||||
"Install with: pip install prompt_toolkit"
|
||||
)
|
||||
while True:
|
||||
if session is not None:
|
||||
raw = await asyncio.to_thread(session.prompt, "benchctl> ")
|
||||
else:
|
||||
raw = await asyncio.to_thread(input, "benchctl> ")
|
||||
cmd = raw.strip()
|
||||
if not cmd:
|
||||
continue
|
||||
try:
|
||||
resp = await _send_command_once(control_socket, cmd)
|
||||
except (FileNotFoundError, ConnectionRefusedError) as e:
|
||||
print(f"Failed to connect to server socket: {e}")
|
||||
return
|
||||
print(resp)
|
||||
if cmd.lower() in ("quit", "exit"):
|
||||
return
|
||||
|
||||
|
||||
async def run_client_oneshot(args: argparse.Namespace) -> None:
|
||||
if not args.cmd:
|
||||
raise ValueError("Client mode with --cmd requires a command string.")
|
||||
control_socket = _control_socket_path()
|
||||
try:
|
||||
resp = await _send_command_once(control_socket, args.cmd)
|
||||
except (FileNotFoundError, ConnectionRefusedError) as e:
|
||||
raise RuntimeError(f"Failed to connect to server socket: {e}") from e
|
||||
print(resp)
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Entry points for cli.py
|
||||
# ============================================================================
|
||||
|
||||
|
||||
def run_interactive_server(args: argparse.Namespace) -> int:
|
||||
"""Entry point for interactive server mode."""
|
||||
try:
|
||||
asyncio.run(run_interactive(args))
|
||||
return 0
|
||||
except Exception as e:
|
||||
logger.error("Interactive server failed: %s", e)
|
||||
return 1
|
||||
|
||||
|
||||
def run_interactive_client(args: argparse.Namespace) -> int:
|
||||
"""Entry point for interactive client mode."""
|
||||
try:
|
||||
if args.cmd:
|
||||
asyncio.run(run_client_oneshot(args))
|
||||
else:
|
||||
asyncio.run(run_client(args))
|
||||
return 0
|
||||
except Exception as e:
|
||||
logger.error("Interactive client failed: %s", e)
|
||||
return 1
|
||||
@@ -0,0 +1,81 @@
|
||||
"""Metrics computation and serialization for the benchmark."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from statistics import mean
|
||||
|
||||
import numpy as np
|
||||
|
||||
from ray.llm._internal.serve.benchmark.models import TurnMetric
|
||||
|
||||
|
||||
def percentile(values: list[float], p: float) -> float:
|
||||
"""Compute the p-th percentile (0-100)."""
|
||||
if not values:
|
||||
return 0.0
|
||||
return float(np.percentile(values, p))
|
||||
|
||||
|
||||
def summarize_metrics(metrics: list[TurnMetric], elapsed_s: float) -> dict:
|
||||
"""Compute aggregate statistics from a list of TurnMetrics.
|
||||
|
||||
ITL (inter-token latency) statistics are computed from raw per-token values
|
||||
flattened across all requests, capturing the full distribution including variance.
|
||||
"""
|
||||
if not metrics:
|
||||
return {"requests": 0, "elapsed_s": round(elapsed_s, 2)}
|
||||
|
||||
ttft = [m.ttft_ms for m in metrics]
|
||||
fc = [m.fc_ms for m in metrics]
|
||||
# Flatten per-token ITL values across all requests for accurate distribution stats
|
||||
itl_all = [v for m in metrics for v in m.itl_ms_list]
|
||||
latency = [m.e2e_latency_ms for m in metrics]
|
||||
out_tok = [m.output_tokens for m in metrics]
|
||||
in_tok = [m.input_tokens for m in metrics]
|
||||
total_output_tokens = sum(out_tok)
|
||||
|
||||
return {
|
||||
"requests": len(metrics),
|
||||
"elapsed_s": round(elapsed_s, 2),
|
||||
"request_rate": round(len(metrics) / elapsed_s, 2) if elapsed_s > 0 else 0.0,
|
||||
"throughput_tok_s": round(total_output_tokens / elapsed_s, 1)
|
||||
if elapsed_s > 0
|
||||
else 0.0,
|
||||
"avg_input_tokens": round(mean(in_tok), 1),
|
||||
"avg_output_tokens": round(mean(out_tok), 1),
|
||||
"avg_ttft_ms": round(mean(ttft), 2),
|
||||
"p50_ttft_ms": round(percentile(ttft, 50), 2),
|
||||
"p90_ttft_ms": round(percentile(ttft, 90), 2),
|
||||
"p99_ttft_ms": round(percentile(ttft, 99), 2),
|
||||
"avg_fc_ms": round(mean(fc), 2),
|
||||
"p50_fc_ms": round(percentile(fc, 50), 2),
|
||||
"p90_fc_ms": round(percentile(fc, 90), 2),
|
||||
"p99_fc_ms": round(percentile(fc, 99), 2),
|
||||
"avg_itl_ms": round(float(np.mean(itl_all)), 2) if itl_all else 0.0,
|
||||
"std_itl_ms": round(float(np.std(itl_all)), 2) if itl_all else 0.0,
|
||||
"p50_itl_ms": round(percentile(itl_all, 50), 2) if itl_all else 0.0,
|
||||
"p90_itl_ms": round(percentile(itl_all, 90), 2) if itl_all else 0.0,
|
||||
"p99_itl_ms": round(percentile(itl_all, 99), 2) if itl_all else 0.0,
|
||||
"avg_e2e_latency_ms": round(mean(latency), 2),
|
||||
"p50_e2e_latency_ms": round(percentile(latency, 50), 2),
|
||||
"p90_e2e_latency_ms": round(percentile(latency, 90), 2),
|
||||
"p99_e2e_latency_ms": round(percentile(latency, 99), 2),
|
||||
}
|
||||
|
||||
|
||||
def serialize_raw_metrics(metrics: list[TurnMetric]) -> list[dict]:
|
||||
"""Serialize TurnMetrics to dicts suitable for JSON output."""
|
||||
return [
|
||||
{
|
||||
"session_id": m.session_id,
|
||||
"turn": m.turn,
|
||||
"ttft_ms": round(m.ttft_ms, 2),
|
||||
"fc_ms": round(m.fc_ms, 2),
|
||||
"itl_ms": round(m.itl_ms, 2),
|
||||
"e2e_latency_ms": round(m.e2e_latency_ms, 2),
|
||||
"input_tokens": m.input_tokens,
|
||||
"output_tokens": m.output_tokens,
|
||||
"start_time_ms": round(m.start_time_ms, 2),
|
||||
}
|
||||
for m in metrics
|
||||
]
|
||||
@@ -0,0 +1,397 @@
|
||||
"""Data models for the multi-turn benchmark."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
from dataclasses import dataclass, field
|
||||
from typing import List, Optional
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@dataclass
|
||||
class TurnResult:
|
||||
"""Result of a single turn's HTTP request."""
|
||||
|
||||
ttft_ms: float # time to first token
|
||||
fc_ms: float # first-chunk latency (time to N-th content chunk)
|
||||
itl_ms: float # mean inter-token latency across output tokens
|
||||
e2e_latency_ms: float # total request latency
|
||||
input_tokens: int # reported by server (usage.prompt_tokens)
|
||||
output_tokens: int # reported by server (usage.completion_tokens)
|
||||
generated_text: str # generated text
|
||||
itl_ms_list: List[float] = field(default_factory=list) # per-token ITL values
|
||||
|
||||
|
||||
@dataclass
|
||||
class TurnMetric:
|
||||
"""Metrics for a single turn."""
|
||||
|
||||
session_id: str
|
||||
turn: int # 0-indexed
|
||||
ttft_ms: float
|
||||
fc_ms: float # first-chunk latency
|
||||
itl_ms: float # mean inter-token latency
|
||||
e2e_latency_ms: float
|
||||
input_tokens: int
|
||||
output_tokens: int
|
||||
start_time_ms: float # relative to benchmark start
|
||||
itl_ms_list: List[float] = field(default_factory=list) # per-token ITL values
|
||||
|
||||
|
||||
@dataclass
|
||||
class WorkloadSpec:
|
||||
"""Workload specification for multi-turn session benchmarks.
|
||||
|
||||
Supports simple mode: specify isl + hit_rate, derive user_tokens and sys_tokens.
|
||||
All parameters are scalar (fixed) values -- no distributions.
|
||||
"""
|
||||
|
||||
# Core parameters
|
||||
num_sessions: Optional[int] = None # total unique sessions (None = duration-based)
|
||||
num_turns: int = 1 # turns per session
|
||||
osl: int = 1 # output sequence length per turn
|
||||
think_time: float = 0.0 # seconds between turns within a session
|
||||
|
||||
# Traffic (use either concurrency or request_rate, not both)
|
||||
concurrency: Optional[int] = None # max concurrent in-flight requests
|
||||
request_rate: Optional[float] = None # requests per second (constant rate mode)
|
||||
ramp_interval: float = -1.0 # seconds between session launches (-1 = auto)
|
||||
|
||||
# Duration-based mode (used with request_rate)
|
||||
duration_s: float = 0.0 # seconds to run benchmark (0 = use num_sessions)
|
||||
|
||||
# Fraction of system prompt shared across all sessions
|
||||
# 1.0 = identical system prompt, 0.0 = all unique
|
||||
shared_system_prompt_ratio: float = 1.0
|
||||
|
||||
# Simple mode inputs (derive user_tokens, sys_tokens)
|
||||
isl: Optional[int] = None
|
||||
hit_rate: Optional[float] = None
|
||||
|
||||
# Resolved values (computed by resolve())
|
||||
_user_tokens: int = field(default=0, init=False, repr=False)
|
||||
_sys_tokens: int = field(default=0, init=False, repr=False)
|
||||
|
||||
def resolve(self) -> "WorkloadSpec":
|
||||
"""Resolve the spec: derive user_tokens and sys_tokens from inputs. Call after init."""
|
||||
if self.isl is None or self.hit_rate is None:
|
||||
raise ValueError("Simple mode requires both --isl and --hit-rate.")
|
||||
self._validate()
|
||||
self._derive_from_simple()
|
||||
return self
|
||||
|
||||
def _derive_from_simple(self) -> None:
|
||||
"""Derive user_tokens and sys_tokens from (ISL, hit_rate, num_turns, OSL, shared_system_prompt_ratio).
|
||||
|
||||
Two equations, two unknowns (u = user_tokens, s = sys_tokens):
|
||||
|
||||
(1) ISL = s + (n+1)/2 · u + (n-1)/2 · a [average input length]
|
||||
(2) (1-h)·ISL = (1-f)·s/n + u [average new-token fraction]
|
||||
|
||||
where n = num_turns, a = osl, f = shared_system_prompt_ratio, h = hit_rate.
|
||||
|
||||
Substituting s from (1) into (2) and solving for u:
|
||||
|
||||
u = [ (1-h)·ISL - (1-f)/n · (ISL - (n-1)·a/2) ]
|
||||
/ [ 1 - (1-f)·(n+1)/(2n) ]
|
||||
|
||||
Then s = ISL - (n+1)/2 · u - (n-1)/2 · a.
|
||||
|
||||
Special case: when n=1 and f=0, equations (1) and (2) collapse to
|
||||
s + u = ISL with h = s/(s+u), giving s = h·ISL and u = (1-h)·ISL.
|
||||
"""
|
||||
isl = self.isl
|
||||
h = self.hit_rate
|
||||
n = self.num_turns
|
||||
a = self.osl
|
||||
f = self.shared_system_prompt_ratio
|
||||
|
||||
denom = 1 - (1 - f) * (n + 1) / (2 * n)
|
||||
if abs(denom) < 1e-9:
|
||||
# n=1, f=0, h=0 (validated earlier): s=0, u=ISL.
|
||||
sys_tokens = 0.0
|
||||
user_tokens = float(isl)
|
||||
else:
|
||||
numer = (1 - h) * isl - (1 - f) / n * (isl - (n - 1) * a / 2)
|
||||
user_tokens = numer / denom
|
||||
sys_tokens = isl - (n + 1) / 2 * user_tokens - (n - 1) / 2 * a
|
||||
|
||||
if user_tokens < 0.5 or sys_tokens < -0.5:
|
||||
suggestions = self._feasibility_suggestions()
|
||||
which = "user_tokens" if user_tokens < 0.5 else "sys_tokens"
|
||||
val = user_tokens if user_tokens < 0.5 else sys_tokens
|
||||
raise ValueError(
|
||||
f"Derived {which} = {val:.1f} is infeasible with "
|
||||
f"(ISL={isl}, hit_rate={h}, num_turns={n}, "
|
||||
f"OSL={a}, shared_system_prompt_ratio={f}).\n"
|
||||
f"To fix, try one of:\n{suggestions}"
|
||||
)
|
||||
|
||||
self._user_tokens = max(1, int(round(user_tokens)))
|
||||
self._sys_tokens = max(0, int(round(sys_tokens)))
|
||||
|
||||
def _feasibility_suggestions(self) -> str:
|
||||
"""Compute feasible boundary values for each parameter and return suggestions.
|
||||
|
||||
For each workload parameter, search for a boundary value that makes
|
||||
the solver yield user_tokens >= 0.5 and sys_tokens >= -0.5 (the
|
||||
minimum values that round to physically meaningful token counts:
|
||||
at least 1 user token and non-negative system tokens).
|
||||
"""
|
||||
isl = self.isl
|
||||
hit_rate = self.hit_rate
|
||||
num_turns = self.num_turns
|
||||
osl = self.osl
|
||||
sharing = self.shared_system_prompt_ratio
|
||||
lines = []
|
||||
|
||||
def _try_solve(isl_, hit_rate_, num_turns_, osl_, sharing_):
|
||||
"""Solve for (user_tokens, sys_tokens) or return None if degenerate."""
|
||||
denom = 1 - (1 - sharing_) * (num_turns_ + 1) / (2 * num_turns_)
|
||||
if abs(denom) < 1e-9:
|
||||
if hit_rate_ > 1e-9:
|
||||
return None
|
||||
return (float(isl_), 0.0)
|
||||
numer = (1 - hit_rate_) * isl_ - (1 - sharing_) / num_turns_ * (
|
||||
isl_ - (num_turns_ - 1) * osl_ / 2
|
||||
)
|
||||
user_tokens = numer / denom
|
||||
sys_tokens = (
|
||||
isl_ - (num_turns_ + 1) / 2 * user_tokens - (num_turns_ - 1) / 2 * osl_
|
||||
)
|
||||
return (user_tokens, sys_tokens)
|
||||
|
||||
def _feasible(isl_, hit_rate_, num_turns_, osl_, sharing_):
|
||||
result = _try_solve(isl_, hit_rate_, num_turns_, osl_, sharing_)
|
||||
# user_tokens >= 0.5 rounds to at least 1 token per turn;
|
||||
# sys_tokens >= -0.5 rounds to at least 0 system prompt tokens.
|
||||
return result is not None and result[0] >= 0.5 and result[1] >= -0.5
|
||||
|
||||
# Min ISL (binary search)
|
||||
lo, hi = isl, isl * 20
|
||||
if _feasible(hi, hit_rate, num_turns, osl, sharing):
|
||||
while hi - lo > 1:
|
||||
mid = (lo + hi) // 2
|
||||
if _feasible(mid, hit_rate, num_turns, osl, sharing):
|
||||
hi = mid
|
||||
else:
|
||||
lo = mid
|
||||
lines.append(f" - ISL >= {hi} (with current params)")
|
||||
|
||||
# Max OSL
|
||||
lo, hi = 1, osl
|
||||
if _feasible(isl, hit_rate, num_turns, lo, sharing):
|
||||
while hi - lo > 1:
|
||||
mid = (lo + hi) // 2
|
||||
if _feasible(isl, hit_rate, num_turns, mid, sharing):
|
||||
lo = mid
|
||||
else:
|
||||
hi = mid
|
||||
lines.append(f" - OSL <= {lo} (with current ISL={isl})")
|
||||
|
||||
# Min hit_rate / max hit_rate (search in 0.01 steps)
|
||||
for h_try in range(0, 100):
|
||||
h_val = h_try / 100.0
|
||||
if _feasible(isl, h_val, num_turns, osl, sharing):
|
||||
if h_val != hit_rate:
|
||||
if h_val > hit_rate:
|
||||
lines.append(
|
||||
f" - hit_rate >= {h_val:.2f} (with current ISL/OSL)"
|
||||
)
|
||||
else:
|
||||
lines.append(
|
||||
f" - hit_rate <= {h_val:.2f} (with current ISL/OSL)"
|
||||
)
|
||||
break
|
||||
|
||||
# Max num_turns
|
||||
for n_try in range(num_turns, 0, -1):
|
||||
if _feasible(isl, hit_rate, n_try, osl, sharing):
|
||||
if n_try != num_turns:
|
||||
lines.append(f" - num_turns <= {n_try} (with current ISL/OSL)")
|
||||
break
|
||||
|
||||
# Min shared_system_prompt_ratio
|
||||
if sharing < 1.0:
|
||||
for f_try in range(int(sharing * 100), 101):
|
||||
f_val = f_try / 100.0
|
||||
if _feasible(isl, hit_rate, num_turns, osl, f_val):
|
||||
if f_val != sharing:
|
||||
lines.append(f" - shared_system_prompt_ratio >= {f_val:.2f}")
|
||||
break
|
||||
|
||||
return "\n".join(lines) if lines else " (no single-parameter fix found)"
|
||||
|
||||
def _validate(self) -> None:
|
||||
"""Validate resolved parameters."""
|
||||
if self.num_turns < 1:
|
||||
raise ValueError("num_turns must be >= 1.")
|
||||
if self.osl < 1:
|
||||
raise ValueError("osl must be >= 1.")
|
||||
if self.num_sessions is not None and self.num_sessions < 1:
|
||||
raise ValueError("num_sessions must be >= 1.")
|
||||
if self.num_sessions is None and self.duration_s <= 0:
|
||||
raise ValueError(
|
||||
"Must specify either --num-sessions or --duration (> 0) for rate-based mode."
|
||||
)
|
||||
if not (0 <= self.shared_system_prompt_ratio <= 1):
|
||||
raise ValueError("shared_system_prompt_ratio must be in [0, 1].")
|
||||
if self.think_time < 0:
|
||||
raise ValueError("think_time must be >= 0.")
|
||||
if (
|
||||
self.num_turns == 1
|
||||
and self.shared_system_prompt_ratio == 0
|
||||
and self.hit_rate is not None
|
||||
and self.hit_rate > 1e-9
|
||||
):
|
||||
raise ValueError(
|
||||
f"Cannot achieve hit_rate={self.hit_rate} with num_turns=1 and "
|
||||
f"shared_system_prompt_ratio=0. There is no caching source "
|
||||
f"(no multi-turn history, no shared prefix). "
|
||||
f"Set shared_system_prompt_ratio > 0 to enable cross-session "
|
||||
f"prefix caching, or use num_turns > 1 for multi-turn caching."
|
||||
)
|
||||
|
||||
if self.concurrency is None and self.request_rate is None:
|
||||
raise ValueError("Must specify either --concurrency or --request-rate.")
|
||||
if self.concurrency is not None and self.request_rate is not None:
|
||||
raise ValueError("Cannot specify both --concurrency and --request-rate.")
|
||||
if self.concurrency is not None and self.concurrency < 1:
|
||||
raise ValueError("concurrency must be >= 1.")
|
||||
if self.request_rate is not None and self.request_rate <= 0:
|
||||
raise ValueError("request_rate must be > 0.")
|
||||
|
||||
if self.ramp_interval < 0:
|
||||
if self.concurrency is not None:
|
||||
if self.think_time > 0:
|
||||
self.ramp_interval = self.think_time / self.concurrency
|
||||
else:
|
||||
self.ramp_interval = 0.0
|
||||
else:
|
||||
self.ramp_interval = 0.0
|
||||
|
||||
if (
|
||||
self.concurrency is not None
|
||||
and self.think_time > 0
|
||||
and self.num_sessions is not None
|
||||
and self.num_sessions < self.concurrency * 2
|
||||
):
|
||||
logger.warning(
|
||||
"num_sessions=%d may be too low to sustain concurrency=%d "
|
||||
"with think_time=%.1f. Consider increasing num_sessions.",
|
||||
self.num_sessions,
|
||||
self.concurrency,
|
||||
self.think_time,
|
||||
)
|
||||
|
||||
@property
|
||||
def user_tokens(self) -> int:
|
||||
return self._user_tokens
|
||||
|
||||
@property
|
||||
def sys_tokens(self) -> int:
|
||||
return self._sys_tokens
|
||||
|
||||
@property
|
||||
def shared_s(self) -> int:
|
||||
return int(round(self._sys_tokens * self.shared_system_prompt_ratio))
|
||||
|
||||
@property
|
||||
def unique_s(self) -> int:
|
||||
return self._sys_tokens - self.shared_s
|
||||
|
||||
def turn_input_tokens(self, k: int) -> int:
|
||||
"""Total input tokens at turn k (1-indexed)."""
|
||||
return self._sys_tokens + k * self._user_tokens + (k - 1) * self.osl
|
||||
|
||||
@property
|
||||
def effective_isl(self) -> float:
|
||||
n = self.num_turns
|
||||
return (
|
||||
self._sys_tokens + self._user_tokens * (n + 1) / 2 + self.osl * (n - 1) / 2
|
||||
)
|
||||
|
||||
@property
|
||||
def effective_h(self) -> float:
|
||||
f = self.shared_system_prompt_ratio
|
||||
n = self.num_turns
|
||||
avg_new = (1 - f) * self._sys_tokens / n + self._user_tokens
|
||||
isl = self.effective_isl
|
||||
return 1.0 - avg_new / isl if isl > 0 else 0.0
|
||||
|
||||
def summary(self) -> dict:
|
||||
per_turn = []
|
||||
for k in range(1, self.num_turns + 1):
|
||||
total = self.turn_input_tokens(k)
|
||||
if k == 1:
|
||||
cached = int(round(self._sys_tokens * self.shared_system_prompt_ratio))
|
||||
else:
|
||||
cached = (
|
||||
self._sys_tokens + (k - 1) * self._user_tokens + (k - 1) * self.osl
|
||||
)
|
||||
new = total - cached
|
||||
h_k = cached / total if total > 0 else 0.0
|
||||
per_turn.append(
|
||||
{
|
||||
"turn": k,
|
||||
"total": total,
|
||||
"cached": cached,
|
||||
"new": new,
|
||||
"hit_rate": round(h_k, 4),
|
||||
}
|
||||
)
|
||||
|
||||
return {
|
||||
"num_sessions": self.num_sessions,
|
||||
"duration_s": self.duration_s,
|
||||
"num_turns": self.num_turns,
|
||||
"osl": self.osl,
|
||||
"think_time": self.think_time,
|
||||
"concurrency": self.concurrency,
|
||||
"request_rate": self.request_rate,
|
||||
"shared_system_prompt_ratio": self.shared_system_prompt_ratio,
|
||||
"user_tokens_per_turn": self._user_tokens,
|
||||
"system_prompt_tokens": self._sys_tokens,
|
||||
"shared_system_prompt": self.shared_s,
|
||||
"unique_system_prompt": self.unique_s,
|
||||
"effective_isl": round(self.effective_isl, 1),
|
||||
"effective_hit_rate": round(self.effective_h, 4),
|
||||
"per_turn": per_turn,
|
||||
}
|
||||
|
||||
def print_summary(self) -> None:
|
||||
s = self.summary()
|
||||
print("=" * 70)
|
||||
print("Workload Spec (resolved)")
|
||||
print("=" * 70)
|
||||
if s["num_sessions"] is not None:
|
||||
print(f" Sessions (N_s): {s['num_sessions']}")
|
||||
else:
|
||||
print(" Sessions (N_s): unlimited (duration-based)")
|
||||
if s["duration_s"] > 0:
|
||||
print(f" Duration: {s['duration_s']}s")
|
||||
print(f" Turns per session (N_t): {s['num_turns']}")
|
||||
print(f" User tokens/turn (u): {s['user_tokens_per_turn']}")
|
||||
print(
|
||||
f" System prompt (s): {s['system_prompt_tokens']} "
|
||||
f"(shared={s['shared_system_prompt']}, unique={s['unique_system_prompt']})"
|
||||
)
|
||||
print(f" Output tokens (o): {s['osl']}")
|
||||
print(f" Think time: {s['think_time']}s")
|
||||
if self.concurrency is not None:
|
||||
print(f" Concurrency (C): {self.concurrency}")
|
||||
print(f" Ramp interval: {self.ramp_interval:.3f}s")
|
||||
if self.request_rate is not None:
|
||||
print(f" Request rate (QPS): {self.request_rate}")
|
||||
print(f" Shared sys prompt ratio: {s['shared_system_prompt_ratio']}")
|
||||
print(f" Effective avg ISL: {s['effective_isl']}")
|
||||
print(f" Effective avg hit rate: {s['effective_hit_rate']:.1%}")
|
||||
print("-" * 70)
|
||||
print(f" {'Turn':<6} {'Total':<8} {'Cached':<8} {'New':<8} {'Hit Rate':<10}")
|
||||
for t in s["per_turn"]:
|
||||
print(
|
||||
f" {t['turn']:<6} {t['total']:<8} {t['cached']:<8} "
|
||||
f"{t['new']:<8} {t['hit_rate']:.1%}"
|
||||
)
|
||||
print("=" * 70)
|
||||
@@ -0,0 +1,154 @@
|
||||
"""Reporting and result persistence for the benchmark."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import logging
|
||||
from pathlib import Path
|
||||
from statistics import mean
|
||||
from typing import Optional
|
||||
|
||||
from ray.llm._internal.serve.benchmark.metrics import (
|
||||
percentile,
|
||||
serialize_raw_metrics,
|
||||
summarize_metrics,
|
||||
)
|
||||
from ray.llm._internal.serve.benchmark.models import TurnMetric, WorkloadSpec
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def report_results(
|
||||
metrics: list[TurnMetric],
|
||||
spec: WorkloadSpec,
|
||||
bench_elapsed_s: float,
|
||||
first_chunk_threshold: int = 16,
|
||||
save_path: Optional[str] = None,
|
||||
warmup_s: float = 0.0,
|
||||
discarded_warmup_requests: int = 0,
|
||||
) -> None:
|
||||
"""Print and optionally save benchmark results."""
|
||||
if not metrics:
|
||||
print("No metrics collected.")
|
||||
return
|
||||
|
||||
all_ttft = [m.ttft_ms for m in metrics]
|
||||
all_fc = [m.fc_ms for m in metrics]
|
||||
all_itl = [v for m in metrics for v in m.itl_ms_list]
|
||||
all_latency = [m.e2e_latency_ms for m in metrics]
|
||||
all_input = [m.input_tokens for m in metrics]
|
||||
all_output = [m.output_tokens for m in metrics]
|
||||
|
||||
total_output_tokens = sum(all_output)
|
||||
throughput = total_output_tokens / bench_elapsed_s if bench_elapsed_s > 0 else 0
|
||||
|
||||
print()
|
||||
print("=" * 70)
|
||||
print("BENCHMARK RESULTS")
|
||||
print("=" * 70)
|
||||
print(f" Total requests: {len(metrics)}")
|
||||
print(f" Unique sessions: {len({m.session_id for m in metrics})}")
|
||||
print(f" Duration: {bench_elapsed_s:.1f}s")
|
||||
if warmup_s > 0:
|
||||
print(f" Warm-up excluded: {warmup_s:.1f}s")
|
||||
if discarded_warmup_requests > 0:
|
||||
print(f" Warm-up requests: {discarded_warmup_requests} (discarded)")
|
||||
print(f" Throughput: {throughput:.1f} output tok/s")
|
||||
print(f" Request rate: {len(metrics) / bench_elapsed_s:.1f} req/s")
|
||||
print(
|
||||
f" Avg input tokens: {mean(all_input):.0f} "
|
||||
f"(target ISL: {spec.effective_isl:.0f})"
|
||||
)
|
||||
print(f" Avg output tokens: {mean(all_output):.0f} (target OSL: {spec.osl})")
|
||||
print()
|
||||
|
||||
fc_label = f"FC({first_chunk_threshold})"
|
||||
print(" Latency Statistics:")
|
||||
for name, values in [
|
||||
("TTFT", all_ttft),
|
||||
(fc_label, all_fc),
|
||||
("ITL", all_itl),
|
||||
("Latency", all_latency),
|
||||
]:
|
||||
if not values:
|
||||
continue
|
||||
print(
|
||||
f" {name:>8}: avg={mean(values):>8.1f}ms "
|
||||
f"P50={percentile(values, 50):>8.1f}ms "
|
||||
f"P90={percentile(values, 90):>8.1f}ms "
|
||||
f"P99={percentile(values, 99):>8.1f}ms"
|
||||
)
|
||||
print()
|
||||
|
||||
print(" Per-Turn Breakdown:")
|
||||
print(
|
||||
f" {'Turn':<6} {'Count':<7} {'Avg ISL':<9} {'Avg TTFT':<10} "
|
||||
f"{'Avg FC':<10} {'Avg ITL':<10} {'Avg Lat':<10}"
|
||||
)
|
||||
for t in range(spec.num_turns):
|
||||
turn_metrics = [m for m in metrics if m.turn == t]
|
||||
if not turn_metrics:
|
||||
continue
|
||||
t_ttft = mean([m.ttft_ms for m in turn_metrics])
|
||||
t_fc = mean([m.fc_ms for m in turn_metrics])
|
||||
t_itl_all = [v for m in turn_metrics for v in m.itl_ms_list]
|
||||
t_itl = mean(t_itl_all) if t_itl_all else 0.0
|
||||
t_lat = mean([m.e2e_latency_ms for m in turn_metrics])
|
||||
t_isl = mean([m.input_tokens for m in turn_metrics])
|
||||
print(
|
||||
f" {t + 1:<6} {len(turn_metrics):<7} {t_isl:<9.0f} "
|
||||
f"{t_ttft:<10.1f} {t_fc:<10.1f} {t_itl:<10.1f} {t_lat:<10.1f}"
|
||||
)
|
||||
print("=" * 70)
|
||||
|
||||
if save_path:
|
||||
stats = summarize_metrics(metrics, bench_elapsed_s)
|
||||
result = {
|
||||
"config": {
|
||||
"concurrency": spec.concurrency,
|
||||
"request_rate": spec.request_rate,
|
||||
},
|
||||
"spec": spec.summary(),
|
||||
"first_chunk_threshold": first_chunk_threshold,
|
||||
"benchmark": {
|
||||
"total_requests": len(metrics),
|
||||
"duration_s": round(bench_elapsed_s, 2),
|
||||
"warmup_s": round(warmup_s, 2),
|
||||
"discarded_warmup_requests": discarded_warmup_requests,
|
||||
},
|
||||
"stats": {
|
||||
("measured_request_rate" if k == "request_rate" else k): v
|
||||
for k, v in stats.items()
|
||||
if k not in ("requests", "elapsed_s")
|
||||
},
|
||||
"per_turn": [],
|
||||
"raw_metrics": serialize_raw_metrics(metrics),
|
||||
}
|
||||
|
||||
for t in range(spec.num_turns):
|
||||
turn_metrics = [m for m in metrics if m.turn == t]
|
||||
if not turn_metrics:
|
||||
continue
|
||||
t_ttft = [m.ttft_ms for m in turn_metrics]
|
||||
t_fc = [m.fc_ms for m in turn_metrics]
|
||||
t_itl = [v for m in turn_metrics for v in m.itl_ms_list]
|
||||
t_isl = [m.input_tokens for m in turn_metrics]
|
||||
result["per_turn"].append(
|
||||
{
|
||||
"turn": t + 1,
|
||||
"count": len(turn_metrics),
|
||||
"avg_isl": round(mean(t_isl), 1),
|
||||
"avg_ttft_ms": round(mean(t_ttft), 2),
|
||||
"avg_fc_ms": round(mean(t_fc), 2),
|
||||
"avg_itl_ms": round(mean(t_itl), 2) if t_itl else 0,
|
||||
"p50_fc_ms": round(percentile(t_fc, 50), 2),
|
||||
"p99_ttft_ms": round(percentile(t_ttft, 99), 2),
|
||||
"p99_fc_ms": round(percentile(t_fc, 99), 2),
|
||||
"p99_itl_ms": (round(percentile(t_itl, 99), 2) if t_itl else 0),
|
||||
}
|
||||
)
|
||||
|
||||
Path(save_path).parent.mkdir(parents=True, exist_ok=True)
|
||||
with open(save_path, "w") as f:
|
||||
json.dump(result, f, indent=2)
|
||||
logger.info("Results saved to %s", save_path)
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,175 @@
|
||||
"""Text generation and conversation management for the benchmark."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
from typing import TYPE_CHECKING, Optional
|
||||
|
||||
import numpy as np
|
||||
|
||||
from ray.llm._internal.serve.benchmark.models import WorkloadSpec
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from transformers import PreTrainedTokenizerBase
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class Conversation:
|
||||
"""A single multi-turn conversation with a unique session ID."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
session_id: str,
|
||||
system_prompt: str,
|
||||
user_messages: list[str],
|
||||
num_turns: int,
|
||||
):
|
||||
self.session_id = session_id
|
||||
self.system_prompt = system_prompt
|
||||
self.user_messages = user_messages
|
||||
self.num_turns = num_turns
|
||||
self._assistant_responses: list[str] = []
|
||||
|
||||
def get_turn_messages(self, turn_idx: int) -> list[dict[str, str]]:
|
||||
"""Build the messages list for turn `turn_idx` (0-indexed)."""
|
||||
messages: list[dict[str, str]] = []
|
||||
if self.system_prompt:
|
||||
messages.append({"role": "system", "content": self.system_prompt})
|
||||
|
||||
for i in range(turn_idx + 1):
|
||||
messages.append({"role": "user", "content": self.user_messages[i]})
|
||||
if i < turn_idx:
|
||||
if i < len(self._assistant_responses):
|
||||
messages.append(
|
||||
{"role": "assistant", "content": self._assistant_responses[i]}
|
||||
)
|
||||
else:
|
||||
messages.append({"role": "assistant", "content": "(placeholder)"})
|
||||
return messages
|
||||
|
||||
def inject_assistant_response(self, turn_idx: int, content: str) -> None:
|
||||
"""Record the server's response for turn `turn_idx`."""
|
||||
if turn_idx == len(self._assistant_responses):
|
||||
self._assistant_responses.append(content)
|
||||
elif turn_idx < len(self._assistant_responses):
|
||||
self._assistant_responses[turn_idx] = content
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Cannot inject response for turn {turn_idx}: "
|
||||
f"only {len(self._assistant_responses)} responses recorded."
|
||||
)
|
||||
|
||||
|
||||
class TextGenerator:
|
||||
"""Generates random text with exact token counts using a tokenizer."""
|
||||
|
||||
def __init__(self, tokenizer: "PreTrainedTokenizerBase"):
|
||||
self._tokenizer = tokenizer
|
||||
self._vocab_size = tokenizer.vocab_size
|
||||
logger.info(
|
||||
"TextGenerator using tokenizer (vocab_size=%d) for exact token counts.",
|
||||
self._vocab_size,
|
||||
)
|
||||
|
||||
def generate(self, num_tokens: int) -> str:
|
||||
if num_tokens <= 0:
|
||||
return ""
|
||||
return self._generate_exact(num_tokens)
|
||||
|
||||
def generate_token_ids(self, num_tokens: int) -> list[int]:
|
||||
if num_tokens <= 0:
|
||||
return []
|
||||
return np.random.randint(0, self._vocab_size, size=num_tokens).tolist()
|
||||
|
||||
def _generate_exact(self, target_tokens: int) -> str:
|
||||
tokenizer = self._tokenizer
|
||||
token_ids = np.random.randint(
|
||||
0, self._vocab_size, size=target_tokens + 20
|
||||
).tolist()
|
||||
|
||||
text = tokenizer.decode(token_ids, skip_special_tokens=True)
|
||||
actual_ids = tokenizer.encode(text, add_special_tokens=False)
|
||||
actual_len = len(actual_ids)
|
||||
|
||||
if actual_len == target_tokens:
|
||||
return text
|
||||
|
||||
if actual_len > target_tokens:
|
||||
trimmed_ids = actual_ids[:target_tokens]
|
||||
text = tokenizer.decode(trimmed_ids, skip_special_tokens=True)
|
||||
final_len = len(tokenizer.encode(text, add_special_tokens=False))
|
||||
if final_len != target_tokens:
|
||||
text = self._binary_search_trim(actual_ids, target_tokens)
|
||||
return text
|
||||
|
||||
deficit = target_tokens - actual_len
|
||||
extra_ids = np.random.randint(0, self._vocab_size, size=deficit + 20).tolist()
|
||||
extra_text = tokenizer.decode(extra_ids, skip_special_tokens=True)
|
||||
combined = text + " " + extra_text
|
||||
combined_ids = tokenizer.encode(combined, add_special_tokens=False)
|
||||
|
||||
if len(combined_ids) >= target_tokens:
|
||||
trimmed = combined_ids[:target_tokens]
|
||||
text = tokenizer.decode(trimmed, skip_special_tokens=True)
|
||||
final_len = len(tokenizer.encode(text, add_special_tokens=False))
|
||||
if final_len != target_tokens:
|
||||
text = self._binary_search_trim(combined_ids, target_tokens)
|
||||
return text
|
||||
|
||||
while len(tokenizer.encode(combined, add_special_tokens=False)) < target_tokens:
|
||||
combined += " hello"
|
||||
combined_ids = tokenizer.encode(combined, add_special_tokens=False)
|
||||
return self._binary_search_trim(combined_ids, target_tokens)
|
||||
|
||||
def _binary_search_trim(self, token_ids: list[int], target: int) -> str:
|
||||
tokenizer = self._tokenizer
|
||||
lo, hi = target, len(token_ids)
|
||||
best_text = tokenizer.decode(token_ids[:target], skip_special_tokens=True)
|
||||
|
||||
while lo <= hi:
|
||||
mid = (lo + hi) // 2
|
||||
text = tokenizer.decode(token_ids[:mid], skip_special_tokens=True)
|
||||
actual = len(tokenizer.encode(text, add_special_tokens=False))
|
||||
if actual == target:
|
||||
return text
|
||||
elif actual < target:
|
||||
lo = mid + 1
|
||||
else:
|
||||
hi = mid - 1
|
||||
best_text = text
|
||||
|
||||
for n in range(target, len(token_ids) + 1):
|
||||
text = tokenizer.decode(token_ids[:n], skip_special_tokens=True)
|
||||
if len(tokenizer.encode(text, add_special_tokens=False)) == target:
|
||||
return text
|
||||
return best_text
|
||||
|
||||
|
||||
def conversation_factory(
|
||||
session_idx: int,
|
||||
spec: WorkloadSpec,
|
||||
shared_system_text: str,
|
||||
text_gen: Optional[TextGenerator],
|
||||
) -> Conversation:
|
||||
"""Create a single conversation on-demand (lazy generation)."""
|
||||
session_id = f"session-{session_idx:06d}"
|
||||
|
||||
if spec.unique_s > 0 and text_gen is not None:
|
||||
unique_text = text_gen.generate(spec.unique_s)
|
||||
system_prompt = shared_system_text + " " + unique_text
|
||||
else:
|
||||
system_prompt = shared_system_text
|
||||
|
||||
user_messages = (
|
||||
[text_gen.generate(spec.user_tokens) for _ in range(spec.num_turns)]
|
||||
if text_gen is not None
|
||||
else ["" for _ in range(spec.num_turns)]
|
||||
)
|
||||
|
||||
return Conversation(
|
||||
session_id=session_id,
|
||||
system_prompt=system_prompt,
|
||||
user_messages=user_messages,
|
||||
num_turns=spec.num_turns,
|
||||
)
|
||||
@@ -0,0 +1,75 @@
|
||||
"""Single-turn execution primitive for the benchmark.
|
||||
|
||||
This module provides the pure core of turn execution: send an HTTP request,
|
||||
build a TurnMetric, and inject the response. It has NO side effects — callers
|
||||
are responsible for inflight tracking, metric recording, and queue management.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import time
|
||||
from dataclasses import dataclass
|
||||
from typing import Optional
|
||||
|
||||
import aiohttp
|
||||
|
||||
from ray.llm._internal.serve.benchmark.http_client import send_chat_completion
|
||||
from ray.llm._internal.serve.benchmark.models import TurnMetric, TurnResult
|
||||
from ray.llm._internal.serve.benchmark.text_gen import Conversation
|
||||
|
||||
|
||||
@dataclass
|
||||
class TurnOutcome:
|
||||
"""Result of executing a single benchmark turn."""
|
||||
|
||||
metric: TurnMetric
|
||||
result: TurnResult
|
||||
|
||||
|
||||
async def execute_single_turn(
|
||||
http_session: aiohttp.ClientSession,
|
||||
conv: Conversation,
|
||||
turn_idx: int,
|
||||
base_url: str,
|
||||
model: str,
|
||||
max_tokens: int,
|
||||
bench_start_ns: int,
|
||||
first_chunk_threshold: int = 16,
|
||||
api_key: Optional[str] = None,
|
||||
) -> TurnOutcome:
|
||||
"""Execute a single benchmark turn: HTTP call, build metric, inject response.
|
||||
|
||||
This is the pure core shared by all three benchmark engines (concurrency,
|
||||
rate-based, interactive). The caller handles inflight tracking, warmup
|
||||
filtering, measurement windows, and queue re-enqueue.
|
||||
"""
|
||||
messages = conv.get_turn_messages(turn_idx)
|
||||
req_start_ns = time.perf_counter_ns()
|
||||
|
||||
result = await send_chat_completion(
|
||||
session=http_session,
|
||||
base_url=base_url,
|
||||
model=model,
|
||||
messages=messages,
|
||||
session_id=conv.session_id,
|
||||
max_tokens=max_tokens,
|
||||
first_chunk_threshold=first_chunk_threshold,
|
||||
api_key=api_key,
|
||||
)
|
||||
|
||||
metric = TurnMetric(
|
||||
session_id=conv.session_id,
|
||||
turn=turn_idx,
|
||||
ttft_ms=result.ttft_ms,
|
||||
fc_ms=result.fc_ms,
|
||||
itl_ms=result.itl_ms,
|
||||
e2e_latency_ms=result.e2e_latency_ms,
|
||||
input_tokens=result.input_tokens,
|
||||
output_tokens=result.output_tokens,
|
||||
start_time_ms=(req_start_ns - bench_start_ns) / 1e6,
|
||||
itl_ms_list=result.itl_ms_list,
|
||||
)
|
||||
|
||||
conv.inject_assistant_response(turn_idx, result.generated_text)
|
||||
|
||||
return TurnOutcome(metric=metric, result=result)
|
||||
Reference in New Issue
Block a user