fix: fix issue #2 - telemetry not showing in graphs

src/metrics.rs

@@ -145,6 +145,7 @@ pub(crate) struct Metrics {
     context_limit: AtomicU32,
     decode_tps: AtomicU32,
     prefill_tps: AtomicU32,
+    prefill_sample: AtomicU32,
     last_runtime_ms: AtomicU64,
     total_runtime_ms: AtomicU64,
     last_prompt_tokens: AtomicU64,
@@ -169,6 +170,8 @@ pub(crate) struct Metrics {
     kv_write_errors: AtomicU64,
     kv_read_bytes: AtomicU64,
     kv_write_bytes: AtomicU64,
+    kv_read_sample_bytes: AtomicU64,
+    kv_write_sample_bytes: AtomicU64,
     last_kv_read_ms: AtomicU64,
     last_kv_write_ms: AtomicU64,
     kv_files: AtomicU64,
@@ -215,6 +218,7 @@ impl Metrics {
             context_limit: AtomicU32::new(0),
             decode_tps: AtomicU32::new(0),
             prefill_tps: AtomicU32::new(0),
+            prefill_sample: AtomicU32::new(0),
             last_runtime_ms: AtomicU64::new(0),
             total_runtime_ms: AtomicU64::new(0),
             last_prompt_tokens: AtomicU64::new(0),
@@ -239,6 +243,8 @@ impl Metrics {
             kv_write_errors: AtomicU64::new(0),
             kv_read_bytes: AtomicU64::new(0),
             kv_write_bytes: AtomicU64::new(0),
+            kv_read_sample_bytes: AtomicU64::new(0),
+            kv_write_sample_bytes: AtomicU64::new(0),
             last_kv_read_ms: AtomicU64::new(0),
             last_kv_write_ms: AtomicU64::new(0),
             kv_files: AtomicU64::new(cache.files),
@@ -282,6 +288,7 @@ impl Metrics {
         self.source.store(source as u8, Ordering::Relaxed);
         self.decode_tps.store(0, Ordering::Relaxed);
         self.prefill_tps.store(0, Ordering::Relaxed);
+        self.prefill_sample.store(0, Ordering::Relaxed);
     }
 
     pub(crate) fn loading(&self) {
@@ -313,6 +320,7 @@ impl Metrics {
         self.context_used.store(used, Ordering::Relaxed);
         self.context_limit.store(limit, Ordering::Relaxed);
         store_f32(&self.prefill_tps, speed);
+        store_f32(&self.prefill_sample, speed);
     }
 
     pub(crate) fn generation_progress(&self, used: u32, limit: u32, speed: f32) {
@@ -376,6 +384,7 @@ impl Metrics {
         self.context_used.store(0, Ordering::Relaxed);
         self.decode_tps.store(0, Ordering::Relaxed);
         self.prefill_tps.store(0, Ordering::Relaxed);
+        self.prefill_sample.store(0, Ordering::Relaxed);
         self.model_unloads.fetch_add(1, Ordering::Relaxed);
     }
 
@@ -418,6 +427,8 @@ impl Metrics {
 
     pub(crate) fn kv_read_bytes(&self, bytes: u64) {
         self.kv_read_bytes.fetch_add(bytes, Ordering::Relaxed);
+        self.kv_read_sample_bytes
+            .fetch_add(bytes, Ordering::Relaxed);
     }
 
     pub(crate) fn kv_read_finished(&self, elapsed: Duration, failed: bool) {
@@ -436,6 +447,8 @@ impl Metrics {
 
     pub(crate) fn kv_write_bytes(&self, bytes: u64) {
         self.kv_write_bytes.fetch_add(bytes, Ordering::Relaxed);
+        self.kv_write_sample_bytes
+            .fetch_add(bytes, Ordering::Relaxed);
     }
 
     pub(crate) fn kv_write_finished(&self, elapsed: Duration, failed: bool) {
@@ -560,6 +573,17 @@ impl Metrics {
         }
     }
 
+    pub(crate) fn take_prefill_sample(&self) -> f32 {
+        f32::from_bits(self.prefill_sample.swap(0, Ordering::Relaxed))
+    }
+
+    pub(crate) fn take_kv_io_sample(&self) -> (u64, u64) {
+        (
+            self.kv_read_sample_bytes.swap(0, Ordering::Relaxed),
+            self.kv_write_sample_bytes.swap(0, Ordering::Relaxed),
+        )
+    }
+
     fn record_checkpoint(&self, source: WorkSource, previous_bytes: Option<u64>, bytes: u64) {
         if bytes == 0 {
             return;
@@ -685,6 +709,8 @@ mod tests {
         metrics.request_started(WorkSource::LocalChat);
         metrics.prefill_progress(100, 1_000, 250.0);
         metrics.generation_progress(120, 1_000, 20.0);
+        assert_eq!(metrics.take_prefill_sample(), 250.0);
+        assert_eq!(metrics.take_prefill_sample(), 0.0);
         metrics.kv_lookup(KvLookup::DiskHit);
         metrics.kv_lookup(KvLookup::MemoryHit);
         metrics.kv_lookup(KvLookup::Miss);
@@ -697,6 +723,8 @@ mod tests {
         metrics.kv_write_started();
         metrics.kv_write_bytes(4_096);
         metrics.kv_write_finished(Duration::from_millis(20), false);
+        assert_eq!(metrics.take_kv_io_sample(), (2_048, 4_096));
+        assert_eq!(metrics.take_kv_io_sample(), (0, 0));
         metrics.request_finished(
             WorkSource::LocalChat,
             Duration::from_millis(250),