RapidsMPF SPMD mode

[madsbk]

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

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

[TomAugspurger]

There's quite a bit of duplication between this and evaluate_pipeline, though I haven't looked closely to see what the difference is.

[madsbk]

I prefer to keep them separate for now. The long-term goal is to remove most of evaluate_pipeline, and eventually drop Dask entirely once this implementation matures.

[TomAugspurger]

Is there any chance that a user could hit this error, say by misconfiguring their engine? If so, I think we'd want a friendlier error message, ideally telling them what they did wrong.

[wence-]

Yeah, I think we shouldn't mention pytest here.

[madsbk]

renamed to:

        raise RuntimeError(
            "spmd_execution() requires the rrun launcher. "
            "Launch your script with `rrun -n <nproc> python your_script.py` "
            "to enable SPMD execution."
        )

[TomAugspurger]

Maybe note that "raise_on_fail", "memory_resource", "executor" are reserved, like you did for exeuctor_options?

And just confirming, we do want to hardcode raise_on_fail here? I understand why the others have to be hardcoded.

[madsbk]

Why do we hardcode raise_on_fail?
I have removed it here

[wence-]

tiny fixes

[wence-]

Yeah, I think we shouldn't mention pytest here.

[madsbk]

/merge