Struct DeviceMesh 

pub struct DeviceMesh {
    pub world_size: usize,
    pub rank: usize,
    pub topology: MeshTopology,
    pub comm: Option<Arc<dyn MeshComm + Send + Sync>>,
    /* private fields */
}

Device Mesh for distributed ML inference

Manages a set of devices arranged in a logical mesh topology for data/tensor/pipeline parallelism. Optimized for memory efficiency.

Fields

world_size: usize

Total number of devices in the mesh

rank: usize

This process’s global rank

topology: MeshTopology

Mesh topology configuration

comm: Option<Arc<dyn MeshComm + Send + Sync>>

Communication backend handle

Implementations

impl DeviceMesh

pub fn new(devices: Vec<Device>) -> Self

Create a new device mesh from a list of devices

pub fn new_with_mock_comm(devices: Vec<Device>, rank: usize) -> Self

Create a new device mesh with a mock communication backend for testing

pub fn with_topology( devices: Vec<Device>, topology: MeshTopology, ) -> Result<Self, String>

Create a mesh with specific topology

pub fn set_rank(&mut self, rank: usize) -> Result<(), String>

Set the rank for this process

pub fn get_device(&self, rank: usize) -> Option<Arc<Device>>

Get device by rank (memory efficient - returns Arc clone)

pub fn local_device(&self) -> Option<Arc<Device>>

Get the local device for this process

pub fn all_devices(&self) -> Arc<Vec<Arc<Device>>>

Get all devices (returns Arc to avoid cloning the entire Vec)

pub fn devices_by_backend(&self, backend: DeviceBackend) -> Vec<Arc<Device>>

Get devices for a specific backend type (memory efficient)

pub fn devices_in_group( &self, group_name: &str, ) -> Result<Vec<Arc<Device>>, String>

Get devices in a specific process group

pub fn add_group( &mut self, name: String, ranks: Vec<usize>, backend: GroupBackend, ) -> Result<(), String>

Add a custom process group

pub fn get_group(&self, name: &str) -> Option<&ProcessGroup>

Get a process group by name

pub fn group_names(&self) -> Vec<String>

List all group names

pub fn in_group(&self, group_name: &str) -> bool

Check if this rank is in a specific group

pub fn group_rank(&self, group_name: &str) -> Option<usize>

Get rank within a group (local rank)

pub fn set_comm(&mut self, comm: Arc<dyn MeshComm + Send + Sync>)

Set the communication backend

pub fn total_memory_mb(&self) -> u64

Get total memory across all devices

pub fn total_compute_units(&self) -> u32

Get total compute units across all devices

pub fn reshape(&mut self, new_topology: MeshTopology) -> Result<(), String>

Reshape the mesh to a different topology

pub fn stats(&self) -> MeshStats

Get mesh statistics

Trait Implementations

impl Clone for DeviceMesh

fn clone(&self) -> DeviceMesh

Returns a duplicate of the value.

fn clone_from(&mut self, source: &Self)

Performs copy-assignment from source.

impl Debug for DeviceMesh

fn fmt(&self, f: &mut Formatter<'_>) -> Result

Formats the value using the given formatter.