923 lines
29 KiB
Rust
923 lines
29 KiB
Rust
//! Process group management for distributed training
|
|
//!
|
|
//! This module provides the core ProcessGroup abstraction that manages
|
|
//! distributed process communication, including initialization, splitting,
|
|
//! merging, and cleanup operations.
|
|
|
|
use crate::backend::{Backend, BackendConfig, BackendImpl};
|
|
#[cfg(any(feature = "rnccl", feature = "nccl"))]
|
|
use crate::comm::CommunicationPrimitive;
|
|
use crate::error::{DistributedError, Result};
|
|
#[cfg(feature = "nccl")]
|
|
use crate::nccl::NcclCommunicator;
|
|
#[cfg(feature = "nccl")]
|
|
use crate::nccl_comm::NcclCommunicationPrimitive;
|
|
#[cfg(feature = "rnccl")]
|
|
use crate::rnccl_backend::RncclBackend;
|
|
use crate::{Tensor, TensorShape};
|
|
use parking_lot::RwLock;
|
|
use serde::{Deserialize, Serialize};
|
|
use std::sync::Arc;
|
|
use std::time::{Instant, SystemTime, UNIX_EPOCH};
|
|
use uuid::Uuid;
|
|
|
|
/// Information about the distributed world
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub struct WorldInfo {
|
|
/// Rank of this process (0-indexed)
|
|
pub rank: usize,
|
|
/// Total number of processes in the world
|
|
pub world_size: usize,
|
|
/// Local rank within the node
|
|
pub local_rank: usize,
|
|
/// Number of processes on local node
|
|
pub local_world_size: usize,
|
|
/// Master address for rendezvous
|
|
pub master_addr: String,
|
|
/// Master port for rendezvous
|
|
pub master_port: u16,
|
|
/// Unique identifier for this process group
|
|
pub group_id: Option<Uuid>,
|
|
/// Backend being used (derived from ProcessGroup)
|
|
#[serde(skip)]
|
|
pub backend: Option<Backend>,
|
|
/// Timestamp when the group was created (seconds since UNIX epoch)
|
|
pub created_at: u64,
|
|
/// Instant when the group was created (not serialized)
|
|
#[serde(skip)]
|
|
pub created_instant: Option<Instant>,
|
|
}
|
|
|
|
impl WorldInfo {
|
|
/// Create new world info for simple cases
|
|
pub fn new(world_size: i32, rank: i32, backend: Backend) -> Self {
|
|
let now = Instant::now();
|
|
let timestamp = SystemTime::now()
|
|
.duration_since(UNIX_EPOCH)
|
|
.unwrap()
|
|
.as_secs();
|
|
|
|
Self {
|
|
rank: rank as usize,
|
|
world_size: world_size as usize,
|
|
local_rank: rank as usize,
|
|
local_world_size: world_size as usize,
|
|
master_addr: "127.0.0.1".to_string(),
|
|
master_port: 29500,
|
|
group_id: Some(Uuid::new_v4()),
|
|
backend: Some(backend),
|
|
created_at: timestamp,
|
|
created_instant: Some(now),
|
|
}
|
|
}
|
|
|
|
/// Check if this process is the root (rank 0)
|
|
pub fn is_root(&self) -> bool {
|
|
self.rank == 0
|
|
}
|
|
|
|
/// Get elapsed time since group creation
|
|
pub fn elapsed(&self) -> std::time::Duration {
|
|
self.created_instant.unwrap_or_else(Instant::now).elapsed()
|
|
}
|
|
}
|
|
|
|
/// Internal state of a process group
|
|
#[derive(Debug)]
|
|
struct ProcessGroupState {
|
|
world_info: WorldInfo,
|
|
config: BackendConfig,
|
|
backend_impl: Option<BackendImpl>,
|
|
is_initialized: bool,
|
|
is_cleaned_up: bool,
|
|
}
|
|
|
|
impl ProcessGroupState {
|
|
fn new(world_size: i32, rank: i32, backend: Backend, config: BackendConfig) -> Self {
|
|
Self {
|
|
world_info: WorldInfo::new(world_size, rank, backend),
|
|
config,
|
|
backend_impl: None,
|
|
is_initialized: false,
|
|
is_cleaned_up: false,
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Process group for distributed training
|
|
///
|
|
/// A ProcessGroup represents a collection of processes that can communicate
|
|
/// with each other using collective operations. It manages the lifetime of
|
|
/// the communication backend and provides a safe, high-level interface.
|
|
#[derive(Debug, Clone)]
|
|
pub struct ProcessGroup {
|
|
state: Arc<RwLock<ProcessGroupState>>,
|
|
}
|
|
|
|
impl ProcessGroup {
|
|
/// Create a new process group with WorldInfo
|
|
pub fn new(backend: Backend, world_info: WorldInfo) -> Result<Self> {
|
|
// Create default configuration for the backend
|
|
let config = match backend {
|
|
Backend::Nccl => BackendConfig::nccl(),
|
|
Backend::Rccl => BackendConfig::rccl(),
|
|
Backend::Rnccl => BackendConfig::nccl(), // Rust NCCL uses same config as NCCL
|
|
Backend::Mpi => BackendConfig::mpi(),
|
|
Backend::Cpu => BackendConfig::cpu(),
|
|
Backend::Tcp => BackendConfig::tcp(),
|
|
};
|
|
|
|
// Update world_info with backend
|
|
let mut updated_world_info = world_info;
|
|
updated_world_info.backend = Some(backend);
|
|
if updated_world_info.created_instant.is_none() {
|
|
updated_world_info.created_instant = Some(Instant::now());
|
|
}
|
|
|
|
let state = ProcessGroupState {
|
|
world_info: updated_world_info,
|
|
config,
|
|
backend_impl: None, // Backend implementation is initialized lazily
|
|
is_initialized: true, // For simplicity, mark as initialized
|
|
is_cleaned_up: false,
|
|
};
|
|
|
|
Ok(Self {
|
|
state: Arc::new(RwLock::new(state)),
|
|
})
|
|
}
|
|
|
|
/// Create a new process group (legacy interface)
|
|
///
|
|
/// # Arguments
|
|
/// * `backend` - Communication backend to use
|
|
/// * `world_size` - Total number of processes
|
|
/// * `rank` - Rank of this process (0-indexed)
|
|
/// * `config` - Backend-specific configuration
|
|
///
|
|
/// # Returns
|
|
/// * `Result<Self>` - New process group or error
|
|
///
|
|
/// # Examples
|
|
/// ```rust,ignore
|
|
/// use rtx_distributed::{Backend, BackendConfig, ProcessGroup};
|
|
///
|
|
/// let config = BackendConfig::nccl();
|
|
/// let pg = ProcessGroup::new(Backend::Nccl, 4, 0, config).await?;
|
|
/// ```
|
|
pub async fn new_with_config(
|
|
backend: Backend,
|
|
world_size: i32,
|
|
rank: i32,
|
|
config: BackendConfig,
|
|
) -> Result<Self> {
|
|
// Validate parameters
|
|
if world_size <= 0 {
|
|
return Err(DistributedError::process_group(
|
|
"world_size must be positive",
|
|
));
|
|
}
|
|
|
|
if rank < 0 || rank >= world_size {
|
|
return Err(DistributedError::process_group(format!(
|
|
"rank {rank} must be in range [0, {world_size})"
|
|
)));
|
|
}
|
|
|
|
// Validate configuration
|
|
config.validate()?;
|
|
|
|
if config.backend != backend {
|
|
return Err(DistributedError::configuration(
|
|
"backend mismatch between parameter and config",
|
|
));
|
|
}
|
|
|
|
// Create process group state
|
|
let state = ProcessGroupState::new(world_size, rank, backend, config);
|
|
let pg = Self {
|
|
state: Arc::new(RwLock::new(state)),
|
|
};
|
|
|
|
// Initialize the backend
|
|
pg.initialize_backend().await?;
|
|
|
|
Ok(pg)
|
|
}
|
|
|
|
/// Initialize the communication backend
|
|
async fn initialize_backend(&self) -> Result<()> {
|
|
let mut state = self.state.write();
|
|
|
|
if state.is_initialized {
|
|
return Ok(());
|
|
}
|
|
|
|
match state.world_info.backend {
|
|
Some(Backend::Nccl) => {
|
|
self.init_nccl(&mut state).await?;
|
|
}
|
|
Some(Backend::Rccl) => {
|
|
self.init_rccl(&mut state).await?;
|
|
}
|
|
Some(Backend::Rnccl) => {
|
|
self.init_rnccl(&mut state).await?;
|
|
}
|
|
Some(Backend::Mpi) => {
|
|
self.init_mpi(&mut state).await?;
|
|
}
|
|
Some(Backend::Cpu) => {
|
|
self.init_cpu(&mut state).await?;
|
|
}
|
|
Some(Backend::Tcp) => {
|
|
self.init_tcp(&mut state).await?;
|
|
}
|
|
None => {
|
|
return Err(DistributedError::process_group("backend not specified"));
|
|
}
|
|
}
|
|
|
|
state.is_initialized = true;
|
|
Ok(())
|
|
}
|
|
|
|
/// Initialize NCCL backend
|
|
async fn init_nccl(&self, state: &mut ProcessGroupState) -> Result<()> {
|
|
tracing::info!(
|
|
"Initializing NCCL backend for rank {} of {}",
|
|
state.world_info.rank,
|
|
state.world_info.world_size
|
|
);
|
|
|
|
#[cfg(feature = "nccl")]
|
|
{
|
|
// Create backend implementation
|
|
let backend_impl = BackendImpl::new(state.config.clone()).await?;
|
|
|
|
// Generate NCCL unique ID (in practice, this would be coordinated across all processes)
|
|
let nccl_id = NcclCommunicator::get_unique_id().map_err(|e| {
|
|
DistributedError::communication("nccl", format!("failed to get NCCL ID: {e}"))
|
|
})?;
|
|
|
|
// Initialize world communicator
|
|
if let BackendImpl::Nccl(nccl_backend) = &backend_impl {
|
|
nccl_backend
|
|
.init_world(
|
|
state.world_info.world_size,
|
|
state.world_info.rank,
|
|
state.world_info.rank as i32,
|
|
nccl_id,
|
|
)
|
|
.await?;
|
|
} else {
|
|
return Err(DistributedError::configuration(
|
|
"expected NCCL backend implementation",
|
|
));
|
|
}
|
|
|
|
state.backend_impl = Some(backend_impl);
|
|
|
|
tracing::info!(
|
|
"NCCL backend successfully initialized for rank {} of {}",
|
|
state.world_info.rank,
|
|
state.world_info.world_size
|
|
);
|
|
}
|
|
|
|
#[cfg(not(feature = "nccl"))]
|
|
{
|
|
return Err(DistributedError::configuration(
|
|
"NCCL backend not available - compile with nccl feature",
|
|
));
|
|
}
|
|
|
|
// Check timeout
|
|
if state.world_info.elapsed() > state.config.timeout {
|
|
return Err(DistributedError::communication(
|
|
"nccl",
|
|
"initialization timeout",
|
|
));
|
|
}
|
|
|
|
Ok(())
|
|
}
|
|
|
|
/// Initialize RCCL backend
|
|
async fn init_rccl(&self, state: &mut ProcessGroupState) -> Result<()> {
|
|
tracing::info!(
|
|
"Initializing RCCL backend for rank {} of {}",
|
|
state.world_info.rank,
|
|
state.world_info.world_size
|
|
);
|
|
|
|
// Simulate RCCL initialization
|
|
tokio::time::sleep(std::time::Duration::from_millis(10)).await;
|
|
|
|
if state.world_info.elapsed() > state.config.timeout {
|
|
return Err(DistributedError::communication(
|
|
"rccl",
|
|
"initialization timeout",
|
|
));
|
|
}
|
|
|
|
Ok(())
|
|
}
|
|
|
|
/// Initialize RNCCL (Rust-native NCCL) backend
|
|
async fn init_rnccl(&self, state: &mut ProcessGroupState) -> Result<()> {
|
|
tracing::info!(
|
|
"Initializing RNCCL backend for rank {} of {}",
|
|
state.world_info.rank,
|
|
state.world_info.world_size
|
|
);
|
|
|
|
#[cfg(feature = "rnccl")]
|
|
{
|
|
// Create RNCCL backend with rank and world_size
|
|
let rnccl_backend = RncclBackend::new(
|
|
state.world_info.rank as i32,
|
|
state.world_info.world_size as i32,
|
|
)?;
|
|
|
|
// Store in BackendImpl
|
|
state.backend_impl = Some(BackendImpl::Rnccl(Arc::new(rnccl_backend)));
|
|
|
|
tracing::info!(
|
|
"RNCCL backend successfully initialized for rank {} of {}",
|
|
state.world_info.rank,
|
|
state.world_info.world_size
|
|
);
|
|
}
|
|
|
|
#[cfg(not(feature = "rnccl"))]
|
|
{
|
|
return Err(DistributedError::configuration(
|
|
"RNCCL backend not available - compile with rnccl feature",
|
|
));
|
|
}
|
|
|
|
// Check timeout
|
|
if state.world_info.elapsed() > state.config.timeout {
|
|
return Err(DistributedError::communication(
|
|
"rnccl",
|
|
"initialization timeout",
|
|
));
|
|
}
|
|
|
|
Ok(())
|
|
}
|
|
|
|
/// Initialize MPI backend
|
|
async fn init_mpi(&self, _state: &mut ProcessGroupState) -> Result<()> {
|
|
#[cfg(feature = "mpi")]
|
|
{
|
|
tracing::info!("Initializing MPI backend");
|
|
// MPI initialization would go here
|
|
Ok(())
|
|
}
|
|
|
|
#[cfg(not(feature = "mpi"))]
|
|
{
|
|
Err(DistributedError::configuration(
|
|
"MPI backend not available - compile with mpi feature",
|
|
))
|
|
}
|
|
}
|
|
|
|
/// Initialize CPU backend (for testing)
|
|
async fn init_cpu(&self, state: &mut ProcessGroupState) -> Result<()> {
|
|
tracing::info!(
|
|
"Initializing CPU backend for rank {} of {}",
|
|
state.world_info.rank,
|
|
state.world_info.world_size
|
|
);
|
|
|
|
// Create CPU backend implementation
|
|
let backend_impl = BackendImpl::new(state.config.clone()).await?;
|
|
state.backend_impl = Some(backend_impl);
|
|
|
|
// CPU backend is always available and fast to initialize
|
|
tokio::time::sleep(std::time::Duration::from_millis(1)).await;
|
|
|
|
tracing::info!(
|
|
"CPU backend successfully initialized for rank {} of {}",
|
|
state.world_info.rank,
|
|
state.world_info.world_size
|
|
);
|
|
|
|
Ok(())
|
|
}
|
|
|
|
/// Initialize TCP fallback backend
|
|
async fn init_tcp(&self, state: &mut ProcessGroupState) -> Result<()> {
|
|
tracing::info!(
|
|
"Initializing TCP fallback backend for rank {} of {}",
|
|
state.world_info.rank,
|
|
state.world_info.world_size
|
|
);
|
|
|
|
// Create TCP backend implementation
|
|
let backend_impl = BackendImpl::new(state.config.clone()).await?;
|
|
state.backend_impl = Some(backend_impl);
|
|
|
|
tracing::info!(
|
|
"TCP fallback backend successfully initialized for rank {} of {}",
|
|
state.world_info.rank,
|
|
state.world_info.world_size
|
|
);
|
|
|
|
Ok(())
|
|
}
|
|
|
|
/// Get world size
|
|
pub fn world_size(&self) -> usize {
|
|
self.state.read().world_info.world_size
|
|
}
|
|
|
|
/// Get rank
|
|
pub fn rank(&self) -> usize {
|
|
self.state.read().world_info.rank
|
|
}
|
|
|
|
/// Get backend
|
|
pub fn backend(&self) -> Backend {
|
|
self.state.read().world_info.backend.unwrap_or(Backend::Cpu)
|
|
}
|
|
|
|
/// Get world info
|
|
pub fn world_info(&self) -> WorldInfo {
|
|
self.state.read().world_info.clone()
|
|
}
|
|
|
|
/// Get communication primitive for this process group
|
|
#[cfg(feature = "nccl")]
|
|
pub fn communication_primitive(&self) -> Result<Option<NcclCommunicationPrimitive>> {
|
|
let state = self.state.read();
|
|
|
|
if let Some(BackendImpl::Nccl(nccl_backend)) = &state.backend_impl {
|
|
Ok(Some(NcclCommunicationPrimitive::new(
|
|
nccl_backend.clone(),
|
|
state.world_info.world_size,
|
|
state.world_info.rank,
|
|
)))
|
|
} else {
|
|
Ok(None)
|
|
}
|
|
}
|
|
|
|
/// Check if NCCL backend is available and initialized
|
|
#[cfg(feature = "nccl")]
|
|
pub fn has_nccl_backend(&self) -> bool {
|
|
let state = self.state.read();
|
|
matches!(&state.backend_impl, Some(BackendImpl::Nccl(_)))
|
|
}
|
|
|
|
/// Check if RNCCL backend is available and initialized
|
|
#[cfg(feature = "rnccl")]
|
|
pub fn has_rnccl_backend(&self) -> bool {
|
|
let state = self.state.read();
|
|
matches!(&state.backend_impl, Some(BackendImpl::Rnccl(_)))
|
|
}
|
|
|
|
/// Get the RNCCL backend reference for direct operations
|
|
#[cfg(feature = "rnccl")]
|
|
pub fn rnccl_backend(&self) -> Option<Arc<RncclBackend>> {
|
|
let state = self.state.read();
|
|
if let Some(BackendImpl::Rnccl(backend)) = &state.backend_impl {
|
|
Some(backend.clone())
|
|
} else {
|
|
None
|
|
}
|
|
}
|
|
|
|
/// AllReduce operation - sum tensors across all processes
|
|
pub async fn all_reduce(&self, tensor: &mut Tensor, op: crate::comm::ReduceOp) -> Result<()> {
|
|
use crate::comm::ReduceOp;
|
|
|
|
// Check initialization
|
|
{
|
|
let state = self.state.read();
|
|
if !state.is_initialized {
|
|
return Err(DistributedError::communication(
|
|
"allreduce",
|
|
"process group not initialized",
|
|
));
|
|
}
|
|
}
|
|
|
|
// Use NCCL backend if available
|
|
#[cfg(feature = "nccl")]
|
|
{
|
|
let nccl_backend_opt = {
|
|
let state = self.state.read();
|
|
if let Some(BackendImpl::Nccl(backend)) = &state.backend_impl {
|
|
Some(backend.clone())
|
|
} else {
|
|
None
|
|
}
|
|
};
|
|
|
|
if let Some(nccl_backend) = nccl_backend_opt {
|
|
let comm_primitive =
|
|
NcclCommunicationPrimitive::new(nccl_backend, self.world_size(), self.rank());
|
|
|
|
return comm_primitive.allreduce(tensor, op).await;
|
|
}
|
|
}
|
|
|
|
// Use RNCCL backend if available
|
|
#[cfg(feature = "rnccl")]
|
|
{
|
|
let rnccl_backend_opt = {
|
|
let state = self.state.read();
|
|
if let Some(BackendImpl::Rnccl(backend)) = &state.backend_impl {
|
|
Some(backend.clone())
|
|
} else {
|
|
None
|
|
}
|
|
};
|
|
|
|
if let Some(rnccl_backend) = rnccl_backend_opt {
|
|
return rnccl_backend.allreduce(tensor, op).await;
|
|
}
|
|
}
|
|
|
|
// Fallback to simulation for CPU backend
|
|
match op {
|
|
ReduceOp::Sum => {
|
|
// In real implementation, this would sum across all ranks
|
|
// For simulation, we modify the tensor to show it was processed
|
|
let scale_factor = self.world_size() as f32;
|
|
*tensor = tensor.mul_scalar(scale_factor)?;
|
|
}
|
|
ReduceOp::Max | ReduceOp::Min => {
|
|
// For min/max, we don't need to scale
|
|
}
|
|
_ => {
|
|
// Other ops would have their own implementations
|
|
}
|
|
}
|
|
|
|
Ok(())
|
|
}
|
|
|
|
/// Broadcast tensor from root process to all others
|
|
pub async fn broadcast(&self, tensor: &mut Tensor, root: usize) -> Result<()> {
|
|
// Check initialization and bounds
|
|
{
|
|
let state = self.state.read();
|
|
if !state.is_initialized {
|
|
return Err(DistributedError::communication(
|
|
"broadcast",
|
|
"process group not initialized",
|
|
));
|
|
}
|
|
|
|
if root >= state.world_info.world_size {
|
|
return Err(DistributedError::communication(
|
|
"broadcast",
|
|
format!(
|
|
"root rank {} out of bounds for world_size {}",
|
|
root, state.world_info.world_size
|
|
),
|
|
));
|
|
}
|
|
}
|
|
|
|
// Use NCCL backend if available
|
|
#[cfg(feature = "nccl")]
|
|
{
|
|
let nccl_backend_opt = {
|
|
let state = self.state.read();
|
|
if let Some(BackendImpl::Nccl(backend)) = &state.backend_impl {
|
|
Some(backend.clone())
|
|
} else {
|
|
None
|
|
}
|
|
};
|
|
|
|
if let Some(nccl_backend) = nccl_backend_opt {
|
|
let comm_primitive =
|
|
NcclCommunicationPrimitive::new(nccl_backend, self.world_size(), self.rank());
|
|
|
|
return comm_primitive.broadcast(tensor, root as i32).await;
|
|
}
|
|
}
|
|
|
|
// Use RNCCL backend if available
|
|
#[cfg(feature = "rnccl")]
|
|
{
|
|
let rnccl_backend_opt = {
|
|
let state = self.state.read();
|
|
if let Some(BackendImpl::Rnccl(backend)) = &state.backend_impl {
|
|
Some(backend.clone())
|
|
} else {
|
|
None
|
|
}
|
|
};
|
|
|
|
if let Some(rnccl_backend) = rnccl_backend_opt {
|
|
return rnccl_backend.broadcast(tensor, root as i32).await;
|
|
}
|
|
}
|
|
|
|
// Simulate broadcast - in real implementation, non-root processes would receive data
|
|
tracing::debug!("Broadcasting tensor from rank {} to all processes", root);
|
|
Ok(())
|
|
}
|
|
|
|
/// AllGather operation - gather tensors from all processes
|
|
pub fn all_gather(&self, tensor: &Tensor) -> Result<Vec<Tensor>> {
|
|
// Check initialization
|
|
let world_size = {
|
|
let state = self.state.read();
|
|
if !state.is_initialized {
|
|
return Err(DistributedError::communication(
|
|
"allgather",
|
|
"process group not initialized",
|
|
));
|
|
}
|
|
state.world_info.world_size
|
|
};
|
|
|
|
// Use RNCCL backend if available (blocking call to async fn)
|
|
#[cfg(feature = "rnccl")]
|
|
{
|
|
let rnccl_backend_opt = {
|
|
let state = self.state.read();
|
|
if let Some(BackendImpl::Rnccl(backend)) = &state.backend_impl {
|
|
Some(backend.clone())
|
|
} else {
|
|
None
|
|
}
|
|
};
|
|
|
|
if let Some(rnccl_backend) = rnccl_backend_opt {
|
|
use crate::comm::AllGatherOutput;
|
|
// Use futures::executor for sync context
|
|
let output = futures::executor::block_on(rnccl_backend.allgather(tensor))?;
|
|
return match output {
|
|
AllGatherOutput::TensorList(tensors) => Ok(tensors),
|
|
AllGatherOutput::Tensor(t) => Ok(vec![t]),
|
|
};
|
|
}
|
|
}
|
|
|
|
// Simulate AllGather by creating copies for each rank
|
|
let mut gathered = Vec::new();
|
|
for _rank in 0..world_size {
|
|
gathered.push(tensor.clone());
|
|
}
|
|
|
|
Ok(gathered)
|
|
}
|
|
|
|
/// ReduceScatter operation - reduce and scatter result
|
|
pub fn reduce_scatter(&self, tensor: &Tensor, op: crate::comm::ReduceOp) -> Result<Tensor> {
|
|
// Check initialization
|
|
let (world_size, rank) = {
|
|
let state = self.state.read();
|
|
if !state.is_initialized {
|
|
return Err(DistributedError::communication(
|
|
"reducescatter",
|
|
"process group not initialized",
|
|
));
|
|
}
|
|
(state.world_info.world_size, state.world_info.rank)
|
|
};
|
|
|
|
// Use RNCCL backend if available (blocking call to async fn)
|
|
#[cfg(feature = "rnccl")]
|
|
{
|
|
let rnccl_backend_opt = {
|
|
let state = self.state.read();
|
|
if let Some(BackendImpl::Rnccl(backend)) = &state.backend_impl {
|
|
Some(backend.clone())
|
|
} else {
|
|
None
|
|
}
|
|
};
|
|
|
|
if let Some(rnccl_backend) = rnccl_backend_opt {
|
|
return futures::executor::block_on(rnccl_backend.reduce_scatter(tensor, op));
|
|
}
|
|
}
|
|
|
|
// Simulate ReduceScatter by returning a portion of the input tensor
|
|
let elements_per_rank = tensor.numel() / world_size;
|
|
let start_idx = rank * elements_per_rank;
|
|
let _end_idx = start_idx + elements_per_rank;
|
|
|
|
// Create a view/slice of the tensor (simplified)
|
|
let result_shape = TensorShape::new(vec![elements_per_rank])?;
|
|
let result = Tensor::zeros(result_shape, &crate::Device::default())?;
|
|
|
|
Ok(result)
|
|
}
|
|
|
|
/// Split the process group
|
|
pub async fn split(&self, _color: i32, key: i32) -> Result<Self> {
|
|
let state = self.state.read();
|
|
|
|
if !state.is_initialized {
|
|
return Err(DistributedError::process_group(
|
|
"cannot split uninitialized process group",
|
|
));
|
|
}
|
|
|
|
// For simplicity, create a new group with modified parameters
|
|
// In a real implementation, this would use backend-specific splitting
|
|
let new_world_size = state.world_info.world_size / 2; // Simplified
|
|
let new_rank = (key as usize) % new_world_size;
|
|
|
|
drop(state); // Release the lock before async call
|
|
|
|
let new_config = self.state.read().config.clone();
|
|
Self::new_with_config(
|
|
self.backend(),
|
|
new_world_size as i32,
|
|
new_rank as i32,
|
|
new_config,
|
|
)
|
|
.await
|
|
}
|
|
|
|
/// Merge multiple process groups
|
|
pub async fn merge(groups: Vec<Self>) -> Result<Self> {
|
|
if groups.is_empty() {
|
|
return Err(DistributedError::process_group(
|
|
"cannot merge empty list of groups",
|
|
));
|
|
}
|
|
|
|
let first_group = &groups[0];
|
|
let backend = first_group.backend();
|
|
let config = first_group.state.read().config.clone();
|
|
|
|
// Calculate merged world size and rank
|
|
let total_world_size: usize = groups.iter().map(ProcessGroup::world_size).sum();
|
|
|
|
let new_rank = groups[0].rank(); // Simplified
|
|
|
|
Self::new_with_config(backend, total_world_size as i32, new_rank as i32, config).await
|
|
}
|
|
|
|
/// Clean up the process group
|
|
pub async fn cleanup(&self) -> Result<()> {
|
|
let mut state = self.state.write();
|
|
|
|
if state.is_cleaned_up {
|
|
return Ok(());
|
|
}
|
|
|
|
tracing::info!(
|
|
"Cleaning up process group for rank {}/{}",
|
|
state.world_info.rank,
|
|
state.world_info.world_size
|
|
);
|
|
|
|
// Clean up backend implementation
|
|
if let Some(backend_impl) = state.backend_impl.take() {
|
|
drop(state); // Release lock before async operation
|
|
backend_impl.cleanup().await?;
|
|
state = self.state.write(); // Re-acquire lock
|
|
}
|
|
|
|
// Additional backend-specific cleanup
|
|
match state.world_info.backend {
|
|
Some(Backend::Nccl | Backend::Rccl | Backend::Rnccl) => {
|
|
// Additional GPU resource cleanup if needed
|
|
tokio::time::sleep(std::time::Duration::from_millis(5)).await;
|
|
}
|
|
Some(Backend::Mpi) => {
|
|
#[cfg(feature = "mpi")]
|
|
{
|
|
// MPI cleanup
|
|
}
|
|
}
|
|
Some(Backend::Cpu | Backend::Tcp) | None => {
|
|
// No special cleanup needed
|
|
}
|
|
}
|
|
|
|
state.is_cleaned_up = true;
|
|
tracing::info!(
|
|
"Process group cleanup completed for rank {}/{}",
|
|
state.world_info.rank,
|
|
state.world_info.world_size
|
|
);
|
|
|
|
Ok(())
|
|
}
|
|
|
|
/// Synchronize all processes in the group (barrier operation)
|
|
pub async fn barrier(&self) -> Result<()> {
|
|
let state = self.state.read();
|
|
|
|
if !state.is_initialized {
|
|
return Err(DistributedError::process_group(
|
|
"cannot perform barrier on uninitialized process group",
|
|
));
|
|
}
|
|
|
|
// For single process groups, barrier is a no-op
|
|
if state.world_info.world_size == 1 {
|
|
return Ok(());
|
|
}
|
|
|
|
// In a real implementation, this would use backend-specific barrier
|
|
// For now, we'll just return success as a placeholder
|
|
tracing::debug!(
|
|
"Barrier operation for rank {} of {} (placeholder implementation)",
|
|
state.world_info.rank,
|
|
state.world_info.world_size
|
|
);
|
|
|
|
Ok(())
|
|
}
|
|
}
|
|
|
|
// Automatic cleanup on drop
|
|
impl Drop for ProcessGroup {
|
|
fn drop(&mut self) {
|
|
let state = self.state.read();
|
|
if state.is_initialized && !state.is_cleaned_up {
|
|
// Log warning about unclean shutdown
|
|
tracing::warn!(
|
|
"ProcessGroup dropped without explicit cleanup for rank {} of {}",
|
|
state.world_info.rank,
|
|
state.world_info.world_size
|
|
);
|
|
}
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
use crate::backend::BackendConfig;
|
|
|
|
#[tokio::test]
|
|
async fn test_process_group_creation() {
|
|
let config = BackendConfig::cpu();
|
|
let pg = ProcessGroup::new_with_config(Backend::Cpu, 2, 0, config)
|
|
.await
|
|
.unwrap();
|
|
|
|
assert_eq!(pg.world_size(), 2);
|
|
assert_eq!(pg.rank(), 0);
|
|
assert_eq!(pg.backend(), Backend::Cpu);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_invalid_parameters() {
|
|
let config = BackendConfig::cpu();
|
|
|
|
// Invalid world size
|
|
let result = ProcessGroup::new_with_config(Backend::Cpu, 0, 0, config.clone()).await;
|
|
assert!(result.is_err());
|
|
|
|
// Invalid rank
|
|
let result = ProcessGroup::new_with_config(Backend::Cpu, 2, 2, config).await;
|
|
assert!(result.is_err());
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_world_info() {
|
|
let config = BackendConfig::cpu();
|
|
let pg = ProcessGroup::new_with_config(Backend::Cpu, 4, 1, config)
|
|
.await
|
|
.unwrap();
|
|
|
|
let info = pg.world_info();
|
|
assert_eq!(info.world_size, 4);
|
|
assert_eq!(info.rank, 1);
|
|
assert_eq!(info.backend, Some(Backend::Cpu));
|
|
assert!(info.group_id.is_some());
|
|
assert!(!info.is_root());
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_cleanup() {
|
|
let config = BackendConfig::cpu();
|
|
let pg = ProcessGroup::new_with_config(Backend::Cpu, 2, 0, config)
|
|
.await
|
|
.unwrap();
|
|
|
|
assert!(pg.cleanup().await.is_ok());
|
|
|
|
// Second cleanup should also work
|
|
assert!(pg.cleanup().await.is_ok());
|
|
}
|
|
|
|
#[test]
|
|
fn test_world_info_properties() {
|
|
let info = WorldInfo::new(4, 0, Backend::Cpu);
|
|
assert!(info.is_root());
|
|
assert!(info.group_id.is_some());
|
|
|
|
let info2 = WorldInfo::new(4, 1, Backend::Cpu);
|
|
assert!(!info2.is_root());
|
|
}
|
|
}
|