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

Question

我有一些使用Numba的代码库达.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 after），它就可以工作了。但是，由于我使用的是CUDA，将submit放在函数前面不起作用。dask文档说你可以针对gpu，但不清楚如何实际设置它。如何设置我的函数以使用目标gpu和库达.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都在另一台计算机上运行你的函数。它不会以任何方式改变或修改你的功能。在