Simplify UCX configs, permitting UCX_TLS=all

[pentschev]

Up until now, we require users to specify what transports should be used by UCX, pushing the configuration burden onto the user, being also error-prone. We can now reduce this configuration burden with just one configuration being added in dask/distributed#5526: DASK_DISTRIBUTED__COMM__UCX__CREATE_CUDA_CONTEXT/distributed.comm.ucx.create_cuda_context, which creates the CUDA context before UCX is initialized.

This is an example of how to setup a cluster with dask-cuda-worker after this change:

# Scheduler
UCX_MEMTYPE_REG_WHOLE_ALLOC_TYPES=cuda DASK_DISTRIBUTED__COMM__UCX__CREATE_CUDA_CONTEXT=True dask-scheduler --protocol ucx --interface ib0

# Workers
UCX_MEMTYPE_REG_WHOLE_ALLOC_TYPES=cuda dask-cuda-worker ucx://${SCHEDULER_IB0_IP}:8786 --interface ib0 --rmm-pool-size 29GiB

# Client
UCX_MEMTYPE_REG_WHOLE_ALLOC_TYPES=cuda DASK_DISTRIBUTED__COMM__UCX__CREATE_CUDA_CONTEXT=True python client.py

Similarly, one can setup: LocalCUDACluster(protocol="ucx", interface="ib0").

Note above how ib0 is intentionally specified. That is mandatory to use RDMACM, as it is necessary to have listeners bind to an InfiniBand interface, but can be left unspecified when using systems without InfiniBand or if RDMACM isn't required (discouraged on systems that have InfiniBand connectivity). The UCX_MEMTYPE_REG_WHOLE_ALLOC_TYPES=cuda option is specified for optimal InfiniBand performance with CUDA and will be default in UCX 1.12, when specifying it won't be necessary anymore.

Changes introduced here are backwards-compatible, meaning the old options such as --enable-nvlink/enable_nvlink=True are still valid. However, if any of those options is specified, the user is responsible to enable/disable all desired transports, which can also be useful for benchmarking specific transports.

Finally, creating a CUDA context may not be necessary by UCX in the future, at a point where it will be possible to remove DASK_DISTRIBUTED__COMM__UCX__CREATE_CUDA_CONTEXT=True from scheduler/client processes entirely.

[pentschev]

cc @VibhuJawa @beckernick @ayushdg @randerzander @quasiben FYI

[codecov-commenter]

Codecov Report

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

@@               Coverage Diff               @@
##             branch-22.02     #792   +/-   ##
===============================================
  Coverage                ?   89.30%           
===============================================
  Files                   ?       16           
  Lines                   ?     2057           
  Branches                ?        0           
===============================================
  Hits                    ?     1837           
  Misses                  ?      220           
  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 f1b0e27...ad669fc. Read the comment docs.

[pentschev]

rerun tests

[pentschev]

Since this is a somewhat large change, I'm targeting 22.02 to avoid any major breakage at this time. If anyone feels differently and think this is a great feature to have in 21.12 and we shouldn't wait, please let me know.

[madsbk]

Nice work @pentschev

[pentschev]

rerun tests

[pentschev]

Thanks @madsbk for the review here!

[pentschev]

@gpucibot merge

[jakirkham]

Thanks Peter! 😄