mstar.communication.tensor_store

Tensor storage, split so the bookkeeping half can have a Rust backend.

TensorStore holds the torch.Tensor objects and stays Python: Rust must never own a tensor, or freeing a record inside a py.allow_threads section would drop a Python object without the GIL.

TensorBookkeeping holds everything else – refcounts, the persist flag, whether the memory has been registered for remote reads – and sees nothing but integer uuids. That is what gets ported, so the Rust graph runtime can adjust refcounts without a hop back to Python.

Uuids are globally unique now (see tensor_uuid), so nothing here is keyed by request. The one exception is _rid_to_uuids, which exists only so a request teardown can find what to free.

Classes

ColumnarTensorInfo(source_session_id, ...[, ...])	One node's output batch, marshalled as columns of primitives.
PythonTensorBookkeeping()
ReferenceInfo([ref_cnt, persist, mem_registered])
RustTensorBookkeeping()	Every method forwards; the batched ones go over in one call.
TensorBookkeeping()	Per-uuid reference state.
TensorStore([bookkeeping])	The tensors themselves, plus which request owns each one.

class mstar.communication.tensor_store.ColumnarTensorInfo(source_session_id, source_entity, source_node_name, source_graph_walk, tp_size, tp_rank, uuids=<factory>, addresses=<factory>, nbytes=<factory>, ndims=<factory>, dims_flat=<factory>, stride_flat=<factory>, dtype_idx=<factory>, dtype_names=<factory>)

Bases: object

One node’s output batch, marshalled as columns of primitives.

The fields above the columns are constant for the whole batch, so they cross into the bookkeeper once and intern once instead of per tensor. tp_size/tp_rank belong there too: the sharding group is a property of the node and this rank, not of the request, so one batch – one node, one graph walk, one worker – cannot be heterogeneous in them.

Parameters:

• source_session_id (str)

• source_entity (str)

• source_node_name (str)

• source_graph_walk (str)

• tp_size (int)

• tp_rank (int)

• uuids (list[int])

• addresses (list[int])

• nbytes (list[int])

• ndims (list[int])

• dims_flat (list[int])

• stride_flat (list[int])

• dtype_idx (list[int])

• dtype_names (list[str])

add_row(uuid, address, nbytes, ndims, dims, stride, dtype_idx)

One tensor whose layout is already known (a row of a template): its entry in every column, in the layout’s order.

Parameters:

• uuid (int)

• address (int)

• nbytes (int)

• ndims (int)

• dtype_idx (int)

Return type:

None

add_tensors_canonical(uuids, canonicals)

Parameters:

• uuids (list[int])

• canonicals (list[Tensor])

addresses: list[int]

dims_flat: list[int]

dtype_idx: list[int]

dtype_index(dtype)

The batch-local index of dtype, registering it on first use.

Parameters:

dtype (dtype)

Return type:

int

dtype_names: list[str]

nbytes: list[int]

ndims: list[int]

source_entity: str

source_graph_walk: str

source_node_name: str

source_session_id: str

stride_flat: list[int]

tp_rank: int

tp_size: int

uuids: list[int]

class mstar.communication.tensor_store.PythonTensorBookkeeping

Bases: TensorBookkeeping

can_gc(uuid)

No references left and not being persisted for the conductor.

Parameters:

uuid (int)

Return type:

bool

dereference(uuid, n=1)

Parameters:

• uuid (int)

• n (int)

forget_tensor(uuid)

Drop the record. The caller frees the tensor itself.

Parameters:

uuid (int)

get_info(uuid)

The descriptor a peer needs to read this tensor.

This is what lets an edge be rebuilt from uuids alone: ingestion and routing carry uuids, and anything that has to put a descriptor back on the wire (a disaggregated loop re-emitting its external inputs, say) looks it up here.

Parameters:

uuid (int)

Return type:

TensorPointerInfo | None

increment_ref(uuid, n=1)

Parameters:

• uuid (int)

• n (int)

is_registered(uuid)

Parameters:

uuid (int)

Return type:

bool

put_tensor(uuid, info)

Parameters:

• uuid (int)

• info (TensorPointerInfo)

register_shm_segment(index, name)

Name one arena segment, once, as the arena grows into it.

Keeps the name off set_shm_placement, which then crosses an index the caller already had.

Parameters:

• index (int)

• name (str)

set_mem_registered(uuid, mem_registered)

Parameters:

• uuid (int)

• mem_registered (bool)

set_persist(uuid, persist)

Parameters:

• uuid (int)

• persist (bool)

set_shm_placement(uuids, segment_idxs, offsets)

Stamp arena placement onto descriptors already held.

register_for_send changes only these two fields, so only they cross – not a rebuilt descriptor. Segments cross as indices from register_shm_segment: the caller reserved by index, and rendering a name per tensor costs a string allocation on each side plus a hash, to recover an integer it already had. -1 is “no segment”, the producer having spilled that tensor to its own file.

Parameters:

• uuids (list[int])

• segment_idxs (list[int])

• offsets (list[int])

class mstar.communication.tensor_store.ReferenceInfo(ref_cnt: 'int' = 0, persist: 'bool' = False, mem_registered: 'bool' = False)

Bases: object

Parameters:

• ref_cnt (int)

• persist (bool)

• mem_registered (bool)

mem_registered: bool = False

persist: bool = False

ref_cnt: int = 0

class mstar.communication.tensor_store.RustTensorBookkeeping

Bases: TensorBookkeeping

Every method forwards; the batched ones go over in one call.

A descriptor handed in is COPIED into Rust, unlike the Python backend which keeps the caller’s object. Anything relying on seeing a later in-place edit (register_for_send stamping shm_segment) has to write it back (set_shm_placement).

can_gc(uuid)

No references left and not being persisted for the conductor.

Parameters:

uuid (int)

Return type:

bool

dereference(uuid, n=1)

Parameters:

• uuid (int)

• n (int)

dereference_batch(uuids, counts, cleanup=False)

Drop counts[i] references from uuids[i], answering the gc question in the same pass.

Returns the uuids that became collectable and whether each was registered for remote reads – the one flag a teardown still needs after cleanup has dropped the record. Per uuid the caller would otherwise pay a crossing each for the decrement, can_gc, is_registered and forget_tensor.

Parameters:

• uuids (list[int])

• counts (list[int])

• cleanup (bool)

Return type:

tuple[list[int], list[bool]]

dereference_batch_uniform(uuids, n=1, cleanup=False)

dereference_batch for a uniform count, which is what a node’s consumed inputs and a streamed output batch always are.

Parameters:

• uuids (list[int])

• n (int)

• cleanup (bool)

Return type:

tuple[list[int], list[bool]]

forget_tensor(uuid)

Drop the record. The caller frees the tensor itself.

Parameters:

uuid (int)

get_info(uuid)

The descriptor a peer needs to read this tensor.

This is what lets an edge be rebuilt from uuids alone: ingestion and routing carry uuids, and anything that has to put a descriptor back on the wire (a disaggregated loop re-emitting its external inputs, say) looks it up here.

Parameters:

uuid (int)

Return type:

TensorPointerInfo | None

property has_put_tensor_batch_columns

increment_ref(uuid, n=1)

Parameters:

• uuid (int)

• n (int)

increment_ref_batch(uuids, counts)

Parameters:

• uuids (list[int])

• counts (list[int])

increment_ref_batch_uniform(uuids, n=1)

increment_ref_batch where every uuid takes the same count.

The safety hold on a freshly stored output batch is always 1, so the caller would otherwise build an n-long list of ones per batch purely to satisfy the parallel-list signature.

Parameters:

• uuids (list[int])

• n (int)

is_registered(uuid)

Parameters:

uuid (int)

Return type:

bool

put_tensor(uuid, info)

Parameters:

• uuid (int)

• info (TensorPointerInfo)

put_tensor_batch(uuids, infos)

Parameters:

• uuids (list[int])

• infos (list[TensorPointerInfo])

put_tensor_batch_columns(infos)

Store a batch already marshalled as columns of primitives.

_to_rust builds a 15-key dict per tensor and pyo3 then allocates a String per string field per tensor, which is where storing a batch actually spends its time – batching the call alone measured zero. Here the caller has filled the columns as it minted, so no TensorPointerInfo is built to be torn apart again, and the fields that are invariant across one node’s output batch cross and intern once rather than per tensor.

Parameters:

infos (ColumnarTensorInfo)

register_shm_segment(index, name)

Name one arena segment, once, as the arena grows into it.

Keeps the name off set_shm_placement, which then crosses an index the caller already had.

Parameters:

• index (int)

• name (str)

set_mem_registered(uuid, mem_registered)

Parameters:

• uuid (int)

• mem_registered (bool)

set_persist(uuid, persist)

Parameters:

• uuid (int)

• persist (bool)

set_shm_placement(uuids, segment_idxs, offsets)

Stamp arena placement onto descriptors already held.

register_for_send changes only these two fields, so only they cross – not a rebuilt descriptor. Segments cross as indices from register_shm_segment: the caller reserved by index, and rendering a name per tensor costs a string allocation on each side plus a hash, to recover an integer it already had. -1 is “no segment”, the producer having spilled that tensor to its own file.

Parameters:

• uuids (list[int])

• segment_idxs (list[int])

• offsets (list[int])

class mstar.communication.tensor_store.TensorBookkeeping

Bases: ABC

Per-uuid reference state. Never sees a tensor or a request id.

The defaults of the batched forms loop, so a backend only overrides where batching buys something. They take two parallel lists rather than a ParallelList because that is what crosses the Rust boundary – pyo3 converts a list of ints directly, where a NamedTuple would need unpacking on every call.

abstractmethod can_gc(uuid)

No references left and not being persisted for the conductor.

Parameters:

uuid (int)

Return type:

bool

abstractmethod dereference(uuid, n=1)

Parameters:

• uuid (int)

• n (int)

dereference_batch(uuids, counts, cleanup=False)

Drop counts[i] references from uuids[i], answering the gc question in the same pass.

Returns the uuids that became collectable and whether each was registered for remote reads – the one flag a teardown still needs after cleanup has dropped the record. Per uuid the caller would otherwise pay a crossing each for the decrement, can_gc, is_registered and forget_tensor.

Parameters:

• uuids (list[int])

• counts (list[int])

• cleanup (bool)

Return type:

tuple[list[int], list[bool]]

dereference_batch_uniform(uuids, n=1, cleanup=False)

dereference_batch for a uniform count, which is what a node’s consumed inputs and a streamed output batch always are.

Parameters:

• uuids (list[int])

• n (int)

• cleanup (bool)

Return type:

tuple[list[int], list[bool]]

abstractmethod forget_tensor(uuid)

Drop the record. The caller frees the tensor itself.

Parameters:

uuid (int)

abstractmethod get_info(uuid)

The descriptor a peer needs to read this tensor.

This is what lets an edge be rebuilt from uuids alone: ingestion and routing carry uuids, and anything that has to put a descriptor back on the wire (a disaggregated loop re-emitting its external inputs, say) looks it up here.

Parameters:

uuid (int)

Return type:

TensorPointerInfo | None

property has_put_tensor_batch_columns

abstractmethod increment_ref(uuid, n=1)

Parameters:

• uuid (int)

• n (int)

increment_ref_batch(uuids, counts)

Parameters:

• uuids (list[int])

• counts (list[int])

increment_ref_batch_uniform(uuids, n=1)

increment_ref_batch where every uuid takes the same count.

The safety hold on a freshly stored output batch is always 1, so the caller would otherwise build an n-long list of ones per batch purely to satisfy the parallel-list signature.

Parameters:

• uuids (list[int])

• n (int)

abstractmethod is_registered(uuid)

Parameters:

uuid (int)

Return type:

bool

abstractmethod put_tensor(uuid, info)

Parameters:

• uuid (int)

• info (TensorPointerInfo)

put_tensor_batch(uuids, infos)

Parameters:

• uuids (list[int])

• infos (list[TensorPointerInfo])

put_tensor_batch_columns(infos)

Columnar put, for a backend that pays per tensor to marshal.

Returns whether it handled the batch. The default says no: a Python bookkeeper stores the descriptor by reference and has nothing to marshal, so flattening it into columns would only add work.

Parameters:

infos (ColumnarTensorInfo)

abstractmethod register_shm_segment(index, name)

Name one arena segment, once, as the arena grows into it.

Keeps the name off set_shm_placement, which then crosses an index the caller already had.

Parameters:

• index (int)

• name (str)

abstractmethod set_mem_registered(uuid, mem_registered)

Parameters:

• uuid (int)

• mem_registered (bool)

abstractmethod set_persist(uuid, persist)

Parameters:

• uuid (int)

• persist (bool)

abstractmethod set_shm_placement(uuids, segment_idxs, offsets)

Stamp arena placement onto descriptors already held.

register_for_send changes only these two fields, so only they cross – not a rebuilt descriptor. Segments cross as indices from register_shm_segment: the caller reserved by index, and rendering a name per tensor costs a string allocation on each side plus a hash, to recover an integer it already had. -1 is “no segment”, the producer having spilled that tensor to its own file.

Parameters:

• uuids (list[int])

• segment_idxs (list[int])

• offsets (list[int])

class mstar.communication.tensor_store.TensorStore(bookkeeping=None)

Bases: object

The tensors themselves, plus which request owns each one.

Reference state lives in bookkeeping; the delegating methods here exist so callers have one object to talk to.

Parameters:

bookkeeping (TensorBookkeeping | None)

can_gc(uuid)

Parameters:

uuid (int)

Return type:

bool

check_uuid_presence(uuid)

Parameters:

uuid (int)

Return type:

bool

dereference(uuid, n=1)

Parameters:

• uuid (int)

• n (int)

dereference_batch(uuids, counts, cleanup=False)

dereference_batch_uniform where the count differs per uuid – an ack that covers several reads of the same tensor, say.

Parameters:

• uuids (list[int])

• counts (list[int])

• cleanup (bool)

Return type:

tuple[list[int], list[bool]]

dereference_batch_uniform(uuids, n=1, cleanup=False)

The uuids that became collectable, and whether each was registered.

cleanup forgets them in the bookkeeper as part of the same call; the tensors themselves stay here until the caller’s teardown removes them, which is what _forgotten covers.

Parameters:

• uuids (list[int])

• n (int)

• cleanup (bool)

Return type:

tuple[list[int], list[bool]]

get_all_uuids(request_id)

Parameters:

request_id (int | str)

Return type:

list[int]

get_cpu_tensor(uuid)

The host copy stored alongside uuid, if the producer gave one.

Parameters:

uuid (int)

Return type:

Tensor | None

get_info(uuid)

Parameters:

uuid (int)

Return type:

TensorPointerInfo | None

get_tensor(uuid)

Parameters:

uuid (int)

Return type:

Tensor

property has_put_tensor_batch_columns: bool

Whether the caller should marshal into ColumnarTensorInfo.

The manager asks before it starts minting, because the two forms are built differently – columnar fills its columns as it goes, where the other materialises a TensorPointerInfo per tensor.

increment_ref(uuid, n=1)

Parameters:

• uuid (int)

• n (int)

increment_ref_batch(uuids, counts)

Parameters:

• uuids (list[int])

• counts (list[int])

increment_ref_batch_uniform(uuids, n=1)

Parameters:

• uuids (list[int])

• n (int)

is_registered(uuid)

Parameters:

uuid (int)

Return type:

bool

mark_forgotten(uuids)

These records are already gone from the bookkeeper, so the teardown that follows must not cross again to drop them. For a caller that went to the bookkeeper directly – the Rust runtime releasing a node’s consumed inputs holds a share of it and never comes through here.

Parameters:

uuids (list[int])

put_tensor(request_id, uuid, tensor, info, cpu_tensor=None)

Parameters:

• request_id (int | str)

• uuid (int)

• tensor (Tensor)

• info (TensorPointerInfo)

• cpu_tensor (Tensor | None)

put_tensor_batch_columns(request_ids, tensors, columns, cpu_tensors=None)

put_tensor_batch_multi for a backend that marshals per tensor.

Takes the columns the caller filled while minting instead of a descriptor per tensor. columns.uuids is the same order as tensors, which is what keeps the ownership maps aligned.

Parameters:

• request_ids (list[int | str])

• tensors (list[Tensor])

• columns (ColumnarTensorInfo)

• cpu_tensors (list[Tensor | None] | None)

put_tensor_batch_multi(request_ids, uuids, tensors, infos, cpu_tensors=None)

put_tensor_batch spanning requests: one bookkeeping call for a whole output batch rather than one per request.

The maps below are plain dicts, so keeping them in a Python loop costs nothing; it is bookkeeping.put_tensor_batch that is a Rust boundary crossing, and taking every request at once collapses it to one.

Parameters:

• request_ids (list[int | str])

• uuids (list[int])

• tensors (list[Tensor])

• infos (list[TensorPointerInfo])

• cpu_tensors (list[Tensor | None] | None)

register_shm_segment(index, name)

Parameters:

• index (int)

• name (str)

remove_tensor(uuid)

Parameters:

uuid (int)

set_metadata(uuid, persist=None, mem_registered=None)

Parameters:

• uuid (int)

• persist (bool | None)

• mem_registered (bool | None)

set_shm_placement(uuids, segment_idxs, offsets)

Parameters:

• uuids (list[int])

• segment_idxs (list[int])

• offsets (list[int])