Move existing PPR sampler work off event loop by mkolodner-sc · Pull Request #708 · Snapchat/GiGL

mkolodner-sc · 2026-07-15T18:54:47Z

Summary

Moves the existing distributed PPR sampler's CPU-heavy C++ work off the asyncio event-loop path.

This PR is limited to the existing/non-typed PPR path. Typed-channel PPR enablement is left for the next PR in the stack.

Stack

Base: Clean up PPR sampler naming #707
Next: Add typed-channel PPR sampler #709

Reasoning for Change

PPR sampling runs several CPU-heavy C++ steps while async graph-store/RPC work is also in flight. When those steps run on the event-loop thread or hold the GIL longer than necessary, they can delay RPC completion callbacks and reduce concurrency across sampler coroutines.

This change keeps the PPR math and returned sequences unchanged while moving construction, queue draining, residual pushes, and extraction through executor/GIL-release boundaries where possible.

Changes

Releases the GIL around PPRForwardPush construction, queue draining, residual push, and extraction in the pybind layer.
Runs PPR state construction, drain_queue, push_residuals, and top-k extraction through run_in_executor from the Python sampler.
Removes a redundant frontier .to(device) before one-hop fetches; extracted tensors are still moved to the sampler device at the existing output boundary.

kmontemayor2-sc

Thanks Matt!

kmontemayor2-sc · 2026-07-15T22:28:42Z

+static std::optional<std::unordered_map<int32_t, torch::Tensor>> drainQueueWrapper(PPRForwardPush& state) {
+    std::optional<std::unordered_map<int32_t, torch::Tensor>> drained;
+    // The drain mutates only this PPRForwardPush instance and returns torch tensors.
+    // No Python objects are touched until pybind converts the result after return.


btw do we know if pybind tensor conversions are cheap? Something else to look into? If we're copying the underlying memory that could be very expensive.

The pybind conversion for torch::Tensor should not copy tensor storage. It creates Python tensor container objects around the existing torch::Tensor handles. The data copy here is earlier, inside drainQueue(), where we materialize the frontier node IDs from C++ sets into CPU tensors. I added comments to make that distinction explicit. In practice, we've observed the memory of these CPU tensors to not be significant on internal benchmarking.

kmontemayor2-sc · 2026-07-15T22:29:40Z

        .def(py::init<torch::Tensor,
                      int32_t,
                      double,
                      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>>(),
+             py::call_guard<py::gil_scoped_release>())


qq does init need the gil release? Or we we just doing this for completion?

It can still be useful to guarantee that the initialization of a PPRState object on a fresh batch isn't blocking the event loop for a different concurrent worker.

is the creation of the object expensive? It's just setting up some empty collections right?

I don't think this is blocking but I am a bit curious if we suspect this is a crucial fix

kmontemayor2-sc · 2026-07-15T22:31:42Z

            edge_type_ids.append(edge_type_id)
            sample_tasks.append(
                self._sample_one_hop(
-                    srcs=nodes_by_edge_type_id[edge_type_id].to(device),


Why do we remove this? Is it just not needed?

Yes it is not needed. These come directly from drainQueue(), which currently creates CPU torch::kLong tensors. _sample_one_hop can consume those CPU ID tensors directly, so the old .to(device) was an avoidable CPU→sampler-device transfer before the fetch.

oh good catch, we do the to(device) in collate_fn right?

kmontemayor2-sc · 2026-07-15T22:38:21Z

+            nodes_by_edge_type_id = await loop.run_in_executor(
+                None, ppr_state.drain_queue
+            )


Sorry for my own understanding, why do we need to run this in the executor? Shouldn't just using the gil release be sufficient here? Have we tested this without run_in_executor but with GIL release? (and also preferable run_in_executor without GIL release?

I guess I'm confused how calling this method and then immediately awaiting gets us a speedup (vs just the gil release?)

run_in_executor and GIL release help at different layers. run_in_executor keeps the asyncio event-loop thread from sitting inside a synchronous C++ call, so other RPC completions/coroutines can be scheduled while PPR work runs in a worker thread. py::gil_scoped_release makes sure that worker thread does not hold the Python GIL while it is doing C++ work, which lets the event-loop thread actually execute Python callbacks/coroutines while the worker is busy.

For example, if push_residuals takes 50ms and we call it directly, the event loop is stuck inside C++ for that full 50ms, even if the GIL is released. If we only use run_in_executor, the event-loop thread is free, but the worker may still hold the GIL and delay Python callbacks. With both, the 50ms C++ push runs in a worker thread without holding the GIL, so the event loop can keep processing RPC completions and scheduling the next PPR work.

hmmm, so the await loop.run_in_executor blocks this thread, but we may have other threads (in the same python process) that would get freed up?

I suspect this pattern may be useful elsewhere in graph store but I'll need to think about that :)

kmontemayor2-sc · 2026-07-15T22:41:11Z

        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()
+            nodes_by_edge_type_id = await loop.run_in_executor(
+                None, ppr_state.drain_queue
+            )
            if nodes_by_edge_type_id is None:
                break

            fetch_budget_remaining = (
                self._max_fetch_iterations is None
                or fetch_iteration_count < self._max_fetch_iterations
            )
            if nodes_by_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,
        )


Something else to look into (not required now) is that it's possible that all of this going into and out of the python / c++ boundaries is expensive. Have we considered the following schema?

Create some new compute_ppr_scores cpp loop that mirrors the logic here

Since we have this logic because of _batch_fetch_neighbors (or specifically sample_one_hop right?) We can pass some object into cpp and we just call sample_one_hop in the cpp code.

it's possible that all of this going into and out of the python / c++ boundaries is expensive

I've done some internal benchmarking on this serialization cost explicitly and the latency is not significant from this. Agreed that this can be investigated in the future though if it becomes an issue -- I've added a TODO here

kmontemayor2-sc

LGTM! Nice work Matt :)

Do you think there's similar low-hanging (but hard to find!) fruit here for the khop sampler?

E.g. is it possible to put this bit in the background? https://github.com/alibaba/graphlearn-for-pytorch/blob/main/graphlearn_torch/python/distributed/dist_neighbor_sampler.py#L297

Or do we need to also write a CPP version of that _sample_from_nodes so we can force-detach the gil? May not be difficult? WDYT?

kmontemayor2-sc · 2026-07-15T23:38:15Z

            edge_type_ids.append(edge_type_id)
            sample_tasks.append(
                self._sample_one_hop(
-                    srcs=nodes_by_edge_type_id[edge_type_id].to(device),


oh good catch, we do the to(device) in collate_fn right?

kmontemayor2-sc · 2026-07-15T23:42:36Z

+            nodes_by_edge_type_id = await loop.run_in_executor(
+                None, ppr_state.drain_queue
+            )


hmmm, so the await loop.run_in_executor blocks this thread, but we may have other threads (in the same python process) that would get freed up?

I suspect this pattern may be useful elsewhere in graph store but I'll need to think about that :)

mkolodner-sc requested review from kmontemayor2-sc, nshah-sc, svij-sc, xgao4-sc, yliu2-sc and zfan3-sc as code owners July 15, 2026 18:54

This was referenced Jul 15, 2026

Add typed-channel PPR sampler #709

Open

Clean up PPR sampler naming #707

Merged

mkolodner-sc force-pushed the mkolodner-sc/pr701-ppr-cleanup branch from 5197b5d to 7f6d9c0 Compare July 15, 2026 20:02

mkolodner-sc force-pushed the mkolodner-sc/pr701-ppr-event-loop branch from fc410ab to 2b85456 Compare July 15, 2026 20:02

mkolodner-sc force-pushed the mkolodner-sc/pr701-ppr-cleanup branch from 7f6d9c0 to 19c0d92 Compare July 15, 2026 21:24

mkolodner-sc force-pushed the mkolodner-sc/pr701-ppr-event-loop branch from 2b85456 to ecb394c Compare July 15, 2026 21:24

mkolodner-sc force-pushed the mkolodner-sc/pr701-ppr-cleanup branch from 19c0d92 to f7aec06 Compare July 15, 2026 21:40

mkolodner-sc force-pushed the mkolodner-sc/pr701-ppr-event-loop branch from ecb394c to 0047b2f Compare July 15, 2026 21:40

kmontemayor2-sc reviewed Jul 15, 2026

View reviewed changes

mkolodner-sc force-pushed the mkolodner-sc/pr701-ppr-event-loop branch 3 times, most recently from 60a0cea to ff2d3ea Compare July 15, 2026 23:04

kmontemayor2-sc approved these changes Jul 15, 2026

View reviewed changes

mkolodner-sc force-pushed the mkolodner-sc/pr701-ppr-event-loop branch from ff2d3ea to d0f9cea Compare July 16, 2026 18:46

mkolodner-sc force-pushed the mkolodner-sc/pr701-ppr-cleanup branch 2 times, most recently from 2dd7085 to 68d6d62 Compare July 16, 2026 19:15

mkolodner-sc force-pushed the mkolodner-sc/pr701-ppr-event-loop branch 2 times, most recently from 33c8dde to 1b0ea2a Compare July 16, 2026 19:21

mkolodner-sc force-pushed the mkolodner-sc/pr701-ppr-cleanup branch from 68d6d62 to 3d34854 Compare July 16, 2026 19:21

mkolodner-sc added 3 commits July 16, 2026 19:57

Clean up PPR sampler naming

f5c430d

Move existing PPR work off event loop

50da876

Avoid redundant PPR frontier device copy

ed7387c

mkolodner-sc force-pushed the mkolodner-sc/pr701-ppr-cleanup branch from 3d34854 to f5c430d Compare July 16, 2026 20:01

mkolodner-sc force-pushed the mkolodner-sc/pr701-ppr-event-loop branch from 1b0ea2a to ed7387c Compare July 16, 2026 20:01

Base automatically changed from mkolodner-sc/pr701-ppr-cleanup to main July 17, 2026 09:13

Uh oh!

Conversation

mkolodner-sc commented Jul 15, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Stack

Reasoning for Change

Changes

Uh oh!

kmontemayor2-sc left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

kmontemayor2-sc left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

mkolodner-sc commented Jul 15, 2026 •

edited

Loading