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

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

Browse Source 
Signed-off-by: Bugen Zhao <i@bugenzhao.com>
This commit is contained in:
Bugen Zhao
2026-07-21 16:53:57 +08:00
committed by GitHub
parent 7a98c7a392
commit eb44b3aaa4
15 changed files with 199 additions and 121 deletions
Download Patch File Download Diff File Expand all files Collapse all files
rust/Cargo.lock
Generated
-1
View File
@@ -3446,7 +3446,6 @@ dependencies = [
"base64 0.22.1",
"bytes",
"encoding_rs",
"futures-channel",
"futures-core",
"futures-util",
"h2",
rust/src/bench/Cargo.toml
+1 -1
View File
@@ -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"] }
rust/src/bench/src/benchmark.rs
+7 -4
View File
@@ -365,7 +365,8 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result<serde_json::Value> {
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<serde_json::Value> {
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<serde_json::Value> {
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<serde_json::Value> {
config.hf_subset.as_deref(),
config.hf_split.as_deref(),
config.num_prompts,
)?;
)
.await?;
crate::datasets::hf_dataset::load_hf_dataset(
tok,
&downloaded_path,
rust/src/bench/src/datasets/custom.rs
+4 -2
View File
@@ -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]
rust/src/bench/src/datasets/hf_dataset.rs
+15 -11
View File
@@ -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<serde_json::Value> {
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.
rust/src/bench/src/datasets/multi_turn.rs
+25 -15
View File
@@ -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
rust/src/bench/src/datasets/prefix_repetition.rs
+4 -2
View File
@@ -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]
rust/src/bench/src/datasets/random.rs
+12 -9
View File
@@ -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) => {
rust/src/bench/src/datasets/random_rerank.rs
+4 -2
View File
@@ -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 {
rust/src/bench/src/datasets/sharegpt.rs
+4 -3
View File
@@ -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<String> {
pub async fn download_sharegpt_dataset() -> Result<String> {
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}"
))
rust/src/bench/src/datasets/speed_bench.rs
+8 -7
View File
@@ -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<String> {
pub async fn download_speed_bench(config: SpeedBenchConfig) -> Result<String> {
let config_name = config.as_str();
let dir = cache_dir();
@@ -42,7 +42,7 @@ pub fn download_speed_bench(config: SpeedBenchConfig) -> Result<String> {
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<String> {
let max_retries = 3;
let mut data: Option<serde_json::Value> = 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<String> {
};
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<String> {
)));
}
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;
rust/src/bench/src/hub.rs
+20 -35
View File
@@ -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, String> {
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, String> {
Self::new(Repo::dataset(repo_id))
}
fn new(repo: Repo) -> Result<Self, String> {
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<PathBuf, String> {
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<PathBuf, String> {
self.repo.get(filename).await.map_err(|e| format!("{e}"))
}
}
rust/src/bench/src/multi_turn.rs
+3 -2
View File
@@ -90,7 +90,8 @@ pub async fn run_multi_turn_benchmark(config: &BenchConfig) -> Result<serde_json
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)
};
@@ -139,7 +140,7 @@ pub async fn run_multi_turn_benchmark(config: &BenchConfig) -> Result<serde_json
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()
}
};
rust/src/bench/src/tiktoken.rs
+21 -15
View File
@@ -209,7 +209,7 @@ pub fn load_builtin_tiktoken(encoding: &str) -> Result<TiktokenTokenizer> {
}
/// Try to load a tiktoken tokenizer from a local directory or HuggingFace model repo.
pub fn try_load_tiktoken(model_id: &str) -> Result<TiktokenTokenizer> {
pub async fn try_load_tiktoken(model_id: &str) -> Result<TiktokenTokenizer> {
// 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<TiktokenTokenizer> {
}
// 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<TiktokenToke
}
/// Load a tiktoken tokenizer from a HuggingFace model repo.
fn try_load_tiktoken_from_hf(model_id: &str) -> Result<TiktokenTokenizer> {
let repo = crate::hub::HubRepo::model(model_id.to_string());
async fn try_load_tiktoken_from_hf(model_id: &str) -> Result<TiktokenTokenizer> {
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<String> {
/// 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<String> {
async fn extract_pat_str_from_repo(repo: &crate::hub::HubRepo) -> Option<String> {
// 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()?;
rust/src/bench/src/tokenizer.rs
+71 -12
View File
@@ -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<Self> {
let client = reqwest::blocking::Client::builder()
pub async fn new(base_url: &str, model: &str) -> Result<Self> {
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<Vec<u32>> {
self.block_on(self.encode_async(text))
}
async fn encode_async(&self, text: &str) -> Result<Vec<u32>> {
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<String> {
self.block_on(self.decode_async(ids))
}
async fn decode_async(&self, ids: &[u32]) -> Result<String> {
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<T>(&self, future: impl Future<Output = Result<T>>) -> Result<T> {
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<Tokenizer> {
async fn try_load_local(model_id: &str) -> Result<Tokenizer> {
// 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<Tokenizer> {
}
// 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);
}
}