entrenar/monitor/gpu/
buffer.rs1use std::collections::VecDeque;
4
5use super::GpuMetrics;
6
7#[derive(Debug)]
9pub struct GpuMetricsBuffer {
10 capacity: usize,
12 buffers: Vec<VecDeque<GpuMetrics>>,
14}
15
16impl GpuMetricsBuffer {
17 pub fn new(capacity: usize, num_devices: usize) -> Self {
19 let buffers = (0..num_devices).map(|_| VecDeque::with_capacity(capacity)).collect();
20 Self { capacity, buffers }
21 }
22
23 pub fn push(&mut self, metrics: &[GpuMetrics]) {
25 for m in metrics {
26 let device_idx = m.device_id as usize;
27 if device_idx >= self.buffers.len() {
28 self.buffers.resize_with(device_idx + 1, || VecDeque::with_capacity(self.capacity));
29 }
30
31 let buffer = &mut self.buffers[device_idx];
32 if buffer.len() >= self.capacity {
33 buffer.pop_front();
34 }
35 buffer.push_back(m.clone());
36 }
37 }
38
39 pub fn last_n(&self, device_id: u32, n: usize) -> Vec<&GpuMetrics> {
41 let device_idx = device_id as usize;
42 if device_idx >= self.buffers.len() {
43 return Vec::new();
44 }
45
46 self.buffers[device_idx].iter().rev().take(n).rev().collect()
47 }
48
49 pub fn utilization_history(&self, device_id: u32) -> Vec<u32> {
51 let device_idx = device_id as usize;
52 if device_idx >= self.buffers.len() {
53 return Vec::new();
54 }
55
56 self.buffers[device_idx].iter().map(|m| m.utilization_percent).collect()
57 }
58
59 pub fn temperature_history(&self, device_id: u32) -> Vec<u32> {
61 let device_idx = device_id as usize;
62 if device_idx >= self.buffers.len() {
63 return Vec::new();
64 }
65
66 self.buffers[device_idx].iter().map(|m| m.temperature_celsius).collect()
67 }
68
69 pub fn memory_history(&self, device_id: u32) -> Vec<f64> {
71 let device_idx = device_id as usize;
72 if device_idx >= self.buffers.len() {
73 return Vec::new();
74 }
75
76 self.buffers[device_idx].iter().map(GpuMetrics::memory_percent).collect()
77 }
78
79 pub fn len(&self, device_id: u32) -> usize {
81 let device_idx = device_id as usize;
82 if device_idx >= self.buffers.len() {
83 return 0;
84 }
85 self.buffers[device_idx].len()
86 }
87
88 pub fn is_empty(&self, device_id: u32) -> bool {
90 self.len(device_id) == 0
91 }
92}
93
94#[cfg(test)]
95mod tests {
96 use super::*;
97
98 #[test]
99 fn test_buffer_push_and_last_n() {
100 let mut buffer = GpuMetricsBuffer::new(10, 1);
101
102 for i in 0..5 {
103 buffer.push(&[GpuMetrics {
104 device_id: 0,
105 utilization_percent: i * 10,
106 ..Default::default()
107 }]);
108 }
109
110 let last3 = buffer.last_n(0, 3);
111 assert_eq!(last3.len(), 3);
112 assert_eq!(last3[0].utilization_percent, 20);
113 assert_eq!(last3[1].utilization_percent, 30);
114 assert_eq!(last3[2].utilization_percent, 40);
115 }
116
117 #[test]
118 fn test_buffer_capacity_limit() {
119 let mut buffer = GpuMetricsBuffer::new(5, 1);
120
121 for i in 0..10 {
122 buffer.push(&[GpuMetrics {
123 device_id: 0,
124 utilization_percent: i,
125 ..Default::default()
126 }]);
127 }
128
129 assert_eq!(buffer.len(0), 5);
130
131 let history = buffer.utilization_history(0);
132 assert_eq!(history, vec![5, 6, 7, 8, 9]);
133 }
134
135 #[test]
136 fn test_buffer_utilization_history() {
137 let mut buffer = GpuMetricsBuffer::new(10, 1);
138
139 for i in 0..5 {
140 buffer.push(&[GpuMetrics {
141 device_id: 0,
142 utilization_percent: i * 20,
143 ..Default::default()
144 }]);
145 }
146
147 let history = buffer.utilization_history(0);
148 assert_eq!(history, vec![0, 20, 40, 60, 80]);
149 }
150
151 #[test]
152 fn test_buffer_temperature_history() {
153 let mut buffer = GpuMetricsBuffer::new(10, 1);
154
155 for i in 0..3 {
156 buffer.push(&[GpuMetrics {
157 device_id: 0,
158 temperature_celsius: 60 + i * 5,
159 ..Default::default()
160 }]);
161 }
162
163 let history = buffer.temperature_history(0);
164 assert_eq!(history, vec![60, 65, 70]);
165 }
166
167 #[test]
168 fn test_buffer_multiple_devices() {
169 let mut buffer = GpuMetricsBuffer::new(10, 2);
170
171 buffer.push(&[
172 GpuMetrics { device_id: 0, utilization_percent: 50, ..Default::default() },
173 GpuMetrics { device_id: 1, utilization_percent: 75, ..Default::default() },
174 ]);
175
176 assert_eq!(buffer.utilization_history(0), vec![50]);
177 assert_eq!(buffer.utilization_history(1), vec![75]);
178 }
179
180 #[test]
181 fn test_buffer_empty_device() {
182 let buffer = GpuMetricsBuffer::new(10, 1);
183 assert!(buffer.is_empty(0));
184 assert!(buffer.utilization_history(5).is_empty()); }
186
187 #[test]
188 fn test_buffer_memory_history() {
189 let mut buffer = GpuMetricsBuffer::new(10, 1);
190
191 for i in 0..3 {
192 buffer.push(&[GpuMetrics {
193 device_id: 0,
194 memory_used_mb: i * 1000,
195 memory_total_mb: 8000,
196 memory_utilization_percent: (i as u32) * 10,
197 ..Default::default()
198 }]);
199 }
200
201 let history = buffer.memory_history(0);
202 assert_eq!(history.len(), 3);
204 assert!((history[0] - 0.0).abs() < 0.1);
205 assert!((history[1] - 12.5).abs() < 0.1);
206 assert!((history[2] - 25.0).abs() < 0.1);
207 }
208
209 #[test]
210 fn test_buffer_last_n_more_than_available() {
211 let mut buffer = GpuMetricsBuffer::new(10, 1);
212
213 buffer.push(&[GpuMetrics { device_id: 0, utilization_percent: 50, ..Default::default() }]);
214
215 let last5 = buffer.last_n(0, 5);
217 assert_eq!(last5.len(), 1);
218 }
219
220 #[test]
221 fn test_buffer_last_nonexistent_device() {
222 let buffer = GpuMetricsBuffer::new(10, 1);
223 let last = buffer.last_n(99, 5);
224 assert!(last.is_empty());
225 }
226}
227
228#[cfg(test)]
229mod property_tests {
230 use super::*;
231 use proptest::prelude::*;
232
233 proptest! {
234 #![proptest_config(ProptestConfig::with_cases(200))]
235
236 #[test]
237 fn prop_buffer_respects_capacity(capacity in 1usize..20, pushes in 1usize..50) {
238 let mut buffer = GpuMetricsBuffer::new(capacity, 1);
239 for i in 0..pushes {
240 buffer.push(&[GpuMetrics {
241 device_id: 0,
242 utilization_percent: i as u32 % 100,
243 ..Default::default()
244 }]);
245 }
246 prop_assert!(buffer.len(0) <= capacity);
247 }
248 }
249}