Add support for CudaAsyncMemoryResource

[charlesbluca]

Closes #565

Adds the --rmm-pool-async/rmm_pool_async option to the CLI and cluster to enable the use of rmm.mr.CudaAsyncMemoryResource in the RMM initialization.

[charlesbluca]

Some small questions/comments:

[charlesbluca]

It looks like CUDA 11.2 has different behavior, in that the memory resource remains an rmm.mr.PoolMemoryResource regardless of if rmm_pool_async is enabled. Is there a way we can check for CUDA version in the tests to handle this?

[pentschev]

cudaMallocAsync was only introduced in CUDA 11.2, https://github.com/rapidsai/rmm/blob/80bfeb2816845da5921551ae5c158e92427bec86/python/rmm/tests/test_rmm.py#L511-L539 will give you some ideas on how to test that async is both active and how you can skip tests for older CUDA versions.

[charlesbluca]

Thanks @pentschev! In that case, I'm a little interested in what's happening for the 11.0 tests, as I would expect a RuntimeError upon trying to set up async there.

EDIT:

I'm assuming that the CudaMemoryResource that the workers return for rmm.mr.get_current_device_resource_type in the 11.0 cases is because RMMSetup is silently failing to allocate a pool, similar to what's happening in #443. With that knowledge, should we try to add some form of check that the user has 11.2 if they enable async, or wait on dask/distributed#4297 to begin handling worker plugin exceptions?

[codecov-io]

Codecov Report

Merging #566 (22e0aa2) into branch-0.20 (cb9bf35) will decrease coverage by 0.47%.
The diff coverage is 52.32%.

❗ Current head 22e0aa2 differs from pull request most recent head fb89979. Consider uploading reports for the commit fb89979 to get more accurate results

@@               Coverage Diff               @@
##           branch-0.20     #566      +/-   ##
===============================================
- Coverage        61.06%   60.59%   -0.48%     
===============================================
  Files               22       22              
  Lines             2571     2644      +73     
===============================================
+ Hits              1570     1602      +32     
- Misses            1001     1042      +41     

Impacted Files	Coverage Δ
dask_cuda/device_host_file.py	67.64% <30.61%> (-20.99%)	⬇️
dask_cuda/cuda_worker.py	78.04% <50.00%> (-0.71%)	⬇️
dask_cuda/local_cuda_cluster.py	78.78% <66.66%> (-2.61%)	⬇️
dask_cuda/utils.py	89.47% <85.71%> (-0.37%)	⬇️
dask_cuda/cli/dask_cuda_worker.py	95.45% <100.00%> (ø)
dask_cuda/explicit_comms/comms.py	98.87% <100.00%> (+0.09%)	⬆️

---

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 3439b83...fb89979. Read the comment docs.

[pentschev]

I'm assuming that the CudaMemoryResource that the workers return for rmm.mr.get_current_device_resource_type in the 11.0 cases is because RMMSetup is silently failing to allocate a pool, similar to what's happening in #443. With that knowledge, should we try to add some form of check that the user has 11.2 if they enable async, or wait on dask/distributed#4297 to begin handling worker plugin exceptions?

It's possible, but some discussion seems to point that we don't want an async pool, but rather just switching to the async allocator, as commented in #565 (comment) . Could you update the PR here to reflect that? I think that will also make the assertion easier, you can probably then just do the same as https://github.com/rapidsai/rmm/blob/80bfeb2816845da5921551ae5c158e92427bec86/python/rmm/tests/test_rmm.py#L519-L521 .

[charlesbluca]

Sure! Does this mean that async can be enabled for RMM pools and managed memory?

[pentschev]

No, that means we're changing the allocator to the async allocator (which is neither the default nor the managed memory allocators), and without a pool. In summary, enabling the async allocator (something like --rmm-async-malloc) is mutually exclusive with both --rmm-pool-size and --rmm-managed-memory and would raise an error if that's attempted.

[charlesbluca]

Updated to reflect this; some questions:

[pentschev]

Since this is the CLI, is this really going to raise an exception or fail with an error instead? I'm not really sure what happens when an exception is raise in the CUDAWorker class when running the CLI tool.

[charlesbluca]

It looks like this fails with an error - I mostly just copied the phrasing used to describe the behavior of using NVLink + managed memory:

dask-cuda/dask_cuda/cli/dask_cuda_worker.py

Lines 112 to 114 in 1526017

	"WARNING: managed memory is currently incompatible with NVLink, "
	"trying to enable both will result in an exception.",
	)

I can change this here and in #561 if it would make more sense to say the CLI will fail outright in these cases.

[pentschev]

Suggested change

	- UCX-Py is not installed or wasn't compiled with hwloc support or
	- UCX-Py is not installed or wasn't compiled with hwloc support; or

[pentschev]

Seems like you accidentally removed a white space from the beginning of each line in this block, which probably will result in a complaint when parsing the docstrings.

[charlesbluca]

Good catch!

[pentschev]

LGTM, thanks @charlesbluca !

[charlesbluca]

@gpucibot merge