Dask波场过滤应用扩缩容故障排查求助

任务函数 #

def filter_wavefield(pos, butter_filter, wave):
    filtered = signal.sosfilt(butter_filter, wave[pos,:]).astype("float32")
    return filtered

数据已成功加载为Dask数组，但扩缩容时遇到系统冻结、内核崩溃问题，以下是尝试过的三种提交方式及问题：

尝试的三种任务提交方式 #

1. 预散射数据并传入Future #

wave_future = client.scatter(wave_on_slice_channel, broadcast=True)

coord_list = [i for i in range(nelem*ngll)]
coord_bag = db.from_sequence(coord_list, npartitions=500)
filtered = coord_bag.map(filter_wavefield, butter_filter, wave_future)
filtered_waves = filtered.compute()

0.5TiB数据时运行正常，但1.35TiB数据执行compute()后系统冻结，仪表盘失效，随后Python内核崩溃。

2. 持久化大型数据 #

提前将数据持久化到分布式内存（500核，总内存>3.5TiB，仪表盘显示1.35TiB已存入）：

wave = wave_on_slice_channel.persist()
wait(wave)

coord_list = [i for i in range(nelem*ngll)]
coord_bag = db.from_sequence(coord_list, npartitions=500)
filtered = coord_bag.map(filter_wavefield, butter_filter, wave)
filtered_waves = filtered.compute()

执行compute()后仍出现系统冻结、内核崩溃问题。

3. For循环+submit提交任务 #

单独执行单个任务signal.sosfilt(butter_filter, wave[120,:])瞬间完成，但批量提交时：

filtered_futures = []
for pos in coord_list:
  filtered_futures.append(client.submit(filter_wavefield, pos, butter_filter, wave))

任务数增长到约2700个时系统冻结，坐标列表总长度为27482400。

问题根源与解决方案 #

核心问题 #

任务粒度太细：单个任务仅处理1个pos，总任务量达2700万，远超Dask调度器的处理能力，导致调度队列爆炸、内存耗尽，最终系统崩溃。

可行解决方案 #

1. 批量处理坐标，合并任务

修改任务函数支持批量处理pos，大幅减少任务总数：

def filter_wavefield_batch(pos_list, butter_filter, wave):
    filtered_batch = []
    for pos in pos_list:
        filtered = signal.sosfilt(butter_filter, wave[pos, :]).astype("float32")
        filtered_batch.append(filtered)
    return filtered_batch

创建Bag时将坐标按批次划分（例如每批次10000个pos）：

batch_size = 10000
coord_batches = [coord_list[i:i+batch_size] for i in range(0, len(coord_list), batch_size)]
coord_bag = db.from_sequence(coord_batches, npartitions=len(coord_batches))
filtered = coord_bag.map(filter_wavefield_batch, butter_filter, wave)
filtered_waves = filtered.compute()

推荐根据集群核数调整批次大小，使任务数保持在核数的2-4倍左右（500核对应1000-2000个任务），平衡并行度与调度开销。

2. 改用Dask数组map_blocks（推荐）

利用Dask数组的原生API处理，自动适配数据分区，无需手动管理坐标：

import numpy as np
from scipy import signal

def filter_row(row, butter_filter):
    return signal.sosfilt(butter_filter, row).astype("float32")

# 对wave的每一行应用过滤，保持原分区结构
filtered_wave = wave.map_blocks(
    lambda x: np.apply_along_axis(filter_row, 1, x, butter_filter),
    chunks=wave.chunks,
    dtype=np.float32
)
filtered_waves = filtered_wave.compute()

此方法充分利用Dask数组的优化逻辑，任务数与wave的分区数一致，调度效率最高，适合大规模数组处理场景。

3. 分批提交任务并流式获取结果

若需保留单pos处理逻辑，可通过as_completed分批提交、流式获取结果，避免客户端内存过载：

from dask.distributed import as_completed

batch_size = 10000
filtered_waves = []

for i in range(0, len(coord_list), batch_size):
    batch_pos = coord_list[i:i+batch_size]
    futures = [client.submit(filter_wavefield, pos, butter_filter, wave) for pos in batch_pos]
    # 逐个获取结果，释放内存
    for future in as_completed(futures):
        filtered_waves.append(future.result())

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