50 lines
1.5 KiB
Rust
50 lines
1.5 KiB
Rust
//! TDD tests for NCCL device handling
|
|
//! These tests define expected behavior for device validation
|
|
|
|
#![cfg(feature = "nccl")]
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use rtx_distributed::nccl_comm::validate_tensor_for_nccl;
|
|
use rtx_tensor::{Device, Shape, Tensor};
|
|
|
|
#[test]
|
|
fn test_validate_cuda_tensor() {
|
|
// Test that CUDA tensors are valid for NCCL
|
|
if let Ok(device) = Device::cuda(0) {
|
|
let shape = Shape::new(vec![10, 10]).unwrap();
|
|
let tensor = Tensor::zeros(shape, &device).unwrap();
|
|
|
|
// CUDA tensor should be valid
|
|
assert!(validate_tensor_for_nccl(&tensor).is_ok());
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn test_validate_cpu_tensor() {
|
|
// Test that CPU tensors are handled properly
|
|
let device = Device::cpu();
|
|
let shape = Shape::new(vec![10, 10]).unwrap();
|
|
let tensor = Tensor::zeros(shape, &device).unwrap();
|
|
|
|
// CPU tensor validation behavior depends on impl
|
|
// It should either succeed (will be moved to GPU) or fail with clear error
|
|
let result = validate_tensor_for_nccl(&tensor);
|
|
assert!(result.is_ok() || result.is_err());
|
|
}
|
|
|
|
#[test]
|
|
fn test_device_pattern_matching() {
|
|
// Test that device matching works correctly
|
|
let device = Device::cpu();
|
|
|
|
let result = match device {
|
|
Device::Cuda(_) => "cuda",
|
|
Device::Cpu => "cpu",
|
|
_ => "other",
|
|
};
|
|
|
|
assert!(result == "cpu" || result == "cuda");
|
|
}
|
|
}
|