diff --git a/rust/src/bench/AGENTS.md b/rust/src/bench/AGENTS.md index 5cf33e0858..62fd963bbf 100644 --- a/rust/src/bench/AGENTS.md +++ b/rust/src/bench/AGENTS.md @@ -50,12 +50,13 @@ cargo test -p vllm-bench -- --ignored - `src/datasets/speed_bench.rs` — NVIDIA SPEED-Bench loader (HF datasets-server API, 6 configs, 11 categories, local cache) - `src/datasets/hf_dataset.rs` — Generic HuggingFace dataset loader (parquet shards via hf-hub into the standard HF hub cache; datasets-server for config/split discovery; native main-branch parquet/JSON fallback for private datasets; column auto-detection) - `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) +- `src/datasets/timed_trace.rs` — Mooncake-style replay trace (`{timestamp, input_length, output_length, hash_ids}` JSONL; `hash_ids` expand to memoized token chunks so prefix sharing reproduces; `--self-timed` fires each request at its recorded offset instead of `--request-rate`; mirrors Python `TimedTrace`) - `src/datasets/prefix_repetition.rs` — Prefix repetition dataset (N shared prefixes × fresh random suffixes, standard prefix-cache stress; mirrors Python `PrefixRepetitionRandomDataset`) - `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`) - `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`) - `src/metrics/mod.rs` — `BenchmarkMetrics` and `MultiTurnMetrics` structs - `src/metrics/calculator.rs` — TTFT/TPOT/ITL/E2EL/throughput stats, goodput SLO checking, peak concurrency, `calculate_multi_turn_metrics` -- `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) +- `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` and off under `--self-timed` (closed-loop only) - `src/output/console.rs` — Terminal output matching Python format + multi-turn per-turn breakdown - `src/output/json.rs` — JSON result file (compatible with Python schema) + multi-turn JSON with `per_turn_metrics` diff --git a/rust/src/bench/src/benchmark.rs b/rust/src/bench/src/benchmark.rs index 161646fe65..5833a3bdce 100644 --- a/rust/src/bench/src/benchmark.rs +++ b/rust/src/bench/src/benchmark.rs @@ -17,9 +17,13 @@ use crate::metrics::calculator::{calculate_embedding_metrics, calculate_metrics} use crate::metrics::steady_state; use crate::output::console::print_results; use crate::output::json::{append_result, build_result_json, compute_result_filename, save_result}; -use crate::rate_control::compute_schedule; +use crate::rate_control::{compute_schedule, trace_schedule}; use crate::ready_checker::{get_first_model, wait_for_endpoint}; +/// Past this first-arrival offset a self-timed run is almost certainly a +/// wrongly scaled or epoch-anchored trace rather than a real idle head. +const TRACE_START_WARN_SECONDS: f64 = 3600.0; + /// Pre-resolve the hostname in `base_url` and pin all resolved IPs on the /// client builder via [`reqwest::ClientBuilder::resolve_to_addrs`]. This /// avoids repeated DNS lookups under high concurrency which can cause @@ -404,6 +408,12 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result { "{} random rerank requests (batch={}, reranker={})", config.num_prompts, config.random_batch_size, config.is_reranker, ), + DatasetName::TimedTrace => format!( + "{} requests from timed trace ({}, chunk_hash_size={})", + config.num_prompts, + config.dataset_path.as_deref().unwrap_or("unknown"), + config.timed_trace_chunk_hash_size, + ), }; tracing::info!( dataset = ?config.dataset_name, @@ -581,6 +591,30 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result { config.is_reranker, )? } + DatasetName::TimedTrace => { + let tok = tokenizer.as_ref().ok_or_else(|| { + BenchError::Config("Timed trace dataset requires a tokenizer".into()) + })?; + let path = config.dataset_path.as_deref().ok_or_else(|| { + BenchError::Config( + "--dataset-path is required for --dataset-name timed_trace".into(), + ) + })?; + crate::datasets::timed_trace::load_timed_trace_dataset( + tok, + path, + config.num_prompts, + &config.request_id_prefix, + &crate::datasets::timed_trace::TimedTraceOptions { + chunk_size: config.timed_trace_chunk_hash_size, + sec_multiplier: config.timed_trace_sec_multiplier, + label_timestamp: &config.timed_trace_label_timestamp, + label_input_length: &config.timed_trace_label_input_length, + label_output_length: &config.timed_trace_label_output_length, + label_hash_ids: &config.timed_trace_label_hash_ids, + }, + )? + } }; let gen_elapsed = gen_start.elapsed(); @@ -800,20 +834,30 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result { tracing::info!("detected speculative decoding; collecting metrics"); } - // Main benchmark - let distribution = if config.burstiness == 1.0 { - "Poisson process" + // Main benchmark. Self-timed runs ignore --request-rate/--burstiness + // entirely (config.rs rejects them), so don't report a rate that isn't used. + if config.self_timed { + tracing::info!( + timing = "trace", + max_concurrency = config.max_concurrency.unwrap_or(config.num_prompts), + prompts = config.num_prompts, + "starting main benchmark run; firing requests at their trace timestamps" + ); } else { - "Gamma distribution" - }; - tracing::info!( - request_rate = config.request_rate, - burstiness = config.burstiness, - distribution, - max_concurrency = config.max_concurrency.unwrap_or(config.num_prompts), - prompts = config.num_prompts, - "starting main benchmark run" - ); + let distribution = if config.burstiness == 1.0 { + "Poisson process" + } else { + "Gamma distribution" + }; + tracing::info!( + request_rate = config.request_rate, + burstiness = config.burstiness, + distribution, + max_concurrency = config.max_concurrency.unwrap_or(config.num_prompts), + prompts = config.num_prompts, + "starting main benchmark run" + ); + } // Pre-assign LoRA adapters to each request (None when --lora-modules not set). let lora_assignments = assign_lora_modules( @@ -832,14 +876,40 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result { ); } - // Compute request schedule - let schedule = compute_schedule( - input_requests.len(), - config.request_rate, - config.burstiness, - config.seed, - config.ramp_up.as_ref(), - ); + // Compute request schedule: the trace's recorded arrival offsets when + // self-timed, synthetic gamma inter-arrivals otherwise. Either way the + // dispatch loop below sleeps to benchmark_start + delay, an absolute + // deadline, so trace replay does not accumulate drift. + let schedule = if config.self_timed { + let schedule = trace_schedule(input_requests.iter().map(|r| r.timestamp)); + let first = schedule.delays.first().copied().unwrap_or(0.0); + let last = schedule.delays.iter().copied().fold(0.0, f64::max); + tracing::info!( + first_arrival_seconds = first, + last_arrival_seconds = last, + "replaying trace arrival schedule" + ); + // Timestamps are absolute offsets from now, never normalized (Python + // does the same). An unscaled millisecond trace or an epoch-anchored + // one parks every request past any plausible run length, so say so + // rather than sitting at 0/N. + if first > TRACE_START_WARN_SECONDS { + tracing::warn!( + first_arrival_seconds = first, + "first trace arrival is far in the future; check \ + --timed-trace-sec-multiplier or re-anchor the trace timestamps" + ); + } + schedule + } else { + compute_schedule( + input_requests.len(), + config.request_rate, + config.burstiness, + config.seed, + config.ramp_up.as_ref(), + ) + }; // Progress bar let pb = if config.disable_tqdm { @@ -1058,7 +1128,10 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result { }; // Attach steady-state metrics when the closed-loop scope gate passes. + // Trace replay is open-loop: an in-flight plateau there reflects a burst + // in the trace, not a saturated server, so it is out of scope. let scope_ok = !config.no_steady_state + && !config.self_timed && config.max_concurrency.is_some() && config.request_rate.is_infinite(); if scope_ok { diff --git a/rust/src/bench/src/cli.rs b/rust/src/bench/src/cli.rs index 4782818ccb..93aa0d9c05 100644 --- a/rust/src/bench/src/cli.rs +++ b/rust/src/bench/src/cli.rs @@ -83,6 +83,8 @@ pub enum DatasetName { PrefixRepetition, #[value(name = "random-rerank")] RandomRerank, + #[value(name = "timed_trace", alias = "timed-trace")] + TimedTrace, } /// Ramp-up strategy for request rate. @@ -321,6 +323,33 @@ pub struct BenchServeArgs { #[arg(long, default_value_t = 128)] pub prefix_repetition_output_len: usize, + // --- Timed-trace dataset --- + /// How many tokens each prefix hash in the trace represents + /// (e.g. 512 for Moonshot traces, 16 for Qwen/Alibaba). + #[arg(long, default_value_t = 16)] + pub timed_trace_chunk_hash_size: usize, + + /// Multiplier converting trace timestamps to seconds + /// (e.g. 0.001 if timestamps are in milliseconds). + #[arg(long, default_value_t = 1.0)] + pub timed_trace_sec_multiplier: f64, + + /// JSON key of the timestamp field in the trace. + #[arg(long, default_value = "timestamp")] + pub timed_trace_label_timestamp: String, + + /// JSON key of the input length field in the trace. + #[arg(long, default_value = "input_length")] + pub timed_trace_label_input_length: String, + + /// JSON key of the output length field in the trace. + #[arg(long, default_value = "output_length")] + pub timed_trace_label_output_length: String, + + /// JSON key of the hash ids field in the trace. + #[arg(long, default_value = "hash_ids")] + pub timed_trace_label_hash_ids: String, + /// Number of prompts to generate. #[arg(long, default_value_t = 1000)] pub num_prompts: usize, @@ -333,6 +362,16 @@ pub struct BenchServeArgs { #[arg(long, default_value_t = 1.0)] pub burstiness: f64, + /// Schedule requests at the timestamps recorded in the trace instead of + /// --request-rate. Defaults to on for --dataset-name timed_trace (the only + /// dataset that carries timestamps); an error for any other dataset. + #[arg(long, overrides_with = "no_self_timed")] + pub self_timed: bool, + + /// Force trace-driven timing off for timed_trace and use --request-rate. + #[arg(long, overrides_with = "self_timed")] + pub no_self_timed: bool, + /// Maximum number of concurrent requests. #[arg(long)] pub max_concurrency: Option, diff --git a/rust/src/bench/src/config.rs b/rust/src/bench/src/config.rs index 197cf08e7d..3e180cfb30 100644 --- a/rust/src/bench/src/config.rs +++ b/rust/src/bench/src/config.rs @@ -198,6 +198,15 @@ pub struct BenchConfig { pub hf_subset: Option, pub hf_output_len: Option, pub hf_text_column: Option, + /// Fire each request at its trace-recorded timestamp instead of a + /// synthetic --request-rate schedule. Only timed_trace can set this. + pub self_timed: bool, + pub timed_trace_chunk_hash_size: usize, + pub timed_trace_sec_multiplier: f64, + pub timed_trace_label_timestamp: String, + pub timed_trace_label_input_length: String, + pub timed_trace_label_output_length: String, + pub timed_trace_label_hash_ids: String, pub reset_prefix_cache: bool, pub prompt_token_ids: bool, // --- Random multimodal dataset --- @@ -376,6 +385,79 @@ impl BenchConfig { let mut multi_turn_min_turns = args.multi_turn_num_turns; let mut multi_turn_max_turns = args.multi_turn_num_turns; + let self_timed = if args.dataset_name == DatasetName::TimedTrace { + if !matches!(args.backend, BackendKind::Vllm | BackendKind::Openai) { + return Err(BenchError::Config( + "timed_trace passes pre-tokenized prompts (list[int]) and requires \ + a completions backend ('vllm' or 'openai')" + .into(), + )); + } + if args.dataset_path.is_none() { + return Err(BenchError::Config( + "--dataset-path is required for --dataset-name timed_trace".into(), + )); + } + // A zero, negative or non-finite scale turns trace timestamps into + // offsets the scheduler cannot sleep on. + if !(args.timed_trace_sec_multiplier.is_finite() + && args.timed_trace_sec_multiplier > 0.0) + { + return Err(BenchError::Config(format!( + "--timed-trace-sec-multiplier must be finite and positive, got {}", + args.timed_trace_sec_multiplier + ))); + } + let self_timed = !args.no_self_timed; + if self_timed { + // Flags that shape a synthetic schedule are unreachable once + // the trace supplies arrival times. Accepting them silently + // would still stamp the result JSON and its filename with a + // request rate that never ran. + let mut unused: Vec<&str> = Vec::new(); + if args.request_rate.is_finite() { + unused.push("--request-rate"); + } + if args.burstiness != 1.0 { + unused.push("--burstiness"); + } + if args.ramp_up_strategy.is_some() { + unused.push("--ramp-up-strategy"); + } + if args.sweep_request_rate.is_some() { + unused.push("--sweep-request-rate"); + } + // A per-point num_prompts truncates the trace to a different + // number of rows, so each sweep point replays a different span. + if args.sweep_max_concurrency.is_some() && args.sweep_num_prompts_factor.is_some() { + return Err(BenchError::Config( + "--sweep-num-prompts-factor truncates the trace to a different \ + number of rows per sweep point, so each point replays a \ + different span; drop it to sweep concurrency over the whole trace" + .into(), + )); + } + if !unused.is_empty() { + return Err(BenchError::Config(format!( + "{} {} no effect under trace-driven timing; pass --no-self-timed \ + to schedule timed_trace by --request-rate instead", + unused.join(" / "), + if unused.len() == 1 { "has" } else { "have" }, + ))); + } + } + self_timed + } else { + if args.self_timed || args.no_self_timed { + return Err(BenchError::Config( + "--self-timed/--no-self-timed is only supported with \ + --dataset-name timed_trace" + .into(), + )); + } + false + }; + // For random datasets with openai-compatible backends, default to ignore_eos. // Exception: multi-turn mode, where ignore_eos causes unbounded context growth // across turns. Multi-turn uses min_tokens instead for output length control. @@ -384,6 +466,8 @@ impl BenchConfig { false } else { args.ignore_eos + // timed_trace: generation must run to the trace's output_length + || args.dataset_name == DatasetName::TimedTrace || ((args.dataset_name == DatasetName::Random || args.dataset_name == DatasetName::RandomMm) && args.backend.is_openai_compatible() @@ -745,6 +829,13 @@ impl BenchConfig { hf_subset: args.hf_subset.clone(), hf_output_len: args.hf_output_len, hf_text_column: args.hf_text_column.clone(), + self_timed, + timed_trace_chunk_hash_size: args.timed_trace_chunk_hash_size, + timed_trace_sec_multiplier: args.timed_trace_sec_multiplier, + timed_trace_label_timestamp: args.timed_trace_label_timestamp.clone(), + timed_trace_label_input_length: args.timed_trace_label_input_length.clone(), + timed_trace_label_output_length: args.timed_trace_label_output_length.clone(), + timed_trace_label_hash_ids: args.timed_trace_label_hash_ids.clone(), reset_prefix_cache: args.reset_prefix_cache, prompt_token_ids: args.prompt_token_ids, random_mm_base_items_per_request: args.random_mm_base_items_per_request, @@ -899,6 +990,98 @@ mod tests { TestCli::parse_from(args).args } + fn base_timed_trace_args() -> Vec<&'static str> { + vec![ + "vllm-bench", + "--model", + "test-model", + "--dataset-name", + "timed_trace", + "--dataset-path", + "trace.jsonl", + ] + } + + #[test] + fn test_timed_trace_defaults_to_self_timed() { + let config = BenchConfig::from_args(&parse_args(base_timed_trace_args())).unwrap(); + assert!(config.self_timed); + assert!(config.ignore_eos); + } + + #[test] + fn test_no_self_timed_restores_rate_scheduling() { + let mut args = base_timed_trace_args(); + args.extend(["--no-self-timed", "--request-rate", "10"]); + let config = BenchConfig::from_args(&parse_args(args)).unwrap(); + assert!(!config.self_timed); + assert_eq!(config.request_rate, 10.0); + } + + /// Rate/burstiness/ramp-up shape a synthetic schedule the trace replaces; + /// accepting them would mislabel the result JSON with a rate that never ran. + #[test] + fn test_self_timed_rejects_synthetic_schedule_flags() { + for flag in [ + vec!["--request-rate", "10"], + vec!["--burstiness", "0.5"], + vec!["--sweep-request-rate", "1,5,10"], + ] { + let mut args = base_timed_trace_args(); + args.extend(flag.iter().copied()); + let msg = BenchConfig::from_args(&parse_args(args)) + .expect_err("self-timed replay should reject synthetic schedule flags") + .to_string(); + assert!( + msg.contains(flag[0]) && msg.contains("--no-self-timed"), + "{flag:?}: {msg}" + ); + } + } + + /// A per-point num_prompts truncates the trace differently at every sweep + /// point, so the replayed span stops being comparable across them. + #[test] + fn test_self_timed_rejects_sweep_num_prompts_factor() { + let mut args = base_timed_trace_args(); + args.extend([ + "--sweep-max-concurrency", + "8,16", + "--sweep-num-prompts-factor", + "4", + ]); + let err = BenchConfig::from_args(&parse_args(args)).expect_err("should reject the factor"); + let msg = err.to_string(); + assert!( + msg.contains("--sweep-num-prompts-factor") && msg.contains("different span"), + "{msg}" + ); + } + + #[test] + fn test_self_timed_allows_concurrency_sweep_over_full_trace() { + let mut args = base_timed_trace_args(); + args.extend(["--sweep-max-concurrency", "8,16"]); + assert!(BenchConfig::from_args(&parse_args(args)).is_ok()); + } + + #[test] + fn test_non_positive_sec_multiplier_fails() { + // `=` form: clap reads a bare "-0.001" as a flag. + for value in [ + "--timed-trace-sec-multiplier=0", + "--timed-trace-sec-multiplier=-0.001", + "--timed-trace-sec-multiplier=inf", + "--timed-trace-sec-multiplier=nan", + ] { + let mut args = base_timed_trace_args(); + args.push(value); + let err = BenchConfig::from_args(&parse_args(args)) + .expect_err("should reject a multiplier that cannot scale timestamps"); + assert!(err.to_string().contains("--timed-trace-sec-multiplier")); + } + } + fn base_multi_turn_args() -> Vec<&'static str> { vec![ "vllm-bench", diff --git a/rust/src/bench/src/datasets/custom.rs b/rust/src/bench/src/datasets/custom.rs index 11ddee02b1..a9fcba3354 100644 --- a/rust/src/bench/src/datasets/custom.rs +++ b/rust/src/bench/src/datasets/custom.rs @@ -43,29 +43,10 @@ pub fn load_custom_dataset( no_oversample: bool, disable_shuffle: bool, ) -> Result> { - let content = std::fs::read_to_string(path) - .map_err(|e| BenchError::Config(format!("Failed to read custom dataset '{path}': {e}")))?; - - let mut lines: Vec = Vec::new(); - for (lineno, line) in content.lines().enumerate() { - let line = line.trim(); - if line.is_empty() { - continue; - } - let parsed: CustomLine = serde_json::from_str(line).map_err(|e| { - BenchError::Config(format!( - "Invalid JSONL at {path}:{}: {e} (each line must be an object \ - with a 'prompt' field)", - lineno + 1 - )) - })?; - lines.push(parsed); - } - if lines.is_empty() { - return Err(BenchError::Config(format!( - "Custom dataset '{path}' contains no entries" - ))); - } + let mut lines: Vec = super::read_jsonl(path, "custom dataset", None)? + .into_iter() + .map(|(_, line)| line) + .collect(); // Python shuffles the loaded data (seeded) before taking num_requests. if !disable_shuffle { @@ -118,26 +99,13 @@ pub fn load_custom_dataset( #[cfg(test)] mod tests { + use super::super::test_support::{test_tokenizer, write_temp_jsonl}; use super::*; - fn write_temp_jsonl(name: &str, content: &str) -> String { - let path = std::env::temp_dir().join(format!("vllm-bench-custom-{name}.jsonl")); - std::fs::write(&path, content).unwrap(); - path.to_string_lossy().into_owned() - } - - /// gpt2 via built-in tiktoken encoding — loads without network access. - fn test_tokenizer() -> TokenizerKind { - TokenizerKind::Tiktoken( - crate::tiktoken::load_builtin_tiktoken("gpt2") - .expect("gpt2 built-in tiktoken should always load without network"), - ) - } - #[test] fn test_load_custom_dataset_basic() { let path = write_temp_jsonl( - "basic", + "custom-basic", r#"{"prompt": "hello world", "output_tokens": 10} {"prompt": "foo bar baz", "output_tokens": 20} "#, @@ -154,7 +122,7 @@ mod tests { #[test] fn test_load_custom_dataset_per_line_output_tokens() { let path = write_temp_jsonl( - "perline", + "custom-perline", r#"{"prompt": "hello", "output_tokens": 10} {"prompt": "world", "output_tokens": 20} "#, @@ -167,7 +135,7 @@ mod tests { #[test] fn test_load_custom_dataset_missing_output_tokens_errors() { - let path = write_temp_jsonl("missing", r#"{"prompt": "hello"}"#); + let path = write_temp_jsonl("custom-missing", r#"{"prompt": "hello"}"#); let err = load_custom_dataset(&test_tokenizer(), &path, 1, -1, 0, "t-", true, true) .expect_err("should fail without output_tokens"); assert!(err.to_string().contains("output_tokens")); @@ -175,7 +143,7 @@ mod tests { #[test] fn test_load_custom_dataset_missing_prompt_errors() { - let path = write_temp_jsonl("noprompt", r#"{"text": "hello"}"#); + let path = write_temp_jsonl("custom-noprompt", r#"{"text": "hello"}"#); assert!( load_custom_dataset(&test_tokenizer(), &path, 1, 256, 0, "t-", true, true).is_err() ); diff --git a/rust/src/bench/src/datasets/mod.rs b/rust/src/bench/src/datasets/mod.rs index 4848ec5303..bfdc4006a6 100644 --- a/rust/src/bench/src/datasets/mod.rs +++ b/rust/src/bench/src/datasets/mod.rs @@ -11,6 +11,7 @@ pub mod random_mm; pub mod random_rerank; pub mod sharegpt; pub mod speed_bench; +pub mod timed_trace; use std::sync::Arc; @@ -54,6 +55,10 @@ pub struct SampleRequest { /// rerank sends `[0]` as the query and `[1..]` as documents (random-rerank). /// Mirrors Python's list-valued `SampleRequest.prompt`. pub prompt_list: Option]>>, + /// Arrival time in seconds relative to benchmark start (timed_trace only). + /// With --self-timed the scheduler fires the request at this offset + /// instead of a synthetic --request-rate schedule. + pub timestamp: Option, } impl Default for SampleRequest { @@ -69,6 +74,7 @@ impl Default for SampleRequest { multi_modal_content: None, chat_messages_json: None, prompt_list: None, + timestamp: None, } } } @@ -137,6 +143,64 @@ pub fn batch_requests( .collect() } +/// Read a JSONL file into `T` values, skipping blank lines. +/// +/// `what` names the dataset in error messages ("custom dataset", "timed +/// trace"); `max_entries` stops the read early. Returns `(line_number, value)` +/// pairs with 0-based line numbers so callers can index or report by row. +pub fn read_jsonl( + path: &str, + what: &str, + max_entries: Option, +) -> crate::error::Result> { + use crate::error::BenchError; + + let content = std::fs::read_to_string(path) + .map_err(|e| BenchError::Config(format!("Failed to read {what} '{path}': {e}")))?; + + let mut entries: Vec<(usize, T)> = Vec::new(); + for (lineno, line) in content.lines().enumerate() { + if max_entries.is_some_and(|max| entries.len() >= max) { + break; + } + let line = line.trim(); + if line.is_empty() { + continue; + } + let value = serde_json::from_str(line).map_err(|e| { + BenchError::Config(format!("Invalid JSONL at {path}:{}: {e}", lineno + 1)) + })?; + entries.push((lineno, value)); + } + + if entries.is_empty() { + return Err(BenchError::Config(format!("No entries in {what} '{path}'"))); + } + Ok(entries) +} + +/// Fixtures shared by the dataset test modules. +#[cfg(test)] +pub(crate) mod test_support { + use crate::tokenizer::TokenizerKind; + + /// gpt2 via built-in tiktoken encoding — loads without network access. + pub(crate) fn test_tokenizer() -> TokenizerKind { + TokenizerKind::Tiktoken( + crate::tiktoken::load_builtin_tiktoken("gpt2") + .expect("gpt2 built-in tiktoken should always load without network"), + ) + } + + /// Write a JSONL fixture to the temp dir. `name` must be unique across the + /// whole test suite — all fixtures share one directory. + pub(crate) fn write_temp_jsonl(name: &str, content: &str) -> String { + let path = std::env::temp_dir().join(format!("vllm-bench-{name}.jsonl")); + std::fs::write(&path, content).unwrap(); + path.to_string_lossy().into_owned() + } +} + #[cfg(test)] mod tests { use super::*; diff --git a/rust/src/bench/src/datasets/prefix_repetition.rs b/rust/src/bench/src/datasets/prefix_repetition.rs index e4dfcd36a3..7eb5fac415 100644 --- a/rust/src/bench/src/datasets/prefix_repetition.rs +++ b/rust/src/bench/src/datasets/prefix_repetition.rs @@ -107,16 +107,9 @@ pub fn generate_prefix_repetition_dataset( #[cfg(test)] mod tests { + use super::super::test_support::test_tokenizer; use super::*; - /// gpt2 via built-in tiktoken encoding — loads without network access. - fn test_tokenizer() -> TokenizerKind { - TokenizerKind::Tiktoken( - crate::tiktoken::load_builtin_tiktoken("gpt2") - .expect("gpt2 built-in tiktoken should always load without network"), - ) - } - #[test] fn test_prefix_repetition_structure() { let tok = test_tokenizer(); diff --git a/rust/src/bench/src/datasets/random_rerank.rs b/rust/src/bench/src/datasets/random_rerank.rs index ad3686cf3f..76a7eb3427 100644 --- a/rust/src/bench/src/datasets/random_rerank.rs +++ b/rust/src/bench/src/datasets/random_rerank.rs @@ -135,16 +135,9 @@ pub fn generate_random_rerank_dataset( #[cfg(test)] mod tests { + use super::super::test_support::test_tokenizer; use super::*; - /// gpt2 via built-in tiktoken encoding — loads without network access. - fn test_tokenizer() -> TokenizerKind { - TokenizerKind::Tiktoken( - crate::tiktoken::load_builtin_tiktoken("gpt2") - .expect("gpt2 built-in tiktoken should always load without network"), - ) - } - fn fixed_ratio() -> RangeRatio { RangeRatio::parse("0.0").unwrap() } diff --git a/rust/src/bench/src/datasets/timed_trace.rs b/rust/src/bench/src/datasets/timed_trace.rs new file mode 100644 index 0000000000..f90a94450a --- /dev/null +++ b/rust/src/bench/src/datasets/timed_trace.rs @@ -0,0 +1,376 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +//! Timed-trace dataset: JSONL replay trace with per-request arrival timestamps. +//! +//! ```jsonl +//! {"timestamp": 12.4, "input_length": 42, "output_length": 8, "hash_ids": [981, 4412, 7]} +//! ``` +//! +//! Mirrors Python's `TimedTrace`. Prompts carry no text: each `hash_id` expands +//! to a pseudo-random token chunk of `chunk_size` tokens (the last chunk may be +//! partial so the prompt length equals `input_length` exactly). Chunks are +//! memoized per (hash_id, size), so requests sharing a hash prefix get +//! byte-identical token prefixes and prefix-cache hits reproduce. Field names +//! are remappable via `--timed-trace-label-*`; timestamps are scaled by +//! `--timed-trace-sec-multiplier`. +//! +//! Deviation from Python: chunk contents are seeded from (hash_id, size) with a +//! fixed mixing function, so they are bit-reproducible across runs. Python +//! seeds via `hash(f"{h}:{size}")`, which is randomized per process unless +//! PYTHONHASHSEED is pinned. + +use std::collections::HashMap; +use std::sync::Arc; + +use rand::rngs::StdRng; +use rand::{Rng, SeedableRng}; +use serde_json::Value; + +use super::SampleRequest; +use crate::error::{BenchError, Result}; +use crate::tokenizer::TokenizerKind; + +pub struct TimedTraceOptions<'a> { + pub chunk_size: usize, + pub sec_multiplier: f64, + pub label_timestamp: &'a str, + pub label_input_length: &'a str, + pub label_output_length: &'a str, + pub label_hash_ids: &'a str, +} + +/// Load a timed-trace JSONL dataset. Row order (and thus the timestamp +/// schedule) is preserved: no shuffle, no oversampling. +pub fn load_timed_trace_dataset( + tokenizer: &TokenizerKind, + path: &str, + num_requests: usize, + request_id_prefix: &str, + opts: &TimedTraceOptions<'_>, +) -> Result> { + let allowed_tokens = tokenizer.get_allowed_tokens(); + if allowed_tokens.is_empty() { + return Err(BenchError::Tokenizer("No allowed tokens found".into())); + } + + let entries: Vec<(usize, Value)> = super::read_jsonl(path, "timed trace", Some(num_requests))?; + + let mut chunk_cache: HashMap<(u64, usize), Arc<[u32]>> = HashMap::new(); + let mut requests: Vec = Vec::with_capacity(entries.len()); + + for (lineno, entry) in &entries { + let input_length = require_u64( + entry, + opts.label_input_length, + "--timed-trace-label-input-length", + )? as usize; + let output_length = require_u64( + entry, + opts.label_output_length, + "--timed-trace-label-output-length", + )? as usize; + let timestamp = require_f64(entry, opts.label_timestamp, "--timed-trace-label-timestamp")? + * opts.sec_multiplier; + // A non-finite offset panics in `Duration::from_secs_f64` once the + // scheduler consumes it; reject here, naming the offending row. + if !timestamp.is_finite() { + return Err(BenchError::Config(format!( + "Field '{}' at {path}:{} scales to a non-finite offset ({timestamp}); \ + check --timed-trace-sec-multiplier", + opts.label_timestamp, + lineno + 1, + ))); + } + let hash_ids = parse_hash_ids(entry, opts.label_hash_ids)?; + + let prompt_ids = expand_prompt( + &hash_ids, + input_length, + opts.chunk_size, + &allowed_tokens, + &mut chunk_cache, + ); + + requests.push(SampleRequest { + prompt_len: prompt_ids.len(), + prompt_token_ids: Some(Arc::from(prompt_ids)), + expected_output_len: output_length, + request_id: Some(format!("{request_id_prefix}{lineno}")), + timestamp: Some(timestamp), + ..Default::default() + }); + } + + Ok(requests) +} + +/// Expand prefix-hash ids into token ids, consuming `chunk_size` tokens per +/// hash until `target_len` is reached (the final chunk may be partial). A +/// repeated (hash_id, size) yields the identical chunk via the cache. If the +/// trace supplies fewer hashes than `target_len / chunk_size`, the prompt +/// comes out shorter — same as Python. +fn expand_prompt( + hash_ids: &[u64], + target_len: usize, + chunk_size: usize, + allowed_tokens: &[u32], + cache: &mut HashMap<(u64, usize), Arc<[u32]>>, +) -> Vec { + // `target_len` is raw trace input, so reserve only what the hashes can + // actually produce: a corrupt row, or one in the wrong unit, would + // otherwise ask the allocator for terabytes before any length filter runs. + let capacity = target_len.min(hash_ids.len().saturating_mul(chunk_size)); + let mut prompt = Vec::with_capacity(capacity); + let mut remaining = target_len; + for &h in hash_ids { + let size = remaining.min(chunk_size); + let chunk = cache.entry((h, size)).or_insert_with(|| { + let mut rng = StdRng::seed_from_u64(chunk_seed(h, size)); + (0..size) + .map(|_| allowed_tokens[rng.random_range(0..allowed_tokens.len())]) + .collect() + }); + prompt.extend_from_slice(chunk); + remaining -= size; + if remaining == 0 { + break; + } + } + prompt +} + +/// Stable seed for a chunk. The golden-ratio multiply is bijective on u64, so +/// distinct `hash_id`s never collide at a fixed `size`. The trailing xor is not +/// injective across sizes: two unrelated hashes seen at different sizes can +/// alias and share a token prefix, at ~2^-64 odds per pair. +fn chunk_seed(hash_id: u64, size: usize) -> u64 { + hash_id.wrapping_mul(0x9E37_79B9_7F4A_7C15) ^ (size as u64) +} + +fn require_field<'v>(entry: &'v Value, label: &str, flag: &str) -> Result<&'v Value> { + entry.get(label).ok_or_else(|| { + let available: Vec<&str> = entry + .as_object() + .map(|o| o.keys().map(String::as_str).collect()) + .unwrap_or_default(); + BenchError::Config(format!( + "Field '{label}' not found in trace entry. Available fields: \ + {available:?}. Use {flag} to specify the correct field name." + )) + }) +} + +/// Float-encoded integers (`512.0`) are accepted, matching Python's `int()`. +fn require_u64(entry: &Value, label: &str, flag: &str) -> Result { + let v = require_field(entry, label, flag)?; + v.as_u64() + .or_else(|| { + v.as_f64() + .filter(|f| f.fract() == 0.0 && (0.0..u64::MAX as f64).contains(f)) + .map(|f| f as u64) + }) + .ok_or_else(|| { + BenchError::Config(format!( + "Field '{label}' must be a non-negative integer, got: {v}" + )) + }) +} + +fn require_f64(entry: &Value, label: &str, flag: &str) -> Result { + let v = require_field(entry, label, flag)?; + v.as_f64() + .ok_or_else(|| BenchError::Config(format!("Field '{label}' must be a number, got: {v}"))) +} + +/// Missing hash_ids field means an empty prompt (mirrors Python's +/// `entry.get(label, [])`); a present field must be an array of integers. +/// +/// Ids are kept as `u64` because traces carry raw unsigned 64-bit block hashes; +/// negative ids keep the two's-complement value `chunk_seed` used before. +fn parse_hash_ids(entry: &Value, label: &str) -> Result> { + let Some(v) = entry.get(label) else { + return Ok(Vec::new()); + }; + let arr = v + .as_array() + .ok_or_else(|| BenchError::Config(format!("Field '{label}' must be an array, got: {v}")))?; + arr.iter() + .map(|x| { + x.as_u64().or_else(|| x.as_i64().map(|i| i as u64)).ok_or_else(|| { + BenchError::Config(format!( + "Field '{label}' must contain only integers, got: {x}" + )) + }) + }) + .collect() +} + +#[cfg(test)] +mod tests { + use super::super::test_support::{test_tokenizer, write_temp_jsonl}; + use super::*; + + fn opts(chunk_size: usize, sec_multiplier: f64) -> TimedTraceOptions<'static> { + TimedTraceOptions { + chunk_size, + sec_multiplier, + label_timestamp: "timestamp", + label_input_length: "input_length", + label_output_length: "output_length", + label_hash_ids: "hash_ids", + } + } + + #[test] + fn test_prompt_len_matches_input_length_and_fields_map() { + let path = write_temp_jsonl( + "timed-trace-basic", + r#"{"timestamp": 1000, "input_length": 40, "output_length": 7, "hash_ids": [1, 2, 3]} +{"timestamp": 2500, "input_length": 32, "output_length": 9, "hash_ids": [1, 4]} +"#, + ); + let reqs = load_timed_trace_dataset(&test_tokenizer(), &path, 10, "t-", &opts(16, 0.001)) + .expect("load should succeed"); + assert_eq!(reqs.len(), 2); + // 40 = 16 + 16 + partial 8 from three hashes + assert_eq!(reqs[0].prompt_len, 40); + assert_eq!(reqs[0].prompt_token_ids.as_ref().unwrap().len(), 40); + assert_eq!(reqs[0].expected_output_len, 7); + assert_eq!(reqs[0].timestamp, Some(1.0)); + assert_eq!(reqs[1].timestamp, Some(2.5)); + assert_eq!(reqs[0].request_id.as_deref(), Some("t-0")); + } + + #[test] + fn test_shared_hash_prefix_gives_identical_tokens() { + let path = write_temp_jsonl( + "timed-trace-prefix", + r#"{"timestamp": 0, "input_length": 32, "output_length": 1, "hash_ids": [7, 8]} +{"timestamp": 1, "input_length": 32, "output_length": 1, "hash_ids": [7, 9]} +"#, + ); + let reqs = load_timed_trace_dataset(&test_tokenizer(), &path, 10, "t-", &opts(16, 1.0)) + .expect("load should succeed"); + let a = reqs[0].prompt_token_ids.as_ref().unwrap(); + let b = reqs[1].prompt_token_ids.as_ref().unwrap(); + assert_eq!(a[..16], b[..16], "shared hash 7 must expand identically"); + assert_ne!(a[16..], b[16..], "distinct hashes 8/9 must differ"); + } + + #[test] + fn test_reload_is_bit_identical() { + let path = write_temp_jsonl( + "timed-trace-determinism", + r#"{"timestamp": 0, "input_length": 40, "output_length": 1, "hash_ids": [11, 12, 13]}"#, + ); + let load = || { + load_timed_trace_dataset(&test_tokenizer(), &path, 10, "t-", &opts(16, 1.0)) + .expect("load should succeed") + }; + assert_eq!(load()[0].prompt_token_ids, load()[0].prompt_token_ids); + } + + #[test] + fn test_too_few_hashes_yields_short_prompt() { + let path = write_temp_jsonl( + "timed-trace-short", + r#"{"timestamp": 0, "input_length": 100, "output_length": 1, "hash_ids": [1]}"#, + ); + let reqs = load_timed_trace_dataset(&test_tokenizer(), &path, 10, "t-", &opts(16, 1.0)) + .expect("load should succeed"); + assert_eq!(reqs[0].prompt_len, 16, "one hash covers only one chunk"); + } + + /// `input_length` is untrusted: the reservation must follow what the hashes + /// can produce, not the raw field, or the allocator aborts the process. + #[test] + fn test_absurd_input_length_does_not_overallocate() { + let path = write_temp_jsonl( + "timed-trace-absurd-len", + r#"{"timestamp": 0, "input_length": 10000000000000, "output_length": 1, "hash_ids": [1]}"#, + ); + let reqs = load_timed_trace_dataset(&test_tokenizer(), &path, 10, "t-", &opts(16, 1.0)) + .expect("load should succeed"); + assert_eq!(reqs[0].prompt_len, 16); + } + + #[test] + fn test_num_requests_caps_rows() { + let path = write_temp_jsonl( + "timed-trace-cap", + r#"{"timestamp": 0, "input_length": 8, "output_length": 1, "hash_ids": [1]} +{"timestamp": 1, "input_length": 8, "output_length": 1, "hash_ids": [2]} +{"timestamp": 2, "input_length": 8, "output_length": 1, "hash_ids": [3]} +"#, + ); + let reqs = load_timed_trace_dataset(&test_tokenizer(), &path, 2, "t-", &opts(16, 1.0)) + .expect("load should succeed"); + assert_eq!(reqs.len(), 2); + } + + #[test] + fn test_missing_field_error_names_label_and_flag() { + let path = write_temp_jsonl( + "timed-trace-missing", + r#"{"timestamp": 0, "input_length": 8, "output_length": 1, "hash_ids": [1]}"#, + ); + let mut o = opts(16, 1.0); + o.label_input_length = "input_len"; + let err = load_timed_trace_dataset(&test_tokenizer(), &path, 1, "t-", &o) + .expect_err("should fail on missing field"); + let msg = err.to_string(); + assert!(msg.contains("input_len") && msg.contains("--timed-trace-label-input-length")); + } + + /// Hashes above `i64::MAX` are ordinary 64-bit block hashes, and Python + /// takes them as unbounded ints. + #[test] + fn test_accepts_unsigned_64bit_hash_ids() { + let path = write_temp_jsonl( + "timed-trace-u64-hash", + r#"{"timestamp": 0, "input_length": 16, "output_length": 1, "hash_ids": [18446744073709551615]}"#, + ); + let reqs = load_timed_trace_dataset(&test_tokenizer(), &path, 1, "t-", &opts(16, 1.0)) + .expect("load should succeed"); + assert_eq!(reqs[0].prompt_len, 16); + } + + /// Exporters that round-trip lengths through floats emit `512.0`; Python's + /// `int()` accepts it. + #[test] + fn test_accepts_float_encoded_lengths() { + let path = write_temp_jsonl( + "timed-trace-float-len", + r#"{"timestamp": 0, "input_length": 16.0, "output_length": 4.0, "hash_ids": [1]}"#, + ); + let reqs = load_timed_trace_dataset(&test_tokenizer(), &path, 1, "t-", &opts(16, 1.0)) + .expect("load should succeed"); + assert_eq!(reqs[0].prompt_len, 16); + assert_eq!(reqs[0].expected_output_len, 4); + } + + #[test] + fn test_non_finite_scaled_timestamp_is_rejected() { + let path = write_temp_jsonl( + "timed-trace-overflow-ts", + r#"{"timestamp": 1e308, "input_length": 16, "output_length": 1, "hash_ids": [1]}"#, + ); + let err = load_timed_trace_dataset(&test_tokenizer(), &path, 1, "t-", &opts(16, 1e10)) + .expect_err("should reject an offset that overflows to inf"); + assert!(err.to_string().contains("non-finite offset")); + } + + /// Negative offsets load unchanged; the scheduler clamps them to fire + /// immediately (Python's `get_request` does the same). + #[test] + fn test_negative_timestamp_loads() { + let path = write_temp_jsonl( + "timed-trace-negative-ts", + r#"{"timestamp": -5.0, "input_length": 16, "output_length": 1, "hash_ids": [1]}"#, + ); + let reqs = load_timed_trace_dataset(&test_tokenizer(), &path, 1, "t-", &opts(16, 1.0)) + .expect("load should succeed"); + assert_eq!(reqs[0].timestamp, Some(-5.0)); + } +} diff --git a/rust/src/bench/src/multi_turn.rs b/rust/src/bench/src/multi_turn.rs index b2b39b14cf..085684d5de 100644 --- a/rust/src/bench/src/multi_turn.rs +++ b/rust/src/bench/src/multi_turn.rs @@ -175,7 +175,10 @@ pub async fn run_multi_turn_benchmark(config: &BenchConfig) -> Result { + DatasetName::Custom + | DatasetName::PrefixRepetition + | DatasetName::RandomRerank + | DatasetName::TimedTrace => { return Err(BenchError::Config( "This dataset does not support multi-turn. Use 'random' or 'sharegpt' with --multi-turn.".into(), )); diff --git a/rust/src/bench/src/rate_control.rs b/rust/src/bench/src/rate_control.rs index aa05f8f8a8..1311cb9354 100644 --- a/rust/src/bench/src/rate_control.rs +++ b/rust/src/bench/src/rate_control.rs @@ -104,10 +104,30 @@ pub fn compute_schedule( } } +/// Schedule built from trace-recorded arrival offsets (`--self-timed`). +/// +/// Offsets are clamped at zero so a row anchored before the trace start fires +/// immediately, as Python's `get_request` does; `Duration::from_secs_f64` in +/// the dispatch loop would otherwise panic. `rates` carries no meaning here. +pub fn trace_schedule(timestamps: impl ExactSizeIterator>) -> RequestSchedule { + let rates = vec![0.0; timestamps.len()]; + RequestSchedule { + delays: timestamps.map(|t| t.unwrap_or(0.0).max(0.0)).collect(), + rates, + } +} + #[cfg(test)] mod tests { use super::*; + #[test] + fn test_trace_schedule_uses_offsets_and_clamps_negatives() { + let sched = trace_schedule([Some(-3.0), Some(0.0), Some(2.5), None].into_iter()); + assert_eq!(sched.delays, vec![0.0, 0.0, 2.5, 0.0]); + assert_eq!(sched.rates.len(), 4); + } + #[test] fn test_infinite_rate_all_zero_delays() { let sched = compute_schedule(100, f64::INFINITY, 1.0, 42, None); diff --git a/rust/src/bench/tests/cli_parity.rs b/rust/src/bench/tests/cli_parity.rs index b0fd09ab9e..b0bb707aa1 100644 --- a/rust/src/bench/tests/cli_parity.rs +++ b/rust/src/bench/tests/cli_parity.rs @@ -36,13 +36,6 @@ const PYTHON_ONLY: &[&str] = &[ // spec_bench dataset "--spec-bench-category", "--spec-bench-output-len", - // timed-trace dataset - "--timed-trace-chunk-hash-size", - "--timed-trace-label-hash-ids", - "--timed-trace-label-input-length", - "--timed-trace-label-output-length", - "--timed-trace-label-timestamp", - "--timed-trace-sec-multiplier", // client-side chat templating / request shaping "--chat-template-kwargs", "--custom-ensure-client-side-data", @@ -53,8 +46,6 @@ const PYTHON_ONLY: &[&str] = &[ "--timeline-itl-thresholds", // misc python-side controls "--hf-name", - "--no-self-timed", - "--self-timed", "--no-stream", "--probe-request-rate", ];