控制 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