cubecl_server/memory_management/drop_queue/
policy.rs

1use cubecl_common::bytes::Bytes;
2
3/// Defines the thresholds that determine when a [`PendingDropQueue`] should be
4/// flushed.
5///
6/// A flush is triggered when **either** limit is exceeded — whichever comes
7/// first. Set a field to `u32::MAX` / `usize::MAX` to effectively disable it.
8#[derive(Debug)]
9pub struct FlushingPolicy {
10    /// Flush when this many allocations have been staged.
11    pub max_bytes_count: u32,
12    /// Flush when the total staged size reaches this many bytes.
13    pub max_bytes_size: u32,
14}
15
16impl Default for FlushingPolicy {
17    fn default() -> Self {
18        Self {
19            max_bytes_count: 64,
20            max_bytes_size: 64 * 1024 * 1024, // 64 MiB
21        }
22    }
23}
24
25/// Tracks staged allocations and evaluates them against a [`FlushingPolicy`].
26#[derive(Default, Debug)]
27pub(crate) struct FlushingPolicyState {
28    bytes_count: u32,
29    bytes_size: u32,
30}
31
32impl FlushingPolicyState {
33    /// Record a newly staged [`Bytes`] allocation.
34    pub(crate) fn register(&mut self, bytes: &Bytes) {
35        self.bytes_count = self.bytes_count.saturating_add(1);
36        self.bytes_size = self
37            .bytes_size
38            .saturating_add(u32::try_from(bytes.len()).unwrap_or(u32::MAX));
39    }
40
41    /// Reset all counters, typically called after a flush.
42    pub(crate) fn reset(&mut self) {
43        self.bytes_count = 0;
44        self.bytes_size = 0;
45    }
46
47    /// Returns `true` if either threshold in `policy` has been reached.
48    pub(crate) fn should_flush(&self, policy: &FlushingPolicy) -> bool {
49        self.bytes_count >= policy.max_bytes_count || self.bytes_size >= policy.max_bytes_size
50    }
51}
52
53#[cfg(test)]
54mod policy_tests {
55    use std::vec;
56
57    use super::*;
58
59    fn policy() -> FlushingPolicy {
60        FlushingPolicy {
61            max_bytes_count: 4,
62            max_bytes_size: 100,
63        }
64    }
65
66    fn state() -> FlushingPolicyState {
67        FlushingPolicyState {
68            bytes_count: 0,
69            bytes_size: 0,
70        }
71    }
72
73    #[test]
74    fn no_flush_when_below_both_thresholds() {
75        let s = state();
76        assert!(!s.should_flush(&policy()));
77    }
78
79    #[test]
80    fn flush_when_count_threshold_reached() {
81        let mut s = state();
82        for _ in 0..4 {
83            s.register(&Bytes::from_elems(vec![0u8]));
84        }
85        assert!(s.should_flush(&policy()));
86    }
87
88    #[test]
89    fn flush_when_size_threshold_reached() {
90        let mut s = state();
91        s.register(&Bytes::from_elems(vec![0u8; 101]));
92        assert!(s.should_flush(&policy()));
93    }
94
95    #[test]
96    fn flush_triggered_by_whichever_limit_comes_first() {
97        let mut s = state();
98        // Only 2 allocations but already over the size limit.
99        s.register(&Bytes::from_elems(vec![0u8; 60]));
100        s.register(&Bytes::from_elems(vec![0u8; 60]));
101        assert!(s.should_flush(&policy()));
102    }
103
104    #[test]
105    fn register_saturates_instead_of_overflowing() {
106        // Regression test for #1359: staging allocations without an intervening
107        // reset must not panic with "attempt to add with overflow".
108        let mut s = FlushingPolicyState {
109            bytes_count: u32::MAX,
110            bytes_size: u32::MAX - 1,
111        };
112        s.register(&Bytes::from_elems(vec![0u8; 8]));
113        assert_eq!(s.bytes_count, u32::MAX);
114        assert_eq!(s.bytes_size, u32::MAX);
115        assert!(s.should_flush(&policy()));
116    }
117
118    #[test]
119    fn reset_clears_state() {
120        let mut s = state();
121        for _ in 0..4 {
122            s.register(&Bytes::from_elems(vec![0u8]));
123        }
124        assert!(s.should_flush(&policy()));
125        s.reset();
126        assert!(!s.should_flush(&policy()));
127    }
128}
cubecl_server/memory_management/drop_queue/policy.rs

cubecl_server/memory_management/drop_queue/
policy.rs
