From eadddbeda774a927b8c828cd9c4fd8db96a781ff Mon Sep 17 00:00:00 2001 From: mkolodner Date: Tue, 19 May 2026 22:02:50 +0000 Subject: [PATCH 1/4] Initial commit --- gigl-core/core/sampling/ppr_forward_push.cpp | 277 ++++++++++---- gigl-core/core/sampling/ppr_forward_push.h | 15 +- .../core/sampling/python_ppr_forward_push.cpp | 10 +- gigl-core/src/gigl_core/ppr_forward_push.pyi | 4 +- gigl/distributed/base_dist_loader.py | 7 - gigl/distributed/dist_ppr_sampler.py | 60 ++- gigl/distributed/dist_sampling_producer.py | 1 + .../shared_dist_sampling_producer.py | 1 + gigl/distributed/utils/dist_sampler.py | 12 +- .../distributed_ppr_weighted_sampling_test.py | 353 ++++++++++++++++++ 10 files changed, 641 insertions(+), 99 deletions(-) create mode 100644 tests/unit/distributed/distributed_ppr_weighted_sampling_test.py diff --git a/gigl-core/core/sampling/ppr_forward_push.cpp b/gigl-core/core/sampling/ppr_forward_push.cpp index 9a2a17f03..c71d81ada 100644 --- a/gigl-core/core/sampling/ppr_forward_push.cpp +++ b/gigl-core/core/sampling/ppr_forward_push.cpp @@ -147,20 +147,33 @@ std::optional> PPRForwardPush::drainQ } void PPRForwardPush::pushResiduals( - const std::unordered_map>& fetchedByEtypeId) { - // Step 1: Unpack the input map into a C++ map keyed by packKey(nodeId, edgeTypeId) + const std::unordered_map>& + fetchedByEtypeId) { + // Step 1: Unpack the input map into C++ maps keyed by packKey(nodeId, edgeTypeId) // for fast lookup during the residual-push loop below. + // fetchedWeights is populated only in weighted mode (_hasWeights becomes true on + // the first call that includes a non-empty flat_weights tensor). std::unordered_map> fetched; + std::unordered_map> fetchedWeights; + for (const auto& [edgeTypeId, neighborTensors] : fetchedByEtypeId) { const auto& nodeIdsTensor = std::get<0>(neighborTensors); const auto& flatNeighborIdsTensor = std::get<1>(neighborTensors); const auto& countsTensor = std::get<2>(neighborTensors); + const auto& flatWeightsTensor = std::get<3>(neighborTensors); + + bool etypeHasWeights = flatWeightsTensor.numel() > 0; + if (etypeHasWeights) { + _hasWeights = true; + } // accessor() gives a bounds-checked, typed 1-D view into // each tensor's data — equivalent to iterating over a NumPy array. auto nodeIdsAccessor = nodeIdsTensor.accessor(); auto flatNeighborIdsAccessor = flatNeighborIdsTensor.accessor(); auto countsAccessor = countsTensor.accessor(); + // Raw pointer for weights avoids a conditional accessor construction. + const double* flatWeightsPtr = etypeHasWeights ? flatWeightsTensor.data_ptr() : nullptr; // Walk the flat neighbor list, slicing out each node's neighbors using // the running offset into the concatenated flat buffer. @@ -168,20 +181,54 @@ void PPRForwardPush::pushResiduals( for (int64_t nodeIdx = 0; nodeIdx < nodeIdsTensor.size(0); ++nodeIdx) { auto nodeId = static_cast(nodeIdsAccessor[nodeIdx]); int64_t count = countsAccessor[nodeIdx]; + uint64_t key = packKey(nodeId, edgeTypeId); + std::vector neighborIds(count); for (int64_t neighborIdx = 0; neighborIdx < count; ++neighborIdx) { neighborIds[neighborIdx] = static_cast(flatNeighborIdsAccessor[offset + neighborIdx]); } - fetched[packKey(nodeId, edgeTypeId)] = std::move(neighborIds); + fetched[key] = std::move(neighborIds); + + if (flatWeightsPtr != nullptr) { + std::vector neighborWeights(count); + for (int64_t neighborIdx = 0; neighborIdx < count; ++neighborIdx) { + neighborWeights[neighborIdx] = flatWeightsPtr[offset + neighborIdx]; + } + fetchedWeights[key] = std::move(neighborWeights); + } + offset += count; } } + // Promote neighbor and weight lists for a newly re-queued node into the persistent cache. + // Called from both uniform and weighted paths — the _hasWeights guard inside handles + // whether _weightCache is populated. + auto promoteToCache = [&](int32_t neighborNodeId, int32_t dstNodeTypeId) { + for (int32_t neighborEdgeTypeId : _nodeTypeToEdgeTypeIds[dstNodeTypeId]) { + uint64_t packedKey = packKey(neighborNodeId, neighborEdgeTypeId); + if (_neighborCache.find(packedKey) == _neighborCache.end()) { + auto fetchedNeighborEntry = fetched.find(packedKey); + if (fetchedNeighborEntry != fetched.end()) { + _neighborCache[packedKey] = fetchedNeighborEntry->second; + if (_hasWeights) { + auto fetchedWeightsNeighborEntry = fetchedWeights.find(packedKey); + if (fetchedWeightsNeighborEntry != fetchedWeights.end()) { + _weightCache[packedKey] = fetchedWeightsNeighborEntry->second; + } + } + } + } + } + }; + // Step 2: For every node that was in the queue (captured in _queuedNodes // by drainQueue()), apply one PPR push step: // a. Absorb residual into the PPR score. - // b. Distribute (1-alpha) * residual equally to each neighbor. - // c. Enqueue any neighbor whose residual now exceeds the requeue threshold. + // b. Compute the normalisation factor (neighbor count or total weight). + // c. Distribute (1-alpha) * residual to each neighbor: uniformly when + // _hasWeights is false; proportionally to edge weight when true. + // d. Enqueue any neighbor whose residual now exceeds the requeue threshold. for (int32_t seedIdx = 0; seedIdx < _batchSize; ++seedIdx) { for (int32_t nodeTypeId = 0; nodeTypeId < _numNodeTypes; ++nodeTypeId) { auto& srcNodeTypeState = _state[seedIdx][nodeTypeId]; @@ -197,89 +244,167 @@ void PPRForwardPush::pushResiduals( srcNodeTypeState.pprScores[sourceNodeId] += sourceResidual; srcNodeTypeState.residuals[sourceNodeId] = 0.0; - // b. Count total fetched/cached neighbors across all edge types for - // this source node. We normalise by the number of neighbors we - // actually retrieved, not the true degree, so residual is fully - // distributed among known neighbors rather than leaking to unfetched - // ones (which matters when num_neighbors_per_hop < true_degree). - int32_t totalFetched = 0; - for (int32_t edgeTypeId : _nodeTypeToEdgeTypeIds[nodeTypeId]) { - auto fetchedEntry = fetched.find(packKey(sourceNodeId, edgeTypeId)); - if (fetchedEntry != fetched.end()) { - totalFetched += static_cast(fetchedEntry->second.size()); - } else { - auto cachedEntry = _neighborCache.find(packKey(sourceNodeId, edgeTypeId)); - if (cachedEntry != _neighborCache.end()) { - totalFetched += static_cast(cachedEntry->second.size()); + if (!_hasWeights) { + // --- Uniform path --- + // b. Count total fetched/cached neighbors across all edge types for + // this source node. We normalise by the number of neighbors we + // actually retrieved, not the true degree, so residual is fully + // distributed among known neighbors rather than leaking to unfetched + // ones (which matters when num_neighbors_per_hop < true_degree). + int32_t totalFetched = 0; + for (int32_t edgeTypeId : _nodeTypeToEdgeTypeIds[nodeTypeId]) { + auto fetchedEntry = fetched.find(packKey(sourceNodeId, edgeTypeId)); + if (fetchedEntry != fetched.end()) { + totalFetched += static_cast(fetchedEntry->second.size()); + } else { + auto cachedEntry = _neighborCache.find(packKey(sourceNodeId, edgeTypeId)); + if (cachedEntry != _neighborCache.end()) { + totalFetched += static_cast(cachedEntry->second.size()); + } } } - } - // Two cases reach here: - // 1. True sink node (no outgoing edges): absorbing the full residual is correct. - // 2. Budget exhausted, no cache entry: the (1-α)·r that should flow to - // neighbors has nowhere to go, so it gets absorbed into src's score instead. - // This overstates src and understates its neighbors. This is expected - // behavior when max_fetch_iterations is set, which intentionally trades - // theoretical PPR correctness for better throughput. - if (totalFetched == 0) { - continue; - } + // Two cases reach here: + // 1. True sink node (no outgoing edges): absorbing the full residual is correct. + // 2. Budget exhausted, no cache entry: the (1-α)·r that should flow to + // neighbors has nowhere to go, so it gets absorbed into src's score instead. + // This overstates src and understates its neighbors. This is expected + // behavior when max_fetch_iterations is set, which intentionally trades + // theoretical PPR correctness for better throughput. + if (totalFetched == 0) { + continue; + } - double residualPerNeighbor = (1.0 - _alpha) * sourceResidual / static_cast(totalFetched); + double residualPerNeighbor = (1.0 - _alpha) * sourceResidual / static_cast(totalFetched); + + for (int32_t edgeTypeId : _nodeTypeToEdgeTypeIds[nodeTypeId]) { + // Invariant: fetched and _neighborCache are mutually exclusive for + // any given (node, etype) key within one iteration. drainQueue() + // only requests a fetch for nodes absent from _neighborCache, so a + // key is in at most one of the two. + // + // Neighbor list for this (src, edgeTypeId) pair, borrowed from whichever + // map holds it. reference_wrapper is used because std::optional cannot + // hold a reference directly, and we want to avoid copying the vector — + // the data already exists in fetched or _neighborCache and both outlive + // this loop body. Access via neighborList->get(). + std::optional>> neighborList; + auto fetchedEntry = fetched.find(packKey(sourceNodeId, edgeTypeId)); + if (fetchedEntry != fetched.end()) { + neighborList = std::cref(fetchedEntry->second); + } else { + auto cachedEntry = _neighborCache.find(packKey(sourceNodeId, edgeTypeId)); + if (cachedEntry != _neighborCache.end()) { + neighborList = std::cref(cachedEntry->second); + } + } + if (!neighborList || neighborList->get().empty()) { + continue; + } - for (int32_t edgeTypeId : _nodeTypeToEdgeTypeIds[nodeTypeId]) { - // Invariant: fetched and _neighborCache are mutually exclusive for - // any given (node, etype) key within one iteration. drainQueue() - // only requests a fetch for nodes absent from _neighborCache, so a - // key is in at most one of the two. - // - // Neighbor list for this (src, edgeTypeId) pair, borrowed from whichever - // map holds it. reference_wrapper is used because std::optional cannot - // hold a reference directly, and we want to avoid copying the vector — - // the data already exists in fetched or _neighborCache and both outlive - // this loop body. Access via neighborList->get(). - std::optional>> neighborList; - auto fetchedEntry = fetched.find(packKey(sourceNodeId, edgeTypeId)); - if (fetchedEntry != fetched.end()) { - neighborList = std::cref(fetchedEntry->second); - } else { - auto cachedEntry = _neighborCache.find(packKey(sourceNodeId, edgeTypeId)); - if (cachedEntry != _neighborCache.end()) { - neighborList = std::cref(cachedEntry->second); + int32_t dstNodeTypeId = _edgeTypeToDstNtypeId[edgeTypeId]; + + // c. Accumulate residual for each neighbor and re-enqueue if threshold + // exceeded. + auto& dstNodeTypeState = _state[seedIdx][dstNodeTypeId]; + for (int32_t neighborNodeId : neighborList->get()) { + dstNodeTypeState.residuals[neighborNodeId] += residualPerNeighbor; + + double threshold = _requeueThresholdFactor * + static_cast(getTotalDegree(neighborNodeId, dstNodeTypeId)); + + if (dstNodeTypeState.queue.find(neighborNodeId) == dstNodeTypeState.queue.end() && + dstNodeTypeState.residuals[neighborNodeId] >= threshold) { + dstNodeTypeState.queue.insert(neighborNodeId); + ++_numNodesInQueue; + promoteToCache(neighborNodeId, dstNodeTypeId); + } + } + } + } else { + // --- Weighted path --- + // b. Sum total weight of fetched/cached neighbors across all edge types. + // We normalise by total fetched weight rather than by true out-weight so + // that the residual is fully distributed among known neighbors, consistent + // with how the uniform path handles truncated neighbor lists. + double totalFetchedWeight = 0.0; + for (int32_t edgeTypeId : _nodeTypeToEdgeTypeIds[nodeTypeId]) { + uint64_t key = packKey(sourceNodeId, edgeTypeId); + auto fetchedWeightsEntry = fetchedWeights.find(key); + if (fetchedWeightsEntry != fetchedWeights.end()) { + for (double w : fetchedWeightsEntry->second) { + totalFetchedWeight += w; + } + } else { + auto cachedWeightsEntry = _weightCache.find(key); + if (cachedWeightsEntry != _weightCache.end()) { + for (double w : cachedWeightsEntry->second) { + totalFetchedWeight += w; + } + } } } - if (!neighborList || neighborList->get().empty()) { + // Sink node or all-zero-weight edges: absorb residual, nothing to distribute. + if (totalFetchedWeight == 0.0) { continue; } - int32_t dstNodeTypeId = _edgeTypeToDstNtypeId[edgeTypeId]; - - // c. Accumulate residual for each neighbor and re-enqueue if threshold - // exceeded. - auto& dstNodeTypeState = _state[seedIdx][dstNodeTypeId]; - for (int32_t neighborNodeId : neighborList->get()) { - dstNodeTypeState.residuals[neighborNodeId] += residualPerNeighbor; - - double threshold = _requeueThresholdFactor * - static_cast(getTotalDegree(neighborNodeId, dstNodeTypeId)); - - if (dstNodeTypeState.queue.find(neighborNodeId) == dstNodeTypeState.queue.end() && - dstNodeTypeState.residuals[neighborNodeId] >= threshold) { - dstNodeTypeState.queue.insert(neighborNodeId); - ++_numNodesInQueue; - - // Promote neighbor lists to the persistent cache: this node will - // be processed next iteration, so caching avoids a re-fetch. - for (int32_t neighborEdgeTypeId : _nodeTypeToEdgeTypeIds[dstNodeTypeId]) { - uint64_t packedKey = packKey(neighborNodeId, neighborEdgeTypeId); - if (_neighborCache.find(packedKey) == _neighborCache.end()) { - auto fetchedNeighborEntry = fetched.find(packedKey); - if (fetchedNeighborEntry != fetched.end()) { - _neighborCache[packedKey] = fetchedNeighborEntry->second; - } + double baseResidual = (1.0 - _alpha) * sourceResidual; + + for (int32_t edgeTypeId : _nodeTypeToEdgeTypeIds[nodeTypeId]) { + uint64_t key = packKey(sourceNodeId, edgeTypeId); + + std::optional>> neighborList; + std::optional>> weightList; + + auto fetchedEntry = fetched.find(key); + if (fetchedEntry != fetched.end()) { + neighborList = std::cref(fetchedEntry->second); + auto fetchedWeightsEntry = fetchedWeights.find(key); + if (fetchedWeightsEntry != fetchedWeights.end()) { + weightList = std::cref(fetchedWeightsEntry->second); + } + } else { + auto cachedEntry = _neighborCache.find(key); + if (cachedEntry != _neighborCache.end()) { + neighborList = std::cref(cachedEntry->second); + auto cachedWeightsEntry = _weightCache.find(key); + if (cachedWeightsEntry != _weightCache.end()) { + weightList = std::cref(cachedWeightsEntry->second); } } } + if (!neighborList || neighborList->get().empty()) { + continue; + } + + int32_t dstNodeTypeId = _edgeTypeToDstNtypeId[edgeTypeId]; + auto& dstNodeTypeState = _state[seedIdx][dstNodeTypeId]; + + // c. Accumulate weight-proportional residual for each neighbor. + // weightList is always populated alongside neighborList in weighted mode: + // fetched[key] and fetchedWeights[key] are set together in Step 1, + // and _neighborCache[key] and _weightCache[key] are promoted together. + TORCH_INTERNAL_ASSERT(weightList.has_value(), + "weightList must be populated alongside neighborList in weighted mode"); + const auto& neighbors = neighborList->get(); + const auto& weights = weightList->get(); + TORCH_INTERNAL_ASSERT(weights.size() == neighbors.size(), + "weightList and neighborList must have the same size"); + for (int32_t i = 0; i < static_cast(neighbors.size()); ++i) { + int32_t neighborNodeId = neighbors[i]; + dstNodeTypeState.residuals[neighborNodeId] += + baseResidual * weights[i] / totalFetchedWeight; + + double threshold = _requeueThresholdFactor * + static_cast(getTotalDegree(neighborNodeId, dstNodeTypeId)); + + if (dstNodeTypeState.queue.find(neighborNodeId) == dstNodeTypeState.queue.end() && + dstNodeTypeState.residuals[neighborNodeId] >= threshold) { + dstNodeTypeState.queue.insert(neighborNodeId); + ++_numNodesInQueue; + promoteToCache(neighborNodeId, dstNodeTypeId); + } + } } } } diff --git a/gigl-core/core/sampling/ppr_forward_push.h b/gigl-core/core/sampling/ppr_forward_push.h index 1c1eef670..2ba811e1a 100644 --- a/gigl-core/core/sampling/ppr_forward_push.h +++ b/gigl-core/core/sampling/ppr_forward_push.h @@ -52,9 +52,12 @@ class PPRForwardPush { std::optional> drainQueue(); // Push residuals given fetched neighbor data. - // fetchedByEtypeId: {etype_id: (node_ids[N], flat_nbrs[sum(counts)], counts[N])} + // fetchedByEtypeId: {etype_id: (node_ids[N], flat_nbrs[sum(counts)], counts[N], flat_weights[sum(counts)])} + // flat_weights is empty (numel()==0) for uniform-residual mode; non-empty for + // weight-proportional mode. _hasWeights is latched true on the first call with a + // non-empty flat_weights and never reset within one PPRForwardPush lifetime. void pushResiduals(const std::unordered_map< - int32_t, std::tuple>& + int32_t, std::tuple>& fetchedByEtypeId); // Return top-k PPR nodes per seed per node type. @@ -103,6 +106,14 @@ class PPRForwardPush { // impractical (contrast with _state above). Populated incrementally; avoids re-fetching. std::unordered_map> _neighborCache; + // True once any pushResiduals call receives a non-empty flat_weights tensor. + // Latched true for the object lifetime; never reset. + bool _hasWeights{false}; + + // Per-edge weights parallel to _neighborCache: _weightCache[packKey(node, etype)][i] + // is the weight of the i-th cached neighbor. Only populated in weighted mode. + std::unordered_map> _weightCache; + }; } // namespace gigl diff --git a/gigl-core/core/sampling/python_ppr_forward_push.cpp b/gigl-core/core/sampling/python_ppr_forward_push.cpp index 22981a48a..4d3674a2f 100644 --- a/gigl-core/core/sampling/python_ppr_forward_push.cpp +++ b/gigl-core/core/sampling/python_ppr_forward_push.cpp @@ -20,15 +20,21 @@ 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. +// +// Each tuple value is (node_ids, flat_nbrs, counts, flat_weights). flat_weights is +// an empty tensor in uniform-residual mode and a non-empty float64 tensor in +// weight-proportional mode. static void pushResidualsWrapper(PPRForwardPush& state, const py::dict& fetchedByEtypeId) { - std::unordered_map> neighborTensorsByEtypeId; + std::unordered_map> + neighborTensorsByEtypeId; // Dict iteration touches Python objects — GIL must be held here. for (auto item : fetchedByEtypeId) { auto edgeTypeId = item.first.cast(); auto neighborTensors = item.second.cast(); neighborTensorsByEtypeId[edgeTypeId] = {neighborTensors[0].cast(), neighborTensors[1].cast(), - neighborTensors[2].cast()}; + neighborTensors[2].cast(), + neighborTensors[3].cast()}; } // C++ push only uses tensor accessor/data_ptr APIs — GIL-safe to release. // Releasing here lets the asyncio event loop process RPC completion callbacks diff --git a/gigl-core/src/gigl_core/ppr_forward_push.pyi b/gigl-core/src/gigl_core/ppr_forward_push.pyi index 0c1ea79af..31b052dd0 100644 --- a/gigl-core/src/gigl_core/ppr_forward_push.pyi +++ b/gigl-core/src/gigl_core/ppr_forward_push.pyi @@ -14,7 +14,9 @@ class PPRForwardPush: def drain_queue(self) -> dict[int, torch.Tensor] | None: ... def push_residuals( self, - fetched_by_etype_id: dict[int, tuple[torch.Tensor, torch.Tensor, torch.Tensor]], + fetched_by_etype_id: dict[ + int, tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor] + ], ) -> None: ... def extract_top_k( self, max_ppr_nodes: int diff --git a/gigl/distributed/base_dist_loader.py b/gigl/distributed/base_dist_loader.py index d993d83ca..4bacf8478 100644 --- a/gigl/distributed/base_dist_loader.py +++ b/gigl/distributed/base_dist_loader.py @@ -348,7 +348,6 @@ def validate_for_weighted_sampling( Raises: ValueError: If ``with_weight=True`` but no edge weights are registered. - NotImplementedError: If ``with_weight=True`` and a PPR sampler is requested. """ if not with_weight: return @@ -362,12 +361,6 @@ def validate_for_weighted_sampling( "with_weight=True requires edge weights to be registered in the dataset. " "Pass weight_edge_feat_name to build_dataset() to register edge weights." ) - # TODO(mkolodner-sc): Implement weight-proportional residual propagation for PPR. - if with_weight and isinstance(sampler_options, PPRSamplerOptions): - raise NotImplementedError( - "Weighted sampling is not yet supported with PPRSamplerOptions. " - "Weight-proportional residual propagation for PPR is planned but not implemented." - ) @staticmethod def create_sampling_config( diff --git a/gigl/distributed/dist_ppr_sampler.py b/gigl/distributed/dist_ppr_sampler.py index 402e381c1..8e61170aa 100644 --- a/gigl/distributed/dist_ppr_sampler.py +++ b/gigl/distributed/dist_ppr_sampler.py @@ -48,6 +48,14 @@ class DistPPRNeighborSampler(BaseDistNeighborSampler): scores are approximated here using the Forward Push algorithm (Andersen et al., 2006). + **Weighted PPR**: when ``edge_weights`` is provided, neighbor fetching during + traversal always uses uniform sampling (all neighbors are fetched without + weight-biased selection). The weighting is applied exclusively in how the PPR + residual is spread: each neighbor receives residual proportional to its edge + weight rather than an equal share. This is the correct formulation — using + weighted sampling during traversal would double-count high-weight edges (once + by over-representing them in the fetch and again by giving them more residual). + This sampler supports both homogeneous and heterogeneous graphs. For heterogeneous graphs, the PPR algorithm traverses across all edge types, switching edge types based on the current node type and the configured edge direction. @@ -90,6 +98,9 @@ def __init__( total_degree_dtype: torch.dtype = torch.int32, degree_tensors: Union[torch.Tensor, dict[EdgeType, torch.Tensor]], max_fetch_iterations: Optional[int] = None, + edge_weights: Optional[ + Union[torch.Tensor, dict[EdgeType, torch.Tensor]] + ] = None, **kwargs, ): super().__init__(*args, **kwargs) @@ -98,6 +109,7 @@ def __init__( self._requeue_threshold_factor = alpha * eps self._num_neighbors_per_hop = num_neighbors_per_hop self._max_fetch_iterations = max_fetch_iterations + self._edge_weights = edge_weights # Build mapping from node type to edge types that can be traversed from that node type. self._node_type_to_edge_types: dict[NodeType, list[EdgeType]] = defaultdict( @@ -251,7 +263,7 @@ async def _batch_fetch_neighbors( self, nodes_by_etype_id: dict[int, torch.Tensor], device: torch.device, - ) -> dict[int, tuple[torch.Tensor, torch.Tensor, torch.Tensor]]: + ) -> dict[int, tuple[torch.Tensor, 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 @@ -265,11 +277,13 @@ async def _batch_fetch_neighbors( 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 etype_id to ``(node_ids, flat_neighbors, counts, flat_weights)`` + as 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]``. + ``node_ids[i]``. ``flat_weights`` is a float64 tensor of the same + shape as ``flat_neighbors`` in weighted mode, or an empty tensor in + uniform mode. Example:: @@ -279,8 +293,8 @@ async def _batch_fetch_neighbors( } # Might return (neighbor lists depend on graph structure): { - 2: (tensor([0, 3]), tensor([5, 9, 2, 1]), tensor([3, 1])), - 5: (tensor([7]), tensor([0, 3]), tensor([2])), + 2: (tensor([0, 3]), tensor([5, 9, 2, 1]), tensor([3, 1]), tensor([])), + 5: (tensor([7]), tensor([0, 3]), tensor([2]), tensor([])), } """ # Fire all per-edge-type RPC calls concurrently. Each _sample_one_hop @@ -299,10 +313,36 @@ async def _batch_fetch_neighbors( ) ) 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) - } + + _empty_weights = torch.empty(0, dtype=torch.float64) + + result: dict[ + int, tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor] + ] = {} + for eid, output in zip(eids, outputs): + if self._edge_weights is not None: + assert output.edge is not None, ( + "output.edge must be set when edge_weights is provided; " + "ensure with_edge=True in SamplingConfig (hardcoded in create_sampling_config)." + ) + if self._is_homogeneous: + assert isinstance(self._edge_weights, torch.Tensor) + flat_weights = self._edge_weights[output.edge].to(torch.float64) + else: + assert isinstance(self._edge_weights, dict) + etype = self._etype_id_to_etype[eid] + flat_weights = self._edge_weights[etype][output.edge].to( + torch.float64 + ) + else: + flat_weights = _empty_weights + result[eid] = ( + nodes_by_etype_id[eid], + output.nbr, + output.nbr_num, + flat_weights, + ) + return result async def _compute_ppr_scores( self, diff --git a/gigl/distributed/dist_sampling_producer.py b/gigl/distributed/dist_sampling_producer.py index 3a51715e2..13c0ff751 100644 --- a/gigl/distributed/dist_sampling_producer.py +++ b/gigl/distributed/dist_sampling_producer.py @@ -103,6 +103,7 @@ def _sampling_worker_loop( sampler_options=sampler_options, degree_tensors=degree_tensors, current_device=current_device, + edge_weights=data.edge_weights, ) dist_sampler.start_loop() diff --git a/gigl/distributed/graph_store/shared_dist_sampling_producer.py b/gigl/distributed/graph_store/shared_dist_sampling_producer.py index 0f7461196..a8a6d6c50 100644 --- a/gigl/distributed/graph_store/shared_dist_sampling_producer.py +++ b/gigl/distributed/graph_store/shared_dist_sampling_producer.py @@ -690,6 +690,7 @@ def _handle_command(command: SharedMpCommand, payload: CommandPayload) -> bool: sampler_options=sampler_options, degree_tensors=degree_tensors, current_device=current_device, + edge_weights=data.edge_weights, ) sampler.start_loop() with state_lock: diff --git a/gigl/distributed/utils/dist_sampler.py b/gigl/distributed/utils/dist_sampler.py index 0333f4138..f218bb8eb 100644 --- a/gigl/distributed/utils/dist_sampler.py +++ b/gigl/distributed/utils/dist_sampler.py @@ -37,6 +37,7 @@ def create_dist_sampler( sampler_options: SamplerOptions, degree_tensors: Optional[Union[torch.Tensor, dict[EdgeType, torch.Tensor]]], current_device: torch.device, + edge_weights: Optional[Union[torch.Tensor, dict[EdgeType, torch.Tensor]]] = None, ) -> SamplerRuntime: """Create a GiGL sampler runtime for one channel on one worker. @@ -49,6 +50,10 @@ def create_dist_sampler( degree_tensors: Pre-computed degree tensors required by PPR sampling. Must not be ``None`` when ``sampler_options`` is :class:`PPRSamplerOptions`. current_device: The device on which sampling will run. + edge_weights: Per-edge weight tensors for this rank's partition. Required when + ``sampler_options`` is :class:`PPRSamplerOptions` and + ``sampling_config.with_weight`` is ``True``. Ignored for k-hop sampling + (GLT handles weight-biased sampling internally via ``with_weight``). Returns: A configured sampler runtime, either :class:`DistNeighborSampler` or @@ -77,8 +82,12 @@ def create_dist_sampler( ) elif isinstance(sampler_options, PPRSamplerOptions): assert degree_tensors is not None + # PPR traversal must always use uniform neighbor sampling: biased selection + # would double-count high-weight edges (once in the fetch, once in residual + # distribution). Weight influence is captured entirely in the residual math. + # Pass edge_weights only when with_weight=True; None disables weighted residuals. sampler = DistPPRNeighborSampler( - **shared_sampler_kwargs, + **{**shared_sampler_kwargs, "with_weight": False}, alpha=sampler_options.alpha, eps=sampler_options.eps, max_ppr_nodes=sampler_options.max_ppr_nodes, @@ -86,6 +95,7 @@ def create_dist_sampler( num_neighbors_per_hop=sampler_options.num_neighbors_per_hop, total_degree_dtype=sampler_options.total_degree_dtype, degree_tensors=degree_tensors, + edge_weights=edge_weights if sampling_config.with_weight else None, ) else: raise NotImplementedError( diff --git a/tests/unit/distributed/distributed_ppr_weighted_sampling_test.py b/tests/unit/distributed/distributed_ppr_weighted_sampling_test.py new file mode 100644 index 000000000..b61680e7b --- /dev/null +++ b/tests/unit/distributed/distributed_ppr_weighted_sampling_test.py @@ -0,0 +1,353 @@ +"""End-to-end correctness tests for PPR weighted sampling. + +Verifies that DistNeighborLoader with PPRSamplerOptions and with_weight=True +never traverses weight=0 edges. The test graph encodes node type in features +(hub=2.0, good=1.0, bad=0.0); any bad node in a sampled subgraph indicates +that a weight=0 edge contributed PPR residual — a test failure. + +With weight-proportional residual distribution, a weight=0 edge contributes +zero to totalFetchedWeight and receives zero residual per push step. Bad +nodes therefore accumulate a PPR score of exactly 0 and are excluded from +every top-k result. +""" + +import torch +import torch.multiprocessing as mp +from absl.testing import absltest +from graphlearn_torch.distributed import shutdown_rpc +from torch_geometric.data import Data, HeteroData + +from gigl.distributed.dist_dataset import DistDataset +from gigl.distributed.distributed_neighborloader import DistNeighborLoader +from gigl.distributed.sampler_options import PPRSamplerOptions +from gigl.src.common.types.graph_data import EdgeType, NodeType, Relation +from gigl.types.graph import ( + FeaturePartitionData, + GraphPartitionData, + PartitionOutput, +) +from tests.test_assets.distributed.utils import create_test_process_group +from tests.test_assets.test_case import TestCase + +_USER = NodeType("user") +_ITEM = NodeType("item") +_USER_TO_ITEM = EdgeType(_USER, Relation("to"), _ITEM) +_ITEM_TO_USER = EdgeType(_ITEM, Relation("to"), _USER) + +# PPR parameters used across all tests. +_PPR_ALPHA = 0.5 +_PPR_EPS = 1e-4 +_PPR_MAX_NODES = 60 +_PPR_NUM_NBRS = 200 + + +# --------------------------------------------------------------------------- +# Graph builders +# --------------------------------------------------------------------------- + + +def _build_homogeneous_bipartite_weight_graph() -> tuple[ + PartitionOutput, int, int, int +]: + """Build a homogeneous graph with hub, good, and bad nodes. + + Graph structure: + - 10 hub nodes (0..9): seed nodes; feature = 2.0 + - 50 good nodes (10..59): reachable via weight=1 edges; feature = 1.0 + - 40 bad nodes (60..99): reachable via weight=0 edges; feature = 0.0 + - Each good node has 5 outgoing weight=1 edges to nearby good nodes (ring). + + Returns: + (partition_output, n_hub, n_good, n_bad) + """ + n_hub = 10 + n_good = 50 + n_bad = 40 + n = n_hub + n_good + n_bad + + hub_ids = torch.arange(n_hub) + good_ids = torch.arange(n_hub, n_hub + n_good) + bad_ids = torch.arange(n_hub + n_good, n) + + hub_good_src = hub_ids.repeat_interleave(n_good) + hub_good_dst = good_ids.repeat(n_hub) + hub_good_w = torch.ones(n_hub * n_good) + + hub_bad_src = hub_ids.repeat_interleave(n_bad) + hub_bad_dst = bad_ids.repeat(n_hub) + hub_bad_w = torch.zeros(n_hub * n_bad) + + connections_per_good = 5 + good_src = good_ids.repeat_interleave(connections_per_good) + good_dst = torch.stack( + [torch.roll(good_ids, -j) for j in range(1, connections_per_good + 1)] + ).T.reshape(-1) + good_w = torch.ones(n_good * connections_per_good) + + edge_src = torch.cat([hub_good_src, hub_bad_src, good_src]) + edge_dst = torch.cat([hub_good_dst, hub_bad_dst, good_dst]) + weights = torch.cat([hub_good_w, hub_bad_w, good_w]) + edge_index = torch.stack([edge_src, edge_dst]) + n_edges = edge_src.shape[0] + + node_feats = torch.cat( + [ + torch.full((n_hub, 1), 2.0), + torch.full((n_good, 1), 1.0), + torch.full((n_bad, 1), 0.0), + ] + ) + + partition_output = PartitionOutput( + node_partition_book=torch.zeros(n), + edge_partition_book=torch.zeros(n_edges), + partitioned_edge_index=GraphPartitionData( + edge_index=edge_index, + edge_ids=None, + weights=weights, + ), + partitioned_node_features=FeaturePartitionData( + feats=node_feats, + ids=torch.arange(n), + ), + partitioned_edge_features=None, + partitioned_positive_labels=None, + partitioned_negative_labels=None, + partitioned_node_labels=None, + ) + return partition_output, n_hub, n_good, n_bad + + +def _build_heterogeneous_bipartite_weight_graph() -> tuple[ + PartitionOutput, int, int, int +]: + """Build a heterogeneous (user/item) graph with good and bad item nodes. + + Graph structure: + - 10 user nodes (0..9): seed nodes; user feature = 2.0 + - 40 good item nodes (0..39): weight=1 from users; feature = 1.0 + - 20 bad item nodes (40..59): weight=0 from users; feature = 0.0 + - Good items also connect back to all users via weight=1 (2nd-hop). + + Returns: + (partition_output, n_user, n_good_item, n_bad_item) + """ + n_user = 10 + n_good_item = 40 + n_bad_item = 20 + n_item = n_good_item + n_bad_item + + user_ids = torch.arange(n_user) + good_item_ids = torch.arange(n_good_item) + bad_item_ids = torch.arange(n_good_item, n_item) + + u2gi_src = user_ids.repeat_interleave(n_good_item) + u2gi_dst = good_item_ids.repeat(n_user) + u2gi_w = torch.ones(n_user * n_good_item) + + u2bi_src = user_ids.repeat_interleave(n_bad_item) + u2bi_dst = bad_item_ids.repeat(n_user) + u2bi_w = torch.zeros(n_user * n_bad_item) + + gi2u_src = good_item_ids.repeat_interleave(n_user) + gi2u_dst = user_ids.repeat(n_good_item) + gi2u_w = torch.ones(n_good_item * n_user) + + u2i_src = torch.cat([u2gi_src, u2bi_src]) + u2i_dst = torch.cat([u2gi_dst, u2bi_dst]) + u2i_w = torch.cat([u2gi_w, u2bi_w]) + n_u2i_edges = u2i_src.shape[0] + + user_feats = torch.full((n_user, 1), 2.0) + item_feats = torch.cat( + [ + torch.full((n_good_item, 1), 1.0), + torch.full((n_bad_item, 1), 0.0), + ] + ) + + partition_output = PartitionOutput( + node_partition_book={ + _USER: torch.zeros(n_user), + _ITEM: torch.zeros(n_item), + }, + edge_partition_book={ + _USER_TO_ITEM: torch.zeros(n_u2i_edges), + _ITEM_TO_USER: torch.zeros(gi2u_src.shape[0]), + }, + partitioned_edge_index={ + _USER_TO_ITEM: GraphPartitionData( + edge_index=torch.stack([u2i_src, u2i_dst]), + edge_ids=None, + weights=u2i_w, + ), + _ITEM_TO_USER: GraphPartitionData( + edge_index=torch.stack([gi2u_src, gi2u_dst]), + edge_ids=None, + weights=gi2u_w, + ), + }, + partitioned_node_features={ + _USER: FeaturePartitionData(feats=user_feats, ids=torch.arange(n_user)), + _ITEM: FeaturePartitionData(feats=item_feats, ids=torch.arange(n_item)), + }, + partitioned_edge_features=None, + partitioned_positive_labels=None, + partitioned_negative_labels=None, + partitioned_node_labels=None, + ) + return partition_output, n_user, n_good_item, n_bad_item + + +# --------------------------------------------------------------------------- +# Subprocess functions +# --------------------------------------------------------------------------- + + +def _run_ppr_weighted_correctness_homogeneous( + _: int, + dataset: DistDataset, + n_hub: int, +) -> None: + """Subprocess: verifies weight=0 edges never contribute PPR residual (homogeneous). + + Seeds are hub nodes only. Node features encode type: hub=2.0, good=1.0, bad=0.0. + Any batch containing a bad node (feature==0.0) means a weight=0 edge contributed + PPR residual — a test failure. + """ + create_test_process_group() + loader = DistNeighborLoader( + dataset=dataset, + input_nodes=torch.arange(n_hub), + num_neighbors=[], + sampler_options=PPRSamplerOptions( + alpha=_PPR_ALPHA, + eps=_PPR_EPS, + max_ppr_nodes=_PPR_MAX_NODES, + num_neighbors_per_hop=_PPR_NUM_NBRS, + ), + with_weight=True, + pin_memory_device=torch.device("cpu"), + batch_size=1, + ) + count = 0 + for datum in loader: + assert isinstance(datum, Data), f"Expected Data, got {type(datum)}" + assert datum.x is not None, "Node features missing from PPR batch" + bad_mask = datum.x[:, 0] == 0.0 + assert not bad_mask.any(), ( + f"weight=0 edge contributed PPR residual: bad node(s) found. " + f"Features of bad nodes: {datum.x[bad_mask].squeeze().tolist()}" + ) + count += 1 + assert count == n_hub, f"Expected {n_hub} batches, got {count}" + shutdown_rpc() + + +def _run_ppr_weighted_correctness_heterogeneous( + _: int, + dataset: DistDataset, + n_user: int, +) -> None: + """Subprocess: verifies weight=0 edges never contribute PPR residual (heterogeneous). + + Seeds are user nodes. Item features encode type: good=1.0, bad=0.0. + Any batch containing a bad item node means a weight=0 edge contributed PPR residual. + """ + create_test_process_group() + node_ids = dataset.node_ids + assert not isinstance(node_ids, torch.Tensor) and node_ids is not None, ( + "Expected heterogeneous dataset with dict node_ids" + ) + loader = DistNeighborLoader( + dataset=dataset, + input_nodes=(_USER, node_ids[_USER]), + num_neighbors=[], + sampler_options=PPRSamplerOptions( + alpha=_PPR_ALPHA, + eps=_PPR_EPS, + max_ppr_nodes=_PPR_MAX_NODES, + num_neighbors_per_hop=_PPR_NUM_NBRS, + ), + with_weight=True, + pin_memory_device=torch.device("cpu"), + batch_size=1, + ) + count = 0 + for datum in loader: + assert isinstance(datum, HeteroData), f"Expected HeteroData, got {type(datum)}" + if _ITEM in datum.node_types: + item_x = datum[_ITEM].x + assert item_x is not None, "Item features missing from PPR batch" + bad_mask = item_x[:, 0] == 0.0 + assert not bad_mask.any(), ( + f"weight=0 edge contributed PPR residual: bad item(s) found. " + f"Features of bad items: {item_x[bad_mask].squeeze().tolist()}" + ) + count += 1 + assert count == n_user, f"Expected {n_user} batches, got {count}" + shutdown_rpc() + + +# --------------------------------------------------------------------------- +# Test class +# --------------------------------------------------------------------------- + + +class PPRWeightedSamplingTest(TestCase): + """End-to-end correctness tests for PPR sampling with weight_proportional_residuals. + + Each test builds a bipartite graph with "good" neighbors (weight=1) and "bad" + neighbors (weight=0) reachable from seed nodes. With weight-proportional PPR, + bad nodes must never appear in any sampled subgraph because weight=0 edges + contribute zero residual per push step. + """ + + def tearDown(self) -> None: + if torch.distributed.is_initialized(): + torch.distributed.destroy_process_group() + super().tearDown() + + def test_ppr_weighted_never_traverses_zero_weight_edges_homogeneous(self) -> None: + """Homogeneous: weight=0 edges to bad nodes never contribute PPR residual. + + Graph: 10 hub seeds, each connected to 50 good nodes (weight=1) and 40 bad + nodes (weight=0). Good nodes have 5 further weight=1 edges for deeper walks. + PPR max_ppr_nodes=60 is larger than the number of good neighbors, so the + sampler must actively filter: correct weighting excludes bad nodes entirely. + """ + partition_output, n_hub, _, _ = _build_homogeneous_bipartite_weight_graph() + dataset = DistDataset(rank=0, world_size=1, edge_dir="out") + dataset.build(partition_output=partition_output) + + self.assertTrue(dataset.has_edge_weights) + + mp.spawn( + fn=_run_ppr_weighted_correctness_homogeneous, + args=(dataset, n_hub), + nprocs=1, + ) + + def test_ppr_weighted_never_traverses_zero_weight_edges_heterogeneous(self) -> None: + """Heterogeneous: weight=0 user→item edges to bad items never contribute PPR residual. + + Graph: 10 user seeds, each connected to 40 good items (weight=1) and 20 bad + items (weight=0). Good items connect back to all users via weight=1 (2nd-hop). + PPR max_ppr_nodes=60 is larger than n_good, so correct weighting is required + to exclude bad items. + """ + partition_output, n_user, _, _ = _build_heterogeneous_bipartite_weight_graph() + dataset = DistDataset(rank=0, world_size=1, edge_dir="out") + dataset.build(partition_output=partition_output) + + self.assertTrue(dataset.has_edge_weights) + + mp.spawn( + fn=_run_ppr_weighted_correctness_heterogeneous, + args=(dataset, n_user), + nprocs=1, + ) + + +if __name__ == "__main__": + absltest.main() From 280875d3b9f99c3107d0090895e762897eaeac23 Mon Sep 17 00:00:00 2001 From: mkolodner Date: Tue, 19 May 2026 22:12:00 +0000 Subject: [PATCH 2/4] Update --- .../distributed/bipartite_weight_graph.py | 188 +++++++++++++++ .../distributed_ppr_weighted_sampling_test.py | 183 +-------------- .../distributed_weighted_sampling_test.py | 219 ++---------------- 3 files changed, 222 insertions(+), 368 deletions(-) create mode 100644 tests/test_assets/distributed/bipartite_weight_graph.py diff --git a/tests/test_assets/distributed/bipartite_weight_graph.py b/tests/test_assets/distributed/bipartite_weight_graph.py new file mode 100644 index 000000000..acf37ba85 --- /dev/null +++ b/tests/test_assets/distributed/bipartite_weight_graph.py @@ -0,0 +1,188 @@ +"""Shared bipartite graph builders for weighted-sampling tests. + +Used by both distributed_weighted_sampling_test and +distributed_ppr_weighted_sampling_test. Each builder returns a single-rank +PartitionOutput that encodes node type as a feature (hub/user=2.0, good=1.0, +bad=0.0) so tests can assert that weight=0 edges never appear in any sampled +subgraph. +""" + +import torch + +from gigl.src.common.types.graph_data import EdgeType, NodeType, Relation +from gigl.types.graph import ( + FeaturePartitionData, + GraphPartitionData, + PartitionOutput, +) + +USER = NodeType("user") +ITEM = NodeType("item") +USER_TO_ITEM = EdgeType(USER, Relation("to"), ITEM) +ITEM_TO_USER = EdgeType(ITEM, Relation("to"), USER) + + +def build_homogeneous_bipartite_weight_graph() -> tuple[PartitionOutput, int]: + """Build a homogeneous graph with hub, good, and bad nodes. + + Graph structure: + - 10 hub nodes (0..9): used as seed nodes; feature value = 2.0 + - 50 good nodes (10..59): reachable from hubs via weight=1 edges; feature = 1.0 + - 40 bad nodes (60..99): reachable from hubs via weight=0 edges; feature = 0.0 + - Each good node also has 5 outgoing weight=1 edges to nearby good nodes + (ring topology, for 2nd-hop sampling). + + With weighted sampling only good nodes should ever appear as sampled + neighbors — weight=0 edges to bad nodes must never be traversed. + + Returns: + (partition_output, n_hub) + """ + n_hub = 10 + n_good = 50 + n_bad = 40 + n = n_hub + n_good + n_bad # 100 + + hub_ids = torch.arange(n_hub) + good_ids = torch.arange(n_hub, n_hub + n_good) + bad_ids = torch.arange(n_hub + n_good, n) + + # Hub → Good: weight=1 + hub_good_src = hub_ids.repeat_interleave(n_good) + hub_good_dst = good_ids.repeat(n_hub) + hub_good_w = torch.ones(n_hub * n_good) + + # Hub → Bad: weight=0 + hub_bad_src = hub_ids.repeat_interleave(n_bad) + hub_bad_dst = bad_ids.repeat(n_hub) + hub_bad_w = torch.zeros(n_hub * n_bad) + + # Good → Good: ring with 5 outgoing edges per node, weight=1 (2nd-hop targets) + connections_per_good = 5 + good_src = good_ids.repeat_interleave(connections_per_good) + # Row i of [connections_per_good, n_good].T gives neighbors of good_ids[i] + good_dst = torch.stack( + [torch.roll(good_ids, -j) for j in range(1, connections_per_good + 1)] + ).T.reshape(-1) + good_w = torch.ones(n_good * connections_per_good) + + edge_src = torch.cat([hub_good_src, hub_bad_src, good_src]) + edge_dst = torch.cat([hub_good_dst, hub_bad_dst, good_dst]) + weights = torch.cat([hub_good_w, hub_bad_w, good_w]) + edge_index = torch.stack([edge_src, edge_dst]) + n_edges = edge_src.shape[0] + + # Feature encodes node type: hub=2.0, good=1.0, bad=0.0 + node_feats = torch.cat( + [ + torch.full((n_hub, 1), 2.0), + torch.full((n_good, 1), 1.0), + torch.full((n_bad, 1), 0.0), + ] + ) + + partition_output = PartitionOutput( + node_partition_book=torch.zeros(n), + edge_partition_book=torch.zeros(n_edges), + partitioned_edge_index=GraphPartitionData( + edge_index=edge_index, + edge_ids=None, + weights=weights, + ), + partitioned_node_features=FeaturePartitionData( + feats=node_feats, + ids=torch.arange(n), + ), + partitioned_edge_features=None, + partitioned_positive_labels=None, + partitioned_negative_labels=None, + partitioned_node_labels=None, + ) + return partition_output, n_hub + + +def build_heterogeneous_bipartite_weight_graph() -> tuple[PartitionOutput, int]: + """Build a heterogeneous (user/item) graph with good and bad item nodes. + + Graph structure: + - 10 user nodes (0..9): seed nodes; user feature = 2.0 + - 60 item nodes total: + - Items 0..39: good, reachable from users via weight=1 edges; feature = 1.0 + - Items 40..59: bad, reachable from users via weight=0 edges; feature = 0.0 + - Good items also have weight=1 edges back to all users (for 2nd-hop). + + With weighted sampling only good item nodes should ever appear as sampled + item neighbors. + + Returns: + (partition_output, n_user) + """ + n_user = 10 + n_good_item = 40 + n_bad_item = 20 + n_item = n_good_item + n_bad_item # 60 + + user_ids = torch.arange(n_user) + good_item_ids = torch.arange(n_good_item) + bad_item_ids = torch.arange(n_good_item, n_item) + + # User → Good Item: weight=1 + u2gi_src = user_ids.repeat_interleave(n_good_item) + u2gi_dst = good_item_ids.repeat(n_user) + u2gi_w = torch.ones(n_user * n_good_item) + + # User → Bad Item: weight=0 + u2bi_src = user_ids.repeat_interleave(n_bad_item) + u2bi_dst = bad_item_ids.repeat(n_user) + u2bi_w = torch.zeros(n_user * n_bad_item) + + # Good Item → User: weight=1 (2nd-hop back to users) + gi2u_src = good_item_ids.repeat_interleave(n_user) + gi2u_dst = user_ids.repeat(n_good_item) + gi2u_w = torch.ones(n_good_item * n_user) + + u2i_src = torch.cat([u2gi_src, u2bi_src]) + u2i_dst = torch.cat([u2gi_dst, u2bi_dst]) + u2i_w = torch.cat([u2gi_w, u2bi_w]) + n_u2i_edges = u2i_src.shape[0] + + user_feats = torch.full((n_user, 1), 2.0) + # Item feature encodes type: good=1.0, bad=0.0 + item_feats = torch.cat( + [ + torch.full((n_good_item, 1), 1.0), + torch.full((n_bad_item, 1), 0.0), + ] + ) + + partition_output = PartitionOutput( + node_partition_book={ + USER: torch.zeros(n_user), + ITEM: torch.zeros(n_item), + }, + edge_partition_book={ + USER_TO_ITEM: torch.zeros(n_u2i_edges), + ITEM_TO_USER: torch.zeros(gi2u_src.shape[0]), + }, + partitioned_edge_index={ + USER_TO_ITEM: GraphPartitionData( + edge_index=torch.stack([u2i_src, u2i_dst]), + edge_ids=None, + weights=u2i_w, + ), + ITEM_TO_USER: GraphPartitionData( + edge_index=torch.stack([gi2u_src, gi2u_dst]), + edge_ids=None, + weights=gi2u_w, + ), + }, + partitioned_node_features={ + USER: FeaturePartitionData(feats=user_feats, ids=torch.arange(n_user)), + ITEM: FeaturePartitionData(feats=item_feats, ids=torch.arange(n_item)), + }, + partitioned_edge_features=None, + partitioned_positive_labels=None, + partitioned_negative_labels=None, + partitioned_node_labels=None, + ) + return partition_output, n_user diff --git a/tests/unit/distributed/distributed_ppr_weighted_sampling_test.py b/tests/unit/distributed/distributed_ppr_weighted_sampling_test.py index b61680e7b..8f666ccb7 100644 --- a/tests/unit/distributed/distributed_ppr_weighted_sampling_test.py +++ b/tests/unit/distributed/distributed_ppr_weighted_sampling_test.py @@ -20,20 +20,15 @@ from gigl.distributed.dist_dataset import DistDataset from gigl.distributed.distributed_neighborloader import DistNeighborLoader from gigl.distributed.sampler_options import PPRSamplerOptions -from gigl.src.common.types.graph_data import EdgeType, NodeType, Relation -from gigl.types.graph import ( - FeaturePartitionData, - GraphPartitionData, - PartitionOutput, +from tests.test_assets.distributed.bipartite_weight_graph import ( + ITEM, + USER, + build_heterogeneous_bipartite_weight_graph, + build_homogeneous_bipartite_weight_graph, ) from tests.test_assets.distributed.utils import create_test_process_group from tests.test_assets.test_case import TestCase -_USER = NodeType("user") -_ITEM = NodeType("item") -_USER_TO_ITEM = EdgeType(_USER, Relation("to"), _ITEM) -_ITEM_TO_USER = EdgeType(_ITEM, Relation("to"), _USER) - # PPR parameters used across all tests. _PPR_ALPHA = 0.5 _PPR_EPS = 1e-4 @@ -41,164 +36,6 @@ _PPR_NUM_NBRS = 200 -# --------------------------------------------------------------------------- -# Graph builders -# --------------------------------------------------------------------------- - - -def _build_homogeneous_bipartite_weight_graph() -> tuple[ - PartitionOutput, int, int, int -]: - """Build a homogeneous graph with hub, good, and bad nodes. - - Graph structure: - - 10 hub nodes (0..9): seed nodes; feature = 2.0 - - 50 good nodes (10..59): reachable via weight=1 edges; feature = 1.0 - - 40 bad nodes (60..99): reachable via weight=0 edges; feature = 0.0 - - Each good node has 5 outgoing weight=1 edges to nearby good nodes (ring). - - Returns: - (partition_output, n_hub, n_good, n_bad) - """ - n_hub = 10 - n_good = 50 - n_bad = 40 - n = n_hub + n_good + n_bad - - hub_ids = torch.arange(n_hub) - good_ids = torch.arange(n_hub, n_hub + n_good) - bad_ids = torch.arange(n_hub + n_good, n) - - hub_good_src = hub_ids.repeat_interleave(n_good) - hub_good_dst = good_ids.repeat(n_hub) - hub_good_w = torch.ones(n_hub * n_good) - - hub_bad_src = hub_ids.repeat_interleave(n_bad) - hub_bad_dst = bad_ids.repeat(n_hub) - hub_bad_w = torch.zeros(n_hub * n_bad) - - connections_per_good = 5 - good_src = good_ids.repeat_interleave(connections_per_good) - good_dst = torch.stack( - [torch.roll(good_ids, -j) for j in range(1, connections_per_good + 1)] - ).T.reshape(-1) - good_w = torch.ones(n_good * connections_per_good) - - edge_src = torch.cat([hub_good_src, hub_bad_src, good_src]) - edge_dst = torch.cat([hub_good_dst, hub_bad_dst, good_dst]) - weights = torch.cat([hub_good_w, hub_bad_w, good_w]) - edge_index = torch.stack([edge_src, edge_dst]) - n_edges = edge_src.shape[0] - - node_feats = torch.cat( - [ - torch.full((n_hub, 1), 2.0), - torch.full((n_good, 1), 1.0), - torch.full((n_bad, 1), 0.0), - ] - ) - - partition_output = PartitionOutput( - node_partition_book=torch.zeros(n), - edge_partition_book=torch.zeros(n_edges), - partitioned_edge_index=GraphPartitionData( - edge_index=edge_index, - edge_ids=None, - weights=weights, - ), - partitioned_node_features=FeaturePartitionData( - feats=node_feats, - ids=torch.arange(n), - ), - partitioned_edge_features=None, - partitioned_positive_labels=None, - partitioned_negative_labels=None, - partitioned_node_labels=None, - ) - return partition_output, n_hub, n_good, n_bad - - -def _build_heterogeneous_bipartite_weight_graph() -> tuple[ - PartitionOutput, int, int, int -]: - """Build a heterogeneous (user/item) graph with good and bad item nodes. - - Graph structure: - - 10 user nodes (0..9): seed nodes; user feature = 2.0 - - 40 good item nodes (0..39): weight=1 from users; feature = 1.0 - - 20 bad item nodes (40..59): weight=0 from users; feature = 0.0 - - Good items also connect back to all users via weight=1 (2nd-hop). - - Returns: - (partition_output, n_user, n_good_item, n_bad_item) - """ - n_user = 10 - n_good_item = 40 - n_bad_item = 20 - n_item = n_good_item + n_bad_item - - user_ids = torch.arange(n_user) - good_item_ids = torch.arange(n_good_item) - bad_item_ids = torch.arange(n_good_item, n_item) - - u2gi_src = user_ids.repeat_interleave(n_good_item) - u2gi_dst = good_item_ids.repeat(n_user) - u2gi_w = torch.ones(n_user * n_good_item) - - u2bi_src = user_ids.repeat_interleave(n_bad_item) - u2bi_dst = bad_item_ids.repeat(n_user) - u2bi_w = torch.zeros(n_user * n_bad_item) - - gi2u_src = good_item_ids.repeat_interleave(n_user) - gi2u_dst = user_ids.repeat(n_good_item) - gi2u_w = torch.ones(n_good_item * n_user) - - u2i_src = torch.cat([u2gi_src, u2bi_src]) - u2i_dst = torch.cat([u2gi_dst, u2bi_dst]) - u2i_w = torch.cat([u2gi_w, u2bi_w]) - n_u2i_edges = u2i_src.shape[0] - - user_feats = torch.full((n_user, 1), 2.0) - item_feats = torch.cat( - [ - torch.full((n_good_item, 1), 1.0), - torch.full((n_bad_item, 1), 0.0), - ] - ) - - partition_output = PartitionOutput( - node_partition_book={ - _USER: torch.zeros(n_user), - _ITEM: torch.zeros(n_item), - }, - edge_partition_book={ - _USER_TO_ITEM: torch.zeros(n_u2i_edges), - _ITEM_TO_USER: torch.zeros(gi2u_src.shape[0]), - }, - partitioned_edge_index={ - _USER_TO_ITEM: GraphPartitionData( - edge_index=torch.stack([u2i_src, u2i_dst]), - edge_ids=None, - weights=u2i_w, - ), - _ITEM_TO_USER: GraphPartitionData( - edge_index=torch.stack([gi2u_src, gi2u_dst]), - edge_ids=None, - weights=gi2u_w, - ), - }, - partitioned_node_features={ - _USER: FeaturePartitionData(feats=user_feats, ids=torch.arange(n_user)), - _ITEM: FeaturePartitionData(feats=item_feats, ids=torch.arange(n_item)), - }, - partitioned_edge_features=None, - partitioned_positive_labels=None, - partitioned_negative_labels=None, - partitioned_node_labels=None, - ) - return partition_output, n_user, n_good_item, n_bad_item - - # --------------------------------------------------------------------------- # Subprocess functions # --------------------------------------------------------------------------- @@ -261,7 +98,7 @@ def _run_ppr_weighted_correctness_heterogeneous( ) loader = DistNeighborLoader( dataset=dataset, - input_nodes=(_USER, node_ids[_USER]), + input_nodes=(USER, node_ids[USER]), num_neighbors=[], sampler_options=PPRSamplerOptions( alpha=_PPR_ALPHA, @@ -276,8 +113,8 @@ def _run_ppr_weighted_correctness_heterogeneous( count = 0 for datum in loader: assert isinstance(datum, HeteroData), f"Expected HeteroData, got {type(datum)}" - if _ITEM in datum.node_types: - item_x = datum[_ITEM].x + if ITEM in datum.node_types: + item_x = datum[ITEM].x assert item_x is not None, "Item features missing from PPR batch" bad_mask = item_x[:, 0] == 0.0 assert not bad_mask.any(), ( @@ -316,7 +153,7 @@ def test_ppr_weighted_never_traverses_zero_weight_edges_homogeneous(self) -> Non PPR max_ppr_nodes=60 is larger than the number of good neighbors, so the sampler must actively filter: correct weighting excludes bad nodes entirely. """ - partition_output, n_hub, _, _ = _build_homogeneous_bipartite_weight_graph() + partition_output, n_hub = build_homogeneous_bipartite_weight_graph() dataset = DistDataset(rank=0, world_size=1, edge_dir="out") dataset.build(partition_output=partition_output) @@ -336,7 +173,7 @@ def test_ppr_weighted_never_traverses_zero_weight_edges_heterogeneous(self) -> N PPR max_ppr_nodes=60 is larger than n_good, so correct weighting is required to exclude bad items. """ - partition_output, n_user, _, _ = _build_heterogeneous_bipartite_weight_graph() + partition_output, n_user = build_heterogeneous_bipartite_weight_graph() dataset = DistDataset(rank=0, world_size=1, edge_dir="out") dataset.build(partition_output=partition_output) diff --git a/tests/unit/distributed/distributed_weighted_sampling_test.py b/tests/unit/distributed/distributed_weighted_sampling_test.py index 659c23f5e..eff598a80 100644 --- a/tests/unit/distributed/distributed_weighted_sampling_test.py +++ b/tests/unit/distributed/distributed_weighted_sampling_test.py @@ -22,12 +22,19 @@ from gigl.distributed.dist_range_partitioner import DistRangePartitioner from gigl.distributed.distributed_neighborloader import DistNeighborLoader from gigl.distributed.utils.networking import get_free_port -from gigl.src.common.types.graph_data import EdgeType, NodeType, Relation from gigl.types.graph import ( FeaturePartitionData, GraphPartitionData, PartitionOutput, ) +from tests.test_assets.distributed.bipartite_weight_graph import ( + ITEM, + ITEM_TO_USER, + USER, + USER_TO_ITEM, + build_heterogeneous_bipartite_weight_graph, + build_homogeneous_bipartite_weight_graph, +) from tests.test_assets.distributed.constants import ( MOCKED_NUM_PARTITIONS, MOCKED_U2U_EDGE_INDEX_ON_RANK_ONE, @@ -43,189 +50,13 @@ from tests.test_assets.distributed.utils import create_test_process_group from tests.test_assets.test_case import TestCase -_USER = NodeType("user") -_ITEM = NodeType("item") -_USER_TO_ITEM = EdgeType(_USER, Relation("to"), _ITEM) -_ITEM_TO_USER = EdgeType(_ITEM, Relation("to"), _USER) - - # --------------------------------------------------------------------------- # Graph builders # --------------------------------------------------------------------------- -def _build_homogeneous_bipartite_weight_graph() -> tuple[ - PartitionOutput, int, int, int -]: - """Build a homogeneous graph with hub, good, and bad nodes. - - Graph structure: - - 10 hub nodes (0..9): used as seed nodes; feature value = 2.0 - - 50 good nodes (10..59): reachable from hubs via weight=1 edges; feature = 1.0 - - 40 bad nodes (60..99): reachable from hubs via weight=0 edges; feature = 0.0 - - Each good node also has 5 outgoing weight=1 edges to nearby good nodes - (ring topology, for 2nd-hop sampling). - - With weighted sampling only good nodes should ever appear as sampled - neighbors — weight=0 edges to bad nodes must never be traversed. - - Returns: - (partition_output, n_hub, n_good, n_bad) - """ - n_hub = 10 - n_good = 50 - n_bad = 40 - n = n_hub + n_good + n_bad # 100 - - hub_ids = torch.arange(n_hub) - good_ids = torch.arange(n_hub, n_hub + n_good) - bad_ids = torch.arange(n_hub + n_good, n) - - # Hub → Good: weight=1 - hub_good_src = hub_ids.repeat_interleave(n_good) - hub_good_dst = good_ids.repeat(n_hub) - hub_good_w = torch.ones(n_hub * n_good) - - # Hub → Bad: weight=0 - hub_bad_src = hub_ids.repeat_interleave(n_bad) - hub_bad_dst = bad_ids.repeat(n_hub) - hub_bad_w = torch.zeros(n_hub * n_bad) - - # Good → Good: ring with 5 outgoing edges per node, weight=1 (2nd-hop targets) - connections_per_good = 5 - good_src = good_ids.repeat_interleave(connections_per_good) - # Row i of [connections_per_good, n_good].T gives neighbors of good_ids[i] - good_dst = torch.stack( - [torch.roll(good_ids, -j) for j in range(1, connections_per_good + 1)] - ).T.reshape(-1) - good_w = torch.ones(n_good * connections_per_good) - - edge_src = torch.cat([hub_good_src, hub_bad_src, good_src]) - edge_dst = torch.cat([hub_good_dst, hub_bad_dst, good_dst]) - weights = torch.cat([hub_good_w, hub_bad_w, good_w]) - edge_index = torch.stack([edge_src, edge_dst]) - n_edges = edge_src.shape[0] - - # Feature encodes node type: hub=2.0, good=1.0, bad=0.0 - node_feats = torch.cat( - [ - torch.full((n_hub, 1), 2.0), - torch.full((n_good, 1), 1.0), - torch.full((n_bad, 1), 0.0), - ] - ) - - partition_output = PartitionOutput( - node_partition_book=torch.zeros(n), - edge_partition_book=torch.zeros(n_edges), - partitioned_edge_index=GraphPartitionData( - edge_index=edge_index, - edge_ids=None, - weights=weights, - ), - partitioned_node_features=FeaturePartitionData( - feats=node_feats, - ids=torch.arange(n), - ), - partitioned_edge_features=None, - partitioned_positive_labels=None, - partitioned_negative_labels=None, - partitioned_node_labels=None, - ) - return partition_output, n_hub, n_good, n_bad - - -def _build_heterogeneous_bipartite_weight_graph() -> tuple[ - PartitionOutput, int, int, int -]: - """Build a heterogeneous (user/item) graph with good and bad item nodes. - - Graph structure: - - 10 user nodes (0..9): seed nodes; user feature = 2.0 - - 60 item nodes total: - - Items 0..39: good, reachable from users via weight=1 edges; feature = 1.0 - - Items 40..59: bad, reachable from users via weight=0 edges; feature = 0.0 - - Good items also have weight=1 edges back to all users (for 2nd-hop). - - With weighted sampling only good item nodes should ever appear as sampled - item neighbors. - - Returns: - (partition_output, n_user, n_good_item, n_bad_item) - """ - n_user = 10 - n_good_item = 40 - n_bad_item = 20 - n_item = n_good_item + n_bad_item # 60 - - user_ids = torch.arange(n_user) - good_item_ids = torch.arange(n_good_item) - bad_item_ids = torch.arange(n_good_item, n_item) - - # User → Good Item: weight=1 - u2gi_src = user_ids.repeat_interleave(n_good_item) - u2gi_dst = good_item_ids.repeat(n_user) - u2gi_w = torch.ones(n_user * n_good_item) - - # User → Bad Item: weight=0 - u2bi_src = user_ids.repeat_interleave(n_bad_item) - u2bi_dst = bad_item_ids.repeat(n_user) - u2bi_w = torch.zeros(n_user * n_bad_item) - - # Good Item → User: weight=1 (2nd-hop back to users) - gi2u_src = good_item_ids.repeat_interleave(n_user) - gi2u_dst = user_ids.repeat(n_good_item) - gi2u_w = torch.ones(n_good_item * n_user) - - u2i_src = torch.cat([u2gi_src, u2bi_src]) - u2i_dst = torch.cat([u2gi_dst, u2bi_dst]) - u2i_w = torch.cat([u2gi_w, u2bi_w]) - n_u2i_edges = u2i_src.shape[0] - - user_feats = torch.full((n_user, 1), 2.0) - # Item feature encodes type: good=1.0, bad=0.0 - item_feats = torch.cat( - [ - torch.full((n_good_item, 1), 1.0), - torch.full((n_bad_item, 1), 0.0), - ] - ) - - partition_output = PartitionOutput( - node_partition_book={ - _USER: torch.zeros(n_user), - _ITEM: torch.zeros(n_item), - }, - edge_partition_book={ - _USER_TO_ITEM: torch.zeros(n_u2i_edges), - _ITEM_TO_USER: torch.zeros(gi2u_src.shape[0]), - }, - partitioned_edge_index={ - _USER_TO_ITEM: GraphPartitionData( - edge_index=torch.stack([u2i_src, u2i_dst]), - edge_ids=None, - weights=u2i_w, - ), - _ITEM_TO_USER: GraphPartitionData( - edge_index=torch.stack([gi2u_src, gi2u_dst]), - edge_ids=None, - weights=gi2u_w, - ), - }, - partitioned_node_features={ - _USER: FeaturePartitionData(feats=user_feats, ids=torch.arange(n_user)), - _ITEM: FeaturePartitionData(feats=item_feats, ids=torch.arange(n_item)), - }, - partitioned_edge_features=None, - partitioned_positive_labels=None, - partitioned_negative_labels=None, - partitioned_node_labels=None, - ) - return partition_output, n_user, n_good_item, n_bad_item - - def _build_heterogeneous_bipartite_partial_weight_graph() -> tuple[ - PartitionOutput, int, int, int + PartitionOutput, int ]: """Same graph as _build_heterogeneous_bipartite_weight_graph but ITEM_TO_USER is unweighted. @@ -268,35 +99,35 @@ def _build_heterogeneous_bipartite_partial_weight_graph() -> tuple[ partition_output = PartitionOutput( node_partition_book={ - _USER: torch.zeros(n_user), - _ITEM: torch.zeros(n_item), + USER: torch.zeros(n_user), + ITEM: torch.zeros(n_item), }, edge_partition_book={ - _USER_TO_ITEM: torch.zeros(n_u2i_edges), - _ITEM_TO_USER: torch.zeros(gi2u_src.shape[0]), + USER_TO_ITEM: torch.zeros(n_u2i_edges), + ITEM_TO_USER: torch.zeros(gi2u_src.shape[0]), }, partitioned_edge_index={ - _USER_TO_ITEM: GraphPartitionData( + USER_TO_ITEM: GraphPartitionData( edge_index=torch.stack([u2i_src, u2i_dst]), edge_ids=None, weights=u2i_w, ), - _ITEM_TO_USER: GraphPartitionData( + ITEM_TO_USER: GraphPartitionData( edge_index=torch.stack([gi2u_src, gi2u_dst]), edge_ids=None, weights=None, # unweighted — samples uniformly ), }, partitioned_node_features={ - _USER: FeaturePartitionData(feats=user_feats, ids=torch.arange(n_user)), - _ITEM: FeaturePartitionData(feats=item_feats, ids=torch.arange(n_item)), + USER: FeaturePartitionData(feats=user_feats, ids=torch.arange(n_user)), + ITEM: FeaturePartitionData(feats=item_feats, ids=torch.arange(n_item)), }, partitioned_edge_features=None, partitioned_positive_labels=None, partitioned_negative_labels=None, partitioned_node_labels=None, ) - return partition_output, n_user, n_good_item, n_bad_item + return partition_output, n_user # --------------------------------------------------------------------------- @@ -355,7 +186,7 @@ def _run_weighted_sampling_correctness_heterogeneous( ) loader = DistNeighborLoader( dataset=dataset, - input_nodes=(_USER, node_ids[_USER]), + input_nodes=(USER, node_ids[USER]), num_neighbors=[10, 5], with_weight=True, pin_memory_device=torch.device("cpu"), @@ -363,8 +194,8 @@ def _run_weighted_sampling_correctness_heterogeneous( count = 0 for datum in loader: assert isinstance(datum, HeteroData), f"Expected HeteroData, got {type(datum)}" - if _ITEM in datum.node_types: - item_x = datum[_ITEM].x + if ITEM in datum.node_types: + item_x = datum[ITEM].x assert item_x is not None, "Item features missing from sampled subgraph" bad_mask = item_x[:, 0] == 0.0 assert not bad_mask.any(), ( @@ -802,7 +633,7 @@ def test_weighted_sampling_never_traverses_zero_weight_edges_homogeneous( sampling. Fanout [10, 5] samples fewer neighbors than available good ones, so the weighted sampler actively selects from the pool each hop. """ - partition_output, n_hub, _, _ = _build_homogeneous_bipartite_weight_graph() + partition_output, n_hub = build_homogeneous_bipartite_weight_graph() assert isinstance(partition_output.partitioned_edge_index, GraphPartitionData) expected_weights = partition_output.partitioned_edge_index.weights @@ -829,7 +660,7 @@ def test_weighted_sampling_never_traverses_zero_weight_edges_heterogeneous( Fanout [10, 5] is smaller than the 40 available good items, so the sampler actively selects. """ - partition_output, n_user, _, _ = _build_heterogeneous_bipartite_weight_graph() + partition_output, n_user = build_heterogeneous_bipartite_weight_graph() dataset = DistDataset(rank=0, world_size=1, edge_dir="out") dataset.build(partition_output=partition_output) @@ -846,9 +677,7 @@ def test_weighted_sampling_partial_weights_heterogeneous(self) -> None: Verifies that mixing weighted and unweighted edge types in one heterogeneous graph does not crash and that weighted edges still behave correctly. """ - partition_output, n_user, _, _ = ( - _build_heterogeneous_bipartite_partial_weight_graph() - ) + partition_output, n_user = _build_heterogeneous_bipartite_partial_weight_graph() dataset = DistDataset(rank=0, world_size=1, edge_dir="out") dataset.build(partition_output=partition_output) From 319e839d247147f3d8a193cbded08702325d516f Mon Sep 17 00:00:00 2001 From: mkolodner Date: Thu, 28 May 2026 18:14:01 +0000 Subject: [PATCH 3/4] Update --- gigl-core/core/sampling/ppr_forward_push.cpp | 7 +- gigl-core/tests/ppr_forward_push_test.cpp | 48 +++++- gigl/distributed/dist_ppr_sampler.py | 156 +++++++++++++++++- .../distributed/bipartite_weight_graph.py | 6 +- .../distributed_ppr_weighted_sampling_test.py | 18 ++ 5 files changed, 218 insertions(+), 17 deletions(-) diff --git a/gigl-core/core/sampling/ppr_forward_push.cpp b/gigl-core/core/sampling/ppr_forward_push.cpp index c71d81ada..a63075e94 100644 --- a/gigl-core/core/sampling/ppr_forward_push.cpp +++ b/gigl-core/core/sampling/ppr_forward_push.cpp @@ -192,7 +192,9 @@ void PPRForwardPush::pushResiduals( if (flatWeightsPtr != nullptr) { std::vector neighborWeights(count); for (int64_t neighborIdx = 0; neighborIdx < count; ++neighborIdx) { - neighborWeights[neighborIdx] = flatWeightsPtr[offset + neighborIdx]; + double weight = flatWeightsPtr[offset + neighborIdx]; + TORCH_CHECK(weight >= 0.0, "PPR edge weights must be non-negative."); + neighborWeights[neighborIdx] = weight; } fetchedWeights[key] = std::move(neighborWeights); } @@ -391,6 +393,9 @@ void PPRForwardPush::pushResiduals( TORCH_INTERNAL_ASSERT(weights.size() == neighbors.size(), "weightList and neighborList must have the same size"); for (int32_t i = 0; i < static_cast(neighbors.size()); ++i) { + if (weights[i] == 0.0) { + continue; + } int32_t neighborNodeId = neighbors[i]; dstNodeTypeState.residuals[neighborNodeId] += baseResidual * weights[i] / totalFetchedWeight; diff --git a/gigl-core/tests/ppr_forward_push_test.cpp b/gigl-core/tests/ppr_forward_push_test.cpp index 604763deb..e61c2dd2b 100644 --- a/gigl-core/tests/ppr_forward_push_test.cpp +++ b/gigl-core/tests/ppr_forward_push_test.cpp @@ -3,6 +3,8 @@ using gigl::PPRForwardPush; +using FetchedByEtypeId = std::unordered_map>; + // Builds a single-edge-type, single-node-type PPRForwardPush. static PPRForwardPush makeState( const std::vector& seeds, @@ -20,16 +22,21 @@ static PPRForwardPush makeState( } // Convenience wrapper: build the fetchedByEtypeId argument for pushResiduals -// from flat vectors, keeping test call sites readable. -static std::unordered_map> +// from flat vectors, keeping test call sites readable. Empty weights select +// uniform-residual mode. +static FetchedByEtypeId makeFetched(int32_t edgeTypeId, const std::vector& nodeIds, const std::vector& flatNeighborIds, - const std::vector& counts) { + const std::vector& counts, + const std::vector& flatWeights = {}) { + auto weightsTensor = + flatWeights.empty() ? torch::empty({0}, torch::kDouble) : torch::tensor(flatWeights, torch::kDouble); return {{edgeTypeId, {torch::tensor(nodeIds, torch::kLong), torch::tensor(flatNeighborIds, torch::kLong), - torch::tensor(counts, torch::kLong)}}}; + torch::tensor(counts, torch::kLong), + weightsTensor}}}; } // After construction, drainQueue() returns the seed node under etype 0. @@ -89,6 +96,39 @@ TEST(PPRForwardPush, ResidualDistributedToNeighbor) { EXPECT_NEAR(weights[1].item(), static_cast((1.0 - alpha) * alpha), 1e-5F); } +// In weighted mode, zero-weight edges must not enqueue a zero-residual neighbor. +TEST(PPRForwardPush, WeightedResidualSkipsZeroWeightNeighbor) { + const double alpha = 0.5; + auto state = makeState(/*seeds=*/{0}, alpha, /*requeueThresholdFactor=*/1e-6, /*degrees=*/{2, 0, 0}); + + state.drainQueue(); + state.pushResiduals(makeFetched( + /*edgeTypeId=*/0, + /*nodeIds=*/{0}, + /*flatNeighborIds=*/{1, 2}, + /*counts=*/{2}, + /*flatWeights=*/{0.0, 2.0})); + + auto iter2 = state.drainQueue(); + ASSERT_TRUE(iter2.has_value()); + const auto& iter2Map = iter2.value(); + ASSERT_NE(iter2Map.find(0), iter2Map.end()); + ASSERT_EQ(iter2Map.at(0).size(0), 1); + EXPECT_EQ(iter2Map.at(0)[0].item(), 2); + + state.pushResiduals({}); + EXPECT_FALSE(state.drainQueue().has_value()); + + auto topk = state.extractTopK(10); + ASSERT_NE(topk.find(0), topk.end()); + const auto& [ids, weights, counts] = topk.at(0); + ASSERT_EQ(counts[0].item(), 2); + EXPECT_EQ(ids[0].item(), 0); + EXPECT_EQ(ids[1].item(), 2); + EXPECT_NEAR(weights[0].item(), static_cast(alpha), 1e-5F); + EXPECT_NEAR(weights[1].item(), static_cast((1.0 - alpha) * alpha), 1e-5F); +} + // Two seeds (0 and 1) both push residual to sink node 2. The neighbor-lookup // request must deduplicate to one entry for node 2, yet both seeds must still // accumulate a PPR score for it. diff --git a/gigl/distributed/dist_ppr_sampler.py b/gigl/distributed/dist_ppr_sampler.py index 8e61170aa..f4cb9c270 100644 --- a/gigl/distributed/dist_ppr_sampler.py +++ b/gigl/distributed/dist_ppr_sampler.py @@ -7,6 +7,7 @@ # TODO: Once gigl_core has a stable Python interface, re-export PPRForwardPush # under a gigl.core namespace rather than importing directly from the C++ extension. from gigl_core import PPRForwardPush +from graphlearn_torch.data import Graph from graphlearn_torch.sampler import ( HeteroSamplerOutput, NeighborOutput, @@ -36,6 +37,8 @@ _PPR_HOMOGENEOUS_NODE_TYPE, ) +EdgeWeightLookup = tuple[torch.Tensor, torch.Tensor] + class DistPPRNeighborSampler(BaseDistNeighborSampler): """Personalized PageRank (PPR) based distributed neighbor sampler. @@ -202,6 +205,9 @@ def __init__( self._node_type_to_total_degree.get(nt, torch.zeros(0, dtype=torch.int32)) for nt in all_node_types ] + self._edge_weight_lookup: Optional[ + Union[EdgeWeightLookup, dict[EdgeType, EdgeWeightLookup]] + ] = self._build_edge_weight_lookup() def _build_total_degree_tensors( self, @@ -259,6 +265,146 @@ def _get_destination_type(self, edge_type: EdgeType) -> NodeType: """Get the node type at the destination end of an edge type.""" return edge_type[0] if self.edge_dir == "in" else edge_type[-1] + def _build_edge_weight_lookup( + self, + ) -> Optional[Union[EdgeWeightLookup, dict[EdgeType, EdgeWeightLookup]]]: + """Build edge-id keyed weight lookup tables for weighted PPR residuals. + + GLT's ``sample_with_edge`` returns graph edge IDs, not guaranteed local + tensor positions. These lookup tables map sampled edge IDs back to the + aligned edge weights stored in the local graph topology. + + Returns: + A homogeneous lookup tuple, a heterogeneous lookup dict, or ``None`` + when weighted residuals are disabled. + + Raises: + ValueError: If edge weights were requested but graph topology does + not contain aligned edge IDs and weights. + """ + if self._edge_weights is None: + return None + + graph = self.data.graph + if graph is None: + raise ValueError("Cannot build PPR edge-weight lookup without graph data.") + + if self._is_homogeneous: + assert isinstance(graph, Graph) + return self._build_sorted_edge_weight_lookup_from_graph(graph) + + assert isinstance(graph, dict) + assert isinstance(self._edge_weights, dict) + return { + edge_type: self._build_sorted_edge_weight_lookup_from_graph( + graph[edge_type] + ) + for edge_type in self._edge_weights + } + + @staticmethod + def _build_sorted_edge_weight_lookup_from_graph( + graph: Graph, + ) -> EdgeWeightLookup: + """Build a sorted edge-id lookup from a GLT graph's local topology.""" + edge_ids = graph.topo.edge_ids + edge_weights = graph.topo.edge_weights + if edge_ids is None or edge_weights is None: + raise ValueError( + "Weighted PPR requires graph topology to contain both edge IDs " + "and edge weights." + ) + return DistPPRNeighborSampler._build_sorted_edge_weight_lookup( + edge_ids=edge_ids, + edge_weights=edge_weights, + ) + + @staticmethod + def _build_sorted_edge_weight_lookup( + edge_ids: torch.Tensor, + edge_weights: torch.Tensor, + ) -> EdgeWeightLookup: + """Build a sorted lookup tuple from edge IDs to weights. + + Args: + edge_ids: Edge IDs aligned with ``edge_weights``. + edge_weights: Per-edge weights aligned with ``edge_ids``. + + Returns: + ``(sorted_edge_ids, sorted_weights)``. + + Raises: + ValueError: If IDs and weights are not aligned or IDs are duplicated. + """ + if edge_ids.numel() != edge_weights.numel(): + raise ValueError( + f"Edge IDs and weights must have the same length, got " + f"{edge_ids.numel()} and {edge_weights.numel()}." + ) + + sorted_edge_ids, sort_indices = torch.sort(edge_ids.to(torch.long)) + if ( + sorted_edge_ids.numel() > 1 + and (sorted_edge_ids[1:] == sorted_edge_ids[:-1]).any().item() + ): + raise ValueError("Edge IDs must be unique within each edge type.") + + return sorted_edge_ids, edge_weights.to(torch.float64)[sort_indices] + + @staticmethod + def _lookup_edge_weights_by_id( + edge_ids: torch.Tensor, + edge_weight_lookup: EdgeWeightLookup, + ) -> torch.Tensor: + """Resolve sampled edge IDs to weights without assuming local ID offsets.""" + sorted_edge_ids, sorted_weights = edge_weight_lookup + if edge_ids.numel() == 0: + return torch.empty(0, dtype=torch.float64, device=edge_ids.device) + if sorted_edge_ids.numel() == 0: + raise ValueError("Cannot resolve non-empty edge IDs against empty weights.") + + requested_edge_ids = edge_ids.to( + device=sorted_edge_ids.device, dtype=torch.long + ) + insertion_indices = torch.searchsorted(sorted_edge_ids, requested_edge_ids) + clamped_indices = insertion_indices.clamp(max=sorted_edge_ids.numel() - 1) + found_mask = (insertion_indices < sorted_edge_ids.numel()) & ( + sorted_edge_ids[clamped_indices] == requested_edge_ids + ) + if not found_mask.all().item(): + missing_edge_ids = requested_edge_ids[~found_mask][:10].tolist() + raise ValueError( + f"Sampled edge IDs were not found in the edge-weight lookup. " + f"First missing IDs: {missing_edge_ids}." + ) + + return sorted_weights[clamped_indices].to( + device=edge_ids.device, dtype=torch.float64 + ) + + def _lookup_edge_weights( + self, + edge_ids: torch.Tensor, + edge_type_id: int, + ) -> torch.Tensor: + """Resolve sampled edge IDs to edge weights for one PPR edge type.""" + if self._edge_weight_lookup is None: + raise ValueError("PPR edge-weight lookup is not initialized.") + + if self._is_homogeneous: + assert isinstance(self._edge_weight_lookup, tuple) + return self._lookup_edge_weights_by_id( + edge_ids=edge_ids, + edge_weight_lookup=self._edge_weight_lookup, + ) + + assert isinstance(self._edge_weight_lookup, dict) + edge_type = self._etype_id_to_etype[edge_type_id] + return self._lookup_edge_weights_by_id( + edge_ids=edge_ids, + edge_weight_lookup=self._edge_weight_lookup[edge_type], + ) + async def _batch_fetch_neighbors( self, nodes_by_etype_id: dict[int, torch.Tensor], @@ -325,15 +471,7 @@ async def _batch_fetch_neighbors( "output.edge must be set when edge_weights is provided; " "ensure with_edge=True in SamplingConfig (hardcoded in create_sampling_config)." ) - if self._is_homogeneous: - assert isinstance(self._edge_weights, torch.Tensor) - flat_weights = self._edge_weights[output.edge].to(torch.float64) - else: - assert isinstance(self._edge_weights, dict) - etype = self._etype_id_to_etype[eid] - flat_weights = self._edge_weights[etype][output.edge].to( - torch.float64 - ) + flat_weights = self._lookup_edge_weights(output.edge, eid) else: flat_weights = _empty_weights result[eid] = ( diff --git a/tests/test_assets/distributed/bipartite_weight_graph.py b/tests/test_assets/distributed/bipartite_weight_graph.py index acf37ba85..b2f8950fc 100644 --- a/tests/test_assets/distributed/bipartite_weight_graph.py +++ b/tests/test_assets/distributed/bipartite_weight_graph.py @@ -86,7 +86,7 @@ def build_homogeneous_bipartite_weight_graph() -> tuple[PartitionOutput, int]: edge_partition_book=torch.zeros(n_edges), partitioned_edge_index=GraphPartitionData( edge_index=edge_index, - edge_ids=None, + edge_ids=torch.arange(n_edges) + 1_000, weights=weights, ), partitioned_node_features=FeaturePartitionData( @@ -167,12 +167,12 @@ def build_heterogeneous_bipartite_weight_graph() -> tuple[PartitionOutput, int]: partitioned_edge_index={ USER_TO_ITEM: GraphPartitionData( edge_index=torch.stack([u2i_src, u2i_dst]), - edge_ids=None, + edge_ids=torch.arange(n_u2i_edges) + 1_000, weights=u2i_w, ), ITEM_TO_USER: GraphPartitionData( edge_index=torch.stack([gi2u_src, gi2u_dst]), - edge_ids=None, + edge_ids=torch.arange(gi2u_src.shape[0]) + 10_000, weights=gi2u_w, ), }, diff --git a/tests/unit/distributed/distributed_ppr_weighted_sampling_test.py b/tests/unit/distributed/distributed_ppr_weighted_sampling_test.py index 8f666ccb7..11b74c774 100644 --- a/tests/unit/distributed/distributed_ppr_weighted_sampling_test.py +++ b/tests/unit/distributed/distributed_ppr_weighted_sampling_test.py @@ -18,6 +18,7 @@ from torch_geometric.data import Data, HeteroData from gigl.distributed.dist_dataset import DistDataset +from gigl.distributed.dist_ppr_sampler import DistPPRNeighborSampler from gigl.distributed.distributed_neighborloader import DistNeighborLoader from gigl.distributed.sampler_options import PPRSamplerOptions from tests.test_assets.distributed.bipartite_weight_graph import ( @@ -165,6 +166,23 @@ def test_ppr_weighted_never_traverses_zero_weight_edges_homogeneous(self) -> Non nprocs=1, ) + def test_ppr_weight_lookup_uses_edge_ids_not_tensor_offsets(self) -> None: + """Sampled edge IDs may be non-contiguous and must resolve by ID.""" + lookup = DistPPRNeighborSampler._build_sorted_edge_weight_lookup( + edge_ids=torch.tensor([42, 7, 100]), + edge_weights=torch.tensor([4.2, 0.7, 10.0]), + ) + + actual = DistPPRNeighborSampler._lookup_edge_weights_by_id( + edge_ids=torch.tensor([7, 100, 42]), + edge_weight_lookup=lookup, + ) + + torch.testing.assert_close( + actual, + torch.tensor([0.7, 10.0, 4.2], dtype=torch.float64), + ) + def test_ppr_weighted_never_traverses_zero_weight_edges_heterogeneous(self) -> None: """Heterogeneous: weight=0 user→item edges to bad items never contribute PPR residual. From 9a5e730c723edefacdfb34b8cbc828e0f98ead29 Mon Sep 17 00:00:00 2001 From: mkolodner Date: Tue, 16 Jun 2026 00:02:48 +0000 Subject: [PATCH 4/4] Fix weighted PPR edge weight lookup --- gigl-core/core/sampling/ppr_forward_push.cpp | 7 +- gigl-core/tests/ppr_forward_push_test.cpp | 84 ++-- gigl/distributed/dist_ppr_sampler.py | 373 +++++++++++------- .../utils/weighted_neighbor_sampling.py | 233 +++++++++++ .../distributed_ppr_weighted_sampling_test.py | 47 ++- tests/unit/distributed/utils/degree_test.py | 4 +- 6 files changed, 564 insertions(+), 184 deletions(-) create mode 100644 gigl/distributed/utils/weighted_neighbor_sampling.py diff --git a/gigl-core/core/sampling/ppr_forward_push.cpp b/gigl-core/core/sampling/ppr_forward_push.cpp index a63075e94..f79c34037 100644 --- a/gigl-core/core/sampling/ppr_forward_push.cpp +++ b/gigl-core/core/sampling/ppr_forward_push.cpp @@ -397,8 +397,11 @@ void PPRForwardPush::pushResiduals( continue; } int32_t neighborNodeId = neighbors[i]; - dstNodeTypeState.residuals[neighborNodeId] += - baseResidual * weights[i] / totalFetchedWeight; + double residualContribution = baseResidual * weights[i] / totalFetchedWeight; + if (residualContribution == 0.0) { + continue; + } + dstNodeTypeState.residuals[neighborNodeId] += residualContribution; double threshold = _requeueThresholdFactor * static_cast(getTotalDegree(neighborNodeId, dstNodeTypeId)); diff --git a/gigl-core/tests/ppr_forward_push_test.cpp b/gigl-core/tests/ppr_forward_push_test.cpp index e61c2dd2b..79bbccc0c 100644 --- a/gigl-core/tests/ppr_forward_push_test.cpp +++ b/gigl-core/tests/ppr_forward_push_test.cpp @@ -3,33 +3,31 @@ using gigl::PPRForwardPush; -using FetchedByEtypeId = std::unordered_map>; +using FetchedByEtypeId = + std::unordered_map>; // Builds a single-edge-type, single-node-type PPRForwardPush. -static PPRForwardPush makeState( - const std::vector& seeds, - double alpha, - double requeueThresholdFactor, - const std::vector& degrees) { - return PPRForwardPush( - torch::tensor(seeds, torch::kLong), - /*seedNodeTypeId=*/0, - alpha, - requeueThresholdFactor, - /*nodeTypeToEdgeTypeIds=*/{{0}}, - /*edgeTypeToDstNtypeId=*/{0}, - {torch::tensor(degrees, torch::kInt)}); +static PPRForwardPush makeState(const std::vector& seeds, + double alpha, + double requeueThresholdFactor, + const std::vector& degrees) { + return PPRForwardPush(torch::tensor(seeds, torch::kLong), + /*seedNodeTypeId=*/0, + alpha, + requeueThresholdFactor, + /*nodeTypeToEdgeTypeIds=*/{{0}}, + /*edgeTypeToDstNtypeId=*/{0}, + {torch::tensor(degrees, torch::kInt)}); } // Convenience wrapper: build the fetchedByEtypeId argument for pushResiduals // from flat vectors, keeping test call sites readable. Empty weights select // uniform-residual mode. -static FetchedByEtypeId -makeFetched(int32_t edgeTypeId, - const std::vector& nodeIds, - const std::vector& flatNeighborIds, - const std::vector& counts, - const std::vector& flatWeights = {}) { +static FetchedByEtypeId makeFetched(int32_t edgeTypeId, + const std::vector& nodeIds, + const std::vector& flatNeighborIds, + const std::vector& counts, + const std::vector& flatWeights = {}) { auto weightsTensor = flatWeights.empty() ? torch::empty({0}, torch::kDouble) : torch::tensor(flatWeights, torch::kDouble); return {{edgeTypeId, @@ -136,13 +134,14 @@ TEST(PPRForwardPush, DeduplicatesNodesAcrossSeeds) { auto state = makeState(/*seeds=*/{0, 1}, /*alpha=*/0.15, /*requeueThresholdFactor=*/1e-6, /*degrees=*/{1, 1, 0}); state.drainQueue(); - state.pushResiduals(makeFetched(/*edgeTypeId=*/0, /*nodeIds=*/{0, 1}, /*flatNeighborIds=*/{2, 2}, /*counts=*/{1, 1})); + state.pushResiduals( + makeFetched(/*edgeTypeId=*/0, /*nodeIds=*/{0, 1}, /*flatNeighborIds=*/{2, 2}, /*counts=*/{1, 1})); auto iter2 = state.drainQueue(); ASSERT_TRUE(iter2.has_value()); const auto& iter2Map = iter2.value(); ASSERT_NE(iter2Map.find(0), iter2Map.end()); - EXPECT_EQ(iter2Map.at(0).size(0), 1); // node 2 deduplicated in the lookup request + EXPECT_EQ(iter2Map.at(0).size(0), 1); // node 2 deduplicated in the lookup request state.pushResiduals({}); EXPECT_FALSE(state.drainQueue().has_value()); @@ -151,12 +150,45 @@ TEST(PPRForwardPush, DeduplicatesNodesAcrossSeeds) { ASSERT_NE(topk.find(0), topk.end()); const auto& [ids, weights, counts] = topk.at(0); // Each seed (batch indices 0 and 1) should have 2 nodes in its top-k. - EXPECT_EQ(counts[0].item(), 2); // seed 0: nodes {0, 2} - EXPECT_EQ(counts[1].item(), 2); // seed 1: nodes {1, 2} + EXPECT_EQ(counts[0].item(), 2); // seed 0: nodes {0, 2} + EXPECT_EQ(counts[1].item(), 2); // seed 1: nodes {1, 2} // The flat id layout is [seed0_top1, seed0_top2, seed1_top1, seed1_top2]. // Within each seed the highest scorer comes first, so seed-node beats node 2. - EXPECT_EQ(ids[1].item(), 2); // seed 0's second node is node 2 - EXPECT_EQ(ids[3].item(), 2); // seed 1's second node is node 2 + EXPECT_EQ(ids[1].item(), 2); // seed 0's second node is node 2 + EXPECT_EQ(ids[3].item(), 2); // seed 1's second node is node 2 +} + +// In weighted mode, a neighbor receiving zero residual must not be re-queued. +// Node 2 is intentionally outside the degree tensor; the test would fail if the +// zero-weight edge tried to compute its requeue threshold. +TEST(PPRForwardPush, WeightedZeroContributionNeighborDoesNotRequeue) { + auto state = makeState(/*seeds=*/{0}, /*alpha=*/0.15, /*requeueThresholdFactor=*/1e-6, /*degrees=*/{2, 0}); + + state.drainQueue(); + state.pushResiduals(makeFetched(/*edgeTypeId=*/0, + /*nodeIds=*/{0}, + /*flatNeighborIds=*/{1, 2}, + /*counts=*/{2}, + /*flatWeights=*/{1.0, 0.0})); + + auto iter2 = state.drainQueue(); + ASSERT_TRUE(iter2.has_value()); + const auto& iter2Map = iter2.value(); + ASSERT_NE(iter2Map.find(0), iter2Map.end()); + ASSERT_EQ(iter2Map.at(0).size(0), 1); + EXPECT_EQ(iter2Map.at(0)[0].item(), 1); + + state.pushResiduals({}); + EXPECT_FALSE(state.drainQueue().has_value()); + + auto topk = state.extractTopK(10); + ASSERT_NE(topk.find(0), topk.end()); + const auto& result = topk.at(0); + const auto& ids = std::get<0>(result); + const auto& counts = std::get<2>(result); + EXPECT_EQ(counts[0].item(), 2); + EXPECT_EQ(ids[0].item(), 0); + EXPECT_EQ(ids[1].item(), 1); } // extractTopK respects the maxPprNodes limit. diff --git a/gigl/distributed/dist_ppr_sampler.py b/gigl/distributed/dist_ppr_sampler.py index f4cb9c270..d51388666 100644 --- a/gigl/distributed/dist_ppr_sampler.py +++ b/gigl/distributed/dist_ppr_sampler.py @@ -7,7 +7,8 @@ # TODO: Once gigl_core has a stable Python interface, re-export PPRForwardPush # under a gigl.core namespace rather than importing directly from the C++ extension. from gigl_core import PPRForwardPush -from graphlearn_torch.data import Graph +from graphlearn_torch.distributed.event_loop import wrap_torch_future +from graphlearn_torch.distributed.rpc import rpc_register, rpc_request_async from graphlearn_torch.sampler import ( HeteroSamplerOutput, NeighborOutput, @@ -18,6 +19,14 @@ from graphlearn_torch.utils import merge_dict from gigl.distributed.base_sampler import BaseDistNeighborSampler +from gigl.distributed.utils.weighted_neighbor_sampling import ( + LocalEdgeWeightLookup, + PartialWeightedNeighborOutput, + WeightedNeighborOutput, + WeightedSamplingCallee, + empty_weighted_neighbor_output, + sample_local_one_hop_with_edge_weights, +) from gigl.types.graph import is_label_edge_type # Trailing "." is an intentional separator. These constants are used both to @@ -37,8 +46,6 @@ _PPR_HOMOGENEOUS_NODE_TYPE, ) -EdgeWeightLookup = tuple[torch.Tensor, torch.Tensor] - class DistPPRNeighborSampler(BaseDistNeighborSampler): """Personalized PageRank (PPR) based distributed neighbor sampler. @@ -112,7 +119,8 @@ def __init__( self._requeue_threshold_factor = alpha * eps self._num_neighbors_per_hop = num_neighbors_per_hop self._max_fetch_iterations = max_fetch_iterations - self._edge_weights = edge_weights + self._edge_weight_lookup: Optional[LocalEdgeWeightLookup] = None + self._ppr_weighted_sample_callee_id: Optional[int] = None # Build mapping from node type to edge types that can be traversed from that node type. self._node_type_to_edge_types: dict[NodeType, list[EdgeType]] = defaultdict( @@ -205,9 +213,18 @@ def __init__( self._node_type_to_total_degree.get(nt, torch.zeros(0, dtype=torch.int32)) for nt in all_node_types ] - self._edge_weight_lookup: Optional[ - Union[EdgeWeightLookup, dict[EdgeType, EdgeWeightLookup]] - ] = self._build_edge_weight_lookup() + + if edge_weights is not None: + self._edge_weight_lookup = LocalEdgeWeightLookup( + self.dist_graph.local_graph + ) + self._ppr_weighted_sample_callee_id = rpc_register( + WeightedSamplingCallee( + sampler=self.sampler, + edge_weight_lookup=self._edge_weight_lookup, + device=self.device, + ) + ) def _build_total_degree_tensors( self, @@ -265,145 +282,209 @@ def _get_destination_type(self, edge_type: EdgeType) -> NodeType: """Get the node type at the destination end of an edge type.""" return edge_type[0] if self.edge_dir == "in" else edge_type[-1] - def _build_edge_weight_lookup( + def _attach_edge_weights_to_neighbor_output( self, - ) -> Optional[Union[EdgeWeightLookup, dict[EdgeType, EdgeWeightLookup]]]: - """Build edge-id keyed weight lookup tables for weighted PPR residuals. - - GLT's ``sample_with_edge`` returns graph edge IDs, not guaranteed local - tensor positions. These lookup tables map sampled edge IDs back to the - aligned edge weights stored in the local graph topology. - - Returns: - A homogeneous lookup tuple, a heterogeneous lookup dict, or ``None`` - when weighted residuals are disabled. - - Raises: - ValueError: If edge weights were requested but graph topology does - not contain aligned edge IDs and weights. + output: Union[NeighborOutput, torch.Tensor], + srcs: torch.Tensor, + edge_type: Optional[EdgeType], + ) -> WeightedNeighborOutput: + """Attach local edge weights to a base one-hop ``NeighborOutput``. + + Used when the inherited GLT ``_sample_one_hop`` can do all routing for + us (graph-cached or single-partition sampling). The only extra behavior + here is resolving sampled edge ids through the local weight lookup. """ - if self._edge_weights is None: - return None - - graph = self.data.graph - if graph is None: - raise ValueError("Cannot build PPR edge-weight lookup without graph data.") - - if self._is_homogeneous: - assert isinstance(graph, Graph) - return self._build_sorted_edge_weight_lookup_from_graph(graph) - - assert isinstance(graph, dict) - assert isinstance(self._edge_weights, dict) - return { - edge_type: self._build_sorted_edge_weight_lookup_from_graph( - graph[edge_type] - ) - for edge_type in self._edge_weights - } - - @staticmethod - def _build_sorted_edge_weight_lookup_from_graph( - graph: Graph, - ) -> EdgeWeightLookup: - """Build a sorted edge-id lookup from a GLT graph's local topology.""" - edge_ids = graph.topo.edge_ids - edge_weights = graph.topo.edge_weights - if edge_ids is None or edge_weights is None: - raise ValueError( - "Weighted PPR requires graph topology to contain both edge IDs " - "and edge weights." + if self._edge_weight_lookup is None: + raise RuntimeError("Cannot attach edge weights without a local lookup.") + if not isinstance(output, NeighborOutput): + return empty_weighted_neighbor_output(srcs=srcs, device=self.device) + if output.edge is None: + raise RuntimeError( + "Weighted PPR requires GLT one-hop sampling to return edge ids." ) - return DistPPRNeighborSampler._build_sorted_edge_weight_lookup( - edge_ids=edge_ids, - edge_weights=edge_weights, + return WeightedNeighborOutput( + nbr=output.nbr, + nbr_num=output.nbr_num, + edge=output.edge, + weight=self._edge_weight_lookup.lookup(output.edge, edge_type), ) - @staticmethod - def _build_sorted_edge_weight_lookup( - edge_ids: torch.Tensor, - edge_weights: torch.Tensor, - ) -> EdgeWeightLookup: - """Build a sorted lookup tuple from edge IDs to weights. + def _stitch_edge_weights( + self, + input_seeds: torch.Tensor, + results: list[PartialWeightedNeighborOutput], + ) -> torch.Tensor: + """Stitch partitioned sampled edge weights in original source-node order. Args: - edge_ids: Edge IDs aligned with ``edge_weights``. - edge_weights: Per-edge weights aligned with ``edge_ids``. + input_seeds: The source nodes passed to the one-hop sampler. + results: Per-partition weighted one-hop outputs. Returns: - ``(sorted_edge_ids, sorted_weights)``. - - Raises: - ValueError: If IDs and weights are not aligned or IDs are duplicated. + A flat weight tensor parallel to the stitched neighbor and edge-id tensors. """ - if edge_ids.numel() != edge_weights.numel(): - raise ValueError( - f"Edge IDs and weights must have the same length, got " - f"{edge_ids.numel()} and {edge_weights.numel()}." - ) - - sorted_edge_ids, sort_indices = torch.sort(edge_ids.to(torch.long)) - if ( - sorted_edge_ids.numel() > 1 - and (sorted_edge_ids[1:] == sorted_edge_ids[:-1]).any().item() - ): - raise ValueError("Edge IDs must be unique within each edge type.") - - return sorted_edge_ids, edge_weights.to(torch.float64)[sort_indices] + # GLT's C++ stitcher only accepts neighbor ids, neighbor counts, and + # edge ids. This mirrors that row-ordering logic for the one extra + # tensor we carry in weighted PPR: edge weights. + if not results: + return torch.empty(0, dtype=torch.float32, device=self.device) + + empty = results[0].output.weight.new_empty((0,)) + weight_chunks: list[torch.Tensor] = [empty] * input_seeds.size(0) + + for result in results: + output = result.output + counts = output.nbr_num.to(torch.device("cpu")) + start = 0 + for local_row, input_row in enumerate(result.index.cpu().tolist()): + count = int(counts[local_row]) + if count > 0: + weight_chunks[input_row] = output.weight[start : start + count] + start += count + + return torch.cat(weight_chunks) if weight_chunks else empty + + def _stitch_weighted_sample_results( + self, + input_seeds: torch.Tensor, + results: list[PartialWeightedNeighborOutput], + ) -> WeightedNeighborOutput: + """Stitch partitioned weighted one-hop outputs. - @staticmethod - def _lookup_edge_weights_by_id( - edge_ids: torch.Tensor, - edge_weight_lookup: EdgeWeightLookup, - ) -> torch.Tensor: - """Resolve sampled edge IDs to weights without assuming local ID offsets.""" - sorted_edge_ids, sorted_weights = edge_weight_lookup - if edge_ids.numel() == 0: - return torch.empty(0, dtype=torch.float64, device=edge_ids.device) - if sorted_edge_ids.numel() == 0: - raise ValueError("Cannot resolve non-empty edge IDs against empty weights.") - - requested_edge_ids = edge_ids.to( - device=sorted_edge_ids.device, dtype=torch.long - ) - insertion_indices = torch.searchsorted(sorted_edge_ids, requested_edge_ids) - clamped_indices = insertion_indices.clamp(max=sorted_edge_ids.numel() - 1) - found_mask = (insertion_indices < sorted_edge_ids.numel()) & ( - sorted_edge_ids[clamped_indices] == requested_edge_ids - ) - if not found_mask.all().item(): - missing_edge_ids = requested_edge_ids[~found_mask][:10].tolist() - raise ValueError( - f"Sampled edge IDs were not found in the edge-weight lookup. " - f"First missing IDs: {missing_edge_ids}." - ) + Reuses GLT's optimized neighbor/edge-id stitcher, then stitches the + parallel weight tensor using the same per-source counts. - return sorted_weights[clamped_indices].to( - device=edge_ids.device, dtype=torch.float64 + Difference from GLT ``_stitch_sample_results``: the inherited stitcher + still owns ``nbr``, ``nbr_num``, and ``edge`` ordering; this wrapper only + adds the parallel ``weight`` tensor. + """ + stitched_output = self._stitch_sample_results(input_seeds, results) + if stitched_output.edge is None: + raise RuntimeError("Weighted PPR stitching expected sampled edge ids.") + return WeightedNeighborOutput( + nbr=stitched_output.nbr, + nbr_num=stitched_output.nbr_num, + edge=stitched_output.edge, + weight=self._stitch_edge_weights(input_seeds, results), ) - def _lookup_edge_weights( + async def _sample_one_hop_with_edge_weights( self, - edge_ids: torch.Tensor, - edge_type_id: int, - ) -> torch.Tensor: - """Resolve sampled edge IDs to edge weights for one PPR edge type.""" + srcs: torch.Tensor, + num_nbr: int, + etype: Optional[EdgeType], + ) -> WeightedNeighborOutput: + """Sample one hop and return edge weights with the sampled edges. + + Remote workers attach weights before returning, so weighted PPR does not + need a second distributed lookup after one-hop sampling. + + This method intentionally tracks GLT + ``DistNeighborSampler._sample_one_hop``. The shared pieces are device + placement, partition routing, ``id_select`` conversion, remote futures, + and final stitching. The differences are limited to: + + - graph-cached/single-partition paths call the inherited method and then + attach local weights; + - multi-partition local work calls ``sample_local_one_hop_with_edge_weights``; + - remote work calls ``WeightedSamplingCallee`` instead of GLT's sampling + callee so weights come back in the same RPC response; + - final stitching adds a parallel weight tensor. + """ if self._edge_weight_lookup is None: - raise ValueError("PPR edge-weight lookup is not initialized.") + raise RuntimeError( + "_sample_one_hop_with_edge_weights called without edge weights." + ) - if self._is_homogeneous: - assert isinstance(self._edge_weight_lookup, tuple) - return self._lookup_edge_weights_by_id( - edge_ids=edge_ids, - edge_weight_lookup=self._edge_weight_lookup, + device = self.device + srcs = srcs.to(device) + if self.data.graph_caching or self.data.num_partitions == 1: + # Same as GLT for routing and sampling; the local weight lookup + # below is the only weighted-PPR addition in this branch. + output = await self._sample_one_hop( + srcs=srcs, + num_nbr=num_nbr, + etype=etype, ) + return self._attach_edge_weights_to_neighbor_output(output, srcs, etype) - assert isinstance(self._edge_weight_lookup, dict) - edge_type = self._etype_id_to_etype[edge_type_id] - return self._lookup_edge_weights_by_id( - edge_ids=edge_ids, - edge_weight_lookup=self._edge_weight_lookup[edge_type], - ) + orders = torch.arange(srcs.size(0), dtype=torch.long, device=device) + if self.edge_dir == "out": + src_ntype = etype[0] if etype is not None else None + elif self.edge_dir == "in": + src_ntype = etype[-1] if etype is not None else None + else: + raise ValueError(f"Unsupported edge_dir: {self.edge_dir}") + + partition_ids = self.dist_graph.get_node_partitions(srcs, src_ntype).to(device) + partition_results: list[PartialWeightedNeighborOutput] = [] + remote_orders_list: list[torch.Tensor] = [] + futures: list[torch.futures.Future] = [] + + for i in range(self.data.num_partitions): + partition_idx = (self.data.partition_idx + i) % self.data.num_partitions + partition_mask = partition_ids == partition_idx + if isinstance(self.dist_graph.node_pb, dict): + assert src_ntype is not None + partition_node_ids = self.data.id_select( + srcs, partition_mask, self.dist_graph.node_pb[src_ntype] + ) + else: + partition_node_ids = self.data.id_select( + srcs, partition_mask, self.dist_graph.node_pb + ) + if partition_node_ids.shape[0] == 0: + continue + + partition_orders = torch.masked_select(orders, partition_mask) + if partition_idx == self.data.partition_idx: + # Same local branch as GLT, except the helper attaches weights + # to the ``sample_one_hop`` result before returning it. + partition_output = sample_local_one_hop_with_edge_weights( + sampler=self.sampler, + edge_weight_lookup=self._edge_weight_lookup, + device=device, + srcs=partition_node_ids, + num_nbr=num_nbr, + edge_type=etype, + ) + partition_results.append( + PartialWeightedNeighborOutput( + index=partition_orders, + output=partition_output, + ) + ) + else: + if self._ppr_weighted_sample_callee_id is None: + raise RuntimeError( + "Weighted PPR sampling RPC callee was not registered." + ) + # Same remote branch as GLT, except the callee id points to a + # weighted callee that returns edge weights with NeighborOutput. + remote_orders_list.append(partition_orders) + to_worker = self.rpc_router.get_to_worker(partition_idx) + futures.append( + rpc_request_async( + to_worker, + self._ppr_weighted_sample_callee_id, + args=(partition_node_ids.cpu(), num_nbr, etype), + ) + ) + + if not remote_orders_list: + if partition_results: + return partition_results[0].output + return empty_weighted_neighbor_output(srcs=srcs, device=device) + + result_futures = await wrap_torch_future(torch.futures.collect_all(futures)) + for i, result_future in enumerate(result_futures): + partition_results.append( + PartialWeightedNeighborOutput( + index=remote_orders_list[i], + output=result_future.wait().to(device), + ) + ) + return self._stitch_weighted_sample_results(srcs, partition_results) async def _batch_fetch_neighbors( self, @@ -450,15 +531,27 @@ async def _batch_fetch_neighbors( sample_tasks = [] for eid in eids: etype = self._etype_id_to_etype[eid] - sample_tasks.append( - self._sample_one_hop( - srcs=nodes_by_etype_id[eid].to(device), - num_nbr=self._num_neighbors_per_hop, - # _sample_one_hop expects None for homogeneous graphs, not the PPR sentinel. - etype=None if etype == _PPR_HOMOGENEOUS_EDGE_TYPE else etype, + sample_etype = None if etype == _PPR_HOMOGENEOUS_EDGE_TYPE else etype + if self._edge_weight_lookup is not None: + sample_tasks.append( + self._sample_one_hop_with_edge_weights( + srcs=nodes_by_etype_id[eid].to(device), + num_nbr=self._num_neighbors_per_hop, + etype=sample_etype, + ) ) - ) - outputs: list[NeighborOutput] = await asyncio.gather(*sample_tasks) + else: + sample_tasks.append( + self._sample_one_hop( + srcs=nodes_by_etype_id[eid].to(device), + num_nbr=self._num_neighbors_per_hop, + # _sample_one_hop expects None for homogeneous graphs, not the PPR sentinel. + etype=sample_etype, + ) + ) + outputs: list[ + Union[NeighborOutput, WeightedNeighborOutput] + ] = await asyncio.gather(*sample_tasks) _empty_weights = torch.empty(0, dtype=torch.float64) @@ -466,13 +559,11 @@ async def _batch_fetch_neighbors( int, tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor] ] = {} for eid, output in zip(eids, outputs): - if self._edge_weights is not None: - assert output.edge is not None, ( - "output.edge must be set when edge_weights is provided; " - "ensure with_edge=True in SamplingConfig (hardcoded in create_sampling_config)." - ) - flat_weights = self._lookup_edge_weights(output.edge, eid) + if self._edge_weight_lookup is not None: + assert isinstance(output, WeightedNeighborOutput) + flat_weights = output.weight.to(torch.float64) else: + assert isinstance(output, NeighborOutput) flat_weights = _empty_weights result[eid] = ( nodes_by_etype_id[eid], diff --git a/gigl/distributed/utils/weighted_neighbor_sampling.py b/gigl/distributed/utils/weighted_neighbor_sampling.py new file mode 100644 index 000000000..491e45e5e --- /dev/null +++ b/gigl/distributed/utils/weighted_neighbor_sampling.py @@ -0,0 +1,233 @@ +"""Helpers for one-hop neighbor sampling that returns edge weights. + +These helpers intentionally mirror GLT's one-hop sampling RPC shape. The only +behavioral difference is that the worker that performs ``sample_one_hop`` also +looks up weights for the returned edge ids and includes them in the response. +""" + +from dataclasses import dataclass +from typing import Optional, cast + +import torch +from graphlearn_torch.distributed.rpc import RpcCalleeBase +from graphlearn_torch.sampler import NeighborSampler +from graphlearn_torch.typing import EdgeType +from graphlearn_torch.utils import ensure_device + + +@dataclass(frozen=True) +class SortedEdgeWeightTable: + """Local edge-weight table keyed by sampled global edge id.""" + + edge_ids: torch.Tensor + weights: torch.Tensor + + @classmethod + def from_graph_topology(cls, graph: object) -> "SortedEdgeWeightTable": + """Build a sorted edge-id table from a GLT Graph's local topology. + + Args: + graph: A ``graphlearn_torch.data.Graph`` instance. + + Returns: + Sorted edge ids and correspondingly permuted weights. + + Raises: + ValueError: If the graph has no edge weights or has malformed topology. + """ + topology = getattr(graph, "topo") + edge_ids = topology.edge_ids + weights = topology.edge_weights + if weights is None: + raise ValueError( + "Weighted sampling requires local graph edge weights, but none were found." + ) + if edge_ids is None: + edge_ids = torch.arange(weights.numel(), dtype=torch.long) + if edge_ids.numel() != weights.numel(): + raise ValueError( + "Local graph edge_ids and edge_weights must have the same length " + f"(got {edge_ids.numel()} ids and {weights.numel()} weights)." + ) + + sorted_edge_ids, order = torch.sort(edge_ids.to(torch.long).cpu()) + sorted_weights = weights.cpu()[order] + return cls(edge_ids=sorted_edge_ids, weights=sorted_weights) + + def lookup(self, sampled_edge_ids: torch.Tensor) -> torch.Tensor: + """Return local weights for sampled global edge ids. + + Args: + sampled_edge_ids: 1-D tensor of sampled edge ids returned by GLT. + + Returns: + A tensor of edge weights parallel to ``sampled_edge_ids``. + + Raises: + KeyError: If a sampled edge id is not present in the local topology. + """ + if sampled_edge_ids.numel() == 0: + return self.weights.new_empty((0,)).to(sampled_edge_ids.device) + if self.edge_ids.numel() == 0: + raise KeyError("Cannot look up sampled edge ids in an empty edge table.") + + query = sampled_edge_ids.to(dtype=torch.long, device=self.edge_ids.device) + positions = torch.searchsorted(self.edge_ids, query) + in_range = positions < self.edge_ids.numel() + matched = torch.zeros_like(in_range, dtype=torch.bool) + if bool(in_range.any()): + matched[in_range] = self.edge_ids[positions[in_range]] == query[in_range] + if not bool(matched.all()): + missing_ids = query[~matched][:10].tolist() + raise KeyError( + f"Sampled edge ids not found in local weighted graph: {missing_ids}" + ) + + return self.weights[positions].to(sampled_edge_ids.device) + + +class LocalEdgeWeightLookup: + """Local worker lookup for edge weights keyed by sampled edge ids.""" + + def __init__(self, local_graph: object): + self._heterogeneous_graphs: Optional[dict[EdgeType, object]] = None + self._heterogeneous_tables: Optional[dict[EdgeType, SortedEdgeWeightTable]] = ( + None + ) + self._homogeneous_table: Optional[SortedEdgeWeightTable] = None + if isinstance(local_graph, dict): + self._heterogeneous_graphs = cast(dict[EdgeType, object], local_graph) + self._heterogeneous_tables = {} + else: + self._homogeneous_table = SortedEdgeWeightTable.from_graph_topology( + local_graph + ) + + def lookup( + self, + sampled_edge_ids: torch.Tensor, + edge_type: Optional[EdgeType], + ) -> torch.Tensor: + """Return local edge weights for sampled edge ids.""" + if self._heterogeneous_tables is not None: + if edge_type is None: + raise ValueError( + "edge_type is required for heterogeneous weight lookup." + ) + if edge_type not in self._heterogeneous_tables: + if self._heterogeneous_graphs is None: + raise RuntimeError( + "heterogeneous graph table is not available for lookup." + ) + self._heterogeneous_tables[edge_type] = ( + SortedEdgeWeightTable.from_graph_topology( + self._heterogeneous_graphs[edge_type] + ) + ) + return self._heterogeneous_tables[edge_type].lookup(sampled_edge_ids) + if self._homogeneous_table is None: + raise RuntimeError("homogeneous graph table is not available for lookup.") + return self._homogeneous_table.lookup(sampled_edge_ids) + + +@dataclass +class WeightedNeighborOutput: + """One-hop sampled neighbors plus sampled edge weights.""" + + nbr: torch.Tensor + nbr_num: torch.Tensor + edge: torch.Tensor + weight: torch.Tensor + + def to(self, device: torch.device) -> "WeightedNeighborOutput": + """Move all output tensors to ``device``.""" + return WeightedNeighborOutput( + nbr=self.nbr.to(device), + nbr_num=self.nbr_num.to(device), + edge=self.edge.to(device), + weight=self.weight.to(device), + ) + + +@dataclass +class PartialWeightedNeighborOutput: + """Weighted one-hop result for a subset of the requested source rows.""" + + index: torch.Tensor + output: WeightedNeighborOutput + + +def empty_weighted_neighbor_output( + srcs: torch.Tensor, + device: torch.device, +) -> WeightedNeighborOutput: + """Create an empty weighted one-hop output for ``srcs``.""" + return WeightedNeighborOutput( + nbr=torch.empty(0, dtype=torch.long, device=device), + nbr_num=torch.zeros(srcs.size(0), dtype=torch.long, device=device), + edge=torch.empty(0, dtype=torch.long, device=device), + weight=torch.empty(0, dtype=torch.float32, device=device), + ) + + +def sample_local_one_hop_with_edge_weights( + sampler: NeighborSampler, + edge_weight_lookup: LocalEdgeWeightLookup, + device: torch.device, + srcs: torch.Tensor, + num_nbr: int, + edge_type: Optional[EdgeType], +) -> WeightedNeighborOutput: + """Sample one hop locally and attach local edge weights in the same call. + + This is the local analogue of GLT's direct ``sampler.sample_one_hop`` call. + The only added step is ``edge_weight_lookup.lookup(output.edge, edge_type)``; + neighbor ids, neighbor counts, and sampled edge ids are otherwise returned + exactly as GLT produced them. + """ + output = sampler.sample_one_hop(srcs, num_nbr, edge_type) + if output is None: + return empty_weighted_neighbor_output(srcs=srcs, device=device) + if output.edge is None: + raise RuntimeError("Weighted sampling requires GLT to return edge ids.") + weights = edge_weight_lookup.lookup(output.edge, edge_type) + return WeightedNeighborOutput( + nbr=output.nbr, + nbr_num=output.nbr_num, + edge=output.edge, + weight=weights, + ) + + +class WeightedSamplingCallee(RpcCalleeBase): + """RPC callee that samples one hop and returns edge weights with the result. + + This mirrors GLT's ``RpcSamplingCallee``. The routing, device handling, and + remote ``sample_one_hop`` call are the same; the delta is that the callee + converts ``NeighborOutput`` into ``WeightedNeighborOutput`` by attaching + local edge weights before moving the result back to CPU for the RPC reply. + """ + + def __init__( + self, + sampler: NeighborSampler, + edge_weight_lookup: LocalEdgeWeightLookup, + device: torch.device, + ) -> None: + super().__init__() + self.sampler = sampler + self.edge_weight_lookup = edge_weight_lookup + self.device = device + + def call(self, *args, **kwargs) -> WeightedNeighborOutput: + """Sample one hop on this worker and attach local edge weights.""" + ensure_device(self.device) + output = sample_local_one_hop_with_edge_weights( + sampler=self.sampler, + edge_weight_lookup=self.edge_weight_lookup, + device=self.device, + srcs=args[0].to(self.device), + num_nbr=args[1], + edge_type=args[2], + ) + return output.to(torch.device("cpu")) diff --git a/tests/unit/distributed/distributed_ppr_weighted_sampling_test.py b/tests/unit/distributed/distributed_ppr_weighted_sampling_test.py index 11b74c774..56b52f8b2 100644 --- a/tests/unit/distributed/distributed_ppr_weighted_sampling_test.py +++ b/tests/unit/distributed/distributed_ppr_weighted_sampling_test.py @@ -11,6 +11,8 @@ every top-k result. """ +from dataclasses import replace + import torch import torch.multiprocessing as mp from absl.testing import absltest @@ -18,9 +20,9 @@ from torch_geometric.data import Data, HeteroData from gigl.distributed.dist_dataset import DistDataset -from gigl.distributed.dist_ppr_sampler import DistPPRNeighborSampler from gigl.distributed.distributed_neighborloader import DistNeighborLoader from gigl.distributed.sampler_options import PPRSamplerOptions +from gigl.types.graph import GraphPartitionData from tests.test_assets.distributed.bipartite_weight_graph import ( ITEM, USER, @@ -166,21 +168,40 @@ def test_ppr_weighted_never_traverses_zero_weight_edges_homogeneous(self) -> Non nprocs=1, ) - def test_ppr_weight_lookup_uses_edge_ids_not_tensor_offsets(self) -> None: - """Sampled edge IDs may be non-contiguous and must resolve by ID.""" - lookup = DistPPRNeighborSampler._build_sorted_edge_weight_lookup( - edge_ids=torch.tensor([42, 7, 100]), - edge_weights=torch.tensor([4.2, 0.7, 10.0]), - ) + def test_ppr_weighted_handles_non_contiguous_edge_ids_homogeneous(self) -> None: + """Homogeneous: sampled global edge ids resolve to local edge weights. - actual = DistPPRNeighborSampler._lookup_edge_weights_by_id( - edge_ids=torch.tensor([7, 100, 42]), - edge_weight_lookup=lookup, + Partitioned weighted graphs store local weights compactly, while GLT + sampling returns the graph edge ids. This test gives every local edge a + large non-contiguous edge id; indexing the local weight tensor directly + with sampled edge ids would go out of bounds. + """ + partition_output, n_hub = build_homogeneous_bipartite_weight_graph() + partitioned_edge_index = partition_output.partitioned_edge_index + self.assertIsInstance(partitioned_edge_index, GraphPartitionData) + assert isinstance(partitioned_edge_index, GraphPartitionData) + num_edges = partitioned_edge_index.edge_index.size(1) + partition_output = replace( + partition_output, + partitioned_edge_index=replace( + partitioned_edge_index, + edge_ids=torch.arange( + 10_000, + 10_000 + num_edges, + dtype=torch.long, + ), + ), ) - torch.testing.assert_close( - actual, - torch.tensor([0.7, 10.0, 4.2], dtype=torch.float64), + dataset = DistDataset(rank=0, world_size=1, edge_dir="out") + dataset.build(partition_output=partition_output) + + self.assertTrue(dataset.has_edge_weights) + + mp.spawn( + fn=_run_ppr_weighted_correctness_homogeneous, + args=(dataset, n_hub), + nprocs=1, ) def test_ppr_weighted_never_traverses_zero_weight_edges_heterogeneous(self) -> None: diff --git a/tests/unit/distributed/utils/degree_test.py b/tests/unit/distributed/utils/degree_test.py index f61210f5e..ffcb6e5a4 100644 --- a/tests/unit/distributed/utils/degree_test.py +++ b/tests/unit/distributed/utils/degree_test.py @@ -119,10 +119,10 @@ def test_heterogeneous_graph_with_missing_topology(self): self.assertEqual(set(result.keys()), set(edge_types)) # Edge type with topology should have computed degrees - self.assert_tensor_equality(result[edge_type_with_topo], expected_degrees) # ty: ignore[invalid-argument-type] TODO(ty-torch-keyed-access): fix ty false positives for torch-backed keyed container access. + self.assert_tensor_equality(result[edge_type_with_topo], expected_degrees) # Edge type without topology should have empty tensor - self.assertEqual(result[edge_type_without_topo].numel(), 0) # ty: ignore[invalid-argument-type] TODO(ty-torch-keyed-access): fix ty false positives for torch-backed keyed container access. + self.assertEqual(result[edge_type_without_topo].numel(), 0) def _run_local_world_size_correction_homogeneous(