确定 dask 计算某事的次数

Question

问题

我想知道 dask （特别是 dask 数组）是否可以知道是否以及何时计算了某些内容。我在考虑单元测试想知道 dask 计算了多少次数组。类似于知道它们被调用多少次的模拟对象。这样的东西已经存在了吗？如果没有，有没有比制作自定义回调更好的方法？如果这不存在，那么 dask 核心开发人员是否有兴趣将其添加到核心 dask 以进行测试？

任何帮助深表感谢。

细节

假设我有一个函数，它接受一个 xarray DataArray，对它做一些事情，然后返回它。在某些情况下，dask 数组被隐式转换为 numpy 数组，包括新的 dask 用户不知道做某事的最佳 dask 友好方式。我想编写我的单元测试，以确保我或其他贡献者不会意外地损害函数的性能。考虑到测试数据通常是真实世界案例的简化/小型版本，并且在这些情况下可能看不到多次计算 dask 数组的性能损失，这一点尤为重要。

编辑：解决方案

这是我最终根据 MRocklin 的答案作为一个简单的解决方案所做的。

class CustomScheduler(object):

    def __init__(self, max_computes=1):

        self.max_computes = max_computes

        self.total_computes = 0

    def __call__(self, dsk, keys, **kwargs):

        self.total_computes += 1

        if self.total_computes > self.max_computes:

            raise RuntimeError("Too many dask computations were scheduled: {}".format(self.total_computes))

        return dask.get(dsk, keys, **kwargs)

然后我像这样使用它：

with dask.config.set(scheduler=CustomScheduler(0)):

    # dask array stuff

Answer 1

有多种方法可以在执行时触发。一种是指定一个自定义调度程序：def my_scheduler(dsk, keys, **kwargs):    print('computing!')    return dask.get(dsk, keys, **kwargs)with dask.config.set(scheduler=my_scheduler):    ...自定义回调，就像您建议的那样，也很容易实现。如果您专门使用 dask 数组，那么您可以查看数组插件测试套件中还使用了多种其他方法。