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/gigl/distributed/dist_partitioner.py b/gigl/distributed/dist_partitioner.py index 2198557ba..375dcbabc 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,216 @@ 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 +1704,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 +1726,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 +1810,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 +1985,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/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() diff --git a/tests/unit/distributed/distributed_partitioner_test.py b/tests/unit/distributed/distributed_partitioner_test.py index 0b02b8e2b..233d93ed3 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,351 @@ 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()