271 lines
7.7 KiB
Rust
271 lines
7.7 KiB
Rust
//! Comprehensive tests for NCCL distributed training functionality
|
|
//!
|
|
//! These tests are currently disabled because the APIs they test
|
|
//! (NCCLBackend, NCCLConfig, CommunicationGroup) are not yet implemented
|
|
//! in the current nccl.rs module which uses cudarc's NCCL bindings.
|
|
|
|
#![cfg(feature = "nccl")]
|
|
#![cfg(disabled)] // Disable entire file until APIs are implemented
|
|
|
|
#[allow(unused_imports)]
|
|
use rtx_distributed::{
|
|
DistributedError,
|
|
nccl::{CommunicationGroup, NCCLBackend, NCCLConfig},
|
|
};
|
|
#[allow(unused_imports)]
|
|
use std::sync::Arc;
|
|
#[allow(unused_imports)]
|
|
use std::time::Duration;
|
|
|
|
#[cfg(test)]
|
|
mod nccl_tests {
|
|
use super::*;
|
|
|
|
fn create_test_config() -> NCCLConfig {
|
|
NCCLConfig {
|
|
socket_ifname: "lo".to_string(),
|
|
debug_level: "WARN".to_string(),
|
|
tree_threshold: 0,
|
|
num_gpus: 1,
|
|
rank: 0,
|
|
world_size: 1,
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn test_nccl_config_creation() {
|
|
let config = create_test_config();
|
|
assert_eq!(config.socket_ifname, "lo");
|
|
assert_eq!(config.debug_level, "WARN");
|
|
assert_eq!(config.tree_threshold, 0);
|
|
assert_eq!(config.num_gpus, 1);
|
|
assert_eq!(config.rank, 0);
|
|
assert_eq!(config.world_size, 1);
|
|
}
|
|
|
|
#[test]
|
|
fn test_nccl_backend_initialization() {
|
|
let config = create_test_config();
|
|
let backend = NCCLBackend::new(config);
|
|
assert!(backend.is_ok());
|
|
}
|
|
|
|
#[test]
|
|
fn test_environment_setup() {
|
|
let config = NCCLConfig {
|
|
socket_ifname: "eth0".to_string(),
|
|
debug_level: "INFO".to_string(),
|
|
tree_threshold: 100,
|
|
num_gpus: 2,
|
|
rank: 0,
|
|
world_size: 2,
|
|
};
|
|
|
|
let result = NCCLBackend::setup_environment(&config);
|
|
assert!(result.is_ok());
|
|
|
|
// Verify environment variables were set (in unsafe block)
|
|
unsafe {
|
|
assert_eq!(std::env::var("NCCL_SOCKET_IFNAME").unwrap(), "eth0");
|
|
assert_eq!(std::env::var("NCCL_DEBUG").unwrap(), "INFO");
|
|
assert_eq!(std::env::var("NCCL_TREE_THRESHOLD").unwrap(), "100");
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn test_communication_group_creation() {
|
|
let config = create_test_config();
|
|
let backend = NCCLBackend::new(config).unwrap();
|
|
|
|
let group = backend.create_communication_group(vec![0]);
|
|
assert!(group.is_ok());
|
|
|
|
let group = group.unwrap();
|
|
assert_eq!(group.ranks(), vec![0]);
|
|
assert_eq!(group.size(), 1);
|
|
}
|
|
|
|
#[test]
|
|
fn test_multi_gpu_configuration() {
|
|
let config = NCCLConfig {
|
|
socket_ifname: "ib0".to_string(),
|
|
debug_level: "WARN".to_string(),
|
|
tree_threshold: 0,
|
|
num_gpus: 4,
|
|
rank: 0,
|
|
world_size: 1,
|
|
};
|
|
|
|
let backend = NCCLBackend::new(config);
|
|
assert!(backend.is_ok());
|
|
|
|
let backend = backend.unwrap();
|
|
assert_eq!(backend.num_gpus(), 4);
|
|
}
|
|
|
|
#[test]
|
|
fn test_distributed_configuration() {
|
|
let config = NCCLConfig {
|
|
socket_ifname: "lo".to_string(),
|
|
debug_level: "ERROR".to_string(),
|
|
tree_threshold: 50,
|
|
num_gpus: 2,
|
|
rank: 1,
|
|
world_size: 4,
|
|
};
|
|
|
|
let backend = NCCLBackend::new(config);
|
|
assert!(backend.is_ok());
|
|
|
|
let backend = backend.unwrap();
|
|
assert_eq!(backend.rank(), 1);
|
|
assert_eq!(backend.world_size(), 4);
|
|
}
|
|
|
|
#[test]
|
|
fn test_all_reduce_operation() {
|
|
let config = create_test_config();
|
|
let mut backend = NCCLBackend::new(config).unwrap();
|
|
|
|
// Create test tensor
|
|
let data = vec![1.0f32, 2.0, 3.0, 4.0];
|
|
let result = backend.all_reduce(data.clone(), "sum");
|
|
|
|
assert!(result.is_ok());
|
|
// In single-node mode, all_reduce should return the same data
|
|
assert_eq!(result.unwrap(), data);
|
|
}
|
|
|
|
#[test]
|
|
fn test_broadcast_operation() {
|
|
let config = create_test_config();
|
|
let mut backend = NCCLBackend::new(config).unwrap();
|
|
|
|
let data = vec![5.0f32, 6.0, 7.0, 8.0];
|
|
let result = backend.broadcast(data.clone(), 0);
|
|
|
|
assert!(result.is_ok());
|
|
assert_eq!(result.unwrap(), data);
|
|
}
|
|
|
|
#[test]
|
|
fn test_reduce_operation() {
|
|
let config = create_test_config();
|
|
let mut backend = NCCLBackend::new(config).unwrap();
|
|
|
|
let data = vec![1.0f32, 1.0, 1.0, 1.0];
|
|
let result = backend.reduce(data.clone(), 0, "sum");
|
|
|
|
assert!(result.is_ok());
|
|
// In single-node mode, reduce should return the same data
|
|
assert_eq!(result.unwrap(), data);
|
|
}
|
|
|
|
#[test]
|
|
fn test_all_gather_operation() {
|
|
let config = create_test_config();
|
|
let mut backend = NCCLBackend::new(config).unwrap();
|
|
|
|
let data = vec![9.0f32, 10.0];
|
|
let result = backend.all_gather(data.clone());
|
|
|
|
assert!(result.is_ok());
|
|
// In single-node mode, all_gather should return the same data
|
|
assert_eq!(result.unwrap(), vec![data]);
|
|
}
|
|
|
|
#[test]
|
|
fn test_invalid_rank() {
|
|
let config = NCCLConfig {
|
|
socket_ifname: "lo".to_string(),
|
|
debug_level: "WARN".to_string(),
|
|
tree_threshold: 0,
|
|
num_gpus: 1,
|
|
rank: 5,
|
|
world_size: 4, // rank >= world_size is invalid
|
|
};
|
|
|
|
let backend = NCCLBackend::new(config);
|
|
assert!(backend.is_err());
|
|
}
|
|
|
|
#[test]
|
|
fn test_barrier_synchronization() {
|
|
let config = create_test_config();
|
|
let mut backend = NCCLBackend::new(config).unwrap();
|
|
|
|
let result = backend.barrier();
|
|
assert!(result.is_ok());
|
|
}
|
|
|
|
#[test]
|
|
fn test_scatter_operation() {
|
|
let config = create_test_config();
|
|
let mut backend = NCCLBackend::new(config).unwrap();
|
|
|
|
let data = vec![vec![1.0f32, 2.0], vec![3.0, 4.0]];
|
|
let result = backend.scatter(data, 0);
|
|
|
|
assert!(result.is_ok());
|
|
}
|
|
|
|
#[test]
|
|
fn test_gather_operation() {
|
|
let config = create_test_config();
|
|
let mut backend = NCCLBackend::new(config).unwrap();
|
|
|
|
let data = vec![11.0f32, 12.0];
|
|
let result = backend.gather(data.clone(), 0);
|
|
|
|
assert!(result.is_ok());
|
|
// In single-node mode, gather returns data wrapped in vec
|
|
assert_eq!(result.unwrap(), vec![data]);
|
|
}
|
|
|
|
#[test]
|
|
fn test_reduce_scatter_operation() {
|
|
let config = create_test_config();
|
|
let mut backend = NCCLBackend::new(config).unwrap();
|
|
|
|
let data = vec![1.0f32, 2.0, 3.0, 4.0];
|
|
let result = backend.reduce_scatter(data.clone(), "sum");
|
|
|
|
assert!(result.is_ok());
|
|
}
|
|
|
|
#[test]
|
|
fn test_point_to_point_send() {
|
|
let config = create_test_config();
|
|
let mut backend = NCCLBackend::new(config).unwrap();
|
|
|
|
let data = vec![13.0f32, 14.0, 15.0];
|
|
// In single-node, send to self should work
|
|
let result = backend.send(data, 0);
|
|
|
|
assert!(result.is_ok());
|
|
}
|
|
|
|
#[test]
|
|
fn test_point_to_point_recv() {
|
|
let config = create_test_config();
|
|
let mut backend = NCCLBackend::new(config).unwrap();
|
|
|
|
// In single-node mode, this will likely timeout or return empty
|
|
let result = backend.recv::<Vec<f32>>(0, 4);
|
|
|
|
// This might fail in single-node, which is expected
|
|
assert!(result.is_ok() || result.is_err());
|
|
}
|
|
|
|
#[test]
|
|
fn test_communicator_cleanup() {
|
|
let config = create_test_config();
|
|
let backend = NCCLBackend::new(config).unwrap();
|
|
|
|
// Test that dropping the backend properly cleans up resources
|
|
drop(backend);
|
|
// If we get here without panicking, cleanup worked
|
|
assert!(true);
|
|
}
|
|
}
|