Initial commit
This commit is contained in:
redclawsystems
@@ -0,0 +1,586 @@
|
||||
//! ClusterViz - Real-Time Cluster Monitor
|
||||
//!
|
||||
//! This crate provides a complete system for monitoring and visualizing
|
||||
//! Thunderbolt 5 Mac cluster topology and data flow in real-time.
|
||||
//!
|
||||
//! # Features
|
||||
//!
|
||||
//! - **Topology Discovery**: Automatic detection of cluster nodes and links
|
||||
//! - **Real-Time Metrics**: GPU utilization, memory usage, bandwidth monitoring
|
||||
//! - **Bottleneck Detection**: Identify performance issues in real-time
|
||||
//! - **Collective Tracing**: Visualize distributed operations (AllReduce, Broadcast, etc.)
|
||||
//!
|
||||
//! # Example
|
||||
//!
|
||||
//! ```rust,no_run
|
||||
//! use rtx_clusterviz_demo::ClusterViz;
|
||||
//!
|
||||
//! #[tokio::main]
|
||||
//! async fn main() -> anyhow::Result<()> {
|
||||
//! let mut clusterviz = ClusterViz::new();
|
||||
//!
|
||||
//! // Discover cluster topology
|
||||
//! clusterviz.discover_topology().await?;
|
||||
//!
|
||||
//! // Collect metrics
|
||||
//! let metrics = clusterviz.collect_metrics().await?;
|
||||
//! println!("Average GPU utilization: {:.1}%", metrics.average_gpu_utilization() * 100.0);
|
||||
//!
|
||||
//! // Check for bottlenecks
|
||||
//! let bottlenecks = clusterviz.detect_bottlenecks().await?;
|
||||
//! for bn in bottlenecks {
|
||||
//! println!("Bottleneck: {} - {}", bn.bottleneck_type, bn.description);
|
||||
//! }
|
||||
//!
|
||||
//! Ok(())
|
||||
//! }
|
||||
//! ```
|
||||
|
||||
pub mod metrics;
|
||||
pub mod sample_data;
|
||||
pub mod topology;
|
||||
|
||||
use std::sync::Arc;
|
||||
use std::time::Duration;
|
||||
|
||||
use anyhow::{Context, Result};
|
||||
use clusterviz_shared::{
|
||||
AlertConfig, BottleneckInfo, BottleneckType, ClusterMetrics, ClusterTopology, CollectiveTrace,
|
||||
Severity,
|
||||
};
|
||||
use tokio::sync::RwLock;
|
||||
use tracing::{debug, info, warn};
|
||||
|
||||
use crate::metrics::MetricsCollector;
|
||||
use crate::topology::TopologyDiscoverer;
|
||||
|
||||
/// Error types specific to ClusterViz
|
||||
#[derive(Debug, thiserror::Error)]
|
||||
pub enum ClusterVizError {
|
||||
/// Topology discovery failed
|
||||
#[error("Failed to discover cluster topology: {0}")]
|
||||
DiscoveryFailed(String),
|
||||
|
||||
/// Metrics collection failed
|
||||
#[error("Failed to collect metrics: {0}")]
|
||||
MetricsCollectionFailed(String),
|
||||
|
||||
/// No topology available
|
||||
#[error("No cluster topology available - run discover_topology() first")]
|
||||
NoTopology,
|
||||
|
||||
/// Node not found
|
||||
#[error("Node not found: {0}")]
|
||||
NodeNotFound(String),
|
||||
|
||||
/// Link not found
|
||||
#[error("Link not found: {0}")]
|
||||
LinkNotFound(String),
|
||||
|
||||
/// Configuration error
|
||||
#[error("Configuration error: {0}")]
|
||||
ConfigError(String),
|
||||
}
|
||||
|
||||
/// Main ClusterViz system for real-time cluster monitoring
|
||||
pub struct ClusterViz {
|
||||
/// Current cluster topology
|
||||
topology: Option<Arc<RwLock<ClusterTopology>>>,
|
||||
/// Metrics collector
|
||||
metrics_collector: MetricsCollector,
|
||||
/// Topology discoverer
|
||||
topology_discoverer: TopologyDiscoverer,
|
||||
/// Alert configuration
|
||||
alert_config: AlertConfig,
|
||||
/// Historical metrics (for trend analysis)
|
||||
metrics_history: Vec<ClusterMetrics>,
|
||||
/// Maximum history size
|
||||
max_history_size: usize,
|
||||
/// Active collective traces
|
||||
active_traces: Vec<CollectiveTrace>,
|
||||
/// Detected bottlenecks
|
||||
detected_bottlenecks: Vec<BottleneckInfo>,
|
||||
}
|
||||
|
||||
impl ClusterViz {
|
||||
/// Create a new ClusterViz instance
|
||||
#[must_use]
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
topology: None,
|
||||
metrics_collector: MetricsCollector::new(),
|
||||
topology_discoverer: TopologyDiscoverer::new(),
|
||||
alert_config: AlertConfig::new(),
|
||||
metrics_history: Vec::new(),
|
||||
max_history_size: 1000,
|
||||
active_traces: Vec::new(),
|
||||
detected_bottlenecks: Vec::new(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Create with custom alert configuration
|
||||
#[must_use]
|
||||
pub fn with_alert_config(mut self, config: AlertConfig) -> Self {
|
||||
self.alert_config = config;
|
||||
self
|
||||
}
|
||||
|
||||
/// Set the maximum history size
|
||||
#[must_use]
|
||||
pub fn with_max_history(mut self, max_size: usize) -> Self {
|
||||
self.max_history_size = max_size;
|
||||
self
|
||||
}
|
||||
|
||||
/// Discover the cluster topology automatically
|
||||
///
|
||||
/// This method probes the network to find all available nodes
|
||||
/// and their interconnections. Returns a clone of the discovered topology.
|
||||
pub async fn discover_topology(&mut self) -> Result<ClusterTopology> {
|
||||
info!("Starting cluster topology discovery");
|
||||
|
||||
let topology = self
|
||||
.topology_discoverer
|
||||
.discover()
|
||||
.await
|
||||
.context("Failed to discover topology")?;
|
||||
|
||||
info!(
|
||||
"Discovered {} nodes and {} links",
|
||||
topology.node_count(),
|
||||
topology.link_count()
|
||||
);
|
||||
|
||||
let topology_clone = topology.clone();
|
||||
self.topology = Some(Arc::new(RwLock::new(topology)));
|
||||
|
||||
Ok(topology_clone)
|
||||
}
|
||||
|
||||
/// Set the topology manually (for testing or pre-configured clusters)
|
||||
pub fn set_topology(&mut self, topology: ClusterTopology) {
|
||||
self.topology = Some(Arc::new(RwLock::new(topology)));
|
||||
}
|
||||
|
||||
/// Get a reference to the current topology
|
||||
#[must_use]
|
||||
pub fn get_topology(&self) -> Option<&Arc<RwLock<ClusterTopology>>> {
|
||||
self.topology.as_ref()
|
||||
}
|
||||
|
||||
/// Collect real-time metrics from the cluster
|
||||
pub async fn collect_metrics(&mut self) -> Result<ClusterMetrics> {
|
||||
let topology = self.topology.as_ref().ok_or(ClusterVizError::NoTopology)?;
|
||||
|
||||
let topology_read = topology.read().await;
|
||||
|
||||
debug!(
|
||||
"Collecting metrics for {} nodes",
|
||||
topology_read.node_count()
|
||||
);
|
||||
|
||||
let metrics = self
|
||||
.metrics_collector
|
||||
.collect(&topology_read)
|
||||
.await
|
||||
.context("Failed to collect metrics")?;
|
||||
|
||||
// Store in history
|
||||
self.metrics_history.push(metrics.clone());
|
||||
if self.metrics_history.len() > self.max_history_size {
|
||||
self.metrics_history.remove(0);
|
||||
}
|
||||
|
||||
Ok(metrics)
|
||||
}
|
||||
|
||||
/// Detect bottlenecks in the cluster
|
||||
pub async fn detect_bottlenecks(&mut self) -> Result<Vec<BottleneckInfo>> {
|
||||
// Collect metrics first (requires mutable borrow)
|
||||
let metrics = self.collect_metrics().await?;
|
||||
|
||||
// Then get topology for analysis
|
||||
let topology = self.topology.as_ref().ok_or(ClusterVizError::NoTopology)?;
|
||||
|
||||
let topology_read = topology.read().await;
|
||||
|
||||
let mut bottlenecks = Vec::new();
|
||||
|
||||
// Check GPU utilization
|
||||
for (node_id, util) in &metrics.gpu_utilization {
|
||||
if *util >= self.alert_config.gpu_util_critical {
|
||||
let mut bn = BottleneckInfo::new(
|
||||
BottleneckType::ComputeBottleneck,
|
||||
Severity::Critical,
|
||||
format!("GPU utilization at {:.1}%", util * 100.0),
|
||||
)
|
||||
.with_metrics(*util, self.alert_config.gpu_util_critical);
|
||||
|
||||
bn.add_affected_node(*node_id);
|
||||
bn.remediation =
|
||||
Some("Consider redistributing workload or adding more nodes".to_string());
|
||||
bottlenecks.push(bn);
|
||||
} else if *util >= self.alert_config.gpu_util_warning {
|
||||
let mut bn = BottleneckInfo::new(
|
||||
BottleneckType::ComputeBottleneck,
|
||||
Severity::Warning,
|
||||
format!("GPU utilization at {:.1}%", util * 100.0),
|
||||
)
|
||||
.with_metrics(*util, self.alert_config.gpu_util_warning);
|
||||
|
||||
bn.add_affected_node(*node_id);
|
||||
bottlenecks.push(bn);
|
||||
}
|
||||
}
|
||||
|
||||
// Check memory usage
|
||||
for (node_id, mem_used) in &metrics.memory_used {
|
||||
if let Some(node) = topology_read.get_node(*node_id) {
|
||||
let usage_ratio = mem_used / node.memory_gb as f64;
|
||||
if usage_ratio >= self.alert_config.memory_critical {
|
||||
let mut bn = BottleneckInfo::new(
|
||||
BottleneckType::MemoryPressure,
|
||||
Severity::Critical,
|
||||
format!("Memory usage at {:.1}%", usage_ratio * 100.0),
|
||||
)
|
||||
.with_metrics(usage_ratio, self.alert_config.memory_critical);
|
||||
|
||||
bn.add_affected_node(*node_id);
|
||||
bn.remediation =
|
||||
Some("Reduce batch size or enable gradient checkpointing".to_string());
|
||||
bottlenecks.push(bn);
|
||||
} else if usage_ratio >= self.alert_config.memory_warning {
|
||||
let mut bn = BottleneckInfo::new(
|
||||
BottleneckType::MemoryPressure,
|
||||
Severity::Warning,
|
||||
format!("Memory usage at {:.1}%", usage_ratio * 100.0),
|
||||
)
|
||||
.with_metrics(usage_ratio, self.alert_config.memory_warning);
|
||||
|
||||
bn.add_affected_node(*node_id);
|
||||
bottlenecks.push(bn);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Check bandwidth saturation
|
||||
for (link_id, bandwidth) in &metrics.bandwidth_usage {
|
||||
if let Some(link) = topology_read.links.iter().find(|l| l.id == *link_id) {
|
||||
let saturation = bandwidth / link.transport_type.max_bandwidth_gbps();
|
||||
if saturation >= self.alert_config.bandwidth_critical {
|
||||
let mut bn = BottleneckInfo::new(
|
||||
BottleneckType::BandwidthSaturation,
|
||||
Severity::Critical,
|
||||
format!("Link bandwidth at {:.1}%", saturation * 100.0),
|
||||
)
|
||||
.with_metrics(saturation, self.alert_config.bandwidth_critical);
|
||||
|
||||
bn.add_affected_link(*link_id);
|
||||
bn.remediation =
|
||||
Some("Add parallel links or optimize communication patterns".to_string());
|
||||
bottlenecks.push(bn);
|
||||
} else if saturation >= self.alert_config.bandwidth_warning {
|
||||
let mut bn = BottleneckInfo::new(
|
||||
BottleneckType::BandwidthSaturation,
|
||||
Severity::Warning,
|
||||
format!("Link bandwidth at {:.1}%", saturation * 100.0),
|
||||
)
|
||||
.with_metrics(saturation, self.alert_config.bandwidth_warning);
|
||||
|
||||
bn.add_affected_link(*link_id);
|
||||
bottlenecks.push(bn);
|
||||
}
|
||||
|
||||
// Check latency
|
||||
if link.latency_us >= self.alert_config.latency_critical_us {
|
||||
let mut bn = BottleneckInfo::new(
|
||||
BottleneckType::HighLatency,
|
||||
Severity::Critical,
|
||||
format!("Link latency at {:.1}us", link.latency_us),
|
||||
)
|
||||
.with_metrics(link.latency_us, self.alert_config.latency_critical_us);
|
||||
|
||||
bn.add_affected_link(*link_id);
|
||||
bottlenecks.push(bn);
|
||||
} else if link.latency_us >= self.alert_config.latency_warning_us {
|
||||
let mut bn = BottleneckInfo::new(
|
||||
BottleneckType::HighLatency,
|
||||
Severity::Warning,
|
||||
format!("Link latency at {:.1}us", link.latency_us),
|
||||
)
|
||||
.with_metrics(link.latency_us, self.alert_config.latency_warning_us);
|
||||
|
||||
bn.add_affected_link(*link_id);
|
||||
bottlenecks.push(bn);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Check load imbalance
|
||||
if !metrics.gpu_utilization.is_empty() {
|
||||
let utils: Vec<f64> = metrics.gpu_utilization.iter().map(|(_, u)| *u).collect();
|
||||
let mean: f64 = utils.iter().sum::<f64>() / utils.len() as f64;
|
||||
let variance: f64 =
|
||||
utils.iter().map(|u| (u - mean).powi(2)).sum::<f64>() / utils.len() as f64;
|
||||
let std_dev = variance.sqrt();
|
||||
let coefficient_of_variation = if mean > 0.0 { std_dev / mean } else { 0.0 };
|
||||
|
||||
if coefficient_of_variation >= self.alert_config.load_imbalance_warning {
|
||||
let mut bn = BottleneckInfo::new(
|
||||
BottleneckType::LoadImbalance,
|
||||
Severity::Warning,
|
||||
format!(
|
||||
"Load imbalance detected (CV: {:.2})",
|
||||
coefficient_of_variation
|
||||
),
|
||||
)
|
||||
.with_metrics(
|
||||
coefficient_of_variation,
|
||||
self.alert_config.load_imbalance_warning,
|
||||
);
|
||||
|
||||
bn.remediation =
|
||||
Some("Review data partitioning and workload distribution".to_string());
|
||||
bottlenecks.push(bn);
|
||||
}
|
||||
}
|
||||
|
||||
// Store detected bottlenecks
|
||||
self.detected_bottlenecks = bottlenecks.clone();
|
||||
|
||||
if !bottlenecks.is_empty() {
|
||||
warn!("Detected {} bottlenecks", bottlenecks.len());
|
||||
}
|
||||
|
||||
Ok(bottlenecks)
|
||||
}
|
||||
|
||||
/// Add an active collective trace
|
||||
pub fn add_trace(&mut self, trace: CollectiveTrace) {
|
||||
self.active_traces.push(trace);
|
||||
}
|
||||
|
||||
/// Get active traces
|
||||
#[must_use]
|
||||
pub fn get_active_traces(&self) -> &[CollectiveTrace] {
|
||||
&self.active_traces
|
||||
}
|
||||
|
||||
/// Clear completed traces
|
||||
pub fn clear_traces(&mut self) {
|
||||
self.active_traces.clear();
|
||||
}
|
||||
|
||||
/// Get metrics history
|
||||
#[must_use]
|
||||
pub fn get_metrics_history(&self) -> &[ClusterMetrics] {
|
||||
&self.metrics_history
|
||||
}
|
||||
|
||||
/// Get detected bottlenecks
|
||||
#[must_use]
|
||||
pub fn get_bottlenecks(&self) -> &[BottleneckInfo] {
|
||||
&self.detected_bottlenecks
|
||||
}
|
||||
|
||||
/// Get current alert configuration
|
||||
#[must_use]
|
||||
pub fn get_alert_config(&self) -> &AlertConfig {
|
||||
&self.alert_config
|
||||
}
|
||||
|
||||
/// Update alert configuration
|
||||
pub fn set_alert_config(&mut self, config: AlertConfig) {
|
||||
self.alert_config = config;
|
||||
}
|
||||
|
||||
/// Run the demo with sample data
|
||||
pub async fn run_demo(&mut self) -> Result<()> {
|
||||
info!("Starting ClusterViz Demo");
|
||||
|
||||
// Use sample topology
|
||||
let topology = sample_data::four_node_cluster();
|
||||
self.set_topology(topology);
|
||||
|
||||
info!("Loaded 4-node cluster topology");
|
||||
|
||||
// Collect metrics
|
||||
let metrics = self.collect_metrics().await?;
|
||||
info!(
|
||||
"Collected metrics - Average GPU utilization: {:.1}%",
|
||||
metrics.average_gpu_utilization() * 100.0
|
||||
);
|
||||
info!(
|
||||
"Total memory used: {:.1} GB",
|
||||
metrics.total_memory_used_gb()
|
||||
);
|
||||
|
||||
// Add sample traces
|
||||
let node_ids: Vec<_> = {
|
||||
let topology = self.topology.as_ref().unwrap().read().await;
|
||||
topology.nodes.iter().map(|n| n.id).collect()
|
||||
};
|
||||
|
||||
let allreduce_trace = clusterviz_shared::sample_allreduce_trace(&node_ids);
|
||||
self.add_trace(allreduce_trace);
|
||||
|
||||
let broadcast_trace = clusterviz_shared::sample_broadcast_trace(&node_ids, 0);
|
||||
self.add_trace(broadcast_trace);
|
||||
|
||||
info!("Added {} collective traces", self.active_traces.len());
|
||||
|
||||
// Detect bottlenecks
|
||||
let bottlenecks = self.detect_bottlenecks().await?;
|
||||
if bottlenecks.is_empty() {
|
||||
info!("No bottlenecks detected");
|
||||
} else {
|
||||
for bn in &bottlenecks {
|
||||
warn!(
|
||||
"Bottleneck: {} - {} (severity: {})",
|
||||
bn.bottleneck_type, bn.description, bn.severity
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
// Run monitoring loop (3 iterations for demo)
|
||||
for i in 1..=3 {
|
||||
tokio::time::sleep(Duration::from_millis(500)).await;
|
||||
|
||||
let metrics = self.collect_metrics().await?;
|
||||
debug!(
|
||||
"Iteration {}: GPU util {:.1}%, Memory {:.1} GB",
|
||||
i,
|
||||
metrics.average_gpu_utilization() * 100.0,
|
||||
metrics.total_memory_used_gb()
|
||||
);
|
||||
}
|
||||
|
||||
info!("ClusterViz Demo completed successfully");
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for ClusterViz {
|
||||
fn default() -> Self {
|
||||
Self::new()
|
||||
}
|
||||
}
|
||||
|
||||
/// Run the demo
|
||||
pub async fn run_demo() -> Result<()> {
|
||||
let mut clusterviz = ClusterViz::new();
|
||||
clusterviz.run_demo().await
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_clusterviz_creation() {
|
||||
let clusterviz = ClusterViz::new();
|
||||
assert!(clusterviz.topology.is_none());
|
||||
assert!(clusterviz.metrics_history.is_empty());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_set_topology() {
|
||||
let mut clusterviz = ClusterViz::new();
|
||||
let topology = sample_data::four_node_cluster();
|
||||
|
||||
clusterviz.set_topology(topology);
|
||||
|
||||
assert!(clusterviz.topology.is_some());
|
||||
let topo = clusterviz.topology.as_ref().unwrap().read().await;
|
||||
assert_eq!(topo.node_count(), 4);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_collect_metrics() {
|
||||
let mut clusterviz = ClusterViz::new();
|
||||
let topology = sample_data::four_node_cluster();
|
||||
clusterviz.set_topology(topology);
|
||||
|
||||
let metrics = clusterviz.collect_metrics().await.unwrap();
|
||||
|
||||
assert!(!metrics.gpu_utilization.is_empty());
|
||||
assert!(!metrics.memory_used.is_empty());
|
||||
assert_eq!(clusterviz.metrics_history.len(), 1);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_collect_metrics_no_topology() {
|
||||
let mut clusterviz = ClusterViz::new();
|
||||
let result = clusterviz.collect_metrics().await;
|
||||
|
||||
assert!(result.is_err());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_detect_bottlenecks() {
|
||||
let mut clusterviz = ClusterViz::new();
|
||||
let topology = sample_data::four_node_cluster();
|
||||
clusterviz.set_topology(topology);
|
||||
|
||||
let bottlenecks = clusterviz.detect_bottlenecks().await.unwrap();
|
||||
|
||||
// With default sample data, we might or might not have bottlenecks
|
||||
// depending on the random utilization values
|
||||
assert!(clusterviz.detected_bottlenecks.len() == bottlenecks.len());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_alert_config() {
|
||||
let clusterviz = ClusterViz::new().with_alert_config(AlertConfig::strict());
|
||||
|
||||
assert!(clusterviz.alert_config.gpu_util_warning < 0.80);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_add_trace() {
|
||||
let mut clusterviz = ClusterViz::new();
|
||||
let nodes = vec![uuid::Uuid::new_v4(), uuid::Uuid::new_v4()];
|
||||
let trace = clusterviz_shared::sample_allreduce_trace(&nodes);
|
||||
|
||||
clusterviz.add_trace(trace);
|
||||
|
||||
assert_eq!(clusterviz.active_traces.len(), 1);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_clear_traces() {
|
||||
let mut clusterviz = ClusterViz::new();
|
||||
let nodes = vec![uuid::Uuid::new_v4(), uuid::Uuid::new_v4()];
|
||||
let trace = clusterviz_shared::sample_allreduce_trace(&nodes);
|
||||
|
||||
clusterviz.add_trace(trace);
|
||||
assert_eq!(clusterviz.active_traces.len(), 1);
|
||||
|
||||
clusterviz.clear_traces();
|
||||
assert!(clusterviz.active_traces.is_empty());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_max_history() {
|
||||
let mut clusterviz = ClusterViz::new().with_max_history(3);
|
||||
let topology = sample_data::four_node_cluster();
|
||||
clusterviz.set_topology(topology);
|
||||
|
||||
// Collect more metrics than history size
|
||||
for _ in 0..5 {
|
||||
let _ = clusterviz.collect_metrics().await.unwrap();
|
||||
}
|
||||
|
||||
assert_eq!(clusterviz.metrics_history.len(), 3);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_run_demo() {
|
||||
let mut clusterviz = ClusterViz::new();
|
||||
let result = clusterviz.run_demo().await;
|
||||
|
||||
assert!(result.is_ok());
|
||||
assert!(clusterviz.topology.is_some());
|
||||
assert!(!clusterviz.metrics_history.is_empty());
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user