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 (or until /dev/shm cannot back another segment: segments commit their pages when created); 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)

Bases: SharedMemoryCommunicationManager

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

Parameters:

• my_entity_id (str)

• hostname (str)

• device (str)

• communicator (BaseCommunicator)

• shm_dir (str | None)

• enable_prof (bool)

cleanup_collectable(collectable, registered)

_cleanup_by_uuid for a batch, inlined: a decode step frees one tensor per request here, and the three-level call chain per uuid was most of the cost. Same steps in the same order as the per-uuid form.

Parameters:

• collectable (list[int])

• registered (list[bool])

close()

Stop the wake watcher and unregister every pinned segment (the segments themselves unmap with the arena’s Drop once the manager is garbage collected). Segments mapped after this stay unpinned.

Return type:

None

register_for_send(request_id, tensor_infos, skip_cuda_sync=False)

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:

• request_id (int | str)

• tensor_infos (list[TensorPointerInfo])

• skip_cuda_sync (bool)

register_for_send_uuids(per_request, skip_cuda_sync=False)

Uuid-driven staging: no descriptor is rebuilt to be read for its uuid, and placement goes back as two columns rather than a whole descriptor. The whole batch stages inside one stream context and syncs once at the end, rather than one host-blocking sync per request.

Parameters:

• per_request (ParallelList[int | str, list[int]])

• skip_cuda_sync (bool)

stamps_shm_placement: bool = True

Whether this transport writes placement onto a descriptor AFTER routing has already copied it. Only the SHM arena does (shm_segment / shm_offset); the per-uuid file and Mooncake paths derive everything a consumer needs from uuid and address, so they leave this off and refresh_shm_placement stays free for them.

start_read_tensors(request_id, graph_edges, graph_walk=None)

Parameters:

• request_id (int | str)

• graph_edges (list[GraphEdge])

• graph_walk (str | None)

stats_summary()

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

unlink_owned_files()

The spill files, then this arena’s own segments, for a process that is stopping: its exit-time finalizer never runs when a hung teardown ends in SIGKILL or the watchdog’s hard exit. Peers keep the mappings they have; the finalizer still runs at a normal exit and finds them gone, and unlinks any segment grown after this call.

Return type:

None