Невозможно передать функцию устройства Numba cuda в удаленную функцию Python Ray для распределенных вычислений - PullRequest
Новая
       13

Невозможно передать функцию устройства Numba cuda в удаленную функцию Python Ray для распределенных вычислений

0 голосов
/ 20 июня 2019

Я использую Python Ray и Numba cuda для распределенных вычислений на GPU.

Я бы хотел передать функцию устройства Numba cuda функции удаленного управления Ray, но она не работает.

Та же версия прекрасно работает для многопроцессорной обработки.

Вот код:

Эта часть работает. Он определяет ядро ​​cuda, а затем использует многопроцессорную обработку для завершения вычислений. (Тем не менее, та же версия не работает на Ray) 

import multiprocessing
import numpy as np
import os
from numba import cuda

# cuda kernel function
def cuda_kernel(fun):
    result = cuda.device_array(10000,dtype=np.float64)

    @cuda.jit
    def kernel(result):
        thread_id = cuda.grid(1)
        if thread_id < 10000:
            result[thread_id] = fun(thread_id)

    # Configure the blocks
    threadsperblock = 16
    blockspergrid = (10000 + (threadsperblock - 1)) // threadsperblock

    # Start the kernel 
    kernel[blockspergrid, threadsperblock](result)

    result = result.copy_to_host()

    return result

# multiprocessing to start the cuda kernel function
def multi_processing():
    os.environ["CUDA_VISIBLE_DEVICES"] = '0'
    result = cuda_kernel(fun)
    np.save(os.getcwd()+'/result', result)

# start multi-processing to allocate     
p = multiprocessing.Process(target = multi_processing)
p.daemon = True
p.start()

# define a device function that should be used by cuda kernel
@cuda.jit(device=True)
def fun(x):
    return x**2

np.load(os.getcwd()+'/result'+'.npy')

Затем я изменяю код на версию Ray следующим образом:

Он имеет те же функции, что и предыдущий, но использует Ray. 

import multiprocessing
import numpy as np
import os
from numba import cuda
import ray

# initialize ray
ray.shutdown()
ray.init()

# cuda kernel function
@ray.remote(num_gpus=1)
def cuda_kernel(fun):
    result = cuda.device_array(10000,dtype=np.float64)

    @cuda.jit
    def kernel(result):
        thread_id = cuda.grid(1)
        if thread_id < 10000:
            result[thread_id] = fun(thread_id)

    # Configure the blocks
    threadsperblock = 16
    blockspergrid = (10000 + (threadsperblock - 1)) // threadsperblock

    # Start the kernel 
    kernel[blockspergrid, threadsperblock](result)

    result = result.copy_to_host()

    return result

# use ray instead of multiprocessing
def ray_process():
    os.environ["CUDA_VISIBLE_DEVICES"] = '0'
    result = ray.get(cuda_kernel.remote(fun))
    np.save(os.getcwd()+'/result', result)

ray_process()

# define a device function that should be used by cuda kernel
@cuda.jit(device=True)
def fun(x):
    return x**2

np.load(os.getcwd()+'/result'+'.npy')

Это выдает ошибку: 

numba.errors.TypingError: Failed in nopython mode pipeline (step: nopython frontend)
Untyped global name 'fun': cannot determine Numba type of <class 'numba.cuda.compiler.DeviceFunctionTemplate'>

File "<ipython-input-6-4c4236d93522>", line 20:
<source missing, REPL/exec in use?>

Это немного странно, и я полагаю, что проблема связана с другой функциональностью Ray и многопроцессорностью. У кого-нибудь есть какие-то решения (я должен передать функцию ядру cuda)?

...