[Bugfix][Benchmark] Check readiness before tokenizer init in rust vllm-bench (#51863)
Signed-off-by: Tyler Michael Smith <tlrmchlsmth@gmail.com> Co-authored-by: OpenAI Codex <codex@openai.com>
Tyler Michael Smith committed
Aug 18, 2026 at 22:52 UTC
aabc1a0a0b55c768ab1f8962157f5fa086932b88
4 files changed
+159
-81
rust/src/bench/src/benchmark.rs
+13
-38
@@ -18,7 +18,7 @@ use crate::metrics::steady_state;
18
use crate::output::console::print_results;
19
use crate::output::json::{append_result, build_result_json, compute_result_filename, save_result};
20
use crate::rate_control::compute_schedule;
21
-use crate::ready_checker::wait_for_endpoint;
21
+use crate::ready_checker::{get_first_model, wait_for_endpoint};
22
23
/// Pre-resolve the hostname in `base_url` and pin all resolved IPs on the
24
/// client builder via [`reqwest::ClientBuilder::resolve_to_addrs`]. This
@@ -281,40 +281,6 @@ pub(crate) fn assign_lora_modules(
281
Some(out)
282
}
283
284
-/// Fetch the first model from the server's /v1/models endpoint.
285
-async fn get_first_model_from_server(
286
- base_url: &str,
287
- client: &reqwest::Client,
288
- extra_headers: &Option<std::collections::HashMap<String, String>>,
289
-) -> Result<(String, String)> {
290
- let url = format!("{base_url}/v1/models");
291
- let mut request = client.get(&url);
292
- if let Some(headers) = extra_headers {
293
- for (k, v) in headers {
294
- request = request.header(k, v);
295
- }
296
- }
297
- // Add API key from environment
298
- if let Ok(api_key) = std::env::var("OPENAI_API_KEY") {
299
- request = request.header("Authorization", format!("Bearer {api_key}"));
300
- }
301
-
302
- let response = request.send().await?;
303
- let data: serde_json::Value = response.json().await?;
304
-
305
- if let Some(models) = data.get("data").and_then(|d| d.as_array())
306
- && let Some(first) = models.first()
307
- {
308
- let id = first.get("id").and_then(|v| v.as_str()).unwrap_or_default().to_string();
309
- let root = first.get("root").and_then(|v| v.as_str()).unwrap_or(&id).to_string();
310
- return Ok((id, root));
311
- }
312
-
313
- Err(BenchError::Config(format!(
314
- "No models found on the server at {base_url}"
315
- )))
316
-}
317
-
284
/// Run the complete benchmark.
285
///
286
/// This is the core orchestrator matching Python's benchmark() + main_async().
@@ -348,8 +314,13 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result<serde_json::Value> {
314
(m.clone(), config.model_name.clone())
315
} else {
316
tracing::info!(base_url = %config.base_url, "fetching first model from server");
351
- let (name, id) =
352
- get_first_model_from_server(&config.base_url, &client, &config.extra_headers).await?;
317
+ let (name, id) = get_first_model(
318
+ &config.base_url,
319
+ &client,
320
+ &config.extra_headers,
321
+ config.ready_check_timeout_sec,
322
+ )
323
+ .await?;
324
tracing::info!(
325
model_name = name,
326
model_id = id,
@@ -364,7 +335,11 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result<serde_json::Value> {
335
} else {
336
let tid = config.tokenizer_id.as_deref().unwrap_or(&model_id);
337
tracing::info!(tokenizer = tid, "loading tokenizer");
367
- let server_info = Some((config.base_url.as_str(), model_id.as_str()));
338
+ let server_info = Some((
339
+ config.base_url.as_str(),
340
+ model_id.as_str(),
341
+ config.ready_check_timeout_sec,
342
+ ));
343
let t =
344
crate::tokenizer::load_tokenizer(tid, config.trust_remote_code, server_info).await?;
345
Some(t)
rust/src/bench/src/multi_turn.rs
+14
-36
@@ -23,6 +23,7 @@ use crate::output::console::print_multi_turn_results;
23
use crate::output::json::{
24
append_result, build_multi_turn_result_json, compute_result_filename, save_result,
25
};
26
+use crate::ready_checker::{get_first_model, wait_for_endpoint};
27
28
/// Output from a single turn within a conversation.
29
#[derive(Debug, Clone)]
@@ -74,7 +75,13 @@ pub async fn run_multi_turn_benchmark(config: &BenchConfig) -> Result<serde_json
75
(m.clone(), config.model_name.clone())
76
} else {
77
tracing::info!(base_url = %config.base_url, "fetching first model from server");
77
- let (name, id) = get_first_model(&config.base_url, &client, &config.extra_headers).await?;
78
+ let (name, id) = get_first_model(
79
+ &config.base_url,
80
+ &client,
81
+ &config.extra_headers,
82
+ config.ready_check_timeout_sec,
83
+ )
84
+ .await?;
85
tracing::info!(
86
model_name = name,
87
model_id = id,
@@ -89,7 +96,11 @@ pub async fn run_multi_turn_benchmark(config: &BenchConfig) -> Result<serde_json
96
} else {
97
let tid = config.tokenizer_id.as_deref().unwrap_or(&model_id);
98
tracing::info!(tokenizer = tid, "loading tokenizer");
92
- let server_info = Some((config.base_url.as_str(), model_id.as_str()));
99
+ let server_info = Some((
100
+ config.base_url.as_str(),
101
+ model_id.as_str(),
102
+ config.ready_check_timeout_sec,
103
+ ));
104
let t =
105
crate::tokenizer::load_tokenizer(tid, config.trust_remote_code, server_info).await?;
106
Some(t)
@@ -263,7 +274,7 @@ pub async fn run_multi_turn_benchmark(config: &BenchConfig) -> Result<serde_json
274
};
275
276
tracing::info!("starting initial single-prompt test run");
266
- let test_output = crate::ready_checker::wait_for_endpoint(
277
+ let test_output = wait_for_endpoint(
278
config.backend,
279
&client,
280
&test_input,
@@ -748,39 +759,6 @@ async fn run_conversation(
759
}
760
}
761
751
-/// Fetch the first model from the server's /v1/models endpoint.
752
-async fn get_first_model(
753
- base_url: &str,
754
- client: &reqwest::Client,
755
- extra_headers: &Option<HashMap<String, String>>,
756
-) -> Result<(String, String)> {
757
- let url = format!("{base_url}/v1/models");
758
- let mut request = client.get(&url);
759
- if let Some(headers) = extra_headers {
760
- for (k, v) in headers {
761
- request = request.header(k, v);
762
- }
763
- }
764
- if let Ok(api_key) = std::env::var("OPENAI_API_KEY") {
765
- request = request.header("Authorization", format!("Bearer {api_key}"));
766
- }
767
-
768
- let response = request.send().await?;
769
- let data: serde_json::Value = response.json().await?;
770
-
771
- if let Some(models) = data.get("data").and_then(|d| d.as_array())
772
- && let Some(first) = models.first()
773
- {
774
- let id = first.get("id").and_then(|v| v.as_str()).unwrap_or_default().to_string();
775
- let root = first.get("root").and_then(|v| v.as_str()).unwrap_or(&id).to_string();
776
- return Ok((id, root));
777
- }
778
-
779
- Err(BenchError::Config(format!(
780
- "No models found on the server at {base_url}"
781
- )))
782
-}
783
-
762
#[cfg(test)]
763
mod tests {
764
use std::sync::Arc;
rust/src/bench/src/ready_checker.rs
+123
-1
@@ -1,14 +1,96 @@
1
// SPDX-License-Identifier: Apache-2.0
2
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
+use std::future::Future;
5
use std::time::Instant;
6
7
use indicatif::{ProgressBar, ProgressStyle};
8
+use thiserror_ext::AsReport as _;
9
10
use crate::backends::{RequestFuncInput, RequestFuncOutput, get_backend};
11
use crate::cli::BackendKind;
12
use crate::error::{BenchError, Result};
13
14
+/// Retry an operation until it succeeds or the readiness timeout expires.
15
+pub(crate) async fn retry_with_timeout<T, F, Fut>(
16
+ timeout_seconds: u64,
17
+ retry_interval: u64,
18
+ mut operation: F,
19
+) -> Result<T>
20
+where
21
+ F: FnMut() -> Fut,
22
+ Fut: Future<Output = Result<T>>,
23
+{
24
+ if timeout_seconds == 0 {
25
+ return operation().await;
26
+ }
27
+
28
+ let deadline = Instant::now() + std::time::Duration::from_secs(timeout_seconds);
29
+
30
+ loop {
31
+ let error = match operation().await {
32
+ Ok(value) => return Ok(value),
33
+ Err(error) => error,
34
+ };
35
+ tracing::warn!(error = %error.as_report(), "endpoint is not ready");
36
+
37
+ let remaining = deadline.saturating_duration_since(Instant::now());
38
+ if remaining.is_zero() {
39
+ return Err(BenchError::EndpointTimeout(
40
+ timeout_seconds,
41
+ format!("{}", error.as_report()),
42
+ ));
43
+ }
44
+
45
+ let sleep_duration =
46
+ std::cmp::min(std::time::Duration::from_secs(retry_interval), remaining);
47
+ if !sleep_duration.is_zero() {
48
+ tokio::time::sleep(sleep_duration).await;
49
+ }
50
+ }
51
+}
52
+
53
+/// Fetch the first model from the server, retrying while it starts.
54
+pub(crate) async fn get_first_model(
55
+ base_url: &str,
56
+ client: &reqwest::Client,
57
+ extra_headers: &Option<std::collections::HashMap<String, String>>,
58
+ timeout_seconds: u64,
59
+) -> Result<(String, String)> {
60
+ let url = format!("{base_url}/v1/models");
61
+ retry_with_timeout(timeout_seconds, 5, || async {
62
+ let mut request = client.get(&url);
63
+ if let Some(headers) = extra_headers {
64
+ for (key, value) in headers {
65
+ request = request.header(key, value);
66
+ }
67
+ }
68
+ // Add API key from environment
69
+ if let Ok(api_key) = std::env::var("OPENAI_API_KEY") {
70
+ request = request.header("Authorization", format!("Bearer {api_key}"));
71
+ }
72
+
73
+ let response = request.send().await?.error_for_status()?;
74
+ let data: serde_json::Value = response.json().await?;
75
+ if let Some(model) = data
76
+ .get("data")
77
+ .and_then(|value| value.as_array())
78
+ .and_then(|models| models.first())
79
+ {
80
+ let id =
81
+ model.get("id").and_then(|value| value.as_str()).unwrap_or_default().to_string();
82
+ let root =
83
+ model.get("root").and_then(|value| value.as_str()).unwrap_or(&id).to_string();
84
+ return Ok((id, root));
85
+ }
86
+
87
+ Err(BenchError::Config(format!(
88
+ "No models found on the server at {base_url}"
89
+ )))
90
+ })
91
+ .await
92
+}
93
+
94
/// Wait for the serving endpoint to become available.
95
///
96
/// Sends test requests with retry until success or timeout.
@@ -63,7 +145,7 @@ pub async fn wait_for_endpoint(
145
last_error = err;
146
}
147
Err(e) => {
66
- last_error = e.to_string();
148
+ last_error = format!("{}", e.as_report());
149
}
150
}
151
@@ -76,3 +158,43 @@ pub async fn wait_for_endpoint(
158
159
Err(BenchError::EndpointTimeout(timeout_seconds, last_error))
160
}
161
+
162
+#[cfg(test)]
163
+mod tests {
164
+ use std::sync::atomic::{AtomicUsize, Ordering};
165
+
166
+ use super::*;
167
+
168
+ #[tokio::test]
169
+ async fn retry_succeeds_after_transient_failures() {
170
+ let attempts = AtomicUsize::new(0);
171
+
172
+ let value = retry_with_timeout(1, 0, || async {
173
+ if attempts.fetch_add(1, Ordering::Relaxed) < 2 {
174
+ Err(BenchError::Backend("not ready".into()))
175
+ } else {
176
+ Ok(42)
177
+ }
178
+ })
179
+ .await
180
+ .unwrap();
181
+
182
+ assert_eq!(value, 42);
183
+ assert_eq!(attempts.load(Ordering::Relaxed), 3);
184
+ }
185
+
186
+ #[tokio::test]
187
+ async fn zero_timeout_does_not_retry() {
188
+ let attempts = AtomicUsize::new(0);
189
+
190
+ let error = retry_with_timeout(0, 0, || async {
191
+ attempts.fetch_add(1, Ordering::Relaxed);
192
+ Err::<(), _>(BenchError::Backend("not ready".into()))
193
+ })
194
+ .await
195
+ .unwrap_err();
196
+
197
+ assert!(matches!(error, BenchError::Backend(_)));
198
+ assert_eq!(attempts.load(Ordering::Relaxed), 1);
199
+ }
200
+}
rust/src/bench/src/tokenizer.rs
+9
-6
@@ -30,7 +30,7 @@ pub struct ServerTokenizer {
30
31
impl ServerTokenizer {
32
/// Create a new server tokenizer and verify connectivity.
33
- pub async fn new(base_url: &str, model: &str) -> Result<Self> {
33
+ pub async fn new(base_url: &str, model: &str, timeout_seconds: u64) -> Result<Self> {
34
let client = reqwest::Client::builder()
35
.timeout(std::time::Duration::from_secs(30))
36
.build()
@@ -49,7 +49,10 @@ impl ServerTokenizer {
49
};
50
51
// Probe the endpoint to verify it works and discover vocab size
52
- let test_tokens = st.encode_async("test").await?;
52
+ let test_tokens = crate::ready_checker::retry_with_timeout(timeout_seconds, 5, || {
53
+ st.encode_async("test")
54
+ })
55
+ .await?;
56
let max_id = test_tokens.iter().copied().max().unwrap_or(0);
57
let estimated_vocab = (max_id * 2).max(131072);
58
@@ -225,11 +228,11 @@ impl TokenizerKind {
228
/// 2. Tiktoken model file (for Kimi, Qwen, etc.)
229
/// 3. Server-side /tokenize + /detokenize endpoints
230
///
228
-/// `server_info` is `Some((base_url, model))` to enable server-side fallback.
231
+/// `server_info` is `Some((base_url, model, timeout_seconds))` to enable server-side fallback.
232
pub async fn load_tokenizer(
233
model_id: &str,
234
_trust_remote_code: bool,
232
- server_info: Option<(&str, &str)>,
235
+ server_info: Option<(&str, &str, u64)>,
236
) -> Result<TokenizerKind> {
237
// 0. Check for built-in tiktoken encoding names (no HF download needed). These are useful for
238
// consistent cross-model token counting (e.g. Artificial Analysis).
@@ -275,13 +278,13 @@ pub async fn load_tokenizer(
278
}
279
Err(tiktoken_err) => {
280
// 3. Try server-side fallback
278
- if let Some((base_url, model)) = server_info {
281
+ if let Some((base_url, model, timeout_seconds)) = server_info {
282
tracing::info!(
283
model = model_id,
284
error = %tiktoken_err.as_report(),
285
"tiktoken unavailable; trying server-side tokenization"
286
);
284
- match ServerTokenizer::new(base_url, model).await {
287
+ match ServerTokenizer::new(base_url, model, timeout_seconds).await {
288
Ok(srv) => {
289
tracing::info!(
290
model = model_id,