[Multi-GPU Polars] Introduce StreamingEngine base class and SPMDEngine (rapidsai#21867)

Introduces `StreamingEngine`, a new base class that all multi-GPU execution frontends (SPMD, Ray, Dask, etc.) will return. This PR provides the initial implementation via `SPMDEngine`; Ray and Dask frontends will follow in separate PRs.

* Add `StreamingEngine` base class and `SPMDEngine` implementation.
* Update tests to use a shared fixture for engine creation instead of constructing engines inline.
* Move tests specific to the legacy tasks/distributed runtime to `tests/experimental/legacy/`.
* Add `pytest.mark.spmd` marker and an autouse fixture that skips non-SPMD tests when running under `rrun -n N` with `N > 1`.
* Separate regular cudf-polars tests from experimental tests, we now have:
  - `run_cudf_polars_pytests.sh`  (CI-blocking).
  - `run_cudf_polars_experimental_pytests.sh` (requires rapidsmpf).

This is the first step in migrating away from the legacy tasks and Dask runtimes toward a unified streaming engine abstraction.

Authors:
  - Mads R. B. Kristensen (https://github.com/madsbk)

Approvers:
  - Bradley Dice (https://github.com/bdice)
  - Tom Augspurger (https://github.com/TomAugspurger)
  - Niranda Perera (https://github.com/nirandaperera)
  - Lawrence Mitchell (https://github.com/wence-)

URL: rapidsai#21867

Author: madsbk

.github/workflows/pr.yaml

@@ -490,7 +490,7 @@ jobs:
       # (rapidsmpf compatibility already validated in rapidsmpf CI)
       matrix_filter: map(select(.ARCH == "amd64")) | group_by(.CUDA_VER|split(".")|map(tonumber)|.[0]) | map(max_by([(.PY_VER|split(".")|map(tonumber)), (.CUDA_VER|split(".")|map(tonumber))]))
       build_type: pull-request
-      script: "ci/test_cudf_polars_with_rapidsmpf.sh"
+      script: "ci/test_cudf_polars_experimental.sh"
   cudf-polars-polars-tests:
     needs: [wheel-build-cudf-polars, changed-files]
     secrets: inherit

.github/workflows/test.yaml

@@ -164,7 +164,7 @@ jobs:
       branch: ${{ inputs.branch }}
       date: ${{ inputs.date }}
       sha: ${{ inputs.sha }}
-      script: "ci/test_cudf_polars_with_rapidsmpf.sh"
+      script: "ci/test_cudf_polars_experimental.sh"
       continue-on-error: true
   cudf-polars-polars-tests:
     secrets: inherit

…un_cudf_polars_with_rapidsmpf_pytests.sh …/run_cudf_polars_experimental_pytests.shci/run_cudf_polars_with_rapidsmpf_pytests.sh renamed to ci/run_cudf_polars_experimental_pytests.sh ci/run_cudf_polars_with_rapidsmpf_pytests.sh renamed to ci/run_cudf_polars_experimental_pytests.sh

@@ -4,24 +4,23 @@
 
 set -euo pipefail
 
-# Test cudf_polars with rapidsmpf integration
-# This script runs experimental tests with single cluster mode and the rapidsmpf runtime
+# Test cudf_polars experimental.
 
 # 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
+# Support invoking outside the script directory
 cd "$(dirname "$(realpath "${BASH_SOURCE[0]}")")"/../python/cudf_polars/
 
-# Run experimental tests with the "single" cluster mode and the "rapidsmpf" runtime
-rapids-logger "Running experimental tests with the 'rapidsmpf' runtime and a 'single' cluster"
-timeout 10m python -m pytest --cache-clear "$@" "tests/experimental" \
+rapids-logger "Running experimental legacy tests with the 'rapidsmpf' runtime and a 'single' cluster"
+timeout 10m python -m pytest --cache-clear "$@" "tests/experimental/legacy" \
     --executor streaming \
     --cluster single \
     --runtime rapidsmpf
 
-# Run experimental tests with the "distributed" cluster mode and the "rapidsmpf" runtime
-rapids-logger "Running experimental tests with the 'rapidsmpf' runtime and a 'distributed' cluster"
-timeout 10m python -m pytest --cache-clear "$@" "tests/experimental" \
+rapids-logger "Running experimental legacy tests with the 'rapidsmpf' runtime and a 'distributed' cluster"
+timeout 10m python -m pytest --cache-clear "$@" "tests/experimental/legacy" \
     --executor streaming \
     --cluster distributed \
     --runtime rapidsmpf
+
+rapids-logger "Running experimental tests"
+timeout 10m python -m pytest --cache-clear "$@" tests/experimental --ignore=tests/experimental/legacy

ci/run_cudf_polars_pytests.sh

@@ -1,25 +1,16 @@
 #!/bin/bash
-# SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION.
+# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION.
 # SPDX-License-Identifier: Apache-2.0
 
 set -euo pipefail
 
 # 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_pytests.sh outside the script directory
 cd "$(dirname "$(realpath "${BASH_SOURCE[0]}")")"/../python/cudf_polars/
 
-# Test the "in-memory" executor
-python -m pytest --cache-clear "$@" tests --executor in-memory
-
-# Test the default "streaming" executor
-python -m pytest --cache-clear "$@" tests --executor streaming
-
-# Test the "streaming" executor with small blocksize
-python -m pytest --cache-clear "$@" tests --executor streaming --blocksize-mode small
-
-# Run experimental tests with Distributed cluster
-python -m pytest --cache-clear "$@" "tests/experimental" \
-    --executor streaming \
-    --cluster distributed \
-    --cov-fail-under=0  # No code-coverage requirement for these tests.
+# Run all non-experimental tests using both the in-memory and streaming executor.
+IGNORE_EXPERIMENTAL="--ignore=tests/experimental/"
+python -m pytest --cache-clear "$@" tests $IGNORE_EXPERIMENTAL --executor in-memory
+python -m pytest --cache-clear "$@" tests $IGNORE_EXPERIMENTAL --executor streaming
+python -m pytest --cache-clear "$@" tests $IGNORE_EXPERIMENTAL --executor streaming \
+    --blocksize-mode small

ci/test_cudf_polars_with_rapidsmpf.sh ci/test_cudf_polars_experimental.shci/test_cudf_polars_with_rapidsmpf.sh renamed to ci/test_cudf_polars_experimental.sh

@@ -49,10 +49,8 @@ EXITCODE=0
 trap set_exitcode ERR
 set +e
 
-rapids-logger "Running cudf_polars tests with rapidsmpf"
-
-# Run cudf_polars tests with rapidsmpf using dedicated test runner
-timeout 15m ./ci/run_cudf_polars_with_rapidsmpf_pytests.sh \
+rapids-logger "Running cudf_polars experimental tests (non-ci-blocking)"
+timeout 15m ./ci/run_cudf_polars_experimental_pytests.sh \
     --no-cov \
     --numprocesses=8 \
     --dist=worksteal \

python/cudf_polars/cudf_polars/dsl/expressions/unary.py

@@ -421,7 +421,7 @@ def do_evaluate(
                     1, dtype=column.dtype.plc_type, stream=df.stream
                 )
             else:
-                assert_never(strategy)  # pragma: no cover
+                assert_never(strategy)
 
             if strategy == "mean":
                 return Column(

python/cudf_polars/cudf_polars/dsl/ir.py

@@ -2317,7 +2317,7 @@ def _joiners(
                 plc.copying.OutOfBoundsPolicy.DONT_CHECK,
                 None,
             )
-        assert_never(how)  # pragma: no cover
+        assert_never(how)
 
     @staticmethod
     def _reorder_maps(

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

@@ -27,9 +27,9 @@
 from typing import TYPE_CHECKING, Any, Literal, assert_never
 
 import nvtx
+from rapidsmpf.config import Options, get_environment_variables
 
 import polars as pl
-import polars.testing
 
 import rmm.statistics
 
@@ -768,7 +768,6 @@ def initialize_dask_cluster(run_config: RunConfig, args: argparse.Namespace):  #
 
     if run_config.shuffle != "tasks":
         try:
-            from rapidsmpf.config import Options
             from rapidsmpf.integrations.dask import bootstrap_dask_cluster
 
             bootstrap_dask_cluster(
@@ -1863,12 +1862,13 @@ def run_polars_spmd(
         rmm.mr.CudaAsyncMemoryResource(release_threshold=args.rmm_release_threshold)
     )
     with spmd_execution(
+        rapidsmpf_options=Options(get_environment_variables()),
         executor_options=executor_options,
         engine_options={
             "parquet_options": parquet_options,
             "cuda_stream_policy": run_config.stream_policy,
         },
-    ) as (comm, ctx, engine):
+    ) as engine:
         from cudf_polars.experimental.rapidsmpf.collectives.common import reserve_op_id
         from cudf_polars.experimental.rapidsmpf.frontend.spmd import (
             allgather_polars_dataframe,
@@ -1877,14 +1877,13 @@ def run_polars_spmd(
         def _allgather_result(df: pl.DataFrame) -> pl.DataFrame:
             with reserve_op_id() as op_id:
                 return allgather_polars_dataframe(
-                    comm=comm,
-                    ctx=ctx,
+                    engine=engine,
                     local_df=df,
                     op_id=op_id,
                 )
 
-        rank = comm.rank
-        run_config = dataclasses.replace(run_config, n_workers=comm.nranks)
+        rank = engine.rank
+        run_config = dataclasses.replace(run_config, n_workers=engine.nranks)
         records, plans, validation_failures, query_failures = _run_query_loop(
             benchmark,
             args,

python/cudf_polars/cudf_polars/experimental/rapidsmpf/frontend/core.py

@@ -0,0 +1,94 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES.
+# SPDX-License-Identifier: Apache-2.0
+"""Multi-GPU frontend core."""
+
+from __future__ import annotations
+
+import contextlib
+from typing import Any, Self
+
+import polars as pl
+
+
+class StreamingEngine(pl.GPUEngine):
+    """
+    Base class for multi-GPU Polars engines.
+
+    The engine manages the lifecycle of a streaming execution and can
+    be used as a context manager. On exit, :meth:`shutdown` is called.
+
+    Notes
+    -----
+    The engine must be created and shut down on the same thread. In particular,
+    destruction and context manager exit must occur on the thread that created
+    the instance.
+
+    Parameters
+    ----------
+    nranks
+        Number of ranks (workers or GPUs) in the cluster.
+    executor_options
+        Key/value options forwarded to the streaming executor.
+    engine_options
+        Additional keyword arguments forwarded to
+        :class:`~polars.lazyframe.engine_config.GPUEngine`.
+    exit_stack
+        A :class:`contextlib.ExitStack` whose registered contexts are closed
+        when :meth:`shutdown` is called. If ``None``, an empty stack is created.
+    """
+
+    def __init__(
+        self,
+        *,
+        nranks: int,
+        executor_options: dict[str, Any],
+        engine_options: dict[str, Any],
+        exit_stack: contextlib.ExitStack | None = None,
+    ):
+        self._nranks = nranks
+        self._exit_stack: contextlib.ExitStack | None = (
+            exit_stack or contextlib.ExitStack()
+        )
+        super().__init__(
+            executor="streaming",
+            executor_options=executor_options,
+            **engine_options,
+        )
+
+    @property
+    def nranks(self) -> int:
+        """
+        Number of ranks (for example GPUs or workers) in the cluster.
+
+        Local execution without a cluster returns 1.
+
+        Returns
+        -------
+        Number of ranks.
+        """
+        return self._nranks
+
+    def shutdown(self) -> None:
+        """
+        Shut down engine and release all owned resources.
+
+        Idempotent: safe to call more than once. Must be called on the same
+        thread that created the engine.
+        """
+        if self._exit_stack is None:
+            return  # already shut down
+        try:
+            self._exit_stack.close()
+        finally:
+            self._exit_stack = None
+            self.device = None
+            self.memory_resource = None
+            self.config = {}
+
+    def __enter__(self) -> Self:
+        """Enter the context manager, returning ``self``."""
+        return self
+
+    def __exit__(self, *_: object) -> None:
+        """Exit the context manager, calling :meth:`shutdown`."""
+        self.shutdown()