diff --git a/rust/Cargo.lock b/rust/Cargo.lock index 0709ef8fc3a..e7a157923f8 100644 --- a/rust/Cargo.lock +++ b/rust/Cargo.lock @@ -5560,6 +5560,7 @@ dependencies = [ "tracing", "tracing-subscriber", "uuid", + "vllm-bench", "vllm-chat", "vllm-engine-core-client", "vllm-managed-engine", diff --git a/rust/Cargo.toml b/rust/Cargo.toml index 492b1c5e260..5fc9ea6edb7 100644 --- a/rust/Cargo.toml +++ b/rust/Cargo.toml @@ -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" } diff --git a/rust/src/bench/src/cli.rs b/rust/src/bench/src/cli.rs index 9aa8f388633..06e32e14ce3 100644 --- a/rust/src/bench/src/cli.rs +++ b/rust/src/bench/src/cli.rs @@ -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 { diff --git a/rust/src/bench/src/config.rs b/rust/src/bench/src/config.rs index 9696612c30f..ec132d11969 100644 --- a/rust/src/bench/src/config.rs +++ b/rust/src/bench/src/config.rs @@ -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 { - if cli.burstiness <= 0.0 { + pub fn from_args(args: &BenchServeArgs) -> Result { + 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 = 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>) -> Result { Ok(config) } -fn parse_ramp_up(cli: &Cli) -> Result> { - let strategy = match cli.ramp_up_strategy { +fn parse_ramp_up(args: &BenchServeArgs) -> Result> { + 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(args: I) -> BenchServeArgs + where + I: IntoIterator, + T: Into + 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() { diff --git a/rust/src/bench/src/hub.rs b/rust/src/bench/src/hub.rs index 0e1b3e962c8..ab51ccdcf08 100644 --- a/rust/src/bench/src/hub.rs +++ b/rust/src/bench/src/hub.rs @@ -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}")) }) }) diff --git a/rust/src/bench/src/lib.rs b/rust/src/bench/src/lib.rs new file mode 100644 index 00000000000..13f252cee33 --- /dev/null +++ b/rust/src/bench/src/lib.rs @@ -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") +} diff --git a/rust/src/bench/src/main.rs b/rust/src/bench/src/main.rs index dd37ec56db3..919e1e0b369 100644 --- a/rust/src/bench/src/main.rs +++ b/rust/src/bench/src/main.rs @@ -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)) } diff --git a/rust/src/cmd/Cargo.toml b/rust/src/cmd/Cargo.toml index 030d4c6d116..a6955059c26 100644 --- a/rust/src/cmd/Cargo.toml +++ b/rust/src/cmd/Cargo.toml @@ -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 diff --git a/rust/src/cmd/src/cli.rs b/rust/src/cmd/src/cli.rs index 314765bab26..454a397ad05 100644 --- a/rust/src/cmd/src/cli.rs +++ b/rust/src/cmd/src/cli.rs @@ -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 diff --git a/rust/src/cmd/src/cli/tests.rs b/rust/src/cmd/src/cli/tests.rs index ee0d3ebf048..d8ec9792b7e 100644 --- a/rust/src/cmd/src/cli/tests.rs +++ b/rust/src/cmd/src/cli/tests.rs @@ -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() { diff --git a/rust/src/cmd/src/main.rs b/rust/src/cmd/src/main.rs index 9637a3be578..69314678103 100644 --- a/rust/src/cmd/src/main.rs +++ b/rust/src/cmd/src/main.rs @@ -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()?; diff --git a/vllm/entrypoints/cli/benchmark/serve.py b/vllm/entrypoints/cli/benchmark/serve.py index 188afd6c703..41a65273ba8 100644 --- a/vllm/entrypoints/cli/benchmark/serve.py +++ b/vllm/entrypoints/cli/benchmark/serve.py @@ -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)