Explicit-comms-shuffle: fine control of task scheduling

[madsbk]

In shuffle, use Client.submit() to control where tasks are executed and release temporary dataframes ASAP.

Context

In the final step in explicit-comms shuffle, we call getitem() to extract the final dataframe partitions from the result of the local shuffle. In some cases, these getitem() tasks would not run on the worker that ran the local shuffle, which would result in extra communication and spilling.
We now use submit(..., worker=...) to make sure that the worker running the local shuffle also runs the getitem() task afterwards.

Is it possible to do this without the use of submit() to avoid the overhead of creating a Future for each dataframe partition?

[pentschev]

Only added a minor code duplication reduction, apart from that LGTM. Thanks @madsbk !

[wence-]

My only comment was similar to @pentschev's; don't know a better way to force location of the compute

[codecov-commenter]

Codecov Report

Base: 0.00% // Head: 0.00% // No change to project coverage 👍

Coverage data is based on head (658c3da) compared to base (99894b1).
Patch coverage: 0.00% of modified lines in pull request are covered.

Additional details and impacted files

@@              Coverage Diff               @@
##           branch-22.12   #1025     +/-   ##
==============================================
  Coverage          0.00%   0.00%             
==============================================
  Files                25      17      -8     
  Lines              3315    2216   -1099     
==============================================
+ Misses             3315    2216   -1099     

Impacted Files	Coverage Δ
dask_cuda/explicit_comms/dataframe/shuffle.py	0.00% <0.00%> (ø)
dask_cuda/benchmarks/local_cupy.py
dask_cuda/_version.py
dask_cuda/benchmarks/local_cudf_merge.py
dask_cuda/benchmarks/local_cudf_shuffle.py
dask_cuda/benchmarks/common.py
dask_cuda/benchmarks/local_cupy_map_overlap.py
dask_cuda/benchmarks/local_cudf_groupby.py
dask_cuda/benchmarks/utils.py

Help us with your feedback. Take ten seconds to tell us how you rate us. Have a feature suggestion? Share it here.

☔ View full report at Codecov.
📢 Do you have feedback about the report comment? Let us know in this issue.

[pentschev]

rerun tests

[pentschev]

Thanks @madsbk !

[pentschev]

@gpucibot merge