diff --git a/rust/Cargo.lock b/rust/Cargo.lock index 795426f67d..5d2faecdb3 100644 --- a/rust/Cargo.lock +++ b/rust/Cargo.lock @@ -3446,7 +3446,6 @@ dependencies = [ "base64 0.22.1", "bytes", "encoding_rs", - "futures-channel", "futures-core", "futures-util", "h2", diff --git a/rust/src/bench/Cargo.toml b/rust/src/bench/Cargo.toml index 22f5608884..2da0f13f1c 100644 --- a/rust/src/bench/Cargo.toml +++ b/rust/src/bench/Cargo.toml @@ -20,7 +20,7 @@ mimalloc.workspace = true rand.workspace = true rand_distr.workspace = true rayon.workspace = true -reqwest = { workspace = true, features = ["json", "stream", "blocking", "http2"] } +reqwest = { workspace = true, features = ["json", "stream", "http2"] } rlimit.workspace = true rustc-hash.workspace = true serde = { workspace = true, features = ["rc"] } diff --git a/rust/src/bench/src/benchmark.rs b/rust/src/bench/src/benchmark.rs index 3210033ba9..cc4acc4f06 100644 --- a/rust/src/bench/src/benchmark.rs +++ b/rust/src/bench/src/benchmark.rs @@ -365,7 +365,8 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result { let tid = config.tokenizer_id.as_deref().unwrap_or(&model_id); 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)?; + let t = + crate::tokenizer::load_tokenizer(tid, config.trust_remote_code, server_info).await?; Some(t) }; let has_tokenizer = tokenizer.is_some(); @@ -481,7 +482,7 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result { let path = match config.dataset_path.as_deref() { Some(p) => p, None => { - downloaded = crate::datasets::sharegpt::download_sharegpt_dataset()?; + downloaded = crate::datasets::sharegpt::download_sharegpt_dataset().await?; downloaded.as_str() } }; @@ -521,7 +522,8 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result { None => { downloaded = crate::datasets::speed_bench::download_speed_bench( config.speed_bench_config, - )?; + ) + .await?; downloaded.as_str() } }; @@ -552,7 +554,8 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result { config.hf_subset.as_deref(), config.hf_split.as_deref(), config.num_prompts, - )?; + ) + .await?; crate::datasets::hf_dataset::load_hf_dataset( tok, &downloaded_path, diff --git a/rust/src/bench/src/datasets/custom.rs b/rust/src/bench/src/datasets/custom.rs index 04bbfcfdba..11ddee02b1 100644 --- a/rust/src/bench/src/datasets/custom.rs +++ b/rust/src/bench/src/datasets/custom.rs @@ -128,8 +128,10 @@ mod tests { /// gpt2 via built-in tiktoken encoding — loads without network access. fn test_tokenizer() -> TokenizerKind { - crate::tokenizer::load_tokenizer("gpt2", false, None) - .expect("gpt2 built-in tiktoken should always load without network") + TokenizerKind::Tiktoken( + crate::tiktoken::load_builtin_tiktoken("gpt2") + .expect("gpt2 built-in tiktoken should always load without network"), + ) } #[test] diff --git a/rust/src/bench/src/datasets/hf_dataset.rs b/rust/src/bench/src/datasets/hf_dataset.rs index 499f54000f..93c9f4cdbd 100644 --- a/rust/src/bench/src/datasets/hf_dataset.rs +++ b/rust/src/bench/src/datasets/hf_dataset.rs @@ -51,18 +51,19 @@ enum ColumnFormat { /// Make a GET request with retry logic (3 retries with exponential backoff). /// Returns the parsed JSON response. -fn get_with_retry( - client: &reqwest::blocking::Client, +async fn get_with_retry( + client: &reqwest::Client, url: &str, label: &str, ) -> Result { let max_retries = 3; for attempt in 0..=max_retries { - let resp = match client.get(url).send() { + let resp = match client.get(url).send().await { Ok(r) => r, Err(e) => { if attempt < max_retries { - std::thread::sleep(std::time::Duration::from_secs(2 * (attempt as u64 + 1))); + tokio::time::sleep(std::time::Duration::from_secs(2 * (attempt as u64 + 1))) + .await; continue; } return Err(BenchError::Config(format!( @@ -81,7 +82,7 @@ fn get_with_retry( } if status.is_server_error() && attempt < max_retries { - std::thread::sleep(std::time::Duration::from_secs(2 * (attempt as u64 + 1))); + tokio::time::sleep(std::time::Duration::from_secs(2 * (attempt as u64 + 1))).await; continue; } @@ -93,6 +94,7 @@ fn get_with_retry( let data: serde_json::Value = resp .json() + .await .map_err(|e| BenchError::Config(format!("Failed to parse {label} response: {e}")))?; return Ok(data); } @@ -106,7 +108,7 @@ fn get_with_retry( /// If both `subset` and `split` are provided, the `/info` call is skipped as an optimization. /// Paginated download fetches rows in pages of 100 until `num_rows_needed` are collected /// or the dataset is exhausted. -pub fn download_hf_dataset( +pub async fn download_hf_dataset( dataset: &str, subset: Option<&str>, split: Option<&str>, @@ -116,7 +118,7 @@ pub fn download_hf_dataset( url::form_urlencoded::byte_serialize(dataset.as_bytes()).collect(); let mut client_builder = - reqwest::blocking::Client::builder().timeout(std::time::Duration::from_secs(120)); + reqwest::Client::builder().timeout(std::time::Duration::from_secs(120)); // Add HF_TOKEN auth header if available if let Ok(token) = std::env::var("HF_TOKEN") { @@ -139,7 +141,7 @@ pub fn download_hf_dataset( // Call /info to discover available configs and splits let info_url = format!("https://datasets-server.huggingface.co/info?dataset={encoded_dataset}"); - let info = get_with_retry(&client, &info_url, "HF dataset /info")?; + let info = get_with_retry(&client, &info_url, "HF dataset /info").await?; let dataset_info = info.get("dataset_info").and_then(|d| d.as_object()).ok_or_else(|| { @@ -252,7 +254,7 @@ pub fn download_hf_dataset( &length={page_size}" ); - let data = get_with_retry(&client, &url, "HF dataset /rows")?; + let data = get_with_retry(&client, &url, "HF dataset /rows").await?; let rows = data["rows"] .as_array() @@ -1035,8 +1037,10 @@ mod tests { /// Build a gpt2 tokenizer using built-in tiktoken encoding (no network required). fn builtin_tokenizer() -> crate::tokenizer::TokenizerKind { - crate::tokenizer::load_tokenizer("gpt2", false, None) - .expect("gpt2 built-in tiktoken should always load without network") + crate::tokenizer::TokenizerKind::Tiktoken( + crate::tiktoken::load_builtin_tiktoken("gpt2") + .expect("gpt2 built-in tiktoken should always load without network"), + ) } /// Write JSON data to a unique temp file and return the path string. diff --git a/rust/src/bench/src/datasets/multi_turn.rs b/rust/src/bench/src/datasets/multi_turn.rs index 513302f640..97ec522e29 100644 --- a/rust/src/bench/src/datasets/multi_turn.rs +++ b/rust/src/bench/src/datasets/multi_turn.rs @@ -526,10 +526,12 @@ mod tests { len } - #[test] + #[tokio::test] #[ignore] - fn test_prefix_sharing_structure() { - let tok = crate::tokenizer::load_tokenizer("nvidia/Kimi-K2.5-NVFP4", false, None).unwrap(); + async fn test_prefix_sharing_structure() { + let tok = crate::tokenizer::load_tokenizer("nvidia/Kimi-K2.5-NVFP4", false, None) + .await + .unwrap(); let cfg = MultiTurnRandomConfig { num_conversations: 5, @@ -611,10 +613,12 @@ mod tests { println!("All prefix sharing checks passed!"); } - #[test] + #[tokio::test] #[ignore] - fn test_per_turn_input_len_default_mode() { - let tok = crate::tokenizer::load_tokenizer("nvidia/Kimi-K2.5-NVFP4", false, None).unwrap(); + async fn test_per_turn_input_len_default_mode() { + let tok = crate::tokenizer::load_tokenizer("nvidia/Kimi-K2.5-NVFP4", false, None) + .await + .unwrap(); let cfg = MultiTurnRandomConfig { num_conversations: 4, @@ -651,10 +655,12 @@ mod tests { println!("per_turn_input_len default-mode checks passed!"); } - #[test] + #[tokio::test] #[ignore] - fn test_variable_turns_range() { - let tok = crate::tokenizer::load_tokenizer("nvidia/Kimi-K2.5-NVFP4", false, None).unwrap(); + async fn test_variable_turns_range() { + let tok = crate::tokenizer::load_tokenizer("nvidia/Kimi-K2.5-NVFP4", false, None) + .await + .unwrap(); let cfg = MultiTurnRandomConfig { num_conversations: 50, @@ -685,10 +691,12 @@ mod tests { println!("variable_turns_range checks passed! counts: {distinct_counts:?}"); } - #[test] + #[tokio::test] #[ignore] - fn test_variable_turns_fixed() { - let tok = crate::tokenizer::load_tokenizer("nvidia/Kimi-K2.5-NVFP4", false, None).unwrap(); + async fn test_variable_turns_fixed() { + let tok = crate::tokenizer::load_tokenizer("nvidia/Kimi-K2.5-NVFP4", false, None) + .await + .unwrap(); let cfg = MultiTurnRandomConfig { num_conversations: 10, @@ -710,10 +718,12 @@ mod tests { println!("variable_turns_fixed checks passed!"); } - #[test] + #[tokio::test] #[ignore] - fn test_per_turn_input_len_prefix_sharing() { - let tok = crate::tokenizer::load_tokenizer("nvidia/Kimi-K2.5-NVFP4", false, None).unwrap(); + async fn test_per_turn_input_len_prefix_sharing() { + let tok = crate::tokenizer::load_tokenizer("nvidia/Kimi-K2.5-NVFP4", false, None) + .await + .unwrap(); // Turn 0 input_len=1000, turns 1+ per_turn_input_len=600 // global_len ≈ 100 (10%), conv_len ≈ 800 (80%), unique ≈ 100 diff --git a/rust/src/bench/src/datasets/prefix_repetition.rs b/rust/src/bench/src/datasets/prefix_repetition.rs index bb0a2a9432..e4dfcd36a3 100644 --- a/rust/src/bench/src/datasets/prefix_repetition.rs +++ b/rust/src/bench/src/datasets/prefix_repetition.rs @@ -111,8 +111,10 @@ mod tests { /// gpt2 via built-in tiktoken encoding — loads without network access. fn test_tokenizer() -> TokenizerKind { - crate::tokenizer::load_tokenizer("gpt2", false, None) - .expect("gpt2 built-in tiktoken should always load without network") + TokenizerKind::Tiktoken( + crate::tiktoken::load_builtin_tiktoken("gpt2") + .expect("gpt2 built-in tiktoken should always load without network"), + ) } #[test] diff --git a/rust/src/bench/src/datasets/random.rs b/rust/src/bench/src/datasets/random.rs index d2068077b6..132dd8420e 100644 --- a/rust/src/bench/src/datasets/random.rs +++ b/rust/src/bench/src/datasets/random.rs @@ -308,7 +308,8 @@ mod tests { #[test] #[ignore] fn test_generate_random_dataset_token_ids() { - let tokenizer = tokenizer::load_tokenizer("gpt2", false, None).unwrap(); + let tokenizer = + TokenizerKind::Tiktoken(crate::tiktoken::load_builtin_tiktoken("gpt2").unwrap()); let requests = generate_random_dataset( &tokenizer, 10, // num_requests @@ -340,7 +341,8 @@ mod tests { #[test] #[ignore] fn test_generate_random_dataset_text() { - let tokenizer = tokenizer::load_tokenizer("gpt2", false, None).unwrap(); + let tokenizer = + TokenizerKind::Tiktoken(crate::tiktoken::load_builtin_tiktoken("gpt2").unwrap()); let requests = generate_random_dataset( &tokenizer, 10, // num_requests @@ -374,7 +376,8 @@ mod tests { #[test] #[ignore] fn test_token_length_exact_local() { - let tokenizer = tokenizer::load_tokenizer("gpt2", false, None).unwrap(); + let tokenizer = + TokenizerKind::Tiktoken(crate::tiktoken::load_builtin_tiktoken("gpt2").unwrap()); let target_len = 512; let requests = generate_random_dataset( &tokenizer, @@ -408,11 +411,11 @@ mod tests { } /// Test that tiktoken tokenizer produces exact target token lengths (token ID mode). - #[test] + #[tokio::test] #[ignore] - fn test_token_length_exact_tiktoken() { + async fn test_token_length_exact_tiktoken() { // Use Qwen2.5 which has a tiktoken-format tokenizer - let tokenizer = tokenizer::load_tokenizer("Qwen/Qwen2.5-0.5B", false, None); + let tokenizer = tokenizer::load_tokenizer("Qwen/Qwen2.5-0.5B", false, None).await; let tokenizer = match tokenizer { Ok(t) => t, Err(e) => { @@ -456,10 +459,10 @@ mod tests { /// Test encode/decode roundtrip stability for tiktoken. /// After one decode→encode cycle with UTF-8-safe tokens, length must not drift. - #[test] + #[tokio::test] #[ignore] - fn test_tiktoken_roundtrip_stability() { - let tokenizer = tokenizer::load_tokenizer("Qwen/Qwen2.5-0.5B", false, None); + async fn test_tiktoken_roundtrip_stability() { + let tokenizer = tokenizer::load_tokenizer("Qwen/Qwen2.5-0.5B", false, None).await; let tokenizer = match tokenizer { Ok(t) => t, Err(e) => { diff --git a/rust/src/bench/src/datasets/random_rerank.rs b/rust/src/bench/src/datasets/random_rerank.rs index 9c335e7e44..ad3686cf3f 100644 --- a/rust/src/bench/src/datasets/random_rerank.rs +++ b/rust/src/bench/src/datasets/random_rerank.rs @@ -139,8 +139,10 @@ mod tests { /// gpt2 via built-in tiktoken encoding — loads without network access. fn test_tokenizer() -> TokenizerKind { - crate::tokenizer::load_tokenizer("gpt2", false, None) - .expect("gpt2 built-in tiktoken should always load without network") + TokenizerKind::Tiktoken( + crate::tiktoken::load_builtin_tiktoken("gpt2") + .expect("gpt2 built-in tiktoken should always load without network"), + ) } fn fixed_ratio() -> RangeRatio { diff --git a/rust/src/bench/src/datasets/sharegpt.rs b/rust/src/bench/src/datasets/sharegpt.rs index d293ba3c54..b8f7543b21 100644 --- a/rust/src/bench/src/datasets/sharegpt.rs +++ b/rust/src/bench/src/datasets/sharegpt.rs @@ -22,14 +22,15 @@ 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 { +pub async fn download_sharegpt_dataset() -> Result { 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| { + let repo = crate::hub::HubRepo::dataset(DEFAULT_SHAREGPT_REPO.to_string()) + .map_err(BenchError::Config)?; + let path = repo.get(DEFAULT_SHAREGPT_FILE).await.map_err(|e| { BenchError::Config(format!( "Failed to download ShareGPT dataset from '{DEFAULT_SHAREGPT_REPO}': {e}" )) diff --git a/rust/src/bench/src/datasets/speed_bench.rs b/rust/src/bench/src/datasets/speed_bench.rs index 6be07164db..d615bb4903 100644 --- a/rust/src/bench/src/datasets/speed_bench.rs +++ b/rust/src/bench/src/datasets/speed_bench.rs @@ -26,7 +26,7 @@ fn cache_dir() -> std::path::PathBuf { /// Download SPEED-Bench dataset from HuggingFace datasets-server API. /// Results are cached as JSON locally for subsequent runs. -pub fn download_speed_bench(config: SpeedBenchConfig) -> Result { +pub async fn download_speed_bench(config: SpeedBenchConfig) -> Result { let config_name = config.as_str(); let dir = cache_dir(); @@ -42,7 +42,7 @@ pub fn download_speed_bench(config: SpeedBenchConfig) -> Result { tracing::info!(config = config_name, "downloading SPEED-Bench dataset"); - let client = reqwest::blocking::Client::builder() + let client = reqwest::Client::builder() .timeout(std::time::Duration::from_secs(120)) .build() .map_err(|e| BenchError::Config(format!("Failed to build HTTP client: {e}")))?; @@ -66,13 +66,14 @@ pub fn download_speed_bench(config: SpeedBenchConfig) -> Result { let max_retries = 3; let mut data: Option = None; for attempt in 0..=max_retries { - let resp = match client.get(&url).send() { + let resp = match client.get(&url).send().await { Ok(r) => r, Err(e) => { if attempt < max_retries { - std::thread::sleep(std::time::Duration::from_secs( + tokio::time::sleep(std::time::Duration::from_secs( 2 * (attempt as u64 + 1), - )); + )) + .await; continue; } return Err(BenchError::Config(format!( @@ -82,7 +83,7 @@ pub fn download_speed_bench(config: SpeedBenchConfig) -> Result { }; if resp.status().is_server_error() && attempt < max_retries { - std::thread::sleep(std::time::Duration::from_secs(2 * (attempt as u64 + 1))); + tokio::time::sleep(std::time::Duration::from_secs(2 * (attempt as u64 + 1))).await; continue; } @@ -93,7 +94,7 @@ pub fn download_speed_bench(config: SpeedBenchConfig) -> Result { ))); } - data = Some(resp.json().map_err(|e| { + data = Some(resp.json().await.map_err(|e| { BenchError::Config(format!("Failed to parse SPEED-Bench API response: {e}")) })?); break; diff --git a/rust/src/bench/src/hub.rs b/rust/src/bench/src/hub.rs index ab51ccdcf0..0b446450fb 100644 --- a/rust/src/bench/src/hub.rs +++ b/rust/src/bench/src/hub.rs @@ -1,54 +1,39 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright contributors to the vLLM project -//! Sync facade over the async `hf_hub` API. -//! -//! The workspace bans rustls (`rust/deny.toml`), but hf-hub's sync `ureq` -//! backend unconditionally pulls ureq's default rustls feature. So we use the -//! reqwest/native-tls tokio API instead, and bridge blocking callers (dataset -//! loaders, tokenizer fallback in rayon threads) by running each download on a -//! dedicated thread with its own single-threaded runtime. - use std::path::PathBuf; +use hf_hub::Repo; +use hf_hub::api::tokio::{ApiBuilder, ApiRepo}; + /// A handle to a HuggingFace Hub repo, downloading via hf-hub's on-disk cache. pub struct HubRepo { - repo: hf_hub::Repo, + repo: ApiRepo, } impl HubRepo { - pub fn model(model_id: String) -> Self { - Self { - repo: hf_hub::Repo::model(model_id), - } + pub fn model(model_id: String) -> Result { + Self::new(Repo::model(model_id)) } - pub fn dataset(repo_id: String) -> Self { - Self { - repo: hf_hub::Repo::dataset(repo_id), + pub fn dataset(repo_id: String) -> Result { + Self::new(Repo::dataset(repo_id)) + } + + fn new(repo: Repo) -> Result { + let mut builder = ApiBuilder::from_env(); + if let Ok(token) = std::env::var("HF_TOKEN") { + builder = builder.with_token(Some(token)); } + let api = builder.build().map_err(|e| format!("Failed to init HF API: {e}"))?; + Ok(Self { + repo: api.repo(repo), + }) } /// Download (or fetch from cache) a single file from the repo. /// Auth is handled by hf-hub via HF_TOKEN / the cached login token. - pub fn get(&self, filename: &str) -> Result { - let repo = self.repo.clone(); - let filename = filename.to_string(); - std::thread::spawn(move || { - let rt = tokio::runtime::Builder::new_current_thread() - .enable_all() - .build() - .map_err(|e| format!("Failed to build download runtime: {e}"))?; - rt.block_on(async move { - let mut builder = hf_hub::api::tokio::ApiBuilder::from_env(); - if let Ok(token) = std::env::var("HF_TOKEN") { - builder = builder.with_token(Some(token)); - } - let api = builder.build().map_err(|e| format!("Failed to init HF API: {e}"))?; - api.repo(repo).get(&filename).await.map_err(|e| format!("{e}")) - }) - }) - .join() - .map_err(|_| "HF Hub download thread panicked".to_string())? + pub async fn get(&self, filename: &str) -> Result { + self.repo.get(filename).await.map_err(|e| format!("{e}")) } } diff --git a/rust/src/bench/src/multi_turn.rs b/rust/src/bench/src/multi_turn.rs index a80bae5528..4b962b879c 100644 --- a/rust/src/bench/src/multi_turn.rs +++ b/rust/src/bench/src/multi_turn.rs @@ -90,7 +90,8 @@ pub async fn run_multi_turn_benchmark(config: &BenchConfig) -> Result Result p, None => { - downloaded = crate::datasets::sharegpt::download_sharegpt_dataset()?; + downloaded = crate::datasets::sharegpt::download_sharegpt_dataset().await?; downloaded.as_str() } }; diff --git a/rust/src/bench/src/tiktoken.rs b/rust/src/bench/src/tiktoken.rs index 342a6f7264..3eda0874b6 100644 --- a/rust/src/bench/src/tiktoken.rs +++ b/rust/src/bench/src/tiktoken.rs @@ -209,7 +209,7 @@ pub fn load_builtin_tiktoken(encoding: &str) -> Result { } /// Try to load a tiktoken tokenizer from a local directory or HuggingFace model repo. -pub fn try_load_tiktoken(model_id: &str) -> Result { +pub async fn try_load_tiktoken(model_id: &str) -> Result { // Phase 1: If model_id is a local directory, look for tiktoken files there let local_dir = Path::new(model_id); if local_dir.is_dir() { @@ -217,7 +217,7 @@ pub fn try_load_tiktoken(model_id: &str) -> Result { } // Phase 2: Fall back to HuggingFace Hub download - try_load_tiktoken_from_hf(model_id) + try_load_tiktoken_from_hf(model_id).await } /// Common tiktoken model filenames to search for. @@ -252,25 +252,28 @@ fn try_load_tiktoken_from_dir(dir: &Path, model_id: &str) -> Result Result { - let repo = crate::hub::HubRepo::model(model_id.to_string()); +async fn try_load_tiktoken_from_hf(model_id: &str) -> Result { + let repo = crate::hub::HubRepo::model(model_id.to_string()).map_err(BenchError::Tokenizer)?; - let model_path = repo - .get("tiktoken.model") - .or_else(|_| repo.get("qwen.tiktoken")) - .or_else(|_| repo.get("vocab.tiktoken")) - .map_err(|_| { - BenchError::Tokenizer(format!("No tiktoken model file found for '{model_id}'")) - })?; + let mut model_path = None; + for filename in TIKTOKEN_MODEL_FILENAMES { + if let Ok(path) = repo.get(filename).await { + model_path = Some(path); + break; + } + } + let model_path = model_path.ok_or_else(|| { + BenchError::Tokenizer(format!("No tiktoken model file found for '{model_id}'")) + })?; let num_base_tokens = count_base_tokens(&model_path)?; - let config = match repo.get("tokenizer_config.json") { + let config = match repo.get("tokenizer_config.json").await { Ok(config_path) => read_tokenizer_config(&config_path), Err(_) => None, }; - let pattern = extract_pat_str_from_repo(&repo); + let pattern = extract_pat_str_from_repo(&repo).await; build_tiktoken(model_id, &model_path, config, pattern, num_base_tokens) } @@ -403,9 +406,12 @@ fn extract_pat_str_from_local_dir(dir: &Path) -> Option { /// Try to download the Python tokenizer source file and extract pat_str via regex. /// Returns None if unavailable or unparsable. -fn extract_pat_str_from_repo(repo: &crate::hub::HubRepo) -> Option { +async fn extract_pat_str_from_repo(repo: &crate::hub::HubRepo) -> Option { // Try common Python tokenizer filenames - let py_path = repo.get("tokenization_kimi.py").or_else(|_| repo.get("tokenizer.py")).ok()?; + let py_path = match repo.get("tokenization_kimi.py").await { + Ok(path) => path, + Err(_) => repo.get("tokenizer.py").await.ok()?, + }; let source = std::fs::read_to_string(&py_path).ok()?; diff --git a/rust/src/bench/src/tokenizer.rs b/rust/src/bench/src/tokenizer.rs index 1cbb665b4d..add890c391 100644 --- a/rust/src/bench/src/tokenizer.rs +++ b/rust/src/bench/src/tokenizer.rs @@ -2,6 +2,7 @@ // SPDX-FileCopyrightText: Copyright contributors to the vLLM project use std::collections::HashSet; +use std::future::Future; use std::path::Path; use thiserror_ext::AsReport as _; @@ -19,7 +20,8 @@ pub enum TokenizerKind { /// Server-side tokenizer using vLLM's /tokenize and /detokenize endpoints. pub struct ServerTokenizer { - client: reqwest::blocking::Client, + client: reqwest::Client, + runtime: tokio::runtime::Handle, tokenize_url: String, detokenize_url: String, model: String, @@ -28,8 +30,8 @@ pub struct ServerTokenizer { impl ServerTokenizer { /// Create a new server tokenizer and verify connectivity. - pub fn new(base_url: &str, model: &str) -> Result { - let client = reqwest::blocking::Client::builder() + pub async fn new(base_url: &str, model: &str) -> Result { + let client = reqwest::Client::builder() .timeout(std::time::Duration::from_secs(30)) .build() .map_err(|e| BenchError::Tokenizer(format!("Failed to build HTTP client: {e}")))?; @@ -39,6 +41,7 @@ impl ServerTokenizer { let st = Self { client, + runtime: tokio::runtime::Handle::current(), tokenize_url, detokenize_url, model: model.to_string(), @@ -46,7 +49,7 @@ impl ServerTokenizer { }; // Probe the endpoint to verify it works and discover vocab size - let test_tokens = st.encode_inner("test")?; + let test_tokens = st.encode_async("test").await?; let max_id = test_tokens.iter().copied().max().unwrap_or(0); let estimated_vocab = (max_id * 2).max(131072); @@ -57,6 +60,10 @@ impl ServerTokenizer { } fn encode_inner(&self, text: &str) -> Result> { + self.block_on(self.encode_async(text)) + } + + async fn encode_async(&self, text: &str) -> Result> { let payload = serde_json::json!({ "model": self.model, "prompt": text, @@ -67,6 +74,7 @@ impl ServerTokenizer { .post(&self.tokenize_url) .json(&payload) .send() + .await .map_err(|e| BenchError::Tokenizer(format!("Server tokenize failed: {e}")))?; if !resp.status().is_success() { @@ -76,7 +84,7 @@ impl ServerTokenizer { ))); } - let data: serde_json::Value = resp.json().map_err(|e| { + let data: serde_json::Value = resp.json().await.map_err(|e| { BenchError::Tokenizer(format!("Failed to parse tokenize response: {e}")) })?; @@ -96,6 +104,10 @@ impl ServerTokenizer { } fn decode_inner(&self, ids: &[u32]) -> Result { + self.block_on(self.decode_async(ids)) + } + + async fn decode_async(&self, ids: &[u32]) -> Result { let payload = serde_json::json!({ "model": self.model, "tokens": ids, @@ -106,6 +118,7 @@ impl ServerTokenizer { .post(&self.detokenize_url) .json(&payload) .send() + .await .map_err(|e| BenchError::Tokenizer(format!("Server detokenize failed: {e}")))?; if !resp.status().is_success() { @@ -115,7 +128,7 @@ impl ServerTokenizer { ))); } - let data: serde_json::Value = resp.json().map_err(|e| { + let data: serde_json::Value = resp.json().await.map_err(|e| { BenchError::Tokenizer(format!("Failed to parse detokenize response: {e}")) })?; @@ -124,6 +137,26 @@ impl ServerTokenizer { .map(|s| s.to_string()) .ok_or_else(|| BenchError::Tokenizer("Missing 'prompt' in detokenize response".into())) } + + fn block_on(&self, future: impl Future>) -> Result { + if matches!( + self.runtime.runtime_flavor(), + tokio::runtime::RuntimeFlavor::CurrentThread + ) { + return Err(BenchError::Tokenizer( + "Server tokenizer fallback requires a multi-thread Tokio runtime".into(), + )); + } + + // Sync tokenizer calls can come from a Tokio worker or a Rayon worker. + // Tokio workers must enter a blocking region before re-entering the runtime; + // Rayon workers can drive the future directly with the saved runtime handle. + if tokio::runtime::Handle::try_current().is_ok() { + tokio::task::block_in_place(|| self.runtime.block_on(future)) + } else { + self.runtime.block_on(future) + } + } } // --- TokenizerKind methods --- @@ -193,7 +226,7 @@ impl TokenizerKind { /// 3. Server-side /tokenize + /detokenize endpoints /// /// `server_info` is `Some((base_url, model))` to enable server-side fallback. -pub fn load_tokenizer( +pub async fn load_tokenizer( model_id: &str, _trust_remote_code: bool, server_info: Option<(&str, &str)>, @@ -213,7 +246,7 @@ pub fn load_tokenizer( } // 1. Try local HuggingFace tokenizer (tokenizer.json) - match try_load_local(model_id) { + match try_load_local(model_id).await { Ok(tok) => { tracing::info!( model = model_id, @@ -230,7 +263,7 @@ pub fn load_tokenizer( error = %local_err.as_report(), "local tokenizer unavailable; trying tiktoken" ); - match crate::tiktoken::try_load_tiktoken(model_id) { + match crate::tiktoken::try_load_tiktoken(model_id).await { Ok(tok) => { tracing::info!( model = model_id, @@ -248,7 +281,7 @@ pub fn load_tokenizer( error = %tiktoken_err.as_report(), "tiktoken unavailable; trying server-side tokenization" ); - match ServerTokenizer::new(base_url, model) { + match ServerTokenizer::new(base_url, model).await { Ok(srv) => { tracing::info!( model = model_id, @@ -282,7 +315,7 @@ pub fn load_tokenizer( } /// Try loading tokenizer.json from local path or HuggingFace Hub. -fn try_load_local(model_id: &str) -> Result { +async fn try_load_local(model_id: &str) -> Result { // 1. Try local directory with tokenizer.json let local_path = Path::new(model_id).join("tokenizer.json"); if local_path.exists() { @@ -308,11 +341,37 @@ fn try_load_local(model_id: &str) -> Result { } // 4. Download from HuggingFace Hub (hf-hub handles auth via HF_TOKEN / cached token) - let repo = crate::hub::HubRepo::model(model_id.to_string()); + let repo = crate::hub::HubRepo::model(model_id.to_string()).map_err(BenchError::Tokenizer)?; let tokenizer_path = repo .get("tokenizer.json") + .await .map_err(|e| BenchError::Tokenizer(format!("No tokenizer.json for '{model_id}': {e}")))?; Tokenizer::from_file(&tokenizer_path) .map_err(|e| BenchError::Tokenizer(format!("Failed to load downloaded tokenizer: {e}"))) } + +#[cfg(test)] +mod tests { + use super::*; + + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn test_server_tokenizer_sync_bridge() { + let tokenizer = std::sync::Arc::new(ServerTokenizer { + client: reqwest::Client::new(), + runtime: tokio::runtime::Handle::current(), + tokenize_url: String::new(), + detokenize_url: String::new(), + model: String::new(), + cached_vocab_size: 0, + }); + + assert_eq!(tokenizer.block_on(async { Ok(1) }).unwrap(), 1); + + let (tx, rx) = tokio::sync::oneshot::channel(); + rayon::spawn(move || { + let _ = tx.send(tokenizer.block_on(async { Ok(2) })); + }); + assert_eq!(rx.await.unwrap().unwrap(), 2); + } +}