Advanced / Utility API

Advanced / Utility API#

2026-03-31

18 min read time

Applies to Linux

GPU utilities#

dask_cuda.utils.get_gpu_count()#

dask_cuda.utils.get_gpu_handle(device_id=0)#

Get GPU handle from device index or UUID.

Parameters:

device_id (int or str) – The index or UUID of the device from which to obtain the handle.

Raises:

ValueError – If acquiring the device handle for the device specified failed.
pynvml.NVMLError – If any NVML error occurred while initializing.

Examples

>>> get_gpu_handle(device_id=0)

>>> get_gpu_handle(device_id="GPU-9fb42d6f-7d6b-368f-f79c-3c3e784c93f6")

dask_cuda.utils.get_gpu_uuid(device_index=0)#

Get GPU UUID from CUDA device index.

Parameters:: device_index (int or str) – The index or UUID of the device from which to obtain the UUID.

Examples

>>> get_gpu_uuid()
'GPU-9baca7f5-0f2f-01ac-6b05-8da14d6e9005'

>>> get_gpu_uuid(3)
'GPU-9fb42d6f-7d6b-368f-f79c-3c3e784c93f6'

>>> get_gpu_uuid("GPU-9fb42d6f-7d6b-368f-f79c-3c3e784c93f6")
'GPU-9fb42d6f-7d6b-368f-f79c-3c3e784c93f6'

dask_cuda.utils.get_n_gpus()#

dask_cuda.utils.get_device_total_memory(device_index=0)#

Return total memory of CUDA device with index or with device identifier UUID.

Parameters:

device_index (int or str) – The index or UUID of the device from which to obtain the CPU affinity.

Returns:

The total memory of the CUDA Device in bytes, or None for devices that do not
have a dedicated memory resource, as is usually the case for system on a chip (SoC)
devices.

dask_cuda.utils.has_device_memory_resource(device_index=0)#

Determine wheter CUDA device has dedicated memory resource.

Certain devices have no dedicated memory resource, such as system on a chip (SoC) devices.

Parameters:: device_index (int or str) – The index or UUID of the device from which to obtain the CPU affinity.
Return type:: Whether the device has a dedicated memory resource.

dask_cuda.utils.get_gpu_count_mig(return_uuids=False)#

Return the number of MIG instances available

Parameters:: return_uuids (bool) – Returns the uuids of the MIG instances available optionally

CPU affinity#

dask_cuda.utils.get_cpu_affinity(device_index=None)#

Get a list containing the CPU indices to which a GPU is directly connected. Use either the device index or the specified device identifier UUID.

Parameters:: device_index (int or str) – The index or UUID of the device from which to obtain the CPU affinity.

Examples

>>> from dask_cuda.utils import get_cpu_affinity
>>> get_cpu_affinity(0)  # DGX-1 has GPUs 0-3 connected to CPUs [0-19, 20-39]
[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19,
 40, 41, 42, 43, 44, 45, 46, 47, 48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59]
>>> get_cpu_affinity(5)  # DGX-1 has GPUs 5-7 connected to CPUs [20-39, 60-79]
[20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39,
 60, 61, 62, 63, 64, 65, 66, 67, 68, 69, 70, 71, 72, 73, 74, 75, 76, 77, 78, 79]
>>> get_cpu_affinity(1000)  # DGX-1 has no device on index 1000
dask_cuda/utils.py:96: UserWarning: Cannot get CPU affinity for device with index
1000, setting default affinity
[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19,
 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39,
 40, 41, 42, 43, 44, 45, 46, 47, 48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59,
 60, 61, 62, 63, 64, 65, 66, 67, 68, 69, 70, 71, 72, 73, 74, 75, 76, 77, 78, 79]

dask_cuda.utils.get_cpu_count()#

dask_cuda.utils.unpack_bitmask(x, mask_bits=64)#

Unpack a list of integers containing bitmasks.

Parameters:

x (list of int) – A list of integers
mask_bits (int) – An integer determining the bitwidth of x

Examples

>>> from dask_cuda.utils import unpack_bitmaps
>>> unpack_bitmask([1 + 2 + 8])
[0, 1, 3]
>>> unpack_bitmask([1 + 2 + 16])
[0, 1, 4]
>>> unpack_bitmask([1 + 2 + 16, 2 + 4])
[0, 1, 4, 65, 66]
>>> unpack_bitmask([1 + 2 + 16, 2 + 4], mask_bits=32)
[0, 1, 4, 33, 34]

Device enumeration#

dask_cuda.utils.cuda_visible_devices(i, visible=None)#

Cycling values for CUDA_VISIBLE_DEVICES environment variable

Examples

>>> cuda_visible_devices(0, range(4))
'0,1,2,3'
>>> cuda_visible_devices(3, range(8))
'3,4,5,6,7,0,1,2'

dask_cuda.utils.nvml_device_index(i, CUDA_VISIBLE_DEVICES)#

Get the device index for NVML addressing

NVML expects the index of the physical device, unlike CUDA runtime which expects the address relative to CUDA_VISIBLE_DEVICES. This function returns the i-th device index from the CUDA_VISIBLE_DEVICES comma-separated string of devices or list.

Examples

>>> nvml_device_index(1, "0,1,2,3")
1
>>> nvml_device_index(1, "1,2,3,0")
2
>>> nvml_device_index(1, [0,1,2,3])
1
>>> nvml_device_index(1, [1,2,3,0])
2
>>> nvml_device_index(1, ["GPU-84fd49f2-48ad-50e8-9f2e-3bf0dfd47ccb",
                          "GPU-d6ac2d46-159b-5895-a854-cb745962ef0f",
                          "GPU-158153b7-51d0-5908-a67c-f406bc86be17"])
"MIG-d6ac2d46-159b-5895-a854-cb745962ef0f"
>>> nvml_device_index(2, ["MIG-41b3359c-e721-56e5-8009-12e5797ed514",
                          "MIG-65b79fff-6d3c-5490-a288-b31ec705f310",
                          "MIG-c6e2bae8-46d4-5a7e-9a68-c6cf1f680ba0"])
"MIG-c6e2bae8-46d4-5a7e-9a68-c6cf1f680ba0"
>>> nvml_device_index(1, 2)
Traceback (most recent call last):
...
ValueError: CUDA_VISIBLE_DEVICES must be `str` or `list`

dask_cuda.utils.parse_cuda_visible_device(dev)#

Parses a single CUDA device identifier

A device identifier must either be an integer, a string containing an integer or a string containing the device’s UUID, beginning with prefix ‘GPU-’ or ‘MIG-‘.

>>> parse_cuda_visible_device(2)
2
>>> parse_cuda_visible_device('2')
2
>>> parse_cuda_visible_device('GPU-9baca7f5-0f2f-01ac-6b05-8da14d6e9005')
'GPU-9baca7f5-0f2f-01ac-6b05-8da14d6e9005'
>>> parse_cuda_visible_device('Foo')
Traceback (most recent call last):
...
ValueError: Devices in CUDA_VISIBLE_DEVICES must be comma-separated integers or
strings beginning with 'GPU-' or 'MIG-' prefixes.

Memory parsing#

dask_cuda.utils.parse_device_bytes(device_bytes, device_index=0, alignment_size=1)#

Parse bytes relative to a specific CUDA device.

Parameters:

device_bytes (float, int, str or None) – Can be an integer (bytes), float (fraction of total device memory), string (like "5GB" or "5000M"), 0 and None are special cases returning None.
device_index (int or str) – The index or UUID of the device from which to obtain the total memory amount. Default: 0.
alignment_size (int) – Number of bytes of alignment to use, i.e., allocation must be a multiple of that size. RMM pool requires 256 bytes alignment.

Returns:

The parsed bytes value relative to the CUDA devices, or None as convenience if
device_bytes is None or any value that would evaluate to 0.

Examples

>>> # On a 32GB CUDA device
>>> parse_device_bytes(None)
None
>>> parse_device_bytes(0)
None
>>> parse_device_bytes(0.0)
None
>>> parse_device_bytes("0 MiB")
None
>>> parse_device_bytes(1.0)
34089730048
>>> parse_device_bytes(0.8)
27271784038
>>> parse_device_bytes(1000000000)
1000000000
>>> parse_device_bytes("1GB")
1000000000
>>> parse_device_bytes("1GB")
1000000000

dask_cuda.utils.parse_device_memory_limit(device_memory_limit, device_index=0, alignment_size=1)#

Parse memory limit to be used by a CUDA device.

Parameters:

device_memory_limit (float, int, str or None) – Can be an integer (bytes), float (fraction of total device memory), string (like "5GB" or "5000M"), "auto", 0 or None to disable spilling to host (i.e. allow full device memory usage). Another special value "default" is also available and returns the recommended Dask-CUDA’s defaults and means 80% of the total device memory (analogous to 0.8), and disabled spilling (analogous to auto/0/None) on devices without a dedicated memory resource, such as system on a chip (SoC) devices.
device_index (int or str) – The index or UUID of the device from which to obtain the total memory amount. Default: 0.
alignment_size (int) – Number of bytes of alignment to use, i.e., allocation must be a multiple of that size. RMM pool requires 256 bytes alignment.

Returns:

The parsed memory limit in bytes, or None as convenience if
device_memory_limit is None or any value that would evaluate to 0.

Examples

>>> # On a 32GB CUDA device
>>> parse_device_memory_limit(None)
None
>>> parse_device_memory_limit(0)
None
>>> parse_device_memory_limit(0.0)
None
>>> parse_device_memory_limit("0 MiB")
None
>>> parse_device_memory_limit(1.0)
34089730048
>>> parse_device_memory_limit(0.8)
27271784038
>>> parse_device_memory_limit(1000000000)
1000000000
>>> parse_device_memory_limit("1GB")
1000000000
>>> parse_device_memory_limit("1GB")
1000000000
>>> parse_device_memory_limit("auto") == (
...    parse_device_memory_limit(1.0)
...    if has_device_memory_resource()
...    else None
... )
True
>>> parse_device_memory_limit("default") == (
...    parse_device_memory_limit(0.8)
...    if has_device_memory_resource()
...    else None
... )
True

dask_cuda.utils.get_rmm_device_memory_usage()#

Get current bytes allocated on current device through RMM

Check the current RMM resource stack for resources such as StatisticsResourceAdaptor and TrackingResourceAdaptor that can report the current allocated bytes. Returns None, if no such resources exist.

Returns:: nbytes – Number of bytes allocated on device through RMM or None
Return type:: int or None

UCX configuration#

dask_cuda.utils.get_ucx_config(enable_tcp_over_ucx=None, enable_infiniband=None, enable_rocm_ipc=None, enable_rdmacm=None, enable_nvlink=None)#

dask_cuda.utils.get_preload_options(protocol=None, create_cuda_context=None, enable_tcp_over_ucx=None, enable_infiniband=None, enable_rocm_ipc=None, enable_rdmacm=None, enable_nvlink=None)#

Return a dictionary with the preload and preload_argv options required to create CUDA context and enabling UCX communication.

Parameters:

protocol (None or str, default None) – If “ucx”, options related to UCX (enable_tcp_over_ucx, enable_infiniband, enable_rocm_ipc) are added to preload_argv.
create_cuda_context (bool, default None) – Ensure the CUDA context gets created at initialization, generally needed by Dask workers.
enable_tcp (bool, default None) – Set environment variables to enable TCP over UCX, even when InfiniBand or NVLink support are disabled.
enable_infiniband (bool, default None) – Set environment variables to enable UCX InfiniBand support. Implies enable_tcp=True.
enable_rdmacm (bool, default None) – Set environment variables to enable UCX RDMA connection manager support. Currently requires enable_infiniband=True.
enable_rocm_ipc (bool, default None) – Set environment variables to enable UCX rocm_ipc support. Implies enable_tcp=True.

Example

>>> from dask_cuda.utils import get_preload_options
>>> get_preload_options()
{'preload': ['dask_cuda.initialize'], 'preload_argv': []}
>>> get_preload_options(protocol="ucx",
...                     create_cuda_context=True,
...                     enable_infiniband=True)
{'preload': ['dask_cuda.initialize'],
 'preload_argv': ['--create-cuda-context',
  '--enable-infiniband']}

Cluster introspection#

dask_cuda.utils.wait_workers(client, min_timeout=10, seconds_per_gpu=2, n_gpus=None, timeout_callback=None)#

Wait for workers to be available. When a timeout occurs, a callback is executed if specified. Generally used for tests.

Parameters:

client (distributed.Client) – Instance of client, used to query for number of workers connected.
min_timeout (float) – Minimum number of seconds to wait before timeout. This value may be overridden by setting the DASK_CUDA_WAIT_WORKERS_MIN_TIMEOUT with a positive integer.
seconds_per_gpu (float) – Seconds to wait for each GPU on the system. For example, if its value is 2 and there is a total of 8 GPUs (workers) being started, a timeout will occur after 16 seconds. Note that this value is only used as timeout when larger than min_timeout.
n_gpus (None or int) – If specified, will wait for a that amount of GPUs (i.e., Dask workers) to come online, else waits for a total of get_n_gpus workers.
timeout_callback (None or callable) – A callback function to be executed if a timeout occurs, ignored if None.

Return type:

True if all workers were started, False if a timeout occurs.

dask_cuda.utils.all_to_all(client)#

dask_cuda.utils.get_worker_config(dask_worker)#

dask_cuda.utils.get_cluster_configuration(client)#

dask_cuda.utils.print_cluster_config(client)#

print current Dask cluster configuration

Return type:: None

class dask_cuda.utils.CommaSeparatedChoice(choices, case_sensitive=True)#

Bases: click.Choice

The choice type allows a value to be checked against a fixed set of supported values.

You may pass any iterable value which will be converted to a tuple and thus will only be iterated once.

The resulting value will always be one of the originally passed choices. See normalize_choice() for more info on the mapping of strings to choices. See choice-opts for an example.

Parameters:

case_sensitive (bool) – Set to false to make choices case insensitive. Defaults to true.
choices (collections.abc.Iterable[ParamTypeValue])

Changed in version 8.2.0: Non-str choices are now supported. It can additionally be any iterable. Before you were not recommended to pass anything but a list or tuple.

Added in version 8.2.0: Choice normalization can be overridden via normalize_choice().

convert(value, param, ctx)#: For a given value from the parser, normalize it and find its matching normalized value in the list of choices. Then return the matched “original” choice.

Worker plugins#

class dask_cuda.plugins.CPUAffinity(cores)#

cores: _typeshed.Incomplete#

setup(worker=None)#

Return type:: None

class dask_cuda.plugins.CUDFSetup(spill, spill_stats)#

setup(worker=None)#

Return type:: None

spill: _typeshed.Incomplete#

spill_stats: _typeshed.Incomplete#

class dask_cuda.plugins.RMMSetup(initial_pool_size, maximum_pool_size, managed_memory, async_alloc, release_threshold, log_directory, track_allocations, external_lib_list)#

async_alloc: _typeshed.Incomplete#

external_lib_list: _typeshed.Incomplete#

initial_pool_size: _typeshed.Incomplete#

log_directory: _typeshed.Incomplete#

logging: _typeshed.Incomplete#

managed_memory: _typeshed.Incomplete#

maximum_pool_size: _typeshed.Incomplete#

release_threshold: _typeshed.Incomplete#

rmm_track_allocations: _typeshed.Incomplete#

setup(worker=None)#

Return type:: None

class dask_cuda.plugins.PreImport(libraries)#

libraries: _typeshed.Incomplete#

setup(worker=None)#

Return type:: None

dask_cuda.plugins.enable_rmm_memory_for_library(lib_name)#

Enable RMM memory pool support for a specified third-party library.

This function allows the given library to utilize RMM’s memory pool if it supports integration with RMM. The library name is passed as a string argument, and if the library is compatible, its memory allocator will be configured to use RMM.

Parameters:

lib_name (str) – The name of the third-party library to enable RMM memory pool support for. Supported libraries are “cupy” and “torch”.

Raises:

ValueError – If the library name is not supported or does not have RMM integration.
ImportError – If the required library is not installed.

Return type:

None

Proxy objects and JIT-Unspilling#

dask_cuda.proxy_object.asproxy(obj, serializers=None, subclass=None)#

Wrap obj in a ProxyObject object if it isn’t already.

Parameters:

obj (object) – Object to wrap in a ProxyObject object.
serializers (Iterable[str], optional) – Serializers to use to serialize obj. If None, no serialization is done.
subclass (class, optional) – Specify a subclass of ProxyObject to create instead of ProxyObject. subclass must be pickable.

Return type:

The ProxyObject proxying ``obj``

class dask_cuda.proxy_object.ProxyObject(detail)#

Object wrapper/proxy for serializable objects

This is used by ProxifyHostFile to delay deserialization of returned objects.

Objects proxied by an instance of this class will be JIT-deserialized when accessed. The instance behaves as the proxied object and can be accessed/used just like the proxied object.

ProxyObject has some limitations and doesn’t mimic the proxied object perfectly. Thus, if encountering problems remember that it is always possible to use unproxy() to access the proxied object directly or disable JIT deserialization completely with jit_unspill=False.

Type checking using instance() works as expected but direct type checking doesn’t: >>> import numpy as np >>> from dask_cuda.proxy_object import asproxy >>> x = np.arange(3) >>> isinstance(asproxy(x), type(x)) True >>> type(asproxy(x)) is type(x) False

_pxy#

Details of all proxy information of the underlying proxied object. Access to _pxy is not pass-through to the proxied object, which is the case for most other access to the ProxyObject.

Type:: ProxyDetail

_pxy_cache#

A dictionary used for caching attributes

Type:: dict

Parameters:: detail (ProxyDetail) – The Any kind of object to be proxied.

dask_cuda.proxify_device_objects.proxify_device_objects(obj, proxied_id_to_proxy=None, found_proxies=None, excl_proxies=False, mark_as_explicit_proxies=False)#

Wrap device objects in ProxyObject

Search through obj and wraps all CUDA device objects in ProxyObject. It uses proxied_id_to_proxy to make sure that identical CUDA device objects found in obj are wrapped by the same ProxyObject.

Parameters:

obj (Any) – Object to search through or wrap in a ProxyObject.
proxied_id_to_proxy (MutableMapping[int, ProxyObject]) – Dict mapping the id() of proxied objects (CUDA device objects) to their proxy and is updated with all new proxied objects found in obj. If None, use an empty dict.
found_proxies (List[ProxyObject]) – List of found proxies in obj. Notice, this includes all proxies found, including those already in proxied_id_to_proxy. If None, use an empty list.
excl_proxies (bool) – Don’t add found objects that are already ProxyObject to found_proxies.
mark_as_explicit_proxies (bool) – Mark found proxies as “explicit”, which means that the user allows them as input arguments to dask tasks even in compatibility-mode.

Returns:

ret – A copy of obj where all CUDA device objects are wrapped in ProxyObject

Return type:

Any

dask_cuda.proxify_device_objects.unproxify_device_objects(obj, skip_explicit_proxies=False, only_incompatible_types=False)#

Unproxify device objects

Search through obj and un-wraps all CUDA device objects.

Parameters:

obj (Any) – Object to search through or unproxify.
skip_explicit_proxies (bool) – When True, skipping proxy objects marked as explicit proxies.
only_incompatible_types (bool) – When True, ONLY unproxify incompatible type. The skip_explicit_proxies argument is ignored.

Returns:

ret – A copy of obj where all CUDA device objects are unproxify

Return type:

Any

Storage#

class dask_cuda.device_host_file.DeviceHostFile(worker_local_directory, *, device_memory_limit=None, memory_limit=None, log_spilling=False)#

Manages serialization/deserialization of objects.

Three LRU cache levels are controlled, for device, host and disk. Each level takes care of serializing objects once its limit has been reached and pass it to the subsequent level. Similarly, each cache may deserialize the object, but storing it back in the appropriate cache, depending on the type of object being deserialized.

Parameters:

worker_local_directory (path) – Path where to store serialized objects on disk
device_memory_limit (int or None) – Number of bytes of CUDA device memory for device LRU cache, spills to host cache once filled. Setting this 0 or None means unlimited device memory, implies no spilling to host.
memory_limit (int or None) – Number of bytes of host memory for host LRU cache, spills to disk once filled. Setting this to 0 or None means unlimited host memory, implies no spilling to disk.
log_spilling (bool) – If True, all spilling operations will be logged directly to distributed.worker with an INFO loglevel. This will eventually be replaced by a Dask configuration flag.

device: _typeshed.Incomplete#

device_buffer: _typeshed.Incomplete#

device_func: _typeshed.Incomplete#

device_host_func: _typeshed.Incomplete#

device_keys: _typeshed.Incomplete#

disk: _typeshed.Incomplete#

disk_func: _typeshed.Incomplete#

disk_func_path: _typeshed.Incomplete#

evict()#

Evicts least recently used host buffer (aka, CPU or system memory)

Implements distributed.spill.ManualEvictProto interface

fast: _typeshed.Incomplete#

get_total_spilling_time()#

host: _typeshed.Incomplete#

host_buffer: _typeshed.Incomplete#

host_func: _typeshed.Incomplete#

others: _typeshed.Incomplete#

set_address(addr)#

Return type:: None

class dask_cuda.device_host_file.LoggedBuffer(*args, fast_name='Fast', slow_name='Slow', addr=None, **kwargs)#

Extends zict.Buffer with logging capabilities

Two arguments fast_name and slow_name are passed to constructor that identify a user-friendly name for logging of where spilling is going from/to. For example, their names can be “Device” and “Host” to identify that spilling is happening from a CUDA device into system memory.

Parameters:

fast_name (str)
slow_name (str)

addr: _typeshed.Incomplete#

fast_name: _typeshed.Incomplete#

fast_to_slow(key, value)#

get_total_spilling_time()#

logger: _typeshed.Incomplete#

msg_template: str#

set_address(addr)#

Return type:: None

slow_name: _typeshed.Incomplete#

slow_to_fast(key)#

total_time_fast_to_slow: float#

total_time_slow_to_fast: float#

RMM utilities#

dask_cuda.utils.get_rmm_log_file_name(dask_worker, logging=False, log_directory=None)#

Parameters:: logging (bool)

Explicit-comms helpers#

dask_cuda.explicit_comms.comms.worker_state(sessionId=None)#

Retrieve the state(s) of the current worker

Parameters:: sessionId (int, optional) – Worker session state ID. If None, all states of the worker are returned.
Returns:: state – Either a single state dict or a dict of state dict
Return type:: dict

dask_cuda.explicit_comms.comms.pop_staging_area(session_state, name)#

Pop the staging area called name

This function must be called within a running explicit-comms task.

Parameters:

session_state (dict) – Worker session state
name (str) – Name for the staging area

Returns:

The staging area, which is a dict that maps keys to their data.

Return type:

dict

dask_cuda.explicit_comms.dataframe.shuffle.patch_shuffle_expression()#

Patch Dasks Shuffle expression.

Notice, this is monkey patched into Dask at dask_cuda import, and it changes Shuffle._layer to lower into an ECShuffle expression when the ‘explicit-comms’ config is set to True.

Return type:: None

Advanced / Utility API

Contents

Advanced / Utility API#

GPU utilities#

CPU affinity#

Device enumeration#

Memory parsing#

UCX configuration#

Cluster introspection#

Worker plugins#

Proxy objects and JIT-Unspilling#

Storage#

RMM utilities#

Explicit-comms helpers#