Update explicit-comms for dask-expr support

[rjzamora]

Makes a few small changes to explicit-comms to support dask-expr.

EDIT: The changes are no longer "small".

[pentschev]

Apart from the linting issues, looks good to me. I think it should be good to go as soon as CI passes. Also tagging @madsbk for awareness.

[quasiben]

I gave this a try and I'm seeing errors. It's been a while since I've run this query so it could be user error:

python local_cudf_merge.py -p ucx -d 0,1,2,3,4,5,6,7 -c 10_000_000

2024-03-20 12:49:49,925 - distributed.shuffle._scheduler_plugin - WARNING - Shuffle 637a8bdfdc46dc6083f254fce7023e81 initialized by task ('shuffle-transfer-637a8bdfdc46dc6083f254fce7023e81', 0) executed on worker ucx://127.0.0.1:50385
2024-03-20 12:49:50,535 - distributed.worker - WARNING - Compute Failed
Key:       ('shuffle-transfer-637a8bdfdc46dc6083f254fce7023e81', 6)
Function:  shuffle_transfer
args:      (              key  shuffle  payload  _partitions
0        60000000        1  9459992            1
1        60000001        5  5005508            5
2        60000002        3  3183193            3
3        60000003        1   536845            1
4        60000004        3  8484645            3
...           ...      ...      ...          ...
9999995  69999995        3  9745632            3
9999996  69999996        5  8371927            5
9999997  69999997        3  8066013            3
9999998  69999998        1  3762504            1
9999999  69999999        7  9411065            7

[10000000 rows x 4 columns], '637a8bdfdc46dc6083f254fce7023e81', 6, 8, '_partitions', Empty DataFrame
Columns: [key, shuffle, payload, _partitions]
Index: [], {0, 1, 2, 3, 4, 5, 6, 7}, True, True)
kwargs:    {}
Exception: "RuntimeError('P2P shuffling 637a8bdfdc46dc6083f254fce7023e81 failed during transfer phase')"

And with explicit-comms:

python local_cudf_merge.py -p ucx -d 0,1,2,3,4,5,6,7 -c 10_000_000 -b explicit-comms

2024-03-20 12:51:47,067 - distributed.worker.state_machine - WARNING - Async instruction for <Task cancelled name="execute(('shuffle-transfer-637a8bdfdc46dc6083f254fce7023e81', 5))" coro=<Worker.execute() done, defined at /datasets/bzaitlen/miniconda3/envs/dask-expr-rapid
s-24.06/lib/python3.10/site-packages/distributed/worker_state_machine.py:3615>> ended with CancelledError
Traceback (most recent call last):
  File "/datasets/bzaitlen/miniconda3/envs/dask-expr-rapids-24.06/lib/python3.10/site-packages/distributed/shuffle/_core.py", line 494, in handle_transfer_errors
    yield
  File "/datasets/bzaitlen/miniconda3/envs/dask-expr-rapids-24.06/lib/python3.10/site-packages/distributed/shuffle/_shuffle.py", line 79, in shuffle_transfer
    return get_worker_plugin().add_partition(
  File "/datasets/bzaitlen/miniconda3/envs/dask-expr-rapids-24.06/lib/python3.10/site-packages/distributed/shuffle/_worker_plugin.py", line 346, in add_partition
    return shuffle_run.add_partition(
  File "/datasets/bzaitlen/miniconda3/envs/dask-expr-rapids-24.06/lib/python3.10/site-packages/distributed/shuffle/_core.py", line 343, in add_partition
    shards = self._shard_partition(data, partition_id)
  File "/datasets/bzaitlen/miniconda3/envs/dask-expr-rapids-24.06/lib/python3.10/site-packages/distributed/shuffle/_shuffle.py", line 521, in _shard_partition
    out = split_by_worker(
  File "/datasets/bzaitlen/miniconda3/envs/dask-expr-rapids-24.06/lib/python3.10/site-packages/distributed/shuffle/_shuffle.py", line 339, in split_by_worker
    t = to_pyarrow_table_dispatch(df, preserve_index=True)
  File "/datasets/bzaitlen/miniconda3/envs/dask-expr-rapids-24.06/lib/python3.10/site-packages/dask/utils.py", line 772, in __call__
    meth = self.dispatch(type(arg))
  File "/datasets/bzaitlen/miniconda3/envs/dask-expr-rapids-24.06/lib/python3.10/site-packages/dask/utils.py", line 766, in dispatch
    raise TypeError(f"No dispatch for {cls}")
TypeError: No dispatch for <class 'cudf.core.dataframe.DataFrame'>

[rjzamora]

I gave this a try and I'm seeing errors.

Hmm - I actually am expecting errors in this PR. (I just realized that I never addressed the fact that we literally patch/wrap a bunch of deprecated shuffling code to automatically dispatch to explicit-comms). That said, I wasn't expecting a local-cudf-merge problem, so I'll try to reproduce.

[rjzamora]

It looks like explicit-comms is going to be a pretty-big headache to migrate over to dask-expr :/

I have the basic benchmarks working locally, but there is no longer a clean way to support the "explicit-comms" config that we used to support by monkey-patching dask.dataframe. Dask-expr no longer uses a monolithic function (like rearrange_by_coumn) to inject a shuffle operation for sorting and joining. Instead, these operations are now a composition of abstract expressions. It may be necessary to deprecate/prohibit this behavior when query-planning is enabled (for now).

[pentschev]

It looks like explicit-comms is going to be a pretty-big headache to migrate over to dask-expr :/

I have the basic benchmarks working locally, but there is no longer a clean way to support the "explicit-comms" config that we used to support by monkey-patching dask.dataframe. Dask-expr no longer uses a monolithic function (like rearrange_by_coumn) to inject a shuffle operation for sorting and joining. Instead, these operations are now a composition of abstract expressions. It may be necessary to deprecate/prohibit this behavior when query-planning is enabled (for now).

Thanks for debugging this Rick. I must admit I'm not of any help when it comes to either Dask Dataframe or explicit-comms, on the other hand @madsbk may have great ideas here.

[madsbk]

Hmm, maybe it is time to deprecate the monkey patching of dask and tell users to call shuffle() directly? And maybe create a basic merge() function that uses shuffle() people can use?

Do you guys know many are using explicit-comms?

[rjzamora]

Hmm, maybe it is time to deprecate the monkey patching of dask and tell users to call shuffle() directly? And maybe create a basic merge() function that uses shuffle() people can use?

Yes, I think this route makes the most sense for now. I have a feeling that there are clean ways to introduce an "official" explicit-comms shuffle using the new API, but I'd like to take some time to get it right.

Do you guys know many are using explicit-comms?

The only "official" user that I personally know of is nemo curator. However, (1) they are not using the monkey patching approach, and (2) they do not develop against the latest RAPIDS version. cc @VibhuJawa @ayushdg

[quasiben]

I'm not aware of any external usage. It is useful to have for benchmarking and continue exploration of shuffling generally. With that said, I agree with @rjzamora we should take our time getting it right

[rjzamora]

Generating a collection in this way puts us in a strange edge case where dask_cudf may never get imported on the worker, and so some type-based dispatching may fail. The only way to avoid this possibility is to import dask_cudf before we deploy the cluster, or to import dask_cudf here.

[rjzamora]

CI failures without the hacky test-skip logic are due to a bad/old dask-expr version being pulled (dask/distributed#8574)

[rjzamora]

@pentschev - The upstream nightly situation should be resolved now. You approved this PR a while back when it looked pretty different. Did you have time for another pass?

[rjzamora]

Note on these changes: For 24.06, we want to explicitly test with both DASK_DATAFRAME__QUERY_PLANNING=True and DASK_DATAFRAME__QUERY_PLANNING=False.

[pentschev]

Thanks @rjzamora , left a few minor requests but otherwise it looks good.

[pentschev]

It seems you're not addressing this for local_cudf_groupby and local_cudf_merge, are explicit-comms still supported for those? If not, I'd suggest moving this into parse_benchmark_args just before the function returns, and as such you'd cover all three cases.

[rjzamora]

Right - sorry. I'll use the same message for local_cudf_groupby. This is local_cudf_merge, so I assume you are also asking why I am not checking for this in local_cudf_shuffle? If so, it's because that benchmark does not use the dask config (it uses dask_cuda.explicit_comms.dataframe.shuffle.shuffle directly).

[pentschev]

Right - sorry. I'll use the same message for local_cudf_groupby.

I'd still suggest simply moving this to parse_benchmarks_args, this will cover all benchmarks and avoid code duplication, which is essentially why we wrote it that way to begin with.

This is local_cudf_merge, so I assume you are also asking why I am not checking for this in local_cudf_shuffle?

Yes, sorry, that is what I meant. 😅

If so, it's because that benchmark does not use the dask config (it uses dask_cuda.explicit_comms.dataframe.shuffle.shuffle directly).

But would using it directly still work with query-planning enabled? If not then moving the exception above to parse_benchmarks_args would also cover it and avoid more boilerplate code.

[rjzamora]

I'd still suggest simply moving this to parse_benchmarks_args, this will cover all benchmarks and avoid code duplication, which is essentially why we wrote it that way to begin with.

Ah - I see what you mean. Using dask_cuda.explicit_comms.dataframe.shuffle.shuffle directly works fine, so it would still be good to allow that option for the shuffle benchmark somehow.

[pentschev]

Ok, then how about still moving this to parse_benchmarks_args but adding another kwarg to it that would still allow running select benchmarks (e.g., local_cudf_shuffle) with planning-query enabled. Something like below:

def parse_benchmark_args(description="Generic dask-cuda Benchmark", args_list=[], *, force_explicit_comms=False):
...

    if args.multi_node and len(args.hosts.split(",")) < 2:
        raise ValueError("--multi-node requires at least 2 hosts")

    # Raise error early if "explicit-comms" is not allowed
    if (
        args.backend == "explicit-comms"
        and dask.config.get(
            "dataframe.query-planning",
            None,
        )
        is not False
        and force_explicit_comms is False
    ):
        raise NotImplementedError(
            "The 'explicit-comms' config is not yet supported when "
            "query-planning is enabled in dask. Please use the legacy "
            "dask-dataframe API (set the 'dataframe.query-planning' "
            "config to `False` before executing).",
        )

    return args

Then local_cudf_{groupby,merge} remain as they were and for local_cudf_shuffle you simply add force_explicit_comms=True to this line.

WDYT?

[rjzamora]

Okay, sounds good. Pushed something very similar to your solution just now (just used a check_explicit_comms arg instead of force_explicit_comms).

[pentschev]

Thanks Rick, everything looks good. I'll let you hit merge if there's nothing else here, or in case someone else still wants to review this.

[pentschev]

Unfortunately, 15 minutes was insufficient, explicit comms tests are some of the longer running ones. It also took me a few minutes to figure out what exactly happened, and to simplify that in the future I've opened #1330 to make this debugging process easier in the future.

[rjzamora]

/merge