Dask-ml ParallelPostFit未使用分布式导致本地内存错误及预测结果保存问题

• 你直接使用原生scikit-learn模型对Dask数组执行predict时，预测逻辑会默认在客户端本地执行，需要把全量x拉取到客户端内存，根本没有提交到分布式集群运行。数据量小时内存可以容纳所以能正常执行，数据量超过客户端内存上限就会触发OOM，和集群配置无关。
• 你调用client.persist(t)时，t本身是本地生成的numpy数组，不是分布式Dask对象，所以persist操作不会触发集群计算，后续转DataFrame写parquet依然需要操作本地全量数据，进一步加剧内存压力。

1. 用Dask-ML的ParallelPostFit包装训练完成的模型 #

ParallelPostFit不需要重新训练模型，仅用于把预测逻辑分发到集群worker节点并行执行，无需把全量数据拉到客户端：

from dask_ml.wrappers import ParallelPostFit

# model为你已经训练完成的RandomForest模型
parallel_model = ParallelPostFit(model)

2. 执行分布式预测 #

用包装后的模型执行预测，此时输出的t是分布式Dask数组，计算会分散到各个worker节点执行：

x = da.from_array(i, chunks=100000) # 可根据worker内存调整chunks大小，单chunk控制在100MB左右最优
t = parallel_model.predict(x)

3. 直接保存预测结果到共享存储 #

不要提前persist全量预测结果，直接转Dask DataFrame后写入parquet，Dask会边计算边写入，不需要全量数据加载到内存：

df = dd.from_array(t, columns=["prediction"])
# 注意存储路径必须是所有worker节点都能访问的共享存储（如共享挂载目录、HDFS、对象存储路径等）
df.to_parquet("共享存储路径/xy.parquet", write_index=False, engine="pyarrow")

• 调整输入数组的分块大小：避免分块过大导致worker单任务内存不足，也避免分块过小导致调度开销过高，单分块大小控制在64MB~256MB区间最优。
• 不要使用本地路径作为输出路径：如果输出路径是客户端本地路径，Dask会把所有worker的计算结果拉回客户端再写入，依然会触发客户端OOM，必须使用所有节点可访问的共享存储路径。
• 如遇写入冲突，可指定partition_on参数或者关闭overwrite避免多个worker同时写同一文件。

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