From 8688a06d676bda6e9e3ec3a87233dce91c40bb4f Mon Sep 17 00:00:00 2001 From: Bugen Zhao Date: Tue, 21 Jul 2026 13:18:13 +0800 Subject: [PATCH] [Rust][Benchmark] Use `tracing` for logs (#48937) Signed-off-by: Bugen Zhao --- rust/Cargo.lock | 3 + rust/src/bench/Cargo.toml | 3 + rust/src/bench/src/backends/pooling.rs | 15 +- rust/src/bench/src/benchmark.rs | 187 ++++++++++-------- rust/src/bench/src/config.rs | 27 ++- rust/src/bench/src/datasets/hf_dataset.rs | 85 +++++--- rust/src/bench/src/datasets/mod.rs | 15 +- rust/src/bench/src/datasets/multi_turn.rs | 7 +- .../bench/src/datasets/prefix_repetition.rs | 12 +- rust/src/bench/src/datasets/progress.rs | 77 ++++++++ rust/src/bench/src/datasets/random.rs | 9 +- rust/src/bench/src/datasets/sharegpt.rs | 24 ++- rust/src/bench/src/datasets/speed_bench.rs | 38 ++-- rust/src/bench/src/lib.rs | 2 +- rust/src/bench/src/main.rs | 12 ++ rust/src/bench/src/metrics/calculator.rs | 34 ++-- rust/src/bench/src/multi_turn.rs | 87 ++++---- rust/src/bench/src/output/json.rs | 4 +- rust/src/bench/src/ready_checker.rs | 10 +- rust/src/bench/src/sweep.rs | 2 +- rust/src/bench/src/tiktoken.rs | 24 ++- rust/src/bench/src/tokenizer.rs | 36 +++- rust/src/cmd/src/logging.rs | 24 ++- rust/src/cmd/src/main.rs | 10 +- 24 files changed, 498 insertions(+), 249 deletions(-) create mode 100644 rust/src/bench/src/datasets/progress.rs diff --git a/rust/Cargo.lock b/rust/Cargo.lock index 0902898f36..795426f67d 100644 --- a/rust/Cargo.lock +++ b/rust/Cargo.lock @@ -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", ] diff --git a/rust/src/bench/Cargo.toml b/rust/src/bench/Cargo.toml index 22c8f84418..22f5608884 100644 --- a/rust/src/bench/Cargo.toml +++ b/rust/src/bench/Cargo.toml @@ -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 diff --git a/rust/src/bench/src/backends/pooling.rs b/rust/src/bench/src/backends/pooling.rs index 6a1cdf4fe1..8f9d4deab9 100644 --- a/rust/src/bench/src/backends/pooling.rs +++ b/rust/src/bench/src/backends/pooling.rs @@ -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!({ diff --git a/rust/src/bench/src/benchmark.rs b/rust/src/bench/src/benchmark.rs index 8c834641ac..3210033ba9 100644 --- a/rust/src/bench/src/benchmark.rs +++ b/rust/src/bench/src/benchmark.rs @@ -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 { 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 { 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 { 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 { }; 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 { // 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 { 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 { 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 { .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 { .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 { // 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 { 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 { 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 { ); 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 { 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>, ) { - 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>, ) { - 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); } diff --git a/rust/src/bench/src/config.rs b/rust/src/bench/src/config.rs index ec132d1196..0700cf1fad 100644 --- a/rust/src/bench/src/config.rs +++ b/rust/src/bench/src/config.rs @@ -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" ); } diff --git a/rust/src/bench/src/datasets/hf_dataset.rs b/rust/src/bench/src/datasets/hf_dataset.rs index 35932a3983..499f54000f 100644 --- a/rust/src/bench/src/datasets/hf_dataset.rs +++ b/rust/src/bench/src/datasets/hf_dataset.rs @@ -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 = 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" ); } } diff --git a/rust/src/bench/src/datasets/mod.rs b/rust/src/bench/src/datasets/mod.rs index 3917b4d833..8f81f8d3c7 100644 --- a/rust/src/bench/src/datasets/mod.rs +++ b/rust/src/bench/src/datasets/mod.rs @@ -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" ); } diff --git a/rust/src/bench/src/datasets/multi_turn.rs b/rust/src/bench/src/datasets/multi_turn.rs index 43098b9725..513302f640 100644 --- a/rust/src/bench/src/datasets/multi_turn.rs +++ b/rust/src/bench/src/datasets/multi_turn.rs @@ -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" ); } diff --git a/rust/src/bench/src/datasets/prefix_repetition.rs b/rust/src/bench/src/datasets/prefix_repetition.rs index 7b06ef6077..bb0a2a9432 100644 --- a/rust/src/bench/src/datasets/prefix_repetition.rs +++ b/rust/src/bench/src/datasets/prefix_repetition.rs @@ -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" ); } diff --git a/rust/src/bench/src/datasets/progress.rs b/rust/src/bench/src/datasets/progress.rs new file mode 100644 index 0000000000..fb2df327e6 --- /dev/null +++ b/rust/src/bench/src/datasets/progress.rs @@ -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))); + } +} diff --git a/rust/src/bench/src/datasets/random.rs b/rust/src/bench/src/datasets/random.rs index 7cf5311278..d2068077b6 100644 --- a/rust/src/bench/src/datasets/random.rs +++ b/rust/src/bench/src/datasets/random.rs @@ -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" ); } diff --git a/rust/src/bench/src/datasets/sharegpt.rs b/rust/src/bench/src/datasets/sharegpt.rs index 2d3901820d..d293ba3c54 100644 --- a/rust/src/bench/src/datasets/sharegpt.rs +++ b/rust/src/bench/src/datasets/sharegpt.rs @@ -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 { - 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 { )) })?; 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" ); } } diff --git a/rust/src/bench/src/datasets/speed_bench.rs b/rust/src/bench/src/datasets/speed_bench.rs index a6490b76e8..6be07164db 100644 --- a/rust/src/bench/src/datasets/speed_bench.rs +++ b/rust/src/bench/src/datasets/speed_bench.rs @@ -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 { // 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 { let mut all_rows: Vec = 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 { 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 { 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 = 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) } diff --git a/rust/src/bench/src/lib.rs b/rust/src/bench/src/lib.rs index 13f252cee3..0a4f19c0d6 100644 --- a/rust/src/bench/src/lib.rs +++ b/rust/src/bench/src/lib.rs @@ -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"); } } diff --git a/rust/src/bench/src/main.rs b/rust/src/bench/src/main.rs index 919e1e0b36..b9ebfc2a0a 100644 --- a/rust/src/bench/src/main.rs +++ b/rust/src/bench/src/main.rs @@ -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(); diff --git a/rust/src/bench/src/metrics/calculator.rs b/rust/src/bench/src/metrics/calculator.rs index 195dad93de..3bce8dd1a7 100644 --- a/rust/src/bench/src/metrics/calculator.rs +++ b/rust/src/bench/src/metrics/calculator.rs @@ -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> = diff --git a/rust/src/bench/src/multi_turn.rs b/rust/src/bench/src/multi_turn.rs index 8f0aedf313..a80bae5528 100644 --- a/rust/src/bench/src/multi_turn.rs +++ b/rust/src/bench/src/multi_turn.rs @@ -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 Result Result 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 Result Result Result Result Result Result Result = 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 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 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(()) } diff --git a/rust/src/bench/src/ready_checker.rs b/rust/src/bench/src/ready_checker.rs index 86af8663d1..144a1ebd99 100644 --- a/rust/src/bench/src/ready_checker.rs +++ b/rust/src/bench/src/ready_checker.rs @@ -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) => { diff --git a/rust/src/bench/src/sweep.rs b/rust/src/bench/src/sweep.rs index 68cee15825..2d695356dc 100644 --- a/rust/src/bench/src/sweep.rs +++ b/rust/src/bench/src/sweep.rs @@ -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(); diff --git a/rust/src/bench/src/tiktoken.rs b/rust/src/bench/src/tiktoken.rs index b0cfbfeae4..342a6f7264 100644 --- a/rust/src/bench/src/tiktoken.rs +++ b/rust/src/bench/src/tiktoken.rs @@ -199,7 +199,12 @@ pub fn load_builtin_tiktoken(encoding: &str) -> Result { } }; 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 { 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); } diff --git a/rust/src/bench/src/tokenizer.rs b/rust/src/bench/src/tokenizer.rs index 3d0d4c9c31..1cbb665b4d 100644 --- a/rust/src/bench/src/tokenizer.rs +++ b/rust/src/bench/src/tokenizer.rs @@ -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)); } diff --git a/rust/src/cmd/src/logging.rs b/rust/src/cmd/src/logging.rs index 4d31b92ee2..936b369244 100644 --- a/rust/src/cmd/src/logging.rs +++ b/rust/src/cmd/src/logging.rs @@ -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")); diff --git a/rust/src/cmd/src/main.rs b/rust/src/cmd/src/main.rs index 6931467810..0806f7b75d 100644 --- a/rust/src/cmd/src/main.rs +++ b/rust/src/cmd/src/main.rs @@ -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();