fix(backend): sélection GGUF prioritaire au fichier fusionné sur les shards
Quand un repo HuggingFace héberge à la fois le fichier GGUF fusionné et ses shards pour une même quantisation, select_gguf_file (tri alphabétique naïf) pouvait retenir un shard partiel plutôt que le fichier fusionné ou l'ensemble complet des shards, causant un échec immédiat au démarrage du serveur llama.cpp (agent Context). select_gguf_files renvoie désormais soit le fichier fusionné seul (priorité), soit tous les shards correspondants triés par index — jamais un shard isolé. Le téléchargement, le cache et la reprise gèrent le cas multi-fichiers (manifest de shards, chemins locaux préservant le nom distant pour l'autoload de llama.cpp). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
@ -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());
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user