在 Slurm 上运行 Jupyter

1. 背景介绍

现在的大模型训练越来越深入每个组了，大规模集群系统也应用的愈发广泛。一般的slurm系统提交作业分为2种，一种是srun，这种所见即所得的申请方式一般适用于短期的调试使用，大概一般允许的时间从几个小时到1天左右，很多集群分组都会限制运行时长。而另一种sbatch，则是批量提交作业，当srun调试程序能够成功运行的时候，就可以使用sbatch提交。如何使用slurm可以参考之前写的《slurm初相识》以及《查询slurm集群各个节点的运行情况》。

那么有没有一种方式可以以sbatch提交作业，但是又能有srun的交互式体验呢？有的，那就是使用sbatch提交jupyter作业。这样就可以在jupyter里自由的使用申请到的资源了。闲话不多说，下面介绍一下具体的实现步骤。

2. 实现

2.1 代码准备

#!/bin/bash
#SBATCH -J [job name]
#SBATCH -p [partition name]
#SBATCH --nodes=1
#SBATCH --gres=gpu:4 # 需要使用多少GPU，n是需要的数量
#SBATCH --cpus-per-gpu 4
#SBATCH -t 5-00:00:00 # 运行总时间，天数-小时数-分钟， D-HH:MM
#SBATCH -o logging_sbatch_8_3.o # 把输出结果STDOUT保存在哪一个文件
#SBATCH -e logging_sbatch_8_3.e # 把报错结果STDERR保存在哪一个文件
#SBATCH --mail-user=[notified E-mail address]
#SBATCH --mail-type=ALL
#SBATCH --nodelist=[Hostname]

export XDG_RUNTIME_DIR=""
node=$(hostname -s)
user=$(whoami)
submit_host=${SLURM_SUBMIT_HOST}
port=[端口号]

echo $node pinned to port $port
# print tunneling instructions jupyter-log

echo -e "
To connect to the compute node ${node} on sribd running your jupyter notebook server,
you need to run following two commands in a terminal
1. Command to create ssh tunnel from you workstation/laptop to cs-login:
ssh -N -f -L ${port}:${node}:${port} ${user}@[登陆机ip]
Copy the link provided below by jupyter-server and replace the NODENAME with localhost before pasting it in your browser on your workstation/laptop
"

# Run Jupyter
jupyter lab --no-browser --port=${port}  --ip=${node} 

当然，由于各种实现方式和资源分配方案，里面的SBATCH可能还需要下面的一些代码：

#SBATCH -A [指定账户]
#SBATCH --reservation=[保留代号]

2.2 在服务器上运行上述代码

当按照自己的需求，直接修改后，保存为sh命令行文件，例如sbatch_run.sh在主节点直接输入下面的命令：

sbatch sbatch_run.sh

如果有报错提示，建议提前安装一下Jupyter：

pip install jupyter notebook

然后，去客户端上执行以下步骤。

2.3 客户端执行代码

在本地电脑上打开一个新的终端窗口，复制并执行 ./logging_sbatch_8_3.o 第一行生成的命令，创建一个SSH隧道：

windows 通过 cmd (命令提示符) 执行 ，Linux 或 Mac 通过 terminal 执行下面代码：

ssh -N -f -L [Port]:[HostName]:[Port] [username]@[Server IP]  # 输入密码登录成功后，请勿关闭此终端窗口。

在通过本地电脑打开浏览器，访问 ./logging_sbatch_8_3.e 最后一行生成的网址。

类似如下的命令:

http://127.0.0.1:8820/lab?token=5ab3861accf57e32f1351ab895cb456c30a20cf9cd10f86c

然后进去以后会长这样:

选择终端[Terminal]即可。这样就可以以sbatch的时长享用srun了。

3. 道德声明

这种做法其实是不对的，本来slurm创建作业系统，就是为了避免资源空闲。一旦这样做，我们就会独享一个服务器的资源，无论我们是否使用，它都已经被占用了。所以建议非必要紧急，不要使用这种方式去申请资源。如果有的话，建议使用sbatch正规提交作业。提高资源利用率。