oxicuda_memory/
pool_pressure.rs

1//! Memory pressure monitoring with threshold callbacks and eviction hooks.
2//!
3//! A long-running allocator benefits from reacting *before* the device runs out
4//! of memory: trimming idle pool pages, evicting cached buffers, or throttling
5//! new work as free VRAM drops.  This module models that control loop on the
6//! host side.
7//!
8//! [`MemoryPressureMonitor`] holds configurable warning / critical thresholds
9//! (expressed as a *used fraction* of total device memory) and, on each sample,
10//! classifies the current state into a [`PressureLevel`].  Transitions into a
11//! more severe level fire user callbacks; entering the critical level fires an
12//! eviction hook that returns the number of bytes reclaimed so the loop can
13//! account for the relief.
14//!
15//! Samples are supplied as [`crate::memory_info::MemoryInfo`] values.  In
16//! production the [`MemoryPressureMonitor::poll`] method fetches a fresh sample
17//! via [`crate::memory_info::memory_info`] (requires a GPU); in tests the
18//! [`MemoryPressureMonitor::observe`] method feeds synthetic samples so the
19//! whole state machine is deterministic and hardware-free.
20//!
21//! # Example
22//!
23//! ```rust
24//! # use oxicuda_memory::pool_pressure::*;
25//! # use oxicuda_memory::memory_info::MemoryInfo;
26//! let mut monitor = MemoryPressureMonitor::new(0.80, 0.95).expect("thresholds");
27//!
28//! // 70% used -> nominal.
29//! let lvl = monitor.observe(MemoryInfo { free: 30, total: 100 });
30//! assert_eq!(lvl, PressureLevel::Nominal);
31//!
32//! // 96% used -> critical.
33//! let lvl = monitor.observe(MemoryInfo { free: 4, total: 100 });
34//! assert_eq!(lvl, PressureLevel::Critical);
35//! ```
36
37use crate::memory_info::MemoryInfo;
38use oxicuda_driver::error::{CudaError, CudaResult};
39
40// ---------------------------------------------------------------------------
41// PressureLevel
42// ---------------------------------------------------------------------------
43
44/// The classified memory-pressure state.
45///
46/// Ordered by severity so levels can be compared directly
47/// (`Nominal < Warning < Critical`).
48#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
49pub enum PressureLevel {
50    /// Used fraction below the warning threshold.
51    Nominal,
52    /// Used fraction at/above the warning threshold but below critical.
53    Warning,
54    /// Used fraction at/above the critical threshold.
55    Critical,
56}
57
58impl PressureLevel {
59    /// Returns `true` if this is the [`Critical`](PressureLevel::Critical) level.
60    #[inline]
61    #[must_use]
62    pub fn is_critical(self) -> bool {
63        matches!(self, Self::Critical)
64    }
65
66    /// Returns `true` if this level is at least [`Warning`](PressureLevel::Warning).
67    #[inline]
68    #[must_use]
69    pub fn is_elevated(self) -> bool {
70        self >= Self::Warning
71    }
72}
73
74impl std::fmt::Display for PressureLevel {
75    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
76        match self {
77            Self::Nominal => write!(f, "nominal"),
78            Self::Warning => write!(f, "warning"),
79            Self::Critical => write!(f, "critical"),
80        }
81    }
82}
83
84// ---------------------------------------------------------------------------
85// PressureSample
86// ---------------------------------------------------------------------------
87
88/// The outcome of classifying one memory sample.
89#[derive(Debug, Clone, Copy, PartialEq)]
90pub struct PressureSample {
91    /// The classified level.
92    pub level: PressureLevel,
93    /// The used fraction (0.0–1.0) that produced the classification.
94    pub used_fraction: f64,
95    /// The previous level before this sample (for transition detection).
96    pub previous: PressureLevel,
97}
98
99impl PressureSample {
100    /// Returns `true` if this sample represents a transition into a strictly
101    /// more severe level than the previous sample.
102    #[inline]
103    #[must_use]
104    pub fn escalated(&self) -> bool {
105        self.level > self.previous
106    }
107
108    /// Returns `true` if this sample represents a transition into a strictly
109    /// less severe level than the previous sample.
110    #[inline]
111    #[must_use]
112    pub fn de_escalated(&self) -> bool {
113        self.level < self.previous
114    }
115}
116
117// ---------------------------------------------------------------------------
118// MemoryPressureMonitor
119// ---------------------------------------------------------------------------
120
121/// Type alias for the eviction hook: called on entering the critical level,
122/// returns the number of bytes reclaimed.
123type EvictionHook = Box<dyn FnMut() -> usize + Send>;
124
125/// Type alias for a level-transition callback.
126type TransitionHook = Box<dyn FnMut(PressureSample) + Send>;
127
128/// Monitors device memory pressure and drives reactive trim / eviction.
129///
130/// Construct with [`new`](Self::new) supplying the warning and critical used
131/// fractions.  Optionally attach an eviction hook with
132/// [`on_eviction`](Self::on_eviction) and a transition callback with
133/// [`on_transition`](Self::on_transition).  Feed samples via
134/// [`observe`](Self::observe) (synthetic) or [`poll`](Self::poll) (live GPU).
135pub struct MemoryPressureMonitor {
136    warning_fraction: f64,
137    critical_fraction: f64,
138    current: PressureLevel,
139    /// Total bytes the eviction hook has reclaimed over the monitor's lifetime.
140    total_evicted: usize,
141    /// Number of times the eviction hook has fired.
142    eviction_count: u64,
143    eviction_hook: Option<EvictionHook>,
144    transition_hook: Option<TransitionHook>,
145}
146
147impl std::fmt::Debug for MemoryPressureMonitor {
148    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
149        f.debug_struct("MemoryPressureMonitor")
150            .field("warning_fraction", &self.warning_fraction)
151            .field("critical_fraction", &self.critical_fraction)
152            .field("current", &self.current)
153            .field("total_evicted", &self.total_evicted)
154            .field("eviction_count", &self.eviction_count)
155            .field("has_eviction_hook", &self.eviction_hook.is_some())
156            .field("has_transition_hook", &self.transition_hook.is_some())
157            .finish()
158    }
159}
160
161impl MemoryPressureMonitor {
162    /// Creates a monitor with the given warning and critical used-fraction
163    /// thresholds.
164    ///
165    /// Both fractions must lie in `(0.0, 1.0]` and `warning_fraction` must be
166    /// strictly less than `critical_fraction`.
167    ///
168    /// # Errors
169    ///
170    /// * [`CudaError::InvalidValue`] if either threshold is non-finite, outside
171    ///   `(0.0, 1.0]`, or if `warning_fraction >= critical_fraction`.
172    pub fn new(warning_fraction: f64, critical_fraction: f64) -> CudaResult<Self> {
173        let valid = |x: f64| x.is_finite() && x > 0.0 && x <= 1.0;
174        if !valid(warning_fraction)
175            || !valid(critical_fraction)
176            || warning_fraction >= critical_fraction
177        {
178            return Err(CudaError::InvalidValue);
179        }
180        Ok(Self {
181            warning_fraction,
182            critical_fraction,
183            current: PressureLevel::Nominal,
184            total_evicted: 0,
185            eviction_count: 0,
186            eviction_hook: None,
187            transition_hook: None,
188        })
189    }
190
191    /// Attaches an eviction hook fired whenever the monitor *enters* the
192    /// critical level.  The hook returns the number of bytes it reclaimed,
193    /// which is accumulated into [`total_evicted`](Self::total_evicted).
194    ///
195    /// Builder-style; returns `self` for chaining.
196    #[must_use]
197    pub fn on_eviction<F>(mut self, hook: F) -> Self
198    where
199        F: FnMut() -> usize + Send + 'static,
200    {
201        self.eviction_hook = Some(Box::new(hook));
202        self
203    }
204
205    /// Attaches a callback fired on every level transition (escalation or
206    /// de-escalation).
207    ///
208    /// Builder-style; returns `self` for chaining.
209    #[must_use]
210    pub fn on_transition<F>(mut self, hook: F) -> Self
211    where
212        F: FnMut(PressureSample) + Send + 'static,
213    {
214        self.transition_hook = Some(Box::new(hook));
215        self
216    }
217
218    /// The warning used-fraction threshold.
219    #[inline]
220    #[must_use]
221    pub fn warning_fraction(&self) -> f64 {
222        self.warning_fraction
223    }
224
225    /// The critical used-fraction threshold.
226    #[inline]
227    #[must_use]
228    pub fn critical_fraction(&self) -> f64 {
229        self.critical_fraction
230    }
231
232    /// The current pressure level (state from the most recent sample).
233    #[inline]
234    #[must_use]
235    pub fn level(&self) -> PressureLevel {
236        self.current
237    }
238
239    /// Total bytes reclaimed by the eviction hook over this monitor's lifetime.
240    #[inline]
241    #[must_use]
242    pub fn total_evicted(&self) -> usize {
243        self.total_evicted
244    }
245
246    /// Number of times the eviction hook has fired.
247    #[inline]
248    #[must_use]
249    pub fn eviction_count(&self) -> u64 {
250        self.eviction_count
251    }
252
253    /// Classifies a `used_fraction` (0.0–1.0) into a [`PressureLevel`] using
254    /// this monitor's thresholds.
255    ///
256    /// The comparison is inclusive at the threshold (a fraction exactly equal
257    /// to the critical threshold is [`Critical`](PressureLevel::Critical)).
258    #[must_use]
259    pub fn classify(&self, used_fraction: f64) -> PressureLevel {
260        if used_fraction >= self.critical_fraction {
261            PressureLevel::Critical
262        } else if used_fraction >= self.warning_fraction {
263            PressureLevel::Warning
264        } else {
265            PressureLevel::Nominal
266        }
267    }
268
269    /// Feeds a synthetic memory sample through the monitor and returns the new
270    /// level.
271    ///
272    /// On a transition the transition callback (if any) fires.  On *entering*
273    /// the critical level the eviction hook (if any) fires and its return value
274    /// is accumulated.  Re-observing the same level does not re-fire hooks.
275    pub fn observe(&mut self, info: MemoryInfo) -> PressureLevel {
276        let used_fraction = info.usage_fraction();
277        let previous = self.current;
278        let level = self.classify(used_fraction);
279        self.current = level;
280
281        if level != previous {
282            let sample = PressureSample {
283                level,
284                used_fraction,
285                previous,
286            };
287            if let Some(hook) = self.transition_hook.as_mut() {
288                hook(sample);
289            }
290            // Fire eviction only on an *upward* transition into Critical.
291            if level == PressureLevel::Critical && previous != PressureLevel::Critical {
292                if let Some(hook) = self.eviction_hook.as_mut() {
293                    let reclaimed = hook();
294                    self.total_evicted = self.total_evicted.saturating_add(reclaimed);
295                    self.eviction_count = self.eviction_count.saturating_add(1);
296                }
297            }
298        }
299        level
300    }
301
302    /// Polls live device memory via [`crate::memory_info::memory_info`] and
303    /// feeds the sample through [`observe`](Self::observe).
304    ///
305    /// # Errors
306    ///
307    /// Forwards any error from the memory-info query (e.g. no current context).
308    pub fn poll(&mut self) -> CudaResult<PressureLevel> {
309        let info = crate::memory_info::memory_info()?;
310        Ok(self.observe(info))
311    }
312
313    /// Resets the monitor's level to [`Nominal`](PressureLevel::Nominal) and
314    /// clears eviction accounting, without changing thresholds or hooks.
315    pub fn reset(&mut self) {
316        self.current = PressureLevel::Nominal;
317        self.total_evicted = 0;
318        self.eviction_count = 0;
319    }
320}
321
322// ---------------------------------------------------------------------------
323// Validation helper
324// ---------------------------------------------------------------------------
325
326/// Validates a pair of `(warning, critical)` used-fraction thresholds without
327/// constructing a monitor.
328///
329/// # Errors
330///
331/// * [`CudaError::InvalidValue`] under the same conditions as
332///   [`MemoryPressureMonitor::new`].
333pub fn validate_thresholds(warning_fraction: f64, critical_fraction: f64) -> CudaResult<()> {
334    MemoryPressureMonitor::new(warning_fraction, critical_fraction).map(|_| ())
335}
336
337// ---------------------------------------------------------------------------
338// Tests
339// ---------------------------------------------------------------------------
340
341#[cfg(test)]
342mod tests {
343    use super::*;
344    use std::sync::Arc;
345    use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering};
346
347    fn info(free: usize, total: usize) -> MemoryInfo {
348        MemoryInfo { free, total }
349    }
350
351    #[test]
352    fn pressure_level_ordering() {
353        assert!(PressureLevel::Nominal < PressureLevel::Warning);
354        assert!(PressureLevel::Warning < PressureLevel::Critical);
355        assert!(PressureLevel::Critical.is_critical());
356        assert!(!PressureLevel::Warning.is_critical());
357        assert!(PressureLevel::Warning.is_elevated());
358        assert!(!PressureLevel::Nominal.is_elevated());
359        assert!(PressureLevel::Critical.is_elevated());
360    }
361
362    #[test]
363    fn pressure_level_display() {
364        assert_eq!(format!("{}", PressureLevel::Nominal), "nominal");
365        assert_eq!(format!("{}", PressureLevel::Warning), "warning");
366        assert_eq!(format!("{}", PressureLevel::Critical), "critical");
367    }
368
369    #[test]
370    fn new_rejects_bad_thresholds() {
371        // warning >= critical
372        assert!(MemoryPressureMonitor::new(0.9, 0.9).is_err());
373        assert!(MemoryPressureMonitor::new(0.95, 0.8).is_err());
374        // out of range
375        assert!(MemoryPressureMonitor::new(0.0, 0.9).is_err());
376        assert!(MemoryPressureMonitor::new(0.8, 1.5).is_err());
377        // non-finite
378        assert!(MemoryPressureMonitor::new(f64::NAN, 0.9).is_err());
379    }
380
381    #[test]
382    fn new_accepts_valid_thresholds() {
383        let m = MemoryPressureMonitor::new(0.8, 0.95).expect("ok");
384        assert!((m.warning_fraction() - 0.8).abs() < 1e-12);
385        assert!((m.critical_fraction() - 0.95).abs() < 1e-12);
386        assert_eq!(m.level(), PressureLevel::Nominal);
387    }
388
389    #[test]
390    fn validate_thresholds_helper() {
391        assert!(validate_thresholds(0.7, 0.9).is_ok());
392        assert!(validate_thresholds(0.9, 0.7).is_err());
393    }
394
395    #[test]
396    fn classify_boundaries() {
397        let m = MemoryPressureMonitor::new(0.80, 0.95).expect("ok");
398        // Below warning.
399        assert_eq!(m.classify(0.50), PressureLevel::Nominal);
400        assert_eq!(m.classify(0.799), PressureLevel::Nominal);
401        // Exactly warning -> Warning (inclusive).
402        assert_eq!(m.classify(0.80), PressureLevel::Warning);
403        assert_eq!(m.classify(0.90), PressureLevel::Warning);
404        // Just below critical.
405        assert_eq!(m.classify(0.949), PressureLevel::Warning);
406        // Exactly critical -> Critical (inclusive).
407        assert_eq!(m.classify(0.95), PressureLevel::Critical);
408        assert_eq!(m.classify(1.00), PressureLevel::Critical);
409    }
410
411    #[test]
412    fn observe_transitions_through_levels() {
413        let mut m = MemoryPressureMonitor::new(0.80, 0.95).expect("ok");
414        // 50% used (free 50 / total 100) -> nominal.
415        assert_eq!(m.observe(info(50, 100)), PressureLevel::Nominal);
416        // 85% used -> warning.
417        assert_eq!(m.observe(info(15, 100)), PressureLevel::Warning);
418        // 97% used -> critical.
419        assert_eq!(m.observe(info(3, 100)), PressureLevel::Critical);
420        assert_eq!(m.level(), PressureLevel::Critical);
421        // back down to 40% -> nominal.
422        assert_eq!(m.observe(info(60, 100)), PressureLevel::Nominal);
423    }
424
425    #[test]
426    fn eviction_hook_fires_on_entering_critical() {
427        let evicted = Arc::new(AtomicUsize::new(0));
428        let calls = Arc::new(AtomicU64::new(0));
429        let e2 = Arc::clone(&evicted);
430        let c2 = Arc::clone(&calls);
431        let mut m = MemoryPressureMonitor::new(0.80, 0.95)
432            .expect("ok")
433            .on_eviction(move || {
434                e2.fetch_add(1, Ordering::SeqCst);
435                c2.fetch_add(1, Ordering::SeqCst);
436                4096 // bytes reclaimed
437            });
438
439        // Nominal -> no eviction.
440        m.observe(info(50, 100));
441        assert_eq!(calls.load(Ordering::SeqCst), 0);
442
443        // Enter critical -> eviction fires once.
444        m.observe(info(2, 100));
445        assert_eq!(calls.load(Ordering::SeqCst), 1);
446        assert_eq!(m.total_evicted(), 4096);
447        assert_eq!(m.eviction_count(), 1);
448
449        // Stay critical -> no re-fire.
450        m.observe(info(1, 100));
451        assert_eq!(calls.load(Ordering::SeqCst), 1);
452        assert_eq!(m.total_evicted(), 4096);
453
454        // Drop to nominal, then re-enter critical -> fires again.
455        m.observe(info(80, 100));
456        m.observe(info(1, 100));
457        assert_eq!(calls.load(Ordering::SeqCst), 2);
458        assert_eq!(m.total_evicted(), 8192);
459        assert_eq!(m.eviction_count(), 2);
460    }
461
462    #[test]
463    fn eviction_does_not_fire_when_warning_to_critical_skips_nominal() {
464        // Going Warning -> Critical (not via Nominal) should still fire once.
465        let calls = Arc::new(AtomicU64::new(0));
466        let c2 = Arc::clone(&calls);
467        let mut m = MemoryPressureMonitor::new(0.80, 0.95)
468            .expect("ok")
469            .on_eviction(move || {
470                c2.fetch_add(1, Ordering::SeqCst);
471                100
472            });
473        m.observe(info(15, 100)); // warning
474        assert_eq!(calls.load(Ordering::SeqCst), 0);
475        m.observe(info(2, 100)); // critical
476        assert_eq!(calls.load(Ordering::SeqCst), 1);
477    }
478
479    #[test]
480    fn transition_hook_records_escalation_and_de_escalation() {
481        let escalations = Arc::new(AtomicU64::new(0));
482        let de_escalations = Arc::new(AtomicU64::new(0));
483        let up = Arc::clone(&escalations);
484        let down = Arc::clone(&de_escalations);
485        let mut m = MemoryPressureMonitor::new(0.80, 0.95)
486            .expect("ok")
487            .on_transition(move |sample: PressureSample| {
488                if sample.escalated() {
489                    up.fetch_add(1, Ordering::SeqCst);
490                }
491                if sample.de_escalated() {
492                    down.fetch_add(1, Ordering::SeqCst);
493                }
494            });
495
496        m.observe(info(50, 100)); // nominal, no transition (already nominal)
497        assert_eq!(escalations.load(Ordering::SeqCst), 0);
498        m.observe(info(15, 100)); // -> warning (escalate)
499        m.observe(info(2, 100)); //  -> critical (escalate)
500        assert_eq!(escalations.load(Ordering::SeqCst), 2);
501        m.observe(info(50, 100)); // -> nominal (de-escalate)
502        assert_eq!(de_escalations.load(Ordering::SeqCst), 1);
503    }
504
505    #[test]
506    fn observe_same_level_does_not_transition() {
507        let calls = Arc::new(AtomicU64::new(0));
508        let c2 = Arc::clone(&calls);
509        let mut m = MemoryPressureMonitor::new(0.80, 0.95)
510            .expect("ok")
511            .on_transition(move |_| {
512                c2.fetch_add(1, Ordering::SeqCst);
513            });
514        m.observe(info(50, 100)); // nominal == initial, no transition
515        m.observe(info(60, 100)); // still nominal, no transition
516        assert_eq!(calls.load(Ordering::SeqCst), 0);
517    }
518
519    #[test]
520    fn pressure_sample_transition_flags() {
521        let escalate = PressureSample {
522            level: PressureLevel::Critical,
523            used_fraction: 0.97,
524            previous: PressureLevel::Warning,
525        };
526        assert!(escalate.escalated());
527        assert!(!escalate.de_escalated());
528
529        let de = PressureSample {
530            level: PressureLevel::Nominal,
531            used_fraction: 0.20,
532            previous: PressureLevel::Critical,
533        };
534        assert!(de.de_escalated());
535        assert!(!de.escalated());
536    }
537
538    #[test]
539    fn reset_clears_state_keeps_thresholds() {
540        let mut m = MemoryPressureMonitor::new(0.80, 0.95)
541            .expect("ok")
542            .on_eviction(|| 1000);
543        m.observe(info(1, 100)); // critical -> evict
544        assert_eq!(m.total_evicted(), 1000);
545        assert_eq!(m.level(), PressureLevel::Critical);
546        m.reset();
547        assert_eq!(m.level(), PressureLevel::Nominal);
548        assert_eq!(m.total_evicted(), 0);
549        assert_eq!(m.eviction_count(), 0);
550        // thresholds intact
551        assert!((m.warning_fraction() - 0.80).abs() < 1e-12);
552    }
553
554    #[test]
555    fn observe_zero_total_is_nominal() {
556        // usage_fraction returns 0.0 when total is 0.
557        let mut m = MemoryPressureMonitor::new(0.80, 0.95).expect("ok");
558        assert_eq!(m.observe(info(0, 0)), PressureLevel::Nominal);
559    }
560
561    #[test]
562    fn poll_signature_compiles() {
563        let _: fn(&mut MemoryPressureMonitor) -> CudaResult<PressureLevel> =
564            MemoryPressureMonitor::poll;
565    }
566
567    #[test]
568    fn debug_does_not_panic() {
569        let m = MemoryPressureMonitor::new(0.8, 0.95)
570            .expect("ok")
571            .on_eviction(|| 0);
572        let s = format!("{m:?}");
573        assert!(s.contains("MemoryPressureMonitor"));
574        assert!(s.contains("has_eviction_hook"));
575    }
576}
oxicuda_memory/pool_pressure.rs

oxicuda_memory/
pool_pressure.rs
