如何使用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 放在我的函数前面（连同后面的 gather），它就可以工作了。但是，因为我使用的是 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.

没有特别的原因会出现这种情况。 All Dask 在另一台计算机上完成运行你的功能。它不会以任何方式改变或修改您的功能。

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

How to use Dask to run python code on the GPU?

python

dask

dask-distributed