cherenkov/qwen4_exp/gpu/activity/
reads.rs1use serde::{Deserialize, Serialize};
4use std::sync::{
5 Arc, Mutex,
6 atomic::{AtomicBool, Ordering},
7};
8use std::time::Instant;
9
10#[derive(Debug, Default, Clone, Copy, Serialize, Deserialize)]
11pub struct ReadStats {
12 pub requested_reads: u64,
13 pub requested_bytes: u64,
14 pub completed_reads: u64,
15 pub failed_reads: u64,
16 pub completed_bytes: u64,
17 pub read_seconds: f64,
18 pub max_read_seconds: f64,
19}
20
21impl ReadStats {
22 pub fn add(&mut self, other: Self) {
23 self.requested_reads += other.requested_reads;
24 self.requested_bytes += other.requested_bytes;
25 self.completed_reads += other.completed_reads;
26 self.failed_reads += other.failed_reads;
27 self.completed_bytes += other.completed_bytes;
28 self.read_seconds += other.read_seconds;
29 self.max_read_seconds = self.max_read_seconds.max(other.max_read_seconds);
30 }
31}
32
33#[derive(Debug, Clone, Copy)]
34pub enum ReadSource {
35 Prefill,
36 Demand,
37 Prefetch,
38}
39
40#[derive(Debug, Default, Clone, Copy, Serialize, Deserialize)]
41pub struct ReadSources {
42 pub prefill: ReadStats,
43 pub demand: ReadStats,
44 pub prefetch: ReadStats,
45}
46
47impl ReadSources {
48 fn source(&mut self, source: ReadSource) -> &mut ReadStats {
49 match source {
50 ReadSource::Prefill => &mut self.prefill,
51 ReadSource::Demand => &mut self.demand,
52 ReadSource::Prefetch => &mut self.prefetch,
53 }
54 }
55
56 pub fn add(&mut self, other: Self) {
57 self.prefill.add(other.prefill);
58 self.demand.add(other.demand);
59 self.prefetch.add(other.prefetch);
60 }
61
62 pub fn total(self) -> ReadStats {
63 let mut total = self.prefill;
64
65 total.add(self.demand);
66 total.add(self.prefetch);
67
68 total
69 }
70}
71
72#[derive(Debug, Clone)]
73pub struct ReadTracker(Arc<Mutex<Vec<[ReadSources; 3]>>>);
74
75impl ReadTracker {
76 pub fn new(layers: usize) -> Self {
77 Self(Arc::new(Mutex::new(vec![
78 [ReadSources::default(); 3];
79 layers
80 ])))
81 }
82
83 pub fn snapshot(&self) -> Vec<[ReadSources; 3]> {
84 self.0.lock().unwrap().clone()
85 }
86
87 pub fn ticket(&self, layer: usize, kind: u8, source: ReadSource, bytes: usize) -> ReadTicket {
88 let mut layers = self.0.lock().unwrap();
89 let stats = layers[layer][kind as usize].source(source);
90 stats.requested_reads += 1;
91 stats.requested_bytes += bytes as u64;
92
93 ReadTicket {
94 tracker: self.clone(),
95 layer,
96 kind,
97 source,
98 bytes,
99 finished: Arc::new(AtomicBool::new(false)),
100 }
101 }
102}
103
104#[derive(Debug, Clone)]
105pub struct ReadTicket {
106 tracker: ReadTracker,
107 layer: usize,
108 kind: u8,
109 source: ReadSource,
110 bytes: usize,
111 finished: Arc<AtomicBool>,
112}
113
114impl ReadTicket {
115 pub fn done(&self) -> bool {
116 self.finished.load(Ordering::Acquire)
117 }
118
119 pub fn measure(&self, read: impl FnOnce() -> usize) {
120 let started = Instant::now();
121 let bytes = read();
122
123 self.finish(bytes, started.elapsed().as_secs_f64());
124 }
125
126 fn finish(&self, bytes: usize, seconds: f64) {
127 let mut layers = self.tracker.0.lock().unwrap();
128 let stats = layers[self.layer][self.kind as usize].source(self.source);
129 stats.completed_reads += u64::from(bytes == self.bytes);
130 stats.failed_reads += u64::from(bytes != self.bytes);
131 stats.completed_bytes += bytes as u64;
132 stats.read_seconds += seconds;
133 stats.max_read_seconds = stats.max_read_seconds.max(seconds);
134
135 self.finished.store(true, Ordering::Release);
136 }
137}
138
139#[cfg(test)]
140#[path = "../../../../tests/unit/qwen4_exp/gpu/read_stats.rs"]
141mod tests;