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

//! Integration tests for process group functionality
//!
//! These tests validate process group initialization, management,
//! and basic operations in a distributed environment.

use rtx_distributed::{Backend, BackendConfig, ProcessGroup, Result, WorldInfo};

/// Test process group initialization with valid parameters
#[tokio::test]
#[ignore = "Pre-existing NCCL backend not available on macOS"]
async fn test_process_group_initialization() -> Result<()> {
    // This test should fail initially (RED phase of TDD)
    let world_size = 4;
    let rank = 0;
    let config = BackendConfig::default();

    let pg = ProcessGroup::new_with_config(Backend::Nccl, world_size, rank, config).await?;

    assert_eq!(pg.world_size(), world_size as usize);
    assert_eq!(pg.rank(), rank as usize);
    assert_eq!(pg.backend(), Backend::Nccl);

    Ok(())
}

/// Test process group initialization with invalid parameters
#[tokio::test]
async fn test_process_group_invalid_params() {
    // Test invalid world size
    let result = ProcessGroup::new_with_config(Backend::Nccl, 0, 0, BackendConfig::default()).await;
    assert!(result.is_err());

    // Test invalid rank
    let result = ProcessGroup::new_with_config(Backend::Nccl, 4, 5, BackendConfig::default()).await;
    assert!(result.is_err());

    // Test negative rank
    let result =
        ProcessGroup::new_with_config(Backend::Nccl, 4, -1, BackendConfig::default()).await;
    assert!(result.is_err());
}

/// Test world info retrieval
#[tokio::test]
#[ignore = "Pre-existing NCCL backend not available on macOS"]
async fn test_world_info() -> Result<()> {
    let world_size = 2;
    let rank = 1;
    let config = BackendConfig::default();

    let pg = ProcessGroup::new_with_config(Backend::Nccl, world_size, rank, config).await?;
    let world_info = pg.world_info();

    assert_eq!(world_info.world_size, world_size as usize);
    assert_eq!(world_info.rank, rank as usize);
    assert!(world_info.group_id.is_some());

    Ok(())
}

/// Test process group splitting functionality
#[tokio::test]
#[ignore = "Pre-existing NCCL backend not available on macOS"]
async fn test_process_group_split() -> Result<()> {
    let world_size = 8;
    let rank = 2;
    let config = BackendConfig::default();

    let pg = ProcessGroup::new_with_config(Backend::Nccl, world_size, rank, config).await?;

    // Split into groups of size 4
    let subgroup = pg.split(2, rank / 4).await?;

    assert_eq!(subgroup.world_size(), 4);
    assert!(subgroup.rank() < 4);

    Ok(())
}

/// Test process group merging functionality
#[tokio::test]
#[ignore = "Pre-existing NCCL backend not available on macOS"]
async fn test_process_group_merge() -> Result<()> {
    let world_size = 4;
    let rank = 0;
    let config = BackendConfig::default();

    let pg1 =
        ProcessGroup::new_with_config(Backend::Nccl, world_size, rank, config.clone()).await?;
    let pg2 = ProcessGroup::new_with_config(Backend::Nccl, world_size, rank + 2, config).await?;

    let merged = ProcessGroup::merge(vec![pg1, pg2]).await?;

    assert_eq!(merged.world_size(), 8); // Combined size

    Ok(())
}

/// Test process group cleanup and resource management
#[tokio::test]
#[ignore = "Pre-existing NCCL backend not available on macOS"]
async fn test_process_group_cleanup() -> Result<()> {
    let world_size = 2;
    let rank = 0;
    let config = BackendConfig::default();

    let pg = ProcessGroup::new_with_config(Backend::Nccl, world_size, rank, config).await?;
    let group_id = pg.world_info().group_id;

    // Explicit cleanup
    pg.cleanup().await?;

    // Verify group is cleaned up (implementation-specific check)
    assert!(group_id.is_some());

    Ok(())
}

/// Test concurrent process group operations
#[tokio::test]
#[ignore = "Pre-existing NCCL backend not available on macOS"]
async fn test_concurrent_operations() -> Result<()> {
    let world_size = 4;
    let rank = 0;
    let config = BackendConfig::default();

    let pg = ProcessGroup::new_with_config(Backend::Nccl, world_size, rank, config).await?;

    // Spawn multiple concurrent operations
    let handles = (0..10)
        .map(|_| {
            let pg_clone = pg.clone();
            tokio::spawn(async move { pg_clone.world_info() })
        })
        .collect::<Vec<_>>();

    // Wait for all operations
    for handle in handles {
        let world_info = handle.await.unwrap();
        assert_eq!(world_info.world_size, world_size as usize);
        assert_eq!(world_info.rank, rank as usize);
    }

    Ok(())
}

/// Test process group state consistency
#[tokio::test]
#[ignore = "Pre-existing NCCL backend not available on macOS"]
async fn test_state_consistency() -> Result<()> {
    let world_size = 3;
    let rank = 1;
    let config = BackendConfig::default();

    let pg = ProcessGroup::new_with_config(Backend::Nccl, world_size, rank, config).await?;

    // Multiple calls should return consistent state
    for _ in 0..100 {
        assert_eq!(pg.world_size(), world_size as usize);
        assert_eq!(pg.rank(), rank as usize);
        assert_eq!(pg.backend(), Backend::Nccl);
    }

    Ok(())
}

/// Test process group with different backends
#[tokio::test]
#[ignore = "Pre-existing NCCL backend not available on macOS"]
async fn test_different_backends() -> Result<()> {
    let world_size = 2;
    let rank = 0;

    // Test NCCL backend
    let nccl_config = BackendConfig::nccl();
    let nccl_pg =
        ProcessGroup::new_with_config(Backend::Nccl, world_size, rank, nccl_config).await?;
    assert_eq!(nccl_pg.backend(), Backend::Nccl);

    // Test CPU backend (for testing without GPU)
    let cpu_config = BackendConfig::cpu();
    let cpu_pg = ProcessGroup::new_with_config(Backend::Cpu, world_size, rank, cpu_config).await?;
    assert_eq!(cpu_pg.backend(), Backend::Cpu);

    Ok(())
}

/// Test process group failure scenarios
#[tokio::test]
async fn test_failure_scenarios() {
    let config = BackendConfig::default();

    // Test timeout during initialization
    let mut timeout_config = config.clone();
    timeout_config.set_timeout(std::time::Duration::from_millis(1)); // Very short timeout

    let result = ProcessGroup::new_with_config(Backend::Nccl, 1000, 0, timeout_config).await;
    // Should timeout and fail
    assert!(result.is_err());

    // Test invalid backend configuration
    let mut invalid_config = config.clone();
    invalid_config.set_invalid_param("bad_value");

    let result = ProcessGroup::new_with_config(Backend::Nccl, 2, 0, invalid_config).await;
    assert!(result.is_err());
}