Files
rustytorch/demos/rtx-clusterviz-demo/src/sample_data.rs
T
2026-03-04 00:08:42 +00:00

589 lines
19 KiB
Rust

//! Sample Data Module
//!
//! This module provides sample cluster configurations, workloads, and traces
//! for demonstration and testing purposes.
use chrono::Utc;
use clusterviz_shared::{
ClusterTopology, CollectiveOp, CollectiveTrace, DataFlowStep, LinkInfo, NodeInfo, NodeStatus,
TransportType,
};
use uuid::Uuid;
/// Create a 4-node M3 Max cluster configuration
///
/// This represents a typical Mac Studio cluster connected via Thunderbolt 5.
#[must_use]
pub fn four_node_cluster() -> ClusterTopology {
let mut topology = ClusterTopology::new("M3-Max-Training-Cluster");
// Create 4 Mac Studio nodes with M3 Max
let node1 = NodeInfo::new("mac-studio-alpha", "Apple M3 Max", 128).as_coordinator();
let node2 = NodeInfo::new("mac-studio-beta", "Apple M3 Max", 128);
let node3 = NodeInfo::new("mac-studio-gamma", "Apple M3 Max", 96);
let node4 = NodeInfo::new("mac-studio-delta", "Apple M3 Max", 96);
let 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 mesh topology
// Full mesh for maximum bandwidth
for i in 0..4 {
for j in (i + 1)..4 {
let mut link = LinkInfo::new(ids[i], ids[j], TransportType::Thunderbolt);
// Slightly vary bandwidth to simulate real conditions
link.bandwidth_gbps = 96.0 + (i as f64 * 2.0); // 96-102 Gbps
link.latency_us = 2.5 + (j as f64 * 0.1); // 2.5-2.8 us
topology.add_link(link);
}
}
topology
}
/// Create an 8-node M4 Ultra cluster configuration
///
/// This represents a high-end cluster with next-generation hardware.
#[must_use]
pub fn eight_node_cluster() -> ClusterTopology {
let mut topology = ClusterTopology::new("M4-Ultra-HPC-Cluster");
// Create 8 Mac Pro nodes with M4 Ultra (theoretical)
let gpu_configs = [
("mac-pro-01", "Apple M4 Ultra", 512),
("mac-pro-02", "Apple M4 Ultra", 512),
("mac-pro-03", "Apple M4 Ultra", 384),
("mac-pro-04", "Apple M4 Ultra", 384),
("mac-pro-05", "Apple M4 Max", 192),
("mac-pro-06", "Apple M4 Max", 192),
("mac-pro-07", "Apple M4 Max", 128),
("mac-pro-08", "Apple M4 Max", 128),
];
let mut nodes = Vec::new();
for (i, (name, gpu, mem)) in gpu_configs.iter().enumerate() {
let mut node = NodeInfo::new(*name, *gpu, *mem);
if i == 0 {
node = node.as_coordinator();
}
nodes.push(node);
}
let ids: Vec<Uuid> = nodes.iter().map(|n| n.id).collect();
for node in nodes {
topology.add_node(node);
}
// Fat-tree topology for 8 nodes
// Tier 1: Pairs (4 pairs)
for i in (0..8).step_by(2) {
let mut link = LinkInfo::new(ids[i], ids[i + 1], TransportType::Thunderbolt);
link.bandwidth_gbps = 100.0;
link.latency_us = 2.0;
topology.add_link(link);
}
// Tier 2: Connect adjacent pairs
for i in 0..4 {
let src = i * 2;
let dst = ((i + 1) % 4) * 2;
let mut link = LinkInfo::new(ids[src], ids[dst], TransportType::Thunderbolt);
link.bandwidth_gbps = 96.0;
link.latency_us = 2.5;
topology.add_link(link);
}
// Cross-tier links for redundancy
topology.add_link(LinkInfo::new(ids[0], ids[4], TransportType::Thunderbolt));
topology.add_link(LinkInfo::new(ids[2], ids[6], TransportType::Thunderbolt));
topology.add_link(LinkInfo::new(ids[1], ids[5], TransportType::Thunderbolt));
topology.add_link(LinkInfo::new(ids[3], ids[7], TransportType::Thunderbolt));
topology
}
/// Create a mixed-generation cluster
///
/// Demonstrates heterogeneous hardware configurations.
#[must_use]
pub fn mixed_generation_cluster() -> ClusterTopology {
let mut topology = ClusterTopology::new("Mixed-Generation-Cluster");
// Mix of M3 and M4 generation hardware
let configs = [
("workstation-01", "Apple M4 Max", 128, true),
("workstation-02", "Apple M3 Max", 128, false),
("workstation-03", "Apple M3 Pro", 64, false),
("workstation-04", "Apple M4 Pro", 96, false),
];
let mut nodes = Vec::new();
for (name, gpu, mem, is_coord) in &configs {
let mut node = NodeInfo::new(*name, *gpu, *mem);
if *is_coord {
node = node.as_coordinator();
}
nodes.push(node);
}
let ids: Vec<Uuid> = nodes.iter().map(|n| n.id).collect();
for node in nodes {
topology.add_node(node);
}
// Star topology with coordinator at center
for i in 1..4 {
let mut link = LinkInfo::new(ids[0], ids[i], TransportType::Thunderbolt);
// Bandwidth varies by GPU generation
link.bandwidth_gbps = if i <= 2 { 96.0 } else { 80.0 };
link.latency_us = 2.5;
topology.add_link(link);
}
// Additional ring connection for redundancy
for i in 1..4 {
let next = if i == 3 { 1 } else { i + 1 };
let mut link = LinkInfo::new(ids[i], ids[next], TransportType::Thunderbolt);
link.bandwidth_gbps = 80.0;
link.latency_us = 3.0;
topology.add_link(link);
}
topology
}
/// Configure a topology for training workload
///
/// Sets utilization patterns typical for distributed training.
pub fn training_workload(topology: &mut ClusterTopology) {
for node in &mut topology.nodes {
// Training typically uses high GPU and memory
node.utilization = 0.85 + (rand::random::<f64>() * 0.10);
node.status = NodeStatus::Healthy;
node.last_heartbeat = Utc::now();
}
for link in &mut topology.links {
// Training uses moderate bandwidth for gradient sync
link.bandwidth_gbps = link.transport_type.max_bandwidth_gbps() * 0.60;
link.is_active = true;
}
}
/// Configure a topology for inference workload
///
/// Sets utilization patterns typical for distributed inference.
pub fn inference_workload(topology: &mut ClusterTopology) {
for (i, node) in topology.nodes.iter_mut().enumerate() {
// Inference has variable utilization based on request load
let base_util = 0.30 + (i as f64 * 0.05);
node.utilization = base_util + (rand::random::<f64>() * 0.20);
node.status = NodeStatus::Healthy;
node.last_heartbeat = Utc::now();
}
for link in &mut topology.links {
// Inference uses lower bandwidth (mostly input data)
link.bandwidth_gbps = link.transport_type.max_bandwidth_gbps() * 0.25;
link.is_active = rand::random::<bool>();
}
}
/// Create sample collective traces for a set of nodes
#[must_use]
pub fn sample_collective_traces(node_ids: &[Uuid]) -> Vec<CollectiveTrace> {
let mut traces = Vec::new();
// AllReduce trace (gradient synchronization)
let mut allreduce = CollectiveTrace::new(CollectiveOp::AllReduce, node_ids.to_vec());
allreduce.description = Some("Gradient synchronization - FP16 gradients".to_string());
let chunk_size: u64 = 128 * 1024 * 1024; // 128 MB
let n = node_ids.len();
// Ring all-reduce pattern
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 = 1000 + (phase as u64 * 50);
allreduce.add_step(step);
}
}
allreduce.duration_us = (n - 1) as u64 * 1200;
traces.push(allreduce);
// Broadcast trace (model distribution)
let mut broadcast = CollectiveTrace::new(CollectiveOp::Broadcast, node_ids.to_vec());
broadcast.description = Some("Model weight broadcast - BF16 weights".to_string());
let model_size: u64 = 2 * 1024 * 1024 * 1024; // 2 GB
if !node_ids.is_empty() {
let root = node_ids[0];
for (i, &dst) in node_ids.iter().enumerate().skip(1) {
let mut step = DataFlowStep::new(root, dst, model_size);
step.sequence = i as u32;
step.duration_us = 15000; // ~15ms at 120 Gbps
broadcast.add_step(step);
}
}
broadcast.duration_us = 15000;
traces.push(broadcast);
// AllGather trace (activation gathering)
let mut allgather = CollectiveTrace::new(CollectiveOp::AllGather, node_ids.to_vec());
allgather.description = Some("Activation gathering for tensor parallelism".to_string());
let activation_size: u64 = 64 * 1024 * 1024; // 64 MB per node
for (i, &src) in node_ids.iter().enumerate() {
for (j, &dst) in node_ids.iter().enumerate() {
if i != j {
let mut step = DataFlowStep::new(src, dst, activation_size);
step.sequence = (i * n + j) as u32;
step.duration_us = 500;
allgather.add_step(step);
}
}
}
allgather.duration_us = 500 * (n - 1) as u64;
traces.push(allgather);
// ReduceScatter trace
let mut reduce_scatter = CollectiveTrace::new(CollectiveOp::ReduceScatter, node_ids.to_vec());
reduce_scatter.description = Some("Reduce-scatter for pipeline parallelism".to_string());
let scatter_size: u64 = 32 * 1024 * 1024; // 32 MB
for phase in 0..(n - 1) {
for i in 0..n {
let src = node_ids[i];
let dst = node_ids[(i + n - 1) % n];
let mut step = DataFlowStep::new(src, dst, scatter_size);
step.sequence = (phase * n + i) as u32;
step.duration_us = 300;
reduce_scatter.add_step(step);
}
}
reduce_scatter.duration_us = 300 * (n - 1) as u64;
traces.push(reduce_scatter);
traces
}
/// Create traces for a training step
#[must_use]
pub fn training_step_traces(node_ids: &[Uuid], batch_size: usize) -> Vec<CollectiveTrace> {
let mut traces = Vec::new();
// Forward pass activation sync
let mut forward = CollectiveTrace::new(CollectiveOp::AllGather, node_ids.to_vec());
forward.description = Some(format!("Forward pass - batch size {}", batch_size));
let activation_per_layer: u64 = (batch_size * 4096 * 4) as u64; // 4 bytes per element
for &src in node_ids {
for &dst in node_ids {
if src != dst {
let step = DataFlowStep::new(src, dst, activation_per_layer);
forward.add_step(step);
}
}
}
forward.duration_us = 800;
traces.push(forward);
// Backward pass gradient sync
let mut backward = CollectiveTrace::new(CollectiveOp::AllReduce, node_ids.to_vec());
backward.description = Some("Backward pass - gradient synchronization".to_string());
let gradient_size: u64 = 512 * 1024 * 1024; // 512 MB total gradients
let n = node_ids.len();
for i in 0..n {
let src = node_ids[i];
let dst = node_ids[(i + 1) % n];
let mut step = DataFlowStep::new(src, dst, gradient_size / n as u64);
step.duration_us = 2000;
backward.add_step(step);
}
backward.duration_us = 2000;
traces.push(backward);
// Parameter update broadcast
let mut update = CollectiveTrace::new(CollectiveOp::Broadcast, node_ids.to_vec());
update.description = Some("Parameter update broadcast".to_string());
if !node_ids.is_empty() {
let root = node_ids[0];
let param_size: u64 = 100 * 1024 * 1024; // 100 MB updated params
for &dst in &node_ids[1..] {
let step = DataFlowStep::new(root, dst, param_size);
update.add_step(step);
}
}
update.duration_us = 1000;
traces.push(update);
traces
}
/// Create traces for inference batching
#[must_use]
pub fn inference_batch_traces(
node_ids: &[Uuid],
requests_per_batch: usize,
) -> Vec<CollectiveTrace> {
let mut traces = Vec::new();
// Input distribution
let mut input_dist = CollectiveTrace::new(CollectiveOp::Scatter, node_ids.to_vec());
input_dist.description = Some(format!("Input scatter - {} requests", requests_per_batch));
if !node_ids.is_empty() {
let root = node_ids[0];
let input_size: u64 = (requests_per_batch * 1024 * 4) as u64; // 4KB per request
for &dst in &node_ids[1..] {
let step = DataFlowStep::new(root, dst, input_size / node_ids.len() as u64);
input_dist.add_step(step);
}
}
input_dist.duration_us = 100;
traces.push(input_dist);
// Output gathering
let mut output_gather = CollectiveTrace::new(CollectiveOp::AllGather, node_ids.to_vec());
output_gather.description = Some("Output gathering".to_string());
let output_size: u64 = (requests_per_batch * 4096 * 4) as u64; // Larger output
for &src in node_ids {
for &dst in node_ids {
if src != dst {
let mut step = DataFlowStep::new(src, dst, output_size / node_ids.len() as u64);
step.duration_us = 200;
output_gather.add_step(step);
}
}
}
output_gather.duration_us = 200;
traces.push(output_gather);
traces
}
/// Generate a stress test topology with high utilization
#[must_use]
pub fn stress_test_cluster() -> ClusterTopology {
let mut topology = four_node_cluster();
// Set all nodes to high utilization
for node in &mut topology.nodes {
node.utilization = 0.95 + (rand::random::<f64>() * 0.05);
node.status = NodeStatus::HighLoad;
}
// Set all links to high bandwidth usage
for link in &mut topology.links {
link.bandwidth_gbps = link.transport_type.max_bandwidth_gbps() * 0.90;
link.latency_us = link.transport_type.typical_latency_us() * 1.5;
link.is_active = true;
}
topology
}
/// Generate a cluster with some degraded nodes
#[must_use]
pub fn degraded_cluster() -> ClusterTopology {
let mut topology = four_node_cluster();
// Mark one node as degraded
if let Some(node) = topology.nodes.get_mut(2) {
node.status = NodeStatus::Degraded;
node.utilization = 0.3;
}
// Mark one node as high load
if let Some(node) = topology.nodes.get_mut(3) {
node.status = NodeStatus::HighLoad;
node.utilization = 0.98;
}
topology
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_four_node_cluster() {
let topology = four_node_cluster();
assert_eq!(topology.node_count(), 4);
assert_eq!(topology.link_count(), 6); // Full mesh: 4*3/2 = 6 links
assert_eq!(topology.total_memory_gb(), 448); // 128+128+96+96
// Verify coordinator
let coordinator_count = topology.nodes.iter().filter(|n| n.is_coordinator).count();
assert_eq!(coordinator_count, 1);
}
#[test]
fn test_eight_node_cluster() {
let topology = eight_node_cluster();
assert_eq!(topology.node_count(), 8);
assert!(topology.link_count() >= 8); // At least pair + tier connections
assert_eq!(topology.total_memory_gb(), 2432); // 512*2 + 384*2 + 192*2 + 128*2
}
#[test]
fn test_mixed_generation_cluster() {
let topology = mixed_generation_cluster();
assert_eq!(topology.node_count(), 4);
assert!(topology.link_count() >= 4);
// Verify we have different GPU models
let gpu_models: Vec<&str> = topology
.nodes
.iter()
.map(|n| n.gpu_model.as_str())
.collect();
assert!(gpu_models.contains(&"Apple M4 Max"));
assert!(gpu_models.contains(&"Apple M3 Max"));
}
#[test]
fn test_training_workload() {
let mut topology = four_node_cluster();
training_workload(&mut topology);
for node in &topology.nodes {
assert!(node.utilization >= 0.85);
assert!(node.utilization <= 0.95);
assert_eq!(node.status, NodeStatus::Healthy);
}
for link in &topology.links {
assert!(link.is_active);
}
}
#[test]
fn test_inference_workload() {
let mut topology = four_node_cluster();
inference_workload(&mut topology);
for node in &topology.nodes {
// Inference workload: base 0.30 + (index * 0.05) + random(0..0.20)
// For 4 nodes, max is 0.30 + 0.15 + 0.20 = 0.65
assert!(node.utilization >= 0.30);
assert!(node.utilization <= 0.70);
}
}
#[test]
fn test_sample_collective_traces() {
let node_ids: Vec<Uuid> = (0..4).map(|_| Uuid::new_v4()).collect();
let traces = sample_collective_traces(&node_ids);
assert_eq!(traces.len(), 4);
// Verify trace types
let ops: Vec<CollectiveOp> = traces.iter().map(|t| t.operation).collect();
assert!(ops.contains(&CollectiveOp::AllReduce));
assert!(ops.contains(&CollectiveOp::Broadcast));
assert!(ops.contains(&CollectiveOp::AllGather));
assert!(ops.contains(&CollectiveOp::ReduceScatter));
}
#[test]
fn test_training_step_traces() {
let node_ids: Vec<Uuid> = (0..4).map(|_| Uuid::new_v4()).collect();
let traces = training_step_traces(&node_ids, 32);
assert_eq!(traces.len(), 3); // Forward, backward, update
}
#[test]
fn test_inference_batch_traces() {
let node_ids: Vec<Uuid> = (0..4).map(|_| Uuid::new_v4()).collect();
let traces = inference_batch_traces(&node_ids, 64);
assert_eq!(traces.len(), 2); // Input scatter, output gather
}
#[test]
fn test_stress_test_cluster() {
let topology = stress_test_cluster();
for node in &topology.nodes {
assert!(node.utilization >= 0.95);
assert_eq!(node.status, NodeStatus::HighLoad);
}
for link in &topology.links {
assert!(link.is_active);
assert!(link.bandwidth_gbps >= link.transport_type.max_bandwidth_gbps() * 0.85);
}
}
#[test]
fn test_degraded_cluster() {
let topology = degraded_cluster();
let degraded_count = topology
.nodes
.iter()
.filter(|n| n.status == NodeStatus::Degraded)
.count();
assert_eq!(degraded_count, 1);
let high_load_count = topology
.nodes
.iter()
.filter(|n| n.status == NodeStatus::HighLoad)
.count();
assert_eq!(high_load_count, 1);
}
#[test]
fn test_trace_data_flow_populated() {
let node_ids: Vec<Uuid> = (0..4).map(|_| Uuid::new_v4()).collect();
let traces = sample_collective_traces(&node_ids);
for trace in &traces {
assert!(!trace.data_flow.is_empty());
assert!(trace.total_bytes > 0);
assert!(trace.duration_us > 0);
}
}
#[test]
fn test_cluster_link_bandwidth_reasonable() {
let topology = four_node_cluster();
for link in &topology.links {
// All links should be Thunderbolt with reasonable bandwidth
assert_eq!(link.transport_type, TransportType::Thunderbolt);
assert!(link.bandwidth_gbps > 0.0);
assert!(link.bandwidth_gbps <= TransportType::Thunderbolt.max_bandwidth_gbps());
assert!(link.latency_us > 0.0);
}
}
}