Skip to main content

kernel/profiles/
fit.rs

1//! Whether a model fits in a machine's memory: a coarse runs-well / tight-fit /
2//! too-large verdict from the model's footprint and the total RAM.
3
4/// How well a model is expected to run given available memory.
5#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
6pub enum FitVerdict {
7    /// Comfortable — well under the runs-well fraction of memory.
8    RunsWell,
9    /// Fits, but with little headroom.
10    TightFit,
11    /// Won't fit comfortably.
12    TooLarge,
13}
14
15/// A fit verdict plus the memory the model is estimated to require.
16#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
17pub struct FitAssessment {
18    /// The verdict.
19    pub verdict: FitVerdict,
20    /// The estimated required bytes (footprint × overhead).
21    pub required_bytes: i64,
22}
23
24impl FitVerdict {
25    /// The stable string form: `runs_well`, `tight_fit`, or `too_large`.
26    pub fn as_str(&self) -> &'static str {
27        match self {
28            FitVerdict::RunsWell => "runs_well",
29            FitVerdict::TightFit => "tight_fit",
30            FitVerdict::TooLarge => "too_large",
31        }
32    }
33
34    /// Weights need working memory beyond the raw footprint.
35    const MEMORY_OVERHEAD_FACTOR: f64 = 1.25;
36    /// Below this share of memory a model runs well.
37    const RUNS_WELL_FRACTION: f64 = 0.75;
38    /// Below this share it's a tight fit; at or above it's too large.
39    const TIGHT_FIT_FRACTION: f64 = 0.95;
40
41    /// Assess a `footprint_mb` model against `total_memory_bytes`. Returns `None`
42    /// when the footprint is unknown/non-positive or the memory total is zero.
43    pub fn assess(footprint_mb: Option<i64>, total_memory_bytes: u64) -> Option<FitAssessment> {
44        let footprint_mb = footprint_mb.filter(|&mb| mb > 0)?;
45        if total_memory_bytes == 0 {
46            return None;
47        }
48        let required_bytes =
49            (footprint_mb as f64 * (1i64 << 20) as f64 * Self::MEMORY_OVERHEAD_FACTOR) as i64;
50        let share = required_bytes as f64 / total_memory_bytes as f64;
51        let verdict = if share < Self::RUNS_WELL_FRACTION {
52            FitVerdict::RunsWell
53        } else if share < Self::TIGHT_FIT_FRACTION {
54            FitVerdict::TightFit
55        } else {
56            FitVerdict::TooLarge
57        };
58        Some(FitAssessment {
59            verdict,
60            required_bytes,
61        })
62    }
63}
64
65#[cfg(test)]
66mod tests {
67    use super::*;
68
69    const GIB: u64 = 1 << 30;
70
71    #[test]
72    fn an_unknown_or_empty_footprint_is_unassessable() {
73        assert!(FitVerdict::assess(None, 16 * GIB).is_none());
74        assert!(FitVerdict::assess(Some(0), 16 * GIB).is_none());
75        assert!(FitVerdict::assess(Some(-5), 16 * GIB).is_none());
76        assert!(FitVerdict::assess(Some(1000), 0).is_none());
77    }
78
79    #[test]
80    fn the_verdict_tracks_the_memory_share() {
81        // ~1 GiB footprint × 1.25 = ~1.25 GiB of 16 GiB → well under 0.75 → runs well.
82        let assessment = FitVerdict::assess(Some(1024), 16 * GIB).unwrap();
83        assert_eq!(assessment.verdict, FitVerdict::RunsWell);
84
85        // 12 GiB × 1.25 = 15 GiB of 16 GiB → share 0.9375 → tight fit.
86        let tight = FitVerdict::assess(Some(12 * 1024), 16 * GIB).unwrap();
87        assert_eq!(tight.verdict, FitVerdict::TightFit);
88
89        // 16 GiB × 1.25 = 20 GiB of 16 GiB → share 1.25 → too large.
90        let too_large = FitVerdict::assess(Some(16 * 1024), 16 * GIB).unwrap();
91        assert_eq!(too_large.verdict, FitVerdict::TooLarge);
92    }
93
94    #[test]
95    fn each_verdict_has_a_stable_slug() {
96        assert_eq!(FitVerdict::RunsWell.as_str(), "runs_well");
97        assert_eq!(FitVerdict::TightFit.as_str(), "tight_fit");
98        assert_eq!(FitVerdict::TooLarge.as_str(), "too_large");
99    }
100
101    #[test]
102    fn required_bytes_includes_the_overhead_factor() {
103        let assessment = FitVerdict::assess(Some(1024), 64 * GIB).unwrap();
104        // 1024 MiB × 1.25 = 1280 MiB.
105        assert_eq!(assessment.required_bytes, 1280 * (1 << 20));
106    }
107}