Scope collective ID pool per query

docs/cudf/source/cudf_polars/api.md

@@ -67,6 +67,8 @@ Returned by `engine.execute()` to keep query results GPU-resident (see {doc}`exe
 .. autofunction:: cudf_polars.engine.spmd.use_gpu
 
 .. autofunction:: cudf_polars.streaming.actor_graph.collectives.common.reserve_op_id
+
+.. autofunction:: cudf_polars.streaming.actor_graph.collectives.common.collective_id_scope
 ```
 
 ## PyTorch interoperability

python/cudf_polars/cudf_polars/engine/core.py

@@ -37,7 +37,10 @@
     prefetch_parquet_file_metadata_for_ir,
 )
 from cudf_polars.quent._plan import build_quent_operator_map, emit_plan
-from cudf_polars.streaming.actor_graph.collectives import ReserveOpIDs
+from cudf_polars.streaming.actor_graph.collectives import (
+    ReserveOpIDs,
+    collective_id_scope,
+)
 from cudf_polars.streaming.actor_graph.collectives.common import reserve_op_id
 from cudf_polars.streaming.actor_graph.core import generate_network
 from cudf_polars.streaming.actor_graph.utils import empty_table_chunk
@@ -908,6 +911,31 @@ def evaluate_on_rank(
     metadata
         Collected channel metadata.
     """
+    # 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,
+        )
+
+
+def _evaluate_on_rank(
+    ctx: Context,
+    comm: Communicator,
+    py_executor: ThreadPoolExecutor,
+    ir: IR,
+    config_options: ConfigOptions[StreamingExecutor],
+    *,
+    quent_query_worker_state: QuentQueryWorkerState | None,
+    query_id: uuid.UUID,
+) -> tuple[DataFrame, list[ChannelMetadata]]:
     stats = allgather_stats(comm, ctx.br(), ir, config_options, py_executor)
     # ``get_stable_plan_id`` is a deterministic function of the IR
     # structure, so every rank derives the same logical plan ID for a

python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/__init__.py

@@ -1,12 +1,13 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES.
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 """Collective operations for building a RapidsMPF actor graph."""
 
 from __future__ import annotations
 
 from cudf_polars.streaming.actor_graph.collectives.common import (
     ReserveOpIDs,
+    collective_id_scope,
     reserve_op_id,
 )
 
-__all__ = ["ReserveOpIDs", "reserve_op_id"]
+__all__ = ["ReserveOpIDs", "collective_id_scope", "reserve_op_id"]

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

@@ -6,6 +6,7 @@
 
 import threading
 from contextlib import contextmanager
+from contextvars import ContextVar
 from typing import TYPE_CHECKING, Literal
 
 from rapidsmpf.shuffler import Shuffler
@@ -34,33 +35,70 @@
     from cudf_polars.utils.config import ConfigOptions, StreamingExecutor
 
 
-# Set of available collective IDs
-_collective_id_vacancy: dict[int, None] = dict.fromkeys(
-    range(Shuffler.max_concurrent_shuffles)
+class _CollectiveIDPool:
+    """Pool of vacant collective IDs."""
+
+    def __init__(self) -> None:
+        # Since 3.7, dict.popitem() is guaranteed LIFO.
+        self.vacancy: dict[int, None] = dict.fromkeys(
+            range(Shuffler.max_concurrent_shuffles)
+        )
+        self.lock = threading.Lock()
+
+    def acquire_unsafe(self) -> int:
+        # Not thread safe, must be called with self.lock held
+        try:
+            # All ranks must choose the same collective IDs during lowering.
+            collective_id, _ = self.vacancy.popitem()
+        except KeyError:
+            raise ValueError(
+                f"Cannot shuffle more than {Shuffler.max_concurrent_shuffles} "
+                "times in a single query."
+            ) from None
+        return collective_id
+
+    def release_unsafe(self, collective_id: int) -> None:
+        """Release a collective ID back to the vacancy set."""
+        # Not thread safe, must be called with self.lock held
+        if collective_id in self.vacancy:
+            raise ValueError("Restoring ID that exists")
+        self.vacancy[collective_id] = None
+
+
+# Collective IDs are only meaningful relative to a communicator, and all ranks
+# must agree on the IDs chosen. A process-global pool breaks that agreement as
+# soon as the same process takes part in queries on different communicators,
+# so the pool is scoped to the current context (see ``collective_id_scope``).
+_collective_id_pool: ContextVar[_CollectiveIDPool] = ContextVar(
+    "cudf_polars_collective_id_pool"
 )
-_collective_id_vacancy_lock: threading.Lock = threading.Lock()
 
 
-def _get_new_collective_id_unsafe() -> int:
-    # Not thread safe, must be called with _collective_id_vacancy_lock held
+def _current_pool() -> _CollectiveIDPool:
+    """Return the collective ID pool of the current context, creating it if unset."""
     try:
-        # All ranks must choose the same collective IDs during lowering.
-        # Since 3.7, dict.popitem() is guaranteed LIFO.
-        collective_id, _ = _collective_id_vacancy.popitem()
-    except KeyError:
-        raise ValueError(
-            f"Cannot shuffle more than {Shuffler.max_concurrent_shuffles} "
-            "times in a single query."
-        ) from None
-    return collective_id
+        return _collective_id_pool.get()
+    except LookupError:
+        pool = _CollectiveIDPool()
+        _collective_id_pool.set(pool)
+        return pool
+
 
+@contextmanager
+def collective_id_scope() -> Iterator[None]:
+    """
+    Use a fresh collective ID pool for the duration of the block.
 
-def _release_collective_id_unsafe(collective_id: int) -> None:
-    """Release a collective ID back to the vacancy set."""
-    # Not thread safe, must be called with _collective_id_vacancy_lock held
-    if collective_id in _collective_id_vacancy:
-        raise ValueError("Restoring ID that exists")
-    _collective_id_vacancy[collective_id] = None
+    Every query should run in its own scope so that the collective IDs it
+    reserves depend only on the query itself, and not on collectives that
+    other queries in the same process (possibly on other communicators)
+    currently hold.
+    """
+    token = _collective_id_pool.set(_CollectiveIDPool())
+    try:
+        yield
+    finally:
+        _collective_id_pool.reset(token)
 
 
 class ReserveOpIDs:
@@ -78,7 +116,8 @@ class ReserveOpIDs:
     -----
     This context manager:
     1. Identifies all IR nodes that may require collective operations
-    2. Reserves collective IDs from the vacancy pool
+    2. Reserves collective IDs from the current context's pool
+       (see :func:`collective_id_scope`)
     3. Creates a mapping from IR nodes to their reserved IDs
     4. Releases all IDs back to the pool on __exit__
 
@@ -130,6 +169,7 @@ def __init__(
             node for node in traversal([ir]) if isinstance(node, collective_types)
         ]
         self.collective_id_map: dict[IR, list[int]] = {}
+        self._pool: _CollectiveIDPool | None = None
 
     def __enter__(self) -> dict[IR, list[int]]:
         """
@@ -143,37 +183,38 @@ def __enter__(self) -> dict[IR, list[int]]:
             collective operations per node.
         """
         # Reserve IDs and map nodes to a list of IDs
-        with _collective_id_vacancy_lock:
+        self._pool = pool = _current_pool()
+        with pool.lock:
             for node in self.collective_nodes:
                 if (
                     isinstance(node, (GroupBy, Distinct))
                     and self.dynamic_planning_enabled
                 ):
                     # GroupBy/Distinct need 2 IDs: one for size allgather, one for shuffle
                     self.collective_id_map[node] = [
-                        _get_new_collective_id_unsafe(),
-                        _get_new_collective_id_unsafe(),
+                        pool.acquire_unsafe(),
+                        pool.acquire_unsafe(),
                     ]
                 elif isinstance(node, Join) and self.dynamic_planning_enabled:
                     # Join needs 3 IDs: allgather, left shuffle, right shuffle.
                     self.collective_id_map[node] = [
-                        _get_new_collective_id_unsafe(),
-                        _get_new_collective_id_unsafe(),
-                        _get_new_collective_id_unsafe(),
+                        pool.acquire_unsafe(),
+                        pool.acquire_unsafe(),
+                        pool.acquire_unsafe(),
                     ]
                 elif isinstance(node, Sort):
                     if self.dynamic_planning_enabled:
                         # 3 IDs: size-estimate allgather, boundary allgather, shuffle
                         self.collective_id_map[node] = [
-                            _get_new_collective_id_unsafe(),
-                            _get_new_collective_id_unsafe(),
-                            _get_new_collective_id_unsafe(),
+                            pool.acquire_unsafe(),
+                            pool.acquire_unsafe(),
+                            pool.acquire_unsafe(),
                         ]
                     else:
                         # 2 IDs: boundary allgather, shuffle
                         self.collective_id_map[node] = [
-                            _get_new_collective_id_unsafe(),
-                            _get_new_collective_id_unsafe(),
+                            pool.acquire_unsafe(),
+                            pool.acquire_unsafe(),
                         ]
                 elif isinstance(node, Over) and not node.is_scalar:
                     # Non-scalar Over needs 2 IDs: one for the size AllGather +
@@ -182,11 +223,11 @@ def __enter__(self) -> dict[IR, list[int]]:
                     # the return shuffle (which overlaps with the forward shuffle
                     # during extract+insert).
                     self.collective_id_map[node] = [
-                        _get_new_collective_id_unsafe(),
-                        _get_new_collective_id_unsafe(),
+                        pool.acquire_unsafe(),
+                        pool.acquire_unsafe(),
                     ]
                 else:
-                    self.collective_id_map[node] = [_get_new_collective_id_unsafe()]
+                    self.collective_id_map[node] = [pool.acquire_unsafe()]
 
         return self.collective_id_map
 
@@ -197,10 +238,12 @@ def __exit__(
         exc_tb: TracebackType | None,
     ) -> Literal[False]:
         """Release all reserved collective IDs back to the vacancy pool."""
-        with _collective_id_vacancy_lock:
+        pool = self._pool
+        assert pool is not None
+        with pool.lock:
             for collective_ids in self.collective_id_map.values():
                 for collective_id in collective_ids:
-                    _release_collective_id_unsafe(collective_id)
+                    pool.release_unsafe(collective_id)
         return False
 
 
@@ -217,12 +260,14 @@ def reserve_op_id() -> Iterator[int]:
     Yields
     ------
     collective_id : int
-        A vacant collective ID reserved from the global vacancy pool.
+        A vacant collective ID reserved from the current context's pool
+        (see :func:`collective_id_scope`).
     """
-    with _collective_id_vacancy_lock:
-        collective_id = _get_new_collective_id_unsafe()
+    pool = _current_pool()
+    with pool.lock:
+        collective_id = pool.acquire_unsafe()
     try:
         yield collective_id
     finally:
-        with _collective_id_vacancy_lock:
-            _release_collective_id_unsafe(collective_id)
+        with pool.lock:
+            pool.release_unsafe(collective_id)

python/cudf_polars/tests/streaming/test_collective_ids.py

@@ -0,0 +1,98 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+import threading
+import uuid
+from typing import Any
+
+import polars as pl
+
+import cudf_polars.engine.core
+from cudf_polars.dsl.translate import Translator
+from cudf_polars.streaming.actor_graph.collectives import (
+    ReserveOpIDs,
+    collective_id_scope,
+    reserve_op_id,
+)
+
+
+def _sort_ir():
+    q = pl.LazyFrame({"a": [3, 1, 2]}).sort("a")
+    return Translator(q._ldf.visit(), pl.GPUEngine()).translate_ir()
+
+
+def _reserved_ids(ir) -> list[int]:
+    with ReserveOpIDs(ir) as collective_id_map:
+        return [i for ids in collective_id_map.values() for i in ids]
+
+
+def test_collective_id_scope_ignores_held_ids() -> None:
+    with collective_id_scope(), reserve_op_id() as expected:
+        pass
+    # Another query (possibly on another communicator) holds an ID.
+    with (
+        collective_id_scope(),
+        reserve_op_id(),
+        collective_id_scope(),
+        reserve_op_id() as op_id,
+    ):
+        assert op_id == expected
+
+
+def test_collective_id_scope_ignores_history() -> None:
+    ir = _sort_ir()
+    with collective_id_scope():
+        expected = _reserved_ids(ir)
+    assert len(expected) > 1
+
+    with collective_id_scope():
+        # IDs are not released in reverse order of reservation, so earlier
+        # queries permute the pool.
+        _reserved_ids(ir)
+        assert _reserved_ids(ir) != expected
+        with collective_id_scope():
+            assert _reserved_ids(ir) == expected
+
+
+def test_collective_id_scope_per_thread() -> None:
+    with collective_id_scope(), reserve_op_id() as expected:
+        pass
+    held = threading.Event()
+    release = threading.Event()
+
+    def hold() -> None:
+        with collective_id_scope(), reserve_op_id():
+            held.set()
+            release.wait()
+
+    thread = threading.Thread(target=hold)
+    thread.start()
+    try:
+        held.wait()
+        with collective_id_scope(), reserve_op_id() as op_id:
+            assert op_id == expected
+    finally:
+        release.set()
+        thread.join()
+
+
+def test_evaluate_on_rank_uses_fresh_scope(monkeypatch) -> None:
+    with collective_id_scope(), reserve_op_id() as expected:
+        pass
+    seen: list[int] = []
+
+    def fake_evaluate_on_rank(*args, **kwargs):
+        with reserve_op_id() as op_id:
+            seen.append(op_id)
+
+    monkeypatch.setattr(
+        cudf_polars.engine.core, "_evaluate_on_rank", fake_evaluate_on_rank
+    )
+    unused: Any = None
+    with reserve_op_id():
+        cudf_polars.engine.core.evaluate_on_rank(
+            unused, unused, unused, unused, unused, query_id=uuid.uuid4()
+        )
+    assert seen == [expected]