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

时间:2019-05-17 23:15:11

标签: python dask dask-distributed

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

示例代码

#!/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,因此无法在函数之前放置Submit。简短的文档说,您可以针对gpu,但是尚不清楚如何在实践中进行实际设置。我该如何设置我的函数以在目标gpu和cuda.jit中使用dask?

1 个答案:

答案 0 :(得分:2)

您可能需要浏览Dask's documentation on GPUs

  

但是,由于我使用的是CUDA,因此无法将提交放在函数前面。

没有特殊原因可以说明这种情况。所有Dask都会在另一台计算机上运行您的功能。它不会以任何方式更改或修改您的功能。