Skip to main content

trustformers_debug/profiler/
io_monitor.rs

1//! I/O operation monitoring and bandwidth tracking
2// reason: debug/profiling scaffolding — structs are constructed and their fields/methods
3// are retained for the data model, serialization completeness, and future consumers that
4// do not yet read every member. Consolidated from many item-level #[allow(dead_code)].
5#![allow(dead_code)]
6
7use serde::{Deserialize, Serialize};
8use std::collections::HashMap;
9use std::time::{Duration, Instant, SystemTime};
10use uuid::Uuid;
11
12/// I/O operation profiling
13#[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    /// How long the operation waited in the device queue before starting.
21    ///
22    /// Always `None` from [`IoMonitor`]: it observes only the start and end of
23    /// each operation, never the wait behind other requests. It used to be
24    /// `queue_depth * 10ms`, an invented per-slot service time.
25    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/// Layer-wise latency analysis
50#[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/// I/O operation monitor
82#[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            // Record bandwidth sample
155            self.bandwidth_history.push(BandwidthSample {
156                timestamp: SystemTime::now(),
157                bandwidth_mb_s,
158                device_type: device_type.clone(),
159            });
160
161            // Keep only recent samples
162            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                // The recorded operation carries no path; callers that know one
169                // set it on the returned profile.
170                file_path: None,
171                bytes_transferred,
172                duration,
173                bandwidth_mb_s,
174                // No real queue-wait measurement is available: the monitor sees
175                // only start/end of the operation, never how long it waited
176                // behind others. This used to be `queue_depth * 10ms`, an
177                // invented per-slot service time published as a measurement.
178                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}