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

//! Standalone NCCL functionality test
//!
//! This test verifies that cudarc NCCL integration works correctly
//! without depending on the main rtx-distributed library.
#![cfg(feature = "disabled_tests")]

use std::sync::Arc;

// Only test if the nccl feature is enabled
#[cfg(feature = "nccl")]
mod nccl_tests {
    use super::*;
    use cudarc::driver::CudaContext;
    use cudarc::nccl::{Comm, Id, ReduceOp as NcclReduceOp};

    /// Test that we can create a CUDA context
    #[test]
    fn test_cuda_context_creation() {
        match CudaContext::new(0) {
            Ok(ctx) => {
                assert_eq!(ctx.ordinal(), 0);
                println!("✅ CUDA context creation: PASSED");
            }
            Err(e) => {
                println!("⚠️  CUDA not available, skipping test: {:?}", e);
            }
        }
    }

    /// Test NCCL ID generation
    #[test]
    fn test_nccl_id_generation() {
        match Id::new() {
            Ok(id1) => {
                let id2 = Id::new().expect("Second ID generation should work");
                // IDs should be different (extremely unlikely to be the same)
                assert_ne!(id1.internal(), id2.internal());
                println!("✅ NCCL ID generation: PASSED");
            }
            Err(e) => {
                println!("⚠️  NCCL not available, skipping test: {:?}", e);
            }
        }
    }

    /// Test NCCL communicator creation with single rank
    #[test]
    fn test_nccl_communicator_single_rank() {
        let ctx = match CudaContext::new(0) {
            Ok(ctx) => Arc::new(ctx),
            Err(e) => {
                println!("⚠️  CUDA not available, skipping test: {:?}", e);
                return;
            }
        };

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

        let stream = ctx.default_stream();
        match Comm::from_rank(stream, 0, 1, id) {
            Ok(comm) => {
                assert_eq!(comm.rank(), 0);
                assert_eq!(comm.world_size(), 1);
                println!("✅ NCCL communicator creation: PASSED");
            }
            Err(e) => {
                println!("❌ NCCL communicator creation failed: {:?}", e);
                panic!("NCCL communicator creation should work in single-rank mode");
            }
        }
    }

    /// Test basic NCCL all-reduce operation
    #[test]
    fn test_nccl_allreduce_single_rank() {
        let ctx = match CudaContext::new(0) {
            Ok(ctx) => Arc::new(ctx),
            Err(e) => {
                println!("⚠️  CUDA not available, skipping test: {:?}", e);
                return;
            }
        };

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

        let stream = ctx.default_stream();
        let comm = match Comm::from_rank(stream.clone(), 0, 1, id) {
            Ok(comm) => comm,
            Err(e) => {
                println!(
                    "⚠️  NCCL communicator creation failed, skipping test: {:?}",
                    e
                );
                return;
            }
        };

        // Create test data
        let input_data = vec![1.0f32, 2.0, 3.0, 4.0];
        let expected_data = input_data.clone(); // Single rank, so no change expected

        // Test the complete memory allocation and operation pipeline
        let input_slice = stream
            .memcpy_stod(&input_data)
            .expect("Input memory allocation should work");

        let mut output_slice = stream
            .alloc_zeros::<f32>(input_data.len())
            .expect("Output memory allocation should work");

        // Perform AllReduce with Sum operation
        let reduce_op = NcclReduceOp::Sum;
        comm.all_reduce(&input_slice, &mut output_slice, &reduce_op)
            .expect("AllReduce operation should work");

        // Synchronize to ensure operation completes
        stream
            .synchronize()
            .expect("Stream synchronization should work");

        // Copy result back to host memory
        let result_data = stream
            .memcpy_dtov(&output_slice)
            .expect("Memory copy to host should work");

        // Verify the results
        assert_eq!(result_data.len(), expected_data.len());
        for (i, (&actual, &expected)) in result_data.iter().zip(expected_data.iter()).enumerate() {
            assert!(
                (actual - expected).abs() < 1e-6,
                "Mismatch at index {}: expected {}, got {}",
                i,
                expected,
                actual
            );
        }

        println!("✅ NCCL AllReduce single rank: PASSED");
        println!("   Input:  {:?}", input_data);
        println!("   Output: {:?}", result_data);
    }

    /// Test NCCL broadcast operation
    #[test]
    fn test_nccl_broadcast_single_rank() {
        let ctx = match CudaContext::new(0) {
            Ok(ctx) => Arc::new(ctx),
            Err(e) => {
                println!("⚠️  CUDA not available, skipping test: {:?}", e);
                return;
            }
        };

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

        let stream = ctx.default_stream();
        let comm = match Comm::from_rank(stream.clone(), 0, 1, id) {
            Ok(comm) => comm,
            Err(e) => {
                println!(
                    "⚠️  NCCL communicator creation failed, skipping test: {:?}",
                    e
                );
                return;
            }
        };

        // Create test data
        let input_data = vec![5.0f32, 6.0, 7.0, 8.0];
        let expected_data = input_data.clone();

        let input_slice = stream
            .memcpy_stod(&input_data)
            .expect("Input memory allocation should work");

        let mut output_slice = stream
            .alloc_zeros::<f32>(input_data.len())
            .expect("Output memory allocation should work");

        // For single rank, root is always 0 and sendbuff is Some
        comm.broadcast(Some(&input_slice), &mut output_slice, 0)
            .expect("Broadcast operation should work");

        stream
            .synchronize()
            .expect("Stream synchronization should work");

        let result_data = stream
            .memcpy_dtov(&output_slice)
            .expect("Memory copy to host should work");

        // Verify the results
        assert_eq!(result_data.len(), expected_data.len());
        for (i, (&actual, &expected)) in result_data.iter().zip(expected_data.iter()).enumerate() {
            assert!(
                (actual - expected).abs() < 1e-6,
                "Mismatch at index {}: expected {}, got {}",
                i,
                expected,
                actual
            );
        }

        println!("✅ NCCL Broadcast single rank: PASSED");
    }
}

// If nccl feature is not enabled, provide a dummy test
#[cfg(not(feature = "nccl"))]
#[test]
fn test_nccl_feature_disabled() {
    println!("⚠️  NCCL feature is disabled. Enable with --features nccl");
}