RapidsMPF SPMD mode in PDSH benchmarks - #21652
Conversation
madsbk
added
improvement
github-actions
Bot
added
Python
This PR adds an `spmd` cluster type to the streaming executor, enabling multi-GPU execution via the `rrun` launcher without requiring Dask.
With this mode, users can run the same Python program on multiple ranks (SPMD style) and use RapidsMPF collectives directly from Python. Each rank processes its own slice of data and coordinates through collectives such as `allgather`.
### Example
Launch with:
```
rrun -n 4 python script.py
```
```python
from cudf_polars.experimental.rapidsmpf.spmd import (
allgather_polars_dataframe,
spmd_execution,
)
from cudf_polars.experimental.rapidsmpf.collectives.common import reserve_op_id
with spmd_execution() as (ctx, engine):
rank = ctx.comm().rank
# Each rank holds its own slice of the data
local = pl.LazyFrame({"a": [rank], "b": [rank * 10]})
local_result = (
local.group_by("a")
.agg(pl.col("b").sum())
.collect(engine=engine)
)
# Gather results from all ranks
with reserve_op_id() as op_id:
full_result = allgather_polars_dataframe(
ctx=ctx, local_df=local_result, op_id=op_id
) # all ranks gets an identical copy of `full_result`
```
### Running the pdsh benchmarks
Update: the benchmark update has been moved to a follow-up PR: #21652
Authors:
- Mads R. B. Kristensen (https://github.com/madsbk)
Approvers:
- Lawrence Mitchell (https://github.com/wence-)
URL: #21619
6805d3a to
e9463fe
Compare
e9463fe to
55825f0
Compare
| ) -> None: | ||
| """Run benchmark queries using SPMD execution via the ``rrun`` launcher.""" | ||
| if run_config.collect_traces: | ||
| warnings.warn( |
There was a problem hiding this comment.
I'd say error rather than warn.
More generally, is there something that prevents us from making --collect-traces (and I'm assuming cudf-polars' structured logging) work with spmd? Single node / process, everything is easy. For multi-node with Dask, we need to call some logging config at the start and collect the logs at the end (we could also configure logging to stream the logs, but that's immaterial). Presumably we can do something similar with spmd, using rank 0 to collect the logs?
There was a problem hiding this comment.
Yes, I think this is possible, but let’s wait until we have the full picture related to the new Ray frontend.
I have changed the warning to a NotImplementedError.
TomAugspurger
left a comment
TomAugspurger
left a comment
There was a problem hiding this comment.
The diff here looks pretty large, but IIUC, the primary changes are
- Refactoring parts of
run_polarsinto_run_query_loop, essentially unchanged - Adding a new
run_polars_smpd, with a branch to take that fromrun_polarsif specified - Some minor adjustments (using ternary expressions, moving some things like dask client initialization around)
If that understanding is correct, then I think my only question (non-blocking) is around the structured logging.
Yes, your understanding is correct. 🙂 Most of the differences are just refactoring so that at least some of the code can be reused in
|
|
/merge |
3 participants
Update the PDSH benchmarks to support
--cluster=spmdruns.Run using something like:
rrun -n 2 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 spmd \ 1