330 lines
11 KiB
Rust
330 lines
11 KiB
Rust
#![cfg(feature = "nccl")]
|
|
#![cfg(disabled)] // Disable until NCCL communicator APIs are fully implemented
|
|
|
|
#[allow(unused_imports)]
|
|
use cudarc::driver::CudaContext;
|
|
#[allow(unused_imports)]
|
|
use rtx_distributed::{Device, Result, comm::ReduceOp, error::DistributedError, nccl};
|
|
#[allow(unused_imports)]
|
|
use rtx_tensor::{DType, Shape, Tensor};
|
|
#[allow(unused_imports)]
|
|
use std::sync::Arc;
|
|
#[allow(unused_imports)]
|
|
use tokio::time::{Duration, timeout};
|
|
|
|
/// Integration tests for NCCL distributed operations
|
|
/// These tests define the expected behavior and drive implementation (TDD)
|
|
///
|
|
/// NOTE: These tests are currently disabled because they require the full
|
|
/// NCCL communicator API to be implemented.
|
|
#[cfg(test)]
|
|
mod nccl_integration_tests {
|
|
use super::*;
|
|
|
|
/// Test basic NCCL communicator creation and initialization
|
|
#[tokio::test]
|
|
async fn test_nccl_communicator_creation() -> Result<()> {
|
|
// Skip if no CUDA devices available
|
|
if !cuda_available() {
|
|
println!("CUDA not available, skipping NCCL tests");
|
|
return Ok(());
|
|
}
|
|
|
|
let world_size = 1;
|
|
let rank = 0;
|
|
|
|
// Create CUDA context
|
|
let device = Arc::new(CudaContext::new(0).map_err(|e| {
|
|
DistributedError::runtime(format!("Failed to create CUDA context: {:?}", e))
|
|
})?);
|
|
|
|
// Create NCCL ID for coordination
|
|
let nccl_id = nccl::NcclCommunicator::get_unique_id()?;
|
|
|
|
// Create communicator
|
|
let config = nccl::NcclConfig::default();
|
|
let comm = nccl::NcclCommunicator::new(world_size, rank, &nccl_id, device, config)?;
|
|
|
|
// Verify communicator properties
|
|
assert_eq!(comm.world_size(), world_size);
|
|
assert_eq!(comm.rank(), rank);
|
|
|
|
Ok(())
|
|
}
|
|
|
|
/// Test NCCL AllReduce operation with sum reduction
|
|
#[tokio::test]
|
|
async fn test_nccl_allreduce_sum() -> Result<()> {
|
|
if !cuda_available() {
|
|
println!("CUDA not available, skipping NCCL tests");
|
|
return Ok(());
|
|
}
|
|
|
|
let device = Device::Cuda(0);
|
|
let comm = create_test_communicator(device.clone())?;
|
|
|
|
// Create test tensor
|
|
let shape = vec![4, 4];
|
|
let mut tensor = Tensor::ones(&shape, &Device::Cuda(0))?;
|
|
|
|
// Perform AllReduce with Sum operation
|
|
let result = timeout(
|
|
Duration::from_secs(5),
|
|
comm.allreduce(&mut tensor, ReduceOp::Sum),
|
|
)
|
|
.await;
|
|
|
|
match result {
|
|
Ok(Ok(())) => {
|
|
// Verify tensor values (with world_size=1, values should remain 1.0)
|
|
let data = tensor.to_vec1::<f32>()?;
|
|
for &val in data.iter() {
|
|
assert!((val - 1.0).abs() < 1e-6, "Expected 1.0, got {}", val);
|
|
}
|
|
}
|
|
Ok(Err(e)) => panic!("AllReduce failed: {:?}", e),
|
|
Err(_) => panic!("AllReduce timed out"),
|
|
}
|
|
|
|
Ok(())
|
|
}
|
|
|
|
/// Test NCCL Broadcast operation
|
|
#[tokio::test]
|
|
async fn test_nccl_broadcast() -> Result<()> {
|
|
if !cuda_available() {
|
|
println!("CUDA not available, skipping NCCL tests");
|
|
return Ok(());
|
|
}
|
|
|
|
let device = Device::Cuda(0);
|
|
let comm = create_test_communicator(device.clone())?;
|
|
|
|
// Create test tensor with specific values
|
|
let shape = Shape::new(vec![3, 3])?;
|
|
let mut tensor = Tensor::zeros(shape, DType::F32, &device)?;
|
|
|
|
// Set root rank values
|
|
if comm.rank() == 0 {
|
|
let data: Vec<f32> = (0..9).map(|i| i as f32).collect();
|
|
tensor.copy_from_cpu(&data)?;
|
|
}
|
|
|
|
// Perform Broadcast from rank 0
|
|
let result = timeout(Duration::from_secs(5), comm.broadcast(&mut tensor, 0)).await;
|
|
|
|
match result {
|
|
Ok(Ok(())) => {
|
|
// Verify all ranks have the broadcasted values
|
|
let data = tensor.to_vec1::<f32>()?;
|
|
for (i, &val) in data.iter().enumerate() {
|
|
assert!(
|
|
(val - i as f32).abs() < 1e-6,
|
|
"Expected {}, got {} at index {}",
|
|
i,
|
|
val,
|
|
i
|
|
);
|
|
}
|
|
}
|
|
Ok(Err(e)) => panic!("Broadcast failed: {:?}", e),
|
|
Err(_) => panic!("Broadcast timed out"),
|
|
}
|
|
|
|
Ok(())
|
|
}
|
|
|
|
/// Test NCCL AllGather operation
|
|
#[tokio::test]
|
|
async fn test_nccl_allgather() -> Result<()> {
|
|
if !cuda_available() {
|
|
println!("CUDA not available, skipping NCCL tests");
|
|
return Ok(());
|
|
}
|
|
|
|
let device = Device::Cuda(0);
|
|
let comm = create_test_communicator(device.clone())?;
|
|
|
|
// Create input tensor with rank-specific values
|
|
let input_shape = Shape::new(vec![2, 2])?;
|
|
let mut input_tensor = Tensor::zeros(input_shape, DType::F32, &device)?;
|
|
|
|
// Fill with rank-specific data
|
|
let rank_data: Vec<f32> = (0..4)
|
|
.map(|i| (comm.rank() as f32) * 10.0 + i as f32)
|
|
.collect();
|
|
input_tensor.copy_from_cpu(&rank_data)?;
|
|
|
|
// Perform AllGather
|
|
let result = timeout(Duration::from_secs(5), comm.allgather(&input_tensor)).await;
|
|
|
|
match result {
|
|
Ok(Ok(output_tensor)) => {
|
|
// Verify output tensor shape (should be [world_size * 2, 2])
|
|
let expected_shape = Shape::new(vec![comm.world_size() * 2, 2])?;
|
|
assert_eq!(output_tensor.shape(), &expected_shape);
|
|
|
|
// Verify gathered data
|
|
let output_data = output_tensor.to_vec1::<f32>()?;
|
|
assert_eq!(output_data.len(), comm.world_size() * 4);
|
|
}
|
|
Ok(Err(e)) => panic!("AllGather failed: {:?}", e),
|
|
Err(_) => panic!("AllGather timed out"),
|
|
}
|
|
|
|
Ok(())
|
|
}
|
|
|
|
/// Test NCCL ReduceScatter operation
|
|
#[tokio::test]
|
|
async fn test_nccl_reduce_scatter() -> Result<()> {
|
|
if !cuda_available() {
|
|
println!("CUDA not available, skipping NCCL tests");
|
|
return Ok(());
|
|
}
|
|
|
|
let device = Device::Cuda(0);
|
|
let comm = create_test_communicator(device.clone())?;
|
|
|
|
// Create input tensor
|
|
let input_shape = Shape::new(vec![4, 2])?; // Will be scattered to [2, 2] per rank
|
|
let mut input_tensor = Tensor::ones(input_shape, DType::F32, &device)?;
|
|
|
|
// Perform ReduceScatter with Sum operation
|
|
let result = timeout(
|
|
Duration::from_secs(5),
|
|
comm.reduce_scatter(&input_tensor, ReduceOp::Sum),
|
|
)
|
|
.await;
|
|
|
|
match result {
|
|
Ok(Ok(output_tensor)) => {
|
|
// Verify output tensor shape
|
|
let expected_shape = Shape::new(vec![2, 2])?;
|
|
assert_eq!(output_tensor.shape(), &expected_shape);
|
|
|
|
// Verify reduced values
|
|
let output_data = output_tensor.to_vec1::<f32>()?;
|
|
for &val in output_data.iter() {
|
|
assert!(
|
|
(val - comm.world_size() as f32).abs() < 1e-6,
|
|
"Expected {}, got {}",
|
|
comm.world_size(),
|
|
val
|
|
);
|
|
}
|
|
}
|
|
Ok(Err(e)) => panic!("ReduceScatter failed: {:?}", e),
|
|
Err(_) => panic!("ReduceScatter timed out"),
|
|
}
|
|
|
|
Ok(())
|
|
}
|
|
|
|
/// Test NCCL point-to-point send/recv operations
|
|
#[tokio::test]
|
|
async fn test_nccl_send_recv() -> Result<()> {
|
|
if !cuda_available() {
|
|
println!("CUDA not available, skipping NCCL tests");
|
|
return Ok(());
|
|
}
|
|
|
|
let device = Device::Cuda(0);
|
|
let comm = create_test_communicator(device.clone())?;
|
|
|
|
if comm.world_size() < 2 {
|
|
println!("Need at least 2 ranks for send/recv test, skipping");
|
|
return Ok(());
|
|
}
|
|
|
|
let shape = Shape::new(vec![3, 3])?;
|
|
|
|
if comm.rank() == 0 {
|
|
// Sender: create tensor with specific values and send to rank 1
|
|
let mut send_tensor = Tensor::zeros(shape, DType::F32, &device)?;
|
|
let send_data: Vec<f32> = (0..9).map(|i| i as f32 + 100.0).collect();
|
|
send_tensor.copy_from_cpu(&send_data)?;
|
|
|
|
let result = timeout(Duration::from_secs(5), comm.send(&send_tensor, 1)).await;
|
|
|
|
match result {
|
|
Ok(Ok(())) => (),
|
|
Ok(Err(e)) => panic!("Send failed: {:?}", e),
|
|
Err(_) => panic!("Send timed out"),
|
|
}
|
|
} else if comm.rank() == 1 {
|
|
// Receiver: receive tensor from rank 0
|
|
let mut recv_tensor = Tensor::zeros(shape, DType::F32, &device)?;
|
|
|
|
let result = timeout(Duration::from_secs(5), comm.recv(&mut recv_tensor, 0)).await;
|
|
|
|
match result {
|
|
Ok(Ok(())) => {
|
|
// Verify received data
|
|
let recv_data = recv_tensor.to_vec1::<f32>()?;
|
|
for (i, &val) in recv_data.iter().enumerate() {
|
|
assert!(
|
|
(val - (i as f32 + 100.0)).abs() < 1e-6,
|
|
"Expected {}, got {} at index {}",
|
|
i as f32 + 100.0,
|
|
val,
|
|
i
|
|
);
|
|
}
|
|
}
|
|
Ok(Err(e)) => panic!("Recv failed: {:?}", e),
|
|
Err(_) => panic!("Recv timed out"),
|
|
}
|
|
}
|
|
|
|
Ok(())
|
|
}
|
|
|
|
/// Test NCCL communicator synchronization
|
|
#[tokio::test]
|
|
async fn test_nccl_synchronization() -> Result<()> {
|
|
if !cuda_available() {
|
|
println!("CUDA not available, skipping NCCL tests");
|
|
return Ok(());
|
|
}
|
|
|
|
let device = Device::Cuda(0);
|
|
let comm = create_test_communicator(device.clone())?;
|
|
|
|
// Test synchronization
|
|
let result = timeout(Duration::from_secs(5), comm.synchronize()).await;
|
|
|
|
match result {
|
|
Ok(Ok(())) => (),
|
|
Ok(Err(e)) => panic!("Synchronization failed: {:?}", e),
|
|
Err(_) => panic!("Synchronization timed out"),
|
|
}
|
|
|
|
Ok(())
|
|
}
|
|
|
|
// Helper functions
|
|
|
|
/// Check if CUDA is available for testing
|
|
fn cuda_available() -> bool {
|
|
// Try to create a CUDA context to check availability
|
|
CudaContext::new(0).is_ok()
|
|
}
|
|
|
|
/// Create a test NCCL communicator for single-GPU testing
|
|
fn create_test_communicator(_device: Device) -> Result<nccl::NcclCommunicator> {
|
|
let world_size = 1;
|
|
let rank = 0;
|
|
|
|
// Create CUDA context
|
|
let device = Arc::new(CudaContext::new(0).map_err(|e| {
|
|
DistributedError::runtime(format!("Failed to create CUDA context: {:?}", e))
|
|
})?);
|
|
|
|
let nccl_id = nccl::NcclCommunicator::get_unique_id()?;
|
|
let config = nccl::NcclConfig::default();
|
|
|
|
nccl::NcclCommunicator::new(world_size, rank, &nccl_id, device, config)
|
|
}
|
|
}
|