From f5c430dfe780f08493e902d16aaa475159f9b985 Mon Sep 17 00:00:00 2001 From: mkolodner Date: Wed, 15 Jul 2026 18:29:37 +0000 Subject: [PATCH 1/3] Clean up PPR sampler naming --- gigl/distributed/dist_ppr_sampler.py | 282 ++++++++++++++------------- 1 file changed, 150 insertions(+), 132 deletions(-) diff --git a/gigl/distributed/dist_ppr_sampler.py b/gigl/distributed/dist_ppr_sampler.py index 204b1c9b4..1ba78c9fc 100644 --- a/gigl/distributed/dist_ppr_sampler.py +++ b/gigl/distributed/dist_ppr_sampler.py @@ -71,8 +71,8 @@ class DistPPRNeighborSampler(BaseDistNeighborSampler): **Heterogeneous (HeteroData)** — one PPR edge type per ``(seed_type, neighbor_type)`` pair, with ``"ppr"`` as the relation: - - ``data[(seed_type, "ppr", ntype)].edge_index``: same format as above. - - ``data[(seed_type, "ppr", ntype)].edge_attr``: same format as above. + - ``data[(seed_type, "ppr", neighbor_type)].edge_index``: same format as above. + - ``data[(seed_type, "ppr", neighbor_type)].edge_attr``: same format as above. Args: alpha: Restart probability (teleport probability back to seed). Higher values @@ -130,17 +130,17 @@ def __init__( # Label edge types (injected by ABLP for supervision) are excluded: including # them would let PPR walks cross label edges, leaking ground-truth targets into # the sampled neighborhood. - for etype in self.edge_types: - if is_label_edge_type(etype): + for edge_type in self.edge_types: + if is_label_edge_type(edge_type): continue if self.edge_dir == "in": # For incoming edges, we traverse FROM the destination node type - anchor_type = etype[-1] + anchor_type = edge_type[-1] else: # "out" # For outgoing edges, we traverse FROM the source node type - anchor_type = etype[0] + anchor_type = edge_type[0] - self._node_type_to_edge_types[anchor_type].append(etype) + self._node_type_to_edge_types[anchor_type].append(edge_type) else: self._node_type_to_edge_types[DEFAULT_HOMOGENEOUS_NODE_TYPE] = [ _PPR_HOMOGENEOUS_EDGE_TYPE @@ -163,42 +163,50 @@ def __init__( # so the kernel can index residual/ppr_score tables for any node it sees. source_node_types: set[NodeType] = set(self._node_type_to_edge_types.keys()) destination_node_types: set[NodeType] = { - self._get_destination_type(et) - for etypes in self._node_type_to_edge_types.values() - for et in etypes + self._get_destination_type(edge_type) + for edge_types in self._node_type_to_edge_types.values() + for edge_type in edge_types } all_node_types: list[NodeType] = sorted( source_node_types | destination_node_types ) all_edge_types: list[EdgeType] = sorted( - {et for etypes in self._node_type_to_edge_types.values() for et in etypes} + { + edge_type + for edge_types in self._node_type_to_edge_types.values() + for edge_type in edge_types + } ) self._node_type_to_id: dict[NodeType, int] = { - nt: i for i, nt in enumerate(all_node_types) + node_type: node_type_id + for node_type_id, node_type in enumerate(all_node_types) } self._ntype_id_to_ntype: list[NodeType] = all_node_types self._etype_to_etype_id: dict[EdgeType, int] = { - et: i for i, et in enumerate(all_edge_types) + edge_type: edge_type_id + for edge_type_id, edge_type in enumerate(all_edge_types) } self._etype_id_to_etype: list[EdgeType] = all_edge_types self._node_type_id_to_edge_type_ids: list[list[int]] = [ [ - self._etype_to_etype_id[et] - for et in self._node_type_to_edge_types.get(nt, []) + self._etype_to_etype_id[edge_type] + for edge_type in self._node_type_to_edge_types.get(node_type, []) ] - for nt in all_node_types + for node_type in all_node_types ] self._edge_type_id_to_dst_ntype_id: list[int] = [ - self._node_type_to_id[self._get_destination_type(et)] - for et in all_edge_types + self._node_type_to_id[self._get_destination_type(edge_type)] + for edge_type in all_edge_types ] # Degree tensors indexed by ntype_id. Destination-only types get an empty # tensor; the C++ kernel returns 0 for those, matching _get_total_degree. self._degree_tensors_for_cpp: list[torch.Tensor] = [ - self._node_type_to_total_degree.get(nt, torch.zeros(0, dtype=torch.int32)) - for nt in all_node_types + self._node_type_to_total_degree.get( + node_type, torch.zeros(0, dtype=torch.int32) + ) + for node_type in all_node_types ] def _convert_degree_tensors_to_dict( @@ -228,33 +236,32 @@ def _get_destination_type(self, edge_type: EdgeType) -> NodeType: async def _batch_fetch_neighbors( self, - nodes_by_etype_id: dict[int, torch.Tensor], + nodes_by_edge_type_id: dict[int, torch.Tensor], device: torch.device, ) -> dict[int, tuple[torch.Tensor, torch.Tensor, torch.Tensor]]: """Batch fetch neighbors for nodes grouped by integer edge type ID. - Issues one ``_sample_one_hop`` call per edge type (not per node), so all - nodes of the same edge type are fetched in a single RPC round-trip. Each - node's neighbor list is capped at ``self._num_neighbors_per_hop``. + Issues one one-hop request per edge type in the frontier. Each node's + neighbor list is capped at ``self._num_neighbors_per_hop``. Args: - nodes_by_etype_id: Dict mapping integer edge type ID to a 1-D int64 + nodes_by_edge_type_id: Dict mapping integer edge type ID to a 1-D int64 tensor of node IDs to fetch neighbors for. Comes directly from ``drain_queue()``; node IDs are already deduplicated. device: Torch device for intermediate tensor creation. Returns: - Dict mapping etype_id to ``(node_ids, flat_neighbors, counts)`` as - int64 tensors, ready to pass directly to ``push_residuals``. + Dict mapping edge type ID to ``(node_ids, flat_neighbors, counts)`` + as int64 tensors, ready to pass directly to ``push_residuals``. ``flat_neighbors`` is the flat concatenation of all neighbor lists for that edge type; ``counts[i]`` is the neighbor count for ``node_ids[i]``. Example:: - nodes_by_etype_id = { - 2: tensor([0, 3]), # etype_id 2 → nodes 0 and 3 - 5: tensor([7]), # etype_id 5 → node 7 + nodes_by_edge_type_id = { + 2: tensor([0, 3]), # edge type ID 2 -> nodes 0 and 3 + 5: tensor([7]), # edge type ID 5 -> node 7 } # Might return (neighbor lists depend on graph structure): { @@ -262,32 +269,35 @@ async def _batch_fetch_neighbors( 5: (tensor([7]), tensor([0, 3]), tensor([2])), } """ - # Fire all per-edge-type RPC calls concurrently. Each _sample_one_hop - # issues a single RPC round-trip; doing them in parallel rather than - # sequentially cuts fetch latency from O(num_edge_types) to O(1). - eids = list(nodes_by_etype_id.keys()) + edge_type_ids: list[int] = [] sample_tasks = [] - for eid in eids: - etype = self._etype_id_to_etype[eid] + for edge_type_id in nodes_by_edge_type_id: + edge_type = self._etype_id_to_etype[edge_type_id] + # _sample_one_hop expects None only for true homogeneous graphs. + # Labeled homogeneous ABLP graphs are hetero-backed because label + # edges are represented as separate edge types, so they still need + # the explicit default edge type here. + rpc_edge_type = ( + None + if self._is_homogeneous and edge_type == _PPR_HOMOGENEOUS_EDGE_TYPE + else edge_type + ) + edge_type_ids.append(edge_type_id) sample_tasks.append( self._sample_one_hop( - srcs=nodes_by_etype_id[eid].to(device), + srcs=nodes_by_edge_type_id[edge_type_id].to(device), num_nbr=self._num_neighbors_per_hop, - # _sample_one_hop expects None only for true homogeneous graphs. - # Labeled homogeneous ABLP graphs are hetero-backed because label - # edges are represented as separate edge types, so they still need - # the explicit default edge type here. - etype=( - None - if self._is_homogeneous and etype == _PPR_HOMOGENEOUS_EDGE_TYPE - else etype - ), + etype=rpc_edge_type, ) ) outputs: list[NeighborOutput] = await asyncio.gather(*sample_tasks) return { - eid: (nodes_by_etype_id[eid], output.nbr, output.nbr_num) - for eid, output in zip(eids, outputs) + edge_type_id: ( + nodes_by_edge_type_id[edge_type_id], + output.nbr, + output.nbr_num, + ) + for edge_type_id, output in zip(edge_type_ids, outputs) } def _extract_ppr_state_top_k( @@ -307,10 +317,9 @@ def _extract_ppr_state_top_k( preserves the homogeneous return shape expected by the rest of the sampler. - ``max_ppr_nodes`` is the combined per-seed cap across finalized PPR and - residual top-up candidates. If residual top-up is enabled, C++ derives - the residual candidate budget from this cap after selecting finalized - PPR nodes. + ``max_ppr_nodes`` is the maximum number of nodes returned for each + source node. If residual top-up is enabled, residual candidates count + against this cap. Returns: ``(flat_ids, flat_weights, valid_counts)`` for homogeneous graphs, @@ -318,42 +327,44 @@ def _extract_ppr_state_top_k( ``flat_ids`` and ``flat_weights`` are concatenated across seeds; ``valid_counts`` stores how many selected nodes belong to each seed. """ - # Translate ntype_id integer keys back to NodeType strings for the rest + # Translate integer node-type IDs back to NodeType strings for the rest # of the pipeline, and move tensors to the correct device. - ntype_to_flat_ids: dict[NodeType, torch.Tensor] = {} - ntype_to_flat_weights: dict[NodeType, torch.Tensor] = {} - ntype_to_valid_counts: dict[NodeType, torch.Tensor] = {} + node_type_to_flat_ids: dict[NodeType, torch.Tensor] = {} + node_type_to_flat_weights: dict[NodeType, torch.Tensor] = {} + node_type_to_valid_counts: dict[NodeType, torch.Tensor] = {} extracted_results = ppr_state.extract_top_k_with_residual_top_up( max_ppr_nodes, self._enable_residual_topup, ) - for ntype_id, ( + for node_type_id, ( flat_ids, flat_weights, valid_counts, ) in extracted_results.items(): - ntype = self._ntype_id_to_ntype[ntype_id] - ntype_to_flat_ids[ntype] = flat_ids.to(device) - ntype_to_flat_weights[ntype] = flat_weights.to(device) - ntype_to_valid_counts[ntype] = valid_counts.to(device) + node_type = self._ntype_id_to_ntype[node_type_id] + # TODO: If these copies become a bottleneck, evaluate + # non_blocking=True together with pinned extraction output tensors. + node_type_to_flat_ids[node_type] = flat_ids.to(device) + node_type_to_flat_weights[node_type] = flat_weights.to(device) + node_type_to_valid_counts[node_type] = valid_counts.to(device) if self._is_homogeneous: assert ( - len(ntype_to_flat_ids) == 1 - and DEFAULT_HOMOGENEOUS_NODE_TYPE in ntype_to_flat_ids + len(node_type_to_flat_ids) == 1 + and DEFAULT_HOMOGENEOUS_NODE_TYPE in node_type_to_flat_ids ) return ( - ntype_to_flat_ids[DEFAULT_HOMOGENEOUS_NODE_TYPE], - ntype_to_flat_weights[DEFAULT_HOMOGENEOUS_NODE_TYPE], - ntype_to_valid_counts[DEFAULT_HOMOGENEOUS_NODE_TYPE], + node_type_to_flat_ids[DEFAULT_HOMOGENEOUS_NODE_TYPE], + node_type_to_flat_weights[DEFAULT_HOMOGENEOUS_NODE_TYPE], + node_type_to_valid_counts[DEFAULT_HOMOGENEOUS_NODE_TYPE], ) else: return ( - ntype_to_flat_ids, - ntype_to_flat_weights, - ntype_to_valid_counts, + node_type_to_flat_ids, + node_type_to_flat_weights, + node_type_to_valid_counts, ) async def _compute_ppr_scores( @@ -435,26 +446,26 @@ async def _compute_ppr_scores( # means all drained nodes either had cached neighbors or no outgoing # edges — we still call push_residuals to flush their residuals into # ppr_scores_. - nodes_by_etype_id = ppr_state.drain_queue() - if nodes_by_etype_id is None: + nodes_by_edge_type_id = ppr_state.drain_queue() + if nodes_by_edge_type_id is None: break fetch_budget_remaining = ( self._max_fetch_iterations is None or fetch_iteration_count < self._max_fetch_iterations ) - if nodes_by_etype_id and fetch_budget_remaining: - fetched_by_etype_id = await self._batch_fetch_neighbors( - nodes_by_etype_id, device + if nodes_by_edge_type_id and fetch_budget_remaining: + fetched_by_edge_type_id = await self._batch_fetch_neighbors( + nodes_by_edge_type_id, device ) fetch_iteration_count += 1 else: # Fetch budget exhausted; push_residuals will use the existing neighbor cache. - fetched_by_etype_id = {} + fetched_by_edge_type_id = {} # Run in executor so the C++ push doesn't block the asyncio event loop. await asyncio.get_running_loop().run_in_executor( - None, ppr_state.push_residuals, fetched_by_etype_id + None, ppr_state.push_residuals, fetched_by_edge_type_id ) # Regular sampling uses residual top-up as fill-only under the @@ -484,7 +495,7 @@ async def _sample_from_nodes( ``edge_attr`` fields on the output Data/HeteroData). Local indices are produced by the inducer (see below), so row 1 of - ``edge_index`` directly indexes into ``data[ntype].x`` without any + ``edge_index`` directly indexes into ``data[node_type].x`` without any additional global→local remapping. The inducer is GLT's C++ data structure (backed by a per-node-type hash map) @@ -494,7 +505,7 @@ async def _sample_from_nodes( 1. **Consistency across seed types.** For heterogeneous ABLP inputs, ``_compute_ppr_scores`` is called once per seed type (anchors, supervision nodes, …). A node reachable from multiple seed types must receive the - *same* local index in ``node_dict[ntype]`` regardless of which seed type + *same* local index in ``node_dict[node_type]`` regardless of which seed type discovered it. The inducer is shared across all those calls, so it guarantees this automatically. @@ -506,9 +517,9 @@ async def _sample_from_nodes( - ``inducer.init_node(seeds)`` registers seed nodes and returns their global IDs (local indices 0, 1, … are assigned internally). - - ``inducer.induce_next(srcs, flat_nbrs, counts)`` (homo) or - ``inducer.induce_next(nbr_dict)`` (hetero) deduplicates neighbors against - all previously seen nodes and returns: + - ``inducer.induce_next(source_nodes, flat_neighbors, counts)`` (homogeneous) + or ``inducer.induce_next(neighbor_dict)`` (heterogeneous) deduplicates + neighbors against all previously seen nodes and returns: - ``new_nodes``: global IDs of nodes not previously registered with the inducer (i.e., not seeds and not returned by a prior @@ -530,7 +541,7 @@ async def _sample_from_nodes( # # 1. Deduplication: when the same global node ID appears from multiple # seeds or seed types, induce_next assigns it a single local index. - # This ensures node_dict[ntype] has no duplicates. + # This ensures node_dict[node_type] has no duplicates. # # 2. Local index assignment: init_node registers seeds at local indices # 0..N-1. induce_next then assigns the next available indices to @@ -545,18 +556,20 @@ async def _sample_from_nodes( if is_hetero: assert isinstance(nodes_to_sample, dict) assert input_type is not None + nodes_by_seed_type = nodes_to_sample # Register all seeds (anchors + supervision nodes for ABLP) with the - # inducer first, so they occupy the lowest local indices. src_dict maps + # inducer first, so they occupy the lowest local indices. source_dict maps # NodeType -> global IDs (same values as nodes_to_sample). - src_dict = inducer.init_node(nodes_to_sample) + source_dict: dict[NodeType, torch.Tensor] = inducer.init_node( + nodes_by_seed_type + ) # Compute PPR for all seed types concurrently, collecting flat global - # neighbor IDs, weights, and per-seed counts. Build nbr_dict for a + # neighbor IDs, weights, and per-seed counts. Build neighbor_dict for a # single inducer.induce_next call using PPR edge types - # (seed_type, 'ppr', ntype) — the inducer only cares about etype[0] - # and etype[-1] as source/dest node types, so the relation name is - # arbitrary. + # (seed_type, 'ppr', neighbor_type). The inducer only cares about + # source and destination node types, so the relation name is arbitrary. # # Each seed type's PPR computation is entirely independent: it creates # its own PPRForwardPush and only reads shared sampler attributes @@ -564,44 +577,47 @@ async def _sample_from_nodes( # Running them with asyncio.gather allows their fetch phases to overlap, # which is most beneficial when there are 2+ distinct seed node types # (e.g. cross-type supervision edges like user→story). - seed_types = list(nodes_to_sample.keys()) + seed_types = list(nodes_by_seed_type.keys()) ppr_results = await asyncio.gather( *[ - self._compute_ppr_scores(nodes_to_sample[seed_type], seed_type) # ty: ignore[invalid-argument-type] TODO(ty-torch-keyed-access): fix ty false positives for torch-backed keyed container access. + self._compute_ppr_scores( + nodes_by_seed_type[seed_type], # ty: ignore[invalid-argument-type] TODO(ty-torch-keyed-access): fix ty false positives for torch-backed keyed container access. + seed_type, + ) for seed_type in seed_types ] ) - nbr_dict: dict[EdgeType, list[torch.Tensor]] = {} + neighbor_dict: dict[EdgeType, list[torch.Tensor]] = {} ppr_edge_type_to_flat_weights: dict[EdgeType, torch.Tensor] = {} for seed_type, ( - ntype_to_flat_ids, - ntype_to_flat_weights, - ntype_to_valid_counts, + node_type_to_flat_ids, + node_type_to_flat_weights, + node_type_to_valid_counts, ) in zip(seed_types, ppr_results): - assert isinstance(ntype_to_flat_ids, dict) - assert isinstance(ntype_to_flat_weights, dict) - assert isinstance(ntype_to_valid_counts, dict) + assert isinstance(node_type_to_flat_ids, dict) + assert isinstance(node_type_to_flat_weights, dict) + assert isinstance(node_type_to_valid_counts, dict) - for ntype, flat_ids in ntype_to_flat_ids.items(): - ppr_edge_type: EdgeType = (seed_type, "ppr", ntype) - valid_counts = ntype_to_valid_counts[ntype] # ty: ignore[invalid-argument-type] TODO(ty-torch-keyed-access): fix ty false positives for torch-backed keyed container access. + for node_type, flat_ids in node_type_to_flat_ids.items(): + ppr_edge_type: EdgeType = (seed_type, "ppr", node_type) + valid_counts = node_type_to_valid_counts[node_type] # ty: ignore[invalid-argument-type] TODO(ty-torch-keyed-access): fix ty false positives for torch-backed keyed container access. ppr_edge_type_to_flat_weights[ppr_edge_type] = ( - ntype_to_flat_weights[ntype] # ty: ignore[invalid-argument-type] TODO(ty-torch-keyed-access): fix ty false positives for torch-backed keyed container access. + node_type_to_flat_weights[node_type] # ty: ignore[invalid-argument-type] TODO(ty-torch-keyed-access): fix ty false positives for torch-backed keyed container access. ) # Skip empty pairs; induce_next handles deduplication across # seed types so a neighbor reachable from multiple seed types - # gets one consistent local index in node_dict[ntype]. + # gets one consistent local index in node_dict[node_type]. if flat_ids.numel() > 0: # ty: ignore[unresolved-attribute] TODO(ty-torch-keyed-access): fix ty false positives for torch-backed keyed container access. - nbr_dict[ppr_edge_type] = [ - src_dict[seed_type], + neighbor_dict[ppr_edge_type] = [ # ty: ignore[invalid-assignment] TODO(ty-torch-keyed-access): fix ty false positives for torch-backed keyed container access. + source_dict[seed_type], flat_ids, valid_counts, - ] # ty: ignore[invalid-assignment] TODO(ty-torch-container-shapes): fix ty false positives for torch container and return shapes. + ] - # induce_next processes all PPR edge types in nbr_dict in one + # induce_next processes all PPR edge types in neighbor_dict in one # pass, assigning local indices to neighbors not yet registered and # deduplicating nodes seen from multiple seed types. Returns: # new_nodes_dict[NodeType] -> global IDs of nodes not previously @@ -610,18 +626,20 @@ async def _sample_from_nodes( # edge type, expanded to match flat_ids # cols_dict[EdgeType] -> flat local destination indices, one # per neighbor in the same order as the - # flat_ids passed in nbr_dict - new_nodes_dict, rows_dict, cols_dict = inducer.induce_next(nbr_dict) + # flat_ids passed in neighbor_dict + new_nodes_dict, rows_dict, cols_dict = inducer.induce_next(neighbor_dict) - # node_dict = seeds (already in src_dict) + PPR neighbors not + # node_dict = seeds (already in source_dict) + PPR neighbors not # previously registered. merge_dict appends tensors into lists; # cat collapses them. - out_nodes_hetero: dict[NodeType, list[torch.Tensor]] = defaultdict(list) - merge_dict(src_dict, out_nodes_hetero) - merge_dict(new_nodes_dict, out_nodes_hetero) + heterogeneous_output_nodes: dict[NodeType, list[torch.Tensor]] = ( + defaultdict(list) + ) + merge_dict(source_dict, heterogeneous_output_nodes) + merge_dict(new_nodes_dict, heterogeneous_output_nodes) node_dict = { - ntype: torch.cat(nodes) - for ntype, nodes in out_nodes_hetero.items() + node_type: torch.cat(nodes) + for node_type, nodes in heterogeneous_output_nodes.items() if nodes } @@ -642,9 +660,9 @@ async def _sample_from_nodes( flat_weights = torch.zeros( 0, dtype=torch.double, device=self.device ) - etype_str = repr(ppr_edge_type) - metadata[f"{PPR_EDGE_INDEX_METADATA_KEY}{etype_str}"] = edge_index - metadata[f"{PPR_WEIGHT_METADATA_KEY}{etype_str}"] = flat_weights + edge_type_repr = repr(ppr_edge_type) + metadata[f"{PPR_EDGE_INDEX_METADATA_KEY}{edge_type_repr}"] = edge_index + metadata[f"{PPR_WEIGHT_METADATA_KEY}{edge_type_repr}"] = flat_weights sample_output = HeteroSamplerOutput( node=node_dict, @@ -658,7 +676,7 @@ async def _sample_from_nodes( edge={}, # Empty dict — GLT SampleQueue requires all values to be tensors batch={input_type: input_seeds}, num_sampled_nodes={ - ntype: [nodes.size(0)] for ntype, nodes in node_dict.items() + node_type: [nodes.size(0)] for node_type, nodes in node_dict.items() }, num_sampled_edges={}, input_type=input_type, @@ -684,34 +702,34 @@ async def _sample_from_nodes( ) # Register seeds; local indices 0..N-1 are assigned internally. - # srcs holds their global IDs (same values as nodes_to_sample). - srcs = inducer.init_node(homogeneous_nodes_to_sample) + # source_nodes holds their global IDs (same values as nodes_to_sample). + source_nodes = inducer.init_node(homogeneous_nodes_to_sample) ( - homo_flat_ids, - homo_flat_weights, - homo_valid_counts, + homogeneous_flat_ids, + homogeneous_flat_weights, + homogeneous_valid_counts, ) = await self._compute_ppr_scores(homogeneous_nodes_to_sample, None) - assert isinstance(homo_flat_ids, torch.Tensor) - assert isinstance(homo_flat_weights, torch.Tensor) - assert isinstance(homo_valid_counts, torch.Tensor) + assert isinstance(homogeneous_flat_ids, torch.Tensor) + assert isinstance(homogeneous_flat_weights, torch.Tensor) + assert isinstance(homogeneous_valid_counts, torch.Tensor) - # induce_next deduplicates homo_flat_ids against already-seen nodes + # induce_next deduplicates homogeneous_flat_ids against already-seen nodes # (the seeds registered above) and returns: # new_nodes: global IDs of nodes not previously registered # with the inducer. # rows: flat local source indices (one per neighbor, expanded). # cols: flat local destination indices for every neighbor, in the - # same order as homo_flat_ids. + # same order as homogeneous_flat_ids. new_nodes, rows, cols = inducer.induce_next( - srcs, homo_flat_ids, homo_valid_counts + source_nodes, homogeneous_flat_ids, homogeneous_valid_counts ) - all_nodes = torch.cat([srcs, new_nodes]) + all_nodes = torch.cat([source_nodes, new_nodes]) ppr_edge_index = torch.stack([rows, cols]) metadata["edge_index"] = ppr_edge_index - metadata["edge_attr"] = homo_flat_weights + metadata["edge_attr"] = homogeneous_flat_weights sample_output = SamplerOutput( node=all_nodes, @@ -722,7 +740,7 @@ async def _sample_from_nodes( [], dtype=torch.long, device=self.device ), # Empty tensor — GLT SampleQueue requires all values to be tensors batch=input_seeds, - num_sampled_nodes=[srcs.size(0), new_nodes.size(0)], + num_sampled_nodes=[source_nodes.size(0), new_nodes.size(0)], num_sampled_edges=[], metadata=metadata, ) From 50da87627892f6d41fa3e57ee7a5aa0bb2134b99 Mon Sep 17 00:00:00 2001 From: mkolodner Date: Wed, 15 Jul 2026 18:32:26 +0000 Subject: [PATCH 2/3] Move existing PPR work off event loop --- .../core/sampling/python_ppr_forward_push.cpp | 31 +++++++++++++++++-- gigl/distributed/dist_ppr_sampler.py | 17 +++++++--- 2 files changed, 40 insertions(+), 8 deletions(-) diff --git a/gigl-core/core/sampling/python_ppr_forward_push.cpp b/gigl-core/core/sampling/python_ppr_forward_push.cpp index 27271559c..5f1f474f6 100644 --- a/gigl-core/core/sampling/python_ppr_forward_push.cpp +++ b/gigl-core/core/sampling/python_ppr_forward_push.cpp @@ -8,6 +8,7 @@ #include #include +#include #include #include @@ -42,6 +43,29 @@ static void pushResidualsWrapper(PPRForwardPush& state, const py::dict& fetchedB } } +static std::optional> drainQueueWrapper(PPRForwardPush& state) { + std::optional> drained; + // The drain mutates only this PPRForwardPush instance and returns torch tensors. + // No Python objects are touched until pybind converts the result after return. + { + py::gil_scoped_release release; + drained = state.drainQueue(); + } + return drained; +} + +static std::unordered_map> +extractTopKWithResidualTopUpWrapper(PPRForwardPush& state, int32_t maxPPRNodes, bool enableResidualTopUp) { + std::unordered_map> result; + // Extraction walks C++ state and builds torch tensors; Python only sees the + // result after pybind converts it on return. + { + py::gil_scoped_release release; + result = state.extractTopKWithResidualTopUp(maxPPRNodes, enableResidualTopUp); + } + return result; +} + } // namespace gigl // TORCH_EXTENSION_NAME is set by PyTorch's build system to match the Python @@ -54,11 +78,12 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { double, std::vector>, std::vector, - std::vector>()) - .def("drain_queue", &gigl::PPRForwardPush::drainQueue) + std::vector>(), + py::call_guard()) + .def("drain_queue", gigl::drainQueueWrapper) .def("push_residuals", gigl::pushResidualsWrapper) .def("extract_top_k_with_residual_top_up", - &gigl::PPRForwardPush::extractTopKWithResidualTopUp, + gigl::extractTopKWithResidualTopUpWrapper, py::arg("max_ppr_nodes"), py::arg("enable_residual_topup")); } diff --git a/gigl/distributed/dist_ppr_sampler.py b/gigl/distributed/dist_ppr_sampler.py index 1ba78c9fc..8701a1dcc 100644 --- a/gigl/distributed/dist_ppr_sampler.py +++ b/gigl/distributed/dist_ppr_sampler.py @@ -427,8 +427,11 @@ async def _compute_ppr_scores( if seed_node_type is None: seed_node_type = DEFAULT_HOMOGENEOUS_NODE_TYPE device = seed_nodes.device + loop = asyncio.get_running_loop() - ppr_state = PPRForwardPush( + ppr_state = await loop.run_in_executor( + None, + PPRForwardPush, seed_nodes, self._node_type_to_id[seed_node_type], self._alpha, @@ -446,7 +449,9 @@ async def _compute_ppr_scores( # means all drained nodes either had cached neighbors or no outgoing # edges — we still call push_residuals to flush their residuals into # ppr_scores_. - nodes_by_edge_type_id = ppr_state.drain_queue() + nodes_by_edge_type_id = await loop.run_in_executor( + None, ppr_state.drain_queue + ) if nodes_by_edge_type_id is None: break @@ -464,17 +469,19 @@ async def _compute_ppr_scores( fetched_by_edge_type_id = {} # Run in executor so the C++ push doesn't block the asyncio event loop. - await asyncio.get_running_loop().run_in_executor( + await loop.run_in_executor( None, ppr_state.push_residuals, fetched_by_edge_type_id ) # Regular sampling uses residual top-up as fill-only under the # max_ppr_nodes cap. If finalized PPR scores already fill that cap, # there is no remaining budget for residual candidates. - return self._extract_ppr_state_top_k( + return await loop.run_in_executor( + None, + self._extract_ppr_state_top_k, ppr_state, device, - max_ppr_nodes=self._max_ppr_nodes, + self._max_ppr_nodes, ) async def _sample_from_nodes( From ed7387c64aa2b39d82f7d839a504960a2ffd1adf Mon Sep 17 00:00:00 2001 From: mkolodner Date: Wed, 15 Jul 2026 18:49:12 +0000 Subject: [PATCH 3/3] Avoid redundant PPR frontier device copy --- .../core/sampling/python_ppr_forward_push.cpp | 16 +++++++++------- gigl/distributed/dist_ppr_sampler.py | 15 ++++++++++----- 2 files changed, 19 insertions(+), 12 deletions(-) diff --git a/gigl-core/core/sampling/python_ppr_forward_push.cpp b/gigl-core/core/sampling/python_ppr_forward_push.cpp index 5f1f474f6..b92b21418 100644 --- a/gigl-core/core/sampling/python_ppr_forward_push.cpp +++ b/gigl-core/core/sampling/python_ppr_forward_push.cpp @@ -18,9 +18,8 @@ namespace py = pybind11; namespace gigl { -// pushResiduals: a wrapper is needed solely to release the GIL during the C++ push. -// pybind11/stl.h handles all type conversions automatically; the other methods use -// direct member function pointers for the same reason. +// pushResiduals receives Python-owned containers, so convert them while the GIL +// is held and release only around the C++ state update. static void pushResidualsWrapper(PPRForwardPush& state, const py::dict& fetchedByEtypeId) { std::unordered_map> neighborTensorsByEtypeId; // Dict iteration touches Python objects — GIL must be held here. @@ -45,8 +44,9 @@ static void pushResidualsWrapper(PPRForwardPush& state, const py::dict& fetchedB static std::optional> drainQueueWrapper(PPRForwardPush& state) { std::optional> drained; - // The drain mutates only this PPRForwardPush instance and returns torch tensors. - // No Python objects are touched until pybind converts the result after return. + // drainQueue mutates only this PPRForwardPush instance and materializes CPU + // tensors for frontier node IDs. pybind converts those tensor handles back + // to Python tensors after return without copying the underlying storage. { py::gil_scoped_release release; drained = state.drainQueue(); @@ -57,8 +57,8 @@ static std::optional> drainQueueWrapp static std::unordered_map> extractTopKWithResidualTopUpWrapper(PPRForwardPush& state, int32_t maxPPRNodes, bool enableResidualTopUp) { std::unordered_map> result; - // Extraction walks C++ state and builds torch tensors; Python only sees the - // result after pybind converts it on return. + // Extraction walks C++ state and builds torch tensors. Returning through + // pybind creates Python container/wrapper objects, not tensor data copies. { py::gil_scoped_release release; result = state.extractTopKWithResidualTopUp(maxPPRNodes, enableResidualTopUp); @@ -79,6 +79,8 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { std::vector>, std::vector, std::vector>(), + // Constructor argument conversion happens before the C++ body; the + // body only initializes PPR state and can run without the GIL. py::call_guard()) .def("drain_queue", gigl::drainQueueWrapper) .def("push_residuals", gigl::pushResidualsWrapper) diff --git a/gigl/distributed/dist_ppr_sampler.py b/gigl/distributed/dist_ppr_sampler.py index 8701a1dcc..78de00436 100644 --- a/gigl/distributed/dist_ppr_sampler.py +++ b/gigl/distributed/dist_ppr_sampler.py @@ -237,7 +237,6 @@ def _get_destination_type(self, edge_type: EdgeType) -> NodeType: async def _batch_fetch_neighbors( self, nodes_by_edge_type_id: dict[int, torch.Tensor], - device: torch.device, ) -> dict[int, tuple[torch.Tensor, torch.Tensor, torch.Tensor]]: """Batch fetch neighbors for nodes grouped by integer edge type ID. @@ -247,8 +246,7 @@ async def _batch_fetch_neighbors( Args: nodes_by_edge_type_id: Dict mapping integer edge type ID to a 1-D int64 tensor of node IDs to fetch neighbors for. Comes directly from - ``drain_queue()``; node IDs are already deduplicated. - device: Torch device for intermediate tensor creation. + ``drain_queue()`` as CPU tensors; node IDs are already deduplicated. Returns: Dict mapping edge type ID to ``(node_ids, flat_neighbors, counts)`` @@ -283,9 +281,11 @@ async def _batch_fetch_neighbors( else edge_type ) edge_type_ids.append(edge_type_id) + # drain_queue materializes CPU frontier tensors; _sample_one_hop can + # consume them directly, so avoid a sampler-device round trip here. sample_tasks.append( self._sample_one_hop( - srcs=nodes_by_edge_type_id[edge_type_id].to(device), + srcs=nodes_by_edge_type_id[edge_type_id], num_nbr=self._num_neighbors_per_hop, etype=rpc_edge_type, ) @@ -443,12 +443,17 @@ async def _compute_ppr_scores( fetch_iteration_count = 0 + # TODO: If Python/C++ boundary overhead in this loop becomes a blocker, + # consider a coarser C++ PPR iteration API while keeping Python + # responsible for async distributed neighbor fetches. while True: # drain_queue returns None when the queue is truly empty (convergence), # or a dict (possibly empty) when nodes were drained. An empty dict # means all drained nodes either had cached neighbors or no outgoing # edges — we still call push_residuals to flush their residuals into # ppr_scores_. + # The pybind wrapper releases the GIL, but a direct call would still + # occupy this coroutine's event-loop thread until C++ returns. nodes_by_edge_type_id = await loop.run_in_executor( None, ppr_state.drain_queue ) @@ -461,7 +466,7 @@ async def _compute_ppr_scores( ) if nodes_by_edge_type_id and fetch_budget_remaining: fetched_by_edge_type_id = await self._batch_fetch_neighbors( - nodes_by_edge_type_id, device + nodes_by_edge_type_id ) fetch_iteration_count += 1 else: