parallel

Dask helper functions.

build_gpu_cycler(workers_per_gpu=1, gpu_ids=None)

Generator to cycle through GPU IDs.

On machine with multiple GPUs and a computation that does not use a full GPU. This will yield GPU IDs in a round-robin fashion.

• If gpu_ids is empty, it will yield 0 indefinitely.
• If workers_per_gpu>0 and gpu_ids=[2], it will yield 2 indefinitely.
• If workers_per_gpu=1 and gpu_ids=[0, 2], it will yield 0, 2 indefinitely.
• If workers_per_gpu=4 and gpu_ids=[0, 2], it will yield 0, 2, 0, 2, 0, 2, 0, 2 indefinitely.

configure_dask_scheduler(scheduler_address, name, workers_per_gpu=0, nproc=1, gpu_backend='opencl')

Configure the Dask scheduler by reusing existing or creating a new cluster.

If scheduler_address is None then creates a local Dask cluster else returns scheduler_address unchanged and the callee is responsible for cluster cleanup.

When creating a local GPU cluster on a machine with multiple GPUs, it will start workers which each can only see a single GPU.

Parameters:

Name	Type	Description	Default
scheduler_address	str | Cluster | None	Address of the Dask scheduler to connect to, or None for local cluster.	required
name	str	Name for the Dask cluster.	required
workers_per_gpu	int	Number of workers per GPU. If > 0, a GPU cluster will be configured otherwise a CPU cluster.	0
nproc	int	Number of processes to use per worker for CPU support.	1
gpu_backend	GpuBackend	GPU backend to use for local GPU cluster setup.	'opencl'

Raises:

Type	Description
ImportError	If GPU support is requested but the selected backend is unavailable.
ValueError	If multiple GPUs are detected but the vendor is unsupported.

Yields:

Type	Description
str | Cluster	A Dask Cluster instance or a string address for the scheduler.

detect_available_gpus(gpu_backend='opencl', cuda_devices=None)

Detect available GPU IDs.

For CUDA backend: CuPy handles CUDA_VISIBLE_DEVICES/ROCR_VISIBLE_DEVICES internally. For OpenCL backend: Checks environment variables before falling back to OpenCL discovery.

Parameters:

Name	Type	Description	Default
gpu_backend	GpuBackend	Backend used for GPU discovery.	'opencl'
cuda_devices	list[int] | None	Optional list of CUDA device IDs to use instead of detecting. If provided, bypasses cuda_available() check. Useful for testing.	None

Returns:

Type	Description
list[int]	List of available GPU IDs. Empty list means no GPU detected.