diff --git a/components/src/dynamo/frontend/sglang_processor.py b/components/src/dynamo/frontend/sglang_processor.py index f5b859a47911..abe31491c118 100644 --- a/components/src/dynamo/frontend/sglang_processor.py +++ b/components/src/dynamo/frontend/sglang_processor.py @@ -622,6 +622,9 @@ async def chat_engine_factory( ) loop = asyncio.get_running_loop() + # TODO(gh-8749): consume mdc.model_info.path()'s parent (slug_dir) + # instead of re-running fetch_model. The MDC cache already has + # blake3-verified copies; this path duplicates the download. source_path = mdc.source_path() if not os.path.exists(source_path): await fetch_model(source_path, ignore_weights=True) diff --git a/components/src/dynamo/frontend/vllm_processor.py b/components/src/dynamo/frontend/vllm_processor.py index f16b15a534a6..d204171e44b7 100644 --- a/components/src/dynamo/frontend/vllm_processor.py +++ b/components/src/dynamo/frontend/vllm_processor.py @@ -760,6 +760,9 @@ async def chat_engine_factory( ) loop = asyncio.get_running_loop() + # TODO(gh-8749): consume mdc.model_info.path()'s parent (slug_dir) + # instead of re-running fetch_model. The MDC cache already has + # blake3-verified copies; this path duplicates the download. source_path = mdc.source_path() if not os.path.exists(source_path): await fetch_model(source_path, ignore_weights=True) diff --git a/lib/bindings/c/src/lib.rs b/lib/bindings/c/src/lib.rs index 405ed1975887..b804326923e3 100644 --- a/lib/bindings/c/src/lib.rs +++ b/lib/bindings/c/src/lib.rs @@ -1498,7 +1498,7 @@ async fn fetch_preprocessor_from_discovery( ); // Download config (tokenizer files) if not local - card.download_config().await?; + card.download_config(None).await?; // Create preprocessor let preprocessor = OpenAIPreprocessor::new(card.clone())?; diff --git a/lib/llm/src/common/checked_file.rs b/lib/llm/src/common/checked_file.rs index 38c8c8474192..a82dfa31f1f5 100644 --- a/lib/llm/src/common/checked_file.rs +++ b/lib/llm/src/common/checked_file.rs @@ -235,6 +235,27 @@ impl Checksum { algorithm, } } + + pub fn hash(&self) -> &str { + &self.hash + } +} + +/// Worker-controlled wire field used as a path component downstream; +/// reject anything that doesn't look like the algorithm's canonical +/// digest. Called from both `Deserialize` and `TryFrom<&str>` so both +/// wire entry points stay in lockstep. +fn validate_checksum_hash(algorithm: CryptographicHashMethods, hash: &str) -> Result<(), String> { + match algorithm { + CryptographicHashMethods::BLAKE3 => { + if hash.len() != 64 || !hash.bytes().all(|b| matches!(b, b'0'..=b'9' | b'a'..=b'f')) { + return Err(format!( + "expected 64-char lowercase hex blake3 digest, got {hash:?}" + )); + } + } + } + Ok(()) } impl Serialize for Checksum { @@ -273,6 +294,9 @@ impl<'de> Deserialize<'de> for Checksum { let algorithm = parts[0].parse().map_err(|_| { de::Error::invalid_value(de::Unexpected::Str(parts[0]), &"invalid algorithm") })?; + validate_checksum_hash(algorithm, parts[1]).map_err(|msg| { + de::Error::invalid_value(de::Unexpected::Str(parts[1]), &msg.as_str()) + })?; Ok(Checksum::new(parts[1], algorithm)) } @@ -297,6 +321,7 @@ impl TryFrom<&str> for Checksum { anyhow::bail!("Unsupported cryptographic hash method: {}", parts[0]); } }; + validate_checksum_hash(algo, parts[1]).map_err(|msg| anyhow::anyhow!("{msg}"))?; Ok(Checksum::new(parts[1], algo)) } @@ -348,26 +373,32 @@ mod tests { assert_eq!(serialized.trim(), "\"blake3:a12c3d4\""); } + const HEX64: &str = "62bc124be974d3a25db05bedc99422660c26715e5bbda0b37d14bd84a0c65ab2"; + #[test] fn test_deserialization_blake3() { - let s = "\"blake3:abcd1234\""; - let deserialized: Checksum = serde_json::from_str(s).unwrap(); - + let s = format!("\"blake3:{HEX64}\""); + let deserialized: Checksum = serde_json::from_str(&s).unwrap(); assert_eq!(deserialized.algorithm, CryptographicHashMethods::BLAKE3); - assert_eq!(deserialized.hash, "abcd1234"); + assert_eq!(deserialized.hash, HEX64); } #[test] fn test_deserialization_invalid_format() { - let s = "\"invalidformat\""; - let result: Result = serde_json::from_str(s); - - assert!(result.is_err()); - - let s = "\"blake3:invalid:format\""; - let result: Result = serde_json::from_str(s); - - assert!(result.is_err()); + for bad in [ + "\"invalidformat\"", + "\"blake3:invalid:format\"", + // Worker-controlled wire field shouldn't survive as a path + // component; reject anything that's not a 64-char lowercase + // hex blake3 digest. + "\"blake3:abcd1234\"", // too short + "\"blake3:../../etc/passwd\"", // path traversal + "\"blake3:62BC124BE974D3A25DB05BEDC99422660C26715E5BBDA0B37D14BD84A0C65AB2\"", // uppercase + &format!("\"blake3:{HEX64}x\""), // 65 chars + ] { + let result: Result = serde_json::from_str(bad); + assert!(result.is_err(), "expected reject: {bad}"); + } } #[test] @@ -378,5 +409,38 @@ mod tests { let expected = Checksum::blake3("62bc124be974d3a25db05bedc99422660c26715e5bbda0b37d14bd84a0c65ab2"); assert_eq!(expected, *cf.checksum()); + assert_eq!(Some(560), cf.size(), "TinyLlama config.json is 560 bytes"); + assert_eq!( + "62bc124be974d3a25db05bedc99422660c26715e5bbda0b37d14bd84a0c65ab2", + cf.checksum().hash() + ); + } + + /// Old MDC bytes (no `size` field) deserialize cleanly with + /// `size = None`, preserving cross-version compat between a + /// pre-size worker and a post-size frontend. + #[test] + fn test_legacy_wire_format_omits_size() { + let legacy = format!(r#"{{"path":"/x/config.json","checksum":"blake3:{HEX64}"}}"#); + let cf: CheckedFile = serde_json::from_str(&legacy).unwrap(); + assert_eq!(None, cf.size()); + assert_eq!("/x/config.json", cf.path().unwrap().to_str().unwrap()); + } + + /// New-format CheckedFiles round-trip with the size field. + #[test] + fn test_new_wire_format_carries_size() { + let cf = CheckedFile { + path: Either::Left(PathBuf::from("/x/tokenizer.json")), + checksum: Checksum::blake3(HEX64), + size: Some(12345), + }; + let json = serde_json::to_string(&cf).unwrap(); + assert!( + json.contains("\"size\":12345"), + "expected size field in serialization, got: {json}" + ); + let back: CheckedFile = serde_json::from_str(&json).unwrap(); + assert_eq!(Some(12345), back.size()); } } diff --git a/lib/llm/src/discovery/watcher.rs b/lib/llm/src/discovery/watcher.rs index 1fe780e3a6d8..6b47c3436ec2 100644 --- a/lib/llm/src/discovery/watcher.rs +++ b/lib/llm/src/discovery/watcher.rs @@ -1,6 +1,7 @@ // SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 +use std::path::PathBuf; use std::sync::Arc; use std::time::Duration; use tokio::sync::Notify; @@ -88,6 +89,10 @@ pub struct ModelWatcher { /// Tracks in-flight `handle_put` tasks by instance path so that `handle_delete` /// can await a racing put before proceeding with cleanup. pending_puts: DashMap>, + /// Frontend's `--model-path`. Threaded into `download_config` so + /// `file://` slots can fall back here when the worker's path is + /// unreachable on this host. + local_model_path: Option, } const ALL_MODEL_TYPES: &[ModelType] = &[ @@ -167,6 +172,7 @@ impl ModelWatcher { registering_worker_sets: DashSet::new(), registration_notify: Notify::new(), pending_puts: DashMap::new(), + local_model_path: None, } } @@ -174,6 +180,10 @@ impl ModelWatcher { self.model_update_tx = Some(tx); } + pub fn set_local_model_path(&mut self, path: Option) { + self.local_model_path = path; + } + /// Wait until we have at least one chat completions model and return it's name. pub async fn wait_for_chat_model(&self) -> String { // Loop in case it gets added and immediately deleted @@ -673,7 +683,8 @@ impl ModelWatcher { mcid: &ModelCardInstanceId, card: &mut ModelDeploymentCard, ) -> anyhow::Result<()> { - card.download_config().await?; + card.download_config(self.local_model_path.as_deref()) + .await?; // Use per-worker-set router config if the worker provided one in its MDC, // otherwise fall back to the frontend-level global config. diff --git a/lib/llm/src/entrypoint/input/common.rs b/lib/llm/src/entrypoint/input/common.rs index 63611566c60c..944d078a2c63 100644 --- a/lib/llm/src/entrypoint/input/common.rs +++ b/lib/llm/src/entrypoint/input/common.rs @@ -101,7 +101,7 @@ pub async fn prepare_engine( let model_manager = Arc::new(ModelManager::new()); // Create metrics for migration tracking (not exposed via /metrics in Dynamic engine mode) let metrics = Arc::new(Metrics::new()); - let watch_obj = Arc::new(ModelWatcher::new( + let mut watcher = ModelWatcher::new( distributed_runtime.clone(), model_manager.clone(), RouterConfig::default(), @@ -110,7 +110,11 @@ pub async fn prepare_engine( None, prefill_load_estimator, metrics, - )); + ); + if !local_model.path().as_os_str().is_empty() { + watcher.set_local_model_path(Some(local_model.path().to_path_buf())); + } + let watch_obj = Arc::new(watcher); let discovery = distributed_runtime.discovery(); let discovery_stream = discovery .list_and_watch( diff --git a/lib/llm/src/entrypoint/input/grpc.rs b/lib/llm/src/entrypoint/input/grpc.rs index 2377dafe1e7b..93fac38d0422 100644 --- a/lib/llm/src/entrypoint/input/grpc.rs +++ b/lib/llm/src/entrypoint/input/grpc.rs @@ -1,6 +1,7 @@ // SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 +use std::path::PathBuf; use std::sync::Arc; use crate::{ @@ -47,6 +48,8 @@ pub async fn run( model.namespace(), model.namespace_prefix(), ); + let local_model_path = + (!model.path().as_os_str().is_empty()).then(|| model.path().to_path_buf()); run_watcher( distributed_runtime.clone(), grpc_service.state().manager_clone(), @@ -55,6 +58,7 @@ pub async fn run( migration_max_seq_len, namespace_filter, prefill_load_estimator.clone(), + local_model_path, ) .await?; grpc_service @@ -112,6 +116,7 @@ pub async fn run( /// Spawns a task that watches for new models in store, /// and registers them with the ModelManager so that the HTTP service can use them. +#[allow(clippy::too_many_arguments)] async fn run_watcher( runtime: DistributedRuntime, model_manager: Arc, @@ -120,10 +125,11 @@ async fn run_watcher( migration_max_seq_len: Option, namespace_filter: NamespaceFilter, prefill_load_estimator: Option>, + local_model_path: Option, ) -> anyhow::Result<()> { // Create metrics for migration tracking (not exposed via /metrics in gRPC mode) let metrics = Arc::new(Metrics::new()); - let watch_obj = ModelWatcher::new( + let mut watch_obj = ModelWatcher::new( runtime.clone(), model_manager, router_config, @@ -133,6 +139,7 @@ async fn run_watcher( prefill_load_estimator, metrics, ); + watch_obj.set_local_model_path(local_model_path); tracing::debug!("Waiting for remote model"); let discovery = runtime.discovery(); let discovery_stream = discovery diff --git a/lib/llm/src/entrypoint/input/http.rs b/lib/llm/src/entrypoint/input/http.rs index a9bdf1c6c09a..09b5d25bd1e1 100644 --- a/lib/llm/src/entrypoint/input/http.rs +++ b/lib/llm/src/entrypoint/input/http.rs @@ -1,6 +1,7 @@ // SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 +use std::path::PathBuf; use std::sync::Arc; use crate::{ @@ -83,6 +84,8 @@ pub async fn run( model.namespace(), model.namespace_prefix(), ); + let local_model_path = + (!model.path().as_os_str().is_empty()).then(|| model.path().to_path_buf()); run_watcher( distributed_runtime.clone(), http_service.state().manager_clone(), @@ -94,6 +97,7 @@ pub async fn run( http_service.state().metrics_clone(), chat_engine_factory.clone(), prefill_load_estimator.clone(), + local_model_path, ) .await?; http_service @@ -173,6 +177,7 @@ async fn run_watcher( metrics: Arc, chat_engine_factory: Option, prefill_load_estimator: Option>, + local_model_path: Option, ) -> anyhow::Result<()> { let mut watch_obj = ModelWatcher::new( runtime.clone(), @@ -184,6 +189,7 @@ async fn run_watcher( prefill_load_estimator, metrics.clone(), ); + watch_obj.set_local_model_path(local_model_path); tracing::debug!("Waiting for remote model"); let discovery = runtime.discovery(); let discovery_stream = discovery diff --git a/lib/llm/src/local_model.rs b/lib/llm/src/local_model.rs index 000f87a300b3..a5fd756ab51d 100644 --- a/lib/llm/src/local_model.rs +++ b/lib/llm/src/local_model.rs @@ -560,13 +560,10 @@ impl LocalModel { let registry = drt.metadata_artifacts(); let mut rewritten = 0usize; - for cf in self.card.iter_metadata_files_mut() { + for (cf, _) in self.card.iter_metadata_files_mut() { let Some(local_path) = cf.path().map(Path::to_path_buf) else { continue; }; - // Filename from the original path — HF cache symlinks would - // otherwise canonicalize to the LFS SHA1 and break downstream - // lookups by literal name (e.g. `parent.join("generation_config.json")`). let Some(filename) = local_path .file_name() .and_then(|f| f.to_str()) @@ -574,14 +571,13 @@ impl LocalModel { else { continue; }; - // Canonicalize so the handler serves the resolved HF blob, not the symlink. - let absolute = match fs::canonicalize(&local_path) { + let absolute = match std::path::absolute(&local_path) { Ok(p) => p, Err(err) => { tracing::warn!( path = %local_path.display(), %err, - "failed to canonicalize self-host metadata path; skipping", + "failed to absolutize self-host metadata path; skipping", ); continue; } diff --git a/lib/llm/src/model_card.rs b/lib/llm/src/model_card.rs index 979f1d861ae3..a2cb817ddb08 100644 --- a/lib/llm/src/model_card.rs +++ b/lib/llm/src/model_card.rs @@ -194,6 +194,343 @@ fn is_exclusively_mistral_model(directory: &Path) -> bool { !directory.join("config.json").exists() && directory.join("params.json").exists() } +/// MDC cache: `blobs/` + `by-slug///`. +/// The `` segment mirrors HF Hub's `snapshots//` and +/// isolates worker sets that share a model name but publish different +/// file content. +fn mdc_cache_root() -> PathBuf { + let home = std::env::var("HOME") + .or_else(|_| std::env::var("USERPROFILE")) + .unwrap_or_else(|_| ".".to_string()); + PathBuf::from(home).join(".cache/dynamo/mdc") +} + +fn mdc_blobs_dir() -> anyhow::Result { + let dir = mdc_cache_root().join("blobs"); + std::fs::create_dir_all(&dir) + .with_context(|| format!("creating MDC blobs dir {}", dir.display()))?; + Ok(dir) +} + +fn mdc_slug_dir(slug: &Slug, mdcsum: &str) -> anyhow::Result { + let dir = mdc_cache_root() + .join("by-slug") + .join(slug.to_string()) + .join(mdcsum); + std::fs::create_dir_all(&dir) + .with_context(|| format!("creating MDC slug dir {}", dir.display()))?; + Ok(dir) +} + +/// Per-call tmp path next to `dest`. pid + uuid suffix keeps tmps +/// disjoint across concurrent callers so `rename(2)` is the only +/// synchronization point — never `create(tmp)`. +fn unique_tmp_path(dest: &Path) -> PathBuf { + let suffix = format!( + "tmp.{}.{}", + std::process::id(), + uuid::Uuid::new_v4().simple() + ); + let name = dest + .file_name() + .map(|f| f.to_string_lossy().into_owned()) + .unwrap_or_default(); + let mut p = dest.to_path_buf(); + p.set_file_name(format!("{name}.{suffix}")); + p +} + +/// RAII guard that unlinks a tmp path on drop. `dismiss()` cancels the +/// cleanup once the tmp has been consumed by a successful rename. Drop +/// ignores ENOENT — safe under double-cleanup races. +struct TmpGuard { + path: Option, +} + +impl TmpGuard { + fn new(path: PathBuf) -> Self { + Self { path: Some(path) } + } + fn dismiss(&mut self) { + self.path = None; + } +} + +impl Drop for TmpGuard { + fn drop(&mut self) { + if let Some(p) = self.path.take() { + let _ = std::fs::remove_file(&p); + } + } +} + +/// Atomic publish: stage via `f(&tmp)`, then `rename(tmp -> dest)`. +/// Concurrent-safe because tmps are per-call (see [`unique_tmp_path`]) +/// and `rename(2)` is atomic + overwrites. Tmp is unlinked on any +/// failure path including async cancellation — see [`TmpGuard`]. +async fn stage_and_rename(dest: &Path, f: F) -> anyhow::Result<()> +where + F: FnOnce(PathBuf) -> Fut, + Fut: std::future::Future>, +{ + let tmp = unique_tmp_path(dest); + let mut guard = TmpGuard::new(tmp.clone()); + f(tmp.clone()).await?; + tokio::fs::rename(&tmp, dest) + .await + .with_context(|| format!("renaming {} -> {}", tmp.display(), dest.display()))?; + guard.dismiss(); + Ok(()) +} + +/// Sync sibling of [`stage_and_rename`] for cheap operations +/// (e.g., creating a symlink). +fn stage_and_rename_sync(dest: &Path, f: F) -> anyhow::Result<()> +where + F: FnOnce(&Path) -> anyhow::Result<()>, +{ + let tmp = unique_tmp_path(dest); + let mut guard = TmpGuard::new(tmp.clone()); + f(&tmp)?; + std::fs::rename(&tmp, dest) + .with_context(|| format!("renaming {} -> {}", tmp.display(), dest.display()))?; + guard.dismiss(); + Ok(()) +} + +/// Concurrent-safe via [`stage_and_rename_sync`]. +fn symlink_force(target: &Path, link: &Path) -> anyhow::Result<()> { + stage_and_rename_sync(link, |tmp| { + #[cfg(unix)] + std::os::unix::fs::symlink(target, tmp) + .with_context(|| format!("symlinking {} -> {}", tmp.display(), target.display()))?; + #[cfg(not(unix))] + std::fs::copy(target, tmp) + .map(|_| ()) + .with_context(|| format!("copying {} -> {}", target.display(), tmp.display()))?; + Ok(()) + }) +} + +/// 1 GiB cap on metadata fetch — realistic files are <20 MiB. Bounds +/// disk usage if a worker advertises a bogus `CheckedFile.size`, and is +/// the fallback when `size` is absent (legacy MDCs). +const ABSOLUTE_MAX_METADATA_BYTES: u64 = 1024 * 1024 * 1024; + +/// Stage `uri` into `dest`, verifying staged bytes against `expected` +/// before publishing. Schemes: `http(s)`, `file`, `hf`. For `hf://`, +/// `hf_snapshots` must already contain the resolved repo path — +/// caller is expected to pre-resolve once per repo. +async fn resolve_uri( + client: &reqwest::Client, + uri: &str, + expected: &CheckedFile, + dest: &Path, + hf_snapshots: &std::collections::HashMap, +) -> anyhow::Result<()> { + let cap = expected + .size() + .unwrap_or(ABSOLUTE_MAX_METADATA_BYTES) + .min(ABSOLUTE_MAX_METADATA_BYTES); + + let parsed = url::Url::parse(uri).with_context(|| format!("parsing artifact uri: {uri}"))?; + + if dest.exists() { + if CheckedFile::from_disk(dest).is_ok_and(|cf| cf.checksum() == expected.checksum()) { + return Ok(()); + } + tracing::warn!(dest = %dest.display(), "MDC cache blob failed re-verification; refetching"); + let _ = std::fs::remove_file(dest); + } + + stage_and_rename(dest, |tmp| async move { + match parsed.scheme() { + "http" | "https" => stream_to_tmp(client, uri, &tmp, cap).await?, + "file" => { + let path = parsed + .to_file_path() + .map_err(|()| anyhow::anyhow!("invalid file:// uri: {uri}"))?; + copy_to_tmp(&path, &tmp, cap).await?; + } + "hf" => { + let (repo, filename) = parse_hf_uri(uri)?; + let snapshot = hf_snapshots + .get(&repo) + .with_context(|| format!("hf snapshot not pre-resolved for {repo}"))?; + copy_to_tmp(&snapshot.join(&filename), &tmp, cap).await?; + } + scheme => anyhow::bail!("unsupported artifact uri scheme: {scheme} (uri: {uri})"), + } + + // Re-blake3 the staged bytes via the same `from_disk` path the + // worker used at registration so the comparison is bit-identical. + let actual = CheckedFile::from_disk(&tmp)?; + if actual.checksum() != expected.checksum() { + anyhow::bail!( + "checksum mismatch for {uri}: expected {}, got {}", + expected.checksum(), + actual.checksum() + ); + } + Ok(()) + }) + .await +} + +/// Stream HTTP body to `tmp`, capped at `cap` bytes. Pre-checks +/// `Content-Length` if present; otherwise the post-write check +/// catches overage. +async fn stream_to_tmp( + client: &reqwest::Client, + uri: &str, + tmp: &Path, + cap: u64, +) -> anyhow::Result<()> { + use futures::TryStreamExt; + use tokio::io::{AsyncReadExt, AsyncWriteExt}; + use tokio_util::io::StreamReader; + + let response = client + .get(uri) + .send() + .await + .with_context(|| format!("fetching {uri}"))?; + if !response.status().is_success() { + anyhow::bail!("fetching {uri} returned status {}", response.status()); + } + if let Some(content_length) = response.content_length() + && content_length > cap + { + anyhow::bail!("{uri} reports {content_length} bytes, exceeds cap {cap}"); + } + + let stream = response.bytes_stream().map_err(std::io::Error::other); + // `cap + 1` so a body of exactly `cap` passes but anything larger + // trips the `written > cap` check below. + let mut reader = StreamReader::new(stream).take(cap + 1); + let mut file = tokio::fs::File::create(tmp) + .await + .with_context(|| format!("creating {}", tmp.display()))?; + let written = tokio::io::copy(&mut reader, &mut file) + .await + .with_context(|| format!("streaming body from {uri} to {}", tmp.display()))?; + file.flush() + .await + .with_context(|| format!("flushing {}", tmp.display()))?; + if written > cap { + anyhow::bail!("{uri} body exceeds cap {cap}"); + } + Ok(()) +} + +async fn copy_to_tmp(src: &Path, tmp: &Path, cap: u64) -> anyhow::Result<()> { + let metadata = tokio::fs::metadata(src) + .await + .with_context(|| format!("reading metadata for {}", src.display()))?; + if metadata.len() > cap { + anyhow::bail!( + "{} is {} bytes, exceeds cap {}", + src.display(), + metadata.len(), + cap + ); + } + tokio::fs::copy(src, tmp) + .await + .with_context(|| format!("copying {} -> {}", src.display(), tmp.display()))?; + Ok(()) +} + +/// Parse `hf://repo[@rev]/filename` into `(repo[@rev], filename)`. +fn parse_hf_uri(uri: &str) -> anyhow::Result<(String, String)> { + let body = uri + .strip_prefix("hf://") + .with_context(|| format!("expected hf:// scheme, got: {uri}"))?; + let (repo, filename) = body + .rsplit_once('/') + .with_context(|| format!("hf:// uri must end in /filename, got: {uri}"))?; + if repo.is_empty() || filename.is_empty() { + anyhow::bail!("malformed hf:// uri: {uri}"); + } + Ok((repo.to_string(), filename.to_string())) +} + +fn checked_file_uri( + cf: &CheckedFile, + source: &str, + local_model_path: Option<&Path>, + is_custom: bool, +) -> anyhow::Result { + use std::borrow::Cow; + + // Coerce path-only into a synthetic file:// URL up front so the + // scheme dispatch below covers both shapes uniformly. `absolute` + // (not `canonicalize`) — worker paths often don't exist here. + let url: Cow = if let Some(u) = cf.url() { + Cow::Borrowed(u) + } else { + let Some(p) = cf.path() else { + anyhow::bail!("CheckedFile has neither path nor url"); + }; + let abs = std::path::absolute(p)?; + Cow::Owned( + url::Url::from_file_path(&abs) + .map_err(|()| anyhow::anyhow!("invalid file path: {}", abs.display()))?, + ) + }; + + match url.scheme() { + "http" | "https" | "hf" => Ok(url.to_string()), + "file" => { + // worker location → --model-path → hf://. Basename + checksum preserved. + // is_custom slots aren't published on HF, so rung 4 errors instead. + let path = url + .to_file_path() + .map_err(|()| anyhow::anyhow!("invalid file uri: {url}"))?; + let filename = path + .file_name() + .and_then(|f| f.to_str()) + .with_context(|| format!("no filename in file uri: {url}"))?; + if path.exists() { + return Ok(url.to_string()); + } + if let Some(prefix) = local_model_path { + let local = prefix.join(filename); + if local.exists() { + return file_uri_for(&local); + } + } + if is_custom { + anyhow::bail!( + "custom file {filename} not reachable on this host \ + (worker path {} missing, no --model-path overlay); \ + custom files aren't published on HF, so ensure the \ + file exists at the same path on every host (shared \ + mount) or pass --model-path with the same basename", + path.display() + ); + } + Ok(format!("hf://{source}/{filename}")) + } + _ => Ok(url.to_string()), + } +} + +fn file_uri_for(p: &Path) -> anyhow::Result { + Ok(url::Url::from_file_path(std::path::absolute(p)?) + .map_err(|()| anyhow::anyhow!("invalid file path: {}", p.display()))? + .to_string()) +} + +fn uri_basename(uri: &str) -> anyhow::Result { + url::Url::parse(uri) + .with_context(|| format!("parsing uri: {uri}"))? + .path_segments() + .and_then(|mut s| s.rfind(|s| !s.is_empty())) + .map(String::from) + .with_context(|| format!("no basename in uri: {uri}")) +} + fn pf_checked_file(p: &PromptFormatterArtifact) -> &CheckedFile { match p { PromptFormatterArtifact::HfTokenizerConfigJson(cf) @@ -538,74 +875,138 @@ impl ModelDeploymentCard { matches!(self.model_input, ModelInput::Tokens) } - /// Walk populated metadata `CheckedFile` slots in deterministic - /// order: model_info, tokenizer, prompt_formatter, - /// chat_template_file, gen_config. - pub fn iter_metadata_files(&self) -> Vec<&CheckedFile> { - let mut out: Vec<&CheckedFile> = Vec::with_capacity(5); - if let Some(m) = self.model_info.as_ref() { - match m { - ModelInfoType::HfConfigJson(cf) => out.push(cf), - } + /// Iterate populated metadata slots in deterministic order: + /// model_info, tokenizer, prompt_formatter, chat_template_file, + /// gen_config. Each entry is `(file, is_custom)` — `is_custom` is + /// only ever true for operator-supplied chat templates, which + /// can't fall back to HF. + /// + /// TODO(gh-8749): external preprocessors (vllm/sglang) read + /// `from_pretrained(slug_dir)` and may expect siblings outside the + /// typed slots — `preprocessor_config.json`, `special_tokens_map.json`, + /// `added_tokens.json`, etc. Add an `extra_files: Vec` + /// MDC field so the worker can advertise everything in its model dir + /// minus weights. Frontend pipeline is already generic over slot count. + pub fn iter_metadata_files(&self) -> Vec<(&CheckedFile, bool)> { + let mut out: Vec<(&CheckedFile, bool)> = Vec::with_capacity(5); + if let Some(ModelInfoType::HfConfigJson(cf)) = self.model_info.as_ref() { + out.push((cf, false)); } - if let Some(t) = self.tokenizer.as_ref() { - match t { - TokenizerKind::HfTokenizerJson(cf) | TokenizerKind::TikTokenModel(cf) => { - out.push(cf) - } - } + if let Some(TokenizerKind::HfTokenizerJson(cf) | TokenizerKind::TikTokenModel(cf)) = + self.tokenizer.as_ref() + { + out.push((cf, false)); } if let Some(p) = self.prompt_formatter.as_ref() { - out.push(pf_checked_file(p)); + out.push((pf_checked_file(p), p.is_custom())); } if let Some(c) = self.chat_template_file.as_ref() { - out.push(pf_checked_file(c)); + out.push((pf_checked_file(c), c.is_custom())); } - if let Some(g) = self.gen_config.as_ref() { - match g { - GenerationConfig::HfGenerationConfigJson(cf) => out.push(cf), - } + if let Some(GenerationConfig::HfGenerationConfigJson(cf)) = self.gen_config.as_ref() { + out.push((cf, false)); } out } - /// Mutable variant of [`Self::iter_metadata_files`]. - pub fn iter_metadata_files_mut(&mut self) -> Vec<&mut CheckedFile> { - let mut out: Vec<&mut CheckedFile> = Vec::with_capacity(5); - if let Some(m) = self.model_info.as_mut() { - match m { - ModelInfoType::HfConfigJson(cf) => out.push(cf), - } + /// Mutable mirror of [`Self::iter_metadata_files`]. + pub fn iter_metadata_files_mut(&mut self) -> Vec<(&mut CheckedFile, bool)> { + let mut out: Vec<(&mut CheckedFile, bool)> = Vec::with_capacity(5); + if let Some(ModelInfoType::HfConfigJson(cf)) = self.model_info.as_mut() { + out.push((cf, false)); } - if let Some(t) = self.tokenizer.as_mut() { - match t { - TokenizerKind::HfTokenizerJson(cf) | TokenizerKind::TikTokenModel(cf) => { - out.push(cf) - } - } + if let Some(TokenizerKind::HfTokenizerJson(cf) | TokenizerKind::TikTokenModel(cf)) = + self.tokenizer.as_mut() + { + out.push((cf, false)); } if let Some(p) = self.prompt_formatter.as_mut() { - out.push(pf_checked_file_mut(p)); + let is_custom = p.is_custom(); + out.push((pf_checked_file_mut(p), is_custom)); } if let Some(c) = self.chat_template_file.as_mut() { - out.push(pf_checked_file_mut(c)); + let is_custom = c.is_custom(); + out.push((pf_checked_file_mut(c), is_custom)); } - if let Some(g) = self.gen_config.as_mut() { - match g { - GenerationConfig::HfGenerationConfigJson(cf) => out.push(cf), - } + if let Some(GenerationConfig::HfGenerationConfigJson(cf)) = self.gen_config.as_mut() { + out.push((cf, false)); } out } - /// Download the files this card needs to work: config.json, tokenizer.json, etc. - pub async fn download_config(&mut self) -> anyhow::Result<()> { - if self.has_local_files() { - tracing::trace!("All model config is local, not downloading"); - return Ok(()); + async fn resolve_metadata_files( + &mut self, + local_model_path: Option<&Path>, + ) -> anyhow::Result<()> { + let source = self.source_path().to_string(); + let mdcsum = self.mdcsum().to_string(); + let blobs = mdc_blobs_dir()?; + let slug_dir = mdc_slug_dir(&self.slug, &mdcsum)?; + + let entries: Vec<(String, CheckedFile)> = self + .iter_metadata_files() + .into_iter() + .map(|(cf, is_custom)| { + Ok(( + checked_file_uri(cf, &source, local_model_path, is_custom)?, + cf.clone(), + )) + }) + .collect::>()?; + + // Pre-resolve hf:// repos once per unique repo; otherwise the + // resolve loop would call hub::from_hf N times for one model. + let mut hf_snapshots: std::collections::HashMap = + std::collections::HashMap::new(); + for (uri, _) in &entries { + if uri.starts_with("hf://") { + let (repo, _) = parse_hf_uri(uri)?; + if let std::collections::hash_map::Entry::Vacant(e) = hf_snapshots.entry(repo) { + let repo_name = e.key().clone(); + let snap = crate::hub::from_hf(&repo_name, /* ignore_weights = */ true) + .await + .with_context(|| format!("hub::from_hf({repo_name})"))?; + e.insert(snap); + } + } } - // For TensorBased models, config files are not used - they handle everything in the backend + let client = reqwest::Client::builder() + .connect_timeout(std::time::Duration::from_secs(10)) + .timeout(std::time::Duration::from_secs(300)) + .build() + .context("building http client for metadata fetch")?; + for (uri, expected) in &entries { + let filename = uri_basename(uri)?; + // Hash is validated at MDC deserialize; safe as a path component. + let blake3_hex = expected.checksum().hash(); + let blob = blobs.join(blake3_hex); + tracing::debug!(filename = %filename, uri = %uri, blake3 = %blake3_hex, "resolving"); + resolve_uri(&client, uri, expected, &blob, &hf_snapshots).await?; + symlink_force(&blob, &slug_dir.join(&filename))?; + } + tracing::debug!( + display_name = %self.display_name, + artifact_count = entries.len(), + cache_root = %mdc_cache_root().display(), + "resolved model metadata files", + ); + + // Pass 3: rewrite cf.path to the cache symlink so downstream + // tokenizer/config loaders read from a verified location. + for (cf, _) in self.iter_metadata_files_mut() { + cf.update_dir(&slug_dir); + } + Ok(()) + } + + /// Resolve every metadata `CheckedFile` through the cache: fetch, + /// blake3-verify, content-address. `local_model_path` (frontend's + /// `--model-path`) supplies a fallback directory for `file://` + /// slots whose worker-published location is unreachable. + pub async fn download_config(&mut self, local_model_path: Option<&Path>) -> anyhow::Result<()> { + // TensorBased models don't use metadata files — backend handles + // everything. if self.model_type.supports_tensor() { tracing::debug!( display_name = %self.display_name, @@ -613,12 +1014,10 @@ impl ModelDeploymentCard { ); return Ok(()); } - - let ignore_weights = true; - let local_path = crate::hub::from_hf(self.source_path(), ignore_weights).await?; - - self.update_dir(&local_path); - Ok(()) + // Single resolve pipeline: every CheckedFile (URL or local + // path, existing or missing) flows through resolve_uri, + // blake3-verifies, lands in the MDC cache. No new/legacy split. + self.resolve_metadata_files(local_model_path).await } /// Re-write all the local disk paths as a URL. Do this before publishing the MDC. @@ -689,63 +1088,6 @@ impl ModelDeploymentCard { Ok(()) } - /// Are all the files we need (tokenizer.json, etc) available locally? - fn has_local_files(&self) -> bool { - let has_model_info = self - .model_info - .as_ref() - .map(|p| p.is_local()) - .unwrap_or(true); - let has_tokenizer = self - .tokenizer - .as_ref() - .map(|p| p.is_local()) - .unwrap_or(true); - let has_prompt_formatter = self - .prompt_formatter - .as_ref() - .map(|p| p.is_local()) - .unwrap_or(true); - let has_chat_template_file = self - .chat_template_file - .as_ref() - .map(|p| p.is_local()) - .unwrap_or(true); - let has_gen_config = self - .gen_config - .as_ref() - .map(|p| p.is_local()) - .unwrap_or(true); - - has_model_info - && has_tokenizer - && has_prompt_formatter - && has_chat_template_file - && has_gen_config - } - - /// Update the directory for files like tokenizer.json be in here. - fn update_dir(&mut self, dir: &Path) { - if let Some(model_info) = self.model_info.as_mut() { - model_info.update_dir(dir); - } - if let Some(tk) = self.tokenizer.as_mut() { - tk.update_dir(dir); - } - if let Some(pf) = self.prompt_formatter.as_mut() { - pf.update_dir(dir); - } - if let Some(gc) = self.gen_config.as_mut() { - gc.update_dir(dir); - } - // If it's a custom chat template we didn't download it, so leave the path untouched - if let Some(ct) = self.chat_template_file.as_mut() - && !ct.is_custom() - { - ct.update_dir(dir); - } - } - /// Creates a ModelDeploymentCard from a local directory path. /// /// Currently HuggingFace format is supported and following files are expected: @@ -1299,7 +1641,7 @@ fn check_valid_local_repo_path(path: impl AsRef) -> Result<()> { mod tests { use super::HFConfig; use std::collections::HashSet; - use std::path::Path; + use std::path::{Path, PathBuf}; #[test] pub fn test_config_json_llama3() -> anyhow::Result<()> { @@ -1353,4 +1695,257 @@ mod tests { ); Ok(()) } + + fn test_cf(uri: &str, size: u64) -> super::CheckedFile { + serde_json::from_value(serde_json::json!({ + "path": uri, + "checksum": format!("blake3:{}", "0".repeat(64)), + "size": size, + })) + .unwrap() + } + + async fn assert_resolve_uri_rejects(body: &[u8], declared_size: u64, expected_err: &str) { + let mut server = mockito::Server::new_async().await; + let _m = server + .mock("GET", "/f") + .with_status(200) + .with_body(body) + .create_async() + .await; + let dir = tempfile::tempdir().unwrap(); + let dest = dir.path().join("f"); + let url = format!("{}/f", server.url()); + let result = super::resolve_uri( + &reqwest::Client::new(), + &url, + &test_cf(&url, declared_size), + &dest, + &std::collections::HashMap::new(), + ) + .await; + let msg = result.expect_err("expected error").to_string(); + assert!( + msg.contains(expected_err), + "want `{expected_err}` in: {msg}" + ); + assert!(!dest.exists(), "no file should be written"); + } + + #[tokio::test] + async fn resolve_uri_http_rejects_checksum_mismatch() { + assert_resolve_uri_rejects(b"hello world", 11, "checksum mismatch").await; + } + + #[tokio::test] + async fn resolve_uri_http_rejects_oversize_body() { + assert_resolve_uri_rejects(b"x".repeat(35).as_slice(), 8, "exceeds cap").await; + } + + /// Cache hit re-verifies the on-disk blob; mismatch → unlink + refetch. + #[tokio::test] + async fn resolve_uri_refetches_on_cache_hit_mismatch() { + let body: &[u8] = b"valid-bytes-for-blob"; + let dir = tempfile::tempdir().unwrap(); + let valid = dir.path().join("valid"); + std::fs::write(&valid, body).unwrap(); + let cf = super::CheckedFile::from_disk(&valid).unwrap(); + + let mut server = mockito::Server::new_async().await; + let _m = server + .mock("GET", "/f") + .with_status(200) + .with_body(body) + .create_async() + .await; + let url = format!("{}/f", server.url()); + + let dest = dir.path().join("blob"); + std::fs::write(&dest, b"corrupt-bytes").unwrap(); + + super::resolve_uri( + &reqwest::Client::new(), + &url, + &cf, + &dest, + &std::collections::HashMap::new(), + ) + .await + .expect("resolve_uri should refetch and succeed"); + + let after = std::fs::read(&dest).unwrap(); + assert_eq!(after, body, "blob should have been replaced"); + } + + /// `parse_hf_uri` round-trip — `hf://repo/filename` parses into + /// `(repo, filename)` and rejects malformed inputs. + #[test] + fn parse_hf_uri_roundtrip() { + let (repo, filename) = super::parse_hf_uri("hf://Qwen/Qwen3-0.6B/tokenizer.json").unwrap(); + assert_eq!(repo, "Qwen/Qwen3-0.6B"); + assert_eq!(filename, "tokenizer.json"); + + assert!(super::parse_hf_uri("hf://just-a-name").is_err()); + assert!(super::parse_hf_uri("https://example.com/x").is_err()); + } + + /// HF-cache-style snapshot of TinyLlama: per-file symlinks into a + /// sibling `blobs/` dir. The symlink layout is what triggers + /// the canonicalize trap `resolve_metadata_files` has to handle. + fn hf_cache_fixture(workspace: &Path) -> anyhow::Result { + use std::hash::{Hash, Hasher}; + let snapshot = workspace.join("snapshots/abc"); + let blobs = workspace.join("blobs"); + std::fs::create_dir_all(&snapshot)?; + std::fs::create_dir_all(&blobs)?; + let src = + Path::new(env!("CARGO_MANIFEST_DIR")).join("tests/data/sample-models/TinyLlama_v1.1"); + for entry in std::fs::read_dir(&src)? { + let entry = entry?; + let name = entry.file_name(); + let mut hasher = std::collections::hash_map::DefaultHasher::new(); + name.to_string_lossy().hash(&mut hasher); + let blob = format!("blob-{:x}", hasher.finish()); + std::fs::copy(entry.path(), blobs.join(&blob))?; + #[cfg(unix)] + std::os::unix::fs::symlink( + Path::new("../..").join("blobs").join(&blob), + snapshot.join(name), + )?; + } + Ok(snapshot) + } + + #[tokio::test] + #[serial_test::serial] + async fn download_config_pipelines_local_files_through_cache() -> anyhow::Result<()> { + let workspace = tempfile::tempdir()?; + let snapshot = hf_cache_fixture(workspace.path())?; + let home = tempfile::tempdir()?; + let home_path = home.path().to_path_buf(); + + temp_env::async_with_vars([("HOME", Some(home.path()))], async { + let mut mdc = super::ModelDeploymentCard::load_from_disk(&snapshot, None)?; + let slug = mdc.slug.clone(); + let mdcsum = mdc.mdcsum().to_string(); + mdc.download_config(None).await?; + + let blobs = home_path.join(".cache/dynamo/mdc/blobs"); + let snap = home_path + .join(".cache/dynamo/mdc/by-slug") + .join(slug.to_string()) + .join(&mdcsum); + + assert!(snap.join("config.json").exists()); + assert!(snap.join("tokenizer.json").exists()); + assert!(snap.join("generation_config.json").exists()); + + for (cf, _) in mdc.iter_metadata_files() { + let path = cf.path().expect("post-download local path"); + assert!(path.starts_with(&snap)); + assert!(std::fs::canonicalize(path)?.starts_with(&blobs)); + } + Ok::<_, anyhow::Error>(()) + }) + .await + } + + /// Build a `CheckedFile` whose wire `path` field is `repr` — parses + /// as a URL when `repr` has a scheme, otherwise as a `PathBuf`. + fn cf_for(repr: &str) -> super::CheckedFile { + serde_json::from_value(serde_json::json!({ + "path": repr, + "checksum": format!("blake3:{}", "0".repeat(64)), + "size": 1u64, + })) + .unwrap() + } + + #[test] + fn checked_file_uri_passes_through_remote_urls() { + let tmp = tempfile::tempdir().unwrap(); + for url in [ + "http://worker:8080/v1/metadata/slug/base/config.json", + "hf://Qwen/Qwen3-0.6B/config.json", + ] { + let got = + super::checked_file_uri(&cf_for(url), "Qwen/Qwen3-0.6B", Some(tmp.path()), false) + .unwrap(); + assert_eq!(got, url); + } + } + + #[test] + fn checked_file_uri_uses_local_model_path_when_worker_path_unreachable() { + let cf = cf_for("/nonexistent/worker/path/config.json"); + let local = tempfile::tempdir().unwrap(); + let local_cfg = local.path().join("config.json"); + std::fs::write(&local_cfg, b"").unwrap(); + let got = + super::checked_file_uri(&cf, "Qwen/Qwen3-0.6B", Some(local.path()), false).unwrap(); + assert_eq!( + got, + url::Url::from_file_path(&local_cfg).unwrap().to_string() + ); + } + + /// Rung 4: when worker path is missing and `--model-path` doesn't + /// supply the basename, synthesize hf://; if `is_custom`, error + /// with a clear operator-action message instead (HF doesn't host + /// custom slots). + #[test] + fn checked_file_uri_rung_4_hf_fallback_or_custom_error() { + let cf = cf_for("/nonexistent/worker/path/template.jinja"); + + let got = super::checked_file_uri(&cf, "Qwen/Qwen3-0.6B", None, false).unwrap(); + assert_eq!(got, "hf://Qwen/Qwen3-0.6B/template.jinja"); + + let err = super::checked_file_uri(&cf, "Qwen/Qwen3-0.6B", None, true) + .expect_err("custom slot must error instead of falling back to HF"); + let msg = err.to_string(); + assert!(msg.contains("template.jinja"), "wrong error: {msg}"); + assert!(msg.contains("custom"), "wrong error: {msg}"); + assert!( + msg.contains("--model-path") || msg.contains("shared mount"), + "wrong error: {msg}" + ); + } + + /// Dropping `stage_and_rename`'s future mid-await (caller cancellation) + /// must still unlink the tmp file via the [`TmpGuard`] drop path. + #[tokio::test] + async fn stage_and_rename_unlinks_tmp_on_cancel() { + let dir = tempfile::tempdir().unwrap(); + let dest = dir.path().join("dest"); + let leaked_before: Vec<_> = std::fs::read_dir(dir.path()) + .unwrap() + .filter_map(Result::ok) + .collect(); + assert!(leaked_before.is_empty(), "tempdir starts empty"); + + let fut = super::stage_and_rename(&dest, |tmp| async move { + std::fs::write(&tmp, b"partial").unwrap(); + // Mimic a worker abort mid-fetch: yield, then never resume. + std::future::pending::<()>().await; + Ok(()) + }); + // Drop the future before it can complete (cancellation). + { + let _ = tokio::time::timeout(std::time::Duration::from_millis(50), fut).await; + } + + // Give Drop a moment to run. + tokio::task::yield_now().await; + + let leaked_after: Vec<_> = std::fs::read_dir(dir.path()) + .unwrap() + .filter_map(Result::ok) + .map(|e| e.path()) + .collect(); + assert!( + leaked_after.is_empty(), + "TmpGuard should have unlinked the tmp on cancel; leaked: {leaked_after:?}" + ); + assert!(!dest.exists(), "dest must not exist after cancel"); + } }