Files
rustytorch/demos/clusterviz-shared/src/lib.rs
T
2026-03-04 00:08:42 +00:00

1068 lines
32 KiB
Rust

//! ClusterViz Shared IPC Types
//!
//! This crate defines the shared types for inter-process communication
//! between the ClusterViz real-time cluster monitor components.
//! Designed for visualizing Thunderbolt 5 Mac cluster topology and data flow.
use chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize};
use uuid::Uuid;
/// Transport type for inter-node communication
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub enum TransportType {
/// Direct Metal GPU-to-GPU transfer (fastest, same node)
MetalDirect,
/// Metal Unified Memory Architecture (shared memory, same SoC)
MetalUMA,
/// Thunderbolt 5 interconnect (120 Gbps bidirectional)
Thunderbolt,
/// Shared memory between processes (same machine)
SharedMem,
}
impl TransportType {
/// Returns the theoretical maximum bandwidth in Gbps
#[must_use]
pub fn max_bandwidth_gbps(&self) -> f64 {
match self {
Self::MetalDirect => 800.0, // M3 Max memory bandwidth
Self::MetalUMA => 400.0, // Unified memory access
Self::Thunderbolt => 120.0, // Thunderbolt 5 bidirectional
Self::SharedMem => 200.0, // PCIe-like shared memory
}
}
/// Returns the typical latency in microseconds
#[must_use]
pub fn typical_latency_us(&self) -> f64 {
match self {
Self::MetalDirect => 0.5,
Self::MetalUMA => 1.0,
Self::Thunderbolt => 2.5,
Self::SharedMem => 1.5,
}
}
}
impl std::fmt::Display for TransportType {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Self::MetalDirect => write!(f, "Metal Direct"),
Self::MetalUMA => write!(f, "Metal UMA"),
Self::Thunderbolt => write!(f, "Thunderbolt 5"),
Self::SharedMem => write!(f, "Shared Memory"),
}
}
}
/// Information about a cluster node (GPU-equipped machine)
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct NodeInfo {
/// Unique node identifier
pub id: Uuid,
/// Human-readable node name
pub name: String,
/// GPU model (e.g., "Apple M3 Ultra", "Apple M4 Max")
pub gpu_model: String,
/// Total GPU memory in gigabytes
pub memory_gb: u32,
/// Current GPU utilization (0.0 - 1.0)
pub utilization: f64,
/// Whether this node is the cluster coordinator
pub is_coordinator: bool,
/// Node status
pub status: NodeStatus,
/// Timestamp of last heartbeat
pub last_heartbeat: DateTime<Utc>,
}
impl NodeInfo {
/// Create a new node info
#[must_use]
pub fn new(name: impl Into<String>, gpu_model: impl Into<String>, memory_gb: u32) -> Self {
Self {
id: Uuid::new_v4(),
name: name.into(),
gpu_model: gpu_model.into(),
memory_gb,
utilization: 0.0,
is_coordinator: false,
status: NodeStatus::Healthy,
last_heartbeat: Utc::now(),
}
}
/// Set this node as the coordinator
#[must_use]
pub fn as_coordinator(mut self) -> Self {
self.is_coordinator = true;
self
}
/// Update utilization
pub fn set_utilization(&mut self, utilization: f64) {
self.utilization = utilization.clamp(0.0, 1.0);
}
}
/// Node operational status
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum NodeStatus {
/// Node is healthy and responsive
Healthy,
/// Node is experiencing high load
HighLoad,
/// Node is experiencing issues
Degraded,
/// Node is unreachable
Unreachable,
/// Node is performing maintenance
Maintenance,
}
impl std::fmt::Display for NodeStatus {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Self::Healthy => write!(f, "Healthy"),
Self::HighLoad => write!(f, "High Load"),
Self::Degraded => write!(f, "Degraded"),
Self::Unreachable => write!(f, "Unreachable"),
Self::Maintenance => write!(f, "Maintenance"),
}
}
}
/// Information about a link between two nodes
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct LinkInfo {
/// Unique link identifier
pub id: Uuid,
/// Source node ID
pub source: Uuid,
/// Target node ID
pub target: Uuid,
/// Current bandwidth in Gbps
pub bandwidth_gbps: f64,
/// Current latency in microseconds
pub latency_us: f64,
/// Transport type for this link
pub transport_type: TransportType,
/// Link health status (0.0 - 1.0, higher is better)
pub health: f64,
/// Whether the link is currently active (transferring data)
pub is_active: bool,
}
impl LinkInfo {
/// Create a new link info
#[must_use]
pub fn new(source: Uuid, target: Uuid, transport_type: TransportType) -> Self {
Self {
id: Uuid::new_v4(),
source,
target,
bandwidth_gbps: transport_type.max_bandwidth_gbps() * 0.8, // 80% efficiency
latency_us: transport_type.typical_latency_us(),
transport_type,
health: 1.0,
is_active: false,
}
}
/// Calculate bandwidth efficiency (actual vs theoretical)
#[must_use]
pub fn bandwidth_efficiency(&self) -> f64 {
self.bandwidth_gbps / self.transport_type.max_bandwidth_gbps()
}
/// Mark the link as active
pub fn set_active(&mut self, active: bool) {
self.is_active = active;
}
}
/// Complete cluster topology with nodes and links
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ClusterTopology {
/// Cluster name
pub name: String,
/// All nodes in the cluster
pub nodes: Vec<NodeInfo>,
/// All links between nodes
pub links: Vec<LinkInfo>,
/// Timestamp of topology discovery
pub discovered_at: DateTime<Utc>,
/// Topology version (incremented on changes)
pub version: u64,
}
impl ClusterTopology {
/// Create a new empty topology
#[must_use]
pub fn new(name: impl Into<String>) -> Self {
Self {
name: name.into(),
nodes: Vec::new(),
links: Vec::new(),
discovered_at: Utc::now(),
version: 1,
}
}
/// Add a node to the topology
pub fn add_node(&mut self, node: NodeInfo) {
self.nodes.push(node);
self.version += 1;
}
/// Add a link to the topology
pub fn add_link(&mut self, link: LinkInfo) {
self.links.push(link);
self.version += 1;
}
/// Get a node by ID
#[must_use]
pub fn get_node(&self, id: Uuid) -> Option<&NodeInfo> {
self.nodes.iter().find(|n| n.id == id)
}
/// Get a mutable reference to a node by ID
#[must_use]
pub fn get_node_mut(&mut self, id: Uuid) -> Option<&mut NodeInfo> {
self.nodes.iter_mut().find(|n| n.id == id)
}
/// Get all links connected to a node
#[must_use]
pub fn get_node_links(&self, node_id: Uuid) -> Vec<&LinkInfo> {
self.links
.iter()
.filter(|l| l.source == node_id || l.target == node_id)
.collect()
}
/// Total number of nodes
#[must_use]
pub fn node_count(&self) -> usize {
self.nodes.len()
}
/// Total number of links
#[must_use]
pub fn link_count(&self) -> usize {
self.links.len()
}
/// Calculate total cluster memory in GB
#[must_use]
pub fn total_memory_gb(&self) -> u32 {
self.nodes.iter().map(|n| n.memory_gb).sum()
}
/// Calculate average GPU utilization across all nodes
#[must_use]
pub fn average_utilization(&self) -> f64 {
if self.nodes.is_empty() {
return 0.0;
}
self.nodes.iter().map(|n| n.utilization).sum::<f64>() / self.nodes.len() as f64
}
}
/// Real-time cluster metrics
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ClusterMetrics {
/// Timestamp of metrics collection
pub timestamp: DateTime<Utc>,
/// Per-node GPU utilization (node_id -> utilization)
pub gpu_utilization: Vec<(Uuid, f64)>,
/// Per-node memory used in GB (node_id -> memory_used_gb)
pub memory_used: Vec<(Uuid, f64)>,
/// Per-link bandwidth usage in Gbps (link_id -> bandwidth)
pub bandwidth_usage: Vec<(Uuid, f64)>,
/// Total cluster TFLOPS
pub total_tflops: f64,
/// Active collective operations count
pub active_collectives: u32,
/// Data transferred in the last second (GB)
pub data_transfer_rate_gb: f64,
}
impl ClusterMetrics {
/// Create empty metrics
#[must_use]
pub fn new() -> Self {
Self {
timestamp: Utc::now(),
gpu_utilization: Vec::new(),
memory_used: Vec::new(),
bandwidth_usage: Vec::new(),
total_tflops: 0.0,
active_collectives: 0,
data_transfer_rate_gb: 0.0,
}
}
/// Average GPU utilization across all nodes
#[must_use]
pub fn average_gpu_utilization(&self) -> f64 {
if self.gpu_utilization.is_empty() {
return 0.0;
}
self.gpu_utilization.iter().map(|(_, u)| u).sum::<f64>() / self.gpu_utilization.len() as f64
}
/// Total memory used across all nodes
#[must_use]
pub fn total_memory_used_gb(&self) -> f64 {
self.memory_used.iter().map(|(_, m)| m).sum()
}
/// Average bandwidth usage across all links
#[must_use]
pub fn average_bandwidth_usage(&self) -> f64 {
if self.bandwidth_usage.is_empty() {
return 0.0;
}
self.bandwidth_usage.iter().map(|(_, b)| b).sum::<f64>() / self.bandwidth_usage.len() as f64
}
}
impl Default for ClusterMetrics {
fn default() -> Self {
Self::new()
}
}
/// Collective operation types
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub enum CollectiveOp {
/// All-reduce: combine and distribute to all
AllReduce,
/// All-gather: gather from all to all
AllGather,
/// Broadcast: one to all
Broadcast,
/// Scatter: one to many (different chunks)
Scatter,
/// Reduce: combine to one
Reduce,
/// Reduce-scatter: reduce then scatter
ReduceScatter,
/// All-to-all: complete exchange
AllToAll,
/// Ring topology transfer
Ring,
}
impl std::fmt::Display for CollectiveOp {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Self::AllReduce => write!(f, "AllReduce"),
Self::AllGather => write!(f, "AllGather"),
Self::Broadcast => write!(f, "Broadcast"),
Self::Scatter => write!(f, "Scatter"),
Self::Reduce => write!(f, "Reduce"),
Self::ReduceScatter => write!(f, "ReduceScatter"),
Self::AllToAll => write!(f, "All-to-All"),
Self::Ring => write!(f, "Ring"),
}
}
}
/// Data flow step in a collective operation
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct DataFlowStep {
/// Source node ID
pub from: Uuid,
/// Target node ID
pub to: Uuid,
/// Data size in bytes
pub size_bytes: u64,
/// Duration of this step in microseconds
pub duration_us: u64,
/// Sequence number within the operation
pub sequence: u32,
}
impl DataFlowStep {
/// Create a new data flow step
#[must_use]
pub fn new(from: Uuid, to: Uuid, size_bytes: u64) -> Self {
Self {
from,
to,
size_bytes,
duration_us: 0,
sequence: 0,
}
}
/// Calculate throughput in Gbps
#[must_use]
pub fn throughput_gbps(&self) -> f64 {
if self.duration_us == 0 {
return 0.0;
}
let bytes_per_us = self.size_bytes as f64 / self.duration_us as f64;
let bits_per_us = bytes_per_us * 8.0;
bits_per_us / 1000.0 // Convert to Gbps
}
}
/// Trace of a collective operation for visualization
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CollectiveTrace {
/// Unique trace identifier
pub id: Uuid,
/// Type of collective operation
pub operation: CollectiveOp,
/// Participating node IDs
pub participants: Vec<Uuid>,
/// Data flow steps (ordered)
pub data_flow: Vec<DataFlowStep>,
/// Total duration in microseconds
pub duration_us: u64,
/// Total data transferred in bytes
pub total_bytes: u64,
/// Start timestamp
pub started_at: DateTime<Utc>,
/// Optional description
pub description: Option<String>,
}
impl CollectiveTrace {
/// Create a new collective trace
#[must_use]
pub fn new(operation: CollectiveOp, participants: Vec<Uuid>) -> Self {
Self {
id: Uuid::new_v4(),
operation,
participants,
data_flow: Vec::new(),
duration_us: 0,
total_bytes: 0,
started_at: Utc::now(),
description: None,
}
}
/// Add a data flow step
pub fn add_step(&mut self, step: DataFlowStep) {
self.total_bytes += step.size_bytes;
self.duration_us = self.duration_us.max(step.duration_us);
self.data_flow.push(step);
}
/// Calculate effective bandwidth in Gbps
#[must_use]
pub fn effective_bandwidth_gbps(&self) -> f64 {
if self.duration_us == 0 {
return 0.0;
}
let bytes_per_us = self.total_bytes as f64 / self.duration_us as f64;
let bits_per_us = bytes_per_us * 8.0;
bits_per_us / 1000.0
}
}
/// Severity level for bottlenecks and alerts
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
pub enum Severity {
/// Informational, no action needed
Info,
/// Warning, should be monitored
Warning,
/// Critical, immediate attention needed
Critical,
}
impl std::fmt::Display for Severity {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Self::Info => write!(f, "Info"),
Self::Warning => write!(f, "Warning"),
Self::Critical => write!(f, "Critical"),
}
}
}
/// Type of bottleneck detected
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub enum BottleneckType {
/// Link bandwidth saturation
BandwidthSaturation,
/// High latency on a link
HighLatency,
/// Memory pressure on a node
MemoryPressure,
/// GPU compute bottleneck
ComputeBottleneck,
/// Network congestion
NetworkCongestion,
/// Load imbalance across nodes
LoadImbalance,
}
impl std::fmt::Display for BottleneckType {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Self::BandwidthSaturation => write!(f, "Bandwidth Saturation"),
Self::HighLatency => write!(f, "High Latency"),
Self::MemoryPressure => write!(f, "Memory Pressure"),
Self::ComputeBottleneck => write!(f, "Compute Bottleneck"),
Self::NetworkCongestion => write!(f, "Network Congestion"),
Self::LoadImbalance => write!(f, "Load Imbalance"),
}
}
}
/// Detected bottleneck information
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct BottleneckInfo {
/// Unique identifier
pub id: Uuid,
/// Type of bottleneck
pub bottleneck_type: BottleneckType,
/// Severity level
pub severity: Severity,
/// Affected node IDs
pub affected_nodes: Vec<Uuid>,
/// Affected link IDs
pub affected_links: Vec<Uuid>,
/// Human-readable description
pub description: String,
/// Suggested remediation
pub remediation: Option<String>,
/// Detected at timestamp
pub detected_at: DateTime<Utc>,
/// Metric value that triggered detection
pub metric_value: f64,
/// Threshold that was exceeded
pub threshold: f64,
}
impl BottleneckInfo {
/// Create a new bottleneck info
#[must_use]
pub fn new(
bottleneck_type: BottleneckType,
severity: Severity,
description: impl Into<String>,
) -> Self {
Self {
id: Uuid::new_v4(),
bottleneck_type,
severity,
affected_nodes: Vec::new(),
affected_links: Vec::new(),
description: description.into(),
remediation: None,
detected_at: Utc::now(),
metric_value: 0.0,
threshold: 0.0,
}
}
/// Add affected node
pub fn add_affected_node(&mut self, node_id: Uuid) {
self.affected_nodes.push(node_id);
}
/// Add affected link
pub fn add_affected_link(&mut self, link_id: Uuid) {
self.affected_links.push(link_id);
}
/// Set metric value and threshold
pub fn with_metrics(mut self, value: f64, threshold: f64) -> Self {
self.metric_value = value;
self.threshold = threshold;
self
}
/// Set remediation suggestion
pub fn with_remediation(mut self, remediation: impl Into<String>) -> Self {
self.remediation = Some(remediation.into());
self
}
}
/// Alert configuration for monitoring
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AlertConfig {
/// Enable alerts
pub enabled: bool,
/// GPU utilization warning threshold (0.0 - 1.0)
pub gpu_util_warning: f64,
/// GPU utilization critical threshold (0.0 - 1.0)
pub gpu_util_critical: f64,
/// Memory usage warning threshold (0.0 - 1.0)
pub memory_warning: f64,
/// Memory usage critical threshold (0.0 - 1.0)
pub memory_critical: f64,
/// Bandwidth saturation warning threshold (0.0 - 1.0)
pub bandwidth_warning: f64,
/// Bandwidth saturation critical threshold (0.0 - 1.0)
pub bandwidth_critical: f64,
/// Latency warning threshold in microseconds
pub latency_warning_us: f64,
/// Latency critical threshold in microseconds
pub latency_critical_us: f64,
/// Load imbalance warning threshold (std dev / mean)
pub load_imbalance_warning: f64,
}
impl AlertConfig {
/// Create a new alert config with default thresholds
#[must_use]
pub fn new() -> Self {
Self {
enabled: true,
gpu_util_warning: 0.85,
gpu_util_critical: 0.95,
memory_warning: 0.80,
memory_critical: 0.90,
bandwidth_warning: 0.75,
bandwidth_critical: 0.90,
latency_warning_us: 5.0,
latency_critical_us: 10.0,
load_imbalance_warning: 0.20,
}
}
/// Create a relaxed config for development/testing
#[must_use]
pub fn relaxed() -> Self {
Self {
enabled: true,
gpu_util_warning: 0.95,
gpu_util_critical: 0.99,
memory_warning: 0.90,
memory_critical: 0.95,
bandwidth_warning: 0.85,
bandwidth_critical: 0.95,
latency_warning_us: 10.0,
latency_critical_us: 20.0,
load_imbalance_warning: 0.30,
}
}
/// Create a strict config for production
#[must_use]
pub fn strict() -> Self {
Self {
enabled: true,
gpu_util_warning: 0.75,
gpu_util_critical: 0.90,
memory_warning: 0.70,
memory_critical: 0.85,
bandwidth_warning: 0.65,
bandwidth_critical: 0.80,
latency_warning_us: 3.0,
latency_critical_us: 7.0,
load_imbalance_warning: 0.15,
}
}
}
impl Default for AlertConfig {
fn default() -> Self {
Self::new()
}
}
// ============================================================================
// Sample Data Functions
// ============================================================================
/// Create a sample 4-node M3 Max cluster topology
#[must_use]
pub fn sample_four_node_topology() -> ClusterTopology {
let mut topology = ClusterTopology::new("M3-Max-Cluster-4");
// Create 4 M3 Max nodes
let node1 = NodeInfo::new("mac-studio-1", "Apple M3 Max", 128).as_coordinator();
let node2 = NodeInfo::new("mac-studio-2", "Apple M3 Max", 128);
let node3 = NodeInfo::new("mac-studio-3", "Apple M3 Max", 128);
let node4 = NodeInfo::new("mac-studio-4", "Apple M3 Max", 128);
let node_ids = [node1.id, node2.id, node3.id, node4.id];
topology.add_node(node1);
topology.add_node(node2);
topology.add_node(node3);
topology.add_node(node4);
// Create Thunderbolt 5 links (ring topology + star from coordinator)
// Ring: 1 -> 2 -> 3 -> 4 -> 1
topology.add_link(LinkInfo::new(
node_ids[0],
node_ids[1],
TransportType::Thunderbolt,
));
topology.add_link(LinkInfo::new(
node_ids[1],
node_ids[2],
TransportType::Thunderbolt,
));
topology.add_link(LinkInfo::new(
node_ids[2],
node_ids[3],
TransportType::Thunderbolt,
));
topology.add_link(LinkInfo::new(
node_ids[3],
node_ids[0],
TransportType::Thunderbolt,
));
// Cross links for better bandwidth
topology.add_link(LinkInfo::new(
node_ids[0],
node_ids[2],
TransportType::Thunderbolt,
));
topology.add_link(LinkInfo::new(
node_ids[1],
node_ids[3],
TransportType::Thunderbolt,
));
topology
}
/// Create a sample 8-node M4 Ultra cluster topology
#[must_use]
pub fn sample_eight_node_topology() -> ClusterTopology {
let mut topology = ClusterTopology::new("M4-Ultra-Cluster-8");
// Create 8 M4 Ultra nodes (simulated future hardware)
let mut nodes = Vec::new();
for i in 0..8 {
let mut node = NodeInfo::new(format!("mac-pro-{}", i + 1), "Apple M4 Ultra", 512);
if i == 0 {
node = node.as_coordinator();
}
nodes.push(node);
}
let node_ids: Vec<Uuid> = nodes.iter().map(|n| n.id).collect();
for node in nodes {
topology.add_node(node);
}
// Create fat-tree topology
// Layer 1: Pairs connected with high bandwidth
for i in (0..8).step_by(2) {
topology.add_link(LinkInfo::new(
node_ids[i],
node_ids[i + 1],
TransportType::Thunderbolt,
));
}
// Layer 2: Connect pairs
for i in 0..4 {
let src = i * 2;
let dst = ((i + 1) % 4) * 2;
topology.add_link(LinkInfo::new(
node_ids[src],
node_ids[dst],
TransportType::Thunderbolt,
));
}
// Cross links for redundancy
topology.add_link(LinkInfo::new(
node_ids[0],
node_ids[4],
TransportType::Thunderbolt,
));
topology.add_link(LinkInfo::new(
node_ids[2],
node_ids[6],
TransportType::Thunderbolt,
));
topology
}
/// Create sample cluster metrics for a topology
#[must_use]
pub fn sample_metrics(topology: &ClusterTopology) -> ClusterMetrics {
let mut metrics = ClusterMetrics::new();
// Generate utilization data for each node
for (i, node) in topology.nodes.iter().enumerate() {
let base_util = 0.6 + (i as f64 * 0.05);
metrics.gpu_utilization.push((node.id, base_util.min(0.95)));
metrics
.memory_used
.push((node.id, node.memory_gb as f64 * (0.4 + i as f64 * 0.05)));
}
// Generate bandwidth data for each link
for link in &topology.links {
let bandwidth = link.bandwidth_gbps * 0.6; // 60% utilization
metrics.bandwidth_usage.push((link.id, bandwidth));
}
// Aggregate metrics
metrics.total_tflops = topology.nodes.len() as f64 * 14.0; // ~14 TFLOPS per M3 Max
metrics.active_collectives = 3;
metrics.data_transfer_rate_gb = 12.5;
metrics
}
/// Create a sample all-reduce collective trace
#[must_use]
pub fn sample_allreduce_trace(node_ids: &[Uuid]) -> CollectiveTrace {
let mut trace = CollectiveTrace::new(CollectiveOp::AllReduce, node_ids.to_vec());
trace.description = Some("Gradient synchronization for distributed training".to_string());
let chunk_size: u64 = 256 * 1024 * 1024; // 256 MB per chunk
// Ring all-reduce pattern
let n = node_ids.len();
for phase in 0..(n - 1) {
for i in 0..n {
let src = node_ids[i];
let dst = node_ids[(i + 1) % n];
let mut step = DataFlowStep::new(src, dst, chunk_size);
step.sequence = (phase * n + i) as u32;
step.duration_us = 2000 + (phase as u64 * 100); // ~2ms per step
trace.add_step(step);
}
}
trace.duration_us = (n - 1) as u64 * 2500; // Total duration
trace
}
/// Create a sample broadcast collective trace
#[must_use]
pub fn sample_broadcast_trace(node_ids: &[Uuid], root_idx: usize) -> CollectiveTrace {
let mut trace = CollectiveTrace::new(CollectiveOp::Broadcast, node_ids.to_vec());
trace.description = Some("Model weight broadcast from coordinator".to_string());
let data_size: u64 = 1024 * 1024 * 1024; // 1 GB
let root = node_ids[root_idx];
for (i, &dst) in node_ids.iter().enumerate() {
if i != root_idx {
let mut step = DataFlowStep::new(root, dst, data_size);
step.sequence = i as u32;
step.duration_us = 8000; // ~8ms for 1GB at 120 Gbps
trace.add_step(step);
}
}
trace.duration_us = 8000; // Parallel broadcast
trace
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_transport_type_properties() {
assert!(
TransportType::MetalDirect.max_bandwidth_gbps()
> TransportType::Thunderbolt.max_bandwidth_gbps()
);
assert!(
TransportType::MetalDirect.typical_latency_us()
< TransportType::Thunderbolt.typical_latency_us()
);
}
#[test]
fn test_node_info_creation() {
let node = NodeInfo::new("test-node", "Apple M3 Max", 128);
assert_eq!(node.name, "test-node");
assert_eq!(node.gpu_model, "Apple M3 Max");
assert_eq!(node.memory_gb, 128);
assert!(!node.is_coordinator);
}
#[test]
fn test_node_coordinator() {
let node = NodeInfo::new("test-node", "Apple M3 Max", 128).as_coordinator();
assert!(node.is_coordinator);
}
#[test]
fn test_link_info_creation() {
let source = Uuid::new_v4();
let target = Uuid::new_v4();
let link = LinkInfo::new(source, target, TransportType::Thunderbolt);
assert_eq!(link.source, source);
assert_eq!(link.target, target);
assert_eq!(link.transport_type, TransportType::Thunderbolt);
assert!(link.bandwidth_efficiency() > 0.0);
assert!(link.bandwidth_efficiency() <= 1.0);
}
#[test]
fn test_cluster_topology() {
let mut topology = ClusterTopology::new("test-cluster");
let node1 = NodeInfo::new("node-1", "M3 Max", 128);
let node2 = NodeInfo::new("node-2", "M3 Max", 128);
let node1_id = node1.id;
let node2_id = node2.id;
topology.add_node(node1);
topology.add_node(node2);
assert_eq!(topology.node_count(), 2);
assert_eq!(topology.total_memory_gb(), 256);
topology.add_link(LinkInfo::new(
node1_id,
node2_id,
TransportType::Thunderbolt,
));
assert_eq!(topology.link_count(), 1);
let links = topology.get_node_links(node1_id);
assert_eq!(links.len(), 1);
}
#[test]
fn test_cluster_metrics() {
let mut metrics = ClusterMetrics::new();
let node_id = Uuid::new_v4();
metrics.gpu_utilization.push((node_id, 0.75));
metrics.memory_used.push((node_id, 64.0));
assert!((metrics.average_gpu_utilization() - 0.75).abs() < f64::EPSILON);
assert!((metrics.total_memory_used_gb() - 64.0).abs() < f64::EPSILON);
}
#[test]
fn test_collective_trace() {
let nodes = vec![Uuid::new_v4(), Uuid::new_v4()];
let mut trace = CollectiveTrace::new(CollectiveOp::AllReduce, nodes.clone());
let step = DataFlowStep::new(nodes[0], nodes[1], 1_000_000);
trace.add_step(step);
assert_eq!(trace.participants.len(), 2);
assert_eq!(trace.data_flow.len(), 1);
assert_eq!(trace.total_bytes, 1_000_000);
}
#[test]
fn test_bottleneck_info() {
let bottleneck = BottleneckInfo::new(
BottleneckType::BandwidthSaturation,
Severity::Warning,
"Link bandwidth at 90%",
)
.with_metrics(0.90, 0.75)
.with_remediation("Consider adding additional interconnects");
assert_eq!(
bottleneck.bottleneck_type,
BottleneckType::BandwidthSaturation
);
assert_eq!(bottleneck.severity, Severity::Warning);
assert!(bottleneck.remediation.is_some());
}
#[test]
fn test_alert_config() {
let default_config = AlertConfig::new();
let strict_config = AlertConfig::strict();
let relaxed_config = AlertConfig::relaxed();
assert!(strict_config.gpu_util_warning < default_config.gpu_util_warning);
assert!(relaxed_config.gpu_util_warning > default_config.gpu_util_warning);
}
#[test]
fn test_sample_four_node_topology() {
let topology = sample_four_node_topology();
assert_eq!(topology.node_count(), 4);
assert_eq!(topology.link_count(), 6); // Ring + cross links
assert_eq!(topology.total_memory_gb(), 512);
// Verify coordinator exists
let coordinator_count = topology.nodes.iter().filter(|n| n.is_coordinator).count();
assert_eq!(coordinator_count, 1);
}
#[test]
fn test_sample_eight_node_topology() {
let topology = sample_eight_node_topology();
assert_eq!(topology.node_count(), 8);
assert!(topology.link_count() >= 6); // At least ring + some cross links
assert_eq!(topology.total_memory_gb(), 4096);
}
#[test]
fn test_sample_metrics() {
let topology = sample_four_node_topology();
let metrics = sample_metrics(&topology);
assert_eq!(metrics.gpu_utilization.len(), 4);
assert_eq!(metrics.memory_used.len(), 4);
assert!(metrics.total_tflops > 0.0);
}
#[test]
fn test_sample_allreduce_trace() {
let node_ids: Vec<Uuid> = (0..4).map(|_| Uuid::new_v4()).collect();
let trace = sample_allreduce_trace(&node_ids);
assert_eq!(trace.operation, CollectiveOp::AllReduce);
assert_eq!(trace.participants.len(), 4);
assert!(!trace.data_flow.is_empty());
assert!(trace.total_bytes > 0);
}
#[test]
fn test_sample_broadcast_trace() {
let node_ids: Vec<Uuid> = (0..4).map(|_| Uuid::new_v4()).collect();
let trace = sample_broadcast_trace(&node_ids, 0);
assert_eq!(trace.operation, CollectiveOp::Broadcast);
assert_eq!(trace.data_flow.len(), 3); // Broadcast from root to 3 others
}
#[test]
fn test_data_flow_step_throughput() {
let mut step = DataFlowStep::new(Uuid::new_v4(), Uuid::new_v4(), 1_000_000_000); // 1 GB
step.duration_us = 8000; // 8 ms
let throughput = step.throughput_gbps();
// 1 GB / 8 ms = 125 GB/s = 1000 Gbps
assert!(throughput > 900.0 && throughput < 1100.0);
}
#[test]
fn test_display_implementations() {
assert_eq!(format!("{}", TransportType::MetalDirect), "Metal Direct");
assert_eq!(format!("{}", NodeStatus::Healthy), "Healthy");
assert_eq!(format!("{}", CollectiveOp::AllReduce), "AllReduce");
assert_eq!(format!("{}", Severity::Critical), "Critical");
assert_eq!(
format!("{}", BottleneckType::MemoryPressure),
"Memory Pressure"
);
}
}