Skip to main content

clawft_kernel/workload_runtime/infer/
roster.rs

1//! Import of the model lab's roster (`~/llm/docs/models/queue.yaml`, card
2//! mesh-placement-20; ADR-101 section 8 step 2) into [`InferenceSpec`]s.
3//!
4//! The roster is read, never written, and `~/llm` stays a separate repo:
5//! the caller names the file. What the roster holds maps directly (role id,
6//! backend, port, `ram_gb`); what it does not hold (co-residency
7//! `excludes`, latency class, stickiness) comes from an operator
8//! [`RosterOverlay`], with the ADR's defaults otherwise: inference is sticky
9//! (warm KV), roles are `batch` unless the overlay or the role's text says
10//! voice. Entries the adapters cannot run (vision, embeddings, STT, TTS
11//! through `mlx-audio`) and aliases of an earlier entry are reported in
12//! [`ImportedRoster::skipped`] with the reason, not dropped silently.
13
14use std::collections::BTreeMap;
15
16use serde::Deserialize;
17
18use super::residency::GB;
19use super::spec::{InferFlavor, InferenceSpec, LatencyClass};
20use crate::workload_pkg::manifest::valid_token;
21
22/// Largest roster file read.
23pub const MAX_ROSTER_BYTES: usize = 1024 * 1024;
24const MAX_ENTRIES: usize = 256;
25
26/// Operator-supplied facts the roster does not carry.
27#[derive(Debug, Clone, Default, Deserialize)]
28#[serde(deny_unknown_fields)]
29pub struct RosterOverlay {
30    /// Roles each role cannot be resident beside (`role:planner` or `planner`).
31    #[serde(default)]
32    pub excludes: BTreeMap<String, Vec<String>>,
33    /// Latency class per role.
34    #[serde(default)]
35    pub latency_class: BTreeMap<String, LatencyClass>,
36    /// Stickiness per role (default true).
37    #[serde(default)]
38    pub sticky: BTreeMap<String, bool>,
39}
40
41/// One roster entry turned into a spec, with what was left behind.
42#[derive(Debug, Clone, PartialEq)]
43pub struct ImportedSpec {
44    /// The inference spec (`model` is the roster's display name: adopt the
45    /// weights under that name in the model registry).
46    pub spec: InferenceSpec,
47    /// Roster `repo` (where the weights came from).
48    pub repo: String,
49    /// Roster `status` (`live`, `ready`, ...).
50    pub status: String,
51}
52
53/// The result of an import.
54#[derive(Debug, Clone, Default, PartialEq)]
55pub struct ImportedRoster {
56    /// Importable entries, in roster order.
57    pub specs: Vec<ImportedSpec>,
58    /// `(id, reason)` for each entry not imported.
59    pub skipped: Vec<(String, String)>,
60}
61
62impl ImportedRoster {
63    /// The spec for roster id `id`.
64    pub fn get(&self, id: &str) -> Option<&ImportedSpec> {
65        self.specs.iter().find(|s| s.spec.role == id)
66    }
67}
68
69#[derive(Debug, Deserialize)]
70struct File {
71    #[serde(default)]
72    roster: Vec<Raw>,
73}
74
75#[derive(Debug, Deserialize)]
76struct Raw {
77    id: String,
78    #[serde(default)]
79    role: String,
80    #[serde(default)]
81    name: String,
82    #[serde(default)]
83    repo: String,
84    #[serde(default)]
85    ram_gb: serde_yaml::Value,
86    #[serde(default)]
87    backend: String,
88    #[serde(default)]
89    port: Option<u16>,
90    #[serde(default)]
91    status: String,
92}
93
94fn ram_bytes(v: &serde_yaml::Value) -> Option<u64> {
95    let gb = match v {
96        serde_yaml::Value::Number(n) => n.as_f64()?,
97        // The roster writes "~10" for unmeasured entries.
98        serde_yaml::Value::String(s) => s.trim().trim_start_matches('~').parse::<f64>().ok()?,
99        _ => return None,
100    };
101    (gb.is_finite() && (0.0..=4096.0).contains(&gb)).then_some((gb * GB as f64) as u64)
102}
103
104/// The adapter flavor the roster's `backend` text names (the first listed
105/// when it offers alternatives).
106pub fn flavor_for_backend(backend: &str) -> Option<InferFlavor> {
107    let first = backend.split_whitespace().next()?.to_ascii_lowercase();
108    match first.as_str() {
109        "mlx_lm" | "mlx-lm" => Some(InferFlavor::MlxLm),
110        "llamacpp" | "llama.cpp" | "llama-server" => Some(InferFlavor::LlamaCpp),
111        "ollama" => Some(InferFlavor::Ollama),
112        _ => None,
113    }
114}
115
116/// Parse roster YAML. `overlay` supplies excludes, latency and stickiness.
117pub fn import_roster(yaml: &str, overlay: &RosterOverlay) -> Result<ImportedRoster, String> {
118    if yaml.len() > MAX_ROSTER_BYTES {
119        return Err(format!("roster is over {MAX_ROSTER_BYTES} bytes"));
120    }
121    let file: File = serde_yaml::from_str(yaml).map_err(|e| format!("roster: {e}"))?;
122    if file.roster.len() > MAX_ENTRIES {
123        return Err(format!("roster lists more than {MAX_ENTRIES} entries"));
124    }
125    let mut out = ImportedRoster::default();
126    // (flavor, port, repo) -> id of the first entry, to spot aliases.
127    let mut seen: BTreeMap<(String, Option<u16>, String), String> = BTreeMap::new();
128    for r in file.roster {
129        let skip = |out: &mut ImportedRoster, why: String| out.skipped.push((r.id.clone(), why));
130        if !valid_token(&r.id, 64) {
131            skip(&mut out, "id is not a plain token".into());
132            continue;
133        }
134        let Some(flavor) = flavor_for_backend(&r.backend) else {
135            skip(&mut out, format!("backend {:?} has no inference adapter", r.backend));
136            continue;
137        };
138        let key = (flavor.id().to_string(), r.port, r.repo.clone());
139        if let Some(first) = seen.get(&key) {
140            skip(&mut out, format!("alias of {first} (same server, port and weights)"));
141            continue;
142        }
143        let Some(weights) = ram_bytes(&r.ram_gb) else {
144            skip(&mut out, "ram_gb is missing or not a number".into());
145            continue;
146        };
147        let mut spec = InferenceSpec::new(&r.id, flavor);
148        spec.model = valid_token(&r.name, 128).then(|| r.name.clone());
149        spec.serve.port = r.port;
150        spec.memory.weights_bytes = weights;
151        spec.sticky = overlay.sticky.get(&r.id).copied().unwrap_or(true);
152        spec.latency_class = overlay.latency_class.get(&r.id).copied().unwrap_or_else(|| {
153            let t = format!("{} {}", r.id, r.role).to_ascii_lowercase();
154            if t.contains("voice") || t.contains("realtime") {
155                LatencyClass::Interactive
156            } else {
157                LatencyClass::Batch
158            }
159        });
160        spec.excludes = overlay.excludes.get(&r.id).cloned().unwrap_or_default();
161        if let Err(e) = spec.validate() {
162            skip(&mut out, format!("not a valid inference spec: {e}"));
163            continue;
164        }
165        seen.insert(key, r.id.clone());
166        out.specs.push(ImportedSpec { spec, repo: r.repo, status: r.status });
167    }
168    Ok(out)
169}