Initial commit

demos/rtx-clusterviz-demo/src/metrics.rs

@@ -0,0 +1,702 @@
+//! Metrics Collection Module
+//!
+//! This module provides real-time metrics collection for cluster monitoring,
+//! including GPU utilization, memory usage, and bandwidth monitoring.
+
+use std::collections::HashMap;
+use std::time::{Duration, Instant};
+
+use anyhow::Result;
+use chrono::Utc;
+use clusterviz_shared::{ClusterMetrics, ClusterTopology, NodeStatus};
+use tracing::{debug, trace};
+use uuid::Uuid;
+
+/// Collects real-time metrics from the cluster
+pub struct MetricsCollector {
+    /// Utilization tracker
+    utilization_tracker: UtilizationTracker,
+    /// Memory tracker
+    memory_tracker: MemoryTracker,
+    /// Bandwidth monitor
+    bandwidth_monitor: BandwidthMonitor,
+    /// Collection interval
+    collection_interval: Duration,
+    /// Last collection time
+    last_collection: Option<Instant>,
+}
+
+impl MetricsCollector {
+    /// Create a new metrics collector
+    #[must_use]
+    pub fn new() -> Self {
+        Self {
+            utilization_tracker: UtilizationTracker::new(),
+            memory_tracker: MemoryTracker::new(),
+            bandwidth_monitor: BandwidthMonitor::new(),
+            collection_interval: Duration::from_millis(100),
+            last_collection: None,
+        }
+    }
+
+    /// Set the collection interval
+    #[must_use]
+    pub fn with_interval(mut self, interval: Duration) -> Self {
+        self.collection_interval = interval;
+        self
+    }
+
+    /// Collect metrics from the topology
+    pub async fn collect(&mut self, topology: &ClusterTopology) -> Result<ClusterMetrics> {
+        debug!("Collecting metrics for {} nodes", topology.nodes.len());
+
+        let mut metrics = ClusterMetrics::new();
+        metrics.timestamp = Utc::now();
+
+        // Collect GPU utilization for each node
+        for node in &topology.nodes {
+            let util = self.utilization_tracker.get_utilization(node.id).await?;
+            metrics.gpu_utilization.push((node.id, util));
+        }
+
+        // Collect memory usage for each node
+        for node in &topology.nodes {
+            let mem_used = self
+                .memory_tracker
+                .get_memory_used(node.id, node.memory_gb)
+                .await?;
+            metrics.memory_used.push((node.id, mem_used));
+        }
+
+        // Collect bandwidth usage for each link
+        for link in &topology.links {
+            let bandwidth = self.bandwidth_monitor.get_bandwidth_usage(link.id).await?;
+            metrics.bandwidth_usage.push((link.id, bandwidth));
+        }
+
+        // Calculate aggregate metrics
+        metrics.total_tflops = self.calculate_total_tflops(topology);
+        metrics.active_collectives = self.count_active_collectives();
+        metrics.data_transfer_rate_gb = self.calculate_transfer_rate(&metrics);
+
+        self.last_collection = Some(Instant::now());
+
+        trace!(
+            "Collected metrics: GPU util {:.1}%, Memory {:.1} GB",
+            metrics.average_gpu_utilization() * 100.0,
+            metrics.total_memory_used_gb()
+        );
+
+        Ok(metrics)
+    }
+
+    /// Calculate total TFLOPS across all nodes
+    fn calculate_total_tflops(&self, topology: &ClusterTopology) -> f64 {
+        // Estimate TFLOPS based on GPU model
+        // M3 Max: ~14 TFLOPS, M4 Ultra: ~28 TFLOPS (estimated)
+        topology
+            .nodes
+            .iter()
+            .filter(|n| n.status == NodeStatus::Healthy)
+            .map(|n| {
+                if n.gpu_model.contains("Ultra") {
+                    28.0
+                } else if n.gpu_model.contains("Max") {
+                    14.0
+                } else if n.gpu_model.contains("Pro") {
+                    8.0
+                } else {
+                    4.0
+                }
+            })
+            .sum()
+    }
+
+    /// Count active collective operations (simulated)
+    fn count_active_collectives(&self) -> u32 {
+        // In a real implementation, this would track actual collectives
+        // For simulation, return a random-ish value
+        (rand::random::<u32>() % 5) + 1
+    }
+
+    /// Calculate data transfer rate from bandwidth metrics
+    fn calculate_transfer_rate(&self, metrics: &ClusterMetrics) -> f64 {
+        let total_bandwidth: f64 = metrics.bandwidth_usage.iter().map(|(_, bw)| bw).sum();
+        // Convert Gbps to GB/s (divide by 8)
+        total_bandwidth / 8.0
+    }
+
+    /// Get the last collection time
+    #[must_use]
+    pub fn last_collection_time(&self) -> Option<Instant> {
+        self.last_collection
+    }
+}
+
+impl Default for MetricsCollector {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+/// Tracks GPU utilization across nodes
+pub struct UtilizationTracker {
+    /// Cached utilization values
+    cache: HashMap<Uuid, f64>,
+    /// Utilization history for trend analysis
+    history: HashMap<Uuid, Vec<f64>>,
+    /// Maximum history entries per node
+    max_history: usize,
+    /// Smoothing factor for exponential moving average
+    smoothing_factor: f64,
+}
+
+impl UtilizationTracker {
+    /// Create a new utilization tracker
+    #[must_use]
+    pub fn new() -> Self {
+        Self {
+            cache: HashMap::new(),
+            history: HashMap::new(),
+            max_history: 100,
+            smoothing_factor: 0.2,
+        }
+    }
+
+    /// Set the maximum history size
+    #[must_use]
+    pub fn with_max_history(mut self, max: usize) -> Self {
+        self.max_history = max;
+        self
+    }
+
+    /// Set the smoothing factor
+    #[must_use]
+    pub fn with_smoothing(mut self, factor: f64) -> Self {
+        self.smoothing_factor = factor.clamp(0.0, 1.0);
+        self
+    }
+
+    /// Get current utilization for a node
+    pub async fn get_utilization(&mut self, node_id: Uuid) -> Result<f64> {
+        // Simulate utilization reading
+        // In real implementation, would query the GPU driver
+        let raw_util = self.read_raw_utilization(node_id).await?;
+
+        // Apply exponential moving average smoothing
+        let smoothed = if let Some(&prev) = self.cache.get(&node_id) {
+            self.smoothing_factor * raw_util + (1.0 - self.smoothing_factor) * prev
+        } else {
+            raw_util
+        };
+
+        // Update cache and history
+        self.cache.insert(node_id, smoothed);
+        let history = self.history.entry(node_id).or_default();
+        history.push(smoothed);
+        if history.len() > self.max_history {
+            history.remove(0);
+        }
+
+        Ok(smoothed)
+    }
+
+    /// Read raw utilization (simulated)
+    async fn read_raw_utilization(&self, node_id: Uuid) -> Result<f64> {
+        // Simulate GPU utilization based on node ID for consistency
+        // In real implementation, would query IOKit or Metal framework
+        let base = (node_id.as_bytes()[0] as f64 / 255.0) * 0.3 + 0.5;
+        let noise = (rand::random::<f64>() - 0.5) * 0.1;
+        Ok((base + noise).clamp(0.0, 1.0))
+    }
+
+    /// Get utilization trend (positive = increasing, negative = decreasing)
+    #[must_use]
+    pub fn get_trend(&self, node_id: Uuid) -> Option<f64> {
+        let history = self.history.get(&node_id)?;
+        if history.len() < 2 {
+            return None;
+        }
+
+        let recent = &history[history.len().saturating_sub(10)..];
+        if recent.len() < 2 {
+            return None;
+        }
+
+        let first_half: f64 =
+            recent[..recent.len() / 2].iter().sum::<f64>() / (recent.len() / 2) as f64;
+        let second_half: f64 =
+            recent[recent.len() / 2..].iter().sum::<f64>() / ((recent.len() + 1) / 2) as f64;
+
+        Some(second_half - first_half)
+    }
+
+    /// Get average utilization over history
+    #[must_use]
+    pub fn get_average(&self, node_id: Uuid) -> Option<f64> {
+        let history = self.history.get(&node_id)?;
+        if history.is_empty() {
+            return None;
+        }
+        Some(history.iter().sum::<f64>() / history.len() as f64)
+    }
+
+    /// Clear cached data for a node
+    pub fn clear_node(&mut self, node_id: Uuid) {
+        self.cache.remove(&node_id);
+        self.history.remove(&node_id);
+    }
+
+    /// Clear all cached data
+    pub fn clear_all(&mut self) {
+        self.cache.clear();
+        self.history.clear();
+    }
+}
+
+impl Default for UtilizationTracker {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+/// Tracks memory usage across nodes
+pub struct MemoryTracker {
+    /// Cached memory usage values (in GB)
+    cache: HashMap<Uuid, f64>,
+    /// Memory usage history
+    history: HashMap<Uuid, Vec<f64>>,
+    /// Maximum history entries
+    max_history: usize,
+}
+
+impl MemoryTracker {
+    /// Create a new memory tracker
+    #[must_use]
+    pub fn new() -> Self {
+        Self {
+            cache: HashMap::new(),
+            history: HashMap::new(),
+            max_history: 100,
+        }
+    }
+
+    /// Set the maximum history size
+    #[must_use]
+    pub fn with_max_history(mut self, max: usize) -> Self {
+        self.max_history = max;
+        self
+    }
+
+    /// Get current memory used for a node
+    pub async fn get_memory_used(&mut self, node_id: Uuid, total_gb: u32) -> Result<f64> {
+        // Simulate memory reading
+        let raw_usage = self.read_raw_memory(node_id, total_gb).await?;
+
+        // Update cache and history
+        self.cache.insert(node_id, raw_usage);
+        let history = self.history.entry(node_id).or_default();
+        history.push(raw_usage);
+        if history.len() > self.max_history {
+            history.remove(0);
+        }
+
+        Ok(raw_usage)
+    }
+
+    /// Read raw memory usage (simulated)
+    async fn read_raw_memory(&self, node_id: Uuid, total_gb: u32) -> Result<f64> {
+        // Simulate memory usage based on node ID
+        let base_ratio = (node_id.as_bytes()[1] as f64 / 255.0) * 0.3 + 0.4;
+        let noise = (rand::random::<f64>() - 0.5) * 0.05;
+        let ratio = (base_ratio + noise).clamp(0.0, 1.0);
+        Ok(ratio * total_gb as f64)
+    }
+
+    /// Get peak memory usage from history
+    #[must_use]
+    pub fn get_peak(&self, node_id: Uuid) -> Option<f64> {
+        self.history.get(&node_id)?.iter().copied().reduce(f64::max)
+    }
+
+    /// Get memory usage trend
+    #[must_use]
+    pub fn get_trend(&self, node_id: Uuid) -> Option<f64> {
+        let history = self.history.get(&node_id)?;
+        if history.len() < 2 {
+            return None;
+        }
+
+        let recent = &history[history.len().saturating_sub(10)..];
+        if recent.len() < 2 {
+            return None;
+        }
+
+        let first = recent.first()?;
+        let last = recent.last()?;
+        Some(last - first)
+    }
+
+    /// Clear cached data
+    pub fn clear_all(&mut self) {
+        self.cache.clear();
+        self.history.clear();
+    }
+}
+
+impl Default for MemoryTracker {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+/// Monitors bandwidth usage across links
+pub struct BandwidthMonitor {
+    /// Current bandwidth readings per link (Gbps)
+    current_bandwidth: HashMap<Uuid, f64>,
+    /// Bandwidth history per link
+    history: HashMap<Uuid, Vec<f64>>,
+    /// Maximum history entries
+    max_history: usize,
+    /// Moving average window size
+    window_size: usize,
+}
+
+impl BandwidthMonitor {
+    /// Create a new bandwidth monitor
+    #[must_use]
+    pub fn new() -> Self {
+        Self {
+            current_bandwidth: HashMap::new(),
+            history: HashMap::new(),
+            max_history: 100,
+            window_size: 10,
+        }
+    }
+
+    /// Set the maximum history size
+    #[must_use]
+    pub fn with_max_history(mut self, max: usize) -> Self {
+        self.max_history = max;
+        self
+    }
+
+    /// Set the moving average window size
+    #[must_use]
+    pub fn with_window_size(mut self, size: usize) -> Self {
+        self.window_size = size.max(1);
+        self
+    }
+
+    /// Get current bandwidth usage for a link
+    pub async fn get_bandwidth_usage(&mut self, link_id: Uuid) -> Result<f64> {
+        let raw_bandwidth = self.read_raw_bandwidth(link_id).await?;
+
+        // Apply moving average
+        let history = self.history.entry(link_id).or_default();
+        history.push(raw_bandwidth);
+        if history.len() > self.max_history {
+            history.remove(0);
+        }
+
+        let window = &history[history.len().saturating_sub(self.window_size)..];
+        let smoothed = window.iter().sum::<f64>() / window.len() as f64;
+
+        self.current_bandwidth.insert(link_id, smoothed);
+        Ok(smoothed)
+    }
+
+    /// Read raw bandwidth (simulated)
+    async fn read_raw_bandwidth(&self, link_id: Uuid) -> Result<f64> {
+        // Simulate bandwidth usage
+        // Thunderbolt 5 max: 120 Gbps, simulate 40-80% utilization
+        let base = (link_id.as_bytes()[0] as f64 / 255.0) * 40.0 + 48.0;
+        let noise = (rand::random::<f64>() - 0.5) * 10.0;
+        Ok((base + noise).clamp(0.0, 120.0))
+    }
+
+    /// Get peak bandwidth for a link
+    #[must_use]
+    pub fn get_peak(&self, link_id: Uuid) -> Option<f64> {
+        self.history.get(&link_id)?.iter().copied().reduce(f64::max)
+    }
+
+    /// Get average bandwidth for a link
+    #[must_use]
+    pub fn get_average(&self, link_id: Uuid) -> Option<f64> {
+        let history = self.history.get(&link_id)?;
+        if history.is_empty() {
+            return None;
+        }
+        Some(history.iter().sum::<f64>() / history.len() as f64)
+    }
+
+    /// Calculate bandwidth efficiency (actual / theoretical max)
+    #[must_use]
+    pub fn get_efficiency(&self, link_id: Uuid, max_bandwidth: f64) -> Option<f64> {
+        let current = self.current_bandwidth.get(&link_id)?;
+        Some(current / max_bandwidth)
+    }
+
+    /// Get all link bandwidths
+    #[must_use]
+    pub fn get_all_bandwidths(&self) -> &HashMap<Uuid, f64> {
+        &self.current_bandwidth
+    }
+
+    /// Clear all data
+    pub fn clear_all(&mut self) {
+        self.current_bandwidth.clear();
+        self.history.clear();
+    }
+}
+
+impl Default for BandwidthMonitor {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+/// Aggregated cluster statistics
+#[derive(Debug, Clone)]
+pub struct ClusterStats {
+    /// Average GPU utilization
+    pub avg_gpu_util: f64,
+    /// Maximum GPU utilization
+    pub max_gpu_util: f64,
+    /// Total memory used (GB)
+    pub total_memory_used: f64,
+    /// Total memory capacity (GB)
+    pub total_memory_capacity: f64,
+    /// Average bandwidth utilization
+    pub avg_bandwidth_util: f64,
+    /// Total compute capacity (TFLOPS)
+    pub total_tflops: f64,
+    /// Number of healthy nodes
+    pub healthy_nodes: usize,
+    /// Total number of nodes
+    pub total_nodes: usize,
+}
+
+impl ClusterStats {
+    /// Calculate from topology and metrics
+    #[must_use]
+    pub fn from_topology_and_metrics(topology: &ClusterTopology, metrics: &ClusterMetrics) -> Self {
+        let healthy_nodes = topology
+            .nodes
+            .iter()
+            .filter(|n| n.status == NodeStatus::Healthy)
+            .count();
+
+        let avg_gpu_util = metrics.average_gpu_utilization();
+        let max_gpu_util = metrics
+            .gpu_utilization
+            .iter()
+            .map(|(_, u)| *u)
+            .fold(0.0, f64::max);
+
+        let total_memory_used = metrics.total_memory_used_gb();
+        let total_memory_capacity = topology.total_memory_gb() as f64;
+
+        let avg_bandwidth_util = metrics.average_bandwidth_usage();
+
+        Self {
+            avg_gpu_util,
+            max_gpu_util,
+            total_memory_used,
+            total_memory_capacity,
+            avg_bandwidth_util,
+            total_tflops: metrics.total_tflops,
+            healthy_nodes,
+            total_nodes: topology.nodes.len(),
+        }
+    }
+
+    /// Calculate memory utilization ratio
+    #[must_use]
+    pub fn memory_utilization(&self) -> f64 {
+        if self.total_memory_capacity > 0.0 {
+            self.total_memory_used / self.total_memory_capacity
+        } else {
+            0.0
+        }
+    }
+
+    /// Calculate cluster health score (0.0 - 1.0)
+    #[must_use]
+    pub fn health_score(&self) -> f64 {
+        if self.total_nodes == 0 {
+            return 0.0;
+        }
+
+        let node_health = self.healthy_nodes as f64 / self.total_nodes as f64;
+        let util_health = 1.0 - (self.max_gpu_util - 0.5).abs() * 0.5; // Optimal around 50-70%
+        let memory_health = 1.0 - self.memory_utilization().max(0.8).powi(2);
+
+        (node_health * 0.4 + util_health * 0.3 + memory_health * 0.3).clamp(0.0, 1.0)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use clusterviz_shared::{LinkInfo, NodeInfo, TransportType};
+
+    fn create_test_topology() -> ClusterTopology {
+        let mut topology = ClusterTopology::new("test-cluster");
+
+        let node1 = NodeInfo::new("node-1", "Apple M3 Max", 128);
+        let node2 = NodeInfo::new("node-2", "Apple M3 Max", 128);
+        let id1 = node1.id;
+        let id2 = node2.id;
+
+        topology.add_node(node1);
+        topology.add_node(node2);
+        topology.add_link(LinkInfo::new(id1, id2, TransportType::Thunderbolt));
+
+        topology
+    }
+
+    #[tokio::test]
+    async fn test_metrics_collector_creation() {
+        let collector = MetricsCollector::new();
+        assert!(collector.last_collection.is_none());
+    }
+
+    #[tokio::test]
+    async fn test_collect_metrics() {
+        let mut collector = MetricsCollector::new();
+        let topology = create_test_topology();
+
+        let metrics = collector.collect(&topology).await.unwrap();
+
+        assert_eq!(metrics.gpu_utilization.len(), 2);
+        assert_eq!(metrics.memory_used.len(), 2);
+        assert_eq!(metrics.bandwidth_usage.len(), 1);
+        assert!(collector.last_collection.is_some());
+    }
+
+    #[tokio::test]
+    async fn test_utilization_tracker() {
+        let mut tracker = UtilizationTracker::new();
+        let node_id = Uuid::new_v4();
+
+        let util = tracker.get_utilization(node_id).await.unwrap();
+
+        assert!(util >= 0.0 && util <= 1.0);
+        assert!(tracker.get_average(node_id).is_some());
+    }
+
+    #[tokio::test]
+    async fn test_utilization_smoothing() {
+        let mut tracker = UtilizationTracker::new().with_smoothing(0.5);
+        let node_id = Uuid::new_v4();
+
+        // Get multiple readings to test smoothing
+        for _ in 0..5 {
+            let _ = tracker.get_utilization(node_id).await.unwrap();
+        }
+
+        let trend = tracker.get_trend(node_id);
+        // Trend might be None if not enough data
+        assert!(trend.is_none() || trend.is_some());
+    }
+
+    #[tokio::test]
+    async fn test_memory_tracker() {
+        let mut tracker = MemoryTracker::new();
+        let node_id = Uuid::new_v4();
+
+        let mem = tracker.get_memory_used(node_id, 128).await.unwrap();
+
+        assert!(mem >= 0.0 && mem <= 128.0);
+    }
+
+    #[tokio::test]
+    async fn test_memory_peak() {
+        let mut tracker = MemoryTracker::new();
+        let node_id = Uuid::new_v4();
+
+        for _ in 0..5 {
+            let _ = tracker.get_memory_used(node_id, 128).await.unwrap();
+        }
+
+        let peak = tracker.get_peak(node_id);
+        assert!(peak.is_some());
+        assert!(peak.unwrap() <= 128.0);
+    }
+
+    #[tokio::test]
+    async fn test_bandwidth_monitor() {
+        let mut monitor = BandwidthMonitor::new();
+        let link_id = Uuid::new_v4();
+
+        let bandwidth = monitor.get_bandwidth_usage(link_id).await.unwrap();
+
+        assert!(bandwidth >= 0.0 && bandwidth <= 120.0);
+    }
+
+    #[tokio::test]
+    async fn test_bandwidth_efficiency() {
+        let mut monitor = BandwidthMonitor::new();
+        let link_id = Uuid::new_v4();
+
+        let _ = monitor.get_bandwidth_usage(link_id).await.unwrap();
+        let efficiency = monitor.get_efficiency(link_id, 120.0);
+
+        assert!(efficiency.is_some());
+        assert!(efficiency.unwrap() >= 0.0 && efficiency.unwrap() <= 1.0);
+    }
+
+    #[test]
+    fn test_cluster_stats() {
+        let topology = create_test_topology();
+        let metrics = clusterviz_shared::sample_metrics(&topology);
+
+        let stats = ClusterStats::from_topology_and_metrics(&topology, &metrics);
+
+        assert_eq!(stats.total_nodes, 2);
+        assert!(stats.avg_gpu_util >= 0.0 && stats.avg_gpu_util <= 1.0);
+        assert!(stats.health_score() >= 0.0 && stats.health_score() <= 1.0);
+    }
+
+    #[test]
+    fn test_cluster_stats_memory_utilization() {
+        let topology = create_test_topology();
+        let metrics = clusterviz_shared::sample_metrics(&topology);
+
+        let stats = ClusterStats::from_topology_and_metrics(&topology, &metrics);
+        let mem_util = stats.memory_utilization();
+
+        assert!(mem_util >= 0.0 && mem_util <= 1.0);
+    }
+
+    #[tokio::test]
+    async fn test_utilization_tracker_clear() {
+        let mut tracker = UtilizationTracker::new();
+        let node_id = Uuid::new_v4();
+
+        let _ = tracker.get_utilization(node_id).await.unwrap();
+        assert!(tracker.cache.contains_key(&node_id));
+
+        tracker.clear_node(node_id);
+        assert!(!tracker.cache.contains_key(&node_id));
+    }
+
+    #[tokio::test]
+    async fn test_bandwidth_moving_average() {
+        let mut monitor = BandwidthMonitor::new().with_window_size(5);
+        let link_id = Uuid::new_v4();
+
+        let mut readings = Vec::new();
+        for _ in 0..10 {
+            let bw = monitor.get_bandwidth_usage(link_id).await.unwrap();
+            readings.push(bw);
+        }
+
+        // Later readings should be smoother due to moving average
+        // Just verify we got valid readings
+        assert_eq!(readings.len(), 10);
+        assert!(readings.iter().all(|&r| r >= 0.0 && r <= 120.0));
+    }
+}