diff --git a/python/cudf_polars/cudf_polars/experimental/benchmarks/utils.py b/python/cudf_polars/cudf_polars/experimental/benchmarks/utils.py index 0ff9ec60e9b..b0d82cda496 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/collectives/__init__.py b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/collectives/__init__.py index 5aa14a2b9eb..b6fc25f335f 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/core.py b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/core.py index d5a9316e3b9..b305fb4f3d4 100644 --- a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/core.py +++ b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/core.py @@ -108,47 +108,62 @@ 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.frontend.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 "ray": + from cudf_polars.experimental.rapidsmpf.frontend.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( + 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 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 00000000000..99eed125221 --- /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/frontend/ray.py b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/frontend/ray.py new file mode 100644 index 00000000000..80138a5473a --- /dev/null +++ b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/frontend/ray.py @@ -0,0 +1,634 @@ +# 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 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 +from cudf_polars.utils.config import RayContext + +if TYPE_CHECKING: + from collections.abc import MutableMapping + + 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 + 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 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 + 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_context`` is ``None``. + """ + if config_options.executor.runtime != "rapidsmpf": + raise RuntimeError("Runtime must be rapidsmpf") + 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_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_context=None), + ) + + # 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)) + # 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( + query_bundle, + actor_config_options, + 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, communicator, streaming context, + etc. Collectively, the actors form an 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`. + 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__( + self, + *, + nranks: int, + rapidsmpf_options_as_bytes: bytes, + py_executor_max_workers: int, + ) -> 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=py_executor_max_workers, + 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 and returns the + serialized root address, which must be passed to :meth:`setup_worker` + on all actors to complete communicator setup. + + 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 + ) + # 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: + """ + Release actor-owned resources and exit the process. + + 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() + + def get_info(self) -> dict[str, Any]: + """ + 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, + query_bundle: tuple[ + IR, + MutableMapping[IR, PartitionInfo], + StatsCollector, + dict[IR, list[int]], + ], + config_options: ConfigOptions[StreamingExecutor], + *, + 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 + ---------- + 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. + 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 + ------ + RuntimeError + If :meth:`setup_worker` has not been called first. + """ + ir, partition_info, stats, collective_id_map = query_bundle + 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 + ) + + 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 + ] + 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, + ) + # 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 + + +class RayClient: + """ + User handle for a RapidsMPF Ray cluster. + + Typically created via :func:`ray_execution`. See that function for usage + and detailed documentation. + + 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``. + 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, + *, + owns_ray: bool, + ) -> None: + self._engine: pl.GPUEngine | None = engine + self._owns_ray: bool = owns_ray + + @property + def rank_actors(self) -> list[ActorHandle[RankActor]]: + """List of Ray rank actor handles.""" + if self._engine is None: + raise RuntimeError("rank_actors is not available after shutdown") + return self._engine.config["executor_options"]["ray_context"].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) + + @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]: + """ + 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]) + + 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. + + Raises + ------ + 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] + for ref in refs: + try: + ray.get(ref) + 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: + # 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 + 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() + + +def ray_execution( + *, + rapidsmpf_options: Options | None = None, + executor_options: 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`. + + The returned client supports both direct use and the context-manager + 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 + 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_options + Additional keyword arguments forwarded to :class:`polars.GPUEngine`. + ray_init_options + Keyword arguments forwarded to :func:`ray.init` when Ray is not + already initialized. + + Returns + ------- + 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 + ------ + 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. + TypeError + If ``executor_options`` contains a reserved key. + TypeError + If ``engine_options`` contains a reserved key. + + 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_options = engine_options or {} + ray_init_options = ray_init_options 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_context"} & executor_options.keys(): + 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 + if rapidsmpf_options is not None + else Options(get_environment_variables()) + ) + rapidsmpf_options.insert_if_absent({"num_streaming_threads": "4"}) + 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). + # 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_options) + + total_gpus = int(ray.cluster_resources().get("GPU", 0.0)) + # 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( + "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. + rank_actors: list[ActorHandle[RankActor]] = [ + RankActor.remote( # type: ignore[attr-defined] + nranks=free_gpus, + 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) + ] + + 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] + ) + + engine = pl.GPUEngine( + memory_resource=None, + executor="streaming", + executor_options={ + **executor_options, + "runtime": "rapidsmpf", + "cluster": "ray", + "ray_context": RayContext(rank_actors), + }, + **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/frontend/spmd.py similarity index 96% rename from python/cudf_polars/cudf_polars/experimental/rapidsmpf/spmd.py rename to python/cudf_polars/cudf_polars/experimental/rapidsmpf/frontend/spmd.py index 6b90860e4ca..20ca4d56425 100644 --- a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/spmd.py +++ b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/frontend/spmd.py @@ -160,6 +160,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 +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. + 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 ------- @@ -286,7 +291,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 @@ -311,10 +316,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 +345,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/cudf_polars/utils/config.py b/python/cudf_polars/cudf_polars/utils/config.py index d047084db81..fbfe4fe6e57 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.frontend.ray import RankActor + __all__ = [ "Cluster", @@ -50,6 +53,7 @@ "DynamicPlanningOptions", "InMemoryExecutor", "ParquetOptions", + "RayContext", "Runtime", "SPMDContext", "Scheduler", # Deprecated, kept for backward compatibility @@ -172,6 +176,7 @@ class Cluster(enum.StrEnum): SINGLE = "single" DISTRIBUTED = "distributed" SPMD = "spmd" + RAY = "ray" class Scheduler(enum.StrEnum): @@ -602,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. @@ -621,6 +626,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.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.frontend.ray.RankActor` + handles, one per GPU in the cluster. + """ + + rank_actors: list[ActorHandle[RankActor]] + + @dataclasses.dataclass(frozen=True, eq=True) class StreamingExecutor: """ @@ -846,6 +873,7 @@ class StreamingExecutor: ) ) spmd: SPMDContext | None = None + ray_context: RayContext | None = None def __post_init__(self) -> None: # noqa: D105 # Check for rapidsmpf runtime 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 00000000000..ad7188f184e --- /dev/null +++ b/python/cudf_polars/docs/cudf-polars-mp.md @@ -0,0 +1,406 @@ +# Multi-GPU Polars + +Multi-GPU Polars extends Polars query execution to multiple GPUs. + +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` | Description | Status | +| --------------- | ---------------------------------------------------- | ----------------- | +| `"single"` | Single-GPU, in-process execution | Stable (legacy) | +| `"distributed"` | Multi-GPU via [Dask Distributed][dask-distributed] | Stable (legacy) | +| `"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: + +* **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. + +* [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 + +`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( + 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"` +* `engine_options`: `"memory_resource"`, `"executor"` + +--- + +## 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 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 + +* 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.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. +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.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): + result = ( + pl.scan_parquet("/data/dataset/*.parquet") + .filter(pl.col("amount") > 100) + .group_by("customer_id") + .agg(pl.col("amount").sum()) + .collect(engine=engine) + ) + + 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: + +* `comm` — [`rapidsmpf.communicator.Communicator`][rapidsmpf-communicator] +* `ctx` — [`rapidsmpf.streaming.core.context.Context`][rapidsmpf-context] +* `engine` — {class}`~polars.lazyframe.engine_config.GPUEngine` + +Pass `engine` to every `LazyFrame.collect()` or `sink*()` call 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()` or `sink*()` calls +* 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 +`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( + comm=comm, + ctx=ctx, + local_df=result, + op_id=op_id, + ) +``` + +`op_id` identifies this collective across ranks — all ranks must pass the same value. +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()`. + +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 + +`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"` +* `engine_options`: `"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 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 00000000000..4ac15a6bd5f --- /dev/null +++ b/python/cudf_polars/tests/experimental/rapidsmpf/test_ray.py @@ -0,0 +1,223 @@ +# 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 +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.frontend.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.""" + 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_options={"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") +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_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_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(TypeError, match="reserved"): + ray_execution(engine_options=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, owns_ray=False) + client.shutdown() + client.shutdown() # must not raise + + +def test_ray_client_post_shutdown_state() -> None: + """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() + 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 + + +def test_raises_inside_rrun() -> None: + """ray_execution() must not be called from within an rrun cluster.""" + with ( + patch( + "rapidsmpf.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 +# --------------------------------------------------------------------------- + + +def test_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_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 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: + """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 isinstance(info["pid"], int) + # Each actor runs in its own process. + assert len({info["pid"] for info in infos}) == ray_client.nranks + + +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) + assert result.shape == (3, 1) + assert sorted(result["a"].to_list()) == [1, 2, 3] + + +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) + assert result.shape == (2, 1) + assert sorted(result["a"].to_list()) == [4, 5] + + +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. + 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") + ) + 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_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_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] diff --git a/python/cudf_polars/tests/experimental/rapidsmpf/test_spmd.py b/python/cudf_polars/tests/experimental/rapidsmpf/test_spmd.py index 585da785833..4b27821a548 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, ) @@ -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,34 +33,34 @@ 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 ( - pytest.raises(ValueError, match="reserved"), + pytest.raises(TypeError, match="reserved"), spmd_execution(executor_options={key: "anything"}), ): 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"} with ( - pytest.raises(ValueError, match="reserved"), + pytest.raises(TypeError, match="reserved"), spmd_execution(**kwargs), ): 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,