如何在不触发即时计算的前提下将Dask delayed对象持久化到内存

当前Dask官方没有内置persist(compute=False)的参数，你需要的「先标记持久化需求、实际使用时再触发对应分支计算」的需求，可以通过以下官方推荐的方案实现：

方案1：按需调用dask.persist批量持久化指定节点（最推荐） #

你不需要在构建链路的时候提前调用.persist()方法，该方法默认会立即触发异步计算。只需要先构建全量delayed链路，确定要使用的分支后，把需要缓存的中间节点和最终节点一起传给dask.persist即可，Dask会自动跳过无关分支的计算，仅计算目标链路的节点并缓存到内存：

# 第一步：仅构建delayed链路，不调用persist，不会触发任何计算
a1 = build_a1()
a2 = build_a2(a1)
a3 = build_a3(a2)

b1 = build_b1()
b2 = build_b2(b1)
b3 = build_b3(b2)

# 第二步：确定使用a分支时，批量持久化需要缓存的节点
from dask import persist
from dask.distributed import wait

# 仅触发a链路计算，同时将a1/a2/a3缓存到内存
a1, a2, a3 = persist(a1, a2, a3)
# 可选：等待所有节点计算完成再执行后续逻辑
wait([a1, a2, a3])

# 后续调用compute直接读取内存缓存，无需重复计算
a3_result = a3.compute()
a1_result = a1.compute()

方案2：封装装饰器实现函数级持久化标记 #

如果你希望用装饰器的形式标记某类函数的输出默认需要持久化，可以自行封装轻量装饰器：

import dask

def delayed_keep_in_memory(**dask_kwargs):
    def decorator(func):
        delayed_func = dask.delayed(func, **dask_kwargs)
        def wrapper(*args, **kwargs):
            res = delayed_func(*args, **kwargs)
            # 给delayed对象打自定义持久化标记
            res.keep_in_memory = True
            return res
        return wrapper
    return decorator

# 用法和你预期的完全一致
@delayed_keep_in_memory
def build_a1():
    # 函数原有逻辑
    return ...

使用时只需要在目标链路构建完成后，遍历依赖节点筛选出所有带keep_in_memory标记的对象，统一调用persist即可。 #

内容的提问来源于stack exchange，提问作者epizut