自定义图中的 Dask 机会性缓存

Question

我有一个自定义 DAG，例如：

dag = {'load': (load, 'myfile.txt'),
       'heavy_comp': (heavy_comp, 'load'),
       'simple_comp_1': (sc_1, 'heavy_comp'),
       'simple_comp_2': (sc_2, 'heavy_comp'),
       'simple_comp_3': (sc_3, 'heavy_comp')}

我正在寻找计算键 simple_comp_1、simple_comp_2 和 simple_comp_3，我执行如下操作，

import dask
from dask.distributed import Client
from dask_yarn import YarnCluster

task_1 = dask.get(dag, 'simple_comp_1')
task_2 = dask.get(dag, 'simple_comp_2')
task_3 = dask.get(dag, 'simple_comp_3')
tasks = [task_1, task_2, task_3]

cluster = YarnCluster()
cluster.scale(3)
client = Client(cluster)
dask.compute(tasks)
cluster.shutdown()

看来，如果没有缓存，这 3 个键的计算将导致 heavy_comp 的计算也 3 次。由于这是一个繁重的计算，我尝试从here 实现机会缓存，如下所示：

from dask.cache import Cache
cache = Cache(2e9)
cache.register()

但是，当我尝试打印缓存的结果时，我什么也没得到：

>>> cache.cache.data
[]
>>> cache.cache.heap.heap
{}
>>> cache.cache.nbytes
{}

我什至尝试将缓存大小增加到 6GB，但没有效果。难道我做错了什么？如何让 Dask 缓存密钥 heavy_comp 的结果？

Answer 1

扩展 MRocklin 的答案并在问题下方的 cmets 中格式化代码。

一次计算整个图表可以按照您的预期进行。 heavy_comp 只会执行一次，这就是你想要的。考虑您在由空函数定义完成的 cmets 中提供的以下代码：

def load(fn):
    print('load')
    return fn

def sc_1(i):
    print('sc_1')
    return i

def sc_2(i):
    print('sc_2')
    return i

def sc_3(i):
    print('sc_3')
    return i

def heavy_comp(i):
    print('heavy_comp')
    return i

def merge(*args):
    print('merge')
    return args

dag = {'load': (load, 'myfile.txt'), 'heavy_comp': (heavy_comp, 'load'), 'simple_comp_1': (sc_1, 'heavy_comp'), 'simple_comp_2': (sc_2, 'heavy_comp'), 'simple_comp_3': (sc_3, 'heavy_comp'), 'merger_comp': (merge, 'sc_1', 'sc_2', 'sc_3')}

import dask
result = dask.get(dag, 'merger_comp')
print('result:', result)

它输出：

load
heavy_comp
sc_1
sc_2
sc_3
merge
result: ('sc_1', 'sc_2', 'sc_3')

如您所见，“heavy_comp”只打印了一次，表明函数heavy_comp只执行了一次。

Answer 2

核心 Dask 库中的机会缓存仅适用于单机调度器，不适用于分布式调度器。

但是，如果您一次只计算整个图形，Dask 将智能地保留中间值。如果您想保留某些值，也可以查看persist 函数。