Scope collective ID pool per query by Jesse-Bakker · Pull Request #24547 · NVIDIA/cudf

Jesse-Bakker · 2026-10-12T08:56:39Z

Description

The pool of vacant collective IDs was process-global. The global pool's
state depends on every query the process has run or is running. So, a
process that first runs a query as part of one cluster and then as part
of another can allocate different collective IDs.
I see two possible solutions:

Return the collective IDs to the global pool in reverse order. This
will make sequential queries work (assuming they always return all their
IDs properly through the context manager)
Scope the collective IDs per query through a ContextVar, making their
ID pools independent

I went for the latter approach, as it seems more robust.

Checklist

I am familiar with the Contributing Guidelines.
New or existing tests cover these changes.
The documentation is up to date with these changes.

copy-pr-bot · 2026-10-12T08:56:43Z

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

coderabbitai · 2026-10-12T09:09:38Z

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Repository: NVIDIA/cudf/.coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: fad455da-e5d4-46f9-b7d9-268d9d1a61ea

📥 Commits

Reviewing files that changed from the base of the PR and between 45b76fc and 94dc15c.

📒 Files selected for processing: 3

docs/cudf/source/cudf_polars/api.md
python/cudf_polars/cudf_polars/engine/core.py
python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/common.py

Included review availability: This review used your included allowance. Your plan provides up to 12 included reviews per hour; 11 remain after this review.

📝 Summary

Summary by CodeRabbit

Bug Fixes
- Queries running concurrently now receive isolated collective operation IDs, preventing reservations from one query or thread from affecting another.
- ID allocation remains consistent when query scopes are nested or restored, preserving the expected sequence of IDs.
Documentation
- Added the collective ID scope helper to the SPMD reference documentation.

Collective ID allocation now uses context-local pools instead of a process-global pool. collective_id_scope() creates and restores a pool. evaluate_on_rank runs its evaluation within that scope. Tests cover nested scopes, separate threads, and rank evaluation.

Changes: Query-scoped collective IDs

Layer / File(s)	Summary
Context-local pool and reservation behavior `python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/*`, `docs/cudf/source/cudf_polars/api.md`	Adds context-local pool selection and `collective_id_scope()`. `ReserveOpIDs` and `reserve_op_id()` acquire and release IDs through the selected pool. The package exports the scope, and the API documentation lists it.
Rank evaluation scope and validation `python/cudf_polars/cudf_polars/engine/core.py`, `python/cudf_polars/tests/streaming/test_collective_ids.py`	`evaluate_on_rank` enters a fresh scope and delegates the existing evaluation pipeline to `_evaluate_on_rank`. Tests cover independent, nested, and threaded scopes, and the evaluation entry point.

Priority: ➖ Normal

Estimated code review effort: 3 (Moderate) | ~20 minutes

Change: Bug fix

Suggested reviewers: wence-, rjzamora

Merge Risk: 🟡 Moderate · up to 94dc1

Concurrent queries on the same SPMD engine may interfere with one another’s collective operations. Resolve or explicitly accept that risk before merging.

Pre-merge checks: 4 · 1

Failed checks: 1 warning

Check name	Status	Explanation	Resolution
Docstring Coverage		Docstring coverage is 36.84% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 19 functions across 4 files. (1 skipped: …	Write docstrings for the functions missing them to satisfy the coverage threshold.

Passed checks: 4 passed

Check name	Status	Explanation
Title check		The title clearly and concisely describes the main change: scoping the collective ID pool per query.
Description check		The description explains the process-global pool problem, the ContextVar-based solution, and the related tests and documentation updates.
Linked Issues check		Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check		Check skipped because no linked issues were found for this pull request.

Full details: Docstring Coverage

Explanation

Docstring coverage is 36.84% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 19 functions across 4 files. (1 skipped: 1 unsupported.)

Fix all pre-merge checks with AI

✨ Finishing Touches

🧪 Generate unit tests (beta)

Create a new PR

Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts

_{Comment @coderabbitai help to get the list of available commands.}

wence- · 2026-10-12T13:19:46Z

/ok to test fe1d521

wence-

Please run pre-commit to fix the minor CI style errors.

Also a question about the safety of using the same communicator with two concurrent queries.

wence- · 2026-10-12T15:44:03Z

+    # Each query reserves its collective IDs from its own pool, so that the
+    # IDs agree across ranks even if this process concurrently takes part in
+    # other queries (possibly on other communicators).
+    with collective_id_scope():
+        return _evaluate_on_rank(
+            ctx,
+            comm,
+            py_executor,
+            ir,
+            config_options,
+            quent_query_worker_state=quent_query_worker_state,
+            query_id=query_id,
+        )


question: What would happen if you ran two queries on the same communicator concurrently with this setup? Because each query generates its own id scope, I think that you would get overlapping ids, and message passing could get mixed up.

Should we key the id scope by the communicator?

Running concurrent queries on the same communicator would probably fail miserably, whereas before this change, it would probably work as long as the queries are started in the same order on all ranks. Keying the scope by communicator would work for that, but again, only if the queries are started in the same order on all ranks. It should be possible for a caller to sequence that using a queue and a barrier though.

Hmm, no, it also depends on release timing and each rank can time that slightly differently, so that can still race and cause mismatched collective ids I think.

The pool of vacant collective IDs was process-global. The global pool's state depends on every query the process has run or is running. So, a process that first runs a query as part of one cluster and then as part of another can allocate different collective IDs. I see two possible solutions: - Return the collective IDs to the global pool in reverse order. This will make sequential queries work (assuming they always return all their IDs properly through the context manager) - Scope the collective IDs per query through a ContextVar, making their ID pools independent I went for the latter approach, as it seems more robust.

TomAugspurger · 2026-10-12T18:33:33Z

/ok to test 45b76fc

TomAugspurger · 2026-10-14T13:56:59Z

/ok to test 94dc15c

github-actions Bot assigned Jesse-Bakker Oct 12, 2026

github-actions Bot added Python Affects Python cuDF API. cudf-polars Issues specific to cudf-polars labels Oct 12, 2026

github-project-automation Bot added this to cuDF Python Oct 12, 2026

Jesse-Bakker force-pushed the collective-id-contextvar branch from afdf8e9 to fe1d521 Compare October 12, 2026 09:05

Jesse-Bakker marked this pull request as ready for review October 12, 2026 09:05

Jesse-Bakker requested a review from a team as a code owner October 12, 2026 09:05

Jesse-Bakker requested a review from TomAugspurger October 12, 2026 09:05

wence- added improvement Improvement / enhancement to an existing function non-breaking Non-breaking change labels Oct 12, 2026

wence- reviewed Oct 12, 2026

View reviewed changes

Jesse-Bakker added 2 commits October 12, 2026 18:54

Add tests for per-query collective ID pools

45b76fc

Jesse-Bakker force-pushed the collective-id-contextvar branch from fe1d521 to 45b76fc Compare October 12, 2026 16:54

TomAugspurger reviewed Oct 13, 2026

View reviewed changes

Comment thread python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/common.py

doc update

a50ef5f

github-actions Bot assigned TomAugspurger Oct 14, 2026

Merge branch 'main' into collective-id-contextvar

94dc15c

Navigation Menu

Repository navigation

Scope collective ID pool per query - #24547

Scope collective ID pool per query#24547
Jesse-Bakker wants to merge 4 commits into
NVIDIA:mainfrom
Jesse-Bakker:collective-id-contextvar

Jesse-Bakker commented Oct 12, 2026

Uh oh!

copy-pr-bot Bot commented Oct 12, 2026

Uh oh!

coderabbitai Bot commented Oct 12, 2026 •

edited

Loading

Summary by CodeRabbit

Uh oh!

wence- commented Oct 12, 2026

Uh oh!

wence- left a comment

Uh oh!

wence- Oct 12, 2026

Uh oh!

Jesse-Bakker Oct 12, 2026

Uh oh!

Jesse-Bakker Oct 12, 2026

Uh oh!

TomAugspurger commented Oct 12, 2026

Uh oh!

Uh oh!

TomAugspurger commented Oct 14, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Conversation

Jesse-Bakker commented Oct 12, 2026

Description

Checklist

Uh oh!

copy-pr-bot Bot commented Oct 12, 2026

Uh oh!

coderabbitai Bot commented Oct 12, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary by CodeRabbit

Uh oh!

wence- commented Oct 12, 2026

Uh oh!

wence- left a comment

Choose a reason for hiding this comment

Uh oh!

wence- Oct 12, 2026

Choose a reason for hiding this comment

Uh oh!

Jesse-Bakker Oct 12, 2026

Choose a reason for hiding this comment

Uh oh!

Jesse-Bakker Oct 12, 2026

Choose a reason for hiding this comment

Uh oh!

TomAugspurger commented Oct 12, 2026

Uh oh!

Uh oh!

TomAugspurger commented Oct 14, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

coderabbitai Bot commented Oct 12, 2026 •

edited

Loading