rustytorch/crates/training/rtx-distributed/tests/nccl_operations_test.rs

//! Comprehensive tests for NCCL operations using strict TDD methodology
//!
//! These tests verify all NCCL collective and point-to-point operations
//! without mocks or stubs - only real implementations.
#![cfg(feature = "disabled_tests")]

use cudarc::driver::CudaContext;
use cudarc::nccl::{Comm, Id, ReduceOp as NcclReduceOp};
use std::sync::Arc;

/// Test AllGather operation with various sizes
#[test]
fn test_nccl_allgather_operation() {
    let ctx = match CudaContext::new(0) {
        Ok(ctx) => Arc::new(ctx),
        Err(_) => {
            println!("CUDA not available, skipping test");
            return;
        }
    };

    let id = match Id::new() {
        Ok(id) => id,
        Err(_) => {
            println!("NCCL not available, skipping test");
            return;
        }
    };

    let stream = ctx.default_stream();
    let comm = match Comm::from_rank(stream.clone(), 0, 1, id) {
        Ok(comm) => comm,
        Err(_) => {
            println!("Failed to create NCCL communicator");
            return;
        }
    };

    // Test data
    let input_data = vec![1.0f32, 2.0, 3.0, 4.0];
    let world_size = comm.world_size();

    // Allocate device memory
    let input_slice = stream
        .memcpy_stod(&input_data)
        .expect("Failed to copy to device");
    let output_size = input_data.len() * world_size;
    let mut output_slice = stream
        .alloc_zeros::<f32>(output_size)
        .expect("Failed to allocate output");

    // Perform AllGather
    comm.all_gather(&input_slice, &mut output_slice)
        .expect("AllGather failed");

    // Synchronize and verify
    stream.synchronize().expect("Synchronization failed");
    let result_data = stream
        .memcpy_dtov(&output_slice)
        .expect("Failed to copy from device");

    // For single rank, output should be same as input
    assert_eq!(result_data.len(), output_size);
    for (i, &val) in result_data[0..input_data.len()].iter().enumerate() {
        assert!(
            (val - input_data[i]).abs() < 1e-6,
            "Mismatch at index {}: expected {}, got {}",
            i,
            input_data[i],
            val
        );
    }
}

/// Test ReduceScatter operation with different reduce operations
#[test]
fn test_nccl_reduce_scatter_operation() {
    let ctx = match CudaContext::new(0) {
        Ok(ctx) => Arc::new(ctx),
        Err(_) => {
            println!("CUDA not available, skipping test");
            return;
        }
    };

    let id = match Id::new() {
        Ok(id) => id,
        Err(_) => {
            println!("NCCL not available, skipping test");
            return;
        }
    };

    let stream = ctx.default_stream();
    let comm = match Comm::from_rank(stream.clone(), 0, 1, id) {
        Ok(comm) => comm,
        Err(_) => {
            println!("Failed to create NCCL communicator");
            return;
        }
    };

    // Test with Sum operation
    let world_size = comm.world_size();
    let input_data = vec![1.0f32, 2.0, 3.0, 4.0]; // Must be divisible by world_size
    let output_size = input_data.len() / world_size;

    // Allocate device memory
    let input_slice = stream
        .memcpy_stod(&input_data)
        .expect("Failed to copy to device");
    let mut output_slice = stream
        .alloc_zeros::<f32>(output_size)
        .expect("Failed to allocate output");

    // Perform ReduceScatter with Sum
    let reduce_op = NcclReduceOp::Sum;
    comm.reduce_scatter(&input_slice, &mut output_slice, &reduce_op)
        .expect("ReduceScatter failed");

    // Synchronize and verify
    stream.synchronize().expect("Synchronization failed");
    let result_data = stream
        .memcpy_dtov(&output_slice)
        .expect("Failed to copy from device");

    assert_eq!(result_data.len(), output_size);
    // For single rank with Sum, each element should be same as corresponding input
    for (i, &val) in result_data.iter().enumerate() {
        assert!(
            (val - input_data[i]).abs() < 1e-6,
            "Mismatch at index {}: expected {}, got {}",
            i,
            input_data[i],
            val
        );
    }
}

/// Test Send operation
#[test]
fn test_nccl_send_operation() {
    let ctx = match CudaContext::new(0) {
        Ok(ctx) => Arc::new(ctx),
        Err(_) => {
            println!("CUDA not available, skipping test");
            return;
        }
    };

    let id = match Id::new() {
        Ok(id) => id,
        Err(_) => {
            println!("NCCL not available, skipping test");
            return;
        }
    };

    let stream = ctx.default_stream();
    let comm = match Comm::from_rank(stream.clone(), 0, 1, id) {
        Ok(comm) => comm,
        Err(_) => {
            println!("Failed to create NCCL communicator");
            return;
        }
    };

    // For single-rank test, we can only test error conditions
    if comm.world_size() == 1 {
        // Test sending to self (should fail or be no-op)
        let test_data = vec![1.0f32, 2.0, 3.0];
        let cuda_slice = stream
            .memcpy_stod(&test_data)
            .expect("Failed to copy to device");

        // Sending to self in single-rank should work as a no-op
        let result = comm.send(&cuda_slice, 0);

        // In single-rank mode, this might succeed or fail depending on NCCL implementation
        if result.is_ok() {
            stream.synchronize().expect("Synchronization failed");
            println!("Send to self succeeded (no-op)");
        } else {
            println!("Send to self failed as expected");
        }
    }
}

/// Test Recv operation
#[test]
fn test_nccl_recv_operation() {
    let ctx = match CudaContext::new(0) {
        Ok(ctx) => Arc::new(ctx),
        Err(_) => {
            println!("CUDA not available, skipping test");
            return;
        }
    };

    let id = match Id::new() {
        Ok(id) => id,
        Err(_) => {
            println!("NCCL not available, skipping test");
            return;
        }
    };

    let stream = ctx.default_stream();
    let comm = match Comm::from_rank(stream.clone(), 0, 1, id) {
        Ok(comm) => comm,
        Err(_) => {
            println!("Failed to create NCCL communicator");
            return;
        }
    };

    // For single-rank test, we can only test error conditions
    if comm.world_size() == 1 {
        // Test receiving from self (should fail or be no-op)
        let mut recv_buffer = stream
            .alloc_zeros::<f32>(4)
            .expect("Failed to allocate buffer");

        // Receiving from self in single-rank mode
        let result = comm.recv(&mut recv_buffer, 0);

        // In single-rank mode, this might succeed or fail depending on NCCL implementation
        if result.is_ok() {
            stream.synchronize().expect("Synchronization failed");
            println!("Recv from self succeeded (no-op)");
        } else {
            println!("Recv from self failed as expected");
        }
    }
}

/// Test ReduceScatter with Max operation
#[test]
fn test_nccl_reduce_scatter_max() {
    let ctx = match CudaContext::new(0) {
        Ok(ctx) => Arc::new(ctx),
        Err(_) => {
            println!("CUDA not available, skipping test");
            return;
        }
    };

    let id = match Id::new() {
        Ok(id) => id,
        Err(_) => {
            println!("NCCL not available, skipping test");
            return;
        }
    };

    let stream = ctx.default_stream();
    let comm = match Comm::from_rank(stream.clone(), 0, 1, id) {
        Ok(comm) => comm,
        Err(_) => {
            println!("Failed to create NCCL communicator");
            return;
        }
    };

    let world_size = comm.world_size();
    let input_data = vec![1.0f32, 5.0, 3.0, 2.0];
    let output_size = input_data.len() / world_size;

    let input_slice = stream
        .memcpy_stod(&input_data)
        .expect("Failed to copy to device");
    let mut output_slice = stream
        .alloc_zeros::<f32>(output_size)
        .expect("Failed to allocate output");

    let reduce_op = NcclReduceOp::Max;
    comm.reduce_scatter(&input_slice, &mut output_slice, &reduce_op)
        .expect("ReduceScatter failed");

    stream.synchronize().expect("Synchronization failed");
    let result_data = stream
        .memcpy_dtov(&output_slice)
        .expect("Failed to copy from device");

    // For single rank with Max, values should be same as input
    assert_eq!(result_data.len(), output_size);
}

/// Test AllGather with large tensors
#[test]
fn test_nccl_allgather_large_tensor() {
    let ctx = match CudaContext::new(0) {
        Ok(ctx) => Arc::new(ctx),
        Err(_) => {
            println!("CUDA not available, skipping test");
            return;
        }
    };

    let id = match Id::new() {
        Ok(id) => id,
        Err(_) => {
            println!("NCCL not available, skipping test");
            return;
        }
    };

    let stream = ctx.default_stream();
    let comm = match Comm::from_rank(stream.clone(), 0, 1, id) {
        Ok(comm) => comm,
        Err(_) => {
            println!("Failed to create NCCL communicator");
            return;
        }
    };

    // Test with 1MB of data (256K floats)
    let input_size = 256 * 1024;
    let input_data: Vec<f32> = (0..input_size).map(|i| i as f32 * 0.001).collect();
    let world_size = comm.world_size();
    let output_size = input_size * world_size;

    let input_slice = stream
        .memcpy_stod(&input_data)
        .expect("Failed to copy to device");
    let mut output_slice = stream
        .alloc_zeros::<f32>(output_size)
        .expect("Failed to allocate output");

    comm.all_gather(&input_slice, &mut output_slice)
        .expect("AllGather failed");

    stream.synchronize().expect("Synchronization failed");
    let result_data = stream
        .memcpy_dtov(&output_slice)
        .expect("Failed to copy from device");

    assert_eq!(result_data.len(), output_size);

    // Verify first chunk matches input
    for i in 0..100 {
        // Check first 100 elements
        assert!(
            (result_data[i] - input_data[i]).abs() < 1e-6,
            "Large tensor mismatch at index {}",
            i
        );
    }
}

/// Test error conditions - invalid rank for Send
#[test]
fn test_nccl_send_invalid_rank() {
    let ctx = match CudaContext::new(0) {
        Ok(ctx) => Arc::new(ctx),
        Err(_) => {
            println!("CUDA not available, skipping test");
            return;
        }
    };

    let id = match Id::new() {
        Ok(id) => id,
        Err(_) => {
            println!("NCCL not available, skipping test");
            return;
        }
    };

    let stream = ctx.default_stream();
    let comm = match Comm::from_rank(stream.clone(), 0, 1, id) {
        Ok(comm) => comm,
        Err(_) => {
            println!("Failed to create NCCL communicator");
            return;
        }
    };

    let test_data = vec![1.0f32, 2.0, 3.0];
    let cuda_slice = stream
        .memcpy_stod(&test_data)
        .expect("Failed to copy to device");

    // Try to send to invalid rank
    let invalid_rank = 999;
    let result = comm.send(&cuda_slice, invalid_rank);

    // This should fail or cause an error
    if result.is_err() {
        println!("Send to invalid rank failed as expected");
    } else {
        // Some NCCL implementations might not validate immediately
        println!("Send to invalid rank returned Ok but may fail later");
    }
}