iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

[Rust][Benchmark] Use tracing for logs (#48937)

Browse Source 
Signed-off-by: Bugen Zhao <i@bugenzhao.com>
This commit is contained in:
Bugen Zhao
2026-07-21 05:18:13 +00:00
committed by GitHub
parent f25953cc59
commit 8688a06d67
24 changed files with 498 additions and 249 deletions
Download Patch File Download Diff File Expand all files Collapse all files
rust/Cargo.lock
Generated
+3
View File
@@ -5498,10 +5498,13 @@ dependencies = [
"serde",
"serde_json",
"thiserror 2.0.18",
"thiserror-ext",
"tiktoken-rs 0.9.1",
"tokenizers",
"tokio",
"tokio-stream",
"tracing",
"tracing-subscriber",
"url",
"uuid",
]
rust/src/bench/Cargo.toml
+3
View File
@@ -26,10 +26,13 @@ rustc-hash.workspace = true
serde = { workspace = true, features = ["rc"] }
serde_json = { workspace = true, features = ["raw_value"] }
thiserror.workspace = true
thiserror-ext.workspace = true
tiktoken-rs.workspace = true
tokenizers.workspace = true
tokio.workspace = true
tokio-stream.workspace = true
tracing.workspace = true
tracing-subscriber.workspace = true
url.workspace = true
uuid.workspace = true
rust/src/bench/src/backends/pooling.rs
+8 -7
View File
@@ -158,9 +158,10 @@ impl PoolingBackend {
// (mirrors Python async_request_vllm_rerank).
if let Some(ref list) = input.prompt_list {
if list.len() < 2 {
eprintln!(
"WARNING: vllm-rerank request has no documents \
(prompt_list needs [query, doc, ...])"
tracing::warn!(
backend = "vllm-rerank",
inputs = list.len(),
"rerank request has no documents"
);
}
let query = list.first().map(|s| s.as_ref()).unwrap_or("");
@@ -175,10 +176,10 @@ impl PoolingBackend {
// Legacy path: text prompt as query, documents via --extra-body.
let query = input.prompt.as_ref();
if query.is_empty() && input.prompt_token_ids.is_some() {
eprintln!(
"WARNING: vllm-rerank received empty query (random dataset uses \
token IDs only). Use --dataset-name random-rerank for meaningful \
rerank benchmarks."
tracing::warn!(
backend = "vllm-rerank",
dataset = "random",
"rerank request has an empty query; use the random-rerank dataset"
);
}
serde_json::json!({
rust/src/bench/src/benchmark.rs
+107 -80
View File
@@ -6,6 +6,7 @@ use std::sync::Arc;
use std::time::Instant;
use indicatif::{ProgressBar, ProgressStyle};
use thiserror_ext::AsReport as _;
use tokio::sync::Semaphore;
use crate::backends::{RequestFuncInput, RequestFuncOutput, get_backend};
@@ -72,12 +73,12 @@ pub fn pre_resolve_dns(
v4.extend(v6);
if !v4.is_empty() {
let ips: Vec<_> = v4.iter().map(|a| a.ip()).collect();
println!("Pre-resolved {host} -> {ips:?}");
tracing::info!(host, addresses = ?ips, "pre-resolved benchmark endpoint DNS");
builder = builder.resolve_to_addrs(host, &v4);
}
}
Err(e) => {
eprintln!("Warning: DNS pre-resolution for '{host}' failed: {e}");
tracing::warn!(host, error = %e.as_report(), "failed to pre-resolve benchmark endpoint DNS");
}
}
@@ -346,10 +347,14 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result<serde_json::Value> {
let (model_id, model_name) = if let Some(ref m) = config.model {
(m.clone(), config.model_name.clone())
} else {
println!("Model not specified, fetching first model from server...");
tracing::info!(base_url = %config.base_url, "fetching first model from server");
let (name, id) =
get_first_model_from_server(&config.base_url, &client, &config.extra_headers).await?;
println!("First model name: {name}, first model id: {id}");
tracing::info!(
model_name = name,
model_id = id,
"selected first model from server"
);
(id, Some(name))
};
@@ -358,10 +363,9 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result<serde_json::Value> {
None
} else {
let tid = config.tokenizer_id.as_deref().unwrap_or(&model_id);
println!("Loading tokenizer: {tid}");
tracing::info!(tokenizer = tid, "loading tokenizer");
let server_info = Some((config.base_url.as_str(), model_id.as_str()));
let t = crate::tokenizer::load_tokenizer(tid, config.trust_remote_code, server_info)?;
println!("Tokenizer loaded successfully.");
Some(t)
};
let has_tokenizer = tokenizer.is_some();
@@ -421,7 +425,12 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result<serde_json::Value> {
config.num_prompts, config.random_batch_size, config.is_reranker,
),
};
println!("Generating {dataset_label}...");
tracing::info!(
dataset = ?config.dataset_name,
prompts = config.num_prompts,
description = %dataset_label,
"generating benchmark dataset"
);
let gen_start = Instant::now();
let mut input_requests = match config.dataset_name {
@@ -608,18 +617,19 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result<serde_json::Value> {
};
let gen_elapsed = gen_start.elapsed();
println!(
"Generated {} prompts in {:.2}s",
input_requests.len(),
gen_elapsed.as_secs_f64()
tracing::info!(
prompts = input_requests.len(),
elapsed_seconds = gen_elapsed.as_secs_f64(),
"generated benchmark dataset"
);
let filtered_count =
filter_requests_by_max_model_len(&mut input_requests, config.max_model_len);
if filtered_count > 0 {
println!(
"Filtered {filtered_count} prompt(s) above --max-model-len {}.",
config.max_model_len.unwrap()
tracing::info!(
filtered_prompts = filtered_count,
max_model_len = config.max_model_len.unwrap(),
"filtered prompts above maximum model length"
);
}
if input_requests.is_empty() {
@@ -670,7 +680,7 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result<serde_json::Value> {
// Ready check
if config.ready_check_timeout_sec > 0 {
println!("Starting initial single prompt test run...");
tracing::info!("starting initial single-prompt test run");
let test_output = wait_for_endpoint(
config.backend,
&client,
@@ -685,7 +695,7 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result<serde_json::Value> {
test_output.error
)));
}
println!("Initial test run completed.");
tracing::info!("initial single-prompt test run completed");
}
// Verify and fix prompt token lengths against the server's /tokenize endpoint.
@@ -703,12 +713,15 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result<serde_json::Value> {
DatasetName::Random | DatasetName::PrefixRepetition
);
if verifiable_dataset && has_token_ids && !config.backend.is_pooling() {
println!("Using prompt_token_ids, skipping server-side tokenizer verification.");
tracing::info!(
reason = "prompt_token_ids",
"skipping server tokenizer verification"
);
}
if verifiable_dataset && !has_token_ids && !config.backend.is_pooling() {
let cache_key = tokenizer_verify_cache_key(&config.base_url, &model_id);
if is_tokenizer_verified(&cache_key) {
println!("Tokenizer verified in previous run (cached), skipping verification.");
tracing::info!(reason = "cached", "skipping server tokenizer verification");
} else {
let num_special =
tokenizer.as_ref().map(|t| t.num_special_tokens_to_add()).unwrap_or(0);
@@ -723,14 +736,17 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result<serde_json::Value> {
.await?
{
SampleVerifyOutcome::Passed => {
println!("Sample verification passed, skipping full verification.");
tracing::info!("tokenizer sample verification passed");
mark_tokenizer_verified(&cache_key);
}
SampleVerifyOutcome::Skipped(reason) => {
println!("Server /tokenize unavailable ({reason}), skipping verification.");
tracing::warn!(
reason = %reason,
"server tokenizer unavailable; skipping prompt verification"
);
}
SampleVerifyOutcome::Mismatch => {
println!("Sample verification found mismatch, running full verify+fix...");
tracing::warn!("tokenizer sample mismatch; verifying and fixing all prompts");
match verify_and_fix_prompt_lengths(
&client,
&config.base_url,
@@ -742,16 +758,16 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result<serde_json::Value> {
.await
{
Ok(()) => {
println!(
"All {} prompts verified: exact token length match.",
input_requests.len()
tracing::info!(
prompts = input_requests.len(),
"verified exact prompt token lengths"
);
mark_tokenizer_verified(&cache_key);
}
Err(BenchError::TokenizeUnavailable(reason)) => {
println!(
"Server /tokenize became unavailable during verification \
({reason}); proceeding with client-side token counts."
tracing::warn!(
reason = %reason,
"server tokenizer became unavailable; using client token counts"
);
}
Err(e) => return Err(e),
@@ -763,7 +779,7 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result<serde_json::Value> {
// Warmup
if config.num_warmups > 0 {
println!("Warming up with {} requests...", config.num_warmups);
tracing::info!(requests = config.num_warmups, "starting benchmark warmup");
run_warmup(
config.backend,
&client,
@@ -776,7 +792,7 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result<serde_json::Value> {
config.disable_tqdm,
)
.await;
println!("Warmup run completed.");
tracing::info!(requests = config.num_warmups, "benchmark warmup completed");
}
// Start profiler if requested (immediate mode — no batch threshold)
@@ -814,28 +830,22 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result<serde_json::Value> {
let spec_decode_before =
fetch_spec_decode_metrics(&config.base_url, &client, &config.extra_headers).await;
if spec_decode_before.is_some() {
println!("Speculative decoding detected, will collect metrics.");
tracing::info!("detected speculative decoding; collecting metrics");
}
// Main benchmark
println!("Starting main benchmark run...");
let distribution = if config.burstiness == 1.0 {
"Poisson process"
} else {
"Gamma distribution"
};
println!(
"Traffic request rate: {}",
if config.request_rate.is_infinite() {
"inf".to_string()
} else {
format!("{}", config.request_rate)
}
);
println!("Burstiness factor: {} ({distribution})", config.burstiness);
println!(
"Maximum request concurrency: {}",
config.max_concurrency.unwrap_or(config.num_prompts)
tracing::info!(
request_rate = config.request_rate,
burstiness = config.burstiness,
distribution,
max_concurrency = config.max_concurrency.unwrap_or(config.num_prompts),
prompts = config.num_prompts,
"starting main benchmark run"
);
// Pre-assign LoRA adapters to each request (None when --lora-modules not set).
@@ -847,11 +857,11 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result<serde_json::Value> {
);
if let (Some(modules), Some(_)) = (config.lora_modules.as_ref(), lora_assignments.as_ref()) {
let names: Vec<&str> = modules.iter().map(|s| s.as_ref()).collect();
println!(
"LoRA adapters ({}): {:?} [assignment={:?}]",
modules.len(),
names,
config.lora_assignment
tracing::info!(
adapters = modules.len(),
names = ?names,
assignment = ?config.lora_assignment,
"assigned LoRA adapters"
);
}
@@ -1125,7 +1135,7 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result<serde_json::Value> {
if let Some((cancel_tx, task)) = profile_task {
let _ = cancel_tx.send(());
if let Err(e) = task.await {
eprintln!("WARNING: Profile background task failed: {e}");
tracing::error!(error = %e.as_report(), "profiler background task failed");
}
}
@@ -1289,12 +1299,14 @@ pub(crate) async fn start_profiler_immediate(
base_url: &str,
extra_headers: &Option<std::collections::HashMap<String, String>>,
) {
println!("Starting profiler...");
let profile_url = format!("{base_url}/start_profile");
tracing::info!(url = %profile_url, "starting profiler");
match send_profile_request(client, &profile_url, extra_headers).await {
Ok(true) => println!("Profiler started"),
Ok(false) => eprintln!("WARNING: Profiler start request returned non-success"),
Err(e) => eprintln!("WARNING: Failed to start profiler: {e}"),
Ok(true) => tracing::info!(url = %profile_url, "profiler started"),
Ok(false) => tracing::warn!(url = %profile_url, "profiler start request was unsuccessful"),
Err(e) => {
tracing::warn!(url = %profile_url, error = %e.as_report(), "failed to start profiler")
}
}
}
@@ -1304,12 +1316,14 @@ pub(crate) async fn stop_profiler_immediate(
base_url: &str,
extra_headers: &Option<std::collections::HashMap<String, String>>,
) {
println!("Stopping profiler...");
let profile_url = format!("{base_url}/stop_profile");
tracing::info!(url = %profile_url, "stopping profiler");
match send_profile_request(client, &profile_url, extra_headers).await {
Ok(true) => println!("Profiler stopped"),
Ok(false) => eprintln!("WARNING: Profiler stop request returned non-success"),
Err(e) => eprintln!("WARNING: Failed to stop profiler: {e}"),
Ok(true) => tracing::info!(url = %profile_url, "profiler stopped"),
Ok(false) => tracing::warn!(url = %profile_url, "profiler stop request was unsuccessful"),
Err(e) => {
tracing::warn!(url = %profile_url, error = %e.as_report(), "failed to stop profiler")
}
}
}
@@ -1371,25 +1385,30 @@ pub(crate) async fn profile_on_batch_threshold(
duration_secs: f64,
mut cancel_rx: tokio::sync::oneshot::Receiver<()>,
) {
println!(
"Waiting for batch size >= {threshold} before starting profiler \
(will capture {duration_secs}s)..."
tracing::info!(
threshold,
duration_seconds = duration_secs,
"waiting for profiler batch threshold"
);
loop {
if let Some(running) = fetch_num_requests_running(client, base_url).await
&& running >= threshold
{
println!("Batch size {running} >= {threshold}, starting profiler...");
tracing::info!(
running_requests = running,
threshold,
"profiler batch threshold reached"
);
break;
}
// Wait 500ms or until the benchmark signals cancellation
tokio::select! {
_ = tokio::time::sleep(std::time::Duration::from_millis(500)) => {}
_ = &mut cancel_rx => {
eprintln!(
"NOTE: Benchmark finished before batch threshold {threshold} was reached; \
profiling skipped."
tracing::warn!(
threshold,
"benchmark finished before profiler batch threshold; skipping profiling"
);
return;
}
@@ -1398,13 +1417,13 @@ pub(crate) async fn profile_on_batch_threshold(
let start_url = format!("{base_url}/start_profile");
match send_profile_request(client, &start_url, extra_headers).await {
Ok(true) => println!("Profiler started"),
Ok(true) => tracing::info!(url = %start_url, "profiler started"),
Ok(false) => {
eprintln!("WARNING: Profiler start request returned non-success");
tracing::warn!(url = %start_url, "profiler start request was unsuccessful");
return;
}
Err(e) => {
eprintln!("WARNING: Failed to start profiler: {e}");
tracing::warn!(url = %start_url, error = %e.as_report(), "failed to start profiler");
return;
}
}
@@ -1413,15 +1432,17 @@ pub(crate) async fn profile_on_batch_threshold(
tokio::select! {
_ = tokio::time::sleep(std::time::Duration::from_secs_f64(duration_secs)) => {}
_ = &mut cancel_rx => {
println!("Benchmark finished, stopping profiler early...");
tracing::info!("benchmark finished; stopping profiler early");
}
}
let stop_url = format!("{base_url}/stop_profile");
match send_profile_request(client, &stop_url, extra_headers).await {
Ok(true) => println!("Profiler stopped after capturing"),
Ok(false) => eprintln!("WARNING: Profiler stop request returned non-success"),
Err(e) => eprintln!("WARNING: Failed to stop profiler: {e}"),
Ok(true) => tracing::info!(url = %stop_url, "profiler stopped after capture"),
Ok(false) => tracing::warn!(url = %stop_url, "profiler stop request was unsuccessful"),
Err(e) => {
tracing::warn!(url = %stop_url, error = %e.as_report(), "failed to stop profiler")
}
}
}
@@ -1502,10 +1523,11 @@ async fn verify_and_fix_prompt_lengths(
let excess = tokens.len().saturating_sub(expected_input_len);
let compensate = if excess > 0 && last_excess == Some(excess) {
if _iter == 1 {
eprintln!(
"Prompt {i}: server consistently adds {excess} extra token(s) \
(likely BOS), compensating target to {}.",
expected_input_len.saturating_sub(excess),
tracing::warn!(
prompt_index = i,
extra_tokens = excess,
adjusted_target = expected_input_len.saturating_sub(excess),
"server consistently adds prompt tokens; compensating verification target"
);
}
excess
@@ -1563,7 +1585,10 @@ async fn verify_and_fix_prompt_lengths(
let fc = fixed_count.load(std::sync::atomic::Ordering::Relaxed);
if fc > 0 {
println!("Fixed {fc} prompt(s) via server tokenize/detokenize convergence.");
tracing::info!(
fixed_prompts = fc,
"fixed prompt lengths using server tokenizer"
);
}
Ok(())
@@ -1818,7 +1843,7 @@ async fn sample_verify_prompts(
let tokenize_url = format!("{base_url}/tokenize");
let api_key = std::env::var("OPENAI_API_KEY").ok();
println!("Sampling {sample_size} prompts for verification...");
tracing::info!(sample_size, "sampling prompts for tokenizer verification");
for (i, request) in requests.iter().enumerate().take(sample_size) {
let tokens = match server_tokenize(
@@ -1841,9 +1866,11 @@ async fn sample_verify_prompts(
let expected = request.prompt_len + num_special;
if tokens.len() != expected {
println!(
"Prompt {i}: expected {expected} tokens, server returned {}",
tokens.len()
tracing::warn!(
prompt_index = i,
expected_tokens = expected,
actual_tokens = tokens.len(),
"tokenizer verification sample mismatch"
);
return Ok(SampleVerifyOutcome::Mismatch);
}
rust/src/bench/src/config.rs
+18 -9
View File
@@ -288,10 +288,18 @@ impl BenchConfig {
}
Some(other) => {
// extra_body was not an object — just use sampling params
eprintln!(
"Warning: --extra-body is not a JSON object, sampling params may be lost"
let value_type = match &other {
serde_json::Value::Null => "null",
serde_json::Value::Bool(_) => "boolean",
serde_json::Value::Number(_) => "number",
serde_json::Value::String(_) => "string",
serde_json::Value::Array(_) => "array",
serde_json::Value::Object(_) => unreachable!(),
};
tracing::warn!(
value_type,
"sampling parameters may be lost because --extra-body is not a JSON object"
);
let _ = other;
sampling_params
}
None => sampling_params,
@@ -489,9 +497,9 @@ impl BenchConfig {
_ => {}
}
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)."
tracing::warn!(
dataset = "custom",
"client-side chat template rendering is unsupported; sending prompts raw"
);
}
}
@@ -570,9 +578,10 @@ impl BenchConfig {
}
if ignore_eos {
eprintln!(
"WARNING: --ignore-eos is set with --multi-turn. The server may not \
respect output length limits, causing unbounded context growth."
tracing::warn!(
ignore_eos,
multi_turn = true,
"output length limits may be ignored, causing unbounded context growth"
);
}
rust/src/bench/src/datasets/hf_dataset.rs
+59 -26
View File
@@ -8,6 +8,7 @@ use rand::seq::SliceRandom;
use rand::{Rng, SeedableRng};
use super::SampleRequest;
use super::progress::RowDownloadReporter;
use crate::error::{BenchError, Result};
use crate::tokenizer::TokenizerKind;
@@ -201,7 +202,12 @@ pub fn download_hf_dataset(
(resolved_config, resolved_split)
};
println!("HF dataset: {dataset} (config={resolved_config}, split={resolved_split})");
tracing::info!(
dataset,
config = resolved_config,
split = resolved_split,
"resolved Hugging Face dataset"
);
// Check cache
let dir = cache_dir();
@@ -215,11 +221,16 @@ pub fn download_hf_dataset(
if cache_path.exists() {
let path_str = cache_path.to_string_lossy().to_string();
println!("HF dataset cached: {path_str}");
tracing::info!(dataset, path = %path_str, "using cached Hugging Face dataset");
return Ok((path_str, resolved_config, resolved_split));
}
println!("Downloading HF dataset '{dataset}' from datasets-server...");
tracing::info!(
dataset,
config = resolved_config,
split = resolved_split,
"downloading Hugging Face dataset"
);
let encoded_config: String =
url::form_urlencoded::byte_serialize(resolved_config.as_bytes()).collect();
@@ -229,6 +240,7 @@ pub fn download_hf_dataset(
let mut all_rows: Vec<serde_json::Value> = Vec::new();
let mut offset = 0usize;
let page_size = 100usize;
let mut progress = RowDownloadReporter::new();
loop {
let url = format!(
@@ -260,14 +272,14 @@ pub fn download_hf_dataset(
offset += fetched;
let total = data["num_rows_total"].as_u64().unwrap_or(0);
eprint!("\r Fetched {offset}/{total} rows...");
progress.update(offset, total);
// Stop if we have enough rows or reached end of dataset
if all_rows.len() >= num_rows_needed || fetched < page_size {
break;
}
}
eprintln!(); // newline after progress
progress.finish();
if all_rows.is_empty() {
return Err(BenchError::Config(format!(
@@ -280,7 +292,12 @@ pub fn download_hf_dataset(
std::fs::write(&cache_path, &json_str)?;
let path_str = cache_path.to_string_lossy().to_string();
println!("HF dataset: {} rows saved to {path_str}", all_rows.len());
tracing::info!(
dataset,
rows = all_rows.len(),
path = %path_str,
"saved Hugging Face dataset"
);
Ok((path_str, resolved_config, resolved_split))
}
@@ -483,21 +500,31 @@ pub fn load_hf_dataset(
// Detect column format from first row
let format = detect_column_format(&entries[0], text_column_override)?;
// Print detected format
match &format {
ColumnFormat::Chat(col) => println!("HF dataset: detected chat column '{col}'"),
ColumnFormat::Chat(col) => {
tracing::info!(
format = "chat",
column = col,
"detected Hugging Face dataset format"
);
}
ColumnFormat::Text {
prompt_col,
output_col,
} => {
let out_msg = output_col.as_deref().unwrap_or("none");
println!("HF dataset: detected text column '{prompt_col}', output column: {out_msg}");
tracing::info!(
format = "text",
prompt_column = prompt_col,
output_column = output_col.as_deref().unwrap_or("none"),
"detected Hugging Face dataset format"
);
}
ColumnFormat::Combined { cols, output_col } => {
let out_msg = output_col.as_deref().unwrap_or("none");
println!(
"HF dataset: detected combined columns {:?}, output column: {out_msg}",
cols
tracing::info!(
format = "combined",
prompt_columns = ?cols,
output_column = output_col.as_deref().unwrap_or("none"),
"detected Hugging Face dataset format"
);
}
}
@@ -608,9 +635,10 @@ pub fn load_hf_dataset(
if len == 0 { 128 } else { len }
} else {
if !warned_no_output {
eprintln!(
"WARNING: No output column detected and --hf-output-len not set. \
Using default output length of 128 tokens."
tracing::warn!(
path = dataset_path,
default_output_tokens = 128,
"no dataset output column or --hf-output-len; using default output length"
);
warned_no_output = true;
}
@@ -618,9 +646,10 @@ pub fn load_hf_dataset(
}
} else {
if !warned_no_output {
eprintln!(
"WARNING: No output column detected and --hf-output-len not set. \
Using default output length of 128 tokens."
tracing::warn!(
path = dataset_path,
default_output_tokens = 128,
"no dataset output column or --hf-output-len; using default output length"
);
warned_no_output = true;
}
@@ -640,9 +669,11 @@ pub fn load_hf_dataset(
// Oversample if needed
if samples.len() < num_requests {
if no_oversample {
println!(
"Skipping oversampling. Total samples: {} (requested: {num_requests})",
samples.len()
tracing::info!(
dataset = "hf",
samples = samples.len(),
requested = num_requests,
"skipping dataset oversampling"
);
} else if !samples.is_empty() {
let original_len = samples.len();
@@ -652,9 +683,11 @@ pub fn load_hf_dataset(
req.request_id = Some(format!("{request_id_prefix}{}", original_len + i));
samples.push(req);
}
println!(
"Oversampled HF dataset from {original_len} to {} total samples.",
samples.len()
tracing::info!(
dataset = "hf",
original_samples = original_len,
samples = samples.len(),
"oversampled dataset"
);
}
}
rust/src/bench/src/datasets/mod.rs
+9 -6
View File
@@ -5,6 +5,7 @@ pub mod custom;
pub mod hf_dataset;
pub mod multi_turn;
pub mod prefix_repetition;
mod progress;
pub mod random;
pub mod random_mm;
pub mod random_rerank;
@@ -90,9 +91,10 @@ pub fn oversample_requests(
return;
}
if no_oversample {
println!(
"Skipping oversampling. Total samples: {} (requested: {num_requests})",
requests.len()
tracing::info!(
samples = requests.len(),
requested = num_requests,
"skipping dataset oversampling"
);
return;
}
@@ -103,9 +105,10 @@ pub fn oversample_requests(
req.request_id = Some(format!("{request_id_prefix}{}", original_len + i));
requests.push(req);
}
println!(
"Oversampled requests from {original_len} to {} total samples.",
requests.len()
tracing::info!(
original_samples = original_len,
samples = requests.len(),
"oversampled dataset"
);
}
rust/src/bench/src/datasets/multi_turn.rs
+4 -3
View File
@@ -445,9 +445,10 @@ pub fn load_sharegpt_multi_turn(
conv.conversation_id = format!("{request_id_prefix}conv-{}", original_len + i);
conversations.push(conv);
}
println!(
"Oversampled multi-turn conversations from {original_len} to {} total.",
conversations.len()
tracing::info!(
original_conversations = original_len,
conversations = conversations.len(),
"oversampled multi-turn conversations"
);
}
rust/src/bench/src/datasets/prefix_repetition.rs
+7 -5
View File
@@ -41,11 +41,13 @@ pub fn generate_prefix_repetition_dataset(
}
let total = prompts_per_prefix * num_prefixes;
if total != num_requests {
println!(
"prefix_repetition: generating {total} requests \
({num_prefixes} prefixes x {prompts_per_prefix} prompts each; \
{} dropped to divide evenly)",
num_requests - total
tracing::info!(
requested = num_requests,
generated = total,
prefixes = num_prefixes,
prompts_per_prefix,
dropped = num_requests - total,
"adjusted prefix-repetition request count"
);
}
rust/src/bench/src/datasets/progress.rs
+77
View File
@@ -0,0 +1,77 @@
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
use std::time::{Duration, Instant};
use indicatif::{ProgressBar, ProgressStyle};
const REPORT_INTERVAL: Duration = Duration::from_secs(10);
/// Reports row download progress to an interactive progress bar, or through
/// periodic tracing events when the progress bar is hidden on a non-TTY.
pub(super) struct RowDownloadReporter {
progress: ProgressBar,
next_report: Instant,
}
impl RowDownloadReporter {
/// Creates a reporter that emits non-TTY updates every 10 seconds.
pub fn new() -> Self {
let progress = ProgressBar::new(0);
progress.set_style(
ProgressStyle::with_template(
"{spinner:.green} Fetching rows [{bar:30.cyan/blue}] {pos}/{len}",
)
.unwrap()
.progress_chars("#>-"),
);
Self {
progress,
next_report: Instant::now() + REPORT_INTERVAL,
}
}
/// Updates the current row count and reports progress when due.
pub fn update(&mut self, rows: usize, total: u64) {
let rows = rows as u64;
let total = total.max(rows);
self.progress.set_length(total);
self.progress.set_position(rows);
if self.should_report(Instant::now()) {
tracing::info!(rows, total, "fetching dataset rows");
}
}
/// Clears the interactive progress bar after the download completes.
pub fn finish(self) {
self.progress.finish_and_clear();
}
fn should_report(&mut self, now: Instant) -> bool {
if !self.progress.is_hidden() || now < self.next_report {
return false;
}
self.next_report = now + REPORT_INTERVAL;
true
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn hidden_reporter_uses_ten_second_deadline() {
let start = Instant::now();
let mut reporter = RowDownloadReporter {
progress: ProgressBar::hidden(),
next_report: start + REPORT_INTERVAL,
};
assert!(!reporter.should_report(start + Duration::from_secs(9)));
assert!(reporter.should_report(start + Duration::from_secs(10)));
assert!(!reporter.should_report(start + Duration::from_secs(19)));
assert!(reporter.should_report(start + Duration::from_secs(20)));
}
}
rust/src/bench/src/datasets/random.rs
+6 -3
View File
@@ -49,9 +49,12 @@ pub fn generate_random_dataset(
let (input_low, input_high) = range_ratio.input_bounds(real_input_len);
let (output_low, output_high) = range_ratio.output_bounds(output_len);
if !range_ratio.is_fixed() {
println!(
"Sampling input_len from [{input_low}, {input_high}] and \
output_len from [{output_low}, {output_high}]"
tracing::info!(
input_low,
input_high,
output_low,
output_high,
"sampling random request lengths"
);
}
rust/src/bench/src/datasets/sharegpt.rs
+15 -9
View File
@@ -23,8 +23,10 @@ const DEFAULT_SHAREGPT_FILE: &str = "ShareGPT_V3_unfiltered_cleaned_split.json";
/// Download the default ShareGPT dataset from HuggingFace Hub.
/// Uses hf-hub's built-in cache — subsequent calls return the cached path instantly.
pub fn download_sharegpt_dataset() -> Result<String> {
println!(
"Downloading ShareGPT dataset from {DEFAULT_SHAREGPT_REPO}/{DEFAULT_SHAREGPT_FILE} ..."
tracing::info!(
repository = DEFAULT_SHAREGPT_REPO,
file = DEFAULT_SHAREGPT_FILE,
"downloading ShareGPT dataset"
);
let repo = crate::hub::HubRepo::dataset(DEFAULT_SHAREGPT_REPO.to_string());
let path = repo.get(DEFAULT_SHAREGPT_FILE).map_err(|e| {
@@ -33,7 +35,7 @@ pub fn download_sharegpt_dataset() -> Result<String> {
))
})?;
let path_str = path.to_string_lossy().to_string();
println!("ShareGPT dataset ready: {path_str}");
tracing::info!(dataset = "sharegpt", path = %path_str, "dataset is ready");
Ok(path_str)
}
@@ -135,9 +137,11 @@ pub fn load_sharegpt_dataset(
// Oversample if dataset is smaller than requested
if samples.len() < num_requests {
if no_oversample {
println!(
"Skipping oversampling. Total samples: {} (requested: {num_requests})",
samples.len()
tracing::info!(
dataset = "sharegpt",
samples = samples.len(),
requested = num_requests,
"skipping dataset oversampling"
);
} else if !samples.is_empty() {
let needed = num_requests - samples.len();
@@ -147,9 +151,11 @@ pub fn load_sharegpt_dataset(
req.request_id = Some(format!("{request_id_prefix}{}", original_len + i));
samples.push(req);
}
println!(
"Oversampled requests from {original_len} to {} total samples.",
samples.len()
tracing::info!(
dataset = "sharegpt",
original_samples = original_len,
samples = samples.len(),
"oversampled dataset"
);
}
}
rust/src/bench/src/datasets/speed_bench.rs
+22 -16
View File
@@ -8,6 +8,7 @@ use rand::seq::SliceRandom;
use rand::{Rng, SeedableRng};
use super::SampleRequest;
use super::progress::RowDownloadReporter;
use crate::cli::SpeedBenchConfig;
use crate::error::{BenchError, Result};
use crate::tokenizer::TokenizerKind;
@@ -35,11 +36,11 @@ pub fn download_speed_bench(config: SpeedBenchConfig) -> Result<String> {
// Return cached file if it exists
if cache_path.exists() {
let path_str = cache_path.to_string_lossy().to_string();
println!("SPEED-Bench ({config_name}) cached: {path_str}");
tracing::info!(config = config_name, path = %path_str, "using cached SPEED-Bench dataset");
return Ok(path_str);
}
println!("Downloading SPEED-Bench ({config_name}) from HuggingFace datasets-server...");
tracing::info!(config = config_name, "downloading SPEED-Bench dataset");
let client = reqwest::blocking::Client::builder()
.timeout(std::time::Duration::from_secs(120))
@@ -49,6 +50,7 @@ pub fn download_speed_bench(config: SpeedBenchConfig) -> Result<String> {
let mut all_rows: Vec<serde_json::Value> = Vec::new();
let mut offset = 0usize;
let page_size = 100usize;
let mut progress = RowDownloadReporter::new();
loop {
let url = format!(
@@ -116,15 +118,14 @@ pub fn download_speed_bench(config: SpeedBenchConfig) -> Result<String> {
let fetched = rows.len();
offset += fetched;
// Print progress
let total = data["num_rows_total"].as_u64().unwrap_or(0);
eprint!("\r Fetched {offset}/{total} rows...");
progress.update(offset, total);
if fetched < page_size {
break;
}
}
eprintln!(); // newline after progress
progress.finish();
if all_rows.is_empty() {
return Err(BenchError::Config(
@@ -137,9 +138,11 @@ pub fn download_speed_bench(config: SpeedBenchConfig) -> Result<String> {
std::fs::write(&cache_path, &json_str)?;
let path_str = cache_path.to_string_lossy().to_string();
println!(
"SPEED-Bench ({config_name}): {} rows saved to {path_str}",
all_rows.len()
tracing::info!(
config = config_name,
rows = all_rows.len(),
path = %path_str,
"saved SPEED-Bench dataset"
);
Ok(path_str)
}
@@ -263,9 +266,11 @@ pub fn load_speed_bench_dataset(
// Oversample if needed
if samples.len() < num_requests {
if no_oversample {
println!(
"Skipping oversampling. Total samples: {} (requested: {num_requests})",
samples.len()
tracing::info!(
dataset = "speed-bench",
samples = samples.len(),
requested = num_requests,
"skipping dataset oversampling"
);
} else if !samples.is_empty() {
let original_len = samples.len();
@@ -275,9 +280,11 @@ pub fn load_speed_bench_dataset(
req.request_id = Some(format!("{request_id_prefix}{}", original_len + i));
samples.push(req);
}
println!(
"Oversampled SPEED-Bench from {original_len} to {} total samples.",
samples.len()
tracing::info!(
dataset = "speed-bench",
original_samples = original_len,
samples = samples.len(),
"oversampled dataset"
);
}
}
@@ -288,7 +295,6 @@ pub fn load_speed_bench_dataset(
));
}
// Print category distribution
let mut cat_counts: std::collections::HashMap<&str, usize> = std::collections::HashMap::new();
for entry in &filtered[..filtered.len().min(samples.len())] {
let cat = entry.get("category").and_then(|c| c.as_str()).unwrap_or("unknown");
@@ -297,7 +303,7 @@ pub fn load_speed_bench_dataset(
let mut cats: Vec<_> = cat_counts.into_iter().collect();
cats.sort_by_key(|b| std::cmp::Reverse(b.1));
let cat_str: Vec<String> = cats.iter().map(|(k, v)| format!("{k}:{v}")).collect();
println!("SPEED-Bench categories: {}", cat_str.join(", "));
tracing::info!(categories = %cat_str.join(", "), "computed SPEED-Bench category distribution");
Ok(samples)
}
rust/src/bench/src/lib.rs
+1 -1
View File
@@ -33,7 +33,7 @@ pub fn prepare_process() {
if let Ok(new) = rlimit::increase_nofile_limit(u64::MAX)
&& new > 1024
{
eprintln!("Open-file limit: {new}");
tracing::info!(soft_limit = new, "raised open-file limit");
}
}
rust/src/bench/src/main.rs
+12
View File
@@ -19,7 +19,19 @@ struct Cli {
args: vllm_bench::BenchServeArgs,
}
// TODO: unify the tracing subscriber used by different binaries.
fn init_tracing() {
let filter = tracing_subscriber::EnvFilter::try_from_default_env()
.unwrap_or_else(|_| tracing_subscriber::EnvFilter::new("info"));
let _ = tracing_subscriber::fmt()
.with_env_filter(filter)
.with_writer(std::io::stderr)
.try_init();
}
fn main() -> anyhow::Result<()> {
init_tracing();
let cli = Cli::parse();
vllm_bench::prepare_process();
rust/src/bench/src/metrics/calculator.rs
+18 -16
View File
@@ -7,6 +7,22 @@ use crate::datasets::SampleRequest;
use crate::metrics::{BenchmarkMetrics, MultiTurnMetrics};
use crate::multi_turn::ConversationOutput;
fn log_failed_requests(outputs: &[RequestFuncOutput]) {
let failed_outputs: Vec<_> = outputs.iter().filter(|output| !output.success).collect();
if failed_outputs.is_empty() {
return;
}
tracing::warn!(
failed_requests = failed_outputs.len(),
displayed_errors = failed_outputs.len().min(10),
"benchmark requests failed"
);
for (index, output) in failed_outputs.into_iter().take(10).enumerate() {
tracing::warn!(index, error = %output.error, "benchmark request failed");
}
}
/// Calculate benchmark metrics from request outputs.
///
/// Mirrors Python's `calculate_metrics()` from serve.py:392-599.
@@ -63,14 +79,7 @@ pub fn calculate_metrics(
let failed = outputs.len() - completed;
// Print failed request errors (capped to 10)
let failed_outputs: Vec<&RequestFuncOutput> = outputs.iter().filter(|o| !o.success).collect();
if !failed_outputs.is_empty() {
eprintln!("Failed requests during benchmark run detected (capping to 10):");
for (i, err) in failed_outputs.iter().take(10).enumerate() {
eprintln!("Error {i}: {}", err.error);
}
}
log_failed_requests(outputs);
// Calculate max output tokens per second and max concurrent requests
let mut max_output_tokens_per_s = 0.0_f64;
@@ -295,14 +304,7 @@ pub fn calculate_embedding_metrics(
let failed = outputs.len() - completed;
// Print failed request errors (capped to 10)
let failed_outputs: Vec<&RequestFuncOutput> = outputs.iter().filter(|o| !o.success).collect();
if !failed_outputs.is_empty() {
eprintln!("Failed requests during benchmark run detected (capping to 10):");
for (i, err) in failed_outputs.iter().take(10).enumerate() {
eprintln!("Error {i}: {}", err.error);
}
}
log_failed_requests(outputs);
// Compute peak concurrent requests from start_time + latency windows
let successful_outputs: Vec<&RequestFuncOutput> =
rust/src/bench/src/multi_turn.rs
+50 -37
View File
@@ -6,6 +6,7 @@ use std::sync::Arc;
use std::time::Instant;
use indicatif::{ProgressBar, ProgressStyle};
use thiserror_ext::AsReport as _;
use tokio::sync::Semaphore;
use crate::backends::{Backend, RequestFuncInput, RequestFuncOutput, get_backend};
@@ -72,9 +73,13 @@ pub async fn run_multi_turn_benchmark(config: &BenchConfig) -> Result<serde_json
let (model_id, model_name) = if let Some(ref m) = config.model {
(m.clone(), config.model_name.clone())
} else {
println!("Model not specified, fetching first model from server...");
tracing::info!(base_url = %config.base_url, "fetching first model from server");
let (name, id) = get_first_model(&config.base_url, &client, &config.extra_headers).await?;
println!("First model name: {name}, first model id: {id}");
tracing::info!(
model_name = name,
model_id = id,
"selected first model from server"
);
(id, Some(name))
};
@@ -83,15 +88,18 @@ pub async fn run_multi_turn_benchmark(config: &BenchConfig) -> Result<serde_json
None
} else {
let tid = config.tokenizer_id.as_deref().unwrap_or(&model_id);
println!("Loading tokenizer: {tid}");
tracing::info!(tokenizer = tid, "loading tokenizer");
let server_info = Some((config.base_url.as_str(), model_id.as_str()));
let t = crate::tokenizer::load_tokenizer(tid, config.trust_remote_code, server_info)?;
println!("Tokenizer loaded successfully.");
Some(t)
};
// Generate/load conversations
println!("Generating multi-turn conversations...");
tracing::info!(
dataset = ?config.dataset_name,
conversations = config.num_prompts,
"generating multi-turn conversations"
);
let gen_start = Instant::now();
let mut conversations = match config.dataset_name {
@@ -179,8 +187,11 @@ pub async fn run_multi_turn_benchmark(config: &BenchConfig) -> Result<serde_json
let (filtered_conversations, filtered_turns) =
filter_turns_by_max_model_len(&mut conversations, max_model_len, no_history);
if filtered_turns > 0 || filtered_conversations > 0 {
println!(
"Filtered {filtered_turns} turn(s) and {filtered_conversations} conversation(s) above --max-model-len {max_model_len}."
tracing::info!(
filtered_turns,
filtered_conversations,
max_model_len,
"filtered conversations above maximum model length"
);
}
if conversations.is_empty() {
@@ -192,11 +203,11 @@ pub async fn run_multi_turn_benchmark(config: &BenchConfig) -> Result<serde_json
let gen_elapsed = gen_start.elapsed();
let total_turns: usize = conversations.iter().map(|c| c.turns.len()).sum();
println!(
"Generated {} conversations ({} total turns) in {:.2}s",
conversations.len(),
tracing::info!(
conversations = conversations.len(),
total_turns,
gen_elapsed.as_secs_f64()
elapsed_seconds = gen_elapsed.as_secs_f64(),
"generated multi-turn conversations"
);
// Log prefix sharing info
@@ -208,18 +219,15 @@ pub async fn run_multi_turn_benchmark(config: &BenchConfig) -> Result<serde_json
let conv_tokens =
(real_input_len as f64 * config.multi_turn_prefix_conversation_ratio).floor() as usize;
let unique_tokens = real_input_len.saturating_sub(global_tokens + conv_tokens);
println!(
"User message prefix sharing: {:.0}% global ({} tokens), {:.0}% per-conversation ({} tokens), {:.0}% unique ({} tokens)",
config.multi_turn_prefix_global_ratio * 100.0,
tracing::info!(
global_ratio = config.multi_turn_prefix_global_ratio,
global_tokens,
config.multi_turn_prefix_conversation_ratio * 100.0,
conv_tokens,
(1.0 - config.multi_turn_prefix_global_ratio
- config.multi_turn_prefix_conversation_ratio)
* 100.0,
conversation_ratio = config.multi_turn_prefix_conversation_ratio,
conversation_tokens = conv_tokens,
unique_tokens,
history_accumulation = false,
"configured multi-turn prefix sharing"
);
println!("No history accumulation: each turn sends fixed-length prompt only.");
}
if config.dry_run {
@@ -253,7 +261,7 @@ pub async fn run_multi_turn_benchmark(config: &BenchConfig) -> Result<serde_json
..Default::default()
};
println!("Starting initial single prompt test run...");
tracing::info!("starting initial single-prompt test run");
let test_output = crate::ready_checker::wait_for_endpoint(
config.backend,
&client,
@@ -268,7 +276,7 @@ pub async fn run_multi_turn_benchmark(config: &BenchConfig) -> Result<serde_json
test_output.error
)));
}
println!("Initial test run completed.");
tracing::info!("initial single-prompt test run completed");
}
// For random datasets in multi-turn mode, auto-set min_tokens to enforce
@@ -283,9 +291,10 @@ pub async fn run_multi_turn_benchmark(config: &BenchConfig) -> Result<serde_json
"min_tokens".to_string(),
serde_json::json!(config.random_output_len),
);
println!(
"Auto-setting min_tokens={} for multi-turn random dataset (use --extra-body to override)",
config.random_output_len
tracing::info!(
min_tokens = config.random_output_len,
dataset = "random",
"set minimum output tokens for multi-turn dataset"
);
}
Some(body)
@@ -297,7 +306,7 @@ pub async fn run_multi_turn_benchmark(config: &BenchConfig) -> Result<serde_json
let spec_decode_before =
fetch_spec_decode_metrics(&config.base_url, &client, &config.extra_headers).await;
if spec_decode_before.is_some() {
println!("Speculative decoding detected, will collect metrics.");
tracing::info!("detected speculative decoding; collecting metrics");
}
// Start profiler if requested (immediate mode — no batch threshold)
@@ -330,10 +339,13 @@ pub async fn run_multi_turn_benchmark(config: &BenchConfig) -> Result<serde_json
};
// Main benchmark
println!("Starting multi-turn benchmark...");
println!("Conversations: {}", conversations.len());
println!("Concurrency: {concurrency}");
println!("Inter-turn delay: {} ms", config.multi_turn_delay_ms);
tracing::info!(
conversations = conversations.len(),
total_turns,
concurrency,
inter_turn_delay_ms = config.multi_turn_delay_ms,
"starting multi-turn benchmark"
);
let max_turn_count = conversations.iter().map(|c| c.turns.len()).max().unwrap_or(0);
@@ -364,11 +376,12 @@ pub async fn run_multi_turn_benchmark(config: &BenchConfig) -> Result<serde_json
);
if let Some(modules) = config.lora_modules.as_ref() {
let names: Vec<&str> = modules.iter().map(|s| s.as_ref()).collect();
println!(
"LoRA adapters ({}): {:?} [assignment={:?}, scope=conversation]",
modules.len(),
names,
config.lora_assignment
tracing::info!(
adapters = modules.len(),
names = ?names,
assignment = ?config.lora_assignment,
scope = "conversation",
"assigned LoRA adapters"
);
}
@@ -433,7 +446,7 @@ pub async fn run_multi_turn_benchmark(config: &BenchConfig) -> Result<serde_json
match handle.await {
Ok(output) => all_outputs.push(output),
Err(e) => {
eprintln!("Conversation task panicked: {e}");
tracing::error!(error = %e.as_report(), "conversation task panicked");
}
}
}
@@ -453,7 +466,7 @@ pub async fn run_multi_turn_benchmark(config: &BenchConfig) -> Result<serde_json
if let Some((cancel_tx, task)) = profile_task {
let _ = cancel_tx.send(());
if let Err(e) = task.await {
eprintln!("WARNING: Profile background task failed: {e}");
tracing::error!(error = %e.as_report(), "profiler background task failed");
}
}
rust/src/bench/src/output/json.rs
+2 -2
View File
@@ -603,7 +603,7 @@ fn add_metric_stats(
pub fn save_result(json: &Value, file_path: &str) -> Result<()> {
let content = serde_json::to_string(json)?;
std::fs::write(file_path, content)?;
println!("Results saved to {file_path}");
tracing::info!(path = file_path, "saved benchmark results");
Ok(())
}
@@ -618,7 +618,7 @@ pub fn append_result(json: &Value, file_path: &str) -> Result<()> {
file.write_all(b"\n")?;
}
file.write_all(content.as_bytes())?;
println!("Results appended to {file_path}");
tracing::info!(path = file_path, "appended benchmark results");
Ok(())
}
rust/src/bench/src/ready_checker.rs
+8 -2
View File
@@ -23,7 +23,11 @@ pub async fn wait_for_endpoint(
let backend = get_backend(backend)?;
let deadline = Instant::now() + std::time::Duration::from_secs(timeout_seconds);
println!("Waiting for endpoint to become up in {timeout_seconds}s");
tracing::info!(
timeout_seconds,
retry_interval,
"waiting for endpoint readiness"
);
let pb = ProgressBar::new(timeout_seconds);
pb.set_style(
@@ -53,7 +57,9 @@ pub async fn wait_for_endpoint(
Ok(output) => {
let err = output.error.clone();
let err_last_line = err.lines().last().unwrap_or(&err);
eprintln!("Endpoint is not ready. Error='{err_last_line}'");
pb.suspend(|| {
tracing::warn!(error = err_last_line, "endpoint is not ready");
});
last_error = err;
}
Err(e) => {
rust/src/bench/src/sweep.rs
+1 -1
View File
@@ -16,7 +16,7 @@ async fn reset_prefix_cache(base_url: &str) -> Result<()> {
.await
.map_err(|e| BenchError::Backend(format!("Failed to reset prefix cache: {e}")))?;
if resp.status().is_success() {
println!("Prefix cache reset successfully.");
tracing::info!(url = %url, "reset prefix cache");
} else {
let status = resp.status();
let body = resp.text().await.unwrap_or_default();
rust/src/bench/src/tiktoken.rs
+15 -9
View File
@@ -199,7 +199,12 @@ pub fn load_builtin_tiktoken(encoding: &str) -> Result<TiktokenTokenizer> {
}
};
let bpe = bpe.map_err(|e| BenchError::Tokenizer(format!("Failed to load {encoding}: {e}")))?;
println!("Tokenizer: Built-in tiktoken {encoding} (vocab_size={vocab_size})");
tracing::info!(
encoding,
kind = "built-in-tiktoken",
vocab_size,
"loaded tokenizer"
);
Ok(TiktokenTokenizer::from_builtin_bpe(bpe, vocab_size))
}
@@ -309,15 +314,16 @@ fn build_tiktoken(
}
}
println!(
"Loading tiktoken model for '{model_id}' (base={}, special={}, pat={})...",
num_base_tokens,
all_special_tokens.len(),
if pattern.is_some() {
tracing::info!(
model = model_id,
base_tokens = num_base_tokens,
special_tokens = all_special_tokens.len(),
pattern = if pattern.is_some() {
"custom"
} else {
"default"
},
"loading tiktoken model"
);
TiktokenTokenizer::from_file(
@@ -438,9 +444,9 @@ fn extract_pat_str_from_source(source: &str) -> Option<String> {
if !fragments.is_empty() {
let pattern = fragments.join("|");
println!(
"Extracted pat_str from Python source: {} fragments",
fragments.len()
tracing::debug!(
fragments = fragments.len(),
"extracted tiktoken pattern from Python source"
);
return Some(pattern);
}
rust/src/bench/src/tokenizer.rs
+27 -9
View File
@@ -4,6 +4,7 @@
use std::collections::HashSet;
use std::path::Path;
use thiserror_ext::AsReport as _;
use tokenizers::Tokenizer;
use crate::error::{BenchError, Result};
@@ -214,29 +215,46 @@ pub fn load_tokenizer(
// 1. Try local HuggingFace tokenizer (tokenizer.json)
match try_load_local(model_id) {
Ok(tok) => {
println!("Tokenizer: Local (vocab_size={})", tok.get_vocab_size(true));
tracing::info!(
model = model_id,
kind = "local",
vocab_size = tok.get_vocab_size(true),
"loaded tokenizer"
);
Ok(TokenizerKind::Local(Box::new(tok)))
}
Err(local_err) => {
// 2. Try tiktoken format
println!("No tokenizer.json for '{model_id}', trying tiktoken format...");
tracing::info!(
model = model_id,
error = %local_err.as_report(),
"local tokenizer unavailable; trying tiktoken"
);
match crate::tiktoken::try_load_tiktoken(model_id) {
Ok(tok) => {
println!("Tokenizer: Tiktoken (vocab_size={})", tok.vocab_size());
tracing::info!(
model = model_id,
kind = "tiktoken",
vocab_size = tok.vocab_size(),
"loaded tokenizer"
);
Ok(TokenizerKind::Tiktoken(tok))
}
Err(tiktoken_err) => {
// 3. Try server-side fallback
if let Some((base_url, model)) = server_info {
println!(
"Tiktoken also not available ({tiktoken_err}), \
trying server-side tokenization..."
tracing::info!(
model = model_id,
error = %tiktoken_err.as_report(),
"tiktoken unavailable; trying server-side tokenization"
);
match ServerTokenizer::new(base_url, model) {
Ok(srv) => {
println!(
"Tokenizer: Server (vocab_size≈{})",
srv.cached_vocab_size
tracing::info!(
model = model_id,
kind = "server",
vocab_size = srv.cached_vocab_size,
"loaded tokenizer"
);
return Ok(TokenizerKind::Server(srv));
}
rust/src/cmd/src/logging.rs
+17 -7
View File
@@ -26,18 +26,21 @@ const RESET: &str = "\x1b[0m";
const VLLM_TIME_FORMAT: &[time::format_description::FormatItem<'static>] =
format_description!("[month]-[day] [hour]:[minute]:[second]");
const PROCESS_LABEL: &str = "RustFrontend";
/// Install the process-wide vLLM-style tracing subscriber for the CLI binary.
pub(crate) fn init_tracing() {
pub(crate) fn init_tracing(process_label: &str) {
let filter = build_targets_filter(
env::var("VLLM_LOGGING_LEVEL").ok().as_deref(),
env::var("RUST_LOG").ok().as_deref(),
);
let formatter = VllmEventFormatter::new();
let formatter = VllmEventFormatter::new(process_label);
let _ = tracing_subscriber::registry()
.with(tracing_subscriber::fmt::layer().event_format(formatter).with_filter(filter))
.with(
tracing_subscriber::fmt::layer()
.event_format(formatter)
.with_writer(std::io::stderr)
.with_filter(filter),
)
.try_init();
}
@@ -94,9 +97,9 @@ struct VllmEventFormatter {
}
impl VllmEventFormatter {
fn new() -> Self {
fn new(process_label: &str) -> Self {
Self {
prefix: format!("({} pid={})", PROCESS_LABEL, process::id()),
prefix: format!("({process_label} pid={})", process::id()),
timer: VllmLocalTimer::default(),
}
}
@@ -291,6 +294,13 @@ fn map_python_log_level(level: &str) -> LevelFilter {
mod tests {
use super::*;
#[test]
fn formatter_prefix_uses_process_label() {
let formatter = VllmEventFormatter::new("Bench");
assert_eq!(formatter.prefix, format!("(Bench pid={})", process::id()));
}
#[test]
fn rust_log_target_overrides_are_merged_with_vllm_default_level() {
let filter = build_targets_filter(Some("DEBUG"), Some("hyper=warn,tower=error"));
rust/src/cmd/src/main.rs
+9 -1
View File
@@ -5,6 +5,7 @@ mod cli;
mod logging;
use std::env;
use std::ffi::OsStr;
use std::process::ExitStatus;
use anyhow::{Context, Result, anyhow, bail};
@@ -82,7 +83,14 @@ fn shutdown_signal() -> CancellationToken {
}
fn main() -> Result<()> {
logging::init_tracing();
let process_label =
match env::args_os().nth(1).as_deref().and_then(OsStr::to_str).unwrap_or_default() {
"bench" => "Bench",
"serve" | "frontend" => "RustFrontend",
_ => "Rust",
};
logging::init_tracing(process_label);
let cli = Cli::parse();
let mut runtime = tokio::runtime::Builder::new_multi_thread();