[Multi-GPU Polars] Ray mode in PDSH benchmarks by madsbk · Pull Request #21811 · NVIDIA/cudf

madsbk · 2026-03-17T08:22:40Z

Update the PDSH benchmarks to support --cluster=ray runs.

Run using something like:

python \
  python/cudf_polars/cudf_polars/experimental/benchmarks/pdsh.py \
  --executor=streaming \
  --suffix="" \
  --spill-device=0.5 \
  --shuffle=rapidsmpf \
  --runtime=rapidsmpf \
  --stream-policy=pool \
  --no-print-results --no-summarize \
  --iterations=10 \
  --path /datasets/datasets/tpch-rs/scale-10 \
  --cluster ray \
  1

TomAugspurger

Trying this out locally, I hit

Traceback (most recent call last):
  File "/home/nfs/toaugspurger/gh/rapidsai/cudf/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsh.py", line 1804, in <module>
    run_polars(PDSHQueries, args)
    ~~~~~~~~~~^^^^^^^^^^^^^^^^^^^
  File "/home/nfs/toaugspurger/gh/rapidsai/cudf/python/cudf_polars/cudf_polars/experimental/benchmarks/utils.py", line 1766, in run_polars
    run_polars_ray(
    ~~~~~~~~~~~~~~^
        benchmark,
        ^^^^^^^^^^
    ...<5 lines>...
        validation_files,
        ^^^^^^^^^^^^^^^^^
    )
    ^
  File "/home/nfs/toaugspurger/gh/rapidsai/cudf/python/cudf_polars/cudf_polars/experimental/benchmarks/utils.py", line 1955, in run_polars_ray
    with ray_execution(
         ~~~~~~~~~~~~~^
        executor_options=executor_options,
        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
        engine_options=engine_options,
        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
    ) as (ray_client, engine):
    ^
  File "/home/nfs/toaugspurger/gh/rapidsai/cudf/python/cudf_polars/cudf_polars/experimental/rapidsmpf/frontend/ray.py", line 615, in ray_execution
    root_ucxx_address_as_bytes = ray.get(rank_actors[0].setup_root.remote())
  File "/raid/toaugspurger/envs/gh/rapidsai/cudf/lib/python3.13/site-packages/ray/_private/auto_init_hook.py", line 22, in auto_init_wrapper
    return fn(*args, **kwargs)
  File "/raid/toaugspurger/envs/gh/rapidsai/cudf/lib/python3.13/site-packages/ray/_private/client_mode_hook.py", line 104, in wrapper
    return func(*args, **kwargs)
  File "/raid/toaugspurger/envs/gh/rapidsai/cudf/lib/python3.13/site-packages/ray/_private/worker.py", line 2981, in get
    values, debugger_breakpoint = worker.get_objects(
                                  ~~~~~~~~~~~~~~~~~~^
        object_refs, timeout, use_object_store=_use_object_store
        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
    )
    ^
  File "/raid/toaugspurger/envs/gh/rapidsai/cudf/lib/python3.13/site-packages/ray/_private/worker.py", line 1014, in get_objects
    raise value
ray.exceptions.ActorDiedError: The actor died because of an error raised in its creation task, ray::RankActor.__init__() (pid=2859867, ip=10.33.227.162, actor_id=dfbdb6136704c13c8ef33ba101000000, repr=<cudf_polars.experimental.rapidsmpf.frontend.ray.FunctionActorManager._create_fake_actor_class.<locals>.TemporaryActor object at 0x7fd68c0df770>)
    ...<4 lines>...
    )
RuntimeError: The actor with name RankActor failed to import on the worker. This may be because needed library dependencies are not installed in the worker environment:

IIRC, you might have hit that as well. Did you come to a conclusion on what caused it?

TomAugspurger · 2026-03-17T13:23:25Z

+        )
+    executor_options = get_executor_options(run_config, benchmark=benchmark)
+    # "runtime", "cluster" are reserved — ray_execution sets them
+    executor_options.pop("runtime", None)


Validate that the user didn't pass something here? I guess the default is tasks, which makes this not straightforward :/

I say we just raise an error? The Ray frontend requires the "rapidsmpf" runtime for sure.

The plan is to remove the runtime option altogether. The runtime is implied by the chosen cluster; you simply set --cluster <single | ray | dask | spmd> (we may rename this to --frontend).

Currently, --cluster <spmd | ray> ignores --runtime, but it still affects performance. For example:

Sampling is only enabled with runtime="rapidsmpf".

The native C++ Parquet reader is only enabled with runtime="rapidsmpf".

For now, let’s ignore the --runtime option, but this is another strong reason to standardize all benchmark options and logging!

rjzamora · 2026-03-17T15:27:57Z

        default=None,
        type=str,
-        choices=["single", "distributed", "spmd"],
+        choices=["single", "distributed", "spmd", "ray"],


Just a note (nothing to do in this PR): I'd like to see these options become ["single", "dask", "spmd", "ray"] for 26.06

rjzamora · 2026-03-17T15:29:33Z

+        raise NotImplementedError(
+            "--collect-traces is not yet supported with --cluster ray."
+        )


Can you describe the blocker for this? I use this option extensively when I debug performance.

The blocker is mainly that I do not yet have the full picture. I want to take this design one step at a time.
Once the new Dask frontend and its benchmark are merged, I plan to work on:

Standardizing all configuration options by defining a clear policy for how cudf-polars options propagate to RapidsMPF, and which options must be configured directly in RapidsMPF rather than through cudf-polars.

Standardizing all benchmark options and logging.

rjzamora · 2026-03-17T15:33:27Z

+        raise NotImplementedError(
+            "--rmm-async is not yet supported with --cluster ray."
+        )


Are we just always using the async allocator with Ray?

Are we accounting for --spill-device anywhere? That option is pretty critical for tuning performance in our nightlies.

Currently, you would have to set RAPIDSMPF_SPILL_DEVICE_LIMIT.

Are we just always using the async allocator with Ray?

yes

rjzamora · 2026-03-17T15:34:15Z

+        )
+    executor_options = get_executor_options(run_config, benchmark=benchmark)
+    # "runtime", "cluster" are reserved — ray_execution sets them
+    executor_options.pop("runtime", None)


I say we just raise an error? The Ray frontend requires the "rapidsmpf" runtime for sure.

rjzamora · 2026-03-17T15:37:31Z

+            args,
+            run_config,
+            engine,
+            None,


This is the client argument, right? Do we need _run_query_loop/etc to handle the Ray client to do collect traces and statistics?

Yes, we would need to either implement a general run method on the ray client, or implement specific logging methods.

…nchmarks

madsbk · 2026-03-18T08:34:39Z

Trying this out locally, I hit
...
IIRC, you might have hit that as well. Did you come to a conclusion on what caused it?

Yes, I ran into that as well when using the pip installation. In a conda environment, it seems to work.
It appears to be related to the ucxx pip package failing to load on the workers. I am still investigating it.

UPDATE: Hopefully, this fixes the issue: rapidsai/rapidsmpf#921

rjzamora

Looks good. Approving since I'm expecting more changes to follow anyway.

My only remaining concern is that the spill-device argument is ignored, and so users may try to tune this value and not understand why it's having no effect.

madsbk · 2026-03-18T14:46:20Z

/merge

Ray mode benchmark

8bc44a2

madsbk self-assigned this Mar 17, 2026

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

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

github-project-automation Bot added this to cuDF Python Mar 17, 2026

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

madsbk mentioned this pull request Mar 17, 2026

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

Closed

11 tasks

madsbk marked this pull request as ready for review March 17, 2026 10:21

madsbk requested a review from a team as a code owner March 17, 2026 10:21

madsbk requested review from galipremsagar and rjzamora March 17, 2026 10:21

TomAugspurger approved these changes Mar 17, 2026

View reviewed changes

rjzamora reviewed Mar 17, 2026

View reviewed changes

Merge branch 'main' of github.com:rapidsai/cudf into rapidsmpf-ray-be…

961fbcb

…nchmarks

NVIDIA deleted a comment from copy-pr-bot Bot Mar 18, 2026

cleaup

7bcf51c

madsbk mentioned this pull request Mar 18, 2026

Remote ray cluster #21838

Closed

style

eef7b87

madsbk requested a review from rjzamora March 18, 2026 13:14

rjzamora approved these changes Mar 18, 2026

View reviewed changes

rapids-bot Bot merged commit 8fc7f3b into NVIDIA:main Mar 18, 2026
163 of 166 checks passed

github-project-automation Bot moved this from In Progress to Done in cuDF Python Mar 18, 2026

madsbk deleted the rapidsmpf-ray-benchmarks branch March 18, 2026 14:46

Uh oh!

Conversation

madsbk commented Mar 17, 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!

madsbk Mar 18, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

madsbk commented Mar 18, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

rjzamora left a comment

Choose a reason for hiding this comment

Uh oh!

madsbk commented Mar 18, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

madsbk Mar 18, 2026 •

edited

Loading

madsbk commented Mar 18, 2026 •

edited

Loading