From d5205fe5e25718b734a222144c3d073dbecca3dd Mon Sep 17 00:00:00 2001 From: kmontemayor Date: Wed, 8 Jul 2026 22:41:12 +0000 Subject: [PATCH 1/3] fix(distributed): skip feature fetch for featureless node/edge types in collate Heterogeneous datasets may carry node/edge features on only a subset of their types. Previously BaseDistNeighborSampler._collate_fn fetched features for every sampled node/edge type, so a featureless-but-reachable type raised a KeyError inside the sampling coroutine. GLT's event loop swallows that exception, drops the batch, and never sends a message, so the loader hangs indefinitely. Guard both hetero feature-fetch loops on membership in DistFeature.local_feature (the per-type feature dict), skipping types with no feature store entry. The all2all node path is left unguarded with a comment, since GiGL never enables use_all2all. Adds regression tests (bounded by join-with-timeout, since the pre-fix bug hangs) for partial edge-feature and node-feature coverage in the standard and ABLP loaders. Co-Authored-By: Claude Fable 5 --- gigl/distributed/base_sampler.py | 19 +++ .../dist_ablp_neighborloader_test.py | 98 +++++++++++++ .../distributed_neighborloader_test.py | 134 ++++++++++++++++++ 3 files changed, 251 insertions(+) diff --git a/gigl/distributed/base_sampler.py b/gigl/distributed/base_sampler.py index 986ba5d58..b7da872c0 100644 --- a/gigl/distributed/base_sampler.py +++ b/gigl/distributed/base_sampler.py @@ -306,6 +306,11 @@ async def _collate_fn( ] if self.dist_node_feature is not None: if self.use_all2all: + # NOTE: The all2all path is intentionally left unguarded against + # featureless node types. GiGL never enables use_all2all (it + # defaults to False and is never set True), so this branch is + # unreachable in practice; get_all2all would raise the same + # swallowed KeyError for a featureless type if it were used. sorted_ntype = sorted(self.dist_node_feature.feature_pb.keys()) nfeat_dict = self.dist_node_feature.get_all2all( output, sorted_ntype @@ -314,7 +319,14 @@ async def _collate_fn( result_map[f"{as_str(ntype)}.nfeats"] = nfeats else: nfeat_fut_dict = {} + # Only fetch features for node types that actually have a feature + # store entry. ``local_feature`` is the per-type dict in the + # heterogeneous case; a type absent from it has no features, and + # fetching it would raise a KeyError swallowed by GLT's event loop. + featured_node_types = self.dist_node_feature.local_feature for ntype, nodes in output.node.items(): + if ntype not in featured_node_types: + continue nodes = nodes.to(torch.long) nfeat_fut_dict[ntype] = self.dist_node_feature.async_get( nodes, ntype @@ -324,7 +336,14 @@ async def _collate_fn( result_map[f"{as_str(ntype)}.nfeats"] = nfeats if self.dist_edge_feature is not None and self.with_edge: efeat_fut_dict = {} + # Only fetch features for edge types that actually have a feature + # store entry. ``local_feature`` is the per-type dict in the + # heterogeneous case; a type absent from it has no features, and + # fetching it would raise a KeyError swallowed by GLT's event loop. + featured_edge_types = self.dist_edge_feature.local_feature for etype in self.edge_types: + if etype not in featured_edge_types: + continue if self.edge_dir == "in": eids = result_map.get( f"{as_str(reverse_edge_type(etype))}.eids", None diff --git a/tests/unit/distributed/dist_ablp_neighborloader_test.py b/tests/unit/distributed/dist_ablp_neighborloader_test.py index 2d8e4c626..d5b7942d0 100644 --- a/tests/unit/distributed/dist_ablp_neighborloader_test.py +++ b/tests/unit/distributed/dist_ablp_neighborloader_test.py @@ -28,6 +28,7 @@ ) from gigl.types.graph import ( DEFAULT_HOMOGENEOUS_EDGE_TYPE, + FeaturePartitionData, GraphPartitionData, PartitionOutput, is_label_edge_type, @@ -1028,5 +1029,102 @@ def test_ablp_dataloader_invalid_inputs( DistABLPLoader(**kwargs) +_USER_TO_STORY_POSITIVE = message_passing_to_positive_label(_USER_TO_STORY) +_USER_TO_USER = EdgeType(_USER, Relation("to"), _USER) + + +def _run_partial_edge_features_ablp( + _, + dataset: DistDataset, + holder, +): + """Iterate a heterogeneous ABLP loader with edge features on only one type. + + ``(user, to, user)`` is reachable from the ``user`` anchors but featureless, so + without the per-type feature-fetch guard it triggers a swallowed ``KeyError`` in + the sampling coroutine and the loader hangs. The caller joins with a timeout and + asserts termination. + """ + create_test_process_group() + loader = DistABLPLoader( + dataset=dataset, + num_neighbors=[2, 2], + # Anchor on the users that carry positive labels (0, 1, 2). Each has a + # (user, to, user) out-edge, so the featureless user->user type is sampled. + input_nodes=(_USER, torch.tensor([0, 1, 2])), + supervision_edge_type=_USER_TO_STORY, + batch_size=3, + pin_memory_device=torch.device("cpu"), + ) + count = 0 + for datum in loader: + assert isinstance(datum, HeteroData) + count += 1 + holder["count"] = count + shutdown_rpc() + + +class TestPartialFeatureCoverageABLP(TestCase): + """ABLP counterpart of the partial-edge-feature-coverage regression test.""" + + def test_partial_edge_features_ablp_yields_batches(self) -> None: + message_passing_edge_index = { + _USER_TO_USER: torch.tensor([[0, 1, 2, 3, 4, 5], [1, 2, 3, 4, 5, 0]]), + _USER_TO_STORY: torch.tensor([[0, 1, 2, 3, 4, 5], [0, 1, 2, 3, 4, 5]]), + _USER_TO_STORY_POSITIVE: torch.tensor([[0, 1, 2], [2, 3, 4]]), + } + partition_output = PartitionOutput( + node_partition_book={ + _USER: torch.zeros(6), + _STORY: torch.zeros(6), + }, + edge_partition_book={ + edge_type: torch.zeros(edge_index.size(1)) + for edge_type, edge_index in message_passing_edge_index.items() + }, + partitioned_edge_index={ + edge_type: GraphPartitionData( + edge_index=edge_index, + edge_ids=torch.arange(edge_index.size(1)), + ) + for edge_type, edge_index in message_passing_edge_index.items() + }, + partitioned_node_features={ + _USER: FeaturePartitionData( + feats=torch.zeros(6, 2), ids=torch.arange(6) + ), + _STORY: FeaturePartitionData( + feats=torch.zeros(6, 2), ids=torch.arange(6) + ), + }, + # Edge features on the user->story type only; user->user is reachable + # from the anchors but featureless. + partitioned_edge_features={ + _USER_TO_STORY: FeaturePartitionData( + feats=torch.zeros(6, 3), ids=torch.arange(6) + ), + }, + partitioned_positive_labels=None, + partitioned_negative_labels=None, + partitioned_node_labels=None, + ) + dataset = DistDataset(rank=0, world_size=1, edge_dir="out") + dataset.build(partition_output=partition_output) + + holder = mp.Manager().dict() + proc = mp.get_context("spawn").Process( + target=_run_partial_edge_features_ablp, + args=(0, dataset, holder), + ) + proc.start() + proc.join(timeout=120) + if proc.is_alive(): + proc.terminate() + proc.join() + self.fail("ABLP loader hung — partial feature coverage not handled") + self.assertEqual(proc.exitcode, 0, "ABLP loader worker exited with an error") + self.assertGreater(holder.get("count", 0), 0) + + if __name__ == "__main__": absltest.main() diff --git a/tests/unit/distributed/distributed_neighborloader_test.py b/tests/unit/distributed/distributed_neighborloader_test.py index f95ca77d9..48bc38d64 100644 --- a/tests/unit/distributed/distributed_neighborloader_test.py +++ b/tests/unit/distributed/distributed_neighborloader_test.py @@ -52,6 +52,7 @@ _STORY = NodeType("story") _USER_TO_STORY = EdgeType(_USER, Relation("to"), _STORY) _STORY_TO_USER = EdgeType(_STORY, Relation("to"), _USER) +_USER_TO_USER = EdgeType(_USER, Relation("to"), _USER) # GLT requires subclasses of DistNeighborLoader to be run in a separate process. Otherwise, we may run into segmentation fault # or other memory issues. Calling these functions in separate proceses also allows us to use shutdown_rpc() to ensure cleanup of @@ -398,6 +399,33 @@ def _run_cora_supervised_node_classification( shutdown_rpc() +def _run_partial_feature_coverage_loader( + _, + dataset: DistDataset, + holder, +): + """Iterate a heterogeneous loader over ``user`` seeds and record the batch count. + + Used by the partial-feature-coverage regression tests. Without the per-type + feature-fetch guard this hangs on a swallowed ``KeyError`` inside the sampling + coroutine, so callers must join the spawned process with a timeout and assert it + terminated. + """ + create_test_process_group() + loader = DistNeighborLoader( + dataset=dataset, + input_nodes=(_USER, dataset.node_ids[_USER]), # ty: ignore[invalid-argument-type, not-subscriptable] TODO(ty-torch-keyed-access): fix ty false positives for torch-backed keyed container access. + num_neighbors=[2, 2], + pin_memory_device=torch.device("cpu"), + ) + count = 0 + for datum in loader: + assert isinstance(datum, HeteroData) + count += 1 + holder["count"] = count + shutdown_rpc() + + class DistributedNeighborLoaderTest(TestCase): def setUp(self): super().setUp() @@ -861,5 +889,111 @@ def test_distributed_neighbor_loader_invalid_inputs_colocated( DistNeighborLoader(**kwargs) +class TestPartialFeatureCoverage(TestCase): + """Regression tests for heterogeneous datasets with features on a type subset. + + Without the per-type feature-fetch guard, ``_collate_fn`` fetched features for + every sampled node/edge type, so a featureless-but-reachable type raised a + ``KeyError`` inside the sampling coroutine that GLT swallowed, hanging the loader + forever. Each test therefore joins the spawned worker with a timeout and asserts + it terminated with batches. + """ + + def _build_partial_edge_feature_dataset(self) -> DistDataset: + """Two edge types both reachable from ``user`` seeds; only U2S has features.""" + partition_output = PartitionOutput( + node_partition_book={ + _USER: torch.zeros(5), + _STORY: torch.zeros(5), + }, + edge_partition_book={ + _USER_TO_USER: torch.zeros(5), + _USER_TO_STORY: torch.zeros(5), + }, + partitioned_edge_index={ + _USER_TO_USER: GraphPartitionData( + edge_index=torch.tensor([[0, 1, 2, 3, 4], [1, 2, 3, 4, 0]]), + edge_ids=None, + ), + _USER_TO_STORY: GraphPartitionData( + edge_index=torch.tensor([[0, 1, 2, 3, 4], [0, 1, 2, 3, 4]]), + edge_ids=None, + ), + }, + partitioned_node_features={ + _USER: FeaturePartitionData( + feats=torch.zeros(5, 2), ids=torch.arange(5) + ), + _STORY: FeaturePartitionData( + feats=torch.zeros(5, 2), ids=torch.arange(5) + ), + }, + # Edge features on U2S only; U2U is reachable but featureless. + partitioned_edge_features={ + _USER_TO_STORY: FeaturePartitionData( + feats=torch.zeros(5, 3), ids=torch.arange(5) + ), + }, + partitioned_positive_labels=None, + partitioned_negative_labels=None, + partitioned_node_labels=None, + ) + dataset = DistDataset(rank=0, world_size=1, edge_dir="out") + dataset.build(partition_output=partition_output) + return dataset + + def _build_partial_node_feature_dataset(self) -> DistDataset: + """Both node types reachable; only ``user`` has node features.""" + partition_output = PartitionOutput( + node_partition_book={ + _USER: torch.zeros(5), + _STORY: torch.zeros(5), + }, + edge_partition_book={ + _USER_TO_STORY: torch.zeros(5), + }, + partitioned_edge_index={ + _USER_TO_STORY: GraphPartitionData( + edge_index=torch.tensor([[0, 1, 2, 3, 4], [0, 1, 2, 3, 4]]), + edge_ids=None, + ), + }, + # Node features on ``user`` only; ``story`` is reachable but featureless. + partitioned_node_features={ + _USER: FeaturePartitionData( + feats=torch.zeros(5, 2), ids=torch.arange(5) + ), + }, + partitioned_edge_features=None, + partitioned_positive_labels=None, + partitioned_negative_labels=None, + partitioned_node_labels=None, + ) + dataset = DistDataset(rank=0, world_size=1, edge_dir="out") + dataset.build(partition_output=partition_output) + return dataset + + def _run_and_assert_batches(self, dataset: DistDataset) -> None: + holder = mp.Manager().dict() + proc = mp.get_context("spawn").Process( + target=_run_partial_feature_coverage_loader, + args=(0, dataset, holder), + ) + proc.start() + proc.join(timeout=120) + if proc.is_alive(): + proc.terminate() + proc.join() + self.fail("loader hung — partial feature coverage not handled") + self.assertEqual(proc.exitcode, 0, "loader worker exited with an error") + self.assertGreater(holder.get("count", 0), 0) + + def test_partial_edge_features_yields_batches(self) -> None: + self._run_and_assert_batches(self._build_partial_edge_feature_dataset()) + + def test_partial_node_features_yields_batches(self) -> None: + self._run_and_assert_batches(self._build_partial_node_feature_dataset()) + + if __name__ == "__main__": absltest.main() From 048014b35ccf21261e97e6f9249843f8e4c78453 Mon Sep 17 00:00:00 2001 From: kmontemayor Date: Wed, 8 Jul 2026 22:41:25 +0000 Subject: [PATCH 2/3] feat(distributed): normalize divergent cross-rank feature registration Direct DistPartitioner users may register different feature/label/weight type sets on different ranks. Because each type is partitioned via its own cross-rank collective, divergent registration desynchronizes those collectives and hangs or corrupts partitioning. Add per-category normalization helpers that all_gather each rank's {type: (trailing_shape, dtype)} schema and per-type local entity counts (via graphlearn_torch.distributed.rpc.all_gather, since torch.distributed is not generally initialized during partitioning), then per type in the global union: backfill an empty (0, dim) tensor on a rank missing it with zero local rows; raise if a rank is missing it but has nonzero local rows (unreconstructable); and raise on trailing-shape/dtype conflicts across ranks. All raise conditions are evaluated from the gathered data so every rank raises symmetrically instead of one raising while others hang. Normalization runs from partition() and from the individual partition_node_features_and_labels() / partition_edge_index_and_edge_features() methods (including the range partitioner override), guarded by per-category flags so the full partition() path does not gather twice. Also fix partition_node_features_and_labels() to partition the union of node feature and label types, so a label-only node type is no longer dropped when some other node type carries features. Adds tests for count-conditional backfill, nonzero-rows and dim-conflict raises, the node feature/label union, a zero-local-rows featured type, and a range-partitioner backfill. Co-Authored-By: Claude Fable 5 --- gigl/distributed/dist_partitioner.py | 249 +++++++++++- gigl/distributed/dist_range_partitioner.py | 4 + .../distributed_partitioner_test.py | 353 +++++++++++++++++- 3 files changed, 597 insertions(+), 9 deletions(-) diff --git a/gigl/distributed/dist_partitioner.py b/gigl/distributed/dist_partitioner.py index 2198557ba..619e3ea25 100644 --- a/gigl/distributed/dist_partitioner.py +++ b/gigl/distributed/dist_partitioner.py @@ -1,7 +1,7 @@ import gc import time from collections import abc, defaultdict -from typing import Callable, Optional, Tuple, Union +from typing import Callable, Optional, Tuple, TypeVar, Union import graphlearn_torch.distributed.rpc as glt_rpc import torch @@ -25,6 +25,9 @@ logger = Logger() +# Node or edge type key for the registration-normalization helpers. +_EntityType = TypeVar("_EntityType", NodeType, EdgeType) + class _DistLinkPredicitonPartitionManager(DistPartitionManager): """ @@ -184,6 +187,13 @@ def __init__( self._is_rpc_initialized: bool = False + # Guards against re-running the cross-rank registration-normalization + # collectives (see ``_normalize_node_feature_and_label_registration`` and + # ``_normalize_edge_feature_and_weight_registration``). Each category is + # normalized at most once per partitioner instance. + self._node_registration_normalized: bool = False + self._edge_registration_normalized: bool = False + self._is_input_homogeneous: Optional[bool] = None self._should_assign_edges_by_src_node: bool = should_assign_edges_by_src_node self._edge_types: list[EdgeType] = [] @@ -1412,6 +1422,214 @@ def _label_pfn(source_node_ids, _): return partitioned_label_edge_index + @staticmethod + def _local_feature_schema( + registered: Optional[dict[_EntityType, torch.Tensor]], + ) -> dict[_EntityType, tuple[tuple[int, ...], torch.dtype]]: + """Summarize a registered feature/label/weight dict as a per-type schema. + + Args: + registered: A registered ``{type: tensor}`` dict, or ``None`` if the + category was never registered on this rank. + + Returns: + A ``{type: (trailing_shape, dtype)}`` dict, where ``trailing_shape`` is + the tensor shape excluding the leading (per-entity) dimension. Empty if + ``registered`` is ``None``. + """ + if registered is None: + return {} + return { + entity_type: (tuple(tensor.shape[1:]), tensor.dtype) + for entity_type, tensor in registered.items() + } + + def _apply_registration_normalization( + self, + gathered: dict[str, tuple], + schema_index: int, + count_index: int, + feat_attr: str, + dim_attr: Optional[str], + category_name: str, + ) -> None: + """Backfill or reject a single feature/label/weight category using gathered schemas. + + For every type in the global union of registered types, validates that all + ranks that registered it agree on trailing shape and dtype, then either + backfills an empty tensor on a rank that is missing it with zero local rows, + or raises when a rank is missing it but has nonzero local rows (the values + cannot be reconstructed). All raise conditions are evaluated from the gathered + (schema, count) tuples, so every rank raises symmetrically rather than one rank + raising while the others hang on the next collective. + + Args: + gathered: Mapping of worker name to the all_gathered payload tuple. + schema_index: Index of this category's ``{type: (shape, dtype)}`` schema + within each payload tuple. + count_index: Index of the relevant per-type local-count dict within each + payload tuple (node counts for node categories, edge counts for edge + categories). + feat_attr: Name of the registration dict attribute to backfill into (e.g. + ``"_edge_feat"``). + dim_attr: Name of the companion dimension-dict attribute to backfill (e.g. + ``"_edge_feat_dim"``), or ``None`` for 1-D categories like edge weights. + category_name: Human-readable category label used in error messages. + """ + union_types: set[Union[NodeType, EdgeType]] = set() + for payload in gathered.values(): + union_types.update(payload[schema_index].keys()) + if not union_types: + return + + local_feats: Optional[dict[Union[NodeType, EdgeType], torch.Tensor]] = getattr( + self, feat_attr + ) + for entity_type in union_types: + distinct_schemas: set[tuple[tuple[int, ...], torch.dtype]] = set() + for payload in gathered.values(): + type_schema = payload[schema_index] + if entity_type in type_schema: + distinct_schemas.add(type_schema[entity_type]) + if len(distinct_schemas) > 1: + raise ValueError( + f"Inconsistent {category_name} schema for type {entity_type} " + f"across ranks: {sorted(str(s) for s in distinct_schemas)}. All " + f"ranks registering a type must agree on trailing shape and dtype." + ) + trailing_shape, dtype = next(iter(distinct_schemas)) + + # A rank missing the tensor but holding local rows cannot be backfilled + # (an empty tensor would desync from its entity ids). Detected from + # gathered counts so every rank raises, not just the offending one. + for payload in gathered.values(): + missing_rank = payload[0] + if entity_type not in payload[schema_index]: + local_count = payload[count_index].get(entity_type, 0) + if local_count > 0: + raise ValueError( + f"Type {entity_type} is globally registered with " + f"{category_name}, but rank {missing_rank} has " + f"{local_count} local rows of it and no {category_name} " + f"tensor; the missing values cannot be reconstructed. " + f"Register {category_name} for {entity_type} on that rank." + ) + + # This rank is missing the type (with zero local rows, per the check + # above) — backfill an empty tensor so its per-type collective still runs. + if local_feats is None or entity_type not in local_feats: + if local_feats is None: + local_feats = {} + setattr(self, feat_attr, local_feats) + local_feats[entity_type] = torch.empty((0, *trailing_shape), dtype=dtype) + if dim_attr is not None: + dim_map: Optional[dict[Union[NodeType, EdgeType], int]] = getattr( + self, dim_attr + ) + if dim_map is None: + dim_map = {} + setattr(self, dim_attr, dim_map) + dim_map[entity_type] = trailing_shape[0] if trailing_shape else 0 + + def _normalize_node_feature_and_label_registration(self) -> None: + """Make every rank hold the union of registered node feature/label types. + + Direct ``DistPartitioner`` users may register different node feature/label + type sets on different ranks. Because each type is partitioned via its own + cross-rank collective, divergent registration desynchronizes those + collectives (hang or corrupt output). This all_gathers each rank's + ``{type: (shape, dtype)}`` node feature and label schemas together with + per-type local node counts, then backfills empty tensors for missing types + with zero local rows and raises on nonzero-rows-without-tensor or shape/dtype + conflicts (see ``_apply_registration_normalization``). + + Idempotent: runs its collective at most once per partitioner instance. + + Uses GLT RPC (``glt_rpc.all_gather``), not ``torch.distributed``, which is not + generally initialized during partitioning. + """ + if self._node_registration_normalized: + return + self._assert_and_get_rpc_setup() + self._node_registration_normalized = True + if self._world_size <= 1: + return + + node_counts = { + node_type: node_ids.size(0) + for node_type, node_ids in (self._node_ids or {}).items() + } + gathered: dict[str, tuple] = glt_rpc.all_gather( + ( + self._rank, + node_counts, + self._local_feature_schema(self._node_feat), + self._local_feature_schema(self._node_labels), + ) + ) + self._apply_registration_normalization( + gathered=gathered, + schema_index=2, + count_index=1, + feat_attr="_node_feat", + dim_attr="_node_feat_dim", + category_name="node features", + ) + self._apply_registration_normalization( + gathered=gathered, + schema_index=3, + count_index=1, + feat_attr="_node_labels", + dim_attr="_node_labels_dim", + category_name="node labels", + ) + + def _normalize_edge_feature_and_weight_registration(self) -> None: + """Make every rank hold the union of registered edge feature/weight types. + + The edge-side counterpart of + ``_normalize_node_feature_and_label_registration``: all_gathers each rank's + edge feature and edge weight schemas plus per-type local edge counts, then + backfills or raises per type (see ``_apply_registration_normalization``). + + Idempotent: runs its collective at most once per partitioner instance. + """ + if self._edge_registration_normalized: + return + self._assert_and_get_rpc_setup() + self._edge_registration_normalized = True + if self._world_size <= 1: + return + + edge_counts = { + edge_type: edge_index.size(1) + for edge_type, edge_index in (self._edge_index or {}).items() + } + gathered: dict[str, tuple] = glt_rpc.all_gather( + ( + self._rank, + edge_counts, + self._local_feature_schema(self._edge_feat), + self._local_feature_schema(self._edge_weights), + ) + ) + self._apply_registration_normalization( + gathered=gathered, + schema_index=2, + count_index=1, + feat_attr="_edge_feat", + dim_attr="_edge_feat_dim", + category_name="edge features", + ) + self._apply_registration_normalization( + gathered=gathered, + schema_index=3, + count_index=1, + feat_attr="_edge_weights", + dim_attr=None, + category_name="edge weights", + ) + def partition_node(self) -> Union[PartitionBook, dict[NodeType, PartitionBook]]: """ Partitions nodes of a graph. If heterogeneous, partitions nodes for all node types. @@ -1484,6 +1702,10 @@ def partition_node_features_and_labels( self._assert_and_get_rpc_setup() + # Normalize divergent cross-rank node feature/label registration before + # partitioning, so each per-type collective below runs on every rank. + self._normalize_node_feature_and_label_registration() + logger.info("Partitioning Node Features and Labels ...") start_time = time.time() @@ -1502,20 +1724,22 @@ def partition_node_features_and_labels( input_entity=self._node_ids, is_node_entity=True, is_subset=False ) + # Partition the UNION of node types that have features or labels. A + # label-only node type must still be partitioned even when other node types + # have features, so we cannot select feature keys in an ``if`` and labels + # only in an ``elif``. node_feature_types: set[NodeType] = set() if self._node_feat is not None: self._assert_data_type_consistency( input_entity=self._node_feat, is_node_entity=True, is_subset=True ) - for node_type in self._node_feat.keys(): - node_feature_types.add(node_type) - elif self._node_labels is not None: + node_feature_types.update(self._node_feat.keys()) + if self._node_labels is not None: self._assert_data_type_consistency( input_entity=self._node_labels, is_node_entity=True, is_subset=True ) - for node_type in self._node_labels.keys(): - node_feature_types.add(node_type) - else: + node_feature_types.update(self._node_labels.keys()) + if not node_feature_types: raise ValueError( "Node features or labels must be registered prior to partitioning." ) @@ -1584,6 +1808,10 @@ def partition_edge_index_and_edge_features( self._assert_and_get_rpc_setup() + # Normalize divergent cross-rank edge feature/weight registration before + # partitioning, so each per-type collective below runs on every rank. + self._normalize_edge_feature_and_weight_registration() + assert ( self._edge_index is not None and self._edge_ids is not None @@ -1755,6 +1983,13 @@ def partition( self._assert_and_get_rpc_setup() + # Normalize divergent cross-rank feature/label/weight registration up front, + # while node ids and edge index (needed for local counts) are still + # registered. The per-category guard flags stop the individual partition_* + # methods below from gathering a second time. + self._normalize_edge_feature_and_weight_registration() + self._normalize_node_feature_and_label_registration() + logger.info(f"Rank {self._rank} starting partitioning ...") start_time = time.time() diff --git a/gigl/distributed/dist_range_partitioner.py b/gigl/distributed/dist_range_partitioner.py index faa1cfb57..64d969ef0 100644 --- a/gigl/distributed/dist_range_partitioner.py +++ b/gigl/distributed/dist_range_partitioner.py @@ -392,6 +392,10 @@ def partition_edge_index_and_edge_features( self._assert_and_get_rpc_setup() + # Normalize divergent cross-rank edge feature/weight registration before + # partitioning, so each per-type collective below runs on every rank. + self._normalize_edge_feature_and_weight_registration() + assert self._edge_index is not None and self._num_edges is not None, ( "Must have registered edges prior to partitioning them" ) diff --git a/tests/unit/distributed/distributed_partitioner_test.py b/tests/unit/distributed/distributed_partitioner_test.py index 0b02b8e2b..eccbdd31e 100644 --- a/tests/unit/distributed/distributed_partitioner_test.py +++ b/tests/unit/distributed/distributed_partitioner_test.py @@ -1,12 +1,13 @@ # Originally taken from https://github.com/alibaba/graphlearn-for-pytorch/blob/main/test/python/test_dist_random_partitioner.py +import traceback from collections import abc, defaultdict from typing import Iterable, Literal, MutableMapping, Optional, Tuple, Type, Union, cast import torch import torch.multiprocessing as mp from absl.testing import absltest -from graphlearn_torch.distributed import init_rpc, init_worker_group +from graphlearn_torch.distributed import init_rpc, init_worker_group, shutdown_rpc from graphlearn_torch.partition import PartitionBook from parameterized import param, parameterized from torch.multiprocessing import Manager @@ -15,7 +16,7 @@ from gigl.distributed.utils import get_process_group_name from gigl.distributed.utils.networking import get_free_port from gigl.distributed.utils.partition_book import get_ids_on_rank -from gigl.src.common.types.graph_data import EdgeType, NodeType +from gigl.src.common.types.graph_data import EdgeType, NodeType, Relation from gigl.types.graph import FeaturePartitionData, GraphPartitionData, PartitionOutput from tests.test_assets.distributed.constants import ( EDGE_TYPE_TO_FEATURE_DIMENSION_MAP, @@ -1433,5 +1434,353 @@ def test_heterogeneous_re_registration(self) -> None: partitioner3.register_node_labels(node_labels=node_labels) +_NORM_USER = NodeType("user") +_NORM_ITEM = NodeType("item") +_NORM_EDGE_A = EdgeType(_NORM_USER, Relation("a"), _NORM_USER) +_NORM_EDGE_B = EdgeType(_NORM_USER, Relation("b"), _NORM_USER) +_NORM_USER_TO_ITEM = EdgeType(_NORM_USER, Relation("to"), _NORM_ITEM) + + +def _run_divergent_edge_feature_registration( + rank: int, + world_size: int, + master_addr: str, + master_port: int, + output_dict: MutableMapping[int, dict], + scenario: str, + partitioner_class: Type[DistPartitioner], +) -> None: + """Partition a graph whose ranks register different edge-feature type sets. + + Both ranks register topology for edge types A and B (topology is assumed + consistent). Edge features on A are registered on both ranks; edge features on + B diverge per ``scenario`` to exercise the cross-rank registration normalization. + """ + init_worker_group( + world_size=world_size, rank=rank, group_name=get_process_group_name(0) + ) + init_rpc(master_addr=master_addr, master_port=master_port, num_rpc_threads=4) + try: + local_node_ids = torch.tensor([0, 1]) if rank == 0 else torch.tensor([2, 3]) + node_ids: dict[NodeType, torch.Tensor] = {_NORM_USER: local_node_ids} + edge_index: dict[EdgeType, torch.Tensor] = { + _NORM_EDGE_A: ( + torch.tensor([[0], [1]]) if rank == 0 else torch.tensor([[2], [3]]) + ), + } + if scenario == "backfill": + # Rank 1 has zero local B edges, so a missing B feature tensor is + # backfillable. Rank 0 holds all B edges, with source nodes spanning + # both ranks' node ranges so that after partitioning each rank owns some + # B edges (range partitioning cannot represent an empty per-rank edge + # range). + edge_index[_NORM_EDGE_B] = ( + torch.tensor([[0, 1, 2, 3], [1, 2, 3, 0]]) + if rank == 0 + else torch.empty((2, 0), dtype=torch.int64) + ) + else: + # Both ranks hold local B edges. + edge_index[_NORM_EDGE_B] = ( + torch.tensor([[0], [1]]) if rank == 0 else torch.tensor([[2], [3]]) + ) + + node_features: dict[NodeType, torch.Tensor] = { + _NORM_USER: torch.rand(local_node_ids.size(0), 2) + } + edge_features: dict[EdgeType, torch.Tensor] = { + _NORM_EDGE_A: torch.rand(edge_index[_NORM_EDGE_A].size(1), 3) + } + if scenario == "backfill": + if rank == 0: + edge_features[_NORM_EDGE_B] = torch.rand( + edge_index[_NORM_EDGE_B].size(1), 4 + ) + # Rank 1 omits B features (zero local B edges -> backfilled). + elif scenario == "nonzero_rows_raises": + if rank == 0: + edge_features[_NORM_EDGE_B] = torch.rand(1, 4) + # Rank 1 omits B features but has one local B edge -> must raise. + elif scenario == "dim_conflict": + edge_features[_NORM_EDGE_B] = ( + torch.rand(1, 4) if rank == 0 else torch.rand(1, 8) + ) + + partitioner = partitioner_class( + should_assign_edges_by_src_node=True, + node_ids=node_ids, + edge_index=edge_index, + node_features=node_features, + edge_features=edge_features, + ) + output = partitioner.partition() + partitioned_edge_features = output.partitioned_edge_features + assert isinstance(partitioned_edge_features, abc.Mapping) + output_dict[rank] = { + "status": "ok", + "edge_feature_types": sorted( + str(edge_type) for edge_type in partitioned_edge_features.keys() + ), + } + except Exception as exception: + output_dict[rank] = { + "status": "error", + "error": f"{type(exception).__name__}: {exception}", + "traceback": traceback.format_exc(), + } + finally: + try: + shutdown_rpc() + except Exception: + pass + + +def _run_node_feature_and_label_union( + rank: int, + master_addr: str, + master_port: int, + output_dict: MutableMapping[int, dict], +) -> None: + """Partition a graph with feature-only and label-only node types. + + ``user`` has node features but no labels; ``item`` has node labels but no + features. Both must be partitioned — a label-only node type cannot be dropped + just because some other node type carries features. + """ + init_worker_group(world_size=1, rank=rank, group_name=get_process_group_name(0)) + init_rpc(master_addr=master_addr, master_port=master_port, num_rpc_threads=4) + try: + partitioner = DistPartitioner( + should_assign_edges_by_src_node=True, + node_ids={_NORM_USER: torch.arange(3), _NORM_ITEM: torch.arange(3)}, + edge_index={_NORM_USER_TO_ITEM: torch.tensor([[0, 1, 2], [0, 1, 2]])}, + node_features={_NORM_USER: torch.rand(3, 2)}, + node_labels={_NORM_ITEM: torch.randint(0, 2, (3, 1))}, + ) + output = partitioner.partition() + partitioned_node_features = output.partitioned_node_features + partitioned_node_labels = output.partitioned_node_labels + assert isinstance(partitioned_node_features, abc.Mapping) + assert isinstance(partitioned_node_labels, abc.Mapping) + output_dict[rank] = { + "status": "ok", + "feature_types": sorted(str(t) for t in partitioned_node_features.keys()), + "label_types": sorted(str(t) for t in partitioned_node_labels.keys()), + } + except Exception as exception: + output_dict[rank] = { + "status": "error", + "error": f"{type(exception).__name__}: {exception}", + "traceback": traceback.format_exc(), + } + finally: + try: + shutdown_rpc() + except Exception: + pass + + +def _run_zero_local_rows_featured_type( + rank: int, + world_size: int, + master_addr: str, + master_port: int, + output_dict: MutableMapping[int, dict], +) -> None: + """Partition where one rank holds zero edges of a consistently-featured type. + + Both ranks register edge features for type B, but rank 1 has zero local B edges + and registers an empty ``(0, dim)`` B feature tensor (mirroring the TFRecord + zero-shard path). The partitioner must still emit B in the partitioned edge + feature dict on every rank, so the featured type stays a feature-store key even + on a rank that received no rows of it — this is the cross-rank consistency the + per-type feature-fetch guard relies on. + """ + init_worker_group( + world_size=world_size, rank=rank, group_name=get_process_group_name(0) + ) + init_rpc(master_addr=master_addr, master_port=master_port, num_rpc_threads=4) + try: + local_node_ids = torch.tensor([0, 1]) if rank == 0 else torch.tensor([2, 3]) + node_ids: dict[NodeType, torch.Tensor] = {_NORM_USER: local_node_ids} + edge_index = { + _NORM_EDGE_A: ( + torch.tensor([[0], [1]]) if rank == 0 else torch.tensor([[2], [3]]) + ), + _NORM_EDGE_B: ( + torch.tensor([[0, 1], [1, 0]]) + if rank == 0 + else torch.empty((2, 0), dtype=torch.int64) + ), + } + edge_features = { + _NORM_EDGE_A: torch.rand(edge_index[_NORM_EDGE_A].size(1), 3), + # Rank 1 registers an empty B feature tensor rather than omitting it. + _NORM_EDGE_B: ( + torch.rand(2, 4) if rank == 0 else torch.empty((0, 4)) + ), + } + partitioner = DistPartitioner( + should_assign_edges_by_src_node=True, + node_ids=node_ids, + edge_index=edge_index, + node_features={_NORM_USER: torch.rand(local_node_ids.size(0), 2)}, + edge_features=edge_features, + ) + output = partitioner.partition() + partitioned_edge_features = output.partitioned_edge_features + assert isinstance(partitioned_edge_features, abc.Mapping) + output_dict[rank] = { + "status": "ok", + "edge_feature_types": sorted( + str(edge_type) for edge_type in partitioned_edge_features.keys() + ), + } + except Exception as exception: + output_dict[rank] = { + "status": "error", + "error": f"{type(exception).__name__}: {exception}", + "traceback": traceback.format_exc(), + } + finally: + try: + shutdown_rpc() + except Exception: + pass + + +class TestRegistrationNormalization(TestCase): + """Cross-rank normalization of divergent feature/label/weight registration.""" + + def setUp(self) -> None: + self._master_ip_address = "localhost" + + def test_missing_edge_feature_type_zero_rows_backfilled(self) -> None: + master_port = get_free_port() + manager = Manager() + output_dict: MutableMapping[int, dict] = manager.dict() + mp.spawn( + _run_divergent_edge_feature_registration, + args=( + MOCKED_NUM_PARTITIONS, + self._master_ip_address, + master_port, + output_dict, + "backfill", + DistPartitioner, + ), + nprocs=MOCKED_NUM_PARTITIONS, + join=True, + ) + for rank in range(MOCKED_NUM_PARTITIONS): + result = output_dict[rank] + self.assertEqual(result["status"], "ok", result.get("error")) + self.assertIn(str(_NORM_EDGE_A), result["edge_feature_types"]) + self.assertIn(str(_NORM_EDGE_B), result["edge_feature_types"]) + + def test_missing_edge_feature_type_zero_rows_backfilled_range(self) -> None: + master_port = get_free_port() + manager = Manager() + output_dict: MutableMapping[int, dict] = manager.dict() + mp.spawn( + _run_divergent_edge_feature_registration, + args=( + MOCKED_NUM_PARTITIONS, + self._master_ip_address, + master_port, + output_dict, + "backfill", + DistRangePartitioner, + ), + nprocs=MOCKED_NUM_PARTITIONS, + join=True, + ) + for rank in range(MOCKED_NUM_PARTITIONS): + result = output_dict[rank] + self.assertEqual(result["status"], "ok", result.get("error")) + self.assertIn(str(_NORM_EDGE_B), result["edge_feature_types"]) + + def test_missing_edge_feature_type_nonzero_rows_raises(self) -> None: + master_port = get_free_port() + manager = Manager() + output_dict: MutableMapping[int, dict] = manager.dict() + mp.spawn( + _run_divergent_edge_feature_registration, + args=( + MOCKED_NUM_PARTITIONS, + self._master_ip_address, + master_port, + output_dict, + "nonzero_rows_raises", + DistPartitioner, + ), + nprocs=MOCKED_NUM_PARTITIONS, + join=True, + ) + for rank in range(MOCKED_NUM_PARTITIONS): + result = output_dict[rank] + self.assertEqual(result["status"], "error") + self.assertIn(str(_NORM_EDGE_B), result["error"]) + + def test_edge_feature_dim_conflict_raises(self) -> None: + master_port = get_free_port() + manager = Manager() + output_dict: MutableMapping[int, dict] = manager.dict() + mp.spawn( + _run_divergent_edge_feature_registration, + args=( + MOCKED_NUM_PARTITIONS, + self._master_ip_address, + master_port, + output_dict, + "dim_conflict", + DistPartitioner, + ), + nprocs=MOCKED_NUM_PARTITIONS, + join=True, + ) + for rank in range(MOCKED_NUM_PARTITIONS): + result = output_dict[rank] + self.assertEqual(result["status"], "error") + self.assertIn(str(_NORM_EDGE_B), result["error"]) + + def test_node_feature_and_label_union_partitioned(self) -> None: + master_port = get_free_port() + manager = Manager() + output_dict: MutableMapping[int, dict] = manager.dict() + mp.spawn( + _run_node_feature_and_label_union, + args=(self._master_ip_address, master_port, output_dict), + nprocs=1, + join=True, + ) + result = output_dict[0] + self.assertEqual(result["status"], "ok", result.get("error")) + self.assertIn(str(_NORM_USER), result["feature_types"]) + self.assertNotIn(str(_NORM_ITEM), result["feature_types"]) + self.assertIn(str(_NORM_ITEM), result["label_types"]) + self.assertNotIn(str(_NORM_USER), result["label_types"]) + + def test_zero_local_rows_featured_type_survives(self) -> None: + master_port = get_free_port() + manager = Manager() + output_dict: MutableMapping[int, dict] = manager.dict() + mp.spawn( + _run_zero_local_rows_featured_type, + args=( + MOCKED_NUM_PARTITIONS, + self._master_ip_address, + master_port, + output_dict, + ), + nprocs=MOCKED_NUM_PARTITIONS, + join=True, + ) + for rank in range(MOCKED_NUM_PARTITIONS): + result = output_dict[rank] + self.assertEqual(result["status"], "ok", result.get("error")) + self.assertIn(str(_NORM_EDGE_B), result["edge_feature_types"]) + + if __name__ == "__main__": absltest.main() From addf7f39d32bc3ed5a19997fb484b3e0ce929640 Mon Sep 17 00:00:00 2001 From: kmontemayor Date: Tue, 14 Jul 2026 17:02:03 +0000 Subject: [PATCH 3/3] style(distributed): apply formatter Co-Authored-By: Claude Fable 5 --- gigl/distributed/dist_partitioner.py | 4 +++- tests/unit/distributed/distributed_partitioner_test.py | 4 +--- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/gigl/distributed/dist_partitioner.py b/gigl/distributed/dist_partitioner.py index 619e3ea25..375dcbabc 100644 --- a/gigl/distributed/dist_partitioner.py +++ b/gigl/distributed/dist_partitioner.py @@ -1521,7 +1521,9 @@ def _apply_registration_normalization( if local_feats is None: local_feats = {} setattr(self, feat_attr, local_feats) - local_feats[entity_type] = torch.empty((0, *trailing_shape), dtype=dtype) + local_feats[entity_type] = torch.empty( + (0, *trailing_shape), dtype=dtype + ) if dim_attr is not None: dim_map: Optional[dict[Union[NodeType, EdgeType], int]] = getattr( self, dim_attr diff --git a/tests/unit/distributed/distributed_partitioner_test.py b/tests/unit/distributed/distributed_partitioner_test.py index eccbdd31e..233d93ed3 100644 --- a/tests/unit/distributed/distributed_partitioner_test.py +++ b/tests/unit/distributed/distributed_partitioner_test.py @@ -1616,9 +1616,7 @@ def _run_zero_local_rows_featured_type( edge_features = { _NORM_EDGE_A: torch.rand(edge_index[_NORM_EDGE_A].size(1), 3), # Rank 1 registers an empty B feature tensor rather than omitting it. - _NORM_EDGE_B: ( - torch.rand(2, 4) if rank == 0 else torch.empty((0, 4)) - ), + _NORM_EDGE_B: (torch.rand(2, 4) if rank == 0 else torch.empty((0, 4))), } partitioner = DistPartitioner( should_assign_edges_by_src_node=True,