关于 Dask 分布式工人的 Numba jit答案

【问题标题】：Numba jit on Dask distrbuted workers关于 Dask 分布式工人的 Numba jit
【发布时间】：2021-03-18 06:31:00
【问题描述】：

我有一个包含几个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 代码的任何指针？

【问题讨论】：

标签： python dask jit numba dask-distributed

【解决方案1】：

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

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

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

【讨论】：

【解决方案2】：

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

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

【讨论】：