使用多处理时避免重新编译 numba 代码

Posted

技术标签:

【中文标题】使用多处理时避免重新编译 numba 代码【英文标题】:Avoid recompilation of numba code when using multiprocessing 【发布时间】:2020-12-31 16:05:28 【问题描述】:

我一直在使用 numba 进行多处理。

唯一的问题 - numba 会分别重新编译每个进程的代码。

(当进程数等于物理CPU数时,问题不大, 但如果不是这种情况,那就是一个巨大的!)

有没有办法让numba编译一次代码, 然后跨进程边界共享编译的工件?

例子-

from multiprocessing import Process
from time import time, sleep
from numba import njit


@njit
def child():
    pass


if __name__ == "__main__":
    ps = [Process(target=child) for _ in range(100)]
    for p in ps:
        p.start()
    s = time()
    for p in ps:
        p.join()
    print("compile time:", time() - s)
compile time: 19.10037922859192

CPU 使用率在所有内核上固定为 100%。 我试过 numba 的 cache=True,但不幸的是我的代码无法缓存。

/Users/dev/PycharmProjects/trading/tradingdo/strategy.py:91: NumbaWarning: Cannot cache compiled function "_strategy1" as it uses dynamic globals (such as ctypes pointers and large global arrays)
  @njit

【问题讨论】:

多处理不能跨进程共享任何内存,所以除非你能以某种方式腌制你的函数,否则它可能是不可能的。也许你必须尝试 Cython 谢谢,我已经尝试过 Cython,但对于我的特定应用来说,它实际上速度较慢! (即使在注释了所有类型并通过其光滑的 html ispector 确保最小的 python 接口之后)。 至于酸洗,numba的缓存真的可以酸洗吗?另外,有没有办法首先访问我的代码中的缓存? 【参考方案1】:

在支持 fork() 的系统上 (Linux),这很容易 -

在启动进程之前只需编译一次函数 - 这将使 numba 缓存编译器输出,就像它通常那样。

但是由于 fork 的写时复制魔法,缓存会自动与子进程共享!

不清楚的是如何在没有适当fork() 支持的系统上执行此操作。 numba的缓存可以腌制吗?

from multiprocessing import Process
from time import time, sleep
from numba import njit


@njit
def child():
    pass


if __name__ == "__main__":
    child() # this will do it

    ps = [Process(target=child) for _ in range(100)]
    for p in ps:
        p.start()
    s = time()
    for p in ps:
        p.join()
    print("compile time:", time() - s)

compile time: 0.011722326278686523

numba 的nogil 也值得一看。这可以消除对进程的需要,并且线程共享 numba 编译缓存就好了

【讨论】:

线程总是让事情变得更简单,但如果你的子方法是 cpu 密集型的,它们不会给你同样的速度。 nogil 让您可以利用多个内核,从而减轻 python 对线程的限制 真棒 - 不错的提示!

以上是关于使用多处理时避免重新编译 numba 代码的主要内容,如果未能解决你的问题,请参考以下文章

如何避免每次重新编译所有文件?

JAX:避免对沿一个轴使用不同数量的元素评估的函数进行即时重新编译

避免使用 makefile 进行不必要的重新编译

确保重新编译时处理定义

强化学习技巧五:numba提速python程序

与纯 Python 代码相比,Numba njit 编译器会导致计算不同的数字?