Edge Feature Quantization [4/5]: Distributed Materialization

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()

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,

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
 

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(

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.