控制 worker 数量

用户可以通过显式设置 CUDA_VISIBLE_DEVICES 来限制活动仅限于特定 GPU;对于 LocalCUDACluster,这可以作为关键字参数提供。例如,要限制活动仅限于前两个索引的 GPU:

from dask_cuda import LocalCUDACluster

cluster = LocalCUDACluster(CUDA_VISIBLE_DEVICES="0,1")

LocalCUDACluster 还可以接受 n_workers 参数,它将限制活动仅限于 CUDA_VISIBLE_DEVICES 中列出的前 N 个 GPU。此参数可以单独使用,也可以与 CUDA_VISIBLE_DEVICES 结合使用:

cluster = LocalCUDACluster(n_workers=2)                                # will use GPUs 0,1
cluster = LocalCUDACluster(CUDA_VISIBLE_DEVICES="3,4,5", n_workers=2)  # will use GPUs 3,4

使用 dask cuda worker 时,必须将 CUDA_VISIBLE_DEVICES 作为环境变量提供。

$ dask scheduler
distributed.scheduler - INFO -   Scheduler at:  tcp://127.0.0.1:8786

$ CUDA_VISIBLE_DEVICES=0,1 dask cuda worker 127.0.0.1:8786

也可以通过 GPU 的 UUID 进行选择,这可以通过 NVIDIA System Management Interface 获取。

$ nvidia-smi -L
GPU 0: Tesla V100-SXM2-32GB (UUID: GPU-dae76d0e-3414-958a-8f3e-fc6682b36f31)
GPU 1: Tesla V100-SXM2-32GB (UUID: GPU-60f2c95a-c564-a078-2a14-b4ff488806ca)

然后可以将这些 UUID 传递给 CUDA_VISIBLE_DEVICES,以代替 GPU 索引。

cluster = LocalCUDACluster(CUDA_VISIBLE_DEVICES="GPU-dae76d0e-3414-958a-8f3e-fc6682b36f31")
$ CUDA_VISIBLE_DEVICES="GPU-dae76d0e-3414-958a-8f3e-fc6682b36f31" \
> dask cuda worker 127.0.0.1:8786