Dask 分布式工作者上的 Numba jit

Question

我有一个包含一些 numba.jit 和 numba.guvectorize 函数的模块。当我导入模块时，编译代码需要几秒钟。

我的问题是，确保代码已编译并可供工作人员进行计算的最佳方法是什么？或者事先编译它真的有什么好处吗？

我目前的解决方案如下，我认为它不是特别优雅，我不确定它是否有用：

import functools
from distributed import Client

import foo # module with numba code

@functools.partial
def callback():
    import foo

client = Client()
client.register_worker_callbacks(callback)

这样，代码在注册回调并准备好计算时编译。

所以这有必要吗？有没有更好的方法来处理这个问题？使用 dask?

处理 numba 代码的任何指示

Answer 1

您的假设是正确的，即在首次访问时将在每个 worker 上编译一个 numba 函数 - 如果您有显式编译，则在 import 中，或者如果隐式调用该函数时。这意味着使用该函数的任务的第一个运行会有额外的延迟。

如果“典型”工作流程没有任何严重缺陷，我会说不要让工作人员回调。通过强制提前编译，您将确保使用该函数的后续任务具有相似的性能。这在极少数情况下可能很有用，因为之前任务的完成时间用于 dask 的任务窃取机制，并且您可能希望仪表板的性能页面不包含编译任务。

您可以使用 worker 启动脚本实现预编译（参见 here）。您的解决方案更简单，但无法在工作人员重新启动后继续存在。

Answer 2

您可以设置 cache=True 将 JIT 编译的字节码存储到 __pycache__。

缓存版本将在第一次使用 JIT 编译函数后使用。参见 numba jit-compilation。

Dask 分布式工作者上的 Numba jit

Numba jit on Dask distrbuted workers

python

jit

numba

dask

dask-distributed