DOCSEARCHbyAlgolia
Submit Docs

Dask CUDA

repository·main·Indexed 43 days ago

https://github.com/rapidsai/dask-cuda

Utilities for the deployment and management of Dask workers on CUDA-enabled systems, specifically targeting multi-GPU environments on Linux. It provides tools like LocalCUDACluster for local GPU cluster setup, CUDAWorker for GPU resource and context management, and the dask-cuda CLI for worker launching and configuration. The library includes advanced GPU memory management via RMM (RAPIDS Memory Manager) and multi-level memory spilling between CUDA device memory, host memory, and disk.

Tokens
15.3K
Snippets
59
Records
71
Agent score
55%

What's inside Dask CUDA

  1. What is Dask-CUDA and when to use it

    main

    Dask-CUDA is an extension of Dask's LocalCluster and Worker designed specifically for distributed GPU workloads. While standard Dask can run GPU-accelerated libraries like cuDF, CuPy, and Numba, Dask-CUDA provides specialized orchestration for multi-GPU nodes that standard Dask Distributed does not handle automatically.

    Key features include:

    • Automatic per-GPU worker instantiation: Automatically launches one worker per available GPU on a node using LocalCUDACluster or the dask cuda worker CLI.
    • Automatic CPU affinity: Sets CPU affinity for each GPU to optimize memory transfer paths.
    • Automatic InfiniBand selection: When using UCX, it selects the optimal InfiniBand device for each GPU.
    • GPU-to-Host memory spilling: Supports spilling GPU memory to host memory when utilization limits are reached.
    • Optimized GPU memory allocation: Integrates with RAPIDS Memory Manager (RMM) to allocate per-GPU memory pools when using UCX, avoiding expensive memory buffer mappings.

  2. Enable cuDF native spilling for ETL workflows

    main

    For ETL workflows using Dask cuDF (Dask DataFrames), it is recommended to use native cuDF spilling instead of Dask-CUDA's default spilling.

    Why use cuDF spilling? Dask-CUDA's default spilling tracks task outputs as whole units, meaning intermediate data inside a task cannot be spilled until the task finishes. In contrast, cuDF spilling allows individual device buffers to be spilled or unspilled while a task is still running.

    To enable this in LocalCUDACluster, use the enable_cudf_spill argument. For the CLI, use the --enable-cudf-spill flag.

    from distributed import Client
from dask_cuda import LocalCUDACluster

cluster = LocalCUDACluster(n_workers=10, enable_cudf_spill=True)
client = Client(cluster)

  3. Understand the scope of Dask CUDA

    main

    Dask CUDA is designed for the deployment and management of Dask workers on CUDA-enabled systems, specifically helping in multi-GPU environments.

    Important Limitations:

    • It does not automatically convert existing Dask code to run on GPUs.
    • It is a management layer for workers; parallelizing GPU libraries like RAPIDS or CuPy with Dask requires additional effort and specific library integration.

  5. Enable UCX communication with LocalCUDACluster (Automatic Configuration)

    main

    For routine deployments, use automatic UCX transport selection. This requires Dask-CUDA >= 22.02 and UCX >= 1.11.1. You only need to specify the protocol="ucx" and an interface. UCX will then decide which transports to use.

    To ensure RDMACM is used on InfiniBand systems, specify interface="ib0". On systems without InfiniBand, you can omit the interface argument or specify a different one.

    from dask.distributed import Client
from dask_cuda import LocalCUDACluster

cluster = LocalCUDACluster(
    protocol="ucx",
    interface="ib0",
    rmm_pool_size="1GB"
)
client = Client(cluster)

  6. Enable UCX communication with LocalCUDACluster (Manual Configuration)

    main

    Use manual configuration when benchmarking specific UCX transports. You must explicitly enable each transport via arguments during LocalCUDACluster construction. This is useful for testing the performance of specific paths like NVLink or InfiniBand.

    from dask.distributed import Client
from dask_cuda import LocalCUDACluster

cluster = LocalCUDACluster(
    protocol="ucx",
    interface="ib0",
    enable_tcp_over_ucx=True,
    enable_nvlink=True,
    enable_infiniband=True,
    enable_rdmacm=True,
    rmm_pool_size="1GB"
)
client = Client(cluster)

  8. Optimize multi-GPU performance by co-locating GPUs

    main

    When setting up multi-GPU environments, choose hardware (like NVIDIA DGX or cloud multi-GPU instances) where GPUs are co-located. This minimizes performance degradation caused by communication overheads and Host-to-Device/Device-to-Host transfers. Co-located GPUs often support accelerated networking like NVLink, which provides significantly higher throughput than traditional networking.

    You can control which GPUs are used by a LocalCUDACluster using the CUDA_VISIBLE_DEVICES environment variable or by specifying the number of workers.

    from dask_cuda import LocalCUDACluster

# Will use GPUs 0 and 1
cluster = LocalCUDACluster(n_workers=2)

# Will use GPUs 3 and 4
cluster = LocalCUDACluster(CUDA_VISIBLE_DEVICES="3,4")

  9. Enable UCX support in Pip installation

    main

    To enable accelerated communications with UCX via Pip, install the distributed-ucxx package. It is recommended to use the version-specific extras to match your CUDA version.

    Generic installation:

    python -m pip install dask-cuda distributed-ucxx

    Version-specific installation:

    # For CUDA 13
python -m pip install 'dask-cuda[cu13]' 'distributed-ucxx[cu13]'

# For CUDA 12
python -m pip install 'dask-cuda[cu12]' 'distributed-ucxx[cu12]'
    python -m pip install 'dask-cuda[cu13]' 'distributed-ucxx[cu13]'

  10. Leverage Accelerated Networking with UCX

    main

    To achieve maximum throughput and avoid costly Host-to-Device/Device-to-Host transfers, use UCX for communication. UCX can leverage accelerated networking hardware like NVLink.

    When initializing a LocalCUDACluster, set protocol="ucx" to allow UCX to automatically select the best available transports (e.g., NVLink) on your system.

    # Example of how the protocol flag changes communication
# Using standard TCP (slower):
# python local_cudf_merge.py -d 0,1 -p tcp -c 50_000_000 --rmm-pool-size 30GB

# Using UCX (faster with NVLink):
# python local_cudf_merge.py -d 0,1 -p ucx -c 50_000_000 --rmm-pool-size 30GB

  11. Enable and collect cuDF spilling statistics

    main

    When enable_cudf_spill is active, you can collect basic spill statistics to understand the performance of memory-intensive workflows.

    • In LocalCUDACluster, set cudf_spill_stats=1.
    • In the dask cuda worker CLI, use --cudf-spill-stats 1.

    To print these statistics during a workflow, you can submit a function that accesses the global spill manager.

    # Enabling statistics in cluster setup
cluster = LocalCUDACluster(n_workers=10, enable_cudf_spill=True, cudf_spill_stats=1)

# Function to print statistics within a workflow
def spill_info():
    from cudf.core.buffer.spill_manager import get_global_manager
    print(get_global_manager().statistics)

client.submit(spill_info)