etha.tensor_bus.batch_state

Batch state management for TensorBus.

A batch represents one call to register_tensors(). Multiple batches can exist for the same pair, each with independent tensors and handlers.

Classes

BatchState

Runtime state for a single batch of registered tensors, held on the Agent.

Module Contents

class etha.tensor_bus.batch_state.BatchState

Runtime state for a single batch of registered tensors, held on the Agent.

Created per register_tensors() call. Holds live tensors, process-group handles, and execution plans (buckets with CUDA work) — in-memory only.

Key design: buckets are FLATTENED across all pairs in the batch, allowing single-pass execution via bucket_comm().

batch_group: torch.distributed.ProcessGroup | None = None

batch_id: str

bucket_size: int | None = None

local_group: torch.distributed.ProcessGroup | None = None

local_leader: int | None = None

pair_names: list[str]

pair_target_dtypes: dict[str, list[torch.dtype]]

pair_tensors: dict[str, list[torch.Tensor]]

recv_buckets: list[etha.comm.ir.Bucket] | None = None

send_buckets: list[etha.comm.ir.Bucket] | None = None