python - 如何使用Dask在GPU上运行python代码？

Question

我有一些使用 Numba cuda.jit 的代码，以便我可以在 GPU 上运行，如果可能的话，我想在其之上分层 dask。

示例代码

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
from numba import cuda, njit
import numpy as np
from dask.distributed import Client, LocalCluster


@cuda.jit()
def addingNumbersCUDA (big_array, big_array2, save_array):
    i = cuda.grid(1)
    if i < big_array.shape[0]:
        for j in range (big_array.shape[1]):
            save_array[i][j] = big_array[i][j] * big_array2[i][j]


if __name__ == "__main__":
    cluster = LocalCluster()
    client = Client(cluster)

    big_array = np.random.random_sample((100, 3000))
    big_array2  = np.random.random_sample((100, 3000))
    save_array = np.zeros(shape=(100, 3000))

    arraysize = 100
    threadsperblock = 64
    blockspergrid = (arraysize + (threadsperblock - 1))

    d_big_array = cuda.to_device(big_array)
    d_big_array2 = cuda.to_device(big_array2)
    d_save_array = cuda.to_device(save_array)

    addingNumbersCUDA[blockspergrid, threadsperblock](d_big_array, d_big_array2, d_save_array)

    save_array = d_save_array.copy_to_host()

如果我的函数 addingNumbersCUDA 没有使用任何 CUDA，我只需将 client.submit 放在我的函数前面(以及后面的收集)，它就会起作用。但是，由于我使用 CUDA，因此将“提交”放在函数前面是行不通的。 dask 文档说您可以以 GPU 为目标，但尚不清楚如何在实践中实际设置它。如果可能的话，我将如何设置我的函数以将 dask 与目标 GPU 一起使用并与 cuda.jit 一起使用？

Answer 1

您可能想查看Dask's documentation on GPUs

But, since I'm using CUDA putting submit in front of the function doesn't work.

没有什么特殊原因可以解释为什么会出现这种情况。所有 Dask 都会在另一台计算机上运行您的功能。它不会以任何方式改变或修改您的功能。