@tduncan / Nvidia-Optimization-Stack-VLLM / commits / c01618fdc8

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

Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> Signed-off-by: Bugen Zhao <i@bugenzhao.com>

Bugen Zhao committed Jul 20, 2026 at 17:31 UTC c01618fdc81524ed4e91652bbf85ed23d6ddc448
12 files changed +421 -280
rust/Cargo.lock
+1
@@ -5560,6 +5560,7 @@ dependencies = [
5560 "tracing",
5561 "tracing-subscriber",
5562 "uuid",
5563 + "vllm-bench",
5564 "vllm-chat",
5565 "vllm-engine-core-client",
5566 "vllm-managed-engine",
rust/Cargo.toml
+1
@@ -132,6 +132,7 @@ trait-set = "0.3.0"
132 url = "2.5.7"
133 uuid = { version = "1.22.0", features = ["v4"] }
134 validator = { version = "0.20.0", features = ["derive"] }
135 +vllm-bench = { path = "src/bench" }
136 vllm-chat = { path = "src/chat" }
137 vllm-engine-core-client = { path = "src/engine-core-client" }
138 vllm-llm = { path = "src/llm" }
rust/src/bench/src/cli.rs
+4 -11
@@ -3,8 +3,6 @@
3
4 use std::fmt;
5
6 -use clap::Parser;
7 -
6 /// Backend type for the benchmark endpoint.
7 #[derive(clap::ValueEnum, Debug, Clone, Copy, PartialEq, Eq)]
8 pub enum BackendKind {
@@ -77,7 +75,7 @@ pub enum DatasetName {
75 ShareGpt,
76 #[value(name = "sonnet")]
77 Sonnet,
80 - #[value(name = "speed-bench")]
78 + #[value(name = "speed-bench", alias = "speed_bench")]
79 SpeedBench,
80 #[value(name = "hf")]
81 Hf,
@@ -144,13 +142,8 @@ impl fmt::Display for SpeedBenchConfig {
142 }
143
144 /// High-performance benchmark client for vLLM serving endpoints.
147 -#[derive(Parser, Debug, Clone)]
148 -#[command(
149 - name = "vllm-bench",
150 - about = "Benchmark online serving throughput",
151 - version
152 -)]
153 -pub struct Cli {
145 +#[derive(clap::Args, Debug, Clone)]
146 +pub struct BenchServeArgs {
147 /// The type of backend or endpoint to use for the benchmark.
148 #[arg(long, default_value = "openai")]
149 pub backend: BackendKind,
@@ -659,7 +652,7 @@ pub struct Cli {
652 pub lora_assignment: LoraAssignment,
653 }
654
662 -impl Cli {
655 +impl BenchServeArgs {
656 /// Resolve the base URL from explicit --base-url or from --host/--port.
657 pub fn resolve_base_url(&self) -> String {
658 if let Some(ref base) = self.base_url {
rust/src/bench/src/config.rs
+212 -188
@@ -4,7 +4,9 @@
4 use std::collections::HashMap;
5 use std::sync::Arc;
6
7 -use crate::cli::{BackendKind, Cli, DatasetName, LoraAssignment, RampUpStrategy, SpeedBenchConfig};
7 +use crate::cli::{
8 + BackendKind, BenchServeArgs, DatasetName, LoraAssignment, RampUpStrategy, SpeedBenchConfig,
9 +};
10 use crate::datasets::random_mm::{MmBucketKey, MmLimitPerPrompt};
11 use crate::error::{BenchError, Result};
12
@@ -215,63 +217,63 @@ pub struct BenchConfig {
217 }
218
219 impl BenchConfig {
218 - pub fn from_cli(cli: &Cli) -> Result<Self> {
219 - if cli.burstiness <= 0.0 {
220 + pub fn from_args(args: &BenchServeArgs) -> Result<Self> {
221 + if args.burstiness <= 0.0 {
222 return Err(BenchError::Config("Burstiness must be positive".into()));
223 }
224
223 - if cli.num_prompts == 0 {
225 + if args.num_prompts == 0 {
226 return Err(BenchError::Config(
227 "--num-prompts must be at least 1".into(),
228 ));
229 }
230
229 - if cli.request_rate <= 0.0 && !cli.request_rate.is_infinite() {
231 + if args.request_rate <= 0.0 && !args.request_rate.is_infinite() {
232 return Err(BenchError::Config(
233 "--request-rate must be positive (or inf)".into(),
234 ));
235 }
234 - if cli.max_model_len == Some(0) {
236 + if args.max_model_len == Some(0) {
237 return Err(BenchError::Config(
238 "--max-model-len must be at least 1".into(),
239 ));
240 }
241
240 - let base_url = cli.resolve_base_url();
241 - let api_url = cli.resolve_api_url();
242 + let base_url = args.resolve_base_url();
243 + let api_url = args.resolve_api_url();
244
243 - let extra_headers = cli.parse_headers()?;
244 - let mut extra_body = cli.parse_extra_body()?;
245 + let extra_headers = args.parse_headers()?;
246 + let mut extra_body = args.parse_extra_body()?;
247
248 // Merge sampling parameters into extra_body (matches Python behavior).
249 // Python collects non-None sampling params and merges them UNDER extra_body,
250 // meaning extra_body keys take precedence over sampling params.
251 {
252 let mut sampling_params = serde_json::Map::new();
251 - if let Some(v) = cli.top_p {
253 + if let Some(v) = args.top_p {
254 sampling_params.insert("top_p".into(), serde_json::json!(v));
255 }
254 - if let Some(v) = cli.top_k {
256 + if let Some(v) = args.top_k {
257 sampling_params.insert("top_k".into(), serde_json::json!(v));
258 }
257 - if let Some(v) = cli.min_p {
259 + if let Some(v) = args.min_p {
260 sampling_params.insert("min_p".into(), serde_json::json!(v));
261 }
260 - if let Some(v) = cli.temperature {
262 + if let Some(v) = args.temperature {
263 sampling_params.insert("temperature".into(), serde_json::json!(v));
264 }
263 - if let Some(v) = cli.frequency_penalty {
265 + if let Some(v) = args.frequency_penalty {
266 sampling_params.insert("frequency_penalty".into(), serde_json::json!(v));
267 }
266 - if let Some(v) = cli.presence_penalty {
268 + if let Some(v) = args.presence_penalty {
269 sampling_params.insert("presence_penalty".into(), serde_json::json!(v));
270 }
269 - if let Some(v) = cli.repetition_penalty {
271 + if let Some(v) = args.repetition_penalty {
272 sampling_params.insert("repetition_penalty".into(), serde_json::json!(v));
273 }
274
275 if !sampling_params.is_empty() {
274 - if !cli.backend.is_openai_compatible() {
276 + if !args.backend.is_openai_compatible() {
277 return Err(BenchError::Config(
278 "Sampling parameters are only supported by openai-compatible backends."
279 .into(),
@@ -299,7 +301,7 @@ impl BenchConfig {
301 }
302
303 // Parse metadata
302 - let metadata = match &cli.metadata {
304 + let metadata = match &args.metadata {
305 None => None,
306 Some(items) => {
307 let mut pairs = Vec::new();
@@ -314,24 +316,24 @@ impl BenchConfig {
316 };
317
318 // Parse goodput SLOs
317 - let goodput = parse_goodput(&cli.goodput)?;
319 + let goodput = parse_goodput(&args.goodput)?;
320
321 // Parse ramp-up config
320 - let ramp_up = parse_ramp_up(cli)?;
322 + let ramp_up = parse_ramp_up(args)?;
323
324 // Default percentile metrics based on backend type
323 - let default_percentile_metrics = if cli.backend.is_pooling() {
325 + let default_percentile_metrics = if args.backend.is_pooling() {
326 "e2el"
327 } else {
328 "ttft,tpot,itl,e2el"
329 };
330 let percentile_metrics_str =
329 - cli.percentile_metrics.as_deref().unwrap_or(default_percentile_metrics);
331 + args.percentile_metrics.as_deref().unwrap_or(default_percentile_metrics);
332 let selected_percentile_metrics: Vec<String> =
333 percentile_metrics_str.split(',').map(|s| s.trim().to_string()).collect();
334
333 - let metric_percentiles = parse_percentiles(&cli.metric_percentiles, false)?;
334 - let sweep_summary_percentiles = cli
335 + let metric_percentiles = parse_percentiles(&args.metric_percentiles, false)?;
336 + let sweep_summary_percentiles = args
337 .sweep_summary_percentiles
338 .as_deref()
339 .map(|raw| parse_percentiles(raw, true))
@@ -344,38 +346,38 @@ impl BenchConfig {
346 selected_percentiles.push(90.0);
347 }
348
347 - let tokenizer_id = if cli.skip_tokenizer_init {
349 + let tokenizer_id = if args.skip_tokenizer_init {
350 None
351 } else {
350 - Some(cli.tokenizer.clone().or_else(|| cli.model.clone()).unwrap_or_default())
352 + args.tokenizer.clone().or_else(|| args.model.clone())
353 };
354
355 // Resolve input/output lengths
354 - let random_input_len = cli.resolved_random_input_len();
355 - let random_output_len = cli.resolved_random_output_len();
356 - let per_turn_input_len = cli.resolved_per_turn_input_len();
356 + let random_input_len = args.resolved_random_input_len();
357 + let random_output_len = args.resolved_random_output_len();
358 + let per_turn_input_len = args.resolved_per_turn_input_len();
359
360 // Normalized multi-turn turn counts (computed in validation block below, defaults
361 // to num_turns if multi-turn mode is not active)
360 - let mut multi_turn_min_turns = cli.multi_turn_num_turns;
361 - let mut multi_turn_max_turns = cli.multi_turn_num_turns;
362 + let mut multi_turn_min_turns = args.multi_turn_num_turns;
363 + let mut multi_turn_max_turns = args.multi_turn_num_turns;
364
365 // For random datasets with openai-compatible backends, default to ignore_eos.
366 // Exception: multi-turn mode, where ignore_eos causes unbounded context growth
367 // across turns. Multi-turn uses min_tokens instead for output length control.
368 // Pooling backends don't generate tokens, so ignore_eos is irrelevant.
367 - let ignore_eos = if cli.backend.is_pooling() {
369 + let ignore_eos = if args.backend.is_pooling() {
370 false
371 } else {
370 - cli.ignore_eos
371 - || ((cli.dataset_name == DatasetName::Random
372 - || cli.dataset_name == DatasetName::RandomMm)
373 - && cli.backend.is_openai_compatible()
374 - && !cli.multi_turn)
372 + args.ignore_eos
373 + || ((args.dataset_name == DatasetName::Random
374 + || args.dataset_name == DatasetName::RandomMm)
375 + && args.backend.is_openai_compatible()
376 + && !args.multi_turn)
377 };
378
379 // Pooling backends don't support multi-turn
378 - if cli.backend.is_pooling() && cli.multi_turn {
380 + if args.backend.is_pooling() && args.multi_turn {
381 return Err(BenchError::Config(
382 "Pooling/embedding backends do not support --multi-turn".into(),
383 ));
@@ -383,7 +385,7 @@ impl BenchConfig {
385
386 // LoRA validation. Adapter names must be non-empty after trim; pooling
387 // backends are out of scope (vLLM LoRA routing is for generative paths).
386 - let lora_modules = match cli.lora_modules.as_ref() {
388 + let lora_modules = match args.lora_modules.as_ref() {
389 None => None,
390 Some(names) => {
391 if names.is_empty() {
@@ -391,7 +393,7 @@ impl BenchConfig {
393 "--lora-modules requires at least one adapter name".into(),
394 ));
395 }
394 - if cli.backend.is_pooling() {
396 + if args.backend.is_pooling() {
397 return Err(BenchError::Config(
398 "--lora-modules is not supported for pooling/embedding backends".into(),
399 ));
@@ -411,18 +413,18 @@ impl BenchConfig {
413 };
414
415 // Random-MM validation and config parsing
414 - let (random_mm_limit, random_mm_buckets) = if cli.dataset_name == DatasetName::RandomMm {
415 - if cli.backend != BackendKind::OpenaiChat {
416 + let (random_mm_limit, random_mm_buckets) = if args.dataset_name == DatasetName::RandomMm {
417 + if args.backend != BackendKind::OpenaiChat {
418 return Err(BenchError::Config(
419 "Multi-modal content (images) is only supported on 'openai-chat' backend."
420 .into(),
421 ));
422 }
423 let limit = crate::datasets::random_mm::parse_limit_mm_per_prompt(
422 - &cli.random_mm_limit_mm_per_prompt,
424 + &args.random_mm_limit_mm_per_prompt,
425 )?;
426 let buckets =
425 - crate::datasets::random_mm::parse_bucket_config(&cli.random_mm_bucket_config)?;
427 + crate::datasets::random_mm::parse_bucket_config(&args.random_mm_bucket_config)?;
428 (limit, buckets)
429 } else {
430 (MmLimitPerPrompt::default(), Vec::new())
@@ -432,18 +434,18 @@ impl BenchConfig {
434 // sonnet (uses built-in Shakespeare's sonnets).
435
436 // Range ratio (Python semantics: [len*(1-r), len*(1+r)], each r in [0,1))
435 - let random_range_ratio = RangeRatio::parse(&cli.random_range_ratio)?;
437 + let random_range_ratio = RangeRatio::parse(&args.random_range_ratio)?;
438
439 // Batched inputs only make sense for pooling backends (the generation
440 // backends send one prompt per request).
439 - if cli.random_batch_size == 0 {
441 + if args.random_batch_size == 0 {
442 return Err(BenchError::Config(
443 "--random-batch-size must be at least 1".into(),
444 ));
445 }
444 - if cli.random_batch_size > 1
445 - && !cli.backend.is_pooling()
446 - && cli.dataset_name != DatasetName::RandomRerank
446 + if args.random_batch_size > 1
447 + && !args.backend.is_pooling()
448 + && args.dataset_name != DatasetName::RandomRerank
449 {
450 return Err(BenchError::Config(
451 "--random-batch-size > 1 is only supported with embeddings/pooling backends".into(),
@@ -451,16 +453,16 @@ impl BenchConfig {
453 }
454
455 // random-rerank validation (mirrors Python RandomDatasetForReranking)
454 - let is_reranker = !cli.no_reranker;
455 - if cli.dataset_name == DatasetName::RandomRerank {
456 - if !cli.backend.is_pooling() {
456 + let is_reranker = !args.no_reranker;
457 + if args.dataset_name == DatasetName::RandomRerank {
458 + if !args.backend.is_pooling() {
459 return Err(BenchError::Config(
460 "--dataset-name random-rerank requires an embeddings/pooling backend \
461 (e.g. --backend vllm-rerank)"
462 .into(),
463 ));
464 }
463 - if !is_reranker && (cli.num_prompts < 2 || cli.random_batch_size < 2) {
465 + if !is_reranker && (args.num_prompts < 2 || args.random_batch_size < 2) {
466 return Err(BenchError::Config(
467 "--no-reranker requires --num-prompts > 1 and --random-batch-size > 1 \
468 (the query is folded into the first batch slot)"
@@ -470,8 +472,8 @@ impl BenchConfig {
472 }
473
474 // Custom dataset validation
473 - if cli.dataset_name == DatasetName::Custom {
474 - match cli.dataset_path.as_deref() {
475 + if args.dataset_name == DatasetName::Custom {
476 + match args.dataset_path.as_deref() {
477 None => {
478 return Err(BenchError::Config(
479 "--dataset-path is required for --dataset-name custom \
@@ -486,7 +488,7 @@ impl BenchConfig {
488 }
489 _ => {}
490 }
489 - if !cli.skip_chat_template {
491 + if !args.skip_chat_template {
492 eprintln!(
493 "NOTE: client-side chat template rendering is not supported; custom \
494 dataset prompts are sent raw (equivalent to --skip-chat-template)."
@@ -495,29 +497,29 @@ impl BenchConfig {
497 }
498
499 // Prefix repetition validation
498 - if cli.dataset_name == DatasetName::PrefixRepetition {
499 - if cli.prefix_repetition_num_prefixes == 0 {
500 + if args.dataset_name == DatasetName::PrefixRepetition {
501 + if args.prefix_repetition_num_prefixes == 0 {
502 return Err(BenchError::Config(
503 "--prefix-repetition-num-prefixes must be at least 1".into(),
504 ));
505 }
504 - if cli.num_prompts < cli.prefix_repetition_num_prefixes {
506 + if args.num_prompts < args.prefix_repetition_num_prefixes {
507 return Err(BenchError::Config(format!(
508 "--num-prompts ({}) must be >= --prefix-repetition-num-prefixes ({})",
507 - cli.num_prompts, cli.prefix_repetition_num_prefixes
509 + args.num_prompts, args.prefix_repetition_num_prefixes
510 )));
511 }
512 }
513
514 // HF dataset validation
513 - if cli.dataset_name == DatasetName::Hf && cli.dataset_path.is_none() {
515 + if args.dataset_name == DatasetName::Hf && args.dataset_path.is_none() {
516 return Err(BenchError::Config(
517 "--dataset-path is required for --dataset-name hf \
518 (set to a HuggingFace dataset ID, e.g. 'allenai/WildChat-4.8M')"
519 .into(),
520 ));
521 }
520 - if let Some(len) = cli.hf_output_len
522 + if let Some(len) = args.hf_output_len
523 && len == 0
524 {
525 return Err(BenchError::Config(
@@ -526,13 +528,13 @@ impl BenchConfig {
528 }
529
530 // Multi-turn validation
529 - if cli.multi_turn {
530 - if cli.backend != BackendKind::OpenaiChat {
531 + if args.multi_turn {
532 + if args.backend != BackendKind::OpenaiChat {
533 return Err(BenchError::Config(
534 "--multi-turn requires --backend openai-chat".into(),
535 ));
536 }
535 - if cli.multi_turn_num_turns == 0 {
537 + if args.multi_turn_num_turns == 0 {
538 return Err(BenchError::Config(
539 "--multi-turn-num-turns must be at least 1".into(),
540 ));
@@ -541,18 +543,18 @@ impl BenchConfig {
543 // Normalize and validate min/max turns. ShareGPT only consumes max_turns
544 // (the loader walks all available turns up to the cap), so the
545 // min/num/max coupling used for synthetic generation does not apply.
544 - if cli.dataset_name == DatasetName::ShareGpt {
545 - if cli.multi_turn_max_turns == 1 {
546 + if args.dataset_name == DatasetName::ShareGpt {
547 + if args.multi_turn_max_turns == 1 {
548 return Err(BenchError::Config(
549 "--multi-turn-max-turns must be at least 2 for ShareGPT multi-turn".into(),
550 ));
551 }
552 } else {
553 (multi_turn_min_turns, multi_turn_max_turns) =
552 - match (cli.multi_turn_min_turns, cli.multi_turn_max_turns) {
553 - (0, 0) => (cli.multi_turn_num_turns, cli.multi_turn_num_turns),
554 - (m, 0) => (m, cli.multi_turn_num_turns),
555 - (0, x) => (cli.multi_turn_num_turns, x),
554 + match (args.multi_turn_min_turns, args.multi_turn_max_turns) {
555 + (0, 0) => (args.multi_turn_num_turns, args.multi_turn_num_turns),
556 + (m, 0) => (m, args.multi_turn_num_turns),
557 + (0, x) => (args.multi_turn_num_turns, x),
558 (m, x) => (m, x),
559 };
560 if multi_turn_min_turns < 1 {
@@ -575,8 +577,8 @@ impl BenchConfig {
577 }
578
579 // Validate prefix sharing ratios
578 - let pg = cli.multi_turn_prefix_global_ratio;
579 - let pc = cli.multi_turn_prefix_conversation_ratio;
580 + let pg = args.multi_turn_prefix_global_ratio;
581 + let pc = args.multi_turn_prefix_conversation_ratio;
582 if !(0.0..=1.0).contains(&pg) {
583 return Err(BenchError::Config(
584 "--multi-turn-prefix-global-ratio must be in [0.0, 1.0]".into(),
@@ -592,20 +594,20 @@ impl BenchConfig {
594 "--multi-turn-prefix-global-ratio + --multi-turn-prefix-conversation-ratio must be < 1.0 (unique suffix required)".into(),
595 ));
596 }
595 - if (pg > 0.0 || pc > 0.0) && cli.dataset_name != DatasetName::Random {
597 + if (pg > 0.0 || pc > 0.0) && args.dataset_name != DatasetName::Random {
598 return Err(BenchError::Config(
599 "Prefix sharing (--multi-turn-prefix-global-ratio / --multi-turn-prefix-conversation-ratio) only works with --dataset-name random".into(),
600 ));
601 }
602 }
603
602 - if !(cli.steady_state_threshold > 0.0 && cli.steady_state_threshold <= 1.0) {
604 + if !(args.steady_state_threshold > 0.0 && args.steady_state_threshold <= 1.0) {
605 return Err(BenchError::Config(format!(
606 "--steady-state-threshold must be in (0.0, 1.0], got {}",
605 - cli.steady_state_threshold
607 + args.steady_state_threshold
608 )));
609 }
608 - if let Some(mw) = cli.steady_state_min_window
610 + if let Some(mw) = args.steady_state_min_window
611 && mw < 0.0
612 {
613 return Err(BenchError::Config(format!(
@@ -613,122 +615,122 @@ impl BenchConfig {
615 )));
616 }
617
616 - if cli.profile_batch_threshold.is_some() && !cli.profile {
618 + if args.profile_batch_threshold.is_some() && !args.profile {
619 return Err(BenchError::Config(
620 "--profile-batch-threshold requires --profile".into(),
621 ));
622 }
621 - if cli.profile_duration <= 0.0 {
623 + if args.profile_duration <= 0.0 {
624 return Err(BenchError::Config(
625 "--profile-duration must be positive".into(),
626 ));
627 }
626 - if cli.profile_batch_threshold.is_none() && cli.profile_duration != 5.0 {
628 + if args.profile_batch_threshold.is_none() && args.profile_duration != 5.0 {
629 return Err(BenchError::Config(
630 "--profile-duration requires --profile-batch-threshold".into(),
631 ));
632 }
633
634 Ok(BenchConfig {
633 - backend: cli.backend,
635 + backend: args.backend,
636 base_url,
637 api_url,
636 - model: cli.model.clone(),
637 - model_name: cli.served_model_name.clone(),
638 + model: args.model.clone(),
639 + model_name: args.served_model_name.clone(),
640 tokenizer_id,
639 - tokenizer_mode: cli.tokenizer_mode.clone(),
640 - trust_remote_code: cli.trust_remote_code,
641 - skip_tokenizer_init: cli.skip_tokenizer_init,
642 - dataset_name: cli.dataset_name,
643 - dataset_path: cli.dataset_path.clone(),
644 - max_model_len: cli.max_model_len,
641 + tokenizer_mode: args.tokenizer_mode.clone(),
642 + trust_remote_code: args.trust_remote_code,
643 + skip_tokenizer_init: args.skip_tokenizer_init,
644 + dataset_name: args.dataset_name,
645 + dataset_path: args.dataset_path.clone(),
646 + max_model_len: args.max_model_len,
647 random_input_len,
648 random_output_len,
647 - random_prefix_len: cli.random_prefix_len,
649 + random_prefix_len: args.random_prefix_len,
650 random_range_ratio,
649 - random_batch_size: cli.random_batch_size,
651 + random_batch_size: args.random_batch_size,
652 is_reranker,
651 - custom_output_len: cli.output_len.map(|v| v as i64).unwrap_or(cli.custom_output_len),
652 - prefix_repetition_prefix_len: cli.prefix_repetition_prefix_len,
653 - prefix_repetition_suffix_len: cli.prefix_repetition_suffix_len,
654 - prefix_repetition_num_prefixes: cli.prefix_repetition_num_prefixes,
655 - prefix_repetition_output_len: cli
653 + custom_output_len: args.output_len.map(|v| v as i64).unwrap_or(args.custom_output_len),
654 + prefix_repetition_prefix_len: args.prefix_repetition_prefix_len,
655 + prefix_repetition_suffix_len: args.prefix_repetition_suffix_len,
656 + prefix_repetition_num_prefixes: args.prefix_repetition_num_prefixes,
657 + prefix_repetition_output_len: args
658 .output_len
657 - .unwrap_or(cli.prefix_repetition_output_len),
658 - random_cache_hit_fraction: cli.random_cache_hit_fraction,
659 - random_cache_ratio: cli.random_cache_ratio,
660 - sharegpt_output_len: cli.sharegpt_output_len,
661 - sonnet_input_len: cli.sonnet_input_len,
662 - sonnet_output_len: cli.sonnet_output_len,
663 - sonnet_prefix_len: cli.sonnet_prefix_len,
664 - no_oversample: cli.no_oversample,
665 - disable_shuffle: cli.disable_shuffle,
666 - num_prompts: cli.num_prompts,
667 - request_rate: cli.request_rate,
668 - burstiness: cli.burstiness,
669 - max_concurrency: cli.max_concurrency,
670 - steady_state_threshold: cli.steady_state_threshold,
671 - steady_state_min_window: cli.steady_state_min_window,
672 - no_steady_state: cli.no_steady_state,
673 - disable_tqdm: cli.disable_tqdm,
674 - num_warmups: cli.num_warmups,
675 - profile: cli.profile,
676 - profile_batch_threshold: cli.profile_batch_threshold,
677 - profile_duration: cli.profile_duration,
678 - save_result: cli.save_result,
679 - save_detailed: cli.save_detailed,
680 - append_result: cli.append_result,
681 - result_dir: cli.result_dir.clone(),
682 - result_filename: cli.result_filename.clone(),
683 - seed: cli.seed,
659 + .unwrap_or(args.prefix_repetition_output_len),
660 + random_cache_hit_fraction: args.random_cache_hit_fraction,
661 + random_cache_ratio: args.random_cache_ratio,
662 + sharegpt_output_len: args.sharegpt_output_len,
663 + sonnet_input_len: args.sonnet_input_len,
664 + sonnet_output_len: args.sonnet_output_len,
665 + sonnet_prefix_len: args.sonnet_prefix_len,
666 + no_oversample: args.no_oversample,
667 + disable_shuffle: args.disable_shuffle,
668 + num_prompts: args.num_prompts,
669 + request_rate: args.request_rate,
670 + burstiness: args.burstiness,
671 + max_concurrency: args.max_concurrency,
672 + steady_state_threshold: args.steady_state_threshold,
673 + steady_state_min_window: args.steady_state_min_window,
674 + no_steady_state: args.no_steady_state,
675 + disable_tqdm: args.disable_tqdm,
676 + num_warmups: args.num_warmups,
677 + profile: args.profile,
678 + profile_batch_threshold: args.profile_batch_threshold,
679 + profile_duration: args.profile_duration,
680 + save_result: args.save_result,
681 + save_detailed: args.save_detailed,
682 + append_result: args.append_result,
683 + result_dir: args.result_dir.clone(),
684 + result_filename: args.result_filename.clone(),
685 + seed: args.seed,
686 ignore_eos,
685 - insecure: cli.insecure,
687 + insecure: args.insecure,
688 selected_percentile_metrics,
689 selected_percentiles,
690 sweep_summary_percentiles,
689 - label: cli.label.clone(),
690 - logprobs: cli.logprobs,
691 - request_id_prefix: cli.get_request_id_prefix(),
692 - ready_check_timeout_sec: cli.ready_check_timeout_sec,
691 + label: args.label.clone(),
692 + logprobs: args.logprobs,
693 + request_id_prefix: args.get_request_id_prefix(),
694 + ready_check_timeout_sec: args.ready_check_timeout_sec,
695 extra_headers,
696 extra_body,
697 metadata,
696 - dry_run: cli.dry_run,
698 + dry_run: args.dry_run,
699 goodput,
700 ramp_up,
699 - multi_turn: cli.multi_turn,
700 - multi_turn_num_turns: cli.multi_turn_num_turns,
701 + multi_turn: args.multi_turn,
702 + multi_turn_num_turns: args.multi_turn_num_turns,
703 multi_turn_min_turns,
704 multi_turn_max_turns,
703 - sharegpt_multi_turn_max_turns: if cli.multi_turn
704 - && cli.dataset_name == DatasetName::ShareGpt
705 - && cli.multi_turn_max_turns != 0
705 + sharegpt_multi_turn_max_turns: if args.multi_turn
706 + && args.dataset_name == DatasetName::ShareGpt
707 + && args.multi_turn_max_turns != 0
708 {
707 - Some(cli.multi_turn_max_turns)
709 + Some(args.multi_turn_max_turns)
710 } else {
711 None
712 },
713 per_turn_input_len,
712 - multi_turn_concurrency: cli.multi_turn_concurrency,
713 - multi_turn_delay_ms: cli.multi_turn_delay_ms,
714 - multi_turn_prefix_global_ratio: cli.multi_turn_prefix_global_ratio,
715 - multi_turn_prefix_conversation_ratio: cli.multi_turn_prefix_conversation_ratio,
716 - speed_bench_config: cli.speed_bench_config,
717 - speed_bench_category: cli.speed_bench_category.clone(),
718 - speed_bench_max_input_len: cli.speed_bench_max_input_len,
719 - hf_split: cli.hf_split.clone(),
720 - hf_subset: cli.hf_subset.clone(),
721 - hf_output_len: cli.hf_output_len,
722 - hf_text_column: cli.hf_text_column.clone(),
723 - reset_prefix_cache: cli.reset_prefix_cache,
724 - prompt_token_ids: cli.prompt_token_ids,
725 - random_mm_base_items_per_request: cli.random_mm_base_items_per_request,
726 - random_mm_num_mm_items_range_ratio: cli.random_mm_num_mm_items_range_ratio,
714 + multi_turn_concurrency: args.multi_turn_concurrency,
715 + multi_turn_delay_ms: args.multi_turn_delay_ms,
716 + multi_turn_prefix_global_ratio: args.multi_turn_prefix_global_ratio,
717 + multi_turn_prefix_conversation_ratio: args.multi_turn_prefix_conversation_ratio,
718 + speed_bench_config: args.speed_bench_config,
719 + speed_bench_category: args.speed_bench_category.clone(),
720 + speed_bench_max_input_len: args.speed_bench_max_input_len,
721 + hf_split: args.hf_split.clone(),
722 + hf_subset: args.hf_subset.clone(),
723 + hf_output_len: args.hf_output_len,
724 + hf_text_column: args.hf_text_column.clone(),
725 + reset_prefix_cache: args.reset_prefix_cache,
726 + prompt_token_ids: args.prompt_token_ids,
727 + random_mm_base_items_per_request: args.random_mm_base_items_per_request,
728 + random_mm_num_mm_items_range_ratio: args.random_mm_num_mm_items_range_ratio,
729 random_mm_limit,
730 random_mm_buckets,
729 - enable_multimodal_chat: cli.enable_multimodal_chat,
731 + enable_multimodal_chat: args.enable_multimodal_chat,
732 lora_modules,
731 - lora_assignment: cli.lora_assignment,
733 + lora_assignment: args.lora_assignment,
734 })
735 }
736 }
@@ -811,17 +813,17 @@ fn parse_goodput(goodput_args: &Option<Vec<String>>) -> Result<GoodputConfig> {
813 Ok(config)
814 }
815
814 -fn parse_ramp_up(cli: &Cli) -> Result<Option<RampUpConfig>> {
815 - let strategy = match cli.ramp_up_strategy {
816 +fn parse_ramp_up(args: &BenchServeArgs) -> Result<Option<RampUpConfig>> {
817 + let strategy = match args.ramp_up_strategy {
818 None => return Ok(None),
819 Some(s) => s,
820 };
821
820 - let start_rps = cli.ramp_up_start_rps.ok_or_else(|| {
822 + let start_rps = args.ramp_up_start_rps.ok_or_else(|| {
823 BenchError::Config("--ramp-up-start-rps is required when --ramp-up-strategy is set".into())
824 })?;
825
824 - let end_rps = cli.ramp_up_end_rps.ok_or_else(|| {
826 + let end_rps = args.ramp_up_end_rps.ok_or_else(|| {
827 BenchError::Config("--ramp-up-end-rps is required when --ramp-up-strategy is set".into())
828 })?;
829
@@ -843,7 +845,21 @@ mod tests {
845 use clap::Parser;
846
847 use super::*;
846 - use crate::cli::Cli;
848 + use crate::cli::BenchServeArgs;
849 +
850 + #[derive(Parser)]
851 + struct TestCli {
852 + #[command(flatten)]
853 + args: BenchServeArgs,
854 + }
855 +
856 + fn parse_args<I, T>(args: I) -> BenchServeArgs
857 + where
858 + I: IntoIterator<Item = T>,
859 + T: Into<std::ffi::OsString> + Clone,
860 + {
861 + TestCli::parse_from(args).args
862 + }
863
864 fn base_multi_turn_args() -> Vec<&'static str> {
865 vec![
@@ -859,8 +875,8 @@ mod tests {
875 #[test]
876 fn test_prefix_sharing_defaults_to_zero() {
877 let args = base_multi_turn_args();
862 - let cli = Cli::parse_from(args);
863 - let config = BenchConfig::from_cli(&cli).unwrap();
878 + let args = parse_args(args);
879 + let config = BenchConfig::from_args(&args).unwrap();
880 assert_eq!(config.multi_turn_prefix_global_ratio, 0.0);
881 assert_eq!(config.multi_turn_prefix_conversation_ratio, 0.0);
882 }
@@ -874,8 +890,8 @@ mod tests {
890 "--multi-turn-prefix-conversation-ratio",
891 "0.8",
892 ]);
877 - let cli = Cli::parse_from(args);
878 - let config = BenchConfig::from_cli(&cli).unwrap();
893 + let args = parse_args(args);
894 + let config = BenchConfig::from_args(&args).unwrap();
895 assert!((config.multi_turn_prefix_global_ratio - 0.1).abs() < 1e-10);
896 assert!((config.multi_turn_prefix_conversation_ratio - 0.8).abs() < 1e-10);
897 }
@@ -889,8 +905,8 @@ mod tests {
905 "--multi-turn-prefix-conversation-ratio",
906 "0.6",
907 ]);
892 - let cli = Cli::parse_from(args);
893 - assert!(BenchConfig::from_cli(&cli).is_err());
908 + let args = parse_args(args);
909 + assert!(BenchConfig::from_args(&args).is_err());
910 }
911
912 #[test]
@@ -902,16 +918,16 @@ mod tests {
918 "--multi-turn-prefix-conversation-ratio",
919 "0.5",
920 ]);
905 - let cli = Cli::parse_from(args);
906 - assert!(BenchConfig::from_cli(&cli).is_err());
921 + let args = parse_args(args);
922 + assert!(BenchConfig::from_args(&args).is_err());
923 }
924
925 #[test]
926 fn test_prefix_sharing_out_of_range_fails() {
927 let mut args = base_multi_turn_args();
928 args.extend(["--multi-turn-prefix-global-ratio", "1.5"]);
913 - let cli = Cli::parse_from(args);
914 - assert!(BenchConfig::from_cli(&cli).is_err());
929 + let args = parse_args(args);
930 + assert!(BenchConfig::from_args(&args).is_err());
931 }
932
933 #[test]
@@ -928,8 +944,8 @@ mod tests {
944 "--multi-turn-prefix-global-ratio",
945 "0.1",
946 ];
931 - let cli = Cli::parse_from(args);
932 - assert!(BenchConfig::from_cli(&cli).is_err());
947 + let args = parse_args(args);
948 + assert!(BenchConfig::from_args(&args).is_err());
949 }
950
951 #[test]
@@ -944,8 +960,8 @@ mod tests {
960 "--dataset-name",
961 "sharegpt",
962 ];
947 - let cli = Cli::parse_from(args);
948 - let config = BenchConfig::from_cli(&cli).unwrap();
963 + let args = parse_args(args);
964 + let config = BenchConfig::from_args(&args).unwrap();
965
966 assert_eq!(config.multi_turn_max_turns, 3);
967 assert_eq!(config.sharegpt_multi_turn_max_turns, None);
@@ -968,8 +984,8 @@ mod tests {
984 "--multi-turn-max-turns",
985 "2",
986 ];
971 - let cli = Cli::parse_from(args);
972 - let config = BenchConfig::from_cli(&cli).unwrap();
987 + let args = parse_args(args);
988 + let config = BenchConfig::from_args(&args).unwrap();
989 assert_eq!(config.sharegpt_multi_turn_max_turns, Some(2));
990 }
991
@@ -987,8 +1003,8 @@ mod tests {
1003 "--multi-turn-max-turns",
1004 "1",
1005 ];
990 - let cli = Cli::parse_from(args);
991 - let err = BenchConfig::from_cli(&cli).unwrap_err().to_string();
1006 + let args = parse_args(args);
1007 + let err = BenchConfig::from_args(&args).unwrap_err().to_string();
1008 assert!(
1009 err.contains("at least 2 for ShareGPT"),
1010 "expected ShareGPT-specific error, got: {err}"
@@ -1009,8 +1025,8 @@ mod tests {
1025 "--multi-turn-max-turns",
1026 "20",
1027 ];
1012 - let cli = Cli::parse_from(args);
1013 - let config = BenchConfig::from_cli(&cli).unwrap();
1028 + let args = parse_args(args);
1029 + let config = BenchConfig::from_args(&args).unwrap();
1030
1031 assert_eq!(config.sharegpt_multi_turn_max_turns, Some(20));
1032 }
@@ -1018,8 +1034,8 @@ mod tests {
1034 #[test]
1035 fn test_sweep_summary_percentiles_default_empty() {
1036 let args = base_multi_turn_args();
1021 - let cli = Cli::parse_from(args);
1022 - let config = BenchConfig::from_cli(&cli).unwrap();
1037 + let args = parse_args(args);
1038 + let config = BenchConfig::from_args(&args).unwrap();
1039
1040 assert!(config.sweep_summary_percentiles.is_empty());
1041 assert_eq!(config.selected_percentiles, vec![99.0, 90.0]);
@@ -1034,8 +1050,8 @@ mod tests {
1050 "--sweep-summary-percentiles",
1051 "90,95,90",
1052 ]);
1037 - let cli = Cli::parse_from(args);
1038 - let config = BenchConfig::from_cli(&cli).unwrap();
1053 + let args = parse_args(args);
1054 + let config = BenchConfig::from_args(&args).unwrap();
1055
1056 assert_eq!(config.sweep_summary_percentiles, vec![90.0, 95.0]);
1057 assert_eq!(config.selected_percentiles, vec![99.0, 95.0, 90.0]);
@@ -1045,8 +1061,8 @@ mod tests {
1061 fn test_invalid_sweep_summary_percentile_fails() {
1062 let mut args = base_multi_turn_args();
1063 args.extend(["--sweep-summary-percentiles", "101"]);
1048 - let cli = Cli::parse_from(args);
1049 - assert!(BenchConfig::from_cli(&cli).is_err());
1064 + let args = parse_args(args);
1065 + assert!(BenchConfig::from_args(&args).is_err());
1066 }
1067
1068 #[test]
@@ -1058,12 +1074,20 @@ mod tests {
1074 "--max-model-len",
1075 "4096",
1076 ];
1061 - let cli = Cli::parse_from(args);
1062 - let config = BenchConfig::from_cli(&cli).unwrap();
1077 + let args = parse_args(args);
1078 + let config = BenchConfig::from_args(&args).unwrap();
1079
1080 assert_eq!(config.max_model_len, Some(4096));
1081 }
1082
1083 + #[test]
1084 + fn test_tokenizer_id_deferred_when_model_is_unspecified() {
1085 + let args = parse_args(["vllm-bench"]);
1086 + let config = BenchConfig::from_args(&args).unwrap();
1087 +
1088 + assert_eq!(config.tokenizer_id, None);
1089 + }
1090 +
1091 #[test]
1092 fn test_zero_max_model_len_fails() {
1093 let args = vec![
@@ -1073,9 +1097,9 @@ mod tests {
1097 "--max-model-len",
1098 "0",
1099 ];
1076 - let cli = Cli::parse_from(args);
1100 + let args = parse_args(args);
1101
1078 - assert!(BenchConfig::from_cli(&cli).is_err());
1102 + assert!(BenchConfig::from_args(&args).is_err());
1103 }
1104 #[test]
1105 fn test_range_ratio_parse_float() {
rust/src/bench/src/hub.rs
+5 -2
@@ -40,8 +40,11 @@ impl HubRepo {
40 .build()
41 .map_err(|e| format!("Failed to build download runtime: {e}"))?;
42 rt.block_on(async move {
43 - let api = hf_hub::api::tokio::Api::new()
44 - .map_err(|e| format!("Failed to init HF API: {e}"))?;
43 + let mut builder = hf_hub::api::tokio::ApiBuilder::from_env();
44 + if let Ok(token) = std::env::var("HF_TOKEN") {
45 + builder = builder.with_token(Some(token));
46 + }
47 + let api = builder.build().map_err(|e| format!("Failed to init HF API: {e}"))?;
48 api.repo(repo).get(&filename).await.map_err(|e| format!("{e}"))
49 })
50 })
rust/src/bench/src/lib.rs new
+86
@@ -0,0 +1,86 @@
1 +// SPDX-License-Identifier: Apache-2.0
2 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +mod backends;
5 +mod benchmark;
6 +mod cli;
7 +mod compare;
8 +mod config;
9 +mod datasets;
10 +mod error;
11 +mod hub;
12 +mod metrics;
13 +mod multi_run;
14 +mod multi_turn;
15 +mod output;
16 +mod rate_control;
17 +mod ready_checker;
18 +mod sweep;
19 +mod tiktoken;
20 +mod tokenizer;
21 +
22 +use anyhow::Context;
23 +
24 +pub use cli::{
25 + BackendKind, BenchServeArgs, DatasetName, LoraAssignment, RampUpStrategy, SpeedBenchConfig,
26 +};
27 +use config::BenchConfig;
28 +
29 +/// Prepare process-wide resources for a benchmark run.
30 +pub fn prepare_process() {
31 + // Raise the open-file soft limit to the hard limit. High-concurrency
32 + // benchmarks (1024+ requests) easily exceed the default 1024 fd soft limit.
33 + if let Ok(new) = rlimit::increase_nofile_limit(u64::MAX)
34 + && new > 1024
35 + {
36 + eprintln!("Open-file limit: {new}");
37 + }
38 +}
39 +
40 +/// Run the online serving benchmark.
41 +pub async fn run(args: BenchServeArgs) -> anyhow::Result<()> {
42 + // --- Compare mode: no server needed, just diff two JSON files ---
43 + if let Some(ref files) = args.compare {
44 + return compare::compare_results(&files[0], &files[1]).context("Comparison failed");
45 + }
46 +
47 + let config = BenchConfig::from_args(&args).context("Configuration error")?;
48 +
49 + async {
50 + if config.multi_turn {
51 + if let Some(ref sweep_mc) = args.sweep_max_concurrency {
52 + // --- Sweep over concurrency in multi-turn mode ---
53 + let values = sweep::parse_concurrency_values(sweep_mc)
54 + .context("Invalid --sweep-max-concurrency")?;
55 + sweep::run_multi_turn_concurrency_sweep(
56 + &config,
57 + &values,
58 + args.sweep_num_prompts_factor,
59 + )
60 + .await?;
61 + } else {
62 + // --- Single multi-turn conversation benchmark ---
63 + multi_turn::run_multi_turn_benchmark(&config).await?;
64 + }
65 + } else if let Some(ref sweep_mc) = args.sweep_max_concurrency {
66 + // --- Sweep over max-concurrency ---
67 + let values = sweep::parse_concurrency_values(sweep_mc)
68 + .context("Invalid --sweep-max-concurrency")?;
69 + sweep::run_concurrency_sweep(&config, &values, args.sweep_num_prompts_factor).await?;
70 + } else if let Some(ref sweep_rate) = args.sweep_request_rate {
71 + // --- Sweep over request-rate ---
72 + let values =
73 + sweep::parse_rate_values(sweep_rate).context("Invalid --sweep-request-rate")?;
74 + sweep::run_rate_sweep(&config, &values).await?;
75 + } else if args.num_runs > 1 {
76 + // --- Multi-run with statistical aggregation ---
77 + multi_run::run_multi(&config, args.num_runs).await?;
78 + } else {
79 + // --- Normal single benchmark ---
80 + benchmark::run_benchmark(&config).await?;
81 + }
82 + anyhow::Ok(())
83 + }
84 + .await
85 + .context("Benchmark failed")
86 +}
rust/src/bench/src/main.rs
+14 -74
@@ -1,92 +1,32 @@
1 // SPDX-License-Identifier: Apache-2.0
2 // SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 -mod backends;
5 -mod benchmark;
6 -mod cli;
7 -mod compare;
8 -mod config;
9 -mod datasets;
10 -mod error;
11 -mod hub;
12 -mod metrics;
13 -mod multi_run;
14 -mod multi_turn;
15 -mod output;
16 -mod rate_control;
17 -mod ready_checker;
18 -mod sweep;
19 -mod tiktoken;
20 -mod tokenizer;
21 -
4 #[cfg(not(target_env = "msvc"))]
5 #[global_allocator]
6 static GLOBAL: mimalloc::MiMalloc = mimalloc::MiMalloc;
7
8 use anyhow::Context;
9 use clap::Parser;
28 -use cli::Cli;
29 -use config::BenchConfig;
10
31 -fn main() -> anyhow::Result<()> {
32 - // Raise the open-file soft limit to the hard limit. High-concurrency
33 - // benchmarks (1024+ requests) easily exceed the default 1024 fd soft limit.
34 - if let Ok(new) = rlimit::increase_nofile_limit(u64::MAX)
35 - && new > 1024
36 - {
37 - eprintln!("Open-file limit: {new}");
38 - }
11 +#[derive(Parser)]
12 +#[command(
13 + name = "vllm-bench",
14 + about = "Benchmark online serving throughput",
15 + version
16 +)]
17 +struct Cli {
18 + #[command(flatten)]
19 + args: vllm_bench::BenchServeArgs,
20 +}
21
22 +fn main() -> anyhow::Result<()> {
23 let cli = Cli::parse();
41 -
42 - // --- Compare mode: no server needed, just diff two JSON files ---
43 - if let Some(ref files) = cli.compare {
44 - return compare::compare_results(&files[0], &files[1]).context("Comparison failed");
45 - }
46 -
47 - let config = BenchConfig::from_cli(&cli).context("Configuration error")?;
24 + vllm_bench::prepare_process();
25
26 let runtime = tokio::runtime::Builder::new_multi_thread()
27 .enable_all()
28 .build()
52 - .expect("Failed to build tokio runtime");
29 + .context("Failed to build tokio runtime")?;
30
54 - runtime
55 - .block_on(async {
56 - if config.multi_turn {
57 - if let Some(ref sweep_mc) = cli.sweep_max_concurrency {
58 - // --- Sweep over concurrency in multi-turn mode ---
59 - let values = sweep::parse_concurrency_values(sweep_mc)
60 - .context("Invalid --sweep-max-concurrency")?;
61 - sweep::run_multi_turn_concurrency_sweep(
62 - &config,
63 - &values,
64 - cli.sweep_num_prompts_factor,
65 - )
66 - .await?;
67 - } else {
68 - // --- Single multi-turn conversation benchmark ---
69 - multi_turn::run_multi_turn_benchmark(&config).await?;
70 - }
71 - } else if let Some(ref sweep_mc) = cli.sweep_max_concurrency {
72 - // --- Sweep over max-concurrency ---
73 - let values = sweep::parse_concurrency_values(sweep_mc)
74 - .context("Invalid --sweep-max-concurrency")?;
75 - sweep::run_concurrency_sweep(&config, &values, cli.sweep_num_prompts_factor)
76 - .await?;
77 - } else if let Some(ref sweep_rate) = cli.sweep_request_rate {
78 - // --- Sweep over request-rate ---
79 - let values =
80 - sweep::parse_rate_values(sweep_rate).context("Invalid --sweep-request-rate")?;
81 - sweep::run_rate_sweep(&config, &values).await?;
82 - } else if cli.num_runs > 1 {
83 - // --- Multi-run with statistical aggregation ---
84 - multi_run::run_multi(&config, cli.num_runs).await?;
85 - } else {
86 - // --- Normal single benchmark ---
87 - benchmark::run_benchmark(&config).await?;
88 - }
89 - anyhow::Ok(())
90 - })
91 - .context("Benchmark failed")
31 + runtime.block_on(vllm_bench::run(cli.args))
32 }
rust/src/cmd/Cargo.toml
+1
@@ -29,6 +29,7 @@ tokio-util.workspace = true
29 tracing.workspace = true
30 tracing-subscriber.workspace = true
31 uuid.workspace = true
32 +vllm-bench.workspace = true
33 vllm-chat.workspace = true
34 vllm-engine-core-client.workspace = true
35 vllm-managed-engine.workspace = true
rust/src/cmd/src/cli.rs
+11 -1
@@ -79,13 +79,23 @@ impl Cli {
79 }
80
81 /// Supported top-level CLI commands.
82 -#[derive(Debug, Subcommand, PartialEq, Eq)]
82 +#[derive(Debug, Subcommand)]
83 pub enum Command {
84 /// Run the Rust OpenAI frontend as a Python-supervised worker.
85 Frontend(FrontendArgs),
86 /// Launch a managed Python headless engine, then run the Rust OpenAI
87 /// frontend.
88 Serve(ServeArgs),
89 + /// Run vLLM benchmarks.
90 + #[command(subcommand)]
91 + Bench(BenchCommand),
92 +}
93 +
94 +/// Supported benchmark commands.
95 +#[derive(Debug, Subcommand)]
96 +pub enum BenchCommand {
97 + /// Benchmark online serving throughput.
98 + Serve(vllm_bench::BenchServeArgs),
99 }
100
101 /// A JSON-encoded list of strings, matching Python's `json.loads` CLI type for
rust/src/cmd/src/cli/tests.rs
+21 -1
@@ -5,7 +5,27 @@ use expect_test::expect;
5 use vllm_engine_core_client::TransportMode;
6 use vllm_server::{Config, HttpListenerMode, ParserSelection, RendererSelection};
7
8 -use super::{Cli, Command};
8 +use super::{BenchCommand, Cli, Command};
9 +
10 +#[test]
11 +fn bench_serve_args_parse_without_managed_engine_repartition() {
12 + let cli = Cli::try_parse_from([
13 + "vllm-rs",
14 + "bench",
15 + "serve",
16 + "--backend",
17 + "openai-chat",
18 + "--request-rate",
19 + "inf",
20 + ])
21 + .unwrap();
22 +
23 + let Command::Bench(BenchCommand::Serve(args)) = cli.command else {
24 + panic!("expected bench serve args");
25 + };
26 + assert_eq!(args.backend, vllm_bench::BackendKind::OpenaiChat);
27 + assert!(args.request_rate.is_infinite());
28 +}
29
30 #[test]
31 fn serve_args_forward_python_flags_with_separator() {
rust/src/cmd/src/main.rs
+5 -1
@@ -12,7 +12,7 @@ use tokio_util::sync::CancellationToken;
12 use tracing::{info, warn};
13 use vllm_managed_engine::ManagedEngineHandle;
14
15 -use crate::cli::{Cli, Command};
15 +use crate::cli::{BenchCommand, Cli, Command};
16
17 #[global_allocator]
18 static GLOBAL: mimalloc::MiMalloc = mimalloc::MiMalloc;
@@ -100,6 +100,10 @@ fn main() -> Result<()> {
100 async fn async_main(cli: Cli) -> Result<()> {
101 match cli.command {
102 Command::Frontend(args) => vllm_server::serve(args.into_config(), shutdown_signal()).await,
103 + Command::Bench(BenchCommand::Serve(bench_args)) => {
104 + vllm_bench::prepare_process();
105 + vllm_bench::run(bench_args).await
106 + }
107 Command::Serve(args) => {
108 let handshake_port = args.managed_engine.resolve_handshake_port()?;
109
vllm/entrypoints/cli/benchmark/serve.py
+60 -2
@@ -1,11 +1,68 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 import argparse
4 +import os
5 +import sys
6 +from pathlib import Path
7
5 -from vllm.benchmarks.serve import add_cli_args, main
8 +from vllm.benchmarks.serve import add_cli_args
9 +from vllm.benchmarks.serve import main as python_main
10 from vllm.entrypoints.cli.benchmark.base import BenchmarkSubcommandBase
11 +from vllm.logger import init_logger
12 from vllm.utils.argparse_utils import FlexibleArgumentParser
13
14 +logger = init_logger(__name__)
15 +_RUST_CLI_PATH = Path(__file__).resolve().parents[3] / "vllm-rs"
16 +_RUST_SUPPORTED_DATASETS = frozenset(
17 + {
18 + "custom",
19 + "hf",
20 + "prefix_repetition",
21 + "random",
22 + "random-mm",
23 + "random-rerank",
24 + "sharegpt",
25 + "sonnet",
26 + "speed_bench",
27 + }
28 +)
29 +_RUST_SUPPORTED_BACKENDS = frozenset(
30 + {
31 + "openai",
32 + "openai-chat",
33 + "openai-embeddings",
34 + "openai-embeddings-chat",
35 + "vllm",
36 + "vllm-pooling",
37 + "vllm-rerank",
38 + }
39 +)
40 +
41 +
42 +def _rust_unsupported_reason(args: argparse.Namespace) -> str | None:
43 + if args.dataset_name not in _RUST_SUPPORTED_DATASETS:
44 + return f"dataset {args.dataset_name!r} is not supported by the Rust benchmark"
45 + if args.backend not in _RUST_SUPPORTED_BACKENDS:
46 + return f"backend {args.backend!r} is not supported by the Rust benchmark"
47 + return None
48 +
49 +
50 +def _maybe_exec_rust_bench(args: argparse.Namespace) -> None:
51 + if reason := _rust_unsupported_reason(args):
52 + logger.info("Using Python benchmark: %s.", reason)
53 + return
54 +
55 + if not _RUST_CLI_PATH.is_file():
56 + logger.warning(
57 + "Rust benchmark binary not found at %s; falling back to Python.",
58 + _RUST_CLI_PATH,
59 + )
60 + return
61 +
62 + rust_cli = str(_RUST_CLI_PATH)
63 + logger.info("Delegating `vllm bench serve` to Rust binary at %s.", rust_cli)
64 + os.execv(rust_cli, [rust_cli, "bench", "serve", *sys.argv[3:]])
65 +
66
67 class BenchmarkServingSubcommand(BenchmarkSubcommandBase):
68 """The `serve` subcommand for `vllm bench`."""
@@ -19,4 +76,5 @@ class BenchmarkServingSubcommand(BenchmarkSubcommandBase):
76
77 @staticmethod
78 def cmd(args: argparse.Namespace) -> None:
22 - main(args)
79 + _maybe_exec_rust_bench(args)
80 + python_main(args)