rustytorch/crates/training/rtx-distributed/src/distributed_context.rs

//! Distributed training context for managing multi-GPU/multi-node state

use crate::backend::Backend;
use crate::error::{DistributedError, Result};
use crate::group::{ProcessGroup, WorldInfo};
use rtx_tensor::{Device, Tensor};
use std::sync::Arc;

/// Distributed training context that encapsulates distributed state and operations
#[derive(Debug, Clone)]
pub struct DistributedContext {
    /// Process group for communication
    process_group: Arc<ProcessGroup>,
    /// World information (rank, size, etc.)
    world_info: WorldInfo,
    /// Backend being used (NCCL, GLOO, etc.)
    backend: Backend,
    /// Whether distributed training is initialized
    initialized: bool,
}

impl DistributedContext {
    /// Create a new distributed context
    pub fn new(process_group: ProcessGroup, world_info: WorldInfo, backend: Backend) -> Self {
        Self {
            process_group: Arc::new(process_group),
            world_info,
            backend,
            initialized: true,
        }
    }

    /// Create an uninitialized context (for single-GPU training)
    pub fn uninitialized() -> Self {
        // Create a simple single-GPU world info
        let world_info = WorldInfo::new(1, 0, Backend::Nccl);

        // Create a simple process group for single GPU
        let process_group =
            ProcessGroup::new(Backend::Nccl, world_info.clone()).unwrap_or_else(|_| {
                // If creation fails, create a minimal dummy group
                ProcessGroup::new(Backend::Cpu, world_info.clone()).unwrap()
            });

        Self {
            process_group: Arc::new(process_group),
            world_info,
            backend: Backend::Nccl,
            initialized: false,
        }
    }

    /// Check if distributed training is initialized
    pub fn is_initialized(&self) -> bool {
        self.initialized
    }

    /// Get the current rank
    pub fn rank(&self) -> usize {
        self.world_info.rank
    }

    /// Get the world size (total number of processes)
    pub fn world_size(&self) -> usize {
        self.world_info.world_size
    }

    /// Check if this is the master process (rank 0)
    pub fn is_master(&self) -> bool {
        self.rank() == 0
    }

    /// Get the local rank (GPU index on this node)
    pub fn local_rank(&self) -> usize {
        self.world_info.local_rank
    }

    /// Get the local world size (number of GPUs on this node)
    pub fn local_world_size(&self) -> usize {
        self.world_info.local_world_size
    }

    /// Get the process group
    pub fn process_group(&self) -> &ProcessGroup {
        &self.process_group
    }

    /// Get the backend type
    pub fn backend(&self) -> Backend {
        self.backend
    }

    /// Get world information
    pub fn world_info(&self) -> &WorldInfo {
        &self.world_info
    }

    /// Synchronize all processes (barrier)
    pub async fn barrier(&self) -> Result<()> {
        if !self.initialized {
            return Ok(()); // No-op for uninitialized context
        }

        self.process_group.barrier().await
    }

    /// Broadcast a tensor from the root process to all other processes
    ///
    /// This is the primary broadcast method for tensor data in distributed training.
    /// The tensor is modified in-place on non-root processes.
    ///
    /// # Arguments
    /// * `tensor` - The tensor to broadcast (modified in-place on non-root)
    /// * `root` - The rank of the process that owns the source data
    pub async fn broadcast_tensor(&self, tensor: &mut Tensor, root: usize) -> Result<()> {
        if !self.initialized {
            return Ok(()); // No-op for uninitialized context
        }

        self.process_group.broadcast(tensor, root).await
    }

    /// Broadcast raw f32 data from the root process to all other processes
    ///
    /// This is useful for broadcasting configuration values, hyperparameters,
    /// or other numeric data that isn't in tensor form.
    ///
    /// # Arguments
    /// * `data` - Slice of f32 values to broadcast (must be same size on all ranks)
    /// * `root` - The rank of the process that owns the source data
    pub async fn broadcast_f32(&self, data: &mut [f32], root: usize) -> Result<()> {
        if !self.initialized {
            return Ok(()); // No-op for uninitialized context
        }

        // Create a tensor from the data, broadcast it, then copy back
        let shape = vec![data.len()];
        let device = Device::cpu();

        let mut tensor = Tensor::from_data(data.to_vec(), shape, &device)
            .map_err(|e| DistributedError::communication("broadcast_f32", e.to_string()))?;

        self.process_group.broadcast(&mut tensor, root).await?;

        // Copy back the broadcasted data
        let result_data = tensor
            .to_vec()
            .map_err(|e| DistributedError::communication("broadcast_f32", e.to_string()))?;

        data.copy_from_slice(&result_data);

        Ok(())
    }

    /// Broadcast serializable data using bincode
    ///
    /// This method serializes the data on the root process, broadcasts the bytes,
    /// and deserializes on all other processes. Useful for configuration structs,
    /// metadata, and other complex types.
    ///
    /// # Arguments
    /// * `data` - The data to broadcast (modified in-place on non-root)
    /// * `root` - The rank of the process that owns the source data
    pub async fn broadcast_bytes<T>(&self, data: &mut T, root: usize) -> Result<()>
    where
        T: serde::Serialize + serde::de::DeserializeOwned,
    {
        if !self.initialized {
            return Ok(()); // No-op for uninitialized context
        }

        let rank = self.rank();

        // Serialize on root
        let serialized = if rank == root {
            bincode::serialize(data).map_err(|e| {
                DistributedError::communication(
                    "broadcast_bytes",
                    format!("serialization failed: {}", e),
                )
            })?
        } else {
            Vec::new()
        };

        // First broadcast the length
        let mut len_buf = [serialized.len() as f32];
        self.broadcast_f32(&mut len_buf, root).await?;
        let expected_len = len_buf[0] as usize;

        // Prepare buffer for the data (pad bytes as f32)
        let num_floats = (expected_len + 3) / 4; // Ceil division
        let mut float_buf = vec![0.0f32; num_floats];

        // On root, pack bytes into f32 buffer
        if rank == root {
            for (i, chunk) in serialized.chunks(4).enumerate() {
                let mut bytes = [0u8; 4];
                bytes[..chunk.len()].copy_from_slice(chunk);
                float_buf[i] = f32::from_le_bytes(bytes);
            }
        }

        // Broadcast the packed data
        self.broadcast_f32(&mut float_buf, root).await?;

        // Unpack and deserialize on non-root
        if rank != root {
            let mut bytes = Vec::with_capacity(expected_len);
            for &f in &float_buf {
                bytes.extend_from_slice(&f.to_le_bytes());
            }
            bytes.truncate(expected_len);

            *data = bincode::deserialize(&bytes).map_err(|e| {
                DistributedError::communication(
                    "broadcast_bytes",
                    format!("deserialization failed: {}", e),
                )
            })?;
        }

        Ok(())
    }

    /// Legacy broadcast method - delegates to broadcast_bytes for backwards compatibility
    #[deprecated(
        since = "1.0.0",
        note = "Use broadcast_tensor, broadcast_f32, or broadcast_bytes instead"
    )]
    pub async fn broadcast<T>(&self, data: &mut T, root: usize) -> Result<()>
    where
        T: Send + Sync + serde::Serialize + serde::de::DeserializeOwned,
    {
        self.broadcast_bytes(data, root).await
    }

    /// Check if the current process is on the specified rank
    pub fn is_rank(&self, rank: usize) -> bool {
        self.rank() == rank
    }

    /// Get a string representation of the distributed setup
    pub fn info_string(&self) -> String {
        if !self.initialized {
            "Distributed: Not initialized (single GPU)".to_string()
        } else {
            format!(
                "Distributed: {:?} backend, rank {}/{}, local rank {}/{}",
                self.backend,
                self.rank(),
                self.world_size(),
                self.local_rank(),
                self.local_world_size()
            )
        }
    }
}

impl Default for DistributedContext {
    fn default() -> Self {
        Self::uninitialized()
    }
}