SLURM管理的HPC系统上Dask优化工作流问询

核心优化思路：让调度器与计算集群独立于提交脚本运行 #

要解决脚本结束后Dashboard关闭、依赖SSH会话的问题，核心是把Dask调度器、Dashboard和计算节点集群从提交脚本中解耦，摆脱对登录节点本地进程和SSH连接的依赖。

方法一：通过SLURM作业启动独立的Dask调度器与集群 #

这是HPC环境下最通用的生产级方案，完全利用SLURM的资源管理能力：

1. 编写SLURM作业脚本启动Dask调度器
   创建start_dask_scheduler.sh脚本，根据你的HPC环境调整配置：

   #!/bin/bash
   #SBATCH --job-name=dask-scheduler
   #SBATCH --nodes=1
   #SBATCH --ntasks=1
   #SBATCH --cpus-per-task=2
   #SBATCH --mem=4G
   #SBATCH --time=24:00:00
   
   # 加载你的Python环境（示例为conda，根据实际调整）
   module load python/3.9
   source /path/to/your/conda/env/bin/activate
   
   # 启动调度器，生成地址文件并保存日志
   dask scheduler --port 8786 --dashboard-address :8787 --scheduler-file scheduler.json > scheduler.log 2>&1
   

   这个脚本会申请一个专用小节点运行调度器，生成的scheduler.json会记录调度器的网络地址，方便后续工作节点和客户端连接。

2. 提交调度器作业
   在登录节点执行：

   sbatch start_dask_scheduler.sh
   

   用squeue查看作业状态，等待scheduler.json文件生成后再进行下一步。

3. 启动Dask工作节点集群
   编写start_dask_workers.sh脚本，按需调整节点数、CPU和内存配置：

   #!/bin/bash
   #SBATCH --job-name=dask-workers
   #SBATCH --nodes=4
   #SBATCH --ntasks-per-node=1
   #SBATCH --cpus-per-task=16
   #SBATCH --mem=64G
   #SBATCH --time=24:00:00
   
   module load python/3.9
   source /path/to/your/conda/env/bin/activate
   
   # 连接到已启动的调度器，启动工作节点
   dask worker --scheduler-file scheduler.json --nthreads 16 --nprocs 1 --memory-limit 64G > worker_%j.log 2>&1
   

   提交作业：

   sbatch start_dask_workers.sh
   

   工作节点会自动连接到调度器，完全由SLURM管理，不依赖登录节点进程。

4. 编写独立的任务提交脚本
   创建submit_tasks.py，仅负责提交任务，不管理集群生命周期：

   from dask.distributed import Client
   
   # 通过scheduler.json连接到已运行的调度器
   client = Client(scheduler_file="scheduler.json")
   
   # 定义你的任务逻辑
   def process_data(x):
       return x * 2 + 1
   
   # 分发并收集任务结果
   futures = client.map(process_data, range(10000))
   final_results = client.gather(futures)
   
   # 任务完成后仅关闭客户端，调度器和集群保持运行
   client.close()
   

   在登录节点运行这个脚本，执行完成后脚本可直接退出，调度器和Dashboard仍会持续运行。

方法二：用tmux/screen维持登录节点会话（临时测试方案） #

如果只是短时间测试，不想编写SLURM作业脚本，可以用会话管理工具让进程在后台持续运行：

1. 登录节点后创建tmux会话：

   tmux new -s dask-session
   

2. 在会话内启动你的原有Dask集群和任务代码：

   from dask_jobqueue import SLURMCluster
   from dask.distributed import Client
   
   cluster = SLURMCluster(cores=16, memory="64GB", walltime="24:00:00")
   cluster.scale(4)
   client = cluster.get_client()
   # 任务代码...
   

3. 按下Ctrl+B再按D脱离会话，此时即使SSH断开，进程仍在登录节点运行。
4. 后续重新登录后，执行tmux attach -t dask-session恢复会话，可继续查看Dashboard或任务状态。
   注意：这种方法占用登录节点资源，长期运行不推荐，仅适合临时测试场景。

持久化访问Dashboard的方法 #

无论用哪种方案，要在本地稳定访问Dashboard，需要建立后台运行的SSH端口转发：

ssh -L 8787:<调度器节点IP>:8787 <登录节点用户名>@<登录节点地址> -N -f

其中<调度器节点IP>可以从scheduler.json文件中获取，-N表示不执行远程命令，-f让SSH在后台运行。即使本地机器重启，重新执行这条命令即可恢复访问。

额外优化建议 #

• 启用自适应集群：在启动工作节点时，用cluster.adapt(minimum=2, maximum=10)替代cluster.scale()，让Dask根据任务负载自动增减工作节点，节省HPC资源。
• 自动清理资源：给SLURM作业设置合理的walltime，任务完成后自动回收节点；也可以编写脚本，在任务全部完成后调用cluster.close()关闭集群。

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