Dask-cuda: Extend CLI and LocalCluster to support CUDA Async Memory Resource

[beckernick]

With the new functionality in CUDA 11.2, RMM now supports a CudaAsyncMemoryResource in addition to the existing resources. It would be nice to extend the dask-cuda CLI to allow starting a cluster with an RMM pool using either of the memory resources.

EDIT: Turns out we don't want to use this with an explicit RMM pool, but we might still want to expose this as it would allow a caching pool allocator that could play more nicely with other GPU libraries that don't use RMM

[pentschev]

This would require extending RMMSetup with something as follows:

pool_async_mr = rmm.mr.PoolMemoryResource(rmm.mr.CudaAsyncMemoryResource())
rmm.mr.set_current_device_resource(pool_async_mr)

# or rmm.mr.set_per_device_resource(device_id, pool_async_mr)

@charlesbluca would you have bandwidth to work on this?

[charlesbluca]

Sure! Out of interest, is this something we only want added into the CLI, or could we also add this to LocalCUDACluster?

[pentschev]

I think we could add to both to keep them both complete. To be backwards compatible, we could add an --rmm-pool-resource/rmm_pool_resource with two options default (which keeps the current behavior, and is also the default option) and async. In either case --rmm-pool-size/rmm_pool_size would enable the pool, and the new flag would be a modifier to that.

[charlesbluca]

Do we expect there to be potentially more pool resources beyond default and async in the future? If not, maybe a flag like --rmm-pool-async/--no-rmm-pool-async would make more sense for this option.

[pentschev]

From a discussion with @shwina , it seems like we don't want an async pool, but rather just a CudaAsyncMemoryResource():

import rmm
cuda_async_mr = rmm.mr.CudaAsyncMemoryResource())
rmm.mr.set_current_device_resource(cuda_async_mr)

Could you update the PR to reflect that @charlesbluca ?

[pentschev]

@beckernick Charles' PR is in, could you try it out and report back if it works for you or if you find any issues?

[beckernick]

Yes, thanks Charles and Peter!