[vllm-bench, feature] Added support for Mooncake-style, timed-traces replay. (#55937)
This commit is contained in:
Rodrigo Garcia
@@ -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`
|
||||
|
||||
|
||||
@@ -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<serde_json::Value> {
|
||||
"{} 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<serde_json::Value> {
|
||||
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<serde_json::Value> {
|
||||
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<serde_json::Value> {
|
||||
);
|
||||
}
|
||||
|
||||
// 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<serde_json::Value> {
|
||||
};
|
||||
|
||||
// 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 {
|
||||
|
||||
@@ -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<usize>,
|
||||
|
||||
@@ -198,6 +198,15 @@ pub struct BenchConfig {
|
||||
pub hf_subset: Option<String>,
|
||||
pub hf_output_len: Option<usize>,
|
||||
pub hf_text_column: Option<String>,
|
||||
/// 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",
|
||||
|
||||
@@ -43,29 +43,10 @@ pub fn load_custom_dataset(
|
||||
no_oversample: bool,
|
||||
disable_shuffle: bool,
|
||||
) -> Result<Vec<SampleRequest>> {
|
||||
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<CustomLine> = 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<CustomLine> = 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()
|
||||
);
|
||||
|
||||
@@ -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<Arc<[Arc<str>]>>,
|
||||
/// 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<f64>,
|
||||
}
|
||||
|
||||
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<T: serde::de::DeserializeOwned>(
|
||||
path: &str,
|
||||
what: &str,
|
||||
max_entries: Option<usize>,
|
||||
) -> crate::error::Result<Vec<(usize, T)>> {
|
||||
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::*;
|
||||
|
||||
@@ -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();
|
||||
|
||||
@@ -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()
|
||||
}
|
||||
|
||||
@@ -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<Vec<SampleRequest>> {
|
||||
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<SampleRequest> = 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<u32> {
|
||||
// `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<u64> {
|
||||
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<f64> {
|
||||
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<Vec<u64>> {
|
||||
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));
|
||||
}
|
||||
}
|
||||
@@ -175,7 +175,10 @@ pub async fn run_multi_turn_benchmark(config: &BenchConfig) -> Result<serde_json
|
||||
"SPEED-Bench multi-turn is not yet supported. Use 'random' or 'sharegpt' with --multi-turn.".into(),
|
||||
));
|
||||
}
|
||||
DatasetName::Custom | DatasetName::PrefixRepetition | DatasetName::RandomRerank => {
|
||||
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(),
|
||||
));
|
||||
|
||||
@@ -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<Item = Option<f64>>) -> 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);
|
||||
|
||||
@@ -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",
|
||||
];
|
||||
|
||||
Reference in New Issue
Block a user