Merge feature/fix-gguf-shard-selection into develop

Fix backend de sélection GGUF (fichier fusionné prioritaire sur les
shards partiels) — QA vert : cargo test -p infrastructure/application,
cargo check --workspace, 3 nouveaux tests unitaires.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
2026-07-23 13:23:59 +02:00

View File

@ -89,6 +89,11 @@ impl HfModelArtifactDownloader {
} }
/// Deterministic cache path for a Hugging Face model reference. /// 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] #[must_use]
pub fn cache_path_for(&self, repo: &HfModelRef) -> PathBuf { pub fn cache_path_for(&self, repo: &HfModelRef) -> PathBuf {
let (base, quant) = split_hf_ref(repo); let (base, quant) = split_hf_ref(repo);
@ -98,7 +103,56 @@ impl HfModelArtifactDownloader {
.join(format!("{}.gguf", safe_cache_component(file_stem))) .join(format!("{}.gguf", safe_cache_component(file_stem)))
} }
async fn resolve_remote_filename(&self, repo: &HfModelRef) -> Result<String, ModelServerError> { 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<Vec<PathBuf>> {
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<PathBuf> = 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<Vec<String>, ModelServerError> {
let (base, quant) = split_hf_ref(repo); let (base, quant) = split_hf_ref(repo);
let url = format!("https://huggingface.co/api/models/{base}"); let url = format!("https://huggingface.co/api/models/{base}");
let response = self let response = self
@ -117,13 +171,20 @@ impl HfModelArtifactDownloader {
.json() .json()
.await .await
.map_err(|e| ModelServerError::Probe(e.to_string()))?; .map_err(|e| ModelServerError::Probe(e.to_string()))?;
select_gguf_file(&metadata.siblings, quant).ok_or_else(|| { let files = select_gguf_files(&metadata.siblings, quant);
ModelServerError::PathNotAccessible(format!( if files.is_empty() {
return Err(ModelServerError::PathNotAccessible(format!(
"no matching .gguf artifact found for {}", "no matching .gguf artifact found for {}",
repo.as_str() repo.as_str()
)) )));
})
} }
Ok(files)
}
}
#[derive(Debug, Default, Serialize, Deserialize)]
struct ShardManifest {
files: Vec<String>,
} }
#[async_trait] #[async_trait]
@ -137,22 +198,48 @@ impl ModelArtifactDownloader for HfModelArtifactDownloader {
if cancel.is_cancelled() { if cancel.is_cancelled() {
return Err(ModelServerError::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 { return Ok(ModelArtifactResolution {
path: model_path_from_pathbuf(path)?, path: model_path_from_pathbuf(merged_path)?,
cache_hit: true, 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() { if cancel.is_cancelled() {
return Err(ModelServerError::Cancelled); return Err(ModelServerError::Cancelled);
} }
let is_shard_set = filenames.len() > 1;
let targets: Vec<PathBuf> = 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 (base, _) = split_hf_ref(repo);
let total_files = filenames.len() as u64;
for (index, (filename, target)) in filenames.iter().zip(targets.iter()).enumerate() {
if cancel.is_cancelled() {
return Err(ModelServerError::Cancelled);
}
let url = format!( let url = format!(
"https://huggingface.co/{base}/resolve/main/{}", "https://huggingface.co/{base}/resolve/main/{}",
url_path_segment(&filename) url_path_segment(filename)
); );
let response = self let response = self
.client .client
@ -166,12 +253,12 @@ impl ModelArtifactDownloader for HfModelArtifactDownloader {
response.status() response.status()
))); )));
} }
if let Some(parent) = path.parent() { if let Some(parent) = target.parent() {
tokio::fs::create_dir_all(parent) tokio::fs::create_dir_all(parent)
.await .await
.map_err(|e| ModelServerError::Store(e.to_string()))?; .map_err(|e| ModelServerError::Store(e.to_string()))?;
} }
let tmp_path = path.with_extension("gguf.part"); let tmp_path = PathBuf::from(format!("{}.part", target.to_string_lossy()));
let mut file = tokio::fs::File::create(&tmp_path) let mut file = tokio::fs::File::create(&tmp_path)
.await .await
.map_err(|e| ModelServerError::Store(e.to_string()))?; .map_err(|e| ModelServerError::Store(e.to_string()))?;
@ -189,8 +276,8 @@ impl ModelArtifactDownloader for HfModelArtifactDownloader {
.map_err(|e| ModelServerError::Store(e.to_string()))?; .map_err(|e| ModelServerError::Store(e.to_string()))?;
downloaded += chunk.len() as u64; downloaded += chunk.len() as u64;
progress(ModelArtifactProgress { progress(ModelArtifactProgress {
downloaded_bytes: Some(downloaded), downloaded_bytes: Some(downloaded + index as u64 * total.unwrap_or(0)),
total_bytes: total, total_bytes: total.map(|t| t * total_files),
source: Some(repo.as_str().to_owned()), source: Some(repo.as_str().to_owned()),
}); });
} }
@ -198,11 +285,18 @@ impl ModelArtifactDownloader for HfModelArtifactDownloader {
.await .await
.map_err(|e| ModelServerError::Store(e.to_string()))?; .map_err(|e| ModelServerError::Store(e.to_string()))?;
drop(file); drop(file);
tokio::fs::rename(&tmp_path, &path) tokio::fs::rename(&tmp_path, target)
.await .await
.map_err(|e| ModelServerError::Store(e.to_string()))?; .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 { Ok(ModelArtifactResolution {
path: model_path_from_pathbuf(path)?, path: model_path_from_pathbuf(targets[0].clone())?,
cache_hit: false, 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))) .map_or((repo.as_str(), None), |(base, quant)| (base, Some(quant)))
} }
fn select_gguf_file(siblings: &[HfSibling], quant: Option<&str>) -> Option<String> { /// Matches the llama.cpp / HF multi-part GGUF naming convention, e.g.
let mut ggufs: Vec<&str> = siblings /// `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<regex::Regex> = 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<u32> {
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<String> {
let ggufs: Vec<&str> = siblings
.iter() .iter()
.map(|sibling| sibling.filename.as_str()) .map(|sibling| sibling.filename.as_str())
.filter(|filename| filename.ends_with(".gguf")) .filter(|filename| filename.ends_with(".gguf"))
.collect(); .collect();
ggufs.sort_unstable();
let Some(quant) = quant else { let matching: Vec<&str> = match quant {
return ggufs.first().map(|filename| (*filename).to_owned()); None => ggufs,
}; Some(quant) => {
let quant = quant.to_ascii_lowercase(); let quant = quant.to_ascii_lowercase();
ggufs ggufs
.into_iter() .into_iter()
.find(|filename| filename.to_ascii_lowercase().contains(&quant)) .filter(|filename| filename.to_ascii_lowercase().contains(&quant))
.map(str::to_owned) .collect()
}
};
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 { fn safe_cache_component(raw: &str) -> String {
@ -644,3 +778,55 @@ impl ModelServerRegistry for FsModelServerRegistry {
self.write_doc(&doc).await 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());
}
}