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.
///
/// 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<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 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<String>,
}
#[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<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 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<String> {
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<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()
.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());
}
}