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

[Rust][Benchmark] Use async HTTP clients (#49295)

Signed-off-by: Bugen Zhao <i@bugenzhao.com>

Bugen Zhao committed Jul 21, 2026 at 16:53 UTC eb44b3aaa40b134ffb9030c19671964c8d0602e3
15 files changed +200 -122
rust/Cargo.lock
-1
@@ -3446,7 +3446,6 @@ dependencies = [
3446 "base64 0.22.1",
3447 "bytes",
3448 "encoding_rs",
3449 - "futures-channel",
3449 "futures-core",
3450 "futures-util",
3451 "h2",
rust/src/bench/Cargo.toml
+1 -1
@@ -20,7 +20,7 @@ mimalloc.workspace = true
20 rand.workspace = true
21 rand_distr.workspace = true
22 rayon.workspace = true
23 -reqwest = { workspace = true, features = ["json", "stream", "blocking", "http2"] }
23 +reqwest = { workspace = true, features = ["json", "stream", "http2"] }
24 rlimit.workspace = true
25 rustc-hash.workspace = true
26 serde = { workspace = true, features = ["rc"] }
rust/src/bench/src/benchmark.rs
+7 -4
@@ -365,7 +365,8 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result<serde_json::Value> {
365 let tid = config.tokenizer_id.as_deref().unwrap_or(&model_id);
366 tracing::info!(tokenizer = tid, "loading tokenizer");
367 let server_info = Some((config.base_url.as_str(), model_id.as_str()));
368 - let t = crate::tokenizer::load_tokenizer(tid, config.trust_remote_code, server_info)?;
368 + let t =
369 + crate::tokenizer::load_tokenizer(tid, config.trust_remote_code, server_info).await?;
370 Some(t)
371 };
372 let has_tokenizer = tokenizer.is_some();
@@ -481,7 +482,7 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result<serde_json::Value> {
482 let path = match config.dataset_path.as_deref() {
483 Some(p) => p,
484 None => {
484 - downloaded = crate::datasets::sharegpt::download_sharegpt_dataset()?;
485 + downloaded = crate::datasets::sharegpt::download_sharegpt_dataset().await?;
486 downloaded.as_str()
487 }
488 };
@@ -521,7 +522,8 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result<serde_json::Value> {
522 None => {
523 downloaded = crate::datasets::speed_bench::download_speed_bench(
524 config.speed_bench_config,
524 - )?;
525 + )
526 + .await?;
527 downloaded.as_str()
528 }
529 };
@@ -552,7 +554,8 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result<serde_json::Value> {
554 config.hf_subset.as_deref(),
555 config.hf_split.as_deref(),
556 config.num_prompts,
555 - )?;
557 + )
558 + .await?;
559 crate::datasets::hf_dataset::load_hf_dataset(
560 tok,
561 &downloaded_path,
rust/src/bench/src/datasets/custom.rs
+4 -2
@@ -128,8 +128,10 @@ mod tests {
128
129 /// gpt2 via built-in tiktoken encoding — loads without network access.
130 fn test_tokenizer() -> TokenizerKind {
131 - crate::tokenizer::load_tokenizer("gpt2", false, None)
132 - .expect("gpt2 built-in tiktoken should always load without network")
131 + TokenizerKind::Tiktoken(
132 + crate::tiktoken::load_builtin_tiktoken("gpt2")
133 + .expect("gpt2 built-in tiktoken should always load without network"),
134 + )
135 }
136
137 #[test]
rust/src/bench/src/datasets/hf_dataset.rs
+15 -11
@@ -51,18 +51,19 @@ enum ColumnFormat {
51
52 /// Make a GET request with retry logic (3 retries with exponential backoff).
53 /// Returns the parsed JSON response.
54 -fn get_with_retry(
55 - client: &reqwest::blocking::Client,
54 +async fn get_with_retry(
55 + client: &reqwest::Client,
56 url: &str,
57 label: &str,
58 ) -> Result<serde_json::Value> {
59 let max_retries = 3;
60 for attempt in 0..=max_retries {
61 - let resp = match client.get(url).send() {
61 + let resp = match client.get(url).send().await {
62 Ok(r) => r,
63 Err(e) => {
64 if attempt < max_retries {
65 - std::thread::sleep(std::time::Duration::from_secs(2 * (attempt as u64 + 1)));
65 + tokio::time::sleep(std::time::Duration::from_secs(2 * (attempt as u64 + 1)))
66 + .await;
67 continue;
68 }
69 return Err(BenchError::Config(format!(
@@ -81,7 +82,7 @@ fn get_with_retry(
82 }
83
84 if status.is_server_error() && attempt < max_retries {
84 - std::thread::sleep(std::time::Duration::from_secs(2 * (attempt as u64 + 1)));
85 + tokio::time::sleep(std::time::Duration::from_secs(2 * (attempt as u64 + 1))).await;
86 continue;
87 }
88
@@ -93,6 +94,7 @@ fn get_with_retry(
94
95 let data: serde_json::Value = resp
96 .json()
97 + .await
98 .map_err(|e| BenchError::Config(format!("Failed to parse {label} response: {e}")))?;
99 return Ok(data);
100 }
@@ -106,7 +108,7 @@ fn get_with_retry(
108 /// If both `subset` and `split` are provided, the `/info` call is skipped as an optimization.
109 /// Paginated download fetches rows in pages of 100 until `num_rows_needed` are collected
110 /// or the dataset is exhausted.
109 -pub fn download_hf_dataset(
111 +pub async fn download_hf_dataset(
112 dataset: &str,
113 subset: Option<&str>,
114 split: Option<&str>,
@@ -116,7 +118,7 @@ pub fn download_hf_dataset(
118 url::form_urlencoded::byte_serialize(dataset.as_bytes()).collect();
119
120 let mut client_builder =
119 - reqwest::blocking::Client::builder().timeout(std::time::Duration::from_secs(120));
121 + reqwest::Client::builder().timeout(std::time::Duration::from_secs(120));
122
123 // Add HF_TOKEN auth header if available
124 if let Ok(token) = std::env::var("HF_TOKEN") {
@@ -139,7 +141,7 @@ pub fn download_hf_dataset(
141 // Call /info to discover available configs and splits
142 let info_url =
143 format!("https://datasets-server.huggingface.co/info?dataset={encoded_dataset}");
142 - let info = get_with_retry(&client, &info_url, "HF dataset /info")?;
144 + let info = get_with_retry(&client, &info_url, "HF dataset /info").await?;
145
146 let dataset_info =
147 info.get("dataset_info").and_then(|d| d.as_object()).ok_or_else(|| {
@@ -252,7 +254,7 @@ pub fn download_hf_dataset(
254 &length={page_size}"
255 );
256
255 - let data = get_with_retry(&client, &url, "HF dataset /rows")?;
257 + let data = get_with_retry(&client, &url, "HF dataset /rows").await?;
258
259 let rows = data["rows"]
260 .as_array()
@@ -1035,8 +1037,10 @@ mod tests {
1037
1038 /// Build a gpt2 tokenizer using built-in tiktoken encoding (no network required).
1039 fn builtin_tokenizer() -> crate::tokenizer::TokenizerKind {
1038 - crate::tokenizer::load_tokenizer("gpt2", false, None)
1039 - .expect("gpt2 built-in tiktoken should always load without network")
1040 + crate::tokenizer::TokenizerKind::Tiktoken(
1041 + crate::tiktoken::load_builtin_tiktoken("gpt2")
1042 + .expect("gpt2 built-in tiktoken should always load without network"),
1043 + )
1044 }
1045
1046 /// Write JSON data to a unique temp file and return the path string.
rust/src/bench/src/datasets/multi_turn.rs
+25 -15
@@ -526,10 +526,12 @@ mod tests {
526 len
527 }
528
529 - #[test]
529 + #[tokio::test]
530 #[ignore]
531 - fn test_prefix_sharing_structure() {
532 - let tok = crate::tokenizer::load_tokenizer("nvidia/Kimi-K2.5-NVFP4", false, None).unwrap();
531 + async fn test_prefix_sharing_structure() {
532 + let tok = crate::tokenizer::load_tokenizer("nvidia/Kimi-K2.5-NVFP4", false, None)
533 + .await
534 + .unwrap();
535
536 let cfg = MultiTurnRandomConfig {
537 num_conversations: 5,
@@ -611,10 +613,12 @@ mod tests {
613 println!("All prefix sharing checks passed!");
614 }
615
614 - #[test]
616 + #[tokio::test]
617 #[ignore]
616 - fn test_per_turn_input_len_default_mode() {
617 - let tok = crate::tokenizer::load_tokenizer("nvidia/Kimi-K2.5-NVFP4", false, None).unwrap();
618 + async fn test_per_turn_input_len_default_mode() {
619 + let tok = crate::tokenizer::load_tokenizer("nvidia/Kimi-K2.5-NVFP4", false, None)
620 + .await
621 + .unwrap();
622
623 let cfg = MultiTurnRandomConfig {
624 num_conversations: 4,
@@ -651,10 +655,12 @@ mod tests {
655 println!("per_turn_input_len default-mode checks passed!");
656 }
657
654 - #[test]
658 + #[tokio::test]
659 #[ignore]
656 - fn test_variable_turns_range() {
657 - let tok = crate::tokenizer::load_tokenizer("nvidia/Kimi-K2.5-NVFP4", false, None).unwrap();
660 + async fn test_variable_turns_range() {
661 + let tok = crate::tokenizer::load_tokenizer("nvidia/Kimi-K2.5-NVFP4", false, None)
662 + .await
663 + .unwrap();
664
665 let cfg = MultiTurnRandomConfig {
666 num_conversations: 50,
@@ -685,10 +691,12 @@ mod tests {
691 println!("variable_turns_range checks passed! counts: {distinct_counts:?}");
692 }
693
688 - #[test]
694 + #[tokio::test]
695 #[ignore]
690 - fn test_variable_turns_fixed() {
691 - let tok = crate::tokenizer::load_tokenizer("nvidia/Kimi-K2.5-NVFP4", false, None).unwrap();
696 + async fn test_variable_turns_fixed() {
697 + let tok = crate::tokenizer::load_tokenizer("nvidia/Kimi-K2.5-NVFP4", false, None)
698 + .await
699 + .unwrap();
700
701 let cfg = MultiTurnRandomConfig {
702 num_conversations: 10,
@@ -710,10 +718,12 @@ mod tests {
718 println!("variable_turns_fixed checks passed!");
719 }
720
713 - #[test]
721 + #[tokio::test]
722 #[ignore]
715 - fn test_per_turn_input_len_prefix_sharing() {
716 - let tok = crate::tokenizer::load_tokenizer("nvidia/Kimi-K2.5-NVFP4", false, None).unwrap();
723 + async fn test_per_turn_input_len_prefix_sharing() {
724 + let tok = crate::tokenizer::load_tokenizer("nvidia/Kimi-K2.5-NVFP4", false, None)
725 + .await
726 + .unwrap();
727
728 // Turn 0 input_len=1000, turns 1+ per_turn_input_len=600
729 // global_len ≈ 100 (10%), conv_len ≈ 800 (80%), unique ≈ 100
rust/src/bench/src/datasets/prefix_repetition.rs
+4 -2
@@ -111,8 +111,10 @@ mod tests {
111
112 /// gpt2 via built-in tiktoken encoding — loads without network access.
113 fn test_tokenizer() -> TokenizerKind {
114 - crate::tokenizer::load_tokenizer("gpt2", false, None)
115 - .expect("gpt2 built-in tiktoken should always load without network")
114 + TokenizerKind::Tiktoken(
115 + crate::tiktoken::load_builtin_tiktoken("gpt2")
116 + .expect("gpt2 built-in tiktoken should always load without network"),
117 + )
118 }
119
120 #[test]
rust/src/bench/src/datasets/random.rs
+12 -9
@@ -308,7 +308,8 @@ mod tests {
308 #[test]
309 #[ignore]
310 fn test_generate_random_dataset_token_ids() {
311 - let tokenizer = tokenizer::load_tokenizer("gpt2", false, None).unwrap();
311 + let tokenizer =
312 + TokenizerKind::Tiktoken(crate::tiktoken::load_builtin_tiktoken("gpt2").unwrap());
313 let requests = generate_random_dataset(
314 &tokenizer,
315 10, // num_requests
@@ -340,7 +341,8 @@ mod tests {
341 #[test]
342 #[ignore]
343 fn test_generate_random_dataset_text() {
343 - let tokenizer = tokenizer::load_tokenizer("gpt2", false, None).unwrap();
344 + let tokenizer =
345 + TokenizerKind::Tiktoken(crate::tiktoken::load_builtin_tiktoken("gpt2").unwrap());
346 let requests = generate_random_dataset(
347 &tokenizer,
348 10, // num_requests
@@ -374,7 +376,8 @@ mod tests {
376 #[test]
377 #[ignore]
378 fn test_token_length_exact_local() {
377 - let tokenizer = tokenizer::load_tokenizer("gpt2", false, None).unwrap();
379 + let tokenizer =
380 + TokenizerKind::Tiktoken(crate::tiktoken::load_builtin_tiktoken("gpt2").unwrap());
381 let target_len = 512;
382 let requests = generate_random_dataset(
383 &tokenizer,
@@ -408,11 +411,11 @@ mod tests {
411 }
412
413 /// Test that tiktoken tokenizer produces exact target token lengths (token ID mode).
411 - #[test]
414 + #[tokio::test]
415 #[ignore]
413 - fn test_token_length_exact_tiktoken() {
416 + async fn test_token_length_exact_tiktoken() {
417 // Use Qwen2.5 which has a tiktoken-format tokenizer
415 - let tokenizer = tokenizer::load_tokenizer("Qwen/Qwen2.5-0.5B", false, None);
418 + let tokenizer = tokenizer::load_tokenizer("Qwen/Qwen2.5-0.5B", false, None).await;
419 let tokenizer = match tokenizer {
420 Ok(t) => t,
421 Err(e) => {
@@ -456,10 +459,10 @@ mod tests {
459
460 /// Test encode/decode roundtrip stability for tiktoken.
461 /// After one decode→encode cycle with UTF-8-safe tokens, length must not drift.
459 - #[test]
462 + #[tokio::test]
463 #[ignore]
461 - fn test_tiktoken_roundtrip_stability() {
462 - let tokenizer = tokenizer::load_tokenizer("Qwen/Qwen2.5-0.5B", false, None);
464 + async fn test_tiktoken_roundtrip_stability() {
465 + let tokenizer = tokenizer::load_tokenizer("Qwen/Qwen2.5-0.5B", false, None).await;
466 let tokenizer = match tokenizer {
467 Ok(t) => t,
468 Err(e) => {
rust/src/bench/src/datasets/random_rerank.rs
+4 -2
@@ -139,8 +139,10 @@ mod tests {
139
140 /// gpt2 via built-in tiktoken encoding — loads without network access.
141 fn test_tokenizer() -> TokenizerKind {
142 - crate::tokenizer::load_tokenizer("gpt2", false, None)
143 - .expect("gpt2 built-in tiktoken should always load without network")
142 + TokenizerKind::Tiktoken(
143 + crate::tiktoken::load_builtin_tiktoken("gpt2")
144 + .expect("gpt2 built-in tiktoken should always load without network"),
145 + )
146 }
147
148 fn fixed_ratio() -> RangeRatio {
rust/src/bench/src/datasets/sharegpt.rs
+4 -3
@@ -22,14 +22,15 @@ const DEFAULT_SHAREGPT_FILE: &str = "ShareGPT_V3_unfiltered_cleaned_split.json";
22
23 /// Download the default ShareGPT dataset from HuggingFace Hub.
24 /// Uses hf-hub's built-in cache — subsequent calls return the cached path instantly.
25 -pub fn download_sharegpt_dataset() -> Result<String> {
25 +pub async fn download_sharegpt_dataset() -> Result<String> {
26 tracing::info!(
27 repository = DEFAULT_SHAREGPT_REPO,
28 file = DEFAULT_SHAREGPT_FILE,
29 "downloading ShareGPT dataset"
30 );
31 - let repo = crate::hub::HubRepo::dataset(DEFAULT_SHAREGPT_REPO.to_string());
32 - let path = repo.get(DEFAULT_SHAREGPT_FILE).map_err(|e| {
31 + let repo = crate::hub::HubRepo::dataset(DEFAULT_SHAREGPT_REPO.to_string())
32 + .map_err(BenchError::Config)?;
33 + let path = repo.get(DEFAULT_SHAREGPT_FILE).await.map_err(|e| {
34 BenchError::Config(format!(
35 "Failed to download ShareGPT dataset from '{DEFAULT_SHAREGPT_REPO}': {e}"
36 ))
rust/src/bench/src/datasets/speed_bench.rs
+8 -7
@@ -26,7 +26,7 @@ fn cache_dir() -> std::path::PathBuf {
26
27 /// Download SPEED-Bench dataset from HuggingFace datasets-server API.
28 /// Results are cached as JSON locally for subsequent runs.
29 -pub fn download_speed_bench(config: SpeedBenchConfig) -> Result<String> {
29 +pub async fn download_speed_bench(config: SpeedBenchConfig) -> Result<String> {
30 let config_name = config.as_str();
31
32 let dir = cache_dir();
@@ -42,7 +42,7 @@ pub fn download_speed_bench(config: SpeedBenchConfig) -> Result<String> {
42
43 tracing::info!(config = config_name, "downloading SPEED-Bench dataset");
44
45 - let client = reqwest::blocking::Client::builder()
45 + let client = reqwest::Client::builder()
46 .timeout(std::time::Duration::from_secs(120))
47 .build()
48 .map_err(|e| BenchError::Config(format!("Failed to build HTTP client: {e}")))?;
@@ -66,13 +66,14 @@ pub fn download_speed_bench(config: SpeedBenchConfig) -> Result<String> {
66 let max_retries = 3;
67 let mut data: Option<serde_json::Value> = None;
68 for attempt in 0..=max_retries {
69 - let resp = match client.get(&url).send() {
69 + let resp = match client.get(&url).send().await {
70 Ok(r) => r,
71 Err(e) => {
72 if attempt < max_retries {
73 - std::thread::sleep(std::time::Duration::from_secs(
73 + tokio::time::sleep(std::time::Duration::from_secs(
74 2 * (attempt as u64 + 1),
75 - ));
75 + ))
76 + .await;
77 continue;
78 }
79 return Err(BenchError::Config(format!(
@@ -82,7 +83,7 @@ pub fn download_speed_bench(config: SpeedBenchConfig) -> Result<String> {
83 };
84
85 if resp.status().is_server_error() && attempt < max_retries {
85 - std::thread::sleep(std::time::Duration::from_secs(2 * (attempt as u64 + 1)));
86 + tokio::time::sleep(std::time::Duration::from_secs(2 * (attempt as u64 + 1))).await;
87 continue;
88 }
89
@@ -93,7 +94,7 @@ pub fn download_speed_bench(config: SpeedBenchConfig) -> Result<String> {
94 )));
95 }
96
96 - data = Some(resp.json().map_err(|e| {
97 + data = Some(resp.json().await.map_err(|e| {
98 BenchError::Config(format!("Failed to parse SPEED-Bench API response: {e}"))
99 })?);
100 break;
rust/src/bench/src/hub.rs
+20 -35
@@ -1,54 +1,39 @@
1 // SPDX-License-Identifier: Apache-2.0
2 // SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 -//! Sync facade over the async `hf_hub` API.
5 -//!
6 -//! The workspace bans rustls (`rust/deny.toml`), but hf-hub's sync `ureq`
7 -//! backend unconditionally pulls ureq's default rustls feature. So we use the
8 -//! reqwest/native-tls tokio API instead, and bridge blocking callers (dataset
9 -//! loaders, tokenizer fallback in rayon threads) by running each download on a
10 -//! dedicated thread with its own single-threaded runtime.
11 -
4 use std::path::PathBuf;
5
6 +use hf_hub::Repo;
7 +use hf_hub::api::tokio::{ApiBuilder, ApiRepo};
8 +
9 /// A handle to a HuggingFace Hub repo, downloading via hf-hub's on-disk cache.
10 pub struct HubRepo {
16 - repo: hf_hub::Repo,
11 + repo: ApiRepo,
12 }
13
14 impl HubRepo {
20 - pub fn model(model_id: String) -> Self {
21 - Self {
22 - repo: hf_hub::Repo::model(model_id),
23 - }
15 + pub fn model(model_id: String) -> Result<Self, String> {
16 + Self::new(Repo::model(model_id))
17 }
18
26 - pub fn dataset(repo_id: String) -> Self {
27 - Self {
28 - repo: hf_hub::Repo::dataset(repo_id),
19 + pub fn dataset(repo_id: String) -> Result<Self, String> {
20 + Self::new(Repo::dataset(repo_id))
21 + }
22 +
23 + fn new(repo: Repo) -> Result<Self, String> {
24 + let mut builder = ApiBuilder::from_env();
25 + if let Ok(token) = std::env::var("HF_TOKEN") {
26 + builder = builder.with_token(Some(token));
27 }
28 + let api = builder.build().map_err(|e| format!("Failed to init HF API: {e}"))?;
29 + Ok(Self {
30 + repo: api.repo(repo),
31 + })
32 }
33
34 /// Download (or fetch from cache) a single file from the repo.
35 /// Auth is handled by hf-hub via HF_TOKEN / the cached login token.
34 - pub fn get(&self, filename: &str) -> Result<PathBuf, String> {
35 - let repo = self.repo.clone();
36 - let filename = filename.to_string();
37 - std::thread::spawn(move || {
38 - let rt = tokio::runtime::Builder::new_current_thread()
39 - .enable_all()
40 - .build()
41 - .map_err(|e| format!("Failed to build download runtime: {e}"))?;
42 - rt.block_on(async move {
43 - let mut builder = hf_hub::api::tokio::ApiBuilder::from_env();
44 - if let Ok(token) = std::env::var("HF_TOKEN") {
45 - builder = builder.with_token(Some(token));
46 - }
47 - let api = builder.build().map_err(|e| format!("Failed to init HF API: {e}"))?;
48 - api.repo(repo).get(&filename).await.map_err(|e| format!("{e}"))
49 - })
50 - })
51 - .join()
52 - .map_err(|_| "HF Hub download thread panicked".to_string())?
36 + pub async fn get(&self, filename: &str) -> Result<PathBuf, String> {
37 + self.repo.get(filename).await.map_err(|e| format!("{e}"))
38 }
39 }
rust/src/bench/src/multi_turn.rs
+3 -2
@@ -90,7 +90,8 @@ pub async fn run_multi_turn_benchmark(config: &BenchConfig) -> Result<serde_json
90 let tid = config.tokenizer_id.as_deref().unwrap_or(&model_id);
91 tracing::info!(tokenizer = tid, "loading tokenizer");
92 let server_info = Some((config.base_url.as_str(), model_id.as_str()));
93 - let t = crate::tokenizer::load_tokenizer(tid, config.trust_remote_code, server_info)?;
93 + let t =
94 + crate::tokenizer::load_tokenizer(tid, config.trust_remote_code, server_info).await?;
95 Some(t)
96 };
97
@@ -139,7 +140,7 @@ pub async fn run_multi_turn_benchmark(config: &BenchConfig) -> Result<serde_json
140 let path = match config.dataset_path.as_deref() {
141 Some(p) => p,
142 None => {
142 - downloaded = crate::datasets::sharegpt::download_sharegpt_dataset()?;
143 + downloaded = crate::datasets::sharegpt::download_sharegpt_dataset().await?;
144 downloaded.as_str()
145 }
146 };
rust/src/bench/src/tiktoken.rs
+22 -16
@@ -209,7 +209,7 @@ pub fn load_builtin_tiktoken(encoding: &str) -> Result<TiktokenTokenizer> {
209 }
210
211 /// Try to load a tiktoken tokenizer from a local directory or HuggingFace model repo.
212 -pub fn try_load_tiktoken(model_id: &str) -> Result<TiktokenTokenizer> {
212 +pub async fn try_load_tiktoken(model_id: &str) -> Result<TiktokenTokenizer> {
213 // Phase 1: If model_id is a local directory, look for tiktoken files there
214 let local_dir = Path::new(model_id);
215 if local_dir.is_dir() {
@@ -217,7 +217,7 @@ pub fn try_load_tiktoken(model_id: &str) -> Result<TiktokenTokenizer> {
217 }
218
219 // Phase 2: Fall back to HuggingFace Hub download
220 - try_load_tiktoken_from_hf(model_id)
220 + try_load_tiktoken_from_hf(model_id).await
221 }
222
223 /// Common tiktoken model filenames to search for.
@@ -252,25 +252,28 @@ fn try_load_tiktoken_from_dir(dir: &Path, model_id: &str) -> Result<TiktokenToke
252 }
253
254 /// Load a tiktoken tokenizer from a HuggingFace model repo.
255 -fn try_load_tiktoken_from_hf(model_id: &str) -> Result<TiktokenTokenizer> {
256 - let repo = crate::hub::HubRepo::model(model_id.to_string());
257 -
258 - let model_path = repo
259 - .get("tiktoken.model")
260 - .or_else(|_| repo.get("qwen.tiktoken"))
261 - .or_else(|_| repo.get("vocab.tiktoken"))
262 - .map_err(|_| {
263 - BenchError::Tokenizer(format!("No tiktoken model file found for '{model_id}'"))
264 - })?;
255 +async fn try_load_tiktoken_from_hf(model_id: &str) -> Result<TiktokenTokenizer> {
256 + let repo = crate::hub::HubRepo::model(model_id.to_string()).map_err(BenchError::Tokenizer)?;
257 +
258 + let mut model_path = None;
259 + for filename in TIKTOKEN_MODEL_FILENAMES {
260 + if let Ok(path) = repo.get(filename).await {
261 + model_path = Some(path);
262 + break;
263 + }
264 + }
265 + let model_path = model_path.ok_or_else(|| {
266 + BenchError::Tokenizer(format!("No tiktoken model file found for '{model_id}'"))
267 + })?;
268
269 let num_base_tokens = count_base_tokens(&model_path)?;
270
268 - let config = match repo.get("tokenizer_config.json") {
271 + let config = match repo.get("tokenizer_config.json").await {
272 Ok(config_path) => read_tokenizer_config(&config_path),
273 Err(_) => None,
274 };
275
273 - let pattern = extract_pat_str_from_repo(&repo);
276 + let pattern = extract_pat_str_from_repo(&repo).await;
277
278 build_tiktoken(model_id, &model_path, config, pattern, num_base_tokens)
279 }
@@ -403,9 +406,12 @@ fn extract_pat_str_from_local_dir(dir: &Path) -> Option<String> {
406
407 /// Try to download the Python tokenizer source file and extract pat_str via regex.
408 /// Returns None if unavailable or unparsable.
406 -fn extract_pat_str_from_repo(repo: &crate::hub::HubRepo) -> Option<String> {
409 +async fn extract_pat_str_from_repo(repo: &crate::hub::HubRepo) -> Option<String> {
410 // Try common Python tokenizer filenames
408 - let py_path = repo.get("tokenization_kimi.py").or_else(|_| repo.get("tokenizer.py")).ok()?;
411 + let py_path = match repo.get("tokenization_kimi.py").await {
412 + Ok(path) => path,
413 + Err(_) => repo.get("tokenizer.py").await.ok()?,
414 + };
415
416 let source = std::fs::read_to_string(&py_path).ok()?;
417
rust/src/bench/src/tokenizer.rs
+71 -12
@@ -2,6 +2,7 @@
2 // SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 use std::collections::HashSet;
5 +use std::future::Future;
6 use std::path::Path;
7
8 use thiserror_ext::AsReport as _;
@@ -19,7 +20,8 @@ pub enum TokenizerKind {
20
21 /// Server-side tokenizer using vLLM's /tokenize and /detokenize endpoints.
22 pub struct ServerTokenizer {
22 - client: reqwest::blocking::Client,
23 + client: reqwest::Client,
24 + runtime: tokio::runtime::Handle,
25 tokenize_url: String,
26 detokenize_url: String,
27 model: String,
@@ -28,8 +30,8 @@ pub struct ServerTokenizer {
30
31 impl ServerTokenizer {
32 /// Create a new server tokenizer and verify connectivity.
31 - pub fn new(base_url: &str, model: &str) -> Result<Self> {
32 - let client = reqwest::blocking::Client::builder()
33 + pub async fn new(base_url: &str, model: &str) -> Result<Self> {
34 + let client = reqwest::Client::builder()
35 .timeout(std::time::Duration::from_secs(30))
36 .build()
37 .map_err(|e| BenchError::Tokenizer(format!("Failed to build HTTP client: {e}")))?;
@@ -39,6 +41,7 @@ impl ServerTokenizer {
41
42 let st = Self {
43 client,
44 + runtime: tokio::runtime::Handle::current(),
45 tokenize_url,
46 detokenize_url,
47 model: model.to_string(),
@@ -46,7 +49,7 @@ impl ServerTokenizer {
49 };
50
51 // Probe the endpoint to verify it works and discover vocab size
49 - let test_tokens = st.encode_inner("test")?;
52 + let test_tokens = st.encode_async("test").await?;
53 let max_id = test_tokens.iter().copied().max().unwrap_or(0);
54 let estimated_vocab = (max_id * 2).max(131072);
55
@@ -57,6 +60,10 @@ impl ServerTokenizer {
60 }
61
62 fn encode_inner(&self, text: &str) -> Result<Vec<u32>> {
63 + self.block_on(self.encode_async(text))
64 + }
65 +
66 + async fn encode_async(&self, text: &str) -> Result<Vec<u32>> {
67 let payload = serde_json::json!({
68 "model": self.model,
69 "prompt": text,
@@ -67,6 +74,7 @@ impl ServerTokenizer {
74 .post(&self.tokenize_url)
75 .json(&payload)
76 .send()
77 + .await
78 .map_err(|e| BenchError::Tokenizer(format!("Server tokenize failed: {e}")))?;
79
80 if !resp.status().is_success() {
@@ -76,7 +84,7 @@ impl ServerTokenizer {
84 )));
85 }
86
79 - let data: serde_json::Value = resp.json().map_err(|e| {
87 + let data: serde_json::Value = resp.json().await.map_err(|e| {
88 BenchError::Tokenizer(format!("Failed to parse tokenize response: {e}"))
89 })?;
90
@@ -96,6 +104,10 @@ impl ServerTokenizer {
104 }
105
106 fn decode_inner(&self, ids: &[u32]) -> Result<String> {
107 + self.block_on(self.decode_async(ids))
108 + }
109 +
110 + async fn decode_async(&self, ids: &[u32]) -> Result<String> {
111 let payload = serde_json::json!({
112 "model": self.model,
113 "tokens": ids,
@@ -106,6 +118,7 @@ impl ServerTokenizer {
118 .post(&self.detokenize_url)
119 .json(&payload)
120 .send()
121 + .await
122 .map_err(|e| BenchError::Tokenizer(format!("Server detokenize failed: {e}")))?;
123
124 if !resp.status().is_success() {
@@ -115,7 +128,7 @@ impl ServerTokenizer {
128 )));
129 }
130
118 - let data: serde_json::Value = resp.json().map_err(|e| {
131 + let data: serde_json::Value = resp.json().await.map_err(|e| {
132 BenchError::Tokenizer(format!("Failed to parse detokenize response: {e}"))
133 })?;
134
@@ -124,6 +137,26 @@ impl ServerTokenizer {
137 .map(|s| s.to_string())
138 .ok_or_else(|| BenchError::Tokenizer("Missing 'prompt' in detokenize response".into()))
139 }
140 +
141 + fn block_on<T>(&self, future: impl Future<Output = Result<T>>) -> Result<T> {
142 + if matches!(
143 + self.runtime.runtime_flavor(),
144 + tokio::runtime::RuntimeFlavor::CurrentThread
145 + ) {
146 + return Err(BenchError::Tokenizer(
147 + "Server tokenizer fallback requires a multi-thread Tokio runtime".into(),
148 + ));
149 + }
150 +
151 + // Sync tokenizer calls can come from a Tokio worker or a Rayon worker.
152 + // Tokio workers must enter a blocking region before re-entering the runtime;
153 + // Rayon workers can drive the future directly with the saved runtime handle.
154 + if tokio::runtime::Handle::try_current().is_ok() {
155 + tokio::task::block_in_place(|| self.runtime.block_on(future))
156 + } else {
157 + self.runtime.block_on(future)
158 + }
159 + }
160 }
161
162 // --- TokenizerKind methods ---
@@ -193,7 +226,7 @@ impl TokenizerKind {
226 /// 3. Server-side /tokenize + /detokenize endpoints
227 ///
228 /// `server_info` is `Some((base_url, model))` to enable server-side fallback.
196 -pub fn load_tokenizer(
229 +pub async fn load_tokenizer(
230 model_id: &str,
231 _trust_remote_code: bool,
232 server_info: Option<(&str, &str)>,
@@ -213,7 +246,7 @@ pub fn load_tokenizer(
246 }
247
248 // 1. Try local HuggingFace tokenizer (tokenizer.json)
216 - match try_load_local(model_id) {
249 + match try_load_local(model_id).await {
250 Ok(tok) => {
251 tracing::info!(
252 model = model_id,
@@ -230,7 +263,7 @@ pub fn load_tokenizer(
263 error = %local_err.as_report(),
264 "local tokenizer unavailable; trying tiktoken"
265 );
233 - match crate::tiktoken::try_load_tiktoken(model_id) {
266 + match crate::tiktoken::try_load_tiktoken(model_id).await {
267 Ok(tok) => {
268 tracing::info!(
269 model = model_id,
@@ -248,7 +281,7 @@ pub fn load_tokenizer(
281 error = %tiktoken_err.as_report(),
282 "tiktoken unavailable; trying server-side tokenization"
283 );
251 - match ServerTokenizer::new(base_url, model) {
284 + match ServerTokenizer::new(base_url, model).await {
285 Ok(srv) => {
286 tracing::info!(
287 model = model_id,
@@ -282,7 +315,7 @@ pub fn load_tokenizer(
315 }
316
317 /// Try loading tokenizer.json from local path or HuggingFace Hub.
285 -fn try_load_local(model_id: &str) -> Result<Tokenizer> {
318 +async fn try_load_local(model_id: &str) -> Result<Tokenizer> {
319 // 1. Try local directory with tokenizer.json
320 let local_path = Path::new(model_id).join("tokenizer.json");
321 if local_path.exists() {
@@ -308,11 +341,37 @@ fn try_load_local(model_id: &str) -> Result<Tokenizer> {
341 }
342
343 // 4. Download from HuggingFace Hub (hf-hub handles auth via HF_TOKEN / cached token)
311 - let repo = crate::hub::HubRepo::model(model_id.to_string());
344 + let repo = crate::hub::HubRepo::model(model_id.to_string()).map_err(BenchError::Tokenizer)?;
345 let tokenizer_path = repo
346 .get("tokenizer.json")
347 + .await
348 .map_err(|e| BenchError::Tokenizer(format!("No tokenizer.json for '{model_id}': {e}")))?;
349
350 Tokenizer::from_file(&tokenizer_path)
351 .map_err(|e| BenchError::Tokenizer(format!("Failed to load downloaded tokenizer: {e}")))
352 }
353 +
354 +#[cfg(test)]
355 +mod tests {
356 + use super::*;
357 +
358 + #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
359 + async fn test_server_tokenizer_sync_bridge() {
360 + let tokenizer = std::sync::Arc::new(ServerTokenizer {
361 + client: reqwest::Client::new(),
362 + runtime: tokio::runtime::Handle::current(),
363 + tokenize_url: String::new(),
364 + detokenize_url: String::new(),
365 + model: String::new(),
366 + cached_vocab_size: 0,
367 + });
368 +
369 + assert_eq!(tokenizer.block_on(async { Ok(1) }).unwrap(), 1);
370 +
371 + let (tx, rx) = tokio::sync::oneshot::channel();
372 + rayon::spawn(move || {
373 + let _ = tx.send(tokenizer.block_on(async { Ok(2) }));
374 + });
375 + assert_eq!(rx.await.unwrap().unwrap(), 2);
376 + }
377 +}