[WIP] New Dask Scheduler for cuDF

[madsbk]

This PR introduce a new task scheduler for cuDF's Frame, which makes it possible to implement specialized optimizations that might be hard to get accepted into upstream Dask.

For now, this PR implements the task ordering optimization from dask/dask#6051.

My hope is that we can use this to implement Dask scheduling and task graph optimizations. Then, when we get to a stable state, we can work on getting the changes into upstream Dask.

cc. @beckernick, @VibhuJawa

[GPUtester]

Please update the changelog in order to start CI tests.

View the gpuCI docs here.

[codecov]

Codecov Report

Merging #5053 into branch-0.15 will decrease coverage by 0.21%.
The diff coverage is 67.00%.

@@               Coverage Diff               @@
##           branch-0.15    #5053      +/-   ##
===============================================
- Coverage        88.50%   88.28%   -0.22%     
===============================================
  Files               54       56       +2     
  Lines            10270    10366      +96     
===============================================
+ Hits              9089     9152      +63     
- Misses            1181     1214      +33     

Impacted Files	Coverage Δ
python/dask_cudf/dask_cudf/core.py	68.55% <29.16%> (-3.60%)	⬇️
python/dask_cudf/dask_cudf/scheduler.py	75.75% <75.75%> (ø)
python/dask_cudf/dask_cudf/order.py	100.00% <100.00%> (ø)

---

Continue to review full report at Codecov.

Legend - Click here to learn more
Δ = absolute <relative> (impact), ø = not affected, ? = missing data
Powered by Codecov. Last update cfb1f6b...46eecf2. Read the comment docs.

[beckernick]

Also cross-listing dask/dask#6059

[mrocklin]

First, I totally understand that you all have to do whatever you have to do. RAPIDS is under constant time pressure it seems.

However, I will say that it would be great to have things fixed upstream. This takes longer, and it means that you'll need to fix more than what is necessary for your specific problem, but I think that it would help other groups a ton.

In the particular situation of task annotations I think that people would really like this feature. It's not that upstream is unhappy about the feature, it's that there are other considerations that would need to be resolved. My guess is that it would take an excellent dev, like @madsbk , weeks to solve correctly. You all may not have that time, which I totally understand, but I wanted to make it clear that you are absolutely welcome to make these changes upstream.

[madsbk]

@mrocklin, what work you think I will have to do to get dask/dask#6059 moving forward?

[mrocklin]

There are probably many approaches here.

Previously you took the approach of using task annotations. I think that task annotations is a general feature that would help out a lot of people (including several other GPU use cases for people who want to use resource restrictions).

You started down this path by annotating the functions themselves. That solution had problems because it wasn't robust to fusion. It also wasn't clear if we wanted to do something more serious, like have a full Task class. I think that a more serious treatment of this problem would be really useful for the project, and potentially have a lot of impact, especially if it can be done efficiently.

As a side effect of that work, it would be easy to annotate the tasks in dask.dataframe.shuffle with a priority, solving your problem here.

Anyway, that's one approach, there might be others, but that one solves enough user issues that have accrued over the years that I think that people would be really excited by it. It does require a lot of work though, both technical and social.

[kkraus14]

Is this just for debugging purposes to make this synchronous?

[madsbk]

Oops yes, I will fix that.

[harrism]

No recent updates, moving to 0.15. @madsbk is this still needed?

[madsbk]

No recent updates, moving to 0.15. @madsbk is this still needed?

Let's close it, I am working on a general solution in Dask/Distributed