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
|
Tensor transport via the Rust shared-memory arena ( |
- class mstar.communication.arena.ArenaShmCommunicationManager(my_entity_id, hostname, device, communicator, shm_dir=None, enable_prof=False)[source]#
Bases:
SharedMemoryCommunicationManagerTensor 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=Trueskips 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 (str)
tensor_infos (list[TensorPointerInfo])
skip_cuda_sync (bool)