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"]