RapidsMPF streaming-engine translation

ci/run_cudf_polars_with_rapidsmpf_pytests.sh

@@ -5,16 +5,16 @@
 set -euo pipefail
 
 # Test cudf_polars with rapidsmpf integration
-# This script runs experimental tests with single cluster mode
+# This script runs experimental tests with single cluster mode and the rapidsmpf runtime
 
 # It is essential to cd into python/cudf_polars as `pytest-xdist` + `coverage` seem to work only at this directory level.
 
 # Support invoking run_cudf_polars_with_rapidsmpf_pytests.sh outside the script directory
 cd "$(dirname "$(realpath "${BASH_SOURCE[0]}")")"/../python/cudf_polars/
 
-# Run experimental tests with the "single" cluster mode
-# These tests use rapidsmpf.integrations.single for single-GPU shuffling
-rapids-logger "Running experimental tests with single cluster mode"
+# Run experimental tests with the "single" cluster mode and the "rapidsmpf" runtime
+rapids-logger "Running experimental tests with the 'rapidsmpf' runtime"
 python -m pytest --cache-clear "$@" "tests/experimental" \
     --executor streaming \
-    --cluster single
+    --cluster single \
+    --runtime rapidsmpf

python/cudf_polars/cudf_polars/experimental/base.py

@@ -402,6 +402,7 @@ class IOPartitionFlavor(IntEnum):
     SINGLE_FILE = enum.auto()  # 1:1 mapping between files and partitions
     SPLIT_FILES = enum.auto()  # Split each file into >1 partition
     FUSED_FILES = enum.auto()  # Fuse multiple files into each partition
+    SINGLE_READ = enum.auto()  # One worker/task reads everything
 
 
 class IOPartitionPlan:
@@ -414,6 +415,7 @@ class IOPartitionPlan:
       - SINGLE_FILE: `factor` must be `1`.
       - SPLIT_FILES: `factor` is the number of partitions per file.
       - FUSED_FILES: `factor` is the number of files per partition.
+      - SINGLE_READ: `factor` is the total number of files.
     """
 
     __slots__ = ("factor", "flavor")

python/cudf_polars/cudf_polars/experimental/benchmarks/utils.py

@@ -227,6 +227,8 @@ class RunConfig:
     queries: list[int]
     suffix: str
     executor: ExecutorType
+    runtime: str
+    stream_policy: str
     cluster: str
     scheduler: str  # Deprecated, kept for backward compatibility
     n_workers: int
@@ -267,6 +269,15 @@ def from_args(cls, args: argparse.Namespace) -> RunConfig:
         executor: ExecutorType = args.executor
         cluster = args.cluster
         scheduler = args.scheduler
+        runtime = args.runtime
+        stream_policy = args.stream_policy
+
+        # Handle "auto" stream policy
+        if stream_policy == "auto":
+            # TODO: Use pool by default for rapidsmpf runtime
+            # once stream-ordering bugs are fixed.
+            # See: https://github.com/rapidsai/cudf/issues/20484
+            stream_policy = "default"
 
         # Deal with deprecated scheduler argument
         # and non-streaming executors
@@ -335,6 +346,8 @@ def from_args(cls, args: argparse.Namespace) -> RunConfig:
             executor=executor,
             cluster=cluster,
             scheduler=scheduler,
+            runtime=runtime,
+            stream_policy=stream_policy,
             n_workers=args.n_workers,
             shuffle=args.shuffle,
             gather_shuffle_stats=args.rapidsmpf_dask_statistics,
@@ -374,7 +387,9 @@ def summarize(self) -> None:
             print(f"path: {self.dataset_path}")
             print(f"scale_factor: {self.scale_factor}")
             print(f"executor: {self.executor}")
+            print(f"stream_policy: {self.stream_policy}")
             if self.executor == "streaming":
+                print(f"runtime: {self.runtime}")
                 print(f"cluster: {self.cluster}")
                 print(f"blocksize: {self.blocksize}")
                 print(f"shuffle_method: {self.shuffle}")
@@ -415,20 +430,25 @@ def get_executor_options(
     """Generate executor_options for GPUEngine."""
     executor_options: dict[str, Any] = {}
 
-    if run_config.blocksize:
-        executor_options["target_partition_size"] = run_config.blocksize
-    if run_config.max_rows_per_partition:
-        executor_options["max_rows_per_partition"] = run_config.max_rows_per_partition
-    if run_config.shuffle:
-        executor_options["shuffle_method"] = run_config.shuffle
-    if run_config.broadcast_join_limit:
-        executor_options["broadcast_join_limit"] = run_config.broadcast_join_limit
-    if run_config.rapidsmpf_spill:
-        executor_options["rapidsmpf_spill"] = run_config.rapidsmpf_spill
-    if run_config.cluster == "distributed":
-        executor_options["cluster"] = "distributed"
-    if run_config.stats_planning:
-        executor_options["stats_planning"] = {"use_reduction_planning": True}
+    if run_config.executor == "streaming":
+        if run_config.blocksize:
+            executor_options["target_partition_size"] = run_config.blocksize
+        if run_config.max_rows_per_partition:
+            executor_options["max_rows_per_partition"] = (
+                run_config.max_rows_per_partition
+            )
+        if run_config.shuffle:
+            executor_options["shuffle_method"] = run_config.shuffle
+        if run_config.broadcast_join_limit:
+            executor_options["broadcast_join_limit"] = run_config.broadcast_join_limit
+        if run_config.rapidsmpf_spill:
+            executor_options["rapidsmpf_spill"] = run_config.rapidsmpf_spill
+        if run_config.cluster == "distributed":
+            executor_options["cluster"] = "distributed"
+        if run_config.stats_planning:
+            executor_options["stats_planning"] = {"use_reduction_planning": True}
+        executor_options["client_device_threshold"] = run_config.spill_device
+        executor_options["runtime"] = run_config.runtime
 
     if (
         benchmark
@@ -467,7 +487,7 @@ def print_query_plan(
         if args.explain_logical:
             print(f"\nQuery {q_id} - Logical plan\n")
             print(explain_query(q, engine, physical=False))
-        if args.explain:
+        if args.explain and run_config.executor == "streaming":
             print(f"\nQuery {q_id} - Physical plan\n")
             print(explain_query(q, engine))
     else:
@@ -671,6 +691,22 @@ def parse_args(
                 - synchronous : Run locally in a single process
                 - distributed : Use Dask for multi-GPU execution"""),
     )
+    parser.add_argument(
+        "--runtime",
+        type=str,
+        choices=["tasks", "rapidsmpf"],
+        default="tasks",
+        help="Runtime to use for the streaming executor (tasks or rapidsmpf).",
+    )
+    parser.add_argument(
+        "--stream-policy",
+        type=str,
+        choices=["auto", "default", "new", "pool"],
+        default="auto",
+        help=textwrap.dedent("""\
+            CUDA stream policy (auto, default, new, pool).
+            Default: auto (use the default policy for the runtime)"""),
+    )
     parser.add_argument(
         "--n-workers",
         default=1,
@@ -869,6 +905,10 @@ def run_polars(
         executor_options = get_executor_options(run_config, benchmark=benchmark)
         engine = pl.GPUEngine(
             raise_on_fail=True,
+            memory_resource=rmm.mr.CudaAsyncMemoryResource()
+            if run_config.rmm_async
+            else None,
+            cuda_stream_policy=run_config.stream_policy,
             executor=run_config.executor,
             executor_options=executor_options,
         )
@@ -928,7 +968,10 @@ def run_polars(
             if args.print_results:
                 print(result)
 
-            print(f"Query {q_id} - Iteration {i} finished in {record.duration:0.4f}s")
+            print(
+                f"Query {q_id} - Iteration {i} finished in {record.duration:0.4f}s",
+                flush=True,
+            )
             records[q_id].append(record)
 
     run_config = dataclasses.replace(run_config, records=dict(records))

python/cudf_polars/cudf_polars/experimental/explain.py

@@ -61,7 +61,17 @@ def explain_query(
     ir = Translator(q._ldf.visit(), engine).translate_ir()
 
     if physical:
-        lowered_ir, partition_info = lower_ir_graph(ir, config)
+        if (
+            config.executor.name == "streaming"
+            and config.executor.runtime == "rapidsmpf"
+        ):  # pragma: no cover; rapidsmpf runtime not tested in CI yet
+            from cudf_polars.experimental.rapidsmpf.core import (
+                lower_ir_graph as rapidsmpf_lower_ir_graph,
+            )
+
+            lowered_ir, partition_info = rapidsmpf_lower_ir_graph(ir, config)
+        else:
+            lowered_ir, partition_info = lower_ir_graph(ir, config)
         return _repr_ir_tree(lowered_ir, partition_info)
     else:
         if config.executor.name == "streaming":

python/cudf_polars/cudf_polars/experimental/join.py

@@ -144,6 +144,7 @@ def _make_bcast_join(
     left: IR,
     right: IR,
     shuffle_method: ShuffleMethod,
+    streaming_runtime: str,
 ) -> tuple[IR, MutableMapping[IR, PartitionInfo]]:
     if ir.options[0] != "Inner":
         left_count = partition_info[left].count
@@ -162,22 +163,23 @@ def _make_bcast_join(
         # - In some cases, we can perform the partial joins
         #   sequentially. However, we are starting with a
         #   catch-all algorithm that works for all cases.
-        if left_count >= right_count:
-            right = _maybe_shuffle_frame(
-                right,
-                ir.right_on,
-                partition_info,
-                shuffle_method,
-                right_count,
-            )
-        else:
-            left = _maybe_shuffle_frame(
-                left,
-                ir.left_on,
-                partition_info,
-                shuffle_method,
-                left_count,
-            )
+        if streaming_runtime == "tasks":
+            if left_count >= right_count:
+                right = _maybe_shuffle_frame(
+                    right,
+                    ir.right_on,
+                    partition_info,
+                    shuffle_method,
+                    right_count,
+                )
+            else:
+                left = _maybe_shuffle_frame(
+                    left,
+                    ir.left_on,
+                    partition_info,
+                    shuffle_method,
+                    left_count,
+                )
 
     new_node = ir.reconstruct([left, right])
     partition_info[new_node] = PartitionInfo(count=output_count)
@@ -288,6 +290,7 @@ def _(
             left,
             right,
             config_options.executor.shuffle_method,
+            config_options.executor.runtime,
         )
     else:
         # Create a hash join

python/cudf_polars/cudf_polars/experimental/parallel.py

@@ -224,6 +224,29 @@ def post_process_task_graph(
     return graph
 
 
+def evaluate_rapidsmpf(
+    ir: IR,
+    config_options: ConfigOptions,
+) -> DataFrame:  # pragma: no cover; rapidsmpf runtime not tested in CI yet
+    """
+    Evaluate with the RapidsMPF streaming runtime.
+
+    Parameters
+    ----------
+    ir
+        Logical plan to evaluate.
+    config_options
+        GPUEngine configuration options.
+
+    Returns
+    -------
+    A cudf-polars DataFrame object.
+    """
+    from cudf_polars.experimental.rapidsmpf.core import evaluate_logical_plan
+
+    return evaluate_logical_plan(ir, config_options)
+
+
 def evaluate_streaming(
     ir: IR,
     config_options: ConfigOptions,
@@ -245,11 +268,19 @@ def evaluate_streaming(
     # Clear source info cache in case data was overwritten
     _clear_source_info_cache()
 
-    ir, partition_info = lower_ir_graph(ir, config_options)
+    assert config_options.executor.name == "streaming", "Executor must be streaming"
+    if (
+        config_options.executor.runtime == "rapidsmpf"
+    ):  # pragma: no cover; rapidsmpf runtime not tested in CI yet
+        # Using the RapidsMPF streaming runtime.
+        return evaluate_rapidsmpf(ir, config_options)
+    else:
+        # Using the default task engine.
+        ir, partition_info = lower_ir_graph(ir, config_options)
 
-    graph, key = task_graph(ir, partition_info, config_options)
+        graph, key = task_graph(ir, partition_info, config_options)
 
-    return get_scheduler(config_options)(graph, key)
+        return get_scheduler(config_options)(graph, key)
 
 
 @generate_ir_tasks.register(IR)

python/cudf_polars/cudf_polars/experimental/rapidsmpf/__init__.py

@@ -0,0 +1,8 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES.
+# SPDX-License-Identifier: Apache-2.0
+
+"""RapidsMPF streaming-engine support."""
+
+from __future__ import annotations
+
+__all__: list[str] = []