Update explicit-comms for dask-expr support #1323

rjzamora · 2024-03-20T19:10:33Z

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

EDIT: The changes are no longer "small".

…it-comms

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.

dask_cuda/utils.py

quasiben · 2024-03-20T19:54:35Z

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 · 2024-03-20T20:16:11Z

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 · 2024-03-21T00:20:07Z

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 · 2024-03-21T06:59:40Z

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 · 2024-03-21T07:15:09Z

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 · 2024-03-21T13:48:07Z

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 · 2024-03-21T14:58:12Z

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 · 2024-03-21T19:23:45Z

dask_cuda/benchmarks/local_cudf_merge.py

    # Setting a seed that triggers max amount of comm in the two-GPU case.
    if gpu:
        import cupy as xp

        import cudf as xdf
+        import dask_cudf  # noqa: F401


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.

dask_cuda/tests/test_explicit_comms.py

dask_cuda/explicit_comms/dataframe/shuffle.py

rjzamora · 2024-03-22T14:23:32Z

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

…it-comms

rjzamora · 2024-04-02T17:41:31Z

@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 · 2024-04-02T18:22:13Z

ci/test_python.sh

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.

ci/test_python.sh

pentschev · 2024-04-02T19:42:06Z

dask_cuda/benchmarks/local_cudf_merge.py

+    # 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
+    ):
+        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).",
+        )


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.

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).

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.

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.

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?

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

…nto expr-explicit-comms

ci/test_python.sh

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 · 2024-04-03T07:29:40Z

ci/test_python.sh

+UCXPY_IFNAME=eth0 \
+UCX_WARN_UNUSED_ENV_VARS=n \
+UCX_MEMTYPE_CACHE=n \
+timeout 30m pytest \


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 · 2024-04-03T23:10:58Z

/merge

rjzamora added 5 commits March 18, 2024 10:21

relax type-check

85c64d4

work toward dask-expr support in dask-cuda

be984be

Merge remote-tracking branch 'upstream/branch-24.06' into expr-explic…

46ae012

…it-comms

remove skip

c11c47c

use futuers_of

dc32127

rjzamora added bug Something isn't working 2 - In Progress Currently a work in progress non-breaking Non-breaking change labels Mar 20, 2024

rjzamora self-assigned this Mar 20, 2024

rjzamora requested a review from a team as a code owner March 20, 2024 19:10

github-actions bot added the python python code needed label Mar 20, 2024

pentschev approved these changes Mar 20, 2024

View reviewed changes

dask_cuda/utils.py Outdated Show resolved Hide resolved

rjzamora marked this pull request as draft March 20, 2024 20:12

rjzamora added 3 commits March 20, 2024 13:53

fix p2p problem

ad2ea62

format

107e31a

test commit

71b9956

update ci and add NotImplementedError

0f8bb29

github-actions bot added the ci label Mar 21, 2024

rjzamora added 5 commits March 21, 2024 08:07

add NotImplementedError to local_cudf_merge.py

d04c755

use from_delayed - TODO: cleanup

1b6ef7d

remove _make_collection - isn't needed

fcc0e04

use prefix kwarg - comming soon to dask-expr anyway

69421df

make sure batchsize is used in key

c79b1f1

rjzamora commented Mar 21, 2024

View reviewed changes

rjzamora added 2 commits March 21, 2024 17:22

skip for dask<2024.2.1

68079bb

more tweaks

af6aa67

rjzamora commented Mar 22, 2024

View reviewed changes

dask_cuda/tests/test_explicit_comms.py Outdated Show resolved Hide resolved

rjzamora commented Mar 22, 2024

View reviewed changes

dask_cuda/explicit_comms/dataframe/shuffle.py Outdated Show resolved Hide resolved

rjzamora added 2 commits April 2, 2024 08:30

Merge remote-tracking branch 'upstream/branch-24.06' into expr-explic…

8c35fba

…it-comms

remove workaround logic

4a93bc9

rjzamora marked this pull request as ready for review April 2, 2024 16:55

rjzamora requested a review from a team as a code owner April 2, 2024 16:55

rjzamora commented Apr 2, 2024

View reviewed changes

Merge branch 'branch-24.06' into expr-explicit-comms

a394334

pentschev requested changes Apr 2, 2024

View reviewed changes

AyodeAwe approved these changes Apr 2, 2024

View reviewed changes

rjzamora added 2 commits April 2, 2024 14:44

respond to code review

92b496c

Merge branch 'expr-explicit-comms' of github.com:rjzamora/dask-cuda i…

206101f

…nto expr-explicit-comms

pentschev reviewed Apr 3, 2024

View reviewed changes

ci/test_python.sh Outdated Show resolved Hide resolved

Increase timeout of explicit comms tests

48e6615

pentschev approved these changes Apr 3, 2024

View reviewed changes

Merge branch 'branch-24.06' into expr-explicit-comms

5dcd2f3

rjzamora added 5 - Ready to Merge Testing and reviews complete, ready to merge and removed 2 - In Progress Currently a work in progress labels Apr 3, 2024

rapids-bot bot merged commit 58e4b95 into rapidsai:branch-24.06 Apr 3, 2024
27 checks passed

rjzamora deleted the expr-explicit-comms branch April 3, 2024 23:11

This was referenced Jul 1, 2024

Incorrect Shuffle results with dask-cuda 24.06 & above NVIDIA/NeMo-Curator#134

Closed

Fix partitioning in explicit-comms shuffle #1356

Merged

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Update explicit-comms for dask-expr support #1323

Update explicit-comms for dask-expr support #1323

rjzamora commented Mar 20, 2024 •

edited

Loading

pentschev left a comment

quasiben commented Mar 20, 2024

rjzamora commented Mar 20, 2024

rjzamora commented Mar 21, 2024

pentschev commented Mar 21, 2024

madsbk commented Mar 21, 2024 •

edited

Loading

rjzamora commented Mar 21, 2024

quasiben commented Mar 21, 2024

rjzamora Mar 21, 2024

rjzamora commented Mar 22, 2024

rjzamora commented Apr 2, 2024

rjzamora Apr 2, 2024

pentschev left a comment

pentschev Apr 2, 2024

rjzamora Apr 2, 2024

pentschev Apr 2, 2024

rjzamora Apr 2, 2024

pentschev Apr 2, 2024

rjzamora Apr 2, 2024 •

edited

Loading

pentschev left a comment

pentschev Apr 3, 2024

rjzamora commented Apr 3, 2024

Update explicit-comms for dask-expr support #1323

Update explicit-comms for dask-expr support #1323

Conversation

rjzamora commented Mar 20, 2024 • edited Loading

pentschev left a comment

Choose a reason for hiding this comment

quasiben commented Mar 20, 2024

rjzamora commented Mar 20, 2024

rjzamora commented Mar 21, 2024

pentschev commented Mar 21, 2024

madsbk commented Mar 21, 2024 • edited Loading

rjzamora commented Mar 21, 2024

quasiben commented Mar 21, 2024

Choose a reason for hiding this comment

rjzamora commented Mar 22, 2024

rjzamora commented Apr 2, 2024

Choose a reason for hiding this comment

pentschev left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

rjzamora Apr 2, 2024 • edited Loading

Choose a reason for hiding this comment

pentschev left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

rjzamora commented Apr 3, 2024

rjzamora commented Mar 20, 2024 •

edited

Loading

madsbk commented Mar 21, 2024 •

edited

Loading

rjzamora Apr 2, 2024 •

edited

Loading