diff --git a/gigl/common/data/load_torch_tensors.py b/gigl/common/data/load_torch_tensors.py index cdc017518..fff631f50 100644 --- a/gigl/common/data/load_torch_tensors.py +++ b/gigl/common/data/load_torch_tensors.py @@ -124,6 +124,47 @@ class SerializedGraphMetadata: ] = 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]]], @@ -482,42 +523,10 @@ def load_torch_tensors_from_tf_record( loaded_graph_tensors (LoadedGraphTensors): Unpartitioned Graph Tensors """ - edge_entity_info = serialized_graph_metadata.edge_entity_info - if weight_edge_feat_name is not None: - 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" - ) - if weight_edge_feat_name not in edge_entity_info.feature_keys: - raise ValueError( - f"Sampling-weight field '{weight_edge_feat_name}' for edge type " - f"{DEFAULT_HOMOGENEOUS_EDGE_TYPE} must be an unquantized raw edge feature." - ) - elif 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())) - if weight_edge_feat_name not in serialized_info.feature_keys: - raise ValueError( - f"Sampling-weight field '{weight_edge_feat_name}' for edge type " - f"{edge_type} must be an unquantized raw edge feature." - ) - 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}" - ) - for edge_type, feature_name in weight_edge_feat_name.items(): - if feature_name not in edge_entity_info[edge_type].feature_keys: - raise ValueError( - f"Sampling-weight field '{feature_name}' for edge type " - f"{edge_type} must be an unquantized raw edge feature." - ) + _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() diff --git a/gigl/distributed/base_dist_loader.py b/gigl/distributed/base_dist_loader.py index 2577cc661..66154b9ef 100644 --- a/gigl/distributed/base_dist_loader.py +++ b/gigl/distributed/base_dist_loader.py @@ -246,6 +246,7 @@ def __init__( self._node_feature_info = dataset_schema.node_feature_info self._edge_feature_info = dataset_schema.edge_feature_info self._node_quantization_metadata = dataset_schema.node_quantization_metadata + self._edge_quantization_metadata = dataset_schema.edge_quantization_metadata self._sampler_options = sampler_options # Sampled-edge PPR output requires a final HeteroData batch so virtual @@ -457,7 +458,10 @@ def create_sampling_config( batch_size=batch_size, shuffle=shuffle, drop_last=drop_last, - with_edge=dataset_schema.edge_feature_info is not None, + with_edge=( + dataset_schema.edge_feature_info is not None + or dataset_schema.edge_quantization_metadata is not None + ), collect_features=True, with_neg=False, with_weight=with_weight, diff --git a/gigl/distributed/base_sampler.py b/gigl/distributed/base_sampler.py index 67ab6d183..5601099e8 100644 --- a/gigl/distributed/base_sampler.py +++ b/gigl/distributed/base_sampler.py @@ -19,6 +19,7 @@ from gigl.common.logger import Logger from gigl.distributed.sampler import ( + EDGE_PACKED_FEATURES_METADATA_KEY, NEGATIVE_LABEL_METADATA_KEY, NODE_PACKED_FEATURES_METADATA_KEY, POSITIVE_LABEL_METADATA_KEY, @@ -116,6 +117,7 @@ def __init__(self, *args, **kwargs) -> None: self._sampling_error_sent: bool = False self.dist_node_quantized_feature: Optional[DistFeature] = None + self.dist_edge_quantized_feature: Optional[DistFeature] = None if ( self.collect_features and data is not None @@ -132,6 +134,20 @@ def __init__(self, *args, **kwargs) -> None: rpc_router=self.rpc_router, device=self.device, ) + if ( + self.collect_features + and data is not None + and getattr(data, "edge_quantized_features", None) is not None + ): + self.dist_edge_quantized_feature = DistFeature( + data.num_partitions, + data.partition_idx, + data.edge_quantized_features, + data.edge_pb, + local_only=False, + rpc_router=self.rpc_router, + device=self.device, + ) def _prepare_sample_loop_inputs( self, @@ -436,6 +452,8 @@ async def _collate_fn( ) if self.dist_edge_feature is not None and self.with_edge: for etype in self.edge_types: + if etype not in self.dist_edge_feature.local_feature: + continue if self.edge_dir == "in": eids = result_map.get( f"{as_str(reverse_edge_type(etype))}.eids", None @@ -451,6 +469,30 @@ async def _collate_fn( futs[result_key] = wrap_torch_future( self.dist_edge_feature.async_get(eids, etype) ) + if self.dist_edge_quantized_feature is not None and self.with_edge: + for etype in self.edge_types: + # Like node features, an edge partition book covers every + # edge type while a feature store may register only some. + if etype not in self.dist_edge_quantized_feature.local_feature: + continue + result_edge_type = ( + reverse_edge_type(etype) if self.edge_dir == "in" else etype + ) + eids = result_map.get(f"{as_str(result_edge_type)}.eids") + if eids is not None: + eids = eids.to(torch.long) + output_edge_type = ( + reverse_edge_type(etype) + if self.edge_dir == "out" + else etype + ) + # GLT maps wire edge types to output stores during collation. + # Metadata bypasses that mapping, so key it by the output store. + futs[ + f"#META.{EDGE_PACKED_FEATURES_METADATA_KEY}.{output_edge_type}" + ] = wrap_torch_future( + self.dist_edge_quantized_feature.async_get(eids, etype) + ) if output.batch is not None: for ntype, batch in output.batch.items(): result_map[f"{as_str(ntype)}.batch"] = batch @@ -490,6 +532,10 @@ async def _collate_fn( futs["efeats"] = wrap_torch_future( self.dist_edge_feature.async_get(eids) ) + if self.dist_edge_quantized_feature is not None: + futs[f"#META.{EDGE_PACKED_FEATURES_METADATA_KEY}"] = wrap_torch_future( + self.dist_edge_quantized_feature.async_get(result_map["eids"]) + ) if output.batch is not None: result_map["batch"] = output.batch diff --git a/gigl/distributed/dataset_factory.py b/gigl/distributed/dataset_factory.py index 14e874be3..0ffa3c462 100644 --- a/gigl/distributed/dataset_factory.py +++ b/gigl/distributed/dataset_factory.py @@ -25,6 +25,7 @@ SerializedGraphMetadata, TFDatasetOptions, load_torch_tensors_from_tf_record, + remove_sampling_weight_from_edge_quantization_metadata, ) from gigl.common.logger import Logger from gigl.common.utils.decorator import tf_on_cpu @@ -232,6 +233,10 @@ def _load_and_build_partitioned_dataset( world_size=world_size, edge_dir=edge_dir, node_quantization_metadata=serialized_graph_metadata.node_quantization_metadata, + edge_quantization_metadata=remove_sampling_weight_from_edge_quantization_metadata( + serialized_graph_metadata=serialized_graph_metadata, + weight_edge_feat_name=weight_edge_feat_name, + ), ) dataset.build( diff --git a/gigl/distributed/dist_ablp_neighborloader.py b/gigl/distributed/dist_ablp_neighborloader.py index 10638c330..269334050 100644 --- a/gigl/distributed/dist_ablp_neighborloader.py +++ b/gigl/distributed/dist_ablp_neighborloader.py @@ -38,6 +38,7 @@ extract_edge_type_metadata, extract_metadata, labeled_to_homogeneous, + materialize_quantized_edge_features, materialize_quantized_node_features, set_missing_features, shard_nodes_by_process, @@ -609,6 +610,7 @@ def _setup_for_colocated( node_feature_info=dataset.node_feature_info, edge_feature_info=dataset.edge_feature_info, node_quantization_metadata=dataset.node_quantization_metadata, + edge_quantization_metadata=dataset.edge_quantization_metadata, edge_dir=dataset.edge_dir, ), ) @@ -799,6 +801,7 @@ def _setup_for_graph_store( node_feature_info=node_feature_info, edge_feature_info=edge_feature_info, node_quantization_metadata=dataset.fetch_node_quantization_metadata(), + edge_quantization_metadata=dataset.fetch_edge_quantization_metadata(), edge_dir=edge_dir, ), backend_key, @@ -972,6 +975,12 @@ def _collate_fn(self, msg: SampleMessage) -> Union[Data, HeteroData]: metadata=metadata, node_quantization_metadata=self._node_quantization_metadata, ) + data, metadata = materialize_quantized_edge_features( + data=data, + metadata=metadata, + edge_quantization_metadata=self._edge_quantization_metadata, + edge_dir=self.edge_dir, + ) # Attach any remaining metadata (e.g. custom user-defined keys) directly onto the # data object so downstream code can access them via attribute lookup. diff --git a/gigl/distributed/dist_dataset.py b/gigl/distributed/dist_dataset.py index a41c0e12a..6d930fd3f 100644 --- a/gigl/distributed/dist_dataset.py +++ b/gigl/distributed/dist_dataset.py @@ -58,6 +58,9 @@ def __init__( node_quantized_feature_partition: Optional[ Union[Feature, dict[NodeType, Feature]] ] = None, + edge_quantized_feature_partition: Optional[ + Union[Feature, dict[EdgeType, Feature]] + ] = None, edge_feature_partition: Optional[ Union[Feature, dict[EdgeType, Feature]] ] = None, @@ -87,6 +90,11 @@ def __init__( dict[NodeType, FeatureQuantizationMetadata], ] ] = None, + edge_quantization_metadata: Optional[ + Union[ + FeatureQuantizationMetadata, dict[EdgeType, FeatureQuantizationMetadata] + ] + ] = None, edge_feature_info: Optional[ Union[FeatureInfo, dict[EdgeType, FeatureInfo]] ] = None, @@ -166,6 +174,8 @@ def __init__( self._node_quantized_features = node_quantized_feature_partition self._node_quantization_metadata = node_quantization_metadata + self._edge_quantized_features = edge_quantized_feature_partition + self._edge_quantization_metadata = edge_quantization_metadata self._degree_tensor: Optional[ Union[torch.Tensor, dict[NodeType, torch.Tensor]] @@ -253,6 +263,13 @@ def edge_features( ): self._edge_features = new_edge_features + @property + def edge_quantized_features( + self, + ) -> Optional[Union[Feature, dict[EdgeType, Feature]]]: + """Packed uint8 main-edge feature sidecar.""" + return self._edge_quantized_features + @property def node_pb( self, @@ -340,6 +357,15 @@ def node_quantization_metadata( ]: return self._node_quantization_metadata + @property + def edge_quantization_metadata( + self, + ) -> Optional[ + Union[FeatureQuantizationMetadata, dict[EdgeType, FeatureQuantizationMetadata]] + ]: + """Metadata required to materialize packed main-edge features.""" + return self._edge_quantization_metadata + @property def edge_feature_info( self, @@ -899,6 +925,42 @@ def _initialize_edge_features( ) logger.info(f"Initialized edge features for homogeneous graph to dataset") + def _initialize_edge_quantized_features( + self, + edge_partition_book: Union[PartitionBook, dict[EdgeType, PartitionBook]], + partitioned_edge_quantized_features: Optional[ + Union[FeaturePartitionData, dict[EdgeType, FeaturePartitionData]] + ], + ) -> None: + """Initialize packed uint8 main-edge feature storage.""" + features, id_to_index = _prepare_feature_data( + partition_book=edge_partition_book, + partitioned_data=partitioned_edge_quantized_features, + ) + if features is None or id_to_index is None: + logger.info("Found no packed quantized edge features to initialize") + return + if isinstance(features, dict): + assert isinstance(id_to_index, dict) + edge_quantized_features: dict[EdgeType, Feature] = {} + for edge_type, features_per_edge_type in features.items(): + assert isinstance(edge_type, EdgeType) + edge_quantized_features[edge_type] = Feature( + feature_tensor=features_per_edge_type, + id2index=id_to_index[edge_type], + with_gpu=False, + dtype=torch.uint8, + ) + self._edge_quantized_features = edge_quantized_features + else: + assert not isinstance(id_to_index, Mapping) + self._edge_quantized_features = Feature( + feature_tensor=features, + id2index=id_to_index, + with_gpu=False, + dtype=torch.uint8, + ) + def build( self, partition_output: PartitionOutput, @@ -1011,6 +1073,13 @@ def build( partition_output.partitioned_edge_features = None gc.collect() + self._initialize_edge_quantized_features( + edge_partition_book=partition_output.edge_partition_book, + partitioned_edge_quantized_features=partition_output.partitioned_edge_quantized_features, + ) + partition_output.partitioned_edge_quantized_features = None + gc.collect() + self._node_partition_book = partition_output.node_partition_book self._edge_partition_book = partition_output.edge_partition_book @@ -1037,6 +1106,7 @@ def share_ipc( Optional[Union[Feature, dict[NodeType, Feature]]], Optional[Union[Feature, dict[NodeType, Feature]]], Optional[Union[Feature, dict[EdgeType, Feature]]], + Optional[Union[Feature, dict[EdgeType, Feature]]], Optional[Union[Feature, dict[NodeType, Feature]]], Optional[Union[PartitionBook, dict[NodeType, PartitionBook]]], Optional[Union[PartitionBook, dict[EdgeType, PartitionBook]]], @@ -1053,6 +1123,12 @@ def share_ipc( dict[NodeType, FeatureQuantizationMetadata], ] ], + Optional[ + Union[ + FeatureQuantizationMetadata, + dict[EdgeType, FeatureQuantizationMetadata], + ] + ], Optional[Union[FeatureInfo, dict[EdgeType, FeatureInfo]]], Optional[Union[torch.Tensor, dict[NodeType, torch.Tensor]]], Optional[int], @@ -1067,6 +1143,7 @@ def share_ipc( Optional[Union[Graph, dict[EdgeType, Graph]]]: Partitioned Graph Data Optional[Union[Feature, dict[NodeType, Feature]]]: Partitioned Node Feature Data Optional[Union[Feature, dict[NodeType, Feature]]]: Partitioned packed uint8 node feature data + Optional[Union[Feature, dict[EdgeType, Feature]]]: Partitioned packed uint8 edge feature data Optional[Union[Feature, dict[EdgeType, Feature]]]: Partitioned Edge Feature Data Optional[Union[Feature, dict[NodeType, Feature]]]: Node labels on the current machine. Will be a dict if heterogeneous. Optional[Union[torch.Tensor, dict[NodeType, torch.Tensor]]]: Node Partition Book Tensor @@ -1079,6 +1156,7 @@ def share_ipc( Optional[Union[int, dict[NodeType, int]]]: Number of test nodes on the current machine. Will be a dict if heterogeneous. Optional[Union[FeatureInfo, dict[NodeType, FeatureInfo]]]: Node feature dim and its data type, will be a dict if heterogeneous Optional[Union[FeatureQuantizationMetadata, dict[NodeType, FeatureQuantizationMetadata]]]: Node quantization metadata. + Optional[Union[FeatureQuantizationMetadata, dict[EdgeType, FeatureQuantizationMetadata]]]: Edge quantization metadata. Optional[Union[FeatureInfo, dict[EdgeType, FeatureInfo]]]: Edge feature dim and its data type, will be a dict if heterogeneous Optional[Union[torch.Tensor, dict[NodeType, torch.Tensor]]]: Degree tensors Optional[int]: Optional per-anchor label cap for ABLP label fetching @@ -1100,6 +1178,7 @@ def share_ipc( self._graph, self._node_features, self._node_quantized_features, + self._edge_quantized_features, self._edge_features, self._node_labels, self._node_partition_book, @@ -1112,6 +1191,7 @@ def share_ipc( self._num_test, # Additional field unique to DistDataset class self._node_feature_info, # Additional field unique to DistDataset class self._node_quantization_metadata, # Additional field unique to DistDataset class + self._edge_quantization_metadata, # Additional field unique to DistDataset class self._edge_feature_info, # Additional field unique to DistDataset class self._degree_tensor, # Additional field unique to DistDataset class self._max_labels_per_anchor_node, # Additional field unique to DistDataset class @@ -1348,6 +1428,9 @@ def _rebuild_distributed_dataset( Optional[ Union[Feature, dict[NodeType, Feature]] ], # Partitioned packed uint8 node feature data + Optional[ + Union[Feature, dict[EdgeType, Feature]] + ], # Partitioned packed uint8 edge feature data Optional[ Union[Feature, dict[EdgeType, Feature]] ], # Partitioned Edge Feature Data @@ -1377,6 +1460,12 @@ def _rebuild_distributed_dataset( dict[NodeType, FeatureQuantizationMetadata], ] ], # Node quantization metadata + Optional[ + Union[ + FeatureQuantizationMetadata, + dict[EdgeType, FeatureQuantizationMetadata], + ] + ], # Edge quantization metadata Optional[ Union[FeatureInfo, dict[EdgeType, FeatureInfo]] ], # Edge feature dim and its data type diff --git a/gigl/distributed/distributed_neighborloader.py b/gigl/distributed/distributed_neighborloader.py index 3effa01c3..270bbfc45 100644 --- a/gigl/distributed/distributed_neighborloader.py +++ b/gigl/distributed/distributed_neighborloader.py @@ -29,6 +29,7 @@ SamplingClusterSetup, extract_metadata, labeled_to_homogeneous, + materialize_quantized_edge_features, materialize_quantized_node_features, set_missing_features, shard_nodes_by_process, @@ -413,6 +414,7 @@ def _setup_for_graph_store( node_feature_info=node_feature_info, edge_feature_info=edge_feature_info, node_quantization_metadata=dataset.fetch_node_quantization_metadata(), + edge_quantization_metadata=dataset.fetch_edge_quantization_metadata(), edge_dir=dataset.fetch_edge_dir(), ), backend_key, @@ -531,6 +533,7 @@ def _setup_for_colocated( node_feature_info=dataset.node_feature_info, edge_feature_info=dataset.edge_feature_info, node_quantization_metadata=dataset.node_quantization_metadata, + edge_quantization_metadata=dataset.edge_quantization_metadata, edge_dir=dataset.edge_dir, ), ) @@ -564,6 +567,12 @@ def _collate_fn(self, msg: SampleMessage) -> Union[Data, HeteroData]: metadata=metadata, node_quantization_metadata=self._node_quantization_metadata, ) + data, metadata = materialize_quantized_edge_features( + data=data, + metadata=metadata, + edge_quantization_metadata=self._edge_quantization_metadata, + edge_dir=self.edge_dir, + ) # Attach any remaining metadata (e.g. custom user-defined keys) directly onto the # data object so downstream code can access them via attribute lookup. diff --git a/gigl/distributed/graph_store/dist_server.py b/gigl/distributed/graph_store/dist_server.py index 0a92d959d..b608be433 100644 --- a/gigl/distributed/graph_store/dist_server.py +++ b/gigl/distributed/graph_store/dist_server.py @@ -418,6 +418,14 @@ def get_node_quantization_metadata( """Get node feature quantization metadata from the dataset.""" return self.dataset.node_quantization_metadata + def get_edge_quantization_metadata( + self, + ) -> Union[ + FeatureQuantizationMetadata, dict[EdgeType, FeatureQuantizationMetadata], None + ]: + """Get main-edge feature quantization metadata from the dataset.""" + return self.dataset.edge_quantization_metadata + def get_edge_feature_info( self, ) -> Union[FeatureInfo, dict[EdgeType, FeatureInfo], None]: diff --git a/gigl/distributed/graph_store/remote_dist_dataset.py b/gigl/distributed/graph_store/remote_dist_dataset.py index 81609961b..127078fca 100644 --- a/gigl/distributed/graph_store/remote_dist_dataset.py +++ b/gigl/distributed/graph_store/remote_dist_dataset.py @@ -80,6 +80,14 @@ def fetch_node_quantization_metadata( DistServer.get_node_quantization_metadata, ) + def fetch_edge_quantization_metadata( + self, + ) -> Union[ + FeatureQuantizationMetadata, dict[EdgeType, FeatureQuantizationMetadata], None + ]: + """Fetch main-edge feature quantization metadata from storage.""" + return request_server(0, DistServer.get_edge_quantization_metadata) + def fetch_edge_feature_info( self, ) -> Union[FeatureInfo, dict[EdgeType, FeatureInfo], None]: diff --git a/gigl/distributed/sampler.py b/gigl/distributed/sampler.py index 7789c6731..1e01ee85f 100644 --- a/gigl/distributed/sampler.py +++ b/gigl/distributed/sampler.py @@ -9,6 +9,7 @@ POSITIVE_LABEL_METADATA_KEY: Final[str] = "gigl_positive_labels." NEGATIVE_LABEL_METADATA_KEY: Final[str] = "gigl_negative_labels." NODE_PACKED_FEATURES_METADATA_KEY: Final[str] = "node_packed_features" +EDGE_PACKED_FEATURES_METADATA_KEY: Final[str] = "edge_packed_features" class ABLPNodeSamplerInput(NodeSamplerInput): diff --git a/gigl/distributed/utils/neighborloader.py b/gigl/distributed/utils/neighborloader.py index 6b81306ee..30d03ff04 100644 --- a/gigl/distributed/utils/neighborloader.py +++ b/gigl/distributed/utils/neighborloader.py @@ -9,14 +9,19 @@ import torch from graphlearn_torch.channel import SampleMessage +from graphlearn_torch.utils import reverse_edge_type from torch_geometric.data import Data, HeteroData -from torch_geometric.data.storage import NodeStorage +from torch_geometric.data.storage import EdgeStorage, NodeStorage from torch_geometric.typing import EdgeType, NodeType from gigl.common.logger import Logger from gigl.common.utils.feature_quantization.torch_ops import dequantize_torch_tensor -from gigl.distributed.sampler import NODE_PACKED_FEATURES_METADATA_KEY +from gigl.distributed.sampler import ( + EDGE_PACKED_FEATURES_METADATA_KEY, + NODE_PACKED_FEATURES_METADATA_KEY, +) from gigl.types.graph import ( + DEFAULT_HOMOGENEOUS_EDGE_TYPE, DEFAULT_HOMOGENEOUS_NODE_TYPE, FeatureInfo, FeatureQuantizationIndexTensors, @@ -59,6 +64,9 @@ class DatasetSchema: node_quantization_metadata: Optional[ Union[FeatureQuantizationMetadata, dict[NodeType, FeatureQuantizationMetadata]] ] = None + edge_quantization_metadata: Optional[ + Union[FeatureQuantizationMetadata, dict[EdgeType, FeatureQuantizationMetadata]] + ] = None def patch_fanout_for_sampling( @@ -337,6 +345,49 @@ def set_missing_features( return data +def _materialize_quantized_features( + store: Union[Data, NodeStorage, EdgeStorage], + packed_features: torch.Tensor, + quantization_metadata: FeatureQuantizationMetadata, + feature_attribute: Literal["x", "edge_attr"], +) -> None: + """Reconstruct and assign quantized features for one PyG feature store. + + Args: + store: PyG data or typed storage receiving the reconstructed features. + packed_features: Quantized feature columns for sampled graph entities. + quantization_metadata: Column layout and dequantization metadata. + feature_attribute: PyG attribute that stores the feature tensor. + + Raises: + ValueError: If expected raw feature columns are absent or have an + unexpected dimension. + """ + dequantized = dequantize_torch_tensor( + packed_features, metadata=quantization_metadata + ) + raw_features = getattr(store, feature_attribute, None) + materialized_features = dequantized.new_empty( + (dequantized.size(0), quantization_metadata.feature_dim) + ) + scatter_idx: FeatureQuantizationIndexTensors = ( + quantization_metadata.scatter_index_tensors(materialized_features.device) + ) + materialized_features[:, scatter_idx.quantized] = dequantized + + if raw_features is None and quantization_metadata.raw_feature_dim: + raise ValueError( + f"Missing {quantization_metadata.raw_feature_dim} unquantized features" + ) + if raw_features is not None: + if raw_features.size(1) != quantization_metadata.raw_feature_dim: + raise ValueError( + f"Expected {quantization_metadata.raw_feature_dim} raw features before dequantization, got {raw_features.size(1)}" + ) + materialized_features[:, scatter_idx.raw] = raw_features + setattr(store, feature_attribute, materialized_features) + + def materialize_quantized_node_features( data: _GraphType, metadata: dict[str, torch.Tensor], @@ -371,48 +422,6 @@ def materialize_quantized_node_features( if node_quantization_metadata is None: return data, metadata - def materialize( - store: Union[Data, NodeStorage], - packed_features: torch.Tensor, - quantization_metadata: FeatureQuantizationMetadata, - ) -> None: - """Reconstruct and assign node features for one PyG node store. - - Args: - store: Node store receiving the reconstructed ``x`` tensor. - packed_features: Quantized feature columns for the sampled nodes. - quantization_metadata: Column layout and dequantization metadata. - - Raises: - ValueError: If expected raw feature columns are absent or have an - unexpected dimension. - """ - dequantized = dequantize_torch_tensor( - packed_features, metadata=quantization_metadata - ) - x = getattr(store, "x", None) - out = dequantized.new_empty( - (dequantized.size(0), quantization_metadata.feature_dim) - ) - scatter_idx: FeatureQuantizationIndexTensors = ( - quantization_metadata.scatter_index_tensors(out.device) - ) - out[:, scatter_idx.quantized] = dequantized - - if x is None and quantization_metadata.raw_feature_dim: - raise ValueError( - f"Missing {quantization_metadata.raw_feature_dim} unquantized features" - ) - if x is not None: - if x.size(1) != quantization_metadata.raw_feature_dim: - raise ValueError( - "Expected " - f"{quantization_metadata.raw_feature_dim} raw node features " - f"before dequantization, got {x.size(1)}" - ) - out[:, scatter_idx.raw] = x - store.x = out - if isinstance(data, Data): if isinstance(node_quantization_metadata, dict): raise ValueError("Expect scalar quantization metadata for homogeneous data") @@ -430,7 +439,12 @@ def materialize( raise ValueError( f"Missing packed quantized features in metadata keys {NODE_PACKED_FEATURES_METADATA_KEY} or {labeled_homogeneous_packed_features_key}" ) - materialize(data, packed_features, node_quantization_metadata) + _materialize_quantized_features( + data, + packed_features, + node_quantization_metadata, + feature_attribute="x", + ) else: if not isinstance(node_quantization_metadata, dict): raise ValueError("Expected per-node-type metadata for heterogeneous data.") @@ -442,7 +456,112 @@ def materialize( packed_features = metadata.pop(metadata_key, None) if packed_features is None: continue - materialize(data[node_type], packed_features, quantization_metadata) + _materialize_quantized_features( + data[node_type], + packed_features, + quantization_metadata, + feature_attribute="x", + ) + + return data, metadata + + +def materialize_quantized_edge_features( + data: _GraphType, + metadata: dict[str, torch.Tensor], + edge_quantization_metadata: Optional[ + Union[FeatureQuantizationMetadata, dict[EdgeType, FeatureQuantizationMetadata]] + ], + edge_dir: Literal["in", "out"] = "in", +) -> tuple[_GraphType, dict[str, torch.Tensor]]: + """Materialize packed quantized edge features into PyG edge feature tensors. + + Reconstructs each edge feature tensor in its original column order by + dequantizing packed features and combining them with any unquantized + feature columns already present in ``data``. Consumed packed-feature + entries are removed from ``metadata``. + + Args: + data: Homogeneous or heterogeneous sampled graph containing raw edge + feature columns. + metadata: Sample metadata containing packed edge feature tensors. + edge_quantization_metadata: Quantization metadata for the graph's edge + features. Homogeneous graphs require a single value; heterogeneous + graphs require metadata for each edge type. + edge_dir: Sampling direction. GLT reverses heterogeneous output edge + stores when sampling outward. + + Returns: + A tuple containing the graph with reconstructed edge features and the + remaining sample metadata. + + Raises: + ValueError: If the graph and quantization metadata shapes do not match, + required packed features are missing, or raw feature dimensions are + inconsistent. + """ + if edge_quantization_metadata is None: + return data, metadata + + if isinstance(data, Data): + if isinstance(edge_quantization_metadata, dict): + raise ValueError("Expect scalar quantization metadata for homogeneous data") + packed_features = metadata.pop(EDGE_PACKED_FEATURES_METADATA_KEY, None) + labeled_homogeneous_output_edge_type = ( + reverse_edge_type(DEFAULT_HOMOGENEOUS_EDGE_TYPE) + if edge_dir == "out" + else DEFAULT_HOMOGENEOUS_EDGE_TYPE + ) + labeled_homogeneous_packed_features_key = ( + f"{EDGE_PACKED_FEATURES_METADATA_KEY}." + f"{labeled_homogeneous_output_edge_type}" + ) + if packed_features is None: + # Labeled homogeneous graphs are sampled as heterogeneous graphs, so + # the transport key uses GLT's direction-dependent output edge type. + packed_features = metadata.pop( + labeled_homogeneous_packed_features_key, None + ) + if packed_features is None: + raise ValueError( + "Missing packed quantized features in metadata keys " + f"{EDGE_PACKED_FEATURES_METADATA_KEY} or " + f"{labeled_homogeneous_packed_features_key}" + ) + _materialize_quantized_features( + data, + packed_features, + edge_quantization_metadata, + feature_attribute="edge_attr", + ) + else: + if not isinstance(edge_quantization_metadata, dict): + raise ValueError("Expected per-edge-type metadata for heterogeneous data.") + edge_quantization_metadata = cast( + dict[EdgeType, FeatureQuantizationMetadata], edge_quantization_metadata + ) + for edge_type, quantization_metadata in edge_quantization_metadata.items(): + output_edge_type = ( + reverse_edge_type(edge_type) if edge_dir == "out" else edge_type + ) + metadata_key = f"{EDGE_PACKED_FEATURES_METADATA_KEY}.{output_edge_type}" + packed_features = metadata.pop(metadata_key, None) + if packed_features is None: + if ( + output_edge_type not in data.edge_types + or data[output_edge_type].num_edges == 0 + ): + continue + raise ValueError( + "Missing packed quantized edge features for sampled edge type " + f"{output_edge_type}" + ) + _materialize_quantized_features( + data[output_edge_type], + packed_features, + quantization_metadata, + feature_attribute="edge_attr", + ) return data, metadata diff --git a/tests/integration/pipeline/data_preprocessor/feature_quantization_transform_test.py b/tests/integration/pipeline/data_preprocessor/feature_quantization_transform_test.py index 540140be9..d2ef8c309 100644 --- a/tests/integration/pipeline/data_preprocessor/feature_quantization_transform_test.py +++ b/tests/integration/pipeline/data_preprocessor/feature_quantization_transform_test.py @@ -5,21 +5,170 @@ import apache_beam as beam import pyarrow as pa import tensorflow as tf +import tensorflow_data_validation as tfdv +import torch from apache_beam.testing.test_pipeline import TestPipeline from apache_beam.testing.util import assert_that, equal_to from parameterized import parameterized from tensorflow_metadata.proto.v0 import schema_pb2 from tensorflow_transform.tf_metadata.dataset_metadata import DatasetMetadata +from torch_geometric.data import Data +from gigl.common.beam.better_tfrecordio import BetterWriteToTFRecord +from gigl.common.data.dataloaders import TFDatasetOptions, TFRecordDataLoader +from gigl.distributed.utils.neighborloader import ( + EDGE_PACKED_FEATURES_METADATA_KEY, + materialize_quantized_edge_features, +) +from gigl.distributed.utils.serialized_graph_metadata_translator import ( + convert_pb_to_serialized_graph_metadata, +) +from gigl.src.common.types.pb_wrappers.graph_metadata import GraphMetadataPbWrapper +from gigl.src.common.types.pb_wrappers.preprocessed_metadata import ( + PreprocessedMetadataPbWrapper, +) from gigl.src.data_preprocessor.lib.transform.feature_quantization import ( + EDGE_PACKED_FEATURE_KEY, NODE_PACKED_FEATURE_KEY, apply_feature_quantization_transform, ) from gigl.src.data_preprocessor.lib.types import FeatureQuantizationSpec +from snapchat.research.gbml import graph_schema_pb2, preprocessed_metadata_pb2 from tests.test_assets.test_case import TestCase class FeatureQuantizationTransformTest(TestCase): + def test_edge_quantization_round_trips_through_storage_and_loading(self) -> None: + with tempfile.TemporaryDirectory() as temp_dir: + metadata_path = os.path.join(temp_dir, "feature_quantization_metadata.json") + tfrecord_prefix = os.path.join(temp_dir, "edges") + schema_path = os.path.join(temp_dir, "schema.pbtxt") + logical_metadata = DatasetMetadata.from_feature_spec( + { + "src": tf.io.FixedLenFeature(shape=[], dtype=tf.int64), + "dst": tf.io.FixedLenFeature(shape=[], dtype=tf.int64), + "quantized": tf.io.FixedLenFeature(shape=[], dtype=tf.float32), + "raw": tf.io.FixedLenFeature(shape=[], dtype=tf.float32), + } + ) + logical_batches = [ + pa.RecordBatch.from_arrays( + [ + pa.array([[0], [1]], type=pa.list_(pa.int64())), + pa.array([[1], [0]], type=pa.list_(pa.int64())), + pa.array([[-2.0], [8.0]], type=pa.list_(pa.float32())), + pa.array([[10.0], [20.0]], type=pa.list_(pa.float32())), + ], + names=["src", "dst", "quantized", "raw"], + ) + ] + + with TestPipeline() as pipeline: + transformed_batches, physical_metadata = ( + apply_feature_quantization_transform( + logical_features=pipeline + | "Create edge RecordBatches" >> beam.Create(logical_batches), + logical_metadata=logical_metadata, + logical_feature_keys=["quantized", "raw"], + quantization_spec=FeatureQuantizationSpec( + feature_keys=["quantized"], bits=2 + ), + quantization_metadata_path=metadata_path, + packed_feature_key=EDGE_PACKED_FEATURE_KEY, + ) + ) + transformed_batches | "Write edge TFRecords" >> BetterWriteToTFRecord( + file_path_prefix=tfrecord_prefix, + transformed_metadata=physical_metadata, + num_shards=1, + ) + + tfdv.write_schema_text(logical_metadata.schema, schema_path) + with tf.io.gfile.GFile(metadata_path) as metadata_file: + quantization_metadata = json.loads(metadata_file.read()) + quantization_metadata_pb = preprocessed_metadata_pb2.PreprocessedMetadata.FeatureQuantizationMetadata( + packed_feature_key=quantization_metadata["packed_feature_key"], + quantized_feature_indices=quantization_metadata[ + "quantized_feature_indices" + ], + ) + quantization_metadata_pb.multi_bit_state.bits = quantization_metadata[ + "bits" + ] + quantization_metadata_pb.multi_bit_state.clip_min = quantization_metadata[ + "clip_min" + ] + quantization_metadata_pb.multi_bit_state.clip_max = quantization_metadata[ + "clip_max" + ] + self.assertEqual( + quantization_metadata_pb.packed_feature_key, "edge_packed_features" + ) + self.assertEqual( + list(quantization_metadata_pb.quantized_feature_indices), [0] + ) + + preprocessed_metadata_pb = preprocessed_metadata_pb2.PreprocessedMetadata() + preprocessed_metadata_pb.condensed_node_type_to_preprocessed_metadata[ + 0 + ].node_id_key = "node_id" + edge_metadata = ( + preprocessed_metadata_pb.condensed_edge_type_to_preprocessed_metadata[0] + ) + edge_metadata.src_node_id_key = "src" + edge_metadata.dst_node_id_key = "dst" + edge_metadata.main_edge_info.CopyFrom( + preprocessed_metadata_pb2.PreprocessedMetadata.EdgeMetadataInfo( + feature_keys=["quantized", "raw"], + feature_dim=2, + tfrecord_uri_prefix=temp_dir, + schema_uri=schema_path, + quantized_feature_metadata=quantization_metadata_pb, + ) + ) + graph_metadata_pb = graph_schema_pb2.GraphMetadata( + node_types=["node"], + edge_types=[ + graph_schema_pb2.EdgeType( + src_node_type="node", relation="connects", dst_node_type="node" + ) + ], + condensed_node_type_map={0: "node"}, + condensed_edge_type_map={ + 0: graph_schema_pb2.EdgeType( + src_node_type="node", relation="connects", dst_node_type="node" + ) + }, + ) + serialized_metadata = convert_pb_to_serialized_graph_metadata( + preprocessed_metadata_pb_wrapper=PreprocessedMetadataPbWrapper( + preprocessed_metadata_pb + ), + graph_metadata_pb_wrapper=GraphMetadataPbWrapper(graph_metadata_pb), + tfrecord_uri_pattern="edges.*\\.tfrecord", + ) + loaded = TFRecordDataLoader(rank=0, world_size=1).load_as_torch_tensors( + serialized_tf_record_info=serialized_metadata.edge_entity_info, + tf_dataset_options=TFDatasetOptions(deterministic=True), + ) + + assert loaded.features is not None + assert loaded.quantized_features is not None + self.assert_tensor_equality(loaded.ids, torch.tensor([[0, 1], [1, 0]])) + self.assert_tensor_equality(loaded.features, torch.tensor([[10.0], [20.0]])) + self.assert_tensor_equality( + loaded.quantized_features, torch.tensor([[0], [192]], dtype=torch.uint8) + ) + materialized, remaining_metadata = materialize_quantized_edge_features( + data=Data(edge_index=loaded.ids, edge_attr=loaded.features), + metadata={EDGE_PACKED_FEATURES_METADATA_KEY: loaded.quantized_features}, + edge_quantization_metadata=serialized_metadata.edge_quantization_metadata, + ) + self.assert_tensor_equality( + materialized.edge_attr, torch.tensor([[-2.0, 10.0], [8.0, 20.0]]) + ) + self.assertEqual(remaining_metadata, {}) + def test_apply_feature_quantization_transform_rejects_reserved_schema_key( self, ) -> None: diff --git a/tests/test_assets/distributed/run_distributed_partitioner.py b/tests/test_assets/distributed/run_distributed_partitioner.py index c16a3469b..c50445fc6 100644 --- a/tests/test_assets/distributed/run_distributed_partitioner.py +++ b/tests/test_assets/distributed/run_distributed_partitioner.py @@ -23,6 +23,9 @@ class InputDataStrategy(Enum): "REGISTER_EDGE_WEIGHTS_WITHOUT_EDGE_FEATURES" ) REGISTER_EDGE_QUANTIZED_FEATURES = "REGISTER_EDGE_QUANTIZED_FEATURES" + REGISTER_EDGE_QUANTIZED_FEATURES_WITHOUT_EDGE_FEATURES = ( + "REGISTER_EDGE_QUANTIZED_FEATURES_WITHOUT_EDGE_FEATURES" + ) def run_distributed_partitioner( @@ -96,7 +99,10 @@ 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 == InputDataStrategy.REGISTER_EDGE_QUANTIZED_FEATURES: + if input_data_strategy in ( + InputDataStrategy.REGISTER_EDGE_QUANTIZED_FEATURES, + InputDataStrategy.REGISTER_EDGE_QUANTIZED_FEATURES_WITHOUT_EDGE_FEATURES, + ): dist_partitioner = partitioner_class( should_assign_edges_by_src_node=should_assign_edges_by_src_node, ) @@ -105,19 +111,35 @@ def run_distributed_partitioner( 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) - assert isinstance(edge_features, dict) - edge_quantized_features = { - edge_type: torch.stack( - (indices[0] * 3 + 17, indices[0] * 5 + 29), dim=1 - ).to(torch.uint8) - for edge_type, indices in edge_index_by_type.items() - if edge_type in edge_features - } + if ( + input_data_strategy + == InputDataStrategy.REGISTER_EDGE_QUANTIZED_FEATURES + ): + assert isinstance(edge_features, dict) + edge_quantized_features = { + edge_type: torch.stack( + (indices[0] * 3 + 17, indices[0] * 5 + 29), dim=1 + ).to(torch.uint8) + for edge_type, indices in edge_index_by_type.items() + if edge_type in edge_features + } + else: + 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 = torch.stack( - (edge_index[0] * 3 + 17, edge_index[0] * 5 + 29), dim=1 - ).to(torch.uint8) - dist_partitioner.register_edge_features(edge_features=edge_features) + if ( + input_data_strategy + == InputDataStrategy.REGISTER_EDGE_QUANTIZED_FEATURES + ): + edge_quantized_features = torch.stack( + (edge_index[0] * 3 + 17, edge_index[0] * 5 + 29), dim=1 + ).to(torch.uint8) + else: + edge_quantized_features = edge_index[0].to(torch.uint8).unsqueeze(1) + if input_data_strategy == InputDataStrategy.REGISTER_EDGE_QUANTIZED_FEATURES: + dist_partitioner.register_edge_features(edge_features=edge_features) dist_partitioner.register_edge_quantized_features( edge_quantized_features=edge_quantized_features ) diff --git a/tests/unit/distributed/dist_ablp_neighborloader_test.py b/tests/unit/distributed/dist_ablp_neighborloader_test.py index a3d25dfeb..e02abe66c 100644 --- a/tests/unit/distributed/dist_ablp_neighborloader_test.py +++ b/tests/unit/distributed/dist_ablp_neighborloader_test.py @@ -200,7 +200,11 @@ def _run_cora_supervised( shutdown_rpc() -def _run_quantized_homogeneous_ablp_loader(_: int, dataset: DistDataset) -> None: +def _run_quantized_homogeneous_ablp_loader( + _: int, + dataset: DistDataset, + expected_edge_features: dict[tuple[int, int], torch.Tensor], +) -> None: """Assert homogeneous ABLP materializes partial packed features.""" create_test_process_group() loader = DistABLPLoader( @@ -218,6 +222,14 @@ def _run_quantized_homogeneous_ablp_loader(_: int, dataset: DistDataset) -> None for batch in loader: assert isinstance(batch, Data) assert_tensor_equality(batch.x, expected_features[batch.node]) + for local_edge_index, edge_feature in zip(batch.edge_index.T, batch.edge_attr): + source, destination = batch.node[local_edge_index] + # Outward sampling reverses the batch edge for message passing while + # preserving the original edge's feature payload. + assert_tensor_equality( + edge_feature, + expected_edge_features[(destination.item(), source.item())], + ) assert _global_pair_set(batch.node, batch.node, batch.y_positive) == [(0, 1)] assert _global_pair_set(batch.node, batch.node, batch.y_negative) == [(0, 2)] batch_count += 1 @@ -770,12 +782,14 @@ def test_homogeneous_ablp_materializes_quantized_features(self) -> None: loaded_graph_tensors = LoadedGraphTensors( node_ids=torch.arange(3), node_features=torch.tensor([[10.0, 20.0], [30.0, 40.0], [50.0, 60.0]]), + # High-order 2-bit codes unpack to [0, 3], [2, 1], and [1, 2]. node_quantized_features=torch.tensor( [[48], [144], [96]], dtype=torch.uint8 ), node_labels=None, edge_index=torch.tensor([[0, 1], [1, 0]]), edge_features=None, + edge_quantized_features=torch.tensor([[48], [144]], dtype=torch.uint8), positive_label=torch.tensor([[0], [1]]), negative_label=torch.tensor([[0], [2]]), ) @@ -792,6 +806,7 @@ def test_homogeneous_ablp_materializes_quantized_features(self) -> None: assert isinstance(loaded_graph_tensors.edge_index, dict) assert isinstance(loaded_graph_tensors.node_features, dict) assert isinstance(loaded_graph_tensors.node_quantized_features, dict) + assert isinstance(loaded_graph_tensors.edge_quantized_features, dict) edge_index = cast(dict[EdgeType, torch.Tensor], loaded_graph_tensors.edge_index) node_features = cast( dict[NodeType, torch.Tensor], loaded_graph_tensors.node_features @@ -799,6 +814,9 @@ def test_homogeneous_ablp_materializes_quantized_features(self) -> None: node_quantized_features = cast( dict[NodeType, torch.Tensor], loaded_graph_tensors.node_quantized_features ) + edge_quantized_features = cast( + dict[EdgeType, torch.Tensor], loaded_graph_tensors.edge_quantized_features + ) partition_output = PartitionOutput( node_partition_book={DEFAULT_HOMOGENEOUS_NODE_TYPE: torch.zeros(3)}, edge_partition_book={edge_type: torch.zeros(3) for edge_type in edge_index}, @@ -820,6 +838,12 @@ def test_homogeneous_ablp_materializes_quantized_features(self) -> None: ids=torch.arange(3), ) }, + partitioned_edge_quantized_features={ + DEFAULT_HOMOGENEOUS_EDGE_TYPE: FeaturePartitionData( + feats=edge_quantized_features[DEFAULT_HOMOGENEOUS_EDGE_TYPE], + ids=torch.arange(2), + ) + }, partitioned_edge_features=None, partitioned_negative_labels=None, partitioned_positive_labels=None, @@ -830,10 +854,26 @@ def test_homogeneous_ablp_materializes_quantized_features(self) -> None: world_size=1, edge_dir="out", node_quantization_metadata=quantization_metadata, + edge_quantization_metadata=FeatureQuantizationMetadata( + bits=2, + feature_dim=2, + quantized_feature_indices=(0, 1), + clip_min=0.0, + clip_max=3.0, + ), ) dataset.build(partition_output=partition_output) - mp.spawn(fn=_run_quantized_homogeneous_ablp_loader, args=(dataset,)) + # The two packed edge bytes decode to [0, 3] for edge 0 -> 1 and + # [2, 1] for edge 1 -> 0. + expected_edge_features = { + (0, 1): torch.tensor([0.0, 3.0]), + (1, 0): torch.tensor([2.0, 1.0]), + } + mp.spawn( + fn=_run_quantized_homogeneous_ablp_loader, + args=(dataset, expected_edge_features), + ) @parameterized.expand( [ diff --git a/tests/unit/distributed/dist_server_test.py b/tests/unit/distributed/dist_server_test.py index c876fcef1..1b0f1675e 100644 --- a/tests/unit/distributed/dist_server_test.py +++ b/tests/unit/distributed/dist_server_test.py @@ -1,5 +1,5 @@ import threading -from unittest.mock import MagicMock, patch +from unittest.mock import MagicMock, call, patch import torch from absl.testing import absltest @@ -65,45 +65,72 @@ def test_get_node_feature_info_with_homogeneous_dataset(self) -> None: # Verify it returns the correct feature info self.assertIsNone(node_feature_info) - def test_get_node_quantization_metadata(self) -> None: - metadata = FeatureQuantizationMetadata( + def test_get_quantization_metadata(self) -> None: + node_metadata = FeatureQuantizationMetadata( bits=2, feature_dim=2, quantized_feature_indices=(0, 1), clip_min=0.0, clip_max=3.0, ) + edge_metadata = { + USER_TO_STORY: FeatureQuantizationMetadata( + bits=4, + feature_dim=3, + quantized_feature_indices=(0, 2), + clip_min=-1.0, + clip_max=1.0, + ) + } dataset = DistDataset( rank=0, world_size=1, edge_dir="out", - node_quantization_metadata=metadata, + node_quantization_metadata=node_metadata, + edge_quantization_metadata=edge_metadata, ) server = dist_server.DistServer(dataset) - self.assertEqual(server.get_node_quantization_metadata(), metadata) + self.assertEqual(server.get_node_quantization_metadata(), node_metadata) + self.assertEqual(server.get_edge_quantization_metadata(), edge_metadata) - def test_remote_dataset_fetches_node_quantization_metadata(self) -> None: - metadata = FeatureQuantizationMetadata( + def test_remote_dataset_fetches_quantization_metadata(self) -> None: + node_metadata = FeatureQuantizationMetadata( bits=2, feature_dim=2, quantized_feature_indices=(0, 1), clip_min=0.0, clip_max=3.0, ) + edge_metadata = { + USER_TO_STORY: FeatureQuantizationMetadata( + bits=4, + feature_dim=3, + quantized_feature_indices=(0, 2), + clip_min=-1.0, + clip_max=1.0, + ) + } with patch( "gigl.distributed.graph_store.remote_dist_dataset.request_server", - return_value=metadata, + side_effect=[node_metadata, edge_metadata], ) as request_server: remote_dataset = RemoteDistDataset(cluster_info=MagicMock(), local_rank=0) self.assertEqual( - remote_dataset.fetch_node_quantization_metadata(), metadata + remote_dataset.fetch_node_quantization_metadata(), node_metadata + ) + self.assertEqual( + remote_dataset.fetch_edge_quantization_metadata(), edge_metadata ) - request_server.assert_called_once_with( - 0, dist_server.DistServer.get_node_quantization_metadata + self.assertEqual( + request_server.call_args_list, + [ + call(0, dist_server.DistServer.get_node_quantization_metadata), + call(0, dist_server.DistServer.get_edge_quantization_metadata), + ], ) def test_get_edge_feature_info_with_heterogeneous_dataset(self) -> None: diff --git a/tests/unit/distributed/distributed_neighborloader_test.py b/tests/unit/distributed/distributed_neighborloader_test.py index e51f10ea0..4ef56a38e 100644 --- a/tests/unit/distributed/distributed_neighborloader_test.py +++ b/tests/unit/distributed/distributed_neighborloader_test.py @@ -5,6 +5,7 @@ import torch.multiprocessing as mp from absl.testing import absltest from graphlearn_torch.distributed import shutdown_rpc +from graphlearn_torch.utils import reverse_edge_type from parameterized import param, parameterized from torch_geometric.data import Data, HeteroData @@ -12,6 +13,7 @@ from gigl.distributed.dataset_factory import build_dataset from gigl.distributed.dist_dataset import DistDataset from gigl.distributed.distributed_neighborloader import DistNeighborLoader +from gigl.distributed.sampler import EDGE_PACKED_FEATURES_METADATA_KEY from gigl.distributed.utils import get_free_port from gigl.distributed.utils.neighborloader import DatasetSchema from gigl.distributed.utils.serialized_graph_metadata_translator import ( @@ -463,7 +465,6 @@ def _run_heterogeneous_partially_quantized_neighbor_loader( batch_size=1, pin_memory_device=torch.device("cpu"), ) - batch = next(iter(loader)) assert isinstance(batch, HeteroData) for node_type, expected_features_for_node_type in expected_features.items(): @@ -1164,17 +1165,40 @@ def test_independent_calls_produce_equal_configs(self) -> None: ) self.assertEqual(first, second) + def test_packed_edge_metadata_requests_sampled_edge_ids(self) -> None: + schema = self._schema() + schema = DatasetSchema( + is_homogeneous_with_labeled_edge_type=False, + edge_types=schema.edge_types, + node_feature_info=schema.node_feature_info, + edge_feature_info=None, + edge_dir=schema.edge_dir, + edge_quantization_metadata=FeatureQuantizationMetadata( + bits=2, + feature_dim=2, + quantized_feature_indices=(0, 1), + clip_min=0.0, + clip_max=3.0, + ), + ) + + config = BaseDistLoader.create_sampling_config( + num_neighbors=[1], dataset_schema=schema + ) + + self.assertTrue(config.with_edge) + # NOTE on the test strategy: GiGL loaders always sample via the multiprocess # producer, which spawns worker subprocesses with a *fresh* interpreter # (`mp.get_context("spawn")`, dist_sampling_producer.py). A `mock.patch` applied in the # loader process therefore never reaches the sampler running in that subprocess, so we # cannot inject a synthetic failure by mocking the sampler. Instead we reproduce a real -# sampler failure end-to-end: a heterogeneous dataset with edge features on only a -# subset of its message-passing edge types. When the featureless type is reached during -# sampling, its feature lookup raises `KeyError` inside the sampling coroutine — the exact -# swallowed-exception case this change surfaces. Without the change this hangs forever, so -# the test uses a bounded join. +# sampler failure end-to-end: a heterogeneous dataset with an incomplete feature store +# for one message-passing edge type. When the missing edge ID is reached during sampling, +# its feature lookup raises inside the sampling coroutine - the exact swallowed-exception +# case this change surfaces. Without the change this hangs forever, so the test uses a +# bounded join. def _run_partial_edge_feature_coverage_raises( @@ -1201,11 +1225,11 @@ def _run_partial_edge_feature_coverage_raises( class TestSamplingErrorPropagation(TestCase): def _build_partial_edge_feature_dataset(self) -> DistDataset: - """Build a hetero dataset with edge features on only one message-passing type. + """Build a hetero dataset with an incomplete edge feature store. - ``user-to-story`` has edge features; ``story-to-user`` does not. Both are - reachable from ``user`` seeds within a 2-hop fanout, so the featureless type is - actually sampled and its edge-feature lookup raises inside the coroutine. + Both edge types are reachable from ``user`` seeds within a 2-hop fanout. + ``story-to-user`` omits the last edge ID, so its feature lookup raises inside + the sampling coroutine. """ n = 5 edge_index = torch.tensor([[0, 1, 2, 3, 4], [0, 1, 2, 3, 4]]) @@ -1235,6 +1259,9 @@ def _build_partial_edge_feature_dataset(self) -> DistDataset: _USER_TO_STORY: FeaturePartitionData( feats=torch.ones(n, 3), ids=torch.arange(n) ), + _STORY_TO_USER: FeaturePartitionData( + feats=torch.ones(n - 1, 3), ids=torch.arange(n - 1) + ), }, partitioned_positive_labels=None, partitioned_negative_labels=None, @@ -1264,7 +1291,169 @@ def test_reachable_sampler_failure_raises_not_hangs(self) -> None: message = error_holder.get("msg", "") # The training process raised with the worker's real traceback embedded. self.assertIn("sampling worker failed", message.lower()) - self.assertIn("story-to-user", message) + self.assertIn("IndexError", message) + self.assertIn("index 4 is out of bounds", message) + + +def _run_heterogeneous_partially_quantized_edge_feature_neighbor_loader( + _, + dataset: DistDataset, + expected_edge_features: dict[EdgeType, torch.Tensor], +) -> None: + create_test_process_group() + loader = DistNeighborLoader( + dataset=dataset, + input_nodes=(_USER, torch.tensor([0])), + num_neighbors=[1, 1], + batch_size=1, + pin_memory_device=torch.device("cpu"), + ) + batch = next(iter(loader)) + assert isinstance(batch, HeteroData) + for edge_type, expected_features in expected_edge_features.items(): + assert_tensor_equality(batch[edge_type].edge_attr, expected_features) + assert not hasattr(batch[_STORY_TO_USER], "edge_attr") + shutdown_rpc() + + +def _run_incoming_heterogeneous_quantized_edge_feature_neighbor_loader( + _, + dataset: DistDataset, + expected_edge_features: torch.Tensor, +) -> None: + create_test_process_group() + loader = DistNeighborLoader( + dataset=dataset, + input_nodes=(_STORY, torch.tensor([0])), + num_neighbors=[1], + batch_size=1, + pin_memory_device=torch.device("cpu"), + ) + + edge_feature_info = loader._edge_feature_info + edge_quantization_metadata = loader._edge_quantization_metadata + assert isinstance(edge_feature_info, dict) + assert isinstance(edge_quantization_metadata, dict) + assert set(edge_feature_info) == {_USER_TO_STORY} + assert set(edge_quantization_metadata) == {_USER_TO_STORY} + assert edge_feature_info[_USER_TO_STORY].dim == 2 + assert edge_quantization_metadata[_USER_TO_STORY].feature_dim == 4 + + batch = next(iter(loader)) + assert isinstance(batch, HeteroData) + assert_tensor_equality(batch[_USER_TO_STORY].edge_attr, expected_edge_features) + assert EDGE_PACKED_FEATURES_METADATA_KEY not in batch + assert f"{EDGE_PACKED_FEATURES_METADATA_KEY}.{_USER_TO_STORY}" not in batch + shutdown_rpc() + + +class HeterogeneousEdgeFeatureLookupTest(TestCase): + def test_heterogeneous_loader_supports_partially_quantized_edge_types( + self, + ) -> None: + # Sampling user reaches both edge types. Only user-to-story has raw and + # packed features, so story-to-user must not be looked up in either store. + expected_edge_features = { + reverse_edge_type(_USER_TO_STORY): torch.tensor([[0.0, 10.0]]) + } + partition_output = PartitionOutput( + node_partition_book={_USER: torch.zeros(1), _STORY: torch.zeros(1)}, + edge_partition_book={_USER_TO_STORY: torch.zeros(1)}, + partitioned_edge_index={ + _USER_TO_STORY: GraphPartitionData( + edge_index=torch.tensor([[0], [0]]), edge_ids=None + ) + }, + partitioned_node_features=None, + partitioned_edge_features={ + _USER_TO_STORY: FeaturePartitionData( + feats=torch.tensor([[10.0]]), ids=torch.tensor([0]) + ), + }, + partitioned_edge_quantized_features={ + _USER_TO_STORY: FeaturePartitionData( + feats=torch.tensor([[0]], dtype=torch.uint8), + ids=torch.tensor([0]), + ) + }, + partitioned_positive_labels=None, + partitioned_negative_labels=None, + partitioned_node_labels=None, + ) + dataset = DistDataset( + rank=0, + world_size=1, + edge_dir="out", + edge_quantization_metadata={ + _USER_TO_STORY: FeatureQuantizationMetadata( + bits=2, + feature_dim=2, + quantized_feature_indices=(0,), + clip_min=0.0, + clip_max=3.0, + ) + }, + ) + dataset.build(partition_output=partition_output) + + mp.spawn( + fn=_run_heterogeneous_partially_quantized_edge_feature_neighbor_loader, + args=(dataset, expected_edge_features), + ) + + def test_incoming_edges_reverse_feature_metadata_and_output_stores(self) -> None: + partition_output = PartitionOutput( + node_partition_book={_USER: torch.zeros(1), _STORY: torch.zeros(1)}, + edge_partition_book={ + _USER_TO_STORY: torch.zeros(1), + _STORY_TO_USER: torch.zeros(1), + }, + partitioned_edge_index={ + _USER_TO_STORY: GraphPartitionData( + edge_index=torch.tensor([[0], [0]]), edge_ids=None + ), + _STORY_TO_USER: GraphPartitionData( + edge_index=torch.tensor([[0], [0]]), edge_ids=None + ), + }, + partitioned_node_features=None, + partitioned_edge_features={ + _USER_TO_STORY: FeaturePartitionData( + feats=torch.tensor([[10.0, 20.0]]), ids=torch.tensor([0]) + ) + }, + partitioned_edge_quantized_features={ + _USER_TO_STORY: FeaturePartitionData( + # The high-order 2-bit codes unpack to [0, 3] in feature + # slots 0 and 2; raw [10, 20] therefore yields [0, 10, 3, 20]. + feats=torch.tensor([[48]], dtype=torch.uint8), + ids=torch.tensor([0]), + ) + }, + partitioned_positive_labels=None, + partitioned_negative_labels=None, + partitioned_node_labels=None, + ) + dataset = DistDataset( + rank=0, + world_size=1, + edge_dir="in", + edge_quantization_metadata={ + _USER_TO_STORY: FeatureQuantizationMetadata( + bits=2, + feature_dim=4, + quantized_feature_indices=(0, 2), + clip_min=0.0, + clip_max=3.0, + ) + }, + ) + dataset.build(partition_output=partition_output) + + mp.spawn( + fn=_run_incoming_heterogeneous_quantized_edge_feature_neighbor_loader, + args=(dataset, torch.tensor([[0.0, 10.0, 3.0, 20.0]])), + ) if __name__ == "__main__": diff --git a/tests/unit/distributed/distributed_partitioner_test.py b/tests/unit/distributed/distributed_partitioner_test.py index 578b8814b..b2b59c474 100644 --- a/tests/unit/distributed/distributed_partitioner_test.py +++ b/tests/unit/distributed/distributed_partitioner_test.py @@ -702,6 +702,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( @@ -775,6 +791,10 @@ def test_partitioning_correctness( has_edge_quantized_features = ( input_data_strategy == InputDataStrategy.REGISTER_EDGE_QUANTIZED_FEATURES ) + 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 @@ -800,7 +820,30 @@ def test_partitioning_correctness( graph.edge_index ) - if has_edge_quantized_features: + if is_packed_edge_only: + self.assertIsNotNone(partition_output.edge_partition_book) + self.assertIsNone(partition_output.partitioned_edge_features) + packed_features = partition_output.partitioned_edge_quantized_features + self.assertIsNotNone(packed_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 has_edge_quantized_features: self.assertIsNotNone(partition_output.edge_partition_book) assert partition_output.partitioned_edge_features is not None self._assert_edge_feature_outputs( @@ -830,10 +873,18 @@ def test_partitioning_correctness( else partitioned_edge_index ) self.assertEqual(edge_type_features.feats.dtype, torch.uint8) + if is_range_based_partition: + self.assertIsNone(edge_type_features.ids) + expected_source_nodes = edge_type_graph.edge_index[0] + else: + assert edge_type_features.ids is not None + expected_source_nodes = MOCKED_UNIFIED_GRAPH.edge_index[ + edge_type + ][0, edge_type_features.ids] expected_features = torch.stack( ( - edge_type_graph.edge_index[0] * 3 + 17, - edge_type_graph.edge_index[0] * 5 + 29, + expected_source_nodes * 3 + 17, + expected_source_nodes * 5 + 29, ), dim=1, ).to(torch.uint8) @@ -841,7 +892,8 @@ def test_partitioning_correctness( tensor_a=edge_type_features.feats, tensor_b=expected_features, ) - if edge_type_features.ids is not None: + if not is_range_based_partition: + assert edge_type_features.ids is not None assert edge_type_graph.edge_ids is not None self.assert_tensor_equality( tensor_a=edge_type_features.ids, diff --git a/tests/unit/distributed/utils/neighborloader_test.py b/tests/unit/distributed/utils/neighborloader_test.py index 20ad8b710..372058e85 100644 --- a/tests/unit/distributed/utils/neighborloader_test.py +++ b/tests/unit/distributed/utils/neighborloader_test.py @@ -7,6 +7,7 @@ from torch_geometric.typing import EdgeType from gigl.distributed.sampler import ( + EDGE_PACKED_FEATURES_METADATA_KEY, NEGATIVE_LABEL_METADATA_KEY, NODE_PACKED_FEATURES_METADATA_KEY, POSITIVE_LABEL_METADATA_KEY, @@ -16,6 +17,7 @@ extract_edge_type_metadata, extract_metadata, labeled_to_homogeneous, + materialize_quantized_edge_features, materialize_quantized_node_features, patch_fanout_for_sampling, set_missing_features, @@ -24,6 +26,7 @@ strip_non_ppr_edge_types, ) from gigl.types.graph import ( + DEFAULT_HOMOGENEOUS_EDGE_TYPE, FeatureInfo, FeatureQuantizationMetadata, message_passing_to_positive_label, @@ -101,6 +104,119 @@ def test_materialize_quantized_node_features_reconstructs_feature_order( self.assertEqual(set(remaining_metadata), {"request_id"}) self.assert_tensor_equality(remaining_metadata["request_id"], torch.tensor([7])) + def test_materialize_quantized_edge_features_reconstructs_feature_order( + self, + ) -> None: + data = Data(edge_attr=torch.tensor([[10.0, 20.0], [30.0, 40.0]])) + # The high-order 2-bit codes unpack into slots 0 and 2: + # 0b00_11_00_00 -> [0, 3], 0b10_01_00_00 -> [2, 1]. + metadata = { + "edge_packed_features": torch.tensor([[48], [144]], dtype=torch.uint8), + "request_id": torch.tensor([7]), + } + + materialized_data, remaining_metadata = materialize_quantized_edge_features( + data=data, + metadata=metadata, + edge_quantization_metadata=FeatureQuantizationMetadata( + bits=2, + feature_dim=4, + quantized_feature_indices=(0, 2), + clip_min=0.0, + clip_max=3.0, + ), + ) + + self.assert_tensor_equality( + materialized_data.edge_attr, + torch.tensor([[0.0, 10.0, 3.0, 20.0], [2.0, 30.0, 1.0, 40.0]]), + ) + self.assertEqual(set(remaining_metadata), {"request_id"}) + + def test_materialize_quantized_edge_features_uses_labeled_homogeneous_key( + self, + ) -> None: + data = Data(edge_index=torch.tensor([[0], [1]])) + typed_key = ( + f"{EDGE_PACKED_FEATURES_METADATA_KEY}.{DEFAULT_HOMOGENEOUS_EDGE_TYPE}" + ) + + # The high-order 2-bit codes in 0b00_11_00_00 unpack to [0, 3]. + materialized_data, remaining_metadata = materialize_quantized_edge_features( + data=data, + metadata={typed_key: torch.tensor([[48]], dtype=torch.uint8)}, + edge_quantization_metadata=FeatureQuantizationMetadata( + bits=2, + feature_dim=2, + quantized_feature_indices=(0, 1), + clip_min=0.0, + clip_max=3.0, + ), + ) + + self.assert_tensor_equality( + materialized_data.edge_attr, torch.tensor([[0.0, 3.0]]) + ) + self.assertEqual(remaining_metadata, {}) + + def test_materialize_quantized_edge_features_maps_outward_sampling_to_output_edge_type( + self, + ) -> None: + source_edge_type = ("user", "to", "item") + output_edge_type = ("item", "rev_to", "user") + data = HeteroData() + data[output_edge_type].edge_attr = torch.tensor([[10.0]]) + # The high-order 2-bit codes in 0b00_11_00_00 unpack to [0, 3]. + metadata = { + f"{EDGE_PACKED_FEATURES_METADATA_KEY}.{output_edge_type}": torch.tensor( + [[48]], dtype=torch.uint8 + ) + } + quantization_metadata = FeatureQuantizationMetadata( + bits=2, + feature_dim=3, + quantized_feature_indices=(0, 2), + clip_min=0.0, + clip_max=3.0, + ) + + materialized_data, remaining_metadata = materialize_quantized_edge_features( + data=data, + metadata=metadata, + edge_quantization_metadata={source_edge_type: quantization_metadata}, + edge_dir="out", + ) + + self.assert_tensor_equality( + materialized_data[output_edge_type].edge_attr, + torch.tensor([[0.0, 10.0, 3.0]]), + ) + self.assertEqual(remaining_metadata, {}) + + def test_materialize_quantized_edge_features_rejects_missing_packed_features_for_sampled_edge_type( + self, + ) -> None: + data = HeteroData() + data[_U2I_EDGE_TYPE].edge_index = torch.tensor([[0], [1]]) + data[_U2I_EDGE_TYPE].edge_attr = torch.tensor([[10.0]]) + + with self.assertRaisesRegex( + ValueError, "Missing packed quantized edge features" + ): + materialize_quantized_edge_features( + data=data, + metadata={}, + edge_quantization_metadata={ + _U2I_EDGE_TYPE: FeatureQuantizationMetadata( + bits=2, + feature_dim=3, + quantized_feature_indices=(0, 2), + clip_min=0.0, + clip_max=3.0, + ) + }, + ) + def test_materialize_quantized_node_features_uses_per_node_type_metadata( self, ) -> None: