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    /// The caveat sentence every printer of this number owes the user.
162    pub fn caveat(&self) -> &'static str {
163        "approximate: ferrox mmaps quantized weights, so their resident cost is the \
164         kernel's page cache to decide -- this charges the whole checkpoint, which is an \
165         upper bound, and the budget itself is a snapshot, not a reservation"
166    }
167}
168
169impl std::fmt::Display for DeviceBudget {
170    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
171        if self.is_unknown() {
172            return write!(
173                f,
174                "{} budget: unknown ({}); no ceiling enforced",
175                self.backend, self.source
176            );
177        }
178        write!(
179            f,
180            "{} budget: {} usable of {} total ({:.0}% held back) via {}",
181            self.backend,
182            human(self.usable_bytes),
183            human(self.total_bytes),
184            self.reserve_fraction * 100.0,
185            self.source
186        )
187    }
188}
189
190fn env_override() -> Option<u64> {
191    std::env::var(BUDGET_ENV)
192        .ok()
193        .and_then(|v| v.trim().parse::<u64>().ok())
194        .filter(|v| *v > 0)
195}
196
197/// `MTLDevice.recommendedMaxWorkingSetSize`. Without `--features metal`
198/// there is no device to ask and no Metal execution either, so this
199/// falls back to host RAM and says so.
200fn metal_budget() -> DeviceBudget {
201    let profile = ferrox_metal::MetalProfile::detect();
202    if profile.available && profile.recommended_working_set_bytes > 0 {
203        return DeviceBudget::new(
204            BudgetBackend::Metal,
205            profile.recommended_working_set_bytes,
206            DEVICE_RESERVE_FRACTION,
207            format!(
208                "Metal recommendedMaxWorkingSetSize on {}",
209                profile.device_name.as_deref().unwrap_or("unnamed device")
210            ),
211        );
212    }
213    let mut fallback = host_ram_budget();
214    fallback.backend = BudgetBackend::Metal;
215    fallback.source = format!(
216        "no Metal device query available; fell back to {}",
217        fallback.source
218    );
219    fallback
220}
221
222/// `cuMemGetInfo`'s free half, not the card's total: another process
223/// may already hold most of it. Compiles without `--features cuda`,
224/// where `HardwareProfile` honestly reports no device and this falls
225/// back to host RAM.
226fn cuda_budget() -> DeviceBudget {
227    let profile = ferrox_cuda::HardwareProfile::detect();
228    if profile.cuda_available && profile.cuda_vram_free_bytes > 0 {
229        return DeviceBudget::new(
230            BudgetBackend::Cuda,
231            profile.cuda_vram_free_bytes,
232            DEVICE_RESERVE_FRACTION,
233            format!(
234                "cuMemGetInfo free VRAM on {} ({} total)",
235                profile.cuda_device_name.as_deref().unwrap_or("device 0"),
236                human(profile.cuda_vram_total_bytes)
237            ),
238        );
239    }
240    let mut fallback = host_ram_budget();
241    fallback.backend = BudgetBackend::Cuda;
242    fallback.source = format!(
243        "no CUDA device query available; fell back to {}",
244        fallback.source
245    );
246    fallback
247}
248
249/// Total physical RAM minus [`CPU_RESERVE_FRACTION`]. Reported as `0`
250/// on a host whose RAM cannot be read (see
251/// `ferrox_cuda::HardwareProfile`), which
252/// [`DeviceBudget::is_unknown`] turns into "do not enforce".
253fn host_ram_budget() -> DeviceBudget {
254    let total = ferrox_cuda::HardwareProfile::detect().host_ram_total_bytes;
255    if total == 0 {
256        return DeviceBudget {
257            backend: BudgetBackend::Cpu,
258            total_bytes: 0,
259            usable_bytes: 0,
260            reserve_fraction: 0.0,
261            source: "host RAM could not be probed on this platform".to_string(),
262            approximate: true,
263        };
264    }
265    DeviceBudget::new(
266        BudgetBackend::Cpu,
267        total,
268        CPU_RESERVE_FRACTION,
269        "total physical host RAM".to_string(),
270    )
271}
272
273pub(crate) fn human(bytes: u64) -> String {
274    const UNITS: [&str; 5] = ["B", "KiB", "MiB", "GiB", "TiB"];
275    let mut v = bytes as f64;
276    let mut u = 0;
277    while v >= 1024.0 && u < UNITS.len() - 1 {
278        v /= 1024.0;
279        u += 1;
280    }
281    format!("{v:.2} {}", UNITS[u])
282}
283
284#[cfg(test)]
285mod tests {
286    use super::*;
287
288    #[test]
289    fn reserve_is_applied_and_reported() {
290        let b = DeviceBudget::new(BudgetBackend::Cpu, 1000, 0.2, "test".into());
291        assert_eq!(b.total_bytes, 1000);
292        assert_eq!(b.usable_bytes, 800);
293        assert_eq!(b.reserve_fraction, 0.2);
294        assert!(!b.is_unknown());
295        // Always approximate: ferrox mmaps its weights.
296        assert!(b.approximate);
297    }
298
299    #[test]
300    fn a_nonsense_reserve_is_clamped_rather_than_producing_a_negative_budget() {
301        let over = DeviceBudget::new(BudgetBackend::Cpu, 1000, 5.0, "test".into());
302        assert_eq!(over.usable_bytes, 0);
303        let under = DeviceBudget::new(BudgetBackend::Cpu, 1000, -1.0, "test".into());
304        assert_eq!(under.usable_bytes, 1000);
305    }
306
307    #[test]
308    fn zero_total_reads_as_unknown_not_as_a_zero_ceiling() {
309        let b = DeviceBudget::new(BudgetBackend::Cpu, 0, 0.2, "nothing to probe".into());
310        assert!(b.is_unknown());
311        assert!(b.to_string().contains("no ceiling enforced"), "{b}");
312    }
313
314    /// Runs in both worlds, like the backend probes themselves: on a
315    /// host that can report RAM the budget must be plausible and
316    /// smaller than the total; on one that cannot it must be unknown.
317    #[test]
318    fn cpu_budget_is_either_unknown_or_a_plausible_fraction_of_real_ram() {
319        let b = DeviceBudget::detect(BudgetBackend::Cpu);
320        assert_eq!(b.backend, BudgetBackend::Cpu);
321        if b.is_unknown() {
322            assert_eq!(b.usable_bytes, 0);
323        } else {
324            assert!(b.total_bytes > 128 * 1024 * 1024);
325            assert!(b.usable_bytes < b.total_bytes);
326            assert!(b.usable_bytes > b.total_bytes / 2);
327            assert!(b.to_string().contains("host RAM"), "{b}");
328        }
329    }
330
331    /// Without `--features metal`/`cuda` these must still resolve (to
332    /// the host-RAM fallback) rather than failing to compile or
333    /// panicking -- the whole point of the honest-zero probe structs.
334    #[test]
335    fn accelerator_budgets_fall_back_to_host_ram_when_no_device_answers() {
336        for backend in [BudgetBackend::Metal, BudgetBackend::Cuda] {
337            let b = DeviceBudget::detect(backend);
338            assert_eq!(b.backend, backend);
339            if b.source.contains("fell back") {
340                assert!(b.source.contains("host RAM"), "{b}");
341            }
342        }
343    }
344
345    #[test]
346    fn human_bytes_are_readable_at_every_scale() {
347        assert_eq!(human(0), "0.00 B");
348        assert_eq!(human(1024), "1.00 KiB");
349        assert_eq!(human(3 * 1024 * 1024 * 1024), "3.00 GiB");
350    }
351}
ferrox_models/device_budget.rs

ferrox_models/
device_budget.rs
