From 381538d876ce98c7627809d730493f5c6fc7ddb8 Mon Sep 17 00:00:00 2001 From: jchmura Date: Sat, 15 Aug 2026 19:25:39 +0000 Subject: [PATCH] Load and partition quantized edge features --- gigl/common/data/dataloaders.py | 21 +-- gigl/common/data/load_torch_tensors.py | 149 ++++++++++++++++-- gigl/distributed/dataset_factory.py | 5 + gigl/distributed/dist_partitioner.py | 146 ++++++++++++++--- gigl/distributed/dist_range_partitioner.py | 105 ++++++++++-- .../serialized_graph_metadata_translator.py | 24 ++- gigl/types/graph.py | 9 ++ .../run_distributed_partitioner.py | 29 +++- tests/unit/common/data/dataloaders_test.py | 87 ++++++++++ .../distributed_partitioner_test.py | 48 +++++- .../distributed_weighted_sampling_test.py | 12 +- 11 files changed, 570 insertions(+), 65 deletions(-) diff --git a/gigl/common/data/dataloaders.py b/gigl/common/data/dataloaders.py index 824e5225d..67152c898 100644 --- a/gigl/common/data/dataloaders.py +++ b/gigl/common/data/dataloaders.py @@ -398,16 +398,6 @@ def load_as_torch_tensors( feature_spec_dict[entity_key] = tf.io.FixedLenFeature( shape=[], dtype=tf.int64 ) - if ( - packed_feature_key is not None - and packed_feature_key not in feature_spec_dict - ): - logger.info( - f"Injecting packed feature key {packed_feature_key} into feature spec dictionary with value `tf.io.FixedLenFeature(shape=[], dtype=tf.string)`" - ) - feature_spec_dict[packed_feature_key] = tf.io.FixedLenFeature( - shape=[], dtype=tf.string - ) else: id_concat_axis = 1 proccess_id_tensor = lambda t: tf.stack( @@ -433,6 +423,17 @@ def load_as_torch_tensors( shape=[], dtype=tf.int64 ) + if ( + packed_feature_key is not None + and packed_feature_key not in feature_spec_dict + ): + logger.info( + f"Injecting packed feature key {packed_feature_key} into feature spec dictionary with value `tf.io.FixedLenFeature(shape=[], dtype=tf.string)`" + ) + feature_spec_dict[packed_feature_key] = tf.io.FixedLenFeature( + shape=[], dtype=tf.string + ) + uris = self._partition_children_uris( serialized_tf_record_info.tfrecord_uri_prefix, serialized_tf_record_info.tfrecord_uri_pattern, diff --git a/gigl/common/data/load_torch_tensors.py b/gigl/common/data/load_torch_tensors.py index 3a1174888..667bc348d 100644 --- a/gigl/common/data/load_torch_tensors.py +++ b/gigl/common/data/load_torch_tensors.py @@ -1,7 +1,7 @@ import time import traceback -from dataclasses import dataclass -from typing import MutableMapping, Optional, Union +from dataclasses import dataclass, replace +from typing import MutableMapping, Optional, Union, cast import torch import torch.multiprocessing as mp @@ -119,6 +119,134 @@ class SerializedGraphMetadata: node_quantization_metadata: Optional[ Union[FeatureQuantizationMetadata, dict[NodeType, FeatureQuantizationMetadata]] ] = None + edge_quantization_metadata: Optional[ + Union[FeatureQuantizationMetadata, dict[EdgeType, FeatureQuantizationMetadata]] + ] = None + + +def _validate_weight_edge_feature_name( + edge_entity_info: Union[ + SerializedTFRecordInfo, dict[EdgeType, SerializedTFRecordInfo] + ], + weight_edge_feat_name: Optional[Union[str, dict[EdgeType, str]]], +) -> None: + if weight_edge_feat_name is None: + return + + configured_weights: list[tuple[EdgeType, str, SerializedTFRecordInfo]] + if isinstance(edge_entity_info, SerializedTFRecordInfo): + if not isinstance(weight_edge_feat_name, str): + raise ValueError("weight_edge_feat_name must be str for homogeneous graph") + edge_type = DEFAULT_HOMOGENEOUS_EDGE_TYPE + configured_weights = [(edge_type, weight_edge_feat_name, edge_entity_info)] + else: + if isinstance(weight_edge_feat_name, str): + if len(edge_entity_info) != 1: + raise ValueError( + "weight_edge_feat_name must be dict[EdgeType, str] for heterogeneous graph with multiple edge types" + ) + edge_type, serialized_info = next(iter(edge_entity_info.items())) + configured_weights = [(edge_type, weight_edge_feat_name, serialized_info)] + else: + unknown_edge_types = set(weight_edge_feat_name) - set(edge_entity_info) + if unknown_edge_types: + raise ValueError( + f"weight_edge_feat_name contains unknown edge types: {unknown_edge_types}" + ) + configured_weights = [ + (edge_type, feature_name, edge_entity_info[edge_type]) + for edge_type, feature_name in weight_edge_feat_name.items() + ] + + for edge_type, feature_name, serialized_info in configured_weights: + if feature_name not in serialized_info.feature_keys: + raise ValueError( + f"Sampling-weight field '{feature_name}' for edge type {edge_type} must be an unquantized raw edge feature." + ) + + +def remove_sampling_weight_from_edge_quantization_metadata( + serialized_graph_metadata: SerializedGraphMetadata, + weight_edge_feat_name: Optional[Union[str, dict[EdgeType, str]]], +) -> Optional[ + Union[FeatureQuantizationMetadata, dict[EdgeType, FeatureQuantizationMetadata]] +]: + """Remove separately stored sampling weights from edge reconstruction metadata. + + TFRecord loading removes the sampling-weight column from raw edge features + before registering it with the weighted sampler. The resulting metadata + must describe the remaining model features so batch reconstruction scatters + raw and dequantized columns into the correct positions. + + Args: + serialized_graph_metadata: Serialized edge schema and quantization metadata. + weight_edge_feat_name: Raw scalar feature configured as sampling weights. + + Returns: + Quantization metadata for the model-facing edge features. + """ + quantization_metadata = serialized_graph_metadata.edge_quantization_metadata + if quantization_metadata is None or weight_edge_feat_name is None: + return quantization_metadata + + if isinstance(serialized_graph_metadata.edge_entity_info, SerializedTFRecordInfo): + assert isinstance(quantization_metadata, FeatureQuantizationMetadata) + assert isinstance(weight_edge_feat_name, str) + edge_info_by_type: dict[EdgeType, SerializedTFRecordInfo] = { + DEFAULT_HOMOGENEOUS_EDGE_TYPE: serialized_graph_metadata.edge_entity_info + } + metadata_by_type: dict[EdgeType, FeatureQuantizationMetadata] = { + DEFAULT_HOMOGENEOUS_EDGE_TYPE: quantization_metadata + } + weight_by_type: dict[EdgeType, str] = { + DEFAULT_HOMOGENEOUS_EDGE_TYPE: weight_edge_feat_name + } + is_homogeneous = True + else: + assert isinstance(quantization_metadata, dict) + edge_info_by_type: dict[EdgeType, SerializedTFRecordInfo] = ( + serialized_graph_metadata.edge_entity_info + ) + metadata_by_type: dict[EdgeType, FeatureQuantizationMetadata] = cast( + dict[EdgeType, FeatureQuantizationMetadata], quantization_metadata + ) + if isinstance(weight_edge_feat_name, str): + edge_type = next(iter(edge_info_by_type)) + weight_by_type: dict[EdgeType, str] = {edge_type: weight_edge_feat_name} + else: + weight_by_type: dict[EdgeType, str] = weight_edge_feat_name + is_homogeneous = False + + adjusted_metadata: dict[EdgeType, FeatureQuantizationMetadata] = {} + for edge_type, metadata in metadata_by_type.items(): + weight_feature_name = weight_by_type.get(edge_type) + if weight_feature_name is None: + adjusted_metadata[edge_type] = metadata + continue + + edge_info = edge_info_by_type[edge_type] + raw_column_offset = 0 + for feature_name in edge_info.feature_keys: + if feature_name == weight_feature_name: + break + feature_spec = edge_info.feature_spec[feature_name] + raw_column_offset += feature_spec.shape[-1] if feature_spec.shape else 1 + weight_logical_index = metadata.raw_feature_indices[raw_column_offset] + adjusted_quantized_feature_indices = tuple( + quantized_feature_index - 1 + if quantized_feature_index > weight_logical_index + else quantized_feature_index + for quantized_feature_index in metadata.quantized_feature_indices + ) + adjusted_metadata[edge_type] = replace( + metadata, + feature_dim=metadata.feature_dim - 1, + quantized_feature_indices=adjusted_quantized_feature_indices, + ) + + if is_homogeneous: + return adjusted_metadata[DEFAULT_HOMOGENEOUS_EDGE_TYPE] + return adjusted_metadata def _data_loading_process( @@ -199,14 +327,6 @@ def _data_loading_process( raise NotImplementedError( "Label keys are not supported for edge entities" ) - if ( - serialized_entity_tf_record_info.packed_feature_key is not None - and not serialized_entity_tf_record_info.is_node_entity - ): - # TODO(quantization): Support feature quantization for edge features. - raise NotImplementedError( - "Packed feature keys are not supported for edge entities" - ) loaded_entity = tf_record_dataloader.load_as_torch_tensors( serialized_tf_record_info=serialized_entity_tf_record_info, tf_dataset_options=tf_dataset_options, @@ -396,6 +516,11 @@ def load_torch_tensors_from_tf_record( loaded_graph_tensors (LoadedGraphTensors): Unpartitioned Graph Tensors """ + _validate_weight_edge_feature_name( + edge_entity_info=serialized_graph_metadata.edge_entity_info, + weight_edge_feat_name=weight_edge_feat_name, + ) + logger.info(f"Rank {rank} starting loading torch tensors from serialized info ...") start_time = time.time() @@ -525,6 +650,9 @@ def load_torch_tensors_from_tf_record( edge_index = edge_output_dict[_ID_FMT.format(entity=_EDGE_KEY)] edge_features = edge_output_dict.get(_FEATURE_FMT.format(entity=_EDGE_KEY), None) + edge_quantized_features = edge_output_dict.get( + _PACKED_FEATURE_FMT.format(entity=_EDGE_KEY), None + ) edge_weights = edge_output_dict.get(_EDGE_WEIGHTS_KEY, None) positive_labels = edge_output_dict.get( @@ -552,6 +680,7 @@ def load_torch_tensors_from_tf_record( node_labels=node_labels, edge_index=edge_index, edge_features=edge_features, + edge_quantized_features=edge_quantized_features, positive_label=positive_labels, negative_label=negative_labels, edge_weights=edge_weights, diff --git a/gigl/distributed/dataset_factory.py b/gigl/distributed/dataset_factory.py index 1a5f859b5..14e874be3 100644 --- a/gigl/distributed/dataset_factory.py +++ b/gigl/distributed/dataset_factory.py @@ -194,6 +194,10 @@ def _load_and_build_partitioned_dataset( partitioner.register_edge_features( edge_features=loaded_graph_tensors.edge_features ) + if loaded_graph_tensors.edge_quantized_features is not None: + partitioner.register_edge_quantized_features( + edge_quantized_features=loaded_graph_tensors.edge_quantized_features + ) if loaded_graph_tensors.positive_label is not None: partitioner.register_labels( label_edge_index=loaded_graph_tensors.positive_label, is_positive=True @@ -212,6 +216,7 @@ def _load_and_build_partitioned_dataset( loaded_graph_tensors.node_quantized_features, loaded_graph_tensors.edge_index, loaded_graph_tensors.edge_features, + loaded_graph_tensors.edge_quantized_features, loaded_graph_tensors.edge_weights, loaded_graph_tensors.positive_label, loaded_graph_tensors.negative_label, diff --git a/gigl/distributed/dist_partitioner.py b/gigl/distributed/dist_partitioner.py index 04de8ce72..9549c9e3f 100644 --- a/gigl/distributed/dist_partitioner.py +++ b/gigl/distributed/dist_partitioner.py @@ -208,6 +208,8 @@ def __init__( self._edge_ids: Optional[dict[EdgeType, tuple[int, int]]] = None self._edge_feat: Optional[dict[EdgeType, torch.Tensor]] = None self._edge_feat_dim: Optional[dict[EdgeType, int]] = None + self._edge_quantized_feat: Optional[dict[EdgeType, torch.Tensor]] = None + self._edge_quantized_feat_dim: Optional[dict[EdgeType, int]] = None self._edge_weights: Optional[dict[EdgeType, torch.Tensor]] = None # TODO (mkolodner-sc): Deprecate the need for explicitly storing labels are part of this class, leveraging @@ -669,6 +671,36 @@ def register_edge_features( for edge_type in input_edge_features: self._edge_feat_dim[edge_type] = input_edge_features[edge_type].shape[1] + def register_edge_quantized_features( + self, edge_quantized_features: Union[torch.Tensor, dict[EdgeType, torch.Tensor]] + ) -> None: + """Register packed uint8 main-edge features for co-partitioning.""" + + self._assert_and_get_rpc_setup() + if self._edge_quantized_feat is not None: + raise ValueError( + "Edge quantized features have already been registered. Cannot re-register edge quantized feature data." + ) + logger.info("Registering Edge Quantized Features ...") + + input_edge_quantized_features = ( + self._convert_edge_entity_to_heterogeneous_format( + input_edge_entity=edge_quantized_features + ) + ) + + assert input_edge_quantized_features, ( + "Edge quantized features is an empty dictionary. Please provide edge quantized features to register." + ) + + self._edge_quantized_feat = convert_to_tensor( + input_edge_quantized_features, dtype=torch.uint8 + ) + self._edge_quantized_feat_dim = { + edge_type: features.shape[1] + for edge_type, features in input_edge_quantized_features.items() + } + def register_edge_weights( self, edge_weights: Union[torch.Tensor, dict[EdgeType, torch.Tensor]] ) -> None: @@ -1201,7 +1233,10 @@ def _partition_edge_index_and_edge_features( node_partition_book: dict[NodeType, PartitionBook], edge_type: EdgeType, ) -> Tuple[ - GraphPartitionData, Optional[FeaturePartitionData], Optional[PartitionBook] + GraphPartitionData, + Optional[FeaturePartitionData], + Optional[FeaturePartitionData], + Optional[PartitionBook], ]: r"""Partition graph topology and edge features of a specific edge type. If there are no edge features for the current edge type, both the returned edge feature and edge partition book will be None. @@ -1213,6 +1248,7 @@ def _partition_edge_index_and_edge_features( Returns: GraphPartitionData: The graph data of the current partition. Optional[FeaturePartitionData]: The edge features on the current partition, will be None if there are no edge features for the current edge type + Optional[FeaturePartitionData]: The quantized edge features on the current partition, will be None if there are no quantized edge features for the current edge type Optional[PartitionBook]: The partition book of graph edges, will be None if there are no edge features for the current edge type """ @@ -1225,11 +1261,17 @@ def _partition_edge_index_and_edge_features( ), "Must have registered edges prior to partitioning them" has_edge_feats = self._edge_feat is not None and edge_type in self._edge_feat + has_edge_quantized_feats = ( + self._edge_quantized_feat is not None + and edge_type in self._edge_quantized_feat + ) has_weights_for_edge_type = ( self._edge_weights is not None and edge_type in self._edge_weights ) # Need a partition book if we have features or weights to reindex. - should_generate_partition_book = has_edge_feats or has_weights_for_edge_type + should_generate_partition_book = ( + has_edge_feats or has_edge_quantized_feats or has_weights_for_edge_type + ) # Partitioning Edge Indices @@ -1283,12 +1325,12 @@ def _edge_pfn(_, chunk_range): gc.collect() - # Partition edge features and weights together in a single pass, + # Partition edge features, packed features, and weights together in a single pass, # mirroring how node features and labels are co-partitioned. - # Input tuple layout: (edge_feat?, edge_weights?, edge_ids) - # IDs are always at r[-1]; features at r[0]; weights at r[1] when - # features are also present, else r[0]. + # Input tuple layout: (edge_feat?, edge_quantized_feat?, edge_weights?, edge_ids) + # IDs are always last; optional tensor indices are recorded when appended. current_feat_part: Optional[FeaturePartitionData] = None + current_quantized_feat_part: Optional[FeaturePartitionData] = None partitioned_weights: Optional[torch.Tensor] = None partitioned_edge_ids: Optional[torch.Tensor] = None @@ -1309,6 +1351,8 @@ def _edge_pfn(_, chunk_range): edge_feat: Optional[torch.Tensor] = None edge_feat_dim: Optional[int] = None edge_weights_tensor: Optional[torch.Tensor] = None + edge_quantized_features: Optional[torch.Tensor] = None + edge_quantized_feature_dim: Optional[int] = None if has_edge_feats: assert self._edge_feat is not None and edge_type in self._edge_feat assert ( @@ -1316,30 +1360,38 @@ def _edge_pfn(_, chunk_range): ) edge_feat = self._edge_feat[edge_type] edge_feat_dim = self._edge_feat_dim[edge_type] + if has_edge_quantized_feats: + assert self._edge_quantized_feat is not None + assert self._edge_quantized_feat_dim is not None + edge_quantized_features = self._edge_quantized_feat[edge_type] + edge_quantized_feature_dim = self._edge_quantized_feat_dim[edge_type] if has_weights_for_edge_type: assert self._edge_weights is not None edge_weights_tensor = self._edge_weights[edge_type] input_parts: list[torch.Tensor] = [] + feat_idx: Optional[int] = None if edge_feat is not None: + feat_idx = len(input_parts) input_parts.append(edge_feat) + quantized_feat_idx: Optional[int] = None + if edge_quantized_features is not None: + quantized_feat_idx = len(input_parts) + input_parts.append(edge_quantized_features) + weight_idx: Optional[int] = None if edge_weights_tensor is not None: + weight_idx = len(input_parts) input_parts.append(edge_weights_tensor) input_parts.append(edge_ids) - # Positional indices: features first, weights next, ids always last. - feat_idx: Optional[int] = 0 if has_edge_feats else None - weight_idx: Optional[int] = None - if has_weights_for_edge_type: - weight_idx = 1 if has_edge_feats else 0 - + # Recorded indices keep result unpacking aligned with optional inputs. def _edge_feat_weight_pfn( ids_chunk: torch.Tensor, _: object ) -> torch.Tensor: assert edge_partition_book is not None return edge_partition_book[ids_chunk] - # Each result tuple contains (edge_feat?, edge_weights?, edge_ids). + # Each result tuple preserves the input tuple layout. feat_weight_res_list, _ = self._partition_by_chunk( input_data=tuple(input_parts), rank_indices=edge_ids, @@ -1360,6 +1412,21 @@ def _edge_feat_weight_pfn( if len(self._edge_feat) == 0 and len(self._edge_feat_dim) == 0: self._edge_feat = None self._edge_feat_dim = None + if has_edge_quantized_feats: + assert edge_quantized_features is not None + assert self._edge_quantized_feat is not None + assert self._edge_quantized_feat_dim is not None + del edge_quantized_features + del ( + self._edge_quantized_feat[edge_type], + self._edge_quantized_feat_dim[edge_type], + ) + if ( + len(self._edge_quantized_feat) == 0 + and len(self._edge_quantized_feat_dim) == 0 + ): + self._edge_quantized_feat = None + self._edge_quantized_feat_dim = None if has_weights_for_edge_type: assert edge_weights_tensor is not None assert self._edge_weights is not None @@ -1377,6 +1444,14 @@ def _edge_feat_weight_pfn( feats=torch.empty(0, edge_feat_dim), ids=partitioned_edge_ids, ) + if has_edge_quantized_feats: + assert edge_quantized_feature_dim is not None + current_quantized_feat_part = FeaturePartitionData( + feats=torch.empty( + 0, edge_quantized_feature_dim, dtype=torch.uint8 + ), + ids=partitioned_edge_ids, + ) if has_weights_for_edge_type: partitioned_weights = torch.empty(0) else: @@ -1387,6 +1462,14 @@ def _edge_feat_weight_pfn( feats=torch.cat([r[feat_idx] for r in feat_weight_res_list]), ids=partitioned_edge_ids, ) + if has_edge_quantized_feats: + assert quantized_feat_idx is not None + current_quantized_feat_part = FeaturePartitionData( + feats=torch.cat( + [r[quantized_feat_idx] for r in feat_weight_res_list] + ), + ids=partitioned_edge_ids, + ) if has_weights_for_edge_type: assert weight_idx is not None partitioned_weights = torch.cat( @@ -1410,7 +1493,12 @@ def _edge_feat_weight_pfn( weights=partitioned_weights, ) - return current_graph_part, current_feat_part, edge_partition_book + return ( + current_graph_part, + current_feat_part, + current_quantized_feat_part, + edge_partition_book, + ) def _partition_label_edge_index( self, @@ -1683,11 +1771,15 @@ def partition_edge_index_and_edge_features( self, node_partition_book: Union[PartitionBook, dict[NodeType, PartitionBook]] ) -> Union[ Tuple[ - GraphPartitionData, Optional[FeaturePartitionData], Optional[PartitionBook] + GraphPartitionData, + Optional[FeaturePartitionData], + Optional[FeaturePartitionData], + Optional[PartitionBook], ], Tuple[ dict[EdgeType, GraphPartitionData], Optional[dict[EdgeType, FeaturePartitionData]], + Optional[dict[EdgeType, FeaturePartitionData]], Optional[dict[EdgeType, PartitionBook]], ], ]: @@ -1698,8 +1790,8 @@ def partition_edge_index_and_edge_features( node_partition_book (Union[PartitionBook, dict[NodeType, PartitionBook]]): The computed Node Partition Book Returns: Union[ - Tuple[GraphPartitionData, FeaturePartitionData, PartitionBook], - Tuple[dict[EdgeType, GraphPartitionData], Optional[dict[EdgeType, FeaturePartitionData]], Optional[dict[EdgeType, PartitionBook]]], + Tuple[GraphPartitionData, Optional[FeaturePartitionData], Optional[FeaturePartitionData], Optional[PartitionBook]], + Tuple[dict[EdgeType, GraphPartitionData], Optional[dict[EdgeType, FeaturePartitionData]], Optional[dict[EdgeType, FeaturePartitionData]], Optional[dict[EdgeType, PartitionBook]]], ]: Partitioned Graph Data, Feature Data, and corresponding edge partition book, is a dictionary if heterogeneous. The second and third elements of this tuple are only present if there are edge features to partition, and are None otherwise. @@ -1748,21 +1840,27 @@ def partition_edge_index_and_edge_features( edge_partition_book: dict[EdgeType, PartitionBook] = {} partitioned_edge_index: dict[EdgeType, GraphPartitionData] = {} partitioned_edge_features: dict[EdgeType, FeaturePartitionData] = {} + partitioned_edge_quantized_features: dict[EdgeType, FeaturePartitionData] = {} for edge_type in self._edge_types: ( partitioned_edge_index_per_edge_type, partitioned_edge_features_per_edge_type, + partitioned_edge_quantized_features_per_edge_type, edge_partition_book_per_edge_type, ) = self._partition_edge_index_and_edge_features( node_partition_book=transformed_node_partition_book, edge_type=edge_type ) partitioned_edge_index[edge_type] = partitioned_edge_index_per_edge_type - if partitioned_edge_features_per_edge_type is not None: - assert edge_partition_book_per_edge_type is not None + if edge_partition_book_per_edge_type is not None: edge_partition_book[edge_type] = edge_partition_book_per_edge_type + if partitioned_edge_features_per_edge_type is not None: partitioned_edge_features[edge_type] = ( partitioned_edge_features_per_edge_type ) + if partitioned_edge_quantized_features_per_edge_type is not None: + partitioned_edge_quantized_features[edge_type] = ( + partitioned_edge_quantized_features_per_edge_type + ) elapsed_time = time.time() - start_time logger.info(f"Edge Partitioning finished, took {elapsed_time:.3f}s") @@ -1784,6 +1882,9 @@ def partition_edge_index_and_edge_features( to_homogeneous(partitioned_edge_features) if partitioned_edge_features else None, + to_homogeneous(partitioned_edge_quantized_features) + if partitioned_edge_quantized_features + else None, to_homogeneous(edge_partition_book) if edge_partition_book else None, ) else: @@ -1791,6 +1892,11 @@ def partition_edge_index_and_edge_features( return ( partitioned_edge_index, partitioned_edge_features if partitioned_edge_features else None, + ( + partitioned_edge_quantized_features + if partitioned_edge_quantized_features + else None + ), edge_partition_book if edge_partition_book else None, ) @@ -1889,6 +1995,7 @@ def partition( ( partitioned_edge_index, partitioned_edge_features, + partitioned_edge_quantized_features, edge_partition_book, ) = self.partition_edge_index_and_edge_features( node_partition_book=node_partition_book @@ -1936,6 +2043,7 @@ def partition( partitioned_node_features=partitioned_node_features, partitioned_node_quantized_features=partitioned_node_quantized_features, partitioned_edge_features=partitioned_edge_features, + partitioned_edge_quantized_features=partitioned_edge_quantized_features, partitioned_positive_labels=partitioned_positive_edge_index, partitioned_negative_labels=partitioned_negative_edge_index, partitioned_node_labels=partitioned_node_labels, diff --git a/gigl/distributed/dist_range_partitioner.py b/gigl/distributed/dist_range_partitioner.py index b7b0754f7..170d5de09 100644 --- a/gigl/distributed/dist_range_partitioner.py +++ b/gigl/distributed/dist_range_partitioner.py @@ -215,7 +215,10 @@ def _partition_edge_index_and_edge_features( node_partition_book: dict[NodeType, PartitionBook], edge_type: EdgeType, ) -> tuple[ - GraphPartitionData, Optional[FeaturePartitionData], Optional[PartitionBook] + GraphPartitionData, + Optional[FeaturePartitionData], + Optional[FeaturePartitionData], + Optional[PartitionBook], ]: """ Partition graph topology of a specific edge type. For range-based partitioning, we partition @@ -232,6 +235,7 @@ def _partition_edge_index_and_edge_features( Returns: GraphPartitionData: The graph data of the current partition. Optional[FeaturePartitionData]: The edge features on the current partition, will be None if there are no edge features for the current edge type + Optional[FeaturePartitionData]: The quantized edge features on the current partition, will be None if there are no quantized edge features for the current edge type Optional[PartitionBook]: The partition book of graph edges, will be None if there are no edge features for the current edge type """ @@ -243,6 +247,10 @@ def _partition_edge_index_and_edge_features( edge_index = self._edge_index[edge_type] has_edge_feats = self._edge_feat is not None and edge_type in self._edge_feat + has_edge_quantized_feats = ( + self._edge_quantized_feat is not None + and edge_type in self._edge_quantized_feat + ) has_edge_weights = ( self._edge_weights is not None and edge_type in self._edge_weights ) @@ -255,24 +263,35 @@ def _partition_edge_index_and_edge_features( edge_feat: Optional[torch.Tensor] = None edge_feat_dim: Optional[int] = None edge_weights_tensor: Optional[torch.Tensor] = None + edge_quantized_features: Optional[torch.Tensor] = None + edge_quantized_feature_dim: Optional[int] = None if has_edge_feats: assert self._edge_feat is not None and self._edge_feat_dim is not None assert edge_type in self._edge_feat_dim edge_feat = self._edge_feat[edge_type] edge_feat_dim = self._edge_feat_dim[edge_type] + if has_edge_quantized_feats: + assert self._edge_quantized_feat is not None + assert self._edge_quantized_feat_dim is not None + edge_quantized_features = self._edge_quantized_feat[edge_type] + edge_quantized_feature_dim = self._edge_quantized_feat_dim[edge_type] if has_edge_weights: assert self._edge_weights is not None edge_weights_tensor = self._edge_weights[edge_type] - # Build input_data tuple: (src, dst[, feat][, weights]) - # Track the index of each optional tensor so we can unpack res_list correctly. + # Build input_data as (src, dst[, feat][, packed feat][, weights]). + # Recorded indices keep result unpacking aligned with optional inputs. input_parts: list[torch.Tensor] = [edge_index[0], edge_index[1]] feat_idx: Optional[int] = None weight_idx: Optional[int] = None if edge_feat is not None: feat_idx = len(input_parts) input_parts.append(edge_feat) + quantized_feat_idx: Optional[int] = None + if edge_quantized_features is not None: + quantized_feat_idx = len(input_parts) + input_parts.append(edge_quantized_features) if edge_weights_tensor is not None: weight_idx = len(input_parts) input_parts.append(edge_weights_tensor) @@ -301,6 +320,15 @@ def edge_partition_fn(rank_indices, _): del self._edge_feat[edge_type], self._edge_feat_dim[edge_type] if self._edge_weights is not None and edge_type in self._edge_weights: del self._edge_weights[edge_type] + if ( + self._edge_quantized_feat is not None + and edge_type in self._edge_quantized_feat + ): + assert self._edge_quantized_feat_dim is not None + del ( + self._edge_quantized_feat[edge_type], + self._edge_quantized_feat_dim[edge_type], + ) # We check if edge_index or edge_feat dict is empty after deleting the tensor. If so, we set these fields to None. if not self._edge_index: @@ -310,6 +338,9 @@ def edge_partition_fn(rank_indices, _): self._edge_feat_dim = None if self._edge_weights is not None and not self._edge_weights: self._edge_weights = None + if self._edge_quantized_feat is not None and not self._edge_quantized_feat: + self._edge_quantized_feat = None + self._edge_quantized_feat_dim = None gc.collect() @@ -319,6 +350,11 @@ def edge_partition_fn(rank_indices, _): torch.empty(0, edge_feat_dim) if edge_feat_dim is not None else None ) partitioned_weights = torch.empty(0) if has_edge_weights else None + partitioned_edge_quantized_features = ( + torch.empty(0, edge_quantized_feature_dim, dtype=torch.uint8) + if edge_quantized_feature_dim is not None + else None + ) else: partitioned_edge_index = torch.stack( ( @@ -337,12 +373,17 @@ def edge_partition_fn(rank_indices, _): if weight_idx is not None else None ) + partitioned_edge_quantized_features = ( + torch.cat([r[quantized_feat_idx] for r in res_list]) + if quantized_feat_idx is not None + else None + ) res_list.clear() gc.collect() - # Generate range-based edge partition book and infer edge IDs. - # Only needed when edge features are present — weights use positional IDs. + # Generate range-based edge partition book and infer edge IDs for every + # sidecar that requires sampled edge lookup. num_edges_on_each_rank: list[tuple[int, int]] = sorted( all_gather((self._rank, partitioned_edge_index.size(1))).values(), key=lambda x: x[0], @@ -354,21 +395,26 @@ def edge_partition_fn(rank_indices, _): partition_ranges.append((start, end)) start = end - if edge_feat_dim is not None: + if ( + edge_feat_dim is not None + or edge_quantized_feature_dim is not None + or has_edge_weights + ): edge_partition_book = RangePartitionBook( partition_ranges=partition_ranges, partition_idx=self._rank ) partitioned_edge_ids = get_ids_on_rank( partition_book=edge_partition_book, rank=self._rank ) - assert partitioned_edge_features is not None current_graph_part = GraphPartitionData( edge_index=partitioned_edge_index, edge_ids=partitioned_edge_ids, weights=partitioned_weights, ) - current_feat_part = FeaturePartitionData( - feats=partitioned_edge_features, ids=None + current_feat_part = ( + FeaturePartitionData(feats=partitioned_edge_features, ids=None) + if partitioned_edge_features is not None + else None ) logger.info( f"Got edge range-based partition book for edge type {edge_type} on rank {self._rank} with partition bounds: {edge_partition_book.partition_bounds}" @@ -386,17 +432,31 @@ def edge_partition_fn(rank_indices, _): f"Edge Index and Feature Partitioning for edge type {edge_type} finished, took {time.time() - start_time:.3f}s" ) - return current_graph_part, current_feat_part, edge_partition_book + current_quantized_feat_part = ( + FeaturePartitionData(feats=partitioned_edge_quantized_features, ids=None) + if partitioned_edge_quantized_features is not None + else None + ) + return ( + current_graph_part, + current_feat_part, + current_quantized_feat_part, + edge_partition_book, + ) def partition_edge_index_and_edge_features( self, node_partition_book: Union[PartitionBook, dict[NodeType, PartitionBook]] ) -> Union[ tuple[ - GraphPartitionData, Optional[FeaturePartitionData], Optional[PartitionBook] + GraphPartitionData, + Optional[FeaturePartitionData], + Optional[FeaturePartitionData], + Optional[PartitionBook], ], tuple[ dict[EdgeType, GraphPartitionData], Optional[dict[EdgeType, FeaturePartitionData]], + Optional[dict[EdgeType, FeaturePartitionData]], Optional[dict[EdgeType, PartitionBook]], ], ]: @@ -408,10 +468,11 @@ def partition_edge_index_and_edge_features( Args: node_partition_book (Union[PartitionBook, dict[NodeType, PartitionBook]]): The computed Node Partition Book + Returns: Union[ - Tuple[GraphPartitionData, Optional[FeaturePartitionData], Optional[PartitionBook]], - Tuple[dict[EdgeType, GraphPartitionData], Optional[dict[EdgeType, FeaturePartitionData]], Optional[dict[EdgeType, PartitionBook]]], + Tuple[GraphPartitionData, Optional[FeaturePartitionData], Optional[FeaturePartitionData], Optional[PartitionBook]], + Tuple[dict[EdgeType, GraphPartitionData], Optional[dict[EdgeType, FeaturePartitionData]], Optional[dict[EdgeType, FeaturePartitionData]], Optional[dict[EdgeType, PartitionBook]]], ]: Partitioned Graph Data, Feature Data, and corresponding edge partition book, is a dictionary if heterogeneous. """ @@ -448,21 +509,27 @@ def partition_edge_index_and_edge_features( edge_partition_book: dict[EdgeType, PartitionBook] = {} partitioned_edge_index: dict[EdgeType, GraphPartitionData] = {} partitioned_edge_features: dict[EdgeType, FeaturePartitionData] = {} + partitioned_edge_quantized_features: dict[EdgeType, FeaturePartitionData] = {} for edge_type in self._edge_types: ( partitioned_edge_index_per_edge_type, partitioned_edge_features_per_edge_type, + partitioned_edge_quantized_features_per_edge_type, edge_partition_book_per_edge_type, ) = self._partition_edge_index_and_edge_features( node_partition_book=transformed_node_partition_book, edge_type=edge_type ) partitioned_edge_index[edge_type] = partitioned_edge_index_per_edge_type - if partitioned_edge_features_per_edge_type is not None: - assert edge_partition_book_per_edge_type is not None + if edge_partition_book_per_edge_type is not None: edge_partition_book[edge_type] = edge_partition_book_per_edge_type + if partitioned_edge_features_per_edge_type is not None: partitioned_edge_features[edge_type] = ( partitioned_edge_features_per_edge_type ) + if partitioned_edge_quantized_features_per_edge_type is not None: + partitioned_edge_quantized_features[edge_type] = ( + partitioned_edge_quantized_features_per_edge_type + ) elapsed_time = time.time() - start_time logger.info(f"Edge Partitioning finished, took {elapsed_time:.3f}s") @@ -481,6 +548,9 @@ def partition_edge_index_and_edge_features( to_homogeneous(partitioned_edge_features) if partitioned_edge_features else None, + to_homogeneous(partitioned_edge_quantized_features) + if partitioned_edge_quantized_features + else None, to_homogeneous(edge_partition_book) if edge_partition_book else None, ) else: @@ -488,5 +558,10 @@ def partition_edge_index_and_edge_features( return ( partitioned_edge_index, partitioned_edge_features if partitioned_edge_features else None, + ( + partitioned_edge_quantized_features + if partitioned_edge_quantized_features + else None + ), edge_partition_book if edge_partition_book else None, ) diff --git a/gigl/distributed/utils/serialized_graph_metadata_translator.py b/gigl/distributed/utils/serialized_graph_metadata_translator.py index 36ad31c52..25fb26882 100644 --- a/gigl/distributed/utils/serialized_graph_metadata_translator.py +++ b/gigl/distributed/utils/serialized_graph_metadata_translator.py @@ -33,18 +33,11 @@ def _build_serialized_tfrecord_entity_info( entity_key (Union[str, Tuple[str, str]]): Entity key to register to SerializedTFRecordInfo, is a str if Node entity or Tuple[str, str] if Edge entity tfrecord_uri_pattern (str): Regex pattern for loading serialized tf records quantization_metadata (Optional[FeatureQuantizationMetadata]): Quantization - metadata for a node entity, when its features are quantized. + metadata for a node or main-edge entity when its features are quantized. Returns: SerializedTFRecordInfo: Stored metadata for current entity """ if quantization_metadata is not None: - if not isinstance( - preprocessed_metadata, PreprocessedMetadata.NodeMetadataOutput - ): - # TODO(quantization): Support edge feature quantization. - raise NotImplementedError( - "Feature quantization is not supported for edge entities." - ) packed_feature_key = ( preprocessed_metadata.quantized_feature_metadata.packed_feature_key ) @@ -146,6 +139,7 @@ def convert_pb_to_serialized_graph_metadata( positive_label_entity_info: dict[EdgeType, SerializedTFRecordInfo] = {} negative_label_entity_info: dict[EdgeType, SerializedTFRecordInfo] = {} node_quantization_metadata: dict[NodeType, FeatureQuantizationMetadata] = {} + edge_quantization_metadata: dict[EdgeType, FeatureQuantizationMetadata] = {} preprocessed_metadata_pb = preprocessed_metadata_pb_wrapper.preprocessed_metadata_pb @@ -202,11 +196,19 @@ def convert_pb_to_serialized_graph_metadata( edge_feature_spec_dict = preprocessed_metadata_pb_wrapper.condensed_edge_type_to_feature_schema_map[ condensed_edge_type ].feature_spec + if edge_metadata.main_edge_info.HasField("quantized_feature_metadata"): + edge_quantization_metadata[edge_type] = ( + _build_feature_quantization_metadata( + quantized_metadata=edge_metadata.main_edge_info.quantized_feature_metadata, + feature_dim=edge_metadata.main_edge_info.feature_dim, + ) + ) edge_entity_info[edge_type] = _build_serialized_tfrecord_entity_info( preprocessed_metadata=edge_metadata.main_edge_info, feature_spec_dict=edge_feature_spec_dict, entity_key=edge_key, tfrecord_uri_pattern=tfrecord_uri_pattern, + quantization_metadata=edge_quantization_metadata.get(edge_type), ) if edge_metadata.HasField("positive_edge_info"): @@ -251,6 +253,9 @@ def convert_pb_to_serialized_graph_metadata( node_quantization_metadata=to_homogeneous(node_quantization_metadata) if len(node_quantization_metadata) > 0 else None, + edge_quantization_metadata=to_homogeneous(edge_quantization_metadata) + if len(edge_quantization_metadata) > 0 + else None, ) else: return SerializedGraphMetadata( @@ -265,4 +270,7 @@ def convert_pb_to_serialized_graph_metadata( node_quantization_metadata=node_quantization_metadata if len(node_quantization_metadata) > 0 else None, + edge_quantization_metadata=edge_quantization_metadata + if len(edge_quantization_metadata) > 0 + else None, ) diff --git a/gigl/types/graph.py b/gigl/types/graph.py index 849f7708a..eb501f0d7 100644 --- a/gigl/types/graph.py +++ b/gigl/types/graph.py @@ -105,6 +105,9 @@ class PartitionOutput: partitioned_node_quantized_features: Optional[ Union[FeaturePartitionData, dict[NodeType, FeaturePartitionData]] ] = None + partitioned_edge_quantized_features: Optional[ + Union[FeaturePartitionData, dict[EdgeType, FeaturePartitionData]] + ] = None @dataclass(frozen=True) @@ -236,6 +239,9 @@ class LoadedGraphTensors: node_quantized_features: Optional[ Union[torch.Tensor, dict[NodeType, torch.Tensor]] ] = None + edge_quantized_features: Optional[ + Union[torch.Tensor, dict[EdgeType, torch.Tensor]] + ] = None def treat_labels_as_edges(self, edge_dir: Literal["in", "out"]) -> None: """ @@ -337,6 +343,9 @@ def treat_labels_as_edges(self, edge_dir: Literal["in", "out"]) -> None: self.node_quantized_features = to_heterogeneous_node( self.node_quantized_features ) + self.edge_quantized_features = to_heterogeneous_edge( + self.edge_quantized_features + ) self.edge_index = edge_index_with_labels self.edge_features = to_heterogeneous_edge(self.edge_features) self.edge_weights = to_heterogeneous_edge(self.edge_weights) diff --git a/tests/test_assets/distributed/run_distributed_partitioner.py b/tests/test_assets/distributed/run_distributed_partitioner.py index 046b8bf49..89c863d07 100644 --- a/tests/test_assets/distributed/run_distributed_partitioner.py +++ b/tests/test_assets/distributed/run_distributed_partitioner.py @@ -22,6 +22,9 @@ class InputDataStrategy(Enum): REGISTER_EDGE_WEIGHTS_WITHOUT_EDGE_FEATURES = ( "REGISTER_EDGE_WEIGHTS_WITHOUT_EDGE_FEATURES" ) + REGISTER_EDGE_QUANTIZED_FEATURES_WITHOUT_EDGE_FEATURES = ( + "REGISTER_EDGE_QUANTIZED_FEATURES_WITHOUT_EDGE_FEATURES" + ) def run_distributed_partitioner( @@ -95,7 +98,29 @@ def run_distributed_partitioner( init_rpc(master_addr=master_addr, master_port=master_port, num_rpc_threads=4) dist_partitioner: DistPartitioner - if input_data_strategy in ( + if ( + input_data_strategy + == InputDataStrategy.REGISTER_EDGE_QUANTIZED_FEATURES_WITHOUT_EDGE_FEATURES + ): + dist_partitioner = partitioner_class( + should_assign_edges_by_src_node=should_assign_edges_by_src_node, + ) + dist_partitioner.register_node_ids(node_ids=node_ids) + dist_partitioner.register_edge_index(edge_index=edge_index) + edge_quantized_features: Union[torch.Tensor, dict[EdgeType, torch.Tensor]] + if isinstance(edge_index, dict): + edge_index_by_type = cast(dict[EdgeType, torch.Tensor], edge_index) + edge_quantized_features = { + edge_type: indices[0].to(torch.uint8).unsqueeze(1) + for edge_type, indices in edge_index_by_type.items() + } + else: + edge_quantized_features = edge_index[0].to(torch.uint8).unsqueeze(1) + dist_partitioner.register_edge_quantized_features( + edge_quantized_features=edge_quantized_features + ) + partition_output = dist_partitioner.partition() + elif input_data_strategy in ( InputDataStrategy.REGISTER_ALL_ENTITIES_SEPARATELY, InputDataStrategy.REGISTER_EDGE_WEIGHTS_WITHOUT_EDGE_FEATURES, ): @@ -119,6 +144,7 @@ def run_distributed_partitioner( ( output_edge_index, output_edge_features, + _, output_edge_partition_book, ) = dist_partitioner.partition_edge_index_and_edge_features( node_partition_book=output_node_partition_book @@ -181,6 +207,7 @@ def run_distributed_partitioner( ( output_graph, output_edge_features, + _, output_edge_partition_book, ) = dist_partitioner.partition_edge_index_and_edge_features( node_partition_book=output_node_partition_book diff --git a/tests/unit/common/data/dataloaders_test.py b/tests/unit/common/data/dataloaders_test.py index 3bfaff851..12b5de2e5 100644 --- a/tests/unit/common/data/dataloaders_test.py +++ b/tests/unit/common/data/dataloaders_test.py @@ -19,6 +19,7 @@ from gigl.common.data.load_torch_tensors import ( SerializedGraphMetadata, load_torch_tensors_from_tf_record, + remove_sampling_weight_from_edge_quantization_metadata, ) from gigl.src.common.types.pb_wrappers.gbml_config import GbmlConfigPbWrapper from gigl.src.data_preprocessor.lib.types import FeatureSpecDict @@ -29,6 +30,7 @@ from gigl.src.mocking.mocking_assets.mocked_datasets_for_pipeline_tests import ( CORA_NODE_CLASSIFICATION_MOCKED_DATASET_INFO, ) +from gigl.types.graph import FeatureQuantizationMetadata from tests.test_assets.test_case import TestCase _FEATURE_SPEC_WITH_ENTITY_KEY: FeatureSpecDict = { @@ -644,6 +646,91 @@ def test_load_edge_weights_from_tf_record(self): torch.tensor(sorted(edge_feature_vals), dtype=torch.float32), ) + def test_load_edge_weights_rejects_non_raw_field_before_loading(self) -> None: + missing_path = UriFactory.create_uri("/does/not/exist") + serialized_graph_metadata = SerializedGraphMetadata( + node_entity_info=SerializedTFRecordInfo( + tfrecord_uri_prefix=missing_path, + feature_spec={"node_id": tf.io.FixedLenFeature([], tf.int64)}, + feature_keys=[], + feature_dim=0, + entity_key="node_id", + ), + edge_entity_info=SerializedTFRecordInfo( + tfrecord_uri_prefix=missing_path, + feature_spec={ + "src_id": tf.io.FixedLenFeature([], tf.int64), + "dst_id": tf.io.FixedLenFeature([], tf.int64), + "edge_packed_features": tf.io.FixedLenFeature([], tf.string), + }, + feature_keys=[], + feature_dim=0, + entity_key=("src_id", "dst_id"), + packed_feature_key="edge_packed_features", + packed_feature_dim=1, + ), + ) + + with self.assertRaises(ValueError): + load_torch_tensors_from_tf_record( + tf_record_dataloader=TFRecordDataLoader(rank=0, world_size=1), + serialized_graph_metadata=serialized_graph_metadata, + should_load_tensors_in_parallel=False, + weight_edge_feat_name="quantized_weight", + ) + + def test_sampling_weight_removal_updates_edge_quantization_metadata( + self, + ) -> None: + missing_path = UriFactory.create_uri("/does/not/exist") + serialized_graph_metadata = SerializedGraphMetadata( + node_entity_info=SerializedTFRecordInfo( + tfrecord_uri_prefix=missing_path, + feature_spec={"node_id": tf.io.FixedLenFeature([], tf.int64)}, + feature_keys=[], + feature_dim=0, + entity_key="node_id", + ), + edge_entity_info=SerializedTFRecordInfo( + tfrecord_uri_prefix=missing_path, + feature_spec={ + "src_id": tf.io.FixedLenFeature([], tf.int64), + "dst_id": tf.io.FixedLenFeature([], tf.int64), + "raw_embedding": tf.io.FixedLenFeature([2], tf.float32), + "weight": tf.io.FixedLenFeature([], tf.float32), + "edge_packed_features": tf.io.FixedLenFeature([], tf.string), + }, + feature_keys=["raw_embedding", "weight"], + feature_dim=3, + entity_key=("src_id", "dst_id"), + packed_feature_key="edge_packed_features", + packed_feature_dim=1, + ), + edge_quantization_metadata=FeatureQuantizationMetadata( + bits=2, + feature_dim=4, + quantized_feature_indices=(3,), + clip_min=0.0, + clip_max=3.0, + ), + ) + + adjusted_metadata = remove_sampling_weight_from_edge_quantization_metadata( + serialized_graph_metadata=serialized_graph_metadata, + weight_edge_feat_name="weight", + ) + + self.assertEqual( + adjusted_metadata, + FeatureQuantizationMetadata( + bits=2, + feature_dim=3, + quantized_feature_indices=(2,), + clip_min=0.0, + clip_max=3.0, + ), + ) + def test_load_edge_weights_multidim_feature(self): """Weight column offset is correct when a preceding feature key is multi-dimensional. diff --git a/tests/unit/distributed/distributed_partitioner_test.py b/tests/unit/distributed/distributed_partitioner_test.py index 0f817bafa..f3aaea091 100644 --- a/tests/unit/distributed/distributed_partitioner_test.py +++ b/tests/unit/distributed/distributed_partitioner_test.py @@ -686,6 +686,22 @@ def _assert_label_outputs( partitioner_class=DistRangePartitioner, expected_pb_dtype=torch.int64, ), + param( + "Homogeneous packed-edge-only tensor partitioning", + is_heterogeneous=False, + input_data_strategy=InputDataStrategy.REGISTER_EDGE_QUANTIZED_FEATURES_WITHOUT_EDGE_FEATURES, + should_assign_edges_by_src_node=True, + partitioner_class=DistPartitioner, + expected_pb_dtype=torch.uint8, + ), + param( + "Homogeneous packed-edge-only range partitioning", + is_heterogeneous=False, + input_data_strategy=InputDataStrategy.REGISTER_EDGE_QUANTIZED_FEATURES_WITHOUT_EDGE_FEATURES, + should_assign_edges_by_src_node=True, + partitioner_class=DistRangePartitioner, + expected_pb_dtype=torch.int64, + ), ] ) def test_partitioning_correctness( @@ -756,6 +772,11 @@ def test_partitioning_correctness( else: expected_edge_feat_types = [USER_TO_USER_EDGE_TYPE] + is_packed_edge_only = ( + input_data_strategy + == InputDataStrategy.REGISTER_EDGE_QUANTIZED_FEATURES_WITHOUT_EDGE_FEATURES + ) + for rank, partition_output in output_dict.items(): partitioned_edge_index = partition_output.partitioned_edge_index assert partitioned_edge_index is not None @@ -780,7 +801,32 @@ def test_partitioning_correctness( graph.edge_index ) - if ( + if is_packed_edge_only: + self.assertIsNotNone(partition_output.edge_partition_book) + self.assertIsNone(partition_output.partitioned_edge_features) + self.assertIsNotNone( + partition_output.partitioned_edge_quantized_features + ) + packed_features = partition_output.partitioned_edge_quantized_features + assert isinstance(packed_features, FeaturePartitionData) + assert isinstance(partitioned_edge_index, GraphPartitionData) + self.assertEqual(packed_features.feats.dtype, torch.uint8) + self.assertEqual( + packed_features.feats.size(0), + partitioned_edge_index.edge_index.size(1), + ) + assert partitioned_edge_index.edge_ids is not None + if packed_features.ids is not None: + self.assert_tensor_equality( + tensor_a=packed_features.ids, + tensor_b=partitioned_edge_index.edge_ids, + ) + for index, edge_id in enumerate(partitioned_edge_index.edge_ids): + self.assert_tensor_equality( + tensor_a=packed_features.feats[index], + tensor_b=edge_id.to(torch.uint8).unsqueeze(0), + ) + elif ( input_data_strategy == InputDataStrategy.REGISTER_MINIMAL_ENTITIES_SEPARATELY ): diff --git a/tests/unit/distributed/distributed_weighted_sampling_test.py b/tests/unit/distributed/distributed_weighted_sampling_test.py index 9b30acdd8..cb9386b6e 100644 --- a/tests/unit/distributed/distributed_weighted_sampling_test.py +++ b/tests/unit/distributed/distributed_weighted_sampling_test.py @@ -553,6 +553,11 @@ def test_weights_only_no_features_partitioned_correctly(self) -> None: ) assert edge_ids is not None + self.assertIsNotNone( + partition_output.edge_partition_book, + msg=f"Rank {rank}: edge partition book must be retained for weights", + ) + self.assertEqual(weights.shape, edge_ids.shape) expected_weights = edge_ids.float() * 0.1 torch.testing.assert_close( @@ -732,7 +737,7 @@ def test_range_partitioner_homogeneous_weights_partitioned_correctly(self) -> No True, # should_assign_edges_by_src_node self._master_ip_address, master_port, - InputDataStrategy.REGISTER_ALL_ENTITIES_SEPARATELY, + InputDataStrategy.REGISTER_EDGE_WEIGHTS_WITHOUT_EDGE_FEATURES, DistRangePartitioner, rank_to_edge_weights, ), @@ -759,6 +764,11 @@ def test_range_partitioner_homogeneous_weights_partitioned_correctly(self) -> No ) assert edge_ids is not None + self.assertIsNotNone( + partition_output.edge_partition_book, + msg=f"Rank {rank}: edge partition book must be retained for weights", + ) + self.assertEqual( weights.shape, edge_ids.shape,