def test_init_options():
ucp.reset()
os.environ["UCX_SEG_SIZE"] = "2M" # Should be ignored
options = {"SEG_SIZE": "3M"}
ucp.init(options)
config = ucp.get_config()
assert config["SEG_SIZE"] == options["SEG_SIZE"]
def test_init_options_and_env():
ucp.reset()
os.environ["UCX_SEG_SIZE"] = "4M"
options = {"SEG_SIZE": "3M"} # Should be ignored
ucp.init(options, env_takes_precedence=True)
config = ucp.get_config()
assert config["SEG_SIZE"] == options["SEG_SIZE"]
def test_init_options():
ucp.reset()
options = {"SEG_SIZE": "3M"}
# environment specification should be ignored
ucp.init(options)
config = ucp.get_config()
assert config["SEG_SIZE"] == options["SEG_SIZE"]
def _scrub_ucx_config():
"""Function to scrub dask config options for valid UCX config options"""
# configuration of UCX can happen in two ways:
# 1) high level on/off flags which correspond to UCX configuration
# 2) explicitly defined UCX configuration flags
# import does not initialize ucp -- this will occur outside this function
from ucp import get_config, get_ucx_version
ucx_110 = get_ucx_version() >= (1, 10, 0)
options = {}
# if any of the high level flags are set, as long as they are not Null/None,
# we assume we should configure basic TLS settings for UCX, otherwise we
# leave UCX to its default configuration
if any([
dask.config.get("ucx.tcp"),
dask.config.get("ucx.nvlink"),
dask.config.get("ucx.infiniband"),
]):
if dask.config.get("ucx.rdmacm"):
tls = "tcp" if ucx_110 else "tcp,rdmacm"
tls_priority = "rdmacm"
else:
tls = "tcp" if ucx_110 else "tcp,sockcm"
tls_priority = "tcp" if ucx_110 else "sockcm"
# CUDA COPY can optionally be used with ucx -- we rely on the user
# to define when messages will include CUDA objects. Note:
# defining only the Infiniband flag will not enable cuda_copy
if any(
[dask.config.get("ucx.nvlink"),
dask.config.get("ucx.cuda_copy")]):
tls = tls + ",cuda_copy"
if dask.config.get("ucx.infiniband"):
tls = "rc," + tls
if dask.config.get("ucx.nvlink"):
tls = tls + ",cuda_ipc"
options = {"TLS": tls, "SOCKADDR_TLS_PRIORITY": tls_priority}
net_devices = dask.config.get("ucx.net-devices")
if net_devices is not None and net_devices != "":
options["NET_DEVICES"] = net_devices
# ANY UCX options defined in config will overwrite high level dask.ucx flags
valid_ucx_keys = list(get_config().keys())
for k, v in dask.config.get("ucx").items():
if k in valid_ucx_keys:
options[k] = v
else:
logger.debug(
"Key: %s with value: %s not a valid UCX configuration option" %
(k, v))
return options
def test_init_options_and_env():
ucp.reset()
options = {"SEG_SIZE": "3M"} # Should be ignored
ucp.init(options, env_takes_precedence=True)
config = ucp.get_config()
assert config["SEG_SIZE"] == os.environ["UCX_SEG_SIZE"]
# Provided options dict was not modified.
assert options == {"SEG_SIZE": "3M"}
def test_get_config():
with patch.dict(os.environ):
# Unset to test default value
if os.environ.get("UCX_TLS") is not None:
del os.environ["UCX_TLS"]
ucp.reset()
config = ucp.get_config()
assert isinstance(config, dict)
assert config["TLS"] == "all"
def _scrub_ucx_config():
"""Function to scrub dask config options for valid UCX config options"""
# configuration of UCX can happen in two ways:
# 1) high level on/off flags which correspond to UCX configuration
# 2) explicitly defined UCX configuration flags
# import does not initialize ucp -- this will occur outside this function
from ucp import get_config
options = {}
# if any of the high level flags are set, as long as they are not Null/None,
# we assume we should configure basic TLS settings for UCX, otherwise we
# leave UCX to its default configuration
if any([
dask.config.get("distributed.comm.ucx.tcp"),
dask.config.get("distributed.comm.ucx.nvlink"),
dask.config.get("distributed.comm.ucx.infiniband"),
]):
if dask.config.get("distributed.comm.ucx.rdmacm"):
tls = "tcp"
tls_priority = "rdmacm"
else:
tls = "tcp"
tls_priority = "tcp"
# CUDA COPY can optionally be used with ucx -- we rely on the user
# to define when messages will include CUDA objects. Note:
# defining only the Infiniband flag will not enable cuda_copy
if any([
dask.config.get("distributed.comm.ucx.nvlink"),
dask.config.get("distributed.comm.ucx.cuda-copy"),
]):
tls = tls + ",cuda_copy"
if dask.config.get("distributed.comm.ucx.infiniband"):
tls = "rc," + tls
if dask.config.get("distributed.comm.ucx.nvlink"):
tls = tls + ",cuda_ipc"
options = {"TLS": tls, "SOCKADDR_TLS_PRIORITY": tls_priority}
# ANY UCX options defined in config will overwrite high level dask.ucx flags
valid_ucx_vars = list(get_config().keys())
for k, v in options.items():
if k not in valid_ucx_vars:
logger.debug(
f"Key: {k} with value: {v} not a valid UCX configuration option"
)
return options
def initialize(
create_cuda_context=True,
enable_tcp_over_ucx=False,
enable_infiniband=False,
enable_nvlink=False,
net_devices="",
):
if create_cuda_context:
try:
numba.cuda.current_context()
except Exception:
logger.error("Unable to start CUDA Context", exc_info=True)
if enable_tcp_over_ucx or enable_infiniband or enable_nvlink:
try:
import ucp
except ImportError:
logger.error(
"UCX protocol requested but ucp module is not available",
exc_info=True)
else:
options = {}
if enable_tcp_over_ucx or enable_infiniband or enable_nvlink:
tls = "tcp,sockcm,cuda_copy"
tls_priority = "sockcm"
if enable_infiniband:
tls = "rc," + tls
if enable_nvlink:
tls = tls + ",cuda_ipc"
options = {"TLS": tls, "SOCKADDR_TLS_PRIORITY": tls_priority}
if net_devices is not None and net_devices != "":
options["NET_DEVICES"] = net_devices
ucp.reset()
ucp.init(options=options)
ucx_env = {}
for k, v in ucp.get_config().items():
# Skip values that aren't actual environment variables (i.e., not strings)
if isinstance(v, str):
ucx_env["UCX_" + k] = v
# Set also UCX environment variables: required by Dask client. It may be best ti
# to have the client asking the scheduler for the proper variables.
os.environ.update(ucx_env)
def _scrub_ucx_config():
"""Function to scrub dask config options for valid UCX config options"""
# configuration of UCX can happen in two ways:
# 1) high level on/off flags which correspond to UCX configuration
# 2) explicity defined UCX configuration flags
# import does not initialize ucp -- this will occur outside this function
from ucp import get_config
options = {}
# if any of the high level flags are set, as long as they are not Null/None,
# we assume we should configure basic TLS settings for UCX
if any([dask.config.get("ucx.nvlink"), dask.config.get("ucx.infiniband")]):
tls = "tcp,sockcm,cuda_copy"
tls_priority = "sockcm"
if dask.config.get("ucx.infiniband"):
tls = "rc," + tls
if dask.config.get("ucx.nvlink"):
tls = tls + ",cuda_ipc"
options = {"TLS": tls, "SOCKADDR_TLS_PRIORITY": tls_priority}
net_devices = dask.config.get("ucx.net-devices")
if net_devices is not None and net_devices != "":
options["NET_DEVICES"] = net_devices
# ANY UCX options defined in config will overwrite high level dask.ucx flags
valid_ucx_keys = list(get_config().keys())
for k, v in dask.config.get("ucx").items():
if k in valid_ucx_keys:
options[k] = v
else:
logger.debug(
"Key: %s with value: %s not a valid UCX configuration option" %
(k, v))
return options
def test_get_config():
ucp.reset()
config = ucp.get_config()
assert isinstance(config, dict)
assert config["MEMTYPE_CACHE"] == "n"
def test_set_env():
ucp.reset()
os.environ["UCX_SEG_SIZE"] = "2M"
config = ucp.get_config()
assert config["SEG_SIZE"] == os.environ["UCX_SEG_SIZE"]
