RapidsMPF SPMD mode

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}")

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}")

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)

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)

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.

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