Dask并行计算效率低下且出现KilledWorker问题求助

核心问题分析 #

1. 小文件IO瓶颈：123000个极小的Parquet文件会导致严重的IO等待——即使Dask并行读取，文件系统的并发IO能力有限，worker多数时间处于等待状态，无法满负荷占用CPU。
2. 分区策略不合理：
   • 初始设置100个分区，每个分区包含上千个小文件，读取时IO开销远超计算开销；
   • 改用partition_size="100MB"后仅生成3个分区，单分区数据量过大，tsfresh计算时生成的中间数据直接撑爆worker内存，导致KilledWorker报错。
3. tsfresh计算的内存开销：tsfresh生成时间序列特征时会产生大量中间变量，单分区数据量超出worker内存限制时，进程会被系统强制杀死。
4. persist时机不当：提前对100个分区执行persist，可能导致集群内存被占用过多，挤占计算资源。

针对性解决方案 #

1. 先合并小Parquet文件

大量小文件是IO瓶颈的根源，先将其合并为少量大文件，从根本上降低IO开销：

• 用Dask读取所有小文件后，重新写入为按固定大小划分的大Parquet文件（比如1GB/个）；
• 写入时开启元数据文件，后续读取效率更高。

2. 优化分区策略

• 合并文件后，根据worker数量设置分区数（建议为worker数的2-4倍，比如10个worker设置20-40个分区），既保证任务并行度，又避免单分区数据量过大；
• 避免直接用partition_size处理极小文件，可根据业务维度（如时间、ID）分区，确保每个分区内的时间序列计算逻辑完整（比如一个分区包含若干完整ID的序列，避免跨分区ID导致的额外处理）。

3. 调整Worker资源配置

• 本机64GB内存，10个worker可将每个worker的内存限制设为6GB（预留10-15GB给系统和其他进程），避免内存过载；
• 保持threads_per_worker=1，规避GIL问题，充分利用多核CPU。

4. 优化tsfresh计算逻辑

• 检查extract_tsfresh_features函数，避免内存泄漏，可考虑分批处理单个ID的时间序列，减少单批次内存占用；
• 在tsfresh的配置中启用efficient=True，关闭冗余特征计算，降低内存开销；
• 只保留业务需要的特征，通过settings参数过滤不必要的特征生成器，减少计算量。

5. 调整persist的使用

• 不要过早对原始小文件数据集执行persist，合并文件后再根据情况决定是否persist；
• 若需persist，添加optimize_graph=True参数，让Dask优化任务图，减少冗余计算。

6. 利用Dask Dashboard定位问题

• 打开Dashboard查看任务流：如果多数任务处于waiting状态，说明是IO瓶颈；如果任务出现failed且伴随内存告警，说明单分区内存过载；
• 测试单个分区的计算：用dask_merged.get_partition(0).compute()验证单分区的计算时间和内存占用，确认分区大小是否合理。

优化后代码示例 #

if __name__ == "__main__":
    # 初始化集群，调整内存限制
    cluster = LocalCluster(n_workers=10, processes=True, threads_per_worker=1, memory_limit="6GB")
    client = Client(cluster)

    # 第一步：合并小文件为大文件
    dask_raw = dd.read_parquet(ROOT_PATH + r"\*\data.parquet",
                               columns=VAR_LIST,
                               engine='pyarrow',
                               ignore_metadata_file=False)
    
    # 按1GB分区合并写入
    dask_raw.repartition(partition_size="1GB").to_parquet(ROOT_PATH + r"\merged_data",
                                                          engine="pyarrow",
                                                          write_metadata_file=True)
    
    # 第二步：读取合并后的文件，设置合理分区数
    dask_merged = dd.read_parquet(ROOT_PATH + r"\merged_data",
                                  columns=VAR_LIST,
                                  engine='pyarrow')
    # 设置分区数为worker数的3倍
    dask_merged = dask_merged.repartition(npartitions=30)
    
    # 计算时间序列特征
    dask_df_features = dask_merged.map_partitions(extract_tsfresh_features,
                                                  column_id=VAR_ID,
                                                  column_sort=VAR_SORT,
                                                  settings=config_tsfresh_extractor)
    
    # 输出结果
    dask_df_features.to_parquet(OUTPUT, engine="pyarrow", append=False)

    client.close()
    cluster.close()

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