mstar.communication.arena#

Tensor transport over a shared-memory arena.

ArenaShmCommunicationManager replaces the file transport’s per-tensor open/write/read/unlink with a Rust segmented /dev/shm arena (persistent mmaps + first-fit coalescing allocator, vendored in rust/). Producer: register_for_send reserves a slot and D2H-copies into the segment on the copy stream; the (segment, offset) rides the existing descriptors. Consumer: start_read_tensors maps the named segment once, reads zero-copy, H2D-copies on the copy stream; the producer reclaims on ACK, gated by a CUDA-event future so an edge is never ACKed before its copies land. Segments are mapped once and never move, so the one-time cudaHostRegister per segment (within MSTAR_SHM_ARENA_PIN_MAX_MB) holds for its lifetime and keeps the side-stream copies truly async.

Capacity degrades in layers: grow by segments up to MSTAR_SHM_ARENA_MAX_SEGMENTS; at the cap, briefly backpressure for consumer ACKs; then spill the tensor to the per-uuid file protocol (MSTAR_SHM_ARENA_SPILL, default on) — slower, never fails, like the old transport at saturation. stats_summary() exposes occupancy and the fragmentation gauge (largest contiguous free block); --log-stats logs it periodically.

Ceilings are PER-ENTITY and multiply across a node: with E entities (workers + the api-server data worker), /dev/shm demand can reach MAX_SEGMENTS x SEGMENT_MB x E and pinned host RAM approx PIN_MAX_MB x E (consumers pin peer segments too, so one process can pin more than its own arena holds). Construction fails fast when one entity’s ceiling already exceeds /dev/shm, and warns when it exceeds current free space or when the pin budget is an outsized share of physical RAM.

Selection: create_tensor_communication_manager picks this manager for the SHM protocol when MSTAR_SHM_ARENA is 1 (require) or AUTO (use if the mstar_rust extension imports); default 0 keeps the file transport. See environment_variables for all knobs.

Classes

ArenaShmCommunicationManager(my_entity_id, ...)

Tensor transport via the Rust shared-memory arena (mstar_rust).

class mstar.communication.arena.ArenaShmCommunicationManager(my_entity_id, hostname, device, communicator, shm_dir=None, enable_prof=False)[source]#

Bases: SharedMemoryCommunicationManager

Tensor transport via the Rust shared-memory arena (mstar_rust).

Parameters:
close()[source]#

Stop the wake watcher (segments/pins release with the arena’s Drop once the manager is garbage collected).

Return type:

None

register_for_send(request_id, tensor_infos, skip_cuda_sync=False)[source]#

Mark these tensors ready for remote consumers to RDMA-read.

skip_cuda_sync=True skips the default-stream sync this call normally issues to ensure the source tensors’ writes are visible before their addresses are shared with peers. Callers must have already synced on their own (e.g. before a batched loop) — meant to cut N serialized syncs to 1 when registering many tensors in a row.

If self.enable_prof is set, this should also update self.req_tx_info

Parameters:
start_read_tensors(request_id, graph_edges, graph_walk=None)[source]#
Parameters:
stats_summary()[source]#

Occupancy/fragmentation snapshot (named apart from the raw SegmentedShmArena.stats tuple). The fragmentation signature is largest_free_block collapsing while free_bytes stays high.

Return type:

dict