如何在Dask同步计算中限制内存使用并保留并行性？

1. 调整Dask资源配置限制单进程/线程内存 #

直接通过Dask的配置参数限制每个工作进程的内存上限，同时控制并行工作的进程/线程数量，让Dask自动在内存阈值内调度任务：

from dask.distributed import Client

# 初始化Client，设置每个进程内存限制与工作进程数
# 示例：每个进程限2GB，开2个进程，总内存占用控制在约4GB（预留系统内存）
client = Client(n_workers=2, memory_limit="2GB")

# 原有数据加载代码
array = da.concatenate([
    da.from_delayed(
        dask.delayed(np.load)(path),
        shape=(size, window_len, vocab_size),
        dtype=np.float32
    )
    for path, size in zip(shard_paths, shard_sizes)
])

result = da.map_blocks(fn, array, drop_axis=[-1]).compute()

如果使用单进程多线程模式，可直接通过配置调整：

import dask

dask.config.set({
    "num_workers": 4,  # 并行线程数
    "memory_limit": "1GB"  # 每个线程可用内存配额
})

2. 手动分批处理分块，兼顾并行与内存控制 #

若需要更精细控制每次加载的分块数量，可将数组拆分为独立分块的Dask数组列表，分批提交计算——每批内部并行处理，批与批之间串行释放内存：

import dask

# 先创建每个分块的Dask数组（不提前concat）
chunk_arrays = [
    da.from_delayed(
        dask.delayed(np.load)(path),
        shape=(size, window_len, vocab_size),
        dtype=np.float32
    )
    for path, size in zip(shard_paths, shard_sizes)
]

# 设置每批处理的分块数量，示例为3
batch_size = 3
results = []

# 分批处理逻辑
for i in range(0, len(chunk_arrays), batch_size):
    batch = chunk_arrays[i:i+batch_size]
    # 处理当前批次并计算
    batch_result = da.map_blocks(fn, da.concatenate(batch), drop_axis=[-1]).compute()
    results.append(batch_result)

# 合并所有批次结果
final_result = np.concatenate(results)

这种方式既保留了单批次内的并行性，又严格控制了同时加载的分块数量，避免内存溢出。

3. 利用Dask Delayed的分组并行 #

通过dask.delayed对分块分组，每组内并行执行处理逻辑，组之间串行执行，完全手动控制并行粒度：

from dask import delayed

# 定义单分块处理函数
def process_single_chunk(path, size):
    arr = np.load(path)
    return fn(arr)

# 按批次分组分块路径与大小
batch_size = 3
batches = [
    list(zip(shard_paths[i:i+batch_size], shard_sizes[i:i+batch_size]))
    for i in range(0, len(shard_paths), batch_size)
]

# 批次内并行、批次间串行处理
delayed_results = []
for batch in batches:
    batch_tasks = [delayed(process_single_chunk)(p, s) for p, s in batch]
    delayed_batch_result = delayed(np.concatenate)(batch_tasks)
    delayed_results.append(delayed_batch_result)

# 计算并合并最终结果
final_result = np.concatenate(dask.compute(*delayed_results))

内容的提问来源于stack exchange，提问作者ondra.cifka