spark_runtime/weights/gguf/
shards.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2
3//! Split-GGUF shard-set resolution for the loader in the parent module.
4
5use anyhow::{Context, Result, bail};
6use std::path::{Path, PathBuf};
7
8use super::sidecar;
9
10/// One resolved split-GGUF file set, in `split.no` order.
11///
12/// A split GGUF is `N` files that share one logical tensor table. llama.cpp
13/// writes the full metadata KV block into shard 0 only; every shard carries
14/// `split.no`, `split.count` and `split.tensors.count`, and each holds its own
15/// slice of the tensors.
16pub struct GgufShardSet {
17    /// Shard paths ordered by `split.no`. Length 1 for an unsplit file.
18    pub paths: Vec<PathBuf>,
19    /// `split.tensors.count` from shard 0: the total across the whole set.
20    pub total_tensors: usize,
21}
22
23/// Resolve every shard of a split GGUF, given the first one.
24///
25/// WHY THIS IS NOT "glob the directory": a directory can legitimately hold more
26/// than one model, an in-progress download, or a leftover shard from a
27/// different quant. The split keys are the only authority on what belongs to
28/// this file set, so the count comes from `split.count` and membership is
29/// checked per shard rather than assumed from the filename.
30///
31/// Refuses, rather than proceeding, when:
32///   * the passed file is not `split.no == 0` (the caller resolved the wrong one)
33///   * any sibling named by the convention is missing or unparseable
34///   * a sibling disagrees about `split.count` or `split.tensors.count`
35///   * the `split.no` values are not exactly `0..count`
36///   * the per-shard tensor counts do not sum to `split.tensors.count`
37///
38/// That last check is the one that matters in practice: a download that lost a
39/// shard still leaves every other shard individually valid and checksum-clean,
40/// so nothing below the tensor-count sum catches it.
41pub fn find_gguf_shards(first: &Path) -> Result<GgufShardSet> {
42    let (_f, _m, g) = sidecar::open_gguf(first)?;
43    let count = match g.get_u64("split.count") {
44        None | Some(0) | Some(1) => {
45            return Ok(GgufShardSet {
46                paths: vec![first.to_path_buf()],
47                total_tensors: g.tensors.len(),
48            });
49        }
50        Some(c) => c as usize,
51    };
52    let no = g.get_u64("split.no").unwrap_or(0) as usize;
53    if no != 0 {
54        bail!(
55            "GGUF split: {} reports split.no={no}, expected shard 0. \
56             Point the loader at the `-00001-of-*` shard.",
57            first.display()
58        );
59    }
60    let total_tensors = g
61        .get_u64("split.tensors.count")
62        .context("GGUF split: shard 0 has 'split.count' but no 'split.tensors.count'")?
63        as usize;
64
65    let name = first
66        .file_name()
67        .and_then(|n| n.to_str())
68        .context("GGUF split: shard path has no filename")?;
69    // llama.cpp convention: `<stem>-00001-of-00007.gguf`, one-based on disk,
70    // zero-based in `split.no`.
71    let marker = format!("-00001-of-{count:05}.gguf");
72    let stem = name.strip_suffix(&marker).with_context(|| {
73        format!(
74            "GGUF split: '{name}' declares split.count={count} but does not end in '{marker}'; \
75             cannot derive sibling shard names"
76        )
77    })?;
78    let dir = first.parent().unwrap_or_else(|| Path::new("."));
79
80    let mut paths = vec![first.to_path_buf()];
81    let mut seen = vec![0usize; count];
82    seen[0] = g.tensors.len();
83    for (i, slot) in seen.iter_mut().enumerate().skip(1) {
84        let p = dir.join(format!("{stem}-{:05}-of-{count:05}.gguf", i + 1));
85        if !p.exists() {
86            bail!(
87                "GGUF split: shard {} of {count} missing: {}",
88                i + 1,
89                p.display()
90            );
91        }
92        let (_f2, _m2, g2) = sidecar::open_gguf(&p)
93            .with_context(|| format!("GGUF split: failed to open shard {}", p.display()))?;
94        let c2 = g2.get_u64("split.count").unwrap_or(0) as usize;
95        if c2 != count {
96            bail!(
97                "GGUF split: {} reports split.count={c2}, shard 0 says {count}",
98                p.display()
99            );
100        }
101        let n2 = g2.get_u64("split.no").unwrap_or(usize::MAX as u64) as usize;
102        if n2 != i {
103            bail!(
104                "GGUF split: {} reports split.no={n2}, expected {i} from its filename",
105                p.display()
106            );
107        }
108        if let Some(t2) = g2.get_u64("split.tensors.count")
109            && t2 as usize != total_tensors
110        {
111            bail!(
112                "GGUF split: {} reports split.tensors.count={t2}, shard 0 says {total_tensors}",
113                p.display()
114            );
115        }
116        *slot = g2.tensors.len();
117        paths.push(p);
118    }
119
120    let summed: usize = seen.iter().sum();
121    if summed != total_tensors {
122        bail!(
123            "GGUF split: shards hold {summed} tensors but 'split.tensors.count' is \
124             {total_tensors} — the file set is incomplete or mismatched (per-shard: {seen:?})"
125        );
126    }
127    Ok(GgufShardSet {
128        paths,
129        total_tensors,
130    })
131}
spark_runtime/weights/gguf/shards.rs

spark_runtime/weights/gguf/
shards.rs
