karylab_agents/vllm
1
0
Fork 0
forked from Karylab-cklius/vllm
Code Pull Requests 2 Actions 1 Packages Activity

[Rust][Benchmark] Integrate vllm-bench to vllm-rs & vllm CLI (#48930)

Browse Source 
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Signed-off-by: Bugen Zhao <i@bugenzhao.com>
This commit is contained in:
Bugen Zhao
2026-07-20 09:31:25 +00:00
committed by GitHub
co-authored by mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
parent 823eaf667d
commit c01618fdc8
12 changed files with 421 additions and 280 deletions
Download Patch File Download Diff File Expand all files Collapse all files
rust/Cargo.lock
Generated
+1
View File
@@ -5560,6 +5560,7 @@ dependencies = [
"tracing",
"tracing-subscriber",
"uuid",
"vllm-bench",
"vllm-chat",
"vllm-engine-core-client",
"vllm-managed-engine",
rust/Cargo.toml
+1
View File
@@ -132,6 +132,7 @@ trait-set = "0.3.0"
url = "2.5.7"
uuid = { version = "1.22.0", features = ["v4"] }
validator = { version = "0.20.0", features = ["derive"] }
vllm-bench = { path = "src/bench" }
vllm-chat = { path = "src/chat" }
vllm-engine-core-client = { path = "src/engine-core-client" }
vllm-llm = { path = "src/llm" }
rust/src/bench/src/cli.rs
+4 -11
View File
@@ -3,8 +3,6 @@
use std::fmt;
use clap::Parser;
/// Backend type for the benchmark endpoint.
#[derive(clap::ValueEnum, Debug, Clone, Copy, PartialEq, Eq)]
pub enum BackendKind {
@@ -77,7 +75,7 @@ pub enum DatasetName {
ShareGpt,
#[value(name = "sonnet")]
Sonnet,
#[value(name = "speed-bench")]
#[value(name = "speed-bench", alias = "speed_bench")]
SpeedBench,
#[value(name = "hf")]
Hf,
@@ -144,13 +142,8 @@ impl fmt::Display for SpeedBenchConfig {
}
/// High-performance benchmark client for vLLM serving endpoints.
#[derive(Parser, Debug, Clone)]
#[command(
name = "vllm-bench",
about = "Benchmark online serving throughput",
version
)]
pub struct Cli {
#[derive(clap::Args, Debug, Clone)]
pub struct BenchServeArgs {
/// The type of backend or endpoint to use for the benchmark.
#[arg(long, default_value = "openai")]
pub backend: BackendKind,
@@ -659,7 +652,7 @@ pub struct Cli {
pub lora_assignment: LoraAssignment,
}
impl Cli {
impl BenchServeArgs {
/// Resolve the base URL from explicit --base-url or from --host/--port.
pub fn resolve_base_url(&self) -> String {
if let Some(ref base) = self.base_url {
rust/src/bench/src/config.rs
+212 -188
View File
@@ -4,7 +4,9 @@
use std::collections::HashMap;
use std::sync::Arc;
use crate::cli::{BackendKind, Cli, DatasetName, LoraAssignment, RampUpStrategy, SpeedBenchConfig};
use crate::cli::{
BackendKind, BenchServeArgs, DatasetName, LoraAssignment, RampUpStrategy, SpeedBenchConfig,
};
use crate::datasets::random_mm::{MmBucketKey, MmLimitPerPrompt};
use crate::error::{BenchError, Result};
@@ -215,63 +217,63 @@ pub struct BenchConfig {
}
impl BenchConfig {
pub fn from_cli(cli: &Cli) -> Result<Self> {
if cli.burstiness <= 0.0 {
pub fn from_args(args: &BenchServeArgs) -> Result<Self> {
if args.burstiness <= 0.0 {
return Err(BenchError::Config("Burstiness must be positive".into()));
}
if cli.num_prompts == 0 {
if args.num_prompts == 0 {
return Err(BenchError::Config(
"--num-prompts must be at least 1".into(),
));
}
if cli.request_rate <= 0.0 && !cli.request_rate.is_infinite() {
if args.request_rate <= 0.0 && !args.request_rate.is_infinite() {
return Err(BenchError::Config(
"--request-rate must be positive (or inf)".into(),
));
}
if cli.max_model_len == Some(0) {
if args.max_model_len == Some(0) {
return Err(BenchError::Config(
"--max-model-len must be at least 1".into(),
));
}
let base_url = cli.resolve_base_url();
let api_url = cli.resolve_api_url();
let base_url = args.resolve_base_url();
let api_url = args.resolve_api_url();
let extra_headers = cli.parse_headers()?;
let mut extra_body = cli.parse_extra_body()?;
let extra_headers = args.parse_headers()?;
let mut extra_body = args.parse_extra_body()?;
// Merge sampling parameters into extra_body (matches Python behavior).
// Python collects non-None sampling params and merges them UNDER extra_body,
// meaning extra_body keys take precedence over sampling params.
{
let mut sampling_params = serde_json::Map::new();
if let Some(v) = cli.top_p {
if let Some(v) = args.top_p {
sampling_params.insert("top_p".into(), serde_json::json!(v));
}
if let Some(v) = cli.top_k {
if let Some(v) = args.top_k {
sampling_params.insert("top_k".into(), serde_json::json!(v));
}
if let Some(v) = cli.min_p {
if let Some(v) = args.min_p {
sampling_params.insert("min_p".into(), serde_json::json!(v));
}
if let Some(v) = cli.temperature {
if let Some(v) = args.temperature {
sampling_params.insert("temperature".into(), serde_json::json!(v));
}
if let Some(v) = cli.frequency_penalty {
if let Some(v) = args.frequency_penalty {
sampling_params.insert("frequency_penalty".into(), serde_json::json!(v));
}
if let Some(v) = cli.presence_penalty {
if let Some(v) = args.presence_penalty {
sampling_params.insert("presence_penalty".into(), serde_json::json!(v));
}
if let Some(v) = cli.repetition_penalty {
if let Some(v) = args.repetition_penalty {
sampling_params.insert("repetition_penalty".into(), serde_json::json!(v));
}
if !sampling_params.is_empty() {
if !cli.backend.is_openai_compatible() {
if !args.backend.is_openai_compatible() {
return Err(BenchError::Config(
"Sampling parameters are only supported by openai-compatible backends."
.into(),
@@ -299,7 +301,7 @@ impl BenchConfig {
}
// Parse metadata
let metadata = match &cli.metadata {
let metadata = match &args.metadata {
None => None,
Some(items) => {
let mut pairs = Vec::new();
@@ -314,24 +316,24 @@ impl BenchConfig {
};
// Parse goodput SLOs
let goodput = parse_goodput(&cli.goodput)?;
let goodput = parse_goodput(&args.goodput)?;
// Parse ramp-up config
let ramp_up = parse_ramp_up(cli)?;
let ramp_up = parse_ramp_up(args)?;
// Default percentile metrics based on backend type
let default_percentile_metrics = if cli.backend.is_pooling() {
let default_percentile_metrics = if args.backend.is_pooling() {
"e2el"
} else {
"ttft,tpot,itl,e2el"
};
let percentile_metrics_str =
cli.percentile_metrics.as_deref().unwrap_or(default_percentile_metrics);
args.percentile_metrics.as_deref().unwrap_or(default_percentile_metrics);
let selected_percentile_metrics: Vec<String> =
percentile_metrics_str.split(',').map(|s| s.trim().to_string()).collect();
let metric_percentiles = parse_percentiles(&cli.metric_percentiles, false)?;
let sweep_summary_percentiles = cli
let metric_percentiles = parse_percentiles(&args.metric_percentiles, false)?;
let sweep_summary_percentiles = args
.sweep_summary_percentiles
.as_deref()
.map(|raw| parse_percentiles(raw, true))
@@ -344,38 +346,38 @@ impl BenchConfig {
selected_percentiles.push(90.0);
}
let tokenizer_id = if cli.skip_tokenizer_init {
let tokenizer_id = if args.skip_tokenizer_init {
None
} else {
Some(cli.tokenizer.clone().or_else(|| cli.model.clone()).unwrap_or_default())
args.tokenizer.clone().or_else(|| args.model.clone())
};
// Resolve input/output lengths
let random_input_len = cli.resolved_random_input_len();
let random_output_len = cli.resolved_random_output_len();
let per_turn_input_len = cli.resolved_per_turn_input_len();
let random_input_len = args.resolved_random_input_len();
let random_output_len = args.resolved_random_output_len();
let per_turn_input_len = args.resolved_per_turn_input_len();
// Normalized multi-turn turn counts (computed in validation block below, defaults
// to num_turns if multi-turn mode is not active)
let mut multi_turn_min_turns = cli.multi_turn_num_turns;
let mut multi_turn_max_turns = cli.multi_turn_num_turns;
let mut multi_turn_min_turns = args.multi_turn_num_turns;
let mut multi_turn_max_turns = args.multi_turn_num_turns;
// 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.
// Pooling backends don't generate tokens, so ignore_eos is irrelevant.
let ignore_eos = if cli.backend.is_pooling() {
let ignore_eos = if args.backend.is_pooling() {
false
} else {
cli.ignore_eos
|| ((cli.dataset_name == DatasetName::Random
|| cli.dataset_name == DatasetName::RandomMm)
&& cli.backend.is_openai_compatible()
&& !cli.multi_turn)
args.ignore_eos
|| ((args.dataset_name == DatasetName::Random
|| args.dataset_name == DatasetName::RandomMm)
&& args.backend.is_openai_compatible()
&& !args.multi_turn)
};
// Pooling backends don't support multi-turn
if cli.backend.is_pooling() && cli.multi_turn {
if args.backend.is_pooling() && args.multi_turn {
return Err(BenchError::Config(
"Pooling/embedding backends do not support --multi-turn".into(),
));
@@ -383,7 +385,7 @@ impl BenchConfig {
// LoRA validation. Adapter names must be non-empty after trim; pooling
// backends are out of scope (vLLM LoRA routing is for generative paths).
let lora_modules = match cli.lora_modules.as_ref() {
let lora_modules = match args.lora_modules.as_ref() {
None => None,
Some(names) => {
if names.is_empty() {
@@ -391,7 +393,7 @@ impl BenchConfig {
"--lora-modules requires at least one adapter name".into(),
));
}
if cli.backend.is_pooling() {
if args.backend.is_pooling() {
return Err(BenchError::Config(
"--lora-modules is not supported for pooling/embedding backends".into(),
));
@@ -411,18 +413,18 @@ impl BenchConfig {
};
// Random-MM validation and config parsing
let (random_mm_limit, random_mm_buckets) = if cli.dataset_name == DatasetName::RandomMm {
if cli.backend != BackendKind::OpenaiChat {
let (random_mm_limit, random_mm_buckets) = if args.dataset_name == DatasetName::RandomMm {
if args.backend != BackendKind::OpenaiChat {
return Err(BenchError::Config(
"Multi-modal content (images) is only supported on 'openai-chat' backend."
.into(),
));
}
let limit = crate::datasets::random_mm::parse_limit_mm_per_prompt(
&cli.random_mm_limit_mm_per_prompt,
&args.random_mm_limit_mm_per_prompt,
)?;
let buckets =
crate::datasets::random_mm::parse_bucket_config(&cli.random_mm_bucket_config)?;
crate::datasets::random_mm::parse_bucket_config(&args.random_mm_bucket_config)?;
(limit, buckets)
} else {
(MmLimitPerPrompt::default(), Vec::new())
@@ -432,18 +434,18 @@ impl BenchConfig {
// sonnet (uses built-in Shakespeare's sonnets).
// Range ratio (Python semantics: [len*(1-r), len*(1+r)], each r in [0,1))
let random_range_ratio = RangeRatio::parse(&cli.random_range_ratio)?;
let random_range_ratio = RangeRatio::parse(&args.random_range_ratio)?;
// Batched inputs only make sense for pooling backends (the generation
// backends send one prompt per request).
if cli.random_batch_size == 0 {
if args.random_batch_size == 0 {
return Err(BenchError::Config(
"--random-batch-size must be at least 1".into(),
));
}
if cli.random_batch_size > 1
&& !cli.backend.is_pooling()
&& cli.dataset_name != DatasetName::RandomRerank
if args.random_batch_size > 1
&& !args.backend.is_pooling()
&& args.dataset_name != DatasetName::RandomRerank
{
return Err(BenchError::Config(
"--random-batch-size > 1 is only supported with embeddings/pooling backends".into(),
@@ -451,16 +453,16 @@ impl BenchConfig {
}
// random-rerank validation (mirrors Python RandomDatasetForReranking)
let is_reranker = !cli.no_reranker;
if cli.dataset_name == DatasetName::RandomRerank {
if !cli.backend.is_pooling() {
let is_reranker = !args.no_reranker;
if args.dataset_name == DatasetName::RandomRerank {
if !args.backend.is_pooling() {
return Err(BenchError::Config(
"--dataset-name random-rerank requires an embeddings/pooling backend \
(e.g. --backend vllm-rerank)"
.into(),
));
}
if !is_reranker && (cli.num_prompts < 2 || cli.random_batch_size < 2) {
if !is_reranker && (args.num_prompts < 2 || args.random_batch_size < 2) {
return Err(BenchError::Config(
"--no-reranker requires --num-prompts > 1 and --random-batch-size > 1 \
(the query is folded into the first batch slot)"
@@ -470,8 +472,8 @@ impl BenchConfig {
}
// Custom dataset validation
if cli.dataset_name == DatasetName::Custom {
match cli.dataset_path.as_deref() {
if args.dataset_name == DatasetName::Custom {
match args.dataset_path.as_deref() {
None => {
return Err(BenchError::Config(
"--dataset-path is required for --dataset-name custom \
@@ -486,7 +488,7 @@ impl BenchConfig {
}
_ => {}
}
if !cli.skip_chat_template {
if !args.skip_chat_template {
eprintln!(
"NOTE: client-side chat template rendering is not supported; custom \
dataset prompts are sent raw (equivalent to --skip-chat-template)."
@@ -495,29 +497,29 @@ impl BenchConfig {
}
// Prefix repetition validation
if cli.dataset_name == DatasetName::PrefixRepetition {
if cli.prefix_repetition_num_prefixes == 0 {
if args.dataset_name == DatasetName::PrefixRepetition {
if args.prefix_repetition_num_prefixes == 0 {
return Err(BenchError::Config(
"--prefix-repetition-num-prefixes must be at least 1".into(),
));
}
if cli.num_prompts < cli.prefix_repetition_num_prefixes {
if args.num_prompts < args.prefix_repetition_num_prefixes {
return Err(BenchError::Config(format!(
"--num-prompts ({}) must be >= --prefix-repetition-num-prefixes ({})",
cli.num_prompts, cli.prefix_repetition_num_prefixes
args.num_prompts, args.prefix_repetition_num_prefixes
)));
}
}
// HF dataset validation
if cli.dataset_name == DatasetName::Hf && cli.dataset_path.is_none() {
if args.dataset_name == DatasetName::Hf && args.dataset_path.is_none() {
return Err(BenchError::Config(
"--dataset-path is required for --dataset-name hf \
(set to a HuggingFace dataset ID, e.g. 'allenai/WildChat-4.8M')"
.into(),
));
}
if let Some(len) = cli.hf_output_len
if let Some(len) = args.hf_output_len
&& len == 0
{
return Err(BenchError::Config(
@@ -526,13 +528,13 @@ impl BenchConfig {
}
// Multi-turn validation
if cli.multi_turn {
if cli.backend != BackendKind::OpenaiChat {
if args.multi_turn {
if args.backend != BackendKind::OpenaiChat {
return Err(BenchError::Config(
"--multi-turn requires --backend openai-chat".into(),
));
}
if cli.multi_turn_num_turns == 0 {
if args.multi_turn_num_turns == 0 {
return Err(BenchError::Config(
"--multi-turn-num-turns must be at least 1".into(),
));
@@ -541,18 +543,18 @@ impl BenchConfig {
// Normalize and validate min/max turns. ShareGPT only consumes max_turns
// (the loader walks all available turns up to the cap), so the
// min/num/max coupling used for synthetic generation does not apply.
if cli.dataset_name == DatasetName::ShareGpt {
if cli.multi_turn_max_turns == 1 {
if args.dataset_name == DatasetName::ShareGpt {
if args.multi_turn_max_turns == 1 {
return Err(BenchError::Config(
"--multi-turn-max-turns must be at least 2 for ShareGPT multi-turn".into(),
));
}
} else {
(multi_turn_min_turns, multi_turn_max_turns) =
match (cli.multi_turn_min_turns, cli.multi_turn_max_turns) {
(0, 0) => (cli.multi_turn_num_turns, cli.multi_turn_num_turns),
(m, 0) => (m, cli.multi_turn_num_turns),
(0, x) => (cli.multi_turn_num_turns, x),
match (args.multi_turn_min_turns, args.multi_turn_max_turns) {
(0, 0) => (args.multi_turn_num_turns, args.multi_turn_num_turns),
(m, 0) => (m, args.multi_turn_num_turns),
(0, x) => (args.multi_turn_num_turns, x),
(m, x) => (m, x),
};
if multi_turn_min_turns < 1 {
@@ -575,8 +577,8 @@ impl BenchConfig {
}
// Validate prefix sharing ratios
let pg = cli.multi_turn_prefix_global_ratio;
let pc = cli.multi_turn_prefix_conversation_ratio;
let pg = args.multi_turn_prefix_global_ratio;
let pc = args.multi_turn_prefix_conversation_ratio;
if !(0.0..=1.0).contains(&pg) {
return Err(BenchError::Config(
"--multi-turn-prefix-global-ratio must be in [0.0, 1.0]".into(),
@@ -592,20 +594,20 @@ impl BenchConfig {
"--multi-turn-prefix-global-ratio + --multi-turn-prefix-conversation-ratio must be < 1.0 (unique suffix required)".into(),
));
}
if (pg > 0.0 || pc > 0.0) && cli.dataset_name != DatasetName::Random {
if (pg > 0.0 || pc > 0.0) && args.dataset_name != DatasetName::Random {
return Err(BenchError::Config(
"Prefix sharing (--multi-turn-prefix-global-ratio / --multi-turn-prefix-conversation-ratio) only works with --dataset-name random".into(),
));
}
}
if !(cli.steady_state_threshold > 0.0 && cli.steady_state_threshold <= 1.0) {
if !(args.steady_state_threshold > 0.0 && args.steady_state_threshold <= 1.0) {
return Err(BenchError::Config(format!(
"--steady-state-threshold must be in (0.0, 1.0], got {}",
cli.steady_state_threshold
args.steady_state_threshold
)));
}
if let Some(mw) = cli.steady_state_min_window
if let Some(mw) = args.steady_state_min_window
&& mw < 0.0
{
return Err(BenchError::Config(format!(
@@ -613,122 +615,122 @@ impl BenchConfig {
)));
}
if cli.profile_batch_threshold.is_some() && !cli.profile {
if args.profile_batch_threshold.is_some() && !args.profile {
return Err(BenchError::Config(
"--profile-batch-threshold requires --profile".into(),
));
}
if cli.profile_duration <= 0.0 {
if args.profile_duration <= 0.0 {
return Err(BenchError::Config(
"--profile-duration must be positive".into(),
));
}
if cli.profile_batch_threshold.is_none() && cli.profile_duration != 5.0 {
if args.profile_batch_threshold.is_none() && args.profile_duration != 5.0 {
return Err(BenchError::Config(
"--profile-duration requires --profile-batch-threshold".into(),
));
}
Ok(BenchConfig {
backend: cli.backend,
backend: args.backend,
base_url,
api_url,
model: cli.model.clone(),
model_name: cli.served_model_name.clone(),
model: args.model.clone(),
model_name: args.served_model_name.clone(),
tokenizer_id,
tokenizer_mode: cli.tokenizer_mode.clone(),
trust_remote_code: cli.trust_remote_code,
skip_tokenizer_init: cli.skip_tokenizer_init,
dataset_name: cli.dataset_name,
dataset_path: cli.dataset_path.clone(),
max_model_len: cli.max_model_len,
tokenizer_mode: args.tokenizer_mode.clone(),
trust_remote_code: args.trust_remote_code,
skip_tokenizer_init: args.skip_tokenizer_init,
dataset_name: args.dataset_name,
dataset_path: args.dataset_path.clone(),
max_model_len: args.max_model_len,
random_input_len,
random_output_len,
random_prefix_len: cli.random_prefix_len,
random_prefix_len: args.random_prefix_len,
random_range_ratio,
random_batch_size: cli.random_batch_size,
random_batch_size: args.random_batch_size,
is_reranker,
custom_output_len: cli.output_len.map(|v| v as i64).unwrap_or(cli.custom_output_len),
prefix_repetition_prefix_len: cli.prefix_repetition_prefix_len,
prefix_repetition_suffix_len: cli.prefix_repetition_suffix_len,
prefix_repetition_num_prefixes: cli.prefix_repetition_num_prefixes,
prefix_repetition_output_len: cli
custom_output_len: args.output_len.map(|v| v as i64).unwrap_or(args.custom_output_len),
prefix_repetition_prefix_len: args.prefix_repetition_prefix_len,
prefix_repetition_suffix_len: args.prefix_repetition_suffix_len,
prefix_repetition_num_prefixes: args.prefix_repetition_num_prefixes,
prefix_repetition_output_len: args
.output_len
.unwrap_or(cli.prefix_repetition_output_len),
random_cache_hit_fraction: cli.random_cache_hit_fraction,
random_cache_ratio: cli.random_cache_ratio,
sharegpt_output_len: cli.sharegpt_output_len,
sonnet_input_len: cli.sonnet_input_len,
sonnet_output_len: cli.sonnet_output_len,
sonnet_prefix_len: cli.sonnet_prefix_len,
no_oversample: cli.no_oversample,
disable_shuffle: cli.disable_shuffle,
num_prompts: cli.num_prompts,
request_rate: cli.request_rate,
burstiness: cli.burstiness,
max_concurrency: cli.max_concurrency,
steady_state_threshold: cli.steady_state_threshold,
steady_state_min_window: cli.steady_state_min_window,
no_steady_state: cli.no_steady_state,
disable_tqdm: cli.disable_tqdm,
num_warmups: cli.num_warmups,
profile: cli.profile,
profile_batch_threshold: cli.profile_batch_threshold,
profile_duration: cli.profile_duration,
save_result: cli.save_result,
save_detailed: cli.save_detailed,
append_result: cli.append_result,
result_dir: cli.result_dir.clone(),
result_filename: cli.result_filename.clone(),
seed: cli.seed,
.unwrap_or(args.prefix_repetition_output_len),
random_cache_hit_fraction: args.random_cache_hit_fraction,
random_cache_ratio: args.random_cache_ratio,
sharegpt_output_len: args.sharegpt_output_len,
sonnet_input_len: args.sonnet_input_len,
sonnet_output_len: args.sonnet_output_len,
sonnet_prefix_len: args.sonnet_prefix_len,
no_oversample: args.no_oversample,
disable_shuffle: args.disable_shuffle,
num_prompts: args.num_prompts,
request_rate: args.request_rate,
burstiness: args.burstiness,
max_concurrency: args.max_concurrency,
steady_state_threshold: args.steady_state_threshold,
steady_state_min_window: args.steady_state_min_window,
no_steady_state: args.no_steady_state,
disable_tqdm: args.disable_tqdm,
num_warmups: args.num_warmups,
profile: args.profile,
profile_batch_threshold: args.profile_batch_threshold,
profile_duration: args.profile_duration,
save_result: args.save_result,
save_detailed: args.save_detailed,
append_result: args.append_result,
result_dir: args.result_dir.clone(),
result_filename: args.result_filename.clone(),
seed: args.seed,
ignore_eos,
insecure: cli.insecure,
insecure: args.insecure,
selected_percentile_metrics,
selected_percentiles,
sweep_summary_percentiles,
label: cli.label.clone(),
logprobs: cli.logprobs,
request_id_prefix: cli.get_request_id_prefix(),
ready_check_timeout_sec: cli.ready_check_timeout_sec,
label: args.label.clone(),
logprobs: args.logprobs,
request_id_prefix: args.get_request_id_prefix(),
ready_check_timeout_sec: args.ready_check_timeout_sec,
extra_headers,
extra_body,
metadata,
dry_run: cli.dry_run,
dry_run: args.dry_run,
goodput,
ramp_up,
multi_turn: cli.multi_turn,
multi_turn_num_turns: cli.multi_turn_num_turns,
multi_turn: args.multi_turn,
multi_turn_num_turns: args.multi_turn_num_turns,
multi_turn_min_turns,
multi_turn_max_turns,
sharegpt_multi_turn_max_turns: if cli.multi_turn
&& cli.dataset_name == DatasetName::ShareGpt
&& cli.multi_turn_max_turns != 0
sharegpt_multi_turn_max_turns: if args.multi_turn
&& args.dataset_name == DatasetName::ShareGpt
&& args.multi_turn_max_turns != 0
{
Some(cli.multi_turn_max_turns)
Some(args.multi_turn_max_turns)
} else {
None
},
per_turn_input_len,
multi_turn_concurrency: cli.multi_turn_concurrency,
multi_turn_delay_ms: cli.multi_turn_delay_ms,
multi_turn_prefix_global_ratio: cli.multi_turn_prefix_global_ratio,
multi_turn_prefix_conversation_ratio: cli.multi_turn_prefix_conversation_ratio,
speed_bench_config: cli.speed_bench_config,
speed_bench_category: cli.speed_bench_category.clone(),
speed_bench_max_input_len: cli.speed_bench_max_input_len,
hf_split: cli.hf_split.clone(),
hf_subset: cli.hf_subset.clone(),
hf_output_len: cli.hf_output_len,
hf_text_column: cli.hf_text_column.clone(),
reset_prefix_cache: cli.reset_prefix_cache,
prompt_token_ids: cli.prompt_token_ids,
random_mm_base_items_per_request: cli.random_mm_base_items_per_request,
random_mm_num_mm_items_range_ratio: cli.random_mm_num_mm_items_range_ratio,
multi_turn_concurrency: args.multi_turn_concurrency,
multi_turn_delay_ms: args.multi_turn_delay_ms,
multi_turn_prefix_global_ratio: args.multi_turn_prefix_global_ratio,
multi_turn_prefix_conversation_ratio: args.multi_turn_prefix_conversation_ratio,
speed_bench_config: args.speed_bench_config,
speed_bench_category: args.speed_bench_category.clone(),
speed_bench_max_input_len: args.speed_bench_max_input_len,
hf_split: args.hf_split.clone(),
hf_subset: args.hf_subset.clone(),
hf_output_len: args.hf_output_len,
hf_text_column: args.hf_text_column.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,
random_mm_num_mm_items_range_ratio: args.random_mm_num_mm_items_range_ratio,
random_mm_limit,
random_mm_buckets,
enable_multimodal_chat: cli.enable_multimodal_chat,
enable_multimodal_chat: args.enable_multimodal_chat,
lora_modules,
lora_assignment: cli.lora_assignment,
lora_assignment: args.lora_assignment,
})
}
}
@@ -811,17 +813,17 @@ fn parse_goodput(goodput_args: &Option<Vec<String>>) -> Result<GoodputConfig> {
Ok(config)
}
fn parse_ramp_up(cli: &Cli) -> Result<Option<RampUpConfig>> {
let strategy = match cli.ramp_up_strategy {
fn parse_ramp_up(args: &BenchServeArgs) -> Result<Option<RampUpConfig>> {
let strategy = match args.ramp_up_strategy {
None => return Ok(None),
Some(s) => s,
};
let start_rps = cli.ramp_up_start_rps.ok_or_else(|| {
let start_rps = args.ramp_up_start_rps.ok_or_else(|| {
BenchError::Config("--ramp-up-start-rps is required when --ramp-up-strategy is set".into())
})?;
let end_rps = cli.ramp_up_end_rps.ok_or_else(|| {
let end_rps = args.ramp_up_end_rps.ok_or_else(|| {
BenchError::Config("--ramp-up-end-rps is required when --ramp-up-strategy is set".into())
})?;
@@ -843,7 +845,21 @@ mod tests {
use clap::Parser;
use super::*;
use crate::cli::Cli;
use crate::cli::BenchServeArgs;
#[derive(Parser)]
struct TestCli {
#[command(flatten)]
args: BenchServeArgs,
}
fn parse_args<I, T>(args: I) -> BenchServeArgs
where
I: IntoIterator<Item = T>,
T: Into<std::ffi::OsString> + Clone,
{
TestCli::parse_from(args).args
}
fn base_multi_turn_args() -> Vec<&'static str> {
vec![
@@ -859,8 +875,8 @@ mod tests {
#[test]
fn test_prefix_sharing_defaults_to_zero() {
let args = base_multi_turn_args();
let cli = Cli::parse_from(args);
let config = BenchConfig::from_cli(&cli).unwrap();
let args = parse_args(args);
let config = BenchConfig::from_args(&args).unwrap();
assert_eq!(config.multi_turn_prefix_global_ratio, 0.0);
assert_eq!(config.multi_turn_prefix_conversation_ratio, 0.0);
}
@@ -874,8 +890,8 @@ mod tests {
"--multi-turn-prefix-conversation-ratio",
"0.8",
]);
let cli = Cli::parse_from(args);
let config = BenchConfig::from_cli(&cli).unwrap();
let args = parse_args(args);
let config = BenchConfig::from_args(&args).unwrap();
assert!((config.multi_turn_prefix_global_ratio - 0.1).abs() < 1e-10);
assert!((config.multi_turn_prefix_conversation_ratio - 0.8).abs() < 1e-10);
}
@@ -889,8 +905,8 @@ mod tests {
"--multi-turn-prefix-conversation-ratio",
"0.6",
]);
let cli = Cli::parse_from(args);
assert!(BenchConfig::from_cli(&cli).is_err());
let args = parse_args(args);
assert!(BenchConfig::from_args(&args).is_err());
}
#[test]
@@ -902,16 +918,16 @@ mod tests {
"--multi-turn-prefix-conversation-ratio",
"0.5",
]);
let cli = Cli::parse_from(args);
assert!(BenchConfig::from_cli(&cli).is_err());
let args = parse_args(args);
assert!(BenchConfig::from_args(&args).is_err());
}
#[test]
fn test_prefix_sharing_out_of_range_fails() {
let mut args = base_multi_turn_args();
args.extend(["--multi-turn-prefix-global-ratio", "1.5"]);
let cli = Cli::parse_from(args);
assert!(BenchConfig::from_cli(&cli).is_err());
let args = parse_args(args);
assert!(BenchConfig::from_args(&args).is_err());
}
#[test]
@@ -928,8 +944,8 @@ mod tests {
"--multi-turn-prefix-global-ratio",
"0.1",
];
let cli = Cli::parse_from(args);
assert!(BenchConfig::from_cli(&cli).is_err());
let args = parse_args(args);
assert!(BenchConfig::from_args(&args).is_err());
}
#[test]
@@ -944,8 +960,8 @@ mod tests {
"--dataset-name",
"sharegpt",
];
let cli = Cli::parse_from(args);
let config = BenchConfig::from_cli(&cli).unwrap();
let args = parse_args(args);
let config = BenchConfig::from_args(&args).unwrap();
assert_eq!(config.multi_turn_max_turns, 3);
assert_eq!(config.sharegpt_multi_turn_max_turns, None);
@@ -968,8 +984,8 @@ mod tests {
"--multi-turn-max-turns",
"2",
];
let cli = Cli::parse_from(args);
let config = BenchConfig::from_cli(&cli).unwrap();
let args = parse_args(args);
let config = BenchConfig::from_args(&args).unwrap();
assert_eq!(config.sharegpt_multi_turn_max_turns, Some(2));
}
@@ -987,8 +1003,8 @@ mod tests {
"--multi-turn-max-turns",
"1",
];
let cli = Cli::parse_from(args);
let err = BenchConfig::from_cli(&cli).unwrap_err().to_string();
let args = parse_args(args);
let err = BenchConfig::from_args(&args).unwrap_err().to_string();
assert!(
err.contains("at least 2 for ShareGPT"),
"expected ShareGPT-specific error, got: {err}"
@@ -1009,8 +1025,8 @@ mod tests {
"--multi-turn-max-turns",
"20",
];
let cli = Cli::parse_from(args);
let config = BenchConfig::from_cli(&cli).unwrap();
let args = parse_args(args);
let config = BenchConfig::from_args(&args).unwrap();
assert_eq!(config.sharegpt_multi_turn_max_turns, Some(20));
}
@@ -1018,8 +1034,8 @@ mod tests {
#[test]
fn test_sweep_summary_percentiles_default_empty() {
let args = base_multi_turn_args();
let cli = Cli::parse_from(args);
let config = BenchConfig::from_cli(&cli).unwrap();
let args = parse_args(args);
let config = BenchConfig::from_args(&args).unwrap();
assert!(config.sweep_summary_percentiles.is_empty());
assert_eq!(config.selected_percentiles, vec![99.0, 90.0]);
@@ -1034,8 +1050,8 @@ mod tests {
"--sweep-summary-percentiles",
"90,95,90",
]);
let cli = Cli::parse_from(args);
let config = BenchConfig::from_cli(&cli).unwrap();
let args = parse_args(args);
let config = BenchConfig::from_args(&args).unwrap();
assert_eq!(config.sweep_summary_percentiles, vec![90.0, 95.0]);
assert_eq!(config.selected_percentiles, vec![99.0, 95.0, 90.0]);
@@ -1045,8 +1061,8 @@ mod tests {
fn test_invalid_sweep_summary_percentile_fails() {
let mut args = base_multi_turn_args();
args.extend(["--sweep-summary-percentiles", "101"]);
let cli = Cli::parse_from(args);
assert!(BenchConfig::from_cli(&cli).is_err());
let args = parse_args(args);
assert!(BenchConfig::from_args(&args).is_err());
}
#[test]
@@ -1058,12 +1074,20 @@ mod tests {
"--max-model-len",
"4096",
];
let cli = Cli::parse_from(args);
let config = BenchConfig::from_cli(&cli).unwrap();
let args = parse_args(args);
let config = BenchConfig::from_args(&args).unwrap();
assert_eq!(config.max_model_len, Some(4096));
}
#[test]
fn test_tokenizer_id_deferred_when_model_is_unspecified() {
let args = parse_args(["vllm-bench"]);
let config = BenchConfig::from_args(&args).unwrap();
assert_eq!(config.tokenizer_id, None);
}
#[test]
fn test_zero_max_model_len_fails() {
let args = vec![
@@ -1073,9 +1097,9 @@ mod tests {
"--max-model-len",
"0",
];
let cli = Cli::parse_from(args);
let args = parse_args(args);
assert!(BenchConfig::from_cli(&cli).is_err());
assert!(BenchConfig::from_args(&args).is_err());
}
#[test]
fn test_range_ratio_parse_float() {
rust/src/bench/src/hub.rs
+5 -2
View File
@@ -40,8 +40,11 @@ impl HubRepo {
.build()
.map_err(|e| format!("Failed to build download runtime: {e}"))?;
rt.block_on(async move {
let api = hf_hub::api::tokio::Api::new()
.map_err(|e| format!("Failed to init HF API: {e}"))?;
let mut builder = hf_hub::api::tokio::ApiBuilder::from_env();
if let Ok(token) = std::env::var("HF_TOKEN") {
builder = builder.with_token(Some(token));
}
let api = builder.build().map_err(|e| format!("Failed to init HF API: {e}"))?;
api.repo(repo).get(&filename).await.map_err(|e| format!("{e}"))
})
})
rust/src/bench/src/lib.rs
+86
View File
@@ -0,0 +1,86 @@
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
mod backends;
mod benchmark;
mod cli;
mod compare;
mod config;
mod datasets;
mod error;
mod hub;
mod metrics;
mod multi_run;
mod multi_turn;
mod output;
mod rate_control;
mod ready_checker;
mod sweep;
mod tiktoken;
mod tokenizer;
use anyhow::Context;
pub use cli::{
BackendKind, BenchServeArgs, DatasetName, LoraAssignment, RampUpStrategy, SpeedBenchConfig,
};
use config::BenchConfig;
/// Prepare process-wide resources for a benchmark run.
pub fn prepare_process() {
// Raise the open-file soft limit to the hard limit. High-concurrency
// benchmarks (1024+ requests) easily exceed the default 1024 fd soft limit.
if let Ok(new) = rlimit::increase_nofile_limit(u64::MAX)
&& new > 1024
{
eprintln!("Open-file limit: {new}");
}
}
/// Run the online serving benchmark.
pub async fn run(args: BenchServeArgs) -> anyhow::Result<()> {
// --- Compare mode: no server needed, just diff two JSON files ---
if let Some(ref files) = args.compare {
return compare::compare_results(&files[0], &files[1]).context("Comparison failed");
}
let config = BenchConfig::from_args(&args).context("Configuration error")?;
async {
if config.multi_turn {
if let Some(ref sweep_mc) = args.sweep_max_concurrency {
// --- Sweep over concurrency in multi-turn mode ---
let values = sweep::parse_concurrency_values(sweep_mc)
.context("Invalid --sweep-max-concurrency")?;
sweep::run_multi_turn_concurrency_sweep(
&config,
&values,
args.sweep_num_prompts_factor,
)
.await?;
} else {
// --- Single multi-turn conversation benchmark ---
multi_turn::run_multi_turn_benchmark(&config).await?;
}
} else if let Some(ref sweep_mc) = args.sweep_max_concurrency {
// --- Sweep over max-concurrency ---
let values = sweep::parse_concurrency_values(sweep_mc)
.context("Invalid --sweep-max-concurrency")?;
sweep::run_concurrency_sweep(&config, &values, args.sweep_num_prompts_factor).await?;
} else if let Some(ref sweep_rate) = args.sweep_request_rate {
// --- Sweep over request-rate ---
let values =
sweep::parse_rate_values(sweep_rate).context("Invalid --sweep-request-rate")?;
sweep::run_rate_sweep(&config, &values).await?;
} else if args.num_runs > 1 {
// --- Multi-run with statistical aggregation ---
multi_run::run_multi(&config, args.num_runs).await?;
} else {
// --- Normal single benchmark ---
benchmark::run_benchmark(&config).await?;
}
anyhow::Ok(())
}
.await
.context("Benchmark failed")
}
rust/src/bench/src/main.rs
+14 -74
View File
@@ -1,92 +1,32 @@
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
mod backends;
mod benchmark;
mod cli;
mod compare;
mod config;
mod datasets;
mod error;
mod hub;
mod metrics;
mod multi_run;
mod multi_turn;
mod output;
mod rate_control;
mod ready_checker;
mod sweep;
mod tiktoken;
mod tokenizer;
#[cfg(not(target_env = "msvc"))]
#[global_allocator]
static GLOBAL: mimalloc::MiMalloc = mimalloc::MiMalloc;
use anyhow::Context;
use clap::Parser;
use cli::Cli;
use config::BenchConfig;
#[derive(Parser)]
#[command(
name = "vllm-bench",
about = "Benchmark online serving throughput",
version
)]
struct Cli {
#[command(flatten)]
args: vllm_bench::BenchServeArgs,
}
fn main() -> anyhow::Result<()> {
// Raise the open-file soft limit to the hard limit. High-concurrency
// benchmarks (1024+ requests) easily exceed the default 1024 fd soft limit.
if let Ok(new) = rlimit::increase_nofile_limit(u64::MAX)
&& new > 1024
{
eprintln!("Open-file limit: {new}");
}
let cli = Cli::parse();
// --- Compare mode: no server needed, just diff two JSON files ---
if let Some(ref files) = cli.compare {
return compare::compare_results(&files[0], &files[1]).context("Comparison failed");
}
let config = BenchConfig::from_cli(&cli).context("Configuration error")?;
vllm_bench::prepare_process();
let runtime = tokio::runtime::Builder::new_multi_thread()
.enable_all()
.build()
.expect("Failed to build tokio runtime");
.context("Failed to build tokio runtime")?;
runtime
.block_on(async {
if config.multi_turn {
if let Some(ref sweep_mc) = cli.sweep_max_concurrency {
// --- Sweep over concurrency in multi-turn mode ---
let values = sweep::parse_concurrency_values(sweep_mc)
.context("Invalid --sweep-max-concurrency")?;
sweep::run_multi_turn_concurrency_sweep(
&config,
&values,
cli.sweep_num_prompts_factor,
)
.await?;
} else {
// --- Single multi-turn conversation benchmark ---
multi_turn::run_multi_turn_benchmark(&config).await?;
}
} else if let Some(ref sweep_mc) = cli.sweep_max_concurrency {
// --- Sweep over max-concurrency ---
let values = sweep::parse_concurrency_values(sweep_mc)
.context("Invalid --sweep-max-concurrency")?;
sweep::run_concurrency_sweep(&config, &values, cli.sweep_num_prompts_factor)
.await?;
} else if let Some(ref sweep_rate) = cli.sweep_request_rate {
// --- Sweep over request-rate ---
let values =
sweep::parse_rate_values(sweep_rate).context("Invalid --sweep-request-rate")?;
sweep::run_rate_sweep(&config, &values).await?;
} else if cli.num_runs > 1 {
// --- Multi-run with statistical aggregation ---
multi_run::run_multi(&config, cli.num_runs).await?;
} else {
// --- Normal single benchmark ---
benchmark::run_benchmark(&config).await?;
}
anyhow::Ok(())
})
.context("Benchmark failed")
runtime.block_on(vllm_bench::run(cli.args))
}
rust/src/cmd/Cargo.toml
+1
View File
@@ -29,6 +29,7 @@ tokio-util.workspace = true
tracing.workspace = true
tracing-subscriber.workspace = true
uuid.workspace = true
vllm-bench.workspace = true
vllm-chat.workspace = true
vllm-engine-core-client.workspace = true
vllm-managed-engine.workspace = true
rust/src/cmd/src/cli.rs
+11 -1
View File
@@ -79,13 +79,23 @@ impl Cli {
}
/// Supported top-level CLI commands.
#[derive(Debug, Subcommand, PartialEq, Eq)]
#[derive(Debug, Subcommand)]
pub enum Command {
/// Run the Rust OpenAI frontend as a Python-supervised worker.
Frontend(FrontendArgs),
/// Launch a managed Python headless engine, then run the Rust OpenAI
/// frontend.
Serve(ServeArgs),
/// Run vLLM benchmarks.
#[command(subcommand)]
Bench(BenchCommand),
}
/// Supported benchmark commands.
#[derive(Debug, Subcommand)]
pub enum BenchCommand {
/// Benchmark online serving throughput.
Serve(vllm_bench::BenchServeArgs),
}
/// A JSON-encoded list of strings, matching Python's `json.loads` CLI type for
rust/src/cmd/src/cli/tests.rs
+21 -1
View File
@@ -5,7 +5,27 @@ use expect_test::expect;
use vllm_engine_core_client::TransportMode;
use vllm_server::{Config, HttpListenerMode, ParserSelection, RendererSelection};
use super::{Cli, Command};
use super::{BenchCommand, Cli, Command};
#[test]
fn bench_serve_args_parse_without_managed_engine_repartition() {
let cli = Cli::try_parse_from([
"vllm-rs",
"bench",
"serve",
"--backend",
"openai-chat",
"--request-rate",
"inf",
])
.unwrap();
let Command::Bench(BenchCommand::Serve(args)) = cli.command else {
panic!("expected bench serve args");
};
assert_eq!(args.backend, vllm_bench::BackendKind::OpenaiChat);
assert!(args.request_rate.is_infinite());
}
#[test]
fn serve_args_forward_python_flags_with_separator() {
rust/src/cmd/src/main.rs
+5 -1
View File
@@ -12,7 +12,7 @@ use tokio_util::sync::CancellationToken;
use tracing::{info, warn};
use vllm_managed_engine::ManagedEngineHandle;
use crate::cli::{Cli, Command};
use crate::cli::{BenchCommand, Cli, Command};
#[global_allocator]
static GLOBAL: mimalloc::MiMalloc = mimalloc::MiMalloc;
@@ -100,6 +100,10 @@ fn main() -> Result<()> {
async fn async_main(cli: Cli) -> Result<()> {
match cli.command {
Command::Frontend(args) => vllm_server::serve(args.into_config(), shutdown_signal()).await,
Command::Bench(BenchCommand::Serve(bench_args)) => {
vllm_bench::prepare_process();
vllm_bench::run(bench_args).await
}
Command::Serve(args) => {
let handshake_port = args.managed_engine.resolve_handshake_port()?;
vllm/entrypoints/cli/benchmark/serve.py
+60 -2
View File
@@ -1,11 +1,68 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import argparse
import os
import sys
from pathlib import Path
from vllm.benchmarks.serve import add_cli_args, main
from vllm.benchmarks.serve import add_cli_args
from vllm.benchmarks.serve import main as python_main
from vllm.entrypoints.cli.benchmark.base import BenchmarkSubcommandBase
from vllm.logger import init_logger
from vllm.utils.argparse_utils import FlexibleArgumentParser
logger = init_logger(__name__)
_RUST_CLI_PATH = Path(__file__).resolve().parents[3] / "vllm-rs"
_RUST_SUPPORTED_DATASETS = frozenset(
{
"custom",
"hf",
"prefix_repetition",
"random",
"random-mm",
"random-rerank",
"sharegpt",
"sonnet",
"speed_bench",
}
)
_RUST_SUPPORTED_BACKENDS = frozenset(
{
"openai",
"openai-chat",
"openai-embeddings",
"openai-embeddings-chat",
"vllm",
"vllm-pooling",
"vllm-rerank",
}
)
def _rust_unsupported_reason(args: argparse.Namespace) -> str | None:
if args.dataset_name not in _RUST_SUPPORTED_DATASETS:
return f"dataset {args.dataset_name!r} is not supported by the Rust benchmark"
if args.backend not in _RUST_SUPPORTED_BACKENDS:
return f"backend {args.backend!r} is not supported by the Rust benchmark"
return None
def _maybe_exec_rust_bench(args: argparse.Namespace) -> None:
if reason := _rust_unsupported_reason(args):
logger.info("Using Python benchmark: %s.", reason)
return
if not _RUST_CLI_PATH.is_file():
logger.warning(
"Rust benchmark binary not found at %s; falling back to Python.",
_RUST_CLI_PATH,
)
return
rust_cli = str(_RUST_CLI_PATH)
logger.info("Delegating `vllm bench serve` to Rust binary at %s.", rust_cli)
os.execv(rust_cli, [rust_cli, "bench", "serve", *sys.argv[3:]])
class BenchmarkServingSubcommand(BenchmarkSubcommandBase):
"""The `serve` subcommand for `vllm bench`."""
@@ -19,4 +76,5 @@ class BenchmarkServingSubcommand(BenchmarkSubcommandBase):
@staticmethod
def cmd(args: argparse.Namespace) -> None:
main(args)
_maybe_exec_rust_bench(args)
python_main(args)