Trait Communicator 

pub trait Communicator: Send + Sync {
    // Required methods
    fn rank(&self) -> usize;
    fn size(&self) -> usize;
    fn all_reduce(
        &self,
        op: AllReduce,
        dtype: DType,
        data: &mut [u8],
    ) -> Result<()>;
    fn all_gather(&self, input: &[u8], output: &mut [u8]) -> Result<()>;
    fn send(&self, dst: usize, data: &[u8]) -> Result<()>;
    fn recv(&self, src: usize, data: &mut [u8]) -> Result<()>;

    // Provided methods
    fn sum_scatter(
        &self,
        dtype: DType,
        input: &[u8],
        output: &mut [u8],
    ) -> Result<()> { ... }
    fn health(&self) -> Result<()> { ... }
    fn split(&self, color: i32, key: i32) -> Result<Arc<dyn Communicator>> { ... }
}

The transport of a Group: every rank of the group calls the same operations in the same order. Buffers are row-contiguous elements; a failure names the peer.

Required Methods

fn rank(&self) -> usize

This process’s rank, 0..size.

fn size(&self) -> usize

The number of processes.

fn all_reduce(&self, op: AllReduce, dtype: DType, data: &mut [u8]) -> Result<()>

Reduces data (elements of dtype) with every rank’s data, in place; every rank gets the same bits.

fn all_gather(&self, input: &[u8], output: &mut [u8]) -> Result<()>

Every rank’s input (equal lengths), concatenated in rank order into output.

fn send(&self, dst: usize, data: &[u8]) -> Result<()>

Sends data to rank dst.

fn recv(&self, src: usize, data: &mut [u8]) -> Result<()>

Receives exactly data.len() bytes from rank src.

Provided Methods

fn sum_scatter( &self, dtype: DType, input: &[u8], output: &mut [u8], ) -> Result<()>

Rank r’s output is the element-wise sum (elements of dtype) of every rank’s r-th chunk of input (input.len() == size · output.len()), as MLX’s sum_scatter.

fn health(&self) -> Result<()>

Whether every peer is still reachable, checked without a collective and without blocking; an error names a lost peer. Backends that cannot tell report Ok.

fn split(&self, color: i32, key: i32) -> Result<Arc<dyn Communicator>>

A subgroup of the ranks with the same color, ordered by key (MLX’s split).

Dyn Compatibility

This trait is dyn compatible.

In older versions of Rust, dyn compatibility was called "object safety".

Implementors