diff --git a/gigl-core/core/sampling/python_ppr_forward_push.cpp b/gigl-core/core/sampling/python_ppr_forward_push.cpp index 27271559c..b92b21418 100644 --- a/gigl-core/core/sampling/python_ppr_forward_push.cpp +++ b/gigl-core/core/sampling/python_ppr_forward_push.cpp @@ -8,6 +8,7 @@ #include #include +#include #include #include @@ -17,9 +18,8 @@ namespace py = pybind11; namespace gigl { -// pushResiduals: a wrapper is needed solely to release the GIL during the C++ push. -// pybind11/stl.h handles all type conversions automatically; the other methods use -// direct member function pointers for the same reason. +// pushResiduals receives Python-owned containers, so convert them while the GIL +// is held and release only around the C++ state update. static void pushResidualsWrapper(PPRForwardPush& state, const py::dict& fetchedByEtypeId) { std::unordered_map> neighborTensorsByEtypeId; // Dict iteration touches Python objects — GIL must be held here. @@ -42,6 +42,30 @@ static void pushResidualsWrapper(PPRForwardPush& state, const py::dict& fetchedB } } +static std::optional> drainQueueWrapper(PPRForwardPush& state) { + std::optional> drained; + // drainQueue mutates only this PPRForwardPush instance and materializes CPU + // tensors for frontier node IDs. pybind converts those tensor handles back + // to Python tensors after return without copying the underlying storage. + { + py::gil_scoped_release release; + drained = state.drainQueue(); + } + return drained; +} + +static std::unordered_map> +extractTopKWithResidualTopUpWrapper(PPRForwardPush& state, int32_t maxPPRNodes, bool enableResidualTopUp) { + std::unordered_map> result; + // Extraction walks C++ state and builds torch tensors. Returning through + // pybind creates Python container/wrapper objects, not tensor data copies. + { + py::gil_scoped_release release; + result = state.extractTopKWithResidualTopUp(maxPPRNodes, enableResidualTopUp); + } + return result; +} + } // namespace gigl // TORCH_EXTENSION_NAME is set by PyTorch's build system to match the Python @@ -54,11 +78,14 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { double, std::vector>, std::vector, - std::vector>()) - .def("drain_queue", &gigl::PPRForwardPush::drainQueue) + std::vector>(), + // Constructor argument conversion happens before the C++ body; the + // body only initializes PPR state and can run without the GIL. + py::call_guard()) + .def("drain_queue", gigl::drainQueueWrapper) .def("push_residuals", gigl::pushResidualsWrapper) .def("extract_top_k_with_residual_top_up", - &gigl::PPRForwardPush::extractTopKWithResidualTopUp, + gigl::extractTopKWithResidualTopUpWrapper, py::arg("max_ppr_nodes"), py::arg("enable_residual_topup")); } diff --git a/gigl/distributed/dist_ppr_sampler.py b/gigl/distributed/dist_ppr_sampler.py index 1ba78c9fc..78de00436 100644 --- a/gigl/distributed/dist_ppr_sampler.py +++ b/gigl/distributed/dist_ppr_sampler.py @@ -237,7 +237,6 @@ def _get_destination_type(self, edge_type: EdgeType) -> NodeType: async def _batch_fetch_neighbors( self, nodes_by_edge_type_id: dict[int, torch.Tensor], - device: torch.device, ) -> dict[int, tuple[torch.Tensor, torch.Tensor, torch.Tensor]]: """Batch fetch neighbors for nodes grouped by integer edge type ID. @@ -247,8 +246,7 @@ async def _batch_fetch_neighbors( Args: nodes_by_edge_type_id: Dict mapping integer edge type ID to a 1-D int64 tensor of node IDs to fetch neighbors for. Comes directly from - ``drain_queue()``; node IDs are already deduplicated. - device: Torch device for intermediate tensor creation. + ``drain_queue()`` as CPU tensors; node IDs are already deduplicated. Returns: Dict mapping edge type ID to ``(node_ids, flat_neighbors, counts)`` @@ -283,9 +281,11 @@ async def _batch_fetch_neighbors( else edge_type ) edge_type_ids.append(edge_type_id) + # drain_queue materializes CPU frontier tensors; _sample_one_hop can + # consume them directly, so avoid a sampler-device round trip here. sample_tasks.append( self._sample_one_hop( - srcs=nodes_by_edge_type_id[edge_type_id].to(device), + srcs=nodes_by_edge_type_id[edge_type_id], num_nbr=self._num_neighbors_per_hop, etype=rpc_edge_type, ) @@ -427,8 +427,11 @@ async def _compute_ppr_scores( if seed_node_type is None: seed_node_type = DEFAULT_HOMOGENEOUS_NODE_TYPE device = seed_nodes.device + loop = asyncio.get_running_loop() - ppr_state = PPRForwardPush( + ppr_state = await loop.run_in_executor( + None, + PPRForwardPush, seed_nodes, self._node_type_to_id[seed_node_type], self._alpha, @@ -440,13 +443,20 @@ async def _compute_ppr_scores( fetch_iteration_count = 0 + # TODO: If Python/C++ boundary overhead in this loop becomes a blocker, + # consider a coarser C++ PPR iteration API while keeping Python + # responsible for async distributed neighbor fetches. while True: # drain_queue returns None when the queue is truly empty (convergence), # or a dict (possibly empty) when nodes were drained. An empty dict # means all drained nodes either had cached neighbors or no outgoing # edges — we still call push_residuals to flush their residuals into # ppr_scores_. - nodes_by_edge_type_id = ppr_state.drain_queue() + # The pybind wrapper releases the GIL, but a direct call would still + # occupy this coroutine's event-loop thread until C++ returns. + nodes_by_edge_type_id = await loop.run_in_executor( + None, ppr_state.drain_queue + ) if nodes_by_edge_type_id is None: break @@ -456,7 +466,7 @@ async def _compute_ppr_scores( ) if nodes_by_edge_type_id and fetch_budget_remaining: fetched_by_edge_type_id = await self._batch_fetch_neighbors( - nodes_by_edge_type_id, device + nodes_by_edge_type_id ) fetch_iteration_count += 1 else: @@ -464,17 +474,19 @@ async def _compute_ppr_scores( fetched_by_edge_type_id = {} # Run in executor so the C++ push doesn't block the asyncio event loop. - await asyncio.get_running_loop().run_in_executor( + await loop.run_in_executor( None, ppr_state.push_residuals, fetched_by_edge_type_id ) # Regular sampling uses residual top-up as fill-only under the # max_ppr_nodes cap. If finalized PPR scores already fill that cap, # there is no remaining budget for residual candidates. - return self._extract_ppr_state_top_k( + return await loop.run_in_executor( + None, + self._extract_ppr_state_top_k, ppr_state, device, - max_ppr_nodes=self._max_ppr_nodes, + self._max_ppr_nodes, ) async def _sample_from_nodes(