trustformers_debug/profiler/
io_monitor.rs1#![allow(dead_code)]
6
7use serde::{Deserialize, Serialize};
8use std::collections::HashMap;
9use std::time::{Duration, Instant, SystemTime};
10use uuid::Uuid;
11
12#[derive(Debug, Clone, Serialize, Deserialize)]
14pub struct IoProfile {
15 pub operation_type: IoOperationType,
16 pub file_path: Option<String>,
17 pub bytes_transferred: usize,
18 pub duration: Duration,
19 pub bandwidth_mb_s: f64,
20 pub queue_time: Option<Duration>,
26 pub device_type: IoDeviceType,
27}
28
29#[derive(Debug, Clone, Serialize, Deserialize)]
30pub enum IoOperationType {
31 FileRead,
32 FileWrite,
33 NetworkRead,
34 NetworkWrite,
35 DatabaseQuery,
36 CacheLoad,
37 CacheStore,
38}
39
40#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Hash)]
41pub enum IoDeviceType {
42 SSD,
43 HDD,
44 Network,
45 Memory,
46 Cache,
47}
48
49#[derive(Debug, Clone, Serialize, Deserialize)]
51pub struct LayerLatencyProfile {
52 pub layer_name: String,
53 pub layer_type: String,
54 pub input_shapes: Vec<Vec<usize>>,
55 pub output_shapes: Vec<Vec<usize>>,
56 pub cpu_time: Duration,
57 pub gpu_time: Duration,
58 pub memory_copy_time: Duration,
59 pub sync_time: Duration,
60 pub parameter_count: usize,
61 pub flops: u64,
62 pub memory_footprint_bytes: usize,
63 pub cache_hit_rate: f64,
64}
65
66#[derive(Debug, Serialize, Deserialize)]
67pub struct IoPerformanceSummary {
68 pub total_operations: usize,
69 pub total_bytes_transferred: usize,
70 pub avg_bandwidth_by_device: HashMap<IoDeviceType, f64>,
71 pub slowest_operations: Vec<String>,
72}
73
74#[derive(Debug, Clone, Serialize, Deserialize)]
75pub struct BandwidthSample {
76 pub timestamp: SystemTime,
77 pub bandwidth_mb_s: f64,
78 pub device_type: IoDeviceType,
79}
80
81#[derive(Debug)]
83pub struct IoMonitor {
84 pub(crate) active_operations: HashMap<Uuid, IoOperation>,
85 pub(crate) bandwidth_history: Vec<BandwidthSample>,
86 pub(crate) io_queue_depth: usize,
87}
88
89#[derive(Debug)]
90pub struct IoOperation {
91 pub(crate) operation_id: Uuid,
92 pub(crate) start_time: Instant,
93 pub(crate) operation_type: IoOperationType,
94 pub(crate) bytes_expected: usize,
95}
96
97impl Default for IoMonitor {
98 fn default() -> Self {
99 Self::new()
100 }
101}
102
103impl IoMonitor {
104 pub fn new() -> Self {
105 Self {
106 active_operations: HashMap::new(),
107 bandwidth_history: Vec::new(),
108 io_queue_depth: 0,
109 }
110 }
111
112 pub fn start_io_operation(
113 &mut self,
114 operation_type: IoOperationType,
115 bytes_expected: usize,
116 ) -> Uuid {
117 let operation_id = Uuid::new_v4();
118 let operation = IoOperation {
119 operation_id,
120 start_time: Instant::now(),
121 operation_type,
122 bytes_expected,
123 };
124
125 self.active_operations.insert(operation_id, operation);
126 self.io_queue_depth += 1;
127 operation_id
128 }
129
130 pub fn finish_io_operation(
131 &mut self,
132 operation_id: Uuid,
133 bytes_transferred: usize,
134 ) -> Option<IoProfile> {
135 if let Some(operation) = self.active_operations.remove(&operation_id) {
136 let duration = operation.start_time.elapsed();
137 let bandwidth_mb_s = if duration.as_secs_f64() > 0.0 {
138 bytes_transferred as f64 / (1024.0 * 1024.0) / duration.as_secs_f64()
139 } else {
140 0.0
141 };
142
143 self.io_queue_depth = self.io_queue_depth.saturating_sub(1);
144
145 let device_type = match operation.operation_type {
146 IoOperationType::FileRead | IoOperationType::FileWrite => IoDeviceType::SSD,
147 IoOperationType::NetworkRead | IoOperationType::NetworkWrite => {
148 IoDeviceType::Network
149 },
150 IoOperationType::CacheLoad | IoOperationType::CacheStore => IoDeviceType::Cache,
151 _ => IoDeviceType::Memory,
152 };
153
154 self.bandwidth_history.push(BandwidthSample {
156 timestamp: SystemTime::now(),
157 bandwidth_mb_s,
158 device_type: device_type.clone(),
159 });
160
161 if self.bandwidth_history.len() > 1000 {
163 self.bandwidth_history.drain(0..500);
164 }
165
166 Some(IoProfile {
167 operation_type: operation.operation_type,
168 file_path: None,
171 bytes_transferred,
172 duration,
173 bandwidth_mb_s,
174 queue_time: None,
179 device_type,
180 })
181 } else {
182 None
183 }
184 }
185
186 pub fn get_average_bandwidth(&self, device_type: &IoDeviceType) -> f64 {
187 let samples: Vec<f64> = self
188 .bandwidth_history
189 .iter()
190 .filter(|s| &s.device_type == device_type)
191 .map(|s| s.bandwidth_mb_s)
192 .collect();
193
194 if samples.is_empty() {
195 0.0
196 } else {
197 samples.iter().sum::<f64>() / samples.len() as f64
198 }
199 }
200}
201
202#[cfg(test)]
203mod tests {
204 use super::*;
205
206 #[test]
207 fn test_io_monitor_new() {
208 let monitor = IoMonitor::new();
209 assert_eq!(monitor.io_queue_depth, 0);
210 assert!(monitor.bandwidth_history.is_empty());
211 }
212
213 #[test]
214 fn test_io_monitor_start_operation() {
215 let mut monitor = IoMonitor::new();
216 let _id = monitor.start_io_operation(IoOperationType::FileRead, 4096);
217 assert_eq!(monitor.io_queue_depth, 1);
218 assert_eq!(monitor.active_operations.len(), 1);
219 }
220
221 #[test]
222 fn test_io_monitor_finish_operation() {
223 let mut monitor = IoMonitor::new();
224 let id = monitor.start_io_operation(IoOperationType::FileWrite, 8192);
225 let profile = monitor.finish_io_operation(id, 8192);
226 assert!(profile.is_some());
227 let p = profile.expect("profile should be Some");
228 assert_eq!(p.bytes_transferred, 8192);
229 assert_eq!(monitor.io_queue_depth, 0);
230 }
231
232 #[test]
233 fn test_io_monitor_finish_nonexistent() {
234 let mut monitor = IoMonitor::new();
235 let profile = monitor.finish_io_operation(Uuid::new_v4(), 100);
236 assert!(profile.is_none());
237 }
238
239 #[test]
240 fn test_io_monitor_average_bandwidth_empty() {
241 let monitor = IoMonitor::new();
242 assert!((monitor.get_average_bandwidth(&IoDeviceType::SSD) - 0.0).abs() < 1e-9);
243 }
244
245 #[test]
246 fn test_io_monitor_device_type_mapping() {
247 let mut monitor = IoMonitor::new();
248 let id = monitor.start_io_operation(IoOperationType::NetworkRead, 1024);
249 let profile = monitor.finish_io_operation(id, 1024);
250 assert!(profile.is_some());
251 let p = profile.expect("profile should be Some");
252 assert_eq!(p.device_type, IoDeviceType::Network);
253 }
254
255 #[test]
256 fn test_io_monitor_cache_device_type() {
257 let mut monitor = IoMonitor::new();
258 let id = monitor.start_io_operation(IoOperationType::CacheLoad, 512);
259 let profile = monitor.finish_io_operation(id, 512);
260 assert!(profile.is_some());
261 let p = profile.expect("profile should be Some");
262 assert_eq!(p.device_type, IoDeviceType::Cache);
263 }
264}