Move existing PPR sampler work off event loop

gigl-core/core/sampling/python_ppr_forward_push.cpp

@@ -8,6 +8,7 @@
 #include <torch/extension.h>
 
 #include <cstdint>
+#include <optional>
 #include <tuple>
 #include <unordered_map>
 
@@ -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<int32_t, std::tuple<torch::Tensor, torch::Tensor, torch::Tensor>> 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<std::unordered_map<int32_t, torch::Tensor>> drainQueueWrapper(PPRForwardPush& state) {
+    std::optional<std::unordered_map<int32_t, torch::Tensor>> 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<int32_t, std::tuple<torch::Tensor, torch::Tensor, torch::Tensor>>
+extractTopKWithResidualTopUpWrapper(PPRForwardPush& state, int32_t maxPPRNodes, bool enableResidualTopUp) {
+    std::unordered_map<int32_t, std::tuple<torch::Tensor, torch::Tensor, torch::Tensor>> 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<int32_t>>,
                       std::vector<int32_t>,
-                      std::vector<torch::Tensor>>())
-        .def("drain_queue", &gigl::PPRForwardPush::drainQueue)
+                      std::vector<torch::Tensor>>(),
+             // Constructor argument conversion happens before the C++ body; the
+             // body only initializes PPR state and can run without the GIL.
+             py::call_guard<py::gil_scoped_release>())
+        .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"));
 }

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,25 +466,27 @@ 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:
                 # Fetch budget exhausted; push_residuals will use the existing neighbor cache.
                 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(