Note that I am seeing errors when I run the benchmark with --shuffle "explicit-comms" --split_out 2 --sort, and the cause seems to be that dask is using a shallow copy in DataFrame.drop, and this "_partitions"-column drop is somehow not working right after an explicit-comms shuffle.

I'm not exactly sure why the "_partitions" column is not getting dropped properly here (I am struggling to create a simple-enough reproducer to properly explore). However, I suspect that the custom round-tripping of of the input dataframe from a persisted list of futures is somehow running us into an undefined edge case.

I see a (possibly different) error:

python -m dask_cuda.benchmarks.local_cudf_groupby --shuffle explicit-comms  --split_out 2 --sort -d 0,1,3,4   -p tcp
=>   File "/root/dask/dask/dataframe/utils.py", line 429, in check_matching_columns
    raise ValueError(
ValueError: The columns in the computed data do not match the columns in the provided metadata
Order of columns does not match

With the columns being:

Index(['key', 'int64___count', 'int64___max', 'float64___count',
       'float64___sum'],
      dtype='object')

Index(['key', 'int64___max', 'int64___count', 'float64___sum',
       'float64___count'],
      dtype='object')

This occurs also with --shuffle tasks

My crystal-ball says that somewhere the columns for aggregation are being constructed through a set, which doesn't provided run-to-run consistent ordering, but I can't find it.

This occurs also with --shuffle tasks

Oh no! You see this same error for the command you shared above, but with --shuffle tasks?

yes (not always), more likely with more devices

Thanks for pointing this out @wence-. It seems like we need to tweak the groupby-aggregation functions (defined in dask_cudf) to output intermediate DataFrame results with sorted column names. Otherwise the ordering is non-deterministic, and we open ourselves up to failed metadata-enforcement checks in down-stream operations :/

Update: rapidsai/cudf#11836 should resolve this problem in dask_cudf.