[Rust Frontend] Honor HF revisions, offline mode, and cache directory (#56386)

Signed-off-by: Bugen Zhao <i@bugenzhao.com>
Co-authored-by: Codex <noreply@openai.com>
This commit is contained in:
Bugen Zhao
2026-09-12 03:39:17 +00:00
committed by GitHub
co-authored by Codex
parent 5f3e4b7447
commit 485421b1c3
11 changed files with 201 additions and 28 deletions
+2 -1
View File
@@ -129,7 +129,7 @@ pub(super) async fn load_model_backends(
model_id: &str,
options: LoadModelBackendsOptions,
) -> Result<LoadedModelBackends> {
let files = ResolvedModelFiles::new(model_id).await?;
let files = ResolvedModelFiles::new(model_id, options.revision.as_deref()).await?;
let text_backend = HfTextBackend::from_resolved_model_files(
files.clone(),
model_id.to_string(),
@@ -231,6 +231,7 @@ mod tests {
resolved_files(config_json, tokenizer_config_json),
"test-model".to_string(),
LoadModelBackendsOptions {
revision: None,
generation_config: Default::default(),
renderer,
language_model_only: false,
+2
View File
@@ -61,6 +61,8 @@ pub type DynChatTextBackend = Arc<dyn ChatTextBackend>;
/// Frontend-side chat backend loading options.
#[derive(Debug, Clone, PartialEq, Eq, Default)]
pub struct LoadModelBackendsOptions {
/// Model revision on the Hugging Face Hub (branch, tag, or commit SHA).
pub revision: Option<String>,
/// Which generation-config sampling defaults to inherit.
pub generation_config: GenerationConfigMode,
/// Which chat renderer implementation to use.
+12
View File
@@ -110,6 +110,9 @@ pub enum BenchCommand {
pub struct RenderArgs {
/// Model identifier or local model directory containing tokenizer files.
model: String,
/// Model revision on the Hugging Face Hub (branch, tag, or commit SHA).
#[arg(long)]
revision: Option<String>,
/// HTTP bind host.
#[arg(long, default_value = "127.0.0.1")]
host: String,
@@ -157,6 +160,7 @@ impl RenderArgs {
RenderConfig {
model: self.model,
revision: self.revision,
served_model_name: self.served_model_name,
host: self.host,
port: self.port,
@@ -195,6 +199,11 @@ pub struct SharedRuntimeArgs {
/// public model ID.
pub model: String,
/// Model revision on the Hugging Face Hub (branch, tag, or commit SHA).
#[arg(long)]
#[serde(default)]
pub revision: Option<String>,
/// The source of generation-config sampling defaults. `"auto"` loads the
/// model's defaults, while `"vllm"` uses vLLM's neutral defaults.
#[arg(long, default_value_t)]
@@ -481,6 +490,7 @@ impl SharedRuntimeArgs {
None => CoordinatorMode::None,
},
model: self.model,
revision: self.revision,
generation_config: self.generation_config,
served_model_name: self.served_model_name,
listener_mode: HttpListenerMode::InheritedFd { fd: listen_fd },
@@ -536,6 +546,7 @@ impl SharedRuntimeArgs {
},
coordinator_mode: CoordinatorMode::MaybeInProc,
model: self.model,
revision: self.revision,
generation_config: self.generation_config,
served_model_name: self.served_model_name,
listener_mode,
@@ -743,6 +754,7 @@ impl ServeArgs {
self.managed_engine.clone().into_config(
self.runtime.model.clone(),
self.runtime.revision.clone(),
self.runtime.max_logprobs,
profiler_config,
reasoning_parser.as_deref(),
+26 -1
View File
@@ -36,6 +36,8 @@ fn render_args_build_config_without_tls() {
"vllm-rs",
"render",
"Qwen/Qwen2.5-0.5B-Instruct",
"--revision",
"release",
"--port",
"8080",
"--max-model-len",
@@ -55,6 +57,7 @@ fn render_args_build_config_without_tls() {
let config = args.into_config();
assert_eq!(config.model, "Qwen/Qwen2.5-0.5B-Instruct");
assert_eq!(config.revision.as_deref(), Some("release"));
assert_eq!(config.host, "127.0.0.1");
assert_eq!(config.port, 8080);
assert_eq!(config.max_model_len, Some(32768));
@@ -173,6 +176,7 @@ fn serve_args_forward_python_flags_with_separator() {
uds: None,
runtime: SharedRuntimeArgs {
model: "Qwen/Qwen3-0.6B",
revision: None,
generation_config: Auto,
engine_ready_timeout_secs: 600,
tool_call_parser: Auto,
@@ -1007,6 +1011,7 @@ fn frontend_args_accept_json() {
data_parallel_size: None,
runtime: SharedRuntimeArgs {
model: "Qwen/Qwen3-0.6B",
revision: None,
generation_config: Auto,
engine_ready_timeout_secs: 600,
tool_call_parser: None,
@@ -1125,7 +1130,7 @@ fn frontend_args_json_accepts_supported_non_default_fields() {
"--output-address",
"ipc:///tmp/output.sock",
"--args-json",
r#"{"model_tag":"Qwen/Qwen3-0.6B","generation_config":"vllm","engine_ready_timeout_secs":42,"tool_call_parser":"hermes","reasoning_parser":"qwen3_thinking","tokenizer_mode":"deepseek_v32","language_model_only":true,"max_logprobs":-1,"shutdown_timeout":3}"#,
r#"{"model_tag":"Qwen/Qwen3-0.6B","generation_config":"vllm","revision":"release","engine_ready_timeout_secs":42,"tool_call_parser":"hermes","reasoning_parser":"qwen3_thinking","tokenizer_mode":"deepseek_v32","language_model_only":true,"max_logprobs":-1,"shutdown_timeout":3}"#,
])
.unwrap();
@@ -1133,6 +1138,7 @@ fn frontend_args_json_accepts_supported_non_default_fields() {
panic!("expected frontend args");
};
assert_eq!(args.runtime.engine_ready_timeout_secs, 42);
assert_eq!(args.runtime.revision.as_deref(), Some("release"));
assert_eq!(args.runtime.generation_config, GenerationConfigMode::Vllm);
assert_eq!(
args.runtime.tool_call_parser,
@@ -1604,6 +1610,7 @@ fn serve_args_accept_handshake_aliases() {
uds: None,
runtime: SharedRuntimeArgs {
model: "Qwen/Qwen3-0.6B",
revision: None,
generation_config: Auto,
engine_ready_timeout_secs: 600,
tool_call_parser: Auto,
@@ -1752,6 +1759,7 @@ fn serve_frontend_config_uses_dp_address_as_advertised_host() {
},
coordinator_mode: MaybeInProc,
model: "Qwen/Qwen3-0.6B",
revision: None,
generation_config: Auto,
served_model_name: [],
listener_mode: BindTcp {
@@ -1839,6 +1847,7 @@ fn serve_frontend_config_keeps_tcp_transport_for_non_local_only_topology() {
},
coordinator_mode: MaybeInProc,
model: "Qwen/Qwen3-0.6B",
revision: None,
generation_config: Auto,
served_model_name: [],
listener_mode: BindTcp {
@@ -1948,6 +1957,7 @@ fn frontend_config_uses_external_coordinator_when_coordinator_address_is_present
address: "tcp://127.0.0.1:7000",
},
model: "Qwen/Qwen3-0.6B",
revision: None,
generation_config: Auto,
served_model_name: [],
listener_mode: InheritedFd {
@@ -2087,3 +2097,18 @@ fn frontend_args_json_disables_profiling_when_profiler_type_is_null() {
let config = args.into_config();
assert_eq!(config.profiler, None);
}
#[test]
fn serve_revision_reaches_frontend_and_managed_engine() {
let cli =
Cli::try_parse_from(["vllm-rs", "serve", "test/model", "--revision", "release"]).unwrap();
let Command::Serve(args) = cli.command else {
panic!("expected serve args");
};
assert_eq!(
args.to_frontend_config("tcp://localhost:1234".into()).revision.as_deref(),
Some("release")
);
let engine = args.to_managed_engine_config(1234);
assert!(engine.python_args.windows(2).any(|args| args == ["--revision", "release"]));
}
+5
View File
@@ -84,6 +84,7 @@ impl ManagedEngineArgs {
pub fn into_config(
self,
model: String,
revision: Option<String>,
max_logprobs: Option<i32>,
profiler_config: Option<String>,
reasoning_parser: Option<&str>,
@@ -95,6 +96,10 @@ impl ManagedEngineArgs {
) -> ManagedEngineConfig {
let mut python_args = self.python_args;
// Manually forward some args to the Python engine.
if let Some(revision) = revision {
python_args.push("--revision".to_string());
python_args.push(revision);
}
if let Some(max_model_len) = self.max_model_len {
python_args.push("--max-model-len".to_string());
python_args.push(max_model_len);
@@ -60,6 +60,7 @@ async fn main() -> Result<()> {
},
coordinator_mode: CoordinatorMode::MaybeInProc,
model: args.model,
revision: None,
generation_config: Default::default(),
served_model_name: vec![],
listener_mode: HttpListenerMode::BindTcp {
+2
View File
@@ -210,6 +210,8 @@ pub struct Config {
pub coordinator_mode: CoordinatorMode,
/// Backend model identifier used for engine-core loading.
pub model: String,
/// Model revision on the Hugging Face Hub (branch, tag, or commit SHA).
pub revision: Option<String>,
/// Which generation-config sampling defaults to inherit.
pub generation_config: GenerationConfigMode,
/// Model name(s) exposed to clients via the OpenAI API. When non-empty,
+1
View File
@@ -98,6 +98,7 @@ async fn build_state(config: &Config) -> Result<Arc<AppState>> {
let loaded = load_model_backends(
&config.model,
LoadModelBackendsOptions {
revision: config.revision.clone(),
generation_config: config.generation_config,
renderer: config.renderer,
language_model_only: config.language_model_only,
+3
View File
@@ -24,6 +24,7 @@ use crate::{
#[derive(Debug)]
pub struct RenderConfig {
pub model: String,
pub revision: Option<String>,
pub served_model_name: Vec<String>,
pub host: String,
pub port: u16,
@@ -67,6 +68,7 @@ async fn build_state(config: &RenderConfig) -> Result<Arc<RenderState>> {
let loaded = load_model_backends(
&config.model,
LoadModelBackendsOptions {
revision: config.revision.clone(),
generation_config: Default::default(),
renderer: config.renderer,
language_model_only: true,
@@ -150,6 +152,7 @@ mod tests {
let error = serve_render(
RenderConfig {
model: "test-model".to_string(),
revision: None,
served_model_name: Vec::new(),
host: "127.0.0.1".to_string(),
port: 8000,
+145 -25
View File
@@ -3,8 +3,8 @@
use std::path::{Path, PathBuf};
use hf_hub::Cache;
use hf_hub::api::tokio::{Api, ApiBuilder, ApiRepo};
use hf_hub::{Cache, Repo, RepoType};
use thiserror_ext::AsReport as _;
use super::config::{HfTokenizerConfig, load_tokenizer_config};
@@ -56,14 +56,31 @@ impl ResolvedModelFiles {
/// Resolve tokenizer/config files from a local model directory first when
/// `model_id` points to one, otherwise consult the local HF cache and
/// finally the Hub.
pub async fn new(model_id: &str) -> Result<Self> {
pub async fn new(model_id: &str, revision: Option<&str>) -> Result<Self> {
if Path::new(model_id).is_dir() {
return resolve_local_model_files(Path::new(model_id));
}
if let Some(files) = resolve_cached_model_files(model_id)? {
let repo = Repo::with_revision(
model_id.to_string(),
RepoType::Model,
revision.unwrap_or("main").to_string(),
);
// TODO: Let hf-hub read HF_HUB_CACHE after upgrading to 1.0.
let cache = std::env::var_os("HF_HUB_CACHE")
.map(|path| Cache::new(path.into()))
.unwrap_or_else(Cache::from_env);
if let Some(files) = resolve_cached_model_files(&cache, &repo)? {
return Ok(files);
}
resolve_remote_model_files(model_id).await
if is_hf_hub_offline() {
return Err(Error::Tokenizer(format!(
"model '{model_id}' at revision '{}' has no supported tokenizer in cache '{}' \
and HF_HUB_OFFLINE is enabled",
repo.revision(),
cache.path().display(),
)));
}
resolve_remote_model_files(model_id, repo, &cache).await
}
}
@@ -87,9 +104,13 @@ fn resolve_local_model_files(model_dir: &Path) -> Result<ResolvedModelFiles> {
})
}
async fn resolve_remote_model_files(model_id: &str) -> Result<ResolvedModelFiles> {
let api = build_api().map_err(|error| Error::Tokenizer(error.to_report_string()))?;
let repo = api.model(model_id.to_string());
async fn resolve_remote_model_files(
model_id: &str,
model_repo: Repo,
cache: &Cache,
) -> Result<ResolvedModelFiles> {
let api = build_api(cache).map_err(|error| Error::Tokenizer(error.to_report_string()))?;
let repo = api.repo(model_repo);
let info = repo.info().await.map_err(|error| {
Error::Tokenizer(format!(
"failed to fetch model '{model_id}': {}",
@@ -146,12 +167,20 @@ async fn resolve_remote_model_files(model_id: &str) -> Result<ResolvedModelFiles
})
}
fn resolve_cached_model_files(model_id: &str) -> Result<Option<ResolvedModelFiles>> {
let cache_repo = Cache::from_env().model(model_id.to_string());
fn resolve_cached_model_files(cache: &Cache, repo: &Repo) -> Result<Option<ResolvedModelFiles>> {
let cache_repo = cache.repo(repo.clone());
// hf-hub 0.5 looks up refs even for commit SHAs; Python caches have no such ref.
let snapshot_dir = (repo.revision().len() == 40
&& repo.revision().bytes().all(|byte| byte.is_ascii_hexdigit()))
.then(|| cache.path().join(repo.folder_name()).join("snapshots").join(repo.revision()));
let get_file = |name: &str| match &snapshot_dir {
Some(dir) => local_file_if_exists(dir, name),
None => cache_repo.get(name),
};
let tokenizer_config_path = cache_repo.get("tokenizer_config.json");
let tokenizer_config_path = get_file("tokenizer_config.json");
let tokenizer_config = load_tokenizer_config(tokenizer_config_path.as_deref())?;
let tokenizer = match resolve_cached_tokenizer_source(&cache_repo, &tokenizer_config)? {
let tokenizer = match resolve_cached_tokenizer_source(&get_file, &tokenizer_config)? {
Some(tokenizer) => tokenizer,
None => return Ok(None),
};
@@ -159,12 +188,12 @@ fn resolve_cached_model_files(model_id: &str) -> Result<Option<ResolvedModelFile
let model_dir = tokenizer.path().parent().ok_or_else(|| {
Error::Tokenizer("resolved tokenizer file has no parent directory".to_string())
})?;
let generation_config_path = cache_repo.get("generation_config.json");
let preprocessor_config_path = cache_repo.get("preprocessor_config.json");
let video_preprocessor_config_path = cache_repo.get("video_preprocessor_config.json");
let processor_config_path = cache_repo.get("processor_config.json");
let generation_config_path = get_file("generation_config.json");
let preprocessor_config_path = get_file("preprocessor_config.json");
let video_preprocessor_config_path = get_file("video_preprocessor_config.json");
let processor_config_path = get_file("processor_config.json");
let chat_template_path = discover_chat_template_in_dir(model_dir);
let config_path = cache_repo.get("config.json");
let config_path = get_file("config.json");
Ok(Some(ResolvedModelFiles {
tokenizer,
@@ -207,20 +236,20 @@ async fn resolve_remote_tokenizer_source(
}
fn resolve_cached_tokenizer_source(
cache_repo: &hf_hub::CacheRepo,
get_file: &impl Fn(&str) -> Option<PathBuf>,
tokenizer_config: &HfTokenizerConfig,
) -> Result<Option<TokenizerSource>> {
let tekken_path = cache_repo.get("tekken.json");
let tekken_path = get_file("tekken.json");
if let Some(tekken_path) = tekken_path {
return Ok(Some(TokenizerSource::Tekken(tekken_path)));
}
let Some(tokenizer_path) = cache_repo.get("tokenizer.json").or_else(|| {
let Some(tokenizer_path) = get_file("tokenizer.json").or_else(|| {
// tiktoken.model is the most common name, try it first.
cache_repo.get("tiktoken.model").or_else(|| {
get_file("tiktoken.model").or_else(|| {
// Scan for any *.tiktoken file in the cache snapshot directory.
let snapshot_dir = cache_repo.get("config.json")?.parent()?.to_path_buf();
let snapshot_dir = get_file("config.json")?.parent()?.to_path_buf();
discover_tiktoken_in_dir(&snapshot_dir)
})
}) else {
@@ -315,8 +344,18 @@ async fn download_known_file(repo: &ApiRepo, model_id: &str, filename: &str) ->
})
}
fn build_api() -> anyhow::Result<Api> {
let mut builder = ApiBuilder::from_env().with_progress(true);
fn is_hf_hub_offline() -> bool {
std::env::var("HF_HUB_OFFLINE").is_ok_and(|value| {
matches!(
value.to_ascii_uppercase().as_str(),
"1" | "ON" | "YES" | "TRUE"
)
})
}
fn build_api(cache: &Cache) -> anyhow::Result<Api> {
let mut builder =
ApiBuilder::from_env().with_cache_dir(cache.path().clone()).with_progress(true);
if let Ok(token) = std::env::var(HF_TOKEN_ENV)
&& !token.is_empty()
{
@@ -388,6 +427,7 @@ fn discover_chat_template_in_dir(dir: &std::path::Path) -> Option<PathBuf> {
#[cfg(test)]
mod tests {
use std::fs;
use std::path::PathBuf;
use tempfile::tempdir;
use vllm_tokenizer::{TiktokenTokenizer, Tokenizer};
@@ -405,7 +445,7 @@ mod tests {
.expect("write tokenizer config");
fs::write(dir.path().join("config.json"), "{}").expect("write config");
let files = ResolvedModelFiles::new(dir.path().to_str().expect("utf8 path"))
let files = ResolvedModelFiles::new(dir.path().to_str().expect("utf8 path"), None)
.await
.expect("resolve local model files");
@@ -422,10 +462,90 @@ mod tests {
);
}
#[test]
fn cached_model_files_respect_branch_tag_and_commit_revisions() {
let dir = tempdir().unwrap();
let cache = hf_hub::Cache::new(dir.path().to_path_buf());
let model = "test/model";
let main_sha = "a".repeat(40);
let pinned_sha = "b".repeat(40);
for (revision, sha) in [("main", &main_sha), ("release/v1", &pinned_sha)] {
let repo =
hf_hub::Repo::with_revision(model.into(), hf_hub::RepoType::Model, revision.into());
cache.repo(repo.clone()).create_ref(sha).unwrap();
let snapshot = dir.path().join(repo.folder_name()).join("snapshots").join(sha);
fs::create_dir_all(&snapshot).unwrap();
fs::write(snapshot.join("tokenizer.json"), "{}").unwrap();
fs::write(snapshot.join("config.json"), "{}").unwrap();
}
for (revision, expected_sha) in [
("main", Some(&main_sha)),
("release/v1", Some(&pinned_sha)),
(pinned_sha.as_str(), Some(&pinned_sha)),
("missing", None),
] {
let repo =
hf_hub::Repo::with_revision(model.into(), hf_hub::RepoType::Model, revision.into());
let files = super::resolve_cached_model_files(&cache, &repo).unwrap();
match expected_sha {
Some(sha) => {
let snapshot = dir.path().join(repo.folder_name()).join("snapshots").join(sha);
let files = files.unwrap();
assert_eq!(files.tokenizer.path(), snapshot.join("tokenizer.json"));
assert_eq!(files.config_path, Some(snapshot.join("config.json")));
}
None => assert!(files.is_none()),
}
}
}
#[tokio::test]
async fn offline_resolution_uses_hub_cache_and_reports_cache_misses() {
const CHILD: &str = "VLLM_TEST_HF_OFFLINE_CHILD";
if std::env::var_os(CHILD).is_some() {
let files = ResolvedModelFiles::new("test/model", Some("release")).await.unwrap();
let cache = PathBuf::from(std::env::var_os("HF_HUB_CACHE").unwrap());
assert_eq!(
files.tokenizer.path(),
cache.join("models--test--model/snapshots/commit/tokenizer.json")
);
let error = ResolvedModelFiles::new("test/model", Some("missing")).await.unwrap_err();
let message = thiserror_ext::AsReport::as_report(&error).to_string();
assert!(message.contains("HF_HUB_OFFLINE is enabled"), "{message}");
assert!(message.contains("revision 'missing'"), "{message}");
return;
}
// Isolate environment variables from other tests and Tokio worker threads.
let dir = tempdir().unwrap();
let cache = dir.path().join("hub-cache");
let repo = cache.join("models--test--model");
fs::create_dir_all(repo.join("refs")).unwrap();
fs::create_dir_all(repo.join("snapshots/commit")).unwrap();
fs::write(repo.join("refs/release"), "commit").unwrap();
fs::write(repo.join("snapshots/commit/tokenizer.json"), "{}").unwrap();
for value in ["1", "ON", "yes", "True"] {
let output = std::process::Command::new(std::env::current_exe().unwrap())
.args(["--exact", "backend::hf::model_files::tests::offline_resolution_uses_hub_cache_and_reports_cache_misses", "--nocapture"])
.env(CHILD, "1")
.env("HF_HOME", dir.path().join("hf-home"))
.env("HF_HUB_CACHE", &cache)
.env("HF_HUB_OFFLINE", value)
.env("HF_ENDPOINT", "http://127.0.0.1:1")
.output().unwrap();
assert!(
output.status.success(),
"{}\n{}",
String::from_utf8_lossy(&output.stdout),
String::from_utf8_lossy(&output.stderr)
);
}
}
#[tokio::test]
#[ignore = "too slow for CI and requires network access to Hugging Face"]
async fn tiktoken_real_kimi_k25_tokenizer_files_load_and_handle_special_tokens() {
let files = ResolvedModelFiles::new("moonshotai/Kimi-K2.5")
let files = ResolvedModelFiles::new("moonshotai/Kimi-K2.5", None)
.await
.expect("resolve real Kimi K2.5 model files");
+2 -1
View File
@@ -802,7 +802,8 @@ mod tests {
#[file_serial(hf_qwen3)]
async fn lower_text_request_uses_real_qwen_generation_defaults() {
let model_id = "Qwen/Qwen3-0.6B";
let files = ResolvedModelFiles::new(model_id).await.expect("resolve qwen model files");
let files =
ResolvedModelFiles::new(model_id, None).await.expect("resolve qwen model files");
let backend = HfTextBackend::from_resolved_model_files(
files,
model_id.to_string(),