Skip to main content

ServerCommunication

cubecl_runtime::server

Trait ServerCommunication 

Source 
pub trait ServerCommunication {
    // Provided methods
    fn sync_collective(
        &mut self,
        stream_id: StreamId,
    ) -> Result<(), ServerError> { ... }
    fn comm_init(
        &mut self,
        device_ids: Vec<DeviceId>,
    ) -> Result<(), ServerError> { ... }
    fn all_reduce(
        &mut self,
        src: BufferBinding,
        dst: BufferBinding,
        reduction: Reduction,
        stream_id: StreamId,
        device_ids: Vec<DeviceId>,
        wait: ComputeWait,
    ) -> Result<(), ServerError> { ... }
    fn all_gather(
        &mut self,
        src: BufferBinding,
        dst: BufferBinding,
        stream_id: StreamId,
        device_ids: Vec<DeviceId>,
    ) -> Result<(), ServerError> { ... }
    fn reduce_scatter(
        &mut self,
        src: BufferBinding,
        dst: BufferBinding,
        reduction: Reduction,
        stream_id: StreamId,
        device_ids: Vec<DeviceId>,
    ) -> Result<(), ServerError> { ... }
    fn broadcast(
        &mut self,
        src: Option<BufferBinding>,
        dst: BufferBinding,
        stream_id: StreamId,
        root: DeviceId,
        device_ids: Vec<DeviceId>,
    ) -> Result<(), ServerError> { ... }
    fn all_to_all(
        &mut self,
        src: BufferBinding,
        dst: BufferBinding,
        stream_id: StreamId,
        device_ids: Vec<DeviceId>,
    ) -> Result<(), ServerError> { ... }
    fn send(
        &mut self,
        desc: CopyDescriptor,
        stream_id: StreamId,
        device_id_dst: DeviceId,
    ) -> Result<(), ServerError> { ... }
    fn recv(
        &mut self,
        handle: Handle,
        stream_id: StreamId,
        device_id_src: DeviceId,
    ) -> Result<(), ServerError> { ... }
}
Expand description

Defines functions for optimized data transfer between servers, supporting custom communication mechanisms such as peer-to-peer communication or specialized implementations.

§Inside the tainted-buffer rules

A collective reads a source buffer and produces a destination one, and owes the same two answers the rest of the server gives: ask whether the source carries a failure on the way in (as read does through FailureStore::ensure_written in cubecl-server), and taint the destination on the way out when the operation fails (as a failed launch does). Skipping either lets a collective reduce stale bytes across every device in the group, or leave a destination that reads back clean when nothing wrote it. A part that fails leaves the group’s other members waiting in it, so the device then refuses every later collective over that group rather than pair it with theirs.

The default methods are for a runtime with no transport between its devices: each returns ServerError::NoDeviceTransport before touching any buffer, so there is no destination to taint.

Provided Methods§

Source

fn sync_collective(&mut self, stream_id: StreamId) -> Result<(), ServerError>

Ensure that all queued collective operations have been executed.

§Arguments
  • stream_id - The StreamId of the stream waiting for the sync.
§Returns

Returns a Result containing an ServerError if the operation fails.

§Errors

The default returns ServerError::NoDeviceTransport.

Source

fn comm_init(&mut self, device_ids: Vec<DeviceId>) -> Result<(), ServerError>

Initialize the communication between the devices in device_ids.

§Arguments
  • device_ids - The IDs of the devices that need communication.
§Returns

Returns a Result containing an ServerError if the operation fails.

§Errors

The default returns ServerError::NoDeviceTransport.

Source

fn all_reduce( &mut self, src: BufferBinding, dst: BufferBinding, reduction: Reduction, stream_id: StreamId, device_ids: Vec<DeviceId>, wait: ComputeWait, ) -> Result<(), ServerError>

Performs an all_reduce operation on the input data and writes it to the output buffer. see https://docs.nvidia.com/deeplearning/nccl/user-guide/docs/usage/collectives.html#allreduce

§Arguments
  • src - The data to be reduced.
  • dst - Where to write the result.
  • reduction - The elements’ type and how they combine.
  • stream_id - The data’s stream id.
  • device_ids - The list of device ids from which to all_reduce.
  • wait - When the compute stream waits for it.
§Returns

Returns a Result containing an ServerError if the operation fails.

§Errors

The default returns ServerError::NoDeviceTransport.

Source

fn all_gather( &mut self, src: BufferBinding, dst: BufferBinding, stream_id: StreamId, device_ids: Vec<DeviceId>, ) -> Result<(), ServerError>

Writes every rank’s src, concatenated in rank order, to dst, which holds one src per rank of device_ids.

§Errors

The default returns ServerError::NoDeviceTransport.

Source

fn reduce_scatter( &mut self, src: BufferBinding, dst: BufferBinding, reduction: Reduction, stream_id: StreamId, device_ids: Vec<DeviceId>, ) -> Result<(), ServerError>

Reduces every rank’s src and writes this rank’s chunk of the result to dst, src holding one dst per rank of device_ids.

§Errors

The default returns ServerError::NoDeviceTransport.

Source

fn broadcast( &mut self, src: Option<BufferBinding>, dst: BufferBinding, stream_id: StreamId, root: DeviceId, device_ids: Vec<DeviceId>, ) -> Result<(), ServerError>

Writes root’s src to dst on every rank of device_ids. Only root has a src.

§Errors

The default returns ServerError::NoDeviceTransport.

Source

fn all_to_all( &mut self, src: BufferBinding, dst: BufferBinding, stream_id: StreamId, device_ids: Vec<DeviceId>, ) -> Result<(), ServerError>

Sends chunk r of src to rank r and writes the chunk received from rank r to chunk r of dst, both holding one chunk per rank of device_ids.

§Errors

The default returns ServerError::NoDeviceTransport.

Source

fn send( &mut self, desc: CopyDescriptor, stream_id: StreamId, device_id_dst: DeviceId, ) -> Result<(), ServerError>

Sends data from this server to a destination server.

§Arguments
  • desc - A descriptor specifying the data to be sent, including shape, strides, and binding.
  • stream_id - The stream ID associated with the server’s operation.
  • device_id_dst - ID of the device receiving the data.
§Returns

Returns a Result containing an ServerError if the operation fails.

§Known limitation

Send and recv are posted fire-and-forget on two devices and block for each other, so a send that refuses — a source whose writer failed, above all — leaves the peer’s already-posted recv waiting on its communication stream with no way to recall it from here. The refusal is still right: completing the send would launder stale bytes onto a handle that carries no claim on the other device. Cross-device failure propagation needs a design pass of its own.

§Errors

The default returns ServerError::NoDeviceTransport.

Source

fn recv( &mut self, handle: Handle, stream_id: StreamId, device_id_src: DeviceId, ) -> Result<(), ServerError>

Receive data from another server.

§Arguments
  • handle - The handle in which the received data is written.
  • stream_id - The stream ID associated with the server’s operation.
  • device_id_src - ID of the device sending the data.
§Returns

Returns a Result containing an ServerError if the operation fails.

§Errors

The default returns ServerError::NoDeviceTransport.

Dyn Compatibility§

This trait is dyn compatible.

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

Implementors§