651 lines
21 KiB
Rust
651 lines
21 KiB
Rust
//! Topology discovery and optimization for distributed training
|
|
//!
|
|
//! This module provides functionality to discover network topology,
|
|
//! optimize communication patterns, and adapt to hardware characteristics.
|
|
|
|
use crate::error::Result;
|
|
use serde::{Deserialize, Serialize};
|
|
use std::collections::HashMap;
|
|
|
|
/// Network topology information
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub struct TopologyInfo {
|
|
/// Number of nodes in the cluster
|
|
pub num_nodes: usize,
|
|
/// Number of GPUs per node
|
|
pub gpus_per_node: usize,
|
|
/// Network interconnect type (e.g., "infiniband", "ethernet")
|
|
pub interconnect: String,
|
|
/// Bandwidth information (Gbps)
|
|
pub bandwidth: TopologyBandwidth,
|
|
/// Node-to-node latency matrix (microseconds)
|
|
pub latency_matrix: Vec<Vec<f32>>,
|
|
/// GPU-to-GPU topology within nodes
|
|
pub intra_node_topology: Vec<GpuTopology>,
|
|
/// Inter-node network topology
|
|
pub inter_node_topology: NetworkTopology,
|
|
}
|
|
|
|
/// Bandwidth information for different communication types
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub struct TopologyBandwidth {
|
|
/// Intra-node GPU-to-GPU bandwidth (NVLink, etc.)
|
|
pub intra_node_gpu: f32,
|
|
/// Intra-node CPU-GPU bandwidth (PCIe)
|
|
pub cpu_gpu: f32,
|
|
/// Inter-node network bandwidth
|
|
pub inter_node: f32,
|
|
/// Memory bandwidth per GPU
|
|
pub memory: f32,
|
|
}
|
|
|
|
/// GPU topology within a single node
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub struct GpuTopology {
|
|
/// Node ID
|
|
pub node_id: usize,
|
|
/// GPU devices in this node
|
|
pub gpus: Vec<GpuDevice>,
|
|
/// GPU-to-GPU connectivity matrix
|
|
pub gpu_connectivity: Vec<Vec<ConnectionType>>,
|
|
}
|
|
|
|
/// GPU device information
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub struct GpuDevice {
|
|
/// Device ID within the node
|
|
pub device_id: usize,
|
|
/// GPU model/architecture
|
|
pub model: String,
|
|
/// Memory size in GB
|
|
pub memory_gb: f32,
|
|
/// Compute capability
|
|
pub compute_capability: String,
|
|
/// CUDA/ROCm device index
|
|
pub device_index: i32,
|
|
}
|
|
|
|
/// Connection type between GPUs
|
|
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
|
|
pub enum ConnectionType {
|
|
/// No direct connection
|
|
None,
|
|
/// PCIe connection
|
|
Pcie,
|
|
/// NVLink connection (NVIDIA)
|
|
NvLink,
|
|
/// Infinity Fabric (AMD)
|
|
InfinityFabric,
|
|
/// Cross-node network connection
|
|
Network,
|
|
}
|
|
|
|
/// Network topology between nodes
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub struct NetworkTopology {
|
|
/// Network type (tree, fat-tree, torus, etc.)
|
|
pub topology_type: String,
|
|
/// Number of network switches/hops between nodes
|
|
pub switch_matrix: Vec<Vec<usize>>,
|
|
/// Network interface details per node
|
|
pub network_interfaces: Vec<NetworkInterface>,
|
|
}
|
|
|
|
impl Default for NetworkTopology {
|
|
fn default() -> Self {
|
|
Self {
|
|
topology_type: "single-node".to_string(),
|
|
switch_matrix: vec![vec![0]],
|
|
network_interfaces: vec![],
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Network interface information
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub struct NetworkInterface {
|
|
/// Node ID
|
|
pub node_id: usize,
|
|
/// Interface name (e.g., "ib0", "eth0")
|
|
pub interface_name: String,
|
|
/// IP address
|
|
pub ip_address: String,
|
|
/// Bandwidth in Gbps
|
|
pub bandwidth_gbps: f32,
|
|
/// MTU size
|
|
pub mtu: usize,
|
|
}
|
|
|
|
/// Topology optimizer for communication patterns
|
|
pub struct TopologyOptimizer {
|
|
topology: TopologyInfo,
|
|
optimization_cache: HashMap<String, OptimizationResult>,
|
|
}
|
|
|
|
/// Result of topology optimization
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub struct OptimizationResult {
|
|
/// Recommended communication pattern
|
|
pub pattern: CommunicationPattern,
|
|
/// Expected performance improvement
|
|
pub improvement_factor: f32,
|
|
/// Memory overhead
|
|
pub memory_overhead_mb: f32,
|
|
}
|
|
|
|
/// Communication pattern recommendation
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub enum CommunicationPattern {
|
|
/// Ring-based AllReduce
|
|
Ring,
|
|
/// Tree-based AllReduce
|
|
Tree,
|
|
/// Recursive doubling
|
|
RecursiveDoubling,
|
|
/// Hierarchical (intra-node + inter-node)
|
|
Hierarchical,
|
|
/// Custom pattern with specific routing
|
|
Custom { routing: Vec<CommunicationRoute> },
|
|
}
|
|
|
|
/// Communication route specification
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub struct CommunicationRoute {
|
|
/// Source rank
|
|
pub src: i32,
|
|
/// Destination rank
|
|
pub dst: i32,
|
|
/// Intermediate hops
|
|
pub hops: Vec<i32>,
|
|
/// Expected bandwidth utilization
|
|
pub bandwidth_utilization: f32,
|
|
}
|
|
|
|
/// AllReduce communication pattern
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub struct AllReducePattern {
|
|
/// Algorithm used (e.g., "ring", "tree", "hierarchical")
|
|
pub algorithm: String,
|
|
/// Communication steps
|
|
pub steps: Vec<CommunicationStep>,
|
|
/// Estimated completion time in milliseconds
|
|
pub estimated_time_ms: f64,
|
|
/// Expected bandwidth in GB/s
|
|
pub expected_bandwidth_gbps: f32,
|
|
}
|
|
|
|
/// Individual communication step
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub struct CommunicationStep {
|
|
/// Operation name
|
|
pub operation: String,
|
|
/// Participating ranks
|
|
pub participants: Vec<usize>,
|
|
/// Estimated time for this step
|
|
pub estimated_time_ms: f64,
|
|
}
|
|
|
|
impl TopologyInfo {
|
|
/// Discover topology information from the current environment
|
|
pub fn discover() -> Result<Self> {
|
|
// For now, return a simulated topology
|
|
// In a real implementation, this would:
|
|
// 1. Query GPU topology via CUDA/ROCm APIs
|
|
// 2. Discover network interfaces
|
|
// 3. Measure inter-node bandwidth/latency
|
|
// 4. Build topology matrices
|
|
|
|
Ok(Self::default_topology())
|
|
}
|
|
|
|
/// Create default topology for testing
|
|
pub fn default_topology() -> Self {
|
|
let num_nodes = 2;
|
|
let gpus_per_node = 4;
|
|
|
|
Self {
|
|
num_nodes,
|
|
gpus_per_node,
|
|
interconnect: "infiniband".to_string(),
|
|
bandwidth: TopologyBandwidth {
|
|
intra_node_gpu: 600.0, // 600 GB/s NVLink
|
|
cpu_gpu: 32.0, // 32 GB/s PCIe 4.0 x16
|
|
inter_node: 200.0, // 200 Gbps InfiniBand
|
|
memory: 2000.0, // 2 TB/s HBM
|
|
},
|
|
latency_matrix: vec![
|
|
vec![0.0, 150.0], // Node 0 to [Node 0, Node 1] in μs
|
|
vec![150.0, 0.0], // Node 1 to [Node 0, Node 1] in μs
|
|
],
|
|
intra_node_topology: (0..num_nodes)
|
|
.map(|node_id| GpuTopology::default_for_node(node_id, gpus_per_node))
|
|
.collect(),
|
|
inter_node_topology: NetworkTopology::default_for_nodes(num_nodes),
|
|
}
|
|
}
|
|
|
|
/// Get total number of GPUs in the cluster
|
|
pub fn total_gpus(&self) -> usize {
|
|
self.num_nodes * self.gpus_per_node
|
|
}
|
|
|
|
/// Check if two ranks are on the same node
|
|
pub fn same_node(&self, rank1: i32, rank2: i32) -> bool {
|
|
let node1 = rank1 as usize / self.gpus_per_node;
|
|
let node2 = rank2 as usize / self.gpus_per_node;
|
|
node1 == node2
|
|
}
|
|
|
|
/// Get node ID for a given rank
|
|
pub fn node_for_rank(&self, rank: i32) -> usize {
|
|
rank as usize / self.gpus_per_node
|
|
}
|
|
|
|
/// Get GPU ID within node for a given rank
|
|
pub fn gpu_for_rank(&self, rank: i32) -> usize {
|
|
rank as usize % self.gpus_per_node
|
|
}
|
|
|
|
/// Get total device count across all nodes
|
|
pub fn device_count(&self) -> usize {
|
|
self.num_nodes * self.gpus_per_node
|
|
}
|
|
|
|
/// Check if topology has NVLink connections
|
|
pub fn has_nvlink_connections(&self) -> bool {
|
|
self.intra_node_topology.iter().any(|node| {
|
|
node.gpu_connectivity
|
|
.iter()
|
|
.any(|row| row.contains(&ConnectionType::NvLink))
|
|
})
|
|
}
|
|
|
|
/// Get bandwidth matrix for all connections
|
|
pub fn bandwidth_matrix(&self) -> Vec<Vec<f32>> {
|
|
let total_devices = self.device_count();
|
|
let mut matrix = vec![vec![0.0; total_devices]; total_devices];
|
|
|
|
// Fill with intra-node bandwidth
|
|
for node in &self.intra_node_topology {
|
|
for (i, row) in node.gpu_connectivity.iter().enumerate() {
|
|
for (j, &conn) in row.iter().enumerate() {
|
|
let global_i = node.node_id * self.gpus_per_node + i;
|
|
let global_j = node.node_id * self.gpus_per_node + j;
|
|
|
|
matrix[global_i][global_j] = match conn {
|
|
ConnectionType::NvLink => self.bandwidth.intra_node_gpu,
|
|
ConnectionType::Pcie => self.bandwidth.cpu_gpu,
|
|
ConnectionType::InfinityFabric => self.bandwidth.intra_node_gpu,
|
|
ConnectionType::Network => self.bandwidth.inter_node,
|
|
ConnectionType::None => 0.0,
|
|
};
|
|
}
|
|
}
|
|
}
|
|
|
|
// Add inter-node bandwidth
|
|
for i in 0..total_devices {
|
|
for j in 0..total_devices {
|
|
if self.node_for_rank(i as i32) != self.node_for_rank(j as i32) {
|
|
matrix[i][j] = self.bandwidth.inter_node;
|
|
}
|
|
}
|
|
}
|
|
|
|
matrix
|
|
}
|
|
|
|
/// Count NVLink connections
|
|
pub fn nvlink_connection_count(&self) -> usize {
|
|
self.intra_node_topology
|
|
.iter()
|
|
.map(|node| {
|
|
node.gpu_connectivity
|
|
.iter()
|
|
.map(|row| {
|
|
row.iter()
|
|
.filter(|&&conn| conn == ConnectionType::NvLink)
|
|
.count()
|
|
})
|
|
.sum::<usize>()
|
|
})
|
|
.sum()
|
|
}
|
|
}
|
|
|
|
impl GpuTopology {
|
|
/// Create default GPU topology for a node
|
|
pub fn default_for_node(node_id: usize, gpus_per_node: usize) -> Self {
|
|
let gpus = (0..gpus_per_node)
|
|
.map(|gpu_id| GpuDevice {
|
|
device_id: gpu_id,
|
|
model: "RTX-5090".to_string(),
|
|
memory_gb: 32.0,
|
|
compute_capability: "9.0".to_string(),
|
|
device_index: (node_id * gpus_per_node + gpu_id) as i32,
|
|
})
|
|
.collect();
|
|
|
|
// Create fully connected NVLink topology
|
|
let mut gpu_connectivity = vec![vec![ConnectionType::None; gpus_per_node]; gpus_per_node];
|
|
for i in 0..gpus_per_node {
|
|
for j in 0..gpus_per_node {
|
|
gpu_connectivity[i][j] = if i == j {
|
|
ConnectionType::None
|
|
} else {
|
|
ConnectionType::NvLink
|
|
};
|
|
}
|
|
}
|
|
|
|
Self {
|
|
node_id,
|
|
gpus,
|
|
gpu_connectivity,
|
|
}
|
|
}
|
|
}
|
|
|
|
impl NetworkTopology {
|
|
/// Create default network topology for nodes
|
|
pub fn default_for_nodes(num_nodes: usize) -> Self {
|
|
// Simple fully-connected topology
|
|
let switch_matrix = vec![vec![1; num_nodes]; num_nodes];
|
|
|
|
let network_interfaces = (0..num_nodes)
|
|
.map(|node_id| NetworkInterface {
|
|
node_id,
|
|
interface_name: format!("ib{node_id}"),
|
|
ip_address: format!("10.0.0.{}", node_id + 1),
|
|
bandwidth_gbps: 200.0,
|
|
mtu: 4096,
|
|
})
|
|
.collect();
|
|
|
|
Self {
|
|
topology_type: "fully_connected".to_string(),
|
|
switch_matrix,
|
|
network_interfaces,
|
|
}
|
|
}
|
|
}
|
|
|
|
impl TopologyOptimizer {
|
|
/// Create new topology optimizer
|
|
pub fn new(topology: TopologyInfo) -> Self {
|
|
Self {
|
|
topology,
|
|
optimization_cache: HashMap::new(),
|
|
}
|
|
}
|
|
|
|
/// Optimize communication pattern for AllReduce
|
|
pub fn optimize_allreduce(
|
|
&mut self,
|
|
world_size: i32,
|
|
tensor_size_mb: f32,
|
|
) -> Result<OptimizationResult> {
|
|
let cache_key = format!("allreduce_{world_size}_{tensor_size_mb}");
|
|
|
|
if let Some(cached) = self.optimization_cache.get(&cache_key) {
|
|
return Ok(cached.clone());
|
|
}
|
|
|
|
let pattern = if world_size <= 8 {
|
|
CommunicationPattern::Ring
|
|
} else if tensor_size_mb < 1.0 {
|
|
CommunicationPattern::Tree
|
|
} else {
|
|
CommunicationPattern::Hierarchical
|
|
};
|
|
|
|
let improvement_factor = self.estimate_improvement(&pattern, world_size, tensor_size_mb);
|
|
let memory_overhead = self.estimate_memory_overhead(&pattern, tensor_size_mb);
|
|
|
|
let result = OptimizationResult {
|
|
pattern,
|
|
improvement_factor,
|
|
memory_overhead_mb: memory_overhead,
|
|
};
|
|
|
|
self.optimization_cache.insert(cache_key, result.clone());
|
|
Ok(result)
|
|
}
|
|
|
|
/// Estimate performance improvement for a communication pattern
|
|
fn estimate_improvement(
|
|
&self,
|
|
pattern: &CommunicationPattern,
|
|
world_size: i32,
|
|
tensor_size_mb: f32,
|
|
) -> f32 {
|
|
match pattern {
|
|
CommunicationPattern::Ring => {
|
|
// Ring is good for large tensors
|
|
if tensor_size_mb > 10.0 { 1.2 } else { 0.9 }
|
|
}
|
|
CommunicationPattern::Tree => {
|
|
// Tree is good for small tensors
|
|
if tensor_size_mb < 1.0 { 1.4 } else { 0.8 }
|
|
}
|
|
CommunicationPattern::Hierarchical => {
|
|
// Hierarchical is good for multi-node
|
|
if world_size as usize > self.topology.gpus_per_node {
|
|
1.5
|
|
} else {
|
|
1.0
|
|
}
|
|
}
|
|
_ => 1.0,
|
|
}
|
|
}
|
|
|
|
/// Estimate memory overhead for a communication pattern
|
|
fn estimate_memory_overhead(&self, pattern: &CommunicationPattern, tensor_size_mb: f32) -> f32 {
|
|
match pattern {
|
|
CommunicationPattern::Ring => tensor_size_mb * 0.1, // 10% overhead
|
|
CommunicationPattern::Tree => tensor_size_mb * 0.2, // 20% overhead
|
|
CommunicationPattern::Hierarchical => tensor_size_mb * 0.15, // 15% overhead
|
|
_ => tensor_size_mb * 0.05, // 5% default overhead
|
|
}
|
|
}
|
|
|
|
/// Create new topology optimizer with default topology
|
|
pub fn with_default() -> Self {
|
|
// Create a default topology for discovery
|
|
let topology = TopologyInfo::default_topology();
|
|
Self {
|
|
topology,
|
|
optimization_cache: HashMap::new(),
|
|
}
|
|
}
|
|
|
|
/// Discover topology from device list
|
|
pub fn discover_topology(&self, devices: &[rtx_runtime::Device]) -> Result<TopologyInfo> {
|
|
// Create topology info based on device list
|
|
let device_count = devices.len();
|
|
let nodes = device_count.div_ceil(8); // Assume 8 GPUs per node
|
|
let gpus_per_node = if nodes == 1 { device_count } else { 8 };
|
|
|
|
// Create GPU devices
|
|
let mut gpus = Vec::new();
|
|
for (i, device) in devices.iter().enumerate() {
|
|
gpus.push(GpuDevice {
|
|
device_id: i,
|
|
model: format!("GPU-{i}"), // Simplified model name
|
|
memory_gb: (device.properties().total_memory / (1024 * 1024 * 1024)) as f32,
|
|
compute_capability: format!(
|
|
"{}.{}",
|
|
device.properties().compute_capability.0,
|
|
device.properties().compute_capability.1
|
|
),
|
|
device_index: i as i32,
|
|
});
|
|
}
|
|
|
|
// Create GPU topology for single node
|
|
let mut intra_node_topology = Vec::new();
|
|
if !gpus.is_empty() {
|
|
let mut gpu_connectivity = vec![vec![ConnectionType::None; device_count]; device_count];
|
|
|
|
// Simulate NVLink connections for RTX 5090
|
|
for i in 0..device_count {
|
|
for j in 0..device_count {
|
|
if i != j && (i as i32 - j as i32).abs() <= 2 {
|
|
gpu_connectivity[i][j] = ConnectionType::NvLink;
|
|
} else if i != j {
|
|
gpu_connectivity[i][j] = ConnectionType::Pcie;
|
|
}
|
|
}
|
|
}
|
|
|
|
intra_node_topology.push(GpuTopology {
|
|
node_id: 0,
|
|
gpus,
|
|
gpu_connectivity,
|
|
});
|
|
}
|
|
|
|
let topology = TopologyInfo {
|
|
num_nodes: nodes,
|
|
gpus_per_node,
|
|
interconnect: "nvlink+infiniband".to_string(),
|
|
bandwidth: TopologyBandwidth {
|
|
intra_node_gpu: 900.0, // RTX 5090 NVLink bandwidth (GB/s)
|
|
cpu_gpu: 32.0, // PCIe Gen4 x16
|
|
inter_node: 200.0, // InfiniBand HDR
|
|
memory: 1000.0, // RTX 5090 memory bandwidth
|
|
},
|
|
latency_matrix: vec![vec![0.1; nodes]; nodes], // 0.1ms latency
|
|
intra_node_topology,
|
|
inter_node_topology: NetworkTopology {
|
|
topology_type: "fat_tree".to_string(),
|
|
switch_matrix: vec![vec![1; nodes]; nodes], // Simple switch matrix
|
|
network_interfaces: vec![
|
|
NetworkInterface {
|
|
node_id: 0,
|
|
interface_name: "ib0".to_string(),
|
|
ip_address: "192.168.1.1".to_string(),
|
|
bandwidth_gbps: 200.0,
|
|
mtu: 9000,
|
|
};
|
|
nodes
|
|
],
|
|
},
|
|
};
|
|
|
|
Ok(topology)
|
|
}
|
|
|
|
/// Optimize AllReduce pattern for given topology
|
|
pub fn optimize_allreduce_pattern(&self, topology: &TopologyInfo) -> Result<AllReducePattern> {
|
|
let total_gpus = topology.total_gpus();
|
|
|
|
// Choose pattern based on topology
|
|
let algorithm = if total_gpus <= 8 && topology.num_nodes == 1 {
|
|
"nvlink_ring".to_string()
|
|
} else if topology.interconnect.contains("nvlink") {
|
|
"hierarchical_ring".to_string()
|
|
} else {
|
|
"recursive_doubling".to_string()
|
|
};
|
|
|
|
// Estimate completion time
|
|
let estimated_time_ms = if topology.num_nodes == 1 {
|
|
5.0 // Single node should be very fast with NVLink
|
|
} else {
|
|
topology.num_nodes as f64 * 10.0 // Scale with node count
|
|
};
|
|
|
|
let steps = if algorithm == "hierarchical_ring" {
|
|
vec![
|
|
CommunicationStep {
|
|
operation: "intra_node_ring".to_string(),
|
|
participants: (0..topology.gpus_per_node).collect(),
|
|
estimated_time_ms: estimated_time_ms * 0.3,
|
|
},
|
|
CommunicationStep {
|
|
operation: "inter_node_allreduce".to_string(),
|
|
participants: vec![0], // Node masters
|
|
estimated_time_ms: estimated_time_ms * 0.7,
|
|
},
|
|
]
|
|
} else {
|
|
vec![CommunicationStep {
|
|
operation: algorithm.clone(),
|
|
participants: (0..total_gpus).collect(),
|
|
estimated_time_ms,
|
|
}]
|
|
};
|
|
|
|
Ok(AllReducePattern {
|
|
algorithm,
|
|
steps,
|
|
estimated_time_ms,
|
|
expected_bandwidth_gbps: topology.bandwidth.intra_node_gpu * 0.8, // 80% efficiency
|
|
})
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
|
|
#[test]
|
|
fn test_topology_discovery() {
|
|
let topology = TopologyInfo::discover().unwrap();
|
|
assert_eq!(topology.num_nodes, 2);
|
|
assert_eq!(topology.gpus_per_node, 4);
|
|
assert_eq!(topology.total_gpus(), 8);
|
|
}
|
|
|
|
#[test]
|
|
fn test_rank_mapping() {
|
|
let topology = TopologyInfo::default_topology();
|
|
|
|
assert_eq!(topology.node_for_rank(0), 0);
|
|
assert_eq!(topology.node_for_rank(3), 0);
|
|
assert_eq!(topology.node_for_rank(4), 1);
|
|
assert_eq!(topology.node_for_rank(7), 1);
|
|
|
|
assert_eq!(topology.gpu_for_rank(0), 0);
|
|
assert_eq!(topology.gpu_for_rank(3), 3);
|
|
assert_eq!(topology.gpu_for_rank(4), 0);
|
|
assert_eq!(topology.gpu_for_rank(7), 3);
|
|
|
|
assert!(topology.same_node(0, 3));
|
|
assert!(!topology.same_node(3, 4));
|
|
}
|
|
|
|
#[test]
|
|
fn test_topology_optimizer() {
|
|
let topology = TopologyInfo::default_topology();
|
|
let mut optimizer = TopologyOptimizer::new(topology);
|
|
|
|
let result = optimizer.optimize_allreduce(8, 10.0).unwrap();
|
|
assert!(result.improvement_factor > 0.0);
|
|
assert!(result.memory_overhead_mb > 0.0);
|
|
|
|
// Test caching
|
|
let result2 = optimizer.optimize_allreduce(8, 10.0).unwrap();
|
|
assert_eq!(result.improvement_factor, result2.improvement_factor);
|
|
}
|
|
|
|
#[test]
|
|
fn test_gpu_topology() {
|
|
let gpu_topo = GpuTopology::default_for_node(0, 4);
|
|
assert_eq!(gpu_topo.gpus.len(), 4);
|
|
assert_eq!(gpu_topo.gpu_connectivity.len(), 4);
|
|
|
|
// Check NVLink connectivity
|
|
assert_eq!(gpu_topo.gpu_connectivity[0][1], ConnectionType::NvLink);
|
|
assert_eq!(gpu_topo.gpu_connectivity[0][0], ConnectionType::None);
|
|
}
|
|
|
|
#[test]
|
|
fn test_connection_type() {
|
|
assert_eq!(ConnectionType::NvLink, ConnectionType::NvLink);
|
|
assert_ne!(ConnectionType::NvLink, ConnectionType::Pcie);
|
|
}
|
|
}
|