Dask-cuda: [FEA]Prevent additional memory allocated when launching a dask-cuda worker with UCX

[VibhuJawa]

[FEA]Prevent additional memory allocated when launching a dask-cuda worker with UCX

In the latest versions of distributed and dask-cuda we seem to be launching a process that takes 254 MB.

According to offline discussion with @pentschev it's probably coming from this change in distributed to create CUDA context before communications are initialized. Comms are always initialized before plugins, so it cant just be moved to Dask-CUDA

Command Run:

CUDA_VISIBLE_DEVICES=0 dask-cuda-worker ucx://10.33.227.151:8786 --enable-nvlink  --rmm-pool-size '5GB' --memory-limit='50000M' &

On Old Distributed+ Dask-Cuda

Libraries:

distributed               2021.4.0         py38h578d9bd_0    conda-forge
dask-cuda                 0.19.0                   py38_0    rapidsai

librmm                    0.19.0          cuda11.0_g7065af3_0    rapidsai
rmm                       0.19.0          cuda_11.0_py38_g7065af3_0    rapidsai

On New Distributed+ Dask-cuda

36790  /raid/vjawa/conda/envs/dask_cuda_new/bin/python /raid/vjawa/conda/envs/dask_cuda_new/bin/dask-cuda-worker ucx://10.33.227.151:8786 --enable-nvlink --rmm-pool-size 5GB --me

36885  /raid/vjawa/conda/envs/dask_cuda_new/bin/python -c from multiprocessing.spawn import spawn_main; spawn_main(tracker_fd=60, pipe_handle=66) --multiprocessing-fork

Libraries:

distributed               2021.5.0         py38h578d9bd_0    conda-forge
dask-cuda                 0.20.0a210517           py38_40    rapidsai-nightly

librmm                    21.08.00a210527 cuda11.0_g139ea5b_9    rapidsai-nightly
rmm                       21.08.00a210527 cuda_11.0_py38_g139ea5b_9    rapidsai-nightly

[pentschev]

Thank you @VibhuJawa for the report. This was introduced in https://github.com/dask/distributed/pull/4787/files#diff-eefe4811faa99fe09003bebb81398442ef60d226dc8d3e4abef629e0c71fad78R59-R70 .

To elaborate on the problem, for various reasons we must ensure we create a CUDA context before initializing UCX, one of them is automatic GPU<->IB device mapping, that's a feature we can use with UCX 1.11 and will allow us then to stop using --ucx-net-devices auto. One of the main advantages is that automatic GPU<->IB mapping works for the DGX A100, whereas --ucx-net-devices auto doesn't. The problem on the Dask side is that comms initialization happens before pretty much anything else, including plugins, which limits us on where we can create the context, i.e., it has to be within Distributed and cannot be moved to Dask-CUDA, which means the context gets created for the Dask scheduler and client when using UCX as well.

After discussing this with other UCX devs, there's not much else we can do now and we indeed need to ensure the CUDA context is created before ucp.init() is called. They will discuss on how to resolve that problem within UCX, but there's not ETA on that.

For now, I see only two options:

1. Live with the additional CUDA context;
2. We could avoid creating the CUDA context for UCX 1.9, but I don't think we should for UCX 1.11 particularly for better user experience in the DGX A100.

If this is a big problem for any current use cases please let us know, in which case we could try option 2 above, or try to elaborate something more complex to suit a specific need, but that will definitely cause other limitations that we need to weigh as well.

[pentschev]

Just for a quick update, I synced yesterday with UCX folks and letting UCX resolve the CUDA context creation matter on its own is a large time investment, thus this is going to be tentatively included in UCX 1.12 (scheduled for release in November 2021).