Skip to main content

entrenar/dashboard/
source.rs

1//! Dashboard data source trait and implementations.
2
3use std::collections::HashMap;
4use std::sync::{Arc, Mutex};
5
6use super::{MetricSnapshot, ResourceSnapshot};
7use crate::run::Run;
8use crate::storage::{ExperimentStorage, RunStatus};
9
10/// Subscription callback type.
11pub type SubscriptionCallback = Box<dyn Fn(&str, f64) + Send>;
12
13/// Dashboard data source trait.
14///
15/// Implement this trait to provide data for real-time dashboards.
16pub trait DashboardSource {
17    /// Get the current run status.
18    fn status(&self) -> RunStatus;
19
20    /// Get recent metrics, limited to `limit` points per metric.
21    fn recent_metrics(&self, limit: usize) -> HashMap<String, MetricSnapshot>;
22
23    /// Subscribe to metric updates.
24    ///
25    /// Returns a receiver that will receive metric snapshots as they arrive.
26    /// The callback is called with the metric key and latest value.
27    fn subscribe(&self, callback: SubscriptionCallback);
28
29    /// Get current resource usage.
30    fn resource_usage(&self) -> ResourceSnapshot;
31}
32
33impl<S: ExperimentStorage> DashboardSource for Run<S> {
34    fn status(&self) -> RunStatus {
35        if self.is_finished() {
36            // Query storage for actual status
37            let storage = self.storage_ref();
38            storage
39                .lock()
40                .ok()
41                .and_then(|s| s.get_run_status(&self.id).ok())
42                .unwrap_or(RunStatus::Success)
43        } else {
44            RunStatus::Running
45        }
46    }
47
48    fn recent_metrics(&self, limit: usize) -> HashMap<String, MetricSnapshot> {
49        let mut result = HashMap::new();
50
51        // Get metrics from storage
52        let storage = self.storage_ref();
53        if let Ok(guard) = storage.lock() {
54            // Get all metric keys we've logged
55            for key in self.metric_keys() {
56                if let Ok(points) = guard.get_metrics(&self.id, &key) {
57                    // Take only the most recent `limit` points
58                    let recent: Vec<_> = if points.len() > limit {
59                        points[points.len() - limit..].to_vec()
60                    } else {
61                        points
62                    };
63
64                    let snapshot = MetricSnapshot::from_points(&key, &recent);
65                    result.insert(key, snapshot);
66                }
67            }
68        }
69
70        result
71    }
72
73    fn subscribe(&self, _callback: SubscriptionCallback) {
74        // Subscriptions are handled at a higher level
75        // This is a placeholder for the trait requirement
76        // In practice, use a channel-based approach or web sockets for WASM
77    }
78
79    fn resource_usage(&self) -> ResourceSnapshot {
80        // Return simulated resource usage
81        // In production, this would query actual system metrics
82        ResourceSnapshot::new().with_cpu_util(0.0).with_gpu_util(0.0).with_memory(0, 0)
83    }
84}
85
86impl<S: ExperimentStorage> Run<S> {
87    /// Get a reference to the storage (for dashboard use).
88    pub(crate) fn storage_ref(&self) -> &Arc<Mutex<S>> {
89        &self.storage
90    }
91
92    /// Get all metric keys that have been logged.
93    pub fn metric_keys(&self) -> Vec<String> {
94        self.step_counters.keys().cloned().collect()
95    }
96}