Я пытаюсь использовать Numba и Dask для ускорения медленных вычислений, аналогичных вычислению ядра оценка плотности огромного набора точек. Мой план состоял в том, чтобы написать вычислительно затратную логику в функции jited, а затем разделить работу между ядрами ЦП, используя dask. Я хотел использовать функцию nogil функции numba.jit, чтобы я мог использовать бэкэнд dask threading, чтобы избежать ненужных копий памяти входных данных (которые очень велики).
К сожалению, Dask не приведет к ускорению, если я не использую планировщик 'processes'. Если вместо этого я использую ThreadPoolExector, то я вижу ожидаемое ускорение.
Вот упрощенный пример моей проблемы:
import os
import numpy as np
import numba
import dask
CPU_COUNT = os.cpu_count()
def render_internal(size, mag):
"""mag is the magnification to apply
generate coordinates internally
"""
coords = np.random.rand(size, 2)
img = np.zeros((mag, mag), dtype=np.int64)
for i in range(len(coords)):
y0, x0 = coords[i] * mag
y1, x1 = int(y0), int(x0)
m = 1
img[y1, x1] += m
jit_render_internal = numba.jit(render_internal, nogil=True, nopython=True)
args = 10000000, 100
print("Linear time:")
%time linear_compute = [jit_render_internal(*args) for i in range(CPU_COUNT)]
delayed_jit_render_internal = dask.delayed(jit_render_internal)
print()
print("Threads time:")
%time dask_compute_threads = dask.compute(*[delayed_jit_render_internal(*args) for i in range(CPU_COUNT)])
print()
print("Processes time:")
%time dask_compute_processes = dask.compute(*[delayed_jit_render_internal(*args) for i in range(CPU_COUNT)], scheduler="processes")
И вот вывод на моей машине:
Linear time:
Wall time: 1min 17s
Threads time:
Wall time: 1min 47s
Processes time:
Wall time: 7.79 s
Как для обработки, так и для обработки потоков я вижу полное использование всех ядер ЦП, как и ожидалось. Но нет ускорения для многопоточного бэкэнда. Я почти уверен, что джиттинговая функция jit_render_internal на самом деле не выпускает GIL.
Мои два вопроса:
- Если ключевое слово
nogilпередается вnumba.jitи GIL не может быть освобожден, почему не возникает ошибка? - Почему код, как я его написал, не выпускает GIL? Все вычисления встроены в функцию, и нет возвращаемого значения.
