@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 9354f22204

[Rust][Benchmark] Port in vllm-bench (#48107)

Co-authored-by: Bugen Zhao <i@bugenzhao.com> Signed-off-by: esmeetu <jasonailu87@gmail.com> Signed-off-by: Bugen Zhao <i@bugenzhao.com>

Roy Wang committed Jul 17, 2026 at 14:25 UTC 9354f222042986addf20709e5274fc26e0d09745
45 files changed +16926 -2
.pre-commit-config.yaml
+1 -1
@@ -30,7 +30,7 @@ repos:
30 - id: markdownlint-cli2
31 language_version: lts
32 args: [--fix]
33 - exclude: ^CLAUDE\.md$
33 + exclude: (^|/)CLAUDE\.md$
34 - repo: https://github.com/rhysd/actionlint
35 rev: v1.7.7
36 hooks:
pyproject.toml
+2 -1
@@ -122,7 +122,8 @@ python = "./.venv"
122 [tool.typos.files]
123 # these files may be written in non english words
124 extend-exclude = ["tests/models/fixtures/*", "tests/prompts/*", "tests/tokenizers_/*",
125 - "benchmarks/sonnet.txt", "tests/lora/data/*", "build/*",
125 + "benchmarks/sonnet.txt", "rust/src/bench/src/datasets/sonnet.txt",
126 + "tests/lora/data/*", "build/*",
127 "examples/pooling/token_embed/*", "tests/models/language/pooling/*",
128 "vllm/third_party/*", "vllm/entrypoints/serve/instrumentator/static/*",
129 "tests/entrypoints/speech_to_text/transcription/test_transcription_validation.py",
rust/Cargo.lock
+53
@@ -2585,6 +2585,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
2585 checksum = "071dfc062690e90b734c0b2273ce72ad0ffa95f0c74596bc250dcfd960262841"
2586 dependencies = [
2587 "autocfg",
2588 + "libm",
2589 ]
2590
2591 [[package]]
@@ -3305,6 +3306,16 @@ dependencies = [
3306 "getrandom 0.3.4",
3307 ]
3308
3309 +[[package]]
3310 +name = "rand_distr"
3311 +version = "0.5.1"
3312 +source = "registry+https://github.com/rust-lang/crates.io-index"
3313 +checksum = "6a8615d50dcf34fa31f7ab52692afec947c4dd0ab803cc87cb3b0b4570ff7463"
3314 +dependencies = [
3315 + "num-traits",
3316 + "rand 0.9.2",
3317 +]
3318 +
3319 [[package]]
3320 name = "rawpointer"
3321 version = "0.2.1"
@@ -3435,6 +3446,7 @@ dependencies = [
3446 "base64 0.22.1",
3447 "bytes",
3448 "encoding_rs",
3449 + "futures-channel",
3450 "futures-core",
3451 "futures-util",
3452 "h2",
@@ -3549,6 +3561,15 @@ dependencies = [
3561 "rustc-hash 2.1.1",
3562 ]
3563
3564 +[[package]]
3565 +name = "rlimit"
3566 +version = "0.11.0"
3567 +source = "registry+https://github.com/rust-lang/crates.io-index"
3568 +checksum = "f35ee2729c56bb610f6dba436bf78135f728b7373bdffae2ec815b2d3eb98cc3"
3569 +dependencies = [
3570 + "libc",
3571 +]
3572 +
3573 [[package]]
3574 name = "rmp"
3575 version = "0.8.15"
@@ -5439,6 +5460,38 @@ version = "0.9.5"
5460 source = "registry+https://github.com/rust-lang/crates.io-index"
5461 checksum = "0b928f33d975fc6ad9f86c8f283853ad26bdd5b10b7f1542aa2fa15e2289105a"
5462
5463 +[[package]]
5464 +name = "vllm-bench"
5465 +version = "0.1.0"
5466 +dependencies = [
5467 + "anyhow",
5468 + "base64 0.22.1",
5469 + "bytes",
5470 + "chrono",
5471 + "clap",
5472 + "dirs",
5473 + "futures",
5474 + "hf-hub",
5475 + "image",
5476 + "indicatif",
5477 + "mimalloc",
5478 + "rand 0.9.2",
5479 + "rand_distr",
5480 + "rayon",
5481 + "reqwest 0.12.28",
5482 + "rlimit",
5483 + "rustc-hash 1.1.0",
5484 + "serde",
5485 + "serde_json",
5486 + "thiserror 2.0.18",
5487 + "tiktoken-rs 0.9.1",
5488 + "tokenizers",
5489 + "tokio",
5490 + "tokio-stream",
5491 + "url",
5492 + "uuid",
5493 +]
5494 +
5495 [[package]]
5496 name = "vllm-chat"
5497 version = "0.1.0"
rust/Cargo.toml
+9
@@ -1,5 +1,6 @@
1 [workspace]
2 members = [
3 + "src/bench",
4 "src/chat",
5 "src/cmd",
6 "src/engine-core-client",
@@ -32,8 +33,10 @@ base64 = "0.22.1"
33 bytemuck = { version = "1.25.0", features = ["extern_crate_alloc"] }
34 byteorder = "1.5.0"
35 bytes = "1.12.0"
36 +chrono = "0.4.42"
37 clap = { version = "4.5.38", features = ["derive", "env"] }
38 criterion = "0.5.1"
39 +dirs = "6.0.0"
40 easy-ext = "1.0.3"
41 educe = "0.6.0"
42 enum-as-inner = "0.7.0"
@@ -50,7 +53,9 @@ hyper-util = { version = "0.1.20", features = [
53 "service",
54 "tokio",
55 ] }
56 +image = { version = "0.25.9", default-features = false, features = ["jpeg"] }
57 indexmap = "2.13.0"
58 +indicatif = "0.18.4"
59 itertools = "0.14.0"
60 libc = "0.2.177"
61 llm-multimodal = { git = "https://github.com/smg-project/llm-multimodal", rev = "5390032d6dc8a3e6fdc83acd320260367eb4b9b5", default-features = false, features = ["native-tls"] }
@@ -71,10 +76,13 @@ prost-types = "0.14.3"
76 pyo3 = "0.28.3"
77 pythonize = "0.28.0"
78 rand = "0.9.2"
79 +rand_distr = "0.5.1"
80 +rayon = "1.11.0"
81 reasoning-parser = "1.2.2"
82 reqwest = { version = "0.12.8", default-features = false, features = ["native-tls"] }
83 reqwest-0-13 = { package = "reqwest", version = "0.13.4", default-features = false, features = ["native-tls"] }
84 riptoken = { version = "0.3.0", default-features = false }
85 +rlimit = "0.11.0"
86 rmp-serde = "1.3.1"
87 rmpv = { version = "1.3.1", features = ["with-serde"] }
88 rustc-hash = "1.1.0"
@@ -121,6 +129,7 @@ tracing = { version = "0.1.44", features = ["release_max_level_debug"] }
129 tracing-futures = { version = "0.2.5", features = ["futures-03"] }
130 tracing-subscriber = { version = "0.3.20", features = ["env-filter", "fmt"] }
131 trait-set = "0.3.0"
132 +url = "2.5.7"
133 uuid = { version = "1.22.0", features = ["v4"] }
134 validator = { version = "0.20.0", features = ["derive"] }
135 vllm-chat = { path = "src/chat" }
rust/src/bench/AGENTS.md new
+181
@@ -0,0 +1,181 @@
1 +# AGENTS.md
2 +
3 +## Project Overview
4 +
5 +Rust rewrite of `vllm bench serve` — a high-performance benchmark client for vLLM serving endpoints. Standalone binary, no Python dependency at runtime.
6 +
7 +Member crate `vllm-bench` of the `rust/` workspace. Uses workspace dependencies and lints; the workspace `[profile.release]` (thin LTO, `panic = "abort"`) applies. Note the workspace bans rustls/ring (`rust/deny.toml`) — all HTTP must stay on native-tls, which is why HF Hub downloads go through `src/hub.rs` (async hf-hub API bridged to sync) instead of hf-hub's ureq backend.
8 +
9 +## Build & Test
10 +
11 +Run from the `rust/` workspace root:
12 +
13 +```bash
14 +# Build release binary (rust/target/release/vllm-bench)
15 +cargo build -p vllm-bench --release
16 +
17 +# Run all tests
18 +cargo test -p vllm-bench
19 +
20 +# Run ignored integration tests (requires network for tokenizer download)
21 +cargo test -p vllm-bench -- --ignored
22 +```
23 +
24 +## Architecture
25 +
26 +- `src/main.rs` — Entry point, mimalloc, tokio runtime, mode dispatch (compare/sweep/multi-run/multi-turn/single)
27 +- `src/cli.rs` — clap derive CLI args (~50+ flags)
28 +- `src/config.rs` — Validated config from CLI; `GoodputConfig`, `RampUpConfig`, sampling param merging
29 +- `src/error.rs` — `BenchError` enum (Http, Json, Tokenizer, Config, EndpointTimeout, Backend, Io)
30 +- `src/benchmark.rs` — Core benchmark orchestrator (spawn-per-request with tokio + Semaphore; fetches speculative decoding metrics from `/metrics`)
31 +- `src/multi_turn.rs` — Multi-turn conversation orchestrator (channel-based worker pool, sequential turns per conversation)
32 +- `src/sweep.rs` — Concurrency/rate parameter sweep (`--sweep-max-concurrency`, `--sweep-request-rate`)
33 +- `src/multi_run.rs` — N-run aggregation with mean/std/min/max/CV (`--num-runs`)
34 +- `src/compare.rs` — Side-by-side diff of two result JSON files (`--compare`)
35 +- `src/tokenizer.rs` — `TokenizerKind` enum: Local(HuggingFace), Tiktoken, OR Server-side `/tokenize`+`/detokenize` fallback
36 +- `src/tiktoken.rs` — Tiktoken BPE loader (`.tiktoken`/`.model` files; built-in encodings o200k_base/cl100k_base; pat_str extraction from Python source)
37 +- `src/hub.rs` — `HubRepo`: sync facade over hf-hub's async (reqwest/native-tls) API — per-download thread with its own runtime; the sync ureq backend is unusable here because it pulls rustls, which `rust/deny.toml` bans
38 +- `src/rate_control.rs` — Gamma/Poisson request scheduling + linear/exponential ramp-up
39 +- `src/ready_checker.rs` — Endpoint readiness with retry
40 +- `src/backends/` — Backend implementations (enum dispatch, not trait objects)
41 + - `mod.rs` — `Backend` enum, `RequestFuncInput`/`RequestFuncOutput` (includes `messages` field for multi-turn)
42 + - `streaming.rs` — SSE parser (`StreamedResponseHandler`) with speculative JSON parse for split TCP segments
43 + - `openai_completions.rs` — `/v1/completions` backend
44 + - `openai_chat.rs` — `/v1/chat/completions` backend (uses `input.messages` when set; zero-copy raw JSON payload for multimodal)
45 + - `pooling.rs` — Non-streaming pooling/embedding backends: `openai-embeddings`, `openai-embeddings-chat`, `vllm-pooling`, `vllm-rerank`
46 +- `src/datasets/random.rs` — Random dataset generation with rayon parallelism
47 +- `src/datasets/random_mm.rs` — Random multimodal dataset (synthetic JPEG images, bucket config sampling, pre-serialized JSON fragments); `--enable-multimodal-chat` pre-builds the chat `messages` array at dataset time (mirrors Python's `apply_multimodal_chat_transformation`)
48 +- `src/datasets/sharegpt.rs` — ShareGPT JSON loader + HuggingFace Hub auto-download with caching
49 +- `src/datasets/sonnet.rs` and `src/datasets/sonnet.txt` — Sonnet dataset (built-in Shakespeare sonnets via `include_str!("sonnet.txt")`; controllable token length + shared prefix; mirrors Python `SonnetDataset`)
50 +- `src/datasets/speed_bench.rs` — NVIDIA SPEED-Bench loader (HF datasets-server API, 6 configs, 11 categories, local cache)
51 +- `src/datasets/hf_dataset.rs` — Generic HuggingFace dataset loader (datasets-server API, column auto-detection)
52 +- `src/datasets/custom.rs` — Custom JSONL dataset (`{"prompt": ..., "output_tokens": ...}` per line; `--custom-output-len -1` uses per-line output_tokens; prompts always sent raw — no client-side chat template)
53 +- `src/datasets/prefix_repetition.rs` — Prefix repetition dataset (N shared prefixes × fresh random suffixes, standard prefix-cache stress; mirrors Python `PrefixRepetitionRandomDataset`)
54 +- `src/datasets/random_rerank.rs` — Random rerank dataset (one query + batched documents per request for `vllm-rerank`; `--no-reranker` for embedding-based scoring; mirrors Python `RandomDatasetForReranking`)
55 +- `src/datasets/multi_turn.rs` — Multi-turn synthetic generator + ShareGPT multi-turn loader (3-tier prefix sharing: global/conversation/unique-suffix; `per_turn_input_len`)
56 +- `src/metrics/mod.rs` — `BenchmarkMetrics` and `MultiTurnMetrics` structs
57 +- `src/metrics/calculator.rs` — TTFT/TPOT/ITL/E2EL/throughput stats, goodput SLO checking, peak concurrency, `calculate_multi_turn_metrics`
58 +- `src/metrics/steady_state.rs` — Steady-state window detection (in-flight concurrency plateau via two-pointer start/end merge) + plateau throughput/TTFT/TPOT; gated on `--max-concurrency` set + `--request-rate inf` (closed-loop)
59 +- `src/output/console.rs` — Terminal output matching Python format + multi-turn per-turn breakdown
60 +- `src/output/json.rs` — JSON result file (compatible with Python schema) + multi-turn JSON with `per_turn_metrics`
61 +
62 +## Key Design Decisions
63 +
64 +- **Enum dispatch** for backends (avoids async trait object issues with `dyn`)
65 +- **reqwest http1_only()** to match Python aiohttp behavior
66 +- **rayon** for parallel dataset generation (key perf win over Python)
67 +- **mimalloc** global allocator to reduce contention at 1400+ concurrency (page-agnostic; works on aarch64 64K-page kernels where jemalloc aborts with `LG_PAGE=12` builds)
68 +- **Arc\<str\> prompts** zero-copy sharing across tokio tasks (~3GB savings at 100k prompts with 8k-token inputs)
69 +- **Spawn-per-request** `tokio::spawn` + `Semaphore` (matches Python asyncio pattern)
70 +- **Speculative JSON parse** in SSE handler — detects complete JSON before `\n\n` arrives, improving TTFT/ITL accuracy when TCP segments split
71 +- **Tokenizer fallback chain**: Local HF → Tiktoken (`.tiktoken`/`.model` + built-in encodings) → Server-side `/tokenize`+`/detokenize`. Blocking HTTP in rayon threads for server fallback.
72 +- **hf-hub** for downloading tokenizers and datasets from HuggingFace Hub
73 +- **Pre-serialized mm fragments** (`Arc<str>`) for multimodal: image content stored as JSON strings, zero-copy concatenated into payload — avoids deep-cloning ~200KB+ base64 per request
74 +- **Steady-state metrics** (default-on in closed-loop): measure throughput/TTFT/TPOT only over the saturated plateau to cut run-to-run variance at high concurrency; `steady_state` is an `Option` in JSON (`#[serde(default)]` for backward compat), null when the scope gate fails or `--no-steady-state`
75 +- **`--prompt-token-ids`** (random dataset only): send token-ID arrays instead of text to skip server-side tokenization; also skips the token-length verification pass (counts exact by construction)
76 +- **`--random-range-ratio`** follows Python semantics: lengths sampled uniformly from `[len*(1-r), len*(1+r)]`, default `0.0` = fixed; accepts a float in `[0,1)` or `'{"input": r1, "output": r2}'`. (The pre-2026-07 Rust-only form `[len*r, len]` with default 1.0 is rejected with a migration hint.)
77 +- **`prompt_list`** (`Arc<[Arc<str>]>` on `SampleRequest`/`RequestFuncInput`): multiple inputs per request for pooling backends — embeddings batches (`--random-batch-size`) send `"input": [...]`, rerank sends `[0]` as query + `[1..]` as documents
78 +- JSON output schema must match Python `vllm bench serve` exactly
79 +
80 +## Common Issues
81 +
82 +- **localhost vs 127.0.0.1**: Some systems resolve `localhost` to IPv6 `::1` while vLLM listens on IPv4 only. Use `127.0.0.1` or the actual hostname.
83 +- **Models without tokenizer.json** (e.g., `nvidia/Kimi-K2.5-NVFP4`): Automatically falls back to server-side tokenization. Can also use `--tokenizer` to point to a model with `tokenizer.json`.
84 +- **usage.completion_tokens parsing**: vLLM sends final usage chunk with `"choices":[]` (empty array). The usage `if` must be separate from the choices `if` (not `else if`).
85 +
86 +## Typical Usage
87 +
88 +```bash
89 +# Embedding benchmark (openai-embeddings, 8 inputs batched per request)
90 +./target/release/vllm-bench \
91 + --backend openai-embeddings \
92 + --base-url http://gb200-10:30000 \
93 + --model BAAI/bge-large-en-v1.5 \
94 + --dataset-name random \
95 + --random-input-len 512 \
96 + --random-batch-size 8 \
97 + --num-prompts 1000 \
98 + --save-result
99 +
100 +# vLLM rerank benchmark (one query + 8 documents per request)
101 +./target/release/vllm-bench \
102 + --backend vllm-rerank \
103 + --base-url http://gb200-10:30000 \
104 + --model BAAI/bge-reranker-v2-m3 \
105 + --dataset-name random-rerank \
106 + --random-input-len 512 \
107 + --random-batch-size 8 \
108 + --num-prompts 500 \
109 + --save-result
110 +
111 +# Prefix-cache stress (10 shared prefixes, 256+256 tokens)
112 +./target/release/vllm-bench \
113 + --backend vllm \
114 + --base-url http://gb200-10:30000 \
115 + --model nvidia/Kimi-K2.5-NVFP4 \
116 + --dataset-name prefix_repetition \
117 + --prefix-repetition-prefix-len 256 \
118 + --prefix-repetition-suffix-len 256 \
119 + --prefix-repetition-num-prefixes 10 \
120 + --num-prompts 1000
121 +
122 +# Custom JSONL workload ({"prompt": ..., "output_tokens": ...} per line)
123 +./target/release/vllm-bench \
124 + --backend openai-chat \
125 + --base-url http://gb200-10:30000 \
126 + --model nvidia/Kimi-K2.5-NVFP4 \
127 + --dataset-name custom \
128 + --dataset-path workload.jsonl \
129 + --custom-output-len -1 \
130 + --num-prompts 1000
131 +
132 +# Random dataset
133 +./target/release/vllm-bench \
134 + --backend vllm \
135 + --base-url http://gb200-10:30000 \
136 + --model nvidia/Kimi-K2.5-NVFP4 \
137 + --dataset-name random \
138 + --random-input-len 8192 \
139 + --random-output-len 1024 \
140 + --ignore-eos \
141 + --num-prompts 4096 \
142 + --percentile-metrics "ttft,tpot,itl,e2el" \
143 + --save-result \
144 + --max-concurrency 1400
145 +
146 +# Random multimodal dataset (VLM benchmark)
147 +./target/release/vllm-bench \
148 + --backend openai-chat \
149 + --base-url http://gb200-10:30000 \
150 + --model Qwen/Qwen2.5-VL-7B-Instruct \
151 + --dataset-name random-mm \
152 + --random-input-len 512 \
153 + --random-output-len 128 \
154 + --num-prompts 100 \
155 + --random-mm-base-items-per-request 1 \
156 + --random-mm-limit-mm-per-prompt '{"image": 1, "video": 0}' \
157 + --random-mm-bucket-config '{(1024, 800, 1): 1.0}'
158 +
159 +# HuggingFace dataset (WildChat)
160 +./target/release/vllm-bench \
161 + --backend openai-chat \
162 + --base-url http://gb200-10:30000 \
163 + --model nvidia/Kimi-K2.5-NVFP4 \
164 + --dataset-name hf \
165 + --dataset-path allenai/WildChat-4.8M \
166 + --hf-split train \
167 + --num-prompts 1000 \
168 + --save-result
169 +
170 +# HuggingFace dataset (LongBench with subset)
171 +./target/release/vllm-bench \
172 + --backend openai-chat \
173 + --base-url http://gb200-10:30000 \
174 + --model nvidia/Kimi-K2.5-NVFP4 \
175 + --dataset-name hf \
176 + --dataset-path THUDM/LongBench \
177 + --hf-subset narrativeqa \
178 + --hf-split test \
179 + --hf-output-len 512 \
180 + --num-prompts 200
181 +```
rust/src/bench/CLAUDE.md new
+1
@@ -0,0 +1 @@
1 +@AGENTS.md
rust/src/bench/Cargo.toml new
+37
@@ -0,0 +1,37 @@
1 +[package]
2 +name = "vllm-bench"
3 +version.workspace = true
4 +edition.workspace = true
5 +description = "High-performance benchmark client for vLLM serving endpoints"
6 +license.workspace = true
7 +
8 +[dependencies]
9 +anyhow.workspace = true
10 +base64.workspace = true
11 +bytes.workspace = true
12 +chrono.workspace = true
13 +clap.workspace = true
14 +dirs.workspace = true
15 +futures.workspace = true
16 +hf-hub.workspace = true
17 +image.workspace = true
18 +indicatif.workspace = true
19 +mimalloc.workspace = true
20 +rand.workspace = true
21 +rand_distr.workspace = true
22 +rayon.workspace = true
23 +reqwest = { workspace = true, features = ["json", "stream", "blocking", "http2"] }
24 +rlimit.workspace = true
25 +rustc-hash.workspace = true
26 +serde = { workspace = true, features = ["rc"] }
27 +serde_json = { workspace = true, features = ["raw_value"] }
28 +thiserror.workspace = true
29 +tiktoken-rs.workspace = true
30 +tokenizers.workspace = true
31 +tokio.workspace = true
32 +tokio-stream.workspace = true
33 +url.workspace = true
34 +uuid.workspace = true
35 +
36 +[lints]
37 +workspace = true
rust/src/bench/README.md new
+810
@@ -0,0 +1,810 @@
1 +# vllm-bench
2 +
3 +High-performance Rust benchmark client for vLLM serving endpoints. A drop-in replacement for `vllm bench serve` with near-instant startup, parallel dataset generation, and a fraction of the memory overhead — and no Python at runtime.
4 +
5 +```bash
6 +vllm-bench --backend vllm --base-url http://127.0.0.1:8000 \
7 + --model <model> --dataset-name random \
8 + --random-input-len 1024 --random-output-len 128 \
9 + --num-prompts 1000 --max-concurrency 200
10 +```
11 +
12 +## Highlights
13 +
14 +- **Fast** — ~7 ms startup, single ~7 MB static binary, no Python imports.
15 +- **Scales** — `Arc<str>` prompt sharing + mimalloc keep memory <100 MB at 1400+ concurrency.
16 +- **Many datasets** — `random`, `random-mm` (VLM), `sharegpt`, `sonnet`, `speed-bench`, and any HuggingFace dataset.
17 +- **Many backends** — completions, chat, embeddings, pooling, and rerank.
18 +- **Beyond a single run** — concurrency/rate **sweeps**, **multi-run** stats, **multi-turn** conversations, **LoRA** multi-adapter, and result **comparison**.
19 +- **Steady-state metrics** — throughput/latency measured over the saturated plateau, excluding ramp-up and drain.
20 +- **Parity** — JSON output schema and timing semantics match Python `vllm bench serve` exactly.
21 +
22 +### Performance vs. Python
23 +
24 +| Metric | Python | Rust |
25 +| -------- | -------- | ------ |
26 +| Startup time | Multi-second (import vllm + numpy + aiohttp) | ~7 ms |
27 +| 100k random prompts (input_len=8192) | Minutes | Seconds (rayon parallelism) |
28 +| Binary size | — | ~7 MB |
29 +| Peak memory at 1400 concurrency | High (GIL + per-object overhead) | <100 MB (`Arc<str>` prompt sharing) |
30 +
31 +## Contents
32 +
33 +- [Install](#install)
34 +- [Quick Start](#quick-start)
35 +- [Usage Examples](#usage-examples)
36 +- [Supported Backends](#supported-backends)
37 +- [Supported Datasets](#supported-datasets)
38 +- [Metrics](#metrics)
39 +- [CLI Reference](#cli-reference)
40 +- [Tokenizer Support](#tokenizer-support)
41 +- [Output Format](#output-format)
42 +- [Architecture](#architecture)
43 +- [Environment Variables](#environment-variables)
44 +
45 +## Install
46 +
47 +### Prebuilt binaries (Linux)
48 +
49 +```bash
50 +curl -fsSL https://github.com/vllm-project/vllm-bench/releases/latest/download/vllm-bench-$(uname -m)-linux-musl -o vllm-bench && chmod +x vllm-bench
51 +```
52 +
53 +### With Cargo
54 +
55 +Install straight from the repository (builds from source; requires [Rust](https://rustup.rs/) stable and a C compiler for the native tokenizer dependency):
56 +
57 +```bash
58 +cargo install --git https://github.com/vllm-project/vllm-bench vllm-bench
59 +```
60 +
61 +The trailing `vllm-bench` selects the package — the repo also ships a `mock-llm-server` binary, so omitting it fails with `multiple packages with binaries found`. The binary is installed to `~/.cargo/bin/`.
62 +
63 +### Build from source
64 +
65 +Requires [Rust](https://rustup.rs/) (stable).
66 +
67 +```bash
68 +git clone https://github.com/vllm-project/vllm-bench.git
69 +cd vllm-bench
70 +./install.sh # builds release and installs to ~/.local/bin
71 +# or: ./install.sh --to ~/bin
72 +```
73 +
74 +## Quick Start
75 +
76 +Point it at a running vLLM server and benchmark with synthetic prompts:
77 +
78 +```bash
79 +vllm-bench \
80 + --backend vllm \
81 + --base-url http://127.0.0.1:8000 \
82 + --model <model-name> \
83 + --dataset-name random \
84 + --random-input-len 1024 \
85 + --random-output-len 128 \
86 + --num-prompts 1000 \
87 + --max-concurrency 200
88 +```
89 +
90 +> **Tip:** prefer `127.0.0.1` over `localhost` — some systems resolve `localhost` to IPv6 `::1` while vLLM listens on IPv4 only.
91 +
92 +Add `--save-result` to write a JSON file, or `--dry-run` to generate and inspect the dataset without sending any requests.
93 +
94 +## Usage Examples
95 +
96 +<details open>
97 +<summary><b>Generation (completions / chat)</b></summary>
98 +
99 +```bash
100 +# Full production-style run with percentile metrics and result file
101 +vllm-bench \
102 + --backend vllm \
103 + --base-url http://127.0.0.1:8000 \
104 + --model nvidia/Kimi-K2.5-NVFP4 \
105 + --dataset-name random \
106 + --random-input-len 8192 \
107 + --random-output-len 1024 \
108 + --ignore-eos \
109 + --num-prompts 4096 \
110 + --percentile-metrics "ttft,tpot,itl,e2el" \
111 + --save-result \
112 + --max-concurrency 1400
113 +
114 +# Send token IDs instead of text (pure vLLM: skips server-side tokenization,
115 +# exact token counts, faster). Random dataset only.
116 +vllm-bench \
117 + --backend vllm \
118 + --base-url http://127.0.0.1:8000 \
119 + --model <model-name> \
120 + --dataset-name random \
121 + --random-input-len 1024 \
122 + --prompt-token-ids \
123 + --num-prompts 1000
124 +```
125 +
126 +</details>
127 +
128 +<details>
129 +<summary><b>Datasets (ShareGPT / Sonnet / HuggingFace / SPEED-Bench)</b></summary>
130 +
131 +```bash
132 +# ShareGPT (auto-downloads from HuggingFace on first run, cached afterwards)
133 +vllm-bench \
134 + --backend openai-chat --base-url http://127.0.0.1:8000 --model <model-name> \
135 + --dataset-name sharegpt --num-prompts 500 --save-result
136 +
137 +# ShareGPT with an explicit local file
138 +vllm-bench \
139 + --backend openai-chat --base-url http://127.0.0.1:8000 --model <model-name> \
140 + --dataset-name sharegpt --dataset-path /path/to/ShareGPT_V3.json \
141 + --num-prompts 500 --save-result
142 +
143 +# Sonnet — built-in Shakespeare sonnets, no dataset file needed.
144 +# Generates prompts of a controllable token length with a shared prefix.
145 +vllm-bench \
146 + --backend openai-chat --base-url http://127.0.0.1:8000 --model <model-name> \
147 + --dataset-name sonnet \
148 + --sonnet-input-len 550 --sonnet-output-len 150 --sonnet-prefix-len 200 \
149 + --num-prompts 500
150 +
151 +# Any public HuggingFace dataset (auto-downloads, auto-detects columns)
152 +vllm-bench \
153 + --backend openai-chat --base-url http://127.0.0.1:8000 --model <model-name> \
154 + --dataset-name hf --dataset-path allenai/WildChat-4.8M \
155 + --hf-split train --num-prompts 1000 --save-result
156 +
157 +# HuggingFace dataset with subset + fixed output length (LongBench)
158 +vllm-bench \
159 + --backend openai-chat --base-url http://127.0.0.1:8000 --model <model-name> \
160 + --dataset-name hf --dataset-path THUDM/LongBench \
161 + --hf-subset narrativeqa --hf-split test --hf-output-len 512 --num-prompts 200
162 +
163 +# Gated HuggingFace dataset (requires HF_TOKEN)
164 +HF_TOKEN=hf_xxx vllm-bench \
165 + --backend openai-chat --base-url http://127.0.0.1:8000 --model <model-name> \
166 + --dataset-name hf --dataset-path lmsys/lmsys-chat-1m \
167 + --hf-split train --hf-output-len 256 --num-prompts 1000
168 +
169 +# SPEED-Bench for speculative decoding evaluation (auto-downloads, cached)
170 +vllm-bench \
171 + --backend openai-chat --base-url http://127.0.0.1:8000 --model <model-name> \
172 + --dataset-name speed-bench --speed-bench-config qualitative \
173 + --num-prompts 200 --output-len 256 --save-result
174 +
175 +# SPEED-Bench throughput split with entropy category filter + input truncation
176 +vllm-bench \
177 + --backend openai-chat --base-url http://127.0.0.1:8000 --model <model-name> \
178 + --dataset-name speed-bench --speed-bench-config throughput_16k \
179 + --speed-bench-max-input-len 10240 --speed-bench-category low_entropy \
180 + --num-prompts 500 --output-len 256 --max-concurrency 200 --save-result
181 +```
182 +
183 +</details>
184 +
185 +<details>
186 +<summary><b>Multimodal (VLM with synthetic images)</b></summary>
187 +
188 +```bash
189 +# One synthetic image per request
190 +vllm-bench \
191 + --backend openai-chat --base-url http://127.0.0.1:8000 \
192 + --model Qwen/Qwen2.5-VL-7B-Instruct \
193 + --dataset-name random-mm \
194 + --random-input-len 512 --random-output-len 128 --num-prompts 100 \
195 + --random-mm-base-items-per-request 1 \
196 + --random-mm-limit-mm-per-prompt '{"image": 1, "video": 0}' \
197 + --random-mm-bucket-config '{(1024, 800, 1): 1.0}'
198 +
199 +# Multiple images per request, mixed resolutions
200 +vllm-bench \
201 + --backend openai-chat --base-url http://127.0.0.1:8000 \
202 + --model Qwen/Qwen2.5-VL-7B-Instruct \
203 + --dataset-name random-mm \
204 + --random-input-len 256 --random-output-len 128 --num-prompts 50 \
205 + --random-mm-base-items-per-request 3 \
206 + --random-mm-limit-mm-per-prompt '{"image": 5, "video": 0}' \
207 + --random-mm-bucket-config '{(256,256,1): 0.5, (720,1280,1): 0.5}'
208 +```
209 +
210 +</details>
211 +
212 +<details>
213 +<summary><b>Embedding / Pooling / Rerank</b></summary>
214 +
215 +```bash
216 +# Text embedding
217 +vllm-bench \
218 + --backend openai-embeddings --base-url http://127.0.0.1:8000 \
219 + --model BAAI/bge-large-en-v1.5 \
220 + --dataset-name random --random-input-len 512 --num-prompts 1000 \
221 + --max-concurrency 200 --save-result
222 +
223 +# Chat-format embedding (supports multimodal content)
224 +vllm-bench \
225 + --backend openai-embeddings-chat --base-url http://127.0.0.1:8000 \
226 + --model BAAI/bge-large-en-v1.5 \
227 + --dataset-name sharegpt --num-prompts 500 --save-result
228 +
229 +# vLLM native pooling endpoint
230 +vllm-bench \
231 + --backend vllm-pooling --base-url http://127.0.0.1:8000 \
232 + --model BAAI/bge-large-en-v1.5 \
233 + --dataset-name random --random-input-len 256 --num-prompts 1000 --save-result
234 +
235 +# Rerank (query from dataset, documents via --extra-body)
236 +vllm-bench \
237 + --backend vllm-rerank --base-url http://127.0.0.1:8000 \
238 + --model BAAI/bge-reranker-v2-m3 \
239 + --dataset-name sharegpt --num-prompts 500 \
240 + --extra-body '{"documents": ["document to rerank"]}' --save-result
241 +```
242 +
243 +</details>
244 +
245 +<details>
246 +<summary><b>Rate control, ramp-up &amp; goodput</b></summary>
247 +
248 +```bash
249 +# Ramp from 10 → 100 RPS with goodput SLO tracking
250 +vllm-bench \
251 + --backend vllm --base-url http://127.0.0.1:8000 --model <model-name> \
252 + --num-prompts 2000 \
253 + --ramp-up-strategy linear --ramp-up-start-rps 10 --ramp-up-end-rps 100 \
254 + --goodput ttft:200 e2el:5000 \
255 + --save-result
256 +
257 +# Fixed Poisson arrival rate at 50 RPS
258 +vllm-bench \
259 + --backend vllm --base-url http://127.0.0.1:8000 --model <model-name> \
260 + --num-prompts 2000 --request-rate 50 --burstiness 1.0
261 +```
262 +
263 +</details>
264 +
265 +<details>
266 +<summary><b>Sweep — find the optimal concurrency / rate</b></summary>
267 +
268 +```bash
269 +# Sweep over concurrency values
270 +vllm-bench \
271 + --backend vllm --base-url http://127.0.0.1:8000 --model <model-name> \
272 + --num-prompts 500 \
273 + --sweep-max-concurrency 1,10,50,100,200,500,1000
274 +
275 +# Sweep over request rates
276 +vllm-bench \
277 + --backend vllm --base-url http://127.0.0.1:8000 --model <model-name> \
278 + --num-prompts 500 \
279 + --sweep-request-rate 1,10,50,100,inf
280 +
281 +# Scale work with concurrency and reset the prefix cache between points
282 +# (--sweep-num-prompts-factor sets num_prompts = concurrency * factor;
283 +# --reset-prefix-cache requires VLLM_SERVER_DEV_MODE=1 on the server)
284 +vllm-bench \
285 + --backend vllm --base-url http://127.0.0.1:8000 --model <model-name> \
286 + --sweep-max-concurrency 1,10,50,100 \
287 + --sweep-num-prompts-factor 20 \
288 + --reset-prefix-cache
289 +```
290 +
291 +</details>
292 +
293 +<details>
294 +<summary><b>Multi-run &amp; comparison</b></summary>
295 +
296 +```bash
297 +# Run 5 times, report mean/std/min/max with coefficient of variation
298 +vllm-bench \
299 + --backend vllm --base-url http://127.0.0.1:8000 --model <model-name> \
300 + --num-prompts 1000 --max-concurrency 200 --num-runs 5
301 +
302 +# Compare two saved result files side-by-side (no server needed)
303 +vllm-bench --compare baseline.json optimized.json
304 +```
305 +
306 +</details>
307 +
308 +<details>
309 +<summary><b>Multi-turn conversations</b></summary>
310 +
311 +```bash
312 +# Synthetic multi-turn (controllable per-turn token lengths)
313 +vllm-bench \
314 + --backend openai-chat --base-url http://127.0.0.1:8000 --model <model-name> \
315 + --dataset-name random --multi-turn --multi-turn-num-turns 5 \
316 + --random-input-len 512 --random-output-len 256 \
317 + --num-prompts 50 --multi-turn-concurrency 10 \
318 + --percentile-metrics "ttft,tpot,itl,e2el" --save-result
319 +
320 +# Variable turn count per conversation + per-turn input length for turns 1+
321 +vllm-bench \
322 + --backend openai-chat --base-url http://127.0.0.1:8000 --model <model-name> \
323 + --dataset-name random --multi-turn \
324 + --multi-turn-min-turns 2 --multi-turn-max-turns 8 \
325 + --random-input-len 2048 --per-turn-input-len 256 --random-output-len 128 \
326 + --num-prompts 100 --multi-turn-concurrency 20
327 +
328 +# ShareGPT conversations (loads all turns, not just the first two)
329 +vllm-bench \
330 + --backend openai-chat --base-url http://127.0.0.1:8000 --model <model-name> \
331 + --dataset-name sharegpt --multi-turn \
332 + --num-prompts 50 --multi-turn-concurrency 10 --save-result
333 +
334 +# Think time between turns
335 +vllm-bench \
336 + --backend openai-chat --base-url http://127.0.0.1:8000 --model <model-name> \
337 + --multi-turn --multi-turn-num-turns 3 --multi-turn-delay-ms 500 \
338 + --num-prompts 100 --multi-turn-concurrency 20
339 +```
340 +
341 +</details>
342 +
343 +<details>
344 +<summary><b>LoRA multi-adapter</b></summary>
345 +
346 +```bash
347 +# Distribute requests across N adapters registered on the server.
348 +# --model stays the BASE model (tokenizer / readiness / /tokenize use it);
349 +# the per-request `model` field is rewritten to one of --lora-modules.
350 +vllm-bench \
351 + --backend openai-chat --base-url http://127.0.0.1:8000 \
352 + --model Qwen/Qwen3-30B-A3B \
353 + --lora-modules sql-lora-1 sql-lora-2 sql-lora-3 sql-lora-4 \
354 + sql-lora-5 sql-lora-6 sql-lora-7 sql-lora-8 \
355 + --lora-assignment random \
356 + --dataset-name random --random-input-len 1024 --random-output-len 256 \
357 + --num-prompts 1000 --max-concurrency 64 --save-result
358 +
359 +# Deterministic round-robin assignment (request i -> adapter[i % N])
360 +vllm-bench \
361 + --backend openai-chat --base-url http://127.0.0.1:8000 \
362 + --model Qwen/Qwen3-30B-A3B \
363 + --lora-modules sql-lora-1 sql-lora-2 sql-lora-3 sql-lora-4 \
364 + --lora-assignment round-robin \
365 + --dataset-name random --num-prompts 1000
366 +```
367 +
368 +Server side — start vLLM with `--enable-lora` and one `name=path` pair per adapter:
369 +
370 +```bash
371 +vllm serve <base-model> \
372 + --enable-lora --max-loras 8 --max-lora-rank 16 \
373 + --lora-modules \
374 + sql-lora-1=jeeejeee/qwen3-moe-text2sql-spider \
375 + sql-lora-2=jeeejeee/qwen3-moe-text2sql-spider \
376 + ...
377 +```
378 +
379 +Set `--max-loras` ≥ number of adapter names to keep them all resident (clean steady-state numbers), or lower to stress the LoRA swap path.
380 +
381 +</details>
382 +
383 +<details>
384 +<summary><b>Profiling &amp; dry-run</b></summary>
385 +
386 +```bash
387 +# Trigger vLLM server-side profiling (start before, stop after the benchmark)
388 +vllm-bench \
389 + --backend vllm --base-url http://127.0.0.1:8000 --model <model-name> \
390 + --num-prompts 100 --profile
391 +
392 +# Defer profiling until the server batch is full, then capture for 10s
393 +vllm-bench \
394 + --backend vllm --base-url http://127.0.0.1:8000 --model <model-name> \
395 + --num-prompts 2000 --max-concurrency 256 \
396 + --profile --profile-batch-threshold 200 --profile-duration 10
397 +
398 +# Dry run: generate dataset, print stats, send nothing
399 +vllm-bench \
400 + --model <model-name> --num-prompts 100000 --random-input-len 8192 --dry-run
401 +```
402 +
403 +</details>
404 +
405 +## Supported Backends
406 +
407 +### Generation
408 +
409 +| Backend | API Endpoint | Description |
410 +| --------- | ------------- | ------------- |
411 +| `vllm` / `openai` | `/v1/completions` | OpenAI-compatible completions (streaming) |
412 +| `openai-chat` | `/v1/chat/completions` | OpenAI-compatible chat completions (streaming, multimodal) |
413 +
414 +### Embedding / Pooling
415 +
416 +| Backend | API Endpoint | Description |
417 +| --------- | ------------- | ------------- |
418 +| `openai-embeddings` | `/v1/embeddings` | Text embedding (accepts text or token IDs) |
419 +| `openai-embeddings-chat` | `/v1/embeddings` | Chat-format embedding (supports multimodal content) |
420 +| `vllm-pooling` | `/v1/pooling` | vLLM native pooling endpoint |
421 +| `vllm-rerank` | `/v1/rerank` | vLLM reranking (query from prompt, documents via `--extra-body`) |
422 +
423 +Pooling backends are non-streaming and report E2EL (end-to-end latency) only. Use `--dataset-name sharegpt`, `sonnet`, or `hf` for text-based embedding/rerank benchmarks, or `random` for token-ID-based embedding benchmarks.
424 +
425 +## Supported Datasets
426 +
427 +| Dataset | Description |
428 +| --------- | ------------- |
429 +| `random` | Synthetic prompts with exact token-length matching (default) |
430 +| `random-mm` | Synthetic multimodal prompts with random JPEG images for VLM benchmarking (requires `openai-chat`) |
431 +| `sharegpt` | Real conversations from ShareGPT (auto-downloads from HuggingFace, or use `--dataset-path`) |
432 +| `sonnet` | Built-in Shakespeare sonnets; controllable token length + shared prefix, no dataset file needed |
433 +| `speed-bench` | NVIDIA SPEED-Bench for speculative decoding evaluation (auto-downloads, 11 categories) |
434 +| `hf` | Any HuggingFace dataset (auto-downloads via datasets-server API, auto-detects chat/text columns) |
435 +
436 +## Metrics
437 +
438 +### Generation backends
439 +
440 +- **TTFT** (Time to First Token) — latency from request send to first token received
441 +- **TPOT** (Time per Output Token) — average time between output tokens
442 +- **ITL** (Inter-Token Latency) — per-token latency distribution
443 +- **E2EL** (End-to-End Latency) — total request latency
444 +- **Throughput** — requests/sec, output tokens/sec, peak output tokens/sec, total tokens/sec
445 +- **Concurrency** — peak concurrent requests
446 +- **Goodput** — requests/sec meeting all specified SLOs (with `--goodput`)
447 +
448 +### Pooling / embedding backends
449 +
450 +- **E2EL** — total request latency (mean, median, std, percentiles)
451 +- **Throughput** — requests/sec, input tokens/sec
452 +- **Concurrency** — peak concurrent requests
453 +
454 +### Steady-state metrics
455 +
456 +When `--max-concurrency` is set and `--request-rate` is `inf` (closed-loop mode), the benchmark automatically reports an additional **Steady-State Metrics** block. It measures throughput and latency only over the window during which in-flight concurrency stays at or above a fraction of `--max-concurrency`, excluding the ramp-up and drain phases. This sharply reduces run-to-run variance at very high concurrency.
457 +
458 +The block reports request/input/output/total token throughput plus TTFT (mean, median, percentiles) and TPOT (mean, median, P90, P99) over the detected plateau, along with the window bounds and how many requests fell inside it. Tune it with `--steady-state-threshold` (default `0.95`) and `--steady-state-min-window`, or disable with `--no-steady-state`. The result JSON carries a `steady_state` object (null when not computed).
459 +
460 +## CLI Reference
461 +
462 +Run `vllm-bench --help` for the authoritative list. Grouped reference below.
463 +
464 +<details>
465 +<summary><b>Server connection</b></summary>
466 +
467 +| Flag | Default | Description |
468 +| ------ | --------- | ------------- |
469 +| `--backend` | `openai` | Backend type (`vllm`, `openai`, `openai-chat`, `openai-embeddings`, `openai-embeddings-chat`, `vllm-pooling`, `vllm-rerank`) |
470 +| `--base-url` | — | Server base URL (overrides `--host`/`--port`) |
471 +| `--host` | `127.0.0.1` | Server host |
472 +| `--port` | `8000` | Server port |
473 +| `--endpoint` | Auto | API endpoint path (auto-selected per backend) |
474 +| `--insecure` | `false` | Disable SSL certificate verification |
475 +
476 +</details>
477 +
478 +<details>
479 +<summary><b>Model &amp; tokenizer</b></summary>
480 +
481 +| Flag | Default | Description |
482 +| ------ | --------- | ------------- |
483 +| `--model` | Auto-detect | Model name (fetched from `/v1/models` if omitted) |
484 +| `--served-model-name` | — | Model name used in API requests |
485 +| `--tokenizer` | Same as model | Tokenizer name or path (supports HF, tiktoken, server fallback) |
486 +| `--tokenizer-mode` | `auto` | Tokenizer mode (`auto`, `hf`, `slow`, `mistral`) |
487 +| `--trust-remote-code` | `false` | Trust remote code for tokenizer |
488 +| `--skip-tokenizer-init` | `false` | Skip tokenizer initialization |
489 +
490 +</details>
491 +
492 +<details>
493 +<summary><b>Dataset</b></summary>
494 +
495 +| Flag | Default | Description |
496 +| ------ | --------- | ------------- |
497 +| `--dataset-name` | `random` | Dataset type (`random`, `random-mm`, `sharegpt`, `sonnet`, `speed-bench`, `hf`) |
498 +| `--dataset-path` | — | Path to dataset file (optional for `sharegpt`/`sonnet`, which auto-source) |
499 +| `--num-prompts` | `1000` | Number of prompts to generate (conversations in multi-turn mode) |
500 +| `--max-model-len` | — | Filter out requests where `prompt_len + output_len` exceeds this context length |
501 +| `--input-len` | — | Override input length (general) |
502 +| `--output-len` | — | Override output length (general) |
503 +| `--no-oversample` | `false` | Don't oversample if dataset is smaller than `--num-prompts` |
504 +| `--disable-shuffle` | `false` | Don't shuffle the dataset |
505 +| `--seed` | `0` | Random seed for reproducibility |
506 +| **Random** | | |
507 +| `--random-input-len` | `1024` | Input token length |
508 +| `--random-output-len` | `128` | Output token length |
509 +| `--random-prefix-len` | `0` | Shared prefix length |
510 +| `--random-range-ratio` | `1.0` | Length jitter, range `(0, 1]`. Lengths sampled from `[ratio × target, target]`; `1.0` = fixed length |
511 +| `--prompt-token-ids` | `false` | Send prompts as token-ID arrays (skips server-side tokenization, exact counts). Random dataset only |
512 +| **Random multimodal** | | |
513 +| `--random-mm-base-items-per-request` | `1` | Base number of multimodal items (images) per request |
514 +| `--random-mm-num-mm-items-range-ratio` | `0.0` | Range ratio for varying item count per request |
515 +| `--random-mm-limit-mm-per-prompt` | `{"image": 255, "video": 1}` | Per-modality hard caps (JSON) |
516 +| `--random-mm-bucket-config` | `{(256,256,1): 0.5, (720,1280,1): 0.5}` | `(height,width,frames)` → probability (Python tuple syntax; frames=1 = image) |
517 +| **ShareGPT** | | |
518 +| `--sharegpt-output-len` | — | Override output length |
519 +| **Sonnet** | | |
520 +| `--sonnet-input-len` | `550` | Input tokens per request |
521 +| `--sonnet-output-len` | `150` | Output tokens per request |
522 +| `--sonnet-prefix-len` | `200` | Prefix tokens shared across requests |
523 +| **SPEED-Bench** | | |
524 +| `--speed-bench-config` | `qualitative` | Split (`qualitative`, `throughput_1k`/`2k`/`8k`/`16k`/`32k`) |
525 +| `--speed-bench-category` | — | Filter by category (`low_entropy`, `high_entropy`, `mixed_entropy`, `coding`, `math`, …) |
526 +| `--speed-bench-max-input-len` | — | Truncate prompts to at most N tokens |
527 +| **HuggingFace** | | |
528 +| `--hf-split` | Auto | Split (`train`, `test`, `validation`); auto-detected if omitted |
529 +| `--hf-subset` | — | Subset/config name (e.g. `narrativeqa` for LongBench) |
530 +| `--hf-output-len` | — | Fixed output length for all requests (overrides dataset-derived length) |
531 +| `--hf-text-column` | Auto | Column containing prompt text; auto-detected from common patterns |
532 +
533 +</details>
534 +
535 +<details>
536 +<summary><b>Rate control</b></summary>
537 +
538 +| Flag | Default | Description |
539 +| ------ | --------- | ------------- |
540 +| `--request-rate` | `inf` | Requests per second (`inf` = all at once) |
541 +| `--burstiness` | `1.0` | Burstiness factor (1.0 = Poisson, >1 = bursty) |
542 +| `--max-concurrency` | `num-prompts` | Maximum concurrent requests (semaphore) |
543 +| `--ramp-up-strategy` | — | Ramp-up mode (`linear` or `exponential`) |
544 +| `--ramp-up-start-rps` | — | Starting request rate for ramp-up |
545 +| `--ramp-up-end-rps` | — | Ending request rate for ramp-up |
546 +
547 +</details>
548 +
549 +<details>
550 +<summary><b>Sampling parameters</b></summary>
551 +
552 +| Flag | Description |
553 +| ------ | ------------- |
554 +| `--temperature` | Temperature (server default if omitted) |
555 +| `--top-p` | Top-p (nucleus) sampling |
556 +| `--top-k` | Top-k sampling |
557 +| `--min-p` | Min-p sampling |
558 +| `--frequency-penalty` | Frequency penalty |
559 +| `--presence-penalty` | Presence penalty |
560 +| `--repetition-penalty` | Repetition penalty |
561 +
562 +Merged into the request body. Only effective with generation backends (`vllm`, `openai`, `openai-chat`); ignored by pooling/embedding backends.
563 +
564 +</details>
565 +
566 +<details>
567 +<summary><b>Output &amp; results</b></summary>
568 +
569 +| Flag | Default | Description |
570 +| ------ | --------- | ------------- |
571 +| `--save-result` | `false` | Save results to JSON file |
572 +| `--save-detailed` | `false` | Include per-request data in JSON (input/output lens, ITLs, texts) |
573 +| `--append-result` | `false` | Append to existing JSON file (JSONL format) |
574 +| `--result-dir` | — | Directory for result files |
575 +| `--result-filename` | Auto | Custom result filename |
576 +| `--percentile-metrics` | `ttft,tpot,itl,e2el` | Metrics for percentile reporting (pooling defaults to `e2el` only) |
577 +| `--metric-percentiles` | `99` | Percentile values to compute |
578 +| `--sweep-summary-percentiles` | — | Extra percentiles for sweep summary tables (auto-added to computed set) |
579 +| `--goodput` | — | SLO pairs for goodput (`ttft:100 tpot:50 e2el:500`, values in ms) |
580 +| `--disable-tqdm` | `false` | Disable progress bar |
581 +| `--label` | — | Label prefix for result files |
582 +| `--metadata` | — | Key-value metadata (`KEY=VALUE`, repeatable) |
583 +
584 +</details>
585 +
586 +<details>
587 +<summary><b>Request options</b></summary>
588 +
589 +| Flag | Default | Description |
590 +| ------ | --------- | ------------- |
591 +| `--ignore-eos` | `false` | Ignore EOS token (force full output length) |
592 +| `--logprobs` | — | Number of logprobs per token |
593 +| `--num-warmups` | `0` | Warmup requests before benchmarking |
594 +| `--ready-check-timeout-sec` | `0` | Endpoint readiness timeout (0 = skip) |
595 +| `--request-id-prefix` | Auto (UUID) | Prefix for request IDs |
596 +| `--header` | — | Extra headers (`KEY=VALUE`, repeatable) |
597 +| `--extra-body` | — | Extra JSON body parameters |
598 +| `--dry-run` | `false` | Generate dataset only, skip benchmark |
599 +
600 +</details>
601 +
602 +<details>
603 +<summary><b>Steady-state metrics</b></summary>
604 +
605 +| Flag | Default | Description |
606 +| ------ | --------- | ------------- |
607 +| `--steady-state-threshold` | `0.95` | Fraction of `--max-concurrency` at which the steady-state window opens, range (0, 1] |
608 +| `--steady-state-min-window` | Auto | Minimum window duration (s) below which a warning is attached. Default `max(10, 0.1 × run_duration)` |
609 +| `--no-steady-state` | `false` | Disable steady-state metrics computation |
610 +
611 +Computed only when `--max-concurrency` is set and `--request-rate` is `inf`.
612 +
613 +</details>
614 +
615 +<details>
616 +<summary><b>Profiling</b></summary>
617 +
618 +| Flag | Default | Description |
619 +| ------ | --------- | ------------- |
620 +| `--profile` | `false` | Trigger vLLM server-side profiling (`/start_profile` before, `/stop_profile` after) |
621 +| `--profile-batch-threshold` | — | Defer profiling until `/metrics` reports ≥ N running requests, then capture. Requires `--profile` |
622 +| `--profile-duration` | `5.0` | Seconds to capture once the batch threshold is reached. Requires `--profile-batch-threshold` |
623 +
624 +</details>
625 +
626 +<details>
627 +<summary><b>Sweep mode</b></summary>
628 +
629 +| Flag | Default | Description |
630 +| ------ | --------- | ------------- |
631 +| `--sweep-max-concurrency` | — | Comma-separated concurrency values to sweep (e.g. `1,10,50,100,500`) |
632 +| `--sweep-request-rate` | — | Comma-separated rate values to sweep, supports `inf` (e.g. `1,10,100,inf`) |
633 +| `--sweep-num-prompts-factor` | — | Set `num_prompts = concurrency × factor` per concurrency sweep point |
634 +| `--reset-prefix-cache` | `false` | Reset the server's prefix cache before each sweep iteration (requires `VLLM_SERVER_DEV_MODE=1`) |
635 +
636 +Runs the benchmark once per value, then prints a summary table comparing throughput and latency across all sweep points and identifies the best-throughput configuration. Works in multi-turn mode too. `--sweep-summary-percentiles` appends extra TTFT/TPOT/E2EL columns to the summary, auto-adding any missing percentiles to the computed set so they also appear in result JSON.
637 +
638 +</details>
639 +
640 +<details>
641 +<summary><b>Multi-turn conversation benchmark</b></summary>
642 +
643 +| Flag | Default | Description |
644 +| ------ | --------- | ------------- |
645 +| `--multi-turn` | `false` | Enable multi-turn conversation mode (requires `--backend openai-chat`) |
646 +| `--multi-turn-num-turns` | `3` | Turns per conversation (synthetic mode) |
647 +| `--multi-turn-min-turns` | `0` | Minimum turns per conversation (0 = use `--multi-turn-num-turns`) |
648 +| `--multi-turn-max-turns` | `0` | Maximum turns per conversation (0 = `--multi-turn-num-turns` synthetic / uncapped ShareGPT) |
649 +| `--multi-turn-concurrency` | — | Concurrent conversations (defaults to `--max-concurrency` or `--num-prompts`) |
650 +| `--multi-turn-delay-ms` | `0` | Delay between turns in ms (simulates user think time) |
651 +| `--per-turn-input-len` | `0` | Input token length for turns 1+ (0 = use `--random-input-len` for all turns) |
652 +| `--multi-turn-prefix-global-ratio` | `0.0` | Fraction of per-turn input shared across all conversations (random dataset only) |
653 +| `--multi-turn-prefix-conversation-ratio` | `0.0` | Fraction shared within each conversation (random dataset only) |
654 +
655 +With `--multi-turn`, `--num-prompts` controls the number of **conversations**, not individual requests.
656 +
657 +**How it works:**
658 +
659 +- Turn 1: send `[user_1]`, get `assistant_1`
660 +- Turn 2: send `[user_1, assistant_1, user_2]`, get `assistant_2`
661 +- Turn N: send full history + `user_N` — measures growing-context performance
662 +
663 +**Data sources:**
664 +
665 +- `--dataset-name random` — synthetic conversations with controllable per-turn token lengths. Auto-sets `min_tokens` to enforce output length without `ignore_eos`.
666 +- `--dataset-name sharegpt` — loads all turns (not just the first two); filters for entries with ≥ 2 real turns.
667 +
668 +**Prefix sharing** (random dataset): when `--multi-turn-prefix-global-ratio` or `--multi-turn-prefix-conversation-ratio` is > 0, each turn sends a fixed-length message (no history accumulation) composed of a global prefix + per-conversation prefix + unique suffix. The two ratios must sum to < 1.0.
669 +
670 +**Router affinity:** every turn sends `X-Session-ID: {conversation_id}` for KV-cache reuse behind a vLLM router.
671 +
672 +**Output:** overall metrics plus a per-turn breakdown (TTFT/TPOT/ITL/E2EL by turn index). Expect TTFT to climb across turns due to growing context. JSON includes a `per_turn_metrics` array.
673 +
674 +</details>
675 +
676 +<details>
677 +<summary><b>LoRA multi-adapter</b></summary>
678 +
679 +| Flag | Default | Description |
680 +| ------ | --------- | ------------- |
681 +| `--lora-modules` | — | Adapter names registered on the server (`vllm serve --lora-modules name=path`). Each request's `model` field is rewritten to one of these. Repeatable |
682 +| `--lora-assignment` | `random` | Distribution: `random` (uniform, seeded by `--seed`) or `round-robin` (deterministic `i % N`) |
683 +
684 +`--model` must stay the **base** model — its tokenizer builds prompts, and `/v1/models`, `/tokenize`, ready check, and warmup all use it. Only the per-request `model` field in completions/chat payloads is rewritten to the assigned adapter (vLLM routes by name).
685 +
686 +**Assignment scope:** per request in single-shot mode; **per conversation** (sticky across all turns) in multi-turn mode, to avoid breaking prefix-cache reuse mid-dialog.
687 +
688 +**Reproducibility:** with `--lora-assignment random`, the same `--seed` + same `--lora-modules` list yields identical request-to-adapter mappings. Pooling/embedding backends are rejected — LoRA routing applies to generative paths only.
689 +
690 +</details>
691 +
692 +<details>
693 +<summary><b>Multi-run &amp; comparison</b></summary>
694 +
695 +| Flag | Default | Description |
696 +| ------ | --------- | ------------- |
697 +| `--num-runs` | `1` | Run benchmark N times; report mean/std/min/max with CV |
698 +| `--compare` | — | Compare two result JSON files side-by-side (skips benchmarking) |
699 +
700 +`--num-runs` aggregates metrics across runs and reports the coefficient of variation (CV) for throughput stability. `--compare` reads two previously-saved result files and prints a diff with delta, % change, and improvement/regression markers.
701 +
702 +</details>
703 +
704 +## Tokenizer Support
705 +
706 +Tokenizers are loaded with a three-tier fallback chain:
707 +
708 +1. **Local HuggingFace** — `tokenizer.json` from a local path or the Hub (fastest)
709 +2. **Tiktoken** — `.tiktoken` / `.model` format for Kimi, Qwen, etc. (auto-extracts `pat_str` from Python source)
710 +3. **Server-side** — falls back to vLLM's `/tokenize` + `/detokenize` endpoints
711 +
712 +For the `random` dataset, prompt token lengths are verified against the server on the first run and cached; subsequent runs with the same model+server skip verification. Verification is also skipped when `--prompt-token-ids` is set (token counts are exact by construction).
713 +
714 +Models without `tokenizer.json` (e.g. `nvidia/Kimi-K2.5-NVFP4`) fall back to server-side tokenization automatically; you can also point `--tokenizer` at a model that ships `tokenizer.json`.
715 +
716 +## Output Format
717 +
718 +JSON output is compatible with the `vllm bench serve` Python schema. Result files are named:
719 +
720 +```text
721 +{label}-{rate}qps-concurrency{max_concurrency}-{model}-{timestamp}.json
722 +```
723 +
724 +Use `--append-result` to append multiple runs to the same file in JSONL format. `--save-detailed` adds per-request arrays (input/output lengths, ITLs, generated text).
725 +
726 +## Architecture
727 +
728 +<details>
729 +<summary><b>Source layout</b></summary>
730 +
731 +```text
732 +src/
733 +├── main.rs # Entry point, mimalloc, tokio runtime, mode dispatch
734 +├── cli.rs # clap CLI argument definitions
735 +├── config.rs # Validated config, goodput/ramp-up parsing
736 +├── benchmark.rs # Core orchestrator (schedule, spawn, collect, verify, profile)
737 +├── multi_turn.rs # Multi-turn conversation orchestrator (channel workers)
738 +├── compare.rs # Result diff (--compare file_a.json file_b.json)
739 +├── sweep.rs # Parameter sweep (--sweep-max-concurrency, --sweep-request-rate)
740 +├── multi_run.rs # Multi-run statistics (--num-runs N)
741 +├── rate_control.rs # Gamma/Poisson scheduling + linear/exponential ramp-up
742 +├── ready_checker.rs # Endpoint readiness with retry
743 +├── tokenizer.rs # Tokenizer abstraction (HF, tiktoken, server)
744 +├── tiktoken.rs # Tiktoken BPE loader with pat_str extraction
745 +├── error.rs # Error types
746 +├── backends/
747 +│ ├── mod.rs # Backend enum dispatch, typed SSE structs
748 +│ ├── streaming.rs # SSE stream parser with speculative JSON parse
749 +│ ├── openai_completions.rs # /v1/completions backend
750 +│ ├── openai_chat.rs # /v1/chat/completions backend
751 +│ └── pooling.rs # Embedding/pooling/rerank backends (non-streaming)
752 +├── datasets/
753 +│ ├── mod.rs # SampleRequest, ConversationTurn, MultiTurnConversation types
754 +│ ├── random.rs # Random dataset with rayon parallelism
755 +│ ├── random_mm.rs # Random multimodal dataset (JPEG generation, bucket sampling)
756 +│ ├── multi_turn.rs # Multi-turn synthetic + ShareGPT conversation generators
757 +│ ├── sharegpt.rs # ShareGPT JSON dataset loader
758 +│ ├── sonnet.rs # Sonnet dataset (built-in Shakespeare sonnets)
759 +│ ├── speed_bench.rs # NVIDIA SPEED-Bench loader (auto-download + cache)
760 +│ └── hf_dataset.rs # Generic HuggingFace dataset (auto-download, column detection)
761 +├── metrics/
762 +│ ├── mod.rs # BenchmarkMetrics, MultiTurnMetrics structs
763 +│ ├── calculator.rs # Percentile/throughput/goodput/peak/multi-turn computation
764 +│ └── steady_state.rs # Steady-state window detection + plateau metrics
765 +└── output/
766 + ├── mod.rs
767 + ├── console.rs # Terminal output (matches Python format)
768 + └── json.rs # JSON result serialization (Python-compatible schema)
769 +```
770 +
771 +</details>
772 +
773 +### Key design decisions
774 +
775 +- **reqwest + tokio** — HTTP client with connection pooling, forced HTTP/1.1, TCP_NODELAY to match Python's aiohttp and avoid Nagle latency inflation on TTFT
776 +- **mimalloc** — global allocator to reduce contention under high concurrency (1400+ tasks); page-agnostic, runs on aarch64 4K- and 64K-page kernels
777 +- **`Arc<str>` prompts** — zero-copy prompt sharing across tokio tasks, eliminating ~3 GB peak memory at 100k requests with 8k-token prompts
778 +- **Spawn-per-request** — `tokio::spawn` per request with a `Semaphore` for concurrency control (matches Python's asyncio pattern)
779 +- **rayon** — parallel dataset generation across CPU cores (200–500× faster than Python for 100k+ prompts)
780 +- **Enum dispatch** — backend variants instead of trait objects (avoids async trait-object limitations)
781 +- **Typed SSE deserialization** — `CompletionChunk`/`ChatChunk` structs skip unused JSON fields (cheaper than `serde_json::Value`)
782 +- **Speculative JSON parse** — SSE handler uses `serde_json::value::RawValue` to detect complete JSON before `\n\n` arrives, improving TTFT/ITL accuracy when TCP segments split
783 +- **Connection error retry** — automatic retry with backoff on connection reset/timeout/refused (up to 3 attempts)
784 +- **Tokenizer verification cache** — server-side token-length verification is cached per model+server pair
785 +
786 +### Behavioral parity with Python
787 +
788 +The Rust implementation matches Python `vllm bench serve` in:
789 +
790 +- SSE streaming protocol handling (including speculative parse for split TCP segments)
791 +- Timing semantics (monotonic `Instant` matching Python's `time.perf_counter()`)
792 +- Chat vs. completions differences (`max_completion_tokens` vs. `max_tokens`, Content-Type, timestamp placement)
793 +- JSON output schema (all fields, key naming, `request_rate` as the string `"inf"`)
794 +- Rate control (Gamma distribution, normalization, burstiness, linear/exponential ramp-up)
795 +- Metrics (TTFT/TPOT/ITL/E2EL percentiles, peak tokens/sec, peak concurrency, goodput)
796 +- Sampling parameters merged into the request body via `extra_body` (same precedence rules)
797 +
798 +## Environment Variables
799 +
800 +| Variable | Description |
801 +| ---------- | ------------- |
802 +| `OPENAI_API_KEY` | API key for authenticated endpoints (cached, not read per-request) |
803 +| `HF_TOKEN` | HuggingFace token for gated model tokenizers and gated datasets |
804 +| `TOKIO_WORKER_THREADS` | Override tokio worker thread count (default: physical cores) |
805 +
806 +## License
807 +
808 +Apache-2.0
809 +</content>
810 +</invoke>
rust/src/bench/src/backends/mod.rs new
+213
@@ -0,0 +1,213 @@
1 +// SPDX-License-Identifier: Apache-2.0
2 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +pub mod openai_chat;
5 +pub mod openai_completions;
6 +pub mod pooling;
7 +pub mod streaming;
8 +
9 +use std::collections::HashMap;
10 +use std::sync::Arc;
11 +
12 +use serde::{Deserialize, Serialize};
13 +
14 +// --- Typed SSE chunk structs for zero-alloc deserialization ---
15 +// Using typed deserialization avoids building a full serde_json::Value tree.
16 +// Only the fields we need are extracted; everything else is skipped by serde.
17 +
18 +/// Completions API streaming chunk (minimal fields).
19 +#[derive(Deserialize)]
20 +pub struct CompletionChunk {
21 + #[serde(default)]
22 + pub choices: Vec<CompletionChoice>,
23 + pub usage: Option<ChunkUsage>,
24 +}
25 +
26 +#[derive(Deserialize)]
27 +pub struct CompletionChoice {
28 + pub text: Option<String>,
29 +}
30 +
31 +/// Chat API streaming chunk (minimal fields).
32 +#[derive(Deserialize)]
33 +pub struct ChatChunk {
34 + #[serde(default)]
35 + pub choices: Vec<ChatChoice>,
36 + pub usage: Option<ChunkUsage>,
37 +}
38 +
39 +#[derive(Deserialize)]
40 +pub struct ChatChoice {
41 + pub delta: Option<ChatDelta>,
42 +}
43 +
44 +#[derive(Deserialize)]
45 +pub struct ChatDelta {
46 + pub content: Option<String>,
47 +}
48 +
49 +#[derive(Deserialize)]
50 +pub struct ChunkUsage {
51 + pub completion_tokens: Option<u64>,
52 +}
53 +
54 +use crate::cli::BackendKind;
55 +use crate::error::Result;
56 +
57 +/// Input for a single benchmark request.
58 +#[derive(Debug, Clone)]
59 +pub struct RequestFuncInput {
60 + pub prompt: Arc<str>,
61 + pub api_url: String,
62 + pub prompt_len: usize,
63 + pub output_len: usize,
64 + pub model: String,
65 + pub model_name: Option<String>,
66 + pub logprobs: Option<usize>,
67 + pub extra_headers: Option<HashMap<String, String>>,
68 + pub extra_body: Option<serde_json::Value>,
69 + pub ignore_eos: bool,
70 + pub request_id: Option<String>,
71 + /// Pre-built messages array for multi-turn conversations.
72 + /// When set, the chat backend uses this instead of building from `prompt`.
73 + pub messages: Option<serde_json::Value>,
74 + /// Pre-computed token IDs for this prompt.
75 + /// When set, the completions backend sends these directly via `prompt_token_ids`
76 + /// instead of the text `prompt`, skipping server-side tokenization.
77 + pub prompt_token_ids: Option<Arc<[u32]>>,
78 + /// Multimodal content as pre-serialized JSON fragments.
79 + /// When set, the chat backend concatenates these directly into the payload bytes,
80 + /// avoiding any parsing or deep-cloning of base64 image data.
81 + pub multi_modal_content: Option<Arc<[Arc<str>]>>,
82 + /// Complete pre-serialized chat `messages` array (--enable-multimodal-chat).
83 + /// When set, the chat backend splices it verbatim into the payload bytes,
84 + /// taking precedence over `messages`, `prompt`, and `multi_modal_content`.
85 + pub chat_messages_json: Option<Arc<str>>,
86 + /// Multiple text inputs for one request (pooling backends only):
87 + /// embeddings batch (`"input": [...]`) or rerank query+documents.
88 + pub prompt_list: Option<Arc<[Arc<str>]>>,
89 +}
90 +
91 +/// Output from a single benchmark request including timing metrics.
92 +#[derive(Debug, Clone, Serialize, Deserialize)]
93 +pub struct RequestFuncOutput {
94 + pub generated_text: String,
95 + pub success: bool,
96 + pub latency: f64,
97 + pub output_tokens: usize,
98 + pub ttft: f64,
99 + pub itl: Vec<f64>,
100 + pub tpot: f64,
101 + pub prompt_len: usize,
102 + pub error: String,
103 + pub start_time: f64,
104 +}
105 +
106 +impl Default for RequestFuncOutput {
107 + fn default() -> Self {
108 + Self {
109 + generated_text: String::new(),
110 + success: false,
111 + latency: 0.0,
112 + output_tokens: 0,
113 + ttft: 0.0,
114 + itl: Vec::new(),
115 + tpot: 0.0,
116 + prompt_len: 0,
117 + error: String::new(),
118 + start_time: 0.0,
119 + }
120 + }
121 +}
122 +
123 +impl Default for RequestFuncInput {
124 + fn default() -> Self {
125 + Self {
126 + prompt: Arc::from(""),
127 + api_url: String::new(),
128 + prompt_len: 0,
129 + output_len: 0,
130 + model: String::new(),
131 + model_name: None,
132 + logprobs: None,
133 + extra_headers: None,
134 + extra_body: None,
135 + ignore_eos: false,
136 + request_id: None,
137 + messages: None,
138 + prompt_token_ids: None,
139 + multi_modal_content: None,
140 + chat_messages_json: None,
141 + prompt_list: None,
142 + }
143 + }
144 +}
145 +
146 +/// Enum dispatch for backend implementations (avoids async trait object issues).
147 +#[derive(Clone)]
148 +pub enum Backend {
149 + OpenAICompletions(openai_completions::OpenAICompletionsBackend),
150 + OpenAIChat(openai_chat::OpenAIChatBackend),
151 + Pooling(pooling::PoolingBackend),
152 +}
153 +
154 +impl Backend {
155 + /// Send a single request and collect timing metrics.
156 + pub async fn send_request(
157 + &self,
158 + input: &RequestFuncInput,
159 + client: &reqwest::Client,
160 + ) -> Result<RequestFuncOutput> {
161 + match self {
162 + Backend::OpenAICompletions(b) => b.send_request(input, client).await,
163 + Backend::OpenAIChat(b) => b.send_request(input, client).await,
164 + Backend::Pooling(b) => b.send_request(input, client).await,
165 + }
166 + }
167 +}
168 +
169 +/// Get a backend by kind.
170 +pub fn get_backend(kind: BackendKind) -> Result<Backend> {
171 + match kind {
172 + BackendKind::Vllm | BackendKind::Openai => Ok(Backend::OpenAICompletions(
173 + openai_completions::OpenAICompletionsBackend,
174 + )),
175 + BackendKind::OpenaiChat => Ok(Backend::OpenAIChat(openai_chat::OpenAIChatBackend)),
176 + kind if kind.is_pooling() => Ok(Backend::Pooling(pooling::PoolingBackend { kind })),
177 + _ => unreachable!(),
178 + }
179 +}
180 +
181 +/// Cached API key to avoid per-request env var syscall.
182 +static API_KEY: std::sync::OnceLock<Option<String>> = std::sync::OnceLock::new();
183 +
184 +fn cached_api_key() -> &'static Option<String> {
185 + API_KEY.get_or_init(|| std::env::var("OPENAI_API_KEY").ok())
186 +}
187 +
188 +/// Build common headers including auth and extras.
189 +pub fn build_headers(
190 + content_type: Option<&str>,
191 + extra_headers: &Option<HashMap<String, String>>,
192 + request_id: &Option<String>,
193 +) -> HashMap<String, String> {
194 + let mut headers = HashMap::new();
195 +
196 + if let Some(ct) = content_type {
197 + headers.insert("Content-Type".to_string(), ct.to_string());
198 + }
199 +
200 + if let Some(api_key) = cached_api_key() {
201 + headers.insert("Authorization".to_string(), format!("Bearer {api_key}"));
202 + }
203 +
204 + if let Some(extra) = extra_headers {
205 + headers.extend(extra.iter().map(|(k, v)| (k.clone(), v.clone())));
206 + }
207 +
208 + if let Some(rid) = request_id {
209 + headers.insert("x-request-id".to_string(), rid.clone());
210 + }
211 +
212 + headers
213 +}
rust/src/bench/src/backends/openai_chat.rs new
+376
@@ -0,0 +1,376 @@
1 +// SPDX-License-Identifier: Apache-2.0
2 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +use std::time::Instant;
5 +
6 +use futures::StreamExt;
7 +
8 +use super::streaming::{StreamedResponseHandler, trim_bytes};
9 +use super::{ChatChunk, RequestFuncInput, RequestFuncOutput, build_headers};
10 +use crate::error::Result;
11 +
12 +/// Backend for OpenAI Chat Completions API (/v1/chat/completions).
13 +#[derive(Clone)]
14 +pub struct OpenAIChatBackend;
15 +
16 +impl OpenAIChatBackend {
17 + pub async fn send_request(
18 + &self,
19 + input: &RequestFuncInput,
20 + client: &reqwest::Client,
21 + ) -> Result<RequestFuncOutput> {
22 + // Content-Type is set below by `.json()` / `.header()`; keep it out of
23 + // headers_map to avoid a duplicate that strict gateways reject.
24 + let headers_map = build_headers(None, &input.extra_headers, &input.request_id);
25 +
26 + let mut output = RequestFuncOutput {
27 + prompt_len: input.prompt_len,
28 + itl: Vec::with_capacity(input.output_len.max(1)),
29 + ..Default::default()
30 + };
31 +
32 + let st = Instant::now();
33 +
34 + let mut most_recent_timestamp = st;
35 + let mut generated_text = String::new();
36 + let mut first_token_received = false;
37 +
38 + // Build request: use zero-copy raw JSON for multimodal, serde_json for text-only
39 + let mut request =
40 + if input.multi_modal_content.is_some() || input.chat_messages_json.is_some() {
41 + let payload_bytes = build_mm_payload(input);
42 + client
43 + .post(&input.api_url)
44 + .header("content-type", "application/json")
45 + .body(payload_bytes)
46 + } else {
47 + let payload = build_text_payload(input);
48 + client.post(&input.api_url).json(&payload)
49 + };
50 + for (k, v) in &headers_map {
51 + request = request.header(k, v);
52 + }
53 +
54 + match request.send().await {
55 + Ok(response) => {
56 + if response.status().is_success() {
57 + let mut handler = StreamedResponseHandler::new();
58 + let mut stream = response.bytes_stream();
59 +
60 + while let Some(chunk_result) = stream.next().await {
61 + let chunk_bytes = match chunk_result {
62 + Ok(b) => b,
63 + Err(e) => {
64 + output.success = false;
65 + output.error = format!("Stream error: {e}");
66 + return Ok(output);
67 + }
68 + };
69 +
70 + let trimmed_bytes = trim_bytes(&chunk_bytes);
71 + if trimmed_bytes.is_empty() {
72 + continue;
73 + }
74 +
75 + let messages = handler.add_chunk(trimmed_bytes);
76 + for message in messages {
77 + // Skip SSE comments
78 + if message.starts_with(':') {
79 + continue;
80 + }
81 +
82 + // Handle multi-field SSE events (e.g., Dynamo sends
83 + // "event: message\ndata: {...}"). Extract the data: line.
84 + let raw = if message.contains('\n') {
85 + match message.lines().find(|l| l.starts_with("data: ")) {
86 + Some(l) => l,
87 + None => continue,
88 + }
89 + } else {
90 + message.as_str()
91 + };
92 +
93 + let chunk = raw.strip_prefix("data: ").unwrap_or(raw);
94 +
95 + if chunk == "[DONE]" {
96 + continue;
97 + }
98 +
99 + // Python chat backend: timestamp is captured for ALL
100 + // non-DONE messages, and most_recent_timestamp is updated
101 + // unconditionally (outside `if choices:`). This differs from
102 + // completions which only timestamps content chunks.
103 + let timestamp = Instant::now();
104 +
105 + let data: ChatChunk = match serde_json::from_str(chunk) {
106 + Ok(d) => d,
107 + Err(_) => continue,
108 + };
109 +
110 + if !data.choices.is_empty() {
111 + let content = data.choices[0]
112 + .delta
113 + .as_ref()
114 + .and_then(|d| d.content.as_deref())
115 + .unwrap_or("");
116 +
117 + if !first_token_received {
118 + first_token_received = true;
119 + output.ttft = timestamp.duration_since(st).as_secs_f64();
120 + } else {
121 + output.itl.push(
122 + timestamp
123 + .duration_since(most_recent_timestamp)
124 + .as_secs_f64(),
125 + );
126 + }
127 +
128 + generated_text.push_str(content);
129 + }
130 + // Separate `if` (not `else if`) — Dynamo may send
131 + // both choices and usage in the same chunk.
132 + if let Some(ref usage) = data.usage
133 + && let Some(ct) = usage.completion_tokens
134 + {
135 + output.output_tokens = ct as usize;
136 + }
137 +
138 + most_recent_timestamp = timestamp;
139 + }
140 + }
141 +
142 + output.generated_text = generated_text;
143 + output.success = true;
144 + output.latency = most_recent_timestamp.duration_since(st).as_secs_f64();
145 + } else {
146 + let status = response.status();
147 + let body = response.text().await.unwrap_or_default();
148 + output.error = if body.is_empty() {
149 + format!("HTTP {status}")
150 + } else {
151 + format!("HTTP {status}: {body}")
152 + };
153 + output.success = false;
154 + }
155 + }
156 + Err(e) => {
157 + output.success = false;
158 + output.error = format!("{e:#}");
159 + }
160 + }
161 +
162 + Ok(output)
163 + }
164 +}
165 +
166 +/// Build a JSON payload for text-only (non-multimodal) requests using serde_json.
167 +fn build_text_payload(input: &RequestFuncInput) -> serde_json::Value {
168 + let model = input.model_name.as_deref().unwrap_or(&input.model);
169 +
170 + let messages = if let Some(ref msgs) = input.messages {
171 + msgs.clone()
172 + } else {
173 + let content = serde_json::json!([
174 + {"type": "text", "text": input.prompt}
175 + ]);
176 + serde_json::json!([{"role": "user", "content": content}])
177 + };
178 +
179 + let mut payload = serde_json::json!({
180 + "model": model,
181 + "messages": messages,
182 + "max_completion_tokens": input.output_len,
183 + "stream": true,
184 + "stream_options": {
185 + "include_usage": true,
186 + },
187 + });
188 +
189 + if input.ignore_eos {
190 + payload["ignore_eos"] = serde_json::json!(true);
191 + }
192 + if let Some(serde_json::Value::Object(map)) = input.extra_body.as_ref() {
193 + for (k, v) in map {
194 + payload[k] = v.clone();
195 + }
196 + }
197 +
198 + payload
199 +}
200 +
201 +/// Build the JSON payload as raw bytes for multimodal requests.
202 +///
203 +/// This is the zero-copy fast path: pre-serialized mm content fragments
204 +/// (each ~200KB+ of base64 image data) are concatenated directly into the
205 +/// output buffer without being parsed, cloned, or re-serialized.
206 +///
207 +/// Saves ~200KB of allocation + copy per image per request compared to
208 +/// the serde_json::Value approach.
209 +fn build_mm_payload(input: &RequestFuncInput) -> Vec<u8> {
210 + let model = input.model_name.as_deref().unwrap_or(&input.model);
211 +
212 + // Estimate total size: JSON overhead (~300 bytes) + prompt + mm fragments
213 + let mm_total: usize = input
214 + .multi_modal_content
215 + .as_ref()
216 + .map(|mm| mm.iter().map(|f| f.len() + 1).sum())
217 + .unwrap_or(0)
218 + + input.chat_messages_json.as_ref().map_or(0, |m| m.len());
219 + let estimated = 512 + input.prompt.len() * 2 + mm_total;
220 + let mut json = String::with_capacity(estimated);
221 +
222 + // {"model": <model>
223 + json.push_str(r#"{"model":"#);
224 + // serde_json::to_string on &str produces a JSON-escaped quoted string
225 + json.push_str(&serde_json::to_string(model).unwrap());
226 +
227 + json.push_str(r#","messages":"#);
228 + if let Some(ref msgs) = input.chat_messages_json {
229 + // --enable-multimodal-chat: the dataset pre-built the full messages
230 + // array (text + mm parts); splice it verbatim.
231 + json.push_str(msgs);
232 + } else {
233 + let mm = input.multi_modal_content.as_ref().unwrap();
234 +
235 + // [{"role":"user","content":[ <text part>
236 + json.push_str(r#"[{"role":"user","content":[{"type":"text","text":""#);
237 + // JSON-escape the prompt text (handles \n, \t, unicode, quotes)
238 + push_json_escaped_str(&mut json, &input.prompt);
239 + json.push_str(r#""}"#);
240 +
241 + // ,<mm fragment 1>,<mm fragment 2>,...
242 + for fragment in mm.iter() {
243 + json.push(',');
244 + json.push_str(fragment);
245 + }
246 +
247 + // Close content, message, messages
248 + json.push_str(r#"]}]"#);
249 + }
250 +
251 + // ,"max_completion_tokens": N, "stream": true, ...
252 + json.push_str(r##","max_completion_tokens":"##);
253 + json.push_str(&input.output_len.to_string());
254 + json.push_str(r##","stream":true,"stream_options":{"include_usage":true}"##);
255 +
256 + if input.ignore_eos {
257 + json.push_str(r#","ignore_eos":true"#);
258 + }
259 +
260 + // Merge extra_body key-value pairs, skipping keys already set above
261 + if let Some(serde_json::Value::Object(map)) = input.extra_body.as_ref() {
262 + for (k, v) in map {
263 + match k.as_str() {
264 + "model"
265 + | "messages"
266 + | "max_completion_tokens"
267 + | "stream"
268 + | "stream_options"
269 + | "ignore_eos" => continue,
270 + _ => {
271 + json.push(',');
272 + json.push_str(&serde_json::to_string(k).unwrap());
273 + json.push(':');
274 + json.push_str(&serde_json::to_string(v).unwrap());
275 + }
276 + }
277 + }
278 + }
279 +
280 + json.push('}');
281 + json.into_bytes()
282 +}
283 +
284 +/// Write a JSON-escaped string (without surrounding quotes) into the buffer.
285 +///
286 +/// Handles: `\n`, `\r`, `\t`, `\\`, `\"`, and control characters.
287 +/// This avoids the allocation of `serde_json::to_string` which produces
288 +/// a new String with surrounding quotes.
289 +fn push_json_escaped_str(buf: &mut String, s: &str) {
290 + use std::fmt::Write;
291 + for ch in s.chars() {
292 + match ch {
293 + '"' => buf.push_str(r#"\""#),
294 + '\\' => buf.push_str(r"\\"),
295 + '\n' => buf.push_str(r"\n"),
296 + '\r' => buf.push_str(r"\r"),
297 + '\t' => buf.push_str(r"\t"),
298 + c if c.is_control() => {
299 + // \uXXXX escape for control characters
300 + for unit in c.encode_utf16(&mut [0; 2]) {
301 + write!(buf, "\\u{unit:04x}").unwrap();
302 + }
303 + }
304 + c => buf.push(c),
305 + }
306 + }
307 +}
308 +
309 +#[cfg(test)]
310 +mod tests {
311 + use std::sync::Arc;
312 +
313 + use super::*;
314 +
315 + fn mm_input() -> RequestFuncInput {
316 + let frag: Arc<str> =
317 + Arc::from(r#"{"type":"image_url","image_url":{"url":"data:image/jpeg;base64,AAAA"}}"#);
318 + RequestFuncInput {
319 + prompt: Arc::from("hello \"world\"\nline2"),
320 + model: "test-model".to_string(),
321 + output_len: 128,
322 + multi_modal_content: Some(Arc::from(vec![frag])),
323 + ..Default::default()
324 + }
325 + }
326 +
327 + /// Regression test: the assembled multimodal payload must be valid JSON
328 + /// (the text part once shipped without its opening quote — a raw-string
329 + /// delimiter eating the trailing `"` in `"text":"`).
330 + #[test]
331 + fn test_build_mm_payload_is_valid_json() {
332 + let payload = build_mm_payload(&mm_input());
333 + let v: serde_json::Value =
334 + serde_json::from_slice(&payload).expect("mm payload must be valid JSON");
335 + assert_eq!(v["model"], "test-model");
336 + assert_eq!(v["messages"][0]["role"], "user");
337 + let content = v["messages"][0]["content"].as_array().unwrap();
338 + assert_eq!(content[0]["type"], "text");
339 + assert_eq!(content[0]["text"], "hello \"world\"\nline2");
340 + assert_eq!(content[1]["type"], "image_url");
341 + assert_eq!(v["max_completion_tokens"], 128);
342 + assert_eq!(v["stream"], true);
343 + assert_eq!(v["stream_options"]["include_usage"], true);
344 + }
345 +
346 + /// --enable-multimodal-chat (dataset pre-built messages) must produce a
347 + /// payload semantically identical to the fragment-assembly path.
348 + #[test]
349 + fn test_chat_messages_json_path_equivalent_to_fragment_path() {
350 + let base = mm_input();
351 + let fragment_payload = build_mm_payload(&base);
352 +
353 + let mut chat = base.clone();
354 + let mm = chat.multi_modal_content.take().unwrap();
355 + let msgs = crate::datasets::random_mm::build_chat_messages_json(&chat.prompt, Some(&mm));
356 + chat.chat_messages_json = Some(Arc::from(msgs.as_str()));
357 + let chat_payload = build_mm_payload(&chat);
358 +
359 + let a: serde_json::Value = serde_json::from_slice(&fragment_payload).unwrap();
360 + let b: serde_json::Value = serde_json::from_slice(&chat_payload).unwrap();
361 + assert_eq!(a, b);
362 + }
363 +
364 + /// ignore_eos and extra_body must survive the raw-splice path.
365 + #[test]
366 + fn test_mm_payload_tail_fields() {
367 + let mut input = mm_input();
368 + input.ignore_eos = true;
369 + input.extra_body = Some(serde_json::json!({"temperature": 0.5, "stream": false}));
370 + let v: serde_json::Value = serde_json::from_slice(&build_mm_payload(&input)).unwrap();
371 + assert_eq!(v["ignore_eos"], true);
372 + assert_eq!(v["temperature"], 0.5);
373 + // keys already set above must not be overridden by extra_body
374 + assert_eq!(v["stream"], true);
375 + }
376 +}
rust/src/bench/src/backends/openai_completions.rs new
+199
@@ -0,0 +1,199 @@
1 +// SPDX-License-Identifier: Apache-2.0
2 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +use std::error::Error as StdError;
5 +use std::time::Instant;
6 +
7 +use futures::StreamExt;
8 +
9 +use super::streaming::{StreamedResponseHandler, trim_bytes};
10 +use super::{CompletionChunk, RequestFuncInput, RequestFuncOutput, build_headers};
11 +use crate::error::Result;
12 +
13 +/// Backend for OpenAI-compatible Completions API (/v1/completions).
14 +/// Used by "vllm" and "openai" backends.
15 +#[derive(Clone)]
16 +pub struct OpenAICompletionsBackend;
17 +
18 +impl OpenAICompletionsBackend {
19 + pub async fn send_request(
20 + &self,
21 + input: &RequestFuncInput,
22 + client: &reqwest::Client,
23 + ) -> Result<RequestFuncOutput> {
24 + let model = input.model_name.as_deref().unwrap_or(&input.model);
25 +
26 + // When prompt_token_ids are available, send them as the `prompt` value
27 + // (JSON array of integers). vLLM's completions API accepts both string
28 + // and token ID array as `prompt`, skipping server-side tokenization.
29 + let prompt_value = if let Some(ref token_ids) = input.prompt_token_ids {
30 + serde_json::json!(token_ids.as_ref())
31 + } else {
32 + serde_json::json!(input.prompt)
33 + };
34 +
35 + let mut payload = serde_json::json!({
36 + "model": model,
37 + "prompt": prompt_value,
38 + "max_tokens": input.output_len,
39 + "stream": true,
40 + "stream_options": {
41 + "include_usage": true,
42 + },
43 + });
44 +
45 + // Always include logprobs (null when not set) — matches Python which
46 + // sends logprobs=None explicitly rather than omitting the key.
47 + payload["logprobs"] = match input.logprobs {
48 + Some(n) => serde_json::json!(n),
49 + None => serde_json::Value::Null,
50 + };
51 +
52 + // Apply ignore_eos and extra_body
53 + if input.ignore_eos {
54 + payload["ignore_eos"] = serde_json::json!(true);
55 + }
56 + if let Some(serde_json::Value::Object(map)) = input.extra_body.as_ref() {
57 + for (k, v) in map {
58 + payload[k] = v.clone();
59 + }
60 + }
61 +
62 + let headers_map = build_headers(None, &input.extra_headers, &input.request_id);
63 +
64 + let mut output = RequestFuncOutput {
65 + prompt_len: input.prompt_len,
66 + itl: Vec::with_capacity(input.output_len.max(1)),
67 + ..Default::default()
68 + };
69 +
70 + let st = Instant::now();
71 + // start_time is overwritten by benchmark.rs with monotonic offset
72 +
73 + let mut most_recent_timestamp = st;
74 + let mut generated_text = String::new();
75 + let mut first_chunk_received = false;
76 +
77 + let mut request = client.post(&input.api_url).json(&payload);
78 + for (k, v) in &headers_map {
79 + request = request.header(k, v);
80 + }
81 +
82 + match request.send().await {
83 + Ok(response) => {
84 + if response.status().is_success() {
85 + let mut handler = StreamedResponseHandler::new();
86 + let mut stream = response.bytes_stream();
87 +
88 + while let Some(chunk_result) = stream.next().await {
89 + let chunk_bytes = match chunk_result {
90 + Ok(b) => b,
91 + Err(e) => {
92 + output.success = false;
93 + output.error = format!("Stream error: {e}");
94 + return Ok(output);
95 + }
96 + };
97 +
98 + let trimmed_bytes = trim_bytes(&chunk_bytes);
99 + if trimmed_bytes.is_empty() {
100 + continue;
101 + }
102 +
103 + let messages = handler.add_chunk(trimmed_bytes);
104 + for message in messages {
105 + // Skip SSE comments
106 + if message.starts_with(':') {
107 + continue;
108 + }
109 +
110 + // Handle multi-field SSE events (e.g., Dynamo sends
111 + // "event: message\ndata: {...}"). Extract the data: line.
112 + let raw = if message.contains('\n') {
113 + match message.lines().find(|l| l.starts_with("data: ")) {
114 + Some(l) => l,
115 + None => continue,
116 + }
117 + } else {
118 + message.as_str()
119 + };
120 +
121 + let chunk = raw.strip_prefix("data: ").unwrap_or(raw);
122 +
123 + if chunk == "[DONE]" {
124 + continue;
125 + }
126 +
127 + // Typed deserialization — avoids allocating a full
128 + // serde_json::Value tree; only extracts needed fields.
129 + let data: CompletionChunk = match serde_json::from_str(chunk) {
130 + Ok(d) => d,
131 + Err(_) => continue,
132 + };
133 +
134 + if !data.choices.is_empty() {
135 + let text = data.choices[0].text.as_deref().unwrap_or("");
136 +
137 + let timestamp = Instant::now();
138 +
139 + if !first_chunk_received {
140 + first_chunk_received = true;
141 + output.ttft = timestamp.duration_since(st).as_secs_f64();
142 + } else {
143 + output.itl.push(
144 + timestamp
145 + .duration_since(most_recent_timestamp)
146 + .as_secs_f64(),
147 + );
148 + }
149 +
150 + most_recent_timestamp = timestamp;
151 + generated_text.push_str(text);
152 + }
153 + // Separate `if` (not `else if`) — Dynamo may send
154 + // both choices and usage in the same chunk.
155 + if let Some(ref usage) = data.usage
156 + && let Some(ct) = usage.completion_tokens
157 + {
158 + output.output_tokens = ct as usize;
159 + }
160 + }
161 + }
162 +
163 + if first_chunk_received {
164 + output.success = true;
165 + } else {
166 + output.success = false;
167 + output.error = "Never received a valid chunk to calculate TTFT. \
168 + This response will be marked as failed!"
169 + .to_string();
170 + }
171 + output.generated_text = generated_text;
172 + output.latency = most_recent_timestamp.duration_since(st).as_secs_f64();
173 + } else {
174 + let status = response.status();
175 + let body = response.text().await.unwrap_or_default();
176 + output.error = if body.is_empty() {
177 + format!("HTTP {status}")
178 + } else {
179 + format!("HTTP {status}: {body}")
180 + };
181 + output.success = false;
182 + }
183 + }
184 + Err(e) => {
185 + output.success = false;
186 + // Capture full error chain for debugging
187 + let mut error_msg = format!("{e}");
188 + let mut source = e.source();
189 + while let Some(cause) = source {
190 + error_msg.push_str(&format!("\n Caused by: {cause}"));
191 + source = cause.source();
192 + }
193 + output.error = error_msg;
194 + }
195 + }
196 +
197 + Ok(output)
198 + }
199 +}
rust/src/bench/src/backends/pooling.rs new
+325
@@ -0,0 +1,325 @@
1 +// SPDX-License-Identifier: Apache-2.0
2 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +//! Pooling/embedding backends: non-streaming HTTP POST for embedding, pooling, and rerank
5 +//! endpoints.
6 +//!
7 +//! Supported variants:
8 +//! - `openai-embeddings`: Standard OpenAI `/v1/embeddings` with text input
9 +//! - `openai-embeddings-chat`: OpenAI `/v1/embeddings` with chat message format (supports
10 +//! multimodal)
11 +//! - `vllm-pooling`: vLLM `/v1/pooling` endpoint
12 +//! - `vllm-rerank`: vLLM `/v1/rerank` endpoint (query + documents)
13 +
14 +use std::time::Instant;
15 +
16 +use crate::backends::{RequestFuncInput, RequestFuncOutput, build_headers};
17 +use crate::cli::BackendKind;
18 +use crate::error::Result;
19 +
20 +/// Response from embedding/pooling endpoints (minimal fields for usage extraction).
21 +#[derive(serde::Deserialize)]
22 +struct PoolingResponse {
23 + usage: Option<PoolingUsage>,
24 +}
25 +
26 +#[derive(serde::Deserialize)]
27 +struct PoolingUsage {
28 + prompt_tokens: Option<u64>,
29 +}
30 +
31 +#[derive(Clone)]
32 +pub struct PoolingBackend {
33 + pub kind: BackendKind,
34 +}
35 +
36 +impl PoolingBackend {
37 + pub async fn send_request(
38 + &self,
39 + input: &RequestFuncInput,
40 + client: &reqwest::Client,
41 + ) -> Result<RequestFuncOutput> {
42 + // Preserve client-side prompt_len as fallback if server doesn't report usage.
43 + let mut output = RequestFuncOutput {
44 + prompt_len: input.prompt_len,
45 + ..Default::default()
46 + };
47 +
48 + let headers = build_headers(
49 + Some("application/json"),
50 + &input.extra_headers,
51 + &input.request_id,
52 + );
53 +
54 + let payload = self.build_payload(input);
55 +
56 + let mut request = client.post(&input.api_url);
57 + for (k, v) in &headers {
58 + request = request.header(k, v);
59 + }
60 +
61 + let st = Instant::now();
62 +
63 + let response = match request.json(&payload).send().await {
64 + Ok(r) => r,
65 + Err(e) => {
66 + output.error = format!("Request failed: {e}");
67 + return Ok(output);
68 + }
69 + };
70 +
71 + if response.status().is_success() {
72 + let latency = st.elapsed().as_secs_f64();
73 + output.latency = latency;
74 + output.ttft = latency;
75 + output.success = true;
76 +
77 + // Parse usage from response; keep client-side prompt_len as fallback.
78 + match response.json::<PoolingResponse>().await {
79 + Ok(data) => {
80 + if let Some(usage) = data.usage
81 + && let Some(tokens) = usage.prompt_tokens
82 + {
83 + output.prompt_len = tokens as usize;
84 + }
85 + }
86 + Err(_) => {
87 + // Response parsed but no usage — keep client-side prompt_len
88 + }
89 + }
90 + } else {
91 + let status = response.status();
92 + let body = response.text().await.unwrap_or_default();
93 + output.error = format!("HTTP {status}: {body}");
94 + }
95 +
96 + Ok(output)
97 + }
98 +
99 + fn build_payload(&self, input: &RequestFuncInput) -> serde_json::Value {
100 + let model = input.model_name.as_deref().unwrap_or(&input.model);
101 +
102 + // For "input" field (openai-embeddings, vllm-pooling): a batched request
103 + // (--random-batch-size) sends the text list; otherwise prefer prompt_token_ids
104 + // when available. The random dataset sets prompt="" and relies on token IDs;
105 + // the OpenAI embeddings API accepts both text strings and token ID arrays.
106 + // Note: embeddings-chat uses text in messages; vllm-rerank uses text as query.
107 + let input_value = if let Some(ref list) = input.prompt_list {
108 + serde_json::json!(list.iter().map(|s| s.as_ref()).collect::<Vec<&str>>())
109 + } else if let Some(ref token_ids) = input.prompt_token_ids {
110 + serde_json::json!(token_ids.as_ref())
111 + } else {
112 + serde_json::json!(input.prompt.as_ref())
113 + };
114 +
115 + let is_vllm_backend = matches!(
116 + self.kind,
117 + BackendKind::VllmPooling | BackendKind::VllmRerank
118 + );
119 +
120 + let mut payload = match self.kind {
121 + BackendKind::OpenaiEmbeddings => {
122 + let mut p = serde_json::json!({
123 + "model": model,
124 + "input": input_value,
125 + });
126 + // truncate_prompt_tokens is vLLM-specific; only include for vLLM backends
127 + // to avoid breaking standard OpenAI providers.
128 + if is_vllm_backend {
129 + p["truncate_prompt_tokens"] = serde_json::json!(-1);
130 + }
131 + p
132 + }
133 + BackendKind::OpenaiEmbeddingsChat => {
134 + // Chat format: uses text prompt in messages array (for multimodal support).
135 + // Python's _get_chat_content always returns a content array.
136 + // Use raw string concatenation for multimodal fragments (zero-copy,
137 + // avoids re-parsing ~200KB+ base64 per image).
138 + let content_json = build_chat_content_json(input);
139 +
140 + let mut p = serde_json::json!({
141 + "model": model,
142 + "messages": [{"role": "user", "content": content_json}],
143 + });
144 + if is_vllm_backend {
145 + p["truncate_prompt_tokens"] = serde_json::json!(-1);
146 + }
147 + p
148 + }
149 + BackendKind::VllmPooling => {
150 + serde_json::json!({
151 + "model": model,
152 + "input": input_value,
153 + "truncate_prompt_tokens": -1,
154 + })
155 + }
156 + BackendKind::VllmRerank => {
157 + // random-rerank dataset: prompt_list = [query, doc1, doc2, ...]
158 + // (mirrors Python async_request_vllm_rerank).
159 + if let Some(ref list) = input.prompt_list {
160 + if list.len() < 2 {
161 + eprintln!(
162 + "WARNING: vllm-rerank request has no documents \
163 + (prompt_list needs [query, doc, ...])"
164 + );
165 + }
166 + let query = list.first().map(|s| s.as_ref()).unwrap_or("");
167 + let documents: Vec<&str> = list.iter().skip(1).map(|s| s.as_ref()).collect();
168 + serde_json::json!({
169 + "model": model,
170 + "query": query,
171 + "documents": documents,
172 + "truncate_prompt_tokens": -1,
173 + })
174 + } else {
175 + // Legacy path: text prompt as query, documents via --extra-body.
176 + let query = input.prompt.as_ref();
177 + if query.is_empty() && input.prompt_token_ids.is_some() {
178 + eprintln!(
179 + "WARNING: vllm-rerank received empty query (random dataset uses \
180 + token IDs only). Use --dataset-name random-rerank for meaningful \
181 + rerank benchmarks."
182 + );
183 + }
184 + serde_json::json!({
185 + "model": model,
186 + "query": query,
187 + "truncate_prompt_tokens": -1,
188 + })
189 + }
190 + }
191 + _ => unreachable!("PoolingBackend with non-pooling kind"),
192 + };
193 +
194 + // Merge extra_body fields into payload
195 + if let Some(ref extra) = input.extra_body
196 + && let (Some(base), Some(extra_obj)) = (payload.as_object_mut(), extra.as_object())
197 + {
198 + for (k, v) in extra_obj {
199 + base.insert(k.clone(), v.clone());
200 + }
201 + }
202 +
203 + payload
204 + }
205 +}
206 +
207 +/// Build the chat content JSON array for embeddings-chat.
208 +/// Uses raw string concatenation for multimodal fragments to avoid
209 +/// re-parsing large base64 image data (matching openai_chat.rs approach).
210 +fn build_chat_content_json(input: &RequestFuncInput) -> serde_json::Value {
211 + if input.multi_modal_content.is_none() {
212 + // Text-only: return content array with single text element
213 + return serde_json::json!([{
214 + "type": "text",
215 + "text": input.prompt.as_ref(),
216 + }]);
217 + }
218 +
219 + // Multimodal: build JSON string manually for zero-copy fragment embedding
220 + let mm = input.multi_modal_content.as_ref().unwrap();
221 + let prompt = input.prompt.as_ref();
222 +
223 + let mm_total: usize = mm.iter().map(|f| f.len() + 1).sum();
224 + let mut json = String::with_capacity(64 + prompt.len() * 2 + mm_total);
225 +
226 + // [{"type":"text","text":"<prompt>"}
227 + json.push_str(r#"[{"type":"text","text":""#);
228 + push_json_escaped_str(&mut json, prompt);
229 + json.push_str(r#""}"#);
230 +
231 + // ,<mm fragment 1>,<mm fragment 2>,...
232 + for fragment in mm.iter() {
233 + json.push(',');
234 + json.push_str(fragment);
235 + }
236 +
237 + json.push(']');
238 +
239 + // Parse the assembled string into a Value for embedding in the payload.
240 + // This parse is O(n) but operates on the pre-built string once, not per-fragment.
241 + serde_json::from_str(&json).unwrap_or_else(|_| {
242 + serde_json::json!([{
243 + "type": "text",
244 + "text": input.prompt.as_ref(),
245 + }])
246 + })
247 +}
248 +
249 +/// Escape a string for safe JSON embedding (matching openai_chat.rs).
250 +fn push_json_escaped_str(buf: &mut String, s: &str) {
251 + use std::fmt::Write;
252 + for ch in s.chars() {
253 + match ch {
254 + '"' => buf.push_str(r#"\""#),
255 + '\\' => buf.push_str(r"\\"),
256 + '\n' => buf.push_str(r"\n"),
257 + '\r' => buf.push_str(r"\r"),
258 + '\t' => buf.push_str(r"\t"),
259 + c if c < '\x20' => {
260 + let _ = write!(buf, "\\u{:04x}", c as u32);
261 + }
262 + c => buf.push(c),
263 + }
264 + }
265 +}
266 +
267 +#[cfg(test)]
268 +mod tests {
269 + use std::sync::Arc;
270 +
271 + use super::*;
272 +
273 + fn list(items: &[&str]) -> Option<Arc<[Arc<str>]>> {
274 + Some(items.iter().map(|s| Arc::from(*s)).collect())
275 + }
276 +
277 + #[test]
278 + fn test_embeddings_payload_batched_input() {
279 + let backend = PoolingBackend {
280 + kind: BackendKind::OpenaiEmbeddings,
281 + };
282 + let input = RequestFuncInput {
283 + model: "bge".to_string(),
284 + prompt_list: list(&["t1", "t2", "t3"]),
285 + ..Default::default()
286 + };
287 + let payload = backend.build_payload(&input);
288 + assert_eq!(payload["input"], serde_json::json!(["t1", "t2", "t3"]));
289 + assert_eq!(payload["model"], "bge");
290 + // truncate_prompt_tokens is vLLM-specific and deliberately omitted for
291 + // the plain OpenAI embeddings backend.
292 + assert!(payload.get("truncate_prompt_tokens").is_none());
293 + }
294 +
295 + #[test]
296 + fn test_rerank_payload_query_and_documents() {
297 + let backend = PoolingBackend {
298 + kind: BackendKind::VllmRerank,
299 + };
300 + let input = RequestFuncInput {
301 + model: "reranker".to_string(),
302 + prompt_list: list(&["the query", "doc a", "doc b"]),
303 + ..Default::default()
304 + };
305 + let payload = backend.build_payload(&input);
306 + assert_eq!(payload["query"], "the query");
307 + assert_eq!(payload["documents"], serde_json::json!(["doc a", "doc b"]));
308 + assert_eq!(payload["truncate_prompt_tokens"], -1);
309 + }
310 +
311 + #[test]
312 + fn test_rerank_payload_legacy_single_prompt() {
313 + let backend = PoolingBackend {
314 + kind: BackendKind::VllmRerank,
315 + };
316 + let input = RequestFuncInput {
317 + model: "reranker".to_string(),
318 + prompt: Arc::from("query text"),
319 + ..Default::default()
320 + };
321 + let payload = backend.build_payload(&input);
322 + assert_eq!(payload["query"], "query text");
323 + assert!(payload.get("documents").is_none());
324 + }
325 +}
rust/src/bench/src/backends/streaming.rs new
+151
@@ -0,0 +1,151 @@
1 +// SPDX-License-Identifier: Apache-2.0
2 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +/// SSE streaming response handler.
5 +///
6 +/// Accumulates incoming byte chunks and extracts complete SSE messages.
7 +/// Mirrors Python's `StreamedResponseHandler` from endpoint_request_func.py:22-60.
8 +pub struct StreamedResponseHandler {
9 + buffer: String,
10 + /// Reusable message buffer — avoids allocating a new Vec per `add_chunk` call.
11 + messages: Vec<String>,
12 +}
13 +
14 +impl StreamedResponseHandler {
15 + pub fn new() -> Self {
16 + Self {
17 + buffer: String::with_capacity(4096),
18 + messages: Vec::with_capacity(4),
19 + }
20 + }
21 +
22 + /// Add a chunk of bytes and return any complete SSE messages.
23 + ///
24 + /// The returned slice borrows from the handler and is valid until the next
25 + /// `add_chunk` call.
26 + pub fn add_chunk(&mut self, chunk_bytes: &[u8]) -> &[String] {
27 + self.messages.clear();
28 +
29 + let chunk_str = String::from_utf8_lossy(chunk_bytes);
30 + self.buffer.push_str(&chunk_str);
31 +
32 + // Split by double newlines (SSE message separator)
33 + while let Some(pos) = self.buffer.find("\n\n") {
34 + let message = self.buffer[..pos].trim().to_string();
35 + // Efficiently remove consumed bytes by shifting remaining data
36 + self.buffer.drain(..pos + 2);
37 + if !message.is_empty() {
38 + self.messages.push(message);
39 + }
40 + }
41 +
42 + // Handle buffered data without trailing `\n\n`.
43 + // Matches Python's speculative json.loads() in StreamedResponseHandler.
44 + // This matters for TTFT/ITL accuracy: when a data message and its `\n\n`
45 + // arrive in separate TCP segments, we want to emit the message at the
46 + // first segment's arrival time, not the second.
47 + //
48 + // Also handles multi-field SSE events where the buffer may start with
49 + // "event: ...\ndata: ..." (Dynamo frontend).
50 + let data_start = if self.buffer.starts_with("data: ") {
51 + Some(0)
52 + } else {
53 + // Look for a "data: " line in multi-field events
54 + self.buffer.find("\ndata: ").map(|p| p + 1)
55 + };
56 + if let Some(offset) = data_start {
57 + let content = self.buffer[offset + 6..].trim();
58 + if content == "[DONE]"
59 + || (!content.is_empty()
60 + && serde_json::from_str::<&serde_json::value::RawValue>(content).is_ok())
61 + {
62 + self.messages.push(self.buffer.trim().to_string());
63 + self.buffer.clear();
64 + }
65 + }
66 +
67 + &self.messages
68 + }
69 +}
70 +
71 +/// Trim leading/trailing ASCII whitespace from a byte slice.
72 +pub fn trim_bytes(bytes: &[u8]) -> &[u8] {
73 + let start = bytes.iter().position(|b| !b.is_ascii_whitespace()).unwrap_or(bytes.len());
74 + let end = bytes
75 + .iter()
76 + .rposition(|b| !b.is_ascii_whitespace())
77 + .map(|p| p + 1)
78 + .unwrap_or(start);
79 + &bytes[start..end]
80 +}
81 +
82 +#[cfg(test)]
83 +mod tests {
84 + use super::*;
85 +
86 + #[test]
87 + fn test_basic_sse() {
88 + let mut handler = StreamedResponseHandler::new();
89 + let msgs =
90 + handler.add_chunk(b"data: {\"choices\":[{\"text\":\"hi\"}]}\n\ndata: [DONE]\n\n");
91 + assert_eq!(msgs.len(), 2);
92 + assert!(msgs[0].contains("choices"));
93 + assert!(msgs[1].contains("[DONE]"));
94 + }
95 +
96 + #[test]
97 + fn test_split_chunks() {
98 + let mut handler = StreamedResponseHandler::new();
99 + let msgs1 = handler.add_chunk(b"data: {\"cho");
100 + assert!(msgs1.is_empty());
101 + let msgs2 = handler.add_chunk(b"ices\":[{\"text\":\"a\"}]}\n\n");
102 + assert_eq!(msgs2.len(), 1);
103 + }
104 +
105 + #[test]
106 + fn test_comment_lines() {
107 + let mut handler = StreamedResponseHandler::new();
108 + let msgs = handler.add_chunk(b": ping\n\ndata: {\"test\":1}\n\n");
109 + assert_eq!(msgs.len(), 2);
110 + assert!(msgs[0].starts_with(":"));
111 + }
112 +
113 + #[test]
114 + fn test_done_without_newlines() {
115 + let mut handler = StreamedResponseHandler::new();
116 + let msgs = handler.add_chunk(b"data: [DONE]");
117 + assert_eq!(msgs.len(), 1);
118 + assert!(msgs[0].contains("[DONE]"));
119 + }
120 +
121 + #[test]
122 + fn test_incomplete_json_in_buffer() {
123 + let mut handler = StreamedResponseHandler::new();
124 + let msgs = handler.add_chunk(b"data: {\"partial\":");
125 + assert!(msgs.is_empty());
126 + // Complete JSON without \n\n — speculative parse emits it
127 + let msgs2 = handler.add_chunk(b"true}");
128 + assert_eq!(msgs2.len(), 1);
129 + assert!(msgs2[0].contains("partial"));
130 + }
131 +
132 + #[test]
133 + fn test_multi_field_sse_event() {
134 + // Dynamo frontend sends "event: message\ndata: {...}\n\n"
135 + let mut handler = StreamedResponseHandler::new();
136 + let msgs =
137 + handler.add_chunk(b"event: message\ndata: {\"choices\":[{\"text\":\"hi\"}]}\n\n");
138 + assert_eq!(msgs.len(), 1);
139 + assert!(msgs[0].contains("choices"));
140 + assert!(msgs[0].contains("event: message"));
141 + }
142 +
143 + #[test]
144 + fn test_multi_field_sse_speculative_parse() {
145 + // Multi-field event without trailing \n\n — speculative parse should emit it
146 + let mut handler = StreamedResponseHandler::new();
147 + let msgs = handler.add_chunk(b"event: message\ndata: {\"choices\":[{\"text\":\"hi\"}]}");
148 + assert_eq!(msgs.len(), 1);
149 + assert!(msgs[0].contains("choices"));
150 + }
151 +}
rust/src/bench/src/benchmark.rs new
+1977
@@ -0,0 +1,1977 @@
1 +// SPDX-License-Identifier: Apache-2.0
2 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +use std::collections::HashMap;
5 +use std::sync::Arc;
6 +use std::time::Instant;
7 +
8 +use indicatif::{ProgressBar, ProgressStyle};
9 +use tokio::sync::Semaphore;
10 +
11 +use crate::backends::{RequestFuncInput, RequestFuncOutput, get_backend};
12 +use crate::cli::{BackendKind, DatasetName, LoraAssignment};
13 +use crate::config::BenchConfig;
14 +use crate::error::{BenchError, Result};
15 +use crate::metrics::calculator::{calculate_embedding_metrics, calculate_metrics};
16 +use crate::metrics::steady_state;
17 +use crate::output::console::print_results;
18 +use crate::output::json::{append_result, build_result_json, compute_result_filename, save_result};
19 +use crate::rate_control::compute_schedule;
20 +use crate::ready_checker::wait_for_endpoint;
21 +
22 +/// Pre-resolve the hostname in `base_url` and pin all resolved IPs on the
23 +/// client builder via [`reqwest::ClientBuilder::resolve_to_addrs`]. This
24 +/// avoids repeated DNS lookups under high concurrency which can cause
25 +/// transient "Temporary failure in name resolution" errors while preserving
26 +/// happy-eyeballs and multi-A failover.
27 +///
28 +/// Skipped when: URL parse fails, host is already an IP, host is a loopback
29 +/// name (resolved from `/etc/hosts`, no DNS pressure and dual-stack ambiguity
30 +/// between `127.0.0.1` and `::1` breaks IPv4-only servers like vLLM), or
31 +/// resolution fails.
32 +pub fn pre_resolve_dns(
33 + base_url: &str,
34 + mut builder: reqwest::ClientBuilder,
35 +) -> reqwest::ClientBuilder {
36 + let parsed = match url::Url::parse(base_url) {
37 + Ok(u) => u,
38 + Err(_) => return builder,
39 + };
40 +
41 + let host = match parsed.host_str() {
42 + Some(h) => h,
43 + None => return builder,
44 + };
45 +
46 + if host.parse::<std::net::IpAddr>().is_ok() {
47 + return builder;
48 + }
49 +
50 + let host_lower = host.to_ascii_lowercase();
51 + if host_lower == "localhost"
52 + || host_lower == "ip6-localhost"
53 + || host_lower.ends_with(".localhost")
54 + {
55 + return builder;
56 + }
57 +
58 + let port = parsed.port_or_known_default().unwrap_or(80);
59 + let addr_str = format!("{host}:{port}");
60 +
61 + match std::net::ToSocketAddrs::to_socket_addrs(&addr_str) {
62 + Ok(addrs) => {
63 + let mut v4 = Vec::new();
64 + let mut v6 = Vec::new();
65 + for addr in addrs {
66 + if addr.is_ipv4() {
67 + v4.push(addr);
68 + } else {
69 + v6.push(addr);
70 + }
71 + }
72 + v4.extend(v6);
73 + if !v4.is_empty() {
74 + let ips: Vec<_> = v4.iter().map(|a| a.ip()).collect();
75 + println!("Pre-resolved {host} -> {ips:?}");
76 + builder = builder.resolve_to_addrs(host, &v4);
77 + }
78 + }
79 + Err(e) => {
80 + eprintln!("Warning: DNS pre-resolution for '{host}' failed: {e}");
81 + }
82 + }
83 +
84 + builder
85 +}
86 +
87 +/// Raw speculative decoding metrics from the server's Prometheus endpoint.
88 +#[derive(Debug, Clone)]
89 +pub(crate) struct SpecDecodeMetrics {
90 + num_drafts: u64,
91 + num_draft_tokens: u64,
92 + num_accepted_tokens: u64,
93 + accepted_per_pos: HashMap<u64, u64>,
94 +}
95 +
96 +/// Computed speculative decoding statistics (delta between before/after benchmark).
97 +#[derive(Debug, Clone)]
98 +pub struct SpecDecodeStats {
99 + pub num_drafts: u64,
100 + pub draft_tokens: u64,
101 + pub accepted_tokens: u64,
102 + pub acceptance_rate: f64,
103 + pub acceptance_length: f64,
104 + pub per_position_acceptance_rates: Vec<f64>,
105 +}
106 +
107 +/// Fetch speculative decoding metrics from the server's Prometheus `/metrics` endpoint.
108 +///
109 +/// Returns None if speculative decoding is not enabled or metrics are not available.
110 +pub(crate) async fn fetch_spec_decode_metrics(
111 + base_url: &str,
112 + client: &reqwest::Client,
113 + extra_headers: &Option<std::collections::HashMap<String, String>>,
114 +) -> Option<SpecDecodeMetrics> {
115 + let metrics_url = format!("{base_url}/metrics");
116 + let mut request = client.get(&metrics_url);
117 + if let Some(headers) = extra_headers {
118 + for (k, v) in headers {
119 + request = request.header(k, v);
120 + }
121 + }
122 + if let Ok(api_key) = std::env::var("OPENAI_API_KEY") {
123 + request = request.header("Authorization", format!("Bearer {api_key}"));
124 + }
125 +
126 + let response = match request.send().await {
127 + Ok(r) if r.status().is_success() => r,
128 + _ => return None,
129 + };
130 +
131 + let text = match response.text().await {
132 + Ok(t) => t,
133 + Err(_) => return None,
134 + };
135 +
136 + let mut num_drafts: u64 = 0;
137 + let mut num_draft_tokens: u64 = 0;
138 + let mut num_accepted_tokens: u64 = 0;
139 + let mut accepted_per_pos: HashMap<u64, u64> = HashMap::new();
140 + let mut found_spec_decode = false;
141 +
142 + for line in text.lines() {
143 + let line = line.trim();
144 + if line.is_empty() || line.starts_with('#') {
145 + continue;
146 + }
147 + if !line.starts_with("vllm:spec_decode") {
148 + continue;
149 + }
150 + found_spec_decode = true;
151 + let parts: Vec<&str> = line.split_whitespace().collect();
152 + if parts.is_empty() {
153 + continue;
154 + }
155 + let val = match parts.last().and_then(|s| s.parse::<f64>().ok()) {
156 + Some(v) => v as u64,
157 + None => continue,
158 + };
159 +
160 + if line.contains("num_drafts") {
161 + num_drafts += val;
162 + } else if line.contains("num_draft_tokens") {
163 + num_draft_tokens += val;
164 + } else if line.contains("num_accepted_tokens_per_pos") {
165 + // Parse position label: position="N"
166 + if let Some(start) = line.find("position=\"") {
167 + let start = start + "position=\"".len();
168 + if let Some(end) = line[start..].find('"')
169 + && let Ok(pos) = line[start..start + end].parse::<u64>()
170 + {
171 + *accepted_per_pos.entry(pos).or_insert(0) += val;
172 + }
173 + }
174 + } else if line.contains("num_accepted_tokens") {
175 + num_accepted_tokens += val;
176 + }
177 + }
178 +
179 + if !found_spec_decode {
180 + return None;
181 + }
182 +
183 + Some(SpecDecodeMetrics {
184 + num_drafts,
185 + num_draft_tokens,
186 + num_accepted_tokens,
187 + accepted_per_pos,
188 + })
189 +}
190 +
191 +/// Compute speculative decoding stats from before/after metrics snapshots.
192 +pub(crate) fn compute_spec_decode_stats(
193 + before: &SpecDecodeMetrics,
194 + after: &SpecDecodeMetrics,
195 +) -> Option<SpecDecodeStats> {
196 + let delta_drafts = after.num_drafts.saturating_sub(before.num_drafts);
197 + let delta_draft_tokens = after.num_draft_tokens.saturating_sub(before.num_draft_tokens);
198 + let delta_accepted = after.num_accepted_tokens.saturating_sub(before.num_accepted_tokens);
199 +
200 + if delta_draft_tokens == 0 {
201 + return None;
202 + }
203 +
204 + let mut per_pos_rates = Vec::new();
205 + if delta_drafts > 0 {
206 + let mut positions: Vec<u64> = before
207 + .accepted_per_pos
208 + .keys()
209 + .chain(after.accepted_per_pos.keys())
210 + .copied()
211 + .collect::<std::collections::HashSet<_>>()
212 + .into_iter()
213 + .collect();
214 + positions.sort();
215 +
216 + for pos in positions {
217 + let before_val = before.accepted_per_pos.get(&pos).copied().unwrap_or(0);
218 + let after_val = after.accepted_per_pos.get(&pos).copied().unwrap_or(before_val);
219 + let delta_pos = after_val.saturating_sub(before_val);
220 + per_pos_rates.push(delta_pos as f64 / delta_drafts as f64);
221 + }
222 + }
223 +
224 + let acceptance_rate = (delta_accepted as f64 / delta_draft_tokens as f64) * 100.0;
225 + let acceptance_length = if delta_drafts > 0 {
226 + 1.0 + delta_accepted as f64 / delta_drafts as f64
227 + } else {
228 + 0.0
229 + };
230 +
231 + Some(SpecDecodeStats {
232 + num_drafts: delta_drafts,
233 + draft_tokens: delta_draft_tokens,
234 + accepted_tokens: delta_accepted,
235 + acceptance_rate,
236 + acceptance_length,
237 + per_position_acceptance_rates: per_pos_rates,
238 + })
239 +}
240 +
241 +/// Pre-assign a LoRA adapter name per item based on the configured strategy.
242 +///
243 +/// Returns `None` when LoRA is not configured. With LoRA enabled, returns a
244 +/// `Vec<Arc<str>>` of length `n` whose i-th entry is the adapter name to use
245 +/// for the i-th item.
246 +///
247 +/// - `RoundRobin` cycles deterministically: `lora_modules[i % N]`.
248 +/// - `Random` uses `StdRng::seed_from_u64(seed)` so the assignment is fully reproducible across
249 +/// runs with the same seed.
250 +///
251 +/// "Item" is a request in single-shot mode and a conversation in multi-turn
252 +/// mode (sticky across all turns of that conversation).
253 +pub(crate) fn assign_lora_modules(
254 + lora_modules: &Option<Vec<Arc<str>>>,
255 + assignment: LoraAssignment,
256 + n: usize,
257 + seed: u64,
258 +) -> Option<Vec<Arc<str>>> {
259 + let modules = lora_modules.as_ref()?;
260 + if modules.is_empty() {
261 + return None;
262 + }
263 + let m = modules.len();
264 + let mut out = Vec::with_capacity(n);
265 + match assignment {
266 + LoraAssignment::RoundRobin => {
267 + for i in 0..n {
268 + out.push(modules[i % m].clone());
269 + }
270 + }
271 + LoraAssignment::Random => {
272 + use rand::rngs::StdRng;
273 + use rand::{Rng, SeedableRng};
274 + let mut rng = StdRng::seed_from_u64(seed);
275 + for _ in 0..n {
276 + out.push(modules[rng.random_range(0..m)].clone());
277 + }
278 + }
279 + }
280 + Some(out)
281 +}
282 +
283 +/// Fetch the first model from the server's /v1/models endpoint.
284 +async fn get_first_model_from_server(
285 + base_url: &str,
286 + client: &reqwest::Client,
287 + extra_headers: &Option<std::collections::HashMap<String, String>>,
288 +) -> Result<(String, String)> {
289 + let url = format!("{base_url}/v1/models");
290 + let mut request = client.get(&url);
291 + if let Some(headers) = extra_headers {
292 + for (k, v) in headers {
293 + request = request.header(k, v);
294 + }
295 + }
296 + // Add API key from environment
297 + if let Ok(api_key) = std::env::var("OPENAI_API_KEY") {
298 + request = request.header("Authorization", format!("Bearer {api_key}"));
299 + }
300 +
301 + let response = request.send().await?;
302 + let data: serde_json::Value = response.json().await?;
303 +
304 + if let Some(models) = data.get("data").and_then(|d| d.as_array())
305 + && let Some(first) = models.first()
306 + {
307 + let id = first.get("id").and_then(|v| v.as_str()).unwrap_or_default().to_string();
308 + let root = first.get("root").and_then(|v| v.as_str()).unwrap_or(&id).to_string();
309 + return Ok((id, root));
310 + }
311 +
312 + Err(BenchError::Config(format!(
313 + "No models found on the server at {base_url}"
314 + )))
315 +}
316 +
317 +/// Run the complete benchmark.
318 +///
319 +/// This is the core orchestrator matching Python's benchmark() + main_async().
320 +pub async fn run_benchmark(config: &BenchConfig) -> Result<serde_json::Value> {
321 + // Validate backend name early
322 + let _ = get_backend(config.backend)?;
323 +
324 + // Build HTTP client with connection pool settings matching Python's aiohttp.TCPConnector
325 + // Force HTTP/1.1 to match aiohttp behavior (avoids HTTP/2 negotiation issues)
326 + let mut client_builder = reqwest::Client::builder()
327 + .pool_max_idle_per_host(config.max_concurrency.unwrap_or(2048).max(256))
328 + .timeout(std::time::Duration::from_secs(6 * 60 * 60))
329 + .connect_timeout(std::time::Duration::from_secs(30))
330 + .tcp_keepalive(std::time::Duration::from_secs(60))
331 + .tcp_nodelay(true)
332 + .http1_only()
333 + .no_proxy();
334 +
335 + if config.insecure {
336 + client_builder = client_builder.danger_accept_invalid_certs(true);
337 + }
338 +
339 + client_builder = pre_resolve_dns(&config.base_url, client_builder);
340 +
341 + let client = client_builder
342 + .build()
343 + .map_err(|e| BenchError::Backend(format!("Failed to build HTTP client: {e}")))?;
344 +
345 + // Resolve model
346 + let (model_id, model_name) = if let Some(ref m) = config.model {
347 + (m.clone(), config.model_name.clone())
348 + } else {
349 + println!("Model not specified, fetching first model from server...");
350 + let (name, id) =
351 + get_first_model_from_server(&config.base_url, &client, &config.extra_headers).await?;
352 + println!("First model name: {name}, first model id: {id}");
353 + (id, Some(name))
354 + };
355 +
356 + // Load tokenizer (if needed)
357 + let tokenizer = if config.skip_tokenizer_init {
358 + None
359 + } else {
360 + let tid = config.tokenizer_id.as_deref().unwrap_or(&model_id);
361 + println!("Loading tokenizer: {tid}");
362 + let server_info = Some((config.base_url.as_str(), model_id.as_str()));
363 + let t = crate::tokenizer::load_tokenizer(tid, config.trust_remote_code, server_info)?;
364 + println!("Tokenizer loaded successfully.");
365 + Some(t)
366 + };
367 + let has_tokenizer = tokenizer.is_some();
368 +
369 + // Generate dataset
370 + let dataset_label = match config.dataset_name {
371 + DatasetName::Random => format!("{} random prompts", config.num_prompts),
372 + DatasetName::RandomMm => format!(
373 + "{} random multimodal prompts ({} items/req)",
374 + config.num_prompts, config.random_mm_base_items_per_request
375 + ),
376 + DatasetName::ShareGpt => format!(
377 + "{} prompts from ShareGPT ({})",
378 + config.num_prompts,
379 + config.dataset_path.as_deref().unwrap_or("auto-download")
380 + ),
381 + DatasetName::Sonnet => format!(
382 + "{} prompts from Sonnet ({}, isl={}, osl={}, prefix={})",
383 + config.num_prompts,
384 + config.dataset_path.as_deref().unwrap_or("built-in"),
385 + config.sonnet_input_len,
386 + config.sonnet_output_len,
387 + config.sonnet_prefix_len,
388 + ),
389 + DatasetName::SpeedBench => {
390 + let truncate_info = config
391 + .speed_bench_max_input_len
392 + .map(|n| format!(", truncated to {n} tokens"))
393 + .unwrap_or_default();
394 + format!(
395 + "{} prompts from SPEED-Bench ({}/{}{})",
396 + config.num_prompts,
397 + config.speed_bench_config,
398 + config.speed_bench_category.as_deref().unwrap_or("all"),
399 + truncate_info,
400 + )
401 + }
402 + DatasetName::Hf => format!(
403 + "{} prompts from HF dataset ({})",
404 + config.num_prompts,
405 + config.dataset_path.as_deref().unwrap_or("unknown"),
406 + ),
407 + DatasetName::Custom => format!(
408 + "{} prompts from custom JSONL ({})",
409 + config.num_prompts,
410 + config.dataset_path.as_deref().unwrap_or("unknown"),
411 + ),
412 + DatasetName::PrefixRepetition => format!(
413 + "{} prefix-repetition prompts ({} prefixes, prefix={}, suffix={})",
414 + config.num_prompts,
415 + config.prefix_repetition_num_prefixes,
416 + config.prefix_repetition_prefix_len,
417 + config.prefix_repetition_suffix_len,
418 + ),
419 + DatasetName::RandomRerank => format!(
420 + "{} random rerank requests (batch={}, reranker={})",
421 + config.num_prompts, config.random_batch_size, config.is_reranker,
422 + ),
423 + };
424 + println!("Generating {dataset_label}...");
425 + let gen_start = Instant::now();
426 +
427 + let mut input_requests = match config.dataset_name {
428 + DatasetName::Random => {
429 + let tok = tokenizer
430 + .as_ref()
431 + .ok_or_else(|| BenchError::Config("Random dataset requires a tokenizer".into()))?;
432 + crate::datasets::random::generate_random_dataset(
433 + tok,
434 + config.num_prompts,
435 + config.random_input_len,
436 + config.random_output_len,
437 + config.random_prefix_len,
438 + config.random_range_ratio,
439 + config.random_cache_hit_fraction,
440 + config.random_cache_ratio,
441 + config.seed,
442 + &config.request_id_prefix,
443 + config.prompt_token_ids,
444 + config.random_batch_size,
445 + )?
446 + }
447 + DatasetName::RandomMm => {
448 + let tok = tokenizer.as_ref().ok_or_else(|| {
449 + BenchError::Config("Random-MM dataset requires a tokenizer".into())
450 + })?;
451 + crate::datasets::random_mm::generate_random_mm_dataset(
452 + tok,
453 + config.num_prompts,
454 + config.random_input_len,
455 + config.random_output_len,
456 + config.random_prefix_len,
457 + config.random_range_ratio,
458 + config.seed,
459 + &config.request_id_prefix,
460 + config.random_mm_base_items_per_request,
461 + config.random_mm_num_mm_items_range_ratio,
462 + &config.random_mm_limit,
463 + &config.random_mm_buckets,
464 + config.enable_multimodal_chat,
465 + )?
466 + }
467 + DatasetName::ShareGpt => {
468 + let tok = tokenizer.as_ref().ok_or_else(|| {
469 + BenchError::Config("ShareGPT dataset requires a tokenizer".into())
470 + })?;
471 + let downloaded;
472 + let path = match config.dataset_path.as_deref() {
473 + Some(p) => p,
474 + None => {
475 + downloaded = crate::datasets::sharegpt::download_sharegpt_dataset()?;
476 + downloaded.as_str()
477 + }
478 + };
479 + crate::datasets::sharegpt::load_sharegpt_dataset(
480 + tok,
481 + path,
482 + config.num_prompts,
483 + config.sharegpt_output_len,
484 + config.seed,
485 + &config.request_id_prefix,
486 + config.no_oversample,
487 + config.disable_shuffle,
488 + )?
489 + }
490 + DatasetName::Sonnet => {
491 + let tok = tokenizer
492 + .as_ref()
493 + .ok_or_else(|| BenchError::Config("Sonnet dataset requires a tokenizer".into()))?;
494 + crate::datasets::sonnet::load_sonnet_dataset(
495 + tok,
496 + config.dataset_path.as_deref(),
497 + config.num_prompts,
498 + config.sonnet_input_len,
499 + config.sonnet_output_len,
500 + config.sonnet_prefix_len,
501 + config.seed,
502 + &config.request_id_prefix,
503 + )?
504 + }
505 + DatasetName::SpeedBench => {
506 + let tok = tokenizer.as_ref().ok_or_else(|| {
507 + BenchError::Config("SPEED-Bench dataset requires a tokenizer".into())
508 + })?;
509 + let downloaded;
510 + let path = match config.dataset_path.as_deref() {
511 + Some(p) => p,
512 + None => {
513 + downloaded = crate::datasets::speed_bench::download_speed_bench(
514 + config.speed_bench_config,
515 + )?;
516 + downloaded.as_str()
517 + }
518 + };
519 + let output_len = config.sharegpt_output_len.unwrap_or(config.random_output_len);
520 + crate::datasets::speed_bench::load_speed_bench_dataset(
521 + tok,
522 + path,
523 + config.num_prompts,
524 + output_len,
525 + config.seed,
526 + &config.request_id_prefix,
527 + config.speed_bench_category.as_deref(),
528 + config.no_oversample,
529 + config.disable_shuffle,
530 + config.speed_bench_max_input_len,
531 + )?
532 + }
533 + DatasetName::Hf => {
534 + let tok = tokenizer
535 + .as_ref()
536 + .ok_or_else(|| BenchError::Config("HF dataset requires a tokenizer".into()))?;
537 + let dataset_id = config.dataset_path.as_deref().ok_or_else(|| {
538 + BenchError::Config("--dataset-path is required for --dataset-name hf".into())
539 + })?;
540 + let (downloaded_path, _config, _split) =
541 + crate::datasets::hf_dataset::download_hf_dataset(
542 + dataset_id,
543 + config.hf_subset.as_deref(),
544 + config.hf_split.as_deref(),
545 + config.num_prompts,
546 + )?;
547 + crate::datasets::hf_dataset::load_hf_dataset(
548 + tok,
549 + &downloaded_path,
550 + config.num_prompts,
551 + config.hf_output_len,
552 + config.seed,
553 + &config.request_id_prefix,
554 + config.hf_text_column.as_deref(),
555 + config.no_oversample,
556 + config.disable_shuffle,
557 + )?
558 + }
559 + DatasetName::Custom => {
560 + let tok = tokenizer
561 + .as_ref()
562 + .ok_or_else(|| BenchError::Config("Custom dataset requires a tokenizer".into()))?;
563 + let path = config.dataset_path.as_deref().ok_or_else(|| {
564 + BenchError::Config("--dataset-path is required for --dataset-name custom".into())
565 + })?;
566 + crate::datasets::custom::load_custom_dataset(
567 + tok,
568 + path,
569 + config.num_prompts,
570 + config.custom_output_len,
571 + config.seed,
572 + &config.request_id_prefix,
573 + config.no_oversample,
574 + config.disable_shuffle,
575 + )?
576 + }
577 + DatasetName::PrefixRepetition => {
578 + let tok = tokenizer.as_ref().ok_or_else(|| {
579 + BenchError::Config("Prefix repetition dataset requires a tokenizer".into())
580 + })?;
581 + crate::datasets::prefix_repetition::generate_prefix_repetition_dataset(
582 + tok,
583 + config.num_prompts,
584 + config.prefix_repetition_prefix_len,
585 + config.prefix_repetition_suffix_len,
586 + config.prefix_repetition_num_prefixes,
587 + config.prefix_repetition_output_len,
588 + config.seed,
589 + &config.request_id_prefix,
590 + config.disable_shuffle,
591 + )?
592 + }
593 + DatasetName::RandomRerank => {
594 + let tok = tokenizer.as_ref().ok_or_else(|| {
595 + BenchError::Config("Random rerank dataset requires a tokenizer".into())
596 + })?;
597 + crate::datasets::random_rerank::generate_random_rerank_dataset(
598 + tok,
599 + config.num_prompts,
600 + config.random_input_len,
601 + config.random_range_ratio,
602 + config.seed,
603 + &config.request_id_prefix,
604 + config.random_batch_size,
605 + config.is_reranker,
606 + )?
607 + }
608 + };
609 +
610 + let gen_elapsed = gen_start.elapsed();
611 + println!(
612 + "Generated {} prompts in {:.2}s",
613 + input_requests.len(),
614 + gen_elapsed.as_secs_f64()
615 + );
616 +
617 + let filtered_count =
618 + filter_requests_by_max_model_len(&mut input_requests, config.max_model_len);
619 + if filtered_count > 0 {
620 + println!(
621 + "Filtered {filtered_count} prompt(s) above --max-model-len {}.",
622 + config.max_model_len.unwrap()
623 + );
624 + }
625 + if input_requests.is_empty() {
626 + return Err(BenchError::Config(
627 + "No requests remain after applying --max-model-len".into(),
628 + ));
629 + }
630 +
631 + // Dry run: just print stats and exit
632 + if config.dry_run {
633 + let total_input_tokens: usize = input_requests.iter().map(|r| r.prompt_len).sum();
634 + let total_output_tokens: usize = input_requests.iter().map(|r| r.expected_output_len).sum();
635 + println!("Dry run stats:");
636 + println!(" Total prompts: {}", input_requests.len());
637 + println!(" Total input tokens: {total_input_tokens}");
638 + println!(" Total expected output tokens: {total_output_tokens}");
639 + println!(
640 + " Avg input tokens: {:.1}",
641 + total_input_tokens as f64 / input_requests.len() as f64
642 + );
643 + println!(
644 + " Avg output tokens: {:.1}",
645 + total_output_tokens as f64 / input_requests.len() as f64
646 + );
647 + return Ok(serde_json::json!({"dry_run": true}));
648 + }
649 +
650 + // Build test input from first request
651 + let first = &input_requests[0];
652 + let test_input = RequestFuncInput {
653 + prompt: first.prompt.clone(),
654 + api_url: config.api_url.clone(),
655 + prompt_len: first.prompt_len,
656 + output_len: first.expected_output_len,
657 + model: model_id.clone(),
658 + model_name: model_name.clone(),
659 + logprobs: config.logprobs,
660 + extra_headers: config.extra_headers.clone(),
661 + extra_body: config.extra_body.clone(),
662 + ignore_eos: config.ignore_eos,
663 + request_id: first.request_id.clone(),
664 + messages: None,
665 + prompt_token_ids: first.prompt_token_ids.clone(),
666 + multi_modal_content: first.multi_modal_content.clone(),
667 + chat_messages_json: first.chat_messages_json.clone(),
668 + prompt_list: first.prompt_list.clone(),
669 + };
670 +
671 + // Ready check
672 + if config.ready_check_timeout_sec > 0 {
673 + println!("Starting initial single prompt test run...");
674 + let test_output = wait_for_endpoint(
675 + config.backend,
676 + &client,
677 + &test_input,
678 + config.ready_check_timeout_sec,
679 + 5,
680 + )
681 + .await?;
682 + if !test_output.success {
683 + return Err(BenchError::Backend(format!(
684 + "Initial test run failed: {}",
685 + test_output.error
686 + )));
687 + }
688 + println!("Initial test run completed.");
689 + }
690 +
691 + // Verify and fix prompt token lengths against the server's /tokenize endpoint.
692 + // Runs after the ready check so a still-starting server isn't mistaken for a
693 + // tokenize failure, and after the dry-run exit so dry runs stay offline.
694 + // Uses a cache: if a previous run with the same model+server already verified OK,
695 + // skip entirely. Otherwise sample 10 prompts first — if all match, cache and skip.
696 + // If any mismatch, do full verify+fix for all prompts.
697 + // Skip verification when prompt_token_ids are set (token counts are exact by construction).
698 + let has_token_ids = input_requests.first().is_some_and(|r| r.prompt_token_ids.is_some());
699 + // Python aligns prompts to the server tokenizer for random AND prefix_repetition
700 + // (both are synthetic exact-length datasets).
701 + let verifiable_dataset = matches!(
702 + config.dataset_name,
703 + DatasetName::Random | DatasetName::PrefixRepetition
704 + );
705 + if verifiable_dataset && has_token_ids && !config.backend.is_pooling() {
706 + println!("Using prompt_token_ids, skipping server-side tokenizer verification.");
707 + }
708 + if verifiable_dataset && !has_token_ids && !config.backend.is_pooling() {
709 + let cache_key = tokenizer_verify_cache_key(&config.base_url, &model_id);
710 + if is_tokenizer_verified(&cache_key) {
711 + println!("Tokenizer verified in previous run (cached), skipping verification.");
712 + } else {
713 + let num_special =
714 + tokenizer.as_ref().map(|t| t.num_special_tokens_to_add()).unwrap_or(0);
715 + match sample_verify_prompts(
716 + &client,
717 + &config.base_url,
718 + &model_id,
719 + &input_requests,
720 + num_special,
721 + &config.extra_headers,
722 + )
723 + .await?
724 + {
725 + SampleVerifyOutcome::Passed => {
726 + println!("Sample verification passed, skipping full verification.");
727 + mark_tokenizer_verified(&cache_key);
728 + }
729 + SampleVerifyOutcome::Skipped(reason) => {
730 + println!("Server /tokenize unavailable ({reason}), skipping verification.");
731 + }
732 + SampleVerifyOutcome::Mismatch => {
733 + println!("Sample verification found mismatch, running full verify+fix...");
734 + match verify_and_fix_prompt_lengths(
735 + &client,
736 + &config.base_url,
737 + &model_id,
738 + &mut input_requests,
739 + num_special,
740 + &config.extra_headers,
741 + )
742 + .await
743 + {
744 + Ok(()) => {
745 + println!(
746 + "All {} prompts verified: exact token length match.",
747 + input_requests.len()
748 + );
749 + mark_tokenizer_verified(&cache_key);
750 + }
751 + Err(BenchError::TokenizeUnavailable(reason)) => {
752 + println!(
753 + "Server /tokenize became unavailable during verification \
754 + ({reason}); proceeding with client-side token counts."
755 + );
756 + }
757 + Err(e) => return Err(e),
758 + }
759 + }
760 + }
761 + }
762 + }
763 +
764 + // Warmup
765 + if config.num_warmups > 0 {
766 + println!("Warming up with {} requests...", config.num_warmups);
767 + run_warmup(
768 + config.backend,
769 + &client,
770 + &test_input,
771 + config.num_warmups,
772 + config.max_concurrency,
773 + config.request_rate,
774 + config.burstiness,
775 + config.seed,
776 + config.disable_tqdm,
777 + )
778 + .await;
779 + println!("Warmup run completed.");
780 + }
781 +
782 + // Start profiler if requested (immediate mode — no batch threshold)
783 + if config.profile && config.profile_batch_threshold.is_none() {
784 + start_profiler_immediate(&client, &config.base_url, &config.extra_headers).await;
785 + }
786 +
787 + // Threshold-based profiling: spawn background task that polls /metrics
788 + // and triggers start/stop profile when batch size is reached.
789 + // A oneshot channel lets us cancel the polling loop when the benchmark ends
790 + // (e.g. if the threshold is never reached).
791 + let profile_task = if let Some(threshold) = config.profile_batch_threshold {
792 + let poll_client = client.clone();
793 + let base_url = config.base_url.clone();
794 + let extra_headers = config.extra_headers.clone();
795 + let duration_secs = config.profile_duration;
796 + let (cancel_tx, cancel_rx) = tokio::sync::oneshot::channel::<()>();
797 + let handle = tokio::spawn(async move {
798 + profile_on_batch_threshold(
799 + &poll_client,
800 + &base_url,
801 + &extra_headers,
802 + threshold,
803 + duration_secs,
804 + cancel_rx,
805 + )
806 + .await;
807 + });
808 + Some((cancel_tx, handle))
809 + } else {
810 + None
811 + };
812 +
813 + // Fetch speculative decoding metrics before benchmark
814 + let spec_decode_before =
815 + fetch_spec_decode_metrics(&config.base_url, &client, &config.extra_headers).await;
816 + if spec_decode_before.is_some() {
817 + println!("Speculative decoding detected, will collect metrics.");
818 + }
819 +
820 + // Main benchmark
821 + println!("Starting main benchmark run...");
822 + let distribution = if config.burstiness == 1.0 {
823 + "Poisson process"
824 + } else {
825 + "Gamma distribution"
826 + };
827 + println!(
828 + "Traffic request rate: {}",
829 + if config.request_rate.is_infinite() {
830 + "inf".to_string()
831 + } else {
832 + format!("{}", config.request_rate)
833 + }
834 + );
835 + println!("Burstiness factor: {} ({distribution})", config.burstiness);
836 + println!(
837 + "Maximum request concurrency: {}",
838 + config.max_concurrency.unwrap_or(config.num_prompts)
839 + );
840 +
841 + // Pre-assign LoRA adapters to each request (None when --lora-modules not set).
842 + let lora_assignments = assign_lora_modules(
843 + &config.lora_modules,
844 + config.lora_assignment,
845 + input_requests.len(),
846 + config.seed,
847 + );
848 + if let (Some(modules), Some(_)) = (config.lora_modules.as_ref(), lora_assignments.as_ref()) {
849 + let names: Vec<&str> = modules.iter().map(|s| s.as_ref()).collect();
850 + println!(
851 + "LoRA adapters ({}): {:?} [assignment={:?}]",
852 + modules.len(),
853 + names,
854 + config.lora_assignment
855 + );
856 + }
857 +
858 + // Compute request schedule
859 + let schedule = compute_schedule(
860 + input_requests.len(),
861 + config.request_rate,
862 + config.burstiness,
863 + config.seed,
864 + config.ramp_up.as_ref(),
865 + );
866 +
867 + // Progress bar
868 + let pb = if config.disable_tqdm {
869 + None
870 + } else {
871 + let bar = ProgressBar::new(input_requests.len() as u64);
872 + bar.set_style(
873 + ProgressStyle::with_template(
874 + "{spinner:.green} [{elapsed_precise}] [{bar:40.cyan/blue}] {pos}/{len} ({eta})",
875 + )
876 + .unwrap()
877 + .progress_chars("#>-"),
878 + );
879 + Some(bar)
880 + };
881 +
882 + // Semaphore for concurrency control
883 + let semaphore = config.max_concurrency.map(|mc| Arc::new(Semaphore::new(mc)));
884 +
885 + let benchmark_start = Instant::now();
886 +
887 + // Arc-wrap shared data to avoid cloning per request
888 + let shared_api_url = Arc::new(config.api_url.clone());
889 + let shared_model_id = Arc::new(model_id.clone());
890 + let shared_model_name = Arc::new(model_name.clone());
891 + let shared_extra_headers = Arc::new(config.extra_headers.clone());
892 + let shared_extra_body = Arc::new(config.extra_body.clone());
893 + let shared_backend = get_backend(config.backend)?;
894 + let shared_logprobs = config.logprobs;
895 + let shared_ignore_eos = config.ignore_eos;
896 +
897 + // Spawn all request tasks. Store prompt_len alongside handle so
898 + // we can preserve it in the panic recovery path.
899 + let mut handles: Vec<(usize, tokio::task::JoinHandle<RequestFuncOutput>)> =
900 + Vec::with_capacity(input_requests.len());
901 +
902 + for (i, (request, delay)) in input_requests.iter().zip(schedule.delays.iter()).enumerate() {
903 + let client = client.clone();
904 + let backend = shared_backend.clone();
905 + let sem = semaphore.clone();
906 + let pb = pb.clone();
907 + let api_url = shared_api_url.clone();
908 + let model = shared_model_id.clone();
909 + let model_name = shared_model_name.clone();
910 + let extra_headers = shared_extra_headers.clone();
911 + let extra_body = shared_extra_body.clone();
912 +
913 + // Per-request LoRA override: the adapter name replaces both `model` and
914 + // `model_name` in the outgoing payload (vLLM routes by name). Tokenizer,
915 + // /v1/models, /tokenize, etc. continue using the base model unchanged.
916 + let lora_name = lora_assignments.as_ref().map(|v| v[i].clone());
917 +
918 + let prompt = request.prompt.clone();
919 + let prompt_len = request.prompt_len;
920 + let output_len = request.expected_output_len;
921 + let request_id = request.request_id.clone();
922 + let prompt_token_ids = request.prompt_token_ids.clone();
923 + let multi_modal_content = request.multi_modal_content.clone();
924 + let chat_messages_json = request.chat_messages_json.clone();
925 + let prompt_list = request.prompt_list.clone();
926 +
927 + let delay_dur = std::time::Duration::from_secs_f64(*delay);
928 + let bench_start = benchmark_start;
929 +
930 + handles.push((
931 + prompt_len,
932 + tokio::spawn(async move {
933 + // Sleep until scheduled time
934 + let target = bench_start + delay_dur;
935 + let now = Instant::now();
936 + if target > now {
937 + tokio::time::sleep(target - now).await;
938 + }
939 +
940 + // Acquire semaphore permit
941 + let _permit = if let Some(ref s) = sem {
942 + Some(s.acquire().await.unwrap())
943 + } else {
944 + None
945 + };
946 +
947 + let (req_model, req_model_name) = match lora_name.as_ref() {
948 + Some(name) => (name.to_string(), Some(name.to_string())),
949 + None => ((*model).clone(), (*model_name).clone()),
950 + };
951 +
952 + let input = RequestFuncInput {
953 + prompt,
954 + api_url: (*api_url).clone(),
955 + prompt_len,
956 + output_len,
957 + model: req_model,
958 + model_name: req_model_name,
959 + logprobs: shared_logprobs,
960 + extra_headers: (*extra_headers).clone(),
961 + extra_body: (*extra_body).clone(),
962 + ignore_eos: shared_ignore_eos,
963 + request_id,
964 + messages: None,
965 + prompt_token_ids,
966 + multi_modal_content,
967 + chat_messages_json,
968 + prompt_list,
969 + };
970 +
971 + // Send request, retry on connection errors
972 + let max_retries = 3;
973 + let mut output = None;
974 +
975 + for attempt in 0..=max_retries {
976 + // Capture monotonic start time right before sending,
977 + // relative to benchmark_start. Python uses perf_counter()
978 + // for both start_time and ttft/itl, keeping them on the
979 + // same clock. We do the same with Instant.
980 + let request_instant = Instant::now();
981 + let result = backend.send_request(&input, &client).await;
982 +
983 + match result {
984 + Ok(mut o) => {
985 + // Override SystemTime-based start_time with monotonic offset
986 + o.start_time =
987 + request_instant.duration_since(bench_start).as_secs_f64();
988 +
989 + // Retry on connection-level failures reported as !success
990 + if !o.success && attempt < max_retries && is_connection_error(&o.error)
991 + {
992 + tokio::time::sleep(std::time::Duration::from_millis(
993 + 500 * (attempt as u64 + 1),
994 + ))
995 + .await;
996 + continue;
997 + }
998 + output = Some(o);
999 + break;
1000 + }
1001 + Err(e) => {
1002 + if attempt < max_retries {
1003 + tokio::time::sleep(std::time::Duration::from_millis(
1004 + 500 * (attempt as u64 + 1),
1005 + ))
1006 + .await;
1007 + continue;
1008 + }
1009 + output = Some(RequestFuncOutput {
1010 + success: false,
1011 + error: e.to_string(),
1012 + prompt_len: input.prompt_len,
1013 + start_time: request_instant
1014 + .duration_since(bench_start)
1015 + .as_secs_f64(),
1016 + ..Default::default()
1017 + });
1018 + break;
1019 + }
1020 + }
1021 + }
1022 +
1023 + if let Some(pb) = pb {
1024 + pb.inc(1);
1025 + }
1026 +
1027 + output.unwrap()
1028 + }),
1029 + ));
1030 + }
1031 +
1032 + // Collect all results, handling task panics gracefully
1033 + let mut outputs = Vec::with_capacity(handles.len());
1034 + for (prompt_len, handle) in handles {
1035 + match handle.await {
1036 + Ok(output) => outputs.push(output),
1037 + Err(e) => {
1038 + outputs.push(RequestFuncOutput {
1039 + success: false,
1040 + error: format!("Task panicked: {e}"),
1041 + prompt_len,
1042 + ..Default::default()
1043 + });
1044 + }
1045 + }
1046 + }
1047 +
1048 + if let Some(ref pb) = pb {
1049 + pb.finish_and_clear();
1050 + }
1051 +
1052 + let benchmark_duration = benchmark_start.elapsed().as_secs_f64();
1053 +
1054 + // Fetch speculative decoding metrics after benchmark and compute stats
1055 + let spec_decode_stats = if spec_decode_before.is_some() {
1056 + let spec_decode_after =
1057 + fetch_spec_decode_metrics(&config.base_url, &client, &config.extra_headers).await;
1058 + match (spec_decode_before.as_ref(), spec_decode_after.as_ref()) {
1059 + (Some(before), Some(after)) => compute_spec_decode_stats(before, after),
1060 + _ => None,
1061 + }
1062 + } else {
1063 + None
1064 + };
1065 +
1066 + // Calculate metrics — pooling uses a dedicated path matching Python's
1067 + // calculate_metrics_for_embeddings (uses server-reported prompt_len, e2el only).
1068 + let (mut metrics, actual_output_lens) = if config.backend.is_pooling() {
1069 + let m =
1070 + calculate_embedding_metrics(&outputs, benchmark_duration, &config.selected_percentiles);
1071 + (m, Vec::new())
1072 + } else {
1073 + calculate_metrics(
1074 + &input_requests,
1075 + &outputs,
1076 + benchmark_duration,
1077 + &config.selected_percentiles,
1078 + has_tokenizer,
1079 + &config.goodput,
1080 + )
1081 + };
1082 +
1083 + // Attach steady-state metrics when the closed-loop scope gate passes.
1084 + let scope_ok = !config.no_steady_state
1085 + && config.max_concurrency.is_some()
1086 + && config.request_rate.is_infinite();
1087 + if scope_ok {
1088 + let target = config.max_concurrency;
1089 + let min_window = config
1090 + .steady_state_min_window
1091 + .unwrap_or_else(|| (0.1 * benchmark_duration).max(10.0));
1092 + if let Some(window) = steady_state::detect_window(
1093 + &outputs,
1094 + target,
1095 + config.steady_state_threshold,
1096 + min_window,
1097 + benchmark_duration,
1098 + ) {
1099 + let ss = steady_state::compute(
1100 + &outputs,
1101 + &input_requests,
1102 + &window,
1103 + &config.selected_percentiles,
1104 + config.backend.is_pooling(),
1105 + );
1106 + metrics.steady_state = Some(ss);
1107 + }
1108 + }
1109 +
1110 + // Print console output
1111 + print_results(
1112 + &metrics,
1113 + benchmark_duration,
1114 + config,
1115 + has_tokenizer,
1116 + spec_decode_stats.as_ref(),
1117 + );
1118 +
1119 + // Stop profiler if requested (immediate mode — no batch threshold)
1120 + if config.profile && config.profile_batch_threshold.is_none() {
1121 + stop_profiler_immediate(&client, &config.base_url, &config.extra_headers).await;
1122 + }
1123 +
1124 + // Signal the threshold-based profile task that the benchmark is done, then wait
1125 + if let Some((cancel_tx, task)) = profile_task {
1126 + let _ = cancel_tx.send(());
1127 + if let Err(e) = task.await {
1128 + eprintln!("WARNING: Profile background task failed: {e}");
1129 + }
1130 + }
1131 +
1132 + // Build result JSON
1133 + let date_iso = chrono::Utc::now().format("%Y-%m-%dT%H:%M:%SZ").to_string();
1134 + let dt_filename = chrono::Local::now().format("%Y%m%d-%H%M%S").to_string();
1135 + let result_json = build_result_json(
1136 + config,
1137 + &metrics,
1138 + &actual_output_lens,
1139 + &outputs,
1140 + benchmark_duration,
1141 + &date_iso,
1142 + spec_decode_stats.as_ref(),
1143 + );
1144 +
1145 + // Save if requested
1146 + if config.save_result || config.append_result {
1147 + let model_for_filename = config.model.as_deref().unwrap_or(&model_id);
1148 + let file_name = compute_result_filename(config, model_for_filename, &dt_filename);
1149 +
1150 + // Create result directory if needed
1151 + if let Some(ref dir) = config.result_dir {
1152 + std::fs::create_dir_all(dir)?;
1153 + }
1154 +
1155 + if config.append_result {
1156 + append_result(&result_json, &file_name)?;
1157 + } else {
1158 + save_result(&result_json, &file_name)?;
1159 + }
1160 + }
1161 +
1162 + Ok(result_json)
1163 +}
1164 +
1165 +fn filter_requests_by_max_model_len(
1166 + requests: &mut Vec<crate::datasets::SampleRequest>,
1167 + max_model_len: Option<usize>,
1168 +) -> usize {
1169 + let Some(max_model_len) = max_model_len else {
1170 + return 0;
1171 + };
1172 +
1173 + let before = requests.len();
1174 + requests.retain(|request| {
1175 + request.prompt_len.saturating_add(request.expected_output_len) <= max_model_len
1176 + });
1177 + before - requests.len()
1178 +}
1179 +
1180 +#[cfg(test)]
1181 +mod max_model_len_tests {
1182 + use std::sync::Arc;
1183 +
1184 + use super::filter_requests_by_max_model_len;
1185 + use crate::datasets::SampleRequest;
1186 +
1187 + fn sample(prompt_len: usize, expected_output_len: usize) -> SampleRequest {
1188 + SampleRequest {
1189 + prompt: Arc::from("prompt"),
1190 + prompt_len,
1191 + expected_output_len,
1192 + request_id: None,
1193 + ..Default::default()
1194 + }
1195 + }
1196 +
1197 + #[test]
1198 + fn test_filter_requests_by_max_model_len_keeps_boundary() {
1199 + let mut requests = vec![sample(80, 20), sample(81, 20), sample(30, 10)];
1200 +
1201 + let filtered = filter_requests_by_max_model_len(&mut requests, Some(100));
1202 +
1203 + assert_eq!(filtered, 1);
1204 + assert_eq!(requests.len(), 2);
1205 + assert_eq!(requests[0].prompt_len, 80);
1206 + assert_eq!(requests[1].prompt_len, 30);
1207 + }
1208 +
1209 + #[test]
1210 + fn test_filter_requests_by_max_model_len_noop_without_limit() {
1211 + let mut requests = vec![sample(80, 20), sample(81, 20)];
1212 +
1213 + let filtered = filter_requests_by_max_model_len(&mut requests, None);
1214 +
1215 + assert_eq!(filtered, 0);
1216 + assert_eq!(requests.len(), 2);
1217 + }
1218 +}
1219 +
1220 +async fn run_warmup(
1221 + backend: BackendKind,
1222 + client: &reqwest::Client,
1223 + test_input: &RequestFuncInput,
1224 + num_warmups: usize,
1225 + max_concurrency: Option<usize>,
1226 + request_rate: f64,
1227 + burstiness: f64,
1228 + seed: u64,
1229 + disable_tqdm: bool,
1230 +) {
1231 + let pb = if disable_tqdm {
1232 + None
1233 + } else {
1234 + let bar = ProgressBar::new(num_warmups as u64);
1235 + bar.set_style(
1236 + ProgressStyle::with_template("{spinner:.green} Warmup [{bar:30}] {pos}/{len}")
1237 + .unwrap()
1238 + .progress_chars("#>-"),
1239 + );
1240 + Some(bar)
1241 + };
1242 +
1243 + let semaphore = max_concurrency.map(|mc| Arc::new(Semaphore::new(mc)));
1244 +
1245 + // Use the same rate-limited scheduling as the main benchmark run
1246 + let schedule = compute_schedule(num_warmups, request_rate, burstiness, seed, None);
1247 + let start = Instant::now();
1248 +
1249 + let mut handles = Vec::with_capacity(num_warmups);
1250 + for i in 0..num_warmups {
1251 + // Wait until scheduled time
1252 + let target = std::time::Duration::from_secs_f64(schedule.delays[i]);
1253 + let elapsed = start.elapsed();
1254 + if target > elapsed {
1255 + tokio::time::sleep(target - elapsed).await;
1256 + }
1257 +
1258 + let client = client.clone();
1259 + let input = test_input.clone();
1260 + let sem = semaphore.clone();
1261 + let pb = pb.clone();
1262 + handles.push(tokio::spawn(async move {
1263 + let _permit = if let Some(ref s) = sem {
1264 + Some(s.acquire().await.unwrap())
1265 + } else {
1266 + None
1267 + };
1268 + if let Ok(b) = get_backend(backend) {
1269 + let _ = b.send_request(&input, &client).await;
1270 + }
1271 + if let Some(pb) = pb {
1272 + pb.inc(1);
1273 + }
1274 + }));
1275 + }
1276 +
1277 + for handle in handles {
1278 + let _ = handle.await;
1279 + }
1280 +
1281 + if let Some(pb) = pb {
1282 + pb.finish_and_clear();
1283 + }
1284 +}
1285 +
1286 +/// Start the profiler (immediate mode — no batch threshold).
1287 +pub(crate) async fn start_profiler_immediate(
1288 + client: &reqwest::Client,
1289 + base_url: &str,
1290 + extra_headers: &Option<std::collections::HashMap<String, String>>,
1291 +) {
1292 + println!("Starting profiler...");
1293 + let profile_url = format!("{base_url}/start_profile");
1294 + match send_profile_request(client, &profile_url, extra_headers).await {
1295 + Ok(true) => println!("Profiler started"),
1296 + Ok(false) => eprintln!("WARNING: Profiler start request returned non-success"),
1297 + Err(e) => eprintln!("WARNING: Failed to start profiler: {e}"),
1298 + }
1299 +}
1300 +
1301 +/// Stop the profiler (immediate mode — no batch threshold).
1302 +pub(crate) async fn stop_profiler_immediate(
1303 + client: &reqwest::Client,
1304 + base_url: &str,
1305 + extra_headers: &Option<std::collections::HashMap<String, String>>,
1306 +) {
1307 + println!("Stopping profiler...");
1308 + let profile_url = format!("{base_url}/stop_profile");
1309 + match send_profile_request(client, &profile_url, extra_headers).await {
1310 + Ok(true) => println!("Profiler stopped"),
1311 + Ok(false) => eprintln!("WARNING: Profiler stop request returned non-success"),
1312 + Err(e) => eprintln!("WARNING: Failed to stop profiler: {e}"),
1313 + }
1314 +}
1315 +
1316 +/// Send a request to the vLLM profiler endpoint (start_profile / stop_profile).
1317 +/// Returns Ok(true) if the server responded with 200.
1318 +pub(crate) async fn send_profile_request(
1319 + client: &reqwest::Client,
1320 + url: &str,
1321 + extra_headers: &Option<std::collections::HashMap<String, String>>,
1322 +) -> Result<bool> {
1323 + let mut request = client.post(url);
1324 + if let Some(headers) = extra_headers {
1325 + for (k, v) in headers {
1326 + request = request.header(k, v);
1327 + }
1328 + }
1329 + if let Ok(api_key) = std::env::var("OPENAI_API_KEY") {
1330 + request = request.header("Authorization", format!("Bearer {api_key}"));
1331 + }
1332 +
1333 + let resp = request
1334 + .send()
1335 + .await
1336 + .map_err(|e| BenchError::Backend(format!("Profile request to {url} failed: {e}")))?;
1337 +
1338 + Ok(resp.status().is_success())
1339 +}
1340 +
1341 +/// Parse `vllm:num_requests_running` from the server's `/metrics` Prometheus endpoint.
1342 +pub(crate) async fn fetch_num_requests_running(
1343 + client: &reqwest::Client,
1344 + base_url: &str,
1345 +) -> Option<usize> {
1346 + let url = format!("{base_url}/metrics");
1347 + let resp = client.get(&url).send().await.ok()?;
1348 + let text = resp.text().await.ok()?;
1349 + for line in text.lines() {
1350 + if let Some(rest) = line.strip_prefix("vllm:num_requests_running") {
1351 + // Line format: `vllm:num_requests_running{model_name="..."} 42`
1352 + // The value is after the last space.
1353 + if let Some(val_str) = rest.rsplit(' ').next()
1354 + && let Ok(v) = val_str.parse::<f64>()
1355 + {
1356 + return Some(v as usize);
1357 + }
1358 + }
1359 + }
1360 + None
1361 +}
1362 +
1363 +/// Poll `/metrics` until `num_requests_running >= threshold`, then start the
1364 +/// profiler, wait `duration_secs`, and stop it. If `cancel_rx` fires before
1365 +/// the threshold is reached (i.e. the benchmark finished), exits early.
1366 +pub(crate) async fn profile_on_batch_threshold(
1367 + client: &reqwest::Client,
1368 + base_url: &str,
1369 + extra_headers: &Option<HashMap<String, String>>,
1370 + threshold: usize,
1371 + duration_secs: f64,
1372 + mut cancel_rx: tokio::sync::oneshot::Receiver<()>,
1373 +) {
1374 + println!(
1375 + "Waiting for batch size >= {threshold} before starting profiler \
1376 + (will capture {duration_secs}s)..."
1377 + );
1378 +
1379 + loop {
1380 + if let Some(running) = fetch_num_requests_running(client, base_url).await
1381 + && running >= threshold
1382 + {
1383 + println!("Batch size {running} >= {threshold}, starting profiler...");
1384 + break;
1385 + }
1386 + // Wait 500ms or until the benchmark signals cancellation
1387 + tokio::select! {
1388 + _ = tokio::time::sleep(std::time::Duration::from_millis(500)) => {}
1389 + _ = &mut cancel_rx => {
1390 + eprintln!(
1391 + "NOTE: Benchmark finished before batch threshold {threshold} was reached; \
1392 + profiling skipped."
1393 + );
1394 + return;
1395 + }
1396 + }
1397 + }
1398 +
1399 + let start_url = format!("{base_url}/start_profile");
1400 + match send_profile_request(client, &start_url, extra_headers).await {
1401 + Ok(true) => println!("Profiler started"),
1402 + Ok(false) => {
1403 + eprintln!("WARNING: Profiler start request returned non-success");
1404 + return;
1405 + }
1406 + Err(e) => {
1407 + eprintln!("WARNING: Failed to start profiler: {e}");
1408 + return;
1409 + }
1410 + }
1411 +
1412 + // Wait for the capture duration, but exit early if the benchmark finishes
1413 + tokio::select! {
1414 + _ = tokio::time::sleep(std::time::Duration::from_secs_f64(duration_secs)) => {}
1415 + _ = &mut cancel_rx => {
1416 + println!("Benchmark finished, stopping profiler early...");
1417 + }
1418 + }
1419 +
1420 + let stop_url = format!("{base_url}/stop_profile");
1421 + match send_profile_request(client, &stop_url, extra_headers).await {
1422 + Ok(true) => println!("Profiler stopped after capturing"),
1423 + Ok(false) => eprintln!("WARNING: Profiler stop request returned non-success"),
1424 + Err(e) => eprintln!("WARNING: Failed to stop profiler: {e}"),
1425 + }
1426 +}
1427 +
1428 +/// Verify and fix prompt token lengths against the server's /tokenize + /detokenize.
1429 +/// For each prompt, if the server's token count doesn't match that prompt's own
1430 +/// target (its generated length plus the special tokens the server adds), adjust
1431 +/// the token sequence (pad/truncate) and detokenize, repeating until exact match.
1432 +/// Runs up to 64 prompts concurrently.
1433 +async fn verify_and_fix_prompt_lengths(
1434 + client: &reqwest::Client,
1435 + base_url: &str,
1436 + model: &str,
1437 + requests: &mut [crate::datasets::SampleRequest],
1438 + num_special: usize,
1439 + extra_headers: &Option<std::collections::HashMap<String, String>>,
1440 +) -> Result<()> {
1441 + let tokenize_url = Arc::new(format!("{base_url}/tokenize"));
1442 + let detokenize_url = Arc::new(format!("{base_url}/detokenize"));
1443 + let api_key = Arc::new(std::env::var("OPENAI_API_KEY").ok());
1444 + let extra_headers = Arc::new(extra_headers.clone());
1445 + let concurrency = 64;
1446 + let sem = Arc::new(Semaphore::new(concurrency));
1447 +
1448 + let pb = Arc::new(ProgressBar::new(requests.len() as u64));
1449 + pb.set_style(
1450 + ProgressStyle::with_template(
1451 + "{spinner:.green} Verifying [{bar:40.cyan/blue}] {pos}/{len} ({per_sec}, {eta})",
1452 + )
1453 + .unwrap()
1454 + .progress_chars("#>-"),
1455 + );
1456 +
1457 + let fixed_count = Arc::new(std::sync::atomic::AtomicUsize::new(0));
1458 +
1459 + // Spawn a task per prompt
1460 + let mut handles: Vec<tokio::task::JoinHandle<Result<(usize, String)>>> =
1461 + Vec::with_capacity(requests.len());
1462 + for (i, req) in requests.iter().enumerate() {
1463 + let client = client.clone();
1464 + let tok_url = tokenize_url.clone();
1465 + let detok_url = detokenize_url.clone();
1466 + let model = model.to_string();
1467 + let prompt = req.prompt.to_string(); // Convert Arc<str> to String for mutation
1468 + let expected_input_len = req.prompt_len + num_special;
1469 + let api_key = api_key.clone();
1470 + let headers = extra_headers.clone();
1471 + let sem = sem.clone();
1472 + let pb = pb.clone();
1473 + let fixed_count = fixed_count.clone();
1474 +
1475 + handles.push(tokio::spawn(async move {
1476 + let _permit = sem.acquire().await.unwrap();
1477 + let mut prompt = prompt;
1478 + let max_iterations = 20;
1479 + let mut needed_fix = false;
1480 + // Track systematic offset (e.g., server auto-adds BOS token).
1481 + // If the same excess appears twice consecutively, compensate.
1482 + let mut last_excess: Option<usize> = None;
1483 +
1484 + for _iter in 0..max_iterations {
1485 + let tokens = server_tokenize(
1486 + &client, &tok_url, &model, &prompt, &api_key, &headers, i,
1487 + ).await?;
1488 +
1489 + if tokens.len() == expected_input_len {
1490 + if needed_fix {
1491 + fixed_count.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
1492 + }
1493 + pb.inc(1);
1494 + return Ok((i, prompt));
1495 + }
1496 +
1497 + needed_fix = true;
1498 +
1499 + // Detect systematic offset: if the server consistently returns
1500 + // expected + N tokens (e.g., auto-prepended BOS), reduce target
1501 + // by N so the server's re-tokenization lands on the expected count.
1502 + let excess = tokens.len().saturating_sub(expected_input_len);
1503 + let compensate = if excess > 0 && last_excess == Some(excess) {
1504 + if _iter == 1 {
1505 + eprintln!(
1506 + "Prompt {i}: server consistently adds {excess} extra token(s) \
1507 + (likely BOS), compensating target to {}.",
1508 + expected_input_len.saturating_sub(excess),
1509 + );
1510 + }
1511 + excess
1512 + } else {
1513 + 0
1514 + };
1515 + last_excess = if excess > 0 { Some(excess) } else { None };
1516 +
1517 + let target = expected_input_len.saturating_sub(compensate);
1518 +
1519 + let mut adjusted = tokens;
1520 + if adjusted.is_empty() {
1521 + return Err(BenchError::Tokenizer(format!(
1522 + "Prompt {i}: server returned no tokens, cannot fix prompt length"
1523 + )));
1524 + }
1525 + if adjusted.len() < target {
1526 + let pad_needed = target - adjusted.len();
1527 + let original_len = adjusted.len();
1528 + for j in 0..pad_needed {
1529 + adjusted.push(adjusted[j % original_len]);
1530 + }
1531 + } else {
1532 + adjusted.truncate(target);
1533 + }
1534 +
1535 + prompt = server_detokenize(
1536 + &client, &detok_url, &model, &adjusted, &api_key, &headers, i,
1537 + ).await?;
1538 + }
1539 +
1540 + Err(BenchError::Tokenizer(format!(
1541 + "Prompt {i}: server verification failed to converge after {max_iterations} iterations"
1542 + )))
1543 + }));
1544 + }
1545 +
1546 + // Collect results and write back
1547 + for handle in handles {
1548 + let result = match handle.await {
1549 + Ok(r) => r,
1550 + Err(e) => {
1551 + return Err(BenchError::Tokenizer(format!(
1552 + "Verification task panicked: {e}"
1553 + )));
1554 + }
1555 + };
1556 + let (i, prompt) = result?;
1557 + requests[i].prompt = Arc::from(prompt);
1558 + // Server-side count the fix loop converged on for this prompt.
1559 + requests[i].prompt_len += num_special;
1560 + }
1561 +
1562 + pb.finish_and_clear();
1563 +
1564 + let fc = fixed_count.load(std::sync::atomic::Ordering::Relaxed);
1565 + if fc > 0 {
1566 + println!("Fixed {fc} prompt(s) via server tokenize/detokenize convergence.");
1567 + }
1568 +
1569 + Ok(())
1570 +}
1571 +
1572 +/// Call the server's /tokenize endpoint and return the token ID list.
1573 +/// Retries up to 3 times on transient errors (5xx, connection errors).
1574 +async fn server_tokenize(
1575 + client: &reqwest::Client,
1576 + url: &str,
1577 + model: &str,
1578 + prompt: &str,
1579 + api_key: &Option<String>,
1580 + extra_headers: &Option<std::collections::HashMap<String, String>>,
1581 + prompt_idx: usize,
1582 +) -> Result<Vec<u64>> {
1583 + let max_retries = 3;
1584 +
1585 + for attempt in 0..=max_retries {
1586 + let payload = serde_json::json!({
1587 + "model": model,
1588 + "prompt": prompt,
1589 + });
1590 +
1591 + let mut request = client.post(url).json(&payload);
1592 + if let Some(headers) = extra_headers {
1593 + for (k, v) in headers {
1594 + request = request.header(k, v);
1595 + }
1596 + }
1597 + if let Some(key) = api_key {
1598 + request = request.header("Authorization", format!("Bearer {key}"));
1599 + }
1600 +
1601 + let resp = match request.send().await {
1602 + Ok(r) => r,
1603 + Err(e) => {
1604 + if attempt < max_retries {
1605 + tokio::time::sleep(std::time::Duration::from_millis(
1606 + 500 * (attempt as u64 + 1),
1607 + ))
1608 + .await;
1609 + continue;
1610 + }
1611 + return Err(BenchError::Tokenizer(format!(
1612 + "Server /tokenize failed for prompt {prompt_idx} after {max_retries} retries: {e}"
1613 + )));
1614 + }
1615 + };
1616 +
1617 + if resp.status().is_server_error() && attempt < max_retries {
1618 + tokio::time::sleep(std::time::Duration::from_millis(500 * (attempt as u64 + 1))).await;
1619 + continue;
1620 + }
1621 +
1622 + if resp.status().is_client_error() {
1623 + // Server doesn't expose /tokenize (404, e.g. Dynamo) or a gateway
1624 + // rejects it with another 4xx (LLM-d/EPP returns 400) — caller
1625 + // should skip verification. 5xx and connection errors stay fatal.
1626 + return Err(BenchError::TokenizeUnavailable(format!(
1627 + "HTTP {}",
1628 + resp.status()
1629 + )));
1630 + }
1631 +
1632 + if !resp.status().is_success() {
1633 + return Err(BenchError::Tokenizer(format!(
1634 + "Server /tokenize returned HTTP {} for prompt {prompt_idx}",
1635 + resp.status()
1636 + )));
1637 + }
1638 +
1639 + let data: serde_json::Value = resp.json().await.map_err(|e| {
1640 + BenchError::Tokenizer(format!(
1641 + "Failed to parse /tokenize response for prompt {prompt_idx}: {e}"
1642 + ))
1643 + })?;
1644 +
1645 + let tokens = data.get("tokens").and_then(|t| t.as_array()).ok_or_else(|| {
1646 + BenchError::Tokenizer(format!(
1647 + "No 'tokens' array in /tokenize response for prompt {prompt_idx}"
1648 + ))
1649 + })?;
1650 +
1651 + return tokens
1652 + .iter()
1653 + .map(|v| {
1654 + v.as_u64().ok_or_else(|| {
1655 + BenchError::Tokenizer("Invalid token ID in /tokenize response".to_string())
1656 + })
1657 + })
1658 + .collect();
1659 + }
1660 +
1661 + unreachable!()
1662 +}
1663 +
1664 +/// Call the server's /detokenize endpoint and return the prompt text.
1665 +/// Retries up to 3 times on transient errors (5xx, connection errors).
1666 +async fn server_detokenize(
1667 + client: &reqwest::Client,
1668 + url: &str,
1669 + model: &str,
1670 + tokens: &[u64],
1671 + api_key: &Option<String>,
1672 + extra_headers: &Option<std::collections::HashMap<String, String>>,
1673 + prompt_idx: usize,
1674 +) -> Result<String> {
1675 + let max_retries = 3;
1676 +
1677 + for attempt in 0..=max_retries {
1678 + let payload = serde_json::json!({
1679 + "model": model,
1680 + "tokens": tokens,
1681 + });
1682 +
1683 + let mut request = client.post(url).json(&payload);
1684 + if let Some(headers) = extra_headers {
1685 + for (k, v) in headers {
1686 + request = request.header(k, v);
1687 + }
1688 + }
1689 + if let Some(key) = api_key {
1690 + request = request.header("Authorization", format!("Bearer {key}"));
1691 + }
1692 +
1693 + let resp = match request.send().await {
1694 + Ok(r) => r,
1695 + Err(e) => {
1696 + if attempt < max_retries {
1697 + tokio::time::sleep(std::time::Duration::from_millis(
1698 + 500 * (attempt as u64 + 1),
1699 + ))
1700 + .await;
1701 + continue;
1702 + }
1703 + return Err(BenchError::Tokenizer(format!(
1704 + "Server /detokenize failed for prompt {prompt_idx} after {max_retries} retries: {e}"
1705 + )));
1706 + }
1707 + };
1708 +
1709 + if resp.status().is_server_error() && attempt < max_retries {
1710 + tokio::time::sleep(std::time::Duration::from_millis(500 * (attempt as u64 + 1))).await;
1711 + continue;
1712 + }
1713 +
1714 + if resp.status().is_client_error() {
1715 + return Err(BenchError::TokenizeUnavailable(format!(
1716 + "HTTP {}",
1717 + resp.status()
1718 + )));
1719 + }
1720 +
1721 + if !resp.status().is_success() {
1722 + return Err(BenchError::Tokenizer(format!(
1723 + "Server /detokenize returned HTTP {} for prompt {prompt_idx}",
1724 + resp.status()
1725 + )));
1726 + }
1727 +
1728 + let data: serde_json::Value = resp.json().await.map_err(|e| {
1729 + BenchError::Tokenizer(format!(
1730 + "Failed to parse /detokenize response for prompt {prompt_idx}: {e}"
1731 + ))
1732 + })?;
1733 +
1734 + return data.get("prompt").and_then(|p| p.as_str()).map(|s| s.to_string()).ok_or_else(
1735 + || {
1736 + BenchError::Tokenizer(format!(
1737 + "No 'prompt' in /detokenize response for prompt {prompt_idx}"
1738 + ))
1739 + },
1740 + );
1741 + }
1742 +
1743 + unreachable!()
1744 +}
1745 +
1746 +/// Check if an error message indicates a transient connection error worth retrying.
1747 +fn is_connection_error(error: &str) -> bool {
1748 + let patterns = [
1749 + "connection reset",
1750 + "Connection reset",
1751 + "connection error",
1752 + "broken pipe",
1753 + "connect error",
1754 + "timed out",
1755 + "connection refused",
1756 + "Connection refused",
1757 + ];
1758 + patterns.iter().any(|p| error.contains(p))
1759 +}
1760 +
1761 +// --- Tokenizer verification cache ---
1762 +
1763 +/// Build a cache key from base_url + model_id.
1764 +fn tokenizer_verify_cache_key(base_url: &str, model_id: &str) -> String {
1765 + use std::collections::hash_map::DefaultHasher;
1766 + use std::hash::{Hash, Hasher};
1767 + let mut h = DefaultHasher::new();
1768 + base_url.hash(&mut h);
1769 + model_id.hash(&mut h);
1770 + format!("{:016x}", h.finish())
1771 +}
1772 +
1773 +/// Cache directory for tokenizer verification markers.
1774 +fn verify_cache_dir() -> std::path::PathBuf {
1775 + dirs::cache_dir()
1776 + .unwrap_or_else(|| std::path::PathBuf::from("/tmp"))
1777 + .join("vllm-bench")
1778 + .join("verified_tokenizers")
1779 +}
1780 +
1781 +/// Check if tokenizer was previously verified for this model+server.
1782 +fn is_tokenizer_verified(cache_key: &str) -> bool {
1783 + verify_cache_dir().join(cache_key).exists()
1784 +}
1785 +
1786 +/// Mark tokenizer as verified for this model+server.
1787 +fn mark_tokenizer_verified(cache_key: &str) {
1788 + let dir = verify_cache_dir();
1789 + let _ = std::fs::create_dir_all(&dir);
1790 + let _ = std::fs::write(dir.join(cache_key), "");
1791 +}
1792 +
1793 +/// Outcome of sample verification against the server's /tokenize.
1794 +enum SampleVerifyOutcome {
1795 + /// All sampled prompts matched their expected token counts.
1796 + Passed,
1797 + /// At least one sampled prompt did not match; full verify+fix is needed.
1798 + Mismatch,
1799 + /// The server cannot tokenize (4xx from /tokenize); verification skipped.
1800 + /// Unlike Passed, this must NOT be cached as "verified".
1801 + Skipped(String),
1802 +}
1803 +
1804 +/// Sample-verify a small number of prompts against the server's /tokenize.
1805 +/// Each prompt is checked against its own generated length (plus the special
1806 +/// tokens the server adds), so variable-length datasets
1807 +/// (--random-range-ratio < 1.0) and shared prefixes (--random-prefix-len)
1808 +/// verify correctly.
1809 +async fn sample_verify_prompts(
1810 + client: &reqwest::Client,
1811 + base_url: &str,
1812 + model: &str,
1813 + requests: &[crate::datasets::SampleRequest],
1814 + num_special: usize,
1815 + extra_headers: &Option<std::collections::HashMap<String, String>>,
1816 +) -> Result<SampleVerifyOutcome> {
1817 + let sample_size = 10.min(requests.len());
1818 + let tokenize_url = format!("{base_url}/tokenize");
1819 + let api_key = std::env::var("OPENAI_API_KEY").ok();
1820 +
1821 + println!("Sampling {sample_size} prompts for verification...");
1822 +
1823 + for (i, request) in requests.iter().enumerate().take(sample_size) {
1824 + let tokens = match server_tokenize(
1825 + client,
1826 + &tokenize_url,
1827 + model,
1828 + &request.prompt,
1829 + &api_key,
1830 + extra_headers,
1831 + i,
1832 + )
1833 + .await
1834 + {
1835 + Ok(t) => t,
1836 + Err(BenchError::TokenizeUnavailable(reason)) => {
1837 + return Ok(SampleVerifyOutcome::Skipped(reason));
1838 + }
1839 + Err(e) => return Err(e),
1840 + };
1841 +
1842 + let expected = request.prompt_len + num_special;
1843 + if tokens.len() != expected {
1844 + println!(
1845 + "Prompt {i}: expected {expected} tokens, server returned {}",
1846 + tokens.len()
1847 + );
1848 + return Ok(SampleVerifyOutcome::Mismatch);
1849 + }
1850 + }
1851 +
1852 + Ok(SampleVerifyOutcome::Passed)
1853 +}
1854 +
1855 +#[cfg(test)]
1856 +mod tests {
1857 + use super::*;
1858 +
1859 + fn make_metrics(
1860 + drafts: u64,
1861 + draft_tokens: u64,
1862 + accepted_tokens: u64,
1863 + per_pos: &[(u64, u64)],
1864 + ) -> SpecDecodeMetrics {
1865 + SpecDecodeMetrics {
1866 + num_drafts: drafts,
1867 + num_draft_tokens: draft_tokens,
1868 + num_accepted_tokens: accepted_tokens,
1869 + accepted_per_pos: per_pos.iter().copied().collect(),
1870 + }
1871 + }
1872 +
1873 + #[test]
1874 + fn test_compute_spec_decode_stats_basic() {
1875 + let before = make_metrics(100, 300, 200, &[(0, 90), (1, 70), (2, 50)]);
1876 + let after = make_metrics(200, 600, 440, &[(0, 180), (1, 150), (2, 120)]);
1877 +
1878 + let stats = compute_spec_decode_stats(&before, &after).unwrap();
1879 +
1880 + assert_eq!(stats.num_drafts, 100);
1881 + assert_eq!(stats.draft_tokens, 300);
1882 + assert_eq!(stats.accepted_tokens, 240);
1883 + assert!((stats.acceptance_rate - 80.0).abs() < 0.01);
1884 + assert!((stats.acceptance_length - 3.4).abs() < 0.01);
1885 + assert_eq!(stats.per_position_acceptance_rates.len(), 3);
1886 + assert!((stats.per_position_acceptance_rates[0] - 0.9).abs() < 0.01);
1887 + assert!((stats.per_position_acceptance_rates[1] - 0.8).abs() < 0.01);
1888 + assert!((stats.per_position_acceptance_rates[2] - 0.7).abs() < 0.01);
1889 + }
1890 +
1891 + #[test]
1892 + fn test_compute_spec_decode_stats_zero_draft_tokens_returns_none() {
1893 + let before = make_metrics(0, 100, 80, &[]);
1894 + let after = make_metrics(0, 100, 80, &[]);
1895 + assert!(compute_spec_decode_stats(&before, &after).is_none());
1896 + }
1897 +
1898 + #[test]
1899 + fn test_compute_spec_decode_stats_zero_drafts() {
1900 + let before = make_metrics(0, 0, 0, &[]);
1901 + let after = make_metrics(0, 100, 80, &[]);
1902 +
1903 + let stats = compute_spec_decode_stats(&before, &after).unwrap();
1904 + assert_eq!(stats.num_drafts, 0);
1905 + assert_eq!(stats.draft_tokens, 100);
1906 + assert_eq!(stats.accepted_tokens, 80);
1907 + assert!((stats.acceptance_rate - 80.0).abs() < 0.01);
1908 + assert!((stats.acceptance_length - 0.0).abs() < 0.01);
1909 + assert!(stats.per_position_acceptance_rates.is_empty());
1910 + }
1911 +
1912 + #[test]
1913 + fn test_compute_spec_decode_stats_new_positions_in_after() {
1914 + let before = make_metrics(50, 150, 100, &[(0, 40)]);
1915 + let after = make_metrics(150, 450, 320, &[(0, 130), (1, 80)]);
1916 +
1917 + let stats = compute_spec_decode_stats(&before, &after).unwrap();
1918 + assert_eq!(stats.num_drafts, 100);
1919 + assert_eq!(stats.per_position_acceptance_rates.len(), 2);
1920 + assert!((stats.per_position_acceptance_rates[0] - 0.9).abs() < 0.01);
1921 + assert!((stats.per_position_acceptance_rates[1] - 0.8).abs() < 0.01);
1922 + }
1923 +
1924 + fn lora_names(v: &[Arc<str>]) -> Vec<&str> {
1925 + v.iter().map(|s| s.as_ref()).collect()
1926 + }
1927 +
1928 + #[test]
1929 + fn test_assign_lora_modules_none_when_unset() {
1930 + assert!(assign_lora_modules(&None, LoraAssignment::Random, 5, 0).is_none());
1931 + }
1932 +
1933 + #[test]
1934 + fn test_assign_lora_modules_round_robin_cycles() {
1935 + let modules = Some(vec![Arc::<str>::from("a"), Arc::from("b"), Arc::from("c")]);
1936 + let out = assign_lora_modules(&modules, LoraAssignment::RoundRobin, 7, 42).unwrap();
1937 + assert_eq!(lora_names(&out), vec!["a", "b", "c", "a", "b", "c", "a"]);
1938 + }
1939 +
1940 + #[test]
1941 + fn test_assign_lora_modules_random_is_seed_reproducible() {
1942 + let modules = Some(vec![Arc::<str>::from("x"), Arc::from("y"), Arc::from("z")]);
1943 + let a = assign_lora_modules(&modules, LoraAssignment::Random, 100, 7).unwrap();
1944 + let b = assign_lora_modules(&modules, LoraAssignment::Random, 100, 7).unwrap();
1945 + assert_eq!(lora_names(&a), lora_names(&b));
1946 +
1947 + // Different seed should (almost certainly) produce a different sequence.
1948 + let c = assign_lora_modules(&modules, LoraAssignment::Random, 100, 8).unwrap();
1949 + assert_ne!(lora_names(&a), lora_names(&c));
1950 + }
1951 +
1952 + #[test]
1953 + fn test_assign_lora_modules_random_covers_all_modules() {
1954 + let modules = Some(vec![
1955 + Arc::<str>::from("a"),
1956 + Arc::from("b"),
1957 + Arc::from("c"),
1958 + Arc::from("d"),
1959 + ]);
1960 + let out = assign_lora_modules(&modules, LoraAssignment::Random, 1000, 0).unwrap();
1961 + let unique: std::collections::HashSet<&str> = out.iter().map(|s| s.as_ref()).collect();
1962 + assert_eq!(
1963 + unique.len(),
1964 + 4,
1965 + "all 4 adapters should be sampled in 1000 draws"
1966 + );
1967 + }
1968 +
1969 + #[test]
1970 + fn test_assign_lora_modules_single_adapter() {
1971 + let modules = Some(vec![Arc::<str>::from("only")]);
1972 + for assignment in [LoraAssignment::Random, LoraAssignment::RoundRobin] {
1973 + let out = assign_lora_modules(&modules, assignment, 5, 0).unwrap();
1974 + assert!(out.iter().all(|s| s.as_ref() == "only"));
1975 + }
1976 + }
1977 +}
rust/src/bench/src/cli.rs new
+750
@@ -0,0 +1,750 @@
1 +// SPDX-License-Identifier: Apache-2.0
2 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +use std::fmt;
5 +
6 +use clap::Parser;
7 +
8 +/// Backend type for the benchmark endpoint.
9 +#[derive(clap::ValueEnum, Debug, Clone, Copy, PartialEq, Eq)]
10 +pub enum BackendKind {
11 + #[value(name = "vllm")]
12 + Vllm,
13 + #[value(name = "openai")]
14 + Openai,
15 + #[value(name = "openai-chat")]
16 + OpenaiChat,
17 + #[value(name = "openai-embeddings")]
18 + OpenaiEmbeddings,
19 + #[value(name = "openai-embeddings-chat")]
20 + OpenaiEmbeddingsChat,
21 + #[value(name = "vllm-pooling")]
22 + VllmPooling,
23 + #[value(name = "vllm-rerank")]
24 + VllmRerank,
25 +}
26 +
27 +impl BackendKind {
28 + pub fn as_str(self) -> &'static str {
29 + match self {
30 + Self::Vllm => "vllm",
31 + Self::Openai => "openai",
32 + Self::OpenaiChat => "openai-chat",
33 + Self::OpenaiEmbeddings => "openai-embeddings",
34 + Self::OpenaiEmbeddingsChat => "openai-embeddings-chat",
35 + Self::VllmPooling => "vllm-pooling",
36 + Self::VllmRerank => "vllm-rerank",
37 + }
38 + }
39 +
40 + /// Return true if the backend is compatible with OpenAI-style API and sampling parameters.
41 + pub fn is_openai_compatible(self) -> bool {
42 + match self {
43 + Self::Vllm | Self::Openai | Self::OpenaiChat => true,
44 + Self::OpenaiEmbeddings
45 + | Self::OpenaiEmbeddingsChat
46 + | Self::VllmPooling
47 + | Self::VllmRerank => false,
48 + }
49 + }
50 +
51 + /// Return true if the backend is a pooling/embedding backend (non-generative).
52 + pub fn is_pooling(self) -> bool {
53 + matches!(
54 + self,
55 + Self::OpenaiEmbeddings
56 + | Self::OpenaiEmbeddingsChat
57 + | Self::VllmPooling
58 + | Self::VllmRerank
59 + )
60 + }
61 +}
62 +
63 +impl fmt::Display for BackendKind {
64 + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
65 + f.write_str(self.as_str())
66 + }
67 +}
68 +
69 +/// Dataset to benchmark with.
70 +#[derive(clap::ValueEnum, Debug, Clone, Copy, PartialEq, Eq)]
71 +pub enum DatasetName {
72 + #[value(name = "random")]
73 + Random,
74 + #[value(name = "random-mm")]
75 + RandomMm,
76 + #[value(name = "sharegpt")]
77 + ShareGpt,
78 + #[value(name = "sonnet")]
79 + Sonnet,
80 + #[value(name = "speed-bench")]
81 + SpeedBench,
82 + #[value(name = "hf")]
83 + Hf,
84 + #[value(name = "custom")]
85 + Custom,
86 + #[value(name = "prefix_repetition", alias = "prefix-repetition")]
87 + PrefixRepetition,
88 + #[value(name = "random-rerank")]
89 + RandomRerank,
90 +}
91 +
92 +/// Ramp-up strategy for request rate.
93 +#[derive(clap::ValueEnum, Debug, Clone, Copy, PartialEq, Eq)]
94 +pub enum RampUpStrategy {
95 + #[value(name = "linear")]
96 + Linear,
97 + #[value(name = "exponential")]
98 + Exponential,
99 +}
100 +
101 +/// Strategy for assigning LoRA modules to requests.
102 +#[derive(clap::ValueEnum, Debug, Clone, Copy, PartialEq, Eq)]
103 +pub enum LoraAssignment {
104 + #[value(name = "random")]
105 + Random,
106 + #[value(name = "round-robin")]
107 + RoundRobin,
108 +}
109 +
110 +/// SPEED-Bench dataset split/config.
111 +#[derive(clap::ValueEnum, Debug, Clone, Copy, PartialEq, Eq)]
112 +pub enum SpeedBenchConfig {
113 + #[value(name = "qualitative")]
114 + Qualitative,
115 + #[value(name = "throughput_1k")]
116 + Throughput1k,
117 + #[value(name = "throughput_2k")]
118 + Throughput2k,
119 + #[value(name = "throughput_8k")]
120 + Throughput8k,
121 + #[value(name = "throughput_16k")]
122 + Throughput16k,
123 + #[value(name = "throughput_32k")]
124 + Throughput32k,
125 +}
126 +
127 +impl SpeedBenchConfig {
128 + pub fn as_str(self) -> &'static str {
129 + match self {
130 + Self::Qualitative => "qualitative",
131 + Self::Throughput1k => "throughput_1k",
132 + Self::Throughput2k => "throughput_2k",
133 + Self::Throughput8k => "throughput_8k",
134 + Self::Throughput16k => "throughput_16k",
135 + Self::Throughput32k => "throughput_32k",
136 + }
137 + }
138 +}
139 +
140 +impl fmt::Display for SpeedBenchConfig {
141 + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
142 + f.write_str(self.as_str())
143 + }
144 +}
145 +
146 +/// High-performance benchmark client for vLLM serving endpoints.
147 +#[derive(Parser, Debug, Clone)]
148 +#[command(
149 + name = "vllm-bench",
150 + about = "Benchmark online serving throughput",
151 + version
152 +)]
153 +pub struct Cli {
154 + /// The type of backend or endpoint to use for the benchmark.
155 + #[arg(long, default_value = "openai")]
156 + pub backend: BackendKind,
157 +
158 + /// Server or API base url if not using http host and port.
159 + #[arg(long)]
160 + pub base_url: Option<String>,
161 +
162 + /// Server host.
163 + #[arg(long, default_value = "127.0.0.1")]
164 + pub host: String,
165 +
166 + /// Server port.
167 + #[arg(long, default_value_t = 8000)]
168 + pub port: u16,
169 +
170 + /// API endpoint. Auto-selected based on --backend if not specified.
171 + #[arg(long)]
172 + pub endpoint: Option<String>,
173 +
174 + /// Name of the model. If not specified, will fetch from server.
175 + #[arg(long)]
176 + pub model: Option<String>,
177 +
178 + /// The model name used in the API (for --served-model-name).
179 + #[arg(long)]
180 + pub served_model_name: Option<String>,
181 +
182 + /// Name or path of the tokenizer.
183 + #[arg(long)]
184 + pub tokenizer: Option<String>,
185 +
186 + /// Tokenizer mode (auto, hf, slow, mistral).
187 + #[arg(long, default_value = "auto")]
188 + pub tokenizer_mode: String,
189 +
190 + /// Skip initialization of tokenizer.
191 + #[arg(long, default_value_t = false)]
192 + pub skip_tokenizer_init: bool,
193 +
194 + /// Trust remote code for tokenizer.
195 + #[arg(long, default_value_t = false)]
196 + pub trust_remote_code: bool,
197 +
198 + /// Dataset name.
199 + #[arg(long, default_value = "random")]
200 + pub dataset_name: DatasetName,
201 +
202 + /// General input length for datasets.
203 + #[arg(long)]
204 + pub input_len: Option<usize>,
205 +
206 + /// General output length for datasets.
207 + #[arg(long)]
208 + pub output_len: Option<usize>,
209 +
210 + /// Maximum model context length. Requests with prompt_len + output_len above this are filtered
211 + /// out.
212 + #[arg(long)]
213 + pub max_model_len: Option<usize>,
214 +
215 + /// Random dataset input length.
216 + #[arg(long, default_value_t = 1024)]
217 + pub random_input_len: usize,
218 +
219 + /// Random dataset output length.
220 + #[arg(long, default_value_t = 128)]
221 + pub random_output_len: usize,
222 +
223 + /// Random dataset prefix length.
224 + #[arg(long, default_value_t = 0)]
225 + pub random_prefix_len: usize,
226 +
227 + /// Per-turn input length for turns 1+ in multi-turn mode.
228 + /// 0 = fallback to --random-input-len for all turns.
229 + /// Mirrors sglang bench_multiturn.py --sub-question-input-length.
230 + #[arg(long, default_value_t = 0)]
231 + pub per_turn_input_len: usize,
232 +
233 + /// Range ratio for sampling input/output lengths, matching Python
234 + /// `vllm bench serve`: lengths are drawn uniformly from
235 + /// [len*(1-r), len*(1+r)]. 0.0 (the default) = exact target lengths.
236 + /// Accepts a single float in [0, 1) or a JSON object
237 + /// '{"input": r1, "output": r2}' for independent control.
238 + /// NOTE: semantics changed — the old Rust-only form sampled [len*r, len]
239 + /// with default 1.0; old values like 1.0 are now rejected.
240 + #[arg(long, default_value = "0.0")]
241 + pub random_range_ratio: String,
242 +
243 + /// Batch multiple generated inputs into one request (embeddings/pooling
244 + /// backends only). E.g. 8 sends "input": [t1..t8] per request. Mirrors
245 + /// Python --random-batch-size. Default 1 = no batching.
246 + #[arg(long, default_value_t = 1)]
247 + pub random_batch_size: usize,
248 +
249 + /// random-rerank: the served model is NOT a reranker (embedding-based
250 + /// scoring). Changes query/document length accounting to mirror Python
251 + /// --no-reranker.
252 + #[arg(long, default_value_t = false)]
253 + pub no_reranker: bool,
254 +
255 + /// Bimodal prefix-cache (random dataset): fraction of prompts that are "warm"
256 + /// and reuse a shared cached prefix. 0.0 = off (default). E.g. 0.8 = 80% warm
257 + /// (prefix-cache hit), 20% cold (full prefill). Requires --random-cache-ratio > 0
258 + /// and --prompt-token-ids. In this mode --random-input-len is the TOTAL length.
259 + #[arg(long, default_value_t = 0.0)]
260 + pub random_cache_hit_fraction: f64,
261 +
262 + /// Bimodal prefix-cache (random dataset): fraction of each WARM prompt's length
263 + /// that is the shared cached prefix. 0.0 = off (default). E.g. 0.95 = 95% cached,
264 + /// 5% unique suffix. Used with --random-cache-hit-fraction.
265 + #[arg(long, default_value_t = 0.0)]
266 + pub random_cache_ratio: f64,
267 +
268 + /// Send prompt as token ID arrays instead of text strings.
269 + /// By default, prompts are decoded to text for maximum
270 + /// compatibility. Enable this for pure vLLM deployments to skip server-side
271 + /// tokenization (faster, exact token counts).
272 + #[arg(long, default_value_t = false)]
273 + pub prompt_token_ids: bool,
274 +
275 + // --- Random multimodal dataset ---
276 + /// Base number of multimodal items (images/videos) per request.
277 + #[arg(long, default_value_t = 1)]
278 + pub random_mm_base_items_per_request: usize,
279 +
280 + /// Range ratio for varying the number of multimodal items per request.
281 + /// Items sampled from [floor(n*(1-r)), ceil(n*(1+r))].
282 + #[arg(long, default_value_t = 0.0)]
283 + pub random_mm_num_mm_items_range_ratio: f64,
284 +
285 + /// Per-modality hard caps as JSON, e.g. '{"image": 3, "video": 0}'.
286 + #[arg(long, default_value = "{\"image\": 255, \"video\": 1}")]
287 + pub random_mm_limit_mm_per_prompt: String,
288 +
289 + /// Bucket config mapping (height,width,num_frames) to probability.
290 + /// Uses Python-style syntax: '{(256,256,1): 0.5, (720,1280,1): 0.5}'.
291 + /// num_frames=1 means image, num_frames>1 means video.
292 + #[arg(long, default_value = "{(256,256,1): 0.5, (720,1280,1): 0.5}")]
293 + pub random_mm_bucket_config: String,
294 +
295 + /// Enable multimodal chat transformation for datasets that support it.
296 + /// The dataset pre-builds the OpenAI chat `messages` array (text part +
297 + /// multimodal items) at generation time, and the request sends it verbatim.
298 + /// Mirrors Python's --enable-multimodal-chat. Currently applies to random-mm.
299 + #[arg(long, default_value_t = false)]
300 + pub enable_multimodal_chat: bool,
301 +
302 + // --- Custom dataset (JSONL) ---
303 + /// Output tokens per request for the custom dataset. Set to -1 to use the
304 + /// per-line "output_tokens" field from the JSONL file instead.
305 + #[arg(long, default_value_t = 256, allow_negative_numbers = true)]
306 + pub custom_output_len: i64,
307 +
308 + /// Skip applying a chat template to custom dataset prompts.
309 + /// NOTE: the Rust client never renders chat templates client-side, so this
310 + /// is always effectively on; passing it silences the informational notice.
311 + #[arg(long, default_value_t = false)]
312 + pub skip_chat_template: bool,
313 +
314 + // --- Prefix repetition dataset ---
315 + /// Shared-prefix token length for the prefix_repetition dataset.
316 + #[arg(long, default_value_t = 256)]
317 + pub prefix_repetition_prefix_len: usize,
318 +
319 + /// Per-request random suffix token length for the prefix_repetition dataset.
320 + #[arg(long, default_value_t = 256)]
321 + pub prefix_repetition_suffix_len: usize,
322 +
323 + /// Number of distinct shared prefixes for the prefix_repetition dataset.
324 + /// Requests are split evenly across prefixes (num-prompts / num-prefixes each).
325 + #[arg(long, default_value_t = 10)]
326 + pub prefix_repetition_num_prefixes: usize,
327 +
328 + /// Output tokens per request for the prefix_repetition dataset.
329 + #[arg(long, default_value_t = 128)]
330 + pub prefix_repetition_output_len: usize,
331 +
332 + /// Number of prompts to generate.
333 + #[arg(long, default_value_t = 1000)]
334 + pub num_prompts: usize,
335 +
336 + /// Number of requests per second. Use "inf" for all at once.
337 + #[arg(long, default_value_t = f64::INFINITY)]
338 + pub request_rate: f64,
339 +
340 + /// Burstiness factor of request generation.
341 + #[arg(long, default_value_t = 1.0)]
342 + pub burstiness: f64,
343 +
344 + /// Maximum number of concurrent requests.
345 + #[arg(long)]
346 + pub max_concurrency: Option<usize>,
347 +
348 + /// Fraction of --max-concurrency at which the steady-state window opens.
349 + /// Range: (0.0, 1.0]. Used only when --max-concurrency is set and
350 + /// --request-rate is inf.
351 + #[arg(long, default_value_t = 0.95)]
352 + pub steady_state_threshold: f64,
353 +
354 + /// Minimum steady-state window duration in seconds. Below this, a warning
355 + /// is attached. If unset, computed as max(10.0, 0.1 * run_duration).
356 + #[arg(long)]
357 + pub steady_state_min_window: Option<f64>,
358 +
359 + /// Disable steady-state metrics computation entirely.
360 + #[arg(long, default_value_t = false)]
361 + pub no_steady_state: bool,
362 +
363 + /// Disable tqdm progress bar.
364 + #[arg(long, default_value_t = false)]
365 + pub disable_tqdm: bool,
366 +
367 + /// Number of warmup requests.
368 + #[arg(long, default_value_t = 0)]
369 + pub num_warmups: usize,
370 +
371 + /// Use vLLM profiling. --profiler-config must be provided on the server.
372 + #[arg(long, default_value_t = false)]
373 + pub profile: bool,
374 +
375 + /// Minimum server batch size (num_requests_running) before starting the
376 + /// profiler. When set, profiling is deferred until the /metrics endpoint
377 + /// reports at least this many running requests, then captures for
378 + /// --profile-duration seconds. Requires --profile.
379 + #[arg(long)]
380 + pub profile_batch_threshold: Option<usize>,
381 +
382 + /// How many seconds to capture once the batch threshold is reached.
383 + /// Defaults to 5. Requires --profile and --profile-batch-threshold.
384 + #[arg(long, default_value_t = 5.0)]
385 + pub profile_duration: f64,
386 +
387 + /// Save benchmark results to a JSON file.
388 + #[arg(long, default_value_t = false)]
389 + pub save_result: bool,
390 +
391 + /// Save detailed per-request results.
392 + #[arg(long, default_value_t = false)]
393 + pub save_detailed: bool,
394 +
395 + /// Directory to save benchmark JSON results.
396 + #[arg(long)]
397 + pub result_dir: Option<String>,
398 +
399 + /// Filename to save benchmark JSON results.
400 + #[arg(long)]
401 + pub result_filename: Option<String>,
402 +
403 + /// Random seed.
404 + #[arg(long, default_value_t = 0)]
405 + pub seed: u64,
406 +
407 + /// Set ignore_eos flag when sending the benchmark request.
408 + #[arg(long, default_value_t = false)]
409 + pub ignore_eos: bool,
410 +
411 + /// Comma-separated list of metrics to report percentiles for.
412 + #[arg(long)]
413 + pub percentile_metrics: Option<String>,
414 +
415 + /// Comma-separated list of percentiles for selected metrics.
416 + #[arg(long, default_value = "99")]
417 + pub metric_percentiles: String,
418 +
419 + /// Comma-separated list of extra percentiles to show in sweep summaries.
420 + #[arg(long)]
421 + pub sweep_summary_percentiles: Option<String>,
422 +
423 + /// The label (prefix) of the benchmark results.
424 + #[arg(long)]
425 + pub label: Option<String>,
426 +
427 + /// Number of logprobs-per-token to compute.
428 + #[arg(long)]
429 + pub logprobs: Option<usize>,
430 +
431 + /// Prefix for request IDs.
432 + #[arg(long)]
433 + pub request_id_prefix: Option<String>,
434 +
435 + /// Maximum time to wait for endpoint readiness in seconds.
436 + #[arg(long, default_value_t = 0)]
437 + pub ready_check_timeout_sec: u64,
438 +
439 + /// Key-value pairs for extra headers (KEY=VALUE).
440 + #[arg(long = "header", num_args = 1..)]
441 + pub headers: Option<Vec<String>>,
442 +
443 + /// JSON string for extra body parameters.
444 + #[arg(long)]
445 + pub extra_body: Option<String>,
446 +
447 + /// Key-value pairs for metadata (KEY=VALUE).
448 + #[arg(long = "metadata", num_args = 1..)]
449 + pub metadata: Option<Vec<String>>,
450 +
451 + /// Dry run: only generate dataset and print stats, don't benchmark.
452 + #[arg(long, default_value_t = false)]
453 + pub dry_run: bool,
454 +
455 + // --- Sampling parameters ---
456 + /// Top-p sampling parameter. Only affects openai-compatible backends.
457 + #[arg(long)]
458 + pub top_p: Option<f64>,
459 +
460 + /// Top-k sampling parameter. Only affects openai-compatible backends.
461 + #[arg(long)]
462 + pub top_k: Option<i64>,
463 +
464 + /// Min-p sampling parameter. Only affects openai-compatible backends.
465 + #[arg(long)]
466 + pub min_p: Option<f64>,
467 +
468 + /// Temperature sampling parameter. Only affects openai-compatible backends.
469 + #[arg(long)]
470 + pub temperature: Option<f64>,
471 +
472 + /// Frequency penalty sampling parameter. Only affects openai-compatible backends.
473 + #[arg(long)]
474 + pub frequency_penalty: Option<f64>,
475 +
476 + /// Presence penalty sampling parameter. Only affects openai-compatible backends.
477 + #[arg(long)]
478 + pub presence_penalty: Option<f64>,
479 +
480 + /// Repetition penalty sampling parameter. Only affects openai-compatible backends.
481 + #[arg(long)]
482 + pub repetition_penalty: Option<f64>,
483 +
484 + // --- SSL ---
485 + /// Disable SSL certificate verification.
486 + #[arg(long, default_value_t = false)]
487 + pub insecure: bool,
488 +
489 + // --- Ramp-up ---
490 + /// Ramp-up strategy for request rate (linear or exponential).
491 + #[arg(long)]
492 + pub ramp_up_strategy: Option<RampUpStrategy>,
493 +
494 + /// Starting request rate for ramp-up (RPS).
495 + #[arg(long)]
496 + pub ramp_up_start_rps: Option<f64>,
497 +
498 + /// Ending request rate for ramp-up (RPS).
499 + #[arg(long)]
500 + pub ramp_up_end_rps: Option<f64>,
501 +
502 + // --- Goodput ---
503 + /// Service level objectives for goodput as "KEY:VALUE" pairs (e.g. ttft:100 tpot:50 e2el:500).
504 + /// Values are in milliseconds.
505 + #[arg(long = "goodput", num_args = 1..)]
506 + pub goodput: Option<Vec<String>>,
507 +
508 + // --- Result ---
509 + /// Append the benchmark result to the existing JSON file.
510 + #[arg(long, default_value_t = false)]
511 + pub append_result: bool,
512 +
513 + // --- ShareGPT dataset ---
514 + /// Path to dataset file (required for sharegpt dataset).
515 + #[arg(long)]
516 + pub dataset_path: Option<String>,
517 +
518 + /// Override output length for ShareGPT dataset.
519 + #[arg(long)]
520 + pub sharegpt_output_len: Option<usize>,
521 +
522 + /// Do not oversample if dataset is smaller than num_prompts.
523 + #[arg(long, default_value_t = false)]
524 + pub no_oversample: bool,
525 +
526 + /// Do not shuffle the dataset.
527 + #[arg(long, default_value_t = false)]
528 + pub disable_shuffle: bool,
529 +
530 + // --- Sonnet dataset ---
531 + /// Number of input tokens per request (sonnet dataset).
532 + #[arg(long, default_value_t = crate::datasets::sonnet::DEFAULT_INPUT_LEN)]
533 + pub sonnet_input_len: usize,
534 +
535 + /// Number of output tokens per request (sonnet dataset).
536 + #[arg(long, default_value_t = crate::datasets::sonnet::DEFAULT_OUTPUT_LEN)]
537 + pub sonnet_output_len: usize,
538 +
539 + /// Number of prefix tokens shared across requests (sonnet dataset).
540 + #[arg(long, default_value_t = crate::datasets::sonnet::DEFAULT_PREFIX_LEN)]
541 + pub sonnet_prefix_len: usize,
542 +
543 + /// SPEED-Bench config/split (qualitative, throughput_1k, throughput_2k, throughput_8k,
544 + /// throughput_16k, throughput_32k).
545 + #[arg(long, default_value = "qualitative")]
546 + pub speed_bench_config: SpeedBenchConfig,
547 +
548 + /// Filter SPEED-Bench by category (e.g. low_entropy, high_entropy, coding, math).
549 + #[arg(long)]
550 + pub speed_bench_category: Option<String>,
551 +
552 + /// Truncate SPEED-Bench prompts to at most this many tokens.
553 + /// Useful for creating custom input lengths from larger splits (e.g. --speed-bench-config
554 + /// throughput_16k --speed-bench-max-input-len 10240).
555 + #[arg(long)]
556 + pub speed_bench_max_input_len: Option<usize>,
557 +
558 + // --- HuggingFace dataset ---
559 + /// HuggingFace dataset split (e.g. train, test, validation).
560 + #[arg(long)]
561 + pub hf_split: Option<String>,
562 +
563 + /// HuggingFace dataset subset/config name.
564 + #[arg(long)]
565 + pub hf_subset: Option<String>,
566 +
567 + /// Fixed output length for HF dataset requests (overrides dataset-derived length).
568 + #[arg(long)]
569 + pub hf_output_len: Option<usize>,
570 +
571 + /// Column name containing the prompt text. Auto-detected if not specified.
572 + #[arg(long)]
573 + pub hf_text_column: Option<String>,
574 +
575 + // --- Compare mode ---
576 + /// Compare two benchmark result JSON files (e.g. --compare a.json b.json).
577 + /// Prints side-by-side metrics with delta and % change. Skips benchmarking.
578 + #[arg(long = "compare", num_args = 2, value_names = ["FILE_A", "FILE_B"])]
579 + pub compare: Option<Vec<String>>,
580 +
581 + // --- Sweep mode ---
582 + /// Sweep over max-concurrency values (comma-separated, e.g. --sweep-max-concurrency
583 + /// 1,10,50,100,500).
584 + #[arg(long)]
585 + pub sweep_max_concurrency: Option<String>,
586 +
587 + /// When sweeping concurrency, set num_prompts = concurrency * this factor for each sweep
588 + /// point.
589 + #[arg(long)]
590 + pub sweep_num_prompts_factor: Option<usize>,
591 +
592 + /// Sweep over request-rate values (comma-separated, supports "inf", e.g. --sweep-request-rate
593 + /// 1,10,100,inf).
594 + #[arg(long)]
595 + pub sweep_request_rate: Option<String>,
596 +
597 + /// Reset the server's prefix cache before each sweep iteration.
598 + /// Requires VLLM_SERVER_DEV_MODE=1 on the vLLM server.
599 + #[arg(long, default_value_t = false)]
600 + pub reset_prefix_cache: bool,
601 +
602 + // --- Multi-run ---
603 + /// Number of benchmark runs for statistical aggregation.
604 + #[arg(long, default_value_t = 1)]
605 + pub num_runs: usize,
606 +
607 + // --- Multi-turn conversation benchmark ---
608 + /// Enable multi-turn conversation benchmark mode.
609 + #[arg(long, default_value_t = false)]
610 + pub multi_turn: bool,
611 +
612 + /// Number of turns per conversation in synthetic multi-turn mode.
613 + #[arg(long, default_value_t = 3)]
614 + pub multi_turn_num_turns: usize,
615 +
616 + /// Minimum turns per conversation. 0 = use --multi-turn-num-turns.
617 + #[arg(long, default_value_t = 0)]
618 + pub multi_turn_min_turns: usize,
619 +
620 + /// Maximum turns per conversation.
621 + /// For synthetic multi-turn, 0 = use --multi-turn-num-turns.
622 + /// For ShareGPT multi-turn, 0 = uncapped.
623 + #[arg(long, default_value_t = 0)]
624 + pub multi_turn_max_turns: usize,
625 +
626 + /// Number of concurrent conversations (defaults to max-concurrency or num-prompts).
627 + #[arg(long)]
628 + pub multi_turn_concurrency: Option<usize>,
629 +
630 + /// Delay between turns in milliseconds (simulates user think time).
631 + #[arg(long, default_value_t = 0)]
632 + pub multi_turn_delay_ms: u64,
633 +
634 + /// Fraction of per-turn input tokens shared across ALL conversations (0.0–1.0).
635 + /// When > 0, enables prefix sharing mode: each turn sends a fixed-length message
636 + /// (no history accumulation). Only works with --dataset-name random.
637 + #[arg(long, default_value_t = 0.0)]
638 + pub multi_turn_prefix_global_ratio: f64,
639 +
640 + /// Fraction of per-turn input tokens shared within each conversation (0.0–1.0).
641 + /// When > 0, enables prefix sharing mode: each turn sends a fixed-length message
642 + /// (no history accumulation). Only works with --dataset-name random.
643 + #[arg(long, default_value_t = 0.0)]
644 + pub multi_turn_prefix_conversation_ratio: f64,
645 +
646 + // --- LoRA ---
647 + /// LoRA adapter names registered on the server (server-side
648 + /// `--lora-modules name=path`). Each request's `model` field is rewritten
649 + /// to one of these names; tokenizer and other endpoints keep using --model.
650 + /// In multi-turn mode, one adapter is assigned per conversation (sticky
651 + /// across turns).
652 + #[arg(long = "lora-modules", num_args = 1..)]
653 + pub lora_modules: Option<Vec<String>>,
654 +
655 + /// Strategy for assigning LoRA adapters to requests.
656 + /// 'random' (default) picks uniformly at random; 'round-robin' cycles
657 + /// through `--lora-modules` deterministically (i % N).
658 + #[arg(long = "lora-assignment", default_value = "random")]
659 + pub lora_assignment: LoraAssignment,
660 +}
661 +
662 +impl Cli {
663 + /// Resolve the base URL from explicit --base-url or from --host/--port.
664 + pub fn resolve_base_url(&self) -> String {
665 + if let Some(ref base) = self.base_url {
666 + base.clone()
667 + } else {
668 + format!("http://{}:{}", self.host, self.port)
669 + }
670 + }
671 +
672 + /// Resolve the API endpoint, auto-selecting based on backend if not explicit.
673 + pub fn resolve_endpoint(&self) -> String {
674 + if let Some(ref ep) = self.endpoint {
675 + return ep.clone();
676 + }
677 + match self.backend {
678 + BackendKind::OpenaiChat => "/v1/chat/completions".to_string(),
679 + BackendKind::Vllm | BackendKind::Openai => "/v1/completions".to_string(),
680 + BackendKind::OpenaiEmbeddings | BackendKind::OpenaiEmbeddingsChat => {
681 + "/v1/embeddings".to_string()
682 + }
683 + BackendKind::VllmPooling => "/v1/pooling".to_string(),
684 + BackendKind::VllmRerank => "/v1/rerank".to_string(),
685 + }
686 + }
687 +
688 + /// Resolve the full API URL.
689 + pub fn resolve_api_url(&self) -> String {
690 + format!("{}{}", self.resolve_base_url(), self.resolve_endpoint())
691 + }
692 +
693 + /// Parse extra headers from KEY=VALUE pairs.
694 + pub fn parse_headers(
695 + &self,
696 + ) -> crate::error::Result<Option<std::collections::HashMap<String, String>>> {
697 + match &self.headers {
698 + None => Ok(None),
699 + Some(items) => {
700 + let mut map = std::collections::HashMap::new();
701 + for item in items {
702 + let (k, v) = item.split_once('=').ok_or_else(|| {
703 + crate::error::BenchError::Config(
704 + "Invalid header format. Use KEY=VALUE".into(),
705 + )
706 + })?;
707 + map.insert(k.trim().to_string(), v.trim().to_string());
708 + }
709 + Ok(Some(map))
710 + }
711 + }
712 + }
713 +
714 + /// Parse extra body JSON.
715 + pub fn parse_extra_body(&self) -> crate::error::Result<Option<serde_json::Value>> {
716 + match &self.extra_body {
717 + None => Ok(None),
718 + Some(s) => {
719 + let v: serde_json::Value = serde_json::from_str(s).map_err(|e| {
720 + crate::error::BenchError::Config(format!("Invalid --extra-body JSON: {e}"))
721 + })?;
722 + Ok(Some(v))
723 + }
724 + }
725 + }
726 +
727 + /// Generate the request ID prefix (auto-generate if not provided).
728 + pub fn get_request_id_prefix(&self) -> String {
729 + self.request_id_prefix
730 + .clone()
731 + .unwrap_or_else(|| format!("bench-{}-", &uuid::Uuid::new_v4().to_string()[..8]))
732 + }
733 +
734 + /// Resolve input/output lengths, applying --input-len/--output-len overrides.
735 + pub fn resolved_random_input_len(&self) -> usize {
736 + self.input_len.unwrap_or(self.random_input_len)
737 + }
738 +
739 + pub fn resolved_random_output_len(&self) -> usize {
740 + self.output_len.unwrap_or(self.random_output_len)
741 + }
742 +
743 + pub fn resolved_per_turn_input_len(&self) -> usize {
744 + if self.per_turn_input_len > 0 {
745 + self.per_turn_input_len
746 + } else {
747 + self.resolved_random_input_len()
748 + }
749 + }
750 +}
rust/src/bench/src/compare.rs new
+302
@@ -0,0 +1,302 @@
1 +// SPDX-License-Identifier: Apache-2.0
2 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +use crate::error::{BenchError, Result};
5 +
6 +/// Metric definition for comparison: name, JSON key, and whether lower is better.
7 +struct MetricDef {
8 + label: &'static str,
9 + key: &'static str,
10 + lower_is_better: bool,
11 +}
12 +
13 +const METRICS: &[MetricDef] = &[
14 + MetricDef {
15 + label: "Request throughput (req/s)",
16 + key: "request_throughput",
17 + lower_is_better: false,
18 + },
19 + MetricDef {
20 + label: "Output throughput (tok/s)",
21 + key: "output_throughput",
22 + lower_is_better: false,
23 + },
24 + MetricDef {
25 + label: "Total token throughput (tok/s)",
26 + key: "total_token_throughput",
27 + lower_is_better: false,
28 + },
29 + MetricDef {
30 + label: "Peak output tokens/s",
31 + key: "max_output_tokens_per_s",
32 + lower_is_better: false,
33 + },
34 + MetricDef {
35 + label: "Peak concurrent requests",
36 + key: "max_concurrent_requests",
37 + lower_is_better: false,
38 + },
39 + MetricDef {
40 + label: "Mean TTFT (ms)",
41 + key: "mean_ttft_ms",
42 + lower_is_better: true,
43 + },
44 + MetricDef {
45 + label: "Median TTFT (ms)",
46 + key: "median_ttft_ms",
47 + lower_is_better: true,
48 + },
49 + MetricDef {
50 + label: "P99 TTFT (ms)",
51 + key: "p99_ttft_ms",
52 + lower_is_better: true,
53 + },
54 + MetricDef {
55 + label: "Mean TPOT (ms)",
56 + key: "mean_tpot_ms",
57 + lower_is_better: true,
58 + },
59 + MetricDef {
60 + label: "Median TPOT (ms)",
61 + key: "median_tpot_ms",
62 + lower_is_better: true,
63 + },
64 + MetricDef {
65 + label: "P99 TPOT (ms)",
66 + key: "p99_tpot_ms",
67 + lower_is_better: true,
68 + },
69 + MetricDef {
70 + label: "Mean ITL (ms)",
71 + key: "mean_itl_ms",
72 + lower_is_better: true,
73 + },
74 + MetricDef {
75 + label: "Median ITL (ms)",
76 + key: "median_itl_ms",
77 + lower_is_better: true,
78 + },
79 + MetricDef {
80 + label: "P99 ITL (ms)",
81 + key: "p99_itl_ms",
82 + lower_is_better: true,
83 + },
84 + MetricDef {
85 + label: "Mean E2EL (ms)",
86 + key: "mean_e2el_ms",
87 + lower_is_better: true,
88 + },
89 + MetricDef {
90 + label: "Median E2EL (ms)",
91 + key: "median_e2el_ms",
92 + lower_is_better: true,
93 + },
94 + MetricDef {
95 + label: "P99 E2EL (ms)",
96 + key: "p99_e2el_ms",
97 + lower_is_better: true,
98 + },
99 + MetricDef {
100 + label: "Completed requests",
101 + key: "completed",
102 + lower_is_better: false,
103 + },
104 + MetricDef {
105 + label: "Failed requests",
106 + key: "failed",
107 + lower_is_better: true,
108 + },
109 + MetricDef {
110 + label: "Duration (s)",
111 + key: "duration",
112 + lower_is_better: true,
113 + },
114 +];
115 +
116 +const STEADY_STATE_METRICS: &[MetricDef] = &[
117 + MetricDef {
118 + label: "SS Request throughput (req/s)",
119 + key: "request_throughput",
120 + lower_is_better: false,
121 + },
122 + MetricDef {
123 + label: "SS Output throughput (tok/s)",
124 + key: "output_throughput",
125 + lower_is_better: false,
126 + },
127 + MetricDef {
128 + label: "SS Input throughput (tok/s)",
129 + key: "input_throughput",
130 + lower_is_better: false,
131 + },
132 + MetricDef {
133 + label: "SS Total token throughput (tok/s)",
134 + key: "total_token_throughput",
135 + lower_is_better: false,
136 + },
137 + MetricDef {
138 + label: "SS Mean TTFT (ms)",
139 + key: "mean_ttft_ms",
140 + lower_is_better: true,
141 + },
142 + MetricDef {
143 + label: "SS Median TTFT (ms)",
144 + key: "median_ttft_ms",
145 + lower_is_better: true,
146 + },
147 + MetricDef {
148 + label: "SS Mean TPOT (ms)",
149 + key: "mean_tpot_ms",
150 + lower_is_better: true,
151 + },
152 + MetricDef {
153 + label: "SS Median TPOT (ms)",
154 + key: "median_tpot_ms",
155 + lower_is_better: true,
156 + },
157 + MetricDef {
158 + label: "SS P90 TPOT (ms)",
159 + key: "p90_tpot_ms",
160 + lower_is_better: true,
161 + },
162 + MetricDef {
163 + label: "SS P99 TPOT (ms)",
164 + key: "p99_tpot_ms",
165 + lower_is_better: true,
166 + },
167 +];
168 +
169 +/// Compare two benchmark result JSON files and print a side-by-side table.
170 +pub fn compare_results(file_a: &str, file_b: &str) -> Result<()> {
171 + let json_a = load_result_json(file_a)?;
172 + let json_b = load_result_json(file_b)?;
173 +
174 + // Print header with file context
175 + let model_a = json_a.get("model_id").and_then(|v| v.as_str()).unwrap_or("?");
176 + let model_b = json_b.get("model_id").and_then(|v| v.as_str()).unwrap_or("?");
177 + let date_a = json_a.get("date").and_then(|v| v.as_str()).unwrap_or("?");
178 + let date_b = json_b.get("date").and_then(|v| v.as_str()).unwrap_or("?");
179 +
180 + println!("{:=^90}", " Benchmark Comparison ");
181 + println!(" A: {} (model: {}, date: {})", file_a, model_a, date_a);
182 + println!(" B: {} (model: {}, date: {})", file_b, model_b, date_b);
183 + println!();
184 +
185 + // Print comparison table
186 + println!(
187 + "{:<35} {:>12} {:>12} {:>10} {:>8}",
188 + "Metric", "A", "B", "Delta", "Change"
189 + );
190 + println!("{:-<35} {:->12} {:->12} {:->10} {:->8}", "", "", "", "", "");
191 +
192 + for metric in METRICS {
193 + let val_a = get_f64(&json_a, metric.key);
194 + let val_b = get_f64(&json_b, metric.key);
195 +
196 + match (val_a, val_b) {
197 + (Some(a), Some(b)) => print_diff_row(metric, a, b),
198 + _ => {
199 + // One or both values missing — skip
200 + }
201 + }
202 + }
203 +
204 + // Steady-state section — both sides must have the block; otherwise render N/A.
205 + let ss_a = json_a.get("steady_state");
206 + let ss_b = json_b.get("steady_state");
207 + let both_present =
208 + matches!(ss_a, Some(v) if !v.is_null()) && matches!(ss_b, Some(v) if !v.is_null());
209 +
210 + println!();
211 + println!("{:=^70}", " Steady-State Comparison ");
212 + if !both_present {
213 + println!("N/A — one or both runs have no steady-state window");
214 + } else {
215 + let ss_a = ss_a.unwrap();
216 + let ss_b = ss_b.unwrap();
217 + for m in STEADY_STATE_METRICS {
218 + let a = ss_a.get(m.key).and_then(|v| v.as_f64());
219 + let b = ss_b.get(m.key).and_then(|v| v.as_f64());
220 + match (a, b) {
221 + (Some(a), Some(b)) => print_diff_row(m, a, b),
222 + _ => println!("{:<35} N/A", m.label),
223 + }
224 + }
225 + }
226 +
227 + println!("{:=<90}", "");
228 + println!();
229 + println!("Legend: + = improvement, - = regression (relative to A → B)");
230 +
231 + Ok(())
232 +}
233 +
234 +fn print_diff_row(metric: &MetricDef, a: f64, b: f64) {
235 + let delta = b - a;
236 + let pct = if a.abs() > 1e-10 {
237 + (delta / a) * 100.0
238 + } else if b.abs() > 1e-10 {
239 + f64::INFINITY
240 + } else {
241 + 0.0
242 + };
243 +
244 + // Determine if change is good/bad/neutral
245 + let marker = if delta.abs() < 1e-10 {
246 + " "
247 + } else if metric.lower_is_better {
248 + if delta < 0.0 { "+" } else { "-" }
249 + } else if delta > 0.0 {
250 + "+"
251 + } else {
252 + "-"
253 + };
254 +
255 + let delta_str = format_delta(delta);
256 + let pct_str = if pct.is_infinite() {
257 + "inf%".to_string()
258 + } else {
259 + format!("{:+.1}%", pct)
260 + };
261 +
262 + println!(
263 + "{:<35} {:>12} {:>12} {:>10} {:>7}{}",
264 + metric.label,
265 + format_value(a),
266 + format_value(b),
267 + delta_str,
268 + pct_str,
269 + marker,
270 + );
271 +}
272 +
273 +fn load_result_json(path: &str) -> Result<serde_json::Value> {
274 + let content = std::fs::read_to_string(path)
275 + .map_err(|e| BenchError::Config(format!("Cannot read result file '{path}': {e}")))?;
276 +
277 + // Support JSONL: take the last line (most recent run)
278 + let json_str = content.lines().rfind(|l| !l.trim().is_empty()).unwrap_or(&content);
279 +
280 + serde_json::from_str(json_str)
281 + .map_err(|e| BenchError::Config(format!("Cannot parse JSON from '{path}': {e}")))
282 +}
283 +
284 +fn get_f64(json: &serde_json::Value, key: &str) -> Option<f64> {
285 + json.get(key).and_then(|v| v.as_f64())
286 +}
287 +
288 +fn format_value(v: f64) -> String {
289 + if v == v.floor() && v.abs() < 1e12 {
290 + format!("{}", v as i64)
291 + } else {
292 + format!("{:.2}", v)
293 + }
294 +}
295 +
296 +fn format_delta(d: f64) -> String {
297 + if d == d.floor() && d.abs() < 1e12 {
298 + format!("{:+}", d as i64)
299 + } else {
300 + format!("{:+.2}", d)
301 + }
302 +}
rust/src/bench/src/config.rs new
+1116
@@ -0,0 +1,1116 @@
1 +// SPDX-License-Identifier: Apache-2.0
2 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +use std::collections::HashMap;
5 +use std::sync::Arc;
6 +
7 +use crate::cli::{BackendKind, Cli, DatasetName, LoraAssignment, RampUpStrategy, SpeedBenchConfig};
8 +use crate::datasets::random_mm::{MmBucketKey, MmLimitPerPrompt};
9 +use crate::error::{BenchError, Result};
10 +
11 +/// Parsed goodput SLO configuration.
12 +#[derive(Debug, Clone, Default)]
13 +pub struct GoodputConfig {
14 + /// TTFT SLO in milliseconds (None = not checked).
15 + pub ttft_ms: Option<f64>,
16 + /// TPOT SLO in milliseconds (None = not checked).
17 + pub tpot_ms: Option<f64>,
18 + /// E2EL SLO in milliseconds (None = not checked).
19 + pub e2el_ms: Option<f64>,
20 +}
21 +
22 +impl GoodputConfig {
23 + pub fn is_empty(&self) -> bool {
24 + self.ttft_ms.is_none() && self.tpot_ms.is_none() && self.e2el_ms.is_none()
25 + }
26 +}
27 +
28 +/// Ramp-up configuration.
29 +#[derive(Debug, Clone)]
30 +pub struct RampUpConfig {
31 + pub strategy: RampUpStrategy,
32 + pub start_rps: f64,
33 + pub end_rps: f64,
34 +}
35 +
36 +/// Range ratio for sampling input/output lengths, matching Python
37 +/// `vllm bench serve`: lengths are drawn uniformly from [len*(1-r), len*(1+r)].
38 +/// A single float applies to both; the JSON form '{"input": r1, "output": r2}'
39 +/// controls them independently. Each ratio must be in [0, 1).
40 +#[derive(Debug, Clone, Copy, PartialEq)]
41 +pub struct RangeRatio {
42 + pub input: f64,
43 + pub output: f64,
44 +}
45 +
46 +impl RangeRatio {
47 + /// Parse `--random-range-ratio`: a bare float or a JSON object with
48 + /// "input" and "output" keys.
49 + pub fn parse(raw: &str) -> Result<Self> {
50 + let trimmed = raw.trim();
51 + let (input, output) = if let Ok(v) = trimmed.parse::<f64>() {
52 + (v, v)
53 + } else {
54 + let v: serde_json::Value = serde_json::from_str(trimmed).map_err(|_| {
55 + BenchError::Config(format!(
56 + "Invalid --random-range-ratio '{raw}': expected a float or \
57 + '{{\"input\": r1, \"output\": r2}}'"
58 + ))
59 + })?;
60 + let obj = v.as_object().ok_or_else(|| {
61 + BenchError::Config(
62 + "--random-range-ratio JSON form must be an object with \
63 + 'input' and 'output' keys"
64 + .into(),
65 + )
66 + })?;
67 + let get = |key: &str| -> Result<f64> {
68 + obj.get(key).and_then(|v| v.as_f64()).ok_or_else(|| {
69 + BenchError::Config(format!(
70 + "--random-range-ratio JSON form must contain a numeric '{key}' key"
71 + ))
72 + })
73 + };
74 + (get("input")?, get("output")?)
75 + };
76 +
77 + for (name, r) in [("input", input), ("output", output)] {
78 + if !(0.0..1.0).contains(&r) {
79 + let hint = if r == 1.0 {
80 + " NOTE: semantics now match Python vllm bench serve — lengths are \
81 + sampled from [len*(1-r), len*(1+r)] and 0.0 means fixed length. \
82 + The old Rust-only default 1.0 ([len*r, len]) is no longer valid."
83 + } else {
84 + ""
85 + };
86 + return Err(BenchError::Config(format!(
87 + "--random-range-ratio {name} ratio must be in [0, 1), got {r}.{hint}"
88 + )));
89 + }
90 + }
91 + Ok(Self { input, output })
92 + }
93 +
94 + /// Sampling interval for input lengths: [floor(len*(1-r)), ceil(len*(1+r))].
95 + pub fn input_bounds(&self, len: usize) -> (usize, usize) {
96 + let low = ((len as f64) * (1.0 - self.input)).floor() as usize;
97 + let high = ((len as f64) * (1.0 + self.input)).ceil() as usize;
98 + (low, high)
99 + }
100 +
101 + /// Sampling interval for output lengths, clamped to at least 1 token.
102 + pub fn output_bounds(&self, len: usize) -> (usize, usize) {
103 + let low = (((len as f64) * (1.0 - self.output)).floor() as usize).max(1);
104 + let high = (((len as f64) * (1.0 + self.output)).ceil() as usize).max(1);
105 + (low, high)
106 + }
107 +
108 + /// True when both ratios are 0 (every request uses the exact target lengths).
109 + pub fn is_fixed(&self) -> bool {
110 + self.input == 0.0 && self.output == 0.0
111 + }
112 +}
113 +
114 +/// Validated benchmark configuration derived from CLI args.
115 +#[derive(Debug, Clone)]
116 +pub struct BenchConfig {
117 + pub backend: BackendKind,
118 + pub base_url: String,
119 + pub api_url: String,
120 + pub model: Option<String>,
121 + pub model_name: Option<String>,
122 + pub tokenizer_id: Option<String>,
123 + #[allow(dead_code)]
124 + pub tokenizer_mode: String,
125 + pub trust_remote_code: bool,
126 + pub skip_tokenizer_init: bool,
127 + pub dataset_name: DatasetName,
128 + pub dataset_path: Option<String>,
129 + pub max_model_len: Option<usize>,
130 + pub random_input_len: usize,
131 + pub random_output_len: usize,
132 + pub random_prefix_len: usize,
133 + pub random_range_ratio: RangeRatio,
134 + pub random_cache_hit_fraction: f64,
135 + pub random_cache_ratio: f64,
136 + /// Inputs per request for embeddings/pooling backends (1 = no batching).
137 + pub random_batch_size: usize,
138 + /// random-rerank: whether the served model is a reranker (default true).
139 + pub is_reranker: bool,
140 + pub custom_output_len: i64,
141 + pub prefix_repetition_prefix_len: usize,
142 + pub prefix_repetition_suffix_len: usize,
143 + pub prefix_repetition_num_prefixes: usize,
144 + pub prefix_repetition_output_len: usize,
145 + pub sharegpt_output_len: Option<usize>,
146 + pub sonnet_input_len: usize,
147 + pub sonnet_output_len: usize,
148 + pub sonnet_prefix_len: usize,
149 + pub no_oversample: bool,
150 + pub disable_shuffle: bool,
151 + pub num_prompts: usize,
152 + pub request_rate: f64,
153 + pub burstiness: f64,
154 + pub max_concurrency: Option<usize>,
155 + pub steady_state_threshold: f64,
156 + pub steady_state_min_window: Option<f64>,
157 + pub no_steady_state: bool,
158 + pub disable_tqdm: bool,
159 + pub num_warmups: usize,
160 + pub profile: bool,
161 + pub profile_batch_threshold: Option<usize>,
162 + pub profile_duration: f64,
163 + pub save_result: bool,
164 + pub save_detailed: bool,
165 + pub append_result: bool,
166 + pub result_dir: Option<String>,
167 + pub result_filename: Option<String>,
168 + pub seed: u64,
169 + pub ignore_eos: bool,
170 + pub insecure: bool,
171 + pub selected_percentile_metrics: Vec<String>,
172 + pub selected_percentiles: Vec<f64>,
173 + pub sweep_summary_percentiles: Vec<f64>,
174 + pub label: Option<String>,
175 + pub logprobs: Option<usize>,
176 + pub request_id_prefix: String,
177 + pub ready_check_timeout_sec: u64,
178 + pub extra_headers: Option<HashMap<String, String>>,
179 + pub extra_body: Option<serde_json::Value>,
180 + pub metadata: Option<Vec<(String, String)>>,
181 + pub dry_run: bool,
182 + pub goodput: GoodputConfig,
183 + pub ramp_up: Option<RampUpConfig>,
184 + pub multi_turn: bool,
185 + pub multi_turn_num_turns: usize,
186 + pub multi_turn_min_turns: usize,
187 + pub multi_turn_max_turns: usize,
188 + pub sharegpt_multi_turn_max_turns: Option<usize>,
189 + pub per_turn_input_len: usize,
190 + pub multi_turn_concurrency: Option<usize>,
191 + pub multi_turn_delay_ms: u64,
192 + pub multi_turn_prefix_global_ratio: f64,
193 + pub multi_turn_prefix_conversation_ratio: f64,
194 + pub speed_bench_config: SpeedBenchConfig,
195 + pub speed_bench_category: Option<String>,
196 + pub speed_bench_max_input_len: Option<usize>,
197 + pub hf_split: Option<String>,
198 + pub hf_subset: Option<String>,
199 + pub hf_output_len: Option<usize>,
200 + pub hf_text_column: Option<String>,
201 + pub reset_prefix_cache: bool,
202 + pub prompt_token_ids: bool,
203 + // --- Random multimodal dataset ---
204 + pub random_mm_base_items_per_request: usize,
205 + pub random_mm_num_mm_items_range_ratio: f64,
206 + pub random_mm_limit: MmLimitPerPrompt,
207 + pub random_mm_buckets: Vec<(MmBucketKey, f64)>,
208 + /// Datasets that support it pre-build the chat `messages` array
209 + /// (text + multimodal parts) instead of prompt + separate mm content.
210 + pub enable_multimodal_chat: bool,
211 + /// LoRA adapter names. None = no LoRA routing (use --model directly).
212 + /// Stored as Arc<str> so per-request override is a cheap clone.
213 + pub lora_modules: Option<Vec<Arc<str>>>,
214 + pub lora_assignment: LoraAssignment,
215 +}
216 +
217 +impl BenchConfig {
218 + pub fn from_cli(cli: &Cli) -> Result<Self> {
219 + if cli.burstiness <= 0.0 {
220 + return Err(BenchError::Config("Burstiness must be positive".into()));
221 + }
222 +
223 + if cli.num_prompts == 0 {
224 + return Err(BenchError::Config(
225 + "--num-prompts must be at least 1".into(),
226 + ));
227 + }
228 +
229 + if cli.request_rate <= 0.0 && !cli.request_rate.is_infinite() {
230 + return Err(BenchError::Config(
231 + "--request-rate must be positive (or inf)".into(),
232 + ));
233 + }
234 + if cli.max_model_len == Some(0) {
235 + return Err(BenchError::Config(
236 + "--max-model-len must be at least 1".into(),
237 + ));
238 + }
239 +
240 + let base_url = cli.resolve_base_url();
241 + let api_url = cli.resolve_api_url();
242 +
243 + let extra_headers = cli.parse_headers()?;
244 + let mut extra_body = cli.parse_extra_body()?;
245 +
246 + // Merge sampling parameters into extra_body (matches Python behavior).
247 + // Python collects non-None sampling params and merges them UNDER extra_body,
248 + // meaning extra_body keys take precedence over sampling params.
249 + {
250 + let mut sampling_params = serde_json::Map::new();
251 + if let Some(v) = cli.top_p {
252 + sampling_params.insert("top_p".into(), serde_json::json!(v));
253 + }
254 + if let Some(v) = cli.top_k {
255 + sampling_params.insert("top_k".into(), serde_json::json!(v));
256 + }
257 + if let Some(v) = cli.min_p {
258 + sampling_params.insert("min_p".into(), serde_json::json!(v));
259 + }
260 + if let Some(v) = cli.temperature {
261 + sampling_params.insert("temperature".into(), serde_json::json!(v));
262 + }
263 + if let Some(v) = cli.frequency_penalty {
264 + sampling_params.insert("frequency_penalty".into(), serde_json::json!(v));
265 + }
266 + if let Some(v) = cli.presence_penalty {
267 + sampling_params.insert("presence_penalty".into(), serde_json::json!(v));
268 + }
269 + if let Some(v) = cli.repetition_penalty {
270 + sampling_params.insert("repetition_penalty".into(), serde_json::json!(v));
271 + }
272 +
273 + if !sampling_params.is_empty() {
274 + if !cli.backend.is_openai_compatible() {
275 + return Err(BenchError::Config(
276 + "Sampling parameters are only supported by openai-compatible backends."
277 + .into(),
278 + ));
279 + }
280 +
281 + // Merge: sampling_params first, then extra_body on top (extra_body wins)
282 + let merged = match extra_body.take() {
283 + Some(serde_json::Value::Object(existing)) => {
284 + sampling_params.extend(existing);
285 + sampling_params
286 + }
287 + Some(other) => {
288 + // extra_body was not an object — just use sampling params
289 + eprintln!(
290 + "Warning: --extra-body is not a JSON object, sampling params may be lost"
291 + );
292 + let _ = other;
293 + sampling_params
294 + }
295 + None => sampling_params,
296 + };
297 + extra_body = Some(serde_json::Value::Object(merged));
298 + }
299 + }
300 +
301 + // Parse metadata
302 + let metadata = match &cli.metadata {
303 + None => None,
304 + Some(items) => {
305 + let mut pairs = Vec::new();
306 + for item in items {
307 + let (k, v) = item.split_once('=').ok_or_else(|| {
308 + BenchError::Config("Invalid metadata format. Use KEY=VALUE".into())
309 + })?;
310 + pairs.push((k.trim().to_string(), v.trim().to_string()));
311 + }
312 + Some(pairs)
313 + }
314 + };
315 +
316 + // Parse goodput SLOs
317 + let goodput = parse_goodput(&cli.goodput)?;
318 +
319 + // Parse ramp-up config
320 + let ramp_up = parse_ramp_up(cli)?;
321 +
322 + // Default percentile metrics based on backend type
323 + let default_percentile_metrics = if cli.backend.is_pooling() {
324 + "e2el"
325 + } else {
326 + "ttft,tpot,itl,e2el"
327 + };
328 + let percentile_metrics_str =
329 + cli.percentile_metrics.as_deref().unwrap_or(default_percentile_metrics);
330 + let selected_percentile_metrics: Vec<String> =
331 + percentile_metrics_str.split(',').map(|s| s.trim().to_string()).collect();
332 +
333 + let metric_percentiles = parse_percentiles(&cli.metric_percentiles, false)?;
334 + let sweep_summary_percentiles = cli
335 + .sweep_summary_percentiles
336 + .as_deref()
337 + .map(|raw| parse_percentiles(raw, true))
338 + .transpose()?
339 + .unwrap_or_default();
340 + let mut selected_percentiles =
341 + merge_percentiles(&metric_percentiles, &sweep_summary_percentiles);
342 + // Always include p90 so it appears in console output and sweep summary
343 + if !selected_percentiles.contains(&90.0) {
344 + selected_percentiles.push(90.0);
345 + }
346 +
347 + let tokenizer_id = if cli.skip_tokenizer_init {
348 + None
349 + } else {
350 + Some(cli.tokenizer.clone().or_else(|| cli.model.clone()).unwrap_or_default())
351 + };
352 +
353 + // Resolve input/output lengths
354 + let random_input_len = cli.resolved_random_input_len();
355 + let random_output_len = cli.resolved_random_output_len();
356 + let per_turn_input_len = cli.resolved_per_turn_input_len();
357 +
358 + // Normalized multi-turn turn counts (computed in validation block below, defaults
359 + // to num_turns if multi-turn mode is not active)
360 + let mut multi_turn_min_turns = cli.multi_turn_num_turns;
361 + let mut multi_turn_max_turns = cli.multi_turn_num_turns;
362 +
363 + // For random datasets with openai-compatible backends, default to ignore_eos.
364 + // Exception: multi-turn mode, where ignore_eos causes unbounded context growth
365 + // across turns. Multi-turn uses min_tokens instead for output length control.
366 + // Pooling backends don't generate tokens, so ignore_eos is irrelevant.
367 + let ignore_eos = if cli.backend.is_pooling() {
368 + false
369 + } else {
370 + cli.ignore_eos
371 + || ((cli.dataset_name == DatasetName::Random
372 + || cli.dataset_name == DatasetName::RandomMm)
373 + && cli.backend.is_openai_compatible()
374 + && !cli.multi_turn)
375 + };
376 +
377 + // Pooling backends don't support multi-turn
378 + if cli.backend.is_pooling() && cli.multi_turn {
379 + return Err(BenchError::Config(
380 + "Pooling/embedding backends do not support --multi-turn".into(),
381 + ));
382 + }
383 +
384 + // LoRA validation. Adapter names must be non-empty after trim; pooling
385 + // backends are out of scope (vLLM LoRA routing is for generative paths).
386 + let lora_modules = match cli.lora_modules.as_ref() {
387 + None => None,
388 + Some(names) => {
389 + if names.is_empty() {
390 + return Err(BenchError::Config(
391 + "--lora-modules requires at least one adapter name".into(),
392 + ));
393 + }
394 + if cli.backend.is_pooling() {
395 + return Err(BenchError::Config(
396 + "--lora-modules is not supported for pooling/embedding backends".into(),
397 + ));
398 + }
399 + let mut out = Vec::with_capacity(names.len());
400 + for n in names {
401 + let trimmed = n.trim();
402 + if trimmed.is_empty() {
403 + return Err(BenchError::Config(
404 + "--lora-modules contains an empty adapter name".into(),
405 + ));
406 + }
407 + out.push(Arc::<str>::from(trimmed));
408 + }
409 + Some(out)
410 + }
411 + };
412 +
413 + // Random-MM validation and config parsing
414 + let (random_mm_limit, random_mm_buckets) = if cli.dataset_name == DatasetName::RandomMm {
415 + if cli.backend != BackendKind::OpenaiChat {
416 + return Err(BenchError::Config(
417 + "Multi-modal content (images) is only supported on 'openai-chat' backend."
418 + .into(),
419 + ));
420 + }
421 + let limit = crate::datasets::random_mm::parse_limit_mm_per_prompt(
422 + &cli.random_mm_limit_mm_per_prompt,
423 + )?;
424 + let buckets =
425 + crate::datasets::random_mm::parse_bucket_config(&cli.random_mm_bucket_config)?;
426 + (limit, buckets)
427 + } else {
428 + (MmLimitPerPrompt::default(), Vec::new())
429 + };
430 +
431 + // Note: --dataset-path is optional for sharegpt (auto-downloads) and
432 + // sonnet (uses built-in Shakespeare's sonnets).
433 +
434 + // Range ratio (Python semantics: [len*(1-r), len*(1+r)], each r in [0,1))
435 + let random_range_ratio = RangeRatio::parse(&cli.random_range_ratio)?;
436 +
437 + // Batched inputs only make sense for pooling backends (the generation
438 + // backends send one prompt per request).
439 + if cli.random_batch_size == 0 {
440 + return Err(BenchError::Config(
441 + "--random-batch-size must be at least 1".into(),
442 + ));
443 + }
444 + if cli.random_batch_size > 1
445 + && !cli.backend.is_pooling()
446 + && cli.dataset_name != DatasetName::RandomRerank
447 + {
448 + return Err(BenchError::Config(
449 + "--random-batch-size > 1 is only supported with embeddings/pooling backends".into(),
450 + ));
451 + }
452 +
453 + // random-rerank validation (mirrors Python RandomDatasetForReranking)
454 + let is_reranker = !cli.no_reranker;
455 + if cli.dataset_name == DatasetName::RandomRerank {
456 + if !cli.backend.is_pooling() {
457 + return Err(BenchError::Config(
458 + "--dataset-name random-rerank requires an embeddings/pooling backend \
459 + (e.g. --backend vllm-rerank)"
460 + .into(),
461 + ));
462 + }
463 + if !is_reranker && (cli.num_prompts < 2 || cli.random_batch_size < 2) {
464 + return Err(BenchError::Config(
465 + "--no-reranker requires --num-prompts > 1 and --random-batch-size > 1 \
466 + (the query is folded into the first batch slot)"
467 + .into(),
468 + ));
469 + }
470 + }
471 +
472 + // Custom dataset validation
473 + if cli.dataset_name == DatasetName::Custom {
474 + match cli.dataset_path.as_deref() {
475 + None => {
476 + return Err(BenchError::Config(
477 + "--dataset-path is required for --dataset-name custom \
478 + (a JSONL file with {\"prompt\": ..., \"output_tokens\": ...} lines)"
479 + .into(),
480 + ));
481 + }
482 + Some(p) if !p.ends_with(".jsonl") => {
483 + return Err(BenchError::Config(
484 + "Only JSONL format is supported for the custom dataset".into(),
485 + ));
486 + }
487 + _ => {}
488 + }
489 + if !cli.skip_chat_template {
490 + eprintln!(
491 + "NOTE: client-side chat template rendering is not supported; custom \
492 + dataset prompts are sent raw (equivalent to --skip-chat-template)."
493 + );
494 + }
495 + }
496 +
497 + // Prefix repetition validation
498 + if cli.dataset_name == DatasetName::PrefixRepetition {
499 + if cli.prefix_repetition_num_prefixes == 0 {
500 + return Err(BenchError::Config(
501 + "--prefix-repetition-num-prefixes must be at least 1".into(),
502 + ));
503 + }
504 + if cli.num_prompts < cli.prefix_repetition_num_prefixes {
505 + return Err(BenchError::Config(format!(
506 + "--num-prompts ({}) must be >= --prefix-repetition-num-prefixes ({})",
507 + cli.num_prompts, cli.prefix_repetition_num_prefixes
508 + )));
509 + }
510 + }
511 +
512 + // HF dataset validation
513 + if cli.dataset_name == DatasetName::Hf && cli.dataset_path.is_none() {
514 + return Err(BenchError::Config(
515 + "--dataset-path is required for --dataset-name hf \
516 + (set to a HuggingFace dataset ID, e.g. 'allenai/WildChat-4.8M')"
517 + .into(),
518 + ));
519 + }
520 + if let Some(len) = cli.hf_output_len
521 + && len == 0
522 + {
523 + return Err(BenchError::Config(
524 + "--hf-output-len must be at least 1".into(),
525 + ));
526 + }
527 +
528 + // Multi-turn validation
529 + if cli.multi_turn {
530 + if cli.backend != BackendKind::OpenaiChat {
531 + return Err(BenchError::Config(
532 + "--multi-turn requires --backend openai-chat".into(),
533 + ));
534 + }
535 + if cli.multi_turn_num_turns == 0 {
536 + return Err(BenchError::Config(
537 + "--multi-turn-num-turns must be at least 1".into(),
538 + ));
539 + }
540 +
541 + // Normalize and validate min/max turns. ShareGPT only consumes max_turns
542 + // (the loader walks all available turns up to the cap), so the
543 + // min/num/max coupling used for synthetic generation does not apply.
544 + if cli.dataset_name == DatasetName::ShareGpt {
545 + if cli.multi_turn_max_turns == 1 {
546 + return Err(BenchError::Config(
547 + "--multi-turn-max-turns must be at least 2 for ShareGPT multi-turn".into(),
548 + ));
549 + }
550 + } else {
551 + (multi_turn_min_turns, multi_turn_max_turns) =
552 + match (cli.multi_turn_min_turns, cli.multi_turn_max_turns) {
553 + (0, 0) => (cli.multi_turn_num_turns, cli.multi_turn_num_turns),
554 + (m, 0) => (m, cli.multi_turn_num_turns),
555 + (0, x) => (cli.multi_turn_num_turns, x),
556 + (m, x) => (m, x),
557 + };
558 + if multi_turn_min_turns < 1 {
559 + return Err(BenchError::Config(
560 + "--multi-turn-min-turns must be at least 1".into(),
561 + ));
562 + }
563 + if multi_turn_min_turns > multi_turn_max_turns {
564 + return Err(BenchError::Config(
565 + "--multi-turn-min-turns must be <= --multi-turn-max-turns".into(),
566 + ));
567 + }
568 + }
569 +
570 + if ignore_eos {
571 + eprintln!(
572 + "WARNING: --ignore-eos is set with --multi-turn. The server may not \
573 + respect output length limits, causing unbounded context growth."
574 + );
575 + }
576 +
577 + // Validate prefix sharing ratios
578 + let pg = cli.multi_turn_prefix_global_ratio;
579 + let pc = cli.multi_turn_prefix_conversation_ratio;
580 + if !(0.0..=1.0).contains(&pg) {
581 + return Err(BenchError::Config(
582 + "--multi-turn-prefix-global-ratio must be in [0.0, 1.0]".into(),
583 + ));
584 + }
585 + if !(0.0..=1.0).contains(&pc) {
586 + return Err(BenchError::Config(
587 + "--multi-turn-prefix-conversation-ratio must be in [0.0, 1.0]".into(),
588 + ));
589 + }
590 + if pg + pc >= 1.0 {
591 + return Err(BenchError::Config(
592 + "--multi-turn-prefix-global-ratio + --multi-turn-prefix-conversation-ratio must be < 1.0 (unique suffix required)".into(),
593 + ));
594 + }
595 + if (pg > 0.0 || pc > 0.0) && cli.dataset_name != DatasetName::Random {
596 + return Err(BenchError::Config(
597 + "Prefix sharing (--multi-turn-prefix-global-ratio / --multi-turn-prefix-conversation-ratio) only works with --dataset-name random".into(),
598 + ));
599 + }
600 + }
601 +
602 + if !(cli.steady_state_threshold > 0.0 && cli.steady_state_threshold <= 1.0) {
603 + return Err(BenchError::Config(format!(
604 + "--steady-state-threshold must be in (0.0, 1.0], got {}",
605 + cli.steady_state_threshold
606 + )));
607 + }
608 + if let Some(mw) = cli.steady_state_min_window
609 + && mw < 0.0
610 + {
611 + return Err(BenchError::Config(format!(
612 + "--steady-state-min-window must be >= 0, got {mw}"
613 + )));
614 + }
615 +
616 + if cli.profile_batch_threshold.is_some() && !cli.profile {
617 + return Err(BenchError::Config(
618 + "--profile-batch-threshold requires --profile".into(),
619 + ));
620 + }
621 + if cli.profile_duration <= 0.0 {
622 + return Err(BenchError::Config(
623 + "--profile-duration must be positive".into(),
624 + ));
625 + }
626 + if cli.profile_batch_threshold.is_none() && cli.profile_duration != 5.0 {
627 + return Err(BenchError::Config(
628 + "--profile-duration requires --profile-batch-threshold".into(),
629 + ));
630 + }
631 +
632 + Ok(BenchConfig {
633 + backend: cli.backend,
634 + base_url,
635 + api_url,
636 + model: cli.model.clone(),
637 + model_name: cli.served_model_name.clone(),
638 + tokenizer_id,
639 + tokenizer_mode: cli.tokenizer_mode.clone(),
640 + trust_remote_code: cli.trust_remote_code,
641 + skip_tokenizer_init: cli.skip_tokenizer_init,
642 + dataset_name: cli.dataset_name,
643 + dataset_path: cli.dataset_path.clone(),
644 + max_model_len: cli.max_model_len,
645 + random_input_len,
646 + random_output_len,
647 + random_prefix_len: cli.random_prefix_len,
648 + random_range_ratio,
649 + random_batch_size: cli.random_batch_size,
650 + is_reranker,
651 + custom_output_len: cli.output_len.map(|v| v as i64).unwrap_or(cli.custom_output_len),
652 + prefix_repetition_prefix_len: cli.prefix_repetition_prefix_len,
653 + prefix_repetition_suffix_len: cli.prefix_repetition_suffix_len,
654 + prefix_repetition_num_prefixes: cli.prefix_repetition_num_prefixes,
655 + prefix_repetition_output_len: cli
656 + .output_len
657 + .unwrap_or(cli.prefix_repetition_output_len),
658 + random_cache_hit_fraction: cli.random_cache_hit_fraction,
659 + random_cache_ratio: cli.random_cache_ratio,
660 + sharegpt_output_len: cli.sharegpt_output_len,
661 + sonnet_input_len: cli.sonnet_input_len,
662 + sonnet_output_len: cli.sonnet_output_len,
663 + sonnet_prefix_len: cli.sonnet_prefix_len,
664 + no_oversample: cli.no_oversample,
665 + disable_shuffle: cli.disable_shuffle,
666 + num_prompts: cli.num_prompts,
667 + request_rate: cli.request_rate,
668 + burstiness: cli.burstiness,
669 + max_concurrency: cli.max_concurrency,
670 + steady_state_threshold: cli.steady_state_threshold,
671 + steady_state_min_window: cli.steady_state_min_window,
672 + no_steady_state: cli.no_steady_state,
673 + disable_tqdm: cli.disable_tqdm,
674 + num_warmups: cli.num_warmups,
675 + profile: cli.profile,
676 + profile_batch_threshold: cli.profile_batch_threshold,
677 + profile_duration: cli.profile_duration,
678 + save_result: cli.save_result,
679 + save_detailed: cli.save_detailed,
680 + append_result: cli.append_result,
681 + result_dir: cli.result_dir.clone(),
682 + result_filename: cli.result_filename.clone(),
683 + seed: cli.seed,
684 + ignore_eos,
685 + insecure: cli.insecure,
686 + selected_percentile_metrics,
687 + selected_percentiles,
688 + sweep_summary_percentiles,
689 + label: cli.label.clone(),
690 + logprobs: cli.logprobs,
691 + request_id_prefix: cli.get_request_id_prefix(),
692 + ready_check_timeout_sec: cli.ready_check_timeout_sec,
693 + extra_headers,
694 + extra_body,
695 + metadata,
696 + dry_run: cli.dry_run,
697 + goodput,
698 + ramp_up,
699 + multi_turn: cli.multi_turn,
700 + multi_turn_num_turns: cli.multi_turn_num_turns,
701 + multi_turn_min_turns,
702 + multi_turn_max_turns,
703 + sharegpt_multi_turn_max_turns: if cli.multi_turn
704 + && cli.dataset_name == DatasetName::ShareGpt
705 + && cli.multi_turn_max_turns != 0
706 + {
707 + Some(cli.multi_turn_max_turns)
708 + } else {
709 + None
710 + },
711 + per_turn_input_len,
712 + multi_turn_concurrency: cli.multi_turn_concurrency,
713 + multi_turn_delay_ms: cli.multi_turn_delay_ms,
714 + multi_turn_prefix_global_ratio: cli.multi_turn_prefix_global_ratio,
715 + multi_turn_prefix_conversation_ratio: cli.multi_turn_prefix_conversation_ratio,
716 + speed_bench_config: cli.speed_bench_config,
717 + speed_bench_category: cli.speed_bench_category.clone(),
718 + speed_bench_max_input_len: cli.speed_bench_max_input_len,
719 + hf_split: cli.hf_split.clone(),
720 + hf_subset: cli.hf_subset.clone(),
721 + hf_output_len: cli.hf_output_len,
722 + hf_text_column: cli.hf_text_column.clone(),
723 + reset_prefix_cache: cli.reset_prefix_cache,
724 + prompt_token_ids: cli.prompt_token_ids,
725 + random_mm_base_items_per_request: cli.random_mm_base_items_per_request,
726 + random_mm_num_mm_items_range_ratio: cli.random_mm_num_mm_items_range_ratio,
727 + random_mm_limit,
728 + random_mm_buckets,
729 + enable_multimodal_chat: cli.enable_multimodal_chat,
730 + lora_modules,
731 + lora_assignment: cli.lora_assignment,
732 + })
733 + }
734 +}
735 +
736 +fn parse_percentiles(raw: &str, dedupe: bool) -> Result<Vec<f64>> {
737 + let mut percentiles = Vec::new();
738 + for s in raw.split(',') {
739 + let p = s
740 + .trim()
741 + .parse::<f64>()
742 + .map_err(|_| BenchError::Config(format!("Invalid percentile: {s}")))?;
743 + if !(0.0..=100.0).contains(&p) {
744 + return Err(BenchError::Config(format!(
745 + "Percentile must be in [0, 100], got: {p}"
746 + )));
747 + }
748 + if !dedupe || !percentiles.contains(&p) {
749 + percentiles.push(p);
750 + }
751 + }
752 + Ok(percentiles)
753 +}
754 +
755 +fn merge_percentiles(metric_percentiles: &[f64], summary_percentiles: &[f64]) -> Vec<f64> {
756 + let mut merged = Vec::with_capacity(metric_percentiles.len() + summary_percentiles.len());
757 + for &p in metric_percentiles {
758 + if !merged.contains(&p) {
759 + merged.push(p);
760 + }
761 + }
762 + for &p in summary_percentiles {
763 + if !merged.contains(&p) {
764 + merged.push(p);
765 + }
766 + }
767 + merged
768 +}
769 +
770 +fn parse_goodput(goodput_args: &Option<Vec<String>>) -> Result<GoodputConfig> {
771 + let items = match goodput_args {
772 + None => return Ok(GoodputConfig::default()),
773 + Some(items) => items,
774 + };
775 +
776 + let valid_names = ["ttft", "tpot", "e2el"];
777 + let mut config = GoodputConfig::default();
778 +
779 + for item in items {
780 + let (name, val_str) = item.split_once(':').ok_or_else(|| {
781 + BenchError::Config(format!(
782 + "Invalid goodput format: '{item}'. Use KEY:VALUE (e.g. ttft:100)"
783 + ))
784 + })?;
785 +
786 + let val: f64 = val_str
787 + .trim()
788 + .parse()
789 + .map_err(|_| BenchError::Config(format!("Invalid goodput value: '{val_str}'")))?;
790 +
791 + if val < 0.0 {
792 + return Err(BenchError::Config(format!(
793 + "Goodput SLO value must be non-negative, got: {name}:{val}"
794 + )));
795 + }
796 +
797 + if !valid_names.contains(&name) {
798 + return Err(BenchError::Config(format!(
799 + "Invalid goodput metric: '{name}'. Valid: {valid_names:?}"
800 + )));
801 + }
802 +
803 + match name {
804 + "ttft" => config.ttft_ms = Some(val),
805 + "tpot" => config.tpot_ms = Some(val),
806 + "e2el" => config.e2el_ms = Some(val),
807 + _ => unreachable!(),
808 + }
809 + }
810 +
811 + Ok(config)
812 +}
813 +
814 +fn parse_ramp_up(cli: &Cli) -> Result<Option<RampUpConfig>> {
815 + let strategy = match cli.ramp_up_strategy {
816 + None => return Ok(None),
817 + Some(s) => s,
818 + };
819 +
820 + let start_rps = cli.ramp_up_start_rps.ok_or_else(|| {
821 + BenchError::Config("--ramp-up-start-rps is required when --ramp-up-strategy is set".into())
822 + })?;
823 +
824 + let end_rps = cli.ramp_up_end_rps.ok_or_else(|| {
825 + BenchError::Config("--ramp-up-end-rps is required when --ramp-up-strategy is set".into())
826 + })?;
827 +
828 + if start_rps <= 0.0 || end_rps <= 0.0 {
829 + return Err(BenchError::Config(
830 + "Ramp-up RPS values must be positive".into(),
831 + ));
832 + }
833 +
834 + Ok(Some(RampUpConfig {
835 + strategy,
836 + start_rps,
837 + end_rps,
838 + }))
839 +}
840 +
841 +#[cfg(test)]
842 +mod tests {
843 + use clap::Parser;
844 +
845 + use super::*;
846 + use crate::cli::Cli;
847 +
848 + fn base_multi_turn_args() -> Vec<&'static str> {
849 + vec![
850 + "vllm-bench",
851 + "--backend",
852 + "openai-chat",
853 + "--multi-turn",
854 + "--model",
855 + "test-model",
856 + ]
857 + }
858 +
859 + #[test]
860 + fn test_prefix_sharing_defaults_to_zero() {
861 + let args = base_multi_turn_args();
862 + let cli = Cli::parse_from(args);
863 + let config = BenchConfig::from_cli(&cli).unwrap();
864 + assert_eq!(config.multi_turn_prefix_global_ratio, 0.0);
865 + assert_eq!(config.multi_turn_prefix_conversation_ratio, 0.0);
866 + }
867 +
868 + #[test]
869 + fn test_prefix_sharing_valid() {
870 + let mut args = base_multi_turn_args();
871 + args.extend([
872 + "--multi-turn-prefix-global-ratio",
873 + "0.1",
874 + "--multi-turn-prefix-conversation-ratio",
875 + "0.8",
876 + ]);
877 + let cli = Cli::parse_from(args);
878 + let config = BenchConfig::from_cli(&cli).unwrap();
879 + assert!((config.multi_turn_prefix_global_ratio - 0.1).abs() < 1e-10);
880 + assert!((config.multi_turn_prefix_conversation_ratio - 0.8).abs() < 1e-10);
881 + }
882 +
883 + #[test]
884 + fn test_prefix_sharing_sum_exceeds_one_fails() {
885 + let mut args = base_multi_turn_args();
886 + args.extend([
887 + "--multi-turn-prefix-global-ratio",
888 + "0.6",
889 + "--multi-turn-prefix-conversation-ratio",
890 + "0.6",
891 + ]);
892 + let cli = Cli::parse_from(args);
893 + assert!(BenchConfig::from_cli(&cli).is_err());
894 + }
895 +
896 + #[test]
897 + fn test_prefix_sharing_sum_equals_one_fails() {
898 + let mut args = base_multi_turn_args();
899 + args.extend([
900 + "--multi-turn-prefix-global-ratio",
901 + "0.5",
902 + "--multi-turn-prefix-conversation-ratio",
903 + "0.5",
904 + ]);
905 + let cli = Cli::parse_from(args);
906 + assert!(BenchConfig::from_cli(&cli).is_err());
907 + }
908 +
909 + #[test]
910 + fn test_prefix_sharing_out_of_range_fails() {
911 + let mut args = base_multi_turn_args();
912 + args.extend(["--multi-turn-prefix-global-ratio", "1.5"]);
913 + let cli = Cli::parse_from(args);
914 + assert!(BenchConfig::from_cli(&cli).is_err());
915 + }
916 +
917 + #[test]
918 + fn test_prefix_sharing_requires_random_dataset() {
919 + let args = vec![
920 + "vllm-bench",
921 + "--backend",
922 + "openai-chat",
923 + "--multi-turn",
924 + "--model",
925 + "test-model",
926 + "--dataset-name",
927 + "sharegpt",
928 + "--multi-turn-prefix-global-ratio",
929 + "0.1",
930 + ];
931 + let cli = Cli::parse_from(args);
932 + assert!(BenchConfig::from_cli(&cli).is_err());
933 + }
934 +
935 + #[test]
936 + fn test_sharegpt_multi_turn_max_turns_default_uncapped() {
937 + let args = vec![
938 + "vllm-bench",
939 + "--backend",
940 + "openai-chat",
941 + "--multi-turn",
942 + "--model",
943 + "test-model",
944 + "--dataset-name",
945 + "sharegpt",
946 + ];
947 + let cli = Cli::parse_from(args);
948 + let config = BenchConfig::from_cli(&cli).unwrap();
949 +
950 + assert_eq!(config.multi_turn_max_turns, 3);
951 + assert_eq!(config.sharegpt_multi_turn_max_turns, None);
952 + }
953 +
954 + #[test]
955 + fn test_sharegpt_multi_turn_max_turns_2_succeeds() {
956 + // Regression: previously the (0, x) normalization arm produced
957 + // (min=multi_turn_num_turns=3, max=2), tripping the min>max check
958 + // with a misleading error about a flag the user never set.
959 + let args = vec![
960 + "vllm-bench",
961 + "--backend",
962 + "openai-chat",
963 + "--multi-turn",
964 + "--model",
965 + "test-model",
966 + "--dataset-name",
967 + "sharegpt",
968 + "--multi-turn-max-turns",
969 + "2",
970 + ];
971 + let cli = Cli::parse_from(args);
972 + let config = BenchConfig::from_cli(&cli).unwrap();
973 + assert_eq!(config.sharegpt_multi_turn_max_turns, Some(2));
974 + }
975 +
976 + #[test]
977 + fn test_sharegpt_multi_turn_max_turns_1_rejected() {
978 + let args = vec![
979 + "vllm-bench",
980 + "--backend",
981 + "openai-chat",
982 + "--multi-turn",
983 + "--model",
984 + "test-model",
985 + "--dataset-name",
986 + "sharegpt",
987 + "--multi-turn-max-turns",
988 + "1",
989 + ];
990 + let cli = Cli::parse_from(args);
991 + let err = BenchConfig::from_cli(&cli).unwrap_err().to_string();
992 + assert!(
993 + err.contains("at least 2 for ShareGPT"),
994 + "expected ShareGPT-specific error, got: {err}"
995 + );
996 + }
997 +
998 + #[test]
999 + fn test_sharegpt_multi_turn_max_turns_explicit_cap() {
1000 + let args = vec![
1001 + "vllm-bench",
1002 + "--backend",
1003 + "openai-chat",
1004 + "--multi-turn",
1005 + "--model",
1006 + "test-model",
1007 + "--dataset-name",
1008 + "sharegpt",
1009 + "--multi-turn-max-turns",
1010 + "20",
1011 + ];
1012 + let cli = Cli::parse_from(args);
1013 + let config = BenchConfig::from_cli(&cli).unwrap();
1014 +
1015 + assert_eq!(config.sharegpt_multi_turn_max_turns, Some(20));
1016 + }
1017 +
1018 + #[test]
1019 + fn test_sweep_summary_percentiles_default_empty() {
1020 + let args = base_multi_turn_args();
1021 + let cli = Cli::parse_from(args);
1022 + let config = BenchConfig::from_cli(&cli).unwrap();
1023 +
1024 + assert!(config.sweep_summary_percentiles.is_empty());
1025 + assert_eq!(config.selected_percentiles, vec![99.0, 90.0]);
1026 + }
1027 +
1028 + #[test]
1029 + fn test_sweep_summary_percentiles_are_deduped_and_merged() {
1030 + let mut args = base_multi_turn_args();
1031 + args.extend([
1032 + "--metric-percentiles",
1033 + "99,95",
1034 + "--sweep-summary-percentiles",
1035 + "90,95,90",
1036 + ]);
1037 + let cli = Cli::parse_from(args);
1038 + let config = BenchConfig::from_cli(&cli).unwrap();
1039 +
1040 + assert_eq!(config.sweep_summary_percentiles, vec![90.0, 95.0]);
1041 + assert_eq!(config.selected_percentiles, vec![99.0, 95.0, 90.0]);
1042 + }
1043 +
1044 + #[test]
1045 + fn test_invalid_sweep_summary_percentile_fails() {
1046 + let mut args = base_multi_turn_args();
1047 + args.extend(["--sweep-summary-percentiles", "101"]);
1048 + let cli = Cli::parse_from(args);
1049 + assert!(BenchConfig::from_cli(&cli).is_err());
1050 + }
1051 +
1052 + #[test]
1053 + fn test_max_model_len_is_configured() {
1054 + let args = vec![
1055 + "vllm-bench",
1056 + "--model",
1057 + "test-model",
1058 + "--max-model-len",
1059 + "4096",
1060 + ];
1061 + let cli = Cli::parse_from(args);
1062 + let config = BenchConfig::from_cli(&cli).unwrap();
1063 +
1064 + assert_eq!(config.max_model_len, Some(4096));
1065 + }
1066 +
1067 + #[test]
1068 + fn test_zero_max_model_len_fails() {
1069 + let args = vec![
1070 + "vllm-bench",
1071 + "--model",
1072 + "test-model",
1073 + "--max-model-len",
1074 + "0",
1075 + ];
1076 + let cli = Cli::parse_from(args);
1077 +
1078 + assert!(BenchConfig::from_cli(&cli).is_err());
1079 + }
1080 + #[test]
1081 + fn test_range_ratio_parse_float() {
1082 + let rr = RangeRatio::parse("0.2").unwrap();
1083 + assert_eq!(rr.input, 0.2);
1084 + assert_eq!(rr.output, 0.2);
1085 + assert!(!rr.is_fixed());
1086 + assert!(RangeRatio::parse("0.0").unwrap().is_fixed());
1087 + }
1088 +
1089 + #[test]
1090 + fn test_range_ratio_parse_dict() {
1091 + let rr = RangeRatio::parse(r#"{"input": 0.1, "output": 0.5}"#).unwrap();
1092 + assert_eq!(rr.input, 0.1);
1093 + assert_eq!(rr.output, 0.5);
1094 + }
1095 +
1096 + #[test]
1097 + fn test_range_ratio_rejects_old_default_with_hint() {
1098 + let err = RangeRatio::parse("1.0").unwrap_err().to_string();
1099 + assert!(err.contains("semantics now match Python"), "got: {err}");
1100 + assert!(RangeRatio::parse("-0.1").is_err());
1101 + assert!(RangeRatio::parse("{\"input\": 0.1}").is_err());
1102 + assert!(RangeRatio::parse("abc").is_err());
1103 + }
1104 +
1105 + #[test]
1106 + fn test_range_ratio_bounds_python_semantics() {
1107 + // Python: [floor(len*(1-r)), ceil(len*(1+r))], output low clamped to 1
1108 + let rr = RangeRatio::parse("0.2").unwrap();
1109 + assert_eq!(rr.input_bounds(100), (80, 120));
1110 + assert_eq!(rr.output_bounds(100), (80, 120));
1111 + let rr = RangeRatio::parse("0.99").unwrap();
1112 + assert_eq!(rr.output_bounds(1).0, 1); // clamped
1113 + let fixed = RangeRatio::parse("0.0").unwrap();
1114 + assert_eq!(fixed.input_bounds(8192), (8192, 8192));
1115 + }
1116 +}
rust/src/bench/src/datasets/custom.rs new
+181
@@ -0,0 +1,181 @@
1 +// SPDX-License-Identifier: Apache-2.0
2 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +//! Custom dataset: JSONL file with one request per line.
5 +//!
6 +//! ```jsonl
7 +//! {"prompt": "What is the capital of India?", "output_tokens": 10}
8 +//! {"prompt": "What is the capital of Iran?", "output_tokens": 1520}
9 +//! ```
10 +//!
11 +//! Mirrors Python's `CustomDataset`. `output_tokens` is optional unless
12 +//! `--custom-output-len -1` is passed. Unlike Python, prompts are always sent
13 +//! raw (no client-side chat template; see `--skip-chat-template`).
14 +
15 +use std::sync::Arc;
16 +
17 +use rand::SeedableRng;
18 +use rand::rngs::StdRng;
19 +use rand::seq::SliceRandom;
20 +use serde::Deserialize;
21 +
22 +use super::SampleRequest;
23 +use crate::error::{BenchError, Result};
24 +use crate::tokenizer::TokenizerKind;
25 +
26 +#[derive(Deserialize)]
27 +struct CustomLine {
28 + prompt: String,
29 + output_tokens: Option<serde_json::Value>,
30 +}
31 +
32 +/// Load the custom JSONL dataset.
33 +///
34 +/// `output_len < 0` means "use the per-line output_tokens field" (Python's
35 +/// `--custom-output-len -1`); otherwise `output_len` applies to every request.
36 +pub fn load_custom_dataset(
37 + tokenizer: &TokenizerKind,
38 + path: &str,
39 + num_requests: usize,
40 + output_len: i64,
41 + seed: u64,
42 + request_id_prefix: &str,
43 + no_oversample: bool,
44 + disable_shuffle: bool,
45 +) -> Result<Vec<SampleRequest>> {
46 + let content = std::fs::read_to_string(path)
47 + .map_err(|e| BenchError::Config(format!("Failed to read custom dataset '{path}': {e}")))?;
48 +
49 + let mut lines: Vec<CustomLine> = Vec::new();
50 + for (lineno, line) in content.lines().enumerate() {
51 + let line = line.trim();
52 + if line.is_empty() {
53 + continue;
54 + }
55 + let parsed: CustomLine = serde_json::from_str(line).map_err(|e| {
56 + BenchError::Config(format!(
57 + "Invalid JSONL at {path}:{}: {e} (each line must be an object \
58 + with a 'prompt' field)",
59 + lineno + 1
60 + ))
61 + })?;
62 + lines.push(parsed);
63 + }
64 + if lines.is_empty() {
65 + return Err(BenchError::Config(format!(
66 + "Custom dataset '{path}' contains no entries"
67 + )));
68 + }
69 +
70 + // Python shuffles the loaded data (seeded) before taking num_requests.
71 + if !disable_shuffle {
72 + let mut rng = StdRng::seed_from_u64(seed);
73 + lines.shuffle(&mut rng);
74 + }
75 +
76 + let mut requests: Vec<SampleRequest> = Vec::with_capacity(num_requests.min(lines.len()));
77 + for (i, item) in lines.iter().enumerate() {
78 + if requests.len() >= num_requests {
79 + break;
80 + }
81 +
82 + let expected_output_len = if output_len < 0 {
83 + let raw = item.output_tokens.as_ref().ok_or_else(|| {
84 + BenchError::Config(
85 + "custom dataset: --custom-output-len -1 requires an \
86 + 'output_tokens' field on every line"
87 + .into(),
88 + )
89 + })?;
90 + raw.as_i64().filter(|v| *v > 0).ok_or_else(|| {
91 + BenchError::Config(format!(
92 + "custom dataset: invalid 'output_tokens' value {raw}: \
93 + must be a positive integer"
94 + ))
95 + })? as usize
96 + } else {
97 + output_len as usize
98 + };
99 +
100 + let prompt_len = tokenizer.encode(&item.prompt, true)?.len();
101 + requests.push(SampleRequest {
102 + prompt: Arc::from(item.prompt.as_str()),
103 + prompt_len,
104 + expected_output_len,
105 + request_id: Some(format!("{request_id_prefix}{i}")),
106 + ..Default::default()
107 + });
108 + }
109 +
110 + super::oversample_requests(
111 + &mut requests,
112 + num_requests,
113 + request_id_prefix,
114 + no_oversample,
115 + );
116 + Ok(requests)
117 +}
118 +
119 +#[cfg(test)]
120 +mod tests {
121 + use super::*;
122 +
123 + fn write_temp_jsonl(name: &str, content: &str) -> String {
124 + let path = std::env::temp_dir().join(format!("vllm-bench-custom-{name}.jsonl"));
125 + std::fs::write(&path, content).unwrap();
126 + path.to_string_lossy().into_owned()
127 + }
128 +
129 + /// gpt2 via built-in tiktoken encoding — loads without network access.
130 + fn test_tokenizer() -> TokenizerKind {
131 + crate::tokenizer::load_tokenizer("gpt2", false, None)
132 + .expect("gpt2 built-in tiktoken should always load without network")
133 + }
134 +
135 + #[test]
136 + fn test_load_custom_dataset_basic() {
137 + let path = write_temp_jsonl(
138 + "basic",
139 + r#"{"prompt": "hello world", "output_tokens": 10}
140 +{"prompt": "foo bar baz", "output_tokens": 20}
141 +"#,
142 + );
143 + let reqs = load_custom_dataset(&test_tokenizer(), &path, 2, 256, 0, "t-", true, true)
144 + .expect("load should succeed");
145 + assert_eq!(reqs.len(), 2);
146 + // Fixed output_len (256) wins over per-line output_tokens by default
147 + assert!(reqs.iter().all(|r| r.expected_output_len == 256));
148 + assert_eq!(&*reqs[0].prompt, "hello world");
149 + assert!(reqs[0].prompt_len > 0);
150 + }
151 +
152 + #[test]
153 + fn test_load_custom_dataset_per_line_output_tokens() {
154 + let path = write_temp_jsonl(
155 + "perline",
156 + r#"{"prompt": "hello", "output_tokens": 10}
157 +{"prompt": "world", "output_tokens": 20}
158 +"#,
159 + );
160 + let reqs = load_custom_dataset(&test_tokenizer(), &path, 2, -1, 0, "t-", true, true)
161 + .expect("load should succeed");
162 + assert_eq!(reqs[0].expected_output_len, 10);
163 + assert_eq!(reqs[1].expected_output_len, 20);
164 + }
165 +
166 + #[test]
167 + fn test_load_custom_dataset_missing_output_tokens_errors() {
168 + let path = write_temp_jsonl("missing", r#"{"prompt": "hello"}"#);
169 + let err = load_custom_dataset(&test_tokenizer(), &path, 1, -1, 0, "t-", true, true)
170 + .expect_err("should fail without output_tokens");
171 + assert!(err.to_string().contains("output_tokens"));
172 + }
173 +
174 + #[test]
175 + fn test_load_custom_dataset_missing_prompt_errors() {
176 + let path = write_temp_jsonl("noprompt", r#"{"text": "hello"}"#);
177 + assert!(
178 + load_custom_dataset(&test_tokenizer(), &path, 1, 256, 0, "t-", true, true).is_err()
179 + );
180 + }
181 +}
rust/src/bench/src/datasets/hf_dataset.rs new
+1239
@@ -0,0 +1,1239 @@
1 +// SPDX-License-Identifier: Apache-2.0
2 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +use std::sync::Arc;
5 +
6 +use rand::rngs::StdRng;
7 +use rand::seq::SliceRandom;
8 +use rand::{Rng, SeedableRng};
9 +
10 +use super::SampleRequest;
11 +use crate::error::{BenchError, Result};
12 +use crate::tokenizer::TokenizerKind;
13 +
14 +/// Cache directory for downloaded HF datasets.
15 +fn cache_dir() -> std::path::PathBuf {
16 + dirs::cache_dir()
17 + .unwrap_or_else(|| std::path::PathBuf::from("/tmp"))
18 + .join("vllm-bench")
19 + .join("datasets")
20 +}
21 +
22 +/// Sanitize a dataset name for use in filenames (replace `/` and other unsafe chars).
23 +fn sanitize_name(name: &str) -> String {
24 + name.chars()
25 + .map(|c| {
26 + if c.is_alphanumeric() || c == '-' || c == '_' || c == '.' {
27 + c
28 + } else {
29 + '_'
30 + }
31 + })
32 + .collect()
33 +}
34 +
35 +/// Detected column format for extracting prompts from HF dataset rows.
36 +enum ColumnFormat {
37 + /// Column contains chat messages (array of {role, content} or {from, value}).
38 + Chat(String),
39 + /// Single text column for prompt, optional output column.
40 + Text {
41 + prompt_col: String,
42 + output_col: Option<String>,
43 + },
44 + /// Multiple columns combined (e.g., context + input for LongBench).
45 + Combined {
46 + cols: Vec<String>,
47 + output_col: Option<String>,
48 + },
49 +}
50 +
51 +/// Make a GET request with retry logic (3 retries with exponential backoff).
52 +/// Returns the parsed JSON response.
53 +fn get_with_retry(
54 + client: &reqwest::blocking::Client,
55 + url: &str,
56 + label: &str,
57 +) -> Result<serde_json::Value> {
58 + let max_retries = 3;
59 + for attempt in 0..=max_retries {
60 + let resp = match client.get(url).send() {
61 + Ok(r) => r,
62 + Err(e) => {
63 + if attempt < max_retries {
64 + std::thread::sleep(std::time::Duration::from_secs(2 * (attempt as u64 + 1)));
65 + continue;
66 + }
67 + return Err(BenchError::Config(format!(
68 + "{label} failed after {max_retries} retries: {e}"
69 + )));
70 + }
71 + };
72 +
73 + // Detect gated/private dataset errors
74 + let status = resp.status();
75 + if status.as_u16() == 401 || status.as_u16() == 403 {
76 + return Err(BenchError::Config(format!(
77 + "{label} returned HTTP {status}. This dataset may be gated or private. \
78 + Try setting the HF_TOKEN environment variable with a valid HuggingFace token."
79 + )));
80 + }
81 +
82 + if status.is_server_error() && attempt < max_retries {
83 + std::thread::sleep(std::time::Duration::from_secs(2 * (attempt as u64 + 1)));
84 + continue;
85 + }
86 +
87 + if !status.is_success() {
88 + return Err(BenchError::Config(format!(
89 + "{label} returned HTTP {status}"
90 + )));
91 + }
92 +
93 + let data: serde_json::Value = resp
94 + .json()
95 + .map_err(|e| BenchError::Config(format!("Failed to parse {label} response: {e}")))?;
96 + return Ok(data);
97 + }
98 + unreachable!()
99 +}
100 +
101 +/// Download an arbitrary HuggingFace dataset via the datasets-server REST API.
102 +///
103 +/// Returns `(cache_path, resolved_config, resolved_split)`.
104 +///
105 +/// If both `subset` and `split` are provided, the `/info` call is skipped as an optimization.
106 +/// Paginated download fetches rows in pages of 100 until `num_rows_needed` are collected
107 +/// or the dataset is exhausted.
108 +pub fn download_hf_dataset(
109 + dataset: &str,
110 + subset: Option<&str>,
111 + split: Option<&str>,
112 + num_rows_needed: usize,
113 +) -> Result<(String, String, String)> {
114 + let encoded_dataset: String =
115 + url::form_urlencoded::byte_serialize(dataset.as_bytes()).collect();
116 +
117 + let mut client_builder =
118 + reqwest::blocking::Client::builder().timeout(std::time::Duration::from_secs(120));
119 +
120 + // Add HF_TOKEN auth header if available
121 + if let Ok(token) = std::env::var("HF_TOKEN") {
122 + let mut headers = reqwest::header::HeaderMap::new();
123 + let header_value = reqwest::header::HeaderValue::from_str(&format!("Bearer {token}"))
124 + .map_err(|e| BenchError::Config(format!("Invalid HF_TOKEN: {e}")))?;
125 + headers.insert(reqwest::header::AUTHORIZATION, header_value);
126 + client_builder = client_builder.default_headers(headers);
127 + }
128 +
129 + let client = client_builder
130 + .build()
131 + .map_err(|e| BenchError::Config(format!("Failed to build HTTP client: {e}")))?;
132 +
133 + // Resolve config and split
134 + let (resolved_config, resolved_split) = if let (Some(cfg), Some(spl)) = (subset, split) {
135 + // Both provided — skip /info call
136 + (cfg.to_string(), spl.to_string())
137 + } else {
138 + // Call /info to discover available configs and splits
139 + let info_url =
140 + format!("https://datasets-server.huggingface.co/info?dataset={encoded_dataset}");
141 + let info = get_with_retry(&client, &info_url, "HF dataset /info")?;
142 +
143 + let dataset_info =
144 + info.get("dataset_info").and_then(|d| d.as_object()).ok_or_else(|| {
145 + BenchError::Config(format!(
146 + "No 'dataset_info' in /info response for '{dataset}'. \
147 + The dataset may not exist or may not be accessible."
148 + ))
149 + })?;
150 +
151 + // Resolve config
152 + let resolved_config = if let Some(user_cfg) = subset {
153 + if !dataset_info.contains_key(user_cfg) {
154 + let available: Vec<&String> = dataset_info.keys().collect();
155 + return Err(BenchError::Config(format!(
156 + "Config '{user_cfg}' not found in dataset '{dataset}'. Available: {available:?}"
157 + )));
158 + }
159 + user_cfg.to_string()
160 + } else if dataset_info.contains_key("default") {
161 + "default".to_string()
162 + } else {
163 + dataset_info.keys().next().cloned().ok_or_else(|| {
164 + BenchError::Config(format!("No configs found in dataset '{dataset}'"))
165 + })?
166 + };
167 +
168 + // Resolve split from the config's splits
169 + let splits_info = dataset_info
170 + .get(&resolved_config)
171 + .and_then(|c| c.get("splits"))
172 + .and_then(|s| s.as_object());
173 +
174 + let resolved_split = if let Some(user_spl) = split {
175 + if let Some(si) = splits_info
176 + && !si.contains_key(user_spl)
177 + {
178 + let available: Vec<&String> = si.keys().collect();
179 + return Err(BenchError::Config(format!(
180 + "Split '{user_spl}' not found in config '{resolved_config}' \
181 + of dataset '{dataset}'. Available: {available:?}"
182 + )));
183 + }
184 + user_spl.to_string()
185 + } else if let Some(si) = splits_info {
186 + // Priority: train > test > validation > first
187 + let priority = ["train", "test", "validation"];
188 + let mut found = None;
189 + for p in &priority {
190 + if si.contains_key(*p) {
191 + found = Some(p.to_string());
192 + break;
193 + }
194 + }
195 + found
196 + .unwrap_or_else(|| si.keys().next().cloned().unwrap_or_else(|| "train".to_string()))
197 + } else {
198 + "train".to_string()
199 + };
200 +
201 + (resolved_config, resolved_split)
202 + };
203 +
204 + println!("HF dataset: {dataset} (config={resolved_config}, split={resolved_split})");
205 +
206 + // Check cache
207 + let dir = cache_dir();
208 + std::fs::create_dir_all(&dir)?;
209 + let cache_path = dir.join(format!(
210 + "hf-{}-{}-{}.json",
211 + sanitize_name(dataset),
212 + sanitize_name(&resolved_config),
213 + sanitize_name(&resolved_split)
214 + ));
215 +
216 + if cache_path.exists() {
217 + let path_str = cache_path.to_string_lossy().to_string();
218 + println!("HF dataset cached: {path_str}");
219 + return Ok((path_str, resolved_config, resolved_split));
220 + }
221 +
222 + println!("Downloading HF dataset '{dataset}' from datasets-server...");
223 +
224 + let encoded_config: String =
225 + url::form_urlencoded::byte_serialize(resolved_config.as_bytes()).collect();
226 + let encoded_split: String =
227 + url::form_urlencoded::byte_serialize(resolved_split.as_bytes()).collect();
228 +
229 + let mut all_rows: Vec<serde_json::Value> = Vec::new();
230 + let mut offset = 0usize;
231 + let page_size = 100usize;
232 +
233 + loop {
234 + let url = format!(
235 + "https://datasets-server.huggingface.co/rows\
236 + ?dataset={encoded_dataset}\
237 + &config={encoded_config}\
238 + &split={encoded_split}\
239 + &offset={offset}\
240 + &length={page_size}"
241 + );
242 +
243 + let data = get_with_retry(&client, &url, "HF dataset /rows")?;
244 +
245 + let rows = data["rows"]
246 + .as_array()
247 + .ok_or_else(|| BenchError::Config("No 'rows' in API response".into()))?;
248 +
249 + if rows.is_empty() {
250 + break;
251 + }
252 +
253 + for row in rows {
254 + if let Some(row_data) = row.get("row") {
255 + all_rows.push(row_data.clone());
256 + }
257 + }
258 +
259 + let fetched = rows.len();
260 + offset += fetched;
261 +
262 + let total = data["num_rows_total"].as_u64().unwrap_or(0);
263 + eprint!("\r Fetched {offset}/{total} rows...");
264 +
265 + // Stop if we have enough rows or reached end of dataset
266 + if all_rows.len() >= num_rows_needed || fetched < page_size {
267 + break;
268 + }
269 + }
270 + eprintln!(); // newline after progress
271 +
272 + if all_rows.is_empty() {
273 + return Err(BenchError::Config(format!(
274 + "HF dataset '{dataset}' download returned no rows"
275 + )));
276 + }
277 +
278 + // Save to cache
279 + let json_str = serde_json::to_string(&all_rows)?;
280 + std::fs::write(&cache_path, &json_str)?;
281 +
282 + let path_str = cache_path.to_string_lossy().to_string();
283 + println!("HF dataset: {} rows saved to {path_str}", all_rows.len());
284 + Ok((path_str, resolved_config, resolved_split))
285 +}
286 +
287 +/// Check if a JSON value looks like a chat message (has role+content or from+value).
288 +fn is_chat_message(val: &serde_json::Value) -> bool {
289 + if let Some(obj) = val.as_object() {
290 + (obj.contains_key("role") && obj.contains_key("content"))
291 + || (obj.contains_key("from") && obj.contains_key("value"))
292 + } else {
293 + false
294 + }
295 +}
296 +
297 +/// Check if a JSON value is an array of chat messages.
298 +fn is_chat_array(val: &serde_json::Value) -> bool {
299 + val.as_array()
300 + .map(|arr| !arr.is_empty() && arr.iter().all(is_chat_message))
301 + .unwrap_or(false)
302 +}
303 +
304 +/// Known column names for chat-format data.
305 +const CHAT_COLUMNS: &[&str] = &["conversation", "conversations", "messages"];
306 +
307 +/// Known column names for single text prompts (in priority order).
308 +const TEXT_COLUMNS: &[&str] = &[
309 + "prompt",
310 + "question",
311 + "problem",
312 + "input",
313 + "text",
314 + "content",
315 + "instruction",
316 +];
317 +
318 +/// Known column names for output/completion data.
319 +const OUTPUT_COLUMNS: &[&str] = &[
320 + "completion",
321 + "response",
322 + "answer",
323 + "output",
324 + "solution",
325 + "answers",
326 +];
327 +
328 +/// Detect the column format from the first row of the dataset.
329 +fn detect_column_format(
330 + row: &serde_json::Value,
331 + text_column_override: Option<&str>,
332 +) -> Result<ColumnFormat> {
333 + let obj = row
334 + .as_object()
335 + .ok_or_else(|| BenchError::Config("HF dataset row is not a JSON object".into()))?;
336 +
337 + // Helper: find first matching output column
338 + let find_output_col = || -> Option<String> {
339 + for col in OUTPUT_COLUMNS {
340 + if obj.contains_key(*col) {
341 + return Some(col.to_string());
342 + }
343 + }
344 + None
345 + };
346 +
347 + // 1. User override via --hf-text-column
348 + if let Some(col_name) = text_column_override {
349 + if !obj.contains_key(col_name) {
350 + let available: Vec<&String> = obj.keys().collect();
351 + return Err(BenchError::Config(format!(
352 + "Column '{col_name}' not found in dataset. Available columns: {available:?}"
353 + )));
354 + }
355 + let val = &obj[col_name];
356 + if is_chat_array(val) {
357 + return Ok(ColumnFormat::Chat(col_name.to_string()));
358 + }
359 + return Ok(ColumnFormat::Text {
360 + prompt_col: col_name.to_string(),
361 + output_col: find_output_col(),
362 + });
363 + }
364 +
365 + // 2. Chat columns
366 + for col in CHAT_COLUMNS {
367 + if let Some(val) = obj.get(*col)
368 + && is_chat_array(val)
369 + {
370 + return Ok(ColumnFormat::Chat(col.to_string()));
371 + }
372 + }
373 +
374 + // 3. "turns" column — array of strings
375 + if let Some(val) = obj.get("turns")
376 + && let Some(arr) = val.as_array()
377 + && !arr.is_empty()
378 + && arr[0].is_string()
379 + {
380 + return Ok(ColumnFormat::Text {
381 + prompt_col: "turns".to_string(),
382 + output_col: find_output_col(),
383 + });
384 + }
385 +
386 + // 4. Combined: context + input
387 + if obj.contains_key("context") && obj.contains_key("input") {
388 + return Ok(ColumnFormat::Combined {
389 + cols: vec!["context".to_string(), "input".to_string()],
390 + output_col: find_output_col(),
391 + });
392 + }
393 +
394 + // 5. Single text columns
395 + for col in TEXT_COLUMNS {
396 + if obj.contains_key(*col) {
397 + return Ok(ColumnFormat::Text {
398 + prompt_col: col.to_string(),
399 + output_col: find_output_col(),
400 + });
401 + }
402 + }
403 +
404 + // Fallback: list available columns
405 + let available: Vec<&String> = obj.keys().collect();
406 + Err(BenchError::Config(format!(
407 + "Could not auto-detect prompt column in HF dataset. \
408 + Available columns: {available:?}. \
409 + Use --hf-text-column to specify the column containing prompts."
410 + )))
411 +}
412 +
413 +/// Extract the first user message from a chat message array.
414 +/// Supports both {role, content} and {from, value} formats.
415 +fn extract_chat_prompt(messages: &[serde_json::Value]) -> Option<String> {
416 + for msg in messages {
417 + let role = msg
418 + .get("role")
419 + .and_then(|r| r.as_str())
420 + .or_else(|| msg.get("from").and_then(|f| f.as_str()));
421 + let content = msg
422 + .get("content")
423 + .and_then(|c| c.as_str())
424 + .or_else(|| msg.get("value").and_then(|v| v.as_str()));
425 +
426 + if let (Some(role), Some(content)) = (role, content)
427 + && (role == "user" || role == "human")
428 + {
429 + return Some(content.to_string());
430 + }
431 + }
432 + None
433 +}
434 +
435 +/// Extract the first assistant message from a chat message array.
436 +fn extract_chat_completion(messages: &[serde_json::Value]) -> Option<String> {
437 + for msg in messages {
438 + let role = msg
439 + .get("role")
440 + .and_then(|r| r.as_str())
441 + .or_else(|| msg.get("from").and_then(|f| f.as_str()));
442 + let content = msg
443 + .get("content")
444 + .and_then(|c| c.as_str())
445 + .or_else(|| msg.get("value").and_then(|v| v.as_str()));
446 +
447 + if let (Some(role), Some(content)) = (role, content)
448 + && (role == "assistant" || role == "gpt")
449 + {
450 + return Some(content.to_string());
451 + }
452 + }
453 + None
454 +}
455 +
456 +/// Load an HF dataset from the cached JSON file and convert to SampleRequests.
457 +pub fn load_hf_dataset(
458 + tokenizer: &TokenizerKind,
459 + dataset_path: &str,
460 + num_requests: usize,
461 + hf_output_len: Option<usize>,
462 + seed: u64,
463 + request_id_prefix: &str,
464 + text_column_override: Option<&str>,
465 + no_oversample: bool,
466 + disable_shuffle: bool,
467 +) -> Result<Vec<SampleRequest>> {
468 + let content = std::fs::read_to_string(dataset_path).map_err(|e| {
469 + BenchError::Config(format!(
470 + "Failed to read HF dataset file '{dataset_path}': {e}"
471 + ))
472 + })?;
473 +
474 + let entries: Vec<serde_json::Value> = serde_json::from_str(&content)
475 + .map_err(|e| BenchError::Config(format!("Invalid JSON in HF dataset file: {e}")))?;
476 +
477 + if entries.is_empty() {
478 + return Err(BenchError::Config(
479 + "HF dataset file contains no rows".into(),
480 + ));
481 + }
482 +
483 + // Detect column format from first row
484 + let format = detect_column_format(&entries[0], text_column_override)?;
485 +
486 + // Print detected format
487 + match &format {
488 + ColumnFormat::Chat(col) => println!("HF dataset: detected chat column '{col}'"),
489 + ColumnFormat::Text {
490 + prompt_col,
491 + output_col,
492 + } => {
493 + let out_msg = output_col.as_deref().unwrap_or("none");
494 + println!("HF dataset: detected text column '{prompt_col}', output column: {out_msg}");
495 + }
496 + ColumnFormat::Combined { cols, output_col } => {
497 + let out_msg = output_col.as_deref().unwrap_or("none");
498 + println!(
499 + "HF dataset: detected combined columns {:?}, output column: {out_msg}",
500 + cols
501 + );
502 + }
503 + }
504 +
505 + // Build shuffled indices
506 + let mut indices: Vec<usize> = (0..entries.len()).collect();
507 + let mut rng = StdRng::seed_from_u64(seed);
508 + if !disable_shuffle {
509 + indices.shuffle(&mut rng);
510 + }
511 +
512 + let mut samples = Vec::new();
513 + let mut idx = 0;
514 + let mut warned_no_output = false;
515 +
516 + for &entry_idx in &indices {
517 + if samples.len() >= num_requests {
518 + break;
519 + }
520 +
521 + let row = &entries[entry_idx];
522 +
523 + // Extract prompt and optional completion based on format
524 + let (prompt, completion) = match &format {
525 + ColumnFormat::Chat(col) => {
526 + let messages =
527 + row.get(col.as_str()).and_then(|v| v.as_array()).cloned().unwrap_or_default();
528 + let prompt = match extract_chat_prompt(&messages) {
529 + Some(p) => p,
530 + None => continue,
531 + };
532 + let completion = extract_chat_completion(&messages);
533 + (prompt, completion)
534 + }
535 + ColumnFormat::Text {
536 + prompt_col,
537 + output_col,
538 + } => {
539 + let prompt_val = match row.get(prompt_col.as_str()) {
540 + Some(v) => v,
541 + None => continue,
542 + };
543 +
544 + // Handle "turns" column (array of strings — take first element)
545 + let prompt = if prompt_col == "turns" {
546 + match prompt_val.as_array().and_then(|arr| arr.first()) {
547 + Some(v) => v.as_str().unwrap_or("").to_string(),
548 + None => continue,
549 + }
550 + } else {
551 + prompt_val.as_str().unwrap_or("").to_string()
552 + };
553 +
554 + let completion = output_col.as_ref().and_then(|col| {
555 + row.get(col.as_str()).and_then(|v| {
556 + // Handle "answers" which may be an array
557 + if let Some(arr) = v.as_array() {
558 + arr.first().and_then(|a| a.as_str()).map(|s| s.to_string())
559 + } else {
560 + v.as_str().map(|s| s.to_string())
561 + }
562 + })
563 + });
564 +
565 + (prompt, completion)
566 + }
567 + ColumnFormat::Combined { cols, output_col } => {
568 + let parts: Vec<String> = cols
569 + .iter()
570 + .filter_map(|col| {
571 + row.get(col.as_str()).and_then(|v| v.as_str()).map(|s| s.to_string())
572 + })
573 + .collect();
574 +
575 + if parts.is_empty() {
576 + continue;
577 + }
578 +
579 + let prompt = parts.join("\n\n");
580 +
581 + let completion = output_col.as_ref().and_then(|col| {
582 + row.get(col.as_str()).and_then(|v| v.as_str()).map(|s| s.to_string())
583 + });
584 +
585 + (prompt, completion)
586 + }
587 + };
588 +
589 + if prompt.is_empty() {
590 + continue;
591 + }
592 +
593 + // Tokenize prompt
594 + let prompt_ids = tokenizer.encode(&prompt, false)?;
595 + let prompt_len = prompt_ids.len();
596 +
597 + if prompt_len < 4 {
598 + continue;
599 + }
600 +
601 + // Determine output length
602 + let output_len = if let Some(fixed_len) = hf_output_len {
603 + fixed_len
604 + } else if let Some(ref comp) = completion {
605 + if !comp.is_empty() {
606 + let comp_ids = tokenizer.encode(comp, false)?;
607 + let len = comp_ids.len();
608 + if len == 0 { 128 } else { len }
609 + } else {
610 + if !warned_no_output {
611 + eprintln!(
612 + "WARNING: No output column detected and --hf-output-len not set. \
613 + Using default output length of 128 tokens."
614 + );
615 + warned_no_output = true;
616 + }
617 + 128
618 + }
619 + } else {
620 + if !warned_no_output {
621 + eprintln!(
622 + "WARNING: No output column detected and --hf-output-len not set. \
623 + Using default output length of 128 tokens."
624 + );
625 + warned_no_output = true;
626 + }
627 + 128
628 + };
629 +
630 + samples.push(SampleRequest {
631 + prompt: Arc::from(prompt.as_str()),
632 + prompt_len,
633 + expected_output_len: output_len,
634 + request_id: Some(format!("{request_id_prefix}{idx}")),
635 + ..Default::default()
636 + });
637 + idx += 1;
638 + }
639 +
640 + // Oversample if needed
641 + if samples.len() < num_requests {
642 + if no_oversample {
643 + println!(
644 + "Skipping oversampling. Total samples: {} (requested: {num_requests})",
645 + samples.len()
646 + );
647 + } else if !samples.is_empty() {
648 + let original_len = samples.len();
649 + let needed = num_requests - original_len;
650 + for i in 0..needed {
651 + let mut req = samples[rng.random_range(0..original_len)].clone();
652 + req.request_id = Some(format!("{request_id_prefix}{}", original_len + i));
653 + samples.push(req);
654 + }
655 + println!(
656 + "Oversampled HF dataset from {original_len} to {} total samples.",
657 + samples.len()
658 + );
659 + }
660 + }
661 +
662 + if samples.is_empty() {
663 + return Err(BenchError::Config(
664 + "No valid samples after processing HF dataset. \
665 + Try a different --hf-text-column or check the dataset format."
666 + .into(),
667 + ));
668 + }
669 +
670 + Ok(samples)
671 +}
672 +
673 +#[cfg(test)]
674 +mod tests {
675 + use super::*;
676 +
677 + #[test]
678 + fn test_detect_chat_conversation() {
679 + let row = serde_json::json!({
680 + "conversation_id": "abc",
681 + "conversation": [
682 + {"role": "user", "content": "Hello"},
683 + {"role": "assistant", "content": "Hi there"}
684 + ]
685 + });
686 + let format = detect_column_format(&row, None).unwrap();
687 + assert!(matches!(format, ColumnFormat::Chat(ref col) if col == "conversation"));
688 + }
689 +
690 + #[test]
691 + fn test_detect_chat_messages() {
692 + let row = serde_json::json!({
693 + "messages": [
694 + {"role": "user", "content": "Hello"},
695 + {"role": "assistant", "content": "Hi"}
696 + ]
697 + });
698 + let format = detect_column_format(&row, None).unwrap();
699 + assert!(matches!(format, ColumnFormat::Chat(ref col) if col == "messages"));
700 + }
701 +
702 + #[test]
703 + fn test_detect_chat_sharegpt_format() {
704 + let row = serde_json::json!({
705 + "conversations": [
706 + {"from": "human", "value": "Hello"},
707 + {"from": "gpt", "value": "Hi!"}
708 + ]
709 + });
710 + let format = detect_column_format(&row, None).unwrap();
711 + assert!(matches!(format, ColumnFormat::Chat(ref col) if col == "conversations"));
712 + }
713 +
714 + #[test]
715 + fn test_detect_plain_prompt() {
716 + let row = serde_json::json!({
717 + "prompt": "What is 2+2?",
718 + "completion": "4"
719 + });
720 + let format = detect_column_format(&row, None).unwrap();
721 + match format {
722 + ColumnFormat::Text {
723 + ref prompt_col,
724 + ref output_col,
725 + } => {
726 + assert_eq!(prompt_col, "prompt");
727 + assert_eq!(output_col.as_deref(), Some("completion"));
728 + }
729 + _ => panic!("Expected Text format"),
730 + }
731 + }
732 +
733 + #[test]
734 + fn test_detect_question_answer() {
735 + let row = serde_json::json!({
736 + "question": "What is AI?",
737 + "answer": "Artificial intelligence"
738 + });
739 + let format = detect_column_format(&row, None).unwrap();
740 + match format {
741 + ColumnFormat::Text {
742 + ref prompt_col,
743 + ref output_col,
744 + } => {
745 + assert_eq!(prompt_col, "question");
746 + assert_eq!(output_col.as_deref(), Some("answer"));
747 + }
748 + _ => panic!("Expected Text format"),
749 + }
750 + }
751 +
752 + #[test]
753 + fn test_detect_combined_context_input() {
754 + let row = serde_json::json!({
755 + "context": "The quick brown fox...",
756 + "input": "What animal was mentioned?",
757 + "answers": ["fox"]
758 + });
759 + let format = detect_column_format(&row, None).unwrap();
760 + match format {
761 + ColumnFormat::Combined {
762 + ref cols,
763 + ref output_col,
764 + } => {
765 + assert_eq!(cols, &["context", "input"]);
766 + assert_eq!(output_col.as_deref(), Some("answers"));
767 + }
768 + _ => panic!("Expected Combined format"),
769 + }
770 + }
771 +
772 + #[test]
773 + fn test_detect_turns_array() {
774 + let row = serde_json::json!({
775 + "turns": ["First turn prompt", "Second turn"],
776 + "answer": "The answer"
777 + });
778 + let format = detect_column_format(&row, None).unwrap();
779 + match format {
780 + ColumnFormat::Text {
781 + ref prompt_col,
782 + ref output_col,
783 + } => {
784 + assert_eq!(prompt_col, "turns");
785 + assert_eq!(output_col.as_deref(), Some("answer"));
786 + }
787 + _ => panic!("Expected Text format"),
788 + }
789 + }
790 +
791 + #[test]
792 + fn test_detect_user_override() {
793 + let row = serde_json::json!({
794 + "my_custom_col": "Hello world",
795 + "answer": "response"
796 + });
797 + let format = detect_column_format(&row, Some("my_custom_col")).unwrap();
798 + match format {
799 + ColumnFormat::Text {
800 + ref prompt_col,
801 + ref output_col,
802 + } => {
803 + assert_eq!(prompt_col, "my_custom_col");
804 + assert_eq!(output_col.as_deref(), Some("answer"));
805 + }
806 + _ => panic!("Expected Text format"),
807 + }
808 + }
809 +
810 + #[test]
811 + fn test_detect_user_override_chat_column() {
812 + let row = serde_json::json!({
813 + "my_chat": [
814 + {"role": "user", "content": "Hi"},
815 + {"role": "assistant", "content": "Hello"}
816 + ]
817 + });
818 + let format = detect_column_format(&row, Some("my_chat")).unwrap();
819 + assert!(matches!(format, ColumnFormat::Chat(ref col) if col == "my_chat"));
820 + }
821 +
822 + #[test]
823 + fn test_detect_user_override_missing_column() {
824 + let row = serde_json::json!({"text": "hello"});
825 + let result = detect_column_format(&row, Some("nonexistent"));
826 + assert!(result.is_err());
827 + }
828 +
829 + #[test]
830 + fn test_detect_no_known_columns() {
831 + let row = serde_json::json!({"id": 1, "label": "positive"});
832 + let result = detect_column_format(&row, None);
833 + assert!(result.is_err());
834 + }
835 +
836 + #[test]
837 + fn test_extract_chat_prompt_role_content() {
838 + let messages = vec![
839 + serde_json::json!({"role": "user", "content": "What is AI?"}),
840 + serde_json::json!({"role": "assistant", "content": "AI is..."}),
841 + ];
842 + assert_eq!(
843 + extract_chat_prompt(&messages),
844 + Some("What is AI?".to_string())
845 + );
846 + }
847 +
848 + #[test]
849 + fn test_extract_chat_prompt_from_value() {
850 + let messages = vec![
851 + serde_json::json!({"from": "human", "value": "Hello"}),
852 + serde_json::json!({"from": "gpt", "value": "Hi!"}),
853 + ];
854 + assert_eq!(extract_chat_prompt(&messages), Some("Hello".to_string()));
855 + }
856 +
857 + #[test]
858 + fn test_extract_chat_prompt_system_first() {
859 + let messages = vec![
860 + serde_json::json!({"role": "system", "content": "You are helpful"}),
861 + serde_json::json!({"role": "user", "content": "Tell me a joke"}),
862 + serde_json::json!({"role": "assistant", "content": "Why did..."}),
863 + ];
864 + assert_eq!(
865 + extract_chat_prompt(&messages),
866 + Some("Tell me a joke".to_string())
867 + );
868 + }
869 +
870 + #[test]
871 + fn test_extract_chat_prompt_no_user() {
872 + let messages = vec![serde_json::json!({"role": "system", "content": "You are helpful"})];
873 + assert_eq!(extract_chat_prompt(&messages), None);
874 + }
875 +
876 + #[test]
877 + fn test_extract_chat_completion_role_content() {
878 + let messages = vec![
879 + serde_json::json!({"role": "user", "content": "Hi"}),
880 + serde_json::json!({"role": "assistant", "content": "Hello!"}),
881 + ];
882 + assert_eq!(
883 + extract_chat_completion(&messages),
884 + Some("Hello!".to_string())
885 + );
886 + }
887 +
888 + #[test]
889 + fn test_extract_chat_completion_gpt() {
890 + let messages = vec![
891 + serde_json::json!({"from": "human", "value": "Hi"}),
892 + serde_json::json!({"from": "gpt", "value": "Hello!"}),
893 + ];
894 + assert_eq!(
895 + extract_chat_completion(&messages),
896 + Some("Hello!".to_string())
897 + );
898 + }
899 +
900 + #[test]
901 + fn test_is_chat_message_valid() {
902 + assert!(is_chat_message(
903 + &serde_json::json!({"role": "user", "content": "hi"})
904 + ));
905 + assert!(is_chat_message(
906 + &serde_json::json!({"from": "human", "value": "hi"})
907 + ));
908 + }
909 +
910 + #[test]
911 + fn test_is_chat_message_invalid() {
912 + assert!(!is_chat_message(&serde_json::json!({"text": "hi"})));
913 + assert!(!is_chat_message(&serde_json::json!("just a string")));
914 + assert!(!is_chat_message(&serde_json::json!(42)));
915 + }
916 +
917 + #[test]
918 + fn test_is_chat_array_valid() {
919 + let val = serde_json::json!([
920 + {"role": "user", "content": "hi"},
921 + {"role": "assistant", "content": "hello"}
922 + ]);
923 + assert!(is_chat_array(&val));
924 + }
925 +
926 + #[test]
927 + fn test_is_chat_array_invalid() {
928 + assert!(!is_chat_array(&serde_json::json!([])));
929 + assert!(!is_chat_array(&serde_json::json!(["a", "b"])));
930 + assert!(!is_chat_array(&serde_json::json!("not an array")));
931 + }
932 +
933 + #[test]
934 + fn test_sanitize_name() {
935 + assert_eq!(
936 + sanitize_name("allenai/WildChat-4.8M"),
937 + "allenai_WildChat-4.8M"
938 + );
939 + assert_eq!(sanitize_name("simple-name"), "simple-name");
940 + assert_eq!(sanitize_name("org/sub/deep"), "org_sub_deep");
941 + }
942 +
943 + // --- extract_chat_prompt edge cases ---
944 +
945 + #[test]
946 + fn test_extract_chat_prompt_empty_messages() {
947 + let messages: Vec<serde_json::Value> = vec![];
948 + assert_eq!(extract_chat_prompt(&messages), None);
949 + }
950 +
951 + #[test]
952 + fn test_extract_chat_prompt_model_role_skipped() {
953 + // "model" role is not "user" or "human" — should be skipped
954 + let messages = vec![
955 + serde_json::json!({"role": "model", "content": "I am the model"}),
956 + serde_json::json!({"role": "assistant", "content": "Hello"}),
957 + ];
958 + assert_eq!(extract_chat_prompt(&messages), None);
959 + }
960 +
961 + #[test]
962 + fn test_extract_chat_prompt_model_role_before_user() {
963 + // "model" role before user: only "user" should be matched
964 + let messages = vec![
965 + serde_json::json!({"role": "model", "content": "intro"}),
966 + serde_json::json!({"role": "user", "content": "actual user message"}),
967 + ];
968 + assert_eq!(
969 + extract_chat_prompt(&messages),
970 + Some("actual user message".to_string())
971 + );
972 + }
973 +
974 + // --- extract_chat_completion edge cases ---
975 +
976 + #[test]
977 + fn test_extract_chat_completion_no_assistant() {
978 + let messages = vec![
979 + serde_json::json!({"role": "user", "content": "What is 2+2?"}),
980 + serde_json::json!({"role": "system", "content": "You are helpful"}),
981 + ];
982 + assert_eq!(extract_chat_completion(&messages), None);
983 + }
984 +
985 + #[test]
986 + fn test_extract_chat_completion_empty_messages() {
987 + let messages: Vec<serde_json::Value> = vec![];
988 + assert_eq!(extract_chat_completion(&messages), None);
989 + }
990 +
991 + #[test]
992 + fn test_extract_chat_completion_model_role_not_matched() {
993 + // "model" role is not "assistant" or "gpt" — should return None
994 + let messages = vec![
995 + serde_json::json!({"role": "user", "content": "Hello"}),
996 + serde_json::json!({"role": "model", "content": "I should not be returned"}),
997 + ];
998 + assert_eq!(extract_chat_completion(&messages), None);
999 + }
1000 +
1001 + // --- load_hf_dataset integration tests (require built-in tiktoken, no network) ---
1002 +
1003 + /// Build a gpt2 tokenizer using built-in tiktoken encoding (no network required).
1004 + fn builtin_tokenizer() -> crate::tokenizer::TokenizerKind {
1005 + crate::tokenizer::load_tokenizer("gpt2", false, None)
1006 + .expect("gpt2 built-in tiktoken should always load without network")
1007 + }
1008 +
1009 + /// Write JSON data to a unique temp file and return the path string.
1010 + fn write_temp_json(name: &str, data: &serde_json::Value) -> String {
1011 + let path = std::env::temp_dir().join(format!("vllm-bench-test-{name}.json"));
1012 + std::fs::write(&path, serde_json::to_string(data).unwrap()).unwrap();
1013 + path.to_string_lossy().to_string()
1014 + }
1015 +
1016 + #[test]
1017 + fn test_load_hf_dataset_chat_format() {
1018 + let tok = builtin_tokenizer();
1019 + let data = serde_json::json!([
1020 + {
1021 + "conversation": [
1022 + {"role": "user", "content": "What is the meaning of life, the universe, and everything?"},
1023 + {"role": "assistant", "content": "The answer is 42, according to Douglas Adams."}
1024 + ]
1025 + },
1026 + {
1027 + "conversation": [
1028 + {"role": "user", "content": "Tell me about quantum computing and its applications in modern science."},
1029 + {"role": "assistant", "content": "Quantum computing uses quantum bits to perform calculations."}
1030 + ]
1031 + }
1032 + ]);
1033 + let path = write_temp_json("chat-format", &data);
1034 +
1035 + let result =
1036 + load_hf_dataset(&tok, &path, 2, Some(50), 42, "test-", None, false, false).unwrap();
1037 +
1038 + assert_eq!(result.len(), 2);
1039 + assert!(
1040 + result.iter().all(|r| r.expected_output_len == 50),
1041 + "all samples should have fixed output len 50"
1042 + );
1043 + }
1044 +
1045 + #[test]
1046 + fn test_load_hf_dataset_plain_text_format() {
1047 + let tok = builtin_tokenizer();
1048 + let data = serde_json::json!([
1049 + {"prompt": "What is the capital of France and why is it important in European history?", "completion": "Paris is the capital of France."},
1050 + {"prompt": "Explain the difference between machine learning and deep learning in simple terms.", "completion": "Machine learning is a subset of AI."},
1051 + {"prompt": "How does photosynthesis work in plants and what role does chlorophyll play?", "completion": "Photosynthesis converts light to energy."}
1052 + ]);
1053 + let path = write_temp_json("plain-text-format", &data);
1054 +
1055 + let result = load_hf_dataset(&tok, &path, 3, None, 0, "req-", None, false, false).unwrap();
1056 +
1057 + assert_eq!(result.len(), 3);
1058 + // Without hf_output_len override, output len is derived from the completion tokens
1059 + assert!(result.iter().all(|r| r.expected_output_len > 0));
1060 + }
1061 +
1062 + #[test]
1063 + fn test_load_hf_dataset_combined_format() {
1064 + let tok = builtin_tokenizer();
1065 + let data = serde_json::json!([
1066 + {
1067 + "context": "The quick brown fox jumped over the lazy dog near the river bank.",
1068 + "input": "What animal jumped over the dog in this sentence?",
1069 + "answers": ["The quick brown fox"]
1070 + },
1071 + {
1072 + "context": "Albert Einstein developed the theory of relativity in the early twentieth century.",
1073 + "input": "Who developed the theory of relativity and when approximately?",
1074 + "answers": ["Albert Einstein"]
1075 + }
1076 + ]);
1077 + let path = write_temp_json("combined-format", &data);
1078 +
1079 + let result =
1080 + load_hf_dataset(&tok, &path, 2, Some(64), 1, "comb-", None, false, false).unwrap();
1081 +
1082 + assert_eq!(result.len(), 2);
1083 + assert!(result.iter().all(|r| r.expected_output_len == 64));
1084 + // Combined format joins context + input with "\n\n"
1085 + assert!(result.iter().all(|r| r.prompt.contains('\n')));
1086 + }
1087 +
1088 + #[test]
1089 + fn test_load_hf_dataset_empty_returns_error() {
1090 + let tok = builtin_tokenizer();
1091 + let data = serde_json::json!([]);
1092 + let path = write_temp_json("empty-dataset", &data);
1093 +
1094 + let result = load_hf_dataset(&tok, &path, 5, None, 0, "test-", None, false, false);
1095 +
1096 + assert!(result.is_err(), "empty dataset should return an error");
1097 + let err = result.unwrap_err().to_string();
1098 + assert!(
1099 + err.contains("no rows") || err.contains("No rows") || err.contains("empty"),
1100 + "error should mention empty/no rows: {err}"
1101 + );
1102 + }
1103 +
1104 + #[test]
1105 + fn test_load_hf_dataset_all_rows_filtered_too_short() {
1106 + let tok = builtin_tokenizer();
1107 + // Very short prompts that will tokenize to fewer than 4 tokens and be filtered out
1108 + let data = serde_json::json!([
1109 + {"prompt": "Hi", "completion": "Ok"},
1110 + {"prompt": "Yes", "completion": "No"},
1111 + {"prompt": "Ok", "completion": "Fine"}
1112 + ]);
1113 + let path = write_temp_json("all-filtered", &data);
1114 +
1115 + let result = load_hf_dataset(&tok, &path, 3, None, 0, "test-", None, false, false);
1116 +
1117 + assert!(
1118 + result.is_err(),
1119 + "all rows being too short should return an error"
1120 + );
1121 + }
1122 +
1123 + #[test]
1124 + fn test_load_hf_dataset_hf_output_len_override() {
1125 + let tok = builtin_tokenizer();
1126 + let data = serde_json::json!([
1127 + {"prompt": "Describe the history of the Roman Empire and its eventual decline over centuries.", "completion": "The Roman Empire fell for many complex reasons."},
1128 + {"prompt": "What are the key differences between supervised and unsupervised machine learning?", "completion": "Supervised learning uses labeled data while unsupervised does not."},
1129 + {"prompt": "Explain how neural networks are inspired by the human brain structure and function.", "completion": "Neural networks mimic brain neurons with layers of nodes."}
1130 + ]);
1131 + let path = write_temp_json("output-len-override", &data);
1132 +
1133 + let fixed_len = 77usize;
1134 + let result = load_hf_dataset(
1135 + &tok,
1136 + &path,
1137 + 3,
1138 + Some(fixed_len),
1139 + 42,
1140 + "t-",
1141 + None,
1142 + false,
1143 + false,
1144 + )
1145 + .unwrap();
1146 +
1147 + assert!(!result.is_empty());
1148 + assert!(
1149 + result.iter().all(|r| r.expected_output_len == fixed_len),
1150 + "all samples must use the fixed output length {fixed_len}"
1151 + );
1152 + }
1153 +
1154 + #[test]
1155 + fn test_load_hf_dataset_no_oversample() {
1156 + let tok = builtin_tokenizer();
1157 + // 2 valid rows, but request 10 with no_oversample=true
1158 + let data = serde_json::json!([
1159 + {"prompt": "Explain how photosynthesis converts sunlight to energy in plant cells.", "completion": "Plants use chlorophyll to absorb sunlight."},
1160 + {"prompt": "What is the difference between a virus and a bacterium in terms of biology?", "completion": "Viruses need host cells while bacteria are self-sufficient."}
1161 + ]);
1162 + let path = write_temp_json("no-oversample", &data);
1163 +
1164 + let result = load_hf_dataset(
1165 + &tok,
1166 + &path,
1167 + 10,
1168 + Some(32),
1169 + 0,
1170 + "t-",
1171 + None,
1172 + true, // no_oversample
1173 + false,
1174 + )
1175 + .unwrap();
1176 +
1177 + // Should have at most 2 samples (the actual dataset size), not 10
1178 + assert!(
1179 + result.len() <= 2,
1180 + "no_oversample should cap result at dataset size, got {}",
1181 + result.len()
1182 + );
1183 + }
1184 +
1185 + #[test]
1186 + fn test_load_hf_dataset_disable_shuffle_preserves_order() {
1187 + let tok = builtin_tokenizer();
1188 + let data = serde_json::json!([
1189 + {"prompt": "Alpha prompt about the first topic in our alphabetically ordered sequence.", "completion": "Alpha answer"},
1190 + {"prompt": "Beta prompt about the second topic continuing from our ordered sequence here.", "completion": "Beta answer"},
1191 + {"prompt": "Gamma prompt about the third item in our clearly ordered alphabetical sequence.", "completion": "Gamma answer"}
1192 + ]);
1193 + let path = write_temp_json("disable-shuffle", &data);
1194 +
1195 + let result = load_hf_dataset(
1196 + &tok,
1197 + &path,
1198 + 3,
1199 + Some(20),
1200 + 99,
1201 + "t-",
1202 + None,
1203 + false,
1204 + true, // disable_shuffle
1205 + )
1206 + .unwrap();
1207 +
1208 + assert_eq!(result.len(), 3);
1209 + // With disable_shuffle, rows come out in original order:
1210 + // Alpha < Beta < Gamma (alphabetical), so first prompt contains "Alpha"
1211 + assert!(
1212 + result[0].prompt.to_lowercase().contains("alpha"),
1213 + "first result should be Alpha prompt with shuffle disabled, got: {}",
1214 + result[0].prompt
1215 + );
1216 + assert!(
1217 + result[1].prompt.to_lowercase().contains("beta"),
1218 + "second result should be Beta prompt, got: {}",
1219 + result[1].prompt
1220 + );
1221 + }
1222 +
1223 + #[test]
1224 + fn test_load_hf_dataset_request_id_prefix() {
1225 + let tok = builtin_tokenizer();
1226 + let data = serde_json::json!([
1227 + {"prompt": "What is the boiling point of water at sea level under standard atmospheric pressure?", "completion": "Water boils at 100 degrees Celsius."},
1228 + {"prompt": "Describe the structure of DNA and how genetic information is encoded within it.", "completion": "DNA is a double helix with base pairs."}
1229 + ]);
1230 + let path = write_temp_json("request-id-prefix", &data);
1231 +
1232 + let result =
1233 + load_hf_dataset(&tok, &path, 2, Some(10), 0, "myprefix-", None, false, false).unwrap();
1234 +
1235 + assert_eq!(result.len(), 2);
1236 + assert_eq!(result[0].request_id.as_deref(), Some("myprefix-0"));
1237 + assert_eq!(result[1].request_id.as_deref(), Some("myprefix-1"));
1238 + }
1239 +}
rust/src/bench/src/datasets/mod.rs new
+207
@@ -0,0 +1,207 @@
1 +// SPDX-License-Identifier: Apache-2.0
2 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +pub mod custom;
5 +pub mod hf_dataset;
6 +pub mod multi_turn;
7 +pub mod prefix_repetition;
8 +pub mod random;
9 +pub mod random_mm;
10 +pub mod random_rerank;
11 +pub mod sharegpt;
12 +pub mod sonnet;
13 +pub mod speed_bench;
14 +
15 +use std::sync::Arc;
16 +
17 +/// Represents a single inference request for benchmarking.
18 +/// Matches Python's SampleRequest dataclass from datasets.py:71-82.
19 +///
20 +/// `prompt` uses `Arc<str>` to avoid expensive String clones when distributing
21 +/// requests across tokio tasks. At 100k prompts with 8k tokens each, this saves
22 +/// ~3GB of peak memory vs cloning String per task.
23 +#[derive(Debug, Clone)]
24 +pub struct SampleRequest {
25 + pub prompt: Arc<str>,
26 + pub prompt_len: usize,
27 + pub expected_output_len: usize,
28 + pub request_id: Option<String>,
29 + /// Pre-computed token IDs for this prompt.
30 + /// When set, the completions backend sends these directly via `prompt_token_ids`
31 + /// instead of the text `prompt`, avoiding server-side re-tokenization.
32 + pub prompt_token_ids: Option<Arc<[u32]>>,
33 + /// Multimodal content items as pre-serialized JSON fragments.
34 + /// Each `Arc<str>` is a complete JSON object string, e.g.
35 + /// `{"type":"image_url","image_url":{"url":"data:image/jpeg;base64,..."}}`
36 + ///
37 + /// Pre-serialized to avoid:
38 + /// 1. `serde_json::Value` tree overhead (3 Maps + keys per image)
39 + /// 2. Deep-cloning ~200KB+ base64 data when building request payloads
40 + ///
41 + /// Double-`Arc` for zero-cost sharing: outer Arc for the slice, inner Arc for each fragment.
42 + pub multi_modal_content: Option<Arc<[Arc<str>]>>,
43 + /// Pre-serialized OpenAI chat `messages` array as a complete JSON string,
44 + /// e.g. `[{"role":"user","content":[{"type":"text","text":"..."},{"type":"image_url",...}]}]`.
45 + ///
46 + /// Set by datasets when `--enable-multimodal-chat` is on (mirrors Python's
47 + /// `apply_multimodal_chat_transformation`: the dataset builds the chat messages
48 + /// and the backend sends them verbatim). When set, `multi_modal_content` is None
49 + /// and the mm items are embedded here instead. `prompt` still holds the text part
50 + /// for token accounting and /tokenize verification.
51 + pub chat_messages_json: Option<Arc<str>>,
52 + /// Multiple text inputs for one request (pooling backends only).
53 + /// Embeddings send it as `"input": [t1, t2, ...]` (--random-batch-size);
54 + /// rerank sends `[0]` as the query and `[1..]` as documents (random-rerank).
55 + /// Mirrors Python's list-valued `SampleRequest.prompt`.
56 + pub prompt_list: Option<Arc<[Arc<str>]>>,
57 +}
58 +
59 +impl Default for SampleRequest {
60 + /// Empty request; struct-update base so dataset builders only spell out the
61 + /// fields they set (new optional fields then don't touch every call site).
62 + fn default() -> Self {
63 + Self {
64 + prompt: Arc::from(""),
65 + prompt_len: 0,
66 + expected_output_len: 0,
67 + request_id: None,
68 + prompt_token_ids: None,
69 + multi_modal_content: None,
70 + chat_messages_json: None,
71 + prompt_list: None,
72 + }
73 + }
74 +}
75 +
76 +/// Oversample `requests` up to `num_requests` by cloning random entries
77 +/// (seeded by list length for determinism), renumbering their request ids.
78 +/// No-op when enough samples exist, `no_oversample` is set, or the list is empty.
79 +/// Mirrors Python `BenchmarkDataset.maybe_oversample_requests`.
80 +pub fn oversample_requests(
81 + requests: &mut Vec<SampleRequest>,
82 + num_requests: usize,
83 + request_id_prefix: &str,
84 + no_oversample: bool,
85 +) {
86 + use rand::rngs::StdRng;
87 + use rand::{Rng, SeedableRng};
88 +
89 + if requests.len() >= num_requests || requests.is_empty() {
90 + return;
91 + }
92 + if no_oversample {
93 + println!(
94 + "Skipping oversampling. Total samples: {} (requested: {num_requests})",
95 + requests.len()
96 + );
97 + return;
98 + }
99 + let original_len = requests.len();
100 + let mut rng = StdRng::seed_from_u64(original_len as u64);
101 + for i in 0..(num_requests - original_len) {
102 + let mut req = requests[rng.random_range(0..original_len)].clone();
103 + req.request_id = Some(format!("{request_id_prefix}{}", original_len + i));
104 + requests.push(req);
105 + }
106 + println!(
107 + "Oversampled requests from {original_len} to {} total samples.",
108 + requests.len()
109 + );
110 +}
111 +
112 +/// Group already-generated single-input requests into batched requests of
113 +/// `batch_size` inputs each (embeddings/pooling only). Mirrors Python
114 +/// `RandomDataset.sample` batching: prompt becomes a list, prompt_len is the
115 +/// sum over the batch, request ids are renumbered per batch.
116 +/// `batch_size <= 1` returns the input unchanged.
117 +pub fn batch_requests(
118 + requests: Vec<SampleRequest>,
119 + batch_size: usize,
120 + request_id_prefix: &str,
121 +) -> Vec<SampleRequest> {
122 + if batch_size <= 1 {
123 + return requests;
124 + }
125 + requests
126 + .chunks(batch_size)
127 + .enumerate()
128 + .map(|(batch_idx, batch)| SampleRequest {
129 + prompt_list: Some(batch.iter().map(|r| r.prompt.clone()).collect()),
130 + prompt_len: batch.iter().map(|r| r.prompt_len).sum(),
131 + expected_output_len: 0,
132 + request_id: Some(format!("{request_id_prefix}{batch_idx}")),
133 + ..Default::default()
134 + })
135 + .collect()
136 +}
137 +
138 +#[cfg(test)]
139 +mod tests {
140 + use super::*;
141 +
142 + fn req(prompt: &str, len: usize) -> SampleRequest {
143 + SampleRequest {
144 + prompt: Arc::from(prompt),
145 + prompt_len: len,
146 + expected_output_len: 128,
147 + ..Default::default()
148 + }
149 + }
150 +
151 + #[test]
152 + fn test_batch_requests_groups_and_sums() {
153 + let reqs = vec![
154 + req("a", 10),
155 + req("b", 20),
156 + req("c", 30),
157 + req("d", 40),
158 + req("e", 50),
159 + ];
160 + let batched = batch_requests(reqs, 2, "t-");
161 + assert_eq!(batched.len(), 3); // 2 + 2 + 1
162 + let first = batched[0].prompt_list.as_ref().unwrap();
163 + assert_eq!(first.len(), 2);
164 + assert_eq!(&*first[0], "a");
165 + assert_eq!(batched[0].prompt_len, 30);
166 + assert_eq!(batched[0].expected_output_len, 0);
167 + assert_eq!(batched[0].request_id.as_deref(), Some("t-0"));
168 + assert_eq!(batched[2].prompt_list.as_ref().unwrap().len(), 1);
169 + assert_eq!(batched[2].prompt_len, 50);
170 + }
171 +
172 + #[test]
173 + fn test_batch_requests_size_one_is_identity() {
174 + let reqs = vec![req("a", 10), req("b", 20)];
175 + let out = batch_requests(reqs, 1, "t-");
176 + assert_eq!(out.len(), 2);
177 + assert!(out[0].prompt_list.is_none());
178 + assert_eq!(&*out[0].prompt, "a");
179 + }
180 +
181 + #[test]
182 + fn test_oversample_requests() {
183 + let mut reqs = vec![req("a", 10), req("b", 20)];
184 + oversample_requests(&mut reqs, 5, "t-", false);
185 + assert_eq!(reqs.len(), 5);
186 + assert_eq!(reqs[4].request_id.as_deref(), Some("t-4"));
187 +
188 + let mut reqs = vec![req("a", 10)];
189 + oversample_requests(&mut reqs, 5, "t-", true); // no_oversample
190 + assert_eq!(reqs.len(), 1);
191 + }
192 +}
193 +
194 +/// A single turn in a multi-turn conversation.
195 +#[derive(Debug, Clone)]
196 +pub struct ConversationTurn {
197 + pub user_message: Arc<str>,
198 + pub user_message_len: usize,
199 + pub expected_output_len: usize,
200 +}
201 +
202 +/// A complete multi-turn conversation with all turns pre-generated.
203 +#[derive(Debug, Clone)]
204 +pub struct MultiTurnConversation {
205 + pub conversation_id: String,
206 + pub turns: Vec<ConversationTurn>,
207 +}
rust/src/bench/src/datasets/multi_turn.rs new
+777
@@ -0,0 +1,777 @@
1 +// SPDX-License-Identifier: Apache-2.0
2 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +use std::sync::Arc;
5 +
6 +use rand::rngs::StdRng;
7 +use rand::seq::SliceRandom;
8 +use rand::{Rng, SeedableRng};
9 +use rayon::prelude::*;
10 +
11 +use super::{ConversationTurn, MultiTurnConversation};
12 +use crate::error::{BenchError, Result};
13 +use crate::tokenizer::TokenizerKind;
14 +
15 +/// Configuration for generating random multi-turn conversations.
16 +#[derive(Debug, Clone)]
17 +pub struct MultiTurnRandomConfig {
18 + pub num_conversations: usize,
19 + pub min_turns: usize,
20 + pub max_turns: usize,
21 + /// Shared prefix length prepended to the conversation.
22 + ///
23 + /// In normal accumulated-history mode this is added to turn 0, so all
24 + /// later turns inherit it through history. In no-history prefix-sharing
25 + /// mode it is added to every independent turn.
26 + pub prefix_len: usize,
27 + /// Input length for turn 0.
28 + pub input_len: usize,
29 + /// Input length for turns 1+. 0 = fallback to input_len.
30 + pub per_turn_input_len: usize,
31 + pub output_len: usize,
32 + pub seed: u64,
33 + pub request_id_prefix: String,
34 + pub prefix_sharing_config: Option<PrefixSharingConfig>,
35 +}
36 +
37 +/// Configuration for 3-tier prefix sharing in multi-turn user messages.
38 +#[derive(Debug, Clone)]
39 +pub struct PrefixSharingConfig {
40 + /// Fraction of per-turn input tokens shared across ALL conversations.
41 + pub global_ratio: f64,
42 + /// Fraction of per-turn input tokens shared within each conversation.
43 + pub conversation_ratio: f64,
44 +}
45 +
46 +/// Generate a deterministic token sequence from allowed tokens using offset+modulo.
47 +fn make_token_seq(allowed_tokens: &[u32], offset: usize, len: usize) -> Vec<u32> {
48 + let at_len = allowed_tokens.len();
49 + (0..len).map(|i| allowed_tokens[(offset + i) % at_len]).collect()
50 +}
51 +
52 +/// Generate synthetic multi-turn conversations with random user messages.
53 +///
54 +/// Each conversation has `num_turns` turns, each with a random user prompt
55 +/// of `input_len` tokens and `output_len` expected output tokens.
56 +pub fn generate_multi_turn_random(
57 + tokenizer: &TokenizerKind,
58 + cfg: &MultiTurnRandomConfig,
59 +) -> Result<Vec<MultiTurnConversation>> {
60 + let num_conversations = cfg.num_conversations;
61 + let min_turns = cfg.min_turns;
62 + let max_turns = cfg.max_turns;
63 + let prefix_len = cfg.prefix_len;
64 + let input_len = cfg.input_len;
65 + let output_len = cfg.output_len;
66 + let seed = cfg.seed;
67 + let request_id_prefix = &cfg.request_id_prefix;
68 + let allowed_tokens = tokenizer.get_allowed_tokens();
69 + if allowed_tokens.is_empty() {
70 + return Err(BenchError::Tokenizer("No allowed tokens found".into()));
71 + }
72 +
73 + let vocab_size = tokenizer.vocab_size() as usize;
74 + let num_special = tokenizer.num_special_tokens_to_add();
75 + let real_input_len = input_len.saturating_sub(num_special);
76 + let real_per_turn_len = if cfg.per_turn_input_len > 0 {
77 + cfg.per_turn_input_len.saturating_sub(num_special)
78 + } else {
79 + real_input_len
80 + };
81 +
82 + if real_input_len < 1 {
83 + return Err(BenchError::Config(format!(
84 + "--random-input-len too small: with {num_special} special tokens, \
85 + effective input length is {real_input_len}"
86 + )));
87 + }
88 + if real_per_turn_len < 1 {
89 + return Err(BenchError::Config(format!(
90 + "--per-turn-input-len too small: with {num_special} special tokens, \
91 + effective per-turn input length is {real_per_turn_len}"
92 + )));
93 + }
94 +
95 + // Prefix sharing mode: generate 3-tier prefixed messages
96 + let mut rng = StdRng::seed_from_u64(seed);
97 + if let Some(ref ps_cfg) = cfg.prefix_sharing_config {
98 + return generate_prefix_sharing_conversations(
99 + tokenizer,
100 + cfg,
101 + ps_cfg,
102 + &allowed_tokens,
103 + &mut rng,
104 + );
105 + }
106 + let shared_prefix_text =
107 + generate_shared_prefix_text(tokenizer, &allowed_tokens, prefix_len, seed)?;
108 +
109 + // Pre-generate per-conversation turn counts and per-turn offsets deterministically.
110 + // Turn counts are drawn first so the RNG sequence is stable regardless of vocab_size.
111 + let conv_turn_counts: Vec<usize> = (0..num_conversations)
112 + .map(|_| {
113 + if min_turns == max_turns {
114 + min_turns
115 + } else {
116 + rng.random_range(min_turns..=max_turns)
117 + }
118 + })
119 + .collect();
120 +
121 + let offsets: Vec<Vec<usize>> = conv_turn_counts
122 + .iter()
123 + .map(|&n| (0..n).map(|_| rng.random_range(0..vocab_size)).collect())
124 + .collect();
125 +
126 + // Parallel generation across conversations
127 + offsets
128 + .par_iter()
129 + .enumerate()
130 + .map(|(conv_idx, conv_offsets)| {
131 + let mut turns = Vec::with_capacity(conv_offsets.len());
132 + for (turn_idx, &offset) in conv_offsets.iter().enumerate() {
133 + let target_len = if turn_idx == 0 {
134 + real_input_len
135 + } else {
136 + real_per_turn_len
137 + };
138 + // Use max_turns stride to keep offsets unique across variable-length convs
139 + let inner_seq = make_token_seq(
140 + &allowed_tokens,
141 + offset + conv_idx * max_turns + turn_idx,
142 + target_len,
143 + );
144 +
145 + let (prompt, adjusted) =
146 + gen_prompt_to_target_len(tokenizer, &inner_seq, target_len)?;
147 + let (prompt, token_len) = if turn_idx == 0 && !shared_prefix_text.is_empty() {
148 + let combined = format!("{}{}", &*shared_prefix_text, prompt);
149 + let token_len = tokenizer.encode(&combined, false)?.len();
150 + (combined, token_len)
151 + } else {
152 + (prompt, adjusted.len())
153 + };
154 +
155 + turns.push(ConversationTurn {
156 + user_message: Arc::from(prompt),
157 + user_message_len: token_len,
158 + expected_output_len: output_len,
159 + });
160 + }
161 +
162 + Ok(MultiTurnConversation {
163 + conversation_id: format!("{request_id_prefix}conv-{conv_idx}"),
164 + turns,
165 + })
166 + })
167 + .collect()
168 +}
169 +
170 +/// Generate conversations with 3-tier prefix sharing.
171 +///
172 +/// Each turn's user message = [global_prefix][conversation_prefix][unique_suffix].
173 +/// No history accumulation — each turn sends only its own fixed-length message.
174 +fn generate_prefix_sharing_conversations(
175 + tokenizer: &TokenizerKind,
176 + cfg: &MultiTurnRandomConfig,
177 + ps_cfg: &PrefixSharingConfig,
178 + allowed_tokens: &[u32],
179 + rng: &mut StdRng,
180 +) -> Result<Vec<MultiTurnConversation>> {
181 + let num_conversations = cfg.num_conversations;
182 + let min_turns = cfg.min_turns;
183 + let max_turns = cfg.max_turns;
184 + let prefix_len = cfg.prefix_len;
185 + let output_len = cfg.output_len;
186 + let request_id_prefix = &cfg.request_id_prefix;
187 +
188 + let num_special = tokenizer.num_special_tokens_to_add();
189 + let real_input_len = cfg.input_len.saturating_sub(num_special);
190 + let real_per_turn_len = if cfg.per_turn_input_len > 0 {
191 + cfg.per_turn_input_len.saturating_sub(num_special)
192 + } else {
193 + real_input_len
194 + };
195 +
196 + // Compute segment lengths from turn-0 (real_input_len) so the shared prefix
197 + // bytes stay byte-identical across all turns regardless of per_turn_input_len.
198 + let global_len = (real_input_len as f64 * ps_cfg.global_ratio).floor() as usize;
199 + let conv_len = (real_input_len as f64 * ps_cfg.conversation_ratio).floor() as usize;
200 + let unique_len = real_input_len.saturating_sub(global_len + conv_len);
201 +
202 + // Validate that turns 1+ still have room for a non-empty unique suffix
203 + if real_per_turn_len <= global_len + conv_len {
204 + return Err(BenchError::Config(format!(
205 + "--per-turn-input-len ({real_per_turn_len} after special tokens) is too small: \
206 + global_len={global_len} + conv_len={conv_len} already fills the budget. \
207 + Increase --per-turn-input-len or reduce prefix ratios."
208 + )));
209 + }
210 +
211 + let at_len = allowed_tokens.len();
212 + let shared_prefix_text =
213 + generate_shared_prefix_text(tokenizer, allowed_tokens, prefix_len, cfg.seed)?;
214 +
215 + // Generate global prefix text once
216 + let global_text: Arc<str> = if global_len > 0 {
217 + let offset: usize = rng.random_range(0..at_len);
218 + let seq = make_token_seq(allowed_tokens, offset, global_len);
219 + let (text, _) = gen_prompt_to_target_len(tokenizer, &seq, global_len)?;
220 + Arc::from(text)
221 + } else {
222 + Arc::from("")
223 + };
224 +
225 + // Generate per-conversation prefix texts
226 + let conv_texts: Vec<Arc<str>> = if conv_len > 0 {
227 + let mut texts = Vec::with_capacity(num_conversations);
228 + for conv_idx in 0..num_conversations {
229 + let offset: usize = rng.random_range(0..at_len);
230 + let seq = make_token_seq(allowed_tokens, offset + conv_idx, conv_len);
231 + let (text, _) = gen_prompt_to_target_len(tokenizer, &seq, conv_len)?;
232 + texts.push(Arc::from(text));
233 + }
234 + texts
235 + } else {
236 + vec![Arc::from(""); num_conversations]
237 + };
238 +
239 + // Pre-generate per-conversation turn counts and unique offsets deterministically.
240 + let vocab_size = tokenizer.vocab_size() as usize;
241 + let conv_turn_counts: Vec<usize> = (0..num_conversations)
242 + .map(|_| {
243 + if min_turns == max_turns {
244 + min_turns
245 + } else {
246 + rng.random_range(min_turns..=max_turns)
247 + }
248 + })
249 + .collect();
250 +
251 + let unique_offsets: Vec<Vec<usize>> = conv_turn_counts
252 + .iter()
253 + .map(|&n| (0..n).map(|_| rng.random_range(0..vocab_size)).collect())
254 + .collect();
255 +
256 + // Parallel generation across conversations
257 + unique_offsets
258 + .par_iter()
259 + .enumerate()
260 + .map(|(conv_idx, conv_offsets)| {
261 + let mut turns = Vec::with_capacity(conv_offsets.len());
262 + for (turn_idx, &offset) in conv_offsets.iter().enumerate() {
263 + // Turn 0 uses unique_len derived from real_input_len;
264 + // turns 1+ use per-turn unique_len (prefix bytes stay identical).
265 + let turn_unique_len = if turn_idx == 0 {
266 + unique_len
267 + } else {
268 + real_per_turn_len.saturating_sub(global_len + conv_len)
269 + };
270 +
271 + // Generate unique suffix
272 + let unique_text = if turn_unique_len > 0 {
273 + let seq = make_token_seq(
274 + allowed_tokens,
275 + offset + conv_idx * max_turns + turn_idx,
276 + turn_unique_len,
277 + );
278 + let (text, _) = gen_prompt_to_target_len(tokenizer, &seq, turn_unique_len)?;
279 + text
280 + } else {
281 + String::new()
282 + };
283 +
284 + // Concatenate: optional random prefix + global + conversation + unique.
285 + // Prefix-sharing mode sends each turn independently, so the random
286 + // prefix must be included on every turn to be present in every request.
287 + let combined = format!(
288 + "{}{}{}{}",
289 + &*shared_prefix_text, &*global_text, &*conv_texts[conv_idx], unique_text
290 + );
291 + // Re-encode to get actual token count (BPE boundary effects)
292 + let token_len = tokenizer.encode(&combined, false)?.len();
293 +
294 + turns.push(ConversationTurn {
295 + user_message: Arc::from(combined),
296 + user_message_len: token_len,
297 + expected_output_len: output_len,
298 + });
299 + }
300 +
301 + Ok(MultiTurnConversation {
302 + conversation_id: format!("{request_id_prefix}conv-{conv_idx}"),
303 + turns,
304 + })
305 + })
306 + .collect()
307 +}
308 +
309 +fn generate_shared_prefix_text(
310 + tokenizer: &TokenizerKind,
311 + allowed_tokens: &[u32],
312 + prefix_len: usize,
313 + seed: u64,
314 +) -> Result<Arc<str>> {
315 + if prefix_len == 0 {
316 + return Ok(Arc::from(""));
317 + }
318 +
319 + let mut rng = StdRng::seed_from_u64(seed.wrapping_add(0xDEAD));
320 + let tokens: Vec<u32> = (0..prefix_len)
321 + .map(|_| allowed_tokens[rng.random_range(0..allowed_tokens.len())])
322 + .collect();
323 + let (text, _) = gen_prompt_to_target_len(tokenizer, &tokens, prefix_len)?;
324 + Ok(Arc::from(text))
325 +}
326 +
327 +/// Load multi-turn conversations from a ShareGPT dataset.
328 +///
329 +/// Walks ALL turns in each entry (not just first 2). Filters entries
330 +/// with at least 4 messages (2 user + 2 assistant = 2 real turns).
331 +pub fn load_sharegpt_multi_turn(
332 + tokenizer: &TokenizerKind,
333 + dataset_path: &str,
334 + num_conversations: usize,
335 + output_len_override: Option<usize>,
336 + max_turns: Option<usize>,
337 + seed: u64,
338 + request_id_prefix: &str,
339 +) -> Result<Vec<MultiTurnConversation>> {
340 + let content = std::fs::read_to_string(dataset_path).map_err(|e| {
341 + BenchError::Config(format!(
342 + "Failed to read ShareGPT file '{dataset_path}': {e}"
343 + ))
344 + })?;
345 +
346 + let data: serde_json::Value = serde_json::from_str(&content)
347 + .map_err(|e| BenchError::Config(format!("Invalid JSON in ShareGPT file: {e}")))?;
348 +
349 + let entries = data
350 + .as_array()
351 + .ok_or_else(|| BenchError::Config("ShareGPT file must contain a JSON array".into()))?;
352 +
353 + // Filter entries with at least 4 messages (2 turns: user+assistant+user+assistant)
354 + let mut filtered: Vec<&serde_json::Value> = entries
355 + .iter()
356 + .filter(|entry| {
357 + entry
358 + .get("conversations")
359 + .and_then(|c| c.as_array())
360 + .map(|a| a.len() >= 4)
361 + .unwrap_or(false)
362 + })
363 + .collect();
364 +
365 + if filtered.is_empty() {
366 + return Err(BenchError::Config(
367 + "No valid multi-turn entries in ShareGPT file (need at least 4 messages per entry)"
368 + .into(),
369 + ));
370 + }
371 +
372 + // Shuffle
373 + let mut rng = StdRng::seed_from_u64(seed);
374 + filtered.shuffle(&mut rng);
375 +
376 + let mut conversations = Vec::new();
377 +
378 + for entry in &filtered {
379 + if conversations.len() >= num_conversations {
380 + break;
381 + }
382 +
383 + let msgs = entry["conversations"].as_array().unwrap();
384 + let mut turns = Vec::new();
385 +
386 + // Walk alternating human/gpt pairs, stopping early once max_turns reached
387 + // to avoid tokenizing turns that would be discarded by truncate().
388 + let mut i = 0;
389 + while i + 1 < msgs.len() {
390 + if let Some(m) = max_turns
391 + && turns.len() >= m
392 + {
393 + break;
394 + }
395 + let from = msgs[i].get("from").and_then(|f| f.as_str()).unwrap_or("");
396 + let user_text = msgs[i].get("value").and_then(|v| v.as_str()).unwrap_or("");
397 + let assistant_text = msgs[i + 1].get("value").and_then(|v| v.as_str()).unwrap_or("");
398 +
399 + // Expect human then gpt
400 + if from != "human" || user_text.is_empty() {
401 + i += 1;
402 + continue;
403 + }
404 +
405 + let user_ids = tokenizer.encode(user_text, false)?;
406 + let user_len = user_ids.len();
407 +
408 + let expected_output_len = if let Some(override_len) = output_len_override {
409 + override_len
410 + } else {
411 + let assistant_ids = tokenizer.encode(assistant_text, false)?;
412 + assistant_ids.len().max(1)
413 + };
414 +
415 + turns.push(ConversationTurn {
416 + user_message: Arc::from(user_text),
417 + user_message_len: user_len,
418 + expected_output_len,
419 + });
420 +
421 + i += 2;
422 + }
423 +
424 + if turns.len() >= 2 {
425 + let conv_idx = conversations.len();
426 + conversations.push(MultiTurnConversation {
427 + conversation_id: format!("{request_id_prefix}conv-{conv_idx}"),
428 + turns,
429 + });
430 + }
431 + }
432 +
433 + if conversations.is_empty() {
434 + return Err(BenchError::Config(
435 + "No valid multi-turn conversations after filtering ShareGPT dataset.".into(),
436 + ));
437 + }
438 +
439 + // Oversample if needed
440 + if conversations.len() < num_conversations {
441 + let original_len = conversations.len();
442 + let needed = num_conversations - original_len;
443 + for i in 0..needed {
444 + let mut conv = conversations[rng.random_range(0..original_len)].clone();
445 + conv.conversation_id = format!("{request_id_prefix}conv-{}", original_len + i);
446 + conversations.push(conv);
447 + }
448 + println!(
449 + "Oversampled multi-turn conversations from {original_len} to {} total.",
450 + conversations.len()
451 + );
452 + }
453 +
454 + Ok(conversations)
455 +}
456 +
457 +/// Ensure decoded-then-encoded prompt length matches the target.
458 +fn gen_prompt_to_target_len(
459 + tokenizer: &TokenizerKind,
460 + token_sequence: &[u32],
461 + target_len: usize,
462 +) -> Result<(String, Vec<u32>)> {
463 + let max_retry = 20;
464 + let mut tokens = token_sequence.to_vec();
465 +
466 + for retry in 0..=max_retry {
467 + let prompt = tokenizer.decode(&tokens, true)?;
468 + tokens = tokenizer.encode(&prompt, false)?;
469 +
470 + if retry >= max_retry {
471 + // BPE tokenizers can oscillate by ±1 on certain boundaries.
472 + // For benchmark random content, accept close-enough and truncate/pad.
473 + if tokens.len() > target_len {
474 + tokens.truncate(target_len);
475 + }
476 + // If still short by 1-2 tokens, accept as-is — negligible for benchmarks.
477 + // Re-decode after truncation to ensure prompt string matches token vector.
478 + let prompt = tokenizer.decode(&tokens, true)?;
479 + return Ok((prompt, tokens));
480 + }
481 +
482 + if tokens.len() == target_len {
483 + return Ok((prompt, tokens));
484 + } else if tokens.len() < target_len {
485 + let allowed = tokenizer.get_allowed_tokens();
486 + let needed = target_len - tokens.len();
487 + if allowed.is_empty() {
488 + let vocab_size = tokenizer.vocab_size() as usize;
489 + for j in 0..needed {
490 + tokens.push(((tokens.len() + j) % vocab_size) as u32);
491 + }
492 + } else {
493 + for j in 0..needed {
494 + tokens.push(allowed[(tokens.len() + j) % allowed.len()]);
495 + }
496 + }
497 + } else {
498 + tokens.truncate(target_len);
499 + }
500 + }
501 +
502 + unreachable!()
503 +}
504 +
505 +#[cfg(test)]
506 +mod tests {
507 + use super::*;
508 +
509 + fn common_prefix_bytes(strings: &[&str]) -> usize {
510 + if strings.is_empty() {
511 + return 0;
512 + }
513 + let first = strings[0].as_bytes();
514 + let mut len = first.len();
515 + for s in &strings[1..] {
516 + let b = s.as_bytes();
517 + len = len.min(b.len());
518 + for i in 0..len {
519 + if first[i] != b[i] {
520 + len = i;
521 + break;
522 + }
523 + }
524 + }
525 + len
526 + }
527 +
528 + #[test]
529 + #[ignore]
530 + fn test_prefix_sharing_structure() {
531 + let tok = crate::tokenizer::load_tokenizer("nvidia/Kimi-K2.5-NVFP4", false, None).unwrap();
532 +
533 + let cfg = MultiTurnRandomConfig {
534 + num_conversations: 5,
535 + min_turns: 3,
536 + max_turns: 3,
537 + prefix_len: 0,
538 + input_len: 1000,
539 + per_turn_input_len: 0,
540 + output_len: 100,
541 + seed: 42,
542 + request_id_prefix: "test-".to_string(),
543 + prefix_sharing_config: Some(PrefixSharingConfig {
544 + global_ratio: 0.1,
545 + conversation_ratio: 0.8,
546 + }),
547 + };
548 +
549 + let conversations = generate_multi_turn_random(&tok, &cfg).unwrap();
550 + assert_eq!(conversations.len(), 5);
551 +
552 + let messages: Vec<Vec<&str>> = conversations
553 + .iter()
554 + .map(|c| c.turns.iter().map(|t| &*t.user_message).collect())
555 + .collect();
556 +
557 + // 1. Global prefix: all messages share a common prefix
558 + let all_msgs: Vec<&str> = messages.iter().flat_map(|v| v.iter().copied()).collect();
559 + let global_prefix = common_prefix_bytes(&all_msgs);
560 + println!("Global prefix bytes: {global_prefix}");
561 + assert!(global_prefix > 0, "Global prefix must be non-empty");
562 +
563 + // 2. Conversation prefix: turns within same conversation share more
564 + for (i, conv_msgs) in messages.iter().enumerate() {
565 + let conv_prefix = common_prefix_bytes(conv_msgs);
566 + println!("Conv {i} prefix bytes: {conv_prefix} (global: {global_prefix})");
567 + assert!(
568 + conv_prefix > global_prefix,
569 + "Conv prefix ({conv_prefix}) must exceed global prefix ({global_prefix})"
570 + );
571 + }
572 +
573 + // 3. Different conversations diverge after global prefix
574 + let cross = common_prefix_bytes(&[messages[0][0], messages[1][0]]);
575 + let within = common_prefix_bytes(&messages[0]);
576 + println!("Cross-conv prefix: {cross}, within-conv prefix: {within}");
577 + assert!(
578 + cross < within,
579 + "Cross-conv ({cross}) must be < within-conv ({within})"
580 + );
581 +
582 + // 4. Turns within same conversation are not identical (unique suffix)
583 + for (i, conv_msgs) in messages.iter().enumerate() {
584 + for a in 0..conv_msgs.len() {
585 + for b in (a + 1)..conv_msgs.len() {
586 + assert_ne!(
587 + conv_msgs[a], conv_msgs[b],
588 + "Conv {i} turn {a} and {b} must differ"
589 + );
590 + }
591 + }
592 + }
593 +
594 + // 5. Token lengths approximately match target
595 + for (i, conv) in conversations.iter().enumerate() {
596 + for (j, turn) in conv.turns.iter().enumerate() {
597 + let diff = (turn.user_message_len as i64 - 1000).abs();
598 + println!(
599 + "Conv {i} turn {j}: {} tokens (diff {diff})",
600 + turn.user_message_len
601 + );
602 + assert!(
603 + diff <= 10,
604 + "Token len {} too far from 1000",
605 + turn.user_message_len
606 + );
607 + }
608 + }
609 +
610 + println!("All prefix sharing checks passed!");
611 + }
612 +
613 + #[test]
614 + #[ignore]
615 + fn test_per_turn_input_len_default_mode() {
616 + let tok = crate::tokenizer::load_tokenizer("nvidia/Kimi-K2.5-NVFP4", false, None).unwrap();
617 +
618 + let cfg = MultiTurnRandomConfig {
619 + num_conversations: 4,
620 + min_turns: 3,
621 + max_turns: 3,
622 + prefix_len: 0,
623 + input_len: 512,
624 + per_turn_input_len: 128,
625 + output_len: 64,
626 + seed: 1,
627 + request_id_prefix: "test-".to_string(),
628 + prefix_sharing_config: None,
629 + };
630 +
631 + let conversations = generate_multi_turn_random(&tok, &cfg).unwrap();
632 + assert_eq!(conversations.len(), 4);
633 +
634 + for (i, conv) in conversations.iter().enumerate() {
635 + assert_eq!(conv.turns.len(), 3);
636 + for (j, turn) in conv.turns.iter().enumerate() {
637 + let expected = if j == 0 { 512usize } else { 128usize };
638 + let diff = (turn.user_message_len as i64 - expected as i64).abs();
639 + println!(
640 + "Conv {i} turn {j}: {} tokens (expected ~{expected}, diff {diff})",
641 + turn.user_message_len
642 + );
643 + assert!(
644 + diff <= 5,
645 + "Conv {i} turn {j}: token len {} too far from {expected}",
646 + turn.user_message_len
647 + );
648 + }
649 + }
650 + println!("per_turn_input_len default-mode checks passed!");
651 + }
652 +
653 + #[test]
654 + #[ignore]
655 + fn test_variable_turns_range() {
656 + let tok = crate::tokenizer::load_tokenizer("nvidia/Kimi-K2.5-NVFP4", false, None).unwrap();
657 +
658 + let cfg = MultiTurnRandomConfig {
659 + num_conversations: 50,
660 + min_turns: 2,
661 + max_turns: 5,
662 + prefix_len: 0,
663 + input_len: 256,
664 + per_turn_input_len: 0,
665 + output_len: 32,
666 + seed: 7,
667 + request_id_prefix: "test-".to_string(),
668 + prefix_sharing_config: None,
669 + };
670 +
671 + let conversations = generate_multi_turn_random(&tok, &cfg).unwrap();
672 + assert_eq!(conversations.len(), 50);
673 +
674 + let mut distinct_counts = std::collections::HashSet::new();
675 + for conv in &conversations {
676 + let n = conv.turns.len();
677 + assert!((2..=5).contains(&n), "turn count {n} out of [2,5]");
678 + distinct_counts.insert(n);
679 + }
680 + assert!(
681 + distinct_counts.len() >= 2,
682 + "expected at least 2 distinct turn counts, got {distinct_counts:?}"
683 + );
684 + println!("variable_turns_range checks passed! counts: {distinct_counts:?}");
685 + }
686 +
687 + #[test]
688 + #[ignore]
689 + fn test_variable_turns_fixed() {
690 + let tok = crate::tokenizer::load_tokenizer("nvidia/Kimi-K2.5-NVFP4", false, None).unwrap();
691 +
692 + let cfg = MultiTurnRandomConfig {
693 + num_conversations: 10,
694 + min_turns: 4,
695 + max_turns: 4,
696 + prefix_len: 0,
697 + input_len: 256,
698 + per_turn_input_len: 0,
699 + output_len: 32,
700 + seed: 42,
701 + request_id_prefix: "test-".to_string(),
702 + prefix_sharing_config: None,
703 + };
704 +
705 + let conversations = generate_multi_turn_random(&tok, &cfg).unwrap();
706 + for conv in &conversations {
707 + assert_eq!(conv.turns.len(), 4, "expected exactly 4 turns");
708 + }
709 + println!("variable_turns_fixed checks passed!");
710 + }
711 +
712 + #[test]
713 + #[ignore]
714 + fn test_per_turn_input_len_prefix_sharing() {
715 + let tok = crate::tokenizer::load_tokenizer("nvidia/Kimi-K2.5-NVFP4", false, None).unwrap();
716 +
717 + // Turn 0 input_len=1000, turns 1+ per_turn_input_len=600
718 + // global_len ≈ 100 (10%), conv_len ≈ 800 (80%), unique ≈ 100
719 + // per-turn unique ≈ 600 - 900 = negative → would error; use smaller ratios
720 + // global=0.05 (50), conv=0.5 (500), unique_t0=450, unique_t1=600-550=50
721 + let cfg = MultiTurnRandomConfig {
722 + num_conversations: 4,
723 + min_turns: 3,
724 + max_turns: 3,
725 + prefix_len: 0,
726 + input_len: 1000,
727 + per_turn_input_len: 600,
728 + output_len: 64,
729 + seed: 3,
730 + request_id_prefix: "test-".to_string(),
731 + prefix_sharing_config: Some(PrefixSharingConfig {
732 + global_ratio: 0.05,
733 + conversation_ratio: 0.50,
734 + }),
735 + };
736 +
737 + let conversations = generate_multi_turn_random(&tok, &cfg).unwrap();
738 + assert_eq!(conversations.len(), 4);
739 +
740 + let messages: Vec<Vec<&str>> = conversations
741 + .iter()
742 + .map(|c| c.turns.iter().map(|t| &*t.user_message).collect())
743 + .collect();
744 +
745 + // Global prefix bytes shared across all turns of all conversations
746 + let all_msgs: Vec<&str> = messages.iter().flat_map(|v| v.iter().copied()).collect();
747 + let global_prefix = common_prefix_bytes(&all_msgs);
748 + assert!(global_prefix > 0, "Global prefix must be non-empty");
749 +
750 + // Within each conversation, prefix grows (conv prefix longer than global)
751 + for (i, conv_msgs) in messages.iter().enumerate() {
752 + let conv_prefix = common_prefix_bytes(conv_msgs);
753 + assert!(
754 + conv_prefix > global_prefix,
755 + "Conv {i}: conv_prefix ({conv_prefix}) must exceed global ({global_prefix})"
756 + );
757 + }
758 +
759 + // Turn 0 length ≈ 1000, turns 1+ ≈ 600
760 + for (i, conv) in conversations.iter().enumerate() {
761 + for (j, turn) in conv.turns.iter().enumerate() {
762 + let expected = if j == 0 { 1000usize } else { 600usize };
763 + let diff = (turn.user_message_len as i64 - expected as i64).abs();
764 + println!(
765 + "Conv {i} turn {j}: {} tokens (expected ~{expected})",
766 + turn.user_message_len
767 + );
768 + assert!(
769 + diff <= 10,
770 + "Conv {i} turn {j}: token len {} too far from {expected}",
771 + turn.user_message_len
772 + );
773 + }
774 + }
775 + println!("per_turn_input_len prefix-sharing checks passed!");
776 + }
777 +}
rust/src/bench/src/datasets/prefix_repetition.rs new
+148
@@ -0,0 +1,148 @@
1 +// SPDX-License-Identifier: Apache-2.0
2 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +//! Prefix repetition dataset: N distinct shared prefixes, each reused by
5 +//! `num_prompts / num_prefixes` requests with a fresh random suffix.
6 +//! The standard prefix-cache stress workload; mirrors Python's
7 +//! `PrefixRepetitionRandomDataset`.
8 +
9 +use std::sync::Arc;
10 +
11 +use rand::rngs::StdRng;
12 +use rand::seq::SliceRandom;
13 +use rand::{Rng, SeedableRng};
14 +use rayon::prelude::*;
15 +
16 +use super::SampleRequest;
17 +use super::random::gen_prompt_decode_to_target_len;
18 +use crate::error::{BenchError, Result};
19 +use crate::tokenizer::TokenizerKind;
20 +
21 +/// Generate the prefix repetition dataset.
22 +///
23 +/// Like Python, `num_requests % num_prefixes` remainder requests are dropped:
24 +/// the total is `(num_requests / num_prefixes) * num_prefixes`.
25 +pub fn generate_prefix_repetition_dataset(
26 + tokenizer: &TokenizerKind,
27 + num_requests: usize,
28 + prefix_len: usize,
29 + suffix_len: usize,
30 + num_prefixes: usize,
31 + output_len: usize,
32 + seed: u64,
33 + request_id_prefix: &str,
34 + disable_shuffle: bool,
35 +) -> Result<Vec<SampleRequest>> {
36 + let prompts_per_prefix = num_requests / num_prefixes;
37 + if prompts_per_prefix == 0 {
38 + return Err(BenchError::Config(format!(
39 + "num_prompts ({num_requests}) must be >= num_prefixes ({num_prefixes})"
40 + )));
41 + }
42 + let total = prompts_per_prefix * num_prefixes;
43 + if total != num_requests {
44 + println!(
45 + "prefix_repetition: generating {total} requests \
46 + ({num_prefixes} prefixes x {prompts_per_prefix} prompts each; \
47 + {} dropped to divide evenly)",
48 + num_requests - total
49 + );
50 + }
51 +
52 + let allowed_tokens = tokenizer.get_allowed_tokens();
53 + if allowed_tokens.is_empty() {
54 + return Err(BenchError::Tokenizer("No allowed tokens found".into()));
55 + }
56 + let allowed_ref = &allowed_tokens;
57 +
58 + // Exact-length random token block: decode -> re-encode -> converge to target.
59 + let gen_block = |target_len: usize, item_seed: u64| -> Result<Vec<u32>> {
60 + let mut rng = StdRng::seed_from_u64(item_seed);
61 + let tokens: Vec<u32> = (0..target_len)
62 + .map(|_| allowed_ref[rng.random_range(0..allowed_ref.len())])
63 + .collect();
64 + let (_, adjusted) =
65 + gen_prompt_decode_to_target_len(tokenizer, &tokens, target_len, false, allowed_ref)?;
66 + Ok(adjusted)
67 + };
68 +
69 + // Generate the shared prefixes (one per group), then suffixes in parallel.
70 + let prefixes: Vec<Vec<u32>> = (0..num_prefixes)
71 + .map(|p| gen_block(prefix_len, seed.wrapping_add(0xF1F0).wrapping_add(p as u64)))
72 + .collect::<Result<Vec<_>>>()?;
73 +
74 + let rid_prefix = request_id_prefix.to_string();
75 + let mut requests: Vec<SampleRequest> = (0..total)
76 + .into_par_iter()
77 + .map(|i| {
78 + let prefix_tokens = &prefixes[i / prompts_per_prefix];
79 + let suffix_tokens = gen_block(suffix_len, seed.wrapping_add(0xBEEF + i as u64))?;
80 +
81 + let mut combined = Vec::with_capacity(prefix_tokens.len() + suffix_tokens.len());
82 + combined.extend_from_slice(prefix_tokens);
83 + combined.extend_from_slice(&suffix_tokens);
84 + let prompt = tokenizer.decode(&combined, true)?;
85 +
86 + Ok(SampleRequest {
87 + prompt: Arc::from(prompt),
88 + prompt_len: combined.len(),
89 + expected_output_len: output_len,
90 + request_id: Some(format!("{rid_prefix}{i}")),
91 + ..Default::default()
92 + })
93 + })
94 + .collect::<Result<Vec<_>>>()?;
95 +
96 + // Interleave prefixes (Python shuffles too) so one prefix group isn't sent
97 + // as a contiguous burst.
98 + if !disable_shuffle {
99 + let mut rng = StdRng::seed_from_u64(seed);
100 + requests.shuffle(&mut rng);
101 + }
102 +
103 + Ok(requests)
104 +}
105 +
106 +#[cfg(test)]
107 +mod tests {
108 + use super::*;
109 +
110 + /// gpt2 via built-in tiktoken encoding — loads without network access.
111 + fn test_tokenizer() -> TokenizerKind {
112 + crate::tokenizer::load_tokenizer("gpt2", false, None)
113 + .expect("gpt2 built-in tiktoken should always load without network")
114 + }
115 +
116 + #[test]
117 + fn test_prefix_repetition_structure() {
118 + let tok = test_tokenizer();
119 + // 7 requests / 3 prefixes -> 2 per prefix, 6 total (remainder dropped like Python)
120 + let reqs = generate_prefix_repetition_dataset(&tok, 7, 32, 16, 3, 64, 0, "t-", true)
121 + .expect("generation should succeed");
122 + assert_eq!(reqs.len(), 6);
123 + assert!(reqs.iter().all(|r| r.expected_output_len == 64));
124 + // Exact-length blocks: prompt_len == prefix + suffix
125 + assert!(
126 + reqs.iter().all(|r| r.prompt_len == 32 + 16),
127 + "lens: {:?}",
128 + reqs.iter().map(|r| r.prompt_len).collect::<Vec<_>>()
129 + );
130 + // Consecutive pairs (shuffle disabled) share a common prefix; requests
131 + // from different groups don't.
132 + let common = |a: &str, b: &str| -> usize {
133 + a.bytes().zip(b.bytes()).take_while(|(x, y)| x == y).count()
134 + };
135 + let same_group = common(&reqs[0].prompt, &reqs[1].prompt);
136 + let diff_group = common(&reqs[0].prompt, &reqs[2].prompt);
137 + assert!(
138 + same_group > diff_group,
139 + "same-group shared prefix ({same_group}) should exceed cross-group ({diff_group})"
140 + );
141 + }
142 +
143 + #[test]
144 + fn test_prefix_repetition_too_few_requests_errors() {
145 + let tok = test_tokenizer();
146 + assert!(generate_prefix_repetition_dataset(&tok, 2, 32, 16, 3, 64, 0, "t-", true).is_err());
147 + }
148 +}
rust/src/bench/src/datasets/random.rs new
+491
@@ -0,0 +1,491 @@
1 +// SPDX-License-Identifier: Apache-2.0
2 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +use std::sync::Arc;
5 +
6 +use rayon::prelude::*;
7 +
8 +use super::SampleRequest;
9 +use crate::config::RangeRatio;
10 +use crate::error::{BenchError, Result};
11 +use crate::tokenizer::TokenizerKind;
12 +
13 +/// Generate random dataset with rayon parallelism.
14 +///
15 +/// This is the key performance win — Python does sequential tokenizer calls
16 +/// while Rust parallelizes across CPU cores with native tokenizer speed.
17 +///
18 +/// Mirrors Python's RandomDataset.sample() from datasets.py:470-560.
19 +pub fn generate_random_dataset(
20 + tokenizer: &TokenizerKind,
21 + num_requests: usize,
22 + input_len: usize,
23 + output_len: usize,
24 + prefix_len: usize,
25 + range_ratio: RangeRatio,
26 + cache_hit_fraction: f64,
27 + cache_ratio: f64,
28 + seed: u64,
29 + request_id_prefix: &str,
30 + use_token_ids: bool,
31 + batch_size: usize,
32 +) -> Result<Vec<SampleRequest>> {
33 + let vocab_size = tokenizer.vocab_size();
34 + let allowed_tokens = tokenizer.get_allowed_tokens();
35 + if allowed_tokens.is_empty() {
36 + return Err(BenchError::Tokenizer("No allowed tokens found".into()));
37 + }
38 +
39 + if batch_size > 1 && use_token_ids {
40 + return Err(BenchError::Config(
41 + "--random-batch-size > 1 is not supported with --prompt-token-ids".into(),
42 + ));
43 + }
44 +
45 + let num_special = tokenizer.num_special_tokens_to_add();
46 + let real_input_len = input_len.saturating_sub(num_special);
47 +
48 + // Python semantics: sample uniformly from [len*(1-r), len*(1+r)].
49 + let (input_low, input_high) = range_ratio.input_bounds(real_input_len);
50 + let (output_low, output_high) = range_ratio.output_bounds(output_len);
51 + if !range_ratio.is_fixed() {
52 + println!(
53 + "Sampling input_len from [{input_low}, {input_high}] and \
54 + output_len from [{output_low}, {output_high}]"
55 + );
56 + }
57 +
58 + // Bimodal prefix-cache mode: a fraction of prompts (warm) reuse a shared cached
59 + // prefix covering `cache_ratio` of their length; the rest (cold) are fully unique.
60 + // Models e.g. "80% of prompts have 95% of input cached" with
61 + // --random-cache-hit-fraction 0.8 --random-cache-ratio 0.95. In this mode
62 + // --random-input-len is the TOTAL prompt length L (the cached prefix is part of L),
63 + // and --random-prefix-len is ignored.
64 + let bimodal = cache_hit_fraction > 0.0 && cache_ratio > 0.0;
65 + if bimodal {
66 + if !use_token_ids {
67 + return Err(BenchError::Config(
68 + "bimodal prefix-cache (--random-cache-hit-fraction) requires --prompt-token-ids \
69 + so warm prompts send identical token IDs and actually hit the prefix cache"
70 + .into(),
71 + ));
72 + }
73 + if cache_hit_fraction > 1.0 || cache_ratio > 1.0 {
74 + return Err(BenchError::Config(
75 + "--random-cache-hit-fraction and --random-cache-ratio must be in [0, 1]".into(),
76 + ));
77 + }
78 + }
79 +
80 + // Length of the shared cached base prefix.
81 + let base_len = if bimodal {
82 + ((input_high as f64) * cache_ratio).ceil() as usize
83 + } else {
84 + prefix_len
85 + };
86 +
87 + // Validate (non-bimodal keeps the original check)
88 + if !bimodal {
89 + let min_total = prefix_len + input_low;
90 + if min_total < 1 {
91 + return Err(BenchError::Config(format!(
92 + "--random-input-len too small: with {num_special} special tokens and \
93 + range_ratio={:?}, minimum total input is {min_total}",
94 + range_ratio
95 + )));
96 + }
97 + }
98 +
99 + // Generate the shared base prefix once (sequential, only happens once).
100 + let prefix_token_ids = if base_len > 0 {
101 + generate_prefix(tokenizer, &allowed_tokens, base_len, seed)?
102 + } else {
103 + Vec::new()
104 + };
105 +
106 + // Pre-generate per-request sampling params using deterministic RNG
107 + use rand::rngs::StdRng;
108 + use rand::{Rng, SeedableRng};
109 + let mut rng = StdRng::seed_from_u64(seed);
110 +
111 + struct RequestParams {
112 + cached_len: usize, // tokens taken from the shared base (cache-hittable)
113 + suffix_len: usize, // unique tokens appended after the cached prefix
114 + output_len: usize,
115 + offset: usize,
116 + }
117 +
118 + let params: Vec<RequestParams> = (0..num_requests)
119 + .map(|_| {
120 + let ol = if output_low == output_high {
121 + output_low
122 + } else {
123 + rng.random_range(output_low..=output_high)
124 + };
125 + let off = rng.random_range(0..vocab_size as usize);
126 + if bimodal {
127 + // Total length L from the input distribution; prefix is part of L.
128 + let l = if input_low == input_high {
129 + input_low
130 + } else {
131 + rng.random_range(input_low..=input_high)
132 + };
133 + let warm = rng.random::<f64>() < cache_hit_fraction;
134 + let cached = if warm {
135 + (((l as f64) * cache_ratio).round() as usize).min(base_len).min(l)
136 + } else {
137 + 0
138 + };
139 + RequestParams {
140 + cached_len: cached,
141 + suffix_len: l - cached,
142 + output_len: ol,
143 + offset: off,
144 + }
145 + } else {
146 + // Original behavior: full shared prefix + variable unique input.
147 + let il = if input_low == input_high {
148 + input_low
149 + } else {
150 + rng.random_range(input_low..=input_high)
151 + };
152 + RequestParams {
153 + cached_len: prefix_len,
154 + suffix_len: il,
155 + output_len: ol,
156 + offset: off,
157 + }
158 + }
159 + })
160 + .collect();
161 +
162 + // Phase 1: Generate all token sequences (parallel, fast — just array ops)
163 + let prefix_ref = &prefix_token_ids;
164 + let allowed_ref = &allowed_tokens;
165 + let rid_prefix = request_id_prefix.to_string();
166 +
167 + let token_sequences: Vec<Vec<u32>> = params
168 + .par_iter()
169 + .enumerate()
170 + .map(|(i, p)| {
171 + let at_len = allowed_ref.len();
172 + let mut seq = Vec::with_capacity(p.cached_len + p.suffix_len);
173 + seq.extend_from_slice(&prefix_ref[..p.cached_len]);
174 + for j in 0..p.suffix_len {
175 + seq.push(allowed_ref[(p.offset + i + j) % at_len]);
176 + }
177 + seq
178 + })
179 + .collect();
180 +
181 + let target_lens: Vec<usize> = params.iter().map(|p| p.cached_len + p.suffix_len).collect();
182 +
183 + if use_token_ids {
184 + // Fast path: store token IDs directly. The completions backend sends
185 + // them as `"prompt": [id1, id2, ...]`, bypassing both client-side decode
186 + // and server-side tokenization. Token counts are exact by construction.
187 + let result: Vec<SampleRequest> = token_sequences
188 + .into_par_iter()
189 + .enumerate()
190 + .map(|(i, tokens)| SampleRequest {
191 + prompt: Arc::from(""),
192 + prompt_len: target_lens[i],
193 + expected_output_len: params[i].output_len,
194 + request_id: Some(format!("{rid_prefix}{i}")),
195 + prompt_token_ids: Some(Arc::from(tokens)),
196 + ..Default::default()
197 + })
198 + .collect();
199 + Ok(result)
200 + } else {
201 + // Default path: decode tokens to text, re-encode,
202 + // truncate to target length, decode again. Sends text prompts for maximum
203 + let result: Vec<SampleRequest> = token_sequences
204 + .into_par_iter()
205 + .enumerate()
206 + .map(|(i, tokens)| {
207 + let target = target_lens[i];
208 + // decode → encode → truncate → decode
209 + let prompt_text = tokenizer.decode(&tokens, true)?;
210 + let mut re_encoded = tokenizer.encode(&prompt_text, false)?;
211 + re_encoded.truncate(target);
212 + let prompt = tokenizer.decode(&re_encoded, true)?;
213 + let prompt_len = re_encoded.len();
214 + Ok(SampleRequest {
215 + prompt: Arc::from(prompt),
216 + prompt_len,
217 + expected_output_len: params[i].output_len,
218 + request_id: Some(format!("{rid_prefix}{i}")),
219 + ..Default::default()
220 + })
221 + })
222 + .collect::<Result<Vec<_>>>()?;
223 + Ok(super::batch_requests(result, batch_size, request_id_prefix))
224 + }
225 +}
226 +
227 +fn generate_prefix(
228 + tokenizer: &TokenizerKind,
229 + allowed_tokens: &[u32],
230 + prefix_len: usize,
231 + seed: u64,
232 +) -> Result<Vec<u32>> {
233 + use rand::rngs::StdRng;
234 + use rand::{Rng, SeedableRng};
235 +
236 + let mut rng = StdRng::seed_from_u64(seed.wrapping_add(0xDEAD));
237 + let tokens: Vec<u32> = (0..prefix_len)
238 + .map(|_| allowed_tokens[rng.random_range(0..allowed_tokens.len())])
239 + .collect();
240 +
241 + let (_, adjusted) =
242 + gen_prompt_decode_to_target_len(tokenizer, &tokens, prefix_len, false, allowed_tokens)?;
243 + Ok(adjusted)
244 +}
245 +
246 +/// Ensure decoded-then-encoded prompt length matches the target.
247 +///
248 +/// Mirrors Python's `gen_prompt_decode_to_target_len` from datasets.py:381-435.
249 +pub(crate) fn gen_prompt_decode_to_target_len(
250 + tokenizer: &TokenizerKind,
251 + token_sequence: &[u32],
252 + target_len: usize,
253 + add_special_tokens: bool,
254 + allowed_tokens: &[u32],
255 +) -> Result<(String, Vec<u32>)> {
256 + let max_retry = 20;
257 + let mut tokens = token_sequence.to_vec();
258 +
259 + for retry in 0..=max_retry {
260 + let prompt = tokenizer.decode(&tokens, true)?;
261 + tokens = tokenizer.encode(&prompt, add_special_tokens)?;
262 +
263 + if retry >= max_retry {
264 + if tokens.len() != target_len {
265 + return Err(BenchError::Tokenizer(format!(
266 + "Token length mismatch after {max_retry} retries: \
267 + target={target_len}, actual={}. \
268 + encode/decode roundtrip cannot converge.",
269 + tokens.len()
270 + )));
271 + }
272 + return Ok((prompt, tokens));
273 + }
274 +
275 + if tokens.len() == target_len {
276 + return Ok((prompt, tokens));
277 + } else if tokens.len() < target_len {
278 + // Pad with tokens from the allowed set (UTF-8-safe for tiktoken)
279 + let needed = target_len - tokens.len();
280 + if allowed_tokens.is_empty() {
281 + let vocab_size = tokenizer.vocab_size() as usize;
282 + for j in 0..needed {
283 + tokens.push(((tokens.len() + j) % vocab_size) as u32);
284 + }
285 + } else {
286 + for j in 0..needed {
287 + tokens.push(allowed_tokens[(tokens.len() + j) % allowed_tokens.len()]);
288 + }
289 + }
290 + } else {
291 + // Truncate
292 + tokens.truncate(target_len);
293 + }
294 + }
295 +
296 + unreachable!()
297 +}
298 +
299 +#[cfg(test)]
300 +mod tests {
301 + use super::*;
302 + use crate::tokenizer;
303 +
304 + // Integration test requires a tokenizer, so only run with --ignored
305 + #[test]
306 + #[ignore]
307 + fn test_generate_random_dataset_token_ids() {
308 + let tokenizer = tokenizer::load_tokenizer("gpt2", false, None).unwrap();
309 + let requests = generate_random_dataset(
310 + &tokenizer,
311 + 10, // num_requests
312 + 128, // input_len
313 + 32, // output_len
314 + 0, // prefix_len
315 + RangeRatio {
316 + input: 0.0,
317 + output: 0.0,
318 + }, // range_ratio (0.0 = fixed length)
319 + 0.0, // cache_hit_fraction (0 = bimodal off)
320 + 0.0, // cache_ratio
321 + 42, // seed
322 + "test-",
323 + true, // use_token_ids
324 + 1, // batch_size
325 + )
326 + .unwrap();
327 +
328 + assert_eq!(requests.len(), 10);
329 + for req in &requests {
330 + assert!(req.prompt_token_ids.is_some());
331 + assert_eq!(req.prompt_token_ids.as_ref().unwrap().len(), req.prompt_len);
332 + assert!(req.prompt_len > 0);
333 + assert_eq!(req.expected_output_len, 32);
334 + }
335 + }
336 +
337 + #[test]
338 + #[ignore]
339 + fn test_generate_random_dataset_text() {
340 + let tokenizer = tokenizer::load_tokenizer("gpt2", false, None).unwrap();
341 + let requests = generate_random_dataset(
342 + &tokenizer,
343 + 10, // num_requests
344 + 128, // input_len
345 + 32, // output_len
346 + 0, // prefix_len
347 + RangeRatio {
348 + input: 0.0,
349 + output: 0.0,
350 + }, // range_ratio (0.0 = fixed length)
351 + 0.0, // cache_hit_fraction (0 = bimodal off)
352 + 0.0, // cache_ratio
353 + 42, // seed
354 + "test-",
355 + false, // use_token_ids = false → text prompts
356 + 1, // batch_size
357 + )
358 + .unwrap();
359 +
360 + assert_eq!(requests.len(), 10);
361 + for req in &requests {
362 + assert!(req.prompt_token_ids.is_none());
363 + assert!(!req.prompt.is_empty());
364 + assert!(req.prompt_len > 0);
365 + assert!(req.prompt_len <= 128);
366 + assert_eq!(req.expected_output_len, 32);
367 + }
368 + }
369 +
370 + /// Test that generated prompts have EXACT target token length (token ID mode).
371 + #[test]
372 + #[ignore]
373 + fn test_token_length_exact_local() {
374 + let tokenizer = tokenizer::load_tokenizer("gpt2", false, None).unwrap();
375 + let target_len = 512;
376 + let requests = generate_random_dataset(
377 + &tokenizer,
378 + 50,
379 + target_len,
380 + 64,
381 + 0,
382 + RangeRatio {
383 + input: 0.0,
384 + output: 0.0,
385 + },
386 + 0.0,
387 + 0.0,
388 + 123,
389 + "len-test-",
390 + true,
391 + 1,
392 + )
393 + .unwrap();
394 +
395 + for (i, req) in requests.iter().enumerate() {
396 + let token_ids = req.prompt_token_ids.as_ref().expect("should have token IDs");
397 + assert_eq!(
398 + token_ids.len(),
399 + target_len,
400 + "Request {i}: expected {target_len} token IDs, got {}",
401 + token_ids.len()
402 + );
403 + assert_eq!(req.prompt_len, target_len);
404 + }
405 + }
406 +
407 + /// Test that tiktoken tokenizer produces exact target token lengths (token ID mode).
408 + #[test]
409 + #[ignore]
410 + fn test_token_length_exact_tiktoken() {
411 + // Use Qwen2.5 which has a tiktoken-format tokenizer
412 + let tokenizer = tokenizer::load_tokenizer("Qwen/Qwen2.5-0.5B", false, None);
413 + let tokenizer = match tokenizer {
414 + Ok(t) => t,
415 + Err(e) => {
416 + eprintln!("Skipping tiktoken test (tokenizer unavailable): {e}");
417 + return;
418 + }
419 + };
420 +
421 + // Verify it's actually a tiktoken tokenizer or local — either way test convergence
422 + let target_len = 256;
423 + let requests = generate_random_dataset(
424 + &tokenizer,
425 + 20,
426 + target_len,
427 + 32,
428 + 0,
429 + RangeRatio {
430 + input: 0.0,
431 + output: 0.0,
432 + },
433 + 0.0,
434 + 0.0,
435 + 42,
436 + "tiktoken-test-",
437 + true,
438 + 1,
439 + )
440 + .unwrap();
441 +
442 + for (i, req) in requests.iter().enumerate() {
443 + let token_ids = req.prompt_token_ids.as_ref().expect("should have token IDs");
444 + assert_eq!(
445 + token_ids.len(),
446 + target_len,
447 + "Request {i}: expected {target_len} token IDs, got {}",
448 + token_ids.len()
449 + );
450 + assert_eq!(req.prompt_len, target_len);
451 + }
452 + }
453 +
454 + /// Test encode/decode roundtrip stability for tiktoken.
455 + /// After one decode→encode cycle with UTF-8-safe tokens, length must not drift.
456 + #[test]
457 + #[ignore]
458 + fn test_tiktoken_roundtrip_stability() {
459 + let tokenizer = tokenizer::load_tokenizer("Qwen/Qwen2.5-0.5B", false, None);
460 + let tokenizer = match tokenizer {
461 + Ok(t) => t,
462 + Err(e) => {
463 + eprintln!("Skipping roundtrip test (tokenizer unavailable): {e}");
464 + return;
465 + }
466 + };
467 +
468 + let allowed = tokenizer.get_allowed_tokens();
469 + assert!(!allowed.is_empty(), "allowed tokens should not be empty");
470 +
471 + // Build a sequence from allowed tokens only
472 + use rand::rngs::StdRng;
473 + use rand::{Rng, SeedableRng};
474 + let mut rng = StdRng::seed_from_u64(99);
475 + let seq: Vec<u32> = (0..512).map(|_| allowed[rng.random_range(0..allowed.len())]).collect();
476 +
477 + let decoded = tokenizer.decode(&seq, true).unwrap();
478 + let re_encoded = tokenizer.encode(&decoded, false).unwrap();
479 + let re_decoded = tokenizer.decode(&re_encoded, true).unwrap();
480 + let re_re_encoded = tokenizer.encode(&re_decoded, false).unwrap();
481 +
482 + // After first cycle, length should stabilize
483 + assert_eq!(
484 + re_encoded.len(),
485 + re_re_encoded.len(),
486 + "Roundtrip should stabilize: first re-encode={}, second re-encode={}",
487 + re_encoded.len(),
488 + re_re_encoded.len()
489 + );
490 + }
491 +}
rust/src/bench/src/datasets/random_mm.rs new
+657
@@ -0,0 +1,657 @@
1 +// SPDX-License-Identifier: Apache-2.0
2 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +use std::io::Cursor;
5 +use std::sync::Arc;
6 +
7 +use base64::Engine as _;
8 +use rand::rngs::StdRng;
9 +use rand::{Rng, SeedableRng};
10 +use rayon::prelude::*;
11 +
12 +use super::SampleRequest;
13 +use crate::error::{BenchError, Result};
14 +use crate::tokenizer::TokenizerKind;
15 +
16 +/// A bucket key: (height, width, num_frames). num_frames=1 means image, >1 means video.
17 +#[derive(Debug, Clone)]
18 +pub struct MmBucketKey {
19 + pub height: u32,
20 + pub width: u32,
21 + pub num_frames: u32,
22 +}
23 +
24 +/// Per-modality hard caps.
25 +#[derive(Debug, Clone)]
26 +pub struct MmLimitPerPrompt {
27 + pub image: usize,
28 + pub video: usize,
29 +}
30 +
31 +impl Default for MmLimitPerPrompt {
32 + fn default() -> Self {
33 + Self {
34 + image: 255,
35 + video: 1,
36 + }
37 + }
38 +}
39 +
40 +/// Parse the limit-mm-per-prompt JSON string, e.g. `{"image": 3, "video": 0}`.
41 +pub fn parse_limit_mm_per_prompt(s: &str) -> Result<MmLimitPerPrompt> {
42 + let v: serde_json::Value = serde_json::from_str(s)
43 + .map_err(|e| BenchError::Config(format!("Invalid --random-mm-limit-mm-per-prompt: {e}")))?;
44 + let obj = v.as_object().ok_or_else(|| {
45 + BenchError::Config("--random-mm-limit-mm-per-prompt must be a JSON object".into())
46 + })?;
47 + let image = obj.get("image").and_then(|v| v.as_u64()).unwrap_or(255) as usize;
48 + let video = obj.get("video").and_then(|v| v.as_u64()).unwrap_or(1) as usize;
49 + Ok(MmLimitPerPrompt { image, video })
50 +}
51 +
52 +/// Parse the bucket config string in Python-style syntax.
53 +///
54 +/// Accepts: `{(256,256,1): 0.5, (720,1280,1): 0.5}`
55 +/// Each key is `(height, width, num_frames)` and value is the probability weight.
56 +pub fn parse_bucket_config(s: &str) -> Result<Vec<(MmBucketKey, f64)>> {
57 + let trimmed = s.trim();
58 + let inner = trimmed
59 + .strip_prefix('{')
60 + .and_then(|s| s.strip_suffix('}'))
61 + .ok_or_else(|| BenchError::Config("Bucket config must be wrapped in {}".into()))?;
62 +
63 + let mut buckets = Vec::new();
64 + let chars: Vec<char> = inner.chars().collect();
65 + let len = chars.len();
66 + let mut i = 0;
67 +
68 + while i < len {
69 + // Skip whitespace and commas
70 + while i < len && (chars[i].is_whitespace() || chars[i] == ',') {
71 + i += 1;
72 + }
73 + if i >= len {
74 + break;
75 + }
76 +
77 + // Expect '('
78 + if chars[i] != '(' {
79 + return Err(BenchError::Config(format!(
80 + "Expected '(' in bucket config at position {i}"
81 + )));
82 + }
83 + i += 1;
84 +
85 + // Read until ')'
86 + let tuple_start = i;
87 + while i < len && chars[i] != ')' {
88 + i += 1;
89 + }
90 + if i >= len {
91 + return Err(BenchError::Config("Unclosed '(' in bucket config".into()));
92 + }
93 + let tuple_str: String = chars[tuple_start..i].iter().collect();
94 + i += 1; // skip ')'
95 +
96 + // Skip whitespace, then expect ':'
97 + while i < len && chars[i].is_whitespace() {
98 + i += 1;
99 + }
100 + if i >= len || chars[i] != ':' {
101 + return Err(BenchError::Config(
102 + "Expected ':' after tuple in bucket config".into(),
103 + ));
104 + }
105 + i += 1;
106 +
107 + // Skip whitespace
108 + while i < len && chars[i].is_whitespace() {
109 + i += 1;
110 + }
111 +
112 + // Read the probability value until ',' or end
113 + let val_start = i;
114 + while i < len && chars[i] != ',' {
115 + i += 1;
116 + }
117 + let val_str: String = chars[val_start..i].iter().collect();
118 +
119 + // Parse tuple
120 + let parts: Vec<&str> = tuple_str.split(',').collect();
121 + if parts.len() != 3 {
122 + return Err(BenchError::Config(format!(
123 + "Bucket key must have 3 values (height,width,num_frames), got: ({tuple_str})"
124 + )));
125 + }
126 +
127 + let height: u32 = parts[0].trim().parse().map_err(|_| {
128 + BenchError::Config(format!(
129 + "Invalid height in bucket config: '{}'",
130 + parts[0].trim()
131 + ))
132 + })?;
133 + let width: u32 = parts[1].trim().parse().map_err(|_| {
134 + BenchError::Config(format!(
135 + "Invalid width in bucket config: '{}'",
136 + parts[1].trim()
137 + ))
138 + })?;
139 + let num_frames: u32 = parts[2].trim().parse().map_err(|_| {
140 + BenchError::Config(format!(
141 + "Invalid num_frames in bucket config: '{}'",
142 + parts[2].trim()
143 + ))
144 + })?;
145 + let prob: f64 = val_str.trim().parse().map_err(|_| {
146 + BenchError::Config(format!(
147 + "Invalid probability in bucket config: '{}'",
148 + val_str.trim()
149 + ))
150 + })?;
151 +
152 + if prob < 0.0 {
153 + return Err(BenchError::Config(format!(
154 + "Bucket probability must be non-negative, got: {prob}"
155 + )));
156 + }
157 +
158 + buckets.push((
159 + MmBucketKey {
160 + height,
161 + width,
162 + num_frames,
163 + },
164 + prob,
165 + ));
166 + }
167 +
168 + if buckets.is_empty() {
169 + return Err(BenchError::Config(
170 + "Bucket config must have at least one entry".into(),
171 + ));
172 + }
173 +
174 + Ok(buckets)
175 +}
176 +
177 +/// JSON fragment prefix/suffix for image content blocks.
178 +const IMG_JSON_PREFIX: &str = r#"{"type":"image_url","image_url":{"url":"data:image/jpeg;base64,"#;
179 +const IMG_JSON_SUFFIX: &str = r#""}}"#;
180 +
181 +/// Generate a synthetic random JPEG image and return it as a pre-serialized JSON fragment.
182 +///
183 +/// Builds the complete JSON string in a single allocation:
184 +/// `{"type":"image_url","image_url":{"url":"data:image/jpeg;base64,<b64>"}}`
185 +///
186 +/// The base64 data is written directly into the final string — no intermediate
187 +/// String or format!() copy.
188 +fn generate_random_image(width: u32, height: u32, rng: &mut StdRng) -> Result<Arc<str>> {
189 + let pixel_count = (width as usize) * (height as usize) * 3;
190 + let mut pixels = vec![0u8; pixel_count];
191 + rng.fill(pixels.as_mut_slice());
192 +
193 + let img = image::RgbImage::from_raw(width, height, pixels)
194 + .ok_or_else(|| BenchError::Config("Failed to create image from random pixels".into()))?;
195 +
196 + // Pre-allocate JPEG buffer (random pixels compress poorly, estimate ~60% of raw)
197 + let estimated_jpeg = pixel_count * 3 / 5;
198 + let mut buf = Cursor::new(Vec::with_capacity(estimated_jpeg));
199 + img.write_to(&mut buf, image::ImageFormat::Jpeg)
200 + .map_err(|e| BenchError::Config(format!("Failed to encode JPEG: {e}")))?;
201 +
202 + let jpeg_bytes = buf.into_inner();
203 +
204 + // Pre-compute exact output size: prefix + base64_len + suffix
205 + let b64_len = jpeg_bytes.len().div_ceil(3) * 4;
206 + let total_len = IMG_JSON_PREFIX.len() + b64_len + IMG_JSON_SUFFIX.len();
207 +
208 + // Single allocation: write base64 directly into the JSON fragment string
209 + let mut json_fragment = String::with_capacity(total_len);
210 + json_fragment.push_str(IMG_JSON_PREFIX);
211 + base64::engine::general_purpose::STANDARD.encode_string(&jpeg_bytes, &mut json_fragment);
212 + json_fragment.push_str(IMG_JSON_SUFFIX);
213 +
214 + Ok(Arc::from(json_fragment))
215 +}
216 +
217 +/// Sample multimodal items for a single request.
218 +///
219 +/// Returns a list of (height, width, num_frames) tuples.
220 +fn sample_mm_items(
221 + rng: &mut StdRng,
222 + min_items: usize,
223 + max_items: usize,
224 + buckets: &[(MmBucketKey, f64)],
225 + limit: &MmLimitPerPrompt,
226 +) -> Vec<MmBucketKey> {
227 + let num_items = if min_items == max_items {
228 + min_items
229 + } else {
230 + rng.random_range(min_items..=max_items)
231 + };
232 +
233 + // Filter to non-zero probability buckets
234 + let active_buckets: Vec<&(MmBucketKey, f64)> =
235 + buckets.iter().filter(|(_, p)| *p > 0.0).collect();
236 + if active_buckets.is_empty() || num_items == 0 {
237 + return Vec::new();
238 + }
239 +
240 + let total_weight: f64 = active_buckets.iter().map(|(_, p)| p).sum();
241 + if total_weight <= 0.0 {
242 + return Vec::new();
243 + }
244 +
245 + let mut result = Vec::with_capacity(num_items);
246 + let mut image_count = 0usize;
247 + let mut video_count = 0usize;
248 +
249 + for _ in 0..num_items {
250 + // Build normalized weights considering remaining capacity
251 + let mut weights: Vec<f64> = Vec::with_capacity(active_buckets.len());
252 + for (key, prob) in &active_buckets {
253 + let is_video = key.num_frames > 1;
254 + let at_limit = if is_video {
255 + video_count >= limit.video
256 + } else {
257 + image_count >= limit.image
258 + };
259 + weights.push(if at_limit { 0.0 } else { *prob });
260 + }
261 +
262 + let w_total: f64 = weights.iter().sum();
263 + if w_total <= 0.0 {
264 + break; // All modalities at limit
265 + }
266 +
267 + // Weighted random selection (strict `<` to avoid selecting zero-weight buckets)
268 + let r = rng.random::<f64>() * w_total;
269 + let mut cumulative = 0.0;
270 + // Default to last non-zero-weight bucket (floating-point accumulation fallback)
271 + let mut selected_idx = weights.iter().rposition(|w| *w > 0.0).unwrap_or(0);
272 + for (i, w) in weights.iter().enumerate() {
273 + cumulative += w;
274 + if r < cumulative {
275 + selected_idx = i;
276 + break;
277 + }
278 + }
279 +
280 + let (key, _) = &active_buckets[selected_idx];
281 + if key.num_frames > 1 {
282 + video_count += 1;
283 + } else {
284 + image_count += 1;
285 + }
286 + result.push(key.clone());
287 + }
288 +
289 + result
290 +}
291 +
292 +/// Generate random multimodal dataset.
293 +///
294 +/// Mirrors Python's RandomMultiModalDataset.sample() from datasets.py.
295 +/// Generates text prompts with exact token lengths and random images/videos.
296 +pub fn generate_random_mm_dataset(
297 + tokenizer: &TokenizerKind,
298 + num_requests: usize,
299 + input_len: usize,
300 + output_len: usize,
301 + prefix_len: usize,
302 + range_ratio: crate::config::RangeRatio,
303 + seed: u64,
304 + request_id_prefix: &str,
305 + base_items_per_request: usize,
306 + num_mm_items_range_ratio: f64,
307 + limit: &MmLimitPerPrompt,
308 + buckets: &[(MmBucketKey, f64)],
309 + enable_multimodal_chat: bool,
310 +) -> Result<Vec<SampleRequest>> {
311 + if !(0.0..=1.0).contains(&num_mm_items_range_ratio) {
312 + return Err(BenchError::Config(
313 + "num_mm_items_range_ratio must be in [0, 1]".into(),
314 + ));
315 + }
316 +
317 + // Check for video buckets with non-zero probability
318 + for (key, prob) in buckets {
319 + if key.num_frames > 1 && *prob > 0.0 {
320 + return Err(BenchError::Config(
321 + "Video generation (num_frames > 1) is not yet supported in Rust. \
322 + Set video bucket probabilities to 0.0."
323 + .into(),
324 + ));
325 + }
326 + }
327 +
328 + // Compute item count bounds
329 + let n = base_items_per_request as f64;
330 + let r = num_mm_items_range_ratio;
331 + let min_items = (n * (1.0 - r)).floor().max(0.0) as usize;
332 + let max_items = (n * (1.0 + r)).ceil() as usize;
333 + // Clamp to total modality limit
334 + let total_limit = limit.image + limit.video;
335 + let max_items = max_items.min(total_limit);
336 + let min_items = min_items.min(max_items);
337 +
338 + let vocab_size = tokenizer.vocab_size();
339 + let allowed_tokens = tokenizer.get_allowed_tokens();
340 + if allowed_tokens.is_empty() {
341 + return Err(BenchError::Tokenizer("No allowed tokens found".into()));
342 + }
343 +
344 + let num_special = tokenizer.num_special_tokens_to_add();
345 + let real_input_len = input_len.saturating_sub(num_special);
346 +
347 + // Python semantics: sample uniformly from [len*(1-r), len*(1+r)].
348 + let (input_low, input_high) = range_ratio.input_bounds(real_input_len);
349 + let (output_low, output_high) = range_ratio.output_bounds(output_len);
350 +
351 + // Pre-generate per-request params
352 + let mut rng = StdRng::seed_from_u64(seed);
353 +
354 + struct RequestParams {
355 + input_len: usize,
356 + output_len: usize,
357 + offset: usize,
358 + }
359 +
360 + let params: Vec<RequestParams> = (0..num_requests)
361 + .map(|_| {
362 + let il = if input_low == input_high {
363 + input_low
364 + } else {
365 + rng.random_range(input_low..=input_high)
366 + };
367 + let ol = if output_low == output_high {
368 + output_low
369 + } else {
370 + rng.random_range(output_low..=output_high)
371 + };
372 + let off = rng.random_range(0..vocab_size as usize);
373 + RequestParams {
374 + input_len: il,
375 + output_len: ol,
376 + offset: off,
377 + }
378 + })
379 + .collect();
380 +
381 + // Pre-generate multimodal item configs per request
382 + let mm_configs: Vec<Vec<MmBucketKey>> = (0..num_requests)
383 + .map(|_| sample_mm_items(&mut rng, min_items, max_items, buckets, limit))
384 + .collect();
385 +
386 + // Generate text prompts (need text for chat backend, not just token IDs)
387 + let prefix_token_ids = if prefix_len > 0 {
388 + generate_prefix(tokenizer, &allowed_tokens, prefix_len, seed)?
389 + } else {
390 + Vec::new()
391 + };
392 +
393 + // Generate token sequences
394 + let prefix_ref = &prefix_token_ids;
395 + let allowed_ref = &allowed_tokens;
396 +
397 + let token_sequences: Vec<Vec<u32>> = params
398 + .par_iter()
399 + .enumerate()
400 + .map(|(i, p)| {
401 + let at_len = allowed_ref.len();
402 + let mut seq = Vec::with_capacity(prefix_ref.len() + p.input_len);
403 + seq.extend_from_slice(prefix_ref);
404 + for j in 0..p.input_len {
405 + seq.push(allowed_ref[(p.offset + i + j) % at_len]);
406 + }
407 + seq
408 + })
409 + .collect();
410 +
411 + let target_lens: Vec<usize> = params.iter().map(|p| prefix_len + p.input_len).collect();
412 +
413 + // Decode tokens to text (chat backend needs text prompts for multimodal)
414 + let prompts: Result<Vec<String>> = token_sequences
415 + .into_par_iter()
416 + .enumerate()
417 + .map(|(i, tokens)| {
418 + let (text, _adjusted) = super::random::gen_prompt_decode_to_target_len(
419 + tokenizer,
420 + &tokens,
421 + target_lens[i],
422 + false,
423 + allowed_ref,
424 + )?;
425 + Ok(text)
426 + })
427 + .collect();
428 + let prompts = prompts?;
429 +
430 + // Generate images for each request (parallel per request)
431 + // Each request gets its own RNG seeded deterministically.
432 + let rid_prefix = request_id_prefix.to_string();
433 + let result: Vec<SampleRequest> = prompts
434 + .into_par_iter()
435 + .enumerate()
436 + .map(|(i, prompt)| {
437 + let mut item_rng =
438 + StdRng::seed_from_u64(seed.wrapping_add(i as u64).wrapping_add(0xBEEF));
439 + let mm_items: Vec<Arc<str>> = mm_configs[i]
440 + .iter()
441 + .map(|key| {
442 + generate_random_image(key.width, key.height, &mut item_rng)
443 + .expect("Image generation should not fail")
444 + })
445 + .collect();
446 +
447 + let mm_content: Option<Arc<[Arc<str>]>> = if mm_items.is_empty() {
448 + None
449 + } else {
450 + Some(Arc::from(mm_items))
451 + };
452 +
453 + // --enable-multimodal-chat: pre-build the full chat `messages` array
454 + // (text part + mm items) at dataset time, mirroring Python's
455 + // apply_multimodal_chat_transformation. mm content moves inside the
456 + // messages string; the backend splices it verbatim.
457 + let (mm_content, chat_messages_json) = if enable_multimodal_chat {
458 + let msgs = build_chat_messages_json(&prompt, mm_content.as_deref());
459 + (None, Some(Arc::from(msgs.as_str())))
460 + } else {
461 + (mm_content, None)
462 + };
463 +
464 + SampleRequest {
465 + prompt: Arc::from(prompt.as_str()),
466 + prompt_len: target_lens[i],
467 + expected_output_len: params[i].output_len,
468 + request_id: Some(format!("{rid_prefix}{i}")),
469 + multi_modal_content: mm_content,
470 + chat_messages_json,
471 + ..Default::default()
472 + }
473 + })
474 + .collect();
475 +
476 + Ok(result)
477 +}
478 +
479 +/// Pre-serialize the OpenAI chat `messages` array for --enable-multimodal-chat.
480 +///
481 +/// Produces `[{"role":"user","content":[{"type":"text","text":"..."},<frag>,...]}]`
482 +/// by concatenating the JSON-escaped prompt with the pre-serialized mm fragments,
483 +/// so the ~200KB+ base64 image data is never parsed or re-serialized.
484 +pub(crate) fn build_chat_messages_json(prompt: &str, mm_items: Option<&[Arc<str>]>) -> String {
485 + let mm_total: usize =
486 + mm_items.map(|items| items.iter().map(|f| f.len() + 1).sum()).unwrap_or(0);
487 + let mut msgs = String::with_capacity(64 + prompt.len() * 2 + mm_total);
488 + msgs.push_str(r#"[{"role":"user","content":[{"type":"text","text":"#);
489 + // serde_json::to_string on &str produces a JSON-escaped quoted string
490 + msgs.push_str(&serde_json::to_string(prompt).unwrap());
491 + msgs.push('}');
492 + for fragment in mm_items.unwrap_or(&[]) {
493 + msgs.push(',');
494 + msgs.push_str(fragment);
495 + }
496 + msgs.push_str("]}]");
497 + msgs
498 +}
499 +
500 +fn generate_prefix(
501 + tokenizer: &TokenizerKind,
502 + allowed_tokens: &[u32],
503 + prefix_len: usize,
504 + seed: u64,
505 +) -> Result<Vec<u32>> {
506 + let mut rng = StdRng::seed_from_u64(seed.wrapping_add(0xDEAD));
507 + let tokens: Vec<u32> = (0..prefix_len)
508 + .map(|_| allowed_tokens[rng.random_range(0..allowed_tokens.len())])
509 + .collect();
510 +
511 + let (_, adjusted) = super::random::gen_prompt_decode_to_target_len(
512 + tokenizer,
513 + &tokens,
514 + prefix_len,
515 + false,
516 + allowed_tokens,
517 + )?;
518 + Ok(adjusted)
519 +}
520 +
521 +#[cfg(test)]
522 +mod tests {
523 + use super::*;
524 +
525 + #[test]
526 + fn test_parse_bucket_config_basic() {
527 + let input = "{(256,256,1): 0.5, (720,1280,1): 0.5}";
528 + let buckets = parse_bucket_config(input).unwrap();
529 + assert_eq!(buckets.len(), 2);
530 + assert_eq!(buckets[0].0.height, 256);
531 + assert_eq!(buckets[0].0.width, 256);
532 + assert_eq!(buckets[0].0.num_frames, 1);
533 + assert!((buckets[0].1 - 0.5).abs() < 1e-10);
534 + assert_eq!(buckets[1].0.height, 720);
535 + assert_eq!(buckets[1].0.width, 1280);
536 + }
537 +
538 + #[test]
539 + fn test_parse_bucket_config_single() {
540 + let input = "{(1024, 800, 1): 1.0}";
541 + let buckets = parse_bucket_config(input).unwrap();
542 + assert_eq!(buckets.len(), 1);
543 + assert_eq!(buckets[0].0.height, 1024);
544 + assert_eq!(buckets[0].0.width, 800);
545 + assert_eq!(buckets[0].0.num_frames, 1);
546 + assert!((buckets[0].1 - 1.0).abs() < 1e-10);
547 + }
548 +
549 + #[test]
550 + fn test_parse_bucket_config_with_video() {
551 + let input = "{(256,256,1): 0.4, (720,1280,1): 0.4, (720,1280,16): 0.2}";
552 + let buckets = parse_bucket_config(input).unwrap();
553 + assert_eq!(buckets.len(), 3);
554 + assert_eq!(buckets[2].0.num_frames, 16);
555 + }
556 +
557 + #[test]
558 + fn test_build_chat_messages_json_valid_and_ordered() {
559 + let frag: Arc<str> =
560 + Arc::from(r#"{"type":"image_url","image_url":{"url":"data:image/jpeg;base64,AAAA"}}"#);
561 + let msgs = build_chat_messages_json("hi \"there\"\nline2", Some(&[frag]));
562 + let v: serde_json::Value = serde_json::from_str(&msgs).expect("must be valid JSON");
563 + assert_eq!(v.as_array().unwrap().len(), 1);
564 + assert_eq!(v[0]["role"], "user");
565 + let content = v[0]["content"].as_array().unwrap();
566 + assert_eq!(content.len(), 2);
567 + assert_eq!(content[0]["type"], "text");
568 + assert_eq!(content[0]["text"], "hi \"there\"\nline2");
569 + assert_eq!(content[1]["type"], "image_url");
570 + }
571 +
572 + #[test]
573 + fn test_build_chat_messages_json_text_only() {
574 + let msgs = build_chat_messages_json("plain", None);
575 + let v: serde_json::Value = serde_json::from_str(&msgs).unwrap();
576 + let content = v[0]["content"].as_array().unwrap();
577 + assert_eq!(content.len(), 1);
578 + assert_eq!(content[0]["text"], "plain");
579 + }
580 +
581 + #[test]
582 + fn test_parse_limit_mm_per_prompt() {
583 + let input = r#"{"image": 3, "video": 0}"#;
584 + let limit = parse_limit_mm_per_prompt(input).unwrap();
585 + assert_eq!(limit.image, 3);
586 + assert_eq!(limit.video, 0);
587 + }
588 +
589 + #[test]
590 + fn test_parse_limit_mm_per_prompt_defaults() {
591 + let input = r#"{}"#;
592 + let limit = parse_limit_mm_per_prompt(input).unwrap();
593 + assert_eq!(limit.image, 255);
594 + assert_eq!(limit.video, 1);
595 + }
596 +
597 + #[test]
598 + fn test_sample_mm_items_basic() {
599 + let mut rng = StdRng::seed_from_u64(42);
600 + let buckets = vec![
601 + (
602 + MmBucketKey {
603 + height: 256,
604 + width: 256,
605 + num_frames: 1,
606 + },
607 + 0.5,
608 + ),
609 + (
610 + MmBucketKey {
611 + height: 720,
612 + width: 1280,
613 + num_frames: 1,
614 + },
615 + 0.5,
616 + ),
617 + ];
618 + let limit = MmLimitPerPrompt { image: 5, video: 0 };
619 + let items = sample_mm_items(&mut rng, 2, 3, &buckets, &limit);
620 + assert!(items.len() >= 2 && items.len() <= 3);
621 + for item in &items {
622 + assert_eq!(item.num_frames, 1);
623 + }
624 + }
625 +
626 + #[test]
627 + fn test_sample_mm_items_respects_limit() {
628 + let mut rng = StdRng::seed_from_u64(42);
629 + let buckets = vec![(
630 + MmBucketKey {
631 + height: 256,
632 + width: 256,
633 + num_frames: 1,
634 + },
635 + 1.0,
636 + )];
637 + let limit = MmLimitPerPrompt { image: 2, video: 0 };
638 + let items = sample_mm_items(&mut rng, 5, 5, &buckets, &limit);
639 + // Should be capped at 2 due to image limit
640 + assert_eq!(items.len(), 2);
641 + }
642 +
643 + #[test]
644 + fn test_generate_random_image() {
645 + let mut rng = StdRng::seed_from_u64(42);
646 + let result = generate_random_image(64, 64, &mut rng).unwrap();
647 + // Result is a pre-serialized JSON fragment
648 + assert!(
649 + result
650 + .starts_with(r#"{"type":"image_url","image_url":{"url":"data:image/jpeg;base64,"#)
651 + );
652 + assert!(result.ends_with(r#""}}"#));
653 + // Verify it's valid JSON
654 + let parsed: serde_json::Value = serde_json::from_str(&result).unwrap();
655 + assert_eq!(parsed["type"], "image_url");
656 + }
657 +}
rust/src/bench/src/datasets/random_rerank.rs new
+185
@@ -0,0 +1,185 @@
1 +// SPDX-License-Identifier: Apache-2.0
2 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +//! Random dataset specialized for scoring/rerank benchmarks: each request is
5 +//! one query plus a batch of documents. Mirrors Python's
6 +//! `RandomDatasetForReranking`.
7 +//!
8 +//! With `is_reranker` (default): the query and each document share the
9 +//! request's token budget (`query + sep + doc ~= input_len`), and every
10 +//! batched request counts the query once per document pair.
11 +//! With `--no-reranker` (embedding-based scoring): the query is just another
12 +//! embedding input occupying the first batch slot.
13 +
14 +use std::sync::Arc;
15 +
16 +use rand::rngs::StdRng;
17 +use rand::{Rng, SeedableRng};
18 +use rayon::prelude::*;
19 +
20 +use super::SampleRequest;
21 +use crate::config::RangeRatio;
22 +use crate::error::{BenchError, Result};
23 +use crate::tokenizer::TokenizerKind;
24 +
25 +pub fn generate_random_rerank_dataset(
26 + tokenizer: &TokenizerKind,
27 + num_requests: usize,
28 + input_len: usize,
29 + range_ratio: RangeRatio,
30 + seed: u64,
31 + request_id_prefix: &str,
32 + batch_size: usize,
33 + is_reranker: bool,
34 +) -> Result<Vec<SampleRequest>> {
35 + let allowed_tokens = tokenizer.get_allowed_tokens();
36 + if allowed_tokens.is_empty() {
37 + return Err(BenchError::Tokenizer("No allowed tokens found".into()));
38 + }
39 + let allowed_ref = &allowed_tokens;
40 +
41 + let num_special = tokenizer.num_special_tokens_to_add();
42 + let real_input_len = input_len.saturating_sub(num_special);
43 +
44 + let n_sep_tokens = usize::from(is_reranker);
45 + let query_len_param = if is_reranker {
46 + (real_input_len / 2).saturating_sub(n_sep_tokens)
47 + } else {
48 + real_input_len
49 + };
50 +
51 + let mut rng = StdRng::seed_from_u64(seed);
52 + let sample = |rng: &mut StdRng, (low, high): (usize, usize)| -> usize {
53 + if low == high {
54 + low
55 + } else {
56 + rng.random_range(low..=high)
57 + }
58 + };
59 +
60 + // One query length for the whole run, like Python.
61 + let query_len = sample(&mut rng, range_ratio.input_bounds(query_len_param));
62 +
63 + // --no-reranker folds the query into the first batch slot.
64 + let (num_docs, docs_per_batch, doc_len_param) = if is_reranker {
65 + let doc_len = real_input_len.saturating_sub(query_len).saturating_sub(n_sep_tokens);
66 + (num_requests, batch_size, doc_len)
67 + } else {
68 + (num_requests - 1, batch_size - 1, real_input_len)
69 + };
70 + if doc_len_param == 0 {
71 + return Err(BenchError::Config(format!(
72 + "random-rerank: --random-input-len {input_len} leaves no budget for documents \
73 + (query_len={query_len})"
74 + )));
75 + }
76 +
77 + // Pre-sample per-document lengths and offsets deterministically.
78 + let doc_bounds = range_ratio.input_bounds(doc_len_param);
79 + let doc_params: Vec<(usize, usize)> = (0..num_docs)
80 + .map(|_| {
81 + (
82 + sample(&mut rng, doc_bounds),
83 + rng.random_range(0..allowed_ref.len()),
84 + )
85 + })
86 + .collect();
87 + let query_offset = rng.random_range(0..allowed_ref.len());
88 +
89 + // Exact-length text: token sequence -> decode -> re-encode -> truncate -> decode.
90 + let gen_text = |target: usize, offset: usize, index: usize| -> Result<(Arc<str>, usize)> {
91 + let at_len = allowed_ref.len();
92 + let tokens: Vec<u32> =
93 + (0..target).map(|j| allowed_ref[(offset + index + j) % at_len]).collect();
94 + let text = tokenizer.decode(&tokens, true)?;
95 + let mut re_encoded = tokenizer.encode(&text, false)?;
96 + re_encoded.truncate(target);
97 + let final_text = tokenizer.decode(&re_encoded, true)?;
98 + Ok((Arc::from(final_text), re_encoded.len()))
99 + };
100 +
101 + let (query_prompt, query_input_len) = gen_text(query_len, query_offset, 0)?;
102 +
103 + let docs: Vec<(Arc<str>, usize)> = doc_params
104 + .par_iter()
105 + .enumerate()
106 + .map(|(i, (len, offset))| gen_text(*len, *offset, i + 1))
107 + .collect::<Result<Vec<_>>>()?;
108 +
109 + // Batch documents; every request is [query, doc1, doc2, ...].
110 + let rid_prefix = request_id_prefix.to_string();
111 + let requests = docs
112 + .chunks(docs_per_batch)
113 + .enumerate()
114 + .map(|(batch_idx, batch)| {
115 + let query_contrib = if is_reranker {
116 + (query_input_len + n_sep_tokens) * batch.len()
117 + } else {
118 + query_input_len
119 + };
120 + let mut prompt_list: Vec<Arc<str>> = Vec::with_capacity(batch.len() + 1);
121 + prompt_list.push(query_prompt.clone());
122 + prompt_list.extend(batch.iter().map(|(text, _)| text.clone()));
123 + SampleRequest {
124 + prompt_list: Some(Arc::from(prompt_list)),
125 + prompt_len: query_contrib + batch.iter().map(|(_, len)| len).sum::<usize>(),
126 + expected_output_len: 0,
127 + request_id: Some(format!("{rid_prefix}{batch_idx}")),
128 + ..Default::default()
129 + }
130 + })
131 + .collect();
132 +
133 + Ok(requests)
134 +}
135 +
136 +#[cfg(test)]
137 +mod tests {
138 + use super::*;
139 +
140 + /// gpt2 via built-in tiktoken encoding — loads without network access.
141 + fn test_tokenizer() -> TokenizerKind {
142 + crate::tokenizer::load_tokenizer("gpt2", false, None)
143 + .expect("gpt2 built-in tiktoken should always load without network")
144 + }
145 +
146 + fn fixed_ratio() -> RangeRatio {
147 + RangeRatio::parse("0.0").unwrap()
148 + }
149 +
150 + #[test]
151 + fn test_random_rerank_reranker_mode() {
152 + let tok = test_tokenizer();
153 + // 6 docs in batches of 3 -> 2 requests of [query, d1, d2, d3]
154 + let reqs = generate_random_rerank_dataset(&tok, 6, 128, fixed_ratio(), 0, "t-", 3, true)
155 + .expect("generation should succeed");
156 + assert_eq!(reqs.len(), 2);
157 + for r in &reqs {
158 + let list = r.prompt_list.as_ref().expect("prompt_list must be set");
159 + assert_eq!(list.len(), 4);
160 + assert_eq!(r.expected_output_len, 0);
161 + assert!(r.prompt_len > 0);
162 + }
163 + // Same query shared across requests
164 + assert_eq!(
165 + reqs[0].prompt_list.as_ref().unwrap()[0],
166 + reqs[1].prompt_list.as_ref().unwrap()[0]
167 + );
168 + // Reranker budget: query+sep+doc pairs stay near input_len per pair
169 + // (query ~63, doc ~64 for input_len=128, gpt2 has no special tokens)
170 + let list = reqs[0].prompt_list.as_ref().unwrap();
171 + let query_tokens = tok.encode(&list[0], false).unwrap().len();
172 + assert!(query_tokens <= 64, "query too long: {query_tokens}");
173 + }
174 +
175 + #[test]
176 + fn test_random_rerank_no_reranker_mode() {
177 + let tok = test_tokenizer();
178 + // no-reranker: query occupies first slot; 5 non-query docs in batches of 2
179 + let reqs = generate_random_rerank_dataset(&tok, 6, 64, fixed_ratio(), 0, "t-", 3, false)
180 + .expect("generation should succeed");
181 + // 6-1=5 docs, batches of 3-1=2 -> 3 requests
182 + assert_eq!(reqs.len(), 3);
183 + assert_eq!(reqs[0].prompt_list.as_ref().unwrap().len(), 3);
184 + }
185 +}
rust/src/bench/src/datasets/sharegpt.rs new
+209
@@ -0,0 +1,209 @@
1 +// SPDX-License-Identifier: Apache-2.0
2 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +use std::sync::Arc;
5 +
6 +use rand::rngs::StdRng;
7 +use rand::seq::SliceRandom;
8 +use rand::{Rng, SeedableRng};
9 +
10 +use super::SampleRequest;
11 +use crate::error::{BenchError, Result};
12 +use crate::tokenizer::TokenizerKind;
13 +
14 +/// Default validation bounds matching Python's is_valid_sequence() defaults.
15 +const MIN_LEN: usize = 4;
16 +const MAX_PROMPT_LEN: usize = 1024;
17 +const MAX_TOTAL_LEN: usize = 2048;
18 +
19 +/// Default HuggingFace dataset repo and filename for ShareGPT.
20 +const DEFAULT_SHAREGPT_REPO: &str = "anon8231489123/ShareGPT_Vicuna_unfiltered";
21 +const DEFAULT_SHAREGPT_FILE: &str = "ShareGPT_V3_unfiltered_cleaned_split.json";
22 +
23 +/// Download the default ShareGPT dataset from HuggingFace Hub.
24 +/// Uses hf-hub's built-in cache — subsequent calls return the cached path instantly.
25 +pub fn download_sharegpt_dataset() -> Result<String> {
26 + println!(
27 + "Downloading ShareGPT dataset from {DEFAULT_SHAREGPT_REPO}/{DEFAULT_SHAREGPT_FILE} ..."
28 + );
29 + let repo = crate::hub::HubRepo::dataset(DEFAULT_SHAREGPT_REPO.to_string());
30 + let path = repo.get(DEFAULT_SHAREGPT_FILE).map_err(|e| {
31 + BenchError::Config(format!(
32 + "Failed to download ShareGPT dataset from '{DEFAULT_SHAREGPT_REPO}': {e}"
33 + ))
34 + })?;
35 + let path_str = path.to_string_lossy().to_string();
36 + println!("ShareGPT dataset ready: {path_str}");
37 + Ok(path_str)
38 +}
39 +
40 +/// Load and sample from a ShareGPT-format JSON dataset.
41 +///
42 +/// Mirrors Python's ShareGPTDataset from datasets.py:1230-1313.
43 +pub fn load_sharegpt_dataset(
44 + tokenizer: &TokenizerKind,
45 + dataset_path: &str,
46 + num_requests: usize,
47 + output_len_override: Option<usize>,
48 + seed: u64,
49 + request_id_prefix: &str,
50 + no_oversample: bool,
51 + disable_shuffle: bool,
52 +) -> Result<Vec<SampleRequest>> {
53 + // Load JSON file
54 + let content = std::fs::read_to_string(dataset_path).map_err(|e| {
55 + BenchError::Config(format!(
56 + "Failed to read ShareGPT file '{dataset_path}': {e}"
57 + ))
58 + })?;
59 +
60 + let data: serde_json::Value = serde_json::from_str(&content)
61 + .map_err(|e| BenchError::Config(format!("Invalid JSON in ShareGPT file: {e}")))?;
62 +
63 + let entries = data
64 + .as_array()
65 + .ok_or_else(|| BenchError::Config("ShareGPT file must contain a JSON array".into()))?;
66 +
67 + // Filter entries with at least 2 conversation turns
68 + let mut filtered: Vec<&serde_json::Value> = entries
69 + .iter()
70 + .filter(|entry| {
71 + entry
72 + .get("conversations")
73 + .and_then(|c| c.as_array())
74 + .map(|a| a.len() >= 2)
75 + .unwrap_or(false)
76 + })
77 + .collect();
78 +
79 + if filtered.is_empty() {
80 + return Err(BenchError::Config(
81 + "No valid entries in ShareGPT file (need at least 2 conversation turns)".into(),
82 + ));
83 + }
84 +
85 + // Shuffle (unless disabled)
86 + let mut rng = StdRng::seed_from_u64(seed);
87 + if !disable_shuffle {
88 + filtered.shuffle(&mut rng);
89 + }
90 +
91 + // Sample requests
92 + let mut samples = Vec::new();
93 + let mut ind = 0;
94 +
95 + for entry in &filtered {
96 + if samples.len() >= num_requests {
97 + break;
98 + }
99 +
100 + let conversations = entry["conversations"].as_array().unwrap();
101 + let prompt = conversations[0]["value"].as_str().unwrap_or("");
102 + let completion = conversations[1]["value"].as_str().unwrap_or("");
103 +
104 + if prompt.is_empty() {
105 + continue;
106 + }
107 +
108 + // Tokenize prompt and completion
109 + let prompt_ids = tokenizer.encode(prompt, false)?;
110 + let prompt_len = prompt_ids.len();
111 +
112 + let new_output_len = if let Some(override_len) = output_len_override {
113 + override_len
114 + } else {
115 + let completion_ids = tokenizer.encode(completion, false)?;
116 + completion_ids.len()
117 + };
118 +
119 + // Validate sequence lengths (matching Python's is_valid_sequence)
120 + let skip_min_output = output_len_override.is_some();
121 + if !is_valid_sequence(prompt_len, new_output_len, skip_min_output) {
122 + continue;
123 + }
124 +
125 + samples.push(SampleRequest {
126 + prompt: Arc::from(prompt),
127 + prompt_len,
128 + expected_output_len: new_output_len,
129 + request_id: Some(format!("{request_id_prefix}{ind}")),
130 + ..Default::default()
131 + });
132 + ind += 1;
133 + }
134 +
135 + // Oversample if dataset is smaller than requested
136 + if samples.len() < num_requests {
137 + if no_oversample {
138 + println!(
139 + "Skipping oversampling. Total samples: {} (requested: {num_requests})",
140 + samples.len()
141 + );
142 + } else if !samples.is_empty() {
143 + let needed = num_requests - samples.len();
144 + let original_len = samples.len();
145 + for i in 0..needed {
146 + let mut req = samples[rng.random_range(0..original_len)].clone();
147 + req.request_id = Some(format!("{request_id_prefix}{}", original_len + i));
148 + samples.push(req);
149 + }
150 + println!(
151 + "Oversampled requests from {original_len} to {} total samples.",
152 + samples.len()
153 + );
154 + }
155 + }
156 +
157 + if samples.is_empty() {
158 + return Err(BenchError::Config(
159 + "No valid samples after filtering ShareGPT dataset. \
160 + Try relaxing constraints or using a larger dataset."
161 + .into(),
162 + ));
163 + }
164 +
165 + Ok(samples)
166 +}
167 +
168 +/// Validate a sequence based on prompt and output lengths.
169 +/// Mirrors Python's is_valid_sequence() from datasets.py:260-284.
170 +fn is_valid_sequence(
171 + prompt_len: usize,
172 + output_len: usize,
173 + skip_min_output_len_check: bool,
174 +) -> bool {
175 + if prompt_len < MIN_LEN {
176 + return false;
177 + }
178 + if !skip_min_output_len_check && output_len < MIN_LEN {
179 + return false;
180 + }
181 + if prompt_len > MAX_PROMPT_LEN {
182 + return false;
183 + }
184 + if prompt_len + output_len > MAX_TOTAL_LEN {
185 + return false;
186 + }
187 + true
188 +}
189 +
190 +#[cfg(test)]
191 +mod tests {
192 + use super::*;
193 +
194 + #[test]
195 + fn test_is_valid_sequence() {
196 + // Valid
197 + assert!(is_valid_sequence(100, 50, false));
198 + // Prompt too short
199 + assert!(!is_valid_sequence(3, 50, false));
200 + // Output too short
201 + assert!(!is_valid_sequence(100, 3, false));
202 + // Output too short but skip check
203 + assert!(is_valid_sequence(100, 1, true));
204 + // Prompt too long
205 + assert!(!is_valid_sequence(1025, 50, false));
206 + // Combined too long
207 + assert!(!is_valid_sequence(1024, 1025, false));
208 + }
209 +}
rust/src/bench/src/datasets/sonnet.rs new
+131
@@ -0,0 +1,131 @@
1 +// SPDX-License-Identifier: Apache-2.0
2 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +use std::sync::Arc;
5 +
6 +use rand::SeedableRng;
7 +use rand::rngs::StdRng;
8 +use rand::seq::IndexedRandom;
9 +
10 +use super::SampleRequest;
11 +use crate::error::{BenchError, Result};
12 +use crate::tokenizer::TokenizerKind;
13 +
14 +/// Default values mirror Python's SonnetDataset defaults (datasets.py).
15 +pub const DEFAULT_PREFIX_LEN: usize = 200;
16 +pub const DEFAULT_INPUT_LEN: usize = 550;
17 +pub const DEFAULT_OUTPUT_LEN: usize = 150;
18 +
19 +const BASE_PROMPT: &str = "Pick as many lines as you can from these poem lines:\n";
20 +
21 +/// Shakespeare's sonnets, public domain. Bundled so `--dataset-name sonnet` works
22 +/// out of the box without `--dataset-path`. Source:
23 +/// https://raw.githubusercontent.com/vllm-project/vllm/main/benchmarks/sonnet.txt
24 +const BUILTIN_SONNET: &str = include_str!("sonnet.txt");
25 +
26 +/// Load the sonnet dataset and generate `num_requests` prompts targeting `input_len`
27 +/// total prompt tokens.
28 +///
29 +/// Mirrors Python's `SonnetDataset.sample()` from `vllm/benchmarks/datasets/datasets.py`.
30 +/// The Rust port skips `apply_chat_template` (no Jinja runtime here): `base_offset`
31 +/// and `prompt_len` are computed from the raw text. The resulting prompts are slightly
32 +/// shorter than Python's chat-template-formatted version (off by the chat scaffolding
33 +/// tokens, typically <20).
34 +pub fn load_sonnet_dataset(
35 + tokenizer: &TokenizerKind,
36 + dataset_path: Option<&str>,
37 + num_requests: usize,
38 + input_len: usize,
39 + output_len: usize,
40 + prefix_len: usize,
41 + seed: u64,
42 + request_id_prefix: &str,
43 +) -> Result<Vec<SampleRequest>> {
44 + let content = match dataset_path {
45 + Some(path) => std::fs::read_to_string(path)
46 + .map_err(|e| BenchError::Config(format!("Failed to read sonnet file '{path}': {e}")))?,
47 + None => BUILTIN_SONNET.to_string(),
48 + };
49 +
50 + // Match Python's f.readlines(): keep trailing newlines so that joining lines
51 + // reconstructs the original text without inserting extra separators.
52 + let lines: Vec<String> = content.split_inclusive('\n').map(|s| s.to_string()).collect();
53 +
54 + if lines.is_empty() {
55 + let src = dataset_path.unwrap_or("<built-in>");
56 + return Err(BenchError::Config(format!("Sonnet file '{src}' is empty")));
57 + }
58 +
59 + // Average tokens per line (used to estimate how many lines to draw).
60 + let mut total_tokens: usize = 0;
61 + for line in &lines {
62 + let ids = tokenizer.encode(line, false)?;
63 + total_tokens += ids.len();
64 + }
65 + let avg_len = total_tokens as f64 / lines.len() as f64;
66 + if avg_len <= 0.0 {
67 + return Err(BenchError::Config(
68 + "Sonnet lines tokenized to zero tokens on average".into(),
69 + ));
70 + }
71 +
72 + let base_ids = tokenizer.encode(BASE_PROMPT, false)?;
73 + let base_offset = base_ids.len();
74 + if input_len <= base_offset {
75 + return Err(BenchError::Config(format!(
76 + "--sonnet-input-len ({input_len}) must be larger than the base prompt length ({base_offset})"
77 + )));
78 + }
79 +
80 + let num_input_lines = ((input_len - base_offset) as f64 / avg_len).round() as i64;
81 + let num_prefix_lines =
82 + (((prefix_len as i64 - base_offset as i64) as f64) / avg_len).round().max(0.0) as i64;
83 + let num_input_lines = num_input_lines.max(0) as usize;
84 + let num_prefix_lines = (num_prefix_lines as usize).min(lines.len());
85 + let num_input_lines = num_input_lines.max(num_prefix_lines);
86 +
87 + let prefix_lines: &[String] = &lines[..num_prefix_lines];
88 + let extras_per_request = num_input_lines - num_prefix_lines;
89 +
90 + let mut rng = StdRng::seed_from_u64(seed);
91 + let mut samples = Vec::with_capacity(num_requests);
92 + let mut ind = 0usize;
93 + let mut attempts = 0usize;
94 + let max_attempts = num_requests.saturating_mul(20).max(1000);
95 +
96 + while samples.len() < num_requests {
97 + if attempts >= max_attempts {
98 + return Err(BenchError::Config(format!(
99 + "Could not assemble {num_requests} sonnet prompts under input_len={input_len} \
100 + after {attempts} attempts. Try increasing --sonnet-input-len."
101 + )));
102 + }
103 + attempts += 1;
104 +
105 + let mut prompt = String::with_capacity(BASE_PROMPT.len() + 256 * num_input_lines);
106 + prompt.push_str(BASE_PROMPT);
107 + for line in prefix_lines {
108 + prompt.push_str(line);
109 + }
110 + for _ in 0..extras_per_request {
111 + // random.choices with replacement — duplicates are allowed.
112 + let line = lines.choose(&mut rng).unwrap();
113 + prompt.push_str(line);
114 + }
115 +
116 + let prompt_ids = tokenizer.encode(&prompt, false)?;
117 + let prompt_len = prompt_ids.len();
118 + if prompt_len <= input_len {
119 + samples.push(SampleRequest {
120 + prompt: Arc::from(prompt),
121 + prompt_len,
122 + expected_output_len: output_len,
123 + request_id: Some(format!("{request_id_prefix}{ind}")),
124 + ..Default::default()
125 + });
126 + ind += 1;
127 + }
128 + }
129 +
130 + Ok(samples)
131 +}
rust/src/bench/src/datasets/sonnet.txt new
+518
@@ -0,0 +1,518 @@
1 +FROM fairest creatures we desire increase,
2 +That thereby beauty's rose might never die,
3 +But as the riper should by time decease,
4 +His tender heir might bear his memory:
5 +But thou, contracted to thine own bright eyes,
6 +Feed'st thy light'st flame with self-substantial fuel,
7 +Making a famine where abundance lies,
8 +Thyself thy foe, to thy sweet self too cruel.
9 +Thou that art now the world's fresh ornament
10 +And only herald to the gaudy spring,
11 +Within thine own bud buriest thy content
12 +And, tender churl, makest waste in niggarding.
13 +Pity the world, or else this glutton be,
14 +To eat the world's due, by the grave and thee.
15 +When forty winters shall beseige thy brow,
16 +And dig deep trenches in thy beauty's field,
17 +Thy youth's proud livery, so gazed on now,
18 +Will be a tatter'd weed, of small worth held:
19 +Then being ask'd where all thy beauty lies,
20 +Where all the treasure of thy lusty days,
21 +To say, within thine own deep-sunken eyes,
22 +Were an all-eating shame and thriftless praise.
23 +How much more praise deserved thy beauty's use,
24 +If thou couldst answer 'This fair child of mine
25 +Shall sum my count and make my old excuse,'
26 +Proving his beauty by succession thine!
27 +This were to be new made when thou art old,
28 +And see thy blood warm when thou feel'st it cold.
29 +Look in thy glass, and tell the face thou viewest
30 +Now is the time that face should form another;
31 +Whose fresh repair if now thou not renewest,
32 +Thou dost beguile the world, unbless some mother.
33 +For where is she so fair whose unear'd womb
34 +Disdains the tillage of thy husbandry?
35 +Or who is he so fond will be the tomb
36 +Of his self-love, to stop posterity?
37 +Thou art thy mother's glass, and she in thee
38 +Calls back the lovely April of her prime:
39 +So thou through windows of thine age shall see
40 +Despite of wrinkles this thy golden time.
41 +But if thou live, remember'd not to be,
42 +Die single, and thine image dies with thee.
43 +Unthrifty loveliness, why dost thou spend
44 +Upon thyself thy beauty's legacy?
45 +Nature's bequest gives nothing but doth lend,
46 +And being frank she lends to those are free.
47 +Then, beauteous niggard, why dost thou abuse
48 +The bounteous largess given thee to give?
49 +Profitless usurer, why dost thou use
50 +So great a sum of sums, yet canst not live?
51 +For having traffic with thyself alone,
52 +Thou of thyself thy sweet self dost deceive.
53 +Then how, when nature calls thee to be gone,
54 +What acceptable audit canst thou leave?
55 +Thy unused beauty must be tomb'd with thee,
56 +Which, used, lives th' executor to be.
57 +Those hours, that with gentle work did frame
58 +The lovely gaze where every eye doth dwell,
59 +Will play the tyrants to the very same
60 +And that unfair which fairly doth excel:
61 +For never-resting time leads summer on
62 +To hideous winter and confounds him there;
63 +Sap cheque'd with frost and lusty leaves quite gone,
64 +Beauty o'ersnow'd and bareness every where:
65 +Then, were not summer's distillation left,
66 +A liquid prisoner pent in walls of glass,
67 +Beauty's effect with beauty were bereft,
68 +Nor it nor no remembrance what it was:
69 +But flowers distill'd though they with winter meet,
70 +Leese but their show; their substance still lives sweet.
71 +Then let not winter's ragged hand deface
72 +In thee thy summer, ere thou be distill'd:
73 +Make sweet some vial; treasure thou some place
74 +With beauty's treasure, ere it be self-kill'd.
75 +That use is not forbidden usury,
76 +Which happies those that pay the willing loan;
77 +That's for thyself to breed another thee,
78 +Or ten times happier, be it ten for one;
79 +Ten times thyself were happier than thou art,
80 +If ten of thine ten times refigured thee:
81 +Then what could death do, if thou shouldst depart,
82 +Leaving thee living in posterity?
83 +Be not self-will'd, for thou art much too fair
84 +To be death's conquest and make worms thine heir.
85 +Lo! in the orient when the gracious light
86 +Lifts up his burning head, each under eye
87 +Doth homage to his new-appearing sight,
88 +Serving with looks his sacred majesty;
89 +And having climb'd the steep-up heavenly hill,
90 +Resembling strong youth in his middle age,
91 +yet mortal looks adore his beauty still,
92 +Attending on his golden pilgrimage;
93 +But when from highmost pitch, with weary car,
94 +Like feeble age, he reeleth from the day,
95 +The eyes, 'fore duteous, now converted are
96 +From his low tract and look another way:
97 +So thou, thyself out-going in thy noon,
98 +Unlook'd on diest, unless thou get a son.
99 +Music to hear, why hear'st thou music sadly?
100 +Sweets with sweets war not, joy delights in joy.
101 +Why lovest thou that which thou receivest not gladly,
102 +Or else receivest with pleasure thine annoy?
103 +If the true concord of well-tuned sounds,
104 +By unions married, do offend thine ear,
105 +They do but sweetly chide thee, who confounds
106 +In singleness the parts that thou shouldst bear.
107 +Mark how one string, sweet husband to another,
108 +Strikes each in each by mutual ordering,
109 +Resembling sire and child and happy mother
110 +Who all in one, one pleasing note do sing:
111 +Whose speechless song, being many, seeming one,
112 +Sings this to thee: 'thou single wilt prove none.'
113 +Is it for fear to wet a widow's eye
114 +That thou consumest thyself in single life?
115 +Ah! if thou issueless shalt hap to die.
116 +The world will wail thee, like a makeless wife;
117 +The world will be thy widow and still weep
118 +That thou no form of thee hast left behind,
119 +When every private widow well may keep
120 +By children's eyes her husband's shape in mind.
121 +Look, what an unthrift in the world doth spend
122 +Shifts but his place, for still the world enjoys it;
123 +But beauty's waste hath in the world an end,
124 +And kept unused, the user so destroys it.
125 +No love toward others in that bosom sits
126 +That on himself such murderous shame commits.
127 +For shame! deny that thou bear'st love to any,
128 +Who for thyself art so unprovident.
129 +Grant, if thou wilt, thou art beloved of many,
130 +But that thou none lovest is most evident;
131 +For thou art so possess'd with murderous hate
132 +That 'gainst thyself thou stick'st not to conspire.
133 +Seeking that beauteous roof to ruinate
134 +Which to repair should be thy chief desire.
135 +O, change thy thought, that I may change my mind!
136 +Shall hate be fairer lodged than gentle love?
137 +Be, as thy presence is, gracious and kind,
138 +Or to thyself at least kind-hearted prove:
139 +Make thee another self, for love of me,
140 +That beauty still may live in thine or thee.
141 +As fast as thou shalt wane, so fast thou growest
142 +In one of thine, from that which thou departest;
143 +And that fresh blood which youngly thou bestowest
144 +Thou mayst call thine when thou from youth convertest.
145 +Herein lives wisdom, beauty and increase:
146 +Without this, folly, age and cold decay:
147 +If all were minded so, the times should cease
148 +And threescore year would make the world away.
149 +Let those whom Nature hath not made for store,
150 +Harsh featureless and rude, barrenly perish:
151 +Look, whom she best endow'd she gave the more;
152 +Which bounteous gift thou shouldst in bounty cherish:
153 +She carved thee for her seal, and meant thereby
154 +Thou shouldst print more, not let that copy die.
155 +When I do count the clock that tells the time,
156 +And see the brave day sunk in hideous night;
157 +When I behold the violet past prime,
158 +And sable curls all silver'd o'er with white;
159 +When lofty trees I see barren of leaves
160 +Which erst from heat did canopy the herd,
161 +And summer's green all girded up in sheaves
162 +Borne on the bier with white and bristly beard,
163 +Then of thy beauty do I question make,
164 +That thou among the wastes of time must go,
165 +Since sweets and beauties do themselves forsake
166 +And die as fast as they see others grow;
167 +And nothing 'gainst Time's scythe can make defence
168 +Save breed, to brave him when he takes thee hence.
169 +O, that you were yourself! but, love, you are
170 +No longer yours than you yourself here live:
171 +Against this coming end you should prepare,
172 +And your sweet semblance to some other give.
173 +So should that beauty which you hold in lease
174 +Find no determination: then you were
175 +Yourself again after yourself's decease,
176 +When your sweet issue your sweet form should bear.
177 +Who lets so fair a house fall to decay,
178 +Which husbandry in honour might uphold
179 +Against the stormy gusts of winter's day
180 +And barren rage of death's eternal cold?
181 +O, none but unthrifts! Dear my love, you know
182 +You had a father: let your son say so.
183 +Not from the stars do I my judgment pluck;
184 +And yet methinks I have astronomy,
185 +But not to tell of good or evil luck,
186 +Of plagues, of dearths, or seasons' quality;
187 +Nor can I fortune to brief minutes tell,
188 +Pointing to each his thunder, rain and wind,
189 +Or say with princes if it shall go well,
190 +By oft predict that I in heaven find:
191 +But from thine eyes my knowledge I derive,
192 +And, constant stars, in them I read such art
193 +As truth and beauty shall together thrive,
194 +If from thyself to store thou wouldst convert;
195 +Or else of thee this I prognosticate:
196 +Thy end is truth's and beauty's doom and date.
197 +When I consider every thing that grows
198 +Holds in perfection but a little moment,
199 +That this huge stage presenteth nought but shows
200 +Whereon the stars in secret influence comment;
201 +When I perceive that men as plants increase,
202 +Cheered and cheque'd even by the self-same sky,
203 +Vaunt in their youthful sap, at height decrease,
204 +And wear their brave state out of memory;
205 +Then the conceit of this inconstant stay
206 +Sets you most rich in youth before my sight,
207 +Where wasteful Time debateth with Decay,
208 +To change your day of youth to sullied night;
209 +And all in war with Time for love of you,
210 +As he takes from you, I engraft you new.
211 +But wherefore do not you a mightier way
212 +Make war upon this bloody tyrant, Time?
213 +And fortify yourself in your decay
214 +With means more blessed than my barren rhyme?
215 +Now stand you on the top of happy hours,
216 +And many maiden gardens yet unset
217 +With virtuous wish would bear your living flowers,
218 +Much liker than your painted counterfeit:
219 +So should the lines of life that life repair,
220 +Which this, Time's pencil, or my pupil pen,
221 +Neither in inward worth nor outward fair,
222 +Can make you live yourself in eyes of men.
223 +To give away yourself keeps yourself still,
224 +And you must live, drawn by your own sweet skill.
225 +Who will believe my verse in time to come,
226 +If it were fill'd with your most high deserts?
227 +Though yet, heaven knows, it is but as a tomb
228 +Which hides your life and shows not half your parts.
229 +If I could write the beauty of your eyes
230 +And in fresh numbers number all your graces,
231 +The age to come would say 'This poet lies:
232 +Such heavenly touches ne'er touch'd earthly faces.'
233 +So should my papers yellow'd with their age
234 +Be scorn'd like old men of less truth than tongue,
235 +And your true rights be term'd a poet's rage
236 +And stretched metre of an antique song:
237 +But were some child of yours alive that time,
238 +You should live twice; in it and in my rhyme.
239 +Shall I compare thee to a summer's day?
240 +Thou art more lovely and more temperate:
241 +Rough winds do shake the darling buds of May,
242 +And summer's lease hath all too short a date:
243 +Sometime too hot the eye of heaven shines,
244 +And often is his gold complexion dimm'd;
245 +And every fair from fair sometime declines,
246 +By chance or nature's changing course untrimm'd;
247 +But thy eternal summer shall not fade
248 +Nor lose possession of that fair thou owest;
249 +Nor shall Death brag thou wander'st in his shade,
250 +When in eternal lines to time thou growest:
251 +So long as men can breathe or eyes can see,
252 +So long lives this and this gives life to thee.
253 +Devouring Time, blunt thou the lion's paws,
254 +And make the earth devour her own sweet brood;
255 +Pluck the keen teeth from the fierce tiger's jaws,
256 +And burn the long-lived phoenix in her blood;
257 +Make glad and sorry seasons as thou fleets,
258 +And do whate'er thou wilt, swift-footed Time,
259 +To the wide world and all her fading sweets;
260 +But I forbid thee one most heinous crime:
261 +O, carve not with thy hours my love's fair brow,
262 +Nor draw no lines there with thine antique pen;
263 +Him in thy course untainted do allow
264 +For beauty's pattern to succeeding men.
265 +Yet, do thy worst, old Time: despite thy wrong,
266 +My love shall in my verse ever live young.
267 +A woman's face with Nature's own hand painted
268 +Hast thou, the master-mistress of my passion;
269 +A woman's gentle heart, but not acquainted
270 +With shifting change, as is false women's fashion;
271 +An eye more bright than theirs, less false in rolling,
272 +Gilding the object whereupon it gazeth;
273 +A man in hue, all 'hues' in his controlling,
274 +Much steals men's eyes and women's souls amazeth.
275 +And for a woman wert thou first created;
276 +Till Nature, as she wrought thee, fell a-doting,
277 +And by addition me of thee defeated,
278 +By adding one thing to my purpose nothing.
279 +But since she prick'd thee out for women's pleasure,
280 +Mine be thy love and thy love's use their treasure.
281 +So is it not with me as with that Muse
282 +Stirr'd by a painted beauty to his verse,
283 +Who heaven itself for ornament doth use
284 +And every fair with his fair doth rehearse
285 +Making a couplement of proud compare,
286 +With sun and moon, with earth and sea's rich gems,
287 +With April's first-born flowers, and all things rare
288 +That heaven's air in this huge rondure hems.
289 +O' let me, true in love, but truly write,
290 +And then believe me, my love is as fair
291 +As any mother's child, though not so bright
292 +As those gold candles fix'd in heaven's air:
293 +Let them say more than like of hearsay well;
294 +I will not praise that purpose not to sell.
295 +My glass shall not persuade me I am old,
296 +So long as youth and thou are of one date;
297 +But when in thee time's furrows I behold,
298 +Then look I death my days should expiate.
299 +For all that beauty that doth cover thee
300 +Is but the seemly raiment of my heart,
301 +Which in thy breast doth live, as thine in me:
302 +How can I then be elder than thou art?
303 +O, therefore, love, be of thyself so wary
304 +As I, not for myself, but for thee will;
305 +Bearing thy heart, which I will keep so chary
306 +As tender nurse her babe from faring ill.
307 +Presume not on thy heart when mine is slain;
308 +Thou gavest me thine, not to give back again.
309 +As an unperfect actor on the stage
310 +Who with his fear is put besides his part,
311 +Or some fierce thing replete with too much rage,
312 +Whose strength's abundance weakens his own heart.
313 +So I, for fear of trust, forget to say
314 +The perfect ceremony of love's rite,
315 +And in mine own love's strength seem to decay,
316 +O'ercharged with burden of mine own love's might.
317 +O, let my books be then the eloquence
318 +And dumb presagers of my speaking breast,
319 +Who plead for love and look for recompense
320 +More than that tongue that more hath more express'd.
321 +O, learn to read what silent love hath writ:
322 +To hear with eyes belongs to love's fine wit.
323 +Mine eye hath play'd the painter and hath stell'd
324 +Thy beauty's form in table of my heart;
325 +My body is the frame wherein 'tis held,
326 +And perspective it is the painter's art.
327 +For through the painter must you see his skill,
328 +To find where your true image pictured lies;
329 +Which in my bosom's shop is hanging still,
330 +That hath his windows glazed with thine eyes.
331 +Now see what good turns eyes for eyes have done:
332 +Mine eyes have drawn thy shape, and thine for me
333 +Are windows to my breast, where-through the sun
334 +Delights to peep, to gaze therein on thee;
335 +Yet eyes this cunning want to grace their art;
336 +They draw but what they see, know not the heart.
337 +Let those who are in favour with their stars
338 +Of public honour and proud titles boast,
339 +Whilst I, whom fortune of such triumph bars,
340 +Unlook'd for joy in that I honour most.
341 +Great princes' favourites their fair leaves spread
342 +But as the marigold at the sun's eye,
343 +And in themselves their pride lies buried,
344 +For at a frown they in their glory die.
345 +The painful warrior famoused for fight,
346 +After a thousand victories once foil'd,
347 +Is from the book of honour razed quite,
348 +And all the rest forgot for which he toil'd:
349 +Then happy I, that love and am beloved
350 +Where I may not remove nor be removed.
351 +Lord of my love, to whom in vassalage
352 +Thy merit hath my duty strongly knit,
353 +To thee I send this written embassage,
354 +To witness duty, not to show my wit:
355 +Duty so great, which wit so poor as mine
356 +May make seem bare, in wanting words to show it,
357 +But that I hope some good conceit of thine
358 +In thy soul's thought, all naked, will bestow it;
359 +Till whatsoever star that guides my moving
360 +Points on me graciously with fair aspect
361 +And puts apparel on my tatter'd loving,
362 +To show me worthy of thy sweet respect:
363 +Then may I dare to boast how I do love thee;
364 +Till then not show my head where thou mayst prove me.
365 +Weary with toil, I haste me to my bed,
366 +The dear repose for limbs with travel tired;
367 +But then begins a journey in my head,
368 +To work my mind, when body's work's expired:
369 +For then my thoughts, from far where I abide,
370 +Intend a zealous pilgrimage to thee,
371 +And keep my drooping eyelids open wide,
372 +Looking on darkness which the blind do see
373 +Save that my soul's imaginary sight
374 +Presents thy shadow to my sightless view,
375 +Which, like a jewel hung in ghastly night,
376 +Makes black night beauteous and her old face new.
377 +Lo! thus, by day my limbs, by night my mind,
378 +For thee and for myself no quiet find.
379 +How can I then return in happy plight,
380 +That am debarr'd the benefit of rest?
381 +When day's oppression is not eased by night,
382 +But day by night, and night by day, oppress'd?
383 +And each, though enemies to either's reign,
384 +Do in consent shake hands to torture me;
385 +The one by toil, the other to complain
386 +How far I toil, still farther off from thee.
387 +I tell the day, to please them thou art bright
388 +And dost him grace when clouds do blot the heaven:
389 +So flatter I the swart-complexion'd night,
390 +When sparkling stars twire not thou gild'st the even.
391 +But day doth daily draw my sorrows longer
392 +And night doth nightly make grief's strength seem stronger.
393 +When, in disgrace with fortune and men's eyes,
394 +I all alone beweep my outcast state
395 +And trouble deal heaven with my bootless cries
396 +And look upon myself and curse my fate,
397 +Wishing me like to one more rich in hope,
398 +Featured like him, like him with friends possess'd,
399 +Desiring this man's art and that man's scope,
400 +With what I most enjoy contented least;
401 +Yet in these thoughts myself almost despising,
402 +Haply I think on thee, and then my state,
403 +Like to the lark at break of day arising
404 +From sullen earth, sings hymns at heaven's gate;
405 +For thy sweet love remember'd such wealth brings
406 +That then I scorn to change my state with kings.
407 +When to the sessions of sweet silent thought
408 +I summon up remembrance of things past,
409 +I sigh the lack of many a thing I sought,
410 +And with old woes new wail my dear time's waste:
411 +Then can I drown an eye, unused to flow,
412 +For precious friends hid in death's dateless night,
413 +And weep afresh love's long since cancell'd woe,
414 +And moan the expense of many a vanish'd sight:
415 +Then can I grieve at grievances foregone,
416 +And heavily from woe to woe tell o'er
417 +The sad account of fore-bemoaned moan,
418 +Which I new pay as if not paid before.
419 +But if the while I think on thee, dear friend,
420 +All losses are restored and sorrows end.
421 +Thy bosom is endeared with all hearts,
422 +Which I by lacking have supposed dead,
423 +And there reigns love and all love's loving parts,
424 +And all those friends which I thought buried.
425 +How many a holy and obsequious tear
426 +Hath dear religious love stol'n from mine eye
427 +As interest of the dead, which now appear
428 +But things removed that hidden in thee lie!
429 +Thou art the grave where buried love doth live,
430 +Hung with the trophies of my lovers gone,
431 +Who all their parts of me to thee did give;
432 +That due of many now is thine alone:
433 +Their images I loved I view in thee,
434 +And thou, all they, hast all the all of me.
435 +If thou survive my well-contented day,
436 +When that churl Death my bones with dust shall cover,
437 +And shalt by fortune once more re-survey
438 +These poor rude lines of thy deceased lover,
439 +Compare them with the bettering of the time,
440 +And though they be outstripp'd by every pen,
441 +Reserve them for my love, not for their rhyme,
442 +Exceeded by the height of happier men.
443 +O, then vouchsafe me but this loving thought:
444 +'Had my friend's Muse grown with this growing age,
445 +A dearer birth than this his love had brought,
446 +To march in ranks of better equipage:
447 +But since he died and poets better prove,
448 +Theirs for their style I'll read, his for his love.'
449 +Full many a glorious morning have I seen
450 +Flatter the mountain-tops with sovereign eye,
451 +Kissing with golden face the meadows green,
452 +Gilding pale streams with heavenly alchemy;
453 +Anon permit the basest clouds to ride
454 +With ugly rack on his celestial face,
455 +And from the forlorn world his visage hide,
456 +Stealing unseen to west with this disgrace:
457 +Even so my sun one early morn did shine
458 +With all triumphant splendor on my brow;
459 +But out, alack! he was but one hour mine;
460 +The region cloud hath mask'd him from me now.
461 +Yet him for this my love no whit disdaineth;
462 +Suns of the world may stain when heaven's sun staineth.
463 +Why didst thou promise such a beauteous day,
464 +And make me travel forth without my cloak,
465 +To let base clouds o'ertake me in my way,
466 +Hiding thy bravery in their rotten smoke?
467 +'Tis not enough that through the cloud thou break,
468 +To dry the rain on my storm-beaten face,
469 +For no man well of such a salve can speak
470 +That heals the wound and cures not the disgrace:
471 +Nor can thy shame give physic to my grief;
472 +Though thou repent, yet I have still the loss:
473 +The offender's sorrow lends but weak relief
474 +To him that bears the strong offence's cross.
475 +Ah! but those tears are pearl which thy love sheds,
476 +And they are rich and ransom all ill deeds.
477 +No more be grieved at that which thou hast done:
478 +Roses have thorns, and silver fountains mud;
479 +Clouds and eclipses stain both moon and sun,
480 +And loathsome canker lives in sweetest bud.
481 +All men make faults, and even I in this,
482 +Authorizing thy trespass with compare,
483 +Myself corrupting, salving thy amiss,
484 +Excusing thy sins more than thy sins are;
485 +For to thy sensual fault I bring in sense--
486 +Thy adverse party is thy advocate--
487 +And 'gainst myself a lawful plea commence:
488 +Such civil war is in my love and hate
489 +That I an accessary needs must be
490 +To that sweet thief which sourly robs from me.
491 +Let me confess that we two must be twain,
492 +Although our undivided loves are one:
493 +So shall those blots that do with me remain
494 +Without thy help by me be borne alone.
495 +In our two loves there is but one respect,
496 +Though in our lives a separable spite,
497 +Which though it alter not love's sole effect,
498 +Yet doth it steal sweet hours from love's delight.
499 +I may not evermore acknowledge thee,
500 +Lest my bewailed guilt should do thee shame,
501 +Nor thou with public kindness honour me,
502 +Unless thou take that honour from thy name:
503 +But do not so; I love thee in such sort
504 +As, thou being mine, mine is thy good report.
505 +As a decrepit father takes delight
506 +To see his active child do deeds of youth,
507 +So I, made lame by fortune's dearest spite,
508 +Take all my comfort of thy worth and truth.
509 +For whether beauty, birth, or wealth, or wit,
510 +Or any of these all, or all, or more,
511 +Entitled in thy parts do crowned sit,
512 +I make my love engrafted to this store:
513 +So then I am not lame, poor, nor despised,
514 +Whilst that this shadow doth such substance give
515 +That I in thy abundance am sufficed
516 +And by a part of all thy glory live.
517 +Look, what is best, that best I wish in thee:
518 +This wish I have; then ten times happy me!
\ No newline at end of file
rust/src/bench/src/datasets/speed_bench.rs new
+303
@@ -0,0 +1,303 @@
1 +// SPDX-License-Identifier: Apache-2.0
2 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +use std::sync::Arc;
5 +
6 +use rand::rngs::StdRng;
7 +use rand::seq::SliceRandom;
8 +use rand::{Rng, SeedableRng};
9 +
10 +use super::SampleRequest;
11 +use crate::cli::SpeedBenchConfig;
12 +use crate::error::{BenchError, Result};
13 +use crate::tokenizer::TokenizerKind;
14 +
15 +/// Marker text for masked entries that need external fetch.
16 +const MASKED_PREFIX: &str = "FULL BENCHMARK DATA SHOULD BE FETCHED";
17 +
18 +/// Cache directory for downloaded SPEED-Bench datasets.
19 +fn cache_dir() -> std::path::PathBuf {
20 + dirs::cache_dir()
21 + .unwrap_or_else(|| std::path::PathBuf::from("/tmp"))
22 + .join("vllm-bench")
23 + .join("datasets")
24 +}
25 +
26 +/// Download SPEED-Bench dataset from HuggingFace datasets-server API.
27 +/// Results are cached as JSON locally for subsequent runs.
28 +pub fn download_speed_bench(config: SpeedBenchConfig) -> Result<String> {
29 + let config_name = config.as_str();
30 +
31 + let dir = cache_dir();
32 + std::fs::create_dir_all(&dir)?;
33 + let cache_path = dir.join(format!("speed-bench-{config_name}.json"));
34 +
35 + // Return cached file if it exists
36 + if cache_path.exists() {
37 + let path_str = cache_path.to_string_lossy().to_string();
38 + println!("SPEED-Bench ({config_name}) cached: {path_str}");
39 + return Ok(path_str);
40 + }
41 +
42 + println!("Downloading SPEED-Bench ({config_name}) from HuggingFace datasets-server...");
43 +
44 + let client = reqwest::blocking::Client::builder()
45 + .timeout(std::time::Duration::from_secs(120))
46 + .build()
47 + .map_err(|e| BenchError::Config(format!("Failed to build HTTP client: {e}")))?;
48 +
49 + let mut all_rows: Vec<serde_json::Value> = Vec::new();
50 + let mut offset = 0usize;
51 + let page_size = 100usize;
52 +
53 + loop {
54 + let url = format!(
55 + "https://datasets-server.huggingface.co/rows\
56 + ?dataset=nvidia/SPEED-Bench\
57 + &config={config_name}\
58 + &split=test\
59 + &offset={offset}\
60 + &length={page_size}"
61 + );
62 +
63 + // Retry on transient errors (502, 503, timeouts)
64 + let max_retries = 3;
65 + let mut data: Option<serde_json::Value> = None;
66 + for attempt in 0..=max_retries {
67 + let resp = match client.get(&url).send() {
68 + Ok(r) => r,
69 + Err(e) => {
70 + if attempt < max_retries {
71 + std::thread::sleep(std::time::Duration::from_secs(
72 + 2 * (attempt as u64 + 1),
73 + ));
74 + continue;
75 + }
76 + return Err(BenchError::Config(format!(
77 + "SPEED-Bench download failed after {max_retries} retries: {e}"
78 + )));
79 + }
80 + };
81 +
82 + if resp.status().is_server_error() && attempt < max_retries {
83 + std::thread::sleep(std::time::Duration::from_secs(2 * (attempt as u64 + 1)));
84 + continue;
85 + }
86 +
87 + if !resp.status().is_success() {
88 + return Err(BenchError::Config(format!(
89 + "SPEED-Bench API returned HTTP {}",
90 + resp.status()
91 + )));
92 + }
93 +
94 + data = Some(resp.json().map_err(|e| {
95 + BenchError::Config(format!("Failed to parse SPEED-Bench API response: {e}"))
96 + })?);
97 + break;
98 + }
99 +
100 + let data = data.unwrap();
101 +
102 + let rows = data["rows"]
103 + .as_array()
104 + .ok_or_else(|| BenchError::Config("No 'rows' in API response".into()))?;
105 +
106 + if rows.is_empty() {
107 + break;
108 + }
109 +
110 + for row in rows {
111 + if let Some(row_data) = row.get("row") {
112 + all_rows.push(row_data.clone());
113 + }
114 + }
115 +
116 + let fetched = rows.len();
117 + offset += fetched;
118 +
119 + // Print progress
120 + let total = data["num_rows_total"].as_u64().unwrap_or(0);
121 + eprint!("\r Fetched {offset}/{total} rows...");
122 +
123 + if fetched < page_size {
124 + break;
125 + }
126 + }
127 + eprintln!(); // newline after progress
128 +
129 + if all_rows.is_empty() {
130 + return Err(BenchError::Config(
131 + "SPEED-Bench download returned no rows".into(),
132 + ));
133 + }
134 +
135 + // Save to cache
136 + let json_str = serde_json::to_string(&all_rows)?;
137 + std::fs::write(&cache_path, &json_str)?;
138 +
139 + let path_str = cache_path.to_string_lossy().to_string();
140 + println!(
141 + "SPEED-Bench ({config_name}): {} rows saved to {path_str}",
142 + all_rows.len()
143 + );
144 + Ok(path_str)
145 +}
146 +
147 +/// Load SPEED-Bench dataset and convert to SampleRequests.
148 +///
149 +/// Filters out masked entries and optionally filters by category.
150 +/// Requires an output length override since SPEED-Bench has no reference outputs.
151 +pub fn load_speed_bench_dataset(
152 + tokenizer: &TokenizerKind,
153 + dataset_path: &str,
154 + num_requests: usize,
155 + output_len: usize,
156 + seed: u64,
157 + request_id_prefix: &str,
158 + category_filter: Option<&str>,
159 + no_oversample: bool,
160 + disable_shuffle: bool,
161 + max_input_len: Option<usize>,
162 +) -> Result<Vec<SampleRequest>> {
163 + let content = std::fs::read_to_string(dataset_path).map_err(|e| {
164 + BenchError::Config(format!(
165 + "Failed to read SPEED-Bench file '{dataset_path}': {e}"
166 + ))
167 + })?;
168 +
169 + let entries: Vec<serde_json::Value> = serde_json::from_str(&content)
170 + .map_err(|e| BenchError::Config(format!("Invalid JSON in SPEED-Bench file: {e}")))?;
171 +
172 + // Filter entries
173 + let mut filtered: Vec<&serde_json::Value> = entries
174 + .iter()
175 + .filter(|entry| {
176 + // Must have turns array with at least one non-empty entry
177 + let turns = match entry.get("turns").and_then(|t| t.as_array()) {
178 + Some(t) if !t.is_empty() => t,
179 + _ => return false,
180 + };
181 +
182 + // Skip masked entries
183 + let first_turn = turns[0].as_str().unwrap_or("");
184 + if first_turn.starts_with(MASKED_PREFIX) || first_turn.is_empty() {
185 + return false;
186 + }
187 +
188 + // Single-turn only: skip multi-turn entries
189 + let is_multiturn = entry.get("multiturn").and_then(|m| m.as_bool()).unwrap_or(false);
190 + if is_multiturn {
191 + return false;
192 + }
193 +
194 + // Category filter
195 + if let Some(cat) = category_filter {
196 + let entry_cat = entry.get("category").and_then(|c| c.as_str()).unwrap_or("");
197 + if entry_cat != cat {
198 + return false;
199 + }
200 + }
201 +
202 + true
203 + })
204 + .collect();
205 +
206 + if filtered.is_empty() {
207 + let cat_msg = category_filter.map(|c| format!(" with category '{c}'")).unwrap_or_default();
208 + return Err(BenchError::Config(format!(
209 + "No valid single-turn entries in SPEED-Bench{cat_msg}. \
210 + Try a different --speed-bench-config or remove --speed-bench-category filter."
211 + )));
212 + }
213 +
214 + // Shuffle
215 + let mut rng = StdRng::seed_from_u64(seed);
216 + if !disable_shuffle {
217 + filtered.shuffle(&mut rng);
218 + }
219 +
220 + // Build SampleRequests
221 + let mut samples = Vec::new();
222 + let mut idx = 0;
223 +
224 + for entry in &filtered {
225 + if samples.len() >= num_requests {
226 + break;
227 + }
228 +
229 + let turns = entry["turns"].as_array().unwrap();
230 + let prompt = turns[0].as_str().unwrap_or("");
231 +
232 + // Tokenize to get prompt length
233 + let prompt_ids = tokenizer.encode(prompt, false)?;
234 + let prompt_len = prompt_ids.len();
235 +
236 + if prompt_len < 4 {
237 + continue;
238 + }
239 +
240 + // Truncate if max_input_len is set
241 + let (final_prompt, final_len) = if let Some(max_len) = max_input_len {
242 + if prompt_len > max_len {
243 + let truncated_ids = &prompt_ids[..max_len];
244 + let truncated_text = tokenizer.decode(truncated_ids, true)?;
245 + (Arc::from(truncated_text.as_str()), max_len)
246 + } else {
247 + (Arc::from(prompt), prompt_len)
248 + }
249 + } else {
250 + (Arc::from(prompt), prompt_len)
251 + };
252 +
253 + samples.push(SampleRequest {
254 + prompt: final_prompt,
255 + prompt_len: final_len,
256 + expected_output_len: output_len,
257 + request_id: Some(format!("{request_id_prefix}{idx}")),
258 + ..Default::default()
259 + });
260 + idx += 1;
261 + }
262 +
263 + // Oversample if needed
264 + if samples.len() < num_requests {
265 + if no_oversample {
266 + println!(
267 + "Skipping oversampling. Total samples: {} (requested: {num_requests})",
268 + samples.len()
269 + );
270 + } else if !samples.is_empty() {
271 + let original_len = samples.len();
272 + let needed = num_requests - original_len;
273 + for i in 0..needed {
274 + let mut req = samples[rng.random_range(0..original_len)].clone();
275 + req.request_id = Some(format!("{request_id_prefix}{}", original_len + i));
276 + samples.push(req);
277 + }
278 + println!(
279 + "Oversampled SPEED-Bench from {original_len} to {} total samples.",
280 + samples.len()
281 + );
282 + }
283 + }
284 +
285 + if samples.is_empty() {
286 + return Err(BenchError::Config(
287 + "No valid samples after filtering SPEED-Bench dataset.".into(),
288 + ));
289 + }
290 +
291 + // Print category distribution
292 + let mut cat_counts: std::collections::HashMap<&str, usize> = std::collections::HashMap::new();
293 + for entry in &filtered[..filtered.len().min(samples.len())] {
294 + let cat = entry.get("category").and_then(|c| c.as_str()).unwrap_or("unknown");
295 + *cat_counts.entry(cat).or_insert(0) += 1;
296 + }
297 + let mut cats: Vec<_> = cat_counts.into_iter().collect();
298 + cats.sort_by_key(|b| std::cmp::Reverse(b.1));
299 + let cat_str: Vec<String> = cats.iter().map(|(k, v)| format!("{k}:{v}")).collect();
300 + println!("SPEED-Bench categories: {}", cat_str.join(", "));
301 +
302 + Ok(samples)
303 +}
rust/src/bench/src/error.rs new
+37
@@ -0,0 +1,37 @@
1 +// SPDX-License-Identifier: Apache-2.0
2 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +use thiserror::Error;
5 +
6 +#[derive(Error, Debug)]
7 +pub enum BenchError {
8 + #[error("HTTP request failed: {0}")]
9 + Http(#[from] reqwest::Error),
10 +
11 + #[error("JSON error: {0}")]
12 + Json(#[from] serde_json::Error),
13 +
14 + #[error("Tokenizer error: {0}")]
15 + Tokenizer(String),
16 +
17 + /// The server's /tokenize//detokenize endpoint is not usable (4xx status:
18 + /// not exposed, or rejected by a gateway such as LLM-d/EPP that returns
19 + /// 400 instead of 404). Callers treat this as "skip verification", unlike
20 + /// `Tokenizer` errors which are genuine failures.
21 + #[error("tokenize endpoint unavailable: {0}")]
22 + TokenizeUnavailable(String),
23 +
24 + #[error("Configuration error: {0}")]
25 + Config(String),
26 +
27 + #[error("Endpoint not ready after {0}s: {1}")]
28 + EndpointTimeout(u64, String),
29 +
30 + #[error("Backend error: {0}")]
31 + Backend(String),
32 +
33 + #[error("IO error: {0}")]
34 + Io(#[from] std::io::Error),
35 +}
36 +
37 +pub type Result<T> = std::result::Result<T, BenchError>;
rust/src/bench/src/hub.rs new
+51
@@ -0,0 +1,51 @@
1 +// SPDX-License-Identifier: Apache-2.0
2 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +//! Sync facade over the async `hf_hub` API.
5 +//!
6 +//! The workspace bans rustls (`rust/deny.toml`), but hf-hub's sync `ureq`
7 +//! backend unconditionally pulls ureq's default rustls feature. So we use the
8 +//! reqwest/native-tls tokio API instead, and bridge blocking callers (dataset
9 +//! loaders, tokenizer fallback in rayon threads) by running each download on a
10 +//! dedicated thread with its own single-threaded runtime.
11 +
12 +use std::path::PathBuf;
13 +
14 +/// A handle to a HuggingFace Hub repo, downloading via hf-hub's on-disk cache.
15 +pub struct HubRepo {
16 + repo: hf_hub::Repo,
17 +}
18 +
19 +impl HubRepo {
20 + pub fn model(model_id: String) -> Self {
21 + Self {
22 + repo: hf_hub::Repo::model(model_id),
23 + }
24 + }
25 +
26 + pub fn dataset(repo_id: String) -> Self {
27 + Self {
28 + repo: hf_hub::Repo::dataset(repo_id),
29 + }
30 + }
31 +
32 + /// Download (or fetch from cache) a single file from the repo.
33 + /// Auth is handled by hf-hub via HF_TOKEN / the cached login token.
34 + pub fn get(&self, filename: &str) -> Result<PathBuf, String> {
35 + let repo = self.repo.clone();
36 + let filename = filename.to_string();
37 + std::thread::spawn(move || {
38 + let rt = tokio::runtime::Builder::new_current_thread()
39 + .enable_all()
40 + .build()
41 + .map_err(|e| format!("Failed to build download runtime: {e}"))?;
42 + rt.block_on(async move {
43 + let api = hf_hub::api::tokio::Api::new()
44 + .map_err(|e| format!("Failed to init HF API: {e}"))?;
45 + api.repo(repo).get(&filename).await.map_err(|e| format!("{e}"))
46 + })
47 + })
48 + .join()
49 + .map_err(|_| "HF Hub download thread panicked".to_string())?
50 + }
51 +}
rust/src/bench/src/main.rs new
+92
@@ -0,0 +1,92 @@
1 +// SPDX-License-Identifier: Apache-2.0
2 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +mod backends;
5 +mod benchmark;
6 +mod cli;
7 +mod compare;
8 +mod config;
9 +mod datasets;
10 +mod error;
11 +mod hub;
12 +mod metrics;
13 +mod multi_run;
14 +mod multi_turn;
15 +mod output;
16 +mod rate_control;
17 +mod ready_checker;
18 +mod sweep;
19 +mod tiktoken;
20 +mod tokenizer;
21 +
22 +#[cfg(not(target_env = "msvc"))]
23 +#[global_allocator]
24 +static GLOBAL: mimalloc::MiMalloc = mimalloc::MiMalloc;
25 +
26 +use anyhow::Context;
27 +use clap::Parser;
28 +use cli::Cli;
29 +use config::BenchConfig;
30 +
31 +fn main() -> anyhow::Result<()> {
32 + // Raise the open-file soft limit to the hard limit. High-concurrency
33 + // benchmarks (1024+ requests) easily exceed the default 1024 fd soft limit.
34 + if let Ok(new) = rlimit::increase_nofile_limit(u64::MAX)
35 + && new > 1024
36 + {
37 + eprintln!("Open-file limit: {new}");
38 + }
39 +
40 + let cli = Cli::parse();
41 +
42 + // --- Compare mode: no server needed, just diff two JSON files ---
43 + if let Some(ref files) = cli.compare {
44 + return compare::compare_results(&files[0], &files[1]).context("Comparison failed");
45 + }
46 +
47 + let config = BenchConfig::from_cli(&cli).context("Configuration error")?;
48 +
49 + let runtime = tokio::runtime::Builder::new_multi_thread()
50 + .enable_all()
51 + .build()
52 + .expect("Failed to build tokio runtime");
53 +
54 + runtime
55 + .block_on(async {
56 + if config.multi_turn {
57 + if let Some(ref sweep_mc) = cli.sweep_max_concurrency {
58 + // --- Sweep over concurrency in multi-turn mode ---
59 + let values = sweep::parse_concurrency_values(sweep_mc)
60 + .context("Invalid --sweep-max-concurrency")?;
61 + sweep::run_multi_turn_concurrency_sweep(
62 + &config,
63 + &values,
64 + cli.sweep_num_prompts_factor,
65 + )
66 + .await?;
67 + } else {
68 + // --- Single multi-turn conversation benchmark ---
69 + multi_turn::run_multi_turn_benchmark(&config).await?;
70 + }
71 + } else if let Some(ref sweep_mc) = cli.sweep_max_concurrency {
72 + // --- Sweep over max-concurrency ---
73 + let values = sweep::parse_concurrency_values(sweep_mc)
74 + .context("Invalid --sweep-max-concurrency")?;
75 + sweep::run_concurrency_sweep(&config, &values, cli.sweep_num_prompts_factor)
76 + .await?;
77 + } else if let Some(ref sweep_rate) = cli.sweep_request_rate {
78 + // --- Sweep over request-rate ---
79 + let values =
80 + sweep::parse_rate_values(sweep_rate).context("Invalid --sweep-request-rate")?;
81 + sweep::run_rate_sweep(&config, &values).await?;
82 + } else if cli.num_runs > 1 {
83 + // --- Multi-run with statistical aggregation ---
84 + multi_run::run_multi(&config, cli.num_runs).await?;
85 + } else {
86 + // --- Normal single benchmark ---
87 + benchmark::run_benchmark(&config).await?;
88 + }
89 + anyhow::Ok(())
90 + })
91 + .context("Benchmark failed")
92 +}
rust/src/bench/src/metrics/calculator.rs new
+503
@@ -0,0 +1,503 @@
1 +// SPDX-License-Identifier: Apache-2.0
2 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +use crate::backends::RequestFuncOutput;
5 +use crate::config::GoodputConfig;
6 +use crate::datasets::SampleRequest;
7 +use crate::metrics::{BenchmarkMetrics, MultiTurnMetrics};
8 +use crate::multi_turn::ConversationOutput;
9 +
10 +/// Calculate benchmark metrics from request outputs.
11 +///
12 +/// Mirrors Python's `calculate_metrics()` from serve.py:392-599.
13 +pub fn calculate_metrics(
14 + input_requests: &[SampleRequest],
15 + outputs: &[RequestFuncOutput],
16 + dur_s: f64,
17 + selected_percentiles: &[f64],
18 + _has_tokenizer: bool,
19 + goodput_config: &GoodputConfig,
20 +) -> (BenchmarkMetrics, Vec<usize>) {
21 + let mut actual_output_lens: Vec<usize> = Vec::with_capacity(outputs.len());
22 + let mut total_input: usize = 0;
23 + let mut completed: usize = 0;
24 + let mut itls: Vec<f64> = Vec::new();
25 + let mut tpots: Vec<f64> = Vec::new();
26 + // Per-request TPOT for goodput SLO checking (parallel to ttfts/e2els).
27 + let mut all_tpots: Vec<f64> = Vec::new();
28 + let mut ttfts: Vec<f64> = Vec::new();
29 + let mut e2els: Vec<f64> = Vec::new();
30 +
31 + for (i, output) in outputs.iter().enumerate() {
32 + if output.success {
33 + let output_len = if output.output_tokens > 0 {
34 + output.output_tokens
35 + } else {
36 + // Fallback: first token + ITL entries.
37 + // Python re-encodes generated_text when tokenizer is available,
38 + // but with vLLM's stream_options.include_usage=true,
39 + // output_tokens is always set so this path is rarely hit.
40 + 1 + output.itl.len()
41 + };
42 +
43 + actual_output_lens.push(output_len);
44 + total_input += input_requests[i].prompt_len;
45 +
46 + if output_len > 1 {
47 + let latency_minus_ttft = output.latency - output.ttft;
48 + let tpot = latency_minus_ttft / (output_len as f64 - 1.0);
49 + tpots.push(tpot);
50 + all_tpots.push(tpot);
51 + } else {
52 + all_tpots.push(0.0);
53 + }
54 +
55 + itls.extend_from_slice(&output.itl);
56 + ttfts.push(output.ttft);
57 + e2els.push(output.latency);
58 + completed += 1;
59 + } else {
60 + actual_output_lens.push(0);
61 + }
62 + }
63 +
64 + let failed = outputs.len() - completed;
65 +
66 + // Print failed request errors (capped to 10)
67 + let failed_outputs: Vec<&RequestFuncOutput> = outputs.iter().filter(|o| !o.success).collect();
68 + if !failed_outputs.is_empty() {
69 + eprintln!("Failed requests during benchmark run detected (capping to 10):");
70 + for (i, err) in failed_outputs.iter().take(10).enumerate() {
71 + eprintln!("Error {i}: {}", err.error);
72 + }
73 + }
74 +
75 + // Calculate max output tokens per second and max concurrent requests
76 + let mut max_output_tokens_per_s = 0.0_f64;
77 + let mut max_concurrent_requests: usize = 0;
78 +
79 + let successful_outputs: Vec<&RequestFuncOutput> =
80 + outputs.iter().filter(|o| o.success).collect();
81 +
82 + if !successful_outputs.is_empty() {
83 + let min_start_time =
84 + successful_outputs.iter().map(|o| o.start_time).fold(f64::INFINITY, f64::min);
85 + let max_end_time = successful_outputs
86 + .iter()
87 + .map(|o| o.start_time + o.latency)
88 + .fold(f64::NEG_INFINITY, f64::max);
89 +
90 + let raw_duration = (max_end_time - min_start_time).ceil() as usize + 1;
91 + // Cap at 24 hours to prevent OOM from corrupted timing data
92 + let duration_seconds = raw_duration.min(86_400);
93 + let mut tokens_per_second = vec![0.0_f64; duration_seconds];
94 + let mut concurrent_per_second = vec![0usize; duration_seconds];
95 +
96 + for output in &successful_outputs {
97 + // Calculate token generation timestamps
98 + let mut token_times = vec![output.start_time + output.ttft];
99 + let mut current_time = token_times[0];
100 + for itl_value in &output.itl {
101 + current_time += itl_value;
102 + token_times.push(current_time);
103 + }
104 +
105 + // Add tokens to second buckets
106 + for token_time in &token_times {
107 + let bucket = (token_time - min_start_time) as usize;
108 + if bucket < duration_seconds {
109 + tokens_per_second[bucket] += 1.0;
110 + }
111 + }
112 +
113 + // Track concurrent requests
114 + let start_second = (output.start_time - min_start_time) as usize;
115 + let end_second = ((output.start_time + output.latency) - min_start_time) as usize;
116 + for slot in concurrent_per_second
117 + .iter_mut()
118 + .take(end_second.min(duration_seconds - 1) + 1)
119 + .skip(start_second)
120 + {
121 + *slot += 1;
122 + }
123 + }
124 +
125 + max_output_tokens_per_s = tokens_per_second.iter().cloned().fold(0.0_f64, f64::max);
126 + max_concurrent_requests = *concurrent_per_second.iter().max().unwrap_or(&0);
127 + }
128 +
129 + let total_output: usize = actual_output_lens.iter().sum();
130 +
131 + // Compute goodput: count requests meeting ALL specified SLOs.
132 + // Mirrors Python serve.py:458-481.
133 + let good_completed = if !goodput_config.is_empty() {
134 + let mut good = 0usize;
135 + // ttfts, all_tpots, e2els are parallel (one per successful request)
136 + for i in 0..ttfts.len() {
137 + let mut is_good = true;
138 + if let Some(slo_ms) = goodput_config.ttft_ms
139 + && ttfts[i] * 1000.0 > slo_ms
140 + {
141 + is_good = false;
142 + }
143 + if let Some(slo_ms) = goodput_config.tpot_ms
144 + && all_tpots[i] * 1000.0 > slo_ms
145 + {
146 + is_good = false;
147 + }
148 + if let Some(slo_ms) = goodput_config.e2el_ms
149 + && e2els[i] * 1000.0 > slo_ms
150 + {
151 + is_good = false;
152 + }
153 + if is_good {
154 + good += 1;
155 + }
156 + }
157 + good
158 + } else {
159 + 0
160 + };
161 +
162 + // Sort each metric array once, then compute median + percentiles from sorted data
163 + let sorted_ttfts = sort_clone(&ttfts);
164 + let sorted_tpots = sort_clone(&tpots);
165 + let sorted_itls = sort_clone(&itls);
166 + let sorted_e2els = sort_clone(&e2els);
167 +
168 + let request_goodput = if !goodput_config.is_empty() {
169 + good_completed as f64 / dur_s
170 + } else {
171 + 0.0
172 + };
173 +
174 + let metrics = BenchmarkMetrics {
175 + completed,
176 + failed,
177 + total_input,
178 + total_output,
179 + request_throughput: completed as f64 / dur_s,
180 + request_goodput,
181 + input_throughput: total_input as f64 / dur_s,
182 + output_throughput: total_output as f64 / dur_s,
183 + total_token_throughput: (total_input + total_output) as f64 / dur_s,
184 + mean_ttft_ms: mean(&ttfts) * 1000.0,
185 + median_ttft_ms: median_sorted(&sorted_ttfts) * 1000.0,
186 + std_ttft_ms: std_dev(&ttfts) * 1000.0,
187 + percentiles_ttft_ms: percentiles_from_sorted(&sorted_ttfts, selected_percentiles),
188 + mean_tpot_ms: mean(&tpots) * 1000.0,
189 + median_tpot_ms: median_sorted(&sorted_tpots) * 1000.0,
190 + std_tpot_ms: std_dev(&tpots) * 1000.0,
191 + percentiles_tpot_ms: percentiles_from_sorted(&sorted_tpots, selected_percentiles),
192 + mean_itl_ms: mean(&itls) * 1000.0,
193 + median_itl_ms: median_sorted(&sorted_itls) * 1000.0,
194 + std_itl_ms: std_dev(&itls) * 1000.0,
195 + percentiles_itl_ms: percentiles_from_sorted(&sorted_itls, selected_percentiles),
196 + mean_e2el_ms: mean(&e2els) * 1000.0,
197 + median_e2el_ms: median_sorted(&sorted_e2els) * 1000.0,
198 + std_e2el_ms: std_dev(&e2els) * 1000.0,
199 + percentiles_e2el_ms: percentiles_from_sorted(&sorted_e2els, selected_percentiles),
200 + max_output_tokens_per_s,
201 + max_concurrent_requests,
202 + steady_state: None,
203 + };
204 +
205 + (metrics, actual_output_lens)
206 +}
207 +
208 +fn mean(data: &[f64]) -> f64 {
209 + if data.is_empty() {
210 + return 0.0;
211 + }
212 + data.iter().sum::<f64>() / data.len() as f64
213 +}
214 +
215 +fn sort_clone(data: &[f64]) -> Vec<f64> {
216 + let mut sorted = data.to_vec();
217 + sorted.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));
218 + sorted
219 +}
220 +
221 +fn median_sorted(sorted: &[f64]) -> f64 {
222 + if sorted.is_empty() {
223 + return 0.0;
224 + }
225 + let mid = sorted.len() / 2;
226 + if sorted.len().is_multiple_of(2) {
227 + (sorted[mid - 1] + sorted[mid]) / 2.0
228 + } else {
229 + sorted[mid]
230 + }
231 +}
232 +
233 +fn std_dev(data: &[f64]) -> f64 {
234 + if data.is_empty() {
235 + return 0.0;
236 + }
237 + let m = mean(data);
238 + let variance = data.iter().map(|x| (x - m).powi(2)).sum::<f64>() / data.len() as f64;
239 + variance.sqrt()
240 +}
241 +
242 +fn percentile_sorted(sorted_data: &[f64], p: f64) -> f64 {
243 + if sorted_data.is_empty() {
244 + return 0.0;
245 + }
246 + if sorted_data.len() == 1 {
247 + return sorted_data[0];
248 + }
249 + // Use the same interpolation as numpy (linear)
250 + let idx = p / 100.0 * (sorted_data.len() - 1) as f64;
251 + let lo = idx.floor() as usize;
252 + let hi = idx.ceil() as usize;
253 + let frac = idx - lo as f64;
254 + if lo == hi {
255 + sorted_data[lo]
256 + } else {
257 + sorted_data[lo] * (1.0 - frac) + sorted_data[hi] * frac
258 + }
259 +}
260 +
261 +fn percentiles_from_sorted(sorted: &[f64], percentiles: &[f64]) -> Vec<(f64, f64)> {
262 + if sorted.is_empty() {
263 + return percentiles.iter().map(|&p| (p, 0.0)).collect();
264 + }
265 + percentiles
266 + .iter()
267 + .map(|&p| (p, percentile_sorted(sorted, p) * 1000.0))
268 + .collect()
269 +}
270 +
271 +/// Calculate benchmark metrics for embedding/pooling requests.
272 +///
273 +/// Mirrors Python's `calculate_metrics_for_embeddings()` from serve.py.
274 +/// Key differences from `calculate_metrics()`:
275 +/// - Uses `output.prompt_len` (server-reported) instead of `input_requests[i].prompt_len`
276 +/// - Only computes E2EL (no TTFT/TPOT/ITL since pooling is non-streaming)
277 +/// - `total_output` is 0 (pooling produces no output tokens)
278 +/// - `total_token_throughput` = `total_input / dur_s` (input tokens only)
279 +pub fn calculate_embedding_metrics(
280 + outputs: &[RequestFuncOutput],
281 + dur_s: f64,
282 + selected_percentiles: &[f64],
283 +) -> BenchmarkMetrics {
284 + let mut total_input: usize = 0;
285 + let mut completed: usize = 0;
286 + let mut e2els: Vec<f64> = Vec::new();
287 +
288 + for output in outputs {
289 + if output.success {
290 + e2els.push(output.latency);
291 + completed += 1;
292 + total_input += output.prompt_len;
293 + }
294 + }
295 +
296 + let failed = outputs.len() - completed;
297 +
298 + // Print failed request errors (capped to 10)
299 + let failed_outputs: Vec<&RequestFuncOutput> = outputs.iter().filter(|o| !o.success).collect();
300 + if !failed_outputs.is_empty() {
301 + eprintln!("Failed requests during benchmark run detected (capping to 10):");
302 + for (i, err) in failed_outputs.iter().take(10).enumerate() {
303 + eprintln!("Error {i}: {}", err.error);
304 + }
305 + }
306 +
307 + // Compute peak concurrent requests from start_time + latency windows
308 + let successful_outputs: Vec<&RequestFuncOutput> =
309 + outputs.iter().filter(|o| o.success).collect();
310 + let max_concurrent_requests = if !successful_outputs.is_empty() {
311 + let min_start =
312 + successful_outputs.iter().map(|o| o.start_time).fold(f64::INFINITY, f64::min);
313 + let max_end = successful_outputs
314 + .iter()
315 + .map(|o| o.start_time + o.latency)
316 + .fold(f64::NEG_INFINITY, f64::max);
317 +
318 + let raw_duration = (max_end - min_start).ceil() as usize + 1;
319 + let duration_seconds = raw_duration.min(86_400);
320 + let mut concurrent_per_second = vec![0usize; duration_seconds];
321 +
322 + for output in &successful_outputs {
323 + let start_second = (output.start_time - min_start) as usize;
324 + let end_second = ((output.start_time + output.latency) - min_start) as usize;
325 + for slot in concurrent_per_second
326 + .iter_mut()
327 + .take(end_second.min(duration_seconds - 1) + 1)
328 + .skip(start_second)
329 + {
330 + *slot += 1;
331 + }
332 + }
333 +
334 + *concurrent_per_second.iter().max().unwrap_or(&0)
335 + } else {
336 + 0
337 + };
338 +
339 + let sorted_e2els = sort_clone(&e2els);
340 +
341 + BenchmarkMetrics {
342 + completed,
343 + failed,
344 + total_input,
345 + total_output: 0,
346 + request_throughput: completed as f64 / dur_s,
347 + request_goodput: 0.0,
348 + input_throughput: total_input as f64 / dur_s,
349 + output_throughput: 0.0,
350 + total_token_throughput: total_input as f64 / dur_s,
351 + mean_ttft_ms: 0.0,
352 + median_ttft_ms: 0.0,
353 + std_ttft_ms: 0.0,
354 + percentiles_ttft_ms: Vec::new(),
355 + mean_tpot_ms: 0.0,
356 + median_tpot_ms: 0.0,
357 + std_tpot_ms: 0.0,
358 + percentiles_tpot_ms: Vec::new(),
359 + mean_itl_ms: 0.0,
360 + median_itl_ms: 0.0,
361 + std_itl_ms: 0.0,
362 + percentiles_itl_ms: Vec::new(),
363 + mean_e2el_ms: mean(&e2els) * 1000.0,
364 + median_e2el_ms: median_sorted(&sorted_e2els) * 1000.0,
365 + std_e2el_ms: std_dev(&e2els) * 1000.0,
366 + percentiles_e2el_ms: percentiles_from_sorted(&sorted_e2els, selected_percentiles),
367 + max_output_tokens_per_s: 0.0,
368 + max_concurrent_requests,
369 + steady_state: None,
370 + }
371 +}
372 +
373 +/// Calculate multi-turn benchmark metrics with per-turn breakdown.
374 +///
375 +/// 1. Flatten all turns into SampleRequest/RequestFuncOutput pairs → overall metrics
376 +/// 2. Group by turn_index → per-turn metrics
377 +/// 3. Conversation-level stats
378 +pub fn calculate_multi_turn_metrics(
379 + conversation_outputs: &[ConversationOutput],
380 + dur_s: f64,
381 + selected_percentiles: &[f64],
382 + goodput_config: &GoodputConfig,
383 + max_turn_count: usize,
384 +) -> MultiTurnMetrics {
385 + // Flatten all turns for overall metrics
386 + let mut all_requests: Vec<SampleRequest> = Vec::new();
387 + let mut all_outputs: Vec<RequestFuncOutput> = Vec::new();
388 +
389 + for conv in conversation_outputs {
390 + for turn in &conv.turns {
391 + all_requests.push(SampleRequest {
392 + prompt: std::sync::Arc::from(""),
393 + prompt_len: turn.cumulative_input_tokens,
394 + expected_output_len: turn.request_output.output_tokens,
395 + request_id: None,
396 + ..Default::default()
397 + });
398 + all_outputs.push(turn.request_output.clone());
399 + }
400 + }
401 +
402 + let (overall, _) = calculate_metrics(
403 + &all_requests,
404 + &all_outputs,
405 + dur_s,
406 + selected_percentiles,
407 + true,
408 + goodput_config,
409 + );
410 +
411 + // Per-turn breakdown
412 + let mut per_turn: Vec<BenchmarkMetrics> = Vec::new();
413 + for turn_idx in 0..max_turn_count {
414 + let mut turn_requests: Vec<SampleRequest> = Vec::new();
415 + let mut turn_outputs: Vec<RequestFuncOutput> = Vec::new();
416 +
417 + for conv in conversation_outputs {
418 + if let Some(turn) = conv.turns.get(turn_idx) {
419 + turn_requests.push(SampleRequest {
420 + prompt: std::sync::Arc::from(""),
421 + prompt_len: turn.cumulative_input_tokens,
422 + expected_output_len: turn.request_output.output_tokens,
423 + request_id: None,
424 + ..Default::default()
425 + });
426 + turn_outputs.push(turn.request_output.clone());
427 + }
428 + }
429 +
430 + if !turn_outputs.is_empty() {
431 + let (turn_metrics, _) = calculate_metrics(
432 + &turn_requests,
433 + &turn_outputs,
434 + dur_s,
435 + selected_percentiles,
436 + true,
437 + goodput_config,
438 + );
439 + per_turn.push(turn_metrics);
440 + }
441 + }
442 +
443 + // Conversation-level stats
444 + let conversations_completed = conversation_outputs.iter().filter(|c| c.all_success).count();
445 + let conversations_failed = conversation_outputs.len() - conversations_completed;
446 +
447 + let avg_turns_completed = if conversation_outputs.is_empty() {
448 + 0.0
449 + } else {
450 + conversation_outputs.iter().map(|c| c.turns.len() as f64).sum::<f64>()
451 + / conversation_outputs.len() as f64
452 + };
453 +
454 + let avg_conversation_duration_ms = if conversation_outputs.is_empty() {
455 + 0.0
456 + } else {
457 + conversation_outputs.iter().map(|c| c.total_duration_ms).sum::<f64>()
458 + / conversation_outputs.len() as f64
459 + };
460 +
461 + MultiTurnMetrics {
462 + overall,
463 + per_turn,
464 + conversations_completed,
465 + conversations_failed,
466 + avg_turns_completed,
467 + avg_conversation_duration_ms,
468 + }
469 +}
470 +
471 +#[cfg(test)]
472 +mod tests {
473 + use super::*;
474 +
475 + #[test]
476 + fn test_mean_empty() {
477 + assert_eq!(mean(&[]), 0.0);
478 + }
479 +
480 + #[test]
481 + fn test_mean_values() {
482 + assert!((mean(&[1.0, 2.0, 3.0]) - 2.0).abs() < 1e-10);
483 + }
484 +
485 + #[test]
486 + fn test_median_odd() {
487 + let sorted = sort_clone(&[3.0, 1.0, 2.0]);
488 + assert!((median_sorted(&sorted) - 2.0).abs() < 1e-10);
489 + }
490 +
491 + #[test]
492 + fn test_median_even() {
493 + let sorted = sort_clone(&[1.0, 2.0, 3.0, 4.0]);
494 + assert!((median_sorted(&sorted) - 2.5).abs() < 1e-10);
495 + }
496 +
497 + #[test]
498 + fn test_percentile_99() {
499 + let data: Vec<f64> = (0..100).map(|i| i as f64).collect();
500 + let p99 = percentile_sorted(&data, 99.0);
501 + assert!((p99 - 98.01).abs() < 0.1);
502 + }
503 +}
rust/src/bench/src/metrics/mod.rs new
+54
@@ -0,0 +1,54 @@
1 +// SPDX-License-Identifier: Apache-2.0
2 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +pub mod calculator;
5 +pub mod steady_state;
6 +
7 +use serde::{Deserialize, Serialize};
8 +pub use steady_state::SteadyStateMetrics;
9 +
10 +/// Multi-turn benchmark metrics with overall and per-turn breakdown.
11 +#[derive(Debug, Clone, Serialize, Deserialize)]
12 +pub struct MultiTurnMetrics {
13 + pub overall: BenchmarkMetrics,
14 + pub per_turn: Vec<BenchmarkMetrics>,
15 + pub conversations_completed: usize,
16 + pub conversations_failed: usize,
17 + pub avg_turns_completed: f64,
18 + pub avg_conversation_duration_ms: f64,
19 +}
20 +
21 +/// Full benchmark metrics for generation tasks.
22 +/// Matches Python's BenchmarkMetrics dataclass.
23 +#[derive(Debug, Clone, Serialize, Deserialize)]
24 +pub struct BenchmarkMetrics {
25 + pub completed: usize,
26 + pub failed: usize,
27 + pub total_input: usize,
28 + pub total_output: usize,
29 + pub request_throughput: f64,
30 + pub request_goodput: f64,
31 + pub input_throughput: f64,
32 + pub output_throughput: f64,
33 + pub total_token_throughput: f64,
34 + pub mean_ttft_ms: f64,
35 + pub median_ttft_ms: f64,
36 + pub std_ttft_ms: f64,
37 + pub percentiles_ttft_ms: Vec<(f64, f64)>,
38 + pub mean_tpot_ms: f64,
39 + pub median_tpot_ms: f64,
40 + pub std_tpot_ms: f64,
41 + pub percentiles_tpot_ms: Vec<(f64, f64)>,
42 + pub mean_itl_ms: f64,
43 + pub median_itl_ms: f64,
44 + pub std_itl_ms: f64,
45 + pub percentiles_itl_ms: Vec<(f64, f64)>,
46 + pub mean_e2el_ms: f64,
47 + pub median_e2el_ms: f64,
48 + pub std_e2el_ms: f64,
49 + pub percentiles_e2el_ms: Vec<(f64, f64)>,
50 + pub max_output_tokens_per_s: f64,
51 + pub max_concurrent_requests: usize,
52 + #[serde(default)]
53 + pub steady_state: Option<SteadyStateMetrics>,
54 +}
rust/src/bench/src/metrics/steady_state.rs new
+629
@@ -0,0 +1,629 @@
1 +// SPDX-License-Identifier: Apache-2.0
2 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +//! Steady-state metrics: throughput and TTFT measured over the window
5 +//! during which client-side in-flight concurrency is at or above a
6 +//! configurable fraction of `--max-concurrency`. Excludes ramp-up and
7 +//! drain phases to reduce run-to-run variance at very high concurrency.
8 +
9 +use serde::{Deserialize, Serialize};
10 +
11 +/// Bounds and metadata of the detected steady-state window.
12 +///
13 +/// `observed_peak` is event-exact and may differ by 1 from the
14 +/// bucket-approximated `max_concurrent_requests` on `BenchmarkMetrics`;
15 +/// this is expected.
16 +#[derive(Debug, Clone, Serialize, Deserialize)]
17 +pub struct SteadyStateWindow {
18 + pub start_s: f64,
19 + pub end_s: f64,
20 + pub duration_s: f64,
21 + pub target_concurrency: usize,
22 + pub threshold: f64,
23 + pub threshold_abs: usize,
24 + pub observed_peak: usize,
25 + pub requests_started_in_window: usize,
26 + pub requests_completed_in_window: usize,
27 + pub requests_total: usize,
28 + pub warning: Option<String>,
29 +}
30 +
31 +#[derive(Debug, Clone, Serialize, Deserialize)]
32 +pub struct SteadyStateMetrics {
33 + pub window: SteadyStateWindow,
34 + pub request_throughput: f64,
35 + pub input_throughput: f64,
36 + pub output_throughput: f64,
37 + pub total_token_throughput: f64,
38 + pub mean_ttft_ms: f64,
39 + pub median_ttft_ms: f64,
40 + pub percentiles_ttft_ms: Vec<(f64, f64)>,
41 + pub mean_tpot_ms: f64,
42 + pub median_tpot_ms: f64,
43 + pub p90_tpot_ms: f64,
44 + pub p99_tpot_ms: f64,
45 +}
46 +
47 +use crate::backends::RequestFuncOutput;
48 +
49 +/// Detect the steady-state window over `outputs` under the given target concurrency.
50 +///
51 +/// Returns `None` when closed-loop gate is not met (no target), when concurrency
52 +/// never crosses `threshold_abs`, or when both started and completed in-window
53 +/// request counts are zero.
54 +///
55 +/// Timestamps in the returned window are normalized so `start_s` / `end_s` are
56 +/// "seconds from the earliest successful request's start_time."
57 +///
58 +/// `user_min_window_s`: if `Some`, sets the minimum window duration below which
59 +/// a `warning` is attached. If `None`, the caller is expected to have pre-resolved
60 +/// the default `max(10.0, 0.1 * total_run_duration_s)`.
61 +pub fn detect_window(
62 + outputs: &[RequestFuncOutput],
63 + target_concurrency: Option<usize>,
64 + threshold: f64,
65 + min_window_s: f64,
66 + total_run_duration_s: f64,
67 +) -> Option<SteadyStateWindow> {
68 + let target = target_concurrency?;
69 + let threshold_abs = ((threshold * target as f64).ceil() as usize).max(1);
70 +
71 + // Collect sorted start / end timestamps from successful requests only.
72 + let mut starts: Vec<f64> = outputs.iter().filter(|o| o.success).map(|o| o.start_time).collect();
73 + let mut ends: Vec<f64> =
74 + outputs.iter().filter(|o| o.success).map(|o| o.start_time + o.latency).collect();
75 + if starts.is_empty() {
76 + return None;
77 + }
78 + starts.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));
79 + ends.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));
80 +
81 + let min_start = starts[0];
82 + let requests_total = outputs.iter().filter(|o| o.success).count();
83 +
84 + // Two-pointer merge. Tie rule: process `end` before `start` at equal timestamps.
85 + let mut i = 0usize;
86 + let mut j = 0usize;
87 + let mut concurrency: usize = 0;
88 + let mut observed_peak: usize = 0;
89 + let mut up_crossing: Option<f64> = None;
90 + let mut last_down_crossing: Option<f64> = None;
91 +
92 + while i < starts.len() || j < ends.len() {
93 + let next_is_end = match (starts.get(i), ends.get(j)) {
94 + (Some(&s), Some(&e)) => e <= s, // tie -> end first
95 + (None, Some(_)) => true,
96 + (Some(_), None) => false,
97 + (None, None) => break,
98 + };
99 +
100 + if next_is_end {
101 + let t = ends[j];
102 + j += 1;
103 + let before = concurrency;
104 + concurrency = concurrency.saturating_sub(1);
105 + if before >= threshold_abs && concurrency < threshold_abs {
106 + last_down_crossing = Some(t);
107 + }
108 + } else {
109 + let t = starts[i];
110 + i += 1;
111 + let before = concurrency;
112 + concurrency += 1;
113 + if concurrency > observed_peak {
114 + observed_peak = concurrency;
115 + }
116 + if before < threshold_abs && concurrency >= threshold_abs && up_crossing.is_none() {
117 + up_crossing = Some(t);
118 + }
119 + }
120 + }
121 +
122 + let start_abs = up_crossing?;
123 + // A down-crossing always exists once an up-crossing exists — the final ends
124 + // drain concurrency to 0. Fall back defensively to the last end timestamp.
125 + let end_abs = last_down_crossing.or_else(|| ends.last().copied())?;
126 +
127 + let start_s = start_abs - min_start;
128 + let end_s = end_abs - min_start;
129 + let duration_s = end_s - start_s;
130 +
131 + // Count requests whose start / completion events land in [start_abs, end_abs).
132 + let mut requests_started_in_window = 0usize;
133 + let mut requests_completed_in_window = 0usize;
134 + for o in outputs.iter().filter(|o| o.success) {
135 + if o.start_time >= start_abs && o.start_time < end_abs {
136 + requests_started_in_window += 1;
137 + }
138 + let end_t = o.start_time + o.latency;
139 + if end_t >= start_abs && end_t < end_abs {
140 + requests_completed_in_window += 1;
141 + }
142 + }
143 +
144 + if requests_started_in_window == 0 && requests_completed_in_window == 0 {
145 + return None;
146 + }
147 +
148 + let warning = if duration_s < min_window_s {
149 + Some(format!(
150 + "steady-state window is {:.1}s ({:.1}% of run); may not be representative",
151 + duration_s,
152 + 100.0 * duration_s / total_run_duration_s.max(1e-9)
153 + ))
154 + } else {
155 + None
156 + };
157 +
158 + Some(SteadyStateWindow {
159 + start_s,
160 + end_s,
161 + duration_s,
162 + target_concurrency: target,
163 + threshold,
164 + threshold_abs,
165 + observed_peak,
166 + requests_started_in_window,
167 + requests_completed_in_window,
168 + requests_total,
169 + warning,
170 + })
171 +}
172 +
173 +use crate::datasets::SampleRequest;
174 +
175 +/// Compute steady-state metrics from per-request outputs, the detected window,
176 +/// and the original request list (for input-token attribution).
177 +///
178 +/// `is_pooling` gates off TTFT because pooling backends write `ttft = latency`
179 +/// as a placeholder; reporting TTFT over that would silently be E2EL.
180 +pub fn compute(
181 + outputs: &[RequestFuncOutput],
182 + input_requests: &[SampleRequest],
183 + window: &SteadyStateWindow,
184 + percentiles: &[f64],
185 + is_pooling: bool,
186 +) -> SteadyStateMetrics {
187 + // Re-derive min_start to convert normalized window bounds back to absolute times.
188 + let min_start = outputs
189 + .iter()
190 + .filter(|o| o.success)
191 + .map(|o| o.start_time)
192 + .fold(f64::INFINITY, f64::min);
193 + let start_abs = window.start_s + min_start;
194 + let end_abs = window.end_s + min_start;
195 + let duration = window.duration_s.max(1e-9);
196 +
197 + let mut completed_in_window = 0usize;
198 + let mut input_tokens_in_window: usize = 0;
199 + let mut output_tokens_in_window: usize = 0;
200 + let mut ttfts_in_window: Vec<f64> = Vec::new();
201 + let mut tpots_in_window: Vec<f64> = Vec::new();
202 +
203 + for (idx, o) in outputs.iter().enumerate() {
204 + if !o.success {
205 + continue;
206 + }
207 + let started_in = o.start_time >= start_abs && o.start_time < end_abs;
208 + let end_t = o.start_time + o.latency;
209 + let completed_in = end_t >= start_abs && end_t < end_abs;
210 +
211 + if completed_in {
212 + completed_in_window += 1;
213 + }
214 +
215 + if started_in {
216 + input_tokens_in_window += input_requests[idx].prompt_len;
217 + if !is_pooling {
218 + ttfts_in_window.push(o.ttft);
219 + // Per-request TPOT: (latency - ttft) / (output_tokens - 1).
220 + // Matches calculator.rs:46-50; skip requests with <= 1 token.
221 + if o.output_tokens > 1 {
222 + let tpot = (o.latency - o.ttft) / (o.output_tokens as f64 - 1.0);
223 + tpots_in_window.push(tpot);
224 + }
225 + }
226 + }
227 +
228 + // Per-token emission attribution (same approach as calculator.rs:98-113).
229 + // Reconstruct: first token at start_time + ttft, then cumulative itl.
230 + if o.output_tokens == 0 && o.itl.is_empty() {
231 + continue;
232 + }
233 + let first_token_t = o.start_time + o.ttft;
234 + if first_token_t >= start_abs && first_token_t < end_abs {
235 + output_tokens_in_window += 1;
236 + }
237 + let mut t = first_token_t;
238 + for &dt in &o.itl {
239 + t += dt;
240 + if t >= start_abs && t < end_abs {
241 + output_tokens_in_window += 1;
242 + }
243 + if t >= end_abs {
244 + break; // tokens only move forward in time
245 + }
246 + }
247 + }
248 +
249 + let request_throughput = completed_in_window as f64 / duration;
250 + let input_throughput = input_tokens_in_window as f64 / duration;
251 + let output_throughput = output_tokens_in_window as f64 / duration;
252 + let total_token_throughput = input_throughput + output_throughput;
253 +
254 + let (mean_ttft, median_ttft, pct_ttft) = if is_pooling {
255 + (0.0, 0.0, Vec::new())
256 + } else {
257 + dist_stats(&ttfts_in_window, percentiles)
258 + };
259 +
260 + let (mean_tpot, median_tpot, pct_tpot) = if is_pooling {
261 + (0.0, 0.0, Vec::new())
262 + } else {
263 + dist_stats(&tpots_in_window, &[90.0, 99.0])
264 + };
265 + let p90_tpot = pct_tpot.first().map(|(_, v)| *v).unwrap_or(0.0);
266 + let p99_tpot = pct_tpot.get(1).map(|(_, v)| *v).unwrap_or(0.0);
267 +
268 + SteadyStateMetrics {
269 + window: window.clone(),
270 + request_throughput,
271 + input_throughput,
272 + output_throughput,
273 + total_token_throughput,
274 + mean_ttft_ms: mean_ttft,
275 + median_ttft_ms: median_ttft,
276 + percentiles_ttft_ms: pct_ttft,
277 + mean_tpot_ms: mean_tpot,
278 + median_tpot_ms: median_tpot,
279 + p90_tpot_ms: p90_tpot,
280 + p99_tpot_ms: p99_tpot,
281 + }
282 +}
283 +
284 +/// Returns (mean_ms, median_ms, percentiles_ms). Values are seconds in; output
285 +/// is milliseconds. Empty input yields zeros with a zero-filled percentile vec
286 +/// aligned to the requested percentile list.
287 +fn dist_stats(values: &[f64], percentiles: &[f64]) -> (f64, f64, Vec<(f64, f64)>) {
288 + if values.is_empty() {
289 + let pct = percentiles.iter().map(|&p| (p, 0.0)).collect();
290 + return (0.0, 0.0, pct);
291 + }
292 + let mut sorted = values.to_vec();
293 + sorted.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));
294 + let n = sorted.len();
295 + let mean = sorted.iter().sum::<f64>() / n as f64;
296 + let median = if n.is_multiple_of(2) {
297 + (sorted[n / 2 - 1] + sorted[n / 2]) / 2.0
298 + } else {
299 + sorted[n / 2]
300 + };
301 + let pct = percentiles
302 + .iter()
303 + .map(|&p| {
304 + let idx = (p / 100.0) * (n - 1) as f64;
305 + let lo = idx.floor() as usize;
306 + let hi = idx.ceil() as usize;
307 + let frac = idx - lo as f64;
308 + let v = if lo == hi {
309 + sorted[lo]
310 + } else {
311 + sorted[lo] * (1.0 - frac) + sorted[hi] * frac
312 + };
313 + (p, v * 1000.0)
314 + })
315 + .collect();
316 + (mean * 1000.0, median * 1000.0, pct)
317 +}
318 +
319 +#[cfg(test)]
320 +mod tests {
321 + use super::*;
322 + use crate::backends::RequestFuncOutput;
323 +
324 + /// Build a successful synthetic output with given `start_time` and `latency`.
325 + fn mk(start: f64, latency: f64) -> RequestFuncOutput {
326 + RequestFuncOutput {
327 + success: true,
328 + start_time: start,
329 + latency,
330 + ttft: 0.0,
331 + itl: Vec::new(),
332 + output_tokens: 0,
333 + ..Default::default()
334 + }
335 + }
336 +
337 + #[test]
338 + fn window_plateau_trapezoid() {
339 + // 10 requests ramp in at t=0..10 each lasting 30s.
340 + // Concurrency reaches 10 at t=10 and stays until t=30 when requests start ending.
341 + // threshold = 1.0 means exact-target; threshold_abs = 10.
342 + let mut outputs: Vec<RequestFuncOutput> = (0..10).map(|i| mk(i as f64, 30.0)).collect();
343 + // Add one request fully inside the plateau so it can pick up as "in window"
344 + outputs.push(mk(15.0, 5.0));
345 +
346 + let w = detect_window(&outputs, Some(10), 1.0, 5.0, 30.0).unwrap();
347 + assert_eq!(w.target_concurrency, 10);
348 + assert_eq!(w.threshold_abs, 10);
349 + // Up-crossing at t=9 (the 10th start) — concurrency transitions 9→10.
350 + // Window start reported in normalized time (min_start = 0.0).
351 + assert!((w.start_s - 9.0).abs() < 1e-9, "start_s = {}", w.start_s);
352 + // Last down-crossing: the first end at t=30 (concurrency 10→9, falling below 10).
353 + assert!((w.end_s - 30.0).abs() < 1e-9, "end_s = {}", w.end_s);
354 + assert!(w.duration_s > 0.0);
355 + assert_eq!(w.observed_peak, 11); // 10 ramp-in + 1 extra = 11 peak briefly
356 + assert!(w.warning.is_none());
357 + }
358 +
359 + #[test]
360 + fn window_none_when_target_none() {
361 + let outputs = vec![mk(0.0, 1.0)];
362 + assert!(detect_window(&outputs, None, 0.95, 1.0, 10.0).is_none());
363 + }
364 +
365 + #[test]
366 + fn window_none_when_never_reaches_threshold() {
367 + // Serial requests, concurrency never exceeds 1.
368 + let outputs: Vec<RequestFuncOutput> = (0..5).map(|i| mk(i as f64 * 10.0, 1.0)).collect();
369 + assert!(detect_window(&outputs, Some(5), 0.95, 1.0, 100.0).is_none());
370 + }
371 +
372 + #[test]
373 + fn window_none_when_all_fail() {
374 + let mut o = mk(0.0, 1.0);
375 + o.success = false;
376 + assert!(detect_window(&[o], Some(1), 1.0, 0.1, 10.0).is_none());
377 + }
378 +
379 + #[test]
380 + fn window_threshold_one() {
381 + // threshold=1.0, target=2 -> threshold_abs=2; needs both requests in flight.
382 + let outputs = vec![mk(0.0, 5.0), mk(1.0, 5.0)];
383 + let w = detect_window(&outputs, Some(2), 1.0, 0.1, 10.0).unwrap();
384 + assert_eq!(w.threshold_abs, 2);
385 + assert!((w.start_s - 1.0).abs() < 1e-9);
386 + assert!((w.end_s - 5.0).abs() < 1e-9);
387 + }
388 +
389 + #[test]
390 + fn window_ties_net_to_zero() {
391 + // A request finishes and another starts at exactly the same instant;
392 + // tie rule (end before start) prevents concurrency > 2 blip.
393 + let outputs = vec![
394 + mk(0.0, 5.0), // ends at 5.0
395 + mk(2.0, 3.0), // ends at 5.0
396 + mk(5.0, 2.0), // starts at 5.0 — ties with the two ends
397 + ];
398 + let w = detect_window(&outputs, Some(2), 1.0, 0.1, 20.0).unwrap();
399 + assert_eq!(w.observed_peak, 2);
400 + }
401 +
402 + #[test]
403 + fn window_warning_when_short() {
404 + let outputs: Vec<RequestFuncOutput> = (0..10).map(|i| mk(i as f64 * 0.1, 1.0)).collect();
405 + // Plateau is ~0.1s; min_window_s = 5s => warning.
406 + let w = detect_window(&outputs, Some(10), 1.0, 5.0, 10.0).unwrap();
407 + assert!(w.warning.is_some(), "expected warning");
408 + }
409 +
410 + #[test]
411 + fn window_constant_concurrency_equals_full_run() {
412 + // Concurrency is at target throughout — steady-state window covers the
413 + // entire middle of the run where all 3 requests overlap.
414 + let outputs = vec![mk(0.0, 10.0), mk(0.0, 10.0), mk(0.0, 10.0)];
415 + let w = detect_window(&outputs, Some(3), 1.0, 0.1, 10.0).unwrap();
416 + assert!((w.start_s - 0.0).abs() < 1e-9);
417 + assert!((w.end_s - 10.0).abs() < 1e-9);
418 + }
419 +
420 + use std::sync::Arc;
421 +
422 + use crate::datasets::SampleRequest;
423 +
424 + fn mk_full(
425 + start: f64,
426 + latency: f64,
427 + ttft: f64,
428 + itl: Vec<f64>,
429 + output_tokens: usize,
430 + ) -> RequestFuncOutput {
431 + RequestFuncOutput {
432 + success: true,
433 + start_time: start,
434 + latency,
435 + ttft,
436 + itl,
437 + output_tokens,
438 + prompt_len: 10,
439 + ..Default::default()
440 + }
441 + }
442 +
443 + fn mk_sample(prompt_len: usize) -> SampleRequest {
444 + SampleRequest {
445 + prompt: Arc::from(""),
446 + prompt_len,
447 + expected_output_len: 0,
448 + request_id: None,
449 + ..Default::default()
450 + }
451 + }
452 +
453 + #[test]
454 + fn compute_constant_concurrency_matches_full_run() {
455 + // 3 identical requests, each 10s, each emitting 5 tokens uniformly via itl=2s.
456 + let itl = vec![2.0, 2.0, 2.0, 2.0]; // 5 tokens: first at t=start+ttft, then 4 more
457 + let outputs = vec![
458 + mk_full(0.0, 10.0, 2.0, itl.clone(), 5),
459 + mk_full(0.0, 10.0, 2.0, itl.clone(), 5),
460 + mk_full(0.0, 10.0, 2.0, itl.clone(), 5),
461 + ];
462 + let requests: Vec<SampleRequest> =
463 + outputs.iter().map(|o| mk_sample(o.prompt_len)).collect();
464 +
465 + let w = detect_window(&outputs, Some(3), 1.0, 0.1, 10.0).unwrap();
466 + let m = compute(
467 + &outputs,
468 + &requests,
469 + &w,
470 + &[99.0],
471 + false, // is_pooling
472 + );
473 +
474 + // Window is [0.0, 10.0). All 3 requests start in window (3).
475 + // Tokens per request: first token at ttft=2s, then 4 more at itl=2s each
476 + // -> t = 2, 4, 6, 8, 10. The 10.0 token is AT end_s and is half-open excluded.
477 + // -> 4 tokens per request land in window.
478 + assert_eq!(w.requests_started_in_window, 3);
479 + // Half-open window [0.0, 10.0) excludes the end_t=10.0 completions, so
480 + // request_throughput is exactly 0.0 here — tightens the half-open invariant.
481 + assert!((m.request_throughput - 0.0).abs() < 1e-9);
482 + // output tokens in window = 4 per request * 3 = 12, over 10s = 1.2 tok/s
483 + assert!(
484 + (m.output_throughput - 1.2).abs() < 1e-6,
485 + "got {}",
486 + m.output_throughput
487 + );
488 + // input throughput = 30 input tokens / 10s = 3.0
489 + assert!(
490 + (m.input_throughput - 3.0).abs() < 1e-9,
491 + "got {}",
492 + m.input_throughput
493 + );
494 + }
495 +
496 + #[test]
497 + fn compute_drain_only_tokens_not_counted() {
498 + // 1 "steady" request defining the window, + 1 late-admitted request whose
499 + // 1000 output tokens land after end_s — must not pump output_throughput.
500 + let mut outputs = vec![mk_full(0.0, 10.0, 0.0, vec![1.0; 9], 10)];
501 + // Late request: starts at 9.9 (within window), long output that drains past end_s.
502 + let late_itl = vec![0.02; 999]; // 1000 tokens; emission starts after end_s
503 + let late = mk_full(9.9, 30.0, 0.05, late_itl, 1000);
504 + outputs.push(late);
505 + let requests: Vec<SampleRequest> =
506 + outputs.iter().map(|o| mk_sample(o.prompt_len)).collect();
507 +
508 + // Target = 2 so threshold_abs = 2; window opens when both are in flight.
509 + let w = detect_window(&outputs, Some(2), 1.0, 0.1, 40.0).unwrap();
510 + let m = compute(&outputs, &requests, &w, &[99.0], false);
511 +
512 + // Late request started in window but most of its tokens are emitted AFTER end_s.
513 + // Only a handful (≤3) should count.
514 + // Window end is 10.0 (the first request's completion).
515 + // Late request first token at 9.9 + 0.05 = 9.95 (in window).
516 + // Next tokens at 9.97, 9.99, 10.01... -> only ~3 fit before end_s=10.0.
517 + let expected_max_late_tokens = 6.0; // generous bound
518 + let window_dur = w.duration_s;
519 + // Upper bound on output throughput: 10 (from req 1) + ≤6 (from late) / window_dur.
520 + let upper_bound = (10.0 + expected_max_late_tokens) / window_dur;
521 + assert!(
522 + m.output_throughput <= upper_bound,
523 + "output_throughput {} exceeds upper bound {} — drain tokens are being counted",
524 + m.output_throughput,
525 + upper_bound
526 + );
527 + }
528 +
529 + #[test]
530 + fn compute_pooling_skips_ttft_and_tpot() {
531 + // For pooling backends, ttft == latency is a placeholder; we must NOT emit TTFT stats.
532 + // Pooling has no output tokens so TPOT is also zeroed.
533 + let outputs = vec![
534 + mk_full(0.0, 5.0, 5.0, Vec::new(), 0),
535 + mk_full(0.0, 5.0, 5.0, Vec::new(), 0),
536 + ];
537 + let requests: Vec<SampleRequest> =
538 + outputs.iter().map(|o| mk_sample(o.prompt_len)).collect();
539 + let w = detect_window(&outputs, Some(2), 1.0, 0.1, 10.0).unwrap();
540 + let m = compute(&outputs, &requests, &w, &[99.0], /* is_pooling */ true);
541 + assert_eq!(m.mean_ttft_ms, 0.0);
542 + assert_eq!(m.median_ttft_ms, 0.0);
543 + assert!(m.percentiles_ttft_ms.is_empty());
544 + assert_eq!(m.mean_tpot_ms, 0.0);
545 + assert_eq!(m.median_tpot_ms, 0.0);
546 + assert_eq!(m.p90_tpot_ms, 0.0);
547 + assert_eq!(m.p99_tpot_ms, 0.0);
548 + // Throughput fields still populated.
549 + assert!(m.input_throughput > 0.0);
550 + }
551 +
552 + #[test]
553 + fn compute_tpot_matches_per_request_definition() {
554 + // 3 requests, each 10s latency with ttft=2s, 5 output tokens.
555 + // Per-request TPOT = (10 - 2) / (5 - 1) = 2s = 2000ms.
556 + // All requests start at t=0 so all are in-window; window duration = 10s.
557 + let itl = vec![2.0; 4];
558 + let outputs = vec![
559 + mk_full(0.0, 10.0, 2.0, itl.clone(), 5),
560 + mk_full(0.0, 10.0, 2.0, itl.clone(), 5),
561 + mk_full(0.0, 10.0, 2.0, itl.clone(), 5),
562 + ];
563 + let requests: Vec<SampleRequest> =
564 + outputs.iter().map(|o| mk_sample(o.prompt_len)).collect();
565 + let w = detect_window(&outputs, Some(3), 1.0, 0.1, 10.0).unwrap();
566 + let m = compute(&outputs, &requests, &w, &[99.0], false);
567 +
568 + assert!((m.mean_tpot_ms - 2000.0).abs() < 1e-6);
569 + assert!((m.median_tpot_ms - 2000.0).abs() < 1e-6);
570 + assert!((m.p90_tpot_ms - 2000.0).abs() < 1e-6);
571 + assert!((m.p99_tpot_ms - 2000.0).abs() < 1e-6);
572 + }
573 +
574 + #[test]
575 + fn json_roundtrip_some() {
576 + let outputs = vec![mk_full(0.0, 10.0, 2.0, vec![2.0; 4], 5); 3];
577 + let requests: Vec<SampleRequest> =
578 + outputs.iter().map(|o| mk_sample(o.prompt_len)).collect();
579 + let w = detect_window(&outputs, Some(3), 1.0, 0.1, 10.0).unwrap();
580 + let m = compute(&outputs, &requests, &w, &[99.0], false);
581 + let s = serde_json::to_string(&m).unwrap();
582 + let parsed: SteadyStateMetrics = serde_json::from_str(&s).unwrap();
583 + assert!((parsed.output_throughput - m.output_throughput).abs() < 1e-12);
584 + assert_eq!(parsed.window.target_concurrency, 3);
585 + }
586 +
587 + #[test]
588 + fn json_deserialize_missing_key_as_none() {
589 + // Pre-feature JSON lacks `steady_state` — deserializing a wrapper that
590 + // contains an Option field with #[serde(default)] should yield None.
591 + #[derive(serde::Deserialize)]
592 + struct Wrap {
593 + #[serde(default)]
594 + steady_state: Option<SteadyStateMetrics>,
595 + }
596 + let old_json = r#"{"other":"value"}"#;
597 + let w: Wrap = serde_json::from_str(old_json).unwrap();
598 + assert!(w.steady_state.is_none());
599 + }
600 +
601 + #[test]
602 + fn end_to_end_trapezoid_run() {
603 + // 5 requests ramp in at t=0..5, each lasts 20s, each emits 10 tokens
604 + // at 1s intervals starting at ttft=1s.
605 + let itl = vec![1.0; 9];
606 + let outputs: Vec<RequestFuncOutput> =
607 + (0..5).map(|i| mk_full(i as f64, 20.0, 1.0, itl.clone(), 10)).collect();
608 + let requests: Vec<SampleRequest> = (0..5).map(|_| mk_sample(100)).collect();
609 +
610 + let w = detect_window(&outputs, Some(5), 1.0, 1.0, 25.0).unwrap();
611 + assert_eq!(w.threshold_abs, 5);
612 + assert!((w.start_s - 4.0).abs() < 1e-9);
613 + assert!((w.end_s - 20.0).abs() < 1e-9);
614 +
615 + let m = compute(&outputs, &requests, &w, &[50.0, 99.0], false);
616 +
617 + // Round-trip through JSON.
618 + let s = serde_json::to_string(&m).unwrap();
619 + let back: SteadyStateMetrics = serde_json::from_str(&s).unwrap();
620 + assert!((back.request_throughput - m.request_throughput).abs() < 1e-12);
621 + assert!((back.output_throughput - m.output_throughput).abs() < 1e-12);
622 + assert_eq!(back.window.target_concurrency, 5);
623 +
624 + // TTFT percentiles should have 2 entries matching the input percentiles.
625 + assert_eq!(m.percentiles_ttft_ms.len(), 2);
626 + assert!((m.percentiles_ttft_ms[0].0 - 50.0).abs() < 1e-9);
627 + assert!((m.percentiles_ttft_ms[1].0 - 99.0).abs() < 1e-9);
628 + }
629 +}
rust/src/bench/src/multi_run.rs new
+357
@@ -0,0 +1,357 @@
1 +// SPDX-License-Identifier: Apache-2.0
2 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +use crate::config::BenchConfig;
5 +use crate::error::Result;
6 +
7 +/// Key metrics extracted from a single run's JSON result.
8 +struct RunMetrics {
9 + request_throughput: f64,
10 + output_throughput: f64,
11 + total_token_throughput: f64,
12 + mean_ttft_ms: f64,
13 + median_ttft_ms: f64,
14 + p99_ttft_ms: f64,
15 + mean_tpot_ms: f64,
16 + median_tpot_ms: f64,
17 + p99_tpot_ms: f64,
18 + mean_itl_ms: f64,
19 + mean_e2el_ms: f64,
20 + p99_e2el_ms: f64,
21 + max_output_tokens_per_s: f64,
22 + completed: f64,
23 + failed: f64,
24 + duration: f64,
25 + ss_request_throughput: Option<f64>,
26 + ss_output_throughput: Option<f64>,
27 + ss_input_throughput: Option<f64>,
28 + ss_total_token_throughput: Option<f64>,
29 + ss_mean_ttft_ms: Option<f64>,
30 + ss_median_ttft_ms: Option<f64>,
31 + ss_mean_tpot_ms: Option<f64>,
32 + ss_median_tpot_ms: Option<f64>,
33 + ss_p90_tpot_ms: Option<f64>,
34 + ss_p99_tpot_ms: Option<f64>,
35 +}
36 +
37 +impl RunMetrics {
38 + fn from_json(json: &serde_json::Value) -> Self {
39 + Self {
40 + request_throughput: get(json, "request_throughput"),
41 + output_throughput: get(json, "output_throughput"),
42 + total_token_throughput: get(json, "total_token_throughput"),
43 + mean_ttft_ms: get(json, "mean_ttft_ms"),
44 + median_ttft_ms: get(json, "median_ttft_ms"),
45 + p99_ttft_ms: get(json, "p99_ttft_ms"),
46 + mean_tpot_ms: get(json, "mean_tpot_ms"),
47 + median_tpot_ms: get(json, "median_tpot_ms"),
48 + p99_tpot_ms: get(json, "p99_tpot_ms"),
49 + mean_itl_ms: get(json, "mean_itl_ms"),
50 + mean_e2el_ms: get(json, "mean_e2el_ms"),
51 + p99_e2el_ms: get(json, "p99_e2el_ms"),
52 + max_output_tokens_per_s: get(json, "max_output_tokens_per_s"),
53 + completed: get(json, "completed"),
54 + failed: get(json, "failed"),
55 + duration: get(json, "duration"),
56 + ss_request_throughput: get_ss_opt(json, "request_throughput"),
57 + ss_output_throughput: get_ss_opt(json, "output_throughput"),
58 + ss_input_throughput: get_ss_opt(json, "input_throughput"),
59 + ss_total_token_throughput: get_ss_opt(json, "total_token_throughput"),
60 + ss_mean_ttft_ms: get_ss_opt(json, "mean_ttft_ms"),
61 + ss_median_ttft_ms: get_ss_opt(json, "median_ttft_ms"),
62 + ss_mean_tpot_ms: get_ss_opt(json, "mean_tpot_ms"),
63 + ss_median_tpot_ms: get_ss_opt(json, "median_tpot_ms"),
64 + ss_p90_tpot_ms: get_ss_opt(json, "p90_tpot_ms"),
65 + ss_p99_tpot_ms: get_ss_opt(json, "p99_tpot_ms"),
66 + }
67 + }
68 +}
69 +
70 +fn get(json: &serde_json::Value, key: &str) -> f64 {
71 + json.get(key).and_then(|v| v.as_f64()).unwrap_or(0.0)
72 +}
73 +
74 +fn get_ss_opt(json: &serde_json::Value, key: &str) -> Option<f64> {
75 + json.get("steady_state")
76 + .and_then(|ss| if ss.is_null() { None } else { Some(ss) })
77 + .and_then(|ss| ss.get(key))
78 + .and_then(|v| v.as_f64())
79 +}
80 +
81 +/// Aggregated statistics for a single metric across N runs.
82 +struct MetricStats {
83 + label: &'static str,
84 + mean: f64,
85 + std: f64,
86 + min: f64,
87 + max: f64,
88 +}
89 +
90 +/// Run the benchmark N times and report aggregated statistics.
91 +pub async fn run_multi(config: &BenchConfig, num_runs: usize) -> Result<()> {
92 + println!(
93 + "{:=^70}",
94 + format!(" Multi-Run Benchmark ({num_runs} runs) ")
95 + );
96 + println!();
97 +
98 + let mut all_runs: Vec<RunMetrics> = Vec::with_capacity(num_runs);
99 +
100 + for i in 0..num_runs {
101 + println!("{:-^70}", format!(" Run {}/{} ", i + 1, num_runs));
102 +
103 + let mut run_config = config.clone();
104 + // Suppress per-run save
105 + run_config.save_result = false;
106 + run_config.append_result = false;
107 +
108 + let result = crate::benchmark::run_benchmark(&run_config).await?;
109 + all_runs.push(RunMetrics::from_json(&result));
110 +
111 + println!();
112 + }
113 +
114 + print_multi_run_summary(&all_runs);
115 + Ok(())
116 +}
117 +
118 +fn print_multi_run_summary(runs: &[RunMetrics]) {
119 + let n = runs.len();
120 +
121 + // Collect each metric into a series, compute stats
122 + let stats = vec![
123 + compute_stats("Request throughput (req/s)", runs, |r| r.request_throughput),
124 + compute_stats("Output throughput (tok/s)", runs, |r| r.output_throughput),
125 + compute_stats("Total token throughput (tok/s)", runs, |r| {
126 + r.total_token_throughput
127 + }),
128 + compute_stats("Peak output tokens/s", runs, |r| r.max_output_tokens_per_s),
129 + compute_stats("Mean TTFT (ms)", runs, |r| r.mean_ttft_ms),
130 + compute_stats("Median TTFT (ms)", runs, |r| r.median_ttft_ms),
131 + compute_stats("P99 TTFT (ms)", runs, |r| r.p99_ttft_ms),
132 + compute_stats("Mean TPOT (ms)", runs, |r| r.mean_tpot_ms),
133 + compute_stats("Median TPOT (ms)", runs, |r| r.median_tpot_ms),
134 + compute_stats("P99 TPOT (ms)", runs, |r| r.p99_tpot_ms),
135 + compute_stats("Mean ITL (ms)", runs, |r| r.mean_itl_ms),
136 + compute_stats("Mean E2EL (ms)", runs, |r| r.mean_e2el_ms),
137 + compute_stats("P99 E2EL (ms)", runs, |r| r.p99_e2el_ms),
138 + compute_stats("Completed requests", runs, |r| r.completed),
139 + compute_stats("Failed requests", runs, |r| r.failed),
140 + compute_stats("Duration (s)", runs, |r| r.duration),
141 + ];
142 +
143 + println!("{:=^80}", format!(" Multi-Run Summary ({n} runs) "));
144 + println!(
145 + "{:<35} {:>10} {:>10} {:>10} {:>10}",
146 + "Metric", "Mean", "Std", "Min", "Max"
147 + );
148 + println!(
149 + "{:-<35} {:->10} {:->10} {:->10} {:->10}",
150 + "", "", "", "", ""
151 + );
152 +
153 + for s in &stats {
154 + println!(
155 + "{:<35} {:>10} {:>10} {:>10} {:>10}",
156 + s.label,
157 + fmt(s.mean),
158 + fmt(s.std),
159 + fmt(s.min),
160 + fmt(s.max),
161 + );
162 + }
163 +
164 + println!("{:=<80}", "");
165 +
166 + // Print coefficient of variation for throughput
167 + let tp = &stats[0]; // request_throughput
168 + if tp.mean > 0.0 {
169 + let cv = (tp.std / tp.mean) * 100.0;
170 + println!(
171 + "Throughput CV: {:.1}% (lower = more stable across runs)",
172 + cv
173 + );
174 + }
175 +
176 + // Steady-state summary — aggregate only runs that have a window.
177 + let ss_stats = vec![
178 + compute_stats_opt("Request throughput (req/s)", runs, |r| {
179 + r.ss_request_throughput
180 + }),
181 + compute_stats_opt("Output throughput (tok/s)", runs, |r| {
182 + r.ss_output_throughput
183 + }),
184 + compute_stats_opt("Input throughput (tok/s)", runs, |r| r.ss_input_throughput),
185 + compute_stats_opt("Total token throughput (tok/s)", runs, |r| {
186 + r.ss_total_token_throughput
187 + }),
188 + compute_stats_opt("Mean TTFT (ms)", runs, |r| r.ss_mean_ttft_ms),
189 + compute_stats_opt("Median TTFT (ms)", runs, |r| r.ss_median_ttft_ms),
190 + compute_stats_opt("Mean TPOT (ms)", runs, |r| r.ss_mean_tpot_ms),
191 + compute_stats_opt("Median TPOT (ms)", runs, |r| r.ss_median_tpot_ms),
192 + compute_stats_opt("P90 TPOT (ms)", runs, |r| r.ss_p90_tpot_ms),
193 + compute_stats_opt("P99 TPOT (ms)", runs, |r| r.ss_p99_tpot_ms),
194 + ];
195 +
196 + let k_present = ss_stats[0].n_present;
197 + let n_total = ss_stats[0].n_total;
198 +
199 + if k_present == 0 {
200 + println!();
201 + println!("Steady-state: no runs had a valid window");
202 + } else {
203 + println!();
204 + println!(
205 + "{:=^80}",
206 + format!(" Steady-State Summary ({k_present}/{n_total} runs) ")
207 + );
208 + if k_present < n_total {
209 + println!(
210 + "Aggregated over {k_present}/{n_total} runs ({} runs had no valid window)",
211 + n_total - k_present
212 + );
213 + }
214 + println!(
215 + "{:<35} {:>10} {:>10} {:>10} {:>10}",
216 + "Metric", "Mean", "Std", "Min", "Max"
217 + );
218 + for s in &ss_stats {
219 + println!(
220 + "{:<35} {:>10.2} {:>10.2} {:>10.2} {:>10.2}",
221 + s.label, s.mean, s.std, s.min, s.max
222 + );
223 + }
224 + }
225 +}
226 +
227 +fn compute_stats<F>(label: &'static str, runs: &[RunMetrics], f: F) -> MetricStats
228 +where
229 + F: Fn(&RunMetrics) -> f64,
230 +{
231 + let values: Vec<f64> = runs.iter().map(&f).collect();
232 + let n = values.len() as f64;
233 + let mean = values.iter().sum::<f64>() / n;
234 + let variance = values.iter().map(|x| (x - mean).powi(2)).sum::<f64>() / n;
235 + let std = variance.sqrt();
236 + let min = values.iter().cloned().fold(f64::INFINITY, f64::min);
237 + let max = values.iter().cloned().fold(f64::NEG_INFINITY, f64::max);
238 +
239 + MetricStats {
240 + label,
241 + mean,
242 + std,
243 + min,
244 + max,
245 + }
246 +}
247 +
248 +fn fmt(v: f64) -> String {
249 + if v == 0.0 {
250 + "0".to_string()
251 + } else if v == v.floor() && v.abs() < 1e9 {
252 + format!("{}", v as i64)
253 + } else {
254 + format!("{:.2}", v)
255 + }
256 +}
257 +
258 +/// Aggregated statistics over only the Some-valued runs.
259 +struct MetricStatsOpt {
260 + label: &'static str,
261 + mean: f64,
262 + std: f64,
263 + min: f64,
264 + max: f64,
265 + n_present: usize,
266 + n_total: usize,
267 +}
268 +
269 +fn compute_stats_opt<F>(label: &'static str, runs: &[RunMetrics], f: F) -> MetricStatsOpt
270 +where
271 + F: Fn(&RunMetrics) -> Option<f64>,
272 +{
273 + let values: Vec<f64> = runs.iter().filter_map(&f).collect();
274 + let n_present = values.len();
275 + let n_total = runs.len();
276 + if values.is_empty() {
277 + return MetricStatsOpt {
278 + label,
279 + mean: 0.0,
280 + std: 0.0,
281 + min: 0.0,
282 + max: 0.0,
283 + n_present,
284 + n_total,
285 + };
286 + }
287 + let mean = values.iter().sum::<f64>() / n_present as f64;
288 + let variance = values.iter().map(|x| (x - mean).powi(2)).sum::<f64>() / n_present as f64;
289 + let std = variance.sqrt();
290 + let min = values.iter().cloned().fold(f64::INFINITY, f64::min);
291 + let max = values.iter().cloned().fold(f64::NEG_INFINITY, f64::max);
292 + MetricStatsOpt {
293 + label,
294 + mean,
295 + std,
296 + min,
297 + max,
298 + n_present,
299 + n_total,
300 + }
301 +}
302 +
303 +#[cfg(test)]
304 +mod tests {
305 + use super::*;
306 +
307 + fn mk_run(ss: Option<f64>) -> RunMetrics {
308 + RunMetrics {
309 + request_throughput: 0.0,
310 + output_throughput: 0.0,
311 + total_token_throughput: 0.0,
312 + mean_ttft_ms: 0.0,
313 + median_ttft_ms: 0.0,
314 + p99_ttft_ms: 0.0,
315 + mean_tpot_ms: 0.0,
316 + median_tpot_ms: 0.0,
317 + p99_tpot_ms: 0.0,
318 + mean_itl_ms: 0.0,
319 + mean_e2el_ms: 0.0,
320 + p99_e2el_ms: 0.0,
321 + max_output_tokens_per_s: 0.0,
322 + completed: 0.0,
323 + failed: 0.0,
324 + duration: 0.0,
325 + ss_request_throughput: ss,
326 + ss_output_throughput: None,
327 + ss_input_throughput: None,
328 + ss_total_token_throughput: None,
329 + ss_mean_ttft_ms: None,
330 + ss_median_ttft_ms: None,
331 + ss_mean_tpot_ms: None,
332 + ss_median_tpot_ms: None,
333 + ss_p90_tpot_ms: None,
334 + ss_p99_tpot_ms: None,
335 + }
336 + }
337 +
338 + #[test]
339 + fn compute_stats_opt_excludes_none() {
340 + let runs = vec![mk_run(Some(10.0)), mk_run(None), mk_run(Some(20.0))];
341 + let s = compute_stats_opt("x", &runs, |r| r.ss_request_throughput);
342 + assert_eq!(s.n_present, 2);
343 + assert_eq!(s.n_total, 3);
344 + assert!((s.mean - 15.0).abs() < 1e-9);
345 + assert_eq!(s.min, 10.0);
346 + assert_eq!(s.max, 20.0);
347 + }
348 +
349 + #[test]
350 + fn compute_stats_opt_all_none() {
351 + let runs = vec![mk_run(None), mk_run(None)];
352 + let s = compute_stats_opt("x", &runs, |r| r.ss_request_throughput);
353 + assert_eq!(s.n_present, 0);
354 + assert_eq!(s.n_total, 2);
355 + assert_eq!(s.mean, 0.0);
356 + }
357 +}
rust/src/bench/src/multi_turn.rs new
+827
@@ -0,0 +1,827 @@
1 +// SPDX-License-Identifier: Apache-2.0
2 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +use std::collections::HashMap;
5 +use std::sync::Arc;
6 +use std::time::Instant;
7 +
8 +use indicatif::{ProgressBar, ProgressStyle};
9 +use tokio::sync::Semaphore;
10 +
11 +use crate::backends::{Backend, RequestFuncInput, RequestFuncOutput, get_backend};
12 +use crate::benchmark::{
13 + assign_lora_modules, compute_spec_decode_stats, fetch_spec_decode_metrics, pre_resolve_dns,
14 + profile_on_batch_threshold, start_profiler_immediate, stop_profiler_immediate,
15 +};
16 +use crate::cli::DatasetName;
17 +use crate::config::BenchConfig;
18 +use crate::datasets::MultiTurnConversation;
19 +use crate::error::{BenchError, Result};
20 +use crate::metrics::calculator::calculate_multi_turn_metrics;
21 +use crate::output::console::print_multi_turn_results;
22 +use crate::output::json::{
23 + append_result, build_multi_turn_result_json, compute_result_filename, save_result,
24 +};
25 +
26 +/// Output from a single turn within a conversation.
27 +#[derive(Debug, Clone)]
28 +pub struct TurnOutput {
29 + pub turn_index: usize,
30 + pub request_output: RequestFuncOutput,
31 + pub cumulative_input_tokens: usize,
32 +}
33 +
34 +/// Output from an entire conversation.
35 +#[derive(Debug, Clone)]
36 +pub struct ConversationOutput {
37 + pub conversation_id: String,
38 + pub turns: Vec<TurnOutput>,
39 + pub total_duration_ms: f64,
40 + pub all_success: bool,
41 +}
42 +
43 +/// Run the multi-turn conversation benchmark.
44 +pub async fn run_multi_turn_benchmark(config: &BenchConfig) -> Result<serde_json::Value> {
45 + let _ = get_backend(config.backend)?;
46 +
47 + // Build HTTP client
48 + let concurrency = config
49 + .multi_turn_concurrency
50 + .unwrap_or(config.max_concurrency.unwrap_or(config.num_prompts));
51 +
52 + let mut client_builder = reqwest::Client::builder()
53 + .pool_max_idle_per_host(concurrency.max(256))
54 + .timeout(std::time::Duration::from_secs(6 * 60 * 60))
55 + .connect_timeout(std::time::Duration::from_secs(30))
56 + .tcp_keepalive(std::time::Duration::from_secs(60))
57 + .tcp_nodelay(true)
58 + .http1_only()
59 + .no_proxy();
60 +
61 + if config.insecure {
62 + client_builder = client_builder.danger_accept_invalid_certs(true);
63 + }
64 +
65 + client_builder = pre_resolve_dns(&config.base_url, client_builder);
66 +
67 + let client = client_builder
68 + .build()
69 + .map_err(|e| BenchError::Backend(format!("Failed to build HTTP client: {e}")))?;
70 +
71 + // Resolve model
72 + let (model_id, model_name) = if let Some(ref m) = config.model {
73 + (m.clone(), config.model_name.clone())
74 + } else {
75 + println!("Model not specified, fetching first model from server...");
76 + let (name, id) = get_first_model(&config.base_url, &client, &config.extra_headers).await?;
77 + println!("First model name: {name}, first model id: {id}");
78 + (id, Some(name))
79 + };
80 +
81 + // Load tokenizer
82 + let tokenizer = if config.skip_tokenizer_init {
83 + None
84 + } else {
85 + let tid = config.tokenizer_id.as_deref().unwrap_or(&model_id);
86 + println!("Loading tokenizer: {tid}");
87 + let server_info = Some((config.base_url.as_str(), model_id.as_str()));
88 + let t = crate::tokenizer::load_tokenizer(tid, config.trust_remote_code, server_info)?;
89 + println!("Tokenizer loaded successfully.");
90 + Some(t)
91 + };
92 +
93 + // Generate/load conversations
94 + println!("Generating multi-turn conversations...");
95 + let gen_start = Instant::now();
96 +
97 + let mut conversations = match config.dataset_name {
98 + DatasetName::Random => {
99 + let tok = tokenizer
100 + .as_ref()
101 + .ok_or_else(|| BenchError::Config("Random dataset requires a tokenizer".into()))?;
102 + let prefix_sharing_config = if config.multi_turn_prefix_global_ratio > 0.0
103 + || config.multi_turn_prefix_conversation_ratio > 0.0
104 + {
105 + Some(crate::datasets::multi_turn::PrefixSharingConfig {
106 + global_ratio: config.multi_turn_prefix_global_ratio,
107 + conversation_ratio: config.multi_turn_prefix_conversation_ratio,
108 + })
109 + } else {
110 + None
111 + };
112 + let random_cfg = crate::datasets::multi_turn::MultiTurnRandomConfig {
113 + num_conversations: config.num_prompts,
114 + min_turns: config.multi_turn_min_turns,
115 + max_turns: config.multi_turn_max_turns,
116 + prefix_len: config.random_prefix_len,
117 + input_len: config.random_input_len,
118 + per_turn_input_len: config.per_turn_input_len,
119 + output_len: config.random_output_len,
120 + seed: config.seed,
121 + request_id_prefix: config.request_id_prefix.clone(),
122 + prefix_sharing_config,
123 + };
124 + crate::datasets::multi_turn::generate_multi_turn_random(tok, &random_cfg)?
125 + }
126 + DatasetName::ShareGpt => {
127 + let tok = tokenizer.as_ref().ok_or_else(|| {
128 + BenchError::Config("ShareGPT dataset requires a tokenizer".into())
129 + })?;
130 + let downloaded;
131 + let path = match config.dataset_path.as_deref() {
132 + Some(p) => p,
133 + None => {
134 + downloaded = crate::datasets::sharegpt::download_sharegpt_dataset()?;
135 + downloaded.as_str()
136 + }
137 + };
138 + crate::datasets::multi_turn::load_sharegpt_multi_turn(
139 + tok,
140 + path,
141 + config.num_prompts,
142 + config.sharegpt_output_len,
143 + config.sharegpt_multi_turn_max_turns,
144 + config.seed,
145 + &config.request_id_prefix,
146 + )?
147 + }
148 + DatasetName::RandomMm => {
149 + return Err(BenchError::Config(
150 + "Random-MM multi-turn is not yet supported. Use 'random' or 'sharegpt' with --multi-turn.".into(),
151 + ));
152 + }
153 + DatasetName::Sonnet => {
154 + return Err(BenchError::Config(
155 + "Sonnet multi-turn is not yet supported. Use 'random' or 'sharegpt' with --multi-turn.".into(),
156 + ));
157 + }
158 + DatasetName::SpeedBench => {
159 + return Err(BenchError::Config(
160 + "SPEED-Bench multi-turn is not yet supported. Use 'random' or 'sharegpt' with --multi-turn.".into(),
161 + ));
162 + }
163 + DatasetName::Custom | DatasetName::PrefixRepetition | DatasetName::RandomRerank => {
164 + return Err(BenchError::Config(
165 + "This dataset does not support multi-turn. Use 'random' or 'sharegpt' with --multi-turn.".into(),
166 + ));
167 + }
168 + DatasetName::Hf => {
169 + return Err(BenchError::Config(
170 + "HF dataset multi-turn is not yet supported. Use 'random' or 'sharegpt' with --multi-turn.".into(),
171 + ));
172 + }
173 + };
174 +
175 + let no_history = config.multi_turn_prefix_global_ratio > 0.0
176 + || config.multi_turn_prefix_conversation_ratio > 0.0;
177 +
178 + if let Some(max_model_len) = config.max_model_len {
179 + let (filtered_conversations, filtered_turns) =
180 + filter_turns_by_max_model_len(&mut conversations, max_model_len, no_history);
181 + if filtered_turns > 0 || filtered_conversations > 0 {
182 + println!(
183 + "Filtered {filtered_turns} turn(s) and {filtered_conversations} conversation(s) above --max-model-len {max_model_len}."
184 + );
185 + }
186 + if conversations.is_empty() {
187 + return Err(BenchError::Config(
188 + "No conversations remain after applying --max-model-len".into(),
189 + ));
190 + }
191 + }
192 +
193 + let gen_elapsed = gen_start.elapsed();
194 + let total_turns: usize = conversations.iter().map(|c| c.turns.len()).sum();
195 + println!(
196 + "Generated {} conversations ({} total turns) in {:.2}s",
197 + conversations.len(),
198 + total_turns,
199 + gen_elapsed.as_secs_f64()
200 + );
201 +
202 + // Log prefix sharing info
203 + if no_history {
204 + let num_special = tokenizer.as_ref().map(|t| t.num_special_tokens_to_add()).unwrap_or(0);
205 + let real_input_len = config.random_input_len.saturating_sub(num_special);
206 + let global_tokens =
207 + (real_input_len as f64 * config.multi_turn_prefix_global_ratio).floor() as usize;
208 + let conv_tokens =
209 + (real_input_len as f64 * config.multi_turn_prefix_conversation_ratio).floor() as usize;
210 + let unique_tokens = real_input_len.saturating_sub(global_tokens + conv_tokens);
211 + println!(
212 + "User message prefix sharing: {:.0}% global ({} tokens), {:.0}% per-conversation ({} tokens), {:.0}% unique ({} tokens)",
213 + config.multi_turn_prefix_global_ratio * 100.0,
214 + global_tokens,
215 + config.multi_turn_prefix_conversation_ratio * 100.0,
216 + conv_tokens,
217 + (1.0 - config.multi_turn_prefix_global_ratio
218 + - config.multi_turn_prefix_conversation_ratio)
219 + * 100.0,
220 + unique_tokens,
221 + );
222 + println!("No history accumulation: each turn sends fixed-length prompt only.");
223 + }
224 +
225 + if config.dry_run {
226 + let total_user_tokens: usize = conversations
227 + .iter()
228 + .flat_map(|c| c.turns.iter())
229 + .map(|t| t.user_message_len)
230 + .sum();
231 + println!("Dry run stats:");
232 + println!(" Total conversations: {}", conversations.len());
233 + println!(" Total turns: {total_turns}");
234 + println!(" Total user message tokens: {total_user_tokens}");
235 + return Ok(serde_json::json!({"dry_run": true, "mode": "multi_turn"}));
236 + }
237 +
238 + // Ready check with a simple single request
239 + if config.ready_check_timeout_sec > 0 {
240 + let first_turn = &conversations[0].turns[0];
241 + let test_input = RequestFuncInput {
242 + prompt: first_turn.user_message.clone(),
243 + api_url: config.api_url.clone(),
244 + prompt_len: first_turn.user_message_len,
245 + output_len: first_turn.expected_output_len,
246 + model: model_id.clone(),
247 + model_name: model_name.clone(),
248 + logprobs: config.logprobs,
249 + extra_headers: config.extra_headers.clone(),
250 + extra_body: config.extra_body.clone(),
251 + ignore_eos: config.ignore_eos,
252 + request_id: None,
253 + ..Default::default()
254 + };
255 +
256 + println!("Starting initial single prompt test run...");
257 + let test_output = crate::ready_checker::wait_for_endpoint(
258 + config.backend,
259 + &client,
260 + &test_input,
261 + config.ready_check_timeout_sec,
262 + 5,
263 + )
264 + .await?;
265 + if !test_output.success {
266 + return Err(BenchError::Backend(format!(
267 + "Initial test run failed: {}",
268 + test_output.error
269 + )));
270 + }
271 + println!("Initial test run completed.");
272 + }
273 +
274 + // For random datasets in multi-turn mode, auto-set min_tokens to enforce
275 + // output length without using ignore_eos (which causes unbounded context growth).
276 + // min_tokens + max_completion_tokens together control output length precisely.
277 + let extra_body = if config.dataset_name == DatasetName::Random && !config.ignore_eos {
278 + let mut body = config.extra_body.clone().unwrap_or_else(|| serde_json::json!({}));
279 + if let serde_json::Value::Object(ref mut map) = body
280 + && !map.contains_key("min_tokens")
281 + {
282 + map.insert(
283 + "min_tokens".to_string(),
284 + serde_json::json!(config.random_output_len),
285 + );
286 + println!(
287 + "Auto-setting min_tokens={} for multi-turn random dataset (use --extra-body to override)",
288 + config.random_output_len
289 + );
290 + }
291 + Some(body)
292 + } else {
293 + config.extra_body.clone()
294 + };
295 +
296 + // Fetch speculative decoding metrics before benchmark
297 + let spec_decode_before =
298 + fetch_spec_decode_metrics(&config.base_url, &client, &config.extra_headers).await;
299 + if spec_decode_before.is_some() {
300 + println!("Speculative decoding detected, will collect metrics.");
301 + }
302 +
303 + // Start profiler if requested (immediate mode — no batch threshold)
304 + if config.profile && config.profile_batch_threshold.is_none() {
305 + start_profiler_immediate(&client, &config.base_url, &config.extra_headers).await;
306 + }
307 +
308 + // Threshold-based profiling: spawn background task that polls /metrics
309 + // and triggers start/stop profile when batch size is reached.
310 + let profile_task = if let Some(threshold) = config.profile_batch_threshold {
311 + let poll_client = client.clone();
312 + let base_url = config.base_url.clone();
313 + let extra_headers = config.extra_headers.clone();
314 + let duration_secs = config.profile_duration;
315 + let (cancel_tx, cancel_rx) = tokio::sync::oneshot::channel::<()>();
316 + let handle = tokio::spawn(async move {
317 + profile_on_batch_threshold(
318 + &poll_client,
319 + &base_url,
320 + &extra_headers,
321 + threshold,
322 + duration_secs,
323 + cancel_rx,
324 + )
325 + .await;
326 + });
327 + Some((cancel_tx, handle))
328 + } else {
329 + None
330 + };
331 +
332 + // Main benchmark
333 + println!("Starting multi-turn benchmark...");
334 + println!("Conversations: {}", conversations.len());
335 + println!("Concurrency: {concurrency}");
336 + println!("Inter-turn delay: {} ms", config.multi_turn_delay_ms);
337 +
338 + let max_turn_count = conversations.iter().map(|c| c.turns.len()).max().unwrap_or(0);
339 +
340 + // Progress bar counts total turns
341 + let pb = if config.disable_tqdm {
342 + None
343 + } else {
344 + let bar = ProgressBar::new(total_turns as u64);
345 + bar.set_style(
346 + ProgressStyle::with_template(
347 + "{spinner:.green} [{elapsed_precise}] [{bar:40.cyan/blue}] {pos}/{len} turns ({eta})",
348 + )
349 + .unwrap()
350 + .progress_chars("#>-"),
351 + );
352 + Some(bar)
353 + };
354 +
355 + let benchmark_start = Instant::now();
356 +
357 + // Per-conversation LoRA assignment (sticky: all turns of a conversation
358 + // share the same adapter). None when --lora-modules not set.
359 + let lora_assignments = assign_lora_modules(
360 + &config.lora_modules,
361 + config.lora_assignment,
362 + conversations.len(),
363 + config.seed,
364 + );
365 + if let Some(modules) = config.lora_modules.as_ref() {
366 + let names: Vec<&str> = modules.iter().map(|s| s.as_ref()).collect();
367 + println!(
368 + "LoRA adapters ({}): {:?} [assignment={:?}, scope=conversation]",
369 + modules.len(),
370 + names,
371 + config.lora_assignment
372 + );
373 + }
374 +
375 + // Shared state
376 + let backend = get_backend(config.backend)?;
377 + let api_url = Arc::new(config.api_url.clone());
378 + let model = Arc::new(model_id.clone());
379 + let model_name_arc = Arc::new(model_name.clone());
380 + let extra_body = Arc::new(extra_body);
381 + let base_extra_headers = Arc::new(config.extra_headers.clone());
382 + let ignore_eos = config.ignore_eos;
383 + let logprobs = config.logprobs;
384 + let delay_ms = config.multi_turn_delay_ms;
385 +
386 + // Semaphore controls max in-flight requests (not conversations).
387 + // Each conversation acquires the permit only during the HTTP request,
388 + // releasing it before the inter-turn delay so other conversations can
389 + // immediately fill the slot.
390 + let semaphore = Arc::new(Semaphore::new(concurrency));
391 +
392 + // Spawn one task per conversation (all at once)
393 + let mut handles = Vec::with_capacity(conversations.len());
394 + for (i, conv) in conversations.into_iter().enumerate() {
395 + let client = client.clone();
396 + let backend = backend.clone();
397 + let api_url = api_url.clone();
398 + let model = model.clone();
399 + let model_name = model_name_arc.clone();
400 + let extra_body = extra_body.clone();
401 + let base_extra_headers = base_extra_headers.clone();
402 + let pb = pb.clone();
403 + let bench_start = benchmark_start;
404 + let semaphore = semaphore.clone();
405 + let lora_name = lora_assignments.as_ref().map(|v| v[i].clone());
406 +
407 + handles.push(tokio::spawn(async move {
408 + run_conversation(
409 + &conv,
410 + &backend,
411 + &client,
412 + &api_url,
413 + &model,
414 + &model_name,
415 + lora_name.as_deref(),
416 + &extra_body,
417 + &base_extra_headers,
418 + ignore_eos,
419 + logprobs,
420 + delay_ms,
421 + no_history,
422 + bench_start,
423 + pb.as_ref(),
424 + &semaphore,
425 + )
426 + .await
427 + }));
428 + }
429 +
430 + // Collect all conversation outputs
431 + let mut all_outputs: Vec<ConversationOutput> = Vec::with_capacity(handles.len());
432 + for handle in handles {
433 + match handle.await {
434 + Ok(output) => all_outputs.push(output),
435 + Err(e) => {
436 + eprintln!("Conversation task panicked: {e}");
437 + }
438 + }
439 + }
440 +
441 + if let Some(ref pb) = pb {
442 + pb.finish_and_clear();
443 + }
444 +
445 + let benchmark_duration = benchmark_start.elapsed().as_secs_f64();
446 +
447 + // Stop profiler if requested (immediate mode — no batch threshold)
448 + if config.profile && config.profile_batch_threshold.is_none() {
449 + stop_profiler_immediate(&client, &config.base_url, &config.extra_headers).await;
450 + }
451 +
452 + // Signal the threshold-based profile task that the benchmark is done, then wait
453 + if let Some((cancel_tx, task)) = profile_task {
454 + let _ = cancel_tx.send(());
455 + if let Err(e) = task.await {
456 + eprintln!("WARNING: Profile background task failed: {e}");
457 + }
458 + }
459 +
460 + // Fetch speculative decoding metrics after benchmark and compute stats
461 + let spec_decode_stats = if spec_decode_before.is_some() {
462 + let spec_decode_after =
463 + fetch_spec_decode_metrics(&config.base_url, &client, &config.extra_headers).await;
464 + match (spec_decode_before.as_ref(), spec_decode_after.as_ref()) {
465 + (Some(before), Some(after)) => compute_spec_decode_stats(before, after),
466 + _ => None,
467 + }
468 + } else {
469 + None
470 + };
471 +
472 + // Calculate metrics
473 + let mt_metrics = calculate_multi_turn_metrics(
474 + &all_outputs,
475 + benchmark_duration,
476 + &config.selected_percentiles,
477 + &config.goodput,
478 + max_turn_count,
479 + );
480 +
481 + // Print console output
482 + print_multi_turn_results(
483 + &mt_metrics,
484 + benchmark_duration,
485 + config,
486 + spec_decode_stats.as_ref(),
487 + );
488 +
489 + // Build result JSON
490 + let date_iso = chrono::Utc::now().format("%Y-%m-%dT%H:%M:%SZ").to_string();
491 + let dt_filename = chrono::Local::now().format("%Y%m%d-%H%M%S").to_string();
492 + let result_json = build_multi_turn_result_json(
493 + config,
494 + &mt_metrics,
495 + &all_outputs,
496 + benchmark_duration,
497 + &date_iso,
498 + spec_decode_stats.as_ref(),
499 + );
500 +
501 + // Save if requested
502 + if config.save_result || config.append_result {
503 + let model_for_filename = config.model.as_deref().unwrap_or(&model_id);
504 + let file_name = compute_result_filename(config, model_for_filename, &dt_filename);
505 +
506 + if let Some(ref dir) = config.result_dir {
507 + std::fs::create_dir_all(dir)?;
508 + }
509 +
510 + if config.append_result {
511 + append_result(&result_json, &file_name)?;
512 + } else {
513 + save_result(&result_json, &file_name)?;
514 + }
515 + }
516 +
517 + Ok(result_json)
518 +}
519 +
520 +fn filter_turns_by_max_model_len(
521 + conversations: &mut Vec<MultiTurnConversation>,
522 + max_model_len: usize,
523 + no_history: bool,
524 +) -> (usize, usize) {
525 + let before_conversations = conversations.len();
526 + let before_turns: usize = conversations.iter().map(|c| c.turns.len()).sum();
527 +
528 + for conversation in conversations.iter_mut() {
529 + if no_history {
530 + conversation.turns.retain(|turn| {
531 + turn.user_message_len.saturating_add(turn.expected_output_len) <= max_model_len
532 + });
533 + continue;
534 + }
535 +
536 + let keep_turns = valid_prefix_len_for_max_model_len(conversation, max_model_len);
537 + conversation.turns.truncate(keep_turns);
538 + }
539 +
540 + conversations.retain(|conversation| conversation.turns.len() >= 2);
541 +
542 + let after_conversations = conversations.len();
543 + let after_turns: usize = conversations.iter().map(|c| c.turns.len()).sum();
544 +
545 + (
546 + before_conversations - after_conversations,
547 + before_turns - after_turns,
548 + )
549 +}
550 +
551 +fn valid_prefix_len_for_max_model_len(
552 + conversation: &MultiTurnConversation,
553 + max_model_len: usize,
554 +) -> usize {
555 + let mut cumulative_input_tokens = 0usize;
556 + let mut keep_turns = 0usize;
557 +
558 + for turn in &conversation.turns {
559 + cumulative_input_tokens = cumulative_input_tokens.saturating_add(turn.user_message_len);
560 +
561 + if cumulative_input_tokens.saturating_add(turn.expected_output_len) > max_model_len {
562 + break;
563 + }
564 +
565 + cumulative_input_tokens = cumulative_input_tokens.saturating_add(turn.expected_output_len);
566 + keep_turns += 1;
567 + }
568 +
569 + keep_turns
570 +}
571 +
572 +/// Run a single conversation: sequential turns, building up message history.
573 +///
574 +/// `lora_name`, when set, replaces both `model` and `model_name` in every
575 +/// turn's request payload — sticky for the whole conversation so prefix-cache
576 +/// reuse across turns isn't broken by mid-conversation adapter switches.
577 +async fn run_conversation(
578 + conversation: &MultiTurnConversation,
579 + backend: &Backend,
580 + client: &reqwest::Client,
581 + api_url: &str,
582 + model: &str,
583 + model_name: &Option<String>,
584 + lora_name: Option<&str>,
585 + extra_body: &Option<serde_json::Value>,
586 + base_extra_headers: &Option<HashMap<String, String>>,
587 + ignore_eos: bool,
588 + logprobs: Option<usize>,
589 + delay_ms: u64,
590 + no_history: bool,
591 + bench_start: Instant,
592 + pb: Option<&ProgressBar>,
593 + semaphore: &Semaphore,
594 +) -> ConversationOutput {
595 + let conv_start = Instant::now();
596 + let mut messages: Vec<serde_json::Value> = Vec::new();
597 + let mut cumulative_tokens: usize = 0;
598 + let mut turn_outputs: Vec<TurnOutput> = Vec::new();
599 + let mut all_success = true;
600 +
601 + // Router affinity: all turns share same X-Session-ID
602 + let mut extra_headers = base_extra_headers.clone().unwrap_or_default();
603 + extra_headers.insert(
604 + "X-Session-ID".to_string(),
605 + conversation.conversation_id.clone(),
606 + );
607 +
608 + for (turn_idx, turn) in conversation.turns.iter().enumerate() {
609 + if no_history {
610 + // Prefix sharing mode: reset messages each turn, send only this turn's message
611 + messages.clear();
612 + messages.push(serde_json::json!({
613 + "role": "user",
614 + "content": [{"type": "text", "text": &*turn.user_message}]
615 + }));
616 + cumulative_tokens = turn.user_message_len;
617 + } else {
618 + // Normal mode: accumulate history
619 + messages.push(serde_json::json!({
620 + "role": "user",
621 + "content": [{"type": "text", "text": &*turn.user_message}]
622 + }));
623 + cumulative_tokens += turn.user_message_len;
624 + }
625 +
626 + // Set min_tokens per-request so each turn's output length matches
627 + // the dataset's expected length. Skip if ignore_eos is set (unbounded
628 + // generation) or the user already provided min_tokens via --extra-body.
629 + let turn_extra_body = if !ignore_eos {
630 + let mut body = extra_body.clone().unwrap_or_else(|| serde_json::json!({}));
631 + if let serde_json::Value::Object(ref mut map) = body
632 + && !map.contains_key("min_tokens")
633 + {
634 + map.insert(
635 + "min_tokens".to_string(),
636 + serde_json::json!(turn.expected_output_len),
637 + );
638 + }
639 + Some(body)
640 + } else {
641 + extra_body.clone()
642 + };
643 +
644 + let (req_model, req_model_name) = match lora_name {
645 + Some(name) => (name.to_string(), Some(name.to_string())),
646 + None => (model.to_string(), model_name.clone()),
647 + };
648 +
649 + let input = RequestFuncInput {
650 + prompt: turn.user_message.clone(),
651 + api_url: api_url.to_string(),
652 + prompt_len: cumulative_tokens,
653 + output_len: turn.expected_output_len,
654 + model: req_model,
655 + model_name: req_model_name,
656 + logprobs,
657 + extra_headers: Some(extra_headers.clone()),
658 + extra_body: turn_extra_body,
659 + ignore_eos,
660 + request_id: Some(format!("{}-turn{}", conversation.conversation_id, turn_idx)),
661 + messages: Some(serde_json::json!(messages)),
662 + ..Default::default()
663 + };
664 +
665 + // Acquire semaphore permit before sending the request.
666 + // Released after the response so the slot is free during inter-turn delay.
667 + let _permit = match semaphore.acquire().await {
668 + Ok(p) => p,
669 + Err(_) => {
670 + all_success = false;
671 + break;
672 + }
673 + };
674 + let request_instant = Instant::now();
675 + let result = backend.send_request(&input, client).await;
676 + drop(_permit);
677 +
678 + let output = match result {
679 + Ok(mut o) => {
680 + o.start_time = request_instant.duration_since(bench_start).as_secs_f64();
681 + o
682 + }
683 + Err(e) => RequestFuncOutput {
684 + success: false,
685 + error: e.to_string(),
686 + prompt_len: cumulative_tokens,
687 + start_time: request_instant.duration_since(bench_start).as_secs_f64(),
688 + ..Default::default()
689 + },
690 + };
691 +
692 + if let Some(pb) = pb {
693 + pb.inc(1);
694 + }
695 +
696 + if output.success {
697 + if !no_history {
698 + // Add assistant response to history for next turn
699 + messages.push(serde_json::json!({
700 + "role": "assistant",
701 + "content": &output.generated_text
702 + }));
703 + cumulative_tokens += output.output_tokens;
704 + }
705 +
706 + turn_outputs.push(TurnOutput {
707 + turn_index: turn_idx,
708 + request_output: output,
709 + cumulative_input_tokens: cumulative_tokens,
710 + });
711 + } else {
712 + all_success = false;
713 + turn_outputs.push(TurnOutput {
714 + turn_index: turn_idx,
715 + request_output: output,
716 + cumulative_input_tokens: cumulative_tokens,
717 + });
718 + break; // Conversation stops on failure
719 + }
720 +
721 + // Inter-turn delay (semaphore NOT held — slot is free for others)
722 + if delay_ms > 0 && turn_idx + 1 < conversation.turns.len() {
723 + tokio::time::sleep(std::time::Duration::from_millis(delay_ms)).await;
724 + }
725 + }
726 +
727 + let total_duration_ms = conv_start.elapsed().as_secs_f64() * 1000.0;
728 +
729 + ConversationOutput {
730 + conversation_id: conversation.conversation_id.clone(),
731 + turns: turn_outputs,
732 + total_duration_ms,
733 + all_success,
734 + }
735 +}
736 +
737 +/// Fetch the first model from the server's /v1/models endpoint.
738 +async fn get_first_model(
739 + base_url: &str,
740 + client: &reqwest::Client,
741 + extra_headers: &Option<HashMap<String, String>>,
742 +) -> Result<(String, String)> {
743 + let url = format!("{base_url}/v1/models");
744 + let mut request = client.get(&url);
745 + if let Some(headers) = extra_headers {
746 + for (k, v) in headers {
747 + request = request.header(k, v);
748 + }
749 + }
750 + if let Ok(api_key) = std::env::var("OPENAI_API_KEY") {
751 + request = request.header("Authorization", format!("Bearer {api_key}"));
752 + }
753 +
754 + let response = request.send().await?;
755 + let data: serde_json::Value = response.json().await?;
756 +
757 + if let Some(models) = data.get("data").and_then(|d| d.as_array())
758 + && let Some(first) = models.first()
759 + {
760 + let id = first.get("id").and_then(|v| v.as_str()).unwrap_or_default().to_string();
761 + let root = first.get("root").and_then(|v| v.as_str()).unwrap_or(&id).to_string();
762 + return Ok((id, root));
763 + }
764 +
765 + Err(BenchError::Config(format!(
766 + "No models found on the server at {base_url}"
767 + )))
768 +}
769 +
770 +#[cfg(test)]
771 +mod tests {
772 + use std::sync::Arc;
773 +
774 + use super::{filter_turns_by_max_model_len, valid_prefix_len_for_max_model_len};
775 + use crate::datasets::{ConversationTurn, MultiTurnConversation};
776 +
777 + fn conversation(turns: &[(usize, usize)]) -> MultiTurnConversation {
778 + MultiTurnConversation {
779 + conversation_id: "conv-0".to_string(),
780 + turns: turns
781 + .iter()
782 + .map(|(user_message_len, expected_output_len)| ConversationTurn {
783 + user_message: Arc::from("hello"),
784 + user_message_len: *user_message_len,
785 + expected_output_len: *expected_output_len,
786 + })
787 + .collect(),
788 + }
789 + }
790 +
791 + #[test]
792 + fn test_valid_prefix_len_for_max_model_len_with_history() {
793 + let conv = conversation(&[(40, 10), (45, 10)]);
794 +
795 + assert_eq!(valid_prefix_len_for_max_model_len(&conv, 105), 2);
796 + assert_eq!(valid_prefix_len_for_max_model_len(&conv, 104), 1);
797 + }
798 +
799 + #[test]
800 + fn test_filter_turns_by_max_model_len_truncates_history_conversations() {
801 + let mut conversations = vec![
802 + conversation(&[(40, 10), (45, 10), (1, 1)]),
803 + conversation(&[(100, 10), (1, 1)]),
804 + ];
805 +
806 + let (filtered_conversations, filtered_turns) =
807 + filter_turns_by_max_model_len(&mut conversations, 105, false);
808 +
809 + assert_eq!(filtered_conversations, 1);
810 + assert_eq!(filtered_turns, 3);
811 + assert_eq!(conversations.len(), 1);
812 + assert_eq!(conversations[0].turns.len(), 2);
813 + }
814 +
815 + #[test]
816 + fn test_filter_turns_by_max_model_len_retains_independent_no_history_turns() {
817 + let mut conversations = vec![conversation(&[(40, 10), (95, 10), (1, 1)])];
818 +
819 + let (filtered_conversations, filtered_turns) =
820 + filter_turns_by_max_model_len(&mut conversations, 104, true);
821 +
822 + assert_eq!(filtered_conversations, 0);
823 + assert_eq!(filtered_turns, 1);
824 + assert_eq!(conversations.len(), 1);
825 + assert_eq!(conversations[0].turns.len(), 2);
826 + }
827 +}
rust/src/bench/src/output/console.rs new
+400
@@ -0,0 +1,400 @@
1 +// SPDX-License-Identifier: Apache-2.0
2 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +use crate::benchmark::SpecDecodeStats;
5 +use crate::config::BenchConfig;
6 +use crate::metrics::{BenchmarkMetrics, MultiTurnMetrics};
7 +
8 +/// Print benchmark results to console in the same format as Python.
9 +///
10 +/// Mirrors the output format from serve.py:918-1093.
11 +pub fn print_results(
12 + metrics: &BenchmarkMetrics,
13 + benchmark_duration: f64,
14 + config: &BenchConfig,
15 + has_tokenizer: bool,
16 + spec_decode_stats: Option<&SpecDecodeStats>,
17 +) {
18 + let is_pooling = config.backend.is_pooling();
19 +
20 + if is_pooling {
21 + println!("{:=^60}", " Embedding/Pooling Benchmark Result ");
22 + } else {
23 + println!("{:=^60}", " Serving Benchmark Result ");
24 + }
25 + println!("{:<40} {:<10}", "Successful requests:", metrics.completed);
26 + println!("{:<40} {:<10}", "Failed requests:", metrics.failed);
27 +
28 + if let Some(mc) = config.max_concurrency {
29 + println!("{:<40} {:<10}", "Maximum request concurrency:", mc);
30 + }
31 + if !config.request_rate.is_infinite() {
32 + println!(
33 + "{:<40} {:<10.2}",
34 + "Request rate configured (RPS):", config.request_rate
35 + );
36 + }
37 + println!(
38 + "{:<40} {:<10.2}",
39 + "Benchmark duration (s):", benchmark_duration
40 + );
41 + println!("{:<40} {:<10}", "Total input tokens:", metrics.total_input);
42 +
43 + if has_tokenizer && !is_pooling {
44 + println!(
45 + "{:<40} {:<10}",
46 + "Total generated tokens:", metrics.total_output
47 + );
48 + }
49 +
50 + println!(
51 + "{:<40} {:<10.2}",
52 + "Request throughput (req/s):", metrics.request_throughput
53 + );
54 + if metrics.request_goodput > 0.0 {
55 + println!(
56 + "{:<40} {:<10.2}",
57 + "Request goodput (req/s):", metrics.request_goodput
58 + );
59 + }
60 +
61 + if !is_pooling {
62 + if has_tokenizer {
63 + println!(
64 + "{:<40} {:<10.2}",
65 + "Output token throughput (tok/s):", metrics.output_throughput
66 + );
67 + println!(
68 + "{:<40} {:<10.2}",
69 + "Peak output token throughput (tok/s):", metrics.max_output_tokens_per_s
70 + );
71 + }
72 +
73 + println!(
74 + "{:<40} {:<10.2}",
75 + "Peak concurrent requests:", metrics.max_concurrent_requests as f64
76 + );
77 +
78 + if has_tokenizer {
79 + println!(
80 + "{:<40} {:<10.2}",
81 + "Total token throughput (tok/s):", metrics.total_token_throughput
82 + );
83 + }
84 + } else {
85 + println!(
86 + "{:<40} {:<10.2}",
87 + "Input token throughput (tok/s):", metrics.input_throughput
88 + );
89 + println!(
90 + "{:<40} {:<10.2}",
91 + "Peak concurrent requests:", metrics.max_concurrent_requests as f64
92 + );
93 + }
94 +
95 + // Print per-metric percentiles
96 + if has_tokenizer && !is_pooling {
97 + print_metric_section(
98 + "ttft",
99 + "TTFT",
100 + "Time to First Token",
101 + &config.selected_percentile_metrics,
102 + metrics.mean_ttft_ms,
103 + metrics.median_ttft_ms,
104 + &metrics.percentiles_ttft_ms,
105 + );
106 + print_metric_section(
107 + "tpot",
108 + "TPOT",
109 + "Time per Output Token (excl. 1st token)",
110 + &config.selected_percentile_metrics,
111 + metrics.mean_tpot_ms,
112 + metrics.median_tpot_ms,
113 + &metrics.percentiles_tpot_ms,
114 + );
115 + print_metric_section(
116 + "itl",
117 + "ITL",
118 + "Inter-token Latency",
119 + &config.selected_percentile_metrics,
120 + metrics.mean_itl_ms,
121 + metrics.median_itl_ms,
122 + &metrics.percentiles_itl_ms,
123 + );
124 + }
125 + print_metric_section(
126 + "e2el",
127 + "E2EL",
128 + "End-to-end Latency",
129 + &config.selected_percentile_metrics,
130 + metrics.mean_e2el_ms,
131 + metrics.median_e2el_ms,
132 + &metrics.percentiles_e2el_ms,
133 + );
134 +
135 + if let Some(stats) = spec_decode_stats {
136 + print_spec_decode_section(stats);
137 + }
138 +
139 + println!("{:=<60}", "");
140 +
141 + print_steady_state(metrics);
142 +}
143 +
144 +/// Print multi-turn benchmark results to console.
145 +pub fn print_multi_turn_results(
146 + mt_metrics: &MultiTurnMetrics,
147 + benchmark_duration: f64,
148 + config: &BenchConfig,
149 + spec_decode_stats: Option<&SpecDecodeStats>,
150 +) {
151 + println!("{:=^60}", " Multi-Turn Benchmark Result ");
152 + println!(
153 + "{:<45} {}/{}",
154 + "Conversations completed/total:",
155 + mt_metrics.conversations_completed,
156 + mt_metrics.conversations_completed + mt_metrics.conversations_failed,
157 + );
158 + println!(
159 + "{:<45} {}",
160 + "Turns per conversation (configured):", config.multi_turn_num_turns,
161 + );
162 + println!(
163 + "{:<45} {:.1}",
164 + "Avg turns completed:", mt_metrics.avg_turns_completed,
165 + );
166 + println!(
167 + "{:<45} {:.0}",
168 + "Avg conversation duration (ms):", mt_metrics.avg_conversation_duration_ms,
169 + );
170 + println!(
171 + "{:<45} {}",
172 + "Concurrency:",
173 + config
174 + .multi_turn_concurrency
175 + .unwrap_or(config.max_concurrency.unwrap_or(config.num_prompts)),
176 + );
177 + println!(
178 + "{:<45} {}",
179 + "Inter-turn delay (ms):", config.multi_turn_delay_ms,
180 + );
181 + println!(
182 + "{:<45} {:.2}",
183 + "Benchmark duration (s):", benchmark_duration,
184 + );
185 +
186 + // Overall metrics
187 + println!("{:-^60}", " Overall (All Turns) ");
188 + print_metrics_block(&mt_metrics.overall, config);
189 +
190 + // Per-turn breakdown
191 + for (i, turn_metrics) in mt_metrics.per_turn.iter().enumerate() {
192 + let samples = turn_metrics.completed + turn_metrics.failed;
193 + println!("{:-^60}", format!(" Turn {} ({} samples) ", i + 1, samples));
194 + print_metrics_block(turn_metrics, config);
195 + }
196 +
197 + if let Some(stats) = spec_decode_stats {
198 + print_spec_decode_section(stats);
199 + }
200 +
201 + println!("{:=<60}", "");
202 +}
203 +
204 +/// Print a metrics block (reused for overall and per-turn).
205 +fn print_metrics_block(metrics: &BenchmarkMetrics, config: &BenchConfig) {
206 + println!("{:<45} {:<10}", "Successful requests:", metrics.completed);
207 + println!("{:<45} {:<10}", "Failed requests:", metrics.failed);
208 + println!("{:<45} {:<10}", "Total input tokens:", metrics.total_input);
209 + println!(
210 + "{:<45} {:<10}",
211 + "Total generated tokens:", metrics.total_output
212 + );
213 + println!(
214 + "{:<45} {:<10.2}",
215 + "Request throughput (req/s):", metrics.request_throughput
216 + );
217 + println!(
218 + "{:<45} {:<10.2}",
219 + "Input token throughput (tok/s):", metrics.input_throughput
220 + );
221 + println!(
222 + "{:<45} {:<10.2}",
223 + "Output token throughput (tok/s):", metrics.output_throughput
224 + );
225 + println!(
226 + "{:<45} {:<10.2}",
227 + "Total token throughput (tok/s):", metrics.total_token_throughput
228 + );
229 +
230 + print_metric_section(
231 + "ttft",
232 + "TTFT",
233 + "Time to First Token",
234 + &config.selected_percentile_metrics,
235 + metrics.mean_ttft_ms,
236 + metrics.median_ttft_ms,
237 + &metrics.percentiles_ttft_ms,
238 + );
239 + print_metric_section(
240 + "tpot",
241 + "TPOT",
242 + "Time per Output Token (excl. 1st token)",
243 + &config.selected_percentile_metrics,
244 + metrics.mean_tpot_ms,
245 + metrics.median_tpot_ms,
246 + &metrics.percentiles_tpot_ms,
247 + );
248 + print_metric_section(
249 + "itl",
250 + "ITL",
251 + "Inter-token Latency",
252 + &config.selected_percentile_metrics,
253 + metrics.mean_itl_ms,
254 + metrics.median_itl_ms,
255 + &metrics.percentiles_itl_ms,
256 + );
257 + print_metric_section(
258 + "e2el",
259 + "E2EL",
260 + "End-to-end Latency",
261 + &config.selected_percentile_metrics,
262 + metrics.mean_e2el_ms,
263 + metrics.median_e2el_ms,
264 + &metrics.percentiles_e2el_ms,
265 + );
266 +}
267 +
268 +fn print_metric_section(
269 + attr_name: &str,
270 + short_name: &str,
271 + header: &str,
272 + selected: &[String],
273 + mean_ms: f64,
274 + median_ms: f64,
275 + percentiles: &[(f64, f64)],
276 +) {
277 + if !selected.iter().any(|s| s == attr_name) {
278 + return;
279 + }
280 + println!("{:-^60}", header);
281 + println!(
282 + "{:<40} {:<10.2}",
283 + format!("Mean {short_name} (ms):"),
284 + mean_ms
285 + );
286 + println!(
287 + "{:<40} {:<10.2}",
288 + format!("Median {short_name} (ms):"),
289 + median_ms
290 + );
291 + for (p, value) in percentiles {
292 + let p_str = if *p == p.floor() {
293 + format!("{}", *p as i64)
294 + } else {
295 + format!("{p}")
296 + };
297 + println!(
298 + "{:<40} {:<10.2}",
299 + format!("P{p_str} {short_name} (ms):"),
300 + value
301 + );
302 + }
303 +}
304 +
305 +/// Print speculative decoding metrics section.
306 +fn print_spec_decode_section(stats: &SpecDecodeStats) {
307 + println!("{:-^50}", "Speculative Decoding");
308 + println!(
309 + "{:<40} {:<10.2}",
310 + "Acceptance rate (%):", stats.acceptance_rate
311 + );
312 + println!(
313 + "{:<40} {:<10.2}",
314 + "Acceptance length:", stats.acceptance_length
315 + );
316 + println!("{:<40} {:<10}", "Drafts:", stats.num_drafts);
317 + println!("{:<40} {:<10}", "Draft tokens:", stats.draft_tokens);
318 + println!("{:<40} {:<10}", "Accepted tokens:", stats.accepted_tokens);
319 + if !stats.per_position_acceptance_rates.is_empty() {
320 + println!("Per-position acceptance (%):");
321 + for (i, rate) in stats.per_position_acceptance_rates.iter().enumerate() {
322 + println!("{:<40} {:<10.2}", format!(" Position {i}:"), rate * 100.0);
323 + }
324 + }
325 +}
326 +
327 +fn print_steady_state(metrics: &crate::metrics::BenchmarkMetrics) {
328 + let Some(ss) = metrics.steady_state.as_ref() else {
329 + return;
330 + };
331 +
332 + if let Some(warning) = ss.window.warning.as_ref() {
333 + println!("Warning: {warning}");
334 + }
335 +
336 + let total = ss.window.requests_total.max(1);
337 + let started_pct = 100.0 * ss.window.requests_started_in_window as f64 / total as f64;
338 +
339 + println!("{:=^47}", " Steady-State Metrics ");
340 + println!(
341 + "{:<35} >= {:.2} * {} = {}",
342 + "Concurrency threshold:",
343 + ss.window.threshold,
344 + ss.window.target_concurrency,
345 + ss.window.threshold_abs,
346 + );
347 + println!(
348 + "{:<35} {:.1}s -> {:.1}s ({:.1}s)",
349 + "Window:", ss.window.start_s, ss.window.end_s, ss.window.duration_s
350 + );
351 + println!(
352 + "{:<35} {}",
353 + "Observed peak concurrency:", ss.window.observed_peak
354 + );
355 + println!(
356 + "{:<35} {} / {} ({:.1}%)",
357 + "Requests started in window:",
358 + ss.window.requests_started_in_window,
359 + ss.window.requests_total,
360 + started_pct
361 + );
362 + println!(
363 + "{:<35} {}",
364 + "Requests completed in window:", ss.window.requests_completed_in_window
365 + );
366 + println!("{:-<47}", "");
367 + println!(
368 + "{:<35} {:.2}",
369 + "Request throughput (req/s):", ss.request_throughput
370 + );
371 + println!(
372 + "{:<35} {:.2}",
373 + "Output token throughput (tok/s):", ss.output_throughput
374 + );
375 + println!(
376 + "{:<35} {:.2}",
377 + "Input token throughput (tok/s):", ss.input_throughput
378 + );
379 + println!(
380 + "{:<35} {:.2}",
381 + "Total token throughput (tok/s):", ss.total_token_throughput
382 + );
383 +
384 + if !ss.percentiles_ttft_ms.is_empty() || ss.mean_ttft_ms > 0.0 {
385 + println!("{:-<47}", "");
386 + println!("{:<35} {:.2}", "Mean TTFT (ms):", ss.mean_ttft_ms);
387 + println!("{:<35} {:.2}", "Median TTFT (ms):", ss.median_ttft_ms);
388 + for (p, v) in &ss.percentiles_ttft_ms {
389 + println!("{:<35} {:.2}", format!("P{} TTFT (ms):", *p as u32), v);
390 + }
391 + }
392 + if ss.mean_tpot_ms > 0.0 || ss.median_tpot_ms > 0.0 {
393 + println!("{:-<47}", "");
394 + println!("{:<35} {:.2}", "Mean TPOT (ms):", ss.mean_tpot_ms);
395 + println!("{:<35} {:.2}", "Median TPOT (ms):", ss.median_tpot_ms);
396 + println!("{:<35} {:.2}", "P90 TPOT (ms):", ss.p90_tpot_ms);
397 + println!("{:<35} {:.2}", "P99 TPOT (ms):", ss.p99_tpot_ms);
398 + }
399 + println!("{:=^47}", "");
400 +}
rust/src/bench/src/output/json.rs new
+776
@@ -0,0 +1,776 @@
1 +// SPDX-License-Identifier: Apache-2.0
2 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +use serde_json::Value;
5 +
6 +use crate::backends::RequestFuncOutput;
7 +use crate::benchmark::SpecDecodeStats;
8 +use crate::config::BenchConfig;
9 +use crate::error::Result;
10 +use crate::metrics::{BenchmarkMetrics, MultiTurnMetrics};
11 +use crate::multi_turn::ConversationOutput;
12 +
13 +/// Build the result JSON object matching the Python output schema exactly.
14 +///
15 +/// Mirrors serve.py:1801-1947 result_json construction.
16 +pub fn build_result_json(
17 + config: &BenchConfig,
18 + metrics: &BenchmarkMetrics,
19 + actual_output_lens: &[usize],
20 + outputs: &[RequestFuncOutput],
21 + benchmark_duration: f64,
22 + current_dt: &str,
23 + spec_decode_stats: Option<&SpecDecodeStats>,
24 +) -> serde_json::Value {
25 + let mut result = serde_json::Map::new();
26 +
27 + // Setup
28 + result.insert("date".into(), Value::String(current_dt.to_string()));
29 + result.insert(
30 + "endpoint_type".into(),
31 + Value::String(config.backend.to_string()),
32 + );
33 + result.insert("backend".into(), Value::String(config.backend.to_string()));
34 + result.insert(
35 + "label".into(),
36 + config.label.as_ref().map(|l| Value::String(l.clone())).unwrap_or(Value::Null),
37 + );
38 + result.insert(
39 + "model_id".into(),
40 + Value::String(config.model.clone().unwrap_or_default()),
41 + );
42 + result.insert(
43 + "tokenizer_id".into(),
44 + config
45 + .tokenizer_id
46 + .as_ref()
47 + .map(|t| Value::String(t.clone()))
48 + .unwrap_or(Value::Null),
49 + );
50 + result.insert(
51 + "num_prompts".into(),
52 + Value::Number(config.num_prompts.into()),
53 + );
54 + result.insert(
55 + "max_model_len".into(),
56 + config.max_model_len.map(|v| Value::Number(v.into())).unwrap_or(Value::Null),
57 + );
58 +
59 + // Metadata
60 + if let Some(ref metadata) = config.metadata {
61 + for (k, v) in metadata {
62 + result.insert(k.clone(), Value::String(v.clone()));
63 + }
64 + }
65 +
66 + // Traffic
67 + if config.request_rate.is_infinite() {
68 + result.insert("request_rate".into(), Value::String("inf".to_string()));
69 + } else {
70 + result.insert(
71 + "request_rate".into(),
72 + serde_json::json!(config.request_rate),
73 + );
74 + }
75 + result.insert("burstiness".into(), serde_json::json!(config.burstiness));
76 + result.insert(
77 + "max_concurrency".into(),
78 + config.max_concurrency.map(|v| Value::Number(v.into())).unwrap_or(Value::Null),
79 + );
80 +
81 + let is_pooling = config.backend.is_pooling();
82 +
83 + // Benchmark results
84 + result.insert("duration".into(), serde_json::json!(benchmark_duration));
85 + result.insert("completed".into(), serde_json::json!(metrics.completed));
86 + result.insert("failed".into(), serde_json::json!(metrics.failed));
87 + result.insert(
88 + "total_input_tokens".into(),
89 + serde_json::json!(metrics.total_input),
90 + );
91 + if !is_pooling {
92 + result.insert(
93 + "total_output_tokens".into(),
94 + serde_json::json!(metrics.total_output),
95 + );
96 + }
97 + result.insert(
98 + "request_throughput".into(),
99 + serde_json::json!(metrics.request_throughput),
100 + );
101 + if config.goodput.is_empty() {
102 + result.insert("request_goodput".into(), Value::Null);
103 + } else {
104 + result.insert(
105 + "request_goodput".into(),
106 + serde_json::json!(metrics.request_goodput),
107 + );
108 + }
109 + if is_pooling {
110 + result.insert(
111 + "total_token_throughput".into(),
112 + serde_json::json!(metrics.total_token_throughput),
113 + );
114 + } else {
115 + result.insert(
116 + "input_throughput".into(),
117 + serde_json::json!(metrics.input_throughput),
118 + );
119 + result.insert(
120 + "output_throughput".into(),
121 + serde_json::json!(metrics.output_throughput),
122 + );
123 + result.insert(
124 + "total_token_throughput".into(),
125 + serde_json::json!(metrics.total_token_throughput),
126 + );
127 + result.insert(
128 + "max_output_tokens_per_s".into(),
129 + serde_json::json!(metrics.max_output_tokens_per_s),
130 + );
131 + result.insert(
132 + "max_concurrent_requests".into(),
133 + serde_json::json!(metrics.max_concurrent_requests),
134 + );
135 + // Inverse Real-Time Factor (for ASR benchmarks; 0.0 for generation)
136 + result.insert("rtfx".into(), serde_json::json!(0.0));
137 + }
138 +
139 + // Per-failure log (always emitted when there are failures, regardless of --save-detailed).
140 + // ttft/itl/output_tokens/start_time/latency are best-effort — zero/empty for failures
141 + // that occurred before any tokens were received. Non-zero values indicate a partial
142 + // failure (some tokens streamed before the error).
143 + if metrics.failed > 0 {
144 + result.insert(
145 + "failed_requests".into(),
146 + Value::Array(collect_failed_requests(outputs)),
147 + );
148 + }
149 +
150 + // Per-request data.
151 + // For pooling: Python always includes input_lens and errors (serve.py:1004-1013).
152 + // For generation: all per-request arrays are gated by --save-detailed.
153 + if is_pooling {
154 + result.insert(
155 + "input_lens".into(),
156 + serde_json::json!(outputs.iter().map(|o| o.prompt_len).collect::<Vec<_>>()),
157 + );
158 + result.insert(
159 + "errors".into(),
160 + serde_json::json!(outputs.iter().map(|o| &o.error).collect::<Vec<_>>()),
161 + );
162 + }
163 + if config.save_detailed {
164 + if !is_pooling {
165 + result.insert(
166 + "input_lens".into(),
167 + serde_json::json!(outputs.iter().map(|o| o.prompt_len).collect::<Vec<_>>()),
168 + );
169 + result.insert("output_lens".into(), serde_json::json!(actual_output_lens));
170 + result.insert(
171 + "ttfts".into(),
172 + serde_json::json!(outputs.iter().map(|o| o.ttft).collect::<Vec<_>>()),
173 + );
174 + result.insert(
175 + "itls".into(),
176 + serde_json::json!(outputs.iter().map(|o| &o.itl).collect::<Vec<_>>()),
177 + );
178 + result.insert(
179 + "generated_texts".into(),
180 + serde_json::json!(outputs.iter().map(|o| &o.generated_text).collect::<Vec<_>>()),
181 + );
182 + result.insert(
183 + "errors".into(),
184 + serde_json::json!(outputs.iter().map(|o| &o.error).collect::<Vec<_>>()),
185 + );
186 + }
187 + result.insert(
188 + "latencies".into(),
189 + serde_json::json!(outputs.iter().map(|o| o.latency).collect::<Vec<_>>()),
190 + );
191 + result.insert(
192 + "start_times".into(),
193 + serde_json::json!(outputs.iter().map(|o| o.start_time).collect::<Vec<_>>()),
194 + );
195 + }
196 +
197 + // Speculative decoding stats
198 + if let Some(stats) = spec_decode_stats {
199 + insert_spec_decode_stats(&mut result, stats);
200 + }
201 +
202 + // Per-metric stats (pooling only has e2el)
203 + if !is_pooling {
204 + add_metric_stats(
205 + &mut result,
206 + "ttft",
207 + metrics.mean_ttft_ms,
208 + metrics.median_ttft_ms,
209 + metrics.std_ttft_ms,
210 + &metrics.percentiles_ttft_ms,
211 + &config.selected_percentile_metrics,
212 + );
213 + add_metric_stats(
214 + &mut result,
215 + "tpot",
216 + metrics.mean_tpot_ms,
217 + metrics.median_tpot_ms,
218 + metrics.std_tpot_ms,
219 + &metrics.percentiles_tpot_ms,
220 + &config.selected_percentile_metrics,
221 + );
222 + add_metric_stats(
223 + &mut result,
224 + "itl",
225 + metrics.mean_itl_ms,
226 + metrics.median_itl_ms,
227 + metrics.std_itl_ms,
228 + &metrics.percentiles_itl_ms,
229 + &config.selected_percentile_metrics,
230 + );
231 + }
232 + add_metric_stats(
233 + &mut result,
234 + "e2el",
235 + metrics.mean_e2el_ms,
236 + metrics.median_e2el_ms,
237 + metrics.std_e2el_ms,
238 + &metrics.percentiles_e2el_ms,
239 + &config.selected_percentile_metrics,
240 + );
241 +
242 + // Steady-state metrics (null when not computed or scope gate failed).
243 + result.insert(
244 + "steady_state".into(),
245 + serde_json::to_value(&metrics.steady_state).unwrap_or(Value::Null),
246 + );
247 +
248 + Value::Object(result)
249 +}
250 +
251 +/// Build multi-turn result JSON.
252 +pub fn build_multi_turn_result_json(
253 + config: &BenchConfig,
254 + mt_metrics: &MultiTurnMetrics,
255 + conversation_outputs: &[ConversationOutput],
256 + benchmark_duration: f64,
257 + current_dt: &str,
258 + spec_decode_stats: Option<&SpecDecodeStats>,
259 +) -> serde_json::Value {
260 + let mut result = serde_json::Map::new();
261 +
262 + // Setup
263 + result.insert("date".into(), Value::String(current_dt.to_string()));
264 + result.insert("mode".into(), Value::String("multi_turn".to_string()));
265 + result.insert("backend".into(), Value::String(config.backend.to_string()));
266 + result.insert(
267 + "model_id".into(),
268 + Value::String(config.model.clone().unwrap_or_default()),
269 + );
270 + result.insert(
271 + "num_conversations".into(),
272 + Value::Number(config.num_prompts.into()),
273 + );
274 + result.insert(
275 + "max_model_len".into(),
276 + config.max_model_len.map(|v| Value::Number(v.into())).unwrap_or(Value::Null),
277 + );
278 + result.insert(
279 + "turns_per_conversation".into(),
280 + Value::Number(config.multi_turn_num_turns.into()),
281 + );
282 + result.insert(
283 + "multi_turn_concurrency".into(),
284 + serde_json::json!(
285 + config
286 + .multi_turn_concurrency
287 + .unwrap_or(config.max_concurrency.unwrap_or(config.num_prompts))
288 + ),
289 + );
290 + result.insert(
291 + "inter_turn_delay_ms".into(),
292 + Value::Number(config.multi_turn_delay_ms.into()),
293 + );
294 + if config.multi_turn_prefix_global_ratio > 0.0
295 + || config.multi_turn_prefix_conversation_ratio > 0.0
296 + {
297 + result.insert(
298 + "prefix_global_ratio".into(),
299 + serde_json::json!(config.multi_turn_prefix_global_ratio),
300 + );
301 + result.insert(
302 + "prefix_conversation_ratio".into(),
303 + serde_json::json!(config.multi_turn_prefix_conversation_ratio),
304 + );
305 + result.insert(
306 + "prefix_unique_ratio".into(),
307 + serde_json::json!(
308 + (1.0 - (config.multi_turn_prefix_global_ratio
309 + + config.multi_turn_prefix_conversation_ratio))
310 + .max(0.0)
311 + ),
312 + );
313 + result.insert("no_history_accumulation".into(), serde_json::json!(true));
314 + }
315 +
316 + // Metadata
317 + if let Some(ref metadata) = config.metadata {
318 + for (k, v) in metadata {
319 + result.insert(k.clone(), Value::String(v.clone()));
320 + }
321 + }
322 +
323 + // Conversation-level stats
324 + result.insert("duration".into(), serde_json::json!(benchmark_duration));
325 + result.insert(
326 + "conversations_completed".into(),
327 + serde_json::json!(mt_metrics.conversations_completed),
328 + );
329 + result.insert(
330 + "conversations_failed".into(),
331 + serde_json::json!(mt_metrics.conversations_failed),
332 + );
333 + result.insert(
334 + "avg_turns_completed".into(),
335 + serde_json::json!(mt_metrics.avg_turns_completed),
336 + );
337 + result.insert(
338 + "avg_conversation_duration_ms".into(),
339 + serde_json::json!(mt_metrics.avg_conversation_duration_ms),
340 + );
341 +
342 + // Overall metrics
343 + let overall = &mt_metrics.overall;
344 + result.insert("completed".into(), serde_json::json!(overall.completed));
345 + result.insert("failed".into(), serde_json::json!(overall.failed));
346 +
347 + // Per-failure log (always emitted when there are failures). Each entry is a single failed turn.
348 + if overall.failed > 0 {
349 + result.insert(
350 + "failed_requests".into(),
351 + Value::Array(collect_failed_turns(conversation_outputs)),
352 + );
353 + }
354 + result.insert(
355 + "total_input_tokens".into(),
356 + serde_json::json!(overall.total_input),
357 + );
358 + result.insert(
359 + "total_output_tokens".into(),
360 + serde_json::json!(overall.total_output),
361 + );
362 + result.insert(
363 + "request_throughput".into(),
364 + serde_json::json!(overall.request_throughput),
365 + );
366 + result.insert(
367 + "input_throughput".into(),
368 + serde_json::json!(overall.input_throughput),
369 + );
370 + result.insert(
371 + "output_throughput".into(),
372 + serde_json::json!(overall.output_throughput),
373 + );
374 + result.insert(
375 + "total_token_throughput".into(),
376 + serde_json::json!(overall.total_token_throughput),
377 + );
378 +
379 + // Overall per-metric stats
380 + add_metric_stats(
381 + &mut result,
382 + "ttft",
383 + overall.mean_ttft_ms,
384 + overall.median_ttft_ms,
385 + overall.std_ttft_ms,
386 + &overall.percentiles_ttft_ms,
387 + &config.selected_percentile_metrics,
388 + );
389 + add_metric_stats(
390 + &mut result,
391 + "tpot",
392 + overall.mean_tpot_ms,
393 + overall.median_tpot_ms,
394 + overall.std_tpot_ms,
395 + &overall.percentiles_tpot_ms,
396 + &config.selected_percentile_metrics,
397 + );
398 + add_metric_stats(
399 + &mut result,
400 + "itl",
401 + overall.mean_itl_ms,
402 + overall.median_itl_ms,
403 + overall.std_itl_ms,
404 + &overall.percentiles_itl_ms,
405 + &config.selected_percentile_metrics,
406 + );
407 + add_metric_stats(
408 + &mut result,
409 + "e2el",
410 + overall.mean_e2el_ms,
411 + overall.median_e2el_ms,
412 + overall.std_e2el_ms,
413 + &overall.percentiles_e2el_ms,
414 + &config.selected_percentile_metrics,
415 + );
416 +
417 + // Speculative decoding stats
418 + if let Some(stats) = spec_decode_stats {
419 + insert_spec_decode_stats(&mut result, stats);
420 + }
421 +
422 + // Per-turn metrics array
423 + let per_turn_json: Vec<Value> = mt_metrics
424 + .per_turn
425 + .iter()
426 + .enumerate()
427 + .map(|(i, m)| {
428 + let mut turn = serde_json::Map::new();
429 + turn.insert("turn_index".into(), serde_json::json!(i));
430 + turn.insert(
431 + "num_samples".into(),
432 + serde_json::json!(m.completed + m.failed),
433 + );
434 + turn.insert("completed".into(), serde_json::json!(m.completed));
435 + turn.insert("failed".into(), serde_json::json!(m.failed));
436 + turn.insert(
437 + "total_input_tokens".into(),
438 + serde_json::json!(m.total_input),
439 + );
440 + turn.insert(
441 + "total_output_tokens".into(),
442 + serde_json::json!(m.total_output),
443 + );
444 + turn.insert(
445 + "request_throughput".into(),
446 + serde_json::json!(m.request_throughput),
447 + );
448 + turn.insert(
449 + "input_throughput".into(),
450 + serde_json::json!(m.input_throughput),
451 + );
452 + turn.insert(
453 + "output_throughput".into(),
454 + serde_json::json!(m.output_throughput),
455 + );
456 + turn.insert(
457 + "total_token_throughput".into(),
458 + serde_json::json!(m.total_token_throughput),
459 + );
460 +
461 + add_metric_stats(
462 + &mut turn,
463 + "ttft",
464 + m.mean_ttft_ms,
465 + m.median_ttft_ms,
466 + m.std_ttft_ms,
467 + &m.percentiles_ttft_ms,
468 + &config.selected_percentile_metrics,
469 + );
470 + add_metric_stats(
471 + &mut turn,
472 + "tpot",
473 + m.mean_tpot_ms,
474 + m.median_tpot_ms,
475 + m.std_tpot_ms,
476 + &m.percentiles_tpot_ms,
477 + &config.selected_percentile_metrics,
478 + );
479 + add_metric_stats(
480 + &mut turn,
481 + "itl",
482 + m.mean_itl_ms,
483 + m.median_itl_ms,
484 + m.std_itl_ms,
485 + &m.percentiles_itl_ms,
486 + &config.selected_percentile_metrics,
487 + );
488 + add_metric_stats(
489 + &mut turn,
490 + "e2el",
491 + m.mean_e2el_ms,
492 + m.median_e2el_ms,
493 + m.std_e2el_ms,
494 + &m.percentiles_e2el_ms,
495 + &config.selected_percentile_metrics,
496 + );
497 +
498 + Value::Object(turn)
499 + })
500 + .collect();
501 +
502 + result.insert("per_turn_metrics".into(), Value::Array(per_turn_json));
503 +
504 + Value::Object(result)
505 +}
506 +
507 +fn collect_failed_requests(outputs: &[RequestFuncOutput]) -> Vec<Value> {
508 + outputs
509 + .iter()
510 + .enumerate()
511 + .filter(|(_, o)| !o.success)
512 + .map(|(i, o)| {
513 + serde_json::json!({
514 + "index": i,
515 + "error": o.error,
516 + "prompt_len": o.prompt_len,
517 + "start_time": o.start_time,
518 + "latency": o.latency,
519 + "ttft": o.ttft,
520 + "output_tokens": o.output_tokens,
521 + "itl": o.itl,
522 + })
523 + })
524 + .collect()
525 +}
526 +
527 +fn collect_failed_turns(conversation_outputs: &[ConversationOutput]) -> Vec<Value> {
528 + let mut failed = Vec::new();
529 + for conv in conversation_outputs {
530 + for turn in &conv.turns {
531 + if !turn.request_output.success {
532 + let o = &turn.request_output;
533 + failed.push(serde_json::json!({
534 + "conversation_id": conv.conversation_id,
535 + "turn_index": turn.turn_index,
536 + "error": o.error,
537 + "prompt_len": o.prompt_len,
538 + "start_time": o.start_time,
539 + "latency": o.latency,
540 + "ttft": o.ttft,
541 + "output_tokens": o.output_tokens,
542 + "itl": o.itl,
543 + }));
544 + }
545 + }
546 + }
547 + failed
548 +}
549 +
550 +fn insert_spec_decode_stats(result: &mut serde_json::Map<String, Value>, stats: &SpecDecodeStats) {
551 + result.insert(
552 + "spec_decode_acceptance_rate".into(),
553 + serde_json::json!(stats.acceptance_rate),
554 + );
555 + result.insert(
556 + "spec_decode_acceptance_length".into(),
557 + serde_json::json!(stats.acceptance_length),
558 + );
559 + result.insert(
560 + "spec_decode_num_drafts".into(),
561 + serde_json::json!(stats.num_drafts),
562 + );
563 + result.insert(
564 + "spec_decode_draft_tokens".into(),
565 + serde_json::json!(stats.draft_tokens),
566 + );
567 + result.insert(
568 + "spec_decode_accepted_tokens".into(),
569 + serde_json::json!(stats.accepted_tokens),
570 + );
571 + result.insert(
572 + "spec_decode_per_position_acceptance_rates".into(),
573 + serde_json::json!(stats.per_position_acceptance_rates),
574 + );
575 +}
576 +
577 +fn add_metric_stats(
578 + result: &mut serde_json::Map<String, Value>,
579 + name: &str,
580 + mean: f64,
581 + median: f64,
582 + std: f64,
583 + percentiles: &[(f64, f64)],
584 + selected: &[String],
585 +) {
586 + if !selected.iter().any(|s| s == name) {
587 + return;
588 + }
589 + result.insert(format!("mean_{name}_ms"), serde_json::json!(mean));
590 + result.insert(format!("median_{name}_ms"), serde_json::json!(median));
591 + result.insert(format!("std_{name}_ms"), serde_json::json!(std));
592 + for (p, value) in percentiles {
593 + let p_str = if *p == p.floor() {
594 + format!("{}", *p as i64)
595 + } else {
596 + format!("{p}")
597 + };
598 + result.insert(format!("p{p_str}_{name}_ms"), serde_json::json!(value));
599 + }
600 +}
601 +
602 +/// Save result JSON to file (overwrite mode).
603 +pub fn save_result(json: &Value, file_path: &str) -> Result<()> {
604 + let content = serde_json::to_string(json)?;
605 + std::fs::write(file_path, content)?;
606 + println!("Results saved to {file_path}");
607 + Ok(())
608 +}
609 +
610 +/// Append result JSON to file (JSONL format, matching Python's --append-result).
611 +pub fn append_result(json: &Value, file_path: &str) -> Result<()> {
612 + use std::io::Write;
613 + let content = serde_json::to_string(json)?;
614 + let mut file = std::fs::OpenOptions::new().create(true).append(true).open(file_path)?;
615 + // If file is non-empty, prepend a newline
616 + let meta = file.metadata()?;
617 + if meta.len() > 0 {
618 + file.write_all(b"\n")?;
619 + }
620 + file.write_all(content.as_bytes())?;
621 + println!("Results appended to {file_path}");
622 + Ok(())
623 +}
624 +
625 +/// Compute the result filename matching Python's logic.
626 +pub fn compute_result_filename(config: &BenchConfig, model_id: &str, current_dt: &str) -> String {
627 + let base_model = model_id.split('/').next_back().unwrap_or(model_id);
628 + let max_conc_str =
629 + config.max_concurrency.map(|mc| format!("-concurrency{mc}")).unwrap_or_default();
630 + let label = config.label.as_deref().unwrap_or_else(|| config.backend.as_str());
631 +
632 + let file_name = if config.request_rate.is_infinite() {
633 + format!("{label}-infqps{max_conc_str}-{base_model}-{current_dt}.json")
634 + } else {
635 + format!(
636 + "{label}-{}qps{max_conc_str}-{base_model}-{current_dt}.json",
637 + config.request_rate
638 + )
639 + };
640 +
641 + if let Some(ref explicit) = config.result_filename {
642 + if let Some(ref dir) = config.result_dir {
643 + return format!("{dir}/{explicit}");
644 + }
645 + return explicit.clone();
646 + }
647 +
648 + if let Some(ref dir) = config.result_dir {
649 + format!("{dir}/{file_name}")
650 + } else {
651 + file_name
652 + }
653 +}
654 +
655 +#[cfg(test)]
656 +mod tests {
657 + use super::*;
658 +
659 + #[test]
660 + fn test_insert_spec_decode_stats() {
661 + let stats = SpecDecodeStats {
662 + num_drafts: 1000,
663 + draft_tokens: 3000,
664 + accepted_tokens: 2400,
665 + acceptance_rate: 80.0,
666 + acceptance_length: 3.4,
667 + per_position_acceptance_rates: vec![0.9, 0.8, 0.7],
668 + };
669 +
670 + let mut result = serde_json::Map::new();
671 + insert_spec_decode_stats(&mut result, &stats);
672 +
673 + assert_eq!(result["spec_decode_acceptance_rate"], 80.0);
674 + assert_eq!(result["spec_decode_acceptance_length"], 3.4);
675 + assert_eq!(result["spec_decode_num_drafts"], 1000);
676 + assert_eq!(result["spec_decode_draft_tokens"], 3000);
677 + assert_eq!(result["spec_decode_accepted_tokens"], 2400);
678 + assert_eq!(
679 + result["spec_decode_per_position_acceptance_rates"],
680 + serde_json::json!([0.9, 0.8, 0.7])
681 + );
682 + }
683 +
684 + #[test]
685 + fn test_insert_spec_decode_stats_empty_positions() {
686 + let stats = SpecDecodeStats {
687 + num_drafts: 500,
688 + draft_tokens: 1500,
689 + accepted_tokens: 1200,
690 + acceptance_rate: 80.0,
691 + acceptance_length: 3.4,
692 + per_position_acceptance_rates: vec![],
693 + };
694 +
695 + let mut result = serde_json::Map::new();
696 + insert_spec_decode_stats(&mut result, &stats);
697 +
698 + assert_eq!(
699 + result["spec_decode_per_position_acceptance_rates"],
700 + serde_json::json!([])
701 + );
702 + assert_eq!(result.len(), 6);
703 + }
704 +
705 + #[test]
706 + fn test_collect_failed_requests_empty_when_all_success() {
707 + let outputs = vec![
708 + RequestFuncOutput {
709 + success: true,
710 + ..Default::default()
711 + },
712 + RequestFuncOutput {
713 + success: true,
714 + ..Default::default()
715 + },
716 + ];
717 + assert!(collect_failed_requests(&outputs).is_empty());
718 + }
719 +
720 + #[test]
721 + fn test_collect_failed_requests_preserves_index_and_fields() {
722 + let outputs = vec![
723 + RequestFuncOutput {
724 + success: true,
725 + ..Default::default()
726 + },
727 + // Pre-stream failure: no tokens received.
728 + RequestFuncOutput {
729 + success: false,
730 + error: "connection reset".into(),
731 + prompt_len: 128,
732 + start_time: 1.5,
733 + latency: 0.0,
734 + ttft: 0.0,
735 + output_tokens: 0,
736 + itl: vec![],
737 + ..Default::default()
738 + },
739 + RequestFuncOutput {
740 + success: true,
741 + ..Default::default()
742 + },
743 + // Partial failure: streamed 3 tokens before timing out.
744 + RequestFuncOutput {
745 + success: false,
746 + error: "timeout".into(),
747 + prompt_len: 64,
748 + start_time: 2.25,
749 + latency: 30.0,
750 + ttft: 0.5,
751 + output_tokens: 3,
752 + itl: vec![0.1, 0.12],
753 + ..Default::default()
754 + },
755 + ];
756 +
757 + let failed = collect_failed_requests(&outputs);
758 + assert_eq!(failed.len(), 2);
759 +
760 + assert_eq!(failed[0]["index"], 1);
761 + assert_eq!(failed[0]["error"], "connection reset");
762 + assert_eq!(failed[0]["prompt_len"], 128);
763 + assert_eq!(failed[0]["start_time"], 1.5);
764 + assert_eq!(failed[0]["latency"], 0.0);
765 + assert_eq!(failed[0]["ttft"], 0.0);
766 + assert_eq!(failed[0]["output_tokens"], 0);
767 + assert_eq!(failed[0]["itl"], serde_json::json!([]));
768 +
769 + assert_eq!(failed[1]["index"], 3);
770 + assert_eq!(failed[1]["error"], "timeout");
771 + assert_eq!(failed[1]["latency"], 30.0);
772 + assert_eq!(failed[1]["ttft"], 0.5);
773 + assert_eq!(failed[1]["output_tokens"], 3);
774 + assert_eq!(failed[1]["itl"], serde_json::json!([0.1, 0.12]));
775 + }
776 +}
rust/src/bench/src/output/mod.rs new
+5
@@ -0,0 +1,5 @@
1 +// SPDX-License-Identifier: Apache-2.0
2 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +pub mod console;
5 +pub mod json;
rust/src/bench/src/rate_control.rs new
+171
@@ -0,0 +1,171 @@
1 +// SPDX-License-Identifier: Apache-2.0
2 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +use rand::SeedableRng;
5 +use rand::rngs::StdRng;
6 +use rand_distr::{Distribution, Gamma};
7 +
8 +use crate::cli::RampUpStrategy;
9 +use crate::config::RampUpConfig;
10 +
11 +/// Pre-computed request schedule: cumulative delays from start and per-request rates.
12 +pub struct RequestSchedule {
13 + /// Cumulative absolute delay from start (seconds) for each request.
14 + pub delays: Vec<f64>,
15 + /// Instantaneous rate for each request (for ramp-up logging).
16 + #[allow(dead_code)]
17 + pub rates: Vec<f64>,
18 +}
19 +
20 +/// Compute the per-request rate, accounting for optional ramp-up.
21 +///
22 +/// Mirrors Python's `_get_current_request_rate()` from serve.py:218-240.
23 +fn get_current_request_rate(
24 + ramp_up: Option<&RampUpConfig>,
25 + request_index: usize,
26 + total_requests: usize,
27 + base_rate: f64,
28 +) -> f64 {
29 + let config = match ramp_up {
30 + Some(c) => c,
31 + None => return base_rate,
32 + };
33 +
34 + let progress = request_index as f64 / (total_requests - 1).max(1) as f64;
35 +
36 + match config.strategy {
37 + RampUpStrategy::Linear => {
38 + let increase = (config.end_rps - config.start_rps) * progress;
39 + config.start_rps + increase
40 + }
41 + RampUpStrategy::Exponential => {
42 + let ratio = config.end_rps / config.start_rps;
43 + config.start_rps * ratio.powf(progress)
44 + }
45 + }
46 +}
47 +
48 +/// Compute the request schedule for all requests.
49 +///
50 +/// Ports the Python `get_request()` / `_generate_request_timestamps()` logic
51 +/// from serve.py:243-340.
52 +pub fn compute_schedule(
53 + num_requests: usize,
54 + request_rate: f64,
55 + burstiness: f64,
56 + seed: u64,
57 + ramp_up: Option<&RampUpConfig>,
58 +) -> RequestSchedule {
59 + assert!(burstiness > 0.0, "burstiness must be positive");
60 + assert!(num_requests > 0, "must have at least one request");
61 +
62 + let mut rng = StdRng::seed_from_u64(seed);
63 + let mut delay_ts = Vec::with_capacity(num_requests);
64 + let mut rates = Vec::with_capacity(num_requests);
65 +
66 + for i in 0..num_requests {
67 + let current_rate = get_current_request_rate(ramp_up, i, num_requests, request_rate);
68 + rates.push(current_rate);
69 +
70 + if current_rate.is_infinite() {
71 + delay_ts.push(0.0);
72 + } else if burstiness.is_infinite() {
73 + // When burstiness → ∞, delay becomes constant = 1/rate
74 + delay_ts.push(1.0 / current_rate);
75 + } else {
76 + let theta = 1.0 / (current_rate * burstiness);
77 + let gamma = Gamma::new(burstiness, theta).unwrap();
78 + delay_ts.push(gamma.sample(&mut rng));
79 + }
80 + }
81 +
82 + // Compute cumulative delays
83 + for i in 1..delay_ts.len() {
84 + delay_ts[i] += delay_ts[i - 1];
85 + }
86 +
87 + // Normalize: scale cumulative delays so total matches target.
88 + // Only for fixed-rate (no ramp-up) mode, matching Python behavior.
89 + if ramp_up.is_none()
90 + && let Some(&last) = delay_ts.last()
91 + && last > 0.0
92 + && !request_rate.is_infinite()
93 + {
94 + let target_total = num_requests as f64 / request_rate;
95 + let factor = target_total / last;
96 + for d in &mut delay_ts {
97 + *d *= factor;
98 + }
99 + }
100 +
101 + RequestSchedule {
102 + delays: delay_ts,
103 + rates,
104 + }
105 +}
106 +
107 +#[cfg(test)]
108 +mod tests {
109 + use super::*;
110 +
111 + #[test]
112 + fn test_infinite_rate_all_zero_delays() {
113 + let sched = compute_schedule(100, f64::INFINITY, 1.0, 42, None);
114 + assert_eq!(sched.delays.len(), 100);
115 + for d in &sched.delays {
116 + assert_eq!(*d, 0.0);
117 + }
118 + }
119 +
120 + #[test]
121 + fn test_fixed_rate_monotonic() {
122 + let sched = compute_schedule(50, 10.0, 1.0, 42, None);
123 + for i in 1..sched.delays.len() {
124 + assert!(sched.delays[i] >= sched.delays[i - 1]);
125 + }
126 + // Total should be close to num_requests / rate = 5.0
127 + let last = *sched.delays.last().unwrap();
128 + assert!((last - 5.0).abs() < 0.01, "last delay = {last}");
129 + }
130 +
131 + #[test]
132 + fn test_infinite_burstiness_constant_delay() {
133 + let sched = compute_schedule(10, 5.0, f64::INFINITY, 42, None);
134 + // Each delay should be 0.2s apart (1/5)
135 + for i in 1..sched.delays.len() {
136 + let delta = sched.delays[i] - sched.delays[i - 1];
137 + assert!((delta - 0.2).abs() < 1e-10);
138 + }
139 + }
140 +
141 + #[test]
142 + fn test_linear_ramp_up() {
143 + let ramp = RampUpConfig {
144 + strategy: RampUpStrategy::Linear,
145 + start_rps: 1.0,
146 + end_rps: 10.0,
147 + };
148 + let sched = compute_schedule(10, 5.0, 1.0, 42, Some(&ramp));
149 + // Rates should increase linearly from 1.0 to 10.0
150 + assert!((sched.rates[0] - 1.0).abs() < 1e-10);
151 + assert!((sched.rates[9] - 10.0).abs() < 1e-10);
152 + // Middle should be ~5.5
153 + assert!((sched.rates[4] - 5.0).abs() < 0.5);
154 + }
155 +
156 + #[test]
157 + fn test_exponential_ramp_up() {
158 + let ramp = RampUpConfig {
159 + strategy: RampUpStrategy::Exponential,
160 + start_rps: 1.0,
161 + end_rps: 100.0,
162 + };
163 + let sched = compute_schedule(10, 5.0, 1.0, 42, Some(&ramp));
164 + assert!((sched.rates[0] - 1.0).abs() < 1e-10);
165 + assert!((sched.rates[9] - 100.0).abs() < 1e-10);
166 + // Exponential: rates should be monotonically increasing
167 + for i in 1..sched.rates.len() {
168 + assert!(sched.rates[i] >= sched.rates[i - 1]);
169 + }
170 + }
171 +}
rust/src/bench/src/ready_checker.rs new
+72
@@ -0,0 +1,72 @@
1 +// SPDX-License-Identifier: Apache-2.0
2 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +use std::time::Instant;
5 +
6 +use indicatif::{ProgressBar, ProgressStyle};
7 +
8 +use crate::backends::{RequestFuncInput, RequestFuncOutput, get_backend};
9 +use crate::cli::BackendKind;
10 +use crate::error::{BenchError, Result};
11 +
12 +/// Wait for the serving endpoint to become available.
13 +///
14 +/// Sends test requests with retry until success or timeout.
15 +/// Mirrors Python's `wait_for_endpoint` in ready_checker.py.
16 +pub async fn wait_for_endpoint(
17 + backend: BackendKind,
18 + client: &reqwest::Client,
19 + test_input: &RequestFuncInput,
20 + timeout_seconds: u64,
21 + retry_interval: u64,
22 +) -> Result<RequestFuncOutput> {
23 + let backend = get_backend(backend)?;
24 + let deadline = Instant::now() + std::time::Duration::from_secs(timeout_seconds);
25 +
26 + println!("Waiting for endpoint to become up in {timeout_seconds}s");
27 +
28 + let pb = ProgressBar::new(timeout_seconds);
29 + pb.set_style(
30 + ProgressStyle::with_template("{msg} |{bar:40}| {elapsed} elapsed, {eta} remaining")
31 + .unwrap()
32 + .progress_chars("##-"),
33 + );
34 +
35 + let mut last_error = String::new();
36 +
37 + loop {
38 + let remaining = deadline.saturating_duration_since(Instant::now());
39 + let elapsed = timeout_seconds.saturating_sub(remaining.as_secs());
40 + pb.set_position(elapsed);
41 +
42 + if remaining.is_zero() {
43 + pb.finish_and_clear();
44 + break;
45 + }
46 +
47 + // Ping the endpoint
48 + match backend.send_request(test_input, client).await {
49 + Ok(output) if output.success => {
50 + pb.finish_and_clear();
51 + return Ok(output);
52 + }
53 + Ok(output) => {
54 + let err = output.error.clone();
55 + let err_last_line = err.lines().last().unwrap_or(&err);
56 + eprintln!("Endpoint is not ready. Error='{err_last_line}'");
57 + last_error = err;
58 + }
59 + Err(e) => {
60 + last_error = e.to_string();
61 + }
62 + }
63 +
64 + // Retry after delay
65 + let sleep_dur = std::cmp::min(std::time::Duration::from_secs(retry_interval), remaining);
66 + if !sleep_dur.is_zero() {
67 + tokio::time::sleep(sleep_dur).await;
68 + }
69 + }
70 +
71 + Err(BenchError::EndpointTimeout(timeout_seconds, last_error))
72 +}
rust/src/bench/src/sweep.rs new
+561
@@ -0,0 +1,561 @@
1 +// SPDX-License-Identifier: Apache-2.0
2 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +use crate::config::BenchConfig;
5 +use crate::error::{BenchError, Result};
6 +use crate::output::json::{compute_result_filename, save_result};
7 +
8 +/// Reset the server's prefix cache by calling POST /reset_prefix_cache.
9 +/// Requires VLLM_SERVER_DEV_MODE=1 on the vLLM server.
10 +async fn reset_prefix_cache(base_url: &str) -> Result<()> {
11 + let url = format!("{}/reset_prefix_cache", base_url.trim_end_matches('/'));
12 + let client = reqwest::Client::new();
13 + let resp = client
14 + .post(&url)
15 + .send()
16 + .await
17 + .map_err(|e| BenchError::Backend(format!("Failed to reset prefix cache: {e}")))?;
18 + if resp.status().is_success() {
19 + println!("Prefix cache reset successfully.");
20 + } else {
21 + let status = resp.status();
22 + let body = resp.text().await.unwrap_or_default();
23 + return Err(BenchError::Backend(format!(
24 + "Failed to reset prefix cache: HTTP {status} — {body}. \
25 + Is VLLM_SERVER_DEV_MODE=1 set on the server?"
26 + )));
27 + }
28 + Ok(())
29 +}
30 +
31 +/// Result of a single sweep point.
32 +struct SweepPoint {
33 + label: String,
34 + #[allow(dead_code)]
35 + value: f64,
36 + result_json: serde_json::Value,
37 + ss_request_throughput: Option<f64>,
38 + ss_output_throughput: Option<f64>,
39 +}
40 +
41 +/// Run a sweep over max-concurrency values.
42 +pub async fn run_concurrency_sweep(
43 + base_config: &BenchConfig,
44 + values: &[usize],
45 + num_prompts_factor: Option<usize>,
46 +) -> Result<()> {
47 + println!("{:=^70}", " Concurrency Sweep ");
48 + println!(
49 + "Sweeping --max-concurrency over {} values: {:?}",
50 + values.len(),
51 + values
52 + );
53 + println!();
54 +
55 + let mut points = Vec::with_capacity(values.len());
56 +
57 + let current_dt = chrono::Local::now().format("%Y%m%d-%H%M%S").to_string();
58 +
59 + for (i, &mc) in values.iter().enumerate() {
60 + println!(
61 + "{:-^70}",
62 + format!(" Run {}/{}: max_concurrency={} ", i + 1, values.len(), mc)
63 + );
64 +
65 + if base_config.reset_prefix_cache {
66 + reset_prefix_cache(&base_config.base_url).await?;
67 + }
68 +
69 + let mut config = base_config.clone();
70 + config.max_concurrency = Some(mc);
71 + if let Some(factor) = num_prompts_factor {
72 + config.num_prompts = mc * factor;
73 + }
74 + // Suppress per-run save (we save with concurrency suffix below)
75 + config.save_result = false;
76 + config.append_result = false;
77 +
78 + let result = crate::benchmark::run_benchmark(&config).await?;
79 +
80 + if base_config.save_result {
81 + let model_id = config.model.as_deref().unwrap_or("unknown");
82 + let file_name = compute_result_filename(&config, model_id, &current_dt);
83 + if let Some(ref dir) = config.result_dir {
84 + std::fs::create_dir_all(dir)?;
85 + }
86 + save_result(&result, &file_name)?;
87 + }
88 +
89 + points.push(SweepPoint {
90 + label: format!("concurrency={mc}"),
91 + value: mc as f64,
92 + ss_request_throughput: ss_f64(&result, "request_throughput"),
93 + ss_output_throughput: ss_f64(&result, "output_throughput"),
94 + result_json: result,
95 + });
96 +
97 + println!();
98 + }
99 +
100 + print_sweep_summary(
101 + "Max Concurrency",
102 + &points,
103 + &base_config.sweep_summary_percentiles,
104 + );
105 + Ok(())
106 +}
107 +
108 +/// Run a sweep over multi-turn concurrency values.
109 +pub async fn run_multi_turn_concurrency_sweep(
110 + base_config: &BenchConfig,
111 + values: &[usize],
112 + num_prompts_factor: Option<usize>,
113 +) -> Result<()> {
114 + println!("{:=^70}", " Multi-Turn Concurrency Sweep ");
115 + println!(
116 + "Sweeping --multi-turn-concurrency over {} values: {:?}",
117 + values.len(),
118 + values
119 + );
120 + println!();
121 +
122 + let mut points = Vec::with_capacity(values.len());
123 + let current_dt = chrono::Local::now().format("%Y%m%d-%H%M%S").to_string();
124 +
125 + for (i, &mc) in values.iter().enumerate() {
126 + println!(
127 + "{:-^70}",
128 + format!(
129 + " Run {}/{}: multi_turn_concurrency={} ",
130 + i + 1,
131 + values.len(),
132 + mc
133 + )
134 + );
135 +
136 + if base_config.reset_prefix_cache {
137 + reset_prefix_cache(&base_config.base_url).await?;
138 + }
139 +
140 + let mut config = base_config.clone();
141 + config.multi_turn_concurrency = Some(mc);
142 + // Set max_concurrency so compute_result_filename includes the
143 + // concurrency suffix — without this every sweep point overwrites
144 + // the same file.
145 + config.max_concurrency = Some(mc);
146 + if let Some(factor) = num_prompts_factor {
147 + config.num_prompts = mc * factor;
148 + }
149 + // Suppress per-run save (we save with concurrency suffix below)
150 + config.save_result = false;
151 + config.append_result = false;
152 +
153 + let result = crate::multi_turn::run_multi_turn_benchmark(&config).await?;
154 +
155 + if base_config.save_result {
156 + let model_id = config.model.as_deref().unwrap_or("unknown");
157 + let file_name = compute_result_filename(&config, model_id, &current_dt);
158 + if let Some(ref dir) = config.result_dir {
159 + std::fs::create_dir_all(dir)?;
160 + }
161 + save_result(&result, &file_name)?;
162 + }
163 +
164 + points.push(SweepPoint {
165 + label: format!("concurrency={mc}"),
166 + value: mc as f64,
167 + ss_request_throughput: ss_f64(&result, "request_throughput"),
168 + ss_output_throughput: ss_f64(&result, "output_throughput"),
169 + result_json: result,
170 + });
171 +
172 + println!();
173 + }
174 +
175 + print_sweep_summary(
176 + "MT Concurrency",
177 + &points,
178 + &base_config.sweep_summary_percentiles,
179 + );
180 + Ok(())
181 +}
182 +
183 +/// Run a sweep over request-rate values.
184 +pub async fn run_rate_sweep(base_config: &BenchConfig, values: &[f64]) -> Result<()> {
185 + println!("{:=^70}", " Request Rate Sweep ");
186 + println!(
187 + "Sweeping --request-rate over {} values: {:?}",
188 + values.len(),
189 + values
190 + );
191 + println!();
192 +
193 + let mut points = Vec::with_capacity(values.len());
194 + let current_dt = chrono::Local::now().format("%Y%m%d-%H%M%S").to_string();
195 +
196 + for (i, &rate) in values.iter().enumerate() {
197 + let rate_str = if rate.is_infinite() {
198 + "inf".to_string()
199 + } else {
200 + format!("{rate}")
201 + };
202 + println!(
203 + "{:-^70}",
204 + format!(
205 + " Run {}/{}: request_rate={} ",
206 + i + 1,
207 + values.len(),
208 + rate_str
209 + )
210 + );
211 +
212 + if base_config.reset_prefix_cache {
213 + reset_prefix_cache(&base_config.base_url).await?;
214 + }
215 +
216 + let mut config = base_config.clone();
217 + config.request_rate = rate;
218 + config.save_result = false;
219 + config.append_result = false;
220 +
221 + let result = crate::benchmark::run_benchmark(&config).await?;
222 +
223 + if base_config.save_result {
224 + let model_id = config.model.as_deref().unwrap_or("unknown");
225 + let file_name = compute_result_filename(&config, model_id, &current_dt);
226 + if let Some(ref dir) = config.result_dir {
227 + std::fs::create_dir_all(dir)?;
228 + }
229 + save_result(&result, &file_name)?;
230 + }
231 +
232 + points.push(SweepPoint {
233 + label: format!("rate={rate_str}"),
234 + value: rate,
235 + ss_request_throughput: ss_f64(&result, "request_throughput"),
236 + ss_output_throughput: ss_f64(&result, "output_throughput"),
237 + result_json: result,
238 + });
239 +
240 + println!();
241 + }
242 +
243 + print_sweep_summary(
244 + "Request Rate",
245 + &points,
246 + &base_config.sweep_summary_percentiles,
247 + );
248 + Ok(())
249 +}
250 +
251 +/// Print a summary table after all sweep points complete.
252 +fn print_sweep_summary(param_name: &str, points: &[SweepPoint], summary_percentiles: &[f64]) {
253 + let param_width = param_name.len().max(20);
254 + let columns = build_summary_columns(summary_percentiles);
255 + let total_width = param_width
256 + + columns.iter().map(SummaryColumn::render_width).sum::<usize>()
257 + + columns.len();
258 +
259 + println!("{:=^width$}", " Sweep Summary ", width = total_width);
260 +
261 + print!("{param_name:<param_width$}");
262 + for column in &columns {
263 + print!(" {:>width$}", column.header, width = column.render_width());
264 + }
265 + println!();
266 +
267 + print!("{:-<param_width$}", "");
268 + for column in &columns {
269 + print!(" {:-<width$}", "", width = column.render_width());
270 + }
271 + println!();
272 +
273 + for point in points {
274 + let row = render_summary_row(point, param_width, &columns);
275 + println!("{row}");
276 + }
277 +
278 + println!("{:=<width$}", "", width = total_width);
279 +
280 + // Find the best throughput point
281 + if let Some(best) = points
282 + .iter()
283 + .filter(|p| get_f64(&p.result_json, "request_throughput").is_some())
284 + .max_by(|a, b| {
285 + let ta = get_f64(&a.result_json, "request_throughput").unwrap_or(0.0);
286 + let tb = get_f64(&b.result_json, "request_throughput").unwrap_or(0.0);
287 + ta.partial_cmp(&tb).unwrap_or(std::cmp::Ordering::Equal)
288 + })
289 + {
290 + let tp = get_f64(&best.result_json, "request_throughput").unwrap_or(0.0);
291 + println!("Best throughput: {:.2} req/s at {}", tp, best.label);
292 + }
293 +}
294 +
295 +fn build_summary_columns(summary_percentiles: &[f64]) -> Vec<SummaryColumn> {
296 + let mut columns = vec![
297 + SummaryColumn::new("Req/s", "request_throughput", 10),
298 + SummaryColumn::new("Tok/s", "output_throughput", 10),
299 + SummaryColumn::new("Total tok/s", "total_token_throughput", 12),
300 + SummaryColumn::new("SS req/s", SS_REQUEST_THROUGHPUT_KEY, 10),
301 + SummaryColumn::new("SS out tok/s", SS_OUTPUT_THROUGHPUT_KEY, 12),
302 + SummaryColumn::new("P50 TTFT(ms)", "median_ttft_ms", 12),
303 + SummaryColumn::new("P50 TPOT(ms)", "median_tpot_ms", 12),
304 + SummaryColumn::new("P90 TTFT(ms)", "p90_ttft_ms", 12),
305 + SummaryColumn::new("P90 TPOT(ms)", "p90_tpot_ms", 12),
306 + ];
307 +
308 + for &percentile in summary_percentiles {
309 + if percentile == 50.0 || percentile == 90.0 {
310 + continue;
311 + }
312 + let p_str = format_percentile(percentile);
313 + columns.push(SummaryColumn::new(
314 + &format!("P{p_str} TTFT(ms)"),
315 + &format!("p{p_str}_ttft_ms"),
316 + 12,
317 + ));
318 + columns.push(SummaryColumn::new(
319 + &format!("P{p_str} TPOT(ms)"),
320 + &format!("p{p_str}_tpot_ms"),
321 + 12,
322 + ));
323 + }
324 +
325 + columns
326 +}
327 +
328 +fn render_summary_row(point: &SweepPoint, param_width: usize, columns: &[SummaryColumn]) -> String {
329 + let mut row = format!("{:<param_width$}", point.label, param_width = param_width);
330 + for column in columns {
331 + let value = match column.key.as_str() {
332 + SS_REQUEST_THROUGHPUT_KEY => point.ss_request_throughput,
333 + SS_OUTPUT_THROUGHPUT_KEY => point.ss_output_throughput,
334 + key => get_f64(&point.result_json, key),
335 + };
336 + row.push(' ');
337 + row.push_str(&format!(
338 + "{:>width$}",
339 + fmt_f64(value),
340 + width = column.render_width()
341 + ));
342 + }
343 + row
344 +}
345 +
346 +fn format_percentile(percentile: f64) -> String {
347 + if percentile == percentile.floor() {
348 + format!("{}", percentile as i64)
349 + } else {
350 + format!("{percentile}")
351 + }
352 +}
353 +
354 +struct SummaryColumn {
355 + header: String,
356 + key: String,
357 + width: usize,
358 +}
359 +
360 +impl SummaryColumn {
361 + fn new(header: &str, key: &str, width: usize) -> Self {
362 + Self {
363 + header: header.to_string(),
364 + key: key.to_string(),
365 + width,
366 + }
367 + }
368 +
369 + fn render_width(&self) -> usize {
370 + self.width.max(self.header.len())
371 + }
372 +}
373 +
374 +/// Parse a comma-separated list of concurrency values.
375 +pub fn parse_concurrency_values(s: &str) -> Result<Vec<usize>> {
376 + s.split(',')
377 + .map(|v| {
378 + v.trim().parse::<usize>().map_err(|_| {
379 + BenchError::Config(format!("Invalid concurrency value: '{}'", v.trim()))
380 + })
381 + })
382 + .collect()
383 +}
384 +
385 +/// Parse a comma-separated list of request rate values (supports "inf").
386 +pub fn parse_rate_values(s: &str) -> Result<Vec<f64>> {
387 + s.split(',')
388 + .map(|v| {
389 + let v = v.trim();
390 + if v == "inf" {
391 + Ok(f64::INFINITY)
392 + } else {
393 + v.parse::<f64>()
394 + .map_err(|_| BenchError::Config(format!("Invalid request rate value: '{v}'")))
395 + }
396 + })
397 + .collect()
398 +}
399 +
400 +fn get_f64(json: &serde_json::Value, key: &str) -> Option<f64> {
401 + json.get(key).and_then(|v| v.as_f64())
402 +}
403 +
404 +/// Sentinel keys used on `SummaryColumn` to mark steady-state columns whose
405 +/// values come from `SweepPoint::ss_request_throughput` / `ss_output_throughput`
406 +/// rather than from the result JSON top level.
407 +const SS_REQUEST_THROUGHPUT_KEY: &str = "__ss_request_throughput";
408 +const SS_OUTPUT_THROUGHPUT_KEY: &str = "__ss_output_throughput";
409 +
410 +/// Extract a numeric field from the `steady_state` sub-object of a result JSON.
411 +/// Returns `None` when `steady_state` is missing or null, or when the field is
412 +/// absent / non-numeric.
413 +fn ss_f64(result: &serde_json::Value, key: &str) -> Option<f64> {
414 + result
415 + .get("steady_state")
416 + .and_then(|ss| if ss.is_null() { None } else { Some(ss) })
417 + .and_then(|ss| ss.get(key))
418 + .and_then(|v| v.as_f64())
419 +}
420 +
421 +fn fmt_f64(v: Option<f64>) -> String {
422 + match v {
423 + Some(f) => format!("{:.2}", f),
424 + None => "-".to_string(),
425 + }
426 +}
427 +
428 +#[cfg(test)]
429 +mod tests {
430 + use serde_json::json;
431 +
432 + use super::*;
433 +
434 + fn point(result_json: serde_json::Value) -> SweepPoint {
435 + let ss_request_throughput = ss_f64(&result_json, "request_throughput");
436 + let ss_output_throughput = ss_f64(&result_json, "output_throughput");
437 + SweepPoint {
438 + label: "point".to_string(),
439 + value: 1.0,
440 + result_json,
441 + ss_request_throughput,
442 + ss_output_throughput,
443 + }
444 + }
445 +
446 + #[test]
447 + fn summary_columns_follow_requested_percentile_order() {
448 + let columns = build_summary_columns(&[90.0, 95.0]);
449 + let headers: Vec<&str> = columns.iter().map(|column| column.header.as_str()).collect();
450 +
451 + assert_eq!(
452 + headers,
453 + vec![
454 + "Req/s",
455 + "Tok/s",
456 + "Total tok/s",
457 + "SS req/s",
458 + "SS out tok/s",
459 + "P50 TTFT(ms)",
460 + "P50 TPOT(ms)",
461 + "P90 TTFT(ms)",
462 + "P90 TPOT(ms)",
463 + "P95 TTFT(ms)",
464 + "P95 TPOT(ms)",
465 + ]
466 + );
467 + }
468 +
469 + #[test]
470 + fn summary_row_uses_dash_for_missing_requested_percentile_values() {
471 + let columns = build_summary_columns(&[90.0]);
472 + let row = render_summary_row(
473 + &point(json!({
474 + "request_throughput": 1.0,
475 + "output_throughput": 2.0,
476 + "total_token_throughput": 3.0,
477 + "median_ttft_ms": 10.0,
478 + "median_tpot_ms": 20.0,
479 + "p90_ttft_ms": 40.0,
480 + })),
481 + 20,
482 + &columns,
483 + );
484 +
485 + assert!(row.contains("40.00"));
486 + assert!(row.contains(" -"));
487 + }
488 +
489 + #[test]
490 + fn summary_columns_skip_duplicate_p50_and_p90_percentiles() {
491 + let columns = build_summary_columns(&[50.0, 90.0]);
492 + let headers: Vec<&str> = columns.iter().map(|column| column.header.as_str()).collect();
493 +
494 + assert_eq!(
495 + headers,
496 + vec![
497 + "Req/s",
498 + "Tok/s",
499 + "Total tok/s",
500 + "SS req/s",
501 + "SS out tok/s",
502 + "P50 TTFT(ms)",
503 + "P50 TPOT(ms)",
504 + "P90 TTFT(ms)",
505 + "P90 TPOT(ms)",
506 + ]
507 + );
508 + }
509 +
510 + #[test]
511 + fn summary_row_renders_steady_state_columns() {
512 + let columns = build_summary_columns(&[]);
513 + let row = render_summary_row(
514 + &point(json!({
515 + "request_throughput": 1.0,
516 + "output_throughput": 2.0,
517 + "total_token_throughput": 3.0,
518 + "median_ttft_ms": 10.0,
519 + "median_tpot_ms": 20.0,
520 + "p90_ttft_ms": 40.0,
521 + "p90_tpot_ms": 50.0,
522 + "steady_state": {
523 + "request_throughput": 0.77,
524 + "output_throughput": 88.88,
525 + },
526 + })),
527 + 20,
528 + &columns,
529 + );
530 +
531 + assert!(row.contains("0.77"));
532 + assert!(row.contains("88.88"));
533 + }
534 +
535 + #[test]
536 + fn summary_row_renders_dash_when_steady_state_missing() {
537 + let columns = build_summary_columns(&[]);
538 + let row = render_summary_row(
539 + &point(json!({
540 + "request_throughput": 1.0,
541 + "output_throughput": 2.0,
542 + "total_token_throughput": 3.0,
543 + "median_ttft_ms": 10.0,
544 + "median_tpot_ms": 20.0,
545 + "p90_ttft_ms": 40.0,
546 + "p90_tpot_ms": 50.0,
547 + "steady_state": null,
548 + })),
549 + 20,
550 + &columns,
551 + );
552 +
553 + // Two dashes for the two SS columns (SS req/s, SS out tok/s). Other
554 + // columns have values so they won't emit dashes.
555 + let dash_count = row.matches(" -").count();
556 + assert!(
557 + dash_count >= 2,
558 + "expected at least 2 dashes, got {dash_count}: {row}"
559 + );
560 + }
561 +}
rust/src/bench/src/tiktoken.rs new
+542
@@ -0,0 +1,542 @@
1 +// SPDX-License-Identifier: Apache-2.0
2 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +use std::collections::HashMap;
5 +use std::path::Path;
6 +
7 +use rustc_hash::FxHashMap;
8 +
9 +use crate::error::{BenchError, Result};
10 +
11 +/// Default regex pattern matching cl100k_base (GPT-4, Qwen, etc.)
12 +const DEFAULT_TIKTOKEN_PATTERN: &str = r"(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\r\n\p{L}\p{N}]?\p{L}+|\p{N}{1,3}| ?[^\s\p{L}\p{N}]+[\r\n]*|\s*[\r\n]+|\s+(?!\S)|\s+";
13 +
14 +/// Number of reserved special token slots (matches Python's
15 +/// TikTokenTokenizer.num_reserved_special_tokens)
16 +const NUM_RESERVED_SPECIAL_TOKENS: u32 = 256;
17 +
18 +/// Tiktoken-based tokenizer for models that use tiktoken format (Kimi, Qwen, etc.)
19 +pub struct TiktokenTokenizer {
20 + bpe: tiktoken_rs::CoreBPE,
21 + vocab_size: u32,
22 + #[allow(dead_code)]
23 + num_base_tokens: u32,
24 + /// IDs of all special tokens (for filtering from allowed_tokens)
25 + special_token_ids: Vec<u32>,
26 + /// Reverse mapping: token_id -> byte sequence, for lossy UTF-8 decoding.
27 + /// Empty for built-in encodings (use bpe.decode instead).
28 + decoder: Vec<Vec<u8>>,
29 + /// True for built-in encodings loaded via tiktoken_rs (o200k_base, cl100k_base, etc.)
30 + is_builtin: bool,
31 +}
32 +
33 +impl TiktokenTokenizer {
34 + /// Load from a tiktoken .model file.
35 + ///
36 + /// `special_tokens_from_config`: tokens from tokenizer_config.json's added_tokens_decoder
37 + /// `all_special_tokens`: ALL special tokens including the full 256 reserved slots
38 + pub fn from_file(
39 + model_path: &Path,
40 + all_special_tokens: HashMap<String, u32>,
41 + special_token_ids: Vec<u32>,
42 + pattern: Option<&str>,
43 + ) -> Result<Self> {
44 + let content = std::fs::read_to_string(model_path)
45 + .map_err(|e| BenchError::Tokenizer(format!("Failed to read tiktoken model: {e}")))?;
46 +
47 + let mut encoder: FxHashMap<Vec<u8>, u32> = FxHashMap::default();
48 +
49 + for line in content.lines() {
50 + let line = line.trim();
51 + if line.is_empty() {
52 + continue;
53 + }
54 + let mut parts = line.split_whitespace();
55 + let b64 = match parts.next() {
56 + Some(s) => s,
57 + None => continue,
58 + };
59 + let rank_str = match parts.next() {
60 + Some(s) => s,
61 + None => continue,
62 + };
63 +
64 + use base64::Engine;
65 + let token_bytes = base64::engine::general_purpose::STANDARD
66 + .decode(b64)
67 + .map_err(|e| BenchError::Tokenizer(format!("Invalid base64 in model file: {e}")))?;
68 + let rank: u32 = rank_str
69 + .parse()
70 + .map_err(|e| BenchError::Tokenizer(format!("Invalid rank: {e}")))?;
71 + encoder.insert(token_bytes, rank);
72 + }
73 +
74 + if encoder.is_empty() {
75 + return Err(BenchError::Tokenizer("Empty tiktoken model file".into()));
76 + }
77 +
78 + let num_base_tokens = encoder.len() as u32;
79 +
80 + // Convert special tokens to FxHashMap for tiktoken-rs
81 + let special_fx: FxHashMap<String, u32> =
82 + all_special_tokens.iter().map(|(k, &v)| (k.clone(), v)).collect();
83 +
84 + // vocab_size = base tokens + all reserved special token slots
85 + let vocab_size = num_base_tokens + NUM_RESERVED_SPECIAL_TOKENS;
86 +
87 + // Build reverse mapping for lossy decode
88 + let mut decoder = vec![Vec::new(); vocab_size as usize];
89 + for (bytes, &rank) in &encoder {
90 + if (rank as usize) < decoder.len() {
91 + decoder[rank as usize] = bytes.clone();
92 + }
93 + }
94 + for (text, &rank) in &all_special_tokens {
95 + if (rank as usize) < decoder.len() {
96 + decoder[rank as usize] = text.as_bytes().to_vec();
97 + }
98 + }
99 +
100 + let pat = pattern.unwrap_or(DEFAULT_TIKTOKEN_PATTERN);
101 +
102 + let bpe = tiktoken_rs::CoreBPE::new(encoder, special_fx, pat)
103 + .map_err(|e| BenchError::Tokenizer(format!("Failed to build tiktoken BPE: {e}")))?;
104 +
105 + Ok(Self {
106 + bpe,
107 + vocab_size,
108 + num_base_tokens,
109 + special_token_ids,
110 + decoder,
111 + is_builtin: false,
112 + })
113 + }
114 +
115 + /// Create from a built-in tiktoken encoding (o200k_base, cl100k_base, etc.)
116 + pub fn from_builtin_bpe(bpe: tiktoken_rs::CoreBPE, vocab_size: u32) -> Self {
117 + Self {
118 + bpe,
119 + vocab_size,
120 + num_base_tokens: vocab_size,
121 + special_token_ids: Vec::new(),
122 + decoder: Vec::new(),
123 + is_builtin: true,
124 + }
125 + }
126 +
127 + pub fn encode(&self, text: &str) -> Vec<u32> {
128 + // Use encode_with_special_tokens to match Python's encode(allowed_special="all")
129 + self.bpe.encode_with_special_tokens(text)
130 + }
131 +
132 + /// Decode token IDs to text with lossy UTF-8 handling.
133 + pub fn decode(&self, ids: &[u32]) -> Result<String> {
134 + if self.is_builtin {
135 + // Byte-level decode + lossy UTF-8, matching the file-based path below.
136 + // (CoreBPE::decode errors on invalid UTF-8, which random token
137 + // sequences routinely produce with byte-level BPE vocabularies.)
138 + let bytes: Vec<u8> =
139 + self.bpe._decode_native_and_split(ids.to_vec()).flatten().collect();
140 + return Ok(String::from_utf8_lossy(&bytes).into_owned());
141 + }
142 + let mut bytes = Vec::new();
143 + for &id in ids {
144 + if let Some(token_bytes) = self.decoder.get(id as usize) {
145 + bytes.extend_from_slice(token_bytes);
146 + }
147 + }
148 + Ok(String::from_utf8_lossy(&bytes).into_owned())
149 + }
150 +
151 + pub fn vocab_size(&self) -> u32 {
152 + self.vocab_size
153 + }
154 +
155 + /// Get non-special token IDs whose byte representation is valid UTF-8.
156 + /// Excludes ALL special tokens (like Python's `set(all_tokens) - set(prohibited_tokens)`).
157 + pub fn get_allowed_tokens(&self) -> Vec<u32> {
158 + if self.is_builtin {
159 + // For built-in encodings, return full token range.
160 + // Note: when used with random dataset, these IDs are sent to vLLM as-is;
161 + // ensure the model's tokenizer is compatible (e.g. GPT-4o for o200k_base).
162 + return (0..self.vocab_size).collect();
163 + }
164 + let special_set: std::collections::HashSet<u32> =
165 + self.special_token_ids.iter().copied().collect();
166 +
167 + self.decoder
168 + .iter()
169 + .enumerate()
170 + .filter(|(id, bytes)| {
171 + !bytes.is_empty()
172 + && std::str::from_utf8(bytes).is_ok()
173 + && !special_set.contains(&(*id as u32))
174 + })
175 + .map(|(id, _)| id as u32)
176 + .collect()
177 + }
178 +}
179 +
180 +/// Load a built-in tiktoken encoding by name.
181 +///
182 +/// Supported names: `o200k_base` (GPT-4o), `cl100k_base` (GPT-4/3.5-turbo),
183 +/// `p50k_base`, `r50k_base`, `gpt2`.
184 +///
185 +/// These encodings are bundled with tiktoken-rs — no network download required.
186 +/// Useful for consistent cross-model token counting (e.g. Artificial Analysis methodology).
187 +pub fn load_builtin_tiktoken(encoding: &str) -> Result<TiktokenTokenizer> {
188 + let (bpe, vocab_size) = match encoding {
189 + "o200k_base" => (tiktoken_rs::o200k_base(), 200_275u32),
190 + "cl100k_base" => (tiktoken_rs::cl100k_base(), 100_277u32),
191 + "p50k_base" => (tiktoken_rs::p50k_base(), 50_281u32),
192 + "p50k_edit" => (tiktoken_rs::p50k_edit(), 50_281u32),
193 + "r50k_base" | "gpt2" => (tiktoken_rs::r50k_base(), 50_257u32),
194 + _ => {
195 + return Err(BenchError::Tokenizer(format!(
196 + "Unknown built-in tiktoken encoding: '{encoding}'. \
197 + Supported: o200k_base, cl100k_base, p50k_base, r50k_base, gpt2"
198 + )));
199 + }
200 + };
201 + let bpe = bpe.map_err(|e| BenchError::Tokenizer(format!("Failed to load {encoding}: {e}")))?;
202 + println!("Tokenizer: Built-in tiktoken {encoding} (vocab_size={vocab_size})");
203 + Ok(TiktokenTokenizer::from_builtin_bpe(bpe, vocab_size))
204 +}
205 +
206 +/// Try to load a tiktoken tokenizer from a local directory or HuggingFace model repo.
207 +pub fn try_load_tiktoken(model_id: &str) -> Result<TiktokenTokenizer> {
208 + // Phase 1: If model_id is a local directory, look for tiktoken files there
209 + let local_dir = Path::new(model_id);
210 + if local_dir.is_dir() {
211 + return try_load_tiktoken_from_dir(local_dir, model_id);
212 + }
213 +
214 + // Phase 2: Fall back to HuggingFace Hub download
215 + try_load_tiktoken_from_hf(model_id)
216 +}
217 +
218 +/// Common tiktoken model filenames to search for.
219 +const TIKTOKEN_MODEL_FILENAMES: &[&str] = &["tiktoken.model", "qwen.tiktoken", "vocab.tiktoken"];
220 +
221 +/// Load a tiktoken tokenizer from a local directory.
222 +fn try_load_tiktoken_from_dir(dir: &Path, model_id: &str) -> Result<TiktokenTokenizer> {
223 + let model_path = TIKTOKEN_MODEL_FILENAMES
224 + .iter()
225 + .map(|f| dir.join(f))
226 + .find(|p| p.exists())
227 + .ok_or_else(|| {
228 + BenchError::Tokenizer(format!(
229 + "No tiktoken model file found in local directory '{model_id}' \
230 + (looked for: {})",
231 + TIKTOKEN_MODEL_FILENAMES.join(", ")
232 + ))
233 + })?;
234 +
235 + let num_base_tokens = count_base_tokens(&model_path)?;
236 +
237 + let config_path = dir.join("tokenizer_config.json");
238 + let config = if config_path.exists() {
239 + read_tokenizer_config(&config_path)
240 + } else {
241 + None
242 + };
243 +
244 + let pattern = extract_pat_str_from_local_dir(dir);
245 +
246 + build_tiktoken(model_id, &model_path, config, pattern, num_base_tokens)
247 +}
248 +
249 +/// Load a tiktoken tokenizer from a HuggingFace model repo.
250 +fn try_load_tiktoken_from_hf(model_id: &str) -> Result<TiktokenTokenizer> {
251 + let repo = crate::hub::HubRepo::model(model_id.to_string());
252 +
253 + let model_path = repo
254 + .get("tiktoken.model")
255 + .or_else(|_| repo.get("qwen.tiktoken"))
256 + .or_else(|_| repo.get("vocab.tiktoken"))
257 + .map_err(|_| {
258 + BenchError::Tokenizer(format!("No tiktoken model file found for '{model_id}'"))
259 + })?;
260 +
261 + let num_base_tokens = count_base_tokens(&model_path)?;
262 +
263 + let config = match repo.get("tokenizer_config.json") {
264 + Ok(config_path) => read_tokenizer_config(&config_path),
265 + Err(_) => None,
266 + };
267 +
268 + let pattern = extract_pat_str_from_repo(&repo);
269 +
270 + build_tiktoken(model_id, &model_path, config, pattern, num_base_tokens)
271 +}
272 +
273 +/// Build a TiktokenTokenizer from discovered model file, config, and pattern.
274 +fn build_tiktoken(
275 + model_id: &str,
276 + model_path: &Path,
277 + config: Option<TokenizerConfig>,
278 + pattern: Option<String>,
279 + num_base_tokens: u32,
280 +) -> Result<TiktokenTokenizer> {
281 + // Build the full 256 reserved special tokens map.
282 + // Python: {special_tokens_mapping.get(i, f"<|reserved_token_{i}|>"): i
283 + // for i in range(num_base_tokens, num_base_tokens + 256)}
284 + let config_special = config.as_ref().map(|c| &c.added_tokens).cloned().unwrap_or_default();
285 +
286 + // Invert config_special: id -> content
287 + let id_to_content: HashMap<u32, String> =
288 + config_special.iter().map(|(content, &id)| (id, content.clone())).collect();
289 +
290 + let mut all_special_tokens: HashMap<String, u32> = HashMap::new();
291 + let mut special_token_ids: Vec<u32> = Vec::new();
292 +
293 + for i in 0..NUM_RESERVED_SPECIAL_TOKENS {
294 + let token_id = num_base_tokens + i;
295 + let content = id_to_content
296 + .get(&token_id)
297 + .cloned()
298 + .unwrap_or_else(|| format!("<|reserved_token_{i}|>"));
299 + all_special_tokens.insert(content, token_id);
300 + special_token_ids.push(token_id);
301 + }
302 +
303 + // Also collect special IDs that Python marks as prohibited
304 + // (those with "special": true in added_tokens_decoder)
305 + let prohibited_ids = config.as_ref().map(|c| c.special_ids.clone()).unwrap_or_default();
306 + for id in &prohibited_ids {
307 + if !special_token_ids.contains(id) {
308 + special_token_ids.push(*id);
309 + }
310 + }
311 +
312 + println!(
313 + "Loading tiktoken model for '{model_id}' (base={}, special={}, pat={})...",
314 + num_base_tokens,
315 + all_special_tokens.len(),
316 + if pattern.is_some() {
317 + "custom"
318 + } else {
319 + "default"
320 + },
321 + );
322 +
323 + TiktokenTokenizer::from_file(
324 + model_path,
325 + all_special_tokens,
326 + special_token_ids,
327 + pattern.as_deref(),
328 + )
329 +}
330 +
331 +/// Count base tokens in a tiktoken .model file without fully parsing it.
332 +fn count_base_tokens(model_path: &Path) -> Result<u32> {
333 + let content = std::fs::read_to_string(model_path)
334 + .map_err(|e| BenchError::Tokenizer(format!("Failed to read tiktoken model: {e}")))?;
335 + let count = content
336 + .lines()
337 + .filter(|l| {
338 + let l = l.trim();
339 + !l.is_empty() && l.split_whitespace().count() >= 2
340 + })
341 + .count();
342 + Ok(count as u32)
343 +}
344 +
345 +/// Parsed tokenizer config
346 +struct TokenizerConfig {
347 + /// All tokens from added_tokens_decoder (content -> id)
348 + added_tokens: HashMap<String, u32>,
349 + /// IDs of tokens marked "special": true
350 + special_ids: Vec<u32>,
351 +}
352 +
353 +/// Read and parse tokenizer_config.json
354 +fn read_tokenizer_config(config_path: &Path) -> Option<TokenizerConfig> {
355 + let content = std::fs::read_to_string(config_path).ok()?;
356 + let config: serde_json::Value = serde_json::from_str(&content).ok()?;
357 +
358 + let mut added_tokens = HashMap::new();
359 + let mut special_ids = Vec::new();
360 +
361 + if let Some(added) = config.get("added_tokens_decoder").and_then(|v| v.as_object()) {
362 + for (id_str, token_info) in added {
363 + if let (Ok(id), Some(content)) = (
364 + id_str.parse::<u32>(),
365 + token_info.get("content").and_then(|c| c.as_str()),
366 + ) {
367 + added_tokens.insert(content.to_string(), id);
368 +
369 + let is_special =
370 + token_info.get("special").and_then(|s| s.as_bool()).unwrap_or(false);
371 + if is_special {
372 + special_ids.push(id);
373 + }
374 + }
375 + }
376 + }
377 +
378 + Some(TokenizerConfig {
379 + added_tokens,
380 + special_ids,
381 + })
382 +}
383 +
384 +/// Try to extract pat_str from Python tokenizer source files in a local directory.
385 +/// Returns None if unavailable or unparsable.
386 +fn extract_pat_str_from_local_dir(dir: &Path) -> Option<String> {
387 + ["tokenization_kimi.py", "tokenizer.py"]
388 + .iter()
389 + .map(|f| dir.join(f))
390 + .filter(|p| p.exists())
391 + .find_map(|p| {
392 + std::fs::read_to_string(p)
393 + .ok()
394 + .and_then(|source| extract_pat_str_from_source(&source))
395 + })
396 +}
397 +
398 +/// Try to download the Python tokenizer source file and extract pat_str via regex.
399 +/// Returns None if unavailable or unparsable.
400 +fn extract_pat_str_from_repo(repo: &crate::hub::HubRepo) -> Option<String> {
401 + // Try common Python tokenizer filenames
402 + let py_path = repo.get("tokenization_kimi.py").or_else(|_| repo.get("tokenizer.py")).ok()?;
403 +
404 + let source = std::fs::read_to_string(&py_path).ok()?;
405 +
406 + // Look for pat_str assignment. Common patterns:
407 + // pat_str = "..." or pat_str = '...' or pat_str = "|".join([...])
408 + extract_pat_str_from_source(&source)
409 +}
410 +
411 +/// Parse pat_str from Python source code.
412 +fn extract_pat_str_from_source(source: &str) -> Option<String> {
413 + // Strategy: find `pat_str = "|".join([` and collect the raw string fragments
414 + // This handles the common Kimi/Qwen pattern of joining a list of regex strings.
415 +
416 + // First try: look for pat_str = "|".join([...]) pattern
417 + if let Some(join_start) = source.find("pat_str") {
418 + let after = &source[join_start..];
419 +
420 + // Check for "|".join([ pattern
421 + if let Some(join_pos) = after.find(".join(") {
422 + let after_join = &after[join_pos + 6..]; // skip ".join("
423 + if let Some(bracket_start) = after_join.find('[') {
424 + let inside = &after_join[bracket_start + 1..];
425 + // Collect all string literals inside the list
426 + let mut fragments = Vec::new();
427 + let mut remaining = inside;
428 +
429 + while let Some(frag) = extract_next_python_string(remaining) {
430 + fragments.push(frag.0);
431 + remaining = frag.1;
432 + // Check if we hit the closing bracket
433 + let trimmed = remaining.trim_start();
434 + if trimmed.starts_with(']') {
435 + break;
436 + }
437 + }
438 +
439 + if !fragments.is_empty() {
440 + let pattern = fragments.join("|");
441 + println!(
442 + "Extracted pat_str from Python source: {} fragments",
443 + fragments.len()
444 + );
445 + return Some(pattern);
446 + }
447 + }
448 + }
449 +
450 + // Fallback: simple pat_str = r"..." or pat_str = "..."
451 + if let Some(eq_pos) = after.find('=') {
452 + let after_eq = after[eq_pos + 1..].trim_start();
453 + if let Some(frag) = extract_next_python_string(after_eq) {
454 + return Some(frag.0);
455 + }
456 + }
457 + }
458 +
459 + None
460 +}
461 +
462 +/// Extract the next Python string literal (r"...", "...", r'''...''', etc.)
463 +/// Returns (string_content, remaining_text)
464 +fn extract_next_python_string(s: &str) -> Option<(String, &str)> {
465 + let s = s.trim_start_matches(|c: char| c == ',' || c.is_whitespace());
466 +
467 + // Skip comments
468 + if s.starts_with('#') {
469 + let next_line = s.find('\n').map(|i| i + 1).unwrap_or(s.len());
470 + return extract_next_python_string(&s[next_line..]);
471 + }
472 +
473 + // Check for r""" (triple-quoted raw string)
474 + for prefix in &["r\"\"\"", "r'''"] {
475 + if let Some(inner) = s.strip_prefix(prefix) {
476 + let delim = &prefix[1..]; // """ or '''
477 + if let Some(end) = inner.find(delim) {
478 + let content = &inner[..end];
479 + let rest = &inner[end + delim.len()..];
480 + return Some((content.to_string(), rest));
481 + }
482 + }
483 + }
484 +
485 + // Check for r"..." (raw string)
486 + if let Some(inner) = s.strip_prefix("r\"")
487 + && let Some(end) = inner.find('"')
488 + {
489 + let content = &inner[..end];
490 + let rest = &inner[end + 1..];
491 + return Some((content.to_string(), rest));
492 + }
493 + if let Some(inner) = s.strip_prefix("r'")
494 + && let Some(end) = inner.find('\'')
495 + {
496 + let content = &inner[..end];
497 + let rest = &inner[end + 1..];
498 + return Some((content.to_string(), rest));
499 + }
500 +
501 + // Check for "..." or '...' (regular string — same as raw for regex patterns)
502 + if let Some(inner) = s.strip_prefix('"')
503 + && let Some(end) = find_unescaped(inner, '"')
504 + {
505 + let content = &inner[..end];
506 + let rest = &inner[end + 1..];
507 + return Some((content.to_string(), rest));
508 + }
509 + if let Some(inner) = s.strip_prefix('\'')
510 + && let Some(end) = find_unescaped(inner, '\'')
511 + {
512 + let content = &inner[..end];
513 + let rest = &inner[end + 1..];
514 + return Some((content.to_string(), rest));
515 + }
516 +
517 + // Check for closing bracket — stop
518 + if s.starts_with(']') {
519 + return None;
520 + }
521 +
522 + None
523 +}
524 +
525 +/// Find position of `ch` that is not preceded by a backslash.
526 +fn find_unescaped(s: &str, ch: char) -> Option<usize> {
527 + let mut escaped = false;
528 + for (i, c) in s.char_indices() {
529 + if escaped {
530 + escaped = false;
531 + continue;
532 + }
533 + if c == '\\' {
534 + escaped = true;
535 + continue;
536 + }
537 + if c == ch {
538 + return Some(i);
539 + }
540 + }
541 + None
542 +}
rust/src/bench/src/tokenizer.rs new
+300
@@ -0,0 +1,300 @@
1 +// SPDX-License-Identifier: Apache-2.0
2 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +use std::collections::HashSet;
5 +use std::path::Path;
6 +
7 +use tokenizers::Tokenizer;
8 +
9 +use crate::error::{BenchError, Result};
10 +use crate::tiktoken::TiktokenTokenizer;
11 +
12 +/// Abstraction over local HuggingFace tokenizer, tiktoken, or server-side tokenization.
13 +pub enum TokenizerKind {
14 + Local(Box<Tokenizer>),
15 + Tiktoken(TiktokenTokenizer),
16 + Server(ServerTokenizer),
17 +}
18 +
19 +/// Server-side tokenizer using vLLM's /tokenize and /detokenize endpoints.
20 +pub struct ServerTokenizer {
21 + client: reqwest::blocking::Client,
22 + tokenize_url: String,
23 + detokenize_url: String,
24 + model: String,
25 + cached_vocab_size: u32,
26 +}
27 +
28 +impl ServerTokenizer {
29 + /// Create a new server tokenizer and verify connectivity.
30 + pub fn new(base_url: &str, model: &str) -> Result<Self> {
31 + let client = reqwest::blocking::Client::builder()
32 + .timeout(std::time::Duration::from_secs(30))
33 + .build()
34 + .map_err(|e| BenchError::Tokenizer(format!("Failed to build HTTP client: {e}")))?;
35 +
36 + let tokenize_url = format!("{base_url}/tokenize");
37 + let detokenize_url = format!("{base_url}/detokenize");
38 +
39 + let st = Self {
40 + client,
41 + tokenize_url,
42 + detokenize_url,
43 + model: model.to_string(),
44 + cached_vocab_size: 0,
45 + };
46 +
47 + // Probe the endpoint to verify it works and discover vocab size
48 + let test_tokens = st.encode_inner("test")?;
49 + let max_id = test_tokens.iter().copied().max().unwrap_or(0);
50 + let estimated_vocab = (max_id * 2).max(131072);
51 +
52 + Ok(Self {
53 + cached_vocab_size: estimated_vocab,
54 + ..st
55 + })
56 + }
57 +
58 + fn encode_inner(&self, text: &str) -> Result<Vec<u32>> {
59 + let payload = serde_json::json!({
60 + "model": self.model,
61 + "prompt": text,
62 + });
63 +
64 + let resp = self
65 + .client
66 + .post(&self.tokenize_url)
67 + .json(&payload)
68 + .send()
69 + .map_err(|e| BenchError::Tokenizer(format!("Server tokenize failed: {e}")))?;
70 +
71 + if !resp.status().is_success() {
72 + return Err(BenchError::Tokenizer(format!(
73 + "Server tokenize returned HTTP {}",
74 + resp.status()
75 + )));
76 + }
77 +
78 + let data: serde_json::Value = resp.json().map_err(|e| {
79 + BenchError::Tokenizer(format!("Failed to parse tokenize response: {e}"))
80 + })?;
81 +
82 + let tokens = data
83 + .get("tokens")
84 + .and_then(|t| t.as_array())
85 + .ok_or_else(|| BenchError::Tokenizer("Missing 'tokens' in tokenize response".into()))?;
86 +
87 + tokens
88 + .iter()
89 + .map(|v| {
90 + v.as_u64()
91 + .map(|id| id as u32)
92 + .ok_or_else(|| BenchError::Tokenizer("Invalid token ID in response".into()))
93 + })
94 + .collect()
95 + }
96 +
97 + fn decode_inner(&self, ids: &[u32]) -> Result<String> {
98 + let payload = serde_json::json!({
99 + "model": self.model,
100 + "tokens": ids,
101 + });
102 +
103 + let resp = self
104 + .client
105 + .post(&self.detokenize_url)
106 + .json(&payload)
107 + .send()
108 + .map_err(|e| BenchError::Tokenizer(format!("Server detokenize failed: {e}")))?;
109 +
110 + if !resp.status().is_success() {
111 + return Err(BenchError::Tokenizer(format!(
112 + "Server detokenize returned HTTP {}",
113 + resp.status()
114 + )));
115 + }
116 +
117 + let data: serde_json::Value = resp.json().map_err(|e| {
118 + BenchError::Tokenizer(format!("Failed to parse detokenize response: {e}"))
119 + })?;
120 +
121 + data.get("prompt")
122 + .and_then(|p| p.as_str())
123 + .map(|s| s.to_string())
124 + .ok_or_else(|| BenchError::Tokenizer("Missing 'prompt' in detokenize response".into()))
125 + }
126 +}
127 +
128 +// --- TokenizerKind methods ---
129 +
130 +impl TokenizerKind {
131 + pub fn encode(&self, text: &str, add_special_tokens: bool) -> Result<Vec<u32>> {
132 + match self {
133 + TokenizerKind::Local(tok) => tok
134 + .encode(text, add_special_tokens)
135 + .map(|enc| enc.get_ids().to_vec())
136 + .map_err(|e| BenchError::Tokenizer(format!("Encode failed: {e}"))),
137 + TokenizerKind::Tiktoken(tok) => Ok(tok.encode(text)),
138 + TokenizerKind::Server(srv) => srv.encode_inner(text),
139 + }
140 + }
141 +
142 + pub fn decode(&self, ids: &[u32], skip_special_tokens: bool) -> Result<String> {
143 + match self {
144 + TokenizerKind::Local(tok) => tok
145 + .decode(ids, skip_special_tokens)
146 + .map_err(|e| BenchError::Tokenizer(format!("Decode failed: {e}"))),
147 + TokenizerKind::Tiktoken(tok) => tok.decode(ids),
148 + TokenizerKind::Server(srv) => srv.decode_inner(ids),
149 + }
150 + }
151 +
152 + pub fn vocab_size(&self) -> u32 {
153 + match self {
154 + TokenizerKind::Local(tok) => tok.get_vocab_size(true) as u32,
155 + TokenizerKind::Tiktoken(tok) => tok.vocab_size(),
156 + TokenizerKind::Server(srv) => srv.cached_vocab_size,
157 + }
158 + }
159 +
160 + pub fn num_special_tokens_to_add(&self) -> usize {
161 + match self {
162 + TokenizerKind::Local(tok) => match tok.encode("", true) {
163 + Ok(enc) => enc.get_ids().len(),
164 + Err(_) => 0,
165 + },
166 + TokenizerKind::Tiktoken(_) | TokenizerKind::Server(_) => 0,
167 + }
168 + }
169 +
170 + pub fn get_allowed_tokens(&self) -> Vec<u32> {
171 + match self {
172 + TokenizerKind::Local(tok) => {
173 + let vs = tok.get_vocab_size(true) as u32;
174 + let mut special_ids = HashSet::new();
175 + for (id, token) in tok.get_added_tokens_decoder() {
176 + if token.special {
177 + special_ids.insert(id);
178 + }
179 + }
180 + (0..vs).filter(|id| !special_ids.contains(id)).collect()
181 + }
182 + TokenizerKind::Tiktoken(tok) => tok.get_allowed_tokens(),
183 + TokenizerKind::Server(srv) => (0..srv.cached_vocab_size).collect(),
184 + }
185 + }
186 +}
187 +
188 +/// Load a tokenizer with fallback chain:
189 +/// 0. Built-in tiktoken encoding (o200k_base, cl100k_base, etc.) — no download needed
190 +/// 1. Local tokenizer.json (HuggingFace fast tokenizer)
191 +/// 2. Tiktoken model file (for Kimi, Qwen, etc.)
192 +/// 3. Server-side /tokenize + /detokenize endpoints
193 +///
194 +/// `server_info` is `Some((base_url, model))` to enable server-side fallback.
195 +pub fn load_tokenizer(
196 + model_id: &str,
197 + _trust_remote_code: bool,
198 + server_info: Option<(&str, &str)>,
199 +) -> Result<TokenizerKind> {
200 + // 0. Check for built-in tiktoken encoding names (no HF download needed). These are useful for
201 + // consistent cross-model token counting (e.g. Artificial Analysis).
202 + const BUILTIN_TIKTOKEN: &[&str] = &[
203 + "o200k_base",
204 + "cl100k_base",
205 + "p50k_base",
206 + "p50k_edit",
207 + "r50k_base",
208 + "gpt2",
209 + ];
210 + if BUILTIN_TIKTOKEN.contains(&model_id) {
211 + return crate::tiktoken::load_builtin_tiktoken(model_id).map(TokenizerKind::Tiktoken);
212 + }
213 +
214 + // 1. Try local HuggingFace tokenizer (tokenizer.json)
215 + match try_load_local(model_id) {
216 + Ok(tok) => {
217 + println!("Tokenizer: Local (vocab_size={})", tok.get_vocab_size(true));
218 + Ok(TokenizerKind::Local(Box::new(tok)))
219 + }
220 + Err(local_err) => {
221 + // 2. Try tiktoken format
222 + println!("No tokenizer.json for '{model_id}', trying tiktoken format...");
223 + match crate::tiktoken::try_load_tiktoken(model_id) {
224 + Ok(tok) => {
225 + println!("Tokenizer: Tiktoken (vocab_size={})", tok.vocab_size());
226 + Ok(TokenizerKind::Tiktoken(tok))
227 + }
228 + Err(tiktoken_err) => {
229 + // 3. Try server-side fallback
230 + if let Some((base_url, model)) = server_info {
231 + println!(
232 + "Tiktoken also not available ({tiktoken_err}), \
233 + trying server-side tokenization..."
234 + );
235 + match ServerTokenizer::new(base_url, model) {
236 + Ok(srv) => {
237 + println!(
238 + "Tokenizer: Server (vocab_size≈{})",
239 + srv.cached_vocab_size
240 + );
241 + return Ok(TokenizerKind::Server(srv));
242 + }
243 + Err(srv_err) => {
244 + return Err(BenchError::Tokenizer(format!(
245 + "All tokenizer loading methods failed:\n \
246 + Local: {local_err}\n \
247 + Tiktoken: {tiktoken_err}\n \
248 + Server: {srv_err}\n \
249 + Try --tokenizer with a model that has tokenizer.json."
250 + )));
251 + }
252 + }
253 + }
254 + Err(BenchError::Tokenizer(format!(
255 + "Failed to load tokenizer:\n \
256 + Local: {local_err}\n \
257 + Tiktoken: {tiktoken_err}\n \
258 + Try --tokenizer or provide --base-url for server fallback."
259 + )))
260 + }
261 + }
262 + }
263 + }
264 +}
265 +
266 +/// Try loading tokenizer.json from local path or HuggingFace Hub.
267 +fn try_load_local(model_id: &str) -> Result<Tokenizer> {
268 + // 1. Try local directory with tokenizer.json
269 + let local_path = Path::new(model_id).join("tokenizer.json");
270 + if local_path.exists() {
271 + return Tokenizer::from_file(&local_path).map_err(|e| {
272 + BenchError::Tokenizer(format!(
273 + "Failed to load tokenizer from {}: {e}",
274 + local_path.display()
275 + ))
276 + });
277 + }
278 +
279 + // 2. Try direct path to tokenizer.json
280 + if Path::new(model_id).exists() && model_id.ends_with("tokenizer.json") {
281 + return Tokenizer::from_file(model_id)
282 + .map_err(|e| BenchError::Tokenizer(format!("Failed to load tokenizer: {e}")));
283 + }
284 +
285 + // 3. If model_id is a local directory, don't try HF Hub — let tiktoken fallback handle it
286 + if Path::new(model_id).is_dir() {
287 + return Err(BenchError::Tokenizer(format!(
288 + "No tokenizer.json in local directory '{model_id}'"
289 + )));
290 + }
291 +
292 + // 4. Download from HuggingFace Hub (hf-hub handles auth via HF_TOKEN / cached token)
293 + let repo = crate::hub::HubRepo::model(model_id.to_string());
294 + let tokenizer_path = repo
295 + .get("tokenizer.json")
296 + .map_err(|e| BenchError::Tokenizer(format!("No tokenizer.json for '{model_id}': {e}")))?;
297 +
298 + Tokenizer::from_file(&tokenizer_path)
299 + .map_err(|e| BenchError::Tokenizer(format!("Failed to load downloaded tokenizer: {e}")))
300 +}