trustformers_debug/dashboard_ws/
metrics.rs

1//! Metric history tracking and dashboard message types.
2//!
3//! This module provides:
4//! - [`DashboardMessage`] — a typed enum of all real-time training events, with
5//!   manual JSON serialisation (no `serde` dependency).
6//! - [`MetricHistory`] — a fixed-capacity rolling window with EMA, trend, and
7//!   linear-regression utilities.
8//! - [`DashboardServerExt`] — an extended server that wraps the core SSE server
9//!   and keeps in-memory metric histories plus CSV/JSON summary generation.
10
11use std::fmt::Write as FmtWrite;
12
13// ─────────────────────────────────────────────────────────────────────────────
14// DashboardError
15// ─────────────────────────────────────────────────────────────────────────────
16
17/// Errors that can occur in the dashboard layer.
18#[derive(Debug, Clone, PartialEq)]
19pub enum DashboardError {
20    /// The network connection to a client could not be established.
21    ConnectionFailed,
22    /// Failed to serialise a [`DashboardMessage`] to JSON.
23    SerializationError,
24    /// The internal message buffer is full; the payload size is given.
25    BufferFull(usize),
26    /// A configuration value is invalid.
27    ConfigError(String),
28}
29
30impl std::fmt::Display for DashboardError {
31    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
32        match self {
33            Self::ConnectionFailed => write!(f, "dashboard connection failed"),
34            Self::SerializationError => write!(f, "failed to serialise dashboard message"),
35            Self::BufferFull(n) => write!(f, "dashboard message buffer full (size={n})"),
36            Self::ConfigError(msg) => write!(f, "dashboard configuration error: {msg}"),
37        }
38    }
39}
40
41impl std::error::Error for DashboardError {}
42
43// ─────────────────────────────────────────────────────────────────────────────
44// DashboardMessage
45// ─────────────────────────────────────────────────────────────────────────────
46
47/// Typed messages emitted by a training run to connected dashboard clients.
48///
49/// All variants are serialised to compact JSON via the hand-written
50/// [`DashboardMessage::to_json`] method (no `serde` macro dependency).
51///
52/// # Example
53///
54/// ```
55/// use trustformers_debug::dashboard_ws::metrics::DashboardMessage;
56///
57/// let msg = DashboardMessage::TrainingMetrics {
58///     step: 1,
59///     loss: 2.3,
60///     learning_rate: 1e-4,
61///     throughput_samples_per_sec: 128.0,
62/// };
63/// let json = msg.to_json();
64/// assert!(json.contains("\"type\":\"training_metrics\""));
65/// assert!(json.contains("\"step\":1"));
66/// ```
67#[derive(Debug, Clone, PartialEq)]
68pub enum DashboardMessage {
69    /// Per-step training metrics.
70    TrainingMetrics {
71        step: u64,
72        loss: f32,
73        learning_rate: f32,
74        throughput_samples_per_sec: f32,
75    },
76    /// Per-step validation metrics.
77    ValidationMetrics {
78        step: u64,
79        val_loss: f32,
80        val_accuracy: f32,
81    },
82    /// Global and per-layer gradient norms.
83    GradientNorm {
84        step: u64,
85        global_norm: f32,
86        per_layer: Vec<(String, f32)>,
87    },
88    /// Host and device memory snapshot.
89    MemoryUsage {
90        step: u64,
91        gpu_mb: f32,
92        cpu_mb: f32,
93        peak_gpu_mb: f32,
94    },
95    /// A checkpoint was written to disk.
96    CheckpointSaved { step: u64, path: String },
97    /// Training finished successfully.
98    TrainingComplete {
99        total_steps: u64,
100        final_loss: f32,
101        duration_secs: f64,
102    },
103    /// Periodic keep-alive ping.
104    Heartbeat { timestamp_ms: u64 },
105    /// Server-side error notification.
106    Error { code: u32, message: String },
107}
108
109impl DashboardMessage {
110    /// Returns the `type` field value used in the JSON envelope.
111    pub fn message_type(&self) -> &str {
112        match self {
113            Self::TrainingMetrics { .. } => "training_metrics",
114            Self::ValidationMetrics { .. } => "validation_metrics",
115            Self::GradientNorm { .. } => "gradient_norm",
116            Self::MemoryUsage { .. } => "memory_usage",
117            Self::CheckpointSaved { .. } => "checkpoint_saved",
118            Self::TrainingComplete { .. } => "training_complete",
119            Self::Heartbeat { .. } => "heartbeat",
120            Self::Error { .. } => "error",
121        }
122    }
123
124    /// Serialise the message to a compact JSON string.
125    ///
126    /// Uses manual formatting to avoid pulling in `serde` inside this module.
127    /// The output always contains a `"type"` key plus variant-specific fields.
128    ///
129    /// # Example
130    ///
131    /// ```
132    /// use trustformers_debug::dashboard_ws::metrics::DashboardMessage;
133    ///
134    /// let msg = DashboardMessage::Heartbeat { timestamp_ms: 12345 };
135    /// let json = msg.to_json();
136    /// assert!(json.contains("\"type\":\"heartbeat\""));
137    /// assert!(json.contains("\"timestamp_ms\":12345"));
138    /// ```
139    pub fn to_json(&self) -> String {
140        match self {
141            Self::TrainingMetrics {
142                step,
143                loss,
144                learning_rate,
145                throughput_samples_per_sec,
146            } => {
147                format!(
148                    r#"{{"type":"training_metrics","step":{step},"loss":{loss},"learning_rate":{learning_rate},"throughput_samples_per_sec":{throughput_samples_per_sec}}}"#,
149                )
150            },
151            Self::ValidationMetrics {
152                step,
153                val_loss,
154                val_accuracy,
155            } => {
156                format!(
157                    r#"{{"type":"validation_metrics","step":{step},"val_loss":{val_loss},"val_accuracy":{val_accuracy}}}"#,
158                )
159            },
160            Self::GradientNorm {
161                step,
162                global_norm,
163                per_layer,
164            } => {
165                let layers_json = Self::per_layer_to_json(per_layer);
166                format!(
167                    r#"{{"type":"gradient_norm","step":{step},"global_norm":{global_norm},"per_layer":{layers_json}}}"#,
168                )
169            },
170            Self::MemoryUsage {
171                step,
172                gpu_mb,
173                cpu_mb,
174                peak_gpu_mb,
175            } => {
176                format!(
177                    r#"{{"type":"memory_usage","step":{step},"gpu_mb":{gpu_mb},"cpu_mb":{cpu_mb},"peak_gpu_mb":{peak_gpu_mb}}}"#,
178                )
179            },
180            Self::CheckpointSaved { step, path } => {
181                let escaped = escape_json_string(path);
182                format!(r#"{{"type":"checkpoint_saved","step":{step},"path":"{escaped}"}}"#)
183            },
184            Self::TrainingComplete {
185                total_steps,
186                final_loss,
187                duration_secs,
188            } => {
189                format!(
190                    r#"{{"type":"training_complete","total_steps":{total_steps},"final_loss":{final_loss},"duration_secs":{duration_secs}}}"#,
191                )
192            },
193            Self::Heartbeat { timestamp_ms } => {
194                format!(r#"{{"type":"heartbeat","timestamp_ms":{timestamp_ms}}}"#)
195            },
196            Self::Error { code, message } => {
197                let escaped = escape_json_string(message);
198                format!(r#"{{"type":"error","code":{code},"message":"{escaped}"}}"#)
199            },
200        }
201    }
202
203    // ── private helpers ──────────────────────────────────────────────────────
204
205    fn per_layer_to_json(per_layer: &[(String, f32)]) -> String {
206        let mut out = String::from('[');
207        for (i, (name, norm)) in per_layer.iter().enumerate() {
208            if i > 0 {
209                out.push(',');
210            }
211            let escaped = escape_json_string(name);
212            let _ = write!(out, r#"["{escaped}",{norm}]"#);
213        }
214        out.push(']');
215        out
216    }
217}
218
219/// Escapes a string for safe embedding inside a JSON string literal.
220fn escape_json_string(s: &str) -> String {
221    let mut out = String::with_capacity(s.len());
222    for c in s.chars() {
223        match c {
224            '"' => out.push_str("\\\""),
225            '\\' => out.push_str("\\\\"),
226            '\n' => out.push_str("\\n"),
227            '\r' => out.push_str("\\r"),
228            '\t' => out.push_str("\\t"),
229            c if (c as u32) < 0x20 => {
230                let _ = write!(out, "\\u{:04x}", c as u32);
231            },
232            c => out.push(c),
233        }
234    }
235    out
236}
237
238// ─────────────────────────────────────────────────────────────────────────────
239// DashboardConfig (extended)
240// ─────────────────────────────────────────────────────────────────────────────
241
242/// Configuration for the extended training-metrics dashboard.
243///
244/// Unlike the SSE server configuration ([`super::websocket::DashboardConfig`]),
245/// this struct captures parameters for heartbeat scheduling and metric buffering.
246#[derive(Debug, Clone)]
247pub struct DashboardConfig {
248    pub host: String,
249    pub port: u16,
250    pub max_connections: usize,
251    pub heartbeat_interval_ms: u64,
252    pub message_buffer_size: usize,
253}
254
255impl Default for DashboardConfig {
256    fn default() -> Self {
257        Self {
258            host: "127.0.0.1".to_string(),
259            port: 7878,
260            max_connections: 16,
261            heartbeat_interval_ms: 5_000,
262            message_buffer_size: 512,
263        }
264    }
265}
266
267// ─────────────────────────────────────────────────────────────────────────────
268// MetricHistory
269// ─────────────────────────────────────────────────────────────────────────────
270
271/// A fixed-capacity rolling window of `(step, value)` pairs.
272///
273/// Older entries are silently dropped when `max_history` is exceeded.
274/// Provides EMA smoothing and linear-regression trend estimation.
275///
276/// # Example
277///
278/// ```
279/// use trustformers_debug::dashboard_ws::metrics::MetricHistory;
280///
281/// let mut h = MetricHistory::new(100);
282/// h.push(0, 1.0);
283/// h.push(1, 0.8);
284/// h.push(2, 0.6);
285/// assert_eq!(h.latest(), Some((2, 0.6)));
286/// let trend = h.trend(3);
287/// assert!(trend.is_some());
288/// assert!(trend.unwrap() < 0.0, "loss should be trending down");
289/// ```
290#[derive(Debug, Clone)]
291pub struct MetricHistory {
292    pub steps: Vec<u64>,
293    pub values: Vec<f32>,
294    pub max_history: usize,
295}
296
297impl MetricHistory {
298    /// Creates a new empty history with the given maximum capacity.
299    pub fn new(max_history: usize) -> Self {
300        Self {
301            steps: Vec::with_capacity(max_history.min(4096)),
302            values: Vec::with_capacity(max_history.min(4096)),
303            max_history,
304        }
305    }
306
307    /// Appends a new `(step, value)` pair, evicting the oldest when full.
308    pub fn push(&mut self, step: u64, value: f32) {
309        if self.steps.len() >= self.max_history {
310            self.steps.remove(0);
311            self.values.remove(0);
312        }
313        self.steps.push(step);
314        self.values.push(value);
315    }
316
317    /// Returns the most-recently pushed `(step, value)` pair, or `None` if empty.
318    pub fn latest(&self) -> Option<(u64, f32)> {
319        self.steps.last().copied().zip(self.values.last().copied())
320    }
321
322    /// Estimates the linear-regression slope over the last `window` data points.
323    ///
324    /// Returns `None` when fewer than two points are available in the window.
325    /// A negative slope means the metric is decreasing (good for loss).
326    ///
327    /// Uses the standard ordinary-least-squares formula:
328    /// `slope = (N·Σ(x·y) − Σx·Σy) / (N·Σx² − (Σx)²)`
329    pub fn trend(&self, window: usize) -> Option<f32> {
330        let n = self.values.len();
331        if n < 2 {
332            return None;
333        }
334        let start = n.saturating_sub(window);
335        let xs = &self.steps[start..];
336        let ys = &self.values[start..];
337        let k = xs.len();
338        if k < 2 {
339            return None;
340        }
341
342        let sum_x: f64 = xs.iter().map(|&x| x as f64).sum();
343        let sum_y: f64 = ys.iter().map(|&y| y as f64).sum();
344        let sum_xx: f64 = xs.iter().map(|&x| (x as f64) * (x as f64)).sum();
345        let sum_xy: f64 = xs.iter().zip(ys.iter()).map(|(&x, &y)| (x as f64) * (y as f64)).sum();
346        let kf = k as f64;
347
348        let denom = kf * sum_xx - sum_x * sum_x;
349        if denom.abs() < f64::EPSILON {
350            return None;
351        }
352
353        let slope = (kf * sum_xy - sum_x * sum_y) / denom;
354        Some(slope as f32)
355    }
356
357    /// Computes an exponential moving average (EMA) over all stored values.
358    ///
359    /// `alpha` controls the smoothing factor (0 < alpha ≤ 1).  Higher `alpha`
360    /// gives more weight to recent observations.  Returns an empty vec if the
361    /// history is empty.
362    ///
363    /// # Panics
364    ///
365    /// Does not panic, but clamps `alpha` to `[1e-6, 1.0]` internally.
366    pub fn smooth(&self, alpha: f32) -> Vec<f32> {
367        if self.values.is_empty() {
368            return Vec::new();
369        }
370        let alpha = alpha.max(1e-6_f32).min(1.0_f32);
371        let mut out = Vec::with_capacity(self.values.len());
372        let mut ema = self.values[0];
373        out.push(ema);
374        for &v in &self.values[1..] {
375            ema = alpha * v + (1.0 - alpha) * ema;
376            out.push(ema);
377        }
378        out
379    }
380
381    /// Returns the number of data points currently stored.
382    pub fn len(&self) -> usize {
383        self.steps.len()
384    }
385
386    /// Returns `true` when no data has been recorded yet.
387    pub fn is_empty(&self) -> bool {
388        self.steps.is_empty()
389    }
390}
391
392// ─────────────────────────────────────────────────────────────────────────────
393// DashboardServer (extended)
394// ─────────────────────────────────────────────────────────────────────────────
395
396/// In-process training dashboard that records metric histories and can generate
397/// CSV / JSON summaries suitable for logging or downstream tools.
398///
399/// This is an *in-process* implementation that does not bind network sockets.
400/// For SSE streaming, see [`super::websocket::DashboardServer`].
401///
402/// # Example
403///
404/// ```
405/// use trustformers_debug::dashboard_ws::metrics::{DashboardConfig, DashboardServerExt};
406///
407/// let mut server = DashboardServerExt::new(DashboardConfig::default());
408/// server.record_metrics(10, 1.5, 1e-3, 0.9);
409/// server.record_metrics(20, 1.2, 9e-4, 0.85);
410/// let csv = server.format_metric_csv();
411/// assert!(csv.contains("step"));
412/// ```
413pub struct DashboardServerExt {
414    pub config: DashboardConfig,
415    pub loss_history: MetricHistory,
416    pub lr_history: MetricHistory,
417    pub grad_norm_history: MetricHistory,
418    pub connected_clients: usize,
419    pub messages_sent: u64,
420    /// In-process broadcast buffer (bounded by `config.message_buffer_size`).
421    message_log: Vec<DashboardMessage>,
422}
423
424impl DashboardServerExt {
425    /// Creates a new in-process dashboard server with the given configuration.
426    pub fn new(config: DashboardConfig) -> Self {
427        let cap = config.message_buffer_size;
428        Self {
429            loss_history: MetricHistory::new(cap),
430            lr_history: MetricHistory::new(cap),
431            grad_norm_history: MetricHistory::new(cap),
432            connected_clients: 0,
433            messages_sent: 0,
434            message_log: Vec::with_capacity(cap.min(4096)),
435            config,
436        }
437    }
438
439    /// Records one training step into the internal metric histories.
440    pub fn record_metrics(&mut self, step: u64, loss: f32, lr: f32, grad_norm: f32) {
441        self.loss_history.push(step, loss);
442        self.lr_history.push(step, lr);
443        self.grad_norm_history.push(step, grad_norm);
444    }
445
446    /// Appends a [`DashboardMessage`] to the internal broadcast buffer.
447    ///
448    /// Returns [`DashboardError::BufferFull`] when the buffer is at capacity.
449    pub fn broadcast_message(&mut self, msg: &DashboardMessage) -> Result<(), DashboardError> {
450        if self.message_log.len() >= self.config.message_buffer_size {
451            return Err(DashboardError::BufferFull(self.message_log.len()));
452        }
453        self.message_log.push(msg.clone());
454        self.messages_sent += 1;
455        Ok(())
456    }
457
458    /// Returns all messages currently in the broadcast buffer (in order).
459    pub fn buffered_messages(&self) -> &[DashboardMessage] {
460        &self.message_log
461    }
462
463    /// Clears the broadcast buffer.
464    pub fn clear_buffer(&mut self) {
465        self.message_log.clear();
466    }
467
468    /// Returns a JSON summary of all metric histories.
469    ///
470    /// Includes the latest value, a 10-point trend (slope), and EMA
471    /// (α = 0.1) for each of `loss`, `lr`, and `grad_norm`.
472    ///
473    /// # Example
474    ///
475    /// ```
476    /// use trustformers_debug::dashboard_ws::metrics::{DashboardConfig, DashboardServerExt};
477    ///
478    /// let mut s = DashboardServerExt::new(DashboardConfig::default());
479    /// s.record_metrics(1, 2.0, 1e-3, 0.5);
480    /// let json = s.generate_summary_json();
481    /// assert!(json.contains("\"loss\""));
482    /// ```
483    pub fn generate_summary_json(&self) -> String {
484        let loss_latest = self
485            .loss_history
486            .latest()
487            .map(|(_, v)| format!("{v}"))
488            .unwrap_or_else(|| "null".to_string());
489        let lr_latest = self
490            .lr_history
491            .latest()
492            .map(|(_, v)| format!("{v}"))
493            .unwrap_or_else(|| "null".to_string());
494        let grad_latest = self
495            .grad_norm_history
496            .latest()
497            .map(|(_, v)| format!("{v}"))
498            .unwrap_or_else(|| "null".to_string());
499
500        let loss_trend = self
501            .loss_history
502            .trend(10)
503            .map(|v| format!("{v}"))
504            .unwrap_or_else(|| "null".to_string());
505        let lr_trend = self
506            .lr_history
507            .trend(10)
508            .map(|v| format!("{v}"))
509            .unwrap_or_else(|| "null".to_string());
510        let grad_trend = self
511            .grad_norm_history
512            .trend(10)
513            .map(|v| format!("{v}"))
514            .unwrap_or_else(|| "null".to_string());
515
516        let loss_ema = format_f32_slice(&self.loss_history.smooth(0.1));
517        let lr_ema = format_f32_slice(&self.lr_history.smooth(0.1));
518        let grad_ema = format_f32_slice(&self.grad_norm_history.smooth(0.1));
519
520        format!(
521            r#"{{"connected_clients":{clients},"messages_sent":{sent},"loss":{{"latest":{loss_latest},"trend_slope":{loss_trend},"ema_alpha0.1":{loss_ema}}},"lr":{{"latest":{lr_latest},"trend_slope":{lr_trend},"ema_alpha0.1":{lr_ema}}},"grad_norm":{{"latest":{grad_latest},"trend_slope":{grad_trend},"ema_alpha0.1":{grad_ema}}}}}"#,
522            clients = self.connected_clients,
523            sent = self.messages_sent,
524        )
525    }
526
527    /// Exports all recorded metrics as a CSV string.
528    ///
529    /// Columns: `step,loss,lr,grad_norm`
530    /// Rows are aligned by index — if one history is shorter it writes empty cells.
531    pub fn format_metric_csv(&self) -> String {
532        let mut out = String::from("step,loss,lr,grad_norm\n");
533        let len = self
534            .loss_history
535            .len()
536            .max(self.lr_history.len())
537            .max(self.grad_norm_history.len());
538
539        for i in 0..len {
540            let step = self
541                .loss_history
542                .steps
543                .get(i)
544                .or_else(|| self.lr_history.steps.get(i))
545                .or_else(|| self.grad_norm_history.steps.get(i))
546                .copied()
547                .unwrap_or(i as u64);
548
549            let loss = self.loss_history.values.get(i).map(|v| format!("{v}")).unwrap_or_default();
550            let lr = self.lr_history.values.get(i).map(|v| format!("{v}")).unwrap_or_default();
551            let grad =
552                self.grad_norm_history.values.get(i).map(|v| format!("{v}")).unwrap_or_default();
553
554            let _ = writeln!(out, "{step},{loss},{lr},{grad}");
555        }
556        out
557    }
558}
559
560// ─────────────────────────────────────────────────────────────────────────────
561// Private helpers
562// ─────────────────────────────────────────────────────────────────────────────
563
564fn format_f32_slice(values: &[f32]) -> String {
565    let mut out = String::from('[');
566    for (i, v) in values.iter().enumerate() {
567        if i > 0 {
568            out.push(',');
569        }
570        let _ = write!(out, "{v}");
571    }
572    out.push(']');
573    out
574}
575
576// ─────────────────────────────────────────────────────────────────────────────
577// Tests
578// ─────────────────────────────────────────────────────────────────────────────
579
580#[cfg(test)]
581mod tests {
582    use super::*;
583
584    // ── DashboardMessage::to_json ────────────────────────────────────────────
585
586    #[test]
587    fn test_training_metrics_json() {
588        let msg = DashboardMessage::TrainingMetrics {
589            step: 42,
590            loss: 1.5,
591            learning_rate: 1e-4,
592            throughput_samples_per_sec: 256.0,
593        };
594        let json = msg.to_json();
595        assert!(json.contains("\"type\":\"training_metrics\""));
596        assert!(json.contains("\"step\":42"));
597        assert!(json.contains("\"loss\":1.5"));
598        assert_eq!(msg.message_type(), "training_metrics");
599    }
600
601    #[test]
602    fn test_validation_metrics_json() {
603        let msg = DashboardMessage::ValidationMetrics {
604            step: 10,
605            val_loss: 0.9,
606            val_accuracy: 0.85,
607        };
608        let json = msg.to_json();
609        assert!(json.contains("\"type\":\"validation_metrics\""));
610        assert!(json.contains("\"step\":10"));
611        assert!(json.contains("\"val_loss\":0.9"));
612        assert!(json.contains("\"val_accuracy\":0.85"));
613    }
614
615    #[test]
616    fn test_gradient_norm_json_with_per_layer() {
617        let msg = DashboardMessage::GradientNorm {
618            step: 5,
619            global_norm: 1.2,
620            per_layer: vec![
621                ("layer_0.weight".to_string(), 0.3),
622                ("layer_1.weight".to_string(), 0.9),
623            ],
624        };
625        let json = msg.to_json();
626        assert!(json.contains("\"type\":\"gradient_norm\""));
627        assert!(json.contains("\"global_norm\":1.2"));
628        assert!(json.contains("layer_0.weight"));
629        assert!(json.contains("layer_1.weight"));
630    }
631
632    #[test]
633    fn test_memory_usage_json() {
634        let msg = DashboardMessage::MemoryUsage {
635            step: 100,
636            gpu_mb: 4096.0,
637            cpu_mb: 8192.0,
638            peak_gpu_mb: 5000.0,
639        };
640        let json = msg.to_json();
641        assert!(json.contains("\"type\":\"memory_usage\""));
642        assert!(json.contains("\"gpu_mb\":4096"));
643        assert!(json.contains("\"peak_gpu_mb\":5000"));
644    }
645
646    #[test]
647    fn test_checkpoint_saved_json_with_escape() {
648        let msg = DashboardMessage::CheckpointSaved {
649            step: 50,
650            path: "/tmp/ckpt/step_50/model\"s.bin".to_string(),
651        };
652        let json = msg.to_json();
653        assert!(json.contains("\"type\":\"checkpoint_saved\""));
654        assert!(json.contains("\\\""));
655    }
656
657    #[test]
658    fn test_training_complete_json() {
659        let msg = DashboardMessage::TrainingComplete {
660            total_steps: 1000,
661            final_loss: 0.05,
662            duration_secs: 3600.5,
663        };
664        let json = msg.to_json();
665        assert!(json.contains("\"type\":\"training_complete\""));
666        assert!(json.contains("\"total_steps\":1000"));
667        assert!(json.contains("3600.5"));
668    }
669
670    #[test]
671    fn test_heartbeat_json() {
672        let msg = DashboardMessage::Heartbeat {
673            timestamp_ms: 999_999,
674        };
675        let json = msg.to_json();
676        assert!(json.contains("\"type\":\"heartbeat\""));
677        assert!(json.contains("\"timestamp_ms\":999999"));
678    }
679
680    #[test]
681    fn test_error_json() {
682        let msg = DashboardMessage::Error {
683            code: 500,
684            message: "internal error".to_string(),
685        };
686        let json = msg.to_json();
687        assert!(json.contains("\"type\":\"error\""));
688        assert!(json.contains("\"code\":500"));
689        assert!(json.contains("internal error"));
690    }
691
692    // ── message_type ────────────────────────────────────────────────────────
693
694    #[test]
695    fn test_message_type_variants() {
696        assert_eq!(
697            DashboardMessage::Heartbeat { timestamp_ms: 0 }.message_type(),
698            "heartbeat"
699        );
700        assert_eq!(
701            DashboardMessage::Error {
702                code: 0,
703                message: String::new()
704            }
705            .message_type(),
706            "error"
707        );
708        assert_eq!(
709            DashboardMessage::CheckpointSaved {
710                step: 0,
711                path: String::new()
712            }
713            .message_type(),
714            "checkpoint_saved"
715        );
716    }
717
718    // ── MetricHistory ────────────────────────────────────────────────────────
719
720    #[test]
721    fn test_metric_history_push_and_latest() {
722        let mut h = MetricHistory::new(5);
723        assert!(h.latest().is_none());
724        h.push(1, 2.0);
725        h.push(2, 1.5);
726        h.push(3, 1.0);
727        assert_eq!(h.latest(), Some((3, 1.0)));
728        assert_eq!(h.len(), 3);
729    }
730
731    #[test]
732    fn test_metric_history_evicts_oldest_when_full() {
733        let mut h = MetricHistory::new(3);
734        h.push(1, 1.0);
735        h.push(2, 2.0);
736        h.push(3, 3.0);
737        h.push(4, 4.0); // should evict step=1
738        assert_eq!(h.len(), 3);
739        assert_eq!(h.steps[0], 2);
740        assert_eq!(h.latest(), Some((4, 4.0)));
741    }
742
743    #[test]
744    fn test_metric_history_trend_decreasing() {
745        let mut h = MetricHistory::new(100);
746        for i in 0_u64..20 {
747            h.push(i, 10.0 - i as f32 * 0.5);
748        }
749        let slope = h.trend(10).expect("should compute slope");
750        assert!(
751            slope < 0.0,
752            "loss is decreasing so slope must be negative: {slope}"
753        );
754    }
755
756    #[test]
757    fn test_metric_history_trend_none_with_single_point() {
758        let mut h = MetricHistory::new(100);
759        h.push(0, 1.0);
760        assert!(h.trend(10).is_none());
761    }
762
763    #[test]
764    fn test_metric_history_smooth_ema() {
765        let mut h = MetricHistory::new(100);
766        for i in 0..10_u64 {
767            h.push(i, i as f32);
768        }
769        let smoothed = h.smooth(0.3);
770        assert_eq!(smoothed.len(), 10);
771        // first value equals original first value
772        assert_eq!(smoothed[0], 0.0);
773        // EMA should be less than the raw value for rising series (lagging)
774        assert!(smoothed.last().copied().unwrap() < 9.0);
775    }
776
777    #[test]
778    fn test_metric_history_smooth_alpha_one_equals_original() {
779        let mut h = MetricHistory::new(10);
780        let vals = [1.0f32, 3.0, 2.0, 5.0, 4.0];
781        for (i, &v) in vals.iter().enumerate() {
782            h.push(i as u64, v);
783        }
784        let smoothed = h.smooth(1.0);
785        for (s, &orig) in smoothed.iter().zip(vals.iter()) {
786            assert!(
787                (s - orig).abs() < 1e-5,
788                "alpha=1 should pass through: {s} vs {orig}"
789            );
790        }
791    }
792
793    // ── DashboardServerExt ───────────────────────────────────────────────────
794
795    #[test]
796    fn test_server_record_and_csv() {
797        let mut s = DashboardServerExt::new(DashboardConfig::default());
798        s.record_metrics(0, 2.0, 1e-3, 1.0);
799        s.record_metrics(1, 1.5, 9e-4, 0.8);
800        let csv = s.format_metric_csv();
801        assert!(csv.starts_with("step,loss,lr,grad_norm\n"));
802        assert!(csv.contains("0,2,"));
803        assert!(csv.contains("1,1.5,"));
804    }
805
806    #[test]
807    fn test_server_generate_summary_json_keys() {
808        let mut s = DashboardServerExt::new(DashboardConfig::default());
809        s.record_metrics(0, 1.0, 0.001, 0.5);
810        let json = s.generate_summary_json();
811        assert!(json.contains("\"loss\""));
812        assert!(json.contains("\"lr\""));
813        assert!(json.contains("\"grad_norm\""));
814        assert!(json.contains("\"connected_clients\""));
815        assert!(json.contains("\"messages_sent\""));
816    }
817
818    #[test]
819    fn test_server_broadcast_and_buffer_full() {
820        let config = DashboardConfig {
821            message_buffer_size: 3,
822            ..Default::default()
823        };
824        let mut s = DashboardServerExt::new(config);
825        let msg = DashboardMessage::Heartbeat { timestamp_ms: 1 };
826        assert!(s.broadcast_message(&msg).is_ok());
827        assert!(s.broadcast_message(&msg).is_ok());
828        assert!(s.broadcast_message(&msg).is_ok());
829        let err = s.broadcast_message(&msg);
830        assert!(matches!(err, Err(DashboardError::BufferFull(_))));
831        assert_eq!(s.messages_sent, 3);
832    }
833
834    #[test]
835    fn test_server_clear_buffer() {
836        let mut s = DashboardServerExt::new(DashboardConfig::default());
837        let msg = DashboardMessage::Heartbeat { timestamp_ms: 0 };
838        s.broadcast_message(&msg).unwrap();
839        assert_eq!(s.buffered_messages().len(), 1);
840        s.clear_buffer();
841        assert!(s.buffered_messages().is_empty());
842    }
843
844    #[test]
845    fn test_dashboard_error_display() {
846        assert_eq!(
847            DashboardError::ConnectionFailed.to_string(),
848            "dashboard connection failed"
849        );
850        assert_eq!(
851            DashboardError::SerializationError.to_string(),
852            "failed to serialise dashboard message"
853        );
854        assert!(DashboardError::BufferFull(10).to_string().contains("10"));
855        assert!(DashboardError::ConfigError("bad port".to_string())
856            .to_string()
857            .contains("bad port"));
858    }
859}
trustformers_debug/dashboard_ws/metrics.rs

trustformers_debug/dashboard_ws/
metrics.rs
