Skip to main content

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}