From 485421b1c3572597a4cfaece04836a843537435a Mon Sep 17 00:00:00 2001 From: Bugen Zhao Date: Sat, 12 Sep 2026 11:39:17 +0800 Subject: [PATCH] [Rust Frontend] Honor HF revisions, offline mode, and cache directory (#56386) Signed-off-by: Bugen Zhao Co-authored-by: Codex --- rust/src/chat/src/backend/hf.rs | 3 +- rust/src/chat/src/backend/mod.rs | 2 + rust/src/cmd/src/cli.rs | 12 ++ rust/src/cmd/src/cli/tests.rs | 27 ++- rust/src/managed-engine/src/cli.rs | 5 + .../examples/external_engine_openai_qwen.rs | 1 + rust/src/server/src/config.rs | 2 + rust/src/server/src/lib.rs | 1 + rust/src/server/src/render.rs | 3 + rust/src/text/src/backend/hf/model_files.rs | 170 +++++++++++++++--- rust/src/text/src/lower.rs | 3 +- 11 files changed, 201 insertions(+), 28 deletions(-) diff --git a/rust/src/chat/src/backend/hf.rs b/rust/src/chat/src/backend/hf.rs index 717bde2898..37212a9650 100644 --- a/rust/src/chat/src/backend/hf.rs +++ b/rust/src/chat/src/backend/hf.rs @@ -129,7 +129,7 @@ pub(super) async fn load_model_backends( model_id: &str, options: LoadModelBackendsOptions, ) -> Result { - 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, diff --git a/rust/src/chat/src/backend/mod.rs b/rust/src/chat/src/backend/mod.rs index 77b9114a58..028c0c40c8 100644 --- a/rust/src/chat/src/backend/mod.rs +++ b/rust/src/chat/src/backend/mod.rs @@ -61,6 +61,8 @@ pub type DynChatTextBackend = Arc; /// 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, /// Which generation-config sampling defaults to inherit. pub generation_config: GenerationConfigMode, /// Which chat renderer implementation to use. diff --git a/rust/src/cmd/src/cli.rs b/rust/src/cmd/src/cli.rs index c0b4c6d295..f7fcf30022 100644 --- a/rust/src/cmd/src/cli.rs +++ b/rust/src/cmd/src/cli.rs @@ -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, /// 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, + /// 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(), diff --git a/rust/src/cmd/src/cli/tests.rs b/rust/src/cmd/src/cli/tests.rs index e9411891e0..884da135d7 100644 --- a/rust/src/cmd/src/cli/tests.rs +++ b/rust/src/cmd/src/cli/tests.rs @@ -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"])); +} diff --git a/rust/src/managed-engine/src/cli.rs b/rust/src/managed-engine/src/cli.rs index eadefced23..f0cd7bf171 100644 --- a/rust/src/managed-engine/src/cli.rs +++ b/rust/src/managed-engine/src/cli.rs @@ -84,6 +84,7 @@ impl ManagedEngineArgs { pub fn into_config( self, model: String, + revision: Option, max_logprobs: Option, profiler_config: Option, 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); diff --git a/rust/src/server/examples/external_engine_openai_qwen.rs b/rust/src/server/examples/external_engine_openai_qwen.rs index d07ba242ae..fc9914e3e2 100644 --- a/rust/src/server/examples/external_engine_openai_qwen.rs +++ b/rust/src/server/examples/external_engine_openai_qwen.rs @@ -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 { diff --git a/rust/src/server/src/config.rs b/rust/src/server/src/config.rs index fde53183f5..4928cddf31 100644 --- a/rust/src/server/src/config.rs +++ b/rust/src/server/src/config.rs @@ -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, /// Which generation-config sampling defaults to inherit. pub generation_config: GenerationConfigMode, /// Model name(s) exposed to clients via the OpenAI API. When non-empty, diff --git a/rust/src/server/src/lib.rs b/rust/src/server/src/lib.rs index 098c803805..95ae359f35 100644 --- a/rust/src/server/src/lib.rs +++ b/rust/src/server/src/lib.rs @@ -98,6 +98,7 @@ async fn build_state(config: &Config) -> Result> { 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, diff --git a/rust/src/server/src/render.rs b/rust/src/server/src/render.rs index 6942694afb..fe7033fa8c 100644 --- a/rust/src/server/src/render.rs +++ b/rust/src/server/src/render.rs @@ -24,6 +24,7 @@ use crate::{ #[derive(Debug)] pub struct RenderConfig { pub model: String, + pub revision: Option, pub served_model_name: Vec, pub host: String, pub port: u16, @@ -67,6 +68,7 @@ async fn build_state(config: &RenderConfig) -> Result> { 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, diff --git a/rust/src/text/src/backend/hf/model_files.rs b/rust/src/text/src/backend/hf/model_files.rs index d0b0caecce..de802fb439 100644 --- a/rust/src/text/src/backend/hf/model_files.rs +++ b/rust/src/text/src/backend/hf/model_files.rs @@ -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 { + pub async fn new(model_id: &str, revision: Option<&str>) -> Result { 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 { }) } -async fn resolve_remote_model_files(model_id: &str) -> Result { - 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 { + 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 Result> { - let cache_repo = Cache::from_env().model(model_id.to_string()); +fn resolve_cached_model_files(cache: &Cache, repo: &Repo) -> Result> { + 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, tokenizer_config: &HfTokenizerConfig, ) -> Result> { - 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 { - 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 { + 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 { #[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"); diff --git a/rust/src/text/src/lower.rs b/rust/src/text/src/lower.rs index 0b18fd4238..bff4a6d48b 100644 --- a/rust/src/text/src/lower.rs +++ b/rust/src/text/src/lower.rs @@ -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(),