ferrox_models/
device_budget.rs

1//! How many bytes the selected backend will let this process hold --
2//! the right-hand side of [`crate::kv_budget`]'s inequality.
3//!
4//! One rule throughout: **ask the device, do not model the operating
5//! system**. Each backend has exactly one query and it is the vendor's
6//! own answer:
7//!
8//! | backend | query | crate |
9//! |---|---|---|
10//! | Metal | `MTLDevice.recommendedMaxWorkingSetSize` | `ferrox_metal::MetalProfile` |
11//! | CUDA | `cuMemGetInfo` free bytes | `ferrox_cuda::HardwareProfile` |
12//! | CPU | total physical RAM, minus a reserve | `ferrox_cuda::HardwareProfile` |
13//!
14//! The serving plan explicitly rules out the alternative -- process
15//! `phys_footprint` sampling, wired-memory limits, jetsam avoidance, a
16//! `free + inactive + active * ratio` dynamic ceiling. None of it is
17//! here and none of it should be added: a conservative, explainable
18//! number beats a clever one, and every one of those mechanisms is an
19//! Apple-specific workaround for an allocator ferrox does not have.
20//!
21//! # What this number is not
22//!
23//! It is a **ceiling to plan against, not a reservation**. Nothing here
24//! allocates, nothing holds the memory, and every source is a snapshot:
25//! another process can take the VRAM a moment later, and macOS can
26//! shrink a recommended working set under pressure.
27//!
28//! It is also **approximate for ferrox specifically**, for a reason
29//! that has nothing to do with the probe: ferrox mmaps its quantized
30//! weights. Their pages are owned by the kernel's page cache, not by
31//! ferrox, so a check that charges the full checkpoint against this
32//! budget is charging an upper bound. A model that overruns the budget
33//! may still run, page-faulting; a model that fits may still be evicted
34//! by pressure from elsewhere on the machine. Every path that prints
35//! this number says so.
36
37/// Which pool a budget was drawn from.
38#[derive(Debug, Clone, Copy, PartialEq, Eq)]
39pub enum BudgetBackend {
40    Cpu,
41    Metal,
42    Cuda,
43}
44
45impl BudgetBackend {
46    pub fn as_str(self) -> &'static str {
47        match self {
48            BudgetBackend::Cpu => "cpu",
49            BudgetBackend::Metal => "metal",
50            BudgetBackend::Cuda => "cuda",
51        }
52    }
53}
54
55impl std::fmt::Display for BudgetBackend {
56    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
57        f.write_str(self.as_str())
58    }
59}
60
61/// Lets a CLI flag take a backend name without this crate depending on
62/// clap: clap derives a value parser from `FromStr`.
63impl std::str::FromStr for BudgetBackend {
64    type Err = String;
65
66    fn from_str(value: &str) -> Result<Self, Self::Err> {
67        match value.trim().to_ascii_lowercase().as_str() {
68            "cpu" | "host" => Ok(BudgetBackend::Cpu),
69            "metal" => Ok(BudgetBackend::Metal),
70            "cuda" => Ok(BudgetBackend::Cuda),
71            other => Err(format!("unknown backend `{other}` (cpu, metal, cuda)")),
72        }
73    }
74}
75
76/// Fraction of a *host RAM* budget held back for the OS and everything
77/// else running on the machine. Deliberately blunt: the alternative is
78/// modelling the OS, which the plan rules out.
79pub const CPU_RESERVE_FRACTION: f64 = 0.2;
80
81/// Fraction of a *device* budget (Metal working set, free VRAM) held
82/// back for driver allocations, command buffers and the activation
83/// scratch this module does not itemise.
84pub const DEVICE_RESERVE_FRACTION: f64 = 0.1;
85
86/// Overrides the probe entirely (`FERROX_DEVICE_BUDGET_BYTES`). The
87/// escape hatch for a host whose real ceiling is something ferrox
88/// cannot see -- a container memory limit, a shared GPU, an operator
89/// who simply knows better.
90pub const BUDGET_ENV: &str = "FERROX_DEVICE_BUDGET_BYTES";
91
92/// A probed byte ceiling plus the sentence explaining where it came
93/// from. The sentence is not decoration: a budget a user cannot trace
94/// back to a query is a budget they will disable.
95#[derive(Debug, Clone, PartialEq)]
96pub struct DeviceBudget {
97    pub backend: BudgetBackend,
98    /// What the query returned, before the reserve.
99    pub total_bytes: u64,
100    /// `total_bytes` minus the reserve: what a plan may actually spend.
101    pub usable_bytes: u64,
102    /// Held-back fraction, as applied.
103    pub reserve_fraction: f64,
104    /// Human sentence naming the query, e.g.
105    /// "Metal recommendedMaxWorkingSetSize".
106    pub source: String,
107    /// True whenever the checkpoint is mmap'd -- i.e. always, for GGUF
108    /// -- because resident weight bytes are then the kernel's business,
109    /// not ferrox's. Kept as a field rather than a constant so nothing
110    /// downstream can print the budget without deciding what to say
111    /// about it.
112    pub approximate: bool,
113}
114
115impl DeviceBudget {
116    /// Applies `reserve` to `total` and records the source.
117    pub fn new(backend: BudgetBackend, total_bytes: u64, reserve: f64, source: String) -> Self {
118        let reserve = reserve.clamp(0.0, 1.0);
119        DeviceBudget {
120            backend,
121            total_bytes,
122            usable_bytes: (total_bytes as f64 * (1.0 - reserve)) as u64,
123            reserve_fraction: reserve,
124            source,
125            approximate: true,
126        }
127    }
128
129    /// Probes the backend the process is configured to use, honouring
130    /// `FERROX_DEVICE_BUDGET_BYTES` first.
131    ///
132    /// `backend` is the caller's already-resolved choice (the CLI's
133    /// `--device`, the server's `FERROX_METAL`/`FERROX_CUDA`), not a
134    /// second guess at it -- this module decides how much memory a
135    /// backend has, never which backend runs.
136    pub fn detect(backend: BudgetBackend) -> Self {
137        if let Some(bytes) = env_override() {
138            return DeviceBudget {
139                backend,
140                total_bytes: bytes,
141                usable_bytes: bytes,
142                reserve_fraction: 0.0,
143                source: format!("{BUDGET_ENV} override (no reserve applied)"),
144                approximate: true,
145            };
146        }
147        match backend {
148            BudgetBackend::Metal => metal_budget(),
149            BudgetBackend::Cuda => cuda_budget(),
150            BudgetBackend::Cpu => host_ram_budget(),
151        }
152    }
153
154    /// True when nothing could be probed. Callers must treat this as
155    /// "do not enforce" rather than "reject everything": refusing to
156    /// load because a probe failed would be worse than not checking.
157    pub fn is_unknown(&self) -> bool {
158        self.total_bytes == 0
159    }
160
161    /// Where `usable_bytes` came from, spelled out.
162    ///
163    /// `source` describes `total_bytes`, so printing it next to
164    /// `usable_bytes` -- as the CLI's over-budget warning did -- reads
165    /// as "27487790694 bytes is the total physical host RAM" when the
166    /// machine has 32 GiB and 20% is held back. Same two-values-one-
167    /// label shape as the rest of this repo's bugs.
168    pub fn usable_provenance(&self) -> String {
169        if self.is_unknown() {
170            return self.source.clone();
171        }
172        format!(
173            "{:.0}% of {} {}, {:.0}% held back",
174            (1.0 - self.reserve_fraction) * 100.0,
175            self.total_bytes,
176            self.source,
177            self.reserve_fraction * 100.0,
178        )
179    }
180
181    /// The caveat sentence every printer of this number owes the user.
182    pub fn caveat(&self) -> &'static str {
183        "approximate: ferrox mmaps quantized weights, so how much of them stays resident is \
184         the kernel's page cache to decide; this charges the whole checkpoint, which is an \
185         upper bound, and the budget itself is a snapshot, not a reservation"
186    }
187}
188
189impl std::fmt::Display for DeviceBudget {
190    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
191        if self.is_unknown() {
192            return write!(
193                f,
194                "{} budget: unknown ({}); no ceiling enforced",
195                self.backend, self.source
196            );
197        }
198        write!(
199            f,
200            "{} budget: {} usable of {} total ({:.0}% held back) via {}",
201            self.backend,
202            human(self.usable_bytes),
203            human(self.total_bytes),
204            self.reserve_fraction * 100.0,
205            self.source
206        )
207    }
208}
209
210fn env_override() -> Option<u64> {
211    std::env::var(BUDGET_ENV)
212        .ok()
213        .and_then(|v| v.trim().parse::<u64>().ok())
214        .filter(|v| *v > 0)
215}
216
217/// `MTLDevice.recommendedMaxWorkingSetSize`. Without `--features metal`
218/// there is no device to ask and no Metal execution either, so this
219/// falls back to host RAM and says so.
220fn metal_budget() -> DeviceBudget {
221    let profile = ferrox_metal::MetalProfile::detect();
222    if profile.available && profile.recommended_working_set_bytes > 0 {
223        return DeviceBudget::new(
224            BudgetBackend::Metal,
225            profile.recommended_working_set_bytes,
226            DEVICE_RESERVE_FRACTION,
227            format!(
228                "Metal recommendedMaxWorkingSetSize on {}",
229                profile.device_name.as_deref().unwrap_or("unnamed device")
230            ),
231        );
232    }
233    let mut fallback = host_ram_budget();
234    fallback.backend = BudgetBackend::Metal;
235    fallback.source = format!(
236        "no Metal device query available; fell back to {}",
237        fallback.source
238    );
239    fallback
240}
241
242/// `cuMemGetInfo`'s free half, not the card's total: another process
243/// may already hold most of it. Compiles without `--features cuda`,
244/// where `HardwareProfile` honestly reports no device and this falls
245/// back to host RAM.
246fn cuda_budget() -> DeviceBudget {
247    let profile = ferrox_cuda::HardwareProfile::detect();
248    if profile.cuda_available && profile.cuda_vram_free_bytes > 0 {
249        return DeviceBudget::new(
250            BudgetBackend::Cuda,
251            profile.cuda_vram_free_bytes,
252            DEVICE_RESERVE_FRACTION,
253            format!(
254                "cuMemGetInfo free VRAM on {} ({} total)",
255                profile.cuda_device_name.as_deref().unwrap_or("device 0"),
256                human(profile.cuda_vram_total_bytes)
257            ),
258        );
259    }
260    let mut fallback = host_ram_budget();
261    fallback.backend = BudgetBackend::Cuda;
262    fallback.source = format!(
263        "no CUDA device query available; fell back to {}",
264        fallback.source
265    );
266    fallback
267}
268
269/// Total physical RAM minus [`CPU_RESERVE_FRACTION`]. Reported as `0`
270/// on a host whose RAM cannot be read (see
271/// `ferrox_cuda::HardwareProfile`), which
272/// [`DeviceBudget::is_unknown`] turns into "do not enforce".
273fn host_ram_budget() -> DeviceBudget {
274    let total = ferrox_cuda::HardwareProfile::detect().host_ram_total_bytes;
275    if total == 0 {
276        return DeviceBudget {
277            backend: BudgetBackend::Cpu,
278            total_bytes: 0,
279            usable_bytes: 0,
280            reserve_fraction: 0.0,
281            source: "host RAM could not be probed on this platform".to_string(),
282            approximate: true,
283        };
284    }
285    DeviceBudget::new(
286        BudgetBackend::Cpu,
287        total,
288        CPU_RESERVE_FRACTION,
289        "total physical host RAM".to_string(),
290    )
291}
292
293pub(crate) fn human(bytes: u64) -> String {
294    const UNITS: [&str; 5] = ["B", "KiB", "MiB", "GiB", "TiB"];
295    let mut v = bytes as f64;
296    let mut u = 0;
297    while v >= 1024.0 && u < UNITS.len() - 1 {
298        v /= 1024.0;
299        u += 1;
300    }
301    format!("{v:.2} {}", UNITS[u])
302}
303
304#[cfg(test)]
305mod tests {
306    use super::*;
307
308    #[test]
309    fn reserve_is_applied_and_reported() {
310        let b = DeviceBudget::new(BudgetBackend::Cpu, 1000, 0.2, "test".into());
311        assert_eq!(b.total_bytes, 1000);
312        assert_eq!(b.usable_bytes, 800);
313        assert_eq!(b.reserve_fraction, 0.2);
314        assert!(!b.is_unknown());
315        // Always approximate: ferrox mmaps its weights.
316        assert!(b.approximate);
317    }
318
319    #[test]
320    fn a_nonsense_reserve_is_clamped_rather_than_producing_a_negative_budget() {
321        let over = DeviceBudget::new(BudgetBackend::Cpu, 1000, 5.0, "test".into());
322        assert_eq!(over.usable_bytes, 0);
323        let under = DeviceBudget::new(BudgetBackend::Cpu, 1000, -1.0, "test".into());
324        assert_eq!(under.usable_bytes, 1000);
325    }
326
327    #[test]
328    fn zero_total_reads_as_unknown_not_as_a_zero_ceiling() {
329        let b = DeviceBudget::new(BudgetBackend::Cpu, 0, 0.2, "nothing to probe".into());
330        assert!(b.is_unknown());
331        assert!(b.to_string().contains("no ceiling enforced"), "{b}");
332    }
333
334    /// Runs in both worlds, like the backend probes themselves: on a
335    /// host that can report RAM the budget must be plausible and
336    /// smaller than the total; on one that cannot it must be unknown.
337    #[test]
338    fn cpu_budget_is_either_unknown_or_a_plausible_fraction_of_real_ram() {
339        let b = DeviceBudget::detect(BudgetBackend::Cpu);
340        assert_eq!(b.backend, BudgetBackend::Cpu);
341        if b.is_unknown() {
342            assert_eq!(b.usable_bytes, 0);
343        } else {
344            assert!(b.total_bytes > 128 * 1024 * 1024);
345            assert!(b.usable_bytes < b.total_bytes);
346            assert!(b.usable_bytes > b.total_bytes / 2);
347            assert!(b.to_string().contains("host RAM"), "{b}");
348        }
349    }
350
351    /// Without `--features metal`/`cuda` these must still resolve (to
352    /// the host-RAM fallback) rather than failing to compile or
353    /// panicking -- the whole point of the honest-zero probe structs.
354    #[test]
355    fn accelerator_budgets_fall_back_to_host_ram_when_no_device_answers() {
356        for backend in [BudgetBackend::Metal, BudgetBackend::Cuda] {
357            let b = DeviceBudget::detect(backend);
358            assert_eq!(b.backend, backend);
359            if b.source.contains("fell back") {
360                assert!(b.source.contains("host RAM"), "{b}");
361            }
362        }
363    }
364
365    #[test]
366    fn human_bytes_are_readable_at_every_scale() {
367        assert_eq!(human(0), "0.00 B");
368        assert_eq!(human(1024), "1.00 KiB");
369        assert_eq!(human(3 * 1024 * 1024 * 1024), "3.00 GiB");
370    }
371}
ferrox_models/device_budget.rs

ferrox_models/
device_budget.rs
