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:
@ -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