def get_ucx_config(
enable_tcp_over_ucx=None,
enable_infiniband=None,
enable_nvlink=None,
enable_rdmacm=None,
net_devices=None,
cuda_device_index=None,
):
if net_devices == "auto" and enable_infiniband is False:
raise ValueError(
"Using ucx_net_devices='auto' is currently only "
"supported when enable_infiniband=True."
)
ucx_config = dask.config.get("distributed.comm.ucx")
ucx_config[canonical_name("create-cuda-context", ucx_config)] = True
ucx_config[canonical_name("reuse-endpoints", ucx_config)] = not _ucx_111
# If any transport is explicitly disabled (`False`) by the user, others that
# are not specified should be enabled (`True`). If transports are explicitly
# enabled (`True`), then default (`None`) or an explicit `False` will suffice
# in disabling others. However, if there's a mix of enable (`True`) and
# disable (`False`), then those choices can be assumed as intended by the
# user.
#
# This may be handled more gracefully in Distributed in the future.
opts = [enable_tcp_over_ucx, enable_infiniband, enable_nvlink]
if any(opt is False for opt in opts) and not any(opt is True for opt in opts):
if enable_tcp_over_ucx is None:
enable_tcp_over_ucx = True
if enable_nvlink is None:
enable_nvlink = True
if enable_infiniband is None:
enable_infiniband = True
ucx_config[canonical_name("tcp", ucx_config)] = enable_tcp_over_ucx
ucx_config[canonical_name("infiniband", ucx_config)] = enable_infiniband
ucx_config[canonical_name("nvlink", ucx_config)] = enable_nvlink
ucx_config[canonical_name("rdmacm", ucx_config)] = enable_rdmacm
if enable_tcp_over_ucx or enable_infiniband or enable_nvlink:
ucx_config[canonical_name("cuda-copy", ucx_config)] = True
else:
ucx_config[canonical_name("cuda-copy", ucx_config)] = None
if net_devices is not None and net_devices != "":
ucx_config[canonical_name("net-devices", ucx_config)] = get_ucx_net_devices(
cuda_device_index, net_devices
)
return ucx_config
def test_canonical_name():
c = {"foo-bar": 1, "fizz_buzz": 2}
assert canonical_name("foo-bar", c) == "foo-bar"
assert canonical_name("foo_bar", c) == "foo-bar"
assert canonical_name("fizz-buzz", c) == "fizz_buzz"
assert canonical_name("fizz_buzz", c) == "fizz_buzz"
assert canonical_name("new-key", c) == "new-key"
assert canonical_name("new_key", c) == "new_key"
def test_canonical_name():
c = {'foo-bar': 1, 'fizz_buzz': 2}
assert canonical_name('foo-bar', c) == 'foo-bar'
assert canonical_name('foo_bar', c) == 'foo-bar'
assert canonical_name('fizz-buzz', c) == 'fizz_buzz'
assert canonical_name('fizz_buzz', c) == 'fizz_buzz'
assert canonical_name('new-key', c) == 'new-key'
assert canonical_name('new_key', c) == 'new_key'
def test_get_ucx_config(enable_tcp_over_ucx, enable_infiniband, enable_nvlink,
net_devices):
pytest.importorskip("ucp")
kwargs = {
"enable_tcp_over_ucx": enable_tcp_over_ucx,
"enable_infiniband": enable_infiniband,
"enable_nvlink": enable_nvlink,
"net_devices": net_devices,
"cuda_device_index": 0,
}
if net_devices == "auto" and enable_infiniband is False:
with pytest.raises(ValueError):
get_ucx_config(**kwargs)
return
else:
ucx_config = get_ucx_config(**kwargs)
assert ucx_config[canonical_name("create_cuda_context",
ucx_config)] is True
if enable_tcp_over_ucx is not None:
assert ucx_config[canonical_name("tcp",
ucx_config)] is enable_tcp_over_ucx
else:
if (enable_infiniband is not True and enable_nvlink is not True
and not (enable_infiniband is None and enable_nvlink is None)):
assert ucx_config[canonical_name("tcp", ucx_config)] is True
else:
assert ucx_config[canonical_name("tcp", ucx_config)] is None
if enable_infiniband is not None:
assert ucx_config[canonical_name("infiniband",
ucx_config)] is enable_infiniband
else:
if (enable_tcp_over_ucx is not True and enable_nvlink is not True and
not (enable_tcp_over_ucx is None and enable_nvlink is None)):
assert ucx_config[canonical_name("infiniband", ucx_config)] is True
else:
assert ucx_config[canonical_name("infiniband", ucx_config)] is None
if enable_nvlink is not None:
assert ucx_config[canonical_name("nvlink",
ucx_config)] is enable_nvlink
else:
if (enable_tcp_over_ucx is not True and enable_infiniband is not True
and not (enable_tcp_over_ucx is None
and enable_infiniband is None)):
assert ucx_config[canonical_name("nvlink", ucx_config)] is True
else:
assert ucx_config[canonical_name("nvlink", ucx_config)] is None
if any(
opt is not None
for opt in [enable_tcp_over_ucx, enable_infiniband, enable_nvlink]
) and not all(opt is False for opt in
[enable_tcp_over_ucx, enable_infiniband, enable_nvlink]):
assert ucx_config[canonical_name("cuda-copy", ucx_config)] is True
else:
assert ucx_config[canonical_name("cuda-copy", ucx_config)] is None
if net_devices == "auto":
# Since the actual device is system-dependent, we don't do any
# checks at the moment. If any InfiniBand devices are available,
# that will be the value of "net-devices", otherwise an empty string.
pass
elif net_devices == "eth0":
assert ucx_config[canonical_name("net-devices", ucx_config)] == "eth0"
else:
assert ucx_config[canonical_name("net-devices", ucx_config)] is None
