diff --git a/crates/infrastructure/src/model_server/mod.rs b/crates/infrastructure/src/model_server/mod.rs index a4b8aa8..209dea5 100644 --- a/crates/infrastructure/src/model_server/mod.rs +++ b/crates/infrastructure/src/model_server/mod.rs @@ -89,6 +89,11 @@ impl HfModelArtifactDownloader { } /// Deterministic cache path for a Hugging Face model reference. + /// + /// This is the local path used for a single **merged** (non-sharded) GGUF file, renamed to a + /// synthetic `{quant}.gguf` name. Sharded files are never renamed this way: llama.cpp + /// discovers multi-part siblings by pattern-matching the *original* HF filename in the same + /// directory, so shard files keep their remote name (see `local_path_for_filename`). #[must_use] pub fn cache_path_for(&self, repo: &HfModelRef) -> PathBuf { let (base, quant) = split_hf_ref(repo); @@ -98,7 +103,56 @@ impl HfModelArtifactDownloader { .join(format!("{}.gguf", safe_cache_component(file_stem))) } - async fn resolve_remote_filename(&self, repo: &HfModelRef) -> Result { + fn cache_dir_for(&self, repo: &HfModelRef) -> PathBuf { + let (base, _) = split_hf_ref(repo); + self.cache_dir.join(base.replace('/', "--")) + } + + /// Local path for a remote filename that is part of a multi-shard set: the original name is + /// preserved (sanitized) so llama.cpp's shard autoload can find siblings by pattern. + fn local_path_for_shard(&self, repo: &HfModelRef, remote_filename: &str) -> PathBuf { + self.cache_dir_for(repo) + .join(safe_cache_component(remote_filename)) + } + + fn manifest_path_for(&self, repo: &HfModelRef) -> PathBuf { + let (_, quant) = split_hf_ref(repo); + let file_stem = quant.unwrap_or("model"); + self.cache_dir_for(repo) + .join(format!("{}.manifest.json", safe_cache_component(file_stem))) + } + + /// Reads a manifest of a previously-downloaded shard set, if present and complete on disk. + fn cached_shard_set(&self, repo: &HfModelRef) -> Option> { + let manifest_path = self.manifest_path_for(repo); + let raw = std::fs::read(&manifest_path).ok()?; + let manifest: ShardManifest = serde_json::from_slice(&raw).ok()?; + let dir = self.cache_dir_for(repo); + let paths: Vec = manifest + .files + .iter() + .map(|filename| dir.join(filename)) + .collect(); + if paths.iter().all(|path| path.is_file()) { + Some(paths) + } else { + None + } + } + + fn write_shard_manifest(&self, repo: &HfModelRef, filenames: &[String]) -> std::io::Result<()> { + let manifest_path = self.manifest_path_for(repo); + let manifest = ShardManifest { + files: filenames.to_vec(), + }; + let json = serde_json::to_vec(&manifest)?; + std::fs::write(manifest_path, json) + } + + async fn resolve_remote_filenames( + &self, + repo: &HfModelRef, + ) -> Result, ModelServerError> { let (base, quant) = split_hf_ref(repo); let url = format!("https://huggingface.co/api/models/{base}"); let response = self @@ -117,15 +171,22 @@ impl HfModelArtifactDownloader { .json() .await .map_err(|e| ModelServerError::Probe(e.to_string()))?; - select_gguf_file(&metadata.siblings, quant).ok_or_else(|| { - ModelServerError::PathNotAccessible(format!( + let files = select_gguf_files(&metadata.siblings, quant); + if files.is_empty() { + return Err(ModelServerError::PathNotAccessible(format!( "no matching .gguf artifact found for {}", repo.as_str() - )) - }) + ))); + } + Ok(files) } } +#[derive(Debug, Default, Serialize, Deserialize)] +struct ShardManifest { + files: Vec, +} + #[async_trait] impl ModelArtifactDownloader for HfModelArtifactDownloader { async fn resolve_hf_model( @@ -137,72 +198,105 @@ impl ModelArtifactDownloader for HfModelArtifactDownloader { if cancel.is_cancelled() { return Err(ModelServerError::Cancelled); } - let path = self.cache_path_for(repo); - if path.is_file() { + + // Fast path 1: a single merged file was already downloaded (synthetic name), no network needed. + let merged_path = self.cache_path_for(repo); + if merged_path.is_file() { return Ok(ModelArtifactResolution { - path: model_path_from_pathbuf(path)?, + path: model_path_from_pathbuf(merged_path)?, cache_hit: true, }); } + // Fast path 2: a previously-downloaded shard set is fully present on disk, no network needed. + if let Some(paths) = self.cached_shard_set(repo) { + if let Some(first) = paths.into_iter().next() { + return Ok(ModelArtifactResolution { + path: model_path_from_pathbuf(first)?, + cache_hit: true, + }); + } + } - let filename = self.resolve_remote_filename(repo).await?; + let filenames = self.resolve_remote_filenames(repo).await?; if cancel.is_cancelled() { return Err(ModelServerError::Cancelled); } + let is_shard_set = filenames.len() > 1; + let targets: Vec = if is_shard_set { + filenames + .iter() + .map(|filename| self.local_path_for_shard(repo, filename)) + .collect() + } else { + vec![merged_path.clone()] + }; + let (base, _) = split_hf_ref(repo); - let url = format!( - "https://huggingface.co/{base}/resolve/main/{}", - url_path_segment(&filename) - ); - let response = self - .client - .get(url) - .send() - .await - .map_err(|e| ModelServerError::Probe(e.to_string()))?; - if !response.status().is_success() { - return Err(ModelServerError::Probe(format!( - "huggingface artifact download returned {}", - response.status() - ))); - } - if let Some(parent) = path.parent() { - tokio::fs::create_dir_all(parent) - .await - .map_err(|e| ModelServerError::Store(e.to_string()))?; - } - let tmp_path = path.with_extension("gguf.part"); - let mut file = tokio::fs::File::create(&tmp_path) - .await - .map_err(|e| ModelServerError::Store(e.to_string()))?; - let total = response.content_length(); - let mut downloaded = 0_u64; - let mut stream = response.bytes_stream(); - while let Some(chunk) = stream.next().await { + let total_files = filenames.len() as u64; + for (index, (filename, target)) in filenames.iter().zip(targets.iter()).enumerate() { if cancel.is_cancelled() { - let _ = tokio::fs::remove_file(&tmp_path).await; return Err(ModelServerError::Cancelled); } - let chunk = chunk.map_err(|e| ModelServerError::Probe(e.to_string()))?; - file.write_all(&chunk) + let url = format!( + "https://huggingface.co/{base}/resolve/main/{}", + url_path_segment(filename) + ); + let response = self + .client + .get(url) + .send() + .await + .map_err(|e| ModelServerError::Probe(e.to_string()))?; + if !response.status().is_success() { + return Err(ModelServerError::Probe(format!( + "huggingface artifact download returned {}", + response.status() + ))); + } + if let Some(parent) = target.parent() { + tokio::fs::create_dir_all(parent) + .await + .map_err(|e| ModelServerError::Store(e.to_string()))?; + } + let tmp_path = PathBuf::from(format!("{}.part", target.to_string_lossy())); + let mut file = tokio::fs::File::create(&tmp_path) + .await + .map_err(|e| ModelServerError::Store(e.to_string()))?; + let total = response.content_length(); + let mut downloaded = 0_u64; + let mut stream = response.bytes_stream(); + while let Some(chunk) = stream.next().await { + if cancel.is_cancelled() { + let _ = tokio::fs::remove_file(&tmp_path).await; + return Err(ModelServerError::Cancelled); + } + let chunk = chunk.map_err(|e| ModelServerError::Probe(e.to_string()))?; + file.write_all(&chunk) + .await + .map_err(|e| ModelServerError::Store(e.to_string()))?; + downloaded += chunk.len() as u64; + progress(ModelArtifactProgress { + downloaded_bytes: Some(downloaded + index as u64 * total.unwrap_or(0)), + total_bytes: total.map(|t| t * total_files), + source: Some(repo.as_str().to_owned()), + }); + } + file.flush() + .await + .map_err(|e| ModelServerError::Store(e.to_string()))?; + drop(file); + tokio::fs::rename(&tmp_path, target) .await .map_err(|e| ModelServerError::Store(e.to_string()))?; - downloaded += chunk.len() as u64; - progress(ModelArtifactProgress { - downloaded_bytes: Some(downloaded), - total_bytes: total, - source: Some(repo.as_str().to_owned()), - }); } - file.flush() - .await - .map_err(|e| ModelServerError::Store(e.to_string()))?; - drop(file); - tokio::fs::rename(&tmp_path, &path) - .await - .map_err(|e| ModelServerError::Store(e.to_string()))?; + + if is_shard_set { + self.write_shard_manifest(repo, &filenames) + .map_err(|e| ModelServerError::Store(e.to_string()))?; + } + Ok(ModelArtifactResolution { - path: model_path_from_pathbuf(path)?, + path: model_path_from_pathbuf(targets[0].clone())?, cache_hit: false, }) } @@ -226,21 +320,61 @@ fn split_hf_ref(repo: &HfModelRef) -> (&str, Option<&str>) { .map_or((repo.as_str(), None), |(base, quant)| (base, Some(quant))) } -fn select_gguf_file(siblings: &[HfSibling], quant: Option<&str>) -> Option { - let mut ggufs: Vec<&str> = siblings +/// Matches the llama.cpp / HF multi-part GGUF naming convention, e.g. +/// `model-q5_k_m-00001-of-00002.gguf`. Captures the shard index and total shard count. +fn shard_suffix() -> &'static regex::Regex { + static RE: std::sync::OnceLock = std::sync::OnceLock::new(); + RE.get_or_init(|| regex::Regex::new(r"-(\d{5})-of-(\d{5})\.gguf$").expect("valid regex")) +} + +fn shard_index(filename: &str) -> Option { + shard_suffix() + .captures(filename) + .and_then(|caps| caps.get(1)) + .and_then(|m| m.as_str().parse().ok()) +} + +/// Resolves the ordered set of remote filenames to download for a quantization: the merged +/// (non-sharded) file if one matches, otherwise all matching shards sorted by index. A merged +/// file always takes priority over a shard set when both exist in the repo, since it is what +/// `llama-server --model` expects and a shard set would otherwise be picked first by naive +/// alphabetical sort ('-' sorts before '.' in ASCII). +fn select_gguf_files(siblings: &[HfSibling], quant: Option<&str>) -> Vec { + let ggufs: Vec<&str> = siblings .iter() .map(|sibling| sibling.filename.as_str()) .filter(|filename| filename.ends_with(".gguf")) .collect(); - ggufs.sort_unstable(); - let Some(quant) = quant else { - return ggufs.first().map(|filename| (*filename).to_owned()); + + let matching: Vec<&str> = match quant { + None => ggufs, + Some(quant) => { + let quant = quant.to_ascii_lowercase(); + ggufs + .into_iter() + .filter(|filename| filename.to_ascii_lowercase().contains(&quant)) + .collect() + } }; - let quant = quant.to_ascii_lowercase(); - ggufs - .into_iter() - .find(|filename| filename.to_ascii_lowercase().contains(&quant)) - .map(str::to_owned) + + let (mut merged, mut shards): (Vec<&str>, Vec<&str>) = (Vec::new(), Vec::new()); + for filename in matching { + if shard_index(filename).is_some() { + shards.push(filename); + } else { + merged.push(filename); + } + } + + if let Some(first_merged) = { + merged.sort_unstable(); + merged.first() + } { + return vec![(*first_merged).to_owned()]; + } + + shards.sort_unstable_by_key(|filename| shard_index(filename).unwrap_or(u32::MAX)); + shards.into_iter().map(str::to_owned).collect() } fn safe_cache_component(raw: &str) -> String { @@ -644,3 +778,55 @@ impl ModelServerRegistry for FsModelServerRegistry { self.write_doc(&doc).await } } + +#[cfg(test)] +mod gguf_selection_tests { + use super::{select_gguf_files, HfSibling}; + + fn sibling(name: &str) -> HfSibling { + HfSibling { + filename: name.to_owned(), + } + } + + #[test] + fn prefers_merged_file_over_colliding_shards() { + let siblings = vec![ + sibling("qwen2.5-coder-7b-instruct-q5_k_m-00001-of-00002.gguf"), + sibling("qwen2.5-coder-7b-instruct-q5_k_m-00002-of-00002.gguf"), + sibling("qwen2.5-coder-7b-instruct-q5_k_m.gguf"), + ]; + + let files = select_gguf_files(&siblings, Some("q5_k_m")); + + assert_eq!(files, vec!["qwen2.5-coder-7b-instruct-q5_k_m.gguf"]); + } + + #[test] + fn selects_all_shards_sorted_by_index_when_no_merged_file() { + let siblings = vec![ + sibling("model-q4_k_m-00002-of-00003.gguf"), + sibling("model-q4_k_m-00001-of-00003.gguf"), + sibling("model-q4_k_m-00003-of-00003.gguf"), + sibling("model-q8_0-00001-of-00002.gguf"), + ]; + + let files = select_gguf_files(&siblings, Some("q4_k_m")); + + assert_eq!( + files, + vec![ + "model-q4_k_m-00001-of-00003.gguf", + "model-q4_k_m-00002-of-00003.gguf", + "model-q4_k_m-00003-of-00003.gguf", + ] + ); + } + + #[test] + fn returns_empty_when_no_match() { + let siblings = vec![sibling("model-q4_k_m.gguf")]; + + assert!(select_gguf_files(&siblings, Some("q5_k_m")).is_empty()); + } +}