确定 dask 计算了多少次答案

【问题标题】：Determine how many times dask computed something确定 dask 计算了多少次
【发布时间】：2018-12-15 09:06:22
【问题描述】：

问题

我想知道是否可以使用 dask（特别是 dask 数组）来知道是否以及何时计算了某些内容。我正在考虑单元测试想知道 dask 计算了多少次数组。类似于模拟对象，知道它们被调用了多少次。这样的事情已经存在了吗？如果没有，有没有比制作custom callback 更好的方法？如果这不存在，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

【问题讨论】：

标签： python dask

【解决方案1】：

有多种方法可以在执行时触发。

一种是指定自定义调度程序：

def my_scheduler(dsk, keys, **kwargs):
    print('computing!')
    return dask.get(dsk, keys, **kwargs)

with dask.config.set(scheduler=my_scheduler):
    ...

Custom callbacks，就像你的建议一样，也很容易实现。

如果您只使用 dask 数组，那么您可以查看 array plugins

测试套件中还使用了多种其他方法。

【讨论】：