ferrox_core/host_memory.rs
1//! How much memory the host actually has free, and whether a model fits.
2//!
3//! Lives here rather than in a binary because both `ferrox-cli` and
4//! `ferrox-server` need the same answer, and two copies of a
5//! platform probe drift.
6
7/// Bytes a new allocation can reasonably expect to get.
8///
9/// On macOS this counts free plus INACTIVE pages, because inactive
10/// pages are reclaimable and counting only free ones understates what
11/// is available by many gigabytes on a machine that has been up for a
12/// while. On Linux it is `MemAvailable`, which the kernel computes for
13/// exactly this question. `None` when the platform will not say, which
14/// callers must treat as "unknown", never as zero.
15pub fn available_bytes() -> Option<u64> {
16 #[cfg(target_os = "macos")]
17 {
18 let out = std::process::Command::new("vm_stat").output().ok()?;
19 if !out.status.success() {
20 return None;
21 }
22 parse_vm_stat(&String::from_utf8_lossy(&out.stdout))
23 }
24 #[cfg(target_os = "linux")]
25 {
26 let text = std::fs::read_to_string("/proc/meminfo").ok()?;
27 text.lines()
28 .find(|l| l.starts_with("MemAvailable:"))
29 .and_then(|l| l.split_whitespace().nth(1)?.parse::<u64>().ok())
30 .map(|kb| kb * 1024)
31 }
32 #[cfg(not(any(target_os = "macos", target_os = "linux")))]
33 {
34 None
35 }
36}
37
38/// Free plus inactive pages from `vm_stat` output, in bytes.
39pub fn parse_vm_stat(text: &str) -> Option<u64> {
40 let page = text
41 .lines()
42 .next()?
43 .split("page size of ")
44 .nth(1)?
45 .split(' ')
46 .next()?
47 .parse::<u64>()
48 .ok()?;
49 let field = |name: &str| -> Option<u64> {
50 text.lines()
51 .find(|l| l.starts_with(name))
52 .and_then(|l| l.split(':').nth(1))
53 .and_then(|v| v.trim().trim_end_matches('.').parse::<u64>().ok())
54 };
55 Some((field("Pages free")? + field("Pages inactive")?) * page)
56}
57
58/// What to do about a model whose weights may not fit.
59#[derive(Debug, Clone, Copy, PartialEq, Eq)]
60pub enum FitPlan {
61 /// Load everything resident. The fast path, and the default.
62 Resident,
63 /// Stream experts, keeping at most this many bytes of them cached.
64 Stream { cache_bytes: u64 },
65}
66
67/// Decide whether streaming is needed, given what the model weighs and
68/// what the host has.
69///
70/// The bar is deliberately not "does it fit exactly". A model that
71/// fills memory to the last byte will thrash the page cache and leave
72/// nothing for the KV cache, so `headroom_bytes` is subtracted first.
73///
74/// `available` of `None` means the platform would not say. That
75/// resolves to `Resident`, NOT to streaming: guessing that a machine is
76/// short on memory would silently put every user on the slow path on
77/// any platform without a probe.
78pub fn plan_for(
79 weight_bytes: u64,
80 available: Option<u64>,
81 headroom_bytes: u64,
82 min_cache_bytes: u64,
83) -> FitPlan {
84 let Some(available) = available else {
85 return FitPlan::Resident;
86 };
87 let usable = available.saturating_sub(headroom_bytes);
88 if weight_bytes <= usable {
89 return FitPlan::Resident;
90 }
91 // It does not fit. Spend what is usable on the expert cache, but
92 // never less than the floor: a cache too small to hold one decode
93 // step's experts turns every acquire into a fresh read, which is
94 // correct but pathologically slow.
95 FitPlan::Stream {
96 cache_bytes: usable.max(min_cache_bytes),
97 }
98}
99
100#[cfg(test)]
101mod tests {
102 use super::*;
103
104 const GB: u64 = 1024 * 1024 * 1024;
105
106 #[test]
107 fn a_model_that_fits_stays_resident() {
108 assert_eq!(
109 plan_for(8 * GB, Some(32 * GB), 4 * GB, GB),
110 FitPlan::Resident
111 );
112 }
113
114 /// Filling memory to the last byte leaves nothing for the KV cache,
115 /// so the headroom is subtracted before the comparison.
116 #[test]
117 fn headroom_is_subtracted_before_deciding() {
118 assert_eq!(
119 plan_for(30 * GB, Some(32 * GB), 4 * GB, GB),
120 FitPlan::Stream {
121 cache_bytes: 28 * GB
122 },
123 "30 GiB of weights into 32 GiB with 4 GiB reserved does not fit"
124 );
125 }
126
127 #[test]
128 fn a_model_far_larger_than_ram_streams() {
129 assert!(matches!(
130 plan_for(155 * GB, Some(32 * GB), 4 * GB, GB),
131 FitPlan::Stream { .. }
132 ));
133 }
134
135 /// An unknown probe must not be read as "no memory". Guessing that
136 /// would put every user on a platform without a probe onto the slow
137 /// path silently.
138 #[test]
139 fn an_unknown_amount_of_memory_never_forces_streaming() {
140 assert_eq!(plan_for(155 * GB, None, 4 * GB, GB), FitPlan::Resident);
141 }
142
143 /// A cache too small to hold one decode step's experts makes every
144 /// acquire a fresh read: correct, and pathologically slow.
145 #[test]
146 fn the_cache_never_falls_below_the_floor() {
147 assert_eq!(
148 plan_for(100 * GB, Some(5 * GB), 4 * GB, 2 * GB),
149 FitPlan::Stream {
150 cache_bytes: 2 * GB
151 },
152 "1 GiB usable must be raised to the 2 GiB floor"
153 );
154 }
155
156 #[test]
157 fn vm_stat_counts_reclaimable_pages_not_just_free_ones() {
158 let sample = "Mach Virtual Memory Statistics: (page size of 16384 bytes)\n\
159 Pages free: 100000.\n\
160 Pages active: 900000.\n\
161 Pages inactive: 50000.\n";
162 // (100000 + 50000) * 16384
163 assert_eq!(parse_vm_stat(sample), Some(150_000 * 16_384));
164 }
165}