Spilling on demand

[madsbk]

Use rapidsai/rmm#892 to implement spilling on demand. Requires use of RMM and JIT-unspill enabled.

The device_memory_limit still works as usual -- when known allocations gets to device_memory_limit, Dask-CUDA starts spilling preemptively. However, with this PR it is should be possible to increase device_memory_limit significantly since memory spikes will be handled by spilling on demand.

Closes #755

[pentschev]

The failure is an OOM in the new test_spill_on_demand. Perhaps we should look for a way to do some sort of mock test? Maybe we can find a way to rewrite RMM's allocation function or the callback function to trigger in a programatic way?

[codecov-commenter]

Codecov Report

❗ No coverage uploaded for pull request base (branch-21.12@ee58ad5). Click here to learn what that means.
The diff coverage is n/a.

❗ Current head 8cefb89 differs from pull request most recent head 9ab0f42. Consider uploading reports for the commit 9ab0f42 to get more accurate results

@@               Coverage Diff               @@
##             branch-21.12     #756   +/-   ##
===============================================
  Coverage                ?   58.20%           
===============================================
  Files                   ?       21           
  Lines                   ?     2991           
  Branches                ?        0           
===============================================
  Hits                    ?     1741           
  Misses                  ?     1250           
  Partials                ?        0           

---

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 ee58ad5...9ab0f42. Read the comment docs.

[pentschev]

LGTM @madsbk , thanks for the great work. I added a couple of minor suggestions only.

[VibhuJawa]

I tested this on the spilling related workflow mentioned here and was able to increase the memory limit from 0.6 to 0.9 but some how we take more time 2min 32s with the PR vs 1min 46s on main. Any reason for that to happen ?

CC: @randerzander

[madsbk]

I tested this on the spilling related workflow mentioned here and was able to increase the memory limit from 0.6 to 0.9 but some how we take more time 2min 32s with the PR vs 1min 46s on main. Any reason for that to happen ?

In principle yes, a memory limit of 0.6 could trigger favorable spilling by chance but generally spilling-on-demand should be better. @VibhuJawa how many runs did you do? In my experiment the memory requirement and runtime of a shuffle workflow fluctuate drastically.
Can I ask you to try comparing main vs this PR setting the a memory limit to 0.6 in both cases? If this PR doesn't slowdown the execute by itself, I think we are good to merge it.
Afterwards, we should investigate why setting the memory limit to 0.9 might be a disadvantageous in some cases.

[VibhuJawa]

Can I ask you to try comparing main vs this PR setting the a memory limit to 0.6 in both cases? If this PR doesn't slowdown the execute by itself, I think we are good to merge it.

We get similar performence in this case on PR and main. Agreed that PR itself does not cause a slowdown.

@VibhuJawa how many runs did you do?

I did two runs so not a comprehensive test by any means.

[jakirkham]

@gpucibot merge

[jakirkham]

Thanks Mads for the PR! Also Peter and Vibhu for the review and testing 😄

[EvenOldridge]

Excited to try this out!

[jakirkham]

Nightlies are up. So should be good to go Even 😀