clawft_kernel/workload_runtime/infer/
residency.rs

1//! Unified-memory budget and co-residency (card mesh-placement-20; ADR-101
2//! section 6). One ledger is shared by every managed adapter on a node:
3//! an instance reserves its weights plus KV budget when it starts and
4//! releases them when it stops or is unloaded. A start that would exceed
5//! the budget, or that excludes (or is excluded by) a running role, is
6//! refused as unplaceable, with the reason.
7//!
8//! Memory is counted in decimal gigabytes at the edges (`ram_gb` of the
9//! model lab's roster, the operator's budget) and in bytes here.
10
11use std::collections::BTreeMap;
12use std::sync::Mutex;
13
14use super::spec::InferenceSpec;
15
16/// Bytes in the roster's and the operator's "GB".
17pub const GB: u64 = 1_000_000_000;
18
19/// Why an instance cannot be resident beside the others.
20#[derive(Debug, Clone, PartialEq, Eq)]
21pub enum CoResidency {
22    /// Its memory does not fit in what is left of the budget.
23    OverBudget {
24        /// Role that was refused.
25        role: String,
26        /// Bytes it needs.
27        need: u64,
28        /// Bytes still free.
29        free: u64,
30        /// The whole budget.
31        budget: u64,
32        /// Who holds the rest: (role, bytes).
33        holders: Vec<(String, u64)>,
34    },
35    /// Either side lists the other in `excludes`.
36    Excluded {
37        /// Role that was refused.
38        role: String,
39        /// Running role it cannot sit beside.
40        with: String,
41        /// Which side declared the exclusion.
42        declared_by: String,
43    },
44}
45
46fn gb(b: u64) -> String {
47    format!("{:.1} GB", b as f64 / GB as f64)
48}
49
50impl std::fmt::Display for CoResidency {
51    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
52        match self {
53            Self::OverBudget { role, need, free, budget, holders } => {
54                let held = holders
55                    .iter()
56                    .map(|(r, b)| format!("{r} {}", gb(*b)))
57                    .collect::<Vec<_>>()
58                    .join(", ");
59                write!(
60                    f,
61                    "unplaceable: '{role}' needs {} but only {} of the {} unified-memory budget is free (held by {})",
62                    gb(*need),
63                    gb(*free),
64                    gb(*budget),
65                    if held.is_empty() { "nothing" } else { &held }
66                )
67            }
68            Self::Excluded { role, with, declared_by } => write!(
69                f,
70                "unplaceable: '{role}' cannot be resident beside '{with}' ('{declared_by}' excludes it)"
71            ),
72        }
73    }
74}
75
76/// Bytes a spec holds while resident.
77pub fn footprint(spec: &InferenceSpec) -> u64 {
78    spec.memory.weights_bytes.saturating_add(spec.memory.kv_budget_bytes)
79}
80
81fn excludes(spec: &InferenceSpec, other_role: &str) -> bool {
82    spec.excludes
83        .iter()
84        .any(|x| x == other_role || x.strip_prefix("role:") == Some(other_role))
85}
86
87/// The node's resident instances.
88pub struct ResidencyLedger {
89    budget: Option<u64>,
90    held: Mutex<BTreeMap<String, InferenceSpec>>,
91}
92
93impl ResidencyLedger {
94    /// A ledger with `budget_bytes` of unified memory (`None`: only
95    /// exclusions are checked).
96    pub fn new(budget_bytes: Option<u64>) -> Self {
97        Self { budget: budget_bytes, held: Mutex::new(BTreeMap::new()) }
98    }
99
100    /// The budget, if any.
101    pub fn budget(&self) -> Option<u64> {
102        self.budget
103    }
104
105    /// Reserve `spec`'s memory for `instance`. `Ok(true)` when newly
106    /// reserved, `Ok(false)` when this instance already held it.
107    pub fn reserve(&self, instance: &str, spec: &InferenceSpec) -> Result<bool, CoResidency> {
108        let mut g = self.held.lock().unwrap_or_else(|e| e.into_inner());
109        if g.contains_key(instance) {
110            return Ok(false);
111        }
112        for (id, other) in g.iter() {
113            if id == instance {
114                continue;
115            }
116            if excludes(spec, &other.role) {
117                return Err(CoResidency::Excluded {
118                    role: spec.role.clone(),
119                    with: other.role.clone(),
120                    declared_by: spec.role.clone(),
121                });
122            }
123            if excludes(other, &spec.role) {
124                return Err(CoResidency::Excluded {
125                    role: spec.role.clone(),
126                    with: other.role.clone(),
127                    declared_by: other.role.clone(),
128                });
129            }
130        }
131        if let Some(budget) = self.budget {
132            let used: u64 = g.values().map(footprint).fold(0, u64::saturating_add);
133            let need = footprint(spec);
134            let free = budget.saturating_sub(used);
135            if need > free {
136                return Err(CoResidency::OverBudget {
137                    role: spec.role.clone(),
138                    need,
139                    free,
140                    budget,
141                    holders: g.values().map(|s| (s.role.clone(), footprint(s))).collect(),
142                });
143            }
144        }
145        g.insert(instance.to_string(), spec.clone());
146        Ok(true)
147    }
148
149    /// Release what `instance` held (no-op if nothing).
150    pub fn release(&self, instance: &str) {
151        self.held.lock().unwrap_or_else(|e| e.into_inner()).remove(instance);
152    }
153
154    /// Bytes reserved.
155    pub fn used(&self) -> u64 {
156        self.held
157            .lock()
158            .unwrap_or_else(|e| e.into_inner())
159            .values()
160            .map(footprint)
161            .fold(0, u64::saturating_add)
162    }
163
164    /// Resident roles and their bytes.
165    pub fn snapshot(&self) -> Vec<(String, u64)> {
166        self.held
167            .lock()
168            .unwrap_or_else(|e| e.into_inner())
169            .values()
170            .map(|s| (s.role.clone(), footprint(s)))
171            .collect()
172    }
173}
174
175#[cfg(test)]
176mod tests {
177    use super::*;
178    use crate::workload_runtime::infer::spec::InferFlavor;
179
180    fn spec(role: &str, gb_: u64, excludes: &[&str]) -> InferenceSpec {
181        let mut s = InferenceSpec::new(role, InferFlavor::MlxLm);
182        s.memory.weights_bytes = gb_ * GB;
183        s.excludes = excludes.iter().map(|x| x.to_string()).collect();
184        s
185    }
186
187    #[test]
188    fn fits_until_the_budget_is_spent_then_says_who_holds_it() {
189        let l = ResidencyLedger::new(Some(96 * GB));
190        assert!(l.reserve("a", &spec("coder-daily", 55, &[])).unwrap());
191        assert!(l.reserve("b", &spec("coder-small", 9, &[])).unwrap());
192        let e = l.reserve("c", &spec("planner", 55, &[])).unwrap_err();
193        let text = e.to_string();
194        assert!(text.contains("'planner' needs 55.0 GB"), "{text}");
195        assert!(text.contains("32.0 GB of the 96.0 GB"), "{text}");
196        assert!(text.contains("coder-daily 55.0 GB") && text.contains("coder-small 9.0 GB"), "{text}");
197        assert_eq!(l.used(), 64 * GB, "a refusal reserves nothing");
198        // Freeing makes room.
199        l.release("a");
200        assert!(l.reserve("c", &spec("planner", 55, &[])).unwrap());
201    }
202
203    #[test]
204    fn exactly_the_budget_fits_and_one_byte_more_does_not() {
205        let l = ResidencyLedger::new(Some(10 * GB));
206        let mut s = spec("x", 10, &[]);
207        assert!(l.reserve("x", &s).unwrap());
208        l.release("x");
209        s.memory.kv_budget_bytes = 1;
210        assert!(matches!(l.reserve("x", &s), Err(CoResidency::OverBudget { .. })));
211    }
212
213    #[test]
214    fn exclusions_hold_in_both_directions_and_with_the_role_prefix() {
215        let l = ResidencyLedger::new(None);
216        assert!(l.reserve("a", &spec("planner", 1, &["role:coder-daily"])).unwrap());
217        // The newcomer is the one excluded by a running role.
218        let e = l.reserve("b", &spec("coder-daily", 1, &[])).unwrap_err();
219        assert_eq!(
220            e,
221            CoResidency::Excluded { role: "coder-daily".into(), with: "planner".into(), declared_by: "planner".into() }
222        );
223        assert!(e.to_string().contains("'planner' excludes it"));
224        // The newcomer excludes a running role (bare name form).
225        let e = l.reserve("c", &spec("swarm", 1, &["planner"])).unwrap_err();
226        assert!(matches!(e, CoResidency::Excluded { ref declared_by, .. } if declared_by == "swarm"));
227        // Unrelated roles coexist; no budget means no memory check.
228        assert!(l.reserve("d", &spec("embed", 500, &[])).unwrap());
229    }
230
231    #[test]
232    fn reserving_twice_for_one_instance_counts_once() {
233        let l = ResidencyLedger::new(Some(60 * GB));
234        let s = spec("coder", 55, &[]);
235        assert!(l.reserve("a", &s).unwrap());
236        assert!(!l.reserve("a", &s).unwrap());
237        assert_eq!(l.used(), 55 * GB);
238    }
239}
clawft_kernel/workload_runtime/infer/residency.rs

clawft_kernel/workload_runtime/infer/
residency.rs
