diff --git a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds.py b/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds.py index e48063058811..30cacaf38186 100644 --- a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds.py +++ b/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds.py @@ -190,8 +190,8 @@ class PDSDSDuckDBQueries(PDSDSQueries): args = parse_args(parser=parser) if args.engine == "polars": - run_polars(PDSDSPolarsQueries, args, num_queries=99) + run_polars(PDSDSPolarsQueries, args) elif args.engine == "duckdb": - run_duckdb(PDSDSDuckDBQueries, args, num_queries=99) + run_duckdb(PDSDSDuckDBQueries, args) else: raise ValueError(f"Invalid engine: {args.engine}") diff --git a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsh.py b/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsh.py index 7592c0ab47b3..77b05e8cfb1c 100644 --- a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsh.py +++ b/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsh.py @@ -1796,8 +1796,8 @@ def q22(run_config: RunConfig) -> str: args = parse_args(parser=parser) if args.engine == "polars": - run_polars(PDSHQueries, args, num_queries=22) + run_polars(PDSHQueries, args) elif args.engine == "duckdb": - run_duckdb(PDSHDuckDBQueries, args, num_queries=22) + run_duckdb(PDSHDuckDBQueries, args) else: raise ValueError(f"Invalid engine: {args.engine}") diff --git a/python/cudf_polars/cudf_polars/experimental/benchmarks/utils.py b/python/cudf_polars/cudf_polars/experimental/benchmarks/utils.py index 5b6380123051..70b32f34838f 100644 --- a/python/cudf_polars/cudf_polars/experimental/benchmarks/utils.py +++ b/python/cudf_polars/cudf_polars/experimental/benchmarks/utils.py @@ -1563,7 +1563,6 @@ def run_polars_query( def run_polars( benchmark: Any, args: argparse.Namespace, - num_queries: int = 22, ) -> None: """Run the queries using the given benchmark and executor options.""" vars(args).update({"query_set": benchmark.name}) @@ -1898,9 +1897,7 @@ def execute_duckdb_query( return conn.execute(query).pl() -def run_duckdb( - duckdb_queries_cls: Any, args: argparse.Namespace, *, num_queries: int -) -> None: +def run_duckdb(duckdb_queries_cls: Any, args: argparse.Namespace) -> None: """Run the benchmark with DuckDB.""" vars(args).update({"query_set": duckdb_queries_cls.name}) run_config = RunConfig.from_args(args) diff --git a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/collectives/common.py b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/collectives/common.py index 5a8a0c112d21..8d747ceb6d95 100644 --- a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/collectives/common.py +++ b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/collectives/common.py @@ -5,6 +5,7 @@ from __future__ import annotations import threading +from contextlib import contextmanager from typing import TYPE_CHECKING, Literal from rapidsmpf.shuffler import Shuffler @@ -17,6 +18,7 @@ from cudf_polars.experimental.shuffle import Shuffle if TYPE_CHECKING: + from collections.abc import Iterator from types import TracebackType from cudf_polars.dsl.ir import IR @@ -139,3 +141,25 @@ def __exit__( for collective_id in collective_ids: _release_collective_id(collective_id) return False + + +@contextmanager +def reserve_op_id() -> Iterator[int]: + """ + Reserve a single collective operation ID. + + This function and the ID it yields must only be used **outside** of a + ``run_actor_graph`` call. It is intended for SPMD mode, where operations + such as gathering results across ranks are performed directly rather than + through the actor graph. The contained block _must_ wait for completion of the collective. + + Yields + ------ + collective_id : int + A vacant collective ID reserved from the global vacancy pool. + """ + collective_id = _get_new_collective_id() + try: + yield collective_id + finally: + _release_collective_id(collective_id) diff --git a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/core.py b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/core.py index 82b6ea4bab3c..50df9f4d4ee7 100644 --- a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/core.py +++ b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/core.py @@ -125,6 +125,19 @@ def evaluate_logical_plan( collective_id_map, collect_metadata=collect_metadata, ) + elif config_options.executor.cluster == "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, + ) else: # Single-process execution: Run locally result, metadata_collector = evaluate_pipeline( @@ -299,11 +312,7 @@ def evaluate_pipeline( stream, ) - # We need to materialize the polars dataframe before we drop the rapidsmpf - # context, which keeps the CUDA streams alive. - stream = df.stream result = df.to_polars() - stream.synchronize() # Now we need to drop *all* GPU data. This ensures that no cudaFreeAsync runs # before the Context, which ultimately contains the rmm MR, goes out of scope. diff --git a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/io.py b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/io.py index 8040a7e72f89..81df4c2d38e9 100644 --- a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/io.py +++ b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/io.py @@ -157,6 +157,7 @@ async def dataframescan_node( num_producers: int, rows_per_partition: int, estimated_chunk_bytes: int, + distributed_scan: bool, ) -> None: """ DataFrameScan node for rapidsmpf. @@ -180,14 +181,23 @@ async def dataframescan_node( estimated_chunk_bytes Estimated size of each chunk in bytes. Used for memory reservation with block spilling to avoid thrashing. + distributed_scan + If ``True``, the DataFrame is treated as a shared object and divided + across workers so each rank reads a disjoint subset. This is normally + used in ``Cluster.DISTRIBUTED`` mode. + + If ``False``, the DataFrame is treated as rank-local and each rank + scans its local DataFrame in full. This is normally used in + ``Cluster.SPMD`` mode. """ async with shutdown_on_error(context, ch_out, trace_ir=ir) as tracer: # Find local partition count. nrows = ir.df.shape()[0] global_count = math.ceil(nrows / rows_per_partition) if nrows > 0 else 0 - # For single rank, simplify the logic - if comm.nranks == 1: + # For single rank or when scanning the full local DataFrame, each rank + # uses all partitions with no offset. + if not distributed_scan or comm.nranks == 1: local_count = global_count local_offset = 0 else: @@ -292,10 +302,10 @@ def _( num_producers=num_producers, rows_per_partition=rows_per_partition, estimated_chunk_bytes=estimated_chunk_bytes, + distributed_scan=config_options.executor.cluster != "spmd", ) ] } - return nodes, channels diff --git a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/spmd.py b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/spmd.py new file mode 100644 index 000000000000..6b90860e4ca6 --- /dev/null +++ b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/spmd.py @@ -0,0 +1,383 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. +# SPDX-License-Identifier: Apache-2.0 +"""RapidsMPF streaming-engine using the SPMD Cluster style.""" + +from __future__ import annotations + +from concurrent.futures import ThreadPoolExecutor +from contextlib import contextmanager +from typing import TYPE_CHECKING, Any, cast + +from rapidsmpf import bootstrap +from rapidsmpf.coll import AllGather +from rapidsmpf.config import Options, get_environment_variables +from rapidsmpf.integrations.cudf.partition import unpack_and_concat +from rapidsmpf.memory.packed_data import PackedData +from rapidsmpf.progress_thread import ProgressThread +from rapidsmpf.rmm_resource_adaptor import RmmResourceAdaptor +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 pylibcudf as plc +import rmm.mr +from pylibcudf.contiguous_split import pack + +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 SPMDContext + +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_spmd_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]: + """ + Build and evaluate a RapidsMPF streaming pipeline in SPMD mode. + + In SPMD mode every rank executes the same Python/Polars script + independently. Each rank owns its local DataFrames, which are + treated as rank-local fragments of a larger distributed dataset and + fed directly into the pipeline. Collective operations (shuffles, + all-gathers, etc.) coordinate across ranks to produce a globally + consistent result. + + 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 + ------- + The concatenated output DataFrame and, if ``collect_metadata`` is + True, the list of channel metadata objects; otherwise ``None``. + """ + if config_options.executor.runtime != "rapidsmpf": + raise RuntimeError("Runtime must be rapidsmpf") + if config_options.executor.spmd is None: + raise RuntimeError("spmd must be set for SPMD mode") + comm = config_options.executor.spmd.comm + context = config_options.executor.spmd.context + py_executor = config_options.executor.spmd.py_executor + + ir_context = IRExecutionContext(get_cuda_stream=context.get_stream_from_pool) + + metadata_collector: list[ChannelMetadata] | None = [] if collect_metadata else None + + nodes, output = generate_network( + context, + 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=py_executor) + + messages = output.release() + chunks = [ + TableChunk.from_message(msg).make_available_and_spill( + context.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 correct schema + stream = ir_context.get_cuda_stream() + chunk = empty_table_chunk(ir, context, 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 + + +def allgather_polars_dataframe( + *, + comm: Communicator, + ctx: Context, + local_df: pl.DataFrame, + op_id: int, +) -> pl.DataFrame: + """ + AllGather a rank-local DataFrame so every rank receives the full result. + + Each rank contributes its local ``local_df`` fragment and receives the + concatenation of all ranks' fragments in rank order. This is the SPMD + equivalent of a distributed ``collect``: after the call, every rank holds + the same complete dataset. + + Parameters + ---------- + comm + The RapidsMPF communicator. + ctx + The RapidsMPF context. + local_df + Rank-local DataFrame to contribute. + op_id + Operation ID for this AllGather collective. Must be identical on every + rank. + + Returns + ------- + DataFrame containing rows from all ranks, ordered by rank. + """ + stream = ctx.get_stream_from_pool() + col_names = local_df.columns + + plc_table = plc.Table.from_arrow(local_df.to_arrow()) + + packed_data = PackedData.from_cudf_packed_columns( + pack(plc_table, stream), + stream, + ctx.br(), + ) + + # Bulk AllGather: each rank contributes once (sequence_number=0) + allgather = AllGather(comm, op_id, ctx.br()) + allgather.insert(0, packed_data) + allgather.insert_finished() + results = allgather.wait_and_extract(ordered=True) + + # Deserialize and concatenate all ranks' contributions + plc_result = unpack_and_concat(results, stream, ctx.br()) + + # pylibcudf Table -> pl.DataFrame (restore column names) + ret = pl.from_arrow(plc_result.to_arrow(col_names)) + assert isinstance(ret, pl.DataFrame) + return ret + + +@contextmanager +def spmd_execution( + *, + mr: rmm.mr.DeviceMemoryResource | None = None, + rapidsmpf_options: Options | None = None, + executor_options: dict[str, object] | None = None, + **engine_kwargs: Any, +) -> Iterator[tuple[Communicator, Context, pl.GPUEngine]]: + """ + Context manager that bootstraps a RapidsMPF SPMD context and a matching GPUEngine. + + **SPMD execution model** + + SPMD (Single Program, Multiple Data) is a parallel programming model where each + process runs the *same* Python script independently on its own slice of data. + When launched with the RapidsMPF launcher `rrun`, multiple identical processes + are started. Each process owns a rank-local :class:`~polars.LazyFrame` + representing its fragment of the distributed dataset. Collective operations, + such as shuffles, all-gathers, and joins, coordinate across ranks to produce + a globally consistent result. + + This context manager is the primary entry point for SPMD execution. It: + + - Bootstraps a UCXX communicator connecting all ``N`` ranks. + - Creates a RapidsMPF :class:`~rapidsmpf.streaming.core.context.Context` + that owns GPU memory and a CUDA-stream pool. + - Returns a :class:`~polars.lazyframe.engine_config.GPUEngine` wired to that + context so that ``LazyFrame.collect(engine=engine)`` dispatches through the + RapidsMPF streaming executor. + + All resources (communicator, stream pool, thread-pool) are released on exit. + + **DataFrame and LazyFrame semantics** + + Because every rank runs an independent Python process, a :class:`~polars.DataFrame` + is always *rank-local* i.e. it contains only that rank's fragment of the distributed + dataset. This is true whether the DataFrame originates from a file reader or from + Python literals. + + File-based sources (``scan_parquet``, ``scan_csv``, ...) distribute their work + automatically: the engine assigns disjoint file- or row-group ranges to each rank, + so different ranks produce different data. + + An in-memory ``DataFrame`` (or one produced by a previous ``collect``) is already + rank-local by construction. Each rank processes its own copy in full; the engine + does **not** re-slice it across ranks. In particular, the two patterns below are + equivalent: + + .. code-block:: python + + # One-step: scan and transform in a single pipeline + result = pl.scan_parquet(...).pipe(transform).collect(engine=engine) + + # Two-step: collect an intermediate result, then transform + intermediate = pl.scan_parquet(...).collect(engine=engine) + result = intermediate.lazy().pipe(transform).collect(engine=engine) + + In both cases rank k operates on exactly the data it read from parquet. The + intermediate ``collect`` simply materializes the data in memory; it does not + change which rows belong to which rank. + + **Query symmetry requirement** + + Every rank 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 rank does not, all ranks will deadlock. This means + your driver script must be fully deterministic: avoid rank-conditional + ``collect`` calls, early exits, or any branching that would cause different + ranks to execute different query graphs. + + Must be invoked under the ``rrun`` launcher. Use + :func:`rapidsmpf.bootstrap.is_running_with_rrun` to test this at runtime. + + Parameters + ---------- + mr + RMM device memory resource to use. Defaults to + ``rmm.mr.CudaAsyncMemoryResource()`` when ``None``. + rapidsmpf_options + RapidsMPF options. Defaults to ``Options(get_environment_variables())`` + when ``None``. + executor_options + Extra keyword arguments forwarded to the ``executor_options`` dict of + :class:`~polars.lazyframe.engine_config.GPUEngine`. The keys + ``"runtime"``, ``"cluster"``, and ``"spmd"`` are reserved and may not + be overridden. + **engine_kwargs + Extra keyword arguments forwarded directly to + :class:`~polars.lazyframe.engine_config.GPUEngine`. For example, + pass ``parquet_options={"use_rapidsmpf_native": True}`` to enable + native Parquet reads. The keys ``"memory_resource"`` and + ``"executor"`` are reserved and may not be overridden. + + Yields + ------ + comm : Communicator + The active RapidsMPF communicator. + ctx : Context + The active RapidsMPF context. + engine : pl.GPUEngine + A Polars GPU engine wired to ``comm`` and ``ctx``. Pass it to + ``LazyFrame.collect(engine=engine)`` on each rank. + + Raises + ------ + RuntimeError + If not running under the ``rrun`` launcher (i.e. + :func:`rapidsmpf.bootstrap.is_running_with_rrun` returns ``False``). + ValueError + If ``executor_options`` contains any of the reserved keys + ``"runtime"``, ``"cluster"``, or ``"spmd"``. + ValueError + If ``engine_kwargs`` contains any of the reserved keys + ``"raise_on_fail"``, ``"memory_resource"``, or ``"executor"``. + + Examples + -------- + >>> with spmd_execution() as (comm, ctx, engine): # doctest: +SKIP + ... result = ( + ... df.lazy().group_by("a").agg(pl.col("b").sum()).collect(engine=engine) + ... ) + ... full = allgather_polars_dataframe( + ... comm=comm, ctx=ctx, local_df=result, op_id=0 + ... ) + """ + if not bootstrap.is_running_with_rrun(): + raise RuntimeError( + "spmd_execution() requires the rrun launcher. " + "Launch your script with `rrun -n python your_script.py` " + "to enable SPMD execution." + ) + + executor_options = executor_options or {} + engine_kwargs = engine_kwargs or {} + + # Check for reserved keys. + if bad := {"runtime", "cluster", "spmd"} & 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()) + ) + mr = RmmResourceAdaptor(mr if mr is not None else rmm.mr.CudaAsyncMemoryResource()) + comm = bootstrap.create_ucxx_comm( + progress_thread=ProgressThread(), + type=bootstrap.BackendType.AUTO, + options=rapidsmpf_options, + ) + py_executor = ThreadPoolExecutor( + max_workers=cast( + int, executor_options.get("rapidsmpf_py_executor_max_workers", 1) + ), + thread_name_prefix="spmd-executor", + ) + try: + with Context.from_options(comm.logger, mr, rapidsmpf_options) as ctx: + engine = pl.GPUEngine( + memory_resource=ctx.br().device_mr, + executor="streaming", + executor_options={ + **executor_options, + "runtime": "rapidsmpf", + "cluster": "spmd", + "spmd": SPMDContext( + comm=comm, context=ctx, py_executor=py_executor + ), + }, + **engine_kwargs, + ) + yield comm, ctx, engine + finally: + # The Context has already been exited above, so no work can be + # pending in py_executor at this point; wait=False is safe. + py_executor.shutdown(wait=False) diff --git a/python/cudf_polars/cudf_polars/utils/config.py b/python/cudf_polars/cudf_polars/utils/config.py index 0f3d00a8a8e5..5d6497adda72 100644 --- a/python/cudf_polars/cudf_polars/utils/config.py +++ b/python/cudf_polars/cudf_polars/utils/config.py @@ -34,6 +34,10 @@ if TYPE_CHECKING: from collections.abc import Callable + from concurrent.futures import ThreadPoolExecutor + + from rapidsmpf.communicator.communicator import Communicator + from rapidsmpf.streaming.core.context import Context import polars.lazyframe.engine_config @@ -47,6 +51,7 @@ "InMemoryExecutor", "ParquetOptions", "Runtime", + "SPMDContext", "Scheduler", # Deprecated, kept for backward compatibility "ShuffleMethod", "ShufflerInsertionMethod", @@ -167,6 +172,7 @@ class Cluster(enum.StrEnum): SINGLE = "single" DISTRIBUTED = "distributed" + SPMD = "spmd" class Scheduler(enum.StrEnum): @@ -601,6 +607,35 @@ def __hash__(self) -> int: return hash((self.qualname, json.dumps(self.options, sort_keys=True))) +@dataclasses.dataclass(frozen=True) +class SPMDContext: + """ + Configuration for SPMD (Single Program Multiple Data) execution. + + .. note:: + This dataclass is **not picklable** because :class:`Communicator`, + :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 + pickling is never required. Do not use this class with Dask or any other + framework that serializes executor configuration across process boundaries. + + Parameters + ---------- + comm + The active RapidsMPF communicator. + context + The active RapidsMPF context. + py_executor + Thread-pool executor used to drive the actor network on each rank. + """ + + comm: Communicator + context: Context + py_executor: ThreadPoolExecutor + + @dataclasses.dataclass(frozen=True, eq=True) class StreamingExecutor: """ @@ -835,6 +870,7 @@ class StreamingExecutor: f"{_env_prefix}__RAPIDSMPF_PY_EXECUTOR_MAX_WORKERS", int, default=None ) ) + spmd: SPMDContext | None = None def __post_init__(self) -> None: # noqa: D105 # Check for rapidsmpf runtime diff --git a/python/cudf_polars/tests/experimental/rapidsmpf/test_spmd.py b/python/cudf_polars/tests/experimental/rapidsmpf/test_spmd.py new file mode 100644 index 000000000000..585da785833c --- /dev/null +++ b/python/cudf_polars/tests/experimental/rapidsmpf/test_spmd.py @@ -0,0 +1,178 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. +# SPDX-License-Identifier: Apache-2.0 +"""Tests for SPMD execution mode.""" + +from __future__ import annotations + +from typing import Any + +import pytest +from rapidsmpf.bootstrap import is_running_with_rrun + +import polars as pl + +import rmm.mr + +from cudf_polars.experimental.rapidsmpf.collectives.common import reserve_op_id +from cudf_polars.experimental.rapidsmpf.spmd import ( + allgather_polars_dataframe, + spmd_execution, +) + +pytestmark = pytest.mark.skipif( + not is_running_with_rrun(), + reason="use something like `rrun -n python -m pytest ...` to run SPMD tests", +) + + +def test_spmd_execution_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 + assert ctx is not None + assert isinstance(engine, pl.GPUEngine) + + +def test_spmd_execution_reserved_keys() -> None: + """executor_options rejects reserved keys.""" + for key in ("runtime", "cluster", "spmd"): + with ( + pytest.raises(ValueError, match="reserved"), + spmd_execution(executor_options={key: "anything"}), + ): + pass + + +def test_spmd_execution_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"), + spmd_execution(**kwargs), + ): + pass + + +def test_spmd_execution_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: + """spmd_execution accepts a custom memory resource.""" + mr = rmm.mr.CudaMemoryResource() + with spmd_execution(mr=mr) as (comm, ctx, engine): + result = pl.LazyFrame({"a": [1, 2, 3]}).collect(engine=engine) + assert result.shape == (3, 1) + + +def test_spmd_execution_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 + lf = pl.LazyFrame({"a": [rank], "b": [rank * 10]}) + result = lf.collect(engine=engine) + assert result.shape == (1, 2) + assert result["a"].to_list() == [rank] + assert result["b"].to_list() == [rank * 10] + + +def test_spmd_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 + into a LazyFrame the engine processes that rank's copy in full rather than + re-slicing it across ranks. So ``lf.collect().lazy().op.collect()`` must + produce the same result as ``lf.op.collect()``. + """ + with spmd_execution() as (comm, ctx, engine): + rank = comm.rank + lf = pl.LazyFrame({"a": [rank, rank + 1, rank + 2], "b": [0, 1, 2]}) + + # One-step + one_step = lf.filter(pl.col("b") >= 1).collect(engine=engine) + + # Two-step: materialize then re-wrap + intermediate = lf.collect(engine=engine) + two_step = intermediate.lazy().filter(pl.col("b") >= 1).collect(engine=engine) + + assert one_step.sort("a").equals(two_step.sort("a")) + + +def test_spmd_execution_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 + nranks = comm.nranks + lf = pl.LazyFrame({"a": [rank], "b": [rank * 10]}) + local_result = lf.group_by("a").agg(pl.col("b").sum()).collect(engine=engine) + with reserve_op_id() as op_id: + global_result = allgather_polars_dataframe( + comm=comm, ctx=ctx, local_df=local_result, op_id=op_id + ) + assert global_result.shape == (nranks, 2) + assert global_result.sort("a")["a"].to_list() == list(range(nranks)) + assert global_result.sort("a")["b"].to_list() == [r * 10 for r in range(nranks)] + + +def test_allgather_polars_dataframe() -> None: + """allgather_polars_dataframe collects every rank's contribution in rank order.""" + with spmd_execution() as (comm, ctx, _): + rank = comm.rank + nranks = comm.nranks + local = pl.DataFrame({"rank": [rank], "val": [rank * 2]}) + with reserve_op_id() as op_id: + result = allgather_polars_dataframe( + comm=comm, ctx=ctx, local_df=local, op_id=op_id + ) + assert result.shape == (nranks, 2) + assert result["rank"].to_list() == list(range(nranks)) + assert result["val"].to_list() == [r * 2 for r in range(nranks)] + + +def test_spmd_execution_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, + ctx, + engine, + ): + result = pl.LazyFrame({"a": [1, 2, 3]}).collect(engine=engine) + assert result.shape == (3, 1) + + +def test_allgather_polars_dataframe_empty() -> None: + """allgather handles an empty (zero-row) local DataFrame on every rank.""" + with spmd_execution() as (comm, ctx, _): + local = pl.DataFrame( + {"a": pl.Series([], dtype=pl.Int32), "b": pl.Series([], dtype=pl.Float64)} + ) + with reserve_op_id() as op_id: + result = allgather_polars_dataframe( + comm=comm, ctx=ctx, local_df=local, op_id=op_id + ) + assert result.shape == (0, 2) + assert result.columns == ["a", "b"] + assert result.dtypes == [pl.Int32, pl.Float64] + + +def test_allgather_polars_dataframe_multi_column() -> None: + """allgather preserves column names, count, and dtypes for multi-column DataFrames.""" + with spmd_execution() as (comm, ctx, _): + rank = comm.rank + nranks = comm.nranks + local = pl.DataFrame( + {"rank": [rank], "x": [float(rank)], "label": [f"r{rank}"]} + ) + with reserve_op_id() as op_id: + result = allgather_polars_dataframe( + comm=comm, ctx=ctx, local_df=local, op_id=op_id + ) + assert result.shape == (nranks, 3) + assert result.columns == ["rank", "x", "label"] + sorted_result = result.sort("rank") + assert sorted_result["rank"].to_list() == list(range(nranks)) + assert sorted_result["x"].to_list() == [float(r) for r in range(nranks)] + assert sorted_result["label"].to_list() == [f"r{r}" for r in range(nranks)]