如何加速Slurm集群上Dask Delayed大规模随机森林计算？

架构老林

2026-6-17

问题背景

我需要重复运行封装了随机森林分类器的函数约10000次（每次执行随机采样），目前在Slurm调度的HPC集群上用Dask Delayed实现并行计算。脚本能正常运行但速度偏慢：500次运行耗时22分钟，推算10000次要7小时，而且大部分时间都花在dask.compute步骤。另外，增加Worker数量或者运行次数超过1000次时，会出现未管理内存过高、垃圾回收频繁的警告，求工作流或脚本的优化方案。

我的Python脚本（基于dask.delayed）

import pandas as pd
import numpy as np
from Bio import SeqIO
from sklearn.model_selection import train_test_split
from sklearn.ensemble import RandomForestClassifier
from sklearn import metrics
from sklearn.metrics import confusion_matrix

import dask
from dask.distributed import Client
from dask_mpi import initialize
import gc

import time

start = time.time()
print(start)

# CONNECT TO DISTRIBUTED CLUSTER
initialize()
client = Client()

# THIS FUNCTION IS PARALLELIZED ACROSS MULTIPLE WORKERS
@dask.delayed(nout=2)
def run_the_classifier(data_df_merged_balanced):
    data_df_merged_balanced = data_df_merged_balanced.groupby('Ecosystem_Type').sample(n=127)
    # Import train_test_split function
    
    X=data_df_merged_balanced.iloc[:,0:5787] # Features
    y=data_df_merged_balanced['Ecosystem_Type']  # Labels
    y = pd.get_dummies(y, dtype=int)
    # Split dataset into training set and test set
    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, stratify=y) # 80% training and 20% test

    # Create a Gaussian Classifier
    clf=RandomForestClassifier(n_estimators=10000,max_depth=5, class_weight='balanced_subsample')

    #Train the model using the training sets y_pred=clf.predict(X_test)
    clf.fit(X_train,y_train)

    # test the model using the 30% test dataset
    y_pred=clf.predict(X_test)

    # get accuracy
    accuracy = metrics.accuracy_score(y_test, y_pred)
    feature_imp = pd.Series(clf.feature_importances_,index=data_df_numerical.columns).sort_values(ascending=False)

    return accuracy, feature_imp

#~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
# A bunch of dataframe operations are here that produce data_df_merged_balanced - that is passed to the function
#~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

accuracy_dict = {}
importance_dict = {}

if __name__ == '__main__':
    for i in range(0, 500):
        accuracy, feature_imp = run_the_classifier(data_df_merged_balanced)
        accuracy_dict[i] = accuracy
        importance_dict[i] = feature_imp
        # client.run(gc.collect)
        print(i)
    time_sent = time.time()
    # print(time_sent)
    result_accuracy_dict = dask.compute(accuracy_dict)[0]
    result_imp_dict = dask.compute(importance_dict)[0]
    time_after_computing = time.time()

print(f"elapsed time to compute {i+1}: {time_after_computing - time_sent}")
df_importances = pd.DataFrame.from_dict(result_imp_dict)
df_importances_mean = df_importances.mean(axis=1).sort_values(ascending=False)
df_importances_mean.to_csv('averaged_feature_importances_500.csv')
print(df_importances_mean.head(n=20))
mean_accuracy = np.array(list(result_accuracy_dict.values())).mean()
mean_accuracy_ser = pd.Series(result_accuracy_dict)
mean_accuracy_ser.to_csv('accuracy_scores_500.csv')
end = time.time()
print(f"total elapsed time: {end - start}")   
print(mean_accuracy)


client.close()

提交用的Slurm脚本

#!/bin/bash

#SBATCH --job-name=dask-py           # Job name
#SBATCH --nodes=1                    # Run all processes on a 10 node
#SBATCH --ntasks=15                  # Number of tasks (MPI workers)
#SBATCH --cpus-per-task=1            # Number of CPU cores per task
#SBATCH --mem=200G                   # Job memory request
#SBATCH --partition=standard,short,long,fat
#SBATCH --output=dask-py%j.log       # Standard output and error log
#SBATCH --error=dask-py%j.err       # Standard output and error log


echo "Running Dask-MPI"

module load mpi/openmpi/4.1.1

source conda.sh

conda activate working-env  # Activate conda environment 

mpirun --oversubscribe -np 15 python3 RF_model_biome_predict_loop.py


'''
# mpirun -np argument needs 2 more workers than what you want for your script
# np rank 0 is for the scheduler
# np rank 1 is for python client process (distribution)
# remaining are workers for your script
'''

警告示例

2024-10-21 17:16:21,239 - distributed.worker.memory - WARNING - Unmanaged memory use is high. This may indicate a memory leak or the memory may not be released to the OS; see https://distributed.dask.org/en/latest/worker-memory.html#memory-not-released-back-to-the-os for more information. -- Unmanaged
 memory: 10.06 GiB -- Worker memory limit: 12.50 GiB

或

2024-10-22 12:59:46,569 - distributed.utils_perf - INFO - full garbage collection released 441.47 MiB from 44 reference cycles (threshold: 9.54 MiB)

优化方案建议

减少数据集重复传递，降低内存开销
当前每次调用run_the_classifier都会传递全量数据集，导致集群内大量数据复制。初始化Client后，把数据集广播到所有Worker节点，避免重复传递：
```
data_df_merged_balanced = client.scatter(data_df_merged_balanced, broadcast=True)
```
同时函数内可手动清理大对象，减少内存占用：
```
def run_the_classifier(data_df_merged_balanced):
    # 原有逻辑...
    del clf, X_train, y_train, X_test, y_test
    gc.collect()
    return accuracy, feature_imp
```
合并dask.compute调用，减少调度开销
原代码分别计算两个字典，会让Dask两次遍历任务图，合并成一次计算能直接降低compute步骤耗时：
```
result_accuracy_dict, result_imp_dict = dask.compute(accuracy_dict, importance_dict)
```
调优随机森林参数，缩短单任务耗时
n_estimators=10000会让单模型训练极慢，先测试降低该值（比如1000或2000），若精度下降在可接受范围内，能大幅提升整体速度。max_depth=5可保留，结合n_estimators做平衡性测试。
优化Worker资源配置，避免线程竞争
sklearn的RandomForest默认多线程运行，原Slurm配置给每个Worker只分配1核，会导致内部线程竞争。调整参数：比如--ntasks=4 --cpus-per-task=4，总核心数不变，但每个Worker能高效利用多线程训练模型。同时给Dask Worker设置内存限制：
```
client = Client(memory_limit='12GiB')
```

分批次处理任务，避免任务图过载
一次性生成10000个延迟任务会让任务图过于庞大，调度开销剧增。改成分批次处理，比如每批1000次，计算完成后释放内存再处理下一批：

batch_size = 1000
total_runs = 10000
all_accuracies = []
all_importances = []

for batch in range(0, total_runs, batch_size):
    accuracy_batch = {}
    importance_batch = {}
    for i in range(batch, min(batch+batch_size, total_runs)):
        acc, imp = run_the_classifier(data_df_merged_balanced)
        accuracy_batch[i] = acc
        importance_batch[i] = imp
    res_acc, res_imp = dask.compute(accuracy_batch, importance_batch)
    all_accuracies.extend(res_acc.values())
    all_importances.extend(res_imp.values())
    del accuracy_batch, importance_batch, res_acc, res_imp
    gc.collect()

替换Dask-MPI为Dask Slurm集群（可选）
若对MPI配置不熟悉，可改用dask-slurm自动管理集群资源，减少手动配置成本：
```
from dask_slurm import SLURMCluster
cluster = SLURMCluster(
    cores=4,
    memory='16GiB',
    n_workers=4,
    partition='standard'
)
cluster.scale(4)
client = Client(cluster)
```
提交时直接用sbatch运行Python脚本即可，无需手动写MPIRun命令。

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

如何加速Slurm集群上Dask Delayed大规模随机森林计算？

我的Python脚本（基于dask.delayed） #

提交用的Slurm脚本 #

警告示例 #

我的Python脚本（基于dask.delayed）

提交用的Slurm脚本

警告示例