diff --git a/examples/link_prediction/graph_store/heterogeneous_training.py b/examples/link_prediction/graph_store/heterogeneous_training.py index 2be34e608..2c772236b 100644 --- a/examples/link_prediction/graph_store/heterogeneous_training.py +++ b/examples/link_prediction/graph_store/heterogeneous_training.py @@ -211,6 +211,7 @@ def _setup_dataloaders( channel_size=sampling_worker_shared_channel_size, process_start_gap_seconds=process_start_gap_seconds, shuffle=shuffle, + use_edge_index_output=True, ) print(f"---Rank {rank} finished setting up main loader for split={split}") @@ -294,23 +295,16 @@ def _compute_loss( ) # Extracting local query, random negative, positive, hard_negative, and random_negative indices. - query_node_idx: torch.Tensor = torch.arange( - main_data[query_node_type].batch_size - ).to(device) + query_node_idx = torch.arange(main_data[query_node_type].batch_size, device=device) random_negative_batch_size = random_negative_data[labeled_node_type].batch_size - positive_idx: torch.Tensor = torch.cat(list(main_data.y_positive.values())).to( - device - ) - repeated_query_node_idx = query_node_idx.repeat_interleave( - torch.tensor([len(v) for v in main_data.y_positive.values()]).to(device) - ) + positive_label_edge_index: torch.Tensor = main_data.y_positive + repeated_query_node_idx = query_node_idx[positive_label_edge_index[0]] + positive_idx = positive_label_edge_index[1] if hasattr(main_data, "y_negative"): - hard_negative_idx: torch.Tensor = torch.cat( - list(main_data.y_negative.values()) - ).to(device) + hard_negative_idx: torch.Tensor = main_data.y_negative[1] else: - hard_negative_idx = torch.empty(0, dtype=torch.long).to(device) + hard_negative_idx = torch.empty(0, dtype=torch.long, device=device) # Use local IDs to get the corresponding embeddings in the tensors diff --git a/examples/link_prediction/graph_store/homogeneous_training.py b/examples/link_prediction/graph_store/homogeneous_training.py index 2d4c22788..25ece94a0 100644 --- a/examples/link_prediction/graph_store/homogeneous_training.py +++ b/examples/link_prediction/graph_store/homogeneous_training.py @@ -241,6 +241,7 @@ def _setup_dataloaders( channel_size=sampling_worker_shared_channel_size, process_start_gap_seconds=process_start_gap_seconds, shuffle=shuffle, + use_edge_index_output=True, ) logger.info(f"---Rank {rank} finished setting up main loader for split={split}") @@ -302,21 +303,16 @@ def _compute_loss( main_embeddings = model(data=main_data, device=device) random_negative_embeddings = model(data=random_negative_data, device=device) - query_node_idx: torch.Tensor = torch.arange(main_data.batch_size).to(device) + query_node_idx = torch.arange(main_data.batch_size, device=device) random_negative_batch_size = random_negative_data.batch_size - positive_idx: torch.Tensor = torch.cat(list(main_data.y_positive.values())).to( - device - ) - repeated_query_node_idx = query_node_idx.repeat_interleave( - torch.tensor([len(v) for v in main_data.y_positive.values()]).to(device) - ) + positive_label_edge_index: torch.Tensor = main_data.y_positive + repeated_query_node_idx = query_node_idx[positive_label_edge_index[0]] + positive_idx = positive_label_edge_index[1] if hasattr(main_data, "y_negative"): - hard_negative_idx: torch.Tensor = torch.cat( - list(main_data.y_negative.values()) - ).to(device) + hard_negative_idx: torch.Tensor = main_data.y_negative[1] else: - hard_negative_idx = torch.empty(0, dtype=torch.long).to(device) + hard_negative_idx = torch.empty(0, dtype=torch.long, device=device) # Use local IDs to get the corresponding embeddings in the tensors diff --git a/examples/link_prediction/heterogeneous_training.py b/examples/link_prediction/heterogeneous_training.py index 8ed672b7c..e41a19191 100644 --- a/examples/link_prediction/heterogeneous_training.py +++ b/examples/link_prediction/heterogeneous_training.py @@ -144,6 +144,7 @@ def _setup_dataloaders( # This is done so that each process on the current machine which initializes a `main_loader` doesn't compete for memory, causing potential OOM process_start_gap_seconds=process_start_gap_seconds, shuffle=shuffle, + use_edge_index_output=True, ) logger.info(f"---Rank {rank} finished setting up main loader") @@ -218,25 +219,16 @@ def _compute_loss( # Local in this case refers to the local index in the batch, while global subsequently refers to the node's unique global ID across all nodes in the dataset. # Global ids are stored in data.node, ex. `data.node = [50, 20, 10]` where the `0` is the local index for global id `50`. Note that each global id in data.node is unique. # TODO (mkolodner-sc): If GLT migrates towards using PyG's `.n_id` field instead of `.node`, update this code. - query_node_idx: torch.Tensor = torch.arange( - main_data[query_node_type].batch_size - ).to(device) + query_node_idx = torch.arange(main_data[query_node_type].batch_size, device=device) random_negative_batch_size = random_negative_data[labeled_node_type].batch_size - # main_data.y_positive is a dict[query_node_local_index: int, labeled_node_local_indices: torch.Tensor], even in the heterogeneous setting. - positive_idx: torch.Tensor = torch.cat(list(main_data.y_positive.values())).to( - device - ) - # We also extract a repeated query node index tensor which upsamples each query node based on the number of positives it has - repeated_query_node_idx = query_node_idx.repeat_interleave( - torch.tensor([len(v) for v in main_data.y_positive.values()]).to(device) - ) + positive_label_edge_index: torch.Tensor = main_data.y_positive + repeated_query_node_idx = query_node_idx[positive_label_edge_index[0]] + positive_idx = positive_label_edge_index[1] if hasattr(main_data, "y_negative"): - hard_negative_idx: torch.Tensor = torch.cat( - list(main_data.y_negative.values()) - ).to(device) + hard_negative_idx: torch.Tensor = main_data.y_negative[1] else: - hard_negative_idx = torch.empty(0, dtype=torch.long).to(device) + hard_negative_idx = torch.empty(0, dtype=torch.long, device=device) # Use local IDs to get the corresponding embeddings in the tensors diff --git a/examples/link_prediction/homogeneous_training.py b/examples/link_prediction/homogeneous_training.py index b95a77489..a84b55114 100644 --- a/examples/link_prediction/homogeneous_training.py +++ b/examples/link_prediction/homogeneous_training.py @@ -134,6 +134,7 @@ def _setup_dataloaders( # This is done so that each process on the current machine which initializes a `main_loader` doesn't compete for memory, causing potential OOM process_start_gap_seconds=process_start_gap_seconds, shuffle=shuffle, + use_edge_index_output=True, ) logger.info(f"---Rank {rank} finished setting up main loader") @@ -187,23 +188,17 @@ def _compute_loss( # Extracting local query, random negative, positive, hard_negative, and random_negative indices. # Local in this case refers to the local index in the batch, while global subsequently refers to the node's unique global ID across all nodes in the dataset. # Global ids are stored in data.node, ex. `data.node = [50, 20, 10]` where the `0` is the local index for global id `50`. Note that each global id in data.node is unique. - query_node_idx: torch.Tensor = torch.arange(main_data.batch_size).to(device) + query_node_idx = torch.arange(main_data.batch_size, device=device) random_negative_batch_size = random_negative_data.batch_size - # main_data.y_positive is a dict[query_node_local_index: int, labeled_node_local_indices: torch.Tensor] - positive_idx: torch.Tensor = torch.cat(list(main_data.y_positive.values())).to( - device - ) - # We also extract a repeated query node index tensor which upsamples each query node based on the number of positives it has - repeated_query_node_idx = query_node_idx.repeat_interleave( - torch.tensor([len(v) for v in main_data.y_positive.values()]).to(device) - ) + # Label edge indices are [anchor_local_id, label_local_id] pairs. + positive_label_edge_index: torch.Tensor = main_data.y_positive + repeated_query_node_idx = query_node_idx[positive_label_edge_index[0]] + positive_idx = positive_label_edge_index[1] if hasattr(main_data, "y_negative"): - hard_negative_idx: torch.Tensor = torch.cat( - list(main_data.y_negative.values()) - ).to(device) + hard_negative_idx: torch.Tensor = main_data.y_negative[1] else: - hard_negative_idx = torch.empty(0, dtype=torch.long).to(device) + hard_negative_idx = torch.empty(0, dtype=torch.long, device=device) # Use local IDs to get the corresponding embeddings in the tensors diff --git a/gigl/distributed/dist_ablp_neighborloader.py b/gigl/distributed/dist_ablp_neighborloader.py index 50f42f5a9..8d220b753 100644 --- a/gigl/distributed/dist_ablp_neighborloader.py +++ b/gigl/distributed/dist_ablp_neighborloader.py @@ -1,6 +1,7 @@ -from collections import abc, defaultdict +import warnings +from collections import abc from itertools import count -from typing import Optional, Union +from typing import Literal, Optional, Union, cast import torch from graphlearn_torch.channel import SampleMessage @@ -27,6 +28,10 @@ SamplerOptions, resolve_sampler_options, ) +from gigl.distributed.utils.ablp import ( + label_edge_index_to_dict, + remap_labels_to_local_edge_indices, +) from gigl.distributed.utils.neighborloader import ( DatasetSchema, SamplingClusterSetup, @@ -43,9 +48,6 @@ from gigl.types.graph import ( DEFAULT_HOMOGENEOUS_EDGE_TYPE, DEFAULT_HOMOGENEOUS_NODE_TYPE, - label_edge_type_to_message_passing_edge_type, - message_passing_to_negative_label, - message_passing_to_positive_label, reverse_edge_type, select_label_edge_types, ) @@ -54,6 +56,182 @@ logger = Logger() +_GraphStoreSchemaSummary = tuple[ + NodeType, + tuple[EdgeType, ...], + tuple[EdgeType, ...], + tuple[EdgeType, ...], +] + + +def _is_integral_tensor(tensor: torch.Tensor) -> bool: + return tensor.dtype in ( + torch.uint8, + torch.int8, + torch.int16, + torch.int32, + torch.int64, + ) + + +def _convert_graph_store_ablp_inputs( + input_nodes: dict[int, ABLPInputNodes], + num_storage_nodes: int, + edge_types: Optional[list[EdgeType]], + edge_dir: Literal["in", "out"], +) -> tuple[ + list[ABLPNodeSamplerInput], + NodeType, + list[EdgeType], + list[EdgeType], + list[EdgeType], +]: + if not input_nodes: + raise ValueError("Graph Store ABLP input_nodes must not be empty.") + invalid_server_ranks = [ + rank + for rank in input_nodes + if not isinstance(rank, int) or rank < 0 or rank >= num_storage_nodes + ] + if invalid_server_ranks: + raise ValueError( + f"Graph Store server ranks must be in [0, {num_storage_nodes}), " + f"received {invalid_server_ranks}." + ) + if edge_types is None: + raise ValueError("Graph Store edge types must be available for ABLP.") + if edge_dir not in ("in", "out"): + raise ValueError(f"Unexpected Graph Store edge direction: {edge_dir}.") + + first_input = next(iter(input_nodes.values())) + input_type = first_input.anchor_node_type + supervision_edge_types = list(first_input.labels) + if not supervision_edge_types: + raise ValueError("Graph Store ABLP labels must not be empty.") + + anchor_endpoint = 0 if edge_dir == "out" else 2 + positive_label_edge_types: list[EdgeType] = [] + negative_label_edge_types: list[EdgeType] = [] + resolved_label_edge_types: dict[EdgeType, tuple[EdgeType, Optional[EdgeType]]] = {} + for supervision_edge_type in supervision_edge_types: + if supervision_edge_type[anchor_endpoint] != input_type: + raise ValueError( + f"Supervision edge type {supervision_edge_type} does not use " + f"anchor node type {input_type} at endpoint {anchor_endpoint}." + ) + positive_edge_type, negative_edge_type = select_label_edge_types( + supervision_edge_type, edge_types + ) + resolved_label_edge_types[supervision_edge_type] = ( + positive_edge_type, + negative_edge_type, + ) + positive_label_edge_types.append(positive_edge_type) + if negative_edge_type is not None: + negative_label_edge_types.append(negative_edge_type) + + expected_supervision_edge_types = set(supervision_edge_types) + converted_inputs_by_rank: dict[int, ABLPNodeSamplerInput] = {} + for server_rank, ablp_input in input_nodes.items(): + if ablp_input.anchor_node_type != input_type: + raise ValueError( + "Graph Store ABLP inputs must use one anchor node type, " + f"received {input_type} and {ablp_input.anchor_node_type}." + ) + if set(ablp_input.labels) != expected_supervision_edge_types: + raise ValueError( + "Graph Store ABLP inputs must use the same supervision edge types." + ) + + anchors = ablp_input.anchor_nodes + if ( + not isinstance(anchors, torch.Tensor) + or anchors.dim() != 1 + or not _is_integral_tensor(anchors) + ): + raise ValueError( + "Graph Store ABLP anchor nodes must be a 1-D integral tensor." + ) + + positive_labels_by_edge_type: dict[EdgeType, torch.Tensor] = {} + negative_labels_by_edge_type: dict[EdgeType, torch.Tensor] = {} + for supervision_edge_type in supervision_edge_types: + positive_labels, negative_labels = ablp_input.labels[supervision_edge_type] + if ( + not isinstance(positive_labels, torch.Tensor) + or positive_labels.dim() != 2 + or not _is_integral_tensor(positive_labels) + or positive_labels.size(0) != anchors.numel() + ): + raise ValueError( + "Graph Store positive labels must be a 2-D integral tensor " + "with one row per anchor." + ) + if negative_labels is not None and ( + not isinstance(negative_labels, torch.Tensor) + or negative_labels.dim() != 2 + or not _is_integral_tensor(negative_labels) + or negative_labels.size(0) != anchors.numel() + ): + raise ValueError( + "Graph Store negative labels must be a 2-D integral tensor " + "with one row per anchor." + ) + + positive_edge_type, negative_edge_type = resolved_label_edge_types[ + supervision_edge_type + ] + positive_labels_by_edge_type[positive_edge_type] = positive_labels + if negative_edge_type is None: + if negative_labels is not None: + raise ValueError( + f"Negative labels were provided for {supervision_edge_type} " + "without registered negative label topology." + ) + else: + if negative_labels is None: + if anchors.numel() != 0: + raise ValueError( + f"Negative labels are required for {supervision_edge_type}." + ) + negative_labels = torch.empty((0, 0), dtype=torch.long) + negative_labels_by_edge_type[negative_edge_type] = negative_labels + + converted_inputs_by_rank[server_rank] = ABLPNodeSamplerInput( + node=anchors, + input_type=input_type, + positive_label_by_edge_types=positive_labels_by_edge_type, + negative_label_by_edge_types=negative_labels_by_edge_type, + ) + + input_data: list[ABLPNodeSamplerInput] = [] + for server_rank in range(num_storage_nodes): + if server_rank in converted_inputs_by_rank: + input_data.append(converted_inputs_by_rank[server_rank]) + else: + input_data.append( + ABLPNodeSamplerInput( + node=torch.empty(0, dtype=torch.long), + input_type=input_type, + positive_label_by_edge_types={ + edge_type: torch.empty((0, 0), dtype=torch.long) + for edge_type in positive_label_edge_types + }, + negative_label_by_edge_types={ + edge_type: torch.empty((0, 0), dtype=torch.long) + for edge_type in negative_label_edge_types + }, + ) + ) + + return ( + input_data, + input_type, + supervision_edge_types, + positive_label_edge_types, + negative_label_edge_types, + ) + class DistABLPLoader(BaseDistLoader): # Counts instantiations of this class, per process. @@ -90,17 +268,19 @@ def __init__( local_process_rank: Optional[int] = None, # TODO: (svij) Deprecate this local_process_world_size: Optional[int] = None, # TODO: (svij) Deprecate this non_blocking_transfers: bool = True, + use_edge_index_output: bool = False, ): """ Neighbor loader for Anchor Based Link Prediction (ABLP) tasks. - Note that for this class, the dataset must *always* be heterogeneous, - as we need separate edge types for positive and negative labels. + The dataset must *always* be heterogeneous here, since positive and + negative labels are carried as separate edge types. By default, the loader will return {py:class} `torch_geometric.data.HeteroData` (heterogeneous) objects, but will return a {py:class}`torch_geometric.data.Data` (homogeneous) object if the dataset is "labeled homogeneous". - The following fields may also be present: + The following fields may also be present (this describes the deprecated + default shape; see ``use_edge_index_output`` for the tensor format): - `y_positive`: `dict[int, torch.Tensor]` mapping from local anchor node id to a tensor of positive label node ids. - `y_negative`: (Optional) `dict[int, torch.Tensor]` mapping from local anchor node id to a tensor of negative @@ -133,11 +313,23 @@ def __init__( - `y_positive`: {0: torch.tensor([1])} # 1 is the only positive label for node 0 - `y_negative`: {0: torch.tensor([2])} # 2 is the only negative label for node 0 - NOTE: both label fields will instead be `dict[EdgeType, dict[int, torch.Tensor]]` if multiple supervision edge types are provided. + NOTE: label fields are flattened only when exactly one supervision edge type + is provided. With multiple supervision types, both fields remain keyed by + edge type even when only one type has negative labels. e.g. if there are supervision edge types: (a, to, b) and (a, to, c), then the label fields could be: - `y_positive`: {(a, to, b): {0: torch.tensor([1])}, (a, to, c): {0: torch.tensor([2])}} - `y_negative`: {(a, to, b): {0: torch.tensor([3])}, (a, to, c): {0: torch.tensor([4])}} + With ``use_edge_index_output=True``, each label field is a ``[2, E]`` + tensor. Row 0 contains local anchor indices and row 1 contains local + label-node indices. For the example above: + + - ``y_positive = torch.tensor([[0], [1]])`` + - ``y_negative = torch.tensor([[0], [2]])`` + + Multiple supervision edge types produce ``dict[EdgeType, torch.Tensor]``. + Pair order within an anchor is unspecified. + Args: dataset (Union[DistDataset, RemoteDistDataset]): The dataset to sample from. If this is a `RemoteDistDataset`, then we are in "Graph Store" mode. @@ -213,11 +405,23 @@ def __init__( is used instead. See https://docs.pytorch.org/tutorials/intermediate/pinmem_nonblock.html for background on pinned memory and non-blocking transfers. + use_edge_index_output (bool): Return labels as ``[2, E]`` edge-index + tensors instead of the deprecated ragged dictionaries. Row 0 + contains local anchor indices and row 1 contains local label-node + indices. Defaults to ``False`` for backward compatibility. """ # Set self._shutdowned right away, that way if we throw here, and __del__ is called, # then we can properly clean up and don't get extraneous error messages. self._shutdowned = True + if not use_edge_index_output: + warnings.warn( + "The ragged dictionary ABLP label output is deprecated. Pass " + "use_edge_index_output=True to receive [2, E] label edge indices.", + FutureWarning, + stacklevel=2, + ) + self._use_edge_index_output = use_edge_index_output sampler_options = resolve_sampler_options(num_neighbors, sampler_options) @@ -612,9 +816,79 @@ def _setup_for_graph_store( Tuple of (list[ABLPNodeSamplerInput], RemoteDistSamplingWorkerOptions, DatasetSchema, backend_key). """ + edge_types = dataset.fetch_edge_types() + raw_edge_dir = dataset.fetch_edge_dir() + converted_inputs = None + local_error: Optional[str] = None + local_schema_summary: Optional[_GraphStoreSchemaSummary] = None + try: + if raw_edge_dir not in ("in", "out"): + raise ValueError( + f"Unexpected Graph Store edge direction: {raw_edge_dir}." + ) + edge_dir = cast(Literal["in", "out"], raw_edge_dir) + converted_inputs = _convert_graph_store_ablp_inputs( + input_nodes=input_nodes, + num_storage_nodes=dataset.cluster_info.num_storage_nodes, + edge_types=edge_types, + edge_dir=edge_dir, + ) + ( + _, + input_type, + supervision_edge_types, + positive_label_edge_types, + negative_label_edge_types, + ) = converted_inputs + if len(supervision_edge_types) != 1: + raise ValueError( + "Graph Store mode currently only supports a single supervision " + f"edge type, but received {supervision_edge_types}." + ) + local_schema_summary = ( + input_type, + tuple(supervision_edge_types), + tuple(positive_label_edge_types), + tuple(negative_label_edge_types), + ) + except ValueError as error: + local_error = str(error) + + gathered_validation: list[ + tuple[Optional[str], Optional[_GraphStoreSchemaSummary]] + ] = [(None, None)] * torch.distributed.get_world_size() + torch.distributed.all_gather_object( + gathered_validation, (local_error, local_schema_summary) + ) + rank_errors = [ + f"rank {rank}: {error}" + for rank, (error, _) in enumerate(gathered_validation) + if error is not None + ] + if rank_errors: + raise ValueError( + "Graph Store ABLP validation failed across compute ranks: " + + "; ".join(rank_errors) + ) + if any( + summary != gathered_validation[0][1] + for _, summary in gathered_validation[1:] + ): + raise ValueError( + "Graph Store ABLP schemas differ across compute ranks: " + f"{[(rank, summary) for rank, (_, summary) in enumerate(gathered_validation)]}" + ) + + assert converted_inputs is not None + ( + input_data, + input_type, + self._supervision_edge_types, + self._positive_label_edge_types, + self._negative_label_edge_types, + ) = converted_inputs node_feature_info = dataset.fetch_node_feature_info() edge_feature_info = dataset.fetch_edge_feature_info() - edge_types = dataset.fetch_edge_types() compute_rank = torch.distributed.get_rank() backend_key = f"dist_ablp_loader_{self._instance_count}" worker_key = f"{backend_key}_compute_rank_{compute_rank}" @@ -632,106 +906,20 @@ def _setup_for_graph_store( f"tcp://{worker_options.master_addr}:{worker_options.master_port}" ) - # Validate server ranks - servers = input_nodes.keys() - if len(servers) > 0: - if ( - max(servers) >= dataset.cluster_info.num_storage_nodes - or min(servers) < 0 - ): - raise ValueError( - f"When using Graph Store mode, the server ranks must be in range " - f"[0, {dataset.cluster_info.num_storage_nodes}), " - f"received inputs for servers: {list(servers)}" - ) - - # Extract node type and label edge types from the ABLPInputNodes dataclass. - # All entries should have the same anchor_node_type and edge type keys. - first_input = next(iter(input_nodes.values())) - input_type = first_input.anchor_node_type is_homogeneous_with_labeled_edge_type = ( input_type == DEFAULT_HOMOGENEOUS_NODE_TYPE ) - # Extract supervision edge types and derive label edge types from the - # ABLPInputNodes.labels dict (keyed by supervision edge type). - self._supervision_edge_types = list(first_input.labels.keys()) - has_negatives = False - for ablp_input in input_nodes.values(): - for maybe_negative_labels in ablp_input.labels.values(): - if maybe_negative_labels is not None: - has_negatives = True - break - - self._positive_label_edge_types = [ - message_passing_to_positive_label(et) for et in self._supervision_edge_types - ] - self._negative_label_edge_types = ( - [ - message_passing_to_negative_label(et) - for et in self._supervision_edge_types - ] - if has_negatives - else [] - ) - - # Graph Store mode currently only supports a single supervision edge type, - # so the labels dict must have exactly 1 entry. - if len(self._supervision_edge_types) != 1: - raise ValueError( - f"Graph Store mode currently only supports a single supervision edge type, " - f"but ABLPInputNodes.labels has {len(self._supervision_edge_types)} " - f"entries: {self._supervision_edge_types}" - ) - logger.info(f"Positive label edge types: {self._positive_label_edge_types}") logger.info(f"Negative label edge types: {self._negative_label_edge_types}") - # Convert from ABLPInputNodes to list of ABLPNodeSamplerInput (one per server) - input_data: list[ABLPNodeSamplerInput] = [] - for server_rank in range(dataset.cluster_info.num_storage_nodes): - positive_label_by_edge_type: dict[EdgeType, torch.Tensor] = {} - negative_label_by_edge_type: dict[EdgeType, torch.Tensor] = {} - if server_rank in input_nodes: - ablp_input_nodes = input_nodes[server_rank] - anchors = ablp_input_nodes.anchor_nodes - for supervision_edge_type, ( - positive_labels, - negative_labels, - ) in ablp_input_nodes.labels.items(): - positive_label_by_edge_type[ - message_passing_to_positive_label(supervision_edge_type) - ] = positive_labels - if negative_labels is not None: - negative_label_by_edge_type[ - message_passing_to_negative_label(supervision_edge_type) - ] = negative_labels - else: - # Empty input for servers with no data for this rank - anchors = torch.empty(0, dtype=torch.long) - positive_label_by_edge_type = { - et: torch.empty(0, 0, dtype=torch.long) - for et in self._positive_label_edge_types - } - if has_negatives: - negative_label_by_edge_type = { - et: torch.empty(0, 0, dtype=torch.long) - for et in self._negative_label_edge_types - } - + for server_rank, ablp_input in enumerate(input_data): logger.info( f"Rank: {torch.distributed.get_rank()}! Building ABLPNodeSamplerInput for server rank: {server_rank} " - f"with input type: {input_type}. anchors: {anchors.shape}, " - f"positive_labels edge types: {list(positive_label_by_edge_type.keys())}, " - f"negative_labels edge types: {list(negative_label_by_edge_type.keys())}" - ) - ablp_input = ABLPNodeSamplerInput( - node=anchors, - input_type=input_type, - positive_label_by_edge_types=positive_label_by_edge_type, - negative_label_by_edge_types=negative_label_by_edge_type, + f"with input type: {input_type}. " + f"positive_labels edge types: {list(ablp_input.positive_label_by_edge_types.keys())}, " + f"negative_labels edge types: {list(ablp_input.negative_label_by_edge_types.keys())}" ) - input_data.append(ablp_input) return ( input_data, @@ -741,7 +929,7 @@ def _setup_for_graph_store( edge_types=edge_types, node_feature_info=node_feature_info, edge_feature_info=edge_feature_info, - edge_dir=dataset.fetch_edge_dir(), + edge_dir=edge_dir, ), backend_key, ) @@ -752,82 +940,123 @@ def _set_labels( positive_labels_by_label_edge_type: dict[EdgeType, torch.Tensor], negative_labels_by_label_edge_type: dict[EdgeType, torch.Tensor], ) -> Union[Data, HeteroData]: - """ - Sets the labels and relevant fields in the torch_geometric Data object, converting the global node ids for labels to their - local index. Removes inserted supervision edge type from the data variables, since this is an implementation detail and should not be - exposed in the final HeteroData/Data object. + """Attach ABLP labels to the collated graph, remapped to subgraph-local indices. + + The tensor output uses ``[2, E]`` label edge indices. The compatibility + path converts the same tensors to the deprecated ragged dictionaries. + Singleton supervision retains flattened output; with multiple supervision + types, outputs remain keyed even when only one type has negative labels. + Args: - data (Union[Data, HeteroData]): Graph to provide labels for - positive_labels_by_label_edge_type (dict[EdgeType, torch.Tensor]): Dict[positive label edge type, label ID tensor], - where the ith row of the tensor corresponds to the ith anchor node ID. - negative_labels_by_label_edge_type (dict[EdgeType, torch.Tensor]): Dict[negative label edge type, label ID tensor], - where the ith row of the tensor corresponds to the ith anchor node ID. + data (Union[Data, HeteroData]): Graph to attach labels to. + positive_labels_by_label_edge_type (dict[EdgeType, torch.Tensor]): Per + positive-label edge type, a ``[N_anchors, M]`` tensor whose ``i``-th + row holds the global label ids of the ``i``-th anchor. + negative_labels_by_label_edge_type (dict[EdgeType, torch.Tensor]): As + above, for negative-label edge types. + Returns: - Union[Data, HeteroData]: torch_geometric HeteroData/Data object with the filtered edge fields and labels set as properties of the instance + Union[Data, HeteroData]: The same object with the supervision edge fields + stripped and ``y_positive`` (and ``y_negative`` when present) attached. + + Raises: + ValueError: If no positive labels are found in ``data``. """ - # shape [N], where N is the number of nodes in the subgraph, and local_node_to_global_node[i] gives the global node id for local node id `i` - node_type_to_local_node_to_global_node: dict[NodeType, torch.Tensor] = {} + local_id_to_global_id_by_node_type: dict[NodeType, torch.Tensor] = {} if isinstance(data, HeteroData): - for e_type in self._supervision_edge_types: - node_type_to_local_node_to_global_node[e_type[0]] = data[e_type[0]].node - node_type_to_local_node_to_global_node[e_type[2]] = data[e_type[2]].node + # Key the map off the label edge types rather than + # self._supervision_edge_types: the latter is stored reversed when + # edge_dir="in" (see __init__), so its dst node type is the anchor + # type there, while label edge types always arrive in outward form. + # Deriving from the keys that _remap_labels_by_edge_type will look up + # keeps the two in step regardless of edge direction. + for label_edge_type, label_tensor in ( + *positive_labels_by_label_edge_type.items(), + *negative_labels_by_label_edge_type.items(), + ): + # Mirror the zero-anchor skip in _remap_labels_by_edge_type. An + # absent node type would otherwise be auto-created as an empty + # store by HeteroData.__getitem__, turning a clear KeyError into + # an AttributeError on .node. + if label_tensor.size(0) == 0: + continue + supervision_node_type = label_edge_type[2] + local_id_to_global_id_by_node_type[supervision_node_type] = data[ + supervision_node_type + ].node else: - node_type_to_local_node_to_global_node[DEFAULT_HOMOGENEOUS_NODE_TYPE] = ( + local_id_to_global_id_by_node_type[DEFAULT_HOMOGENEOUS_NODE_TYPE] = ( data.node ) - output_positive_labels: dict[EdgeType, dict[int, torch.Tensor]] = defaultdict( - dict - ) - output_negative_labels: dict[EdgeType, dict[int, torch.Tensor]] = defaultdict( - dict + + ( + positive_label_edge_indices, + negative_label_edge_indices, + ) = remap_labels_to_local_edge_indices( + local_id_to_global_id_by_node_type=local_id_to_global_id_by_node_type, + positive_labels_by_edge_type=positive_labels_by_label_edge_type, + negative_labels_by_edge_type=negative_labels_by_label_edge_type, ) - # We always have supervision edge types of the form (anchor_node_type, to, supervision_node_type) - # So we can index into the edge type accordingly. - edge_index = 2 - for edge_type, label_tensor in positive_labels_by_label_edge_type.items(): - for local_anchor_node_id in range(label_tensor.size(0)): - positive_mask = ( - node_type_to_local_node_to_global_node[ - edge_type[edge_index] - ].unsqueeze(1) - == label_tensor[local_anchor_node_id] - ) # shape [N, P], where N is the number of nodes and P is the number of positive labels for the current anchor node - - # Gets the indexes of the items in local_node_to_global_node which match any of the positive labels for the current anchor node - output_positive_labels[ - label_edge_type_to_message_passing_edge_type(edge_type) - ][local_anchor_node_id] = torch.nonzero(positive_mask)[:, 0].to( - self.to_device + output_positive_labels: dict[ + EdgeType, Union[torch.Tensor, dict[int, torch.Tensor]] + ] = {} + output_negative_labels: dict[ + EdgeType, Union[torch.Tensor, dict[int, torch.Tensor]] + ] = {} + if self._use_edge_index_output: + for edge_type, label_edge_index in positive_label_edge_indices.items(): + output_positive_labels[edge_type] = label_edge_index + for edge_type, label_edge_index in negative_label_edge_indices.items(): + output_negative_labels[edge_type] = label_edge_index + else: + label_edge_types = ( + positive_label_edge_indices.keys() | negative_label_edge_indices.keys() + ) + if isinstance(data, HeteroData): + num_anchors_by_edge_type = { + edge_type: int(data[edge_type[0]].batch_size) + for edge_type in label_edge_types + } + else: + num_anchors_by_edge_type = { + edge_type: int(data.batch_size) for edge_type in label_edge_types + } + output_positive_labels = { + edge_type: label_edge_index_to_dict( + label_edge_index=label_edge_index, + num_anchors=num_anchors_by_edge_type[edge_type], ) - # Shape [X], where X is the number of indexes in the original local_node_to_global_node which match a node in the positive labels for the current anchor node - - for edge_type, label_tensor in negative_labels_by_label_edge_type.items(): - for local_anchor_node_id in range(label_tensor.size(0)): - negative_mask = ( - node_type_to_local_node_to_global_node[ - edge_type[edge_index] - ].unsqueeze(1) - == label_tensor[local_anchor_node_id] - ) # shape [N, M], where N is the number of nodes and M is the number of negative labels for the current anchor node - - # Gets the indexes of the items in local_node_to_global_node which match any of the negative labels for the current anchor node - output_negative_labels[ - label_edge_type_to_message_passing_edge_type(edge_type) - ][local_anchor_node_id] = torch.nonzero(negative_mask)[:, 0].to( - self.to_device + for edge_type, label_edge_index in positive_label_edge_indices.items() + } + output_negative_labels = { + edge_type: label_edge_index_to_dict( + label_edge_index=label_edge_index, + num_anchors=num_anchors_by_edge_type[edge_type], ) - # Shape [X], where X is the number of indexes in the original local_node_to_global_node which match a node in the negative labels for the current anchor node + for edge_type, label_edge_index in negative_label_edge_indices.items() + } if not output_positive_labels: raise ValueError("No positive labels were found in the data!") - elif len(output_positive_labels) == 1: + elif len(self._supervision_edge_types) == 1: + if len(output_positive_labels) != 1: + raise ValueError( + "Expected exactly one positive output for a single " + f"supervision edge type, got {list(output_positive_labels)}." + ) data.y_positive = next(iter(output_positive_labels.values())) else: data.y_positive = output_positive_labels - if len(output_negative_labels) == 1: - data.y_negative = next(iter(output_negative_labels.values())) - elif len(output_negative_labels) > 0: - data.y_negative = output_negative_labels + if output_negative_labels: + if len(self._supervision_edge_types) == 1: + if len(output_negative_labels) != 1: + raise ValueError( + "Expected exactly one negative output for a single " + f"supervision edge type, got {list(output_negative_labels)}." + ) + data.y_negative = next(iter(output_negative_labels.values())) + else: + data.y_negative = output_negative_labels return data def _collate_fn(self, msg: SampleMessage) -> Union[Data, HeteroData]: diff --git a/gigl/distributed/utils/ablp.py b/gigl/distributed/utils/ablp.py new file mode 100644 index 000000000..aeff1e569 --- /dev/null +++ b/gigl/distributed/utils/ablp.py @@ -0,0 +1,160 @@ +"""Utilities for remapping ABLP labels to sampled-subgraph indices.""" + +import torch +from torch_geometric.typing import EdgeType + +from gigl.src.common.types.graph_data import NodeType +from gigl.types.graph import label_edge_type_to_message_passing_edge_type +from gigl.utils.data_splitters import PADDING_NODE + + +def label_edge_index_to_dict( + label_edge_index: torch.Tensor, num_anchors: int +) -> dict[int, torch.Tensor]: + """Convert a label edge index to the deprecated per-anchor dictionary. + + The edge index must be grouped by anchor, as guaranteed by + :func:`remap_labels_to_local_edge_indices`. + + Args: + label_edge_index: A ``[2, E]`` tensor. Row 0 contains local anchor + indices and row 1 contains local label-node indices. + num_anchors: Number of anchors in the sampled batch, including anchors + with no labels. + + Returns: + A mapping from every local anchor index to its local label-node indices. + """ + counts = torch.bincount(label_edge_index[0], minlength=num_anchors) + labels_by_anchor = torch.split(label_edge_index[1], counts.tolist()) + return {anchor: labels_by_anchor[anchor] for anchor in range(num_anchors)} + + +def _remap_label_tensor_to_local_edge_index( + label_tensor: torch.Tensor, + sorted_global_node_ids: torch.Tensor, + local_ids_by_sorted_global_id: torch.Tensor, +) -> torch.Tensor: + """Remap one padded global-label tensor to a local label edge index. + + Args: + label_tensor: A ``[N_anchors, M]`` tensor of global label-node ids, + padded with ``PADDING_NODE``. + sorted_global_node_ids: Sorted global ids for the sampled supervision + nodes. + local_ids_by_sorted_global_id: Local node indices in the order given by + ``sorted_global_node_ids``. + + Returns: + A ``[2, E]`` long tensor on the input device. Row 0 contains local + anchor indices and row 1 contains local label-node indices. Labels not + present in the sampled subgraph are omitted. + + Raises: + ValueError: If the sampled node map contains duplicate global ids. + """ + num_anchors = int(label_tensor.size(0)) + num_nodes = int(sorted_global_node_ids.size(0)) + empty_edge_index = torch.empty((2, 0), dtype=torch.long, device=label_tensor.device) + if num_anchors == 0: + return empty_edge_index + + num_labels = int(label_tensor.size(1)) + candidate_global_label_ids = label_tensor.reshape(-1) + candidate_anchor_indices = torch.arange( + num_anchors, device=label_tensor.device + ).repeat_interleave(num_labels) + + is_not_padding = candidate_global_label_ids != PADDING_NODE + candidate_global_label_ids = candidate_global_label_ids[is_not_padding] + candidate_anchor_indices = candidate_anchor_indices[is_not_padding] + if num_nodes == 0 or candidate_global_label_ids.numel() == 0: + return empty_edge_index + + if not bool((sorted_global_node_ids[1:] > sorted_global_node_ids[:-1]).all()): + raise ValueError( + "Vectorized label remapping requires unique global ids in the " + "sampled node map." + ) + + insertion_indices = torch.searchsorted( + sorted_global_node_ids, candidate_global_label_ids + ) + # An absent label larger than every sampled id inserts at num_nodes. Clamp + # before gathering; the exact-match mask below still discards that label. + insertion_indices = insertion_indices.clamp_(max=num_nodes - 1) + is_in_sampled_subgraph = ( + sorted_global_node_ids[insertion_indices] == candidate_global_label_ids + ) + + local_label_indices = local_ids_by_sorted_global_id[insertion_indices][ + is_in_sampled_subgraph + ] + anchor_indices = candidate_anchor_indices[is_in_sampled_subgraph] + return torch.stack((anchor_indices, local_label_indices)) + + +def _remap_labels_by_edge_type( + labels_by_edge_type: dict[EdgeType, torch.Tensor], + local_id_to_global_id_by_node_type: dict[NodeType, torch.Tensor], + sorted_node_lookup_by_type: dict[NodeType, tuple[torch.Tensor, torch.Tensor]], +) -> dict[EdgeType, torch.Tensor]: + """Remap positive or negative labels for each supervision edge type.""" + output: dict[EdgeType, torch.Tensor] = {} + for label_edge_type, label_tensor in labels_by_edge_type.items(): + if label_tensor.size(0) == 0: + continue + + supervision_node_type = label_edge_type[2] + if supervision_node_type not in sorted_node_lookup_by_type: + sorted_node_lookup_by_type[supervision_node_type] = torch.sort( + local_id_to_global_id_by_node_type[supervision_node_type] + ) + ( + sorted_global_node_ids, + local_ids_by_sorted_global_id, + ) = sorted_node_lookup_by_type[supervision_node_type] + output[label_edge_type_to_message_passing_edge_type(label_edge_type)] = ( + _remap_label_tensor_to_local_edge_index( + label_tensor=label_tensor, + sorted_global_node_ids=sorted_global_node_ids, + local_ids_by_sorted_global_id=local_ids_by_sorted_global_id, + ) + ) + return output + + +def remap_labels_to_local_edge_indices( + local_id_to_global_id_by_node_type: dict[NodeType, torch.Tensor], + positive_labels_by_edge_type: dict[EdgeType, torch.Tensor], + negative_labels_by_edge_type: dict[EdgeType, torch.Tensor], +) -> tuple[dict[EdgeType, torch.Tensor], dict[EdgeType, torch.Tensor]]: + """Remap padded global ABLP labels to local label edge indices. + + Positive and negative labels share one sorted node lookup per supervision + node type. Pair order within an anchor is unspecified. + + Args: + local_id_to_global_id_by_node_type: Per node type, a tensor whose + ``i``-th entry is the global id for local node ``i``. + positive_labels_by_edge_type: Per positive-label edge type, a padded + ``[N_anchors, M]`` tensor of global label-node ids. + negative_labels_by_edge_type: Equivalent negative-label tensors. May + be empty. + + Returns: + Positive and negative mappings keyed by message-passing edge type. Each + value is a ``[2, E]`` local label edge index. + """ + sorted_node_lookup_by_type: dict[NodeType, tuple[torch.Tensor, torch.Tensor]] = {} + positive_label_edge_indices = _remap_labels_by_edge_type( + labels_by_edge_type=positive_labels_by_edge_type, + local_id_to_global_id_by_node_type=local_id_to_global_id_by_node_type, + sorted_node_lookup_by_type=sorted_node_lookup_by_type, + ) + negative_label_edge_indices = _remap_labels_by_edge_type( + labels_by_edge_type=negative_labels_by_edge_type, + local_id_to_global_id_by_node_type=local_id_to_global_id_by_node_type, + sorted_node_lookup_by_type=sorted_node_lookup_by_type, + ) + return positive_label_edge_indices, negative_label_edge_indices diff --git a/tests/unit/distributed/dist_ablp_neighborloader_test.py b/tests/unit/distributed/dist_ablp_neighborloader_test.py index 31d3d1cbc..43c26bc5d 100644 --- a/tests/unit/distributed/dist_ablp_neighborloader_test.py +++ b/tests/unit/distributed/dist_ablp_neighborloader_test.py @@ -1,6 +1,7 @@ import unittest from collections import defaultdict -from typing import Literal, Optional, Union +from typing import Any, Literal, Optional, Union +from unittest.mock import patch import torch import torch.multiprocessing as mp @@ -10,8 +11,13 @@ from parameterized import param, parameterized from torch_geometric.data import Data, HeteroData +from gigl.distributed.base_dist_loader import BaseDistLoader from gigl.distributed.dataset_factory import build_dataset -from gigl.distributed.dist_ablp_neighborloader import DistABLPLoader +from gigl.distributed.dist_ablp_neighborloader import ( + DistABLPLoader, + _convert_graph_store_ablp_inputs, + _is_integral_tensor, +) from gigl.distributed.dist_dataset import DistDataset from gigl.distributed.dist_partitioner import DistPartitioner from gigl.distributed.dist_range_partitioner import DistRangePartitioner @@ -36,9 +42,12 @@ to_homogeneous, ) from gigl.utils.data_splitters import DistNodeAnchorLinkSplitter +from gigl.utils.sampling import ABLPInputNodes from tests.test_assets.distributed.utils import ( + MockRemoteDistDataset, assert_tensor_equality, create_test_process_group, + get_process_group_init_method, ) from tests.test_assets.test_case import TestCase @@ -358,12 +367,13 @@ def _run_distributed_ablp_neighbor_loader_multiple_supervision_edge_types( expected=expected_positive_labels[edge_type], ) if expected_negative_labels is not None: - _assert_labels( - anchor_nodes=datum[edge_type[anchor_index]].node, - supervision_nodes=datum[edge_type[supervision_index]].node, - y=datum.y_negative[edge_type], - expected=expected_negative_labels[edge_type], - ) + for edge_type, expected_labels in expected_negative_labels.items(): + _assert_labels( + anchor_nodes=datum[edge_type[anchor_index]].node, + supervision_nodes=datum[edge_type[supervision_index]].node, + y=datum.y_negative[edge_type], + expected=expected_labels, + ) else: assert not hasattr(datum, "y_negative") @@ -416,6 +426,234 @@ def _run_distributed_ablp_neighbor_loader_multiple_supervision_edge_types( shutdown_rpc() +def _global_pair_set( + anchor_node: torch.Tensor, + label_node: torch.Tensor, + label_dict: dict[int, torch.Tensor], +) -> list[tuple[int, int]]: + """Convert a label dictionary to sorted global-id pairs. + + Anchors and labels are looked up in separate node maps so this works for + heterogeneous graphs, where the two live in different node stores. Pass the + same tensor twice for a homogeneous graph. + """ + pairs: list[tuple[int, int]] = [] + for local_anchor, local_labels in label_dict.items(): + global_anchor = int(anchor_node[local_anchor].item()) + for local_label in local_labels.tolist(): + pairs.append((global_anchor, int(label_node[local_label].item()))) + return sorted(pairs) + + +def _global_pair_set_from_edge_index( + anchor_node: torch.Tensor, + label_node: torch.Tensor, + label_edge_index: torch.Tensor, +) -> list[tuple[int, int]]: + """Convert a label edge index to sorted global-id pairs. + + Takes separate anchor and label node maps for the same reason as + :func:`_global_pair_set`. + """ + return sorted( + ( + int(anchor_node[local_anchor].item()), + int(label_node[local_label].item()), + ) + for local_anchor, local_label in label_edge_index.t().tolist() + ) + + +def _collect_homogeneous_labels( + _, + return_dict, + use_edge_index_output: bool, + dataset: DistDataset, + input_nodes: torch.Tensor, + batch_size: int, + has_negatives: bool, +): + """Run one loader format and return its sorted global-id label pairs.""" + create_test_process_group() + loader = DistABLPLoader( + dataset=dataset, + num_neighbors=[2, 2], + input_nodes=input_nodes, + batch_size=batch_size, + pin_memory_device=torch.device("cpu"), + use_edge_index_output=use_edge_index_output, + ) + positive_pairs: list[tuple[int, int]] = [] + negative_pairs: list[tuple[int, int]] = [] + for datum in loader: + assert isinstance(datum, Data) + node = datum.node + if use_edge_index_output: + assert isinstance(datum.y_positive, torch.Tensor) + assert datum.y_positive.size(0) == 2 + positive_pairs.extend( + _global_pair_set_from_edge_index(node, node, datum.y_positive) + ) + else: + positive_pairs.extend(_global_pair_set(node, node, datum.y_positive)) + if has_negatives: + if use_edge_index_output: + assert isinstance(datum.y_negative, torch.Tensor) + assert datum.y_negative.size(0) == 2 + negative_pairs.extend( + _global_pair_set_from_edge_index(node, node, datum.y_negative) + ) + else: + negative_pairs.extend(_global_pair_set(node, node, datum.y_negative)) + else: + assert not hasattr(datum, "y_negative"), ( + f"expected no negatives, got {getattr(datum, 'y_negative', None)}" + ) + return_dict[use_edge_index_output] = ( + sorted(positive_pairs), + sorted(negative_pairs), + ) + shutdown_rpc() + + +def _edge_type_key(edge_type: EdgeType) -> tuple[str, ...]: + """Canonical dict key for an edge type. + + Label edge types reach collation as plain tuples on some paths and as + ``EdgeType`` on others, and the two stringify differently (``"('a', 'to', + 'b')"`` vs ``'a-to-b'``). Normalizing to plain strings keeps keys comparable + across the process boundary regardless of which arrives. + """ + return tuple(str(part) for part in edge_type) + + +def _accumulate_heterogeneous_pairs( + data: HeteroData, + labels_by_edge_type: dict[EdgeType, Union[torch.Tensor, dict[int, torch.Tensor]]], + use_edge_index_output: bool, + into: dict[tuple[str, ...], list[tuple[int, int]]], +) -> None: + """Accumulate one batch's labels as global (anchor, label) id pairs. + + Anchors and supervision nodes are resolved through the node maps of the edge + type's src and dst stores respectively, so a label remapped against the wrong + node type surfaces as a wrong global id rather than passing silently. + """ + anchor_index = 0 + supervision_index = 2 + for edge_type, labels in labels_by_edge_type.items(): + anchor_node = data[edge_type[anchor_index]].node + label_node = data[edge_type[supervision_index]].node + if use_edge_index_output: + assert isinstance(labels, torch.Tensor), f"{edge_type}: {type(labels)}" + assert labels.size(0) == 2 + pairs = _global_pair_set_from_edge_index(anchor_node, label_node, labels) + else: + assert not isinstance(labels, torch.Tensor), f"{edge_type}: {type(labels)}" + pairs = _global_pair_set(anchor_node, label_node, labels) + into[_edge_type_key(edge_type)].extend(pairs) + + +def _collect_heterogeneous_labels( + _, + return_dict, + use_edge_index_output: bool, + dataset: DistDataset, + input_nodes: tuple[NodeType, torch.Tensor], + supervision_edge_types: list[EdgeType], + batch_size: int, +): + """Run one loader format on a heterogeneous graph and return its label pairs. + + Results are keyed by ``str(edge_type)`` so they cross the process boundary as + plain data. + """ + create_test_process_group() + loader = DistABLPLoader( + dataset=dataset, + num_neighbors=[2, 2], + input_nodes=input_nodes, + batch_size=batch_size, + pin_memory_device=torch.device("cpu"), + supervision_edge_type=supervision_edge_types, + use_edge_index_output=use_edge_index_output, + ) + positive_pairs: dict[tuple[str, ...], list[tuple[int, int]]] = defaultdict(list) + negative_pairs: dict[tuple[str, ...], list[tuple[int, int]]] = defaultdict(list) + for datum in loader: + assert isinstance(datum, HeteroData) + # Several supervision edge types keep y_positive / y_negative in their + # dict-of-edge-type form rather than collapsing to a bare value. + _accumulate_heterogeneous_pairs( + datum, datum.y_positive, use_edge_index_output, positive_pairs + ) + _accumulate_heterogeneous_pairs( + datum, datum.y_negative, use_edge_index_output, negative_pairs + ) + return_dict[use_edge_index_output] = ( + {edge_type: sorted(pairs) for edge_type, pairs in positive_pairs.items()}, + {edge_type: sorted(pairs) for edge_type, pairs in negative_pairs.items()}, + ) + shutdown_rpc() + + +def _run_graph_store_validation_protocol( + rank: int, + init_method: str, + scenario: str, + results: Any, +) -> None: + torch.distributed.init_process_group( + backend="gloo", + rank=rank, + world_size=2, + init_method=init_method, + ) + try: + supervision_edge_type = ( + _A_TO_C if scenario == "schema_mismatch" and rank == 1 else _A_TO_B + ) + ablp_input = ABLPInputNodes( + anchor_nodes=torch.tensor([1]), + anchor_node_type=_A, + labels={supervision_edge_type: (torch.tensor([[2]]), None)}, + ) + if scenario == "malformed_payload" and rank == 1: + object.__setattr__(ablp_input, "anchor_nodes", [1]) + dataset = MockRemoteDistDataset( + num_storage_nodes=1, + num_compute_nodes=2, + edge_types=[message_passing_to_positive_label(supervision_edge_type)], + edge_dir=( + "sideways" if scenario == "invalid_direction" and rank == 1 else "out" + ), + ) + loader = DistABLPLoader.__new__(DistABLPLoader) + loader._instance_count = 0 + loader._shutdowned = True + with patch.object( + BaseDistLoader, "create_graph_store_worker_options" + ) as create_worker_options: + try: + loader._setup_for_graph_store( + input_nodes={0: ablp_input}, + dataset=dataset, + num_workers=1, + worker_concurrency=1, + channel_size="1MB", + prefetch_size=1, + ) + except ValueError as error: + results[rank] = (str(error), create_worker_options.call_count) + else: + results[rank] = ( + "no validation error", + create_worker_options.call_count, + ) + finally: + torch.distributed.destroy_process_group() + + class DistABLPLoaderTest(TestCase): def tearDown(self): if torch.distributed.is_initialized(): @@ -425,6 +663,524 @@ def tearDown(self): torch.distributed.destroy_process_group() super().tearDown() + def test_is_integral_tensor(self): + for dtype in ( + torch.uint8, + torch.int8, + torch.int16, + torch.int32, + torch.int64, + ): + with self.subTest(dtype=dtype): + self.assertTrue(_is_integral_tensor(torch.empty(0, dtype=dtype))) + for dtype in (torch.bool, torch.float16, torch.float32, torch.float64): + with self.subTest(dtype=dtype): + self.assertFalse(_is_integral_tensor(torch.empty(0, dtype=dtype))) + + def test_convert_graph_store_ablp_inputs_outward(self): + positive_a_to_b = message_passing_to_positive_label(_A_TO_B) + negative_a_to_b = message_passing_to_negative_label(_A_TO_B) + positive_a_to_c = message_passing_to_positive_label(_A_TO_C) + anchors = torch.tensor([1, 2]) + positive_b = torch.tensor([[10], [11]]) + negative_b = torch.tensor([[20, 21], [22, 23]]) + positive_c = torch.tensor([[30], [31]]) + input_nodes = { + 0: ABLPInputNodes( + anchor_nodes=anchors, + anchor_node_type=_A, + labels={ + _A_TO_B: (positive_b, negative_b), + _A_TO_C: (positive_c, None), + }, + ) + } + + ( + sampler_inputs, + input_type, + supervision_edge_types, + positive_label_edge_types, + negative_label_edge_types, + ) = _convert_graph_store_ablp_inputs( + input_nodes=input_nodes, + num_storage_nodes=2, + edge_types=[ + positive_a_to_b, + negative_a_to_b, + positive_a_to_c, + ], + edge_dir="out", + ) + + self.assertEqual(input_type, _A) + self.assertEqual(supervision_edge_types, [_A_TO_B, _A_TO_C]) + self.assertEqual(positive_label_edge_types, [positive_a_to_b, positive_a_to_c]) + self.assertEqual(negative_label_edge_types, [negative_a_to_b]) + self.assertLen(sampler_inputs, 2) + + supplied_input = sampler_inputs[0] + self.assertIs(supplied_input.node, anchors) + self.assertEqual(supplied_input.input_type, _A) + self.assertEqual( + set(supplied_input.positive_label_by_edge_types), + {positive_a_to_b, positive_a_to_c}, + ) + self.assertIs( + supplied_input.positive_label_by_edge_types[positive_a_to_b], + positive_b, + ) + self.assertIs( + supplied_input.positive_label_by_edge_types[positive_a_to_c], + positive_c, + ) + self.assertEqual( + set(supplied_input.negative_label_by_edge_types), {negative_a_to_b} + ) + self.assertIs( + supplied_input.negative_label_by_edge_types[negative_a_to_b], + negative_b, + ) + + missing_input = sampler_inputs[1] + self.assertEqual(missing_input.input_type, _A) + self.assertEqual(missing_input.node.dtype, torch.long) + self.assertEqual(tuple(missing_input.node.shape), (0,)) + self.assertEqual( + set(missing_input.positive_label_by_edge_types), + {positive_a_to_b, positive_a_to_c}, + ) + for labels in missing_input.positive_label_by_edge_types.values(): + self.assertEqual(labels.dtype, torch.long) + self.assertEqual(tuple(labels.shape), (0, 0)) + self.assertEqual( + set(missing_input.negative_label_by_edge_types), {negative_a_to_b} + ) + self.assertEqual( + missing_input.negative_label_by_edge_types[negative_a_to_b].dtype, + torch.long, + ) + self.assertEqual( + tuple(missing_input.negative_label_by_edge_types[negative_a_to_b].shape), + (0, 0), + ) + + def test_convert_graph_store_ablp_inputs_normalizes_base_placeholder(self): + positive_edge_type = message_passing_to_positive_label(_A_TO_B) + negative_edge_type = message_passing_to_negative_label(_A_TO_B) + positive_labels = torch.empty((0, 0), dtype=torch.long) + topology_complete_negative = torch.empty((0, 0), dtype=torch.long) + sampler_inputs, *_ = _convert_graph_store_ablp_inputs( + input_nodes={ + 0: ABLPInputNodes( + anchor_nodes=torch.empty(0, dtype=torch.long), + anchor_node_type=_A, + labels={_A_TO_B: (positive_labels, None)}, + ), + 1: ABLPInputNodes( + anchor_nodes=torch.empty(0, dtype=torch.long), + anchor_node_type=_A, + labels={ + _A_TO_B: ( + torch.empty((0, 0), dtype=torch.long), + topology_complete_negative, + ) + }, + ), + }, + num_storage_nodes=2, + edge_types=[positive_edge_type, negative_edge_type], + edge_dir="out", + ) + + self.assertIs( + sampler_inputs[0].positive_label_by_edge_types[positive_edge_type], + positive_labels, + ) + normalized_negative = sampler_inputs[0].negative_label_by_edge_types[ + negative_edge_type + ] + self.assertEqual(normalized_negative.dtype, torch.long) + self.assertEqual(tuple(normalized_negative.shape), (0, 0)) + self.assertIs( + sampler_inputs[1].negative_label_by_edge_types[negative_edge_type], + topology_complete_negative, + ) + + def test_convert_graph_store_ablp_inputs_incoming(self): + positive_b_to_a = message_passing_to_positive_label(_B_TO_A) + negative_b_to_a = message_passing_to_negative_label(_B_TO_A) + positive_c_to_a = message_passing_to_positive_label(_C_TO_A) + ( + sampler_inputs, + input_type, + supervision_types, + positive_types, + negative_types, + ) = _convert_graph_store_ablp_inputs( + input_nodes={ + 0: ABLPInputNodes( + anchor_nodes=torch.tensor([1]), + anchor_node_type=_A, + labels={ + _B_TO_A: (torch.tensor([[10]]), torch.tensor([[20]])), + _C_TO_A: (torch.tensor([[30]]), None), + }, + ) + }, + num_storage_nodes=1, + edge_types=[ + positive_b_to_a, + negative_b_to_a, + positive_c_to_a, + ], + edge_dir="in", + ) + + self.assertEqual(input_type, _A) + self.assertEqual(supervision_types, [_B_TO_A, _C_TO_A]) + self.assertEqual(positive_types, [positive_b_to_a, positive_c_to_a]) + self.assertEqual(negative_types, [negative_b_to_a]) + self.assertEqual( + list(sampler_inputs[0].positive_label_by_edge_types), + [positive_b_to_a, positive_c_to_a], + ) + self.assertEqual( + list(sampler_inputs[0].negative_label_by_edge_types), + [negative_b_to_a], + ) + + def test_convert_graph_store_ablp_inputs_rejects_invalid_inputs(self): + positive_edge_type = message_passing_to_positive_label(_A_TO_B) + negative_edge_type = message_passing_to_negative_label(_A_TO_B) + positive_c_edge_type = message_passing_to_positive_label(_A_TO_C) + topology = [positive_edge_type, negative_edge_type, positive_c_edge_type] + + def make_input( + anchor_nodes: Any, + labels: Any, + anchor_node_type: NodeType = _A, + ) -> ABLPInputNodes: + return ABLPInputNodes( + anchor_nodes=anchor_nodes, + anchor_node_type=anchor_node_type, + labels=labels, + ) + + valid_input = make_input( + torch.tensor([1]), + {_A_TO_B: (torch.tensor([[2]]), torch.tensor([[3]]))}, + ) + cases = [ + ("empty input", {}, 1, topology, "out"), + ( + "empty labels", + {0: make_input(torch.tensor([], dtype=torch.long), {})}, + 1, + topology, + "out", + ), + ("negative rank", {-1: valid_input}, 1, topology, "out"), + ("out of range rank", {1: valid_input}, 1, topology, "out"), + ("missing edge types", {0: valid_input}, 1, None, "out"), + ( + "missing positive topology", + {0: valid_input}, + 1, + [negative_edge_type], + "out", + ), + ( + "wrong outward anchor endpoint", + { + 0: make_input( + torch.tensor([1]), + {_B_TO_A: (torch.tensor([[2]]), None)}, + ) + }, + 1, + [message_passing_to_positive_label(_B_TO_A)], + "out", + ), + ( + "wrong incoming anchor endpoint", + { + 0: make_input( + torch.tensor([1]), + {_A_TO_B: (torch.tensor([[2]]), None)}, + ) + }, + 1, + [positive_edge_type], + "in", + ), + ( + "inconsistent anchor types", + { + 0: valid_input, + 1: make_input( + torch.tensor([2]), + {_A_TO_B: (torch.tensor([[3]]), torch.tensor([[4]]))}, + _B, + ), + }, + 2, + topology, + "out", + ), + ( + "inconsistent supervision keys", + { + 0: valid_input, + 1: make_input( + torch.tensor([2]), + {_A_TO_C: (torch.tensor([[3]]), None)}, + ), + }, + 2, + topology, + "out", + ), + ( + "non tensor anchors", + {0: make_input([1], valid_input.labels)}, + 1, + topology, + "out", + ), + ( + "non 1-D anchors", + {0: make_input(torch.tensor([[1]]), valid_input.labels)}, + 1, + topology, + "out", + ), + ( + "floating anchors", + {0: make_input(torch.tensor([1.0]), valid_input.labels)}, + 1, + topology, + "out", + ), + ( + "boolean anchors", + {0: make_input(torch.tensor([True]), valid_input.labels)}, + 1, + topology, + "out", + ), + ( + "None positive labels", + { + 0: make_input( + torch.tensor([1]), {_A_TO_B: (None, torch.tensor([[3]]))} + ) + }, + 1, + topology, + "out", + ), + ( + "non tensor positive labels", + { + 0: make_input( + torch.tensor([1]), {_A_TO_B: ([2], torch.tensor([[3]]))} + ) + }, + 1, + topology, + "out", + ), + ( + "non tensor negative labels", + { + 0: make_input( + torch.tensor([1]), {_A_TO_B: (torch.tensor([[2]]), [3])} + ) + }, + 1, + topology, + "out", + ), + ( + "non 2-D positive labels", + { + 0: make_input( + torch.tensor([1]), + {_A_TO_B: (torch.tensor([2]), torch.tensor([[3]]))}, + ) + }, + 1, + topology, + "out", + ), + ( + "non 2-D negative labels", + { + 0: make_input( + torch.tensor([1]), + {_A_TO_B: (torch.tensor([[2]]), torch.tensor([3]))}, + ) + }, + 1, + topology, + "out", + ), + ( + "floating positive labels", + { + 0: make_input( + torch.tensor([1]), + {_A_TO_B: (torch.tensor([[2.0]]), torch.tensor([[3]]))}, + ) + }, + 1, + topology, + "out", + ), + ( + "boolean negative labels", + { + 0: make_input( + torch.tensor([1]), + {_A_TO_B: (torch.tensor([[2]]), torch.tensor([[True]]))}, + ) + }, + 1, + topology, + "out", + ), + ( + "positive row mismatch", + { + 0: make_input( + torch.tensor([1, 2]), + {_A_TO_B: (torch.tensor([[3]]), torch.tensor([[4], [5]]))}, + ) + }, + 1, + topology, + "out", + ), + ( + "negative row mismatch", + { + 0: make_input( + torch.tensor([1, 2]), + {_A_TO_B: (torch.tensor([[3], [4]]), torch.tensor([[5]]))}, + ) + }, + 1, + topology, + "out", + ), + ( + "missing required negatives", + { + 0: make_input( + torch.tensor([1]), {_A_TO_B: (torch.tensor([[2]]), None)} + ) + }, + 1, + topology, + "out", + ), + ( + "unexpected negatives", + {0: valid_input}, + 1, + [positive_edge_type], + "out", + ), + ] + for name, input_nodes, num_storage_nodes, edge_types, edge_dir in cases: + with self.subTest(name=name): + with self.assertRaises(ValueError): + _convert_graph_store_ablp_inputs( + input_nodes=input_nodes, + num_storage_nodes=num_storage_nodes, + edge_types=edge_types, + edge_dir=edge_dir, # ty: ignore[invalid-argument-type] + ) + + def test_convert_graph_store_ablp_inputs_rejects_invalid_direction(self): + with self.assertRaises(ValueError): + _convert_graph_store_ablp_inputs( + input_nodes={ + 0: ABLPInputNodes( + anchor_nodes=torch.tensor([1]), + anchor_node_type=_A, + labels={_A_TO_B: (torch.tensor([[2]]), None)}, + ) + }, + num_storage_nodes=1, + edge_types=[message_passing_to_positive_label(_A_TO_B)], + edge_dir="sideways", # ty: ignore[invalid-argument-type] + ) + + def test_graph_store_setup_retains_single_supervision_gate(self): + create_test_process_group() + loader = DistABLPLoader.__new__(DistABLPLoader) + loader._instance_count = 0 + loader._shutdowned = True + dataset = MockRemoteDistDataset( + num_storage_nodes=1, + edge_types=[ + message_passing_to_positive_label(_A_TO_B), + message_passing_to_positive_label(_A_TO_C), + ], + ) + with patch.object( + BaseDistLoader, "create_graph_store_worker_options" + ) as create_worker_options: + with self.assertRaisesRegex( + ValueError, + "Graph Store mode currently only supports a single supervision edge type", + ): + loader._setup_for_graph_store( + input_nodes={ + 0: ABLPInputNodes( + anchor_nodes=torch.tensor([1]), + anchor_node_type=_A, + labels={ + _A_TO_B: (torch.tensor([[2]]), None), + _A_TO_C: (torch.tensor([[3]]), None), + }, + ) + }, + dataset=dataset, + num_workers=1, + worker_concurrency=1, + channel_size="1MB", + prefetch_size=1, + ) + create_worker_options.assert_not_called() + + @parameterized.expand( + [ + param("malformed peer payload", "malformed_payload", "rank 1:"), + param("invalid peer direction", "invalid_direction", "rank 1:"), + param("different peer schema", "schema_mismatch", "schemas differ"), + ] + ) + def test_graph_store_setup_synchronizes_validation( + self, _: str, scenario: str, expected_error: str + ): + manager = mp.Manager() + results = manager.dict() + mp.spawn( + fn=_run_graph_store_validation_protocol, + args=(get_process_group_init_method(), scenario, results), + nprocs=2, + join=True, + ) + + self.assertLen(results, 2) + for rank in range(2): + error, worker_option_calls = results[rank] + self.assertIn(expected_error, error) + self.assertEqual(worker_option_calls, 0) + @parameterized.expand( [ param( @@ -556,6 +1312,209 @@ def test_ablp_dataloader( ), ) + @parameterized.expand( + [ + param( + "positive and negative", + labeled_edges={ + _POSITIVE_EDGE_TYPE: torch.tensor([[10, 15], [15, 16]]), + _NEGATIVE_EDGE_TYPE: torch.tensor( + [[10, 10, 11, 15], [13, 16, 14, 17]] + ), + }, + input_nodes=torch.tensor([10, 15]), + batch_size=2, + has_negatives=True, + ), + param( + "positive only", + labeled_edges={_POSITIVE_EDGE_TYPE: torch.tensor([[10, 15], [15, 16]])}, + input_nodes=torch.tensor([10, 15]), + batch_size=2, + has_negatives=False, + ), + # Anchor 11 has message-passing edges (11 -> {13, 17}) but is the + # source of NO positive-label edge, so its positive-label row is + # all-padding and y_positive[11] is a guaranteed-empty tensor. This + # exercises the empty-anchor branch end-to-end for both outputs. + param( + "guaranteed empty positive anchor", + labeled_edges={ + _POSITIVE_EDGE_TYPE: torch.tensor([[10, 15], [15, 16]]), + _NEGATIVE_EDGE_TYPE: torch.tensor( + [[10, 10, 11, 15], [13, 16, 14, 17]] + ), + }, + input_nodes=torch.tensor([10, 11, 15]), + batch_size=3, + has_negatives=True, + ), + ] + ) + def test_edge_index_output_matches_dict_output( + self, _, labeled_edges, input_nodes, batch_size, has_negatives + ): + """Both output formats contain the same global anchor-label pairs.""" + edge_index = { + DEFAULT_HOMOGENEOUS_EDGE_TYPE: torch.tensor( + [[10, 10, 11, 11, 15, 15, 16, 16], [11, 12, 13, 17, 13, 14, 12, 14]] + ), + } + edge_index.update(labeled_edges) + partition_output = PartitionOutput( + node_partition_book=to_heterogeneous_node(torch.zeros(18)), + edge_partition_book={ + e_type: torch.zeros(int(e_idx.max().item() + 1)) + for e_type, e_idx in edge_index.items() + }, + partitioned_edge_index={ + etype: GraphPartitionData( + edge_index=idx, edge_ids=torch.arange(idx.size(1)) + ) + for etype, idx in edge_index.items() + }, + partitioned_edge_features=None, + partitioned_node_features=None, + partitioned_negative_labels=None, + partitioned_positive_labels=None, + partitioned_node_labels=None, + ) + dataset = DistDataset(rank=0, world_size=1, edge_dir="out") + dataset.build(partition_output=partition_output) + + manager = mp.Manager() + return_dict = manager.dict() + for use_edge_index_output in (False, True): + mp.spawn( + fn=_collect_homogeneous_labels, + args=( + return_dict, + use_edge_index_output, + dataset, + input_nodes, + batch_size, + has_negatives, + ), + ) + self.assertEqual(return_dict[False][0], return_dict[True][0]) + self.assertEqual(return_dict[False][1], return_dict[True][1]) + + @parameterized.expand( + [ + param( + "edge_dir=out", + edge_dir="out", + edge_index={ + _A_TO_B: torch.tensor([[10, 10], [11, 12]]), + message_passing_to_positive_label(_A_TO_B): torch.tensor( + [[10, 10], [13, 14]] + ), + message_passing_to_negative_label(_A_TO_B): torch.tensor( + [[10, 10], [15, 16]] + ), + _A_TO_C: torch.tensor([[10, 10], [20, 21]]), + message_passing_to_positive_label(_A_TO_C): torch.tensor( + [[10, 10], [22, 23]] + ), + }, + ), + # edge_dir="in" stores the supervision edge types reversed, so their + # dst node type is the anchor type rather than the supervision type, + # while the label edge types reaching collation stay outward-facing. + # A loader that resolves supervision nodes off the former instead of + # the latter fails here and passes the edge_dir="out" case above. + param( + "edge_dir=in", + edge_dir="in", + edge_index={ + _B_TO_A: torch.tensor([[11, 12], [10, 10]]), + message_passing_to_positive_label(_B_TO_A): torch.tensor( + [[13, 14], [10, 10]] + ), + message_passing_to_negative_label(_B_TO_A): torch.tensor( + [[15, 16], [10, 10]] + ), + _C_TO_A: torch.tensor([[20, 21], [10, 10]]), + message_passing_to_positive_label(_C_TO_A): torch.tensor( + [[22, 23], [10, 10]] + ), + }, + ), + ] + ) + def test_heterogeneous_edge_index_output_matches_dict_output( + self, + _, + edge_dir: Literal["in", "out"], + edge_index: dict[EdgeType, torch.Tensor], + ): + """Both output formats agree on a heterogeneous graph, for either edge_dir. + + Anchors are node type ``a``; supervision nodes are ``b`` and ``c``. Because + the anchor and supervision node maps are distinct, a label remapped against + the wrong node type yields a wrong global id here instead of going unnoticed. + """ + nodes: dict[NodeType, list[torch.Tensor]] = defaultdict(list) + for edge_type, edge_idx in edge_index.items(): + nodes[edge_type[0]].append(edge_idx[0]) + nodes[edge_type[2]].append(edge_idx[1]) + partition_output = PartitionOutput( + node_partition_book={ + node_type: torch.zeros(int(torch.cat(node_ids).max().item() + 1)) + for node_type, node_ids in nodes.items() + }, + edge_partition_book={ + e_type: torch.zeros(int(e_idx.max().item() + 1)) + for e_type, e_idx in edge_index.items() + }, + partitioned_edge_index={ + etype: GraphPartitionData( + edge_index=idx, edge_ids=torch.arange(idx.size(1)) + ) + for etype, idx in edge_index.items() + }, + partitioned_edge_features=None, + partitioned_node_features=None, + partitioned_negative_labels=None, + partitioned_positive_labels=None, + partitioned_node_labels=None, + ) + dataset = DistDataset(rank=0, world_size=1, edge_dir=edge_dir) + dataset.build(partition_output=partition_output) + + manager = mp.Manager() + return_dict = manager.dict() + for use_edge_index_output in (False, True): + mp.spawn( + fn=_collect_heterogeneous_labels, + args=( + return_dict, + use_edge_index_output, + dataset, + (_A, torch.tensor([10])), + [_A_TO_B, _A_TO_C], + 1, # batch_size + ), + ) + # Both formats resolve to the same global ids... + self.assertEqual(dict(return_dict[False][0]), dict(return_dict[True][0])) + self.assertEqual(dict(return_dict[False][1]), dict(return_dict[True][1])) + # ...and did so over labels that actually exist, so equality above cannot + # hold vacuously. + self.assertEqual( + dict(return_dict[True][0]), + { + _edge_type_key(_A_TO_B): [(10, 13), (10, 14)], + _edge_type_key(_A_TO_C): [(10, 22), (10, 23)], + }, + ) + self.assertEqual( + dict(return_dict[True][1]), + { + _edge_type_key(_A_TO_B): [(10, 15), (10, 16)], + }, + ) + def test_cora_supervised(self): create_test_process_group() cora_supervised_info = get_mocked_dataset_artifact_metadata()[ diff --git a/tests/unit/distributed/label_remap_cuda_device_test.py b/tests/unit/distributed/label_remap_cuda_device_test.py new file mode 100644 index 000000000..53ac1ff07 --- /dev/null +++ b/tests/unit/distributed/label_remap_cuda_device_test.py @@ -0,0 +1,60 @@ +"""CUDA device-placement regression test for ABLP label remapping.""" + +import unittest + +import torch + +from gigl.distributed.utils.ablp import remap_labels_to_local_edge_indices +from gigl.src.common.types.graph_data import EdgeType, NodeType, Relation +from gigl.types.graph import message_passing_to_positive_label +from tests.test_assets.test_case import TestCase + +_USER = NodeType("user") +_STORY = NodeType("story") +_USER_TO_STORY = EdgeType(_USER, Relation("to"), _STORY) + + +def _inputs( + device: torch.device, +) -> tuple[dict[NodeType, torch.Tensor], dict[EdgeType, torch.Tensor]]: + """Build an unsorted node map and padded labels on one device.""" + node_map = {_STORY: torch.tensor([15, 10, 16, 11, 12], device=device)} + positive_labels = { + message_passing_to_positive_label(_USER_TO_STORY): torch.tensor( + [[15, 15], [16, -1], [-1, -1]], + dtype=torch.long, + device=device, + ) + } + return node_map, positive_labels + + +@unittest.skipUnless(torch.cuda.is_available(), "requires a CUDA device") +class LabelRemapCudaDeviceTest(TestCase): + def test_cuda_output_matches_cpu(self) -> None: + cpu_node_map, cpu_positive_labels = _inputs(torch.device("cpu")) + expected_positive, _ = remap_labels_to_local_edge_indices( + local_id_to_global_id_by_node_type=cpu_node_map, + positive_labels_by_edge_type=cpu_positive_labels, + negative_labels_by_edge_type={}, + ) + + cuda_node_map, cuda_positive_labels = _inputs(torch.device("cuda")) + actual_positive, _ = remap_labels_to_local_edge_indices( + local_id_to_global_id_by_node_type=cuda_node_map, + positive_labels_by_edge_type=cuda_positive_labels, + negative_labels_by_edge_type={}, + ) + + self.assertEqual(set(actual_positive.keys()), set(expected_positive.keys())) + for edge_type, expected_edge_index in expected_positive.items(): + actual_edge_index = actual_positive[edge_type] + self.assertEqual(actual_edge_index.device.type, "cuda") + torch.testing.assert_close( + actual_edge_index.cpu(), + expected_edge_index, + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/unit/distributed/utils/ablp_test.py b/tests/unit/distributed/utils/ablp_test.py new file mode 100644 index 000000000..a65726d65 --- /dev/null +++ b/tests/unit/distributed/utils/ablp_test.py @@ -0,0 +1,285 @@ +"""Tests for vectorized ABLP label remapping.""" + +import unittest + +import torch +from parameterized import param, parameterized +from torch_geometric.typing import EdgeType as PyGEdgeType + +from gigl.distributed.utils.ablp import ( + label_edge_index_to_dict, + remap_labels_to_local_edge_indices, +) +from gigl.src.common.types.graph_data import EdgeType, NodeType, Relation +from gigl.types.graph import ( + DEFAULT_HOMOGENEOUS_EDGE_TYPE, + DEFAULT_HOMOGENEOUS_NODE_TYPE, + message_passing_to_negative_label, + message_passing_to_positive_label, +) +from tests.test_assets.test_case import TestCase + +_USER = NodeType("user") +_STORY = NodeType("story") +_USER_TO_STORY = EdgeType(_USER, Relation("to"), _STORY) +_A = NodeType("a") +_B = NodeType("b") +_C = NodeType("c") +_A_TO_B = EdgeType(_A, Relation("to"), _B) +_A_TO_C = EdgeType(_A, Relation("to"), _C) + + +def _positive_label_edge_type(edge_type: EdgeType) -> EdgeType: + return message_passing_to_positive_label(edge_type) + + +def _negative_label_edge_type(edge_type: EdgeType) -> EdgeType: + return message_passing_to_negative_label(edge_type) + + +def _assert_label_sets_equal( + actual: dict[PyGEdgeType, torch.Tensor], + expected: dict[PyGEdgeType, dict[int, list[int]]], +) -> None: + assert set(actual.keys()) == set(expected.keys()) + for edge_type, expected_by_anchor in expected.items(): + actual_by_anchor = label_edge_index_to_dict( + label_edge_index=actual[edge_type], + num_anchors=len(expected_by_anchor), + ) + assert set(actual_by_anchor.keys()) == set(expected_by_anchor.keys()) + for anchor, expected_labels in expected_by_anchor.items(): + actual_labels = actual_by_anchor[anchor] + assert actual_labels.dtype == torch.long + assert sorted(actual_labels.tolist()) == sorted(expected_labels) + + +class LabelEdgeIndexToDictTest(TestCase): + def test_preserves_empty_and_multi_label_anchors(self) -> None: + label_edge_index = torch.tensor([[0, 2, 2], [5, 7, 8]]) + + labels_by_anchor = label_edge_index_to_dict( + label_edge_index=label_edge_index, + num_anchors=3, + ) + + self.assertEqual(set(labels_by_anchor.keys()), {0, 1, 2}) + torch.testing.assert_close( + labels_by_anchor[0], torch.tensor([5], dtype=torch.long) + ) + torch.testing.assert_close( + labels_by_anchor[1], torch.empty(0, dtype=torch.long) + ) + torch.testing.assert_close( + labels_by_anchor[2], torch.tensor([7, 8], dtype=torch.long) + ) + + def test_all_anchors_empty(self) -> None: + labels_by_anchor = label_edge_index_to_dict( + label_edge_index=torch.empty((2, 0), dtype=torch.long), + num_anchors=2, + ) + + self.assertEqual(set(labels_by_anchor.keys()), {0, 1}) + torch.testing.assert_close( + labels_by_anchor[0], torch.empty(0, dtype=torch.long) + ) + torch.testing.assert_close( + labels_by_anchor[1], torch.empty(0, dtype=torch.long) + ) + + +class RemapLabelsToLocalEdgeIndicesTest(TestCase): + @parameterized.expand( + [ + param( + "sorted_present_empty_and_padded", + local_id_to_global_id_by_node_type={ + _STORY: torch.tensor([10, 11, 12, 13, 14, 15, 16, 17]) + }, + positive_labels={ + _positive_label_edge_type(_USER_TO_STORY): torch.tensor( + [[15, -1], [15, 16], [-1, -1], [99, -1]] + ) + }, + negative_labels={}, + expected_positive={_USER_TO_STORY: {0: [5], 1: [5, 6], 2: [], 3: []}}, + expected_negative={}, + ), + param( + "duplicate_labels_preserve_multiplicity", + local_id_to_global_id_by_node_type={ + _STORY: torch.tensor([10, 11, 12, 13, 14, 15]) + }, + positive_labels={ + _positive_label_edge_type(_USER_TO_STORY): torch.tensor( + [[15, 15], [11, 11]] + ) + }, + negative_labels={}, + expected_positive={_USER_TO_STORY: {0: [5, 5], 1: [1, 1]}}, + expected_negative={}, + ), + param( + "default_homogeneous_keying", + local_id_to_global_id_by_node_type={ + DEFAULT_HOMOGENEOUS_NODE_TYPE: torch.tensor([20, 10, 30, 11, 15]) + }, + positive_labels={ + message_passing_to_positive_label( + DEFAULT_HOMOGENEOUS_EDGE_TYPE + ): torch.tensor([[30, 10], [-1, -1]]) + }, + negative_labels={}, + expected_positive={DEFAULT_HOMOGENEOUS_EDGE_TYPE: {0: [2, 1], 1: []}}, + expected_negative={}, + ), + param( + "positive_and_negative_labels", + local_id_to_global_id_by_node_type={ + _STORY: torch.tensor([10, 11, 12, 13, 14, 15, 16, 17]) + }, + positive_labels={ + _positive_label_edge_type(_USER_TO_STORY): torch.tensor( + [[15], [16]] + ) + }, + negative_labels={ + _negative_label_edge_type(_USER_TO_STORY): torch.tensor( + [[13, 16], [17, -1]] + ) + }, + expected_positive={_USER_TO_STORY: {0: [5], 1: [6]}}, + expected_negative={_USER_TO_STORY: {0: [3, 6], 1: [7]}}, + ), + param( + "multiple_supervision_edge_types", + local_id_to_global_id_by_node_type={ + _B: torch.tensor([11, 12, 13, 14, 15, 16]), + _C: torch.tensor([20, 21, 22, 23, 24, 25]), + }, + positive_labels={ + _positive_label_edge_type(_A_TO_B): torch.tensor([[13, 14]]), + _positive_label_edge_type(_A_TO_C): torch.tensor([[22, 23]]), + }, + negative_labels={ + _negative_label_edge_type(_A_TO_B): torch.tensor([[15, 16]]), + _negative_label_edge_type(_A_TO_C): torch.tensor([[24, 25]]), + }, + expected_positive={ + _A_TO_B: {0: [2, 3]}, + _A_TO_C: {0: [2, 3]}, + }, + expected_negative={ + _A_TO_B: {0: [4, 5]}, + _A_TO_C: {0: [4, 5]}, + }, + ), + param( + "all_anchors_empty", + local_id_to_global_id_by_node_type={_STORY: torch.tensor([10, 11, 12])}, + positive_labels={ + _positive_label_edge_type(_USER_TO_STORY): torch.tensor( + [[-1, -1], [99, 98]] + ) + }, + negative_labels={}, + expected_positive={_USER_TO_STORY: {0: [], 1: []}}, + expected_negative={}, + ), + ] + ) + def test_matches_constructed_labels( + self, + _, + local_id_to_global_id_by_node_type: dict[NodeType, torch.Tensor], + positive_labels: dict[EdgeType, torch.Tensor], + negative_labels: dict[EdgeType, torch.Tensor], + expected_positive: dict[PyGEdgeType, dict[int, list[int]]], + expected_negative: dict[PyGEdgeType, dict[int, list[int]]], + ) -> None: + positive_edge_indices, negative_edge_indices = ( + remap_labels_to_local_edge_indices( + local_id_to_global_id_by_node_type=local_id_to_global_id_by_node_type, + positive_labels_by_edge_type=positive_labels, + negative_labels_by_edge_type=negative_labels, + ) + ) + + _assert_label_sets_equal(positive_edge_indices, expected_positive) + _assert_label_sets_equal(negative_edge_indices, expected_negative) + + def test_pins_exact_edge_index(self) -> None: + positive_edge_indices, _ = remap_labels_to_local_edge_indices( + local_id_to_global_id_by_node_type={ + _STORY: torch.tensor([10, 11, 12, 13, 14, 15, 16, 17]) + }, + positive_labels_by_edge_type={ + _positive_label_edge_type(_USER_TO_STORY): torch.tensor( + [[15, -1], [15, 16], [-1, -1]] + ) + }, + negative_labels_by_edge_type={}, + ) + + torch.testing.assert_close( + positive_edge_indices[_USER_TO_STORY], + torch.tensor([[0, 1, 1], [5, 5, 6]], dtype=torch.long), + ) + + def test_unsorted_node_map_returns_local_indices(self) -> None: + positive_edge_indices, _ = remap_labels_to_local_edge_indices( + local_id_to_global_id_by_node_type={_STORY: torch.tensor([15, 10, 16, 11])}, + positive_labels_by_edge_type={ + _positive_label_edge_type(_USER_TO_STORY): torch.tensor([[16, 15]]) + }, + negative_labels_by_edge_type={}, + ) + + labels_by_anchor = label_edge_index_to_dict( + positive_edge_indices[_USER_TO_STORY], num_anchors=1 + ) + self.assertEqual(sorted(labels_by_anchor[0].tolist()), [0, 2]) + + def test_zero_anchor_tensor_yields_no_edge_type_key(self) -> None: + positive_edge_indices, negative_edge_indices = ( + remap_labels_to_local_edge_indices( + local_id_to_global_id_by_node_type={_STORY: torch.tensor([10, 11, 12])}, + positive_labels_by_edge_type={ + _positive_label_edge_type(_USER_TO_STORY): torch.empty((0, 0)) + }, + negative_labels_by_edge_type={}, + ) + ) + + self.assertEqual(positive_edge_indices, {}) + self.assertEqual(negative_edge_indices, {}) + + def test_output_follows_input_device_and_dtype(self) -> None: + positive_edge_indices, _ = remap_labels_to_local_edge_indices( + local_id_to_global_id_by_node_type={ + _STORY: torch.tensor([10, 11, 12, 13, 14, 15]) + }, + positive_labels_by_edge_type={ + _positive_label_edge_type(_USER_TO_STORY): torch.tensor([[15], [11]]) + }, + negative_labels_by_edge_type={}, + ) + + label_edge_index = positive_edge_indices[_USER_TO_STORY] + self.assertEqual(label_edge_index.device.type, "cpu") + self.assertEqual(label_edge_index.dtype, torch.long) + + def test_duplicate_node_map_raises(self) -> None: + with self.assertRaises(ValueError): + remap_labels_to_local_edge_indices( + local_id_to_global_id_by_node_type={_STORY: torch.tensor([10, 10, 11])}, + positive_labels_by_edge_type={ + _positive_label_edge_type(_USER_TO_STORY): torch.tensor([[10, 11]]) + }, + negative_labels_by_edge_type={}, + ) + + +if __name__ == "__main__": + unittest.main()