From b3d633c79d5b633d14489373a06577e27dde7ed6 Mon Sep 17 00:00:00 2001 From: "Mads R. B. Kristensen" Date: Mon, 9 Mar 2026 17:23:34 +0100 Subject: [PATCH 01/33] match config_options.executor.cluster --- .../experimental/rapidsmpf/core.py | 82 ++++++++++--------- 1 file changed, 42 insertions(+), 40 deletions(-) diff --git a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/core.py b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/core.py index d5a9316e3b94..d5a4e4457d0d 100644 --- a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/core.py +++ b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/core.py @@ -108,47 +108,49 @@ def evaluate_logical_plan( # Build and execute the streaming pipeline. # This must be done on all worker processes # for cluster == "distributed". - if ( - config_options.executor.cluster == "distributed" - ): # pragma: no cover; block depends on executor type and Distributed cluster - # Distributed execution: Use client.run - - # NOTE: Distributed execution requires Dask for now - from cudf_polars.experimental.rapidsmpf.dask import evaluate_pipeline_dask - - result, metadata_collector = evaluate_pipeline_dask( - evaluate_pipeline, - ir, - partition_info, - config_options, - stats, - collective_id_map, - collect_metadata=collect_metadata, - ) - elif config_options.executor.cluster == "spmd": - from cudf_polars.experimental.rapidsmpf.spmd import ( - evaluate_pipeline_spmd_mode, - ) + match config_options.executor.cluster: + case "distributed": # pragma: no cover; block depends on executor type and Distributed cluster + # Distributed execution: Use client.run + # NOTE: Distributed execution requires Dask for now + from cudf_polars.experimental.rapidsmpf.dask import ( + evaluate_pipeline_dask, + ) - result, metadata_collector = evaluate_pipeline_spmd_mode( - ir, - partition_info, - config_options, - stats, - collective_id_map, - collect_metadata=collect_metadata, - ) - else: - # Single-process execution: Run locally - result, metadata_collector = evaluate_pipeline( - ir, - partition_info, - config_options, - stats, - collective_id_map, - single_process_communicator(Options(), ProgressThread()), - collect_metadata=collect_metadata, - ) + result, metadata_collector = evaluate_pipeline_dask( + evaluate_pipeline, + ir, + partition_info, + config_options, + stats, + collective_id_map, + collect_metadata=collect_metadata, + ) + case "spmd": + from cudf_polars.experimental.rapidsmpf.spmd import ( + evaluate_pipeline_spmd_mode, + ) + + result, metadata_collector = evaluate_pipeline_spmd_mode( + ir, + partition_info, + config_options, + stats, + collective_id_map, + collect_metadata=collect_metadata, + ) + case "single": + # Single-process execution: Run locally + result, metadata_collector = evaluate_pipeline( + ir, + partition_info, + config_options, + stats, + collective_id_map, + single_process_communicator(Options(), ProgressThread()), + collect_metadata=collect_metadata, + ) + case other: + raise ValueError(f"Unknown cluster mode: {other}") return result, metadata_collector From e2c9aecd9804d244b67cd58e4f39497fefb7adf3 Mon Sep 17 00:00:00 2001 From: "Mads R. B. Kristensen" Date: Sun, 8 Mar 2026 22:15:12 +0100 Subject: [PATCH 02/33] Introduce Ray mode --- .../experimental/rapidsmpf/core.py | 13 + .../cudf_polars/experimental/rapidsmpf/ray.py | 538 ++++++++++++++++++ .../experimental/rapidsmpf/spmd.py | 2 +- .../cudf_polars/cudf_polars/utils/config.py | 4 + .../tests/experimental/rapidsmpf/test_ray.py | 146 +++++ 5 files changed, 702 insertions(+), 1 deletion(-) create mode 100644 python/cudf_polars/cudf_polars/experimental/rapidsmpf/ray.py create mode 100644 python/cudf_polars/tests/experimental/rapidsmpf/test_ray.py diff --git a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/core.py b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/core.py index d5a4e4457d0d..7de650f1d6d6 100644 --- a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/core.py +++ b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/core.py @@ -138,6 +138,19 @@ def evaluate_logical_plan( collective_id_map, collect_metadata=collect_metadata, ) + case "ray": + from cudf_polars.experimental.rapidsmpf.ray import ( + evaluate_pipeline_ray_mode, + ) + + result, metadata_collector = evaluate_pipeline_ray_mode( + ir, + partition_info, + config_options, + stats, + collective_id_map, + collect_metadata=collect_metadata, + ) case "single": # Single-process execution: Run locally result, metadata_collector = evaluate_pipeline( diff --git a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/ray.py b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/ray.py new file mode 100644 index 000000000000..ff4e756e7d65 --- /dev/null +++ b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/ray.py @@ -0,0 +1,538 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# SPDX-License-Identifier: Apache-2.0 +"""RapidsMPF streaming engine running on a Ray cluster.""" + +from __future__ import annotations + +import dataclasses +import os +import socket +from concurrent.futures import ThreadPoolExecutor +from contextlib import contextmanager +from typing import TYPE_CHECKING, Any, cast + +import ray +import ucxx._lib.libucxx as ucx_api +from rapidsmpf import bootstrap +from rapidsmpf.communicator.ucxx import barrier, get_root_ucxx_address, new_communicator +from rapidsmpf.config import ( + Options, + get_environment_variables, +) +from rapidsmpf.progress_thread import ProgressThread +from rapidsmpf.rmm_resource_adaptor import RmmResourceAdaptor +from rapidsmpf.statistics import Statistics +from rapidsmpf.streaming.core.actor import run_actor_network +from rapidsmpf.streaming.core.context import Context +from rapidsmpf.streaming.cudf.table_chunk import TableChunk + +import polars as pl + +import rmm.mr + +from cudf_polars.containers import DataFrame +from cudf_polars.dsl.ir import IRExecutionContext +from cudf_polars.experimental.rapidsmpf.core import generate_network +from cudf_polars.experimental.rapidsmpf.utils import empty_table_chunk +from cudf_polars.experimental.utils import _concat + +if TYPE_CHECKING: + from collections.abc import Iterator, MutableMapping + + from rapidsmpf.communicator.communicator import Communicator + from rapidsmpf.streaming.cudf.channel_metadata import ChannelMetadata + + from cudf_polars.dsl.ir import IR + from cudf_polars.experimental.base import PartitionInfo, StatsCollector + from cudf_polars.experimental.parallel import ConfigOptions + from cudf_polars.utils.config import StreamingExecutor + + +def evaluate_pipeline_ray_mode( + ir: IR, + partition_info: MutableMapping[IR, PartitionInfo], + config_options: ConfigOptions[StreamingExecutor], + stats: StatsCollector, + collective_id_map: dict[IR, list[int]], + *, + collect_metadata: bool = False, +) -> tuple[pl.DataFrame, list[ChannelMetadata] | None]: + """ + Evaluate a RapidsMPF streaming pipeline in Ray mode. + + The query is dispatched concurrently to every :class:`RankActor` in the + Ray cluster. Each actor evaluates the full pipeline on its local GPU and + participates in collective operations through the shared UCXX + communicator. The per-rank outputs are concatenated on the client before + being returned. + + Parameters + ---------- + ir + The IR node. + partition_info + The partition information. + config_options + Executor configuration, including the rapidsmpf context and the + Python thread-pool executor used to drive the actor network. + stats + The statistics collector. + collective_id_map + Mapping from IR nodes to their pre-allocated collective operation + IDs. + collect_metadata + Whether to collect runtime metadata. + + Returns + ------- + result + Concatenated output from all Ray actors as a Polars DataFrame. + metadata + Collected channel metadata if ``collect_metadata`` is ``True``, + otherwise ``None``. + + Raises + ------ + RuntimeError + If the configured executor runtime is not ``"rapidsmpf"``. + RuntimeError + If ``config_options.executor.ray_client`` is ``None``. + """ + if config_options.executor.runtime != "rapidsmpf": + raise RuntimeError("Runtime must be rapidsmpf") + if config_options.executor.ray_client is None: + raise RuntimeError("ray_client must be set when cluster='ray'") + rank_actors = config_options.executor.ray_client._rank_actors + + # Strip ray_client before pickling config_options for remote calls: + # actors don't need the full actor list, and sending actor handles to each + # actor is wasteful. + actor_config_options = dataclasses.replace( + config_options, + executor=dataclasses.replace(config_options.executor, ray_client=None), + ) + + result = ray.get( + [ + rank.evaluate_polars_ir.remote( + ir, + partition_info, + actor_config_options, + stats, + collective_id_map, + collect_metadata=collect_metadata, + ) + for rank in rank_actors + ] + ) + dfs: list[pl.DataFrame] = [] + metadata_collector: list[ChannelMetadata] = [] + for df, md in result: + dfs.append(df) + if md is not None: + metadata_collector.extend(md) + + return pl.concat(dfs), metadata_collector or None + + +@ray.remote( + max_restarts=0, + max_task_retries=0, + num_cpus=0, + num_gpus=1, +) +class RankActor: + """ + Ray actor that owns one GPU and participates in a RapidsMPF cluster. + + Each actor manages its own memory resource, statistics collector, + communicator, and streaming context. Collectively, the actors form a + SPMD execution cluster used by the client-side Ray integration. + + Parameters + ---------- + nranks + Total number of actors, typically one per GPU in the Ray cluster. + rapidsmpf_options_as_bytes + Serialized RapidsMPF options produced by + :meth:`rapidsmpf.config.Options.serialize`. + executor_options + Additional executor options forwarded from the client. + """ + + def __init__( + self, + *, + nranks: int, + rapidsmpf_options_as_bytes: bytes, + executor_options: dict[str, object], + ) -> None: + self._mr = RmmResourceAdaptor(rmm.mr.CudaAsyncMemoryResource()) + self._rapidsmpf_options: Options = Options.deserialize( + rapidsmpf_options_as_bytes + ) + self._statistics: Statistics = Statistics.from_options( + self._mr, self._rapidsmpf_options + ) + self._nranks: int = nranks + self._py_executor = ThreadPoolExecutor( + max_workers=cast( + int, executor_options.get("rapidsmpf_py_executor_max_workers", 1) + ), + thread_name_prefix="ray-executor", + ) + self._comm: Communicator | None = None + self._ctx: Context | None = None + + def setup_root(self) -> bytes: + """ + Initialize this actor as the root rank. + + The root actor creates a new UCXX communicator without an existing + root address. The resulting root address is returned so it can be + distributed to all other actors during bootstrap. + + Returns + ------- + Serialized UCXX root address for communicator bootstrap. + """ + self._comm = new_communicator( + nranks=self._nranks, + ucx_worker=None, + root_ucxx_address=None, + options=self._rapidsmpf_options, + progress_thread=ProgressThread(self._statistics), + ) + return get_root_ucxx_address(self._comm) + + def setup_worker(self, root_ucxx_address_as_bytes: bytes) -> None: + """ + Complete communicator bootstrap and create the streaming context. + + This method must be called concurrently on all actors, including the + root. Non-root actors connect to the root using the provided UCXX + address. Once all ranks have joined, the actors synchronize with a + barrier and create their RapidsMPF streaming contexts. + + Parameters + ---------- + root_ucxx_address_as_bytes + Serialized UCXX root address returned by :meth:`setup_root`. + """ + if self._comm is None: + root_ucxx_address = ucx_api.UCXAddress.create_from_buffer( + root_ucxx_address_as_bytes + ) + self._comm = new_communicator( + nranks=self._nranks, + ucx_worker=None, + root_ucxx_address=root_ucxx_address, + options=self._rapidsmpf_options, + progress_thread=ProgressThread(self._statistics), + ) + barrier(self._comm) + self._ctx = Context.from_options( + self._comm.logger, self._mr, self._rapidsmpf_options + ) + + def shutdown(self) -> None: + """ + Release actor-owned resources and exit the process. + + This shuts down the local Python executor, drops communicator and + memory-resource references, and then terminates the Ray actor process. + """ + self._py_executor.shutdown(wait=True, cancel_futures=True) + self._comm = None + self._mr = None + ray.actor.exit_actor() + + def get_info(self) -> dict: + """ + Return diagnostic information about actor placement. + + Returns + ------- + Diagnostic information about this actor's placement and state. + """ + return { + "pid": os.getpid(), + "hostname": socket.gethostname(), + "cuda_visible_devices": os.environ.get("CUDA_VISIBLE_DEVICES"), + "node_id": ray.get_runtime_context().get_node_id(), + } + + def evaluate_polars_ir( + self, + ir: IR, + partition_info: MutableMapping[IR, PartitionInfo], + config_options: ConfigOptions[StreamingExecutor], + stats: StatsCollector, + collective_id_map: dict[IR, list[int]], + *, + collect_metadata: bool, + ) -> tuple[pl.DataFrame, list[ChannelMetadata] | None]: + """ + Execute a Polars IR query on this actor's GPU. + + The IR is lowered to a RapidsMPF actor network, executed locally, and + the resulting output messages are assembled into a Polars DataFrame. + Collective operations in the network communicate with peer actors + through the shared UCXX communicator. + + Parameters + ---------- + ir + Root IR node describing the query to execute. + partition_info + Per-node partition metadata produced by the planner. + config_options + Executor configuration forwarded from the client. + stats + Statistics collector used during execution. + collective_id_map + Mapping from IR nodes to their pre-allocated collective operation + IDs. + collect_metadata + If ``True``, collect channel metadata during execution. + + Returns + ------- + result + This rank's output fragment as a Polars DataFrame. + metadata + Collected channel metadata if ``collect_metadata`` is ``True``, + otherwise ``None``. + + Raises + ------ + AssertionError + If :meth:`setup_worker` has not been called first. + """ + assert self._ctx is not None, ( + "setup_worker must be called before evaluate_polars_ir" + ) + ir_context = IRExecutionContext(get_cuda_stream=self._ctx.get_stream_from_pool) + metadata_collector: list[ChannelMetadata] | None = ( + [] if collect_metadata else None + ) + + nodes, output = generate_network( + self._ctx, + self._comm, + ir, + partition_info, + config_options, + stats, + ir_context=ir_context, + collective_id_map=collective_id_map, + metadata_collector=metadata_collector, + ) + + run_actor_network(actors=nodes, py_executor=self._py_executor) + + messages = output.release() + chunks = [ + TableChunk.from_message(msg).make_available_and_spill( + self._ctx.br(), allow_overbooking=True + ) + for msg in messages + ] + dfs: list[DataFrame] + if chunks: + dfs = [ + DataFrame.from_table( + chunk.table_view(), + list(ir.schema.keys()), + list(ir.schema.values()), + chunk.stream, + ) + for chunk in chunks + ] + df = _concat(*dfs, context=ir_context) + else: + # No chunks received, create an empty DataFrame with the correct schema. + stream = ir_context.get_cuda_stream() + chunk = empty_table_chunk(ir, self._ctx, stream) + df = DataFrame.from_table( + chunk.table_view(), + list(ir.schema.keys()), + list(ir.schema.values()), + stream, + ) + result = df.to_polars() + return result, metadata_collector + + +class RayClient: + """Client-side handle for the distributed cudf-polars execution.""" + + def __init__(self, rank_actors: list[Any]) -> None: + self._rank_actors: list[Any] = rank_actors + + @property + def nranks(self) -> int: + """ + Number of Ray rank actors. + + Returns + ------- + Number of ranks/nodes in the Ray cluster. + """ + return len(self._rank_actors) + + def gather_cluster_info(self) -> list[dict]: + """ + Collect diagnostic information from every rank actor. + + Returns + ------- + List of info dicts (see :meth:`RankActor.get_info`), one per rank + in rank order. + + Examples + -------- + >>> with ray_execution() as (ray_client, engine): # doctest: +SKIP + ... for i, info in enumerate(ray_client.gather_cluster_info()): + ... print(f"rank {i}: {info}") + """ + return ray.get([rank.get_info.remote() for rank in self._rank_actors]) + + +@contextmanager +def ray_execution( + *, + rapidsmpf_options: Options | None = None, + executor_options: dict[str, object] | None = None, + engine_kwargs: dict[str, Any] | None = None, + ray_init_kwargs: dict[str, object] | None = None, +) -> Iterator[tuple[RayClient, pl.GPUEngine]]: + """ + Create a RapidsMPF Ray cluster and matching Polars GPU engine. + + If Ray is not already initialized, this context manager calls + :func:`ray.init` on entry and :func:`ray.shutdown` on exit. If Ray is + already initialized, cluster lifetime remains managed by the caller. + + Parameters + ---------- + rapidsmpf_options + RapidsMPF options forwarded to every actor. If ``None``, defaults to + ``Options(get_environment_variables())``. + executor_options + Additional key-value pairs forwarded to the Polars executor options. + engine_kwargs + Additional keyword arguments forwarded to :class:`polars.GPUEngine`. + ray_init_kwargs + Keyword arguments forwarded to :func:`ray.init` when Ray is not + already initialized. + + Yields + ------ + ray_client + Client-side handle to the Ray actor cluster. + engine + Polars GPU engine configured to execute through RapidsMPF on Ray. + + Raises + ------ + RuntimeError + If called from within an ``rrun`` cluster. + RuntimeError + If not all GPUs in the Ray cluster are free at startup. + RuntimeError + If no GPUs are available in the Ray cluster. + ValueError + If ``executor_options`` contains a reserved key. + ValueError + If ``engine_kwargs`` contains a reserved key. + + Examples + -------- + >>> with ray_execution() as (ray_client, engine): # doctest: +SKIP + ... result = pl.LazyFrame({"a": [1, 2, 3]}).collect(engine=engine) + """ + executor_options = executor_options or {} + engine_kwargs = engine_kwargs or {} + ray_init_kwargs = ray_init_kwargs or {} + + if bootstrap.is_running_with_rrun(): + raise RuntimeError( + "ray_execution() must not be called from within an rrun cluster. Instead " + "launch the rrun cluster separately and let this client connect to its " + "cluster nodes." + ) + + # Check for reserved keys. + if bad := {"runtime", "cluster", "spmd", "ray_client"} & executor_options.keys(): + raise ValueError(f"executor_options may not contain reserved keys: {bad}") + if bad := {"memory_resource", "executor"} & engine_kwargs.keys(): + raise ValueError(f"engine_kwargs may not contain reserved keys: {bad}") + + rapidsmpf_options = ( + rapidsmpf_options + if rapidsmpf_options is not None + else Options(get_environment_variables()) + ) + rapidsmpf_options_as_bytes = rapidsmpf_options.serialize() + + ray_was_initialized: bool = ray.is_initialized() + if not ray_was_initialized: + # Prevent Ray from overriding CUDA_VISIBLE_DEVICES to "" when a worker + # process starts with zero visible GPUs (e.g. the driver process itself). + # Without this, Ray's accelerator detection resets the variable before our + # actors acquire their GPU assignment, hiding all GPUs from CUDA. + os.environ.setdefault("RAY_ACCEL_ENV_VAR_OVERRIDE_ON_ZERO", "0") + ray.init(**ray_init_kwargs) + + total_gpus = int(ray.cluster_resources().get("GPU", 0.0)) + free_gpus = int(ray.available_resources().get("GPU", 0.0)) + if total_gpus != free_gpus: + raise RuntimeError( + "Ray execution expects all GPUs in the Ray cluster to be available at startup" + ) + if free_gpus == 0: + raise RuntimeError("No available GPUs in the Ray cluster at startup") + + # Create one actor per GPU. Ray adds .remote() dynamically; no type stubs. + rank_actors: list[Any] = [ + RankActor.remote( # type: ignore[attr-defined] + nranks=free_gpus, + executor_options=executor_options, + rapidsmpf_options_as_bytes=rapidsmpf_options_as_bytes, + ) + for _ in range(free_gpus) + ] + + root_ucxx_address_as_bytes = ray.get(rank_actors[0].setup_root.remote()) + # 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) for rank in rank_actors] + ) + + try: + ray_client = RayClient(rank_actors) + engine = pl.GPUEngine( + memory_resource=None, + executor="streaming", + executor_options={ + **executor_options, + "runtime": "rapidsmpf", + "cluster": "ray", + "ray_client": ray_client, + }, + **engine_kwargs, + ) + yield ray_client, engine + finally: + for a in rank_actors: + try: + ray.get(a.shutdown.remote()) + except ray.exceptions.RayActorError: + pass # expected: exit_actor() terminates the process immediately. + except Exception as e: + print(f"shutdown error: {e}") + if not ray_was_initialized: + ray.shutdown() diff --git a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/spmd.py b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/spmd.py index 6b90860e4ca6..abf797a7ee35 100644 --- a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/spmd.py +++ b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/spmd.py @@ -286,7 +286,7 @@ def spmd_execution( when ``None``. executor_options Extra keyword arguments forwarded to the ``executor_options`` dict of - :class:`~polars.lazyframe.engine_config.GPUEngine`. The keys + :class:`~polars.lazyframe.engine_config.GPUEngine`. The keys ``"runtime"``, ``"cluster"``, and ``"spmd"`` are reserved and may not be overridden. **engine_kwargs diff --git a/python/cudf_polars/cudf_polars/utils/config.py b/python/cudf_polars/cudf_polars/utils/config.py index e6a469558ce8..55d80fec5ab8 100644 --- a/python/cudf_polars/cudf_polars/utils/config.py +++ b/python/cudf_polars/cudf_polars/utils/config.py @@ -43,6 +43,8 @@ import rmm.mr + from cudf_polars.experimental.rapidsmpf.ray import RayClient + __all__ = [ "Cluster", @@ -173,6 +175,7 @@ class Cluster(enum.StrEnum): SINGLE = "single" DISTRIBUTED = "distributed" SPMD = "spmd" + RAY = "ray" class Scheduler(enum.StrEnum): @@ -873,6 +876,7 @@ class StreamingExecutor: ) ) spmd: SPMDContext | None = None + ray_client: RayClient | None = None def __post_init__(self) -> None: # noqa: D105 # Check for rapidsmpf runtime diff --git a/python/cudf_polars/tests/experimental/rapidsmpf/test_ray.py b/python/cudf_polars/tests/experimental/rapidsmpf/test_ray.py new file mode 100644 index 000000000000..53622aec1e56 --- /dev/null +++ b/python/cudf_polars/tests/experimental/rapidsmpf/test_ray.py @@ -0,0 +1,146 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# SPDX-License-Identifier: Apache-2.0 +"""Tests for Ray execution mode.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Any + +import pytest + +import polars as pl + +ray = pytest.importorskip("ray") + +from cudf_polars.experimental.rapidsmpf.ray import ( # noqa: E402 + RayClient, + ray_execution, +) + +if TYPE_CHECKING: + from collections.abc import Iterator + + +@pytest.fixture(scope="session") +def _ray_env() -> Iterator[tuple[RayClient, pl.GPUEngine]]: + """Create one Ray cluster + GPU actors shared across the test session.""" + with ray_execution(ray_init_kwargs={"include_dashboard": False}) as ( + ray_client, + engine, + ): + yield ray_client, engine + + +@pytest.fixture(scope="session") +def ray_client(_ray_env: tuple[RayClient, pl.GPUEngine]) -> RayClient: + """Session-scoped Ray cluster client.""" + return _ray_env[0] + + +@pytest.fixture(scope="session") +def engine(_ray_env: tuple[RayClient, pl.GPUEngine]) -> pl.GPUEngine: + """Session-scoped GPU engine backed by the Ray cluster.""" + return _ray_env[1] + + +pytestmark = [ + # Ray's internal subprocess management leaks /dev/null file handles; + # suppress the resulting ResourceWarning noise from its internals. + pytest.mark.filterwarnings("ignore::ResourceWarning"), +] + + +# --------------------------------------------------------------------------- +# Context-manager smoke tests (no GPU required) +# --------------------------------------------------------------------------- + + +def test_ray_execution_reserved_executor_keys() -> None: + """executor_options rejects reserved keys.""" + for key in ("runtime", "cluster", "spmd", "ray_client"): + with ( + pytest.raises(ValueError, match="reserved"), + ray_execution(executor_options={key: "anything"}), + ): + pass + + +def test_ray_execution_reserved_engine_kwargs_keys() -> None: + """engine_kwargs rejects keys that are set explicitly by ray_execution.""" + for key in ("memory_resource", "executor"): + kwargs: dict[str, Any] = {key: "anything"} + with ( + pytest.raises(ValueError, match="reserved"), + ray_execution(engine_kwargs=kwargs), + ): + pass + + +# --------------------------------------------------------------------------- +# GPU tests — share a single Ray cluster + actor set for the whole session +# --------------------------------------------------------------------------- + + +def test_ray_execution_yields_client_and_engine( + ray_client: RayClient, + engine: pl.GPUEngine, +) -> None: + """ray_execution yields a (RayClient, GPUEngine) pair.""" + assert isinstance(ray_client, RayClient) + assert isinstance(engine, pl.GPUEngine) + assert ray_client.nranks >= 1 + + +def test_gather_cluster_info(ray_client: RayClient) -> None: + """gather_cluster_info returns one info dict per rank with expected fields.""" + infos = ray_client.gather_cluster_info() + assert len(infos) == ray_client.nranks + for info in infos: + assert "node_id" in info + assert "hostname" in info + assert "pid" in info + assert "cuda_visible_devices" in info + assert "comm_rank" in info + assert isinstance(info["pid"], int) + # comm_rank is set after setup_worker; verify it is a valid rank index. + assert info["comm_rank"] in range(ray_client.nranks) + # Each actor runs in its own process. + assert len({info["pid"] for info in infos}) == ray_client.nranks + + +def test_ray_execution_scan(engine: pl.GPUEngine) -> None: + """Input rows are partitioned across actors; total output equals input.""" + lf = pl.LazyFrame({"a": [1, 2, 3]}) + result = lf.collect(engine=engine) + assert result.shape == (3, 1) + assert sorted(result["a"].to_list()) == [1, 2, 3] + + +def test_ray_execution_filter(engine: pl.GPUEngine) -> None: + """Filter is applied correctly across all actors.""" + lf = pl.LazyFrame({"a": [1, 2, 3, 4, 5]}) + result = lf.filter(pl.col("a") > 3).collect(engine=engine) + assert result.shape == (2, 1) + assert sorted(result["a"].to_list()) == [4, 5] + + +def test_ray_execution_group_by(engine: pl.GPUEngine) -> None: + """Group-by produces the correct global aggregation across all actors.""" + lf = pl.LazyFrame({"key": ["a", "a", "b"], "val": [1, 2, 3]}) + result = ( + lf.group_by("key").agg(pl.col("val").sum()).collect(engine=engine).sort("key") + ) + assert result.shape == (2, 2) + assert result["key"].to_list() == ["a", "b"] + assert result["val"].to_list() == [3, 3] + + +def test_ray_execution_empty_dataframe(engine: pl.GPUEngine) -> None: + """An empty LazyFrame produces an empty result with the correct schema.""" + lf = pl.LazyFrame( + {"a": pl.Series([], dtype=pl.Int32), "b": pl.Series([], dtype=pl.Float64)} + ) + result = lf.collect(engine=engine) + assert result.shape == (0, 2) + assert result.columns == ["a", "b"] + assert result.dtypes == [pl.Int32, pl.Float64] From 66572ff67cee3e3c626621c2734cfc7d9f0e70bd Mon Sep 17 00:00:00 2001 From: "Mads R. B. Kristensen" Date: Tue, 10 Mar 2026 17:16:35 +0100 Subject: [PATCH 03/33] docs --- .../cudf_polars/docs/cudf-polars-mp-design.md | 427 ++++++++++++++++++ python/cudf_polars/docs/cudf_polars_mp.md | 257 +++++++++++ 2 files changed, 684 insertions(+) create mode 100644 python/cudf_polars/docs/cudf-polars-mp-design.md create mode 100644 python/cudf_polars/docs/cudf_polars_mp.md diff --git a/python/cudf_polars/docs/cudf-polars-mp-design.md b/python/cudf_polars/docs/cudf-polars-mp-design.md new file mode 100644 index 000000000000..a912f5cff63a --- /dev/null +++ b/python/cudf_polars/docs/cudf-polars-mp-design.md @@ -0,0 +1,427 @@ +# cudf-polars Multi-GPU Design + +This document describes the multi-GPU execution architecture of cudf-polars. +For user-facing setup instructions, see +`cudf_polars_mp.md`. + +- [1. Architecture Overview](#1-architecture-overview) +- [2. SPMD Mode](#2-spmd-mode) +- [3. Ray Mode](#3-ray-mode) +- [4. Comparison](#4-comparison) +- [5. Hardware Mapping (GPU Pinning)](#5-hardware-mapping-gpu-pinning) + +--- + +## 1. Architecture Overview + +Both modes share the same bottom two layers (RapidsMPF engine and SPMD +cluster) and differ only in how the user script interacts with them. + +### SPMD mode + +The user script runs on **all N workers simultaneously**. There is no separate +client — every process is both driver and worker. + +``` + rank 0 rank 1 ... rank N-1 +┌─────────────────┐ ┌─────────────────┐ ┌─────────────────┐ +│ User script │ │ User script │ │ User script │ +│ (same code on │ │ (same code on │ │ (same code on │ +│ every rank) │ │ every rank) │ │ every rank) │ +└────────┬────────┘ └────────┬────────┘ └────────┬────────┘ + │ │ │ + │ LazyFrame.collect(engine=engine) │ + ↓ ↓ ↓ +┌─────────────────┐ ┌─────────────────┐ ┌─────────────────┐ +│ run IR │ │ run IR │ │ run IR │ +└────────┬────────┘ └────────┬────────┘ └────────┬────────┘ + │ │ │ + ↓ ↓ ↓ +┌────────────────────────────────────────────────────────────────┐ +│ RapidsMPF streaming engine │ +│ shuffle / all-gather · UCXX communicator · RMM GPU memory │ +└────────────────────────────────────────────────────────────────┘ + ↑ ↑ ↑ + GPU 0 GPU 1 GPU N-1 +``` + +Results are rank-local after `collect`. Call `allgather_polars_dataframe()` to +assemble the full dataset on every rank. + +### Ray mode + +A single driver script dispatches work to N `RankActor` Ray actors (one per +GPU). The driver never touches a GPU directly. + +``` + ┌──────────────────────────────┐ + │ User script │ + │ (single driver process) │ + │ LazyFrame.collect(engine=…) │ + └──────────────┬───────────────┘ + │ IR dispatched to all actors + ┌────────────────|─────────────────┐ + ↓ ↓ ↓ + ┌─────────────┐ ┌─────────────┐ ┌─────────────┐ + │ RankActor │ │ RankActor │ │ RankActor │ + │ rank 0 │ │ rank 1 │ │ rank N-1 │ + │ run IR │ │ run IR │ │ run IR │ + └──────┬──────┘ └──────┬──────┘ └──────┬──────┘ + ↓ ↓ ↓ +┌────────────────────────────────────────────────────────────────┐ +│ RapidsMPF streaming engine │ +│ shuffle / all-gather · UCXX communicator · RMM GPU memory │ +└────────────────────────────────────────────────────────────────┘ + ↑ ↑ ↑ + GPU 0 GPU 1 GPU N-1 +``` + +Per-rank output fragments are concatenated on the driver before being returned. +No `allgather` step is needed. + +### Key insight — why client modes exist + +In SPMD mode every process runs the same code, which is unfamiliar to users +accustomed to single-process or Dask-style driver/worker workflows. Client +frontends such as Ray let users write a normal single-process script while the +cluster handles distribution transparently. The underlying engine is unchanged; +only the dispatch layer differs. Future frontends (Dask, custom clients) can +target the same SPMD cluster without modifying the engine. + +--- + +## 2. SPMD Mode + +### Execution model + +The user script is launched with `rrun -n N python script.py`. `rrun` starts N +identical processes, each pinned to one GPU. There is no separate client +process — every process runs the full script, acting simultaneously as driver +and worker on its rank-local data. + +Because every rank runs independent Python, a `pl.DataFrame` is always +*rank-local*: it holds only that rank's fragment of the distributed dataset. +File-based sources (`scan_parquet`, `scan_csv`) distribute work automatically +— the engine assigns disjoint file- or row-group ranges to each rank. + +### Bootstrapping + +`spmd_execution()` calls `bootstrap.create_ucxx_comm(type=BackendType.AUTO)`. +Under `rrun`, `BackendType.AUTO` resolves to the `rrun`-native bootstrap +mechanism, connecting all N ranks without additional configuration. + +### GPU assignment + +`rrun` sets `CUDA_VISIBLE_DEVICES` for each process automatically. Each rank +sees exactly one GPU. No user action is required. + +### Query symmetry requirement + +All ranks must issue the **same** sequence of Polars queries in the **same** +order. Collective operations (shuffles, all-gathers, joins) are matched across +ranks by a monotonically increasing operation ID. If one rank calls a +collective that another does not, all ranks will deadlock. Driver logic must be +fully deterministic: avoid rank-conditional `collect` calls, early exits, or +branching that causes different ranks to execute different query graphs. + +### Entry point + +```python +from cudf_polars.experimental.rapidsmpf.spmd import ( + spmd_execution, + allgather_polars_dataframe, +) +``` + +`spmd_execution()` is a context manager that yields `(comm, ctx, engine)`: + +| Object | Type | Purpose | +|----------|----------------|---------------------------------------------| +| `comm` | `Communicator` | Active RapidsMPF communicator | +| `ctx` | `Context` | RapidsMPF context (GPU memory, stream pool) | +| `engine` | `pl.GPUEngine` | Polars GPU engine wired to `comm` and `ctx` | + +### Collecting results + +Each `collect(engine=engine)` returns a rank-local `pl.DataFrame`. To +assemble the full result on every rank, call `allgather_polars_dataframe()`: + +```python +full = allgather_polars_dataframe( + comm=comm, ctx=ctx, local_df=result, op_id=0 +) +``` + +`op_id` must be the same on every rank and must be unique across all +`allgather_polars_dataframe` calls in the script. + +### Example + +```python +# Launch with: rrun -n 4 python spmd_script.py +import polars as pl +from cudf_polars.experimental.rapidsmpf.spmd import ( + spmd_execution, + allgather_polars_dataframe, +) + +with spmd_execution() as (comm, ctx, engine): + result = ( + pl.scan_parquet("/data/dataset/*.parquet") + .filter(pl.col("value") > 0) + .group_by("category") + .agg(pl.col("value").sum()) + .collect(engine=engine) + ) + # result is rank-local; gather all fragments onto every rank + full = allgather_polars_dataframe( + comm=comm, ctx=ctx, local_df=result, op_id=0 + ) + # full is now identical on every rank + print(full) +``` + +### Options pass-through + +```python +with spmd_execution( + executor_options={"rapidsmpf_py_executor_max_workers": 2}, + parquet_options={"use_rapidsmpf_native": True}, +) as (comm, ctx, engine): + ... +``` + +Reserved `executor_options` keys: `"runtime"`, `"cluster"`, `"spmd"`. +Reserved `engine_kwargs` keys: `"memory_resource"`, `"executor"`. + +--- + +## 3. Ray Mode + +### Execution model + +The user runs a single driver script. `ray_execution()` creates N `RankActor` +Ray remote actors — one per available GPU. The actors form a private SPMD +cluster; the driver dispatches Polars IR to them and receives concatenated +results directly, with no `allgather` step needed on the client. + +### Bootstrapping + +1. `ray_execution()` creates N `RankActor` instances (each requesting + `num_gpus=1` from Ray's resource scheduler). +2. Root actor (rank 0) calls `setup_root()` → returns its UCXX address. +3. All N actors (including the root) call `setup_worker(root_ucxx_address)` + **concurrently**. Non-root actors connect to the root; the root skips + communicator creation and proceeds directly to the barrier. All ranks must + reach the barrier simultaneously. +4. After `setup_worker` completes, each actor holds a fully initialized + `Communicator` and `Context`. + +### GPU assignment + +Ray's resource scheduler assigns `num_gpus=1` to each `RankActor` before the +actor process starts, setting `CUDA_VISIBLE_DEVICES` automatically. The +`RAY_ACCEL_ENV_VAR_OVERRIDE_ON_ZERO=0` environment variable prevents Ray from +overriding `CUDA_VISIBLE_DEVICES` to empty on the driver process (which has +zero GPUs assigned). + +For hardware placement details, see [Section 5](#5-hardware-mapping-gpu-pinning). + +### Entry point + +```python +from cudf_polars.experimental.rapidsmpf.ray import ray_execution +``` + +`ray_execution()` is a context manager that yields `(ray_client, engine)`: + +| Object | Type | Purpose | +|--------------|----------------|----------------------------------------| +| `ray_client` | `RayClient` | Client handle to the actor cluster | +| `engine` | `pl.GPUEngine` | Polars GPU engine backed by Ray actors | + +### Results + +`collect(engine=engine)` dispatches the query to all actors, concatenates +their per-rank output fragments on the client, and returns a single +`pl.DataFrame`. No `allgather` is needed. + +### Ray lifecycle + +`ray_execution()` calls `ray.init()` on entry if Ray is not already +initialized, and `ray.shutdown()` on exit. If `ray.is_initialized()` returns +`True` before entry, the caller manages the cluster lifetime. + +### Diagnostics + +```python +for i, info in enumerate(ray_client.gather_cluster_info()): + print(f"rank {i}: {info}") +# Each info dict contains: pid, hostname, cuda_visible_devices, node_id +``` + +### Example + +```python +# Launch with: python ray_script.py +import polars as pl +from cudf_polars.experimental.rapidsmpf.ray import ray_execution + +with ray_execution() as (ray_client, engine): + print(ray_client.gather_cluster_info()) # verify actor placement + + result = ( + pl.scan_parquet("/data/dataset/*.parquet") + .filter(pl.col("value") > 0) + .group_by("category") + .agg(pl.col("value").sum()) + .collect(engine=engine) + ) + # result is the full concatenated output, returned directly to the driver + print(result) +``` + +### Options pass-through + +```python +with ray_execution( + executor_options={"rapidsmpf_py_executor_max_workers": 2}, + engine_kwargs={"parquet_options": {"use_rapidsmpf_native": True}}, + ray_init_kwargs={"address": "auto"}, +) as (ray_client, engine): + ... +``` + +Reserved `executor_options` keys: `"runtime"`, `"cluster"`, `"spmd"`, +`"ray_client"`. Reserved `engine_kwargs` keys: `"memory_resource"`, +`"executor"`. + +--- + +## 4. Comparison + +| | SPMD | Ray | Future (Dask / custom) | +|------------------|----------------------------------|-----------------------------------|------------------------| +| Driver | Script runs on **all N workers** | Single client process | Single client process | +| Launch | `rrun -n N python script.py` | `python script.py` | `python script.py` | +| GPU pinning | `rrun` auto-pins each rank | Ray scheduler (one actor per GPU) | Depends on frontend | +| Result delivery | Rank-local; `allgather` needed | Concatenated, returned to client | Returned to client | +| Query symmetry | Required (all ranks same order) | Not required | Not required | +| Extra dependency | `rrun` / RapidsMPF | `ray` | `dask` / none | + +**Note on future frontends.** Dask and custom clients are not yet implemented, +but the architecture supports them. A new frontend only needs to: + +1. Manage a pool of SPMD workers (one per GPU). +2. Bootstrap the UCXX communicator across those workers. +3. Dispatch pickled IR + `partition_info` + `collective_id_map` to all workers + concurrently and collect their output fragments. + +No changes to the underlying RapidsMPF streaming engine are required. + +--- + +## 5. Hardware Mapping (GPU Pinning) + +### Launch model + +Hardware mapping depends on how ranks are **launched**. + +Two launch models are supported: + +* **SPMD mode**, where ranks are started directly by the `rrun` launcher. +* **Ray mode**, where ranks run inside Ray actors scheduled by the Ray runtime. + +`rrun` is a lightweight process launcher designed for GPU workloads. It starts +multiple ranks, assigns GPUs, and optionally applies topology-aware bindings. + +A typical SPMD program is started with: + +```bash +rrun -n 4 python script.py +``` + +This launches four identical Python processes. Each process becomes a **rank** +in the SPMD program and runs the same code independently. + +Because `rrun` performs the hardware setup **before the program starts**, it +must be used to launch the application. It cannot attach to or configure +processes that are already running. + +Other execution modes may require such functionality. For example, in **Ray +mode** the ranks run inside actors created by the Ray runtime rather than by +`rrun`, so GPU assignment is handled by Ray instead. + +--- + +### SPMD mode + +In SPMD mode, `rrun` assigns one GPU to each rank. + +By default it detects all GPUs on the node, but a specific list can be +provided: + +```bash +rrun -n 4 -g 0,1,2,3 python script.py +``` + +Each rank receives a single GPU via `CUDA_VISIBLE_DEVICES`. Inside the process +this GPU always appears as **device 0**, so CUDA programs require no special +configuration. + +If more ranks than GPUs are launched, multiple ranks will share a GPU. + +`rrun` can also apply topology-aware bindings so that each rank runs close to +its GPU. This includes CPU affinity, NUMA memory locality, and network-device +selection. Bindings are enabled by default and can be disabled with: + +```bash +rrun -n 4 --bind-to none python script.py +``` + +If topology discovery is unavailable, bindings are skipped automatically. + +`rrun` also sets environment variables used by RapidsMPF to bootstrap the +communication backend. + +When running under Slurm (for example with `srun`), Slurm launches the ranks +across nodes while `rrun` performs the same local setup for each rank. + +--- + +### Ray mode + +In Ray mode, ranks run as Ray actors scheduled by the Ray runtime rather than +being launched by `rrun`. + +Each rank is defined with: + +```python +@ray.remote(num_gpus=1) +``` + +Ray's scheduler selects a node with a free GPU, launches the actor there, and +sets `CUDA_VISIBLE_DEVICES` before the Python code starts. As in SPMD mode, +each rank therefore sees its assigned GPU as **device 0**. + +Unlike SPMD mode, the processes already exist when RapidsMPF code begins +executing. Ray creates the worker processes and then runs the user code inside +them. This means the `rrun` launcher cannot perform hardware setup ahead of +time. + +#### Future Request for `rrun` +To support such execution models, `rrun` will need a **library API** that can +configure an already-running process. Conceptually, the workflow would look +like: + +1. Ray launches one worker per GPU. +2. Each worker determines which GPU it has been assigned. +3. The worker calls a new `rrun` API to apply the hardware bindings locally. + +This API would configure the process in-place, for example by applying CPU +affinity, NUMA bindings, and other topology-aware settings for the specified +GPU. The GPU could be specified by index or by a stable identifier such as a +GPU UUID. + +This capability is not currently provided by `rrun`, which today only supports +processes that it launches itself, see [Launch model](#launch-model). diff --git a/python/cudf_polars/docs/cudf_polars_mp.md b/python/cudf_polars/docs/cudf_polars_mp.md new file mode 100644 index 000000000000..ba37cff0b318 --- /dev/null +++ b/python/cudf_polars/docs/cudf_polars_mp.md @@ -0,0 +1,257 @@ +# cudf-polars-mp + +`cudf-polars-mp` extends Polars query execution to multiple GPUs. + +Multi-process (mp) execution distributes a query across several GPU workers. Each +worker owns a disjoint fragment of the data and participates in collective operations +(shuffles, all-gathers, joins) to produce a globally correct result. + +The entry point in all cases is the Polars `GPUEngine` with `executor="streaming"`. +The `cluster` option selects the execution model: + +| `cluster` value | Description | Status | +| --------------- | ------------------------------------------- | --------------- | +| `"single"` | Single-GPU, in-process execution | Stable (legacy) | +| `"distributed"` | Multi-GPU via Dask Distributed | Stable (legacy) | +| `"spmd"` | Multi-GPU via SPMD with the `rrun` launcher | Experimental | +| `"ray"` | Multi-GPU via Ray actors | Experimental | + +This document describes the two experimental multi-GPU modes. Both rely on RapidsMPF +for shuffle and collective communication. + +* [SPMD cluster mode](#spmd-cluster-mode) +* [Ray cluster mode](#ray-cluster-mode) + +--- + +# SPMD cluster mode + +In SPMD (Single Program, Multiple Data) execution, the same Python script is launched +multiple times simultaneously, once per GPU, using the `rrun` launcher bundled with +RapidsMPF. Each process is assigned a GPU and receives a **rank**. Ranks communicate +through a UCXX-based communicator established at startup. + +Each rank runs an independent Python process and owns its local data. File-based +sources (`scan_parquet`, `scan_csv`, etc.) are automatically partitioned so that +different ranks read different file or row-group ranges. In-memory `DataFrame` +objects are already rank-local, so each rank processes its own copy. + +## Prerequisites + +* RapidsMPF (`rapidsmpf`) installed +* UCXX available (usually installed with RapidsMPF) +* `rrun` launcher available (`rrun --help` should succeed) + +## Running in SPMD mode + +`spmd_execution()` is the primary entry point for SPMD execution. It is a context +manager imported from `cudf_polars.experimental.rapidsmpf.spmd`. On entry it: + +1. Bootstraps a UCXX communicator connecting all ranks. +2. Creates a RapidsMPF streaming `Context` that owns GPU memory and a CUDA stream pool. +3. Constructs and yields a `pl.GPUEngine` bound to that context. + +All resources are released when the context exits. + +`spmd_execution()` must run inside an `rrun` cluster. It raises `RuntimeError` +if `rapidsmpf.bootstrap.is_running_with_rrun()` returns `False`. + +```python +# launch with: rrun -n 4 python my_script.py +import polars as pl +from cudf_polars.experimental.rapidsmpf.spmd import ( + spmd_execution, + allgather_polars_dataframe, +) + +with spmd_execution() as (comm, ctx, engine): + result = ( + pl.scan_parquet("/data/dataset/*.parquet") + .filter(pl.col("amount") > 100) + .group_by("customer_id") + .agg(pl.col("amount").sum()) + .collect(engine=engine) + ) + + full = allgather_polars_dataframe( + comm=comm, + ctx=ctx, + local_df=result, + op_id=0, + ) +``` + +The context manager yields: + +* `comm` — `rapidsmpf.communicator.Communicator` +* `ctx` — `rapidsmpf.streaming.core.context.Context` +* `engine` — `pl.GPUEngine` configured for SPMD execution + +Pass `engine` to every `LazyFrame.collect()` inside the context block. + +## Collecting distributed results + +`collect()` returns a rank-local result. Use +`allgather_polars_dataframe()` to gather all fragments: + +```python +full = allgather_polars_dataframe( + comm=comm, + ctx=ctx, + local_df=result, + op_id=0, +) +``` + +`op_id` is a unique integer that identifies this collective operation across ranks. +All ranks must call the same collective with the same `op_id`. Otherwise the program +will deadlock. + +The result is a `pl.DataFrame` containing rows from all ranks, ordered by rank. + +## Query symmetry requirement + +All ranks must execute the **same sequence of queries in the same order**. Collective +operations are matched using internal operation IDs. If one rank executes a collective +that another rank does not, the program will deadlock. + +In practice: + +* Avoid rank-conditional `collect()` calls +* Avoid branches that change the query graph +* Keep the driver script deterministic + +## Passing options + +`executor_options` and `engine_kwargs` accept pass-through dictionaries: + +```python +with spmd_execution( + executor_options={ + "max_rows_per_partition": 500_000, + "rapidsmpf_spill": True, + }, + parquet_options={"use_rapidsmpf_native": True}, # forwarded via **engine_kwargs +) as (comm, ctx, engine): + ... +``` + +`executor_options` keys map to `StreamingExecutor` fields. Any additional keyword +arguments to `spmd_execution()` (such as `parquet_options`) are forwarded directly +to `pl.GPUEngine` as `**engine_kwargs`. + +The keys `"runtime"`, `"cluster"`, and `"spmd"` in `executor_options`, and +`"memory_resource"` and `"executor"` in `engine_kwargs`, are reserved. + +--- + +# Ray cluster mode + +Ray mode uses a single client process that drives execution across multiple GPU +workers. Internally, the system uses the concept of **ranks**, similar to MPI ranks. +Each rank corresponds to one GPU worker and participates in collective operations +through a shared UCXX communicator. + +In the Ray implementation, each rank is implemented as a **Ray actor**, with one +actor created per available GPU. Each rank owns its GPU, memory resource, +communicator endpoint, and RapidsMPF streaming context. + +The client sends the query plan to all ranks. The ranks execute the pipeline +collectively through UCXX, and their outputs are streamed back and concatenated on +the client. + +Unlike SPMD mode, the driver script runs as a normal Python program. There is no +`rrun` launcher and no symmetry requirement for the driver code. + +## Prerequisites + +* Ray (`ray`) installed +* RapidsMPF and UCXX available on all GPU nodes + +## Running in Ray mode + +`ray_execution()` is imported from `cudf_polars.experimental.rapidsmpf.ray`. It: + +1. Calls `ray.init()` if Ray is not already running +2. Creates one `RankActor` per GPU +3. Bootstraps a UCXX communicator across the actors +4. Yields a `pl.GPUEngine` and a `RayClient` + +Actors are shut down on exit. If the context started Ray, it also calls +`ray.shutdown()`. + +```python +import polars as pl +from cudf_polars.experimental.rapidsmpf.ray import ray_execution + +with ray_execution() as (ray_client, engine): + result = ( + pl.scan_parquet("/data/dataset/*.parquet") + .filter(pl.col("amount") > 100) + .group_by("customer_id") + .agg(pl.col("amount").sum()) + .collect(engine=engine) + ) + +print(result) +``` + +The context manager yields: + +* `ray_client` — cluster diagnostics and utilities +* `engine` — `pl.GPUEngine` configured for Ray execution + +## Ray lifecycle + +If Ray is already initialized, `ray_execution()` attaches to the existing cluster and +does not call `ray.shutdown()` on exit. + +```python +import ray +import polars as pl +from cudf_polars.experimental.rapidsmpf.ray import ray_execution + +ray.init(address="auto") + +try: + with ray_execution() as (ray_client, engine): + result = pl.scan_parquet(...).collect(engine=engine) +finally: + ray.shutdown() +``` + +`ray_execution()` raises `RuntimeError` if called inside an `rrun` cluster or if no +GPUs are available. + +## Cluster diagnostics + +`RayClient.gather_cluster_info()` returns placement information for all rank actors: + +```python +with ray_execution() as (ray_client, engine): + for i, info in enumerate(ray_client.gather_cluster_info()): + print( + f"rank {i}: hostname={info['hostname']}, pid={info['pid']}, " + f"CUDA_VISIBLE_DEVICES={info['cuda_visible_devices']}" + ) +``` + +Each entry includes `pid`, `hostname`, `cuda_visible_devices`, and `node_id`. + +## Passing options + +`executor_options`, `engine_kwargs`, and `ray_init_kwargs` accept pass-through +dictionaries: + +```python +with ray_execution( + executor_options={"max_rows_per_partition": 500_000}, + engine_kwargs={"raise_on_fail": True}, + ray_init_kwargs={"num_cpus": 4}, +) as (ray_client, engine): + ... +``` + +The keys `"runtime"`, `"cluster"`, `"spmd"`, and `"ray_client"` in +`executor_options`, and `"memory_resource"` and `"executor"` in `engine_kwargs`, +are reserved. From 27369853e9058c9d1e8f6577303a302e8bb4238c Mon Sep 17 00:00:00 2001 From: "Mads R. B. Kristensen" Date: Wed, 11 Mar 2026 08:41:39 +0100 Subject: [PATCH 04/33] cleanup --- python/cudf_polars/tests/experimental/rapidsmpf/test_ray.py | 3 --- 1 file changed, 3 deletions(-) diff --git a/python/cudf_polars/tests/experimental/rapidsmpf/test_ray.py b/python/cudf_polars/tests/experimental/rapidsmpf/test_ray.py index 53622aec1e56..1f22d8fc098d 100644 --- a/python/cudf_polars/tests/experimental/rapidsmpf/test_ray.py +++ b/python/cudf_polars/tests/experimental/rapidsmpf/test_ray.py @@ -100,10 +100,7 @@ def test_gather_cluster_info(ray_client: RayClient) -> None: assert "hostname" in info assert "pid" in info assert "cuda_visible_devices" in info - assert "comm_rank" in info assert isinstance(info["pid"], int) - # comm_rank is set after setup_worker; verify it is a valid rank index. - assert info["comm_rank"] in range(ray_client.nranks) # Each actor runs in its own process. assert len({info["pid"] for info in infos}) == ray_client.nranks From 99833868304bde922a96c0332917683d814d796f Mon Sep 17 00:00:00 2001 From: "Mads R. B. Kristensen" Date: Wed, 11 Mar 2026 09:11:05 +0100 Subject: [PATCH 05/33] docs --- .../cudf_polars/docs/cudf-polars-mp-design.md | 38 +++++++++---------- 1 file changed, 19 insertions(+), 19 deletions(-) diff --git a/python/cudf_polars/docs/cudf-polars-mp-design.md b/python/cudf_polars/docs/cudf-polars-mp-design.md index a912f5cff63a..eff694cbc08a 100644 --- a/python/cudf_polars/docs/cudf-polars-mp-design.md +++ b/python/cudf_polars/docs/cudf-polars-mp-design.md @@ -20,28 +20,28 @@ cluster) and differ only in how the user script interacts with them. ### SPMD mode The user script runs on **all N workers simultaneously**. There is no separate -client — every process is both driver and worker. +client — every process is both client and worker. ``` - rank 0 rank 1 ... rank N-1 + rank 0 rank 1 ... rank N-1 ┌─────────────────┐ ┌─────────────────┐ ┌─────────────────┐ │ User script │ │ User script │ │ User script │ │ (same code on │ │ (same code on │ │ (same code on │ │ every rank) │ │ every rank) │ │ every rank) │ └────────┬────────┘ └────────┬────────┘ └────────┬────────┘ - │ │ │ - │ LazyFrame.collect(engine=engine) │ - ↓ ↓ ↓ + │ │ │ + │ LazyFrame.collect(engine=engine) │ + ↓ ↓ ↓ ┌─────────────────┐ ┌─────────────────┐ ┌─────────────────┐ │ run IR │ │ run IR │ │ run IR │ └────────┬────────┘ └────────┬────────┘ └────────┬────────┘ - │ │ │ - ↓ ↓ ↓ + │ │ │ + ↓ ↓ ↓ ┌────────────────────────────────────────────────────────────────┐ │ RapidsMPF streaming engine │ │ shuffle / all-gather · UCXX communicator · RMM GPU memory │ └────────────────────────────────────────────────────────────────┘ - ↑ ↑ ↑ + ↑ ↑ ↑ GPU 0 GPU 1 GPU N-1 ``` @@ -50,13 +50,13 @@ assemble the full dataset on every rank. ### Ray mode -A single driver script dispatches work to N `RankActor` Ray actors (one per -GPU). The driver never touches a GPU directly. +A single client script dispatches work to N `RankActor` Ray actors (one per +GPU). The client never touches a GPU directly. ``` ┌──────────────────────────────┐ │ User script │ - │ (single driver process) │ + │ (single client process) │ │ LazyFrame.collect(engine=…) │ └──────────────┬───────────────┘ │ IR dispatched to all actors @@ -76,13 +76,13 @@ GPU). The driver never touches a GPU directly. GPU 0 GPU 1 GPU N-1 ``` -Per-rank output fragments are concatenated on the driver before being returned. +Per-rank output fragments are concatenated on the client before being returned. No `allgather` step is needed. ### Key insight — why client modes exist In SPMD mode every process runs the same code, which is unfamiliar to users -accustomed to single-process or Dask-style driver/worker workflows. Client +accustomed to single-process or Dask-style client/worker workflows. Client frontends such as Ray let users write a normal single-process script while the cluster handles distribution transparently. The underlying engine is unchanged; only the dispatch layer differs. Future frontends (Dask, custom clients) can @@ -96,7 +96,7 @@ target the same SPMD cluster without modifying the engine. The user script is launched with `rrun -n N python script.py`. `rrun` starts N identical processes, each pinned to one GPU. There is no separate client -process — every process runs the full script, acting simultaneously as driver +process — every process runs the full script, acting simultaneously as client and worker on its rank-local data. Because every rank runs independent Python, a `pl.DataFrame` is always @@ -200,10 +200,10 @@ Reserved `engine_kwargs` keys: `"memory_resource"`, `"executor"`. ### Execution model -The user runs a single driver script. `ray_execution()` creates N `RankActor` +The user runs a single client script. `ray_execution()` creates N `RankActor` Ray remote actors — one per available GPU. The actors form a private SPMD -cluster; the driver dispatches Polars IR to them and receives concatenated -results directly, with no `allgather` step needed on the client. +cluster; the client dispatches Polars IR to them and receives concatenated +results directly, with no `allgather` step needed. ### Bootstrapping @@ -222,7 +222,7 @@ results directly, with no `allgather` step needed on the client. Ray's resource scheduler assigns `num_gpus=1` to each `RankActor` before the actor process starts, setting `CUDA_VISIBLE_DEVICES` automatically. The `RAY_ACCEL_ENV_VAR_OVERRIDE_ON_ZERO=0` environment variable prevents Ray from -overriding `CUDA_VISIBLE_DEVICES` to empty on the driver process (which has +overriding `CUDA_VISIBLE_DEVICES` to empty on the client process (which has zero GPUs assigned). For hardware placement details, see [Section 5](#5-hardware-mapping-gpu-pinning). @@ -277,7 +277,7 @@ with ray_execution() as (ray_client, engine): .agg(pl.col("value").sum()) .collect(engine=engine) ) - # result is the full concatenated output, returned directly to the driver + # result is the full concatenated output, returned directly to the client print(result) ``` From 0eaf862bfc1095a20bf64e3656212815b3cc2dd1 Mon Sep 17 00:00:00 2001 From: "Mads R. B. Kristensen" Date: Wed, 11 Mar 2026 12:34:41 +0100 Subject: [PATCH 06/33] query_bundle --- .../cudf_polars/experimental/rapidsmpf/ray.py | 35 ++++++++++--------- 1 file changed, 18 insertions(+), 17 deletions(-) diff --git a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/ray.py b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/ray.py index ff4e756e7d65..863a6f813a51 100644 --- a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/ray.py +++ b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/ray.py @@ -112,14 +112,15 @@ def evaluate_pipeline_ray_mode( executor=dataclasses.replace(config_options.executor, ray_client=None), ) + # ir, partition_info, stats, and collective_id_map must be pickled together + # so that the IR-node keys in partition_info / collective_id_map are the + # same objects as the nodes in the ir tree after deserialization. + query_bundle = (ir, partition_info, stats, collective_id_map) result = ray.get( [ rank.evaluate_polars_ir.remote( - ir, - partition_info, + query_bundle, actor_config_options, - stats, - collective_id_map, collect_metadata=collect_metadata, ) for rank in rank_actors @@ -264,11 +265,13 @@ def get_info(self) -> dict: def evaluate_polars_ir( self, - ir: IR, - partition_info: MutableMapping[IR, PartitionInfo], + query_bundle: tuple[ + IR, + MutableMapping[IR, PartitionInfo], + StatsCollector, + dict[IR, list[int]], + ], config_options: ConfigOptions[StreamingExecutor], - stats: StatsCollector, - collective_id_map: dict[IR, list[int]], *, collect_metadata: bool, ) -> tuple[pl.DataFrame, list[ChannelMetadata] | None]: @@ -282,17 +285,14 @@ def evaluate_polars_ir( Parameters ---------- - ir - Root IR node describing the query to execute. - partition_info - Per-node partition metadata produced by the planner. + query_bundle + Tuple of ``(ir, partition_info, stats, collective_id_map)``. + Bundled into a single argument so that all four objects are + pickled together, preserving object identity between IR-node + keys in ``partition_info`` / ``collective_id_map`` and the + nodes in the ``ir`` tree. config_options Executor configuration forwarded from the client. - stats - Statistics collector used during execution. - collective_id_map - Mapping from IR nodes to their pre-allocated collective operation - IDs. collect_metadata If ``True``, collect channel metadata during execution. @@ -309,6 +309,7 @@ def evaluate_polars_ir( AssertionError If :meth:`setup_worker` has not been called first. """ + ir, partition_info, stats, collective_id_map = query_bundle assert self._ctx is not None, ( "setup_worker must be called before evaluate_polars_ir" ) From 12f940b1d058878a77170887aa17d0a1c56f772d Mon Sep 17 00:00:00 2001 From: "Mads R. B. Kristensen" Date: Thu, 12 Mar 2026 12:40:40 +0100 Subject: [PATCH 07/33] set_current_device_resource --- python/cudf_polars/cudf_polars/experimental/rapidsmpf/ray.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/ray.py b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/ray.py index 863a6f813a51..39527d088579 100644 --- a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/ray.py +++ b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/ray.py @@ -235,6 +235,9 @@ def setup_worker(self, root_ucxx_address_as_bytes: bytes) -> None: self._ctx = Context.from_options( self._comm.logger, self._mr, self._rapidsmpf_options ) + # Set the current RMM device resource so all temporary allocations + # in libcudf also use the same memory resource. + rmm.mr.set_current_device_resource(self._ctx.br().device_mr) def shutdown(self) -> None: """ From 9d067e891998a5db459ee4d8f4eb01e53e97b837 Mon Sep 17 00:00:00 2001 From: "Mads R. B. Kristensen" Date: Thu, 12 Mar 2026 12:53:12 +0100 Subject: [PATCH 08/33] num_streaming_threads <- max_io_threads --- python/cudf_polars/cudf_polars/experimental/rapidsmpf/ray.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/ray.py b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/ray.py index 39527d088579..39134767ce5c 100644 --- a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/ray.py +++ b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/ray.py @@ -478,6 +478,9 @@ def ray_execution( if rapidsmpf_options is not None else Options(get_environment_variables()) ) + rapidsmpf_options.insert_if_absent( + {"num_streaming_threads": str(executor_options.get("max_io_threads", 4))} + ) rapidsmpf_options_as_bytes = rapidsmpf_options.serialize() ray_was_initialized: bool = ray.is_initialized() From ccc59efaae230fb8525b669e17908cf3857ed062 Mon Sep 17 00:00:00 2001 From: "Mads R. B. Kristensen" Date: Thu, 12 Mar 2026 13:01:41 +0100 Subject: [PATCH 09/33] ray.put() query_bundle --- .../cudf_polars/experimental/rapidsmpf/ray.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/ray.py b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/ray.py index 39134767ce5c..73ee009a1a72 100644 --- a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/ray.py +++ b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/ray.py @@ -112,10 +112,11 @@ def evaluate_pipeline_ray_mode( executor=dataclasses.replace(config_options.executor, ray_client=None), ) - # ir, partition_info, stats, and collective_id_map must be pickled together - # so that the IR-node keys in partition_info / collective_id_map are the - # same objects as the nodes in the ir tree after deserialization. - query_bundle = (ir, partition_info, stats, collective_id_map) + # Serialize the IR bundle once into the Ray object store. The objects must be + # pickled together so IR-node keys in partition_info / collective_id_map remain + # identical to the nodes in the deserialized IR tree. Actors fetch the bundle + # by reference instead of receiving N copies. + query_bundle = ray.put((ir, partition_info, stats, collective_id_map)) result = ray.get( [ rank.evaluate_polars_ir.remote( From 1a8658bfb4e4f5e86ee23f1b174b8715b307628d Mon Sep 17 00:00:00 2001 From: "Mads R. B. Kristensen" Date: Thu, 12 Mar 2026 13:07:01 +0100 Subject: [PATCH 10/33] rapidsmpf_py_executor_max_workers default None --- python/cudf_polars/cudf_polars/experimental/rapidsmpf/ray.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/ray.py b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/ray.py index 73ee009a1a72..b89851715fe0 100644 --- a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/ray.py +++ b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/ray.py @@ -179,7 +179,8 @@ def __init__( self._nranks: int = nranks self._py_executor = ThreadPoolExecutor( max_workers=cast( - int, executor_options.get("rapidsmpf_py_executor_max_workers", 1) + int | None, + executor_options.get("rapidsmpf_py_executor_max_workers"), ), thread_name_prefix="ray-executor", ) From 47a6c7a57e6c7822fcebf24fcf72c6c5fb37fb55 Mon Sep 17 00:00:00 2001 From: "Mads R. B. Kristensen" Date: Thu, 12 Mar 2026 13:21:21 +0100 Subject: [PATCH 11/33] doc --- .../cudf_polars/experimental/rapidsmpf/ray.py | 11 +++++------ 1 file changed, 5 insertions(+), 6 deletions(-) diff --git a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/ray.py b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/ray.py index b89851715fe0..82a44e9e26c0 100644 --- a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/ray.py +++ b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/ray.py @@ -60,7 +60,7 @@ def evaluate_pipeline_ray_mode( """ Evaluate a RapidsMPF streaming pipeline in Ray mode. - The query is dispatched concurrently to every :class:`RankActor` in the + The query is dispatched in parallel to every :class:`RankActor` in the Ray cluster. Each actor evaluates the full pipeline on its local GPU and participates in collective operations through the shared UCXX communicator. The per-rank outputs are concatenated on the client before @@ -147,9 +147,9 @@ class RankActor: """ Ray actor that owns one GPU and participates in a RapidsMPF cluster. - Each actor manages its own memory resource, statistics collector, - communicator, and streaming context. Collectively, the actors form a - SPMD execution cluster used by the client-side Ray integration. + Each actor manages its own memory resource, communicator, streaming context, + etc. Collectively, the actors form an SPMD execution cluster used by the + client-side Ray integration. Parameters ---------- @@ -245,8 +245,7 @@ def shutdown(self) -> None: """ Release actor-owned resources and exit the process. - This shuts down the local Python executor, drops communicator and - memory-resource references, and then terminates the Ray actor process. + Raises `ray.exceptions.RayActorError` """ self._py_executor.shutdown(wait=True, cancel_futures=True) self._comm = None From 3f344edd63617df53279d5013971c831cc7d348e Mon Sep 17 00:00:00 2001 From: "Mads R. B. Kristensen" Date: Thu, 12 Mar 2026 15:57:42 +0100 Subject: [PATCH 12/33] RayContext --- .../cudf_polars/experimental/rapidsmpf/ray.py | 151 ++++++++++++------ .../cudf_polars/cudf_polars/utils/config.py | 27 +++- .../tests/experimental/rapidsmpf/test_ray.py | 27 ++-- 3 files changed, 144 insertions(+), 61 deletions(-) diff --git a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/ray.py b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/ray.py index 82a44e9e26c0..398defe43520 100644 --- a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/ray.py +++ b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/ray.py @@ -8,7 +8,6 @@ import os import socket from concurrent.futures import ThreadPoolExecutor -from contextlib import contextmanager from typing import TYPE_CHECKING, Any, cast import ray @@ -35,9 +34,10 @@ from cudf_polars.experimental.rapidsmpf.core import generate_network from cudf_polars.experimental.rapidsmpf.utils import empty_table_chunk from cudf_polars.experimental.utils import _concat +from cudf_polars.utils.config import RayContext if TYPE_CHECKING: - from collections.abc import Iterator, MutableMapping + from collections.abc import MutableMapping from rapidsmpf.communicator.communicator import Communicator from rapidsmpf.streaming.cudf.channel_metadata import ChannelMetadata @@ -96,20 +96,20 @@ def evaluate_pipeline_ray_mode( RuntimeError If the configured executor runtime is not ``"rapidsmpf"``. RuntimeError - If ``config_options.executor.ray_client`` is ``None``. + If ``config_options.executor.ray_context`` is ``None``. """ if config_options.executor.runtime != "rapidsmpf": raise RuntimeError("Runtime must be rapidsmpf") - if config_options.executor.ray_client is None: - raise RuntimeError("ray_client must be set when cluster='ray'") - rank_actors = config_options.executor.ray_client._rank_actors + if config_options.executor.ray_context is None: + raise RuntimeError("ray_context must be set when cluster='ray'") + rank_actors = config_options.executor.ray_context.rank_actors - # Strip ray_client before pickling config_options for remote calls: + # Strip ray_context before pickling config_options for remote calls: # actors don't need the full actor list, and sending actor handles to each # actor is wasteful. actor_config_options = dataclasses.replace( config_options, - executor=dataclasses.replace(config_options.executor, ray_client=None), + executor=dataclasses.replace(config_options.executor, ray_context=None), ) # Serialize the IR bundle once into the Ray object store. The objects must be @@ -370,10 +370,36 @@ def evaluate_polars_ir( class RayClient: - """Client-side handle for the distributed cudf-polars execution.""" + """ + User handle for a RapidsMPF Ray cluster. + + Typically created via :func:`ray_execution`. See that function for usage + and detailed documentation. - def __init__(self, rank_actors: list[Any]) -> None: - self._rank_actors: list[Any] = rank_actors + Parameters + ---------- + engine + Polars GPU engine configured to execute queries on the Ray-backed + RapidsMPF cluster. Must have been created with ``"ray_context"`` in + its ``executor_options``. + ray_was_initialized + Indicates whether Ray was already initialized before this client + was created. Used to determine whether the client should shut down + Ray on :meth:`shutdown`. + """ + + def __init__( + self, + engine: pl.GPUEngine, + *, + ray_was_initialized: bool, + ) -> None: + self._engine: pl.GPUEngine = engine + self._ray_was_initialized: bool = ray_was_initialized + # Own copy of actor handles for shutdown; drained to [] on first call. + self._rank_actors: list[Any] = list( + engine.config["executor_options"]["ray_context"].rank_actors + ) @property def nranks(self) -> int: @@ -386,6 +412,11 @@ def nranks(self) -> int: """ return len(self._rank_actors) + @property + def engine(self) -> pl.GPUEngine: + """The Polars GPU engine bound to this cluster.""" + return self._engine + def gather_cluster_info(self) -> list[dict]: """ Collect diagnostic information from every rank actor. @@ -403,20 +434,50 @@ def gather_cluster_info(self) -> list[dict]: """ return ray.get([rank.get_info.remote() for rank in self._rank_actors]) + def shutdown(self) -> None: + """ + Shut down all rank actors. + + If Ray was initialized by this client, also calls :func:`ray.shutdown`. + Safe to call more than once. + """ + actors, self._rank_actors = self._rank_actors, [] + for a in actors: + try: + ray.get(a.shutdown.remote()) + except ray.exceptions.RayActorError: + pass # expected: exit_actor() terminates the process immediately + except Exception as e: + print(f"shutdown error: {e}") + if not self._ray_was_initialized: + self._ray_was_initialized = True + ray.shutdown() + + def __enter__(self) -> tuple[RayClient, pl.GPUEngine]: + """Enter the context manager, returning ``(self, engine)``.""" + return self, self.engine + + def __exit__(self, *_: object) -> None: + """Exit the context manager, calling :meth:`shutdown`.""" + self.shutdown() + -@contextmanager def ray_execution( *, rapidsmpf_options: Options | None = None, executor_options: dict[str, object] | None = None, engine_kwargs: dict[str, Any] | None = None, ray_init_kwargs: dict[str, object] | None = None, -) -> Iterator[tuple[RayClient, pl.GPUEngine]]: +) -> RayClient: """ - Create a RapidsMPF Ray cluster and matching Polars GPU engine. + Create a RapidsMPF Ray cluster and return a :class:`RayClient`. + + The returned client supports both direct use and the context-manager + protocol. Prefer the context-manager form in scripts; use the direct + form in interactive environments such as Jupyter notebooks. - If Ray is not already initialized, this context manager calls - :func:`ray.init` on entry and :func:`ray.shutdown` on exit. If Ray is + If Ray is not already initialized, :func:`ray.init` is called here and + :func:`ray.shutdown` is called by :meth:`RayClient.shutdown`. If Ray is already initialized, cluster lifetime remains managed by the caller. Parameters @@ -432,12 +493,12 @@ def ray_execution( Keyword arguments forwarded to :func:`ray.init` when Ray is not already initialized. - Yields - ------ - ray_client - Client-side handle to the Ray actor cluster. - engine - Polars GPU engine configured to execute through RapidsMPF on Ray. + Returns + ------- + RayClient + A client connected to the newly created Ray actor cluster. + Call :meth:`RayClient.shutdown` (or use it as a context manager) + to release resources when done. Raises ------ @@ -454,8 +515,16 @@ def ray_execution( Examples -------- + Context-manager style: + >>> with ray_execution() as (ray_client, engine): # doctest: +SKIP ... result = pl.LazyFrame({"a": [1, 2, 3]}).collect(engine=engine) + + Jupyter / manual style: + + >>> client, engine = ray_execution() # doctest: +SKIP + >>> result = pl.LazyFrame({"a": [1, 2, 3]}).collect(engine=engine) + >>> client.shutdown() """ executor_options = executor_options or {} engine_kwargs = engine_kwargs or {} @@ -469,7 +538,7 @@ def ray_execution( ) # Check for reserved keys. - if bad := {"runtime", "cluster", "spmd", "ray_client"} & executor_options.keys(): + if bad := {"runtime", "cluster", "spmd", "ray_context"} & executor_options.keys(): raise ValueError(f"executor_options may not contain reserved keys: {bad}") if bad := {"memory_resource", "executor"} & engine_kwargs.keys(): raise ValueError(f"engine_kwargs may not contain reserved keys: {bad}") @@ -520,27 +589,15 @@ def ray_execution( [rank.setup_worker.remote(root_ucxx_address_as_bytes) for rank in rank_actors] ) - try: - ray_client = RayClient(rank_actors) - engine = pl.GPUEngine( - memory_resource=None, - executor="streaming", - executor_options={ - **executor_options, - "runtime": "rapidsmpf", - "cluster": "ray", - "ray_client": ray_client, - }, - **engine_kwargs, - ) - yield ray_client, engine - finally: - for a in rank_actors: - try: - ray.get(a.shutdown.remote()) - except ray.exceptions.RayActorError: - pass # expected: exit_actor() terminates the process immediately. - except Exception as e: - print(f"shutdown error: {e}") - if not ray_was_initialized: - ray.shutdown() + engine = pl.GPUEngine( + memory_resource=None, + executor="streaming", + executor_options={ + **executor_options, + "runtime": "rapidsmpf", + "cluster": "ray", + "ray_context": RayContext(rank_actors), + }, + **engine_kwargs, + ) + return RayClient(engine, ray_was_initialized=ray_was_initialized) diff --git a/python/cudf_polars/cudf_polars/utils/config.py b/python/cudf_polars/cudf_polars/utils/config.py index 55d80fec5ab8..a1f5b1188223 100644 --- a/python/cudf_polars/cudf_polars/utils/config.py +++ b/python/cudf_polars/cudf_polars/utils/config.py @@ -43,8 +43,6 @@ import rmm.mr - from cudf_polars.experimental.rapidsmpf.ray import RayClient - __all__ = [ "Cluster", @@ -52,6 +50,7 @@ "DynamicPlanningOptions", "InMemoryExecutor", "ParquetOptions", + "RayContext", "Runtime", "SPMDContext", "Scheduler", # Deprecated, kept for backward compatibility @@ -639,6 +638,28 @@ class SPMDContext: py_executor: ThreadPoolExecutor +@dataclasses.dataclass(frozen=True) +class RayContext: + """ + Configuration for Ray cluster execution. + + .. note:: + This dataclass holds Ray actor handles, which are only valid within the + Ray session that created them. It is stripped from ``config_options`` + before pickling for remote actor calls in + :func:`~cudf_polars.experimental.rapidsmpf.ray.evaluate_pipeline_ray_mode`. + Do not persist or transfer this object across Ray sessions. + + Parameters + ---------- + rank_actors + List of :class:`~cudf_polars.experimental.rapidsmpf.ray.RankActor` + handles, one per GPU in the cluster. + """ + + rank_actors: list + + @dataclasses.dataclass(frozen=True, eq=True) class StreamingExecutor: """ @@ -876,7 +897,7 @@ class StreamingExecutor: ) ) spmd: SPMDContext | None = None - ray_client: RayClient | None = None + ray_context: RayContext | None = None def __post_init__(self) -> None: # noqa: D105 # Check for rapidsmpf runtime diff --git a/python/cudf_polars/tests/experimental/rapidsmpf/test_ray.py b/python/cudf_polars/tests/experimental/rapidsmpf/test_ray.py index 1f22d8fc098d..7528ac0b9773 100644 --- a/python/cudf_polars/tests/experimental/rapidsmpf/test_ray.py +++ b/python/cudf_polars/tests/experimental/rapidsmpf/test_ray.py @@ -5,6 +5,7 @@ from __future__ import annotations from typing import TYPE_CHECKING, Any +from unittest.mock import MagicMock import pytest @@ -14,6 +15,7 @@ from cudf_polars.experimental.rapidsmpf.ray import ( # noqa: E402 RayClient, + RayContext, ray_execution, ) @@ -57,23 +59,26 @@ def engine(_ray_env: tuple[RayClient, pl.GPUEngine]) -> pl.GPUEngine: def test_ray_execution_reserved_executor_keys() -> None: """executor_options rejects reserved keys.""" - for key in ("runtime", "cluster", "spmd", "ray_client"): - with ( - pytest.raises(ValueError, match="reserved"), - ray_execution(executor_options={key: "anything"}), - ): - pass + for key in ("runtime", "cluster", "spmd", "ray_context"): + with pytest.raises(ValueError, match="reserved"): + ray_execution(executor_options={key: "anything"}) def test_ray_execution_reserved_engine_kwargs_keys() -> None: """engine_kwargs rejects keys that are set explicitly by ray_execution.""" for key in ("memory_resource", "executor"): kwargs: dict[str, Any] = {key: "anything"} - with ( - pytest.raises(ValueError, match="reserved"), - ray_execution(engine_kwargs=kwargs), - ): - pass + with pytest.raises(ValueError, match="reserved"): + ray_execution(engine_kwargs=kwargs) + + +def test_ray_client_shutdown_idempotent() -> None: + """RayClient.shutdown() is safe to call more than once.""" + mock_engine = MagicMock(spec=pl.GPUEngine) + mock_engine.config = {"executor_options": {"ray_context": RayContext([])}} + client = RayClient(mock_engine, ray_was_initialized=True) + client.shutdown() + client.shutdown() # must not raise # --------------------------------------------------------------------------- From 1d1129de6fa2639c1a0cea1aa4cd2b6283c4aad2 Mon Sep 17 00:00:00 2001 From: "Mads R. B. Kristensen" Date: Thu, 12 Mar 2026 16:11:02 +0100 Subject: [PATCH 13/33] cleanup --- .../cudf_polars/experimental/rapidsmpf/ray.py | 46 +++++----- .../cudf_polars/cudf_polars/utils/config.py | 5 +- python/cudf_polars/docs/cudf_polars_mp.md | 2 +- .../tests/experimental/rapidsmpf/test_ray.py | 88 ++++++++++++++++--- 4 files changed, 102 insertions(+), 39 deletions(-) diff --git a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/ray.py b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/ray.py index 398defe43520..4c57e9aa5e79 100644 --- a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/ray.py +++ b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/ray.py @@ -41,6 +41,7 @@ from rapidsmpf.communicator.communicator import Communicator from rapidsmpf.streaming.cudf.channel_metadata import ChannelMetadata + from ray.actor import ActorHandle from cudf_polars.dsl.ir import IR from cudf_polars.experimental.base import PartitionInfo, StatsCollector @@ -245,7 +246,7 @@ def shutdown(self) -> None: """ Release actor-owned resources and exit the process. - Raises `ray.exceptions.RayActorError` + Raises `ray.exceptions.RayActorError`. """ self._py_executor.shutdown(wait=True, cancel_futures=True) self._comm = None @@ -310,13 +311,12 @@ def evaluate_polars_ir( Raises ------ - AssertionError + RuntimeError If :meth:`setup_worker` has not been called first. """ ir, partition_info, stats, collective_id_map = query_bundle - assert self._ctx is not None, ( - "setup_worker must be called before evaluate_polars_ir" - ) + if self._ctx is None: + raise RuntimeError("setup_worker must be called before evaluate_polars_ir") ir_context = IRExecutionContext(get_cuda_stream=self._ctx.get_stream_from_pool) metadata_collector: list[ChannelMetadata] | None = ( [] if collect_metadata else None @@ -343,7 +343,6 @@ def evaluate_polars_ir( ) for msg in messages ] - dfs: list[DataFrame] if chunks: dfs = [ DataFrame.from_table( @@ -382,22 +381,21 @@ class RayClient: Polars GPU engine configured to execute queries on the Ray-backed RapidsMPF cluster. Must have been created with ``"ray_context"`` in its ``executor_options``. - ray_was_initialized - Indicates whether Ray was already initialized before this client - was created. Used to determine whether the client should shut down - Ray on :meth:`shutdown`. + owns_ray + If ``True``, this client initialized Ray and is responsible for + calling :func:`ray.shutdown` in :meth:`shutdown`. """ def __init__( self, engine: pl.GPUEngine, *, - ray_was_initialized: bool, + owns_ray: bool, ) -> None: self._engine: pl.GPUEngine = engine - self._ray_was_initialized: bool = ray_was_initialized + self._owns_ray: bool = owns_ray # Own copy of actor handles for shutdown; drained to [] on first call. - self._rank_actors: list[Any] = list( + self._rank_actors: list[ActorHandle[RankActor]] = list( engine.config["executor_options"]["ray_context"].rank_actors ) @@ -449,8 +447,8 @@ def shutdown(self) -> None: pass # expected: exit_actor() terminates the process immediately except Exception as e: print(f"shutdown error: {e}") - if not self._ray_was_initialized: - self._ray_was_initialized = True + if self._owns_ray: + self._owns_ray = False ray.shutdown() def __enter__(self) -> tuple[RayClient, pl.GPUEngine]: @@ -487,6 +485,8 @@ def ray_execution( ``Options(get_environment_variables())``. executor_options Additional key-value pairs forwarded to the Polars executor options. + If ``"max_io_threads"`` is present, its value is also used as the + default for ``rapidsmpf_options["num_streaming_threads"]``. engine_kwargs Additional keyword arguments forwarded to :class:`polars.GPUEngine`. ray_init_kwargs @@ -495,10 +495,9 @@ def ray_execution( Returns ------- - RayClient - A client connected to the newly created Ray actor cluster. - Call :meth:`RayClient.shutdown` (or use it as a context manager) - to release resources when done. + A client connected to the newly created Ray actor cluster. + Call :meth:`RayClient.shutdown` (or use it as a context manager) + to release resources when done. Raises ------ @@ -557,12 +556,13 @@ def ray_execution( if not ray_was_initialized: # Prevent Ray from overriding CUDA_VISIBLE_DEVICES to "" when a worker # process starts with zero visible GPUs (e.g. the driver process itself). - # Without this, Ray's accelerator detection resets the variable before our - # actors acquire their GPU assignment, hiding all GPUs from CUDA. os.environ.setdefault("RAY_ACCEL_ENV_VAR_OVERRIDE_ON_ZERO", "0") ray.init(**ray_init_kwargs) total_gpus = int(ray.cluster_resources().get("GPU", 0.0)) + # Note: available_resources() is a snapshot and inherently racy. This is a + # best-effort guard; another process could claim GPUs between this check and + # actor creation. free_gpus = int(ray.available_resources().get("GPU", 0.0)) if total_gpus != free_gpus: raise RuntimeError( @@ -572,7 +572,7 @@ def ray_execution( raise RuntimeError("No available GPUs in the Ray cluster at startup") # Create one actor per GPU. Ray adds .remote() dynamically; no type stubs. - rank_actors: list[Any] = [ + rank_actors: list[ActorHandle[RankActor]] = [ RankActor.remote( # type: ignore[attr-defined] nranks=free_gpus, executor_options=executor_options, @@ -600,4 +600,4 @@ def ray_execution( }, **engine_kwargs, ) - return RayClient(engine, ray_was_initialized=ray_was_initialized) + return RayClient(engine, owns_ray=not ray_was_initialized) diff --git a/python/cudf_polars/cudf_polars/utils/config.py b/python/cudf_polars/cudf_polars/utils/config.py index a1f5b1188223..337aeb7f51d0 100644 --- a/python/cudf_polars/cudf_polars/utils/config.py +++ b/python/cudf_polars/cudf_polars/utils/config.py @@ -38,11 +38,14 @@ from rapidsmpf.communicator.communicator import Communicator from rapidsmpf.streaming.core.context import Context + from ray.actor import ActorHandle import polars.lazyframe.engine_config import rmm.mr + from cudf_polars.experimental.rapidsmpf.ray import RankActor + __all__ = [ "Cluster", @@ -657,7 +660,7 @@ class RayContext: handles, one per GPU in the cluster. """ - rank_actors: list + rank_actors: list[ActorHandle[RankActor]] @dataclasses.dataclass(frozen=True, eq=True) diff --git a/python/cudf_polars/docs/cudf_polars_mp.md b/python/cudf_polars/docs/cudf_polars_mp.md index ba37cff0b318..18f8981cec36 100644 --- a/python/cudf_polars/docs/cudf_polars_mp.md +++ b/python/cudf_polars/docs/cudf_polars_mp.md @@ -1,6 +1,6 @@ # cudf-polars-mp -`cudf-polars-mp` extends Polars query execution to multiple GPUs. +cudf-polars-mp extends Polars query execution to multiple GPUs. Multi-process (mp) execution distributes a query across several GPU workers. Each worker owns a disjoint fragment of the data and participates in collective operations diff --git a/python/cudf_polars/tests/experimental/rapidsmpf/test_ray.py b/python/cudf_polars/tests/experimental/rapidsmpf/test_ray.py index 7528ac0b9773..67f83050aa8b 100644 --- a/python/cudf_polars/tests/experimental/rapidsmpf/test_ray.py +++ b/python/cudf_polars/tests/experimental/rapidsmpf/test_ray.py @@ -5,17 +5,18 @@ from __future__ import annotations from typing import TYPE_CHECKING, Any -from unittest.mock import MagicMock +from unittest.mock import MagicMock, patch import pytest import polars as pl +from cudf_polars.utils.config import RayContext + ray = pytest.importorskip("ray") from cudf_polars.experimental.rapidsmpf.ray import ( # noqa: E402 RayClient, - RayContext, ray_execution, ) @@ -26,11 +27,19 @@ @pytest.fixture(scope="session") def _ray_env() -> Iterator[tuple[RayClient, pl.GPUEngine]]: """Create one Ray cluster + GPU actors shared across the test session.""" - with ray_execution(ray_init_kwargs={"include_dashboard": False}) as ( - ray_client, - engine, - ): - yield ray_client, engine + try: + with ray_execution( + # Use a small partition size so tests exercise the multi-partition + # code path deterministically, regardless of input size. + executor_options={"max_rows_per_partition": 10}, + ray_init_kwargs={"include_dashboard": False}, + ) as ( + ray_client, + engine, + ): + yield ray_client, engine + except RuntimeError as e: + pytest.skip(f"Ray GPU cluster unavailable: {e}") @pytest.fixture(scope="session") @@ -76,11 +85,23 @@ def test_ray_client_shutdown_idempotent() -> None: """RayClient.shutdown() is safe to call more than once.""" mock_engine = MagicMock(spec=pl.GPUEngine) mock_engine.config = {"executor_options": {"ray_context": RayContext([])}} - client = RayClient(mock_engine, ray_was_initialized=True) + client = RayClient(mock_engine, owns_ray=False) client.shutdown() client.shutdown() # must not raise +def test_ray_execution_raises_inside_rrun() -> None: + """ray_execution() must not be called from within an rrun cluster.""" + with ( + patch( + "cudf_polars.experimental.rapidsmpf.ray.bootstrap.is_running_with_rrun", + return_value=True, + ), + pytest.raises(RuntimeError, match="rrun"), + ): + ray_execution() + + # --------------------------------------------------------------------------- # GPU tests — share a single Ray cluster + actor set for the whole session # --------------------------------------------------------------------------- @@ -96,6 +117,18 @@ def test_ray_execution_yields_client_and_engine( assert ray_client.nranks >= 1 +def test_ray_execution_executor_options_forwarded( + ray_client: RayClient, + engine: pl.GPUEngine, +) -> None: + """Reserved executor_options keys are injected into the engine config.""" + opts = engine.config["executor_options"] + assert opts["runtime"] == "rapidsmpf" + assert opts["cluster"] == "ray" + assert isinstance(opts["ray_context"], RayContext) + assert len(opts["ray_context"].rank_actors) == ray_client.nranks + + def test_gather_cluster_info(ray_client: RayClient) -> None: """gather_cluster_info returns one info dict per rank with expected fields.""" infos = ray_client.gather_cluster_info() @@ -126,15 +159,42 @@ def test_ray_execution_filter(engine: pl.GPUEngine) -> None: assert sorted(result["a"].to_list()) == [4, 5] -def test_ray_execution_group_by(engine: pl.GPUEngine) -> None: - """Group-by produces the correct global aggregation across all actors.""" - lf = pl.LazyFrame({"key": ["a", "a", "b"], "val": [1, 2, 3]}) +def test_ray_execution_group_by(ray_client: RayClient, engine: pl.GPUEngine) -> None: + """Group-by produces the correct aggregation across all ranks.""" + # max_rows_per_partition=10 (set on the session fixture) gives each rank + # exactly 5 partitions, so the multi-partition path is always exercised. + n, n_keys = ray_client.nranks * 50, 5 + keys = [str(i % n_keys) for i in range(n)] + vals = list(range(n)) + lf = pl.LazyFrame({"key": keys, "val": vals}) result = ( lf.group_by("key").agg(pl.col("val").sum()).collect(engine=engine).sort("key") ) - assert result.shape == (2, 2) - assert result["key"].to_list() == ["a", "b"] - assert result["val"].to_list() == [3, 3] + expected = ( + pl.LazyFrame({"key": keys, "val": vals}) + .group_by("key") + .agg(pl.col("val").sum()) + .collect() + .sort("key") + ) + assert result.shape == expected.shape + assert result["key"].to_list() == expected["key"].to_list() + assert result["val"].to_list() == expected["val"].to_list() + + +def test_ray_execution_join(ray_client: RayClient, engine: pl.GPUEngine) -> None: + """Hash join between two tables produces the correct result across all ranks.""" + # max_rows_per_partition=10 (set on the session fixture) gives each rank + # exactly 5 partitions, so the multi-partition path is always exercised. + n = ray_client.nranks * 50 + lf_left = pl.LazyFrame({"key": list(range(n)), "val_left": list(range(n))}) + lf_right = pl.LazyFrame( + {"key": list(range(n)), "val_right": [x * 2 for x in range(n)]} + ) + result = lf_left.join(lf_right, on="key").collect(engine=engine).sort("key") + assert result.shape == (n, 3) + assert result["val_left"].to_list() == list(range(n)) + assert result["val_right"].to_list() == [x * 2 for x in range(n)] def test_ray_execution_empty_dataframe(engine: pl.GPUEngine) -> None: From 671b6e269a83390999c74231ae88a742fe282b26 Mon Sep 17 00:00:00 2001 From: "Mads R. B. Kristensen" Date: Thu, 12 Mar 2026 16:59:51 +0100 Subject: [PATCH 14/33] Multi-GPU Polars --- python/cudf_polars/docs/cudf-polars-mp-design.md | 2 +- .../cudf_polars/docs/{cudf_polars_mp.md => cudf-polars-mp.md} | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) rename python/cudf_polars/docs/{cudf_polars_mp.md => cudf-polars-mp.md} (98%) diff --git a/python/cudf_polars/docs/cudf-polars-mp-design.md b/python/cudf_polars/docs/cudf-polars-mp-design.md index eff694cbc08a..84d702c5b28f 100644 --- a/python/cudf_polars/docs/cudf-polars-mp-design.md +++ b/python/cudf_polars/docs/cudf-polars-mp-design.md @@ -2,7 +2,7 @@ This document describes the multi-GPU execution architecture of cudf-polars. For user-facing setup instructions, see -`cudf_polars_mp.md`. +`cudf-polars-mp.md`. - [1. Architecture Overview](#1-architecture-overview) - [2. SPMD Mode](#2-spmd-mode) diff --git a/python/cudf_polars/docs/cudf_polars_mp.md b/python/cudf_polars/docs/cudf-polars-mp.md similarity index 98% rename from python/cudf_polars/docs/cudf_polars_mp.md rename to python/cudf_polars/docs/cudf-polars-mp.md index 18f8981cec36..508784d25aa3 100644 --- a/python/cudf_polars/docs/cudf_polars_mp.md +++ b/python/cudf_polars/docs/cudf-polars-mp.md @@ -1,6 +1,6 @@ -# cudf-polars-mp +# Multi-GPU Polars -cudf-polars-mp extends Polars query execution to multiple GPUs. +Multi-GPU Polars extends Polars query execution to multiple GPUs. Multi-process (mp) execution distributes a query across several GPU workers. Each worker owns a disjoint fragment of the data and participates in collective operations From 844544e00ec4db80b9727fd9b355c3bd9fd4e4eb Mon Sep 17 00:00:00 2001 From: "Mads R. B. Kristensen" Date: Fri, 13 Mar 2026 08:45:45 +0100 Subject: [PATCH 15/33] doc --- .../cudf_polars/cudf_polars/experimental/rapidsmpf/ray.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/ray.py b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/ray.py index 4c57e9aa5e79..0e0b71d8da1c 100644 --- a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/ray.py +++ b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/ray.py @@ -192,9 +192,9 @@ def setup_root(self) -> bytes: """ Initialize this actor as the root rank. - The root actor creates a new UCXX communicator without an existing - root address. The resulting root address is returned so it can be - distributed to all other actors during bootstrap. + The root actor creates a new UCXX communicator and returns the + serialized root address, which must be passed to :meth:`setup_worker` + on all actors to complete communicator setup. Returns ------- From 296eb55114672fb1276a1c55f4e5a91ca56a8bb1 Mon Sep 17 00:00:00 2001 From: "Mads R. B. Kristensen" Date: Fri, 13 Mar 2026 08:53:57 +0100 Subject: [PATCH 16/33] RayClient: cleanup --- .../cudf_polars/experimental/rapidsmpf/ray.py | 23 +++++++++++-------- .../tests/experimental/rapidsmpf/test_ray.py | 15 +++++++++++- 2 files changed, 28 insertions(+), 10 deletions(-) diff --git a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/ray.py b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/ray.py index 0e0b71d8da1c..123b1cd63669 100644 --- a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/ray.py +++ b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/ray.py @@ -392,12 +392,15 @@ def __init__( *, owns_ray: bool, ) -> None: - self._engine: pl.GPUEngine = engine + self._engine: pl.GPUEngine | None = engine self._owns_ray: bool = owns_ray - # Own copy of actor handles for shutdown; drained to [] on first call. - self._rank_actors: list[ActorHandle[RankActor]] = list( - engine.config["executor_options"]["ray_context"].rank_actors - ) + + @property + def rank_actors(self) -> list[ActorHandle[RankActor]]: + """List of Ray rank actor handles, or ``[]`` after :meth:`shutdown`.""" + if self._engine is None: + return [] + return self._engine.config["executor_options"]["ray_context"].rank_actors @property def nranks(self) -> int: @@ -408,11 +411,13 @@ def nranks(self) -> int: ------- Number of ranks/nodes in the Ray cluster. """ - return len(self._rank_actors) + return len(self.rank_actors) @property def engine(self) -> pl.GPUEngine: """The Polars GPU engine bound to this cluster.""" + if self._engine is None: + raise RuntimeError("engine is not available after shutdown") return self._engine def gather_cluster_info(self) -> list[dict]: @@ -430,7 +435,7 @@ def gather_cluster_info(self) -> list[dict]: ... for i, info in enumerate(ray_client.gather_cluster_info()): ... print(f"rank {i}: {info}") """ - return ray.get([rank.get_info.remote() for rank in self._rank_actors]) + return ray.get([rank.get_info.remote() for rank in self.rank_actors]) def shutdown(self) -> None: """ @@ -439,14 +444,14 @@ def shutdown(self) -> None: If Ray was initialized by this client, also calls :func:`ray.shutdown`. Safe to call more than once. """ - actors, self._rank_actors = self._rank_actors, [] - for a in actors: + for a in self.rank_actors: try: ray.get(a.shutdown.remote()) except ray.exceptions.RayActorError: pass # expected: exit_actor() terminates the process immediately except Exception as e: print(f"shutdown error: {e}") + self._engine = None if self._owns_ray: self._owns_ray = False ray.shutdown() diff --git a/python/cudf_polars/tests/experimental/rapidsmpf/test_ray.py b/python/cudf_polars/tests/experimental/rapidsmpf/test_ray.py index 67f83050aa8b..f65b69b7686d 100644 --- a/python/cudf_polars/tests/experimental/rapidsmpf/test_ray.py +++ b/python/cudf_polars/tests/experimental/rapidsmpf/test_ray.py @@ -90,6 +90,18 @@ def test_ray_client_shutdown_idempotent() -> None: client.shutdown() # must not raise +def test_ray_client_post_shutdown_state() -> None: + """After shutdown, rank_actors is empty, nranks is 0, and engine raises.""" + mock_engine = MagicMock(spec=pl.GPUEngine) + mock_engine.config = {"executor_options": {"ray_context": RayContext([])}} + client = RayClient(mock_engine, owns_ray=False) + client.shutdown() + assert client.rank_actors == [] + assert client.nranks == 0 + with pytest.raises(RuntimeError, match="shutdown"): + _ = client.engine + + def test_ray_execution_raises_inside_rrun() -> None: """ray_execution() must not be called from within an rrun cluster.""" with ( @@ -126,7 +138,8 @@ def test_ray_execution_executor_options_forwarded( assert opts["runtime"] == "rapidsmpf" assert opts["cluster"] == "ray" assert isinstance(opts["ray_context"], RayContext) - assert len(opts["ray_context"].rank_actors) == ray_client.nranks + assert ray_client.rank_actors == opts["ray_context"].rank_actors + assert len(ray_client.rank_actors) == ray_client.nranks def test_gather_cluster_info(ray_client: RayClient) -> None: From 3bae4a761230b24a8a0e74e99f5a6da228ad8822 Mon Sep 17 00:00:00 2001 From: "Mads R. B. Kristensen" Date: Fri, 13 Mar 2026 09:40:54 +0100 Subject: [PATCH 17/33] removed the design docs, now part of cudf-polars-mp.md --- .../cudf_polars/docs/cudf-polars-mp-design.md | 427 ------------------ python/cudf_polars/docs/cudf-polars-mp.md | 161 +++++-- 2 files changed, 112 insertions(+), 476 deletions(-) delete mode 100644 python/cudf_polars/docs/cudf-polars-mp-design.md diff --git a/python/cudf_polars/docs/cudf-polars-mp-design.md b/python/cudf_polars/docs/cudf-polars-mp-design.md deleted file mode 100644 index 84d702c5b28f..000000000000 --- a/python/cudf_polars/docs/cudf-polars-mp-design.md +++ /dev/null @@ -1,427 +0,0 @@ -# cudf-polars Multi-GPU Design - -This document describes the multi-GPU execution architecture of cudf-polars. -For user-facing setup instructions, see -`cudf-polars-mp.md`. - -- [1. Architecture Overview](#1-architecture-overview) -- [2. SPMD Mode](#2-spmd-mode) -- [3. Ray Mode](#3-ray-mode) -- [4. Comparison](#4-comparison) -- [5. Hardware Mapping (GPU Pinning)](#5-hardware-mapping-gpu-pinning) - ---- - -## 1. Architecture Overview - -Both modes share the same bottom two layers (RapidsMPF engine and SPMD -cluster) and differ only in how the user script interacts with them. - -### SPMD mode - -The user script runs on **all N workers simultaneously**. There is no separate -client — every process is both client and worker. - -``` - rank 0 rank 1 ... rank N-1 -┌─────────────────┐ ┌─────────────────┐ ┌─────────────────┐ -│ User script │ │ User script │ │ User script │ -│ (same code on │ │ (same code on │ │ (same code on │ -│ every rank) │ │ every rank) │ │ every rank) │ -└────────┬────────┘ └────────┬────────┘ └────────┬────────┘ - │ │ │ - │ LazyFrame.collect(engine=engine) │ - ↓ ↓ ↓ -┌─────────────────┐ ┌─────────────────┐ ┌─────────────────┐ -│ run IR │ │ run IR │ │ run IR │ -└────────┬────────┘ └────────┬────────┘ └────────┬────────┘ - │ │ │ - ↓ ↓ ↓ -┌────────────────────────────────────────────────────────────────┐ -│ RapidsMPF streaming engine │ -│ shuffle / all-gather · UCXX communicator · RMM GPU memory │ -└────────────────────────────────────────────────────────────────┘ - ↑ ↑ ↑ - GPU 0 GPU 1 GPU N-1 -``` - -Results are rank-local after `collect`. Call `allgather_polars_dataframe()` to -assemble the full dataset on every rank. - -### Ray mode - -A single client script dispatches work to N `RankActor` Ray actors (one per -GPU). The client never touches a GPU directly. - -``` - ┌──────────────────────────────┐ - │ User script │ - │ (single client process) │ - │ LazyFrame.collect(engine=…) │ - └──────────────┬───────────────┘ - │ IR dispatched to all actors - ┌────────────────|─────────────────┐ - ↓ ↓ ↓ - ┌─────────────┐ ┌─────────────┐ ┌─────────────┐ - │ RankActor │ │ RankActor │ │ RankActor │ - │ rank 0 │ │ rank 1 │ │ rank N-1 │ - │ run IR │ │ run IR │ │ run IR │ - └──────┬──────┘ └──────┬──────┘ └──────┬──────┘ - ↓ ↓ ↓ -┌────────────────────────────────────────────────────────────────┐ -│ RapidsMPF streaming engine │ -│ shuffle / all-gather · UCXX communicator · RMM GPU memory │ -└────────────────────────────────────────────────────────────────┘ - ↑ ↑ ↑ - GPU 0 GPU 1 GPU N-1 -``` - -Per-rank output fragments are concatenated on the client before being returned. -No `allgather` step is needed. - -### Key insight — why client modes exist - -In SPMD mode every process runs the same code, which is unfamiliar to users -accustomed to single-process or Dask-style client/worker workflows. Client -frontends such as Ray let users write a normal single-process script while the -cluster handles distribution transparently. The underlying engine is unchanged; -only the dispatch layer differs. Future frontends (Dask, custom clients) can -target the same SPMD cluster without modifying the engine. - ---- - -## 2. SPMD Mode - -### Execution model - -The user script is launched with `rrun -n N python script.py`. `rrun` starts N -identical processes, each pinned to one GPU. There is no separate client -process — every process runs the full script, acting simultaneously as client -and worker on its rank-local data. - -Because every rank runs independent Python, a `pl.DataFrame` is always -*rank-local*: it holds only that rank's fragment of the distributed dataset. -File-based sources (`scan_parquet`, `scan_csv`) distribute work automatically -— the engine assigns disjoint file- or row-group ranges to each rank. - -### Bootstrapping - -`spmd_execution()` calls `bootstrap.create_ucxx_comm(type=BackendType.AUTO)`. -Under `rrun`, `BackendType.AUTO` resolves to the `rrun`-native bootstrap -mechanism, connecting all N ranks without additional configuration. - -### GPU assignment - -`rrun` sets `CUDA_VISIBLE_DEVICES` for each process automatically. Each rank -sees exactly one GPU. No user action is required. - -### Query symmetry requirement - -All ranks must issue the **same** sequence of Polars queries in the **same** -order. Collective operations (shuffles, all-gathers, joins) are matched across -ranks by a monotonically increasing operation ID. If one rank calls a -collective that another does not, all ranks will deadlock. Driver logic must be -fully deterministic: avoid rank-conditional `collect` calls, early exits, or -branching that causes different ranks to execute different query graphs. - -### Entry point - -```python -from cudf_polars.experimental.rapidsmpf.spmd import ( - spmd_execution, - allgather_polars_dataframe, -) -``` - -`spmd_execution()` is a context manager that yields `(comm, ctx, engine)`: - -| Object | Type | Purpose | -|----------|----------------|---------------------------------------------| -| `comm` | `Communicator` | Active RapidsMPF communicator | -| `ctx` | `Context` | RapidsMPF context (GPU memory, stream pool) | -| `engine` | `pl.GPUEngine` | Polars GPU engine wired to `comm` and `ctx` | - -### Collecting results - -Each `collect(engine=engine)` returns a rank-local `pl.DataFrame`. To -assemble the full result on every rank, call `allgather_polars_dataframe()`: - -```python -full = allgather_polars_dataframe( - comm=comm, ctx=ctx, local_df=result, op_id=0 -) -``` - -`op_id` must be the same on every rank and must be unique across all -`allgather_polars_dataframe` calls in the script. - -### Example - -```python -# Launch with: rrun -n 4 python spmd_script.py -import polars as pl -from cudf_polars.experimental.rapidsmpf.spmd import ( - spmd_execution, - allgather_polars_dataframe, -) - -with spmd_execution() as (comm, ctx, engine): - result = ( - pl.scan_parquet("/data/dataset/*.parquet") - .filter(pl.col("value") > 0) - .group_by("category") - .agg(pl.col("value").sum()) - .collect(engine=engine) - ) - # result is rank-local; gather all fragments onto every rank - full = allgather_polars_dataframe( - comm=comm, ctx=ctx, local_df=result, op_id=0 - ) - # full is now identical on every rank - print(full) -``` - -### Options pass-through - -```python -with spmd_execution( - executor_options={"rapidsmpf_py_executor_max_workers": 2}, - parquet_options={"use_rapidsmpf_native": True}, -) as (comm, ctx, engine): - ... -``` - -Reserved `executor_options` keys: `"runtime"`, `"cluster"`, `"spmd"`. -Reserved `engine_kwargs` keys: `"memory_resource"`, `"executor"`. - ---- - -## 3. Ray Mode - -### Execution model - -The user runs a single client script. `ray_execution()` creates N `RankActor` -Ray remote actors — one per available GPU. The actors form a private SPMD -cluster; the client dispatches Polars IR to them and receives concatenated -results directly, with no `allgather` step needed. - -### Bootstrapping - -1. `ray_execution()` creates N `RankActor` instances (each requesting - `num_gpus=1` from Ray's resource scheduler). -2. Root actor (rank 0) calls `setup_root()` → returns its UCXX address. -3. All N actors (including the root) call `setup_worker(root_ucxx_address)` - **concurrently**. Non-root actors connect to the root; the root skips - communicator creation and proceeds directly to the barrier. All ranks must - reach the barrier simultaneously. -4. After `setup_worker` completes, each actor holds a fully initialized - `Communicator` and `Context`. - -### GPU assignment - -Ray's resource scheduler assigns `num_gpus=1` to each `RankActor` before the -actor process starts, setting `CUDA_VISIBLE_DEVICES` automatically. The -`RAY_ACCEL_ENV_VAR_OVERRIDE_ON_ZERO=0` environment variable prevents Ray from -overriding `CUDA_VISIBLE_DEVICES` to empty on the client process (which has -zero GPUs assigned). - -For hardware placement details, see [Section 5](#5-hardware-mapping-gpu-pinning). - -### Entry point - -```python -from cudf_polars.experimental.rapidsmpf.ray import ray_execution -``` - -`ray_execution()` is a context manager that yields `(ray_client, engine)`: - -| Object | Type | Purpose | -|--------------|----------------|----------------------------------------| -| `ray_client` | `RayClient` | Client handle to the actor cluster | -| `engine` | `pl.GPUEngine` | Polars GPU engine backed by Ray actors | - -### Results - -`collect(engine=engine)` dispatches the query to all actors, concatenates -their per-rank output fragments on the client, and returns a single -`pl.DataFrame`. No `allgather` is needed. - -### Ray lifecycle - -`ray_execution()` calls `ray.init()` on entry if Ray is not already -initialized, and `ray.shutdown()` on exit. If `ray.is_initialized()` returns -`True` before entry, the caller manages the cluster lifetime. - -### Diagnostics - -```python -for i, info in enumerate(ray_client.gather_cluster_info()): - print(f"rank {i}: {info}") -# Each info dict contains: pid, hostname, cuda_visible_devices, node_id -``` - -### Example - -```python -# Launch with: python ray_script.py -import polars as pl -from cudf_polars.experimental.rapidsmpf.ray import ray_execution - -with ray_execution() as (ray_client, engine): - print(ray_client.gather_cluster_info()) # verify actor placement - - result = ( - pl.scan_parquet("/data/dataset/*.parquet") - .filter(pl.col("value") > 0) - .group_by("category") - .agg(pl.col("value").sum()) - .collect(engine=engine) - ) - # result is the full concatenated output, returned directly to the client - print(result) -``` - -### Options pass-through - -```python -with ray_execution( - executor_options={"rapidsmpf_py_executor_max_workers": 2}, - engine_kwargs={"parquet_options": {"use_rapidsmpf_native": True}}, - ray_init_kwargs={"address": "auto"}, -) as (ray_client, engine): - ... -``` - -Reserved `executor_options` keys: `"runtime"`, `"cluster"`, `"spmd"`, -`"ray_client"`. Reserved `engine_kwargs` keys: `"memory_resource"`, -`"executor"`. - ---- - -## 4. Comparison - -| | SPMD | Ray | Future (Dask / custom) | -|------------------|----------------------------------|-----------------------------------|------------------------| -| Driver | Script runs on **all N workers** | Single client process | Single client process | -| Launch | `rrun -n N python script.py` | `python script.py` | `python script.py` | -| GPU pinning | `rrun` auto-pins each rank | Ray scheduler (one actor per GPU) | Depends on frontend | -| Result delivery | Rank-local; `allgather` needed | Concatenated, returned to client | Returned to client | -| Query symmetry | Required (all ranks same order) | Not required | Not required | -| Extra dependency | `rrun` / RapidsMPF | `ray` | `dask` / none | - -**Note on future frontends.** Dask and custom clients are not yet implemented, -but the architecture supports them. A new frontend only needs to: - -1. Manage a pool of SPMD workers (one per GPU). -2. Bootstrap the UCXX communicator across those workers. -3. Dispatch pickled IR + `partition_info` + `collective_id_map` to all workers - concurrently and collect their output fragments. - -No changes to the underlying RapidsMPF streaming engine are required. - ---- - -## 5. Hardware Mapping (GPU Pinning) - -### Launch model - -Hardware mapping depends on how ranks are **launched**. - -Two launch models are supported: - -* **SPMD mode**, where ranks are started directly by the `rrun` launcher. -* **Ray mode**, where ranks run inside Ray actors scheduled by the Ray runtime. - -`rrun` is a lightweight process launcher designed for GPU workloads. It starts -multiple ranks, assigns GPUs, and optionally applies topology-aware bindings. - -A typical SPMD program is started with: - -```bash -rrun -n 4 python script.py -``` - -This launches four identical Python processes. Each process becomes a **rank** -in the SPMD program and runs the same code independently. - -Because `rrun` performs the hardware setup **before the program starts**, it -must be used to launch the application. It cannot attach to or configure -processes that are already running. - -Other execution modes may require such functionality. For example, in **Ray -mode** the ranks run inside actors created by the Ray runtime rather than by -`rrun`, so GPU assignment is handled by Ray instead. - ---- - -### SPMD mode - -In SPMD mode, `rrun` assigns one GPU to each rank. - -By default it detects all GPUs on the node, but a specific list can be -provided: - -```bash -rrun -n 4 -g 0,1,2,3 python script.py -``` - -Each rank receives a single GPU via `CUDA_VISIBLE_DEVICES`. Inside the process -this GPU always appears as **device 0**, so CUDA programs require no special -configuration. - -If more ranks than GPUs are launched, multiple ranks will share a GPU. - -`rrun` can also apply topology-aware bindings so that each rank runs close to -its GPU. This includes CPU affinity, NUMA memory locality, and network-device -selection. Bindings are enabled by default and can be disabled with: - -```bash -rrun -n 4 --bind-to none python script.py -``` - -If topology discovery is unavailable, bindings are skipped automatically. - -`rrun` also sets environment variables used by RapidsMPF to bootstrap the -communication backend. - -When running under Slurm (for example with `srun`), Slurm launches the ranks -across nodes while `rrun` performs the same local setup for each rank. - ---- - -### Ray mode - -In Ray mode, ranks run as Ray actors scheduled by the Ray runtime rather than -being launched by `rrun`. - -Each rank is defined with: - -```python -@ray.remote(num_gpus=1) -``` - -Ray's scheduler selects a node with a free GPU, launches the actor there, and -sets `CUDA_VISIBLE_DEVICES` before the Python code starts. As in SPMD mode, -each rank therefore sees its assigned GPU as **device 0**. - -Unlike SPMD mode, the processes already exist when RapidsMPF code begins -executing. Ray creates the worker processes and then runs the user code inside -them. This means the `rrun` launcher cannot perform hardware setup ahead of -time. - -#### Future Request for `rrun` -To support such execution models, `rrun` will need a **library API** that can -configure an already-running process. Conceptually, the workflow would look -like: - -1. Ray launches one worker per GPU. -2. Each worker determines which GPU it has been assigned. -3. The worker calls a new `rrun` API to apply the hardware bindings locally. - -This API would configure the process in-place, for example by applying CPU -affinity, NUMA bindings, and other topology-aware settings for the specified -GPU. The GPU could be specified by index or by a stable identifier such as a -GPU UUID. - -This capability is not currently provided by `rrun`, which today only supports -processes that it launches itself, see [Launch model](#launch-model). diff --git a/python/cudf_polars/docs/cudf-polars-mp.md b/python/cudf_polars/docs/cudf-polars-mp.md index 508784d25aa3..23267a18d569 100644 --- a/python/cudf_polars/docs/cudf-polars-mp.md +++ b/python/cudf_polars/docs/cudf-polars-mp.md @@ -2,39 +2,74 @@ Multi-GPU Polars extends Polars query execution to multiple GPUs. -Multi-process (mp) execution distributes a query across several GPU workers. Each -worker owns a disjoint fragment of the data and participates in collective operations -(shuffles, all-gathers, joins) to produce a globally correct result. +Multi-GPU execution distributes a query across several GPU workers. Each worker +owns a disjoint fragment of the data and participates in collective operations +such as shuffles, all-gathers, and joins to produce a globally correct result. The entry point in all cases is the Polars `GPUEngine` with `executor="streaming"`. The `cluster` option selects the execution model: -| `cluster` value | Description | Status | -| --------------- | ------------------------------------------- | --------------- | -| `"single"` | Single-GPU, in-process execution | Stable (legacy) | -| `"distributed"` | Multi-GPU via Dask Distributed | Stable (legacy) | -| `"spmd"` | Multi-GPU via SPMD with the `rrun` launcher | Experimental | -| `"ray"` | Multi-GPU via Ray actors | Experimental | +| `cluster` | Description | Status | +| --------------- | --------------------------------------- | ----------------- | +| `"single"` | Single-GPU, in-process execution | Stable | +| `"distributed"` | Multi-GPU via Dask Distributed | Stable (legacy) | +| `"spmd"` | Multi-GPU via SPMD launched with `rrun` | Preview (new API) | +| `"ray"` | Multi-GPU via Ray actors | Preview (new API) | -This document describes the two experimental multi-GPU modes. Both rely on RapidsMPF -for shuffle and collective communication. +Two preview execution modes are available: -* [SPMD cluster mode](#spmd-cluster-mode) -* [Ray cluster mode](#ray-cluster-mode) +* **SPMD mode** — MPI-style execution where the user launches one Python process + per GPU using `rrun`. +* **Ray mode** — a single-client model where a driver program coordinates GPU + workers implemented as Ray actors. + +This document describes these two execution modes. + +* [SPMD execution mode](#spmd-execution-mode) +* [Ray execution mode](#ray-execution-mode) --- -# SPMD cluster mode +# SPMD execution mode In SPMD (Single Program, Multiple Data) execution, the same Python script is launched multiple times simultaneously, once per GPU, using the `rrun` launcher bundled with RapidsMPF. Each process is assigned a GPU and receives a **rank**. Ranks communicate through a UCXX-based communicator established at startup. -Each rank runs an independent Python process and owns its local data. File-based -sources (`scan_parquet`, `scan_csv`, etc.) are automatically partitioned so that -different ranks read different file or row-group ranges. In-memory `DataFrame` -objects are already rank-local, so each rank processes its own copy. +Each rank runs an independent Python process and owns its local data fragment. + +File-based sources (`scan_parquet`, `scan_csv`, etc.) are automatically partitioned +so that different ranks read different file or row-group ranges. In-memory +`DataFrame` objects are already rank-local, so each rank processes its own copy. + +Conceptually the setup looks like this: + +``` + rank 0 rank 1 ... rank N-1 +┌─────────────────┐ ┌─────────────────┐ ┌─────────────────┐ +│ User script │ │ User script │ │ User script │ +│ (same code on │ │ (same code on │ │ (same code on │ +│ every rank) │ │ every rank) │ │ every rank) │ +└────────┬────────┘ └────────┬────────┘ └────────┬────────┘ + │ │ │ + │ LazyFrame.collect(engine=engine) │ + ↓ ↓ ↓ +┌─────────────────┐ ┌─────────────────┐ ┌─────────────────┐ +│ run IR │ │ run IR │ │ run IR │ +└────────┬────────┘ └────────┬────────┘ └────────┬────────┘ + │ │ │ + ↓ ↓ ↓ +┌────────────────────────────────────────────────────────────────┐ +│ RapidsMPF streaming engine │ +│ shuffle / all-gather · UCXX communicator · RMM GPU memory │ +└────────────────────────────────────────────────────────────────┘ + ↑ ↑ ↑ + GPU 0 GPU 1 GPU N-1 +``` + +After `collect`, results are **rank-local**. To assemble the full dataset on +every rank, call `allgather_polars_dataframe()`. ## Prerequisites @@ -89,6 +124,18 @@ The context manager yields: Pass `engine` to every `LazyFrame.collect()` inside the context block. +## Query symmetry requirement + +All ranks must execute the **same sequence of queries in the same order**. Collective +operations are matched using internal operation IDs. If one rank executes a collective +that another rank does not, the program will deadlock. + +In practice: + +* Avoid rank-conditional `collect()` calls +* Avoid branches that change the query graph +* Keep the driver script deterministic + ## Collecting distributed results `collect()` returns a rank-local result. Use @@ -104,23 +151,10 @@ full = allgather_polars_dataframe( ``` `op_id` is a unique integer that identifies this collective operation across ranks. -All ranks must call the same collective with the same `op_id`. Otherwise the program -will deadlock. +All ranks must call the same collective with the same `op_id`. The result is a `pl.DataFrame` containing rows from all ranks, ordered by rank. -## Query symmetry requirement - -All ranks must execute the **same sequence of queries in the same order**. Collective -operations are matched using internal operation IDs. If one rank executes a collective -that another rank does not, the program will deadlock. - -In practice: - -* Avoid rank-conditional `collect()` calls -* Avoid branches that change the query graph -* Keep the driver script deterministic - ## Passing options `executor_options` and `engine_kwargs` accept pass-through dictionaries: @@ -131,34 +165,61 @@ with spmd_execution( "max_rows_per_partition": 500_000, "rapidsmpf_spill": True, }, - parquet_options={"use_rapidsmpf_native": True}, # forwarded via **engine_kwargs + parquet_options={"use_rapidsmpf_native": True}, ) as (comm, ctx, engine): ... ``` -`executor_options` keys map to `StreamingExecutor` fields. Any additional keyword -arguments to `spmd_execution()` (such as `parquet_options`) are forwarded directly -to `pl.GPUEngine` as `**engine_kwargs`. +`executor_options` keys map to `StreamingExecutor` fields. Additional keyword +arguments to `spmd_execution()` are forwarded directly to `pl.GPUEngine`. + +Reserved keys: -The keys `"runtime"`, `"cluster"`, and `"spmd"` in `executor_options`, and -`"memory_resource"` and `"executor"` in `engine_kwargs`, are reserved. +* `executor_options`: `"runtime"`, `"cluster"`, `"spmd"` +* `engine_kwargs`: `"memory_resource"`, `"executor"` --- -# Ray cluster mode +# Ray execution mode Ray mode uses a single client process that drives execution across multiple GPU -workers. Internally, the system uses the concept of **ranks**, similar to MPI ranks. -Each rank corresponds to one GPU worker and participates in collective operations +workers. + +Internally the system uses the concept of **ranks**, similar to MPI ranks. Each +rank corresponds to one GPU worker and participates in collective operations through a shared UCXX communicator. -In the Ray implementation, each rank is implemented as a **Ray actor**, with one -actor created per available GPU. Each rank owns its GPU, memory resource, -communicator endpoint, and RapidsMPF streaming context. +In the Ray implementation each rank is implemented as a **Ray actor**, with one +actor created per available GPU. + +Conceptually the system looks like this: -The client sends the query plan to all ranks. The ranks execute the pipeline +``` + ┌──────────────────────────────┐ + │ User script │ + │ (single client process) │ + │ LazyFrame.collect(engine=…) │ + └──────────────┬───────────────┘ + │ IR dispatched to all actors + ┌────────────────|─────────────────┐ + ↓ ↓ ↓ + ┌─────────────┐ ┌─────────────┐ ┌─────────────┐ + │ RankActor │ │ RankActor │ │ RankActor │ + │ rank 0 │ │ rank 1 │ │ rank N-1 │ + │ run IR │ │ run IR │ │ run IR │ + └──────┬──────┘ └──────┬──────┘ └──────┬──────┘ + ↓ ↓ ↓ +┌────────────────────────────────────────────────────────────────┐ +│ RapidsMPF streaming engine │ +│ shuffle / all-gather · UCXX communicator · RMM GPU memory │ +└────────────────────────────────────────────────────────────────┘ + ↑ ↑ ↑ + GPU 0 GPU 1 GPU N-1 +``` + +The client broadcasts the query plan to all ranks. The ranks execute the pipeline collectively through UCXX, and their outputs are streamed back and concatenated on -the client. +the client process. Unlike SPMD mode, the driver script runs as a normal Python program. There is no `rrun` launcher and no symmetry requirement for the driver code. @@ -229,6 +290,7 @@ GPUs are available. ```python with ray_execution() as (ray_client, engine): + print(f"cluster has {ray_client.nranks} ranks") for i, info in enumerate(ray_client.gather_cluster_info()): print( f"rank {i}: hostname={info['hostname']}, pid={info['pid']}, " @@ -252,6 +314,7 @@ with ray_execution( ... ``` -The keys `"runtime"`, `"cluster"`, `"spmd"`, and `"ray_client"` in -`executor_options`, and `"memory_resource"` and `"executor"` in `engine_kwargs`, -are reserved. +Reserved keys: + +* `executor_options`: `"runtime"`, `"cluster"`, `"spmd"`, `"ray_context"` +* `engine_kwargs`: `"memory_resource"`, `"executor"` From 18bb8e4d31de7a0eb451bcb7d6dc61f06d5353b6 Mon Sep 17 00:00:00 2001 From: "Mads R. B. Kristensen" Date: Fri, 13 Mar 2026 13:58:32 +0100 Subject: [PATCH 18/33] docs --- python/cudf_polars/docs/cudf-polars-mp.md | 74 +++++++++++++++++------ 1 file changed, 57 insertions(+), 17 deletions(-) diff --git a/python/cudf_polars/docs/cudf-polars-mp.md b/python/cudf_polars/docs/cudf-polars-mp.md index 23267a18d569..e44a2ac386d8 100644 --- a/python/cudf_polars/docs/cudf-polars-mp.md +++ b/python/cudf_polars/docs/cudf-polars-mp.md @@ -9,12 +9,12 @@ such as shuffles, all-gathers, and joins to produce a globally correct result. The entry point in all cases is the Polars `GPUEngine` with `executor="streaming"`. The `cluster` option selects the execution model: -| `cluster` | Description | Status | -| --------------- | --------------------------------------- | ----------------- | -| `"single"` | Single-GPU, in-process execution | Stable | -| `"distributed"` | Multi-GPU via Dask Distributed | Stable (legacy) | -| `"spmd"` | Multi-GPU via SPMD launched with `rrun` | Preview (new API) | -| `"ray"` | Multi-GPU via Ray actors | Preview (new API) | +| `cluster` | Description | Status | +| --------------- | ---------------------------------------------------- | ----------------- | +| `"single"` | Single-GPU, in-process execution | Stable | +| `"distributed"` | Multi-GPU via [Dask Distributed][dask-distributed] | Stable (legacy) | +| `"spmd"` | Multi-GPU via [SPMD][spmd-wiki] launched with `rrun` | Preview (new API) | +| `"ray"` | Multi-GPU via [Ray][ray-docs] actors | Preview (new API) | Two preview execution modes are available: @@ -53,7 +53,9 @@ Conceptually the setup looks like this: │ every rank) │ │ every rank) │ │ every rank) │ └────────┬────────┘ └────────┬────────┘ └────────┬────────┘ │ │ │ - │ LazyFrame.collect(engine=engine) │ +┌────────┴────────────────────┴───────────────────────┴────────┐ +│ LazyFrame.collect(engine=engine) │ +└────────┬────────────────────┬───────────────────────┬────────┘ ↓ ↓ ↓ ┌─────────────────┐ ┌─────────────────┐ ┌─────────────────┐ │ run IR │ │ run IR │ │ run IR │ @@ -118,9 +120,9 @@ with spmd_execution() as (comm, ctx, engine): The context manager yields: -* `comm` — `rapidsmpf.communicator.Communicator` -* `ctx` — `rapidsmpf.streaming.core.context.Context` -* `engine` — `pl.GPUEngine` configured for SPMD execution +* `comm` — [`rapidsmpf.communicator.Communicator`][rapidsmpf-communicator] +* `ctx` — [`rapidsmpf.streaming.core.context.Context`][rapidsmpf-context] +* `engine` — [`pl.GPUEngine`][polars-gpuengine] configured for SPMD execution Pass `engine` to every `LazyFrame.collect()` inside the context block. @@ -136,6 +138,32 @@ In practice: * Avoid branches that change the query graph * Keep the driver script deterministic +**Example that works correctly:** + +```python +# Every rank executes the same query in the same order. +with spmd_execution() as (comm, ctx, engine): + result = ( + pl.scan_parquet("/data/*.parquet") + .filter(pl.col("amount") > 100) + .group_by("customer_id") + .agg(pl.col("amount").sum()) + .collect(engine=engine) + ) +``` + +**Example that deadlocks:** + +```python +# Rank 0 executes a group_by collective; other ranks do not. +# The collective IDs go out of sync → deadlock. +with spmd_execution() as (comm, ctx, engine): + df = pl.scan_parquet("/data/*.parquet") + if comm.rank() == 0: # DON'T DO THIS + df = df.group_by("customer_id").agg(pl.col("amount").sum()) + result = df.collect(engine=engine) +``` + ## Collecting distributed results `collect()` returns a rank-local result. Use @@ -185,12 +213,13 @@ Reserved keys: Ray mode uses a single client process that drives execution across multiple GPU workers. -Internally the system uses the concept of **ranks**, similar to MPI ranks. Each -rank corresponds to one GPU worker and participates in collective operations -through a shared UCXX communicator. +Internally the system uses the concept of **ranks** (described in the +[SPMD section](#spmd-execution-mode) above). Each rank corresponds to one GPU +worker and participates in collective operations through a shared UCXX +communicator. -In the Ray implementation each rank is implemented as a **Ray actor**, with one -actor created per available GPU. +In the Ray implementation each rank is implemented as a [**Ray actor**][ray-actors], +with one actor created per available GPU. Conceptually the system looks like this: @@ -221,8 +250,10 @@ The client broadcasts the query plan to all ranks. The ranks execute the pipelin collectively through UCXX, and their outputs are streamed back and concatenated on the client process. -Unlike SPMD mode, the driver script runs as a normal Python program. There is no -`rrun` launcher and no symmetry requirement for the driver code. +Unlike SPMD mode, the driver script runs as a normal Python program with no +`rrun` launcher. Query symmetry is handled automatically: the client serializes +the complete query plan and broadcasts it to all actors, so every rank always +executes the same query. ## Prerequisites @@ -318,3 +349,12 @@ Reserved keys: * `executor_options`: `"runtime"`, `"cluster"`, `"spmd"`, `"ray_context"` * `engine_kwargs`: `"memory_resource"`, `"executor"` + + +[dask-distributed]: https://distributed.dask.org/ +[spmd-wiki]: https://en.wikipedia.org/wiki/Single_program,_multiple_data +[ray-docs]: https://docs.ray.io/ +[ray-actors]: https://docs.ray.io/en/latest/ray-core/actors.html +[rapidsmpf-communicator]: https://docs.rapids.ai/api/rapidsmpf/stable/api/communicator/ +[rapidsmpf-context]: https://docs.rapids.ai/api/rapidsmpf/stable/api/streaming/context/ +[polars-gpuengine]: https://docs.pola.rs/api/python/stable/reference/api/polars.GPUEngine.html From c32c0b45e5feab36e5daa37eee462a6a64147441 Mon Sep 17 00:00:00 2001 From: "Mads R. B. Kristensen" Date: Fri, 13 Mar 2026 14:11:01 +0100 Subject: [PATCH 19/33] @wence- review --- .../cudf_polars/experimental/rapidsmpf/ray.py | 53 ++++++++++++------- python/cudf_polars/docs/cudf-polars-mp.md | 10 ++-- 2 files changed, 37 insertions(+), 26 deletions(-) diff --git a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/ray.py b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/ray.py index 123b1cd63669..44b5c8843ba3 100644 --- a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/ray.py +++ b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/ray.py @@ -249,6 +249,9 @@ def shutdown(self) -> None: Raises `ray.exceptions.RayActorError`. """ self._py_executor.shutdown(wait=True, cancel_futures=True) + # Release resources in dependency order before exit_actor() terminates + # the process. + self._ctx = None self._comm = None self._mr = None ray.actor.exit_actor() @@ -364,6 +367,10 @@ def evaluate_polars_ir( list(ir.schema.values()), stream, ) + # Ray transfers the returned Polars DataFrame back to the client via the + # object store (pickle / Arrow IPC). The DataFrame is already on CPU at + # this point (to_polars() copies the result off-GPU), so no GPU memory + # crosses process boundaries. result = df.to_polars() return result, metadata_collector @@ -443,18 +450,28 @@ def shutdown(self) -> None: If Ray was initialized by this client, also calls :func:`ray.shutdown`. Safe to call more than once. + + Raises + ------ + ExceptionGroup + If one or more actors raise an unexpected exception during shutdown. """ - for a in self.rank_actors: - try: - ray.get(a.shutdown.remote()) - except ray.exceptions.RayActorError: - pass # expected: exit_actor() terminates the process immediately - except Exception as e: - print(f"shutdown error: {e}") - self._engine = None - if self._owns_ray: - self._owns_ray = False - ray.shutdown() + exceptions: list[Exception] = [] + try: + for a in self.rank_actors: + try: + ray.get(a.shutdown.remote()) + except ray.exceptions.RayActorError: + pass # expected: exit_actor() terminates the process immediately + except Exception as e: + exceptions.append(e) + if exceptions: + raise ExceptionGroup("Actor shutdown failed", exceptions) + finally: + self._engine = None + if self._owns_ray: + self._owns_ray = False + ray.shutdown() def __enter__(self) -> tuple[RayClient, pl.GPUEngine]: """Enter the context manager, returning ``(self, engine)``.""" @@ -476,8 +493,10 @@ def ray_execution( Create a RapidsMPF Ray cluster and return a :class:`RayClient`. The returned client supports both direct use and the context-manager - protocol. Prefer the context-manager form in scripts; use the direct - form in interactive environments such as Jupyter notebooks. + protocol. Prefer the context-manager form in scripts: it guarantees that + actors and Ray are shut down even if an exception is raised. In interactive + environments such as Jupyter notebooks, the direct form lets the cluster + persist across multiple cells without tearing it down after every query. If Ray is not already initialized, :func:`ray.init` is called here and :func:`ray.shutdown` is called by :meth:`RayClient.shutdown`. If Ray is @@ -490,8 +509,6 @@ def ray_execution( ``Options(get_environment_variables())``. executor_options Additional key-value pairs forwarded to the Polars executor options. - If ``"max_io_threads"`` is present, its value is also used as the - default for ``rapidsmpf_options["num_streaming_threads"]``. engine_kwargs Additional keyword arguments forwarded to :class:`polars.GPUEngine`. ray_init_kwargs @@ -552,9 +569,7 @@ def ray_execution( if rapidsmpf_options is not None else Options(get_environment_variables()) ) - rapidsmpf_options.insert_if_absent( - {"num_streaming_threads": str(executor_options.get("max_io_threads", 4))} - ) + rapidsmpf_options.insert_if_absent({"num_streaming_threads": "4"}) rapidsmpf_options_as_bytes = rapidsmpf_options.serialize() ray_was_initialized: bool = ray.is_initialized() @@ -576,7 +591,7 @@ def ray_execution( if free_gpus == 0: raise RuntimeError("No available GPUs in the Ray cluster at startup") - # Create one actor per GPU. Ray adds .remote() dynamically; no type stubs. + # Create one actor per GPU. rank_actors: list[ActorHandle[RankActor]] = [ RankActor.remote( # type: ignore[attr-defined] nranks=free_gpus, diff --git a/python/cudf_polars/docs/cudf-polars-mp.md b/python/cudf_polars/docs/cudf-polars-mp.md index e44a2ac386d8..3f11c848bf1e 100644 --- a/python/cudf_polars/docs/cudf-polars-mp.md +++ b/python/cudf_polars/docs/cudf-polars-mp.md @@ -210,13 +210,9 @@ Reserved keys: # Ray execution mode -Ray mode uses a single client process that drives execution across multiple GPU -workers. - -Internally the system uses the concept of **ranks** (described in the -[SPMD section](#spmd-execution-mode) above). Each rank corresponds to one GPU -worker and participates in collective operations through a shared UCXX -communicator. +Ray mode uses a single client process that drives execution across multiple ranks. +Each rank corresponds to one GPU worker and participates in collective operations +through a shared UCXX communicator. In the Ray implementation each rank is implemented as a [**Ray actor**][ray-actors], with one actor created per available GPU. From 68aafbd06e57b545f9ed343cb3fb5b107882325d Mon Sep 17 00:00:00 2001 From: "Mads R. B. Kristensen" Date: Fri, 13 Mar 2026 14:15:04 +0100 Subject: [PATCH 20/33] doc --- python/cudf_polars/docs/cudf-polars-mp.md | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/python/cudf_polars/docs/cudf-polars-mp.md b/python/cudf_polars/docs/cudf-polars-mp.md index 3f11c848bf1e..b37d978d3264 100644 --- a/python/cudf_polars/docs/cudf-polars-mp.md +++ b/python/cudf_polars/docs/cudf-polars-mp.md @@ -18,8 +18,8 @@ The `cluster` option selects the execution model: Two preview execution modes are available: -* **SPMD mode** — MPI-style execution where the user launches one Python process - per GPU using `rrun`. +* **SPMD mode** — each GPU runs the same script as an independent process, + launched with `rrun`. * **Ray mode** — a single-client model where a driver program coordinates GPU workers implemented as Ray actors. From 12060c380210ed10dc4a5c0d598ff1549165063f Mon Sep 17 00:00:00 2001 From: "Mads R. B. Kristensen" Date: Mon, 16 Mar 2026 10:39:01 +0100 Subject: [PATCH 21/33] fix test patch --- python/cudf_polars/tests/experimental/rapidsmpf/test_ray.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/cudf_polars/tests/experimental/rapidsmpf/test_ray.py b/python/cudf_polars/tests/experimental/rapidsmpf/test_ray.py index f65b69b7686d..dd1cb92897de 100644 --- a/python/cudf_polars/tests/experimental/rapidsmpf/test_ray.py +++ b/python/cudf_polars/tests/experimental/rapidsmpf/test_ray.py @@ -106,7 +106,7 @@ def test_ray_execution_raises_inside_rrun() -> None: """ray_execution() must not be called from within an rrun cluster.""" with ( patch( - "cudf_polars.experimental.rapidsmpf.ray.bootstrap.is_running_with_rrun", + "rapidsmpf.bootstrap.is_running_with_rrun", return_value=True, ), pytest.raises(RuntimeError, match="rrun"), From 6ca5a751d725a4a2904e0914a7794f8dd456dde0 Mon Sep 17 00:00:00 2001 From: "Mads R. B. Kristensen" Date: Mon, 16 Mar 2026 13:09:39 +0100 Subject: [PATCH 22/33] Apply suggestions from code review Co-authored-by: Tom Augspurger Co-authored-by: Lawrence Mitchell --- python/cudf_polars/docs/cudf-polars-mp.md | 28 +++++++++++------------ 1 file changed, 14 insertions(+), 14 deletions(-) diff --git a/python/cudf_polars/docs/cudf-polars-mp.md b/python/cudf_polars/docs/cudf-polars-mp.md index b37d978d3264..00d071f45416 100644 --- a/python/cudf_polars/docs/cudf-polars-mp.md +++ b/python/cudf_polars/docs/cudf-polars-mp.md @@ -30,7 +30,7 @@ This document describes these two execution modes. --- -# SPMD execution mode +## SPMD execution mode In SPMD (Single Program, Multiple Data) execution, the same Python script is launched multiple times simultaneously, once per GPU, using the `rrun` launcher bundled with @@ -73,13 +73,13 @@ Conceptually the setup looks like this: After `collect`, results are **rank-local**. To assemble the full dataset on every rank, call `allgather_polars_dataframe()`. -## Prerequisites +### Prerequisites * RapidsMPF (`rapidsmpf`) installed * UCXX available (usually installed with RapidsMPF) * `rrun` launcher available (`rrun --help` should succeed) -## Running in SPMD mode +### Running in SPMD mode `spmd_execution()` is the primary entry point for SPMD execution. It is a context manager imported from `cudf_polars.experimental.rapidsmpf.spmd`. On entry it: @@ -122,11 +122,11 @@ The context manager yields: * `comm` — [`rapidsmpf.communicator.Communicator`][rapidsmpf-communicator] * `ctx` — [`rapidsmpf.streaming.core.context.Context`][rapidsmpf-context] -* `engine` — [`pl.GPUEngine`][polars-gpuengine] configured for SPMD execution +* `engine` — {class}`~polars.lazyframe.engine_config.GPUEngine` Pass `engine` to every `LazyFrame.collect()` inside the context block. -## Query symmetry requirement +### Query symmetry requirement All ranks must execute the **same sequence of queries in the same order**. Collective operations are matched using internal operation IDs. If one rank executes a collective @@ -159,12 +159,12 @@ with spmd_execution() as (comm, ctx, engine): # The collective IDs go out of sync → deadlock. with spmd_execution() as (comm, ctx, engine): df = pl.scan_parquet("/data/*.parquet") - if comm.rank() == 0: # DON'T DO THIS + if comm.rank == 0: # DON'T DO THIS df = df.group_by("customer_id").agg(pl.col("amount").sum()) result = df.collect(engine=engine) ``` -## Collecting distributed results +### Collecting distributed results `collect()` returns a rank-local result. Use `allgather_polars_dataframe()` to gather all fragments: @@ -183,7 +183,7 @@ All ranks must call the same collective with the same `op_id`. The result is a `pl.DataFrame` containing rows from all ranks, ordered by rank. -## Passing options +### Passing options `executor_options` and `engine_kwargs` accept pass-through dictionaries: @@ -208,7 +208,7 @@ Reserved keys: --- -# Ray execution mode +## Ray execution mode Ray mode uses a single client process that drives execution across multiple ranks. Each rank corresponds to one GPU worker and participates in collective operations @@ -251,12 +251,12 @@ Unlike SPMD mode, the driver script runs as a normal Python program with no the complete query plan and broadcasts it to all actors, so every rank always executes the same query. -## Prerequisites +### Prerequisites * Ray (`ray`) installed * RapidsMPF and UCXX available on all GPU nodes -## Running in Ray mode +### Running in Ray mode `ray_execution()` is imported from `cudf_polars.experimental.rapidsmpf.ray`. It: @@ -289,7 +289,7 @@ The context manager yields: * `ray_client` — cluster diagnostics and utilities * `engine` — `pl.GPUEngine` configured for Ray execution -## Ray lifecycle +### Ray lifecycle If Ray is already initialized, `ray_execution()` attaches to the existing cluster and does not call `ray.shutdown()` on exit. @@ -311,7 +311,7 @@ finally: `ray_execution()` raises `RuntimeError` if called inside an `rrun` cluster or if no GPUs are available. -## Cluster diagnostics +### Cluster diagnostics `RayClient.gather_cluster_info()` returns placement information for all rank actors: @@ -327,7 +327,7 @@ with ray_execution() as (ray_client, engine): Each entry includes `pid`, `hostname`, `cuda_visible_devices`, and `node_id`. -## Passing options +### Passing options `executor_options`, `engine_kwargs`, and `ray_init_kwargs` accept pass-through dictionaries: From 43539a28af426103e0e8e5f4fa691d175138a964 Mon Sep 17 00:00:00 2001 From: "Mads R. B. Kristensen" Date: Mon, 16 Mar 2026 13:41:06 +0100 Subject: [PATCH 23/33] changes based on reviews --- .../cudf_polars/experimental/rapidsmpf/ray.py | 36 ++++++++++--------- .../experimental/rapidsmpf/spmd.py | 8 ++--- python/cudf_polars/docs/cudf-polars-mp.md | 6 ++-- .../tests/experimental/rapidsmpf/test_ray.py | 12 +++---- .../tests/experimental/rapidsmpf/test_spmd.py | 4 +-- 5 files changed, 34 insertions(+), 32 deletions(-) diff --git a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/ray.py b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/ray.py index 44b5c8843ba3..d47a1320acf3 100644 --- a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/ray.py +++ b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/ray.py @@ -458,9 +458,10 @@ def shutdown(self) -> None: """ exceptions: list[Exception] = [] try: - for a in self.rank_actors: + refs = [a.shutdown.remote() for a in self.rank_actors] + for ref in refs: try: - ray.get(a.shutdown.remote()) + ray.get(ref) except ray.exceptions.RayActorError: pass # expected: exit_actor() terminates the process immediately except Exception as e: @@ -486,8 +487,8 @@ def ray_execution( *, rapidsmpf_options: Options | None = None, executor_options: dict[str, object] | None = None, - engine_kwargs: dict[str, Any] | None = None, - ray_init_kwargs: dict[str, object] | None = None, + engine_options: dict[str, Any] | None = None, + ray_init_options: dict[str, object] | None = None, ) -> RayClient: """ Create a RapidsMPF Ray cluster and return a :class:`RayClient`. @@ -509,9 +510,9 @@ def ray_execution( ``Options(get_environment_variables())``. executor_options Additional key-value pairs forwarded to the Polars executor options. - engine_kwargs + engine_options Additional keyword arguments forwarded to :class:`polars.GPUEngine`. - ray_init_kwargs + ray_init_options Keyword arguments forwarded to :func:`ray.init` when Ray is not already initialized. @@ -529,10 +530,10 @@ def ray_execution( If not all GPUs in the Ray cluster are free at startup. RuntimeError If no GPUs are available in the Ray cluster. - ValueError + TypeError If ``executor_options`` contains a reserved key. - ValueError - If ``engine_kwargs`` contains a reserved key. + TypeError + If ``engine_options`` contains a reserved key. Examples -------- @@ -548,8 +549,8 @@ def ray_execution( >>> client.shutdown() """ executor_options = executor_options or {} - engine_kwargs = engine_kwargs or {} - ray_init_kwargs = ray_init_kwargs or {} + engine_options = engine_options or {} + ray_init_options = ray_init_options or {} if bootstrap.is_running_with_rrun(): raise RuntimeError( @@ -560,9 +561,9 @@ def ray_execution( # Check for reserved keys. if bad := {"runtime", "cluster", "spmd", "ray_context"} & executor_options.keys(): - raise ValueError(f"executor_options may not contain reserved keys: {bad}") - if bad := {"memory_resource", "executor"} & engine_kwargs.keys(): - raise ValueError(f"engine_kwargs may not contain reserved keys: {bad}") + raise TypeError(f"executor_options may not contain reserved keys: {bad}") + if bad := {"memory_resource", "executor"} & engine_options.keys(): + raise TypeError(f"engine_options may not contain reserved keys: {bad}") rapidsmpf_options = ( rapidsmpf_options @@ -575,9 +576,10 @@ def ray_execution( ray_was_initialized: bool = ray.is_initialized() if not ray_was_initialized: # Prevent Ray from overriding CUDA_VISIBLE_DEVICES to "" when a worker - # process starts with zero visible GPUs (e.g. the driver process itself). + # process starts with zero visible GPUs (e.g., the driver process itself). + # In the future, this behavior will become the default in Ray. os.environ.setdefault("RAY_ACCEL_ENV_VAR_OVERRIDE_ON_ZERO", "0") - ray.init(**ray_init_kwargs) + ray.init(**ray_init_options) total_gpus = int(ray.cluster_resources().get("GPU", 0.0)) # Note: available_resources() is a snapshot and inherently racy. This is a @@ -618,6 +620,6 @@ def ray_execution( "cluster": "ray", "ray_context": RayContext(rank_actors), }, - **engine_kwargs, + **engine_options, ) return RayClient(engine, owns_ray=not ray_was_initialized) diff --git a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/spmd.py b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/spmd.py index abf797a7ee35..aa093ff20e6f 100644 --- a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/spmd.py +++ b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/spmd.py @@ -311,10 +311,10 @@ def spmd_execution( RuntimeError If not running under the ``rrun`` launcher (i.e. :func:`rapidsmpf.bootstrap.is_running_with_rrun` returns ``False``). - ValueError + TypeError If ``executor_options`` contains any of the reserved keys ``"runtime"``, ``"cluster"``, or ``"spmd"``. - ValueError + TypeError If ``engine_kwargs`` contains any of the reserved keys ``"raise_on_fail"``, ``"memory_resource"``, or ``"executor"``. @@ -340,9 +340,9 @@ def spmd_execution( # Check for reserved keys. if bad := {"runtime", "cluster", "spmd"} & executor_options.keys(): - raise ValueError(f"executor_options may not contain reserved keys: {bad}") + raise TypeError(f"executor_options may not contain reserved keys: {bad}") if bad := {"memory_resource", "executor"} & engine_kwargs.keys(): - raise ValueError(f"engine_kwargs may not contain reserved keys: {bad}") + raise TypeError(f"engine_kwargs may not contain reserved keys: {bad}") rapidsmpf_options = ( rapidsmpf_options diff --git a/python/cudf_polars/docs/cudf-polars-mp.md b/python/cudf_polars/docs/cudf-polars-mp.md index 00d071f45416..337af96f7ddb 100644 --- a/python/cudf_polars/docs/cudf-polars-mp.md +++ b/python/cudf_polars/docs/cudf-polars-mp.md @@ -329,14 +329,14 @@ Each entry includes `pid`, `hostname`, `cuda_visible_devices`, and `node_id`. ### Passing options -`executor_options`, `engine_kwargs`, and `ray_init_kwargs` accept pass-through +`executor_options`, `engine_options`, and `ray_init_options` accept pass-through dictionaries: ```python with ray_execution( executor_options={"max_rows_per_partition": 500_000}, - engine_kwargs={"raise_on_fail": True}, - ray_init_kwargs={"num_cpus": 4}, + engine_options={"raise_on_fail": True}, + ray_init_options={"num_cpus": 4}, ) as (ray_client, engine): ... ``` diff --git a/python/cudf_polars/tests/experimental/rapidsmpf/test_ray.py b/python/cudf_polars/tests/experimental/rapidsmpf/test_ray.py index dd1cb92897de..7477d2af7477 100644 --- a/python/cudf_polars/tests/experimental/rapidsmpf/test_ray.py +++ b/python/cudf_polars/tests/experimental/rapidsmpf/test_ray.py @@ -32,7 +32,7 @@ def _ray_env() -> Iterator[tuple[RayClient, pl.GPUEngine]]: # Use a small partition size so tests exercise the multi-partition # code path deterministically, regardless of input size. executor_options={"max_rows_per_partition": 10}, - ray_init_kwargs={"include_dashboard": False}, + ray_init_options={"include_dashboard": False}, ) as ( ray_client, engine, @@ -69,16 +69,16 @@ def engine(_ray_env: tuple[RayClient, pl.GPUEngine]) -> pl.GPUEngine: def test_ray_execution_reserved_executor_keys() -> None: """executor_options rejects reserved keys.""" for key in ("runtime", "cluster", "spmd", "ray_context"): - with pytest.raises(ValueError, match="reserved"): + with pytest.raises(TypeError, match="reserved"): ray_execution(executor_options={key: "anything"}) -def test_ray_execution_reserved_engine_kwargs_keys() -> None: - """engine_kwargs rejects keys that are set explicitly by ray_execution.""" +def test_ray_execution_reserved_engine_options_keys() -> None: + """engine_options rejects keys that are set explicitly by ray_execution.""" for key in ("memory_resource", "executor"): kwargs: dict[str, Any] = {key: "anything"} - with pytest.raises(ValueError, match="reserved"): - ray_execution(engine_kwargs=kwargs) + with pytest.raises(TypeError, match="reserved"): + ray_execution(engine_options=kwargs) def test_ray_client_shutdown_idempotent() -> None: diff --git a/python/cudf_polars/tests/experimental/rapidsmpf/test_spmd.py b/python/cudf_polars/tests/experimental/rapidsmpf/test_spmd.py index 585da785833c..d26f4ca70d69 100644 --- a/python/cudf_polars/tests/experimental/rapidsmpf/test_spmd.py +++ b/python/cudf_polars/tests/experimental/rapidsmpf/test_spmd.py @@ -37,7 +37,7 @@ def test_spmd_execution_reserved_keys() -> None: """executor_options rejects reserved keys.""" for key in ("runtime", "cluster", "spmd"): with ( - pytest.raises(ValueError, match="reserved"), + pytest.raises(TypeError, match="reserved"), spmd_execution(executor_options={key: "anything"}), ): pass @@ -48,7 +48,7 @@ def test_spmd_execution_engine_kwargs_reserved_keys() -> None: for key in ("memory_resource", "executor"): kwargs: dict[str, Any] = {key: "anything"} with ( - pytest.raises(ValueError, match="reserved"), + pytest.raises(TypeError, match="reserved"), spmd_execution(**kwargs), ): pass From fc137fc7f71a489a3f421c14eaa4d8b54900f8d9 Mon Sep 17 00:00:00 2001 From: "Mads R. B. Kristensen" Date: Mon, 16 Mar 2026 14:04:41 +0100 Subject: [PATCH 24/33] changes based on reviews --- .../cudf_polars/experimental/rapidsmpf/ray.py | 16 +++++++++++----- python/cudf_polars/docs/cudf-polars-mp.md | 17 +++++++++++------ 2 files changed, 22 insertions(+), 11 deletions(-) diff --git a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/ray.py b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/ray.py index d47a1320acf3..ca19f0d6a488 100644 --- a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/ray.py +++ b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/ray.py @@ -404,9 +404,9 @@ def __init__( @property def rank_actors(self) -> list[ActorHandle[RankActor]]: - """List of Ray rank actor handles, or ``[]`` after :meth:`shutdown`.""" + """List of Ray rank actor handles.""" if self._engine is None: - return [] + raise RuntimeError("rank_actors is not available after shutdown") return self._engine.config["executor_options"]["ray_context"].rank_actors @property @@ -456,6 +456,8 @@ def shutdown(self) -> None: ExceptionGroup If one or more actors raise an unexpected exception during shutdown. """ + if self._engine is None: + return # already shut down; idempotent exceptions: list[Exception] = [] try: refs = [a.shutdown.remote() for a in self.rank_actors] @@ -469,6 +471,9 @@ def shutdown(self) -> None: if exceptions: raise ExceptionGroup("Actor shutdown failed", exceptions) finally: + # Setting _engine to None serves two purposes: it marks this client as + # shut down (so engine/rank_actors/nranks raise) and releases the + # reference to RayContext, allowing its now-dead actor handles to be GC'd. self._engine = None if self._owns_ray: self._owns_ray = False @@ -582,9 +587,10 @@ def ray_execution( ray.init(**ray_init_options) total_gpus = int(ray.cluster_resources().get("GPU", 0.0)) - # Note: available_resources() is a snapshot and inherently racy. This is a - # best-effort guard; another process could claim GPUs between this check and - # actor creation. + # Note: available_resources() returns a snapshot and is inherently racy. + # This is only a best-effort guard, another process could claim GPUs between + # this check and actor creation. That is fine here, because the snapshot + # simply determines the number of ranks used for this cluster instance. free_gpus = int(ray.available_resources().get("GPU", 0.0)) if total_gpus != free_gpus: raise RuntimeError( diff --git a/python/cudf_polars/docs/cudf-polars-mp.md b/python/cudf_polars/docs/cudf-polars-mp.md index 337af96f7ddb..39078f529d13 100644 --- a/python/cudf_polars/docs/cudf-polars-mp.md +++ b/python/cudf_polars/docs/cudf-polars-mp.md @@ -124,7 +124,7 @@ The context manager yields: * `ctx` — [`rapidsmpf.streaming.core.context.Context`][rapidsmpf-context] * `engine` — {class}`~polars.lazyframe.engine_config.GPUEngine` -Pass `engine` to every `LazyFrame.collect()` inside the context block. +Pass `engine` to every `LazyFrame.collect()` or `sink*()` call inside the context block. ### Query symmetry requirement @@ -134,7 +134,7 @@ that another rank does not, the program will deadlock. In practice: -* Avoid rank-conditional `collect()` calls +* Avoid rank-conditional `collect()` or `sink*()` calls * Avoid branches that change the query graph * Keep the driver script deterministic @@ -198,13 +198,14 @@ with spmd_execution( ... ``` -`executor_options` keys map to `StreamingExecutor` fields. Additional keyword -arguments to `spmd_execution()` are forwarded directly to `pl.GPUEngine`. +`executor_options` is forwarded directly to `pl.GPUEngine` as its `executor_options` +argument; user-supplied keys are merged with reserved entries set by `spmd_execution()`. +Any additional keyword arguments to `spmd_execution()` are also forwarded to `pl.GPUEngine`. Reserved keys: * `executor_options`: `"runtime"`, `"cluster"`, `"spmd"` -* `engine_kwargs`: `"memory_resource"`, `"executor"` +* `engine_options`: `"memory_resource"`, `"executor"` --- @@ -341,10 +342,14 @@ with ray_execution( ... ``` +`executor_options` is forwarded directly to `pl.GPUEngine` as its `executor_options` +argument; user-supplied keys are merged with reserved entries set by `ray_execution()`. +Any additional keyword arguments to `ray_execution()` are also forwarded to `pl.GPUEngine`. + Reserved keys: * `executor_options`: `"runtime"`, `"cluster"`, `"spmd"`, `"ray_context"` -* `engine_kwargs`: `"memory_resource"`, `"executor"` +* `engine_options`: `"memory_resource"`, `"executor"` [dask-distributed]: https://distributed.dask.org/ From ee737a8ad2ccbe23c7bbd92764e17e8a20407fcb Mon Sep 17 00:00:00 2001 From: "Mads R. B. Kristensen" Date: Mon, 16 Mar 2026 14:15:21 +0100 Subject: [PATCH 25/33] docs --- python/cudf_polars/cudf_polars/experimental/rapidsmpf/ray.py | 2 ++ python/cudf_polars/docs/cudf-polars-mp.md | 3 ++- 2 files changed, 4 insertions(+), 1 deletion(-) diff --git a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/ray.py b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/ray.py index ca19f0d6a488..ebb5a197143f 100644 --- a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/ray.py +++ b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/ray.py @@ -118,6 +118,8 @@ def evaluate_pipeline_ray_mode( # identical to the nodes in the deserialized IR tree. Actors fetch the bundle # by reference instead of receiving N copies. query_bundle = ray.put((ir, partition_info, stats, collective_id_map)) + # ray.get() returns results in the same order as the input list of object refs, + # guaranteeing that result[i] corresponds to rank_actors[i] (rank order). result = ray.get( [ rank.evaluate_polars_ir.remote( diff --git a/python/cudf_polars/docs/cudf-polars-mp.md b/python/cudf_polars/docs/cudf-polars-mp.md index 39078f529d13..bdee8f787ebb 100644 --- a/python/cudf_polars/docs/cudf-polars-mp.md +++ b/python/cudf_polars/docs/cudf-polars-mp.md @@ -181,7 +181,8 @@ full = allgather_polars_dataframe( `op_id` is a unique integer that identifies this collective operation across ranks. All ranks must call the same collective with the same `op_id`. -The result is a `pl.DataFrame` containing rows from all ranks, ordered by rank. +The result is guaranteed to be a `pl.DataFrame` containing rows from all ranks in rank order +(rank 0 first, then rank 1, …, rank N-1). ### Passing options From 3a9073260af3da364213d0dfe48f66e5f5727d72 Mon Sep 17 00:00:00 2001 From: "Mads R. B. Kristensen" Date: Mon, 16 Mar 2026 14:24:35 +0100 Subject: [PATCH 26/33] docs --- .../rapidsmpf/collectives/__init__.py | 9 +++-- .../experimental/rapidsmpf/spmd.py | 9 ++++- python/cudf_polars/docs/cudf-polars-mp.md | 37 +++++++++++-------- 3 files changed, 36 insertions(+), 19 deletions(-) diff --git a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/collectives/__init__.py b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/collectives/__init__.py index 5aa14a2b9eb1..b6fc25f335f5 100644 --- a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/collectives/__init__.py +++ b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/collectives/__init__.py @@ -1,9 +1,12 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. # SPDX-License-Identifier: Apache-2.0 """Collective operations for the RapidsMPF streaming runtime.""" from __future__ import annotations -from cudf_polars.experimental.rapidsmpf.collectives.common import ReserveOpIDs +from cudf_polars.experimental.rapidsmpf.collectives.common import ( + ReserveOpIDs, + reserve_op_id, +) -__all__ = ["ReserveOpIDs"] +__all__ = ["ReserveOpIDs", "reserve_op_id"] diff --git a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/spmd.py b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/spmd.py index aa093ff20e6f..f516aa7aad13 100644 --- a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/spmd.py +++ b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/spmd.py @@ -27,6 +27,9 @@ from cudf_polars.containers import DataFrame from cudf_polars.dsl.ir import IRExecutionContext +from cudf_polars.experimental.rapidsmpf.collectives.common import ( + reserve_op_id as reserve_op_id, # noqa: PLC0414 (explicit re-export) +) from cudf_polars.experimental.rapidsmpf.core import generate_network from cudf_polars.experimental.rapidsmpf.utils import empty_table_chunk from cudf_polars.experimental.utils import _concat @@ -160,6 +163,9 @@ def allgather_polars_dataframe( equivalent of a distributed ``collect``: after the call, every rank holds the same complete dataset. + Must be called inside a :func:`spmd_execution` context block while ``comm`` + and ``ctx`` are still alive. + Parameters ---------- comm @@ -170,7 +176,8 @@ def allgather_polars_dataframe( Rank-local DataFrame to contribute. op_id Operation ID for this AllGather collective. Must be identical on every - rank. + rank. Use :func:`reserve_op_id` to obtain a collision-free ID from the + same pool used internally by cudf-polars. Do not pass hardcoded integers. Returns ------- diff --git a/python/cudf_polars/docs/cudf-polars-mp.md b/python/cudf_polars/docs/cudf-polars-mp.md index bdee8f787ebb..2fc370b208e6 100644 --- a/python/cudf_polars/docs/cudf-polars-mp.md +++ b/python/cudf_polars/docs/cudf-polars-mp.md @@ -97,8 +97,9 @@ if `rapidsmpf.bootstrap.is_running_with_rrun()` returns `False`. # launch with: rrun -n 4 python my_script.py import polars as pl from cudf_polars.experimental.rapidsmpf.spmd import ( - spmd_execution, allgather_polars_dataframe, + reserve_op_id, + spmd_execution, ) with spmd_execution() as (comm, ctx, engine): @@ -110,12 +111,13 @@ with spmd_execution() as (comm, ctx, engine): .collect(engine=engine) ) - full = allgather_polars_dataframe( - comm=comm, - ctx=ctx, - local_df=result, - op_id=0, - ) + with reserve_op_id() as op_id: + full = allgather_polars_dataframe( + comm=comm, + ctx=ctx, + local_df=result, + op_id=op_id, + ) ``` The context manager yields: @@ -170,16 +172,21 @@ with spmd_execution() as (comm, ctx, engine): `allgather_polars_dataframe()` to gather all fragments: ```python -full = allgather_polars_dataframe( - comm=comm, - ctx=ctx, - local_df=result, - op_id=0, -) +with spmd_execution() as (comm, ctx, engine): + with reserve_op_id() as op_id: + full = allgather_polars_dataframe( + comm=comm, + ctx=ctx, + local_df=result, + op_id=op_id, + ) ``` -`op_id` is a unique integer that identifies this collective operation across ranks. -All ranks must call the same collective with the same `op_id`. +`op_id` identifies this collective across ranks — all ranks must pass the same value. +Use `reserve_op_id()` (imported from `cudf_polars.experimental.rapidsmpf.collectives.common`) +to obtain a safe ID. It draws from the same pool that cudf-polars uses internally for shuffle +and join collectives, so there is no risk of collision. Do not pass hardcoded integers: they +may silently collide with an ID already reserved by an active collective inside `collect()`. The result is guaranteed to be a `pl.DataFrame` containing rows from all ranks in rank order (rank 0 first, then rank 1, …, rank N-1). From 4b638591778652fcc5cca2da878247ae7bc41f65 Mon Sep 17 00:00:00 2001 From: "Mads R. B. Kristensen" Date: Mon, 16 Mar 2026 14:38:00 +0100 Subject: [PATCH 27/33] move to frontend --- .../cudf_polars/experimental/benchmarks/utils.py | 6 ++++-- .../cudf_polars/experimental/rapidsmpf/core.py | 4 ++-- .../experimental/rapidsmpf/frontend/__init__.py | 7 +++++++ .../experimental/rapidsmpf/{ => frontend}/ray.py | 0 .../experimental/rapidsmpf/{ => frontend}/spmd.py | 0 python/cudf_polars/cudf_polars/utils/config.py | 8 ++++---- python/cudf_polars/docs/cudf-polars-mp.md | 12 ++++++------ .../tests/experimental/rapidsmpf/test_ray.py | 2 +- .../tests/experimental/rapidsmpf/test_spmd.py | 2 +- 9 files changed, 25 insertions(+), 16 deletions(-) create mode 100644 python/cudf_polars/cudf_polars/experimental/rapidsmpf/frontend/__init__.py rename python/cudf_polars/cudf_polars/experimental/rapidsmpf/{ => frontend}/ray.py (100%) rename python/cudf_polars/cudf_polars/experimental/rapidsmpf/{ => frontend}/spmd.py (100%) diff --git a/python/cudf_polars/cudf_polars/experimental/benchmarks/utils.py b/python/cudf_polars/cudf_polars/experimental/benchmarks/utils.py index 0ff9ec60e9b8..b0d82cda4964 100644 --- a/python/cudf_polars/cudf_polars/experimental/benchmarks/utils.py +++ b/python/cudf_polars/cudf_polars/experimental/benchmarks/utils.py @@ -33,7 +33,7 @@ import rmm.statistics -from cudf_polars.experimental.rapidsmpf.spmd import spmd_execution +from cudf_polars.experimental.rapidsmpf.frontend.spmd import spmd_execution # The dtype for count() aggregations depends on the presence # of the polars-runtime-64 package (`polars[rt64]`). @@ -1866,7 +1866,9 @@ def run_polars_spmd( cuda_stream_policy=run_config.stream_policy, ) as (comm, ctx, engine): from cudf_polars.experimental.rapidsmpf.collectives.common import reserve_op_id - from cudf_polars.experimental.rapidsmpf.spmd import allgather_polars_dataframe + from cudf_polars.experimental.rapidsmpf.frontend.spmd import ( + allgather_polars_dataframe, + ) def _allgather_result(df: pl.DataFrame) -> pl.DataFrame: with reserve_op_id() as op_id: diff --git a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/core.py b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/core.py index 7de650f1d6d6..b305fb4f3d4e 100644 --- a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/core.py +++ b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/core.py @@ -126,7 +126,7 @@ def evaluate_logical_plan( collect_metadata=collect_metadata, ) case "spmd": - from cudf_polars.experimental.rapidsmpf.spmd import ( + from cudf_polars.experimental.rapidsmpf.frontend.spmd import ( evaluate_pipeline_spmd_mode, ) @@ -139,7 +139,7 @@ def evaluate_logical_plan( collect_metadata=collect_metadata, ) case "ray": - from cudf_polars.experimental.rapidsmpf.ray import ( + from cudf_polars.experimental.rapidsmpf.frontend.ray import ( evaluate_pipeline_ray_mode, ) diff --git a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/frontend/__init__.py b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/frontend/__init__.py new file mode 100644 index 000000000000..99eed1252217 --- /dev/null +++ b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/frontend/__init__.py @@ -0,0 +1,7 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. +# SPDX-License-Identifier: Apache-2.0 +"""Multi-GPU frontend execution modes for the RapidsMPF streaming engine.""" + +from __future__ import annotations + +__all__: list[str] = [] diff --git a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/ray.py b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/frontend/ray.py similarity index 100% rename from python/cudf_polars/cudf_polars/experimental/rapidsmpf/ray.py rename to python/cudf_polars/cudf_polars/experimental/rapidsmpf/frontend/ray.py diff --git a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/spmd.py b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/frontend/spmd.py similarity index 100% rename from python/cudf_polars/cudf_polars/experimental/rapidsmpf/spmd.py rename to python/cudf_polars/cudf_polars/experimental/rapidsmpf/frontend/spmd.py diff --git a/python/cudf_polars/cudf_polars/utils/config.py b/python/cudf_polars/cudf_polars/utils/config.py index 642e24798464..fbfe4fe6e575 100644 --- a/python/cudf_polars/cudf_polars/utils/config.py +++ b/python/cudf_polars/cudf_polars/utils/config.py @@ -44,7 +44,7 @@ import rmm.mr - from cudf_polars.experimental.rapidsmpf.ray import RankActor + from cudf_polars.experimental.rapidsmpf.frontend.ray import RankActor __all__ = [ @@ -607,7 +607,7 @@ class SPMDContext: :class:`Context`, and :class:`~concurrent.futures.ThreadPoolExecutor` cannot be serialized. In SPMD mode each rank constructs its own ``SPMDContext`` locally inside - :func:`~cudf_polars.experimental.rapidsmpf.spmd.spmd_execution`, so + :func:`~cudf_polars.experimental.rapidsmpf.frontend.spmd.spmd_execution`, so pickling is never required. Do not use this class with Dask or any other framework that serializes executor configuration across process boundaries. @@ -635,13 +635,13 @@ class RayContext: This dataclass holds Ray actor handles, which are only valid within the Ray session that created them. It is stripped from ``config_options`` before pickling for remote actor calls in - :func:`~cudf_polars.experimental.rapidsmpf.ray.evaluate_pipeline_ray_mode`. + :func:`~cudf_polars.experimental.rapidsmpf.frontend.ray.evaluate_pipeline_ray_mode`. Do not persist or transfer this object across Ray sessions. Parameters ---------- rank_actors - List of :class:`~cudf_polars.experimental.rapidsmpf.ray.RankActor` + List of :class:`~cudf_polars.experimental.rapidsmpf.frontend.ray.RankActor` handles, one per GPU in the cluster. """ diff --git a/python/cudf_polars/docs/cudf-polars-mp.md b/python/cudf_polars/docs/cudf-polars-mp.md index 2fc370b208e6..dd92eb0b3908 100644 --- a/python/cudf_polars/docs/cudf-polars-mp.md +++ b/python/cudf_polars/docs/cudf-polars-mp.md @@ -11,7 +11,7 @@ The `cluster` option selects the execution model: | `cluster` | Description | Status | | --------------- | ---------------------------------------------------- | ----------------- | -| `"single"` | Single-GPU, in-process execution | Stable | +| `"single"` | Single-GPU, in-process execution | Stable (legacy) | | `"distributed"` | Multi-GPU via [Dask Distributed][dask-distributed] | Stable (legacy) | | `"spmd"` | Multi-GPU via [SPMD][spmd-wiki] launched with `rrun` | Preview (new API) | | `"ray"` | Multi-GPU via [Ray][ray-docs] actors | Preview (new API) | @@ -82,7 +82,7 @@ every rank, call `allgather_polars_dataframe()`. ### Running in SPMD mode `spmd_execution()` is the primary entry point for SPMD execution. It is a context -manager imported from `cudf_polars.experimental.rapidsmpf.spmd`. On entry it: +manager imported from `cudf_polars.experimental.rapidsmpf.frontend.spmd`. On entry it: 1. Bootstraps a UCXX communicator connecting all ranks. 2. Creates a RapidsMPF streaming `Context` that owns GPU memory and a CUDA stream pool. @@ -96,7 +96,7 @@ if `rapidsmpf.bootstrap.is_running_with_rrun()` returns `False`. ```python # launch with: rrun -n 4 python my_script.py import polars as pl -from cudf_polars.experimental.rapidsmpf.spmd import ( +from cudf_polars.experimental.rapidsmpf.frontend.spmd import ( allgather_polars_dataframe, reserve_op_id, spmd_execution, @@ -267,7 +267,7 @@ executes the same query. ### Running in Ray mode -`ray_execution()` is imported from `cudf_polars.experimental.rapidsmpf.ray`. It: +`ray_execution()` is imported from `cudf_polars.experimental.rapidsmpf.frontend.ray`. It: 1. Calls `ray.init()` if Ray is not already running 2. Creates one `RankActor` per GPU @@ -279,7 +279,7 @@ Actors are shut down on exit. If the context started Ray, it also calls ```python import polars as pl -from cudf_polars.experimental.rapidsmpf.ray import ray_execution +from cudf_polars.experimental.rapidsmpf.frontend.ray import ray_execution with ray_execution() as (ray_client, engine): result = ( @@ -306,7 +306,7 @@ does not call `ray.shutdown()` on exit. ```python import ray import polars as pl -from cudf_polars.experimental.rapidsmpf.ray import ray_execution +from cudf_polars.experimental.rapidsmpf.frontend.ray import ray_execution ray.init(address="auto") diff --git a/python/cudf_polars/tests/experimental/rapidsmpf/test_ray.py b/python/cudf_polars/tests/experimental/rapidsmpf/test_ray.py index 7477d2af7477..4b0fc3fde3f5 100644 --- a/python/cudf_polars/tests/experimental/rapidsmpf/test_ray.py +++ b/python/cudf_polars/tests/experimental/rapidsmpf/test_ray.py @@ -15,7 +15,7 @@ ray = pytest.importorskip("ray") -from cudf_polars.experimental.rapidsmpf.ray import ( # noqa: E402 +from cudf_polars.experimental.rapidsmpf.frontend.ray import ( # noqa: E402 RayClient, ray_execution, ) diff --git a/python/cudf_polars/tests/experimental/rapidsmpf/test_spmd.py b/python/cudf_polars/tests/experimental/rapidsmpf/test_spmd.py index d26f4ca70d69..281bc652ce90 100644 --- a/python/cudf_polars/tests/experimental/rapidsmpf/test_spmd.py +++ b/python/cudf_polars/tests/experimental/rapidsmpf/test_spmd.py @@ -14,7 +14,7 @@ import rmm.mr from cudf_polars.experimental.rapidsmpf.collectives.common import reserve_op_id -from cudf_polars.experimental.rapidsmpf.spmd import ( +from cudf_polars.experimental.rapidsmpf.frontend.spmd import ( allgather_polars_dataframe, spmd_execution, ) From c06f0f93462c947c17cce00d2a709809334a3db8 Mon Sep 17 00:00:00 2001 From: "Mads R. B. Kristensen" Date: Mon, 16 Mar 2026 14:43:12 +0100 Subject: [PATCH 28/33] update test --- .../cudf_polars/tests/experimental/rapidsmpf/test_ray.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/python/cudf_polars/tests/experimental/rapidsmpf/test_ray.py b/python/cudf_polars/tests/experimental/rapidsmpf/test_ray.py index 4b0fc3fde3f5..00a1bc058e3a 100644 --- a/python/cudf_polars/tests/experimental/rapidsmpf/test_ray.py +++ b/python/cudf_polars/tests/experimental/rapidsmpf/test_ray.py @@ -91,13 +91,15 @@ def test_ray_client_shutdown_idempotent() -> None: def test_ray_client_post_shutdown_state() -> None: - """After shutdown, rank_actors is empty, nranks is 0, and engine raises.""" + """After shutdown, rank_actors, nranks, and engine all raise RuntimeError.""" mock_engine = MagicMock(spec=pl.GPUEngine) mock_engine.config = {"executor_options": {"ray_context": RayContext([])}} client = RayClient(mock_engine, owns_ray=False) client.shutdown() - assert client.rank_actors == [] - assert client.nranks == 0 + with pytest.raises(RuntimeError, match="shutdown"): + _ = client.rank_actors + with pytest.raises(RuntimeError, match="shutdown"): + _ = client.nranks with pytest.raises(RuntimeError, match="shutdown"): _ = client.engine From 94bc8cbc95d47491da30fb3c408cf3133f3a3e4f Mon Sep 17 00:00:00 2001 From: "Mads R. B. Kristensen" Date: Mon, 16 Mar 2026 15:32:14 +0100 Subject: [PATCH 29/33] remove test prefix --- .../tests/experimental/rapidsmpf/test_ray.py | 20 +++++++++---------- .../tests/experimental/rapidsmpf/test_spmd.py | 18 ++++++++--------- 2 files changed, 19 insertions(+), 19 deletions(-) diff --git a/python/cudf_polars/tests/experimental/rapidsmpf/test_ray.py b/python/cudf_polars/tests/experimental/rapidsmpf/test_ray.py index 00a1bc058e3a..4ac15a6bd5f0 100644 --- a/python/cudf_polars/tests/experimental/rapidsmpf/test_ray.py +++ b/python/cudf_polars/tests/experimental/rapidsmpf/test_ray.py @@ -66,14 +66,14 @@ def engine(_ray_env: tuple[RayClient, pl.GPUEngine]) -> pl.GPUEngine: # --------------------------------------------------------------------------- -def test_ray_execution_reserved_executor_keys() -> None: +def test_reserved_executor_keys() -> None: """executor_options rejects reserved keys.""" for key in ("runtime", "cluster", "spmd", "ray_context"): with pytest.raises(TypeError, match="reserved"): ray_execution(executor_options={key: "anything"}) -def test_ray_execution_reserved_engine_options_keys() -> None: +def test_reserved_engine_options_keys() -> None: """engine_options rejects keys that are set explicitly by ray_execution.""" for key in ("memory_resource", "executor"): kwargs: dict[str, Any] = {key: "anything"} @@ -104,7 +104,7 @@ def test_ray_client_post_shutdown_state() -> None: _ = client.engine -def test_ray_execution_raises_inside_rrun() -> None: +def test_raises_inside_rrun() -> None: """ray_execution() must not be called from within an rrun cluster.""" with ( patch( @@ -121,7 +121,7 @@ def test_ray_execution_raises_inside_rrun() -> None: # --------------------------------------------------------------------------- -def test_ray_execution_yields_client_and_engine( +def test_yields_client_and_engine( ray_client: RayClient, engine: pl.GPUEngine, ) -> None: @@ -131,7 +131,7 @@ def test_ray_execution_yields_client_and_engine( assert ray_client.nranks >= 1 -def test_ray_execution_executor_options_forwarded( +def test_executor_options_forwarded( ray_client: RayClient, engine: pl.GPUEngine, ) -> None: @@ -158,7 +158,7 @@ def test_gather_cluster_info(ray_client: RayClient) -> None: assert len({info["pid"] for info in infos}) == ray_client.nranks -def test_ray_execution_scan(engine: pl.GPUEngine) -> None: +def test_scan(engine: pl.GPUEngine) -> None: """Input rows are partitioned across actors; total output equals input.""" lf = pl.LazyFrame({"a": [1, 2, 3]}) result = lf.collect(engine=engine) @@ -166,7 +166,7 @@ def test_ray_execution_scan(engine: pl.GPUEngine) -> None: assert sorted(result["a"].to_list()) == [1, 2, 3] -def test_ray_execution_filter(engine: pl.GPUEngine) -> None: +def test_filter(engine: pl.GPUEngine) -> None: """Filter is applied correctly across all actors.""" lf = pl.LazyFrame({"a": [1, 2, 3, 4, 5]}) result = lf.filter(pl.col("a") > 3).collect(engine=engine) @@ -174,7 +174,7 @@ def test_ray_execution_filter(engine: pl.GPUEngine) -> None: assert sorted(result["a"].to_list()) == [4, 5] -def test_ray_execution_group_by(ray_client: RayClient, engine: pl.GPUEngine) -> None: +def test_group_by(ray_client: RayClient, engine: pl.GPUEngine) -> None: """Group-by produces the correct aggregation across all ranks.""" # max_rows_per_partition=10 (set on the session fixture) gives each rank # exactly 5 partitions, so the multi-partition path is always exercised. @@ -197,7 +197,7 @@ def test_ray_execution_group_by(ray_client: RayClient, engine: pl.GPUEngine) -> assert result["val"].to_list() == expected["val"].to_list() -def test_ray_execution_join(ray_client: RayClient, engine: pl.GPUEngine) -> None: +def test_join(ray_client: RayClient, engine: pl.GPUEngine) -> None: """Hash join between two tables produces the correct result across all ranks.""" # max_rows_per_partition=10 (set on the session fixture) gives each rank # exactly 5 partitions, so the multi-partition path is always exercised. @@ -212,7 +212,7 @@ def test_ray_execution_join(ray_client: RayClient, engine: pl.GPUEngine) -> None assert result["val_right"].to_list() == [x * 2 for x in range(n)] -def test_ray_execution_empty_dataframe(engine: pl.GPUEngine) -> None: +def test_empty_dataframe(engine: pl.GPUEngine) -> None: """An empty LazyFrame produces an empty result with the correct schema.""" lf = pl.LazyFrame( {"a": pl.Series([], dtype=pl.Int32), "b": pl.Series([], dtype=pl.Float64)} diff --git a/python/cudf_polars/tests/experimental/rapidsmpf/test_spmd.py b/python/cudf_polars/tests/experimental/rapidsmpf/test_spmd.py index 281bc652ce90..4b27821a5483 100644 --- a/python/cudf_polars/tests/experimental/rapidsmpf/test_spmd.py +++ b/python/cudf_polars/tests/experimental/rapidsmpf/test_spmd.py @@ -25,7 +25,7 @@ ) -def test_spmd_execution_yields_context_and_engine() -> None: +def test_yields_context_and_engine() -> None: """spmd_execution yields a (Communicator, Context, GPUEngine) triple.""" with spmd_execution() as (comm, ctx, engine): assert comm is not None @@ -33,7 +33,7 @@ def test_spmd_execution_yields_context_and_engine() -> None: assert isinstance(engine, pl.GPUEngine) -def test_spmd_execution_reserved_keys() -> None: +def test_reserved_keys() -> None: """executor_options rejects reserved keys.""" for key in ("runtime", "cluster", "spmd"): with ( @@ -43,7 +43,7 @@ def test_spmd_execution_reserved_keys() -> None: pass -def test_spmd_execution_engine_kwargs_reserved_keys() -> None: +def test_engine_kwargs_reserved_keys() -> None: """engine_kwargs rejects keys that are set explicitly by spmd_execution.""" for key in ("memory_resource", "executor"): kwargs: dict[str, Any] = {key: "anything"} @@ -54,13 +54,13 @@ def test_spmd_execution_engine_kwargs_reserved_keys() -> None: pass -def test_spmd_execution_engine_kwargs_parquet_options() -> None: +def test_engine_kwargs_parquet_options() -> None: """engine_kwargs forwards parquet_options to GPUEngine without error.""" with spmd_execution(parquet_options={}) as (comm, ctx, engine): assert isinstance(engine, pl.GPUEngine) -def test_spmd_execution_custom_mr() -> None: +def test_custom_mr() -> None: """spmd_execution accepts a custom memory resource.""" mr = rmm.mr.CudaMemoryResource() with spmd_execution(mr=mr) as (comm, ctx, engine): @@ -68,7 +68,7 @@ def test_spmd_execution_custom_mr() -> None: assert result.shape == (3, 1) -def test_spmd_execution_scan() -> None: +def test_scan() -> None: """Each rank scans its own single-row LazyFrame and gets that row back.""" with spmd_execution() as (comm, ctx, engine): rank = comm.rank @@ -79,7 +79,7 @@ def test_spmd_execution_scan() -> None: assert result["b"].to_list() == [rank * 10] -def test_spmd_collect_then_lazy_equivalent() -> None: +def test_collect_then_lazy_equivalent() -> None: """collect().lazy() preserves SPMD semantics: an intermediate materialize is a no-op. In SPMD mode a DataFrame is always rank-local. When it is wrapped back @@ -101,7 +101,7 @@ def test_spmd_collect_then_lazy_equivalent() -> None: assert one_step.sort("a").equals(two_step.sort("a")) -def test_spmd_execution_group_by() -> None: +def test_group_by() -> None: """Group-by on rank-local data, then allgather to verify the global result.""" with spmd_execution() as (comm, ctx, engine): rank = comm.rank @@ -132,7 +132,7 @@ def test_allgather_polars_dataframe() -> None: assert result["val"].to_list() == [r * 2 for r in range(nranks)] -def test_spmd_execution_max_workers() -> None: +def test_max_workers() -> None: """executor_options forwards rapidsmpf_py_executor_max_workers to the thread pool.""" with spmd_execution(executor_options={"rapidsmpf_py_executor_max_workers": 2}) as ( comm, From 8feec15a4468ae7368c1dc85ae8ea8f33566a12b Mon Sep 17 00:00:00 2001 From: "Mads R. B. Kristensen" Date: Mon, 16 Mar 2026 15:44:46 +0100 Subject: [PATCH 30/33] reorg docs --- python/cudf_polars/docs/cudf-polars-mp.md | 295 +++++++++++----------- 1 file changed, 147 insertions(+), 148 deletions(-) diff --git a/python/cudf_polars/docs/cudf-polars-mp.md b/python/cudf_polars/docs/cudf-polars-mp.md index dd92eb0b3908..0fce3a219bab 100644 --- a/python/cudf_polars/docs/cudf-polars-mp.md +++ b/python/cudf_polars/docs/cudf-polars-mp.md @@ -13,20 +13,163 @@ The `cluster` option selects the execution model: | --------------- | ---------------------------------------------------- | ----------------- | | `"single"` | Single-GPU, in-process execution | Stable (legacy) | | `"distributed"` | Multi-GPU via [Dask Distributed][dask-distributed] | Stable (legacy) | -| `"spmd"` | Multi-GPU via [SPMD][spmd-wiki] launched with `rrun` | Preview (new API) | | `"ray"` | Multi-GPU via [Ray][ray-docs] actors | Preview (new API) | +| `"spmd"` | Multi-GPU via [SPMD][spmd-wiki] launched with `rrun` | Preview (new API) | Two preview execution modes are available: -* **SPMD mode** — each GPU runs the same script as an independent process, - launched with `rrun`. * **Ray mode** — a single-client model where a driver program coordinates GPU workers implemented as Ray actors. +* **SPMD mode** — each GPU runs the same script as an independent process, + launched with `rrun`. This document describes these two execution modes. -* [SPMD execution mode](#spmd-execution-mode) * [Ray execution mode](#ray-execution-mode) +* [SPMD execution mode](#spmd-execution-mode) + +--- + +## Ray execution mode + +Ray mode uses a single client process that drives execution across multiple ranks. +Each rank corresponds to one GPU worker and participates in collective operations +through a shared UCXX communicator. + +In the Ray implementation each rank is implemented as a [**Ray actor**][ray-actors], +with one actor created per available GPU. + +Conceptually the system looks like this: + +``` + ┌──────────────────────────────┐ + │ User script │ + │ (single client process) │ + │ LazyFrame.collect(engine=…) │ + └──────────────┬───────────────┘ + │ IR dispatched to all actors + ┌────────────────|─────────────────┐ + ↓ ↓ ↓ + ┌─────────────┐ ┌─────────────┐ ┌─────────────┐ + │ RankActor │ │ RankActor │ │ RankActor │ + │ rank 0 │ │ rank 1 │ │ rank N-1 │ + │ run IR │ │ run IR │ │ run IR │ + └──────┬──────┘ └──────┬──────┘ └──────┬──────┘ + ↓ ↓ ↓ +┌────────────────────────────────────────────────────────────────┐ +│ RapidsMPF streaming engine │ +│ shuffle / all-gather · UCXX communicator · RMM GPU memory │ +└────────────────────────────────────────────────────────────────┘ + ↑ ↑ ↑ + GPU 0 GPU 1 GPU N-1 +``` + +The client broadcasts the query plan to all ranks. The ranks execute the pipeline +collectively through UCXX, and their outputs are streamed back and concatenated on +the client process. + +The driver script runs as a normal Python program with no `rrun` launcher. Query +symmetry is handled automatically: the client serializes the complete query plan and +broadcasts it to all actors, so every rank always executes the same query. + +### Prerequisites + +* Ray (`ray`) installed +* RapidsMPF and UCXX available on all GPU nodes + +### Running in Ray mode + +`ray_execution()` is imported from `cudf_polars.experimental.rapidsmpf.frontend.ray`. It: + +1. Calls `ray.init()` if Ray is not already running +2. Creates one `RankActor` per GPU +3. Bootstraps a UCXX communicator across the actors +4. Yields a `pl.GPUEngine` and a `RayClient` + +Actors are shut down on exit. If the context started Ray, it also calls +`ray.shutdown()`. + +```python +import polars as pl +from cudf_polars.experimental.rapidsmpf.frontend.ray import ray_execution + +with ray_execution() as (ray_client, engine): + result = ( + pl.scan_parquet("/data/dataset/*.parquet") + .filter(pl.col("amount") > 100) + .group_by("customer_id") + .agg(pl.col("amount").sum()) + .collect(engine=engine) + ) + +print(result) +``` + +The context manager yields: + +* `ray_client` — cluster diagnostics and utilities +* `engine` — `pl.GPUEngine` configured for Ray execution + +### Ray lifecycle + +If Ray is already initialized, `ray_execution()` attaches to the existing cluster and +does not call `ray.shutdown()` on exit. + +```python +import ray +import polars as pl +from cudf_polars.experimental.rapidsmpf.frontend.ray import ray_execution + +ray.init(address="auto") + +try: + with ray_execution() as (ray_client, engine): + result = pl.scan_parquet(...).collect(engine=engine) +finally: + ray.shutdown() +``` + +`ray_execution()` raises `RuntimeError` if called inside an `rrun` cluster or if no +GPUs are available. + +### Cluster diagnostics + +`RayClient.gather_cluster_info()` returns placement information for all rank actors: + +```python +with ray_execution() as (ray_client, engine): + print(f"cluster has {ray_client.nranks} ranks") + for i, info in enumerate(ray_client.gather_cluster_info()): + print( + f"rank {i}: hostname={info['hostname']}, pid={info['pid']}, " + f"CUDA_VISIBLE_DEVICES={info['cuda_visible_devices']}" + ) +``` + +Each entry includes `pid`, `hostname`, `cuda_visible_devices`, and `node_id`. + +### Passing options + +`executor_options`, `engine_options`, and `ray_init_options` accept pass-through +dictionaries: + +```python +with ray_execution( + executor_options={"max_rows_per_partition": 500_000}, + engine_options={"raise_on_fail": True}, + ray_init_options={"num_cpus": 4}, +) as (ray_client, engine): + ... +``` + +`executor_options` is forwarded directly to `pl.GPUEngine` as its `executor_options` +argument; user-supplied keys are merged with reserved entries set by `ray_execution()`. +Any additional keyword arguments to `ray_execution()` are also forwarded to `pl.GPUEngine`. + +Reserved keys: + +* `executor_options`: `"runtime"`, `"cluster"`, `"spmd"`, `"ray_context"` +* `engine_options`: `"memory_resource"`, `"executor"` --- @@ -215,150 +358,6 @@ Reserved keys: * `executor_options`: `"runtime"`, `"cluster"`, `"spmd"` * `engine_options`: `"memory_resource"`, `"executor"` ---- - -## Ray execution mode - -Ray mode uses a single client process that drives execution across multiple ranks. -Each rank corresponds to one GPU worker and participates in collective operations -through a shared UCXX communicator. - -In the Ray implementation each rank is implemented as a [**Ray actor**][ray-actors], -with one actor created per available GPU. - -Conceptually the system looks like this: - -``` - ┌──────────────────────────────┐ - │ User script │ - │ (single client process) │ - │ LazyFrame.collect(engine=…) │ - └──────────────┬───────────────┘ - │ IR dispatched to all actors - ┌────────────────|─────────────────┐ - ↓ ↓ ↓ - ┌─────────────┐ ┌─────────────┐ ┌─────────────┐ - │ RankActor │ │ RankActor │ │ RankActor │ - │ rank 0 │ │ rank 1 │ │ rank N-1 │ - │ run IR │ │ run IR │ │ run IR │ - └──────┬──────┘ └──────┬──────┘ └──────┬──────┘ - ↓ ↓ ↓ -┌────────────────────────────────────────────────────────────────┐ -│ RapidsMPF streaming engine │ -│ shuffle / all-gather · UCXX communicator · RMM GPU memory │ -└────────────────────────────────────────────────────────────────┘ - ↑ ↑ ↑ - GPU 0 GPU 1 GPU N-1 -``` - -The client broadcasts the query plan to all ranks. The ranks execute the pipeline -collectively through UCXX, and their outputs are streamed back and concatenated on -the client process. - -Unlike SPMD mode, the driver script runs as a normal Python program with no -`rrun` launcher. Query symmetry is handled automatically: the client serializes -the complete query plan and broadcasts it to all actors, so every rank always -executes the same query. - -### Prerequisites - -* Ray (`ray`) installed -* RapidsMPF and UCXX available on all GPU nodes - -### Running in Ray mode - -`ray_execution()` is imported from `cudf_polars.experimental.rapidsmpf.frontend.ray`. It: - -1. Calls `ray.init()` if Ray is not already running -2. Creates one `RankActor` per GPU -3. Bootstraps a UCXX communicator across the actors -4. Yields a `pl.GPUEngine` and a `RayClient` - -Actors are shut down on exit. If the context started Ray, it also calls -`ray.shutdown()`. - -```python -import polars as pl -from cudf_polars.experimental.rapidsmpf.frontend.ray import ray_execution - -with ray_execution() as (ray_client, engine): - result = ( - pl.scan_parquet("/data/dataset/*.parquet") - .filter(pl.col("amount") > 100) - .group_by("customer_id") - .agg(pl.col("amount").sum()) - .collect(engine=engine) - ) - -print(result) -``` - -The context manager yields: - -* `ray_client` — cluster diagnostics and utilities -* `engine` — `pl.GPUEngine` configured for Ray execution - -### Ray lifecycle - -If Ray is already initialized, `ray_execution()` attaches to the existing cluster and -does not call `ray.shutdown()` on exit. - -```python -import ray -import polars as pl -from cudf_polars.experimental.rapidsmpf.frontend.ray import ray_execution - -ray.init(address="auto") - -try: - with ray_execution() as (ray_client, engine): - result = pl.scan_parquet(...).collect(engine=engine) -finally: - ray.shutdown() -``` - -`ray_execution()` raises `RuntimeError` if called inside an `rrun` cluster or if no -GPUs are available. - -### Cluster diagnostics - -`RayClient.gather_cluster_info()` returns placement information for all rank actors: - -```python -with ray_execution() as (ray_client, engine): - print(f"cluster has {ray_client.nranks} ranks") - for i, info in enumerate(ray_client.gather_cluster_info()): - print( - f"rank {i}: hostname={info['hostname']}, pid={info['pid']}, " - f"CUDA_VISIBLE_DEVICES={info['cuda_visible_devices']}" - ) -``` - -Each entry includes `pid`, `hostname`, `cuda_visible_devices`, and `node_id`. - -### Passing options - -`executor_options`, `engine_options`, and `ray_init_options` accept pass-through -dictionaries: - -```python -with ray_execution( - executor_options={"max_rows_per_partition": 500_000}, - engine_options={"raise_on_fail": True}, - ray_init_options={"num_cpus": 4}, -) as (ray_client, engine): - ... -``` - -`executor_options` is forwarded directly to `pl.GPUEngine` as its `executor_options` -argument; user-supplied keys are merged with reserved entries set by `ray_execution()`. -Any additional keyword arguments to `ray_execution()` are also forwarded to `pl.GPUEngine`. - -Reserved keys: - -* `executor_options`: `"runtime"`, `"cluster"`, `"spmd"`, `"ray_context"` -* `engine_options`: `"memory_resource"`, `"executor"` - [dask-distributed]: https://distributed.dask.org/ [spmd-wiki]: https://en.wikipedia.org/wiki/Single_program,_multiple_data From e1ac68d6d78dbb787c9d513a601775dce9281b05 Mon Sep 17 00:00:00 2001 From: "Mads R. B. Kristensen" Date: Mon, 16 Mar 2026 19:05:28 +0100 Subject: [PATCH 31/33] cleanup --- .../cudf_polars/experimental/rapidsmpf/frontend/spmd.py | 8 +++----- 1 file changed, 3 insertions(+), 5 deletions(-) diff --git a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/frontend/spmd.py b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/frontend/spmd.py index f516aa7aad13..20ca4d564257 100644 --- a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/frontend/spmd.py +++ b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/frontend/spmd.py @@ -27,9 +27,6 @@ from cudf_polars.containers import DataFrame from cudf_polars.dsl.ir import IRExecutionContext -from cudf_polars.experimental.rapidsmpf.collectives.common import ( - reserve_op_id as reserve_op_id, # noqa: PLC0414 (explicit re-export) -) from cudf_polars.experimental.rapidsmpf.core import generate_network from cudf_polars.experimental.rapidsmpf.utils import empty_table_chunk from cudf_polars.experimental.utils import _concat @@ -176,8 +173,9 @@ def allgather_polars_dataframe( Rank-local DataFrame to contribute. op_id Operation ID for this AllGather collective. Must be identical on every - rank. Use :func:`reserve_op_id` to obtain a collision-free ID from the - same pool used internally by cudf-polars. Do not pass hardcoded integers. + rank. For example, use :func:`reserve_op_id` to obtain a collision-free + ID from the same pool used internally by cudf-polars. Avoid passing + hardcoded integers. Returns ------- From 96619cd6a01310cb738cbc8e041b9efa2c27b24b Mon Sep 17 00:00:00 2001 From: "Mads R. B. Kristensen" Date: Mon, 16 Mar 2026 19:11:59 +0100 Subject: [PATCH 32/33] cleanup --- .../experimental/rapidsmpf/frontend/ray.py | 19 ++++++++++--------- 1 file changed, 10 insertions(+), 9 deletions(-) diff --git a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/frontend/ray.py b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/frontend/ray.py index ebb5a197143f..80138a5473a6 100644 --- a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/frontend/ray.py +++ b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/frontend/ray.py @@ -161,8 +161,9 @@ class RankActor: rapidsmpf_options_as_bytes Serialized RapidsMPF options produced by :meth:`rapidsmpf.config.Options.serialize`. - executor_options - Additional executor options forwarded from the client. + py_executor_max_workers + Maximum number of threads for the actor's Python thread-pool executor. + ``None`` lets :class:`~concurrent.futures.ThreadPoolExecutor` choose. """ def __init__( @@ -170,7 +171,7 @@ def __init__( *, nranks: int, rapidsmpf_options_as_bytes: bytes, - executor_options: dict[str, object], + py_executor_max_workers: int, ) -> None: self._mr = RmmResourceAdaptor(rmm.mr.CudaAsyncMemoryResource()) self._rapidsmpf_options: Options = Options.deserialize( @@ -181,10 +182,7 @@ def __init__( ) self._nranks: int = nranks self._py_executor = ThreadPoolExecutor( - max_workers=cast( - int | None, - executor_options.get("rapidsmpf_py_executor_max_workers"), - ), + max_workers=py_executor_max_workers, thread_name_prefix="ray-executor", ) self._comm: Communicator | None = None @@ -258,7 +256,7 @@ def shutdown(self) -> None: self._mr = None ray.actor.exit_actor() - def get_info(self) -> dict: + def get_info(self) -> dict[str, Any]: """ Return diagnostic information about actor placement. @@ -605,8 +603,11 @@ def ray_execution( rank_actors: list[ActorHandle[RankActor]] = [ RankActor.remote( # type: ignore[attr-defined] nranks=free_gpus, - executor_options=executor_options, rapidsmpf_options_as_bytes=rapidsmpf_options_as_bytes, + py_executor_max_workers=cast( + int, + executor_options.get("rapidsmpf_py_executor_max_workers", 1), + ), ) for _ in range(free_gpus) ] From 5c73829d6e3bd803b9902d16d30f90b5774bbc01 Mon Sep 17 00:00:00 2001 From: "Mads R. B. Kristensen" Date: Mon, 16 Mar 2026 19:29:51 +0100 Subject: [PATCH 33/33] docs --- python/cudf_polars/docs/cudf-polars-mp.md | 52 ++++++++++++++++++++--- 1 file changed, 45 insertions(+), 7 deletions(-) diff --git a/python/cudf_polars/docs/cudf-polars-mp.md b/python/cudf_polars/docs/cudf-polars-mp.md index 0fce3a219bab..ad7188f184ee 100644 --- a/python/cudf_polars/docs/cudf-polars-mp.md +++ b/python/cudf_polars/docs/cudf-polars-mp.md @@ -150,22 +150,37 @@ Each entry includes `pid`, `hostname`, `cuda_visible_devices`, and `node_id`. ### Passing options -`executor_options`, `engine_options`, and `ray_init_options` accept pass-through -dictionaries: +`rapidsmpf_options`, `executor_options`, `engine_options`, and `ray_init_options` accept +pass-through dictionaries: ```python +from rapidsmpf.integrations.cudf_polars import Options + with ray_execution( - executor_options={"max_rows_per_partition": 500_000}, + rapidsmpf_options=Options(num_streaming_threads=8), + executor_options={ + "max_rows_per_partition": 500_000, + "rapidsmpf_py_executor_max_workers": 2, + }, engine_options={"raise_on_fail": True}, ray_init_options={"num_cpus": 4}, ) as (ray_client, engine): ... ``` +`rapidsmpf_options` is an `Options` object passed to the RapidsMPF `Context` on each +worker. If not provided, `ray_execution()` constructs a default `Options` with +`num_streaming_threads=4`. + `executor_options` is forwarded directly to `pl.GPUEngine` as its `executor_options` argument; user-supplied keys are merged with reserved entries set by `ray_execution()`. Any additional keyword arguments to `ray_execution()` are also forwarded to `pl.GPUEngine`. +Notable `executor_options` keys: + +* `"rapidsmpf_py_executor_max_workers"` (default: `1`) — number of threads in the Python + `ThreadPoolExecutor` that drives the RapidsMPF actor network on each worker. + Reserved keys: * `executor_options`: `"runtime"`, `"cluster"`, `"spmd"`, `"ray_context"` @@ -239,9 +254,9 @@ if `rapidsmpf.bootstrap.is_running_with_rrun()` returns `False`. ```python # launch with: rrun -n 4 python my_script.py import polars as pl +from cudf_polars.experimental.rapidsmpf.collectives.common import reserve_op_id from cudf_polars.experimental.rapidsmpf.frontend.spmd import ( allgather_polars_dataframe, - reserve_op_id, spmd_execution, ) @@ -315,6 +330,12 @@ with spmd_execution() as (comm, ctx, engine): `allgather_polars_dataframe()` to gather all fragments: ```python +from cudf_polars.experimental.rapidsmpf.collectives.common import reserve_op_id +from cudf_polars.experimental.rapidsmpf.frontend.spmd import ( + allgather_polars_dataframe, + spmd_execution, +) + with spmd_execution() as (comm, ctx, engine): with reserve_op_id() as op_id: full = allgather_polars_dataframe( @@ -326,8 +347,7 @@ with spmd_execution() as (comm, ctx, engine): ``` `op_id` identifies this collective across ranks — all ranks must pass the same value. -Use `reserve_op_id()` (imported from `cudf_polars.experimental.rapidsmpf.collectives.common`) -to obtain a safe ID. It draws from the same pool that cudf-polars uses internally for shuffle +Use `reserve_op_id()` to obtain a safe ID. It draws from the same pool that cudf-polars uses internally for shuffle and join collectives, so there is no risk of collision. Do not pass hardcoded integers: they may silently collide with an ID already reserved by an active collective inside `collect()`. @@ -336,23 +356,41 @@ The result is guaranteed to be a `pl.DataFrame` containing rows from all ranks i ### Passing options -`executor_options` and `engine_kwargs` accept pass-through dictionaries: +`mr`, `rapidsmpf_options`, `executor_options`, and `engine_kwargs` accept pass-through +arguments: ```python +import rmm +from rapidsmpf.integrations.cudf_polars import Options + with spmd_execution( + mr=rmm.mr.PoolMemoryResource(rmm.mr.CudaMemoryResource()), + rapidsmpf_options=Options(num_streaming_threads=8), executor_options={ "max_rows_per_partition": 500_000, "rapidsmpf_spill": True, + "rapidsmpf_py_executor_max_workers": 2, }, parquet_options={"use_rapidsmpf_native": True}, ) as (comm, ctx, engine): ... ``` +`mr` is an `rmm.mr.DeviceMemoryResource` used as the GPU memory resource for the +RapidsMPF `Context`. Defaults to `None` (uses the current device resource). + +`rapidsmpf_options` is an `Options` object passed to the RapidsMPF `Context`. Defaults +to `None` (uses RapidsMPF defaults). + `executor_options` is forwarded directly to `pl.GPUEngine` as its `executor_options` argument; user-supplied keys are merged with reserved entries set by `spmd_execution()`. Any additional keyword arguments to `spmd_execution()` are also forwarded to `pl.GPUEngine`. +Notable `executor_options` keys: + +* `"rapidsmpf_py_executor_max_workers"` (default: `1`) — number of threads in the Python + `ThreadPoolExecutor` that drives the RapidsMPF actor network. + Reserved keys: * `executor_options`: `"runtime"`, `"cluster"`, `"spmd"`