[Multi-GPU Polars] Unify streaming engine options by madsbk · Pull Request #21930 · rapidsai/cudf

madsbk · 2026-03-25T15:55:37Z

Introduces StreamingOptions, a single typed dataclass that unifies all RapidsMPF, executor, and engine configuration in one place.

SPMD, Ray, and the upcoming Dask frontend now expose a from_options classmethod:

engine = SPMDEngine.from_options(opts)
engine = RayEngine.from_options(opts)
engine = DaskEngine.from_options(opts)  # upcoming

This is the first step toward standardizing the many configuration options we currently support.

This is a single commit intended to be cherry-picked onto `main` once rapidsai#21930 has been merged.

TomAugspurger

Did you consider nesting the options under their category (e.g. streaming_options.rapidsmpf.num_streaming_threads instead of streaming_options.num_streaming_threads?

I like this because

It avoids an risk of a name clash between options from different categories
It simplifies usage of the options later on (you don't need the category when converting to rapidsmpf Options, for example).

TomAugspurger · 2026-03-31T14:51:32Z

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

+    ----------
+    num_streaming_threads
+        Threads used to execute coroutines.
+        Env: ``RAPIDSMPF_NUM_STREAMING_THREADS``.


IMO, any all environment variables should be prefixed by CUDF_POLARS_.

I'd also recommend the convention CUDF_POLARS__<CATEGORY>__<name> to avoid any possibility of a name clash between two options with the same name from different categories.

Do we really want to have duplicate options like CUDF_POLARS_LOG and RAPIDSMPF_LOG? I think it might be better to accept that RapidsMPF-specific options use a different prefix?

Just to be clear, the environment variable would be CUDF_POLARS_RAPIDSMPF_LOG, not CUDF_POLARS_LOG.

My preference would be to prefix everything. Then there's no guessing about the environment variable names, it's always CUDF_POLARS_<name> (or CUDF_POLARS__<CATEGORY>__<NAME> if we nest things).

Special casing rapidsmpf variables to not have the prefix CUDF_POLARS prefix isn't the worst thing, just something to remember.

So, if the user has RAPIDSMPF_NUM_STREAMING_THREADS set in their environment, that value will be used unless they have CUDF_POLARS_RAPIDSMPF_NUM_STREAMING_THREADS set. Right?

I think that would be the ultimate effect, given that (IIUC) it would eventually fall back to the rapidsmpf default, which is to read RAPIDSMPF_NUM_STREAMING_THREADS from the environment.

TomAugspurger · 2026-03-31T14:54:02Z

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

+    # ---- RapidsMPF runtime ----
+    num_streaming_threads: int | _Unspecified = _opt("rapidsmpf")
+    num_streams: int | _Unspecified = _opt("rapidsmpf")
+    log: str | _Unspecified = _opt("rapidsmpf")


This should be the literal levels allowed.

Suggested change

log: str | _Unspecified = _opt("rapidsmpf")

log: Literal["NONE", "PRINT", "WARN", "INFO", "DEBUG", "TRACE"] | _Unspecified = _opt("rapidsmpf")

TomAugspurger · 2026-03-31T14:55:30Z

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

+    spill_device_limit: str | _Unspecified = _opt("rapidsmpf")
+    periodic_spill_check: str | _Unspecified = _opt("rapidsmpf")
+    # ---- Executor ----
+    rapidsmpf_py_executor_max_workers: int | _Unspecified = _opt("executor")


It's a bit strange to have an option with the "rapidsmpf" prefix under the executor category.

TomAugspurger · 2026-03-31T14:56:18Z

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

+        ``RAPIDSMPF_*`` environment variable (if set); otherwise the rapidsmpf
+        C++ library applies its own built-in default.
+        """
+        from rapidsmpf.config import Options, get_environment_variables


We should be able to import rapidsmpf at the top of the file here.

TomAugspurger · 2026-03-31T15:21:23Z

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

+        env = get_environment_variables()
+
+        opts: dict[str, str] = {}
+        for f in dataclasses.fields(self):


There's also dataclasses.asdict(self), which would avoid the need for a getattr later on.

I guess we'd lose the category which is on the field rather than the fields value, but we nest the options then that'd be present (and you'd just do something like dataclasses.asdict(self.rapidsmpf)).

TomAugspurger · 2026-03-31T15:22:31Z

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

+        known = {f.name for f in dataclasses.fields(cls)}
+        unknown = d.keys() - known
+        if unknown:
+            raise TypeError(
+                f"StreamingOptions.from_dict() got unknown field(s): "
+                f"{', '.join(sorted(unknown))}"
+            )


Why not let the dataclass's __init__ handle this for us?

TomAugspurger · 2026-03-31T15:23:20Z

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

+        return cls(**kwargs)
+
+    @classmethod
+    def from_argparse(cls, args: argparse.Namespace) -> StreamingOptions:


This shouldn't be part of the public API IMO.

If we want it (presumably for our benchmarks file) then let's make it private.

I was thinking this could be generally useful. Most applications built on cudf-polars would likely want to support these kinds of command-line arguments, right?

Mmm I'm not sure. In my experience (primarily deploying services on Kubernetes) most configuration is done through some kind of centralized management system that places the configuration files / environment variables before the application starts up.

If we're able to use it internally (e.g. for benchmarks) then great. But I'd caution against advertising this as a stable, public API.

madsbk self-assigned this Mar 25, 2026

madsbk added improvement Improvement / enhancement to an existing function non-breaking Non-breaking change labels Mar 25, 2026

github-actions bot added Python Affects Python cuDF API. cudf-polars Issues specific to cudf-polars labels Mar 25, 2026

github-project-automation bot added this to cuDF Python Mar 25, 2026

GPUtester moved this to In Progress in cuDF Python Mar 25, 2026

madsbk mentioned this pull request Mar 25, 2026

[FEA] New user-facing multi-GPU Polars API #21769

Open

14 tasks

madsbk force-pushed the StreamingOptions branch from 872e49a to 54be0e8 Compare March 25, 2026 21:18

madsbk marked this pull request as ready for review March 26, 2026 06:51

madsbk requested a review from a team as a code owner March 26, 2026 06:51

madsbk requested review from TomAugspurger and rjzamora March 26, 2026 06:51

madsbk force-pushed the StreamingOptions branch from 54be0e8 to cd99655 Compare March 26, 2026 06:51

rapidsai deleted a comment from copy-pr-bot bot Mar 26, 2026

[Multi-GPU Polars] Unify streaming engine options

c995fc1

madsbk force-pushed the StreamingOptions branch from cd99655 to c995fc1 Compare March 26, 2026 06:53

madsbk added a commit to madsbk/cudf that referenced this pull request Mar 26, 2026

[Multi-GPU Polars] Split PDSH utils into legacy and new frontend paths

79fd5f1

This is a single commit intended to be cherry-picked onto `main` once rapidsai#21930 has been merged.

madsbk mentioned this pull request Mar 26, 2026

[Multi-GPU Polars] Split PDSH utils into legacy and new frontend paths #21941

Draft

Merge branch 'main' into StreamingOptions

e9501af

TomAugspurger reviewed Mar 31, 2026

View reviewed changes

madsbk added 2 commits April 2, 2026 00:09

Merge branch 'main' into StreamingOptions

c5802b6

Merge branch 'main' into StreamingOptions

b488f42

	log: str \| _Unspecified = _opt("rapidsmpf")
	log: Literal["NONE", "PRINT", "WARN", "INFO", "DEBUG", "TRACE"] \| _Unspecified = _opt("rapidsmpf")

Conversation

madsbk commented Mar 25, 2026

Uh oh!

TomAugspurger left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants