clawft_kernel/workload_runtime/infer/
roster.rs1use 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
22pub const MAX_ROSTER_BYTES: usize = 1024 * 1024;
24const MAX_ENTRIES: usize = 256;
25
26#[derive(Debug, Clone, Default, Deserialize)]
28#[serde(deny_unknown_fields)]
29pub struct RosterOverlay {
30 #[serde(default)]
32 pub excludes: BTreeMap<String, Vec<String>>,
33 #[serde(default)]
35 pub latency_class: BTreeMap<String, LatencyClass>,
36 #[serde(default)]
38 pub sticky: BTreeMap<String, bool>,
39}
40
41#[derive(Debug, Clone, PartialEq)]
43pub struct ImportedSpec {
44 pub spec: InferenceSpec,
47 pub repo: String,
49 pub status: String,
51}
52
53#[derive(Debug, Clone, Default, PartialEq)]
55pub struct ImportedRoster {
56 pub specs: Vec<ImportedSpec>,
58 pub skipped: Vec<(String, String)>,
60}
61
62impl ImportedRoster {
63 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 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
104pub 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
116pub 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 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}