Dask-cuda: GPU topology

[gram526]

Hi,
I am having issues with splitting a dak-cuda pipeline equally between GPUs. Here is the pseudo code:

import dask.array as da
import numpy as np
from dask.distributed import Client
from dask_cuda import LocalCUDACluster

cluster = LocalCUDACluster(n_workers=2)
with Client(cluster) as client:
    chunk = da.from_zarr(input_file_name, chunks=(512, 512, 512))
    out = da.map_overlap(some_function, chunk, depth={0: 64, 1: 64, 2: 64}, boundary='reflect',
                   trim=False, align_arrays=False, chunks=(512, 512, 512, 3), new_axis=3, meta=np.array((), dtype=np.int32))
    out.to_zarr(output_file_name, overwrite=True)

The function in question uses cupy.

This code works as expected on a server with two K80 gpu by splitting the task between the 2 gpus (CUDA version 11.2). The topology of the server is:

    GPU0    GPU1    CPU Affinity    NUMA Affinity

GPU0 X PIX 1,3,5,7,9,11 1
GPU1 PIX X 1,3,5,7,9,11 1

However, on a workstation with two m6000, the task is never split between the 2 gpus (CUDA version 11.3). The workstation has this topology:

    GPU0       GPU1 CPU Affinity    NUMA Affinity

GPU0 X PHB 0-11,24-35 0
GPU1 PHB X 0-11,24-35 0

The two gpus are visible on dask's server, but only one is used (not always the same one) every time I run the code. Could the difference be the GPU topology?

[pentschev]

I can't see what could be the cause for this, but unfortunately your code is incomplete so I can't test it on my end. I would suggest running one of Dask-CUDA's benchmarks to test that it works fine. For example:

$ python dask_cuda/benchmarks/local_cupy.py -d 0,1
Roundtrip benchmark
--------------------------
Operation          | transpose_sum
User size          | 10000
User second size   | 1000
User chunk-size    | 2500
Compute shape      | (10000, 10000)
Compute chunk-size | (2500, 2500)
Ignore-size        | 1.00 MiB
Protocol           | tcp
Device(s)          | 0,1
Worker Thread(s)   | 1
==========================
Wall-clock         | npartitions
--------------------------
1.35 s             | 16
1.14 s             | 16
1.21 s             | 16
==========================
(w1,w2)            | 25% 50% 75% (total nbytes)
--------------------------
(00,01)            | 131.90 MiB/s 261.46 MiB/s 314.96 MiB/s (524.52 MiB)
(01,00)            | 139.90 MiB/s 260.52 MiB/s 314.43 MiB/s (524.52 MiB)

In the example above, I'm specifying two workers on devices 0 and 1 (-d 0,1), and the last two lines prints the bandwidth between each worker pair, in the case above (00,01) and (01,00). If the above works, it may be that your data is too small or code is too specific and for some reason Dask can't parallelize to both workers.

[gram526]

The benchmarks didn't work because I didn't have rmm module installed. The problem is it doesn't seem to be compatible with CUDA 11.3. I downgraded to 11.2, everything seems to work fine.

[pentschev]

Thanks for the reply @gram526 . That's right, we currently don't officially support and build packages for CUDA 11.3. For future reference, you can always check the versions we support in https://rapids.ai/start.html .