gigl.distributed.utils.fused_feature#

Fetching a raw feature matrix and its packed sidecar with one request per partition.

A node or edge type can store its columns in a raw float matrix, in a packed uint8 sidecar (feature quantization), or in both. Fetching two stores through two DistFeature objects costs two RPCs per remote partition for every batch. FusedFeatureStore serves whichever stores a type has as one row of bytes (the float columns reinterpreted as bytes, followed by the sidecar columns) and FusedDistFeature is a DistFeature over such stores, so one lookup is one request per partition for every type. FusedDistFeature.async_get resolves to the raw matrix and the sidecar rows, the latter None for a type without a sidecar.

All stores are CPU stores: GiGL builds every Feature with with_gpu=False.

Classes#

`FusedDistFeature`	A `DistFeature` whose local stores are `FusedFeatureStore` objects.
`FusedFeatureStore`	A raw float feature store and/or its packed sidecar, served to other partitions as one `uint8` row per entity.

Functions#

build_fused_dist_feature(num_partitions, ...[, device])

Builds a FusedDistFeature over every entity type that has a raw store, a packed sidecar, or both.

Module Contents#

class gigl.distributed.utils.fused_feature.FusedDistFeature(num_partitions, partition_idx, local_feature, feature_pb, rpc_router, device=None)[source]#

Bases: graphlearn_torch.distributed.dist_feature.DistFeature

A DistFeature whose local stores are FusedFeatureStore objects.

GLT’s DistFeature looks up local rows and answers remote requests by indexing its local stores with ids, which FusedFeatureStore supports; only the constructor checks the store type, so it is repeated here without that check.

async_get resolves to the raw float matrix and the packed sidecar rows, stitched straight from the local rows and the received byte rows. The sidecar rows are None for a type without a sidecar.

Parameters:

num_partitions (int) – Number of data partitions.
partition_idx (int) – Data partition of the current process.
local_feature (Union[FusedFeatureStore, dict[Union[graphlearn_torch.typing.NodeType, graphlearn_torch.typing.EdgeType], FusedFeatureStore]]) – A fused store, or one per node/edge type for heterogeneous data.
feature_pb (Union[graphlearn_torch.partition.PartitionBook, dict[Union[graphlearn_torch.typing.NodeType, graphlearn_torch.typing.EdgeType], graphlearn_torch.partition.PartitionBook]]) – Partition book of the entities, or one per type for heterogeneous data.
rpc_router (graphlearn_torch.distributed.rpc.RpcDataPartitionRouter) – Router used to reach the other partitions.
device (Optional[torch.device]) – Device used for computing.

async_get(ids, input_type=None)[source]#

Looks up the raw float rows and packed sidecar rows of ids.

Parameters:

ids (jaxtyping.Int64[torch.Tensor, num_ids]) – Ids to look up.
input_type (Optional[Union[graphlearn_torch.typing.NodeType, graphlearn_torch.typing.EdgeType]]) – Node or edge type of the ids for heterogeneous data.

Returns:

A future of the [len(ids), raw_dim] raw matrix and the [len(ids), packed_dim] sidecar rows, or None for the sidecar rows if the type has no sidecar.

Return type:

torch.futures.Future[tuple[jaxtyping.Float[torch.Tensor, num_ids raw_dim], Optional[jaxtyping.UInt8[torch.Tensor, num_ids packed_dim]]]]

device[source]#

feature_pb[source]#

local_feature[source]#

num_partitions[source]#

partition_idx[source]#

rpc_callee_id = 1[source]#

rpc_router[source]#

class gigl.distributed.utils.fused_feature.FusedFeatureStore(raw, packed)[source]#

A raw float feature store and/or its packed sidecar, served to other partitions as one uint8 row per entity.

A served row holds the raw float columns reinterpreted as bytes, then the sidecar columns, then zero padding up to a multiple of the raw dtype’s size. The padding keeps every row aligned, so the raw columns of a received row can be read as floats without a copy.

A missing or zero-width store contributes no columns. Without a raw store the raw dtype is float32, the dtype the raw matrix is later rebuilt as. Without a sidecar the row is the raw bytes alone, which need no padding. A row is never narrower than one raw element, so a type with no columns at all is served as one zero element.

Both stores must be CPU stores.

Example

A float32 raw store with 3 columns (12 bytes) and a 2-column uint8 sidecar make 14 bytes, padded to 16. Bytes 0-11 are the raw floats, bytes 12-13 the sidecar and bytes 14-15 zero padding.

>>> raw = Feature(torch.ones(4, 3), id2index=torch.arange(4), with_gpu=False)
>>> packed = Feature(
...     torch.full((4, 2), 7, dtype=torch.uint8),
...     id2index=torch.arange(4),
...     with_gpu=False,
...     dtype=torch.uint8,
... )
>>> store = FusedFeatureStore(raw, packed)
>>> rows = store.cpu_get(torch.tensor([0, 2]))
>>> rows.shape
torch.Size([2, 16])
>>> rows[0, 12:].tolist()
[7, 7, 0, 0]
>>> raw_rows, packed_rows = store.split(rows)
>>> raw_rows.shape, packed_rows.shape
(torch.Size([2, 3]), torch.Size([2, 2]))

Parameters:

raw (Optional[graphlearn_torch.data.Feature]) – Raw float feature store, or None if the type has none.
packed (Optional[graphlearn_torch.data.Feature]) – Packed uint8 sidecar store with the same ids and id mapping as raw, or None.

Raises:

ValueError – If both stores are None.

cpu_get(ids)[source]#

Served rows of ids from the CPU tensors, as [len(ids), row_bytes] bytes, for answering remote requests.

Each store is gathered straight into its columns of the output, so every value is copied once.

Parameters:: ids (jaxtyping.Int64[torch.Tensor, num_ids])
Return type:: jaxtyping.UInt8[torch.Tensor, num_ids row_bytes]

lazy_init_with_ipc_handle()[source]#

Materializes the stores after an IPC transfer.

Return type:: None

split(rows)[source]#

Views of served byte rows as the raw float matrix and the packed sidecar rows, without copying.

The views have zero columns for a missing store.

Parameters:: rows (jaxtyping.UInt8[torch.Tensor, num_rows row_bytes])
Return type:: tuple[jaxtyping.Float[torch.Tensor, num_rows raw_dim], jaxtyping.UInt8[torch.Tensor, num_rows packed_dim]]

gigl.distributed.utils.fused_feature.build_fused_dist_feature(num_partitions, partition_idx, raw_features, packed_features, feature_pb, rpc_router, device=None)[source]#

Builds a FusedDistFeature over every entity type that has a raw store, a packed sidecar, or both.

Parameters:

num_partitions (int) – Number of data partitions.
partition_idx (int) – Data partition of the current process.
raw_features (Optional[Union[graphlearn_torch.data.Feature, dict[Union[graphlearn_torch.typing.NodeType, graphlearn_torch.typing.EdgeType], graphlearn_torch.data.Feature]]]) – Raw float store, or one per type.
packed_features (Optional[Union[graphlearn_torch.data.Feature, dict[Union[graphlearn_torch.typing.NodeType, graphlearn_torch.typing.EdgeType], graphlearn_torch.data.Feature]]]) – Packed sidecar store, or one per type.
feature_pb (Union[graphlearn_torch.partition.PartitionBook, dict[Union[graphlearn_torch.typing.NodeType, graphlearn_torch.typing.EdgeType], graphlearn_torch.partition.PartitionBook]]) – Partition book of the entities, or one per type.
rpc_router (graphlearn_torch.distributed.rpc.RpcDataPartitionRouter) – Router used to reach the other partitions.
device (Optional[torch.device]) – Device used for computing.

Returns:

None if there are no stores at all.

Return type:

Optional[FusedDistFeature]