metrale_kernels/
resolve.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: Kernel-target resolution: which compiled target serves a
4//! checkpoint. The rules are pure functions over [`ResolveCandidate`]s, so
5//! they are tested without compiled kernels.
6//!
7//! `(model_type, hidden_size)` alone can collide: `kernels/gb10/qwen3.6-27b`
8//! and `kernels/gb10/qwen3.8-27b` both declare `(qwen3_5, 5120)`.
9//!
10//! The rules:
11//! 1. Exact `(model_type, Some(hidden_size))` declarations are tried before
12//!    wildcard `(model_type, None)` ones.
13//! 2. Within the first tier that matches, one target name wins outright.
14//! 3. Several names are a collision, broken by `match_names`: a candidate
15//!    survives when one of its needles is a case-insensitive substring of a
16//!    checkpoint reference (HF id, `--model-name`, resolved model dir).
17//!    Exactly one surviving name wins.
18//! 4. Zero or several survivors are [`TargetResolveError::Ambiguous`]; the
19//!    wildcard tier is not tried.
20//! 5. `--kernel-target <name>` pins a target and skips the tie-break, but the
21//!    pinned target must declare the checkpoint's `model_type` for this
22//!    `hidden_size` or a wildcard ([`TargetResolveError::PinIncompatible`]).
23//!
24//! Owner: kernels crate.
25//! Invariants:
26//! - Resolution never chooses between different target names by iteration
27//!   order: a collision it cannot break is an error.
28//! - `build.rs` (`validate_collision_match_names`) panics when a target in a
29//!   collision of the build's targets has no `match_names`.
30
31use crate::{ModelTypeMatch, TargetPtxSet};
32
33/// 2026-09-25: The parts of one compiled target that resolution reads,
34/// borrowed so production can wrap `TargetPtxSet`s without copying blobs.
35pub struct ResolveCandidate<'a> {
36    /// 2026-09-25: Kernel-target directory name (`KernelTarget::model`), e.g.
37    /// `"qwen3.8-27b"`. A multi-quant build repeats a name once per quant;
38    /// candidates with the same name never collide, and the first one wins.
39    pub name: &'a str,
40    /// 2026-09-25: MODEL.toml `[[model_types]]`.
41    pub type_matches: &'a [ModelTypeMatch],
42    /// 2026-09-25: MODEL.toml `[model] match_names` needles.
43    pub match_names: &'a [&'a str],
44}
45
46/// 2026-09-25: Why resolution could not choose a target.
47#[derive(Debug, Clone, PartialEq, Eq)]
48pub enum TargetResolveError {
49    /// 2026-09-25: Several target names claim the checkpoint's
50    /// `(model_type, hidden_size)` in one tier, and the reference tie-break
51    /// did not leave exactly one.
52    Ambiguous {
53        model_type: String,
54        hidden_size: usize,
55        /// 2026-09-25: `"exact"` or `"wildcard"`: the tier that collided.
56        tier: &'static str,
57        /// 2026-09-25: Distinct target names in that tier, with their needles.
58        candidates: Vec<(String, Vec<String>)>,
59        /// 2026-09-25: The names whose needles matched a reference.
60        matched: Vec<String>,
61        /// 2026-09-25: The checkpoint references searched.
62        model_refs: Vec<String>,
63    },
64    /// 2026-09-25: `--kernel-target` named a target this binary did not
65    /// compile.
66    PinNotFound { pin: String, available: Vec<String> },
67    /// 2026-09-25: `--kernel-target` named a compiled target that declares
68    /// neither the checkpoint's `(model_type, hidden_size)` nor a wildcard for
69    /// its `model_type`.
70    PinIncompatible {
71        pin: String,
72        model_type: String,
73        hidden_size: usize,
74        declared: Vec<String>,
75    },
76}
77
78impl std::fmt::Display for TargetResolveError {
79    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
80        match self {
81            Self::Ambiguous {
82                model_type,
83                hidden_size,
84                tier,
85                candidates,
86                matched,
87                model_refs,
88            } => {
89                let cands = candidates
90                    .iter()
91                    .map(|(n, needles)| format!("{n} (match_names: {needles:?})"))
92                    .collect::<Vec<_>>()
93                    .join(", ");
94                let outcome = if matched.is_empty() {
95                    "no checkpoint reference names any of them".to_string()
96                } else {
97                    format!("the references match several of them: {matched:?}")
98                };
99                write!(
100                    f,
101                    "AMBIGUOUS kernel target: {} compiled targets declare {tier} support for \
102                     (model_type '{model_type}', hidden_size {hidden_size}) — [{cands}] — and \
103                     {outcome} (references searched: {model_refs:?}). Refusing to pick by build \
104                     order. Fix: serve with a model id/path that contains exactly one target's \
105                     match_names needle, pin explicitly with --kernel-target <name>, or build \
106                     single-target with METRALE_TARGET_MODEL=<name>.",
107                    candidates.len(),
108                )
109            }
110            Self::PinNotFound { pin, available } => write!(
111                f,
112                "--kernel-target '{pin}' does not name a compiled kernel target \
113                 (available: {available:?})"
114            ),
115            Self::PinIncompatible {
116                pin,
117                model_type,
118                hidden_size,
119                declared,
120            } => write!(
121                f,
122                "--kernel-target '{pin}' is compiled but declares no support for this \
123                 checkpoint's (model_type '{model_type}', hidden_size {hidden_size}) — it \
124                 declares {declared:?}. Serving another architecture's kernels would be \
125                 garbage; refusing."
126            ),
127        }
128    }
129}
130
131impl std::error::Error for TargetResolveError {}
132
133/// 2026-09-25: Whether any non-empty needle, lower-cased, is a substring of
134/// any lower-cased reference.
135fn needles_hit(match_names: &[&str], refs_lower: &[String]) -> bool {
136    match_names.iter().any(|needle| {
137        let n = needle.to_lowercase();
138        !n.is_empty() && refs_lower.iter().any(|r| r.contains(&n))
139    })
140}
141
142/// 2026-09-25: Distinct candidate names among `idxs`, in first-seen order.
143fn distinct_names<'a>(candidates: &[ResolveCandidate<'a>], idxs: &[usize]) -> Vec<&'a str> {
144    let mut names: Vec<&str> = Vec::new();
145    for &i in idxs {
146        if !names.contains(&candidates[i].name) {
147            names.push(candidates[i].name);
148        }
149    }
150    names
151}
152
153/// 2026-09-25: Which candidate serves `(model_type, hidden_size)` for a
154/// checkpoint identified by `model_refs`: its index, `Ok(None)` when no tier
155/// declares the pair, or [`TargetResolveError::Ambiguous`] when a collision
156/// does not break to exactly one name.
157pub fn resolve_target(
158    candidates: &[ResolveCandidate<'_>],
159    model_type: &str,
160    hidden_size: usize,
161    model_refs: &[&str],
162) -> Result<Option<usize>, TargetResolveError> {
163    let refs_lower: Vec<String> = model_refs.iter().map(|r| r.to_lowercase()).collect();
164
165    let tiers: [(&'static str, Option<usize>); 2] =
166        [("exact", Some(hidden_size)), ("wildcard", None)];
167    for (tier, want_hidden) in tiers {
168        let idxs: Vec<usize> = candidates
169            .iter()
170            .enumerate()
171            .filter(|(_, c)| {
172                c.type_matches
173                    .iter()
174                    .any(|m| m.model_type == model_type && m.hidden_size == want_hidden)
175            })
176            .map(|(i, _)| i)
177            .collect();
178        if idxs.is_empty() {
179            continue;
180        }
181        let names = distinct_names(candidates, &idxs);
182        if names.len() == 1 {
183            // 2026-09-25: One name, possibly several quant variants: the first
184            // wins.
185            return Ok(Some(idxs[0]));
186        }
187        let matched: Vec<&str> = names
188            .iter()
189            .copied()
190            .filter(|n| {
191                idxs.iter().any(|&i| {
192                    candidates[i].name == *n && needles_hit(candidates[i].match_names, &refs_lower)
193                })
194            })
195            .collect();
196        if let [winner] = matched.as_slice() {
197            let idx = idxs
198                .iter()
199                .copied()
200                .find(|&i| candidates[i].name == *winner)
201                .expect("winner name came from idxs");
202            return Ok(Some(idx));
203        }
204        // 2026-09-25: Zero or several survivors: error, without trying the
205        // wildcard tier.
206        return Err(TargetResolveError::Ambiguous {
207            model_type: model_type.to_string(),
208            hidden_size,
209            tier,
210            candidates: names
211                .iter()
212                .map(|n| {
213                    let needles = idxs
214                        .iter()
215                        .filter(|&&i| candidates[i].name == *n)
216                        .flat_map(|&i| candidates[i].match_names.iter().map(|s| s.to_string()))
217                        .collect();
218                    (n.to_string(), needles)
219                })
220                .collect(),
221            matched: matched.iter().map(|n| n.to_string()).collect(),
222            model_refs: model_refs.iter().map(|r| r.to_string()).collect(),
223        });
224    }
225    Ok(None)
226}
227
228/// 2026-09-25: Resolve a `--kernel-target` pin, matched case-insensitively
229/// against target names. The first same-named candidate that declares the
230/// checkpoint's `model_type` with this `hidden_size` or a wildcard wins; the
231/// tie-break is not consulted.
232pub fn resolve_pinned(
233    candidates: &[ResolveCandidate<'_>],
234    pin: &str,
235    model_type: &str,
236    hidden_size: usize,
237) -> Result<usize, TargetResolveError> {
238    let pinned: Vec<usize> = candidates
239        .iter()
240        .enumerate()
241        .filter(|(_, c)| c.name.eq_ignore_ascii_case(pin))
242        .map(|(i, _)| i)
243        .collect();
244    if pinned.is_empty() {
245        let all: Vec<usize> = (0..candidates.len()).collect();
246        return Err(TargetResolveError::PinNotFound {
247            pin: pin.to_string(),
248            available: distinct_names(candidates, &all)
249                .into_iter()
250                .map(String::from)
251                .collect(),
252        });
253    }
254    let compatible = pinned.iter().copied().find(|&i| {
255        candidates[i].type_matches.iter().any(|m| {
256            m.model_type == model_type
257                && (m.hidden_size.is_none() || m.hidden_size == Some(hidden_size))
258        })
259    });
260    compatible.ok_or_else(|| TargetResolveError::PinIncompatible {
261        pin: pin.to_string(),
262        model_type: model_type.to_string(),
263        hidden_size,
264        declared: pinned
265            .iter()
266            .flat_map(|&i| candidates[i].type_matches.iter())
267            .map(|m| format!("({}, {:?})", m.model_type, m.hidden_size))
268            .collect(),
269    })
270}
271
272/// 2026-09-25: The compiled target that serves a checkpoint, by the module's
273/// rules: [`resolve_pinned`] when `pinned_target` (`--kernel-target`) is set,
274/// else [`resolve_target`] over `all_ptx_sets()`. `Ok(None)` when no compiled
275/// target declares the pair.
276pub fn ptx_for_config(
277    model_type: &str,
278    hidden_size: usize,
279    model_refs: &[&str],
280    pinned_target: Option<&str>,
281) -> Result<Option<TargetPtxSet>, TargetResolveError> {
282    let targets = crate::all_ptx_sets();
283    let candidates: Vec<ResolveCandidate<'_>> = targets
284        .iter()
285        .map(|t| ResolveCandidate {
286            name: t.target.model,
287            type_matches: &t.model_type_matches,
288            match_names: t.match_names,
289        })
290        .collect();
291    let idx = match pinned_target {
292        Some(pin) => Some(resolve_pinned(&candidates, pin, model_type, hidden_size)?),
293        None => resolve_target(&candidates, model_type, hidden_size, model_refs)?,
294    };
295    drop(candidates);
296    Ok(idx.and_then(|i| targets.into_iter().nth(i)))
297}
298
299/// 2026-09-25: The compiled target with exactly this `(model, quant)`
300/// identity, or `None`. For callers that already know the resolved target,
301/// such as the dashboard's kernel table, which looks up the target `serve`
302/// published.
303pub fn ptx_for_exact_target(model: &str, quant: &str) -> Option<TargetPtxSet> {
304    crate::all_ptx_sets()
305        .into_iter()
306        .find(|t| t.target.model == model && t.target.quant == quant)
307}
308
309#[cfg(test)]
310#[path = "resolve_tests.rs"]
311mod tests;
metrale_kernels/resolve.rs

metrale_kernels/
resolve.rs
