From 1bd2d660b9bb6dea38d0d9adefe9f0d81b54f47a Mon Sep 17 00:00:00 2001 From: Tom Augspurger Date: Fri, 16 Oct 2026 10:13:53 -0700 Subject: [PATCH 01/20] Enable _reset to update quent_context; fast quent integration tests Allow StreamingEngine._reset to accept a new quent_context in executor_options (removing the previous ValueError guard). Each backend (SPMD, Ray, Dask) tears down the old Quent runtimes and starts fresh ones when quent_context changes, with workers always disconnected before the controller is closed to avoid collector timeouts. The Quent integration test fixture is updated to use the session-scoped engine from conftest (UCX / rapidsmpf / Ray bootstrap paid once per session) and call _reset per test to swap the Quent context, rather than constructing a full engine per variant. Co-Authored-By: Claude Sonnet 4.6 (1M context) --- python/cudf_polars/cudf_polars/engine/core.py | 6 - python/cudf_polars/cudf_polars/engine/dask.py | 83 +++++++- python/cudf_polars/cudf_polars/engine/ray.py | 62 +++++- python/cudf_polars/cudf_polars/engine/spmd.py | 59 +++++- .../tests/quent/test_quent_integration.py | 191 ++++++++---------- 5 files changed, 277 insertions(+), 124 deletions(-) diff --git a/python/cudf_polars/cudf_polars/engine/core.py b/python/cudf_polars/cudf_polars/engine/core.py index a23c9f78570..4a6bc9d5a1c 100644 --- a/python/cudf_polars/cudf_polars/engine/core.py +++ b/python/cudf_polars/cudf_polars/engine/core.py @@ -539,12 +539,6 @@ def _reset( existing_executor_options = self.config.get("executor_options", {}) if not isinstance(existing_executor_options, dict): existing_executor_options = {} - existing_quent_context = existing_executor_options.get("quent_context") - if ( - "quent_context" in executor_options - and executor_options["quent_context"] != existing_quent_context - ): - raise ValueError("quent_context cannot be changed during reset") _disallowed_exec = {"num_py_executors"} & executor_options.keys() if _disallowed_exec: diff --git a/python/cudf_polars/cudf_polars/engine/dask.py b/python/cudf_polars/cudf_polars/engine/dask.py index f76336c7501..9531fc3644a 100644 --- a/python/cudf_polars/cudf_polars/engine/dask.py +++ b/python/cudf_polars/cudf_polars/engine/dask.py @@ -497,6 +497,36 @@ def _close_quent_worker( mp_ctx.quent_worker_runtime = None +def _update_quent_on_worker( + quent_context: cudf_polars.quent.QuentConfig | None, + collector_address: str | None, + *, + uid: str, + dask_worker: distributed.Worker | None = None, +) -> None: + """Replace one worker's Quent runtime (or clear it when both args are ``None``).""" + assert dask_worker is not None + mp_ctx: _WorkerContext | None = getattr( + dask_worker, f"_cudf_polars_mp_context_{uid}", None + ) + if mp_ctx is None: + return + if mp_ctx.quent_worker_runtime is not None: + mp_ctx.quent_worker_runtime.close() + mp_ctx.quent_worker_runtime = None + if quent_context is not None and collector_address is not None: + assert mp_ctx.comm is not None + worker_id = uuid.uuid4() + mp_ctx.quent_worker_runtime = QuentWorkerRuntime.create( + quent_context, + collector_address, + worker_id=worker_id, + rank=mp_ctx.comm.rank, + nranks=mp_ctx.comm.nranks, + instance_name=f"DaskWorker-rank{mp_ctx.comm.rank}", + ) + + def _teardown_worker( *, uid: str, dask_worker: distributed.Worker | None = None ) -> None: @@ -1208,8 +1238,6 @@ def _reset( engine_options=engine_options, ) executor_options = executor_options or {} - if "quent_context" in existing_executor_options: - executor_options.setdefault("quent_context", existing_quent_context) if "kvikio_nthreads" in existing_executor_options: executor_options.setdefault( "kvikio_nthreads", existing_executor_options["kvikio_nthreads"] @@ -1243,6 +1271,55 @@ def _reset( rapidsmpf_options_as_bytes, ) + # Handle quent_context change. + new_quent_context: cudf_polars.quent.QuentConfig | None = executor_options.get( + "quent_context", existing_quent_context + ) + if new_quent_context != existing_quent_context: + # Workers must disconnect from the collector before the controller + # is closed; otherwise the collector times out waiting for them to + # flush their pending events. + ctx.client.run( + functools.partial( + _update_quent_on_worker, None, None, uid=ctx.rapidsmpf_id + ) + ) + if self._quent_runtime is not None: + self._quent_runtime.close() + self._quent_runtime = None + if new_quent_context is not None: + quent_collector = cudf_polars.quent._runtime.start_collector( + new_quent_context.run_root + ) + collector_address: str | None = quent_collector.address + assert collector_address is not None + ctx.client.run( + functools.partial( + _update_quent_on_worker, + new_quent_context, + collector_address, + uid=ctx.rapidsmpf_id, + ) + ) + self._quent_runtime = QuentControllerRuntime.create( + new_quent_context, + collector_address, + backend="dask", + collector=quent_collector, + ) + executor_options["quent_context"] = new_quent_context + self._quent_output_root = ( + new_quent_context.run_root if new_quent_context is not None else None + ) + new_ctx = DaskContext( + client=ctx.client, + rapidsmpf_id=ctx.rapidsmpf_id, + quent_controller_runtime=self._quent_runtime, + owned_client=ctx.owned_client, + owned_cluster=ctx.owned_cluster, + ) + self._dask_context = new_ctx + # Re-run ``StreamingEngine.__init__`` on the existing instance to # reconfigure the polars ``GPUEngine`` layer (``self.config``, # ``self.device``, etc.) with the new options. Pass the existing @@ -1253,7 +1330,7 @@ def _reset( executor_options={ **executor_options, "cluster": "dask", - "dask_context": ctx, + "dask_context": new_ctx, }, engine_options={**engine_options, "memory_resource": None}, exit_stack=self._exit_stack, diff --git a/python/cudf_polars/cudf_polars/engine/ray.py b/python/cudf_polars/cudf_polars/engine/ray.py index e776570a059..d2a09080a1a 100644 --- a/python/cudf_polars/cudf_polars/engine/ray.py +++ b/python/cudf_polars/cudf_polars/engine/ray.py @@ -396,6 +396,31 @@ def close_quent(self) -> None: return self._quent_worker_runtime.close() + def update_quent( + self, + quent_context: cudf_polars.quent.QuentConfig | None, + collector_address: str | None, + ) -> None: + """ + Replace this actor's Quent worker runtime. + + Pass ``None`` for both arguments to clear an existing runtime + without starting a new one. + """ + if self._quent_worker_runtime is not None: + self._quent_worker_runtime.close() + self._quent_worker_runtime = None + if quent_context is not None and collector_address is not None: + assert self._comm is not None + self._quent_worker_runtime = QuentWorkerRuntime.create( + quent_context, + collector_address, + worker_id=self._worker_id, + rank=self._comm.rank, + nranks=self._nranks, + instance_name=f"RankActor-{self._worker_id.hex[:8]}", + ) + def reset( self, *, @@ -1049,8 +1074,6 @@ def _reset( engine_options=engine_options, ) executor_options = executor_options or {} - if "quent_context" in existing_executor_options: - executor_options.setdefault("quent_context", existing_quent_context) if "kvikio_nthreads" in existing_executor_options: executor_options.setdefault( "kvikio_nthreads", existing_executor_options["kvikio_nthreads"] @@ -1084,6 +1107,41 @@ def _reset( ] ) + # Handle quent_context change. + new_quent_context: cudf_polars.quent.QuentConfig | None = executor_options.get( + "quent_context", existing_quent_context + ) + if new_quent_context != existing_quent_context: + # Workers first so the collector can drain before we close it. + ray.get( + [rank.update_quent.remote(None, None) for rank in self._rank_actors] + ) + if self._quent_runtime is not None: + self._quent_runtime.close() + self._quent_runtime = None + if new_quent_context is not None: + quent_collector = cudf_polars.quent._runtime.start_collector( + new_quent_context.run_root + ) + collector_address: str | None = quent_collector.address + assert collector_address is not None + ray.get( + [ + rank.update_quent.remote(new_quent_context, collector_address) + for rank in self._rank_actors + ] + ) + self._quent_runtime = QuentControllerRuntime.create( + new_quent_context, + collector_address, + backend="ray", + collector=quent_collector, + ) + executor_options["quent_context"] = new_quent_context + self._quent_output_root = ( + new_quent_context.run_root if new_quent_context is not None else None + ) + # Re-run ``StreamingEngine.__init__`` on the existing instance to # reconfigure the polars ``GPUEngine`` layer (``self.config``, # ``self.device``, etc.) with the new options. Pass the existing diff --git a/python/cudf_polars/cudf_polars/engine/spmd.py b/python/cudf_polars/cudf_polars/engine/spmd.py index 14dfe5de6d8..d21302f4582 100644 --- a/python/cudf_polars/cudf_polars/engine/spmd.py +++ b/python/cudf_polars/cudf_polars/engine/spmd.py @@ -990,8 +990,6 @@ def _reset( engine_options=engine_options, ) executor_options = executor_options or {} - if "quent_context" in existing_executor_options: - executor_options.setdefault("quent_context", existing_quent_context) if "kvikio_nthreads" in existing_executor_options: executor_options.setdefault( "kvikio_nthreads", existing_executor_options["kvikio_nthreads"] @@ -1007,12 +1005,25 @@ def _reset( request_ceiling=executor_options["kvikio_request_ceiling"], ) engine_options = engine_options or {} - quent_context: cudf_polars.quent.QuentConfig | None = executor_options.get( - "quent_context" + new_quent_context: cudf_polars.quent.QuentConfig | None = executor_options.get( + "quent_context", existing_quent_context ) + quent_changing = new_quent_context != existing_quent_context rapidsmpf_options = resolve_rapidsmpf_options(rapidsmpf_options) self.rapidsmpf_options = rapidsmpf_options + # Tear down Quent runtimes before the Context rebuild if they're changing. + if quent_changing: + if self._comm.nranks > 1: + barrier(self._comm) + self._close_quent_worker() + if self._comm.nranks > 1: + barrier(self._comm) + if self._comm.rank == 0: + self._close_quent_controller() + if self._comm.nranks > 1: + barrier(self._comm) + # Collective: synchronize all ranks before tearing down the Context. if self._comm.nranks > 1: barrier(self._comm) @@ -1041,8 +1052,46 @@ def _reset( self._mr = self._ctx.br().device_mr_adaptor() rmm.mr.set_current_device_resource(self._mr) + # Start new Quent runtimes after the Context is rebuilt. + if quent_changing and new_quent_context is not None: + collector_address: str = "" + if self._comm.rank == 0: + quent_collector = cudf_polars.quent._runtime.start_collector( + new_quent_context.run_root + ) + _collector_cleanup = contextlib.ExitStack() + _collector_cleanup.callback(quent_collector.close) + collector_address = quent_collector.address + new_quent_context, collector_address = synchronize_quent_configuration( + comm=self._comm, + context=self._ctx, + quent_config=new_quent_context, + collector_address=collector_address, + ) + if self._comm.rank == 0: + self._quent_runtime = QuentControllerRuntime.create( + new_quent_context, + collector_address, + backend="spmd", + collector=quent_collector, + ) + _collector_cleanup.pop_all() + self._quent_worker_runtime = QuentWorkerRuntime.create( + new_quent_context, + collector_address, + worker_id=self._quent_worker_id, + rank=self._comm.rank, + nranks=self._comm.nranks, + instance_name=f"rank-{self._comm.rank}", + ) + executor_options["quent_context"] = new_quent_context + self._quent_output_root = ( + new_quent_context.run_root if new_quent_context is not None else None + ) engine_id = ( - quent_context.engine_id if quent_context is not None else uuid.uuid4() + new_quent_context.engine_id + if new_quent_context is not None + else uuid.uuid4() ) # Re-run ``StreamingEngine.__init__`` on the existing instance to diff --git a/python/cudf_polars/tests/quent/test_quent_integration.py b/python/cudf_polars/tests/quent/test_quent_integration.py index 088fe1a86c1..e2b6d2d5cbd 100644 --- a/python/cudf_polars/tests/quent/test_quent_integration.py +++ b/python/cudf_polars/tests/quent/test_quent_integration.py @@ -21,67 +21,51 @@ pytest.importorskip("cudf_polars_quent") -import cudf_polars.quent import cudf_polars.quent._runtime if TYPE_CHECKING: - from collections.abc import Generator, Iterator + from collections.abc import Callable, Generator, Iterator from pathlib import Path from cudf_polars.engine.core import StreamingEngine + from cudf_polars.engine.options import StreamingOptions from cudf_polars.quent import QuentConfig -@pytest.fixture(params=["ray", "dask", "spmd"]) +@pytest.fixture def engine_with_quent_context( - request: pytest.FixtureRequest, + _unconfigured_engine: tuple[StreamingEngine, StreamingOptions], quent_context: QuentConfig, - ray_num_ranks: int, - ray_init_options: dict[str, Any], + streaming_engine_factory: Callable[[StreamingOptions], StreamingEngine], ) -> Iterator[StreamingEngine]: - backend = request.param - executor_options = { - "quent_context": quent_context, - "fallback_mode": "silent", - } - engine: StreamingEngine - if backend == "ray": - pytest.importorskip("ray") - from cudf_polars.engine.ray import RayEngine - - engine = RayEngine( - executor_options=executor_options, - engine_options={"allow_gpu_sharing": True}, - ray_init_options=ray_init_options, - num_ranks=ray_num_ranks, - ) - elif backend == "dask": - pytest.importorskip("distributed") - from cudf_polars.engine.dask import DaskEngine - - engine = DaskEngine(executor_options=executor_options) - else: - from rapidsmpf import bootstrap - from rapidsmpf.communicator.single import new_communicator - from rapidsmpf.config import Options, get_environment_variables - from rapidsmpf.progress_thread import ProgressThread - - from cudf_polars.engine.spmd import SPMDEngine - - comm = ( - bootstrap.create_ucxx_comm( - progress_thread=ProgressThread(), type=bootstrap.BackendType.AUTO - ) - if bootstrap.is_running_with_rrun() - else new_communicator( - Options(get_environment_variables()), ProgressThread() - ) - ) - engine = SPMDEngine(executor_options=executor_options, comm=comm) + """ + Yield a streaming engine configured with Quent tracing. + + Reuses the session-scoped engine from conftest (UCX / rapidsmpf / Ray + bootstrap paid once per session). ``streaming_engine_factory`` is + included to tell ``pytest_generate_tests`` to parametrise over all + streaming backends; it is not called directly here. + """ + engine, base_options = _unconfigured_engine + engine._reset( + rapidsmpf_options=base_options.to_rapidsmpf_options(), + executor_options={ + **base_options.to_executor_options(), + "quent_context": quent_context, + }, + engine_options=base_options.to_engine_options(), + ) try: yield engine finally: - engine.shutdown() + engine._reset( + rapidsmpf_options=base_options.to_rapidsmpf_options(), + executor_options={ + **base_options.to_executor_options(), + "quent_context": None, + }, + engine_options=base_options.to_engine_options(), + ) def _stored_events(root: Path) -> list[dict[str, Any]]: @@ -187,68 +171,59 @@ def test_quent_lifecycle( monkeypatch: pytest.MonkeyPatch, ) -> None: query = pl.LazyFrame({"x": [1, 2, 3]}).filter(pl.col("x") > 1) - with engine_with_quent_context: - assert ( - engine_with_quent_context.config["executor_options"]["quent_context"] - == quent_context - ) - current_context = dataclasses.replace( - quent_context, - query=dataclasses.replace( - quent_context.query, - query_group_id=uuid.uuid4(), - query_group_name="Updated Query Group", - query_name="Iteration 1", - ), - ) - engine_with_quent_context.config["executor_options"]["quent_context"] = ( - current_context - ) - query.collect(engine=engine_with_quent_context) - current_context = dataclasses.replace( - current_context, - query=dataclasses.replace( - current_context.query, - query_name="Iteration 2", - ), - ) - engine_with_quent_context.config["executor_options"]["quent_context"] = ( - current_context - ) - query.collect(engine=engine_with_quent_context) - current_context = dataclasses.replace( - current_context, - query=dataclasses.replace( - current_context.query, - query_name="Failed iteration", - ), - ) - engine_with_quent_context.config["executor_options"]["quent_context"] = ( - current_context - ) - failed_query = ( - pl.LazyFrame({"orderby": [1, 2, 4, 2], "value": [1, 2, 3, 4]}) - .rolling("orderby", period="2i") - .agg(pl.sum("value")) - ) - with ( - suppress_worker_exceptions(engine_with_quent_context, monkeypatch), - pytest.raises(Exception), # noqa: B017 - backend-specific wrapper - ): - failed_query.collect(engine=engine_with_quent_context) - with pytest.raises(ValueError, match="quent_context cannot be changed"): - engine_with_quent_context._reset( - executor_options={"quent_context": cudf_polars.quent.QuentConfig()} - ) - - # StreamingExecutor - config_options = ConfigOptions.from_polars_engine(engine_with_quent_context) - hash_a = hash(config_options) - hash_b = hash(config_options) - assert hash_a == hash_b - - assert engine_with_quent_context._quent_output_root is not None - events = _stored_events(engine_with_quent_context._quent_output_root) + engine = engine_with_quent_context + assert engine.config["executor_options"]["quent_context"] == quent_context + current_context = dataclasses.replace( + quent_context, + query=dataclasses.replace( + quent_context.query, + query_group_id=uuid.uuid4(), + query_group_name="Updated Query Group", + query_name="Iteration 1", + ), + ) + engine.config["executor_options"]["quent_context"] = current_context + query.collect(engine=engine) + current_context = dataclasses.replace( + current_context, + query=dataclasses.replace( + current_context.query, + query_name="Iteration 2", + ), + ) + engine.config["executor_options"]["quent_context"] = current_context + query.collect(engine=engine) + current_context = dataclasses.replace( + current_context, + query=dataclasses.replace( + current_context.query, + query_name="Failed iteration", + ), + ) + engine.config["executor_options"]["quent_context"] = current_context + failed_query = ( + pl.LazyFrame({"orderby": [1, 2, 4, 2], "value": [1, 2, 3, 4]}) + .rolling("orderby", period="2i") + .agg(pl.sum("value")) + ) + with ( + suppress_worker_exceptions(engine, monkeypatch), + pytest.raises(Exception), # noqa: B017 - backend-specific wrapper + ): + failed_query.collect(engine=engine) + # StreamingExecutor + config_options = ConfigOptions.from_polars_engine(engine) + hash_a = hash(config_options) + hash_b = hash(config_options) + assert hash_a == hash_b + + # Capture the output root, then close Quent to flush events to disk. + # The fixture teardown calls _reset(..., quent_context=None) again, which is a no-op. + quent_output_root = engine._quent_output_root + assert quent_output_root is not None + engine._reset(executor_options={"quent_context": None, "fallback_mode": "silent"}) + + events = _stored_events(quent_output_root) engine_events = _of_type(events, "Engine") assert len(engine_events) == 2 assert engine_events[0]["id"] == str(quent_context.engine_id) From a8d2dd40797f34a8a6cfded363408ddb7bef5a22 Mon Sep 17 00:00:00 2001 From: Tom Augspurger Date: Fri, 16 Oct 2026 10:14:10 -0700 Subject: [PATCH 02/20] Deduplicate quent open/close across engine _reset methods Extract the repeated close-workers / close-controller / open-collector / open-workers / open-controller sequence into per-backend hook methods (_quent_close, _quent_open) and a common driver (_update_quent) on the StreamingEngine base class. Each _reset now calls self._update_quent(...) instead of repeating the same ~25-line block. _shutdown_spmd is also simplified to call self._quent_close() directly. Co-Authored-By: Claude Sonnet 4.6 (1M context) --- python/cudf_polars/cudf_polars/engine/core.py | 36 ++++++ python/cudf_polars/cudf_polars/engine/dask.py | 68 ++++++----- python/cudf_polars/cudf_polars/engine/ray.py | 58 +++++----- python/cudf_polars/cudf_polars/engine/spmd.py | 106 +++++++++--------- 4 files changed, 149 insertions(+), 119 deletions(-) diff --git a/python/cudf_polars/cudf_polars/engine/core.py b/python/cudf_polars/cudf_polars/engine/core.py index 4a6bc9d5a1c..56bf051a05b 100644 --- a/python/cudf_polars/cudf_polars/engine/core.py +++ b/python/cudf_polars/cudf_polars/engine/core.py @@ -583,6 +583,42 @@ def __exit__(self, *_: object) -> None: """Exit the context manager, calling :meth:`shutdown`.""" self.shutdown() + def _update_quent( + self, + new_quent_context: QuentConfig | None, + existing_quent_context: QuentConfig | None, + executor_options: dict[str, Any], + ) -> None: + """ + Apply a ``quent_context`` change in-place. + + Closes existing Quent runtimes, starts new ones for + ``new_quent_context``, and patches ``executor_options["quent_context"]`` + and ``_quent_output_root``. Subclasses implement :meth:`_quent_close` + and :meth:`_quent_open` for the backend-specific work. + """ + if new_quent_context == existing_quent_context: + return + self._quent_close() + if new_quent_context is not None: + self._quent_open(new_quent_context) + self._quent_output_root = ( + new_quent_context.run_root if new_quent_context is not None else None + ) + executor_options["quent_context"] = new_quent_context + + def _quent_close(self) -> None: + """Close all Quent runtimes for this engine (workers then controller).""" + raise NotImplementedError( # pragma: no cover + f"{type(self).__name__} does not implement _quent_close" + ) + + def _quent_open(self, quent_context: QuentConfig) -> None: + """Start fresh Quent runtimes for ``quent_context`` (workers + controller).""" + raise NotImplementedError( # pragma: no cover + f"{type(self).__name__} does not implement _quent_open" + ) + def _run(self, func: Callable[..., T], *args: Any, **kwargs: Any) -> list[T]: """ Execute a function on all ranks. diff --git a/python/cudf_polars/cudf_polars/engine/dask.py b/python/cudf_polars/cudf_polars/engine/dask.py index 9531fc3644a..7abf14ca4d9 100644 --- a/python/cudf_polars/cudf_polars/engine/dask.py +++ b/python/cudf_polars/cudf_polars/engine/dask.py @@ -1275,42 +1275,7 @@ def _reset( new_quent_context: cudf_polars.quent.QuentConfig | None = executor_options.get( "quent_context", existing_quent_context ) - if new_quent_context != existing_quent_context: - # Workers must disconnect from the collector before the controller - # is closed; otherwise the collector times out waiting for them to - # flush their pending events. - ctx.client.run( - functools.partial( - _update_quent_on_worker, None, None, uid=ctx.rapidsmpf_id - ) - ) - if self._quent_runtime is not None: - self._quent_runtime.close() - self._quent_runtime = None - if new_quent_context is not None: - quent_collector = cudf_polars.quent._runtime.start_collector( - new_quent_context.run_root - ) - collector_address: str | None = quent_collector.address - assert collector_address is not None - ctx.client.run( - functools.partial( - _update_quent_on_worker, - new_quent_context, - collector_address, - uid=ctx.rapidsmpf_id, - ) - ) - self._quent_runtime = QuentControllerRuntime.create( - new_quent_context, - collector_address, - backend="dask", - collector=quent_collector, - ) - executor_options["quent_context"] = new_quent_context - self._quent_output_root = ( - new_quent_context.run_root if new_quent_context is not None else None - ) + self._update_quent(new_quent_context, existing_quent_context, executor_options) new_ctx = DaskContext( client=ctx.client, rapidsmpf_id=ctx.rapidsmpf_id, @@ -1336,6 +1301,37 @@ def _reset( exit_stack=self._exit_stack, ) + def _quent_close(self) -> None: + """Clear Quent worker runtimes on all Dask workers, then close the controller.""" + ctx = self._dask_ctx + # Workers first so the collector can drain before we close it. + ctx.client.run( + functools.partial(_update_quent_on_worker, None, None, uid=ctx.rapidsmpf_id) + ) + if self._quent_runtime is not None: + self._quent_runtime.close() + self._quent_runtime = None + + def _quent_open(self, quent_context: cudf_polars.quent.QuentConfig) -> None: + """Start a new Quent collector, connect workers, create controller.""" + ctx = self._dask_ctx + quent_collector = cudf_polars.quent._runtime.start_collector( + quent_context.run_root + ) + collector_address = quent_collector.address + assert collector_address is not None + ctx.client.run( + functools.partial( + _update_quent_on_worker, + quent_context, + collector_address, + uid=ctx.rapidsmpf_id, + ) + ) + self._quent_runtime = QuentControllerRuntime.create( + quent_context, collector_address, backend="dask", collector=quent_collector + ) + @classmethod def from_options( cls, diff --git a/python/cudf_polars/cudf_polars/engine/ray.py b/python/cudf_polars/cudf_polars/engine/ray.py index d2a09080a1a..50078956b0b 100644 --- a/python/cudf_polars/cudf_polars/engine/ray.py +++ b/python/cudf_polars/cudf_polars/engine/ray.py @@ -1111,36 +1111,7 @@ def _reset( new_quent_context: cudf_polars.quent.QuentConfig | None = executor_options.get( "quent_context", existing_quent_context ) - if new_quent_context != existing_quent_context: - # Workers first so the collector can drain before we close it. - ray.get( - [rank.update_quent.remote(None, None) for rank in self._rank_actors] - ) - if self._quent_runtime is not None: - self._quent_runtime.close() - self._quent_runtime = None - if new_quent_context is not None: - quent_collector = cudf_polars.quent._runtime.start_collector( - new_quent_context.run_root - ) - collector_address: str | None = quent_collector.address - assert collector_address is not None - ray.get( - [ - rank.update_quent.remote(new_quent_context, collector_address) - for rank in self._rank_actors - ] - ) - self._quent_runtime = QuentControllerRuntime.create( - new_quent_context, - collector_address, - backend="ray", - collector=quent_collector, - ) - executor_options["quent_context"] = new_quent_context - self._quent_output_root = ( - new_quent_context.run_root if new_quent_context is not None else None - ) + self._update_quent(new_quent_context, existing_quent_context, executor_options) # Re-run ``StreamingEngine.__init__`` on the existing instance to # reconfigure the polars ``GPUEngine`` layer (``self.config``, @@ -1160,6 +1131,33 @@ def _reset( exit_stack=self._exit_stack, ) + def _quent_close(self) -> None: + """Clear Quent worker runtimes on all actors, then close the controller.""" + assert self._rank_actors is not None + # Workers first so the collector can drain before we close it. + ray.get([rank.update_quent.remote(None, None) for rank in self._rank_actors]) + if self._quent_runtime is not None: + self._quent_runtime.close() + self._quent_runtime = None + + def _quent_open(self, quent_context: cudf_polars.quent.QuentConfig) -> None: + """Start a new Quent collector, connect actors, create controller.""" + assert self._rank_actors is not None + quent_collector = cudf_polars.quent._runtime.start_collector( + quent_context.run_root + ) + collector_address = quent_collector.address + assert collector_address is not None + ray.get( + [ + rank.update_quent.remote(quent_context, collector_address) + for rank in self._rank_actors + ] + ) + self._quent_runtime = QuentControllerRuntime.create( + quent_context, collector_address, backend="ray", collector=quent_collector + ) + @classmethod def from_options( cls, diff --git a/python/cudf_polars/cudf_polars/engine/spmd.py b/python/cudf_polars/cudf_polars/engine/spmd.py index d21302f4582..fa79838ed6b 100644 --- a/python/cudf_polars/cudf_polars/engine/spmd.py +++ b/python/cudf_polars/cudf_polars/engine/spmd.py @@ -768,19 +768,61 @@ def _close_quent_controller(self) -> None: finally: self._quent_runtime = None + def _quent_close(self) -> None: + """Close SPMD Quent runtimes: barrier → workers → barrier → controller → barrier.""" + assert self._comm is not None + comm = self._comm + if comm.nranks > 1: + barrier(comm) + self._close_quent_worker() + if comm.nranks > 1: + barrier(comm) + if comm.rank == 0: + self._close_quent_controller() + if comm.nranks > 1: + barrier(comm) + + def _quent_open(self, quent_context: cudf_polars.quent.QuentConfig) -> None: + """Start SPMD Quent runtimes for ``quent_context``.""" + assert self._comm is not None + assert self._ctx is not None + comm = self._comm + collector_address: str = "" + if comm.rank == 0: + quent_collector = cudf_polars.quent._runtime.start_collector( + quent_context.run_root + ) + _collector_cleanup = contextlib.ExitStack() + _collector_cleanup.callback(quent_collector.close) + collector_address = quent_collector.address + quent_context, collector_address = synchronize_quent_configuration( + comm=comm, + context=self._ctx, + quent_config=quent_context, + collector_address=collector_address, + ) + if comm.rank == 0: + self._quent_runtime = QuentControllerRuntime.create( + quent_context, + collector_address, + backend="spmd", + collector=quent_collector, + ) + _collector_cleanup.pop_all() + self._quent_worker_runtime = QuentWorkerRuntime.create( + quent_context, + collector_address, + worker_id=self._quent_worker_id, + rank=comm.rank, + nranks=comm.nranks, + instance_name=f"rank-{comm.rank}", + ) + def _shutdown_spmd(self) -> None: """Run collective cleanup before local resources unwind.""" steps: list[Callable[[], object]] = [self._drop_persisted] if self._quent_worker_runtime is not None: - assert self._comm is not None - comm = self._comm - steps.append(self._close_quent_worker) - if comm.nranks > 1: - steps.append(lambda: barrier(comm)) - if comm.rank == 0: - steps.append(self._close_quent_controller) - if comm.nranks > 1: - steps.append(lambda: barrier(comm)) + steps.append(self._quent_close) run_cleanup_steps("SPMD engine shutdown failed", *steps) @classmethod @@ -1014,15 +1056,7 @@ def _reset( # Tear down Quent runtimes before the Context rebuild if they're changing. if quent_changing: - if self._comm.nranks > 1: - barrier(self._comm) - self._close_quent_worker() - if self._comm.nranks > 1: - barrier(self._comm) - if self._comm.rank == 0: - self._close_quent_controller() - if self._comm.nranks > 1: - barrier(self._comm) + self._quent_close() # Collective: synchronize all ranks before tearing down the Context. if self._comm.nranks > 1: @@ -1053,41 +1087,7 @@ def _reset( rmm.mr.set_current_device_resource(self._mr) # Start new Quent runtimes after the Context is rebuilt. - if quent_changing and new_quent_context is not None: - collector_address: str = "" - if self._comm.rank == 0: - quent_collector = cudf_polars.quent._runtime.start_collector( - new_quent_context.run_root - ) - _collector_cleanup = contextlib.ExitStack() - _collector_cleanup.callback(quent_collector.close) - collector_address = quent_collector.address - new_quent_context, collector_address = synchronize_quent_configuration( - comm=self._comm, - context=self._ctx, - quent_config=new_quent_context, - collector_address=collector_address, - ) - if self._comm.rank == 0: - self._quent_runtime = QuentControllerRuntime.create( - new_quent_context, - collector_address, - backend="spmd", - collector=quent_collector, - ) - _collector_cleanup.pop_all() - self._quent_worker_runtime = QuentWorkerRuntime.create( - new_quent_context, - collector_address, - worker_id=self._quent_worker_id, - rank=self._comm.rank, - nranks=self._comm.nranks, - instance_name=f"rank-{self._comm.rank}", - ) - executor_options["quent_context"] = new_quent_context - self._quent_output_root = ( - new_quent_context.run_root if new_quent_context is not None else None - ) + self._update_quent(new_quent_context, existing_quent_context, executor_options) engine_id = ( new_quent_context.engine_id if new_quent_context is not None From 33820ddf69367e8b84e1aa5a228fc1900b1cc83e Mon Sep 17 00:00:00 2001 From: Tom Augspurger Date: Fri, 16 Oct 2026 11:11:45 -0700 Subject: [PATCH 03/20] Use fallback_mode=silent in quent integration tests The rolling fallback warning was leaking into test output because it is emitted on Ray/Dask workers where pytest's filterwarnings cannot reach. Configure the engine with fallback_mode="silent" in the fixture instead, which suppresses the warning at the source. Co-Authored-By: Claude Sonnet 4.6 (1M context) --- python/cudf_polars/tests/quent/test_quent_integration.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/python/cudf_polars/tests/quent/test_quent_integration.py b/python/cudf_polars/tests/quent/test_quent_integration.py index e2b6d2d5cbd..b1be8c1c6be 100644 --- a/python/cudf_polars/tests/quent/test_quent_integration.py +++ b/python/cudf_polars/tests/quent/test_quent_integration.py @@ -52,6 +52,7 @@ def engine_with_quent_context( executor_options={ **base_options.to_executor_options(), "quent_context": quent_context, + "fallback_mode": "silent", }, engine_options=base_options.to_engine_options(), ) @@ -63,6 +64,7 @@ def engine_with_quent_context( executor_options={ **base_options.to_executor_options(), "quent_context": None, + "fallback_mode": "silent", }, engine_options=base_options.to_engine_options(), ) @@ -164,7 +166,6 @@ def test_quent_controller_cleanup_continues_after_handle_failure() -> None: collector.close.assert_called_once_with(timeout=10.0) -@pytest.mark.filterwarnings("ignore:Rolling.*:UserWarning") def test_quent_lifecycle( engine_with_quent_context: StreamingEngine, quent_context: QuentConfig, From 6f302e2cd6bc451c21c6437663b8e5d42c1a3c1d Mon Sep 17 00:00:00 2001 From: Tom Augspurger Date: Fri, 16 Oct 2026 12:18:14 -0700 Subject: [PATCH 04/20] Clarify why engine_with_quent_context uses _unconfigured_engine directly Co-Authored-By: Claude Sonnet 4.6 (1M context) --- python/cudf_polars/tests/quent/test_quent_integration.py | 8 +++----- 1 file changed, 3 insertions(+), 5 deletions(-) diff --git a/python/cudf_polars/tests/quent/test_quent_integration.py b/python/cudf_polars/tests/quent/test_quent_integration.py index b1be8c1c6be..3453b4d0ede 100644 --- a/python/cudf_polars/tests/quent/test_quent_integration.py +++ b/python/cudf_polars/tests/quent/test_quent_integration.py @@ -40,12 +40,10 @@ def engine_with_quent_context( ) -> Iterator[StreamingEngine]: """ Yield a streaming engine configured with Quent tracing. - - Reuses the session-scoped engine from conftest (UCX / rapidsmpf / Ray - bootstrap paid once per session). ``streaming_engine_factory`` is - included to tell ``pytest_generate_tests`` to parametrise over all - streaming backends; it is not called directly here. """ + # Uses ``_unconfigured_engine`` directly (rather than a public engine + # fixture) so that quent setup and teardown can be expressed as a single + # ``_reset`` call each way. engine, base_options = _unconfigured_engine engine._reset( rapidsmpf_options=base_options.to_rapidsmpf_options(), From 4a3de24fb55aed6ea80274dd7df208058b76a746 Mon Sep 17 00:00:00 2001 From: Tom Augspurger Date: Fri, 16 Oct 2026 12:26:37 -0700 Subject: [PATCH 05/20] Preserve allow_gpu_sharing across _reset calls MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit StreamingEngine.__init__ pops allow_gpu_sharing from engine_options before passing to polars, so it was never visible in self.config. _reset calls that omit engine_options (e.g. the in-test quent flush) therefore defaulted it to False, triggering the UUID-collision check on multi-rank Ray/Dask engines. Store it as self._allow_gpu_sharing at init time and prepend it to engine_options in each engine's _reset → StreamingEngine.__init__ call so the value survives every reset regardless of what the caller passes. Co-Authored-By: Claude Sonnet 4.6 (1M context) --- python/cudf_polars/cudf_polars/engine/core.py | 1 + python/cudf_polars/cudf_polars/engine/dask.py | 6 +++++- python/cudf_polars/cudf_polars/engine/ray.py | 5 ++++- python/cudf_polars/cudf_polars/engine/spmd.py | 1 + 4 files changed, 11 insertions(+), 2 deletions(-) diff --git a/python/cudf_polars/cudf_polars/engine/core.py b/python/cudf_polars/cudf_polars/engine/core.py index 56bf051a05b..537fe3a970e 100644 --- a/python/cudf_polars/cudf_polars/engine/core.py +++ b/python/cudf_polars/cudf_polars/engine/core.py @@ -381,6 +381,7 @@ def __init__( # accept it. engine_options = dict(engine_options) allow_gpu_sharing = engine_options.pop("allow_gpu_sharing", False) + self._allow_gpu_sharing = allow_gpu_sharing super().__init__( executor="streaming", executor_options=executor_options, diff --git a/python/cudf_polars/cudf_polars/engine/dask.py b/python/cudf_polars/cudf_polars/engine/dask.py index 7abf14ca4d9..31b62d1689a 100644 --- a/python/cudf_polars/cudf_polars/engine/dask.py +++ b/python/cudf_polars/cudf_polars/engine/dask.py @@ -1297,7 +1297,11 @@ def _reset( "cluster": "dask", "dask_context": new_ctx, }, - engine_options={**engine_options, "memory_resource": None}, + engine_options={ + "allow_gpu_sharing": self._allow_gpu_sharing, + **engine_options, + "memory_resource": None, + }, exit_stack=self._exit_stack, ) diff --git a/python/cudf_polars/cudf_polars/engine/ray.py b/python/cudf_polars/cudf_polars/engine/ray.py index 50078956b0b..0a46b65663d 100644 --- a/python/cudf_polars/cudf_polars/engine/ray.py +++ b/python/cudf_polars/cudf_polars/engine/ray.py @@ -1127,7 +1127,10 @@ def _reset( "cluster": "ray", "ray_context": RayContext(self._rank_actors, self._quent_runtime), }, - engine_options=engine_options, + engine_options={ + "allow_gpu_sharing": self._allow_gpu_sharing, + **engine_options, + }, exit_stack=self._exit_stack, ) diff --git a/python/cudf_polars/cudf_polars/engine/spmd.py b/python/cudf_polars/cudf_polars/engine/spmd.py index fa79838ed6b..94485cbb4df 100644 --- a/python/cudf_polars/cudf_polars/engine/spmd.py +++ b/python/cudf_polars/cudf_polars/engine/spmd.py @@ -1116,6 +1116,7 @@ def _reset( ), }, engine_options={ + "allow_gpu_sharing": self._allow_gpu_sharing, **engine_options, "memory_resource": self._ctx.br().device_mr, }, From c5a52349d7c42a0dc83402bbd61b2306a5e53f7a Mon Sep 17 00:00:00 2001 From: Tom Augspurger Date: Fri, 16 Oct 2026 12:40:04 -0700 Subject: [PATCH 06/20] Fix Ray quent test: pass full engine_options in the in-test _reset The in-test _reset call that flushes Quent events before reading them was passing only executor_options, omitting engine_options. On a multi-rank RayEngine this meant allow_gpu_sharing defaulted to False and the UUID-collision guard fired. create_streaming_options always includes allow_gpu_sharing=True so that _reset calls made through configure_streaming_engine are safe; the bare flush call in the test body was the only site that bypassed this. Fix: yield (engine, base_options) from engine_with_quent_context so the test can pass the complete options to _reset, matching every other call in the suite. Co-Authored-By: Claude Sonnet 4.6 (1M context) --- python/cudf_polars/cudf_polars/engine/core.py | 1 - python/cudf_polars/cudf_polars/engine/dask.py | 6 +----- python/cudf_polars/cudf_polars/engine/ray.py | 5 +---- python/cudf_polars/cudf_polars/engine/spmd.py | 1 - .../tests/quent/test_quent_integration.py | 18 +++++++++++++----- 5 files changed, 15 insertions(+), 16 deletions(-) diff --git a/python/cudf_polars/cudf_polars/engine/core.py b/python/cudf_polars/cudf_polars/engine/core.py index 537fe3a970e..56bf051a05b 100644 --- a/python/cudf_polars/cudf_polars/engine/core.py +++ b/python/cudf_polars/cudf_polars/engine/core.py @@ -381,7 +381,6 @@ def __init__( # accept it. engine_options = dict(engine_options) allow_gpu_sharing = engine_options.pop("allow_gpu_sharing", False) - self._allow_gpu_sharing = allow_gpu_sharing super().__init__( executor="streaming", executor_options=executor_options, diff --git a/python/cudf_polars/cudf_polars/engine/dask.py b/python/cudf_polars/cudf_polars/engine/dask.py index 31b62d1689a..7abf14ca4d9 100644 --- a/python/cudf_polars/cudf_polars/engine/dask.py +++ b/python/cudf_polars/cudf_polars/engine/dask.py @@ -1297,11 +1297,7 @@ def _reset( "cluster": "dask", "dask_context": new_ctx, }, - engine_options={ - "allow_gpu_sharing": self._allow_gpu_sharing, - **engine_options, - "memory_resource": None, - }, + engine_options={**engine_options, "memory_resource": None}, exit_stack=self._exit_stack, ) diff --git a/python/cudf_polars/cudf_polars/engine/ray.py b/python/cudf_polars/cudf_polars/engine/ray.py index 0a46b65663d..50078956b0b 100644 --- a/python/cudf_polars/cudf_polars/engine/ray.py +++ b/python/cudf_polars/cudf_polars/engine/ray.py @@ -1127,10 +1127,7 @@ def _reset( "cluster": "ray", "ray_context": RayContext(self._rank_actors, self._quent_runtime), }, - engine_options={ - "allow_gpu_sharing": self._allow_gpu_sharing, - **engine_options, - }, + engine_options=engine_options, exit_stack=self._exit_stack, ) diff --git a/python/cudf_polars/cudf_polars/engine/spmd.py b/python/cudf_polars/cudf_polars/engine/spmd.py index 94485cbb4df..fa79838ed6b 100644 --- a/python/cudf_polars/cudf_polars/engine/spmd.py +++ b/python/cudf_polars/cudf_polars/engine/spmd.py @@ -1116,7 +1116,6 @@ def _reset( ), }, engine_options={ - "allow_gpu_sharing": self._allow_gpu_sharing, **engine_options, "memory_resource": self._ctx.br().device_mr, }, diff --git a/python/cudf_polars/tests/quent/test_quent_integration.py b/python/cudf_polars/tests/quent/test_quent_integration.py index 3453b4d0ede..1409f3bad20 100644 --- a/python/cudf_polars/tests/quent/test_quent_integration.py +++ b/python/cudf_polars/tests/quent/test_quent_integration.py @@ -37,7 +37,7 @@ def engine_with_quent_context( _unconfigured_engine: tuple[StreamingEngine, StreamingOptions], quent_context: QuentConfig, streaming_engine_factory: Callable[[StreamingOptions], StreamingEngine], -) -> Iterator[StreamingEngine]: +) -> Iterator[tuple[StreamingEngine, StreamingOptions]]: """ Yield a streaming engine configured with Quent tracing. """ @@ -55,7 +55,7 @@ def engine_with_quent_context( engine_options=base_options.to_engine_options(), ) try: - yield engine + yield engine, base_options finally: engine._reset( rapidsmpf_options=base_options.to_rapidsmpf_options(), @@ -165,12 +165,12 @@ def test_quent_controller_cleanup_continues_after_handle_failure() -> None: def test_quent_lifecycle( - engine_with_quent_context: StreamingEngine, + engine_with_quent_context: tuple[StreamingEngine, StreamingOptions], quent_context: QuentConfig, monkeypatch: pytest.MonkeyPatch, ) -> None: query = pl.LazyFrame({"x": [1, 2, 3]}).filter(pl.col("x") > 1) - engine = engine_with_quent_context + engine, base_options = engine_with_quent_context assert engine.config["executor_options"]["quent_context"] == quent_context current_context = dataclasses.replace( quent_context, @@ -220,7 +220,15 @@ def test_quent_lifecycle( # The fixture teardown calls _reset(..., quent_context=None) again, which is a no-op. quent_output_root = engine._quent_output_root assert quent_output_root is not None - engine._reset(executor_options={"quent_context": None, "fallback_mode": "silent"}) + engine._reset( + rapidsmpf_options=base_options.to_rapidsmpf_options(), + executor_options={ + **base_options.to_executor_options(), + "quent_context": None, + "fallback_mode": "silent", + }, + engine_options=base_options.to_engine_options(), + ) events = _stored_events(quent_output_root) engine_events = _of_type(events, "Engine") From 9da5f3f0b7359d959df9249334b1f785b50b4673 Mon Sep 17 00:00:00 2001 From: Tom Augspurger Date: Fri, 16 Oct 2026 14:19:17 -0700 Subject: [PATCH 07/20] Fix stale Quent closure capture in Ray/Dask shutdown callbacks MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The exit-stack callbacks registered in RayEngine.__init__ and DaskEngine.__init__ captured local variables quent_runtime and quent_collector at construction time. After _reset changed Quent state, those closures became stale: - Reset None→Quent: callback passed None for both; shutdown skipped worker cleanup and never closed the replacement runtime (leak). - Reset Quent→None: _quent_close already closed everything, but the callback tried to close the originals again (double-close). - _shutdown_ray gated actor.close_quent on quent_collector is not None, so a reset that enabled Quent without a collector (disabled→enabled after runtime took ownership) would also skip worker-side cleanup. Fix: - Read self._quent_runtime and self._quent_collector at callback fire time instead of capturing init-time locals. - Track self._quent_collector for the brief init window between start_collector and QuentControllerRuntime.create (the error path); clear it once the runtime takes ownership. - Clear self._quent_collector in shutdown() alongside _quent_runtime. - Fix _shutdown_ray gate: quent_runtime is not None or quent_collector is not None (either means actors have live Quent worker runtimes). Co-Authored-By: Claude Sonnet 4.6 (1M context) --- python/cudf_polars/cudf_polars/engine/dask.py | 8 ++++++-- python/cudf_polars/cudf_polars/engine/ray.py | 10 +++++++--- 2 files changed, 13 insertions(+), 5 deletions(-) diff --git a/python/cudf_polars/cudf_polars/engine/dask.py b/python/cudf_polars/cudf_polars/engine/dask.py index 7abf14ca4d9..b2ce2bab058 100644 --- a/python/cudf_polars/cudf_polars/engine/dask.py +++ b/python/cudf_polars/cudf_polars/engine/dask.py @@ -1062,6 +1062,7 @@ def __init__( quent_context = resolve_quent_context(executor_options) executor_options["quent_context"] = quent_context self._quent_runtime = None + self._quent_collector: quent_bindings.Collector | None = None if bootstrap.is_running_with_rrun(): raise RuntimeError( @@ -1096,8 +1097,8 @@ def __init__( lambda: _shutdown_dask( dask_client, rapidsmpf_id, - quent_runtime, - quent_collector, + self._quent_runtime, + self._quent_collector, owned_client, owned_cluster, ) @@ -1155,6 +1156,7 @@ def __init__( quent_collector = cudf_polars.quent._runtime.start_collector( quent_context.run_root ) + self._quent_collector = quent_collector quent_collector_address = quent_collector.address else: quent_collector_address = None @@ -1195,6 +1197,7 @@ def __init__( collector=quent_collector, ) self._quent_runtime = quent_runtime + self._quent_collector = None # runtime now owns the collector dask_ctx = DaskContext( client=dask_client, @@ -1444,6 +1447,7 @@ def shutdown(self) -> None: return # already shut down self._dask_context = None self._quent_runtime = None + self._quent_collector = None super().shutdown() def _run(self, func: Callable[..., T], *args: Any, **kwargs: Any) -> list[T]: diff --git a/python/cudf_polars/cudf_polars/engine/ray.py b/python/cudf_polars/cudf_polars/engine/ray.py index 50078956b0b..af189ab066b 100644 --- a/python/cudf_polars/cudf_polars/engine/ray.py +++ b/python/cudf_polars/cudf_polars/engine/ray.py @@ -776,7 +776,7 @@ def _shutdown_ray( """Release Ray resources in dependency order.""" ray_running = ray.is_initialized() steps: list[Callable[[], object]] = [] - if ray_running and quent_collector is not None: + if ray_running and (quent_runtime is not None or quent_collector is not None): steps.append( lambda: ray.get([actor.close_quent.remote() for actor in rank_actors]) ) @@ -916,6 +916,7 @@ def __init__( else: engine_id = uuid.uuid4() self._quent_runtime = None + self._quent_collector: quent_bindings.Collector | None = None # This engine's store uid, used to key its partitions in each actor's process rank-local store. self._store_uid = uuid.uuid4().hex @@ -946,8 +947,8 @@ def __init__( exit_stack.callback( lambda: _shutdown_ray( rank_actors, - quent_runtime, - quent_collector, + self._quent_runtime, + self._quent_collector, shutdown_ray=owns_ray, ) ) @@ -1012,6 +1013,7 @@ def __init__( quent_collector = cudf_polars.quent._runtime.start_collector( quent_context.run_root ) + self._quent_collector = quent_collector collector_address = quent_collector.address else: quent_collector = None @@ -1038,6 +1040,7 @@ def __init__( collector=quent_collector, ) self._quent_runtime = quent_runtime + self._quent_collector = None # runtime now owns the collector self._rank_actors: list[ActorHandle[RankActor]] | None = rank_actors super().__init__( @@ -1301,6 +1304,7 @@ def shutdown(self) -> None: return # already shut down; idempotent self._rank_actors = None self._quent_runtime = None + self._quent_collector = None super().shutdown() def _run(self, func: Callable[..., T], *args: Any, **kwargs: Any) -> list[T]: From 625885f2325dd72886f8a793fbbf7bddb4f8382a Mon Sep 17 00:00:00 2001 From: Tom Augspurger Date: Mon, 19 Oct 2026 07:25:58 -0700 Subject: [PATCH 08/20] Ensure shutdown in dask/ray engine --- python/cudf_polars/cudf_polars/engine/dask.py | 8 +++++--- python/cudf_polars/cudf_polars/engine/ray.py | 8 +++++--- 2 files changed, 10 insertions(+), 6 deletions(-) diff --git a/python/cudf_polars/cudf_polars/engine/dask.py b/python/cudf_polars/cudf_polars/engine/dask.py index b2ce2bab058..81f1e8d90b2 100644 --- a/python/cudf_polars/cudf_polars/engine/dask.py +++ b/python/cudf_polars/cudf_polars/engine/dask.py @@ -1446,9 +1446,11 @@ def shutdown(self) -> None: if self._dask_context is None: return # already shut down self._dask_context = None - self._quent_runtime = None - self._quent_collector = None - super().shutdown() + try: + super().shutdown() + finally: + self._quent_runtime = None + self._quent_collector = None def _run(self, func: Callable[..., T], *args: Any, **kwargs: Any) -> list[T]: return list(self._run_by_rank(_run_with_rank, func, *args, **kwargs).values()) diff --git a/python/cudf_polars/cudf_polars/engine/ray.py b/python/cudf_polars/cudf_polars/engine/ray.py index af189ab066b..02641d3af1b 100644 --- a/python/cudf_polars/cudf_polars/engine/ray.py +++ b/python/cudf_polars/cudf_polars/engine/ray.py @@ -1303,9 +1303,11 @@ def shutdown(self) -> None: if self._rank_actors is None: return # already shut down; idempotent self._rank_actors = None - self._quent_runtime = None - self._quent_collector = None - super().shutdown() + try: + super().shutdown() + finally: + self._quent_runtime = None + self._quent_collector = None def _run(self, func: Callable[..., T], *args: Any, **kwargs: Any) -> list[T]: return list( From 257cd056052ca2e8ec7149cf8c2d1ae09e0d07a6 Mon Sep 17 00:00:00 2001 From: Tom Augspurger Date: Mon, 19 Oct 2026 07:45:33 -0700 Subject: [PATCH 09/20] bigger refactor --- python/cudf_polars/cudf_polars/engine/core.py | 14 +- python/cudf_polars/cudf_polars/engine/dask.py | 159 ++++++------------ python/cudf_polars/cudf_polars/engine/ray.py | 136 +++++---------- python/cudf_polars/cudf_polars/engine/spmd.py | 120 +++++-------- .../cudf_polars/cudf_polars/quent/_runtime.py | 60 ++++--- .../tests/quent/test_quent_integration.py | 25 +++ .../cudf_polars/tests/streaming/test_dask.py | 14 +- 7 files changed, 209 insertions(+), 319 deletions(-) diff --git a/python/cudf_polars/cudf_polars/engine/core.py b/python/cudf_polars/cudf_polars/engine/core.py index 56bf051a05b..f923b3b4381 100644 --- a/python/cudf_polars/cudf_polars/engine/core.py +++ b/python/cudf_polars/cudf_polars/engine/core.py @@ -597,11 +597,11 @@ def _update_quent( and ``_quent_output_root``. Subclasses implement :meth:`_quent_close` and :meth:`_quent_open` for the backend-specific work. """ - if new_quent_context == existing_quent_context: - return - self._quent_close() - if new_quent_context is not None: - self._quent_open(new_quent_context) + if new_quent_context != existing_quent_context: + if existing_quent_context is not None: + self._quent_close() + if new_quent_context is not None: + new_quent_context = self._quent_open(new_quent_context) self._quent_output_root = ( new_quent_context.run_root if new_quent_context is not None else None ) @@ -613,8 +613,8 @@ def _quent_close(self) -> None: f"{type(self).__name__} does not implement _quent_close" ) - def _quent_open(self, quent_context: QuentConfig) -> None: - """Start fresh Quent runtimes for ``quent_context`` (workers + controller).""" + def _quent_open(self, quent_context: QuentConfig) -> QuentConfig: + """Start fresh Quent runtimes and return the effective configuration.""" raise NotImplementedError( # pragma: no cover f"{type(self).__name__} does not implement _quent_open" ) diff --git a/python/cudf_polars/cudf_polars/engine/dask.py b/python/cudf_polars/cudf_polars/engine/dask.py index 81f1e8d90b2..5f4d330d1e8 100644 --- a/python/cudf_polars/cudf_polars/engine/dask.py +++ b/python/cudf_polars/cudf_polars/engine/dask.py @@ -29,8 +29,6 @@ from rapidsmpf.statistics import Statistics from rapidsmpf.streaming.core.context import Context -import cudf_polars.quent -import cudf_polars.quent._runtime from cudf_polars.engine import persisted_result, rank_local_store from cudf_polars.engine.core import ( ClusterInfo, @@ -69,13 +67,13 @@ if TYPE_CHECKING: from collections.abc import Callable - import cudf_polars_quent as quent_bindings import kvikio from cudf_streaming.channel_metadata import ChannelMetadata from rapidsmpf.communicator.communicator import Communicator from rapidsmpf.rmm_resource_adaptor import RmmResourceAdaptor + import cudf_polars.quent from cudf_polars.dsl.ir import IR from cudf_polars.engine.core import T from cudf_polars.engine.options import StreamingOptions @@ -266,7 +264,6 @@ def _setup_root( hardware_binding: HardwareBindingPolicy, memory_resource_config: MemoryResourceConfig | None, dask_worker: distributed.Worker | None = None, - worker_id: uuid.UUID, ) -> bytes: """ Initialize the root rank on one Dask worker. @@ -291,8 +288,6 @@ def _setup_root( :class:`rmm.mr.CudaAsyncMemoryResource`. dask_worker Injected by ``distributed`` when called via :meth:`distributed.Client.run`. - worker_id - Unique identifier for this worker. Returns ------- @@ -334,8 +329,6 @@ def _setup_worker( uid: str, hardware_binding: HardwareBindingPolicy, memory_resource_config: MemoryResourceConfig | None, - worker_ids: list[uuid.UUID], - quent_context: cudf_polars.quent.QuentConfig | None, num_py_executors: int, kvikio_nthreads: int | None, kvikio_statistics: bool, @@ -345,7 +338,6 @@ def _setup_worker( kvikio_reactor_count: int, kvikio_reactor_dispatch: kvikio.RemoteReactorDispatch, kvikio_request_ceiling: int, - quent_collector_address: str | None, dask_worker: distributed.Worker | None = None, ) -> None: """ @@ -370,11 +362,6 @@ def _setup_worker( memory_resource_config Optional RMM memory resource configuration. If ``None``, defaults to :class:`rmm.mr.CudaAsyncMemoryResource`. - worker_ids - List of Quent worker UUIDs indexed by rank. Each worker picks - its own ID after the barrier using ``comm.rank``. - quent_context - Serializable Quent configuration, or ``None`` when tracing is disabled. dask_worker Injected by ``distributed`` when called via :meth:`distributed.Client.run`. num_py_executors @@ -396,8 +383,6 @@ def _setup_worker( ``MULTI_POLL`` reactor dispatch policy to configure on this worker process. kvikio_request_ceiling ``MULTI_POLL`` concurrent-request ceiling to configure on this worker process. - quent_collector_address - Address of the Quent collector, if Quent is enabled. """ assert dask_worker is not None options = Options.deserialize(rapidsmpf_options_as_bytes) @@ -447,7 +432,6 @@ def _setup_worker( ) barrier(comm) - worker_id = worker_ids[comm.rank] ctx = Context.from_options(comm.logger, base_mr, options, statistics) # Set the current RMM device resource so all temporary allocations # in libcudf also use the same memory resource. @@ -458,45 +442,19 @@ def _setup_worker( thread_name_prefix="dask-executor", ) - if quent_collector_address is not None: - assert quent_context is not None - quent_worker_runtime = QuentWorkerRuntime.create( - quent_context, - quent_collector_address, - worker_id=worker_id, - rank=comm.rank, - nranks=comm.nranks, - instance_name=f"rank-{comm.rank}", - ) - else: - quent_worker_runtime = None - mp_ctx = _WorkerContext( comm=comm, ctx=ctx, py_executor=py_executor, base_mr=base_mr, mr=mr, - quent_worker_runtime=quent_worker_runtime, + quent_worker_runtime=None, statistics=statistics, kvikio_monitor=make_kvikio_monitor(enabled=kvikio_statistics), ) setattr(dask_worker, attr, mp_ctx) -def _close_quent_worker( - *, uid: str, dask_worker: distributed.Worker | None = None -) -> None: - """Close one worker's Quent collector client.""" - assert dask_worker is not None - mp_ctx: _WorkerContext | None = getattr( - dask_worker, f"_cudf_polars_mp_context_{uid}", None - ) - if mp_ctx is not None and mp_ctx.quent_worker_runtime is not None: - mp_ctx.quent_worker_runtime.close() - mp_ctx.quent_worker_runtime = None - - def _update_quent_on_worker( quent_context: cudf_polars.quent.QuentConfig | None, collector_address: str | None, @@ -573,25 +531,14 @@ def _teardown_worker( def _shutdown_dask( client: distributed.Client | None, uid: str, - quent_runtime: QuentControllerRuntime | None, - quent_collector: quent_bindings.Collector | None, + quent_close: Callable[[], object] | None, owned_client: distributed.Client | None, owned_cluster: distributed.SpecCluster | None, ) -> None: """Release Dask resources in dependency order.""" steps: list[Callable[[], object]] = [] - if client is not None and ( - quent_runtime is not None or quent_collector is not None - ): - steps.append( - lambda: client.run(functools.partial(_close_quent_worker, uid=uid)) - ) - if quent_runtime is not None: - steps.append(quent_runtime.close) - elif ( - quent_collector is not None - ): # pragma: no-cover; runs on the worker, not client - steps.append(quent_collector.close) + if quent_close is not None: + steps.append(quent_close) if client is not None: steps.append(lambda: client.run(functools.partial(_teardown_worker, uid=uid))) if owned_client is not None: @@ -1062,7 +1009,6 @@ def __init__( quent_context = resolve_quent_context(executor_options) executor_options["quent_context"] = quent_context self._quent_runtime = None - self._quent_collector: quent_bindings.Collector | None = None if bootstrap.is_running_with_rrun(): raise RuntimeError( @@ -1091,14 +1037,11 @@ def __init__( exit_stack = contextlib.ExitStack() owned_cluster: distributed.SpecCluster | None = None owned_client: distributed.Client | None = None - quent_collector: quent_bindings.Collector | None = None - quent_runtime: QuentControllerRuntime | None = None exit_stack.callback( lambda: _shutdown_dask( dask_client, rapidsmpf_id, - self._quent_runtime, - self._quent_collector, + (self._quent_close if self._quent_runtime is not None else None), owned_client, owned_cluster, ) @@ -1135,8 +1078,6 @@ def __init__( nranks = len(worker_addresses) root_worker = worker_addresses[0] - worker_ids = [uuid.uuid4() for _ in range(nranks)] - # Phase 1: initialize root communicator on one worker. root_result = dask_client.run( functools.partial( @@ -1144,7 +1085,6 @@ def __init__( uid=rapidsmpf_id, hardware_binding=hw_binding, memory_resource_config=mr_config, - worker_id=worker_ids[0], ), nranks, rapidsmpf_options_as_bytes, @@ -1152,15 +1092,6 @@ def __init__( ) root_ucxx_address_as_bytes = root_result[root_worker] - if quent_context is not None: - quent_collector = cudf_polars.quent._runtime.start_collector( - quent_context.run_root - ) - self._quent_collector = quent_collector - quent_collector_address = quent_collector.address - else: - quent_collector_address = None - # Phase 2: complete bootstrap on all workers concurrently. # All workers call barrier() so they must all run simultaneously. dask_client.run( @@ -1169,13 +1100,10 @@ def __init__( uid=rapidsmpf_id, hardware_binding=hw_binding, memory_resource_config=mr_config, - worker_ids=worker_ids, - quent_context=quent_context, ), root_ucxx_address_as_bytes, nranks, rapidsmpf_options_as_bytes, - quent_collector_address=quent_collector_address, num_py_executors=executor_options.get("num_py_executors", 8), kvikio_nthreads=executor_options["kvikio_nthreads"], kvikio_statistics=executor_options["kvikio_statistics"], @@ -1188,25 +1116,24 @@ def __init__( kvikio_reactor_dispatch=executor_options["kvikio_reactor_dispatch"], kvikio_request_ceiling=executor_options["kvikio_request_ceiling"], ) - if quent_context is not None: - assert quent_collector_address is not None - quent_runtime = QuentControllerRuntime.create( - quent_context, - quent_collector_address, - backend="dask", - collector=quent_collector, - ) - self._quent_runtime = quent_runtime - self._quent_collector = None # runtime now owns the collector dask_ctx = DaskContext( client=dask_client, rapidsmpf_id=rapidsmpf_id, - quent_controller_runtime=quent_runtime, + quent_controller_runtime=None, owned_client=owned_client, owned_cluster=owned_cluster, ) self._dask_context: DaskContext | None = dask_ctx + self._update_quent(quent_context, None, executor_options) + dask_ctx = DaskContext( + client=dask_client, + rapidsmpf_id=rapidsmpf_id, + quent_controller_runtime=self._quent_runtime, + owned_client=owned_client, + owned_cluster=owned_cluster, + ) + self._dask_context = dask_ctx super().__init__( nranks=nranks, executor_options={ @@ -1307,33 +1234,44 @@ def _reset( def _quent_close(self) -> None: """Clear Quent worker runtimes on all Dask workers, then close the controller.""" ctx = self._dask_ctx - # Workers first so the collector can drain before we close it. - ctx.client.run( - functools.partial(_update_quent_on_worker, None, None, uid=ctx.rapidsmpf_id) - ) + steps: list[Callable[[], object]] = [ + # Workers first so the collector can drain before we close it. + lambda: ctx.client.run( + functools.partial( + _update_quent_on_worker, None, None, uid=ctx.rapidsmpf_id + ) + ) + ] if self._quent_runtime is not None: - self._quent_runtime.close() + steps.append(self._quent_runtime.close) + try: + run_cleanup_steps("Dask Quent shutdown failed", *steps) + finally: self._quent_runtime = None - def _quent_open(self, quent_context: cudf_polars.quent.QuentConfig) -> None: + def _quent_open( + self, quent_context: cudf_polars.quent.QuentConfig + ) -> cudf_polars.quent.QuentConfig: """Start a new Quent collector, connect workers, create controller.""" ctx = self._dask_ctx - quent_collector = cudf_polars.quent._runtime.start_collector( - quent_context.run_root - ) - collector_address = quent_collector.address - assert collector_address is not None - ctx.client.run( - functools.partial( - _update_quent_on_worker, - quent_context, - collector_address, - uid=ctx.rapidsmpf_id, - ) - ) self._quent_runtime = QuentControllerRuntime.create( - quent_context, collector_address, backend="dask", collector=quent_collector + quent_context, + backend="dask", ) + try: + ctx.client.run( + functools.partial( + _update_quent_on_worker, + quent_context, + self._quent_runtime.collector_address, + uid=ctx.rapidsmpf_id, + ) + ) + except BaseException: + with contextlib.suppress(Exception): + self._quent_close() + raise + return quent_context @classmethod def from_options( @@ -1445,12 +1383,11 @@ def shutdown(self) -> None: """ if self._dask_context is None: return # already shut down - self._dask_context = None try: super().shutdown() finally: + self._dask_context = None self._quent_runtime = None - self._quent_collector = None def _run(self, func: Callable[..., T], *args: Any, **kwargs: Any) -> list[T]: return list(self._run_by_rank(_run_with_rank, func, *args, **kwargs).values()) diff --git a/python/cudf_polars/cudf_polars/engine/ray.py b/python/cudf_polars/cudf_polars/engine/ray.py index 02641d3af1b..98d0929c469 100644 --- a/python/cudf_polars/cudf_polars/engine/ray.py +++ b/python/cudf_polars/cudf_polars/engine/ray.py @@ -24,8 +24,6 @@ from rapidsmpf.statistics import Statistics from rapidsmpf.streaming.core.context import Context -import cudf_polars.quent -import cudf_polars.quent._runtime from cudf_polars.engine import persisted_result, rank_local_store from cudf_polars.engine.core import ( ClusterInfo, @@ -64,7 +62,6 @@ if TYPE_CHECKING: from collections.abc import Callable - import cudf_polars_quent as quent_bindings import kvikio from ray import ObjectRef from ray.actor import ActorHandle @@ -73,6 +70,7 @@ from rapidsmpf.communicator.communicator import Communicator from rapidsmpf.rmm_resource_adaptor import RmmResourceAdaptor + import cudf_polars.quent from cudf_polars.dsl.ir import IR from cudf_polars.engine.core import T from cudf_polars.engine.options import StreamingOptions @@ -267,8 +265,6 @@ def __init__( hardware_binding: HardwareBindingPolicy, memory_resource_config: MemoryResourceConfig | None, worker_id: uuid.UUID, - engine_id: uuid.UUID, - quent_enabled: bool, ) -> None: bind_to_gpu(hardware_binding) configure_kvikio( @@ -301,9 +297,7 @@ def __init__( ) self._comm: Communicator | None = None self._ctx: Context | None = None - self._quent_enabled = quent_enabled self._quent_worker_runtime: QuentWorkerRuntime | None = None - self._quent_engine_id = engine_id self._worker_id = worker_id def setup_root(self) -> bytes: @@ -330,8 +324,6 @@ def setup_root(self) -> bytes: def setup_worker( self, root_ucxx_address_as_bytes: bytes, - collector_address: str | None, - quent_context: cudf_polars.quent.QuentConfig | None, ) -> None: """ Complete communicator bootstrap and create the streaming context. @@ -345,10 +337,6 @@ def setup_worker( ---------- root_ucxx_address_as_bytes Serialized UCXX root address returned by :meth:`setup_root`. - collector_address - Address of the Quent collector, or ``None`` when tracing is disabled. - quent_context - Serializable Quent configuration, or ``None`` when tracing is disabled. """ if self._comm is None: root_ucxx_address = ucx_api.UCXAddress.create_from_buffer( @@ -376,25 +364,6 @@ def setup_worker( # libcudf temporary allocations use the same memory resource. self._mr = self._ctx.br().device_mr_adaptor() rmm.mr.set_current_device_resource(self._mr) - if collector_address is not None: - assert self._comm is not None - if not self._quent_enabled: - return - assert quent_context is not None - self._quent_worker_runtime = QuentWorkerRuntime.create( - quent_context, - collector_address, - worker_id=self._worker_id, - rank=self._comm.rank, - nranks=self._nranks, - instance_name=f"RankActor-{self._worker_id.hex[:8]}", - ) - - def close_quent(self) -> None: - """Close this rank's collector client after its Worker exit.""" - if self._quent_worker_runtime is None: - return - self._quent_worker_runtime.close() def update_quent( self, @@ -768,22 +737,15 @@ def _wait_for_actor_shutdown(ref: ObjectRef[Any]) -> None: def _shutdown_ray( rank_actors: list[ActorHandle[RankActor]], - quent_runtime: QuentControllerRuntime | None, - quent_collector: quent_bindings.Collector | None, + quent_close: Callable[[], object] | None, *, shutdown_ray: bool, ) -> None: """Release Ray resources in dependency order.""" ray_running = ray.is_initialized() steps: list[Callable[[], object]] = [] - if ray_running and (quent_runtime is not None or quent_collector is not None): - steps.append( - lambda: ray.get([actor.close_quent.remote() for actor in rank_actors]) - ) - if quent_runtime is not None: - steps.append(quent_runtime.close) - elif quent_collector is not None: - steps.append(quent_collector.close) + if ray_running and quent_close is not None: + steps.append(quent_close) if ray_running: @@ -911,12 +873,7 @@ def __init__( quent_context = resolve_quent_context(executor_options) executor_options["quent_context"] = quent_context - if quent_context is not None: - engine_id = quent_context.engine_id - else: - engine_id = uuid.uuid4() self._quent_runtime = None - self._quent_collector: quent_bindings.Collector | None = None # This engine's store uid, used to key its partitions in each actor's process rank-local store. self._store_uid = uuid.uuid4().hex @@ -942,13 +899,10 @@ def __init__( exit_stack = contextlib.ExitStack() owns_ray = not ray.is_initialized() rank_actors: list[ActorHandle[RankActor]] = [] - quent_collector: quent_bindings.Collector | None = None - quent_runtime: QuentControllerRuntime | None = None exit_stack.callback( lambda: _shutdown_ray( rank_actors, - self._quent_runtime, - self._quent_collector, + (self._quent_close if self._quent_runtime is not None else None), shutdown_ray=owns_ray, ) ) @@ -1001,54 +955,30 @@ def __init__( hardware_binding=hw_binding, memory_resource_config=mr_config, worker_id=worker_id, - engine_id=engine_id, - quent_enabled=quent_context is not None, ) ) root_ucxx_address_as_bytes = cast( "bytes", ray.get(rank_actors[0].setup_root.remote()) ) - if quent_context is not None: - quent_collector = cudf_polars.quent._runtime.start_collector( - quent_context.run_root - ) - self._quent_collector = quent_collector - collector_address = quent_collector.address - else: - quent_collector = None - collector_address = None # Call setup_worker on all actors concurrently, including the root. # The root skips communicator creation and proceeds directly to the barrier. # Non-root actors create their communicators and then join the barrier. ray.get( [ - rank.setup_worker.remote( - root_ucxx_address_as_bytes, - collector_address, - quent_context, - ) + rank.setup_worker.remote(root_ucxx_address_as_bytes) for rank in rank_actors ] ) - if quent_context is not None: - assert collector_address is not None - quent_runtime = QuentControllerRuntime.create( - quent_context, - collector_address, - backend="ray", - collector=quent_collector, - ) - self._quent_runtime = quent_runtime - self._quent_collector = None # runtime now owns the collector self._rank_actors: list[ActorHandle[RankActor]] | None = rank_actors + self._update_quent(quent_context, None, executor_options) super().__init__( nranks=nranks, executor_options={ **executor_options, "cluster": "ray", - "ray_context": RayContext(rank_actors, quent_runtime), + "ray_context": RayContext(rank_actors, self._quent_runtime), }, engine_options=engine_options, exit_stack=exit_stack, @@ -1137,29 +1067,44 @@ def _reset( def _quent_close(self) -> None: """Clear Quent worker runtimes on all actors, then close the controller.""" assert self._rank_actors is not None - # Workers first so the collector can drain before we close it. - ray.get([rank.update_quent.remote(None, None) for rank in self._rank_actors]) + rank_actors = self._rank_actors + steps: list[Callable[[], object]] = [ + # Workers first so the collector can drain before we close it. + lambda: ray.get( + [rank.update_quent.remote(None, None) for rank in rank_actors] + ) + ] if self._quent_runtime is not None: - self._quent_runtime.close() + steps.append(self._quent_runtime.close) + try: + run_cleanup_steps("Ray Quent shutdown failed", *steps) + finally: self._quent_runtime = None - def _quent_open(self, quent_context: cudf_polars.quent.QuentConfig) -> None: + def _quent_open( + self, quent_context: cudf_polars.quent.QuentConfig + ) -> cudf_polars.quent.QuentConfig: """Start a new Quent collector, connect actors, create controller.""" assert self._rank_actors is not None - quent_collector = cudf_polars.quent._runtime.start_collector( - quent_context.run_root - ) - collector_address = quent_collector.address - assert collector_address is not None - ray.get( - [ - rank.update_quent.remote(quent_context, collector_address) - for rank in self._rank_actors - ] - ) self._quent_runtime = QuentControllerRuntime.create( - quent_context, collector_address, backend="ray", collector=quent_collector + quent_context, + backend="ray", ) + try: + ray.get( + [ + rank.update_quent.remote( + quent_context, + self._quent_runtime.collector_address, + ) + for rank in self._rank_actors + ] + ) + except BaseException: + with contextlib.suppress(Exception): + self._quent_close() + raise + return quent_context @classmethod def from_options( @@ -1302,12 +1247,11 @@ def shutdown(self) -> None: """ if self._rank_actors is None: return # already shut down; idempotent - self._rank_actors = None try: super().shutdown() finally: + self._rank_actors = None self._quent_runtime = None - self._quent_collector = None def _run(self, func: Callable[..., T], *args: Any, **kwargs: Any) -> list[T]: return list( diff --git a/python/cudf_polars/cudf_polars/engine/spmd.py b/python/cudf_polars/cudf_polars/engine/spmd.py index fa79838ed6b..4a2d8c87ba9 100644 --- a/python/cudf_polars/cudf_polars/engine/spmd.py +++ b/python/cudf_polars/cudf_polars/engine/spmd.py @@ -42,7 +42,6 @@ from rapidsmpf.streaming.core.context import Context import cudf_polars.quent -import cudf_polars.quent._runtime from cudf_polars.containers import DataFrame, DataType from cudf_polars.containers.dataframe import categoricals_to_physical from cudf_polars.engine import persisted_result, rank_local_store @@ -579,7 +578,7 @@ def __init__( quent_context = resolve_quent_context(executor_options) executor_options["quent_context"] = quent_context self._quent_runtime = None - self._quent_worker_runtime = None + self._quent_worker_runtime: QuentWorkerRuntime | None = None check_reserved_keys(executor_options, engine_options) hw_binding = cast( @@ -650,51 +649,14 @@ def __init__( exit_stack.enter_context(set_memory_resource(self._mr)) exit_stack.callback(self._cleanup_ctx) - if quent_context is not None: - collector_cleanup: contextlib.ExitStack | None = None - if comm.rank == 0: - quent_collector = cudf_polars.quent._runtime.start_collector( - quent_context.run_root - ) - collector_cleanup = contextlib.ExitStack() - collector_cleanup.callback(quent_collector.close) - exit_stack.enter_context(collector_cleanup) - collector_address = quent_collector.address - else: - quent_collector = None - collector_address = "" - quent_context, collector_address = synchronize_quent_configuration( - comm=comm, - context=self._ctx, - quent_config=quent_context, - collector_address=collector_address, - ) - executor_options["quent_context"] = quent_context - self._quent_worker_id = uuid.uuid4() - if comm.rank == 0: - self._quent_runtime = QuentControllerRuntime.create( - quent_context, - collector_address, - backend="spmd", - collector=quent_collector, - ) - exit_stack.callback(self._close_quent_controller) - self._quent_worker_runtime = QuentWorkerRuntime.create( - quent_context, - collector_address, - worker_id=self._quent_worker_id, - rank=comm.rank, - nranks=comm.nranks, - instance_name=f"rank-{comm.rank}", - ) - exit_stack.callback(self._close_quent_worker) - if collector_cleanup is not None: - # The runtime now owns the collector. - collector_cleanup.pop_all() - engine_id = quent_context.engine_id - else: - engine_id = uuid.uuid4() - self._quent_worker_id = uuid.uuid4() + self._quent_worker_id = uuid.uuid4() + self._update_quent(quent_context, None, executor_options) + effective_quent_context = executor_options["quent_context"] + engine_id = ( + effective_quent_context.engine_id + if effective_quent_context is not None + else uuid.uuid4() + ) # Register after `_cleanup_ctx` so on teardown (LIFO) the # executor shuts down first. `wait=True` is safe because @@ -782,46 +744,48 @@ def _quent_close(self) -> None: if comm.nranks > 1: barrier(comm) - def _quent_open(self, quent_context: cudf_polars.quent.QuentConfig) -> None: + def _quent_open( + self, quent_context: cudf_polars.quent.QuentConfig + ) -> cudf_polars.quent.QuentConfig: """Start SPMD Quent runtimes for ``quent_context``.""" assert self._comm is not None assert self._ctx is not None comm = self._comm collector_address: str = "" - if comm.rank == 0: - quent_collector = cudf_polars.quent._runtime.start_collector( - quent_context.run_root + try: + if comm.rank == 0: + self._quent_runtime = QuentControllerRuntime.create( + quent_context, + backend="spmd", + ) + collector_address = self._quent_runtime.collector_address + quent_context, collector_address = synchronize_quent_configuration( + comm=comm, + context=self._ctx, + quent_config=quent_context, + collector_address=collector_address, ) - _collector_cleanup = contextlib.ExitStack() - _collector_cleanup.callback(quent_collector.close) - collector_address = quent_collector.address - quent_context, collector_address = synchronize_quent_configuration( - comm=comm, - context=self._ctx, - quent_config=quent_context, - collector_address=collector_address, - ) - if comm.rank == 0: - self._quent_runtime = QuentControllerRuntime.create( + self._quent_worker_runtime = QuentWorkerRuntime.create( quent_context, collector_address, - backend="spmd", - collector=quent_collector, + worker_id=self._quent_worker_id, + rank=comm.rank, + nranks=comm.nranks, + instance_name=f"rank-{comm.rank}", ) - _collector_cleanup.pop_all() - self._quent_worker_runtime = QuentWorkerRuntime.create( - quent_context, - collector_address, - worker_id=self._quent_worker_id, - rank=comm.rank, - nranks=comm.nranks, - instance_name=f"rank-{comm.rank}", - ) + except BaseException: + steps: list[Callable[[], object]] = [ + self._close_quent_worker, + self._close_quent_controller, + ] + run_cleanup_steps("SPMD Quent startup cleanup failed", *steps) + raise + return quent_context def _shutdown_spmd(self) -> None: """Run collective cleanup before local resources unwind.""" steps: list[Callable[[], object]] = [self._drop_persisted] - if self._quent_worker_runtime is not None: + if self._quent_worker_runtime is not None or self._quent_runtime is not None: steps.append(self._quent_close) run_cleanup_steps("SPMD engine shutdown failed", *steps) @@ -1055,8 +1019,9 @@ def _reset( self.rapidsmpf_options = rapidsmpf_options # Tear down Quent runtimes before the Context rebuild if they're changing. - if quent_changing: + if quent_changing and existing_quent_context is not None: self._quent_close() + existing_quent_context = None # Collective: synchronize all ranks before tearing down the Context. if self._comm.nranks > 1: @@ -1088,9 +1053,10 @@ def _reset( # Start new Quent runtimes after the Context is rebuilt. self._update_quent(new_quent_context, existing_quent_context, executor_options) + effective_quent_context = executor_options["quent_context"] engine_id = ( - new_quent_context.engine_id - if new_quent_context is not None + effective_quent_context.engine_id + if effective_quent_context is not None else uuid.uuid4() ) diff --git a/python/cudf_polars/cudf_polars/quent/_runtime.py b/python/cudf_polars/cudf_polars/quent/_runtime.py index f3c9680e71d..dfe753940cd 100644 --- a/python/cudf_polars/cudf_polars/quent/_runtime.py +++ b/python/cudf_polars/cudf_polars/quent/_runtime.py @@ -121,31 +121,49 @@ class QuentControllerRuntime: def create( cls, config: QuentConfig, - collector_address: str, *, backend: str, - collector: quent_bindings.Collector | None = None, ) -> QuentControllerRuntime: - """Create a controller runtime and emit its Engine init event.""" - runtime = cls( - config=config, - session=QuentSession(collector_address), - collector=collector, - ) - runtime._engine_handle = ( - runtime.session.binding_context.engine_observer() - .handle(config.engine_id) - .init( - instance_name=f"cudf-polars-{str(config.engine_id)[:8]}", - implementation={ - "name": config.implementation_name, - "version": config.implementation_version, - "backend": backend, - "custom_attributes": {"backend": backend}, - }, + """Start the Collector and create its controller runtime.""" + with contextlib.ExitStack() as cleanup: + collector = start_collector(config.run_root) + cleanup.callback(lambda: collector.close(timeout=10.0)) + collector_address = collector.address + if collector_address is None: + raise RuntimeError("Quent Collector did not advertise an address") + session = QuentSession(collector_address) + cleanup.callback(session.close) + engine_handle = ( + session.binding_context.engine_observer() + .handle(config.engine_id) + .init( + instance_name=f"cudf-polars-{str(config.engine_id)[:8]}", + implementation={ + "name": config.implementation_name, + "version": config.implementation_version, + "backend": backend, + "custom_attributes": {"backend": backend}, + }, + ) ) - ) - return runtime + runtime = cls( + config=config, + session=session, + collector=collector, + _engine_handle=engine_handle, + ) + cleanup.pop_all() + return runtime + + @property + def collector_address(self) -> str: + """Return the live Collector's advertised address.""" + if self.collector is None: + raise RuntimeError("Quent controller runtime is closed") + address = self.collector.address + if address is None: + raise RuntimeError("Quent Collector did not advertise an address") + return address @contextlib.contextmanager def query( diff --git a/python/cudf_polars/tests/quent/test_quent_integration.py b/python/cudf_polars/tests/quent/test_quent_integration.py index 1409f3bad20..8cf20cb8c7f 100644 --- a/python/cudf_polars/tests/quent/test_quent_integration.py +++ b/python/cudf_polars/tests/quent/test_quent_integration.py @@ -164,6 +164,31 @@ def test_quent_controller_cleanup_continues_after_handle_failure() -> None: collector.close.assert_called_once_with(timeout=10.0) +def test_quent_controller_create_closes_collector_on_failure( + quent_context: QuentConfig, + monkeypatch: pytest.MonkeyPatch, +) -> None: + collector = MagicMock(address="127.0.0.1:1234") + monkeypatch.setattr( + cudf_polars.quent._runtime, + "start_collector", + MagicMock(return_value=collector), + ) + monkeypatch.setattr( + cudf_polars.quent._runtime, + "QuentSession", + MagicMock(side_effect=RuntimeError("session creation failed")), + ) + + with pytest.raises(RuntimeError, match="session creation failed"): + cudf_polars.quent._runtime.QuentControllerRuntime.create( + quent_context, + backend="test", + ) + + collector.close.assert_called_once_with(timeout=10.0) + + def test_quent_lifecycle( engine_with_quent_context: tuple[StreamingEngine, StreamingOptions], quent_context: QuentConfig, diff --git a/python/cudf_polars/tests/streaming/test_dask.py b/python/cudf_polars/tests/streaming/test_dask.py index b258bb70d0e..e3e320c5fc1 100644 --- a/python/cudf_polars/tests/streaming/test_dask.py +++ b/python/cudf_polars/tests/streaming/test_dask.py @@ -60,26 +60,26 @@ def test_shutdown_dask_runs_all_steps_in_order() -> None: def run(func: Any) -> None: events.append(func.func.__name__) - if func.func.__name__ == "_close_quent_worker": - raise RuntimeError("worker Quent close failed") client = SimpleNamespace(run=run) - runtime = SimpleNamespace(close=lambda: events.append("controller")) + + def quent_close() -> None: + events.append("quent") + raise RuntimeError("Quent close failed") + owned_client = SimpleNamespace(close=lambda: events.append("client")) owned_cluster = SimpleNamespace(close=lambda: events.append("cluster")) with pytest.raises(ExceptionGroup, match="Dask engine shutdown failed"): _shutdown_dask( cast("Any", client), "test", - cast("Any", runtime), - None, + quent_close, cast("Any", owned_client), cast("Any", owned_cluster), ) assert events == [ - "_close_quent_worker", - "controller", + "quent", "_teardown_worker", "client", "cluster", From 154e2bd2ef980edd76fbb552fde42a77a88db6c4 Mon Sep 17 00:00:00 2001 From: Tom Augspurger Date: Mon, 19 Oct 2026 11:01:05 -0700 Subject: [PATCH 10/20] coverage --- python/cudf_polars/cudf_polars/engine/dask.py | 8 ++++---- python/cudf_polars/cudf_polars/engine/ray.py | 6 +++--- python/cudf_polars/cudf_polars/engine/spmd.py | 10 +++++----- python/cudf_polars/cudf_polars/quent/_runtime.py | 6 +++--- 4 files changed, 15 insertions(+), 15 deletions(-) diff --git a/python/cudf_polars/cudf_polars/engine/dask.py b/python/cudf_polars/cudf_polars/engine/dask.py index 5f4d330d1e8..d4bfef5466e 100644 --- a/python/cudf_polars/cudf_polars/engine/dask.py +++ b/python/cudf_polars/cudf_polars/engine/dask.py @@ -461,13 +461,13 @@ def _update_quent_on_worker( *, uid: str, dask_worker: distributed.Worker | None = None, -) -> None: +) -> None: # pragma: no cover - runs in a Dask worker process """Replace one worker's Quent runtime (or clear it when both args are ``None``).""" assert dask_worker is not None mp_ctx: _WorkerContext | None = getattr( dask_worker, f"_cudf_polars_mp_context_{uid}", None ) - if mp_ctx is None: + if mp_ctx is None: # pragma: no cover; unreachable return if mp_ctx.quent_worker_runtime is not None: mp_ctx.quent_worker_runtime.close() @@ -538,7 +538,7 @@ def _shutdown_dask( """Release Dask resources in dependency order.""" steps: list[Callable[[], object]] = [] if quent_close is not None: - steps.append(quent_close) + steps.append(quent_close) # pragma: no cover - same cleanup as reset if client is not None: steps.append(lambda: client.run(functools.partial(_teardown_worker, uid=uid))) if owned_client is not None: @@ -1267,7 +1267,7 @@ def _quent_open( uid=ctx.rapidsmpf_id, ) ) - except BaseException: + except BaseException: # pragma: no cover - defensive startup cleanup with contextlib.suppress(Exception): self._quent_close() raise diff --git a/python/cudf_polars/cudf_polars/engine/ray.py b/python/cudf_polars/cudf_polars/engine/ray.py index 98d0929c469..209c69bf11a 100644 --- a/python/cudf_polars/cudf_polars/engine/ray.py +++ b/python/cudf_polars/cudf_polars/engine/ray.py @@ -369,7 +369,7 @@ def update_quent( self, quent_context: cudf_polars.quent.QuentConfig | None, collector_address: str | None, - ) -> None: + ) -> None: # pragma: no cover - runs in a Ray actor process """ Replace this actor's Quent worker runtime. @@ -745,7 +745,7 @@ def _shutdown_ray( ray_running = ray.is_initialized() steps: list[Callable[[], object]] = [] if ray_running and quent_close is not None: - steps.append(quent_close) + steps.append(quent_close) # pragma: no cover - same cleanup as reset if ray_running: @@ -1100,7 +1100,7 @@ def _quent_open( for rank in self._rank_actors ] ) - except BaseException: + except BaseException: # pragma: no cover - defensive startup cleanup with contextlib.suppress(Exception): self._quent_close() raise diff --git a/python/cudf_polars/cudf_polars/engine/spmd.py b/python/cudf_polars/cudf_polars/engine/spmd.py index 4a2d8c87ba9..ee88ab92023 100644 --- a/python/cudf_polars/cudf_polars/engine/spmd.py +++ b/python/cudf_polars/cudf_polars/engine/spmd.py @@ -734,14 +734,14 @@ def _quent_close(self) -> None: """Close SPMD Quent runtimes: barrier → workers → barrier → controller → barrier.""" assert self._comm is not None comm = self._comm - if comm.nranks > 1: + if comm.nranks > 1: # pragma: no cover - requires a multi-rank test process barrier(comm) self._close_quent_worker() - if comm.nranks > 1: + if comm.nranks > 1: # pragma: no cover - requires a multi-rank test process barrier(comm) if comm.rank == 0: self._close_quent_controller() - if comm.nranks > 1: + if comm.nranks > 1: # pragma: no cover - requires a multi-rank test process barrier(comm) def _quent_open( @@ -773,7 +773,7 @@ def _quent_open( nranks=comm.nranks, instance_name=f"rank-{comm.rank}", ) - except BaseException: + except BaseException: # pragma: no cover - defensive startup cleanup steps: list[Callable[[], object]] = [ self._close_quent_worker, self._close_quent_controller, @@ -786,7 +786,7 @@ def _shutdown_spmd(self) -> None: """Run collective cleanup before local resources unwind.""" steps: list[Callable[[], object]] = [self._drop_persisted] if self._quent_worker_runtime is not None or self._quent_runtime is not None: - steps.append(self._quent_close) + steps.append(self._quent_close) # pragma: no cover - same cleanup as reset run_cleanup_steps("SPMD engine shutdown failed", *steps) @classmethod diff --git a/python/cudf_polars/cudf_polars/quent/_runtime.py b/python/cudf_polars/cudf_polars/quent/_runtime.py index dfe753940cd..fa6402e5d48 100644 --- a/python/cudf_polars/cudf_polars/quent/_runtime.py +++ b/python/cudf_polars/cudf_polars/quent/_runtime.py @@ -129,7 +129,7 @@ def create( collector = start_collector(config.run_root) cleanup.callback(lambda: collector.close(timeout=10.0)) collector_address = collector.address - if collector_address is None: + if collector_address is None: # pragma: no cover - Collector invariant raise RuntimeError("Quent Collector did not advertise an address") session = QuentSession(collector_address) cleanup.callback(session.close) @@ -158,10 +158,10 @@ def create( @property def collector_address(self) -> str: """Return the live Collector's advertised address.""" - if self.collector is None: + if self.collector is None: # pragma: no cover - caller lifecycle invariant raise RuntimeError("Quent controller runtime is closed") address = self.collector.address - if address is None: + if address is None: # pragma: no cover - Collector invariant raise RuntimeError("Quent Collector did not advertise an address") return address From 557bbd6cb29ff456c76c25da8e378b52e208ab60 Mon Sep 17 00:00:00 2001 From: Tom Augspurger Date: Mon, 19 Oct 2026 11:58:27 -0700 Subject: [PATCH 11/20] Move the test --- .../tests/quent/test_quent_context.py | 28 +++++++++++++++++++ .../tests/quent/test_quent_integration.py | 25 ----------------- 2 files changed, 28 insertions(+), 25 deletions(-) diff --git a/python/cudf_polars/tests/quent/test_quent_context.py b/python/cudf_polars/tests/quent/test_quent_context.py index d6efb97cbc4..9df2d3678eb 100644 --- a/python/cudf_polars/tests/quent/test_quent_context.py +++ b/python/cudf_polars/tests/quent/test_quent_context.py @@ -14,6 +14,9 @@ import uuid from unittest.mock import MagicMock, call +import pytest + +import cudf_polars.quent._runtime from cudf_polars.quent._context import ( ActorBoundQuentIRExecutionState, QuentConfig, @@ -118,3 +121,28 @@ def test_emit_evaluate_failure() -> None: running_evaluation.failed.assert_called_once_with(error=str(error)) assert evaluate_id not in runtime.session._evaluations + + +def test_quent_controller_create_closes_collector_on_failure( + quent_context: QuentConfig, + monkeypatch: pytest.MonkeyPatch, +) -> None: + collector = MagicMock(address="127.0.0.1:1234") + monkeypatch.setattr( + cudf_polars.quent._runtime, + "start_collector", + MagicMock(return_value=collector), + ) + monkeypatch.setattr( + cudf_polars.quent._runtime, + "QuentSession", + MagicMock(side_effect=RuntimeError("session creation failed")), + ) + + with pytest.raises(RuntimeError, match="session creation failed"): + cudf_polars.quent._runtime.QuentControllerRuntime.create( + quent_context, + backend="test", + ) + + collector.close.assert_called_once_with(timeout=10.0) diff --git a/python/cudf_polars/tests/quent/test_quent_integration.py b/python/cudf_polars/tests/quent/test_quent_integration.py index 8cf20cb8c7f..1409f3bad20 100644 --- a/python/cudf_polars/tests/quent/test_quent_integration.py +++ b/python/cudf_polars/tests/quent/test_quent_integration.py @@ -164,31 +164,6 @@ def test_quent_controller_cleanup_continues_after_handle_failure() -> None: collector.close.assert_called_once_with(timeout=10.0) -def test_quent_controller_create_closes_collector_on_failure( - quent_context: QuentConfig, - monkeypatch: pytest.MonkeyPatch, -) -> None: - collector = MagicMock(address="127.0.0.1:1234") - monkeypatch.setattr( - cudf_polars.quent._runtime, - "start_collector", - MagicMock(return_value=collector), - ) - monkeypatch.setattr( - cudf_polars.quent._runtime, - "QuentSession", - MagicMock(side_effect=RuntimeError("session creation failed")), - ) - - with pytest.raises(RuntimeError, match="session creation failed"): - cudf_polars.quent._runtime.QuentControllerRuntime.create( - quent_context, - backend="test", - ) - - collector.close.assert_called_once_with(timeout=10.0) - - def test_quent_lifecycle( engine_with_quent_context: tuple[StreamingEngine, StreamingOptions], quent_context: QuentConfig, From 436e0fabc545e5bf660d54586e109eb0a17c69e0 Mon Sep 17 00:00:00 2001 From: Tom Augspurger Date: Tue, 20 Oct 2026 06:58:45 -0700 Subject: [PATCH 12/20] sync quent context --- python/cudf_polars/cudf_polars/engine/ray.py | 4 +- python/cudf_polars/cudf_polars/engine/spmd.py | 54 +++++++++++++++++++ .../tests/quent/test_quent_context.py | 44 +++++++++++++++ 3 files changed, 101 insertions(+), 1 deletion(-) diff --git a/python/cudf_polars/cudf_polars/engine/ray.py b/python/cudf_polars/cudf_polars/engine/ray.py index 209c69bf11a..10cf7b7c1d0 100644 --- a/python/cudf_polars/cudf_polars/engine/ray.py +++ b/python/cudf_polars/cudf_polars/engine/ray.py @@ -750,7 +750,9 @@ def _shutdown_ray( if ray_running: def shutdown_actors() -> None: - refs = [actor.shutdown.remote() for actor in rank_actors] + refs: list[ObjectRef[Any]] = [ + actor.shutdown.remote() for actor in rank_actors + ] run_cleanup_steps( "Ray actor shutdown failed", *(lambda ref=ref: _wait_for_actor_shutdown(ref) for ref in refs), diff --git a/python/cudf_polars/cudf_polars/engine/spmd.py b/python/cudf_polars/cudf_polars/engine/spmd.py index ee88ab92023..09fb08e7113 100644 --- a/python/cudf_polars/cudf_polars/engine/spmd.py +++ b/python/cudf_polars/cudf_polars/engine/spmd.py @@ -402,6 +402,54 @@ def synchronize_quent_configuration( ) +def synchronize_quent_reset_contexts( + *, + comm: Communicator, + context: Context, + existing_context: cudf_polars.quent.QuentConfig | None, + new_context: cudf_polars.quent.QuentConfig | None, +) -> tuple[ + cudf_polars.quent.QuentConfig | None, + cudf_polars.quent.QuentConfig | None, +]: + """Broadcast rank 0's current and requested Quent contexts.""" + if comm.nranks == 1: + return existing_context, new_context + if comm.rank == 0: + data = json.dumps( + { + "existing": ( + existing_context._serialize().decode() + if existing_context is not None + else None + ), + "new": ( + new_context._serialize().decode() + if new_context is not None + else None + ), + } + ).encode() + else: + data = b"" + with reserve_op_id() as op_id: + all_data = all_gather_host_data(comm, context.br(), op_id, data) + synchronized = json.loads(all_data[0]) + existing_data = synchronized["existing"] + new_data = synchronized["new"] + deserialized_existing = ( + cudf_polars.quent.QuentConfig._deserialize(existing_data.encode()) + if existing_data is not None + else None + ) + deserialized_new = ( + cudf_polars.quent.QuentConfig._deserialize(new_data.encode()) + if new_data is not None + else None + ) + return (deserialized_existing, deserialized_new) + + def synchronize_quent_query_id( *, comm: Communicator, @@ -1014,6 +1062,12 @@ def _reset( new_quent_context: cudf_polars.quent.QuentConfig | None = executor_options.get( "quent_context", existing_quent_context ) + existing_quent_context, new_quent_context = synchronize_quent_reset_contexts( + comm=self._comm, + context=self._ctx, + existing_context=existing_quent_context, + new_context=new_quent_context, + ) quent_changing = new_quent_context != existing_quent_context rapidsmpf_options = resolve_rapidsmpf_options(rapidsmpf_options) self.rapidsmpf_options = rapidsmpf_options diff --git a/python/cudf_polars/tests/quent/test_quent_context.py b/python/cudf_polars/tests/quent/test_quent_context.py index 9df2d3678eb..8af1cffbd7a 100644 --- a/python/cudf_polars/tests/quent/test_quent_context.py +++ b/python/cudf_polars/tests/quent/test_quent_context.py @@ -11,12 +11,14 @@ from __future__ import annotations +import dataclasses import uuid from unittest.mock import MagicMock, call import pytest import cudf_polars.quent._runtime +from cudf_polars.engine.spmd import synchronize_quent_reset_contexts from cudf_polars.quent._context import ( ActorBoundQuentIRExecutionState, QuentConfig, @@ -41,6 +43,48 @@ def test_context_serialization_roundtrip(tmp_path) -> None: assert QuentConfig._deserialize(context._serialize()) == context +@pytest.mark.parametrize( + "existing_enabled,new_enabled", + [ + pytest.param(False, True, id="enable"), + pytest.param(True, True, id="replace"), + pytest.param(True, False, id="disable"), + ], +) +def test_synchronize_quent_reset_contexts( + quent_context: QuentConfig, + monkeypatch: pytest.MonkeyPatch, + *, + existing_enabled: bool, + new_enabled: bool, +) -> None: + existing_context = quent_context if existing_enabled else None + new_context = ( + dataclasses.replace(quent_context, engine_id=uuid.uuid4()) + if new_enabled + else None + ) + + def fake_all_gather(*args: object) -> list[bytes]: + data = args[-1] + assert isinstance(data, bytes) + return [data, b""] + + monkeypatch.setattr( + "cudf_polars.engine.spmd.all_gather_host_data", + fake_all_gather, + ) + + synchronized = synchronize_quent_reset_contexts( + comm=MagicMock(nranks=2, rank=0), + context=MagicMock(), + existing_context=existing_context, + new_context=new_context, + ) + + assert synchronized == (existing_context, new_context) + + def test_ir_execution_state_bind_actor() -> None: actor_id = uuid.uuid4() state = UnboundQuentIRExecutionState( From b2b0a4fbbb062a99c853f5250f723e9c969f7f23 Mon Sep 17 00:00:00 2001 From: Tom Augspurger Date: Tue, 20 Oct 2026 07:10:24 -0700 Subject: [PATCH 13/20] set_quent --- python/cudf_polars/cudf_polars/engine/core.py | 36 ++----------- python/cudf_polars/cudf_polars/engine/dask.py | 52 +++++++++--------- python/cudf_polars/cudf_polars/engine/ray.py | 49 +++++++++-------- python/cudf_polars/cudf_polars/engine/spmd.py | 53 +++++++++---------- 4 files changed, 82 insertions(+), 108 deletions(-) diff --git a/python/cudf_polars/cudf_polars/engine/core.py b/python/cudf_polars/cudf_polars/engine/core.py index f923b3b4381..a40a3bf5e82 100644 --- a/python/cudf_polars/cudf_polars/engine/core.py +++ b/python/cudf_polars/cudf_polars/engine/core.py @@ -583,40 +583,10 @@ def __exit__(self, *_: object) -> None: """Exit the context manager, calling :meth:`shutdown`.""" self.shutdown() - def _update_quent( - self, - new_quent_context: QuentConfig | None, - existing_quent_context: QuentConfig | None, - executor_options: dict[str, Any], - ) -> None: - """ - Apply a ``quent_context`` change in-place. - - Closes existing Quent runtimes, starts new ones for - ``new_quent_context``, and patches ``executor_options["quent_context"]`` - and ``_quent_output_root``. Subclasses implement :meth:`_quent_close` - and :meth:`_quent_open` for the backend-specific work. - """ - if new_quent_context != existing_quent_context: - if existing_quent_context is not None: - self._quent_close() - if new_quent_context is not None: - new_quent_context = self._quent_open(new_quent_context) - self._quent_output_root = ( - new_quent_context.run_root if new_quent_context is not None else None - ) - executor_options["quent_context"] = new_quent_context - - def _quent_close(self) -> None: - """Close all Quent runtimes for this engine (workers then controller).""" - raise NotImplementedError( # pragma: no cover - f"{type(self).__name__} does not implement _quent_close" - ) - - def _quent_open(self, quent_context: QuentConfig) -> QuentConfig: - """Start fresh Quent runtimes and return the effective configuration.""" + def _set_quent(self, quent_context: QuentConfig | None) -> QuentConfig | None: + """Replace this engine's Quent runtimes and return the effective config.""" raise NotImplementedError( # pragma: no cover - f"{type(self).__name__} does not implement _quent_open" + f"{type(self).__name__} does not implement _set_quent" ) def _run(self, func: Callable[..., T], *args: Any, **kwargs: Any) -> list[T]: diff --git a/python/cudf_polars/cudf_polars/engine/dask.py b/python/cudf_polars/cudf_polars/engine/dask.py index d4bfef5466e..e17cdf138c7 100644 --- a/python/cudf_polars/cudf_polars/engine/dask.py +++ b/python/cudf_polars/cudf_polars/engine/dask.py @@ -1041,7 +1041,11 @@ def __init__( lambda: _shutdown_dask( dask_client, rapidsmpf_id, - (self._quent_close if self._quent_runtime is not None else None), + ( + functools.partial(self._set_quent, None) + if self._quent_runtime is not None + else None + ), owned_client, owned_cluster, ) @@ -1125,7 +1129,7 @@ def __init__( owned_cluster=owned_cluster, ) self._dask_context: DaskContext | None = dask_ctx - self._update_quent(quent_context, None, executor_options) + executor_options["quent_context"] = self._set_quent(quent_context) dask_ctx = DaskContext( client=dask_client, rapidsmpf_id=rapidsmpf_id, @@ -1205,7 +1209,9 @@ def _reset( new_quent_context: cudf_polars.quent.QuentConfig | None = executor_options.get( "quent_context", existing_quent_context ) - self._update_quent(new_quent_context, existing_quent_context, executor_options) + if new_quent_context != existing_quent_context: + new_quent_context = self._set_quent(new_quent_context) + executor_options["quent_context"] = new_quent_context new_ctx = DaskContext( client=ctx.client, rapidsmpf_id=ctx.rapidsmpf_id, @@ -1231,29 +1237,27 @@ def _reset( exit_stack=self._exit_stack, ) - def _quent_close(self) -> None: - """Clear Quent worker runtimes on all Dask workers, then close the controller.""" + def _set_quent( + self, quent_context: cudf_polars.quent.QuentConfig | None + ) -> cudf_polars.quent.QuentConfig | None: + """Replace this engine's Dask Quent runtimes.""" ctx = self._dask_ctx - steps: list[Callable[[], object]] = [ - # Workers first so the collector can drain before we close it. - lambda: ctx.client.run( - functools.partial( - _update_quent_on_worker, None, None, uid=ctx.rapidsmpf_id - ) - ) - ] if self._quent_runtime is not None: + steps: list[Callable[[], object]] = [ + # Workers first so the collector can drain before we close it. + lambda: ctx.client.run( + functools.partial( + _update_quent_on_worker, None, None, uid=ctx.rapidsmpf_id + ) + ) + ] steps.append(self._quent_runtime.close) - try: - run_cleanup_steps("Dask Quent shutdown failed", *steps) - finally: - self._quent_runtime = None - - def _quent_open( - self, quent_context: cudf_polars.quent.QuentConfig - ) -> cudf_polars.quent.QuentConfig: - """Start a new Quent collector, connect workers, create controller.""" - ctx = self._dask_ctx + try: + run_cleanup_steps("Dask Quent shutdown failed", *steps) + finally: + self._quent_runtime = None + if quent_context is None: + return None self._quent_runtime = QuentControllerRuntime.create( quent_context, backend="dask", @@ -1269,7 +1273,7 @@ def _quent_open( ) except BaseException: # pragma: no cover - defensive startup cleanup with contextlib.suppress(Exception): - self._quent_close() + self._set_quent(None) raise return quent_context diff --git a/python/cudf_polars/cudf_polars/engine/ray.py b/python/cudf_polars/cudf_polars/engine/ray.py index 10cf7b7c1d0..1e17b9eb1af 100644 --- a/python/cudf_polars/cudf_polars/engine/ray.py +++ b/python/cudf_polars/cudf_polars/engine/ray.py @@ -5,6 +5,7 @@ from __future__ import annotations import contextlib +import functools import os import uuid from concurrent.futures import ThreadPoolExecutor @@ -904,7 +905,11 @@ def __init__( exit_stack.callback( lambda: _shutdown_ray( rank_actors, - (self._quent_close if self._quent_runtime is not None else None), + ( + functools.partial(self._set_quent, None) + if self._quent_runtime is not None + else None + ), shutdown_ray=owns_ray, ) ) @@ -974,7 +979,7 @@ def __init__( ) self._rank_actors: list[ActorHandle[RankActor]] | None = rank_actors - self._update_quent(quent_context, None, executor_options) + executor_options["quent_context"] = self._set_quent(quent_context) super().__init__( nranks=nranks, executor_options={ @@ -1046,7 +1051,9 @@ def _reset( new_quent_context: cudf_polars.quent.QuentConfig | None = executor_options.get( "quent_context", existing_quent_context ) - self._update_quent(new_quent_context, existing_quent_context, executor_options) + if new_quent_context != existing_quent_context: + new_quent_context = self._set_quent(new_quent_context) + executor_options["quent_context"] = new_quent_context # Re-run ``StreamingEngine.__init__`` on the existing instance to # reconfigure the polars ``GPUEngine`` layer (``self.config``, @@ -1066,28 +1073,26 @@ def _reset( exit_stack=self._exit_stack, ) - def _quent_close(self) -> None: - """Clear Quent worker runtimes on all actors, then close the controller.""" + def _set_quent( + self, quent_context: cudf_polars.quent.QuentConfig | None + ) -> cudf_polars.quent.QuentConfig | None: + """Replace this engine's Ray Quent runtimes.""" assert self._rank_actors is not None rank_actors = self._rank_actors - steps: list[Callable[[], object]] = [ - # Workers first so the collector can drain before we close it. - lambda: ray.get( - [rank.update_quent.remote(None, None) for rank in rank_actors] - ) - ] if self._quent_runtime is not None: + steps: list[Callable[[], object]] = [ + # Workers first so the collector can drain before we close it. + lambda: ray.get( + [rank.update_quent.remote(None, None) for rank in rank_actors] + ) + ] steps.append(self._quent_runtime.close) - try: - run_cleanup_steps("Ray Quent shutdown failed", *steps) - finally: - self._quent_runtime = None - - def _quent_open( - self, quent_context: cudf_polars.quent.QuentConfig - ) -> cudf_polars.quent.QuentConfig: - """Start a new Quent collector, connect actors, create controller.""" - assert self._rank_actors is not None + try: + run_cleanup_steps("Ray Quent shutdown failed", *steps) + finally: + self._quent_runtime = None + if quent_context is None: + return None self._quent_runtime = QuentControllerRuntime.create( quent_context, backend="ray", @@ -1104,7 +1109,7 @@ def _quent_open( ) except BaseException: # pragma: no cover - defensive startup cleanup with contextlib.suppress(Exception): - self._quent_close() + self._set_quent(None) raise return quent_context diff --git a/python/cudf_polars/cudf_polars/engine/spmd.py b/python/cudf_polars/cudf_polars/engine/spmd.py index 09fb08e7113..f159af2eb71 100644 --- a/python/cudf_polars/cudf_polars/engine/spmd.py +++ b/python/cudf_polars/cudf_polars/engine/spmd.py @@ -6,6 +6,7 @@ import contextlib import dataclasses +import functools import json import os import uuid @@ -698,7 +699,7 @@ def __init__( exit_stack.callback(self._cleanup_ctx) self._quent_worker_id = uuid.uuid4() - self._update_quent(quent_context, None, executor_options) + executor_options["quent_context"] = self._set_quent(quent_context) effective_quent_context = executor_options["quent_context"] engine_id = ( effective_quent_context.engine_id @@ -778,27 +779,25 @@ def _close_quent_controller(self) -> None: finally: self._quent_runtime = None - def _quent_close(self) -> None: - """Close SPMD Quent runtimes: barrier → workers → barrier → controller → barrier.""" - assert self._comm is not None - comm = self._comm - if comm.nranks > 1: # pragma: no cover - requires a multi-rank test process - barrier(comm) - self._close_quent_worker() - if comm.nranks > 1: # pragma: no cover - requires a multi-rank test process - barrier(comm) - if comm.rank == 0: - self._close_quent_controller() - if comm.nranks > 1: # pragma: no cover - requires a multi-rank test process - barrier(comm) - - def _quent_open( - self, quent_context: cudf_polars.quent.QuentConfig - ) -> cudf_polars.quent.QuentConfig: - """Start SPMD Quent runtimes for ``quent_context``.""" + def _set_quent( + self, quent_context: cudf_polars.quent.QuentConfig | None + ) -> cudf_polars.quent.QuentConfig | None: + """Replace this engine's SPMD Quent runtimes.""" assert self._comm is not None assert self._ctx is not None comm = self._comm + if self._quent_worker_runtime is not None or self._quent_runtime is not None: + if comm.nranks > 1: # pragma: no cover - requires a multi-rank test process + barrier(comm) + self._close_quent_worker() + if comm.nranks > 1: # pragma: no cover - requires a multi-rank test process + barrier(comm) + if comm.rank == 0: + self._close_quent_controller() + if comm.nranks > 1: # pragma: no cover - requires a multi-rank test process + barrier(comm) + if quent_context is None: + return None collector_address: str = "" try: if comm.rank == 0: @@ -834,7 +833,9 @@ def _shutdown_spmd(self) -> None: """Run collective cleanup before local resources unwind.""" steps: list[Callable[[], object]] = [self._drop_persisted] if self._quent_worker_runtime is not None or self._quent_runtime is not None: - steps.append(self._quent_close) # pragma: no cover - same cleanup as reset + steps.append( # pragma: no cover - same cleanup as reset + functools.partial(self._set_quent, None) + ) run_cleanup_steps("SPMD engine shutdown failed", *steps) @classmethod @@ -1072,14 +1073,10 @@ def _reset( rapidsmpf_options = resolve_rapidsmpf_options(rapidsmpf_options) self.rapidsmpf_options = rapidsmpf_options - # Tear down Quent runtimes before the Context rebuild if they're changing. - if quent_changing and existing_quent_context is not None: - self._quent_close() - existing_quent_context = None + if quent_changing: + new_quent_context = self._set_quent(new_quent_context) + executor_options["quent_context"] = new_quent_context - # Collective: synchronize all ranks before tearing down the Context. - if self._comm.nranks > 1: - barrier(self._comm) # Free persisted partitions before the Context is torn down. self._drop_persisted() # Same-thread shutdown, _reset runs on the thread that built the @@ -1105,8 +1102,6 @@ def _reset( self._mr = self._ctx.br().device_mr_adaptor() rmm.mr.set_current_device_resource(self._mr) - # Start new Quent runtimes after the Context is rebuilt. - self._update_quent(new_quent_context, existing_quent_context, executor_options) effective_quent_context = executor_options["quent_context"] engine_id = ( effective_quent_context.engine_id From 9dab31b9943223f615266732944694d5d60ae1cf Mon Sep 17 00:00:00 2001 From: Tom Augspurger Date: Tue, 20 Oct 2026 07:55:07 -0700 Subject: [PATCH 14/20] ensure closed --- python/cudf_polars/cudf_polars/engine/dask.py | 6 ++++-- python/cudf_polars/cudf_polars/engine/ray.py | 6 ++++-- 2 files changed, 8 insertions(+), 4 deletions(-) diff --git a/python/cudf_polars/cudf_polars/engine/dask.py b/python/cudf_polars/cudf_polars/engine/dask.py index e17cdf138c7..8cc374688cd 100644 --- a/python/cudf_polars/cudf_polars/engine/dask.py +++ b/python/cudf_polars/cudf_polars/engine/dask.py @@ -470,8 +470,10 @@ def _update_quent_on_worker( if mp_ctx is None: # pragma: no cover; unreachable return if mp_ctx.quent_worker_runtime is not None: - mp_ctx.quent_worker_runtime.close() - mp_ctx.quent_worker_runtime = None + try: + mp_ctx.quent_worker_runtime.close() + finally: + mp_ctx.quent_worker_runtime = None if quent_context is not None and collector_address is not None: assert mp_ctx.comm is not None worker_id = uuid.uuid4() diff --git a/python/cudf_polars/cudf_polars/engine/ray.py b/python/cudf_polars/cudf_polars/engine/ray.py index 1e17b9eb1af..68f60dc3c43 100644 --- a/python/cudf_polars/cudf_polars/engine/ray.py +++ b/python/cudf_polars/cudf_polars/engine/ray.py @@ -378,8 +378,10 @@ def update_quent( without starting a new one. """ if self._quent_worker_runtime is not None: - self._quent_worker_runtime.close() - self._quent_worker_runtime = None + try: + self._quent_worker_runtime.close() + finally: + self._quent_worker_runtime = None if quent_context is not None and collector_address is not None: assert self._comm is not None self._quent_worker_runtime = QuentWorkerRuntime.create( From 06b884d2e0929bb546623d9d4643e8706e5b2290 Mon Sep 17 00:00:00 2001 From: Tom Augspurger Date: Tue, 20 Oct 2026 11:01:48 -0700 Subject: [PATCH 15/20] wip --- python/cudf_polars/quent/bridge/pyproject.toml | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/python/cudf_polars/quent/bridge/pyproject.toml b/python/cudf_polars/quent/bridge/pyproject.toml index 24333557945..a83be74ae77 100644 --- a/python/cudf_polars/quent/bridge/pyproject.toml +++ b/python/cudf_polars/quent/bridge/pyproject.toml @@ -7,8 +7,8 @@ build-backend = "maturin" [project] name = "cudf-polars-quent" -version = "0.1.0" -requires-python = ">=3.12" +# rapids-pre-commit-hooks: disable-next-line +version = "26.12.0" license = "Apache-2.0" [tool.maturin] From ddfd1d415edf342d0d55a1ae7293ad59122b2bdf Mon Sep 17 00:00:00 2001 From: Tom Augspurger Date: Tue, 20 Oct 2026 11:02:35 -0700 Subject: [PATCH 16/20] wip --- python/cudf_polars/quent/bridge/pyproject.toml | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/python/cudf_polars/quent/bridge/pyproject.toml b/python/cudf_polars/quent/bridge/pyproject.toml index a83be74ae77..24333557945 100644 --- a/python/cudf_polars/quent/bridge/pyproject.toml +++ b/python/cudf_polars/quent/bridge/pyproject.toml @@ -7,8 +7,8 @@ build-backend = "maturin" [project] name = "cudf-polars-quent" -# rapids-pre-commit-hooks: disable-next-line -version = "26.12.0" +version = "0.1.0" +requires-python = ">=3.12" license = "Apache-2.0" [tool.maturin] From 3d623f55696018e272ea12a73f4f1168d0c818c4 Mon Sep 17 00:00:00 2001 From: Tom Augspurger Date: Tue, 20 Oct 2026 11:09:54 -0700 Subject: [PATCH 17/20] fixup --- python/cudf_polars/cudf_polars/engine/dask.py | 5 ++- python/cudf_polars/cudf_polars/engine/ray.py | 5 ++- .../cudf_polars/cudf_polars/quent/_runtime.py | 35 ++++++++++++------- .../cudf_polars/tests/streaming/test_dask.py | 2 ++ .../cudf_polars/tests/streaming/test_ray.py | 2 +- 5 files changed, 33 insertions(+), 16 deletions(-) diff --git a/python/cudf_polars/cudf_polars/engine/dask.py b/python/cudf_polars/cudf_polars/engine/dask.py index 8cc374688cd..847729a0d89 100644 --- a/python/cudf_polars/cudf_polars/engine/dask.py +++ b/python/cudf_polars/cudf_polars/engine/dask.py @@ -1007,6 +1007,7 @@ def __init__( ) -> None: executor_options = resolve_kvikio_executor_options(executor_options or {}) engine_options = engine_options or {} + self._allow_gpu_sharing = bool(engine_options.get("allow_gpu_sharing", False)) quent_context = resolve_quent_context(executor_options) executor_options["quent_context"] = quent_context @@ -1179,7 +1180,9 @@ def _reset( "kvikio_nthreads", existing_executor_options["kvikio_nthreads"] ) executor_options = resolve_kvikio_executor_options(executor_options) - engine_options = engine_options or {} + engine_options = dict(engine_options or {}) + engine_options.setdefault("allow_gpu_sharing", self._allow_gpu_sharing) + self._allow_gpu_sharing = bool(engine_options["allow_gpu_sharing"]) self.rapidsmpf_options = resolve_rapidsmpf_options(rapidsmpf_options) rapidsmpf_options_as_bytes = self.rapidsmpf_options.serialize() diff --git a/python/cudf_polars/cudf_polars/engine/ray.py b/python/cudf_polars/cudf_polars/engine/ray.py index 68f60dc3c43..6ceb97e0ef5 100644 --- a/python/cudf_polars/cudf_polars/engine/ray.py +++ b/python/cudf_polars/cudf_polars/engine/ray.py @@ -865,6 +865,7 @@ def __init__( ) -> None: executor_options = resolve_kvikio_executor_options(executor_options or {}) engine_options = engine_options or {} + self._allow_gpu_sharing = bool(engine_options.get("allow_gpu_sharing", False)) ray_init_options = ray_init_options or {} if bootstrap.is_running_with_rrun(): @@ -1021,7 +1022,9 @@ def _reset( "kvikio_nthreads", existing_executor_options["kvikio_nthreads"] ) executor_options = resolve_kvikio_executor_options(executor_options) - engine_options = engine_options or {} + engine_options = dict(engine_options or {}) + engine_options.setdefault("allow_gpu_sharing", self._allow_gpu_sharing) + self._allow_gpu_sharing = bool(engine_options["allow_gpu_sharing"]) self.rapidsmpf_options = resolve_rapidsmpf_options(rapidsmpf_options) rapidsmpf_options_as_bytes = self.rapidsmpf_options.serialize() diff --git a/python/cudf_polars/cudf_polars/quent/_runtime.py b/python/cudf_polars/cudf_polars/quent/_runtime.py index fa6402e5d48..c9cc10a9104 100644 --- a/python/cudf_polars/cudf_polars/quent/_runtime.py +++ b/python/cudf_polars/cudf_polars/quent/_runtime.py @@ -242,19 +242,28 @@ def create( from cudf_polars.quent._context import WorkerResources session = QuentSession(collector_address) - worker_handle = ( - session.binding_context.worker_observer() - .handle(worker_id) - .init(instance_name=instance_name, engine=config.engine_id) - ) - resources = WorkerResources.build( - instance_suffix=instance_name, - engine_id=config.engine_id, - worker_id=worker_id, - rank=rank, - nranks=nranks, - ) - resources.declare(session) + worker_handle: quent_bindings.WorkerInitHandle | None = None + try: + worker_handle = ( + session.binding_context.worker_observer() + .handle(worker_id) + .init(instance_name=instance_name, engine=config.engine_id) + ) + resources = WorkerResources.build( + instance_suffix=instance_name, + engine_id=config.engine_id, + worker_id=worker_id, + rank=rank, + nranks=nranks, + ) + resources.declare(session) + except BaseException: # pragma: no cover + if worker_handle is not None: + with contextlib.suppress(Exception): + worker_handle.exit() + with contextlib.suppress(Exception): + session.close() + raise return cls( config=config, session=session, diff --git a/python/cudf_polars/tests/streaming/test_dask.py b/python/cudf_polars/tests/streaming/test_dask.py index e3e320c5fc1..0143570a4be 100644 --- a/python/cudf_polars/tests/streaming/test_dask.py +++ b/python/cudf_polars/tests/streaming/test_dask.py @@ -156,6 +156,7 @@ def reset_engine(dask_client: distributed.Client) -> Iterator[DaskEngine]: """ with DaskEngine( executor_options={"max_rows_per_partition": 10}, + engine_options={"allow_gpu_sharing": True}, dask_client=dask_client, ) as e: yield e @@ -175,6 +176,7 @@ def test_reset_keeps_workers_alive(reset_engine: DaskEngine) -> None: ) pids_after = sorted(reset_engine._run(os.getpid)) + assert reset_engine._allow_gpu_sharing is True # Same worker addresses … assert workers_before == workers_after # … and the workers are running in the same OS processes. diff --git a/python/cudf_polars/tests/streaming/test_ray.py b/python/cudf_polars/tests/streaming/test_ray.py index 1a5e5c298a1..d8cf07caef5 100644 --- a/python/cudf_polars/tests/streaming/test_ray.py +++ b/python/cudf_polars/tests/streaming/test_ray.py @@ -148,12 +148,12 @@ def test_reset_keeps_actors_alive(reset_engine: RayEngine) -> None: reset_engine._reset( executor_options={"max_rows_per_partition": 7}, - engine_options={"allow_gpu_sharing": True}, ) actors_after = list(reset_engine.rank_actors) pids_after = reset_engine._run(os.getpid) + assert reset_engine._allow_gpu_sharing is True # The Python ``ActorHandle`` objects are the same instances … assert all(a is b for a, b in zip(actors_before, actors_after, strict=True)) # … and the actors are running in the same OS processes. From 2eff55f970c9b4a23c784bfa578875692e5b96d2 Mon Sep 17 00:00:00 2001 From: Tom Augspurger Date: Wed, 21 Oct 2026 07:39:42 -0700 Subject: [PATCH 18/20] docstrings --- python/cudf_polars/cudf_polars/engine/core.py | 17 ++++++++++++++++- python/cudf_polars/cudf_polars/engine/spmd.py | 18 +++++++++++++++++- 2 files changed, 33 insertions(+), 2 deletions(-) diff --git a/python/cudf_polars/cudf_polars/engine/core.py b/python/cudf_polars/cudf_polars/engine/core.py index a40a3bf5e82..79caea6c865 100644 --- a/python/cudf_polars/cudf_polars/engine/core.py +++ b/python/cudf_polars/cudf_polars/engine/core.py @@ -584,7 +584,22 @@ def __exit__(self, *_: object) -> None: self.shutdown() def _set_quent(self, quent_context: QuentConfig | None) -> QuentConfig | None: - """Replace this engine's Quent runtimes and return the effective config.""" + """ + Replace this engine's Quent runtimes. + + Closes any running Quent runtimes (on all ranks), then starts new ones for + ``quent_context`` unless it is ``None``. + + Parameters + ---------- + quent_context + Quent configuration to start, or ``None`` to only close. + + Returns + ------- + The effective Quent configuration, which may differ from + ``quent_context``, or ``None`` if Quent is disabled. + """ raise NotImplementedError( # pragma: no cover f"{type(self).__name__} does not implement _set_quent" ) diff --git a/python/cudf_polars/cudf_polars/engine/spmd.py b/python/cudf_polars/cudf_polars/engine/spmd.py index e1f5ebcd7d9..3bc0503a050 100644 --- a/python/cudf_polars/cudf_polars/engine/spmd.py +++ b/python/cudf_polars/cudf_polars/engine/spmd.py @@ -782,7 +782,23 @@ def _close_quent_controller(self) -> None: def _set_quent( self, quent_context: cudf_polars.quent.QuentConfig | None ) -> cudf_polars.quent.QuentConfig | None: - """Replace this engine's SPMD Quent runtimes.""" + """ + Replace this engine's Quent runtimes. + + Must be called collectively on all ranks, with ``quent_context`` either + ``None`` on all ranks or not ``None`` on all ranks. Rank 0 starts the + Collector and controller, and every rank starts a worker runtime. + + Parameters + ---------- + quent_context + Quent configuration to start, or ``None`` to only close. + + Returns + ------- + Rank 0's synchronized Quent configuration, or ``None`` if Quent is + disabled. + """ assert self._comm is not None assert self._ctx is not None comm = self._comm From 5c1bda1b6d0d80bd236981ef0b9c69cb6b610b35 Mon Sep 17 00:00:00 2001 From: Tom Augspurger Date: Wed, 21 Oct 2026 07:43:17 -0700 Subject: [PATCH 19/20] sync before set --- python/cudf_polars/cudf_polars/engine/spmd.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/python/cudf_polars/cudf_polars/engine/spmd.py b/python/cudf_polars/cudf_polars/engine/spmd.py index 3bc0503a050..3465f3ee720 100644 --- a/python/cudf_polars/cudf_polars/engine/spmd.py +++ b/python/cudf_polars/cudf_polars/engine/spmd.py @@ -699,6 +699,12 @@ def __init__( exit_stack.callback(self._cleanup_ctx) self._quent_worker_id = uuid.uuid4() + _, quent_context = synchronize_quent_reset_contexts( + comm=comm, + context=self._ctx, + existing_context=None, + new_context=quent_context, + ) executor_options["quent_context"] = self._set_quent(quent_context) effective_quent_context = executor_options["quent_context"] engine_id = ( From 74ea0e409d0e666e98225640e0b62243a5d1f2b0 Mon Sep 17 00:00:00 2001 From: Tom Augspurger Date: Thu, 22 Oct 2026 20:01:20 -0700 Subject: [PATCH 20/20] Trigger CI