如何在VSCode的Jupyter扩展中使用Slurm计算节点运行代码？

下面提供两种可行的方案，适配你不想用浏览器、计算节点无公网访问的场景：

方案一：配置VSCode SSH多跳连接直接访问计算节点 #

1. 先通过登录节点申请计算节点：

   srun -p cpu -n 4 --pty /bin/bash
   

   记录下分配的节点名（比如node001）。
2. 打开VSCode，按Ctrl+Shift+P调出命令面板，选择Remote-SSH: Open SSH Configuration File，选择你的SSH配置文件（通常是~/.ssh/config）。
3. 在配置文件中添加以下内容，替换占位符为你的实际信息：

   Host login-node
       HostName <登录节点的IP或域名>
       User <你的服务器用户名>
   
   Host compute-node
       HostName node001  # 替换为Slurm分配的实际节点名
       User <你的服务器用户名>
       ProxyJump login-node
   

4. 保存配置后，通过Remote-SSH: Connect to Host...选择compute-node，连接到计算节点。
5. 连接成功后，打开VSCode的Jupyter扩展，创建或打开Notebook，选择计算节点上的Python环境即可运行代码。

方案二：通过反向隧道连接计算节点的Jupyter服务 #

1. 申请并进入计算节点：

   srun -p cpu -n 4 --pty /bin/bash
   

2. 在计算节点启动Jupyter服务，指定端口并禁用浏览器：

   jupyter notebook --no-browser --port=8889
   

   记下输出中的登录token（格式类似http://localhost:8889/?token=xxxxxx）。
3. 保持计算节点的终端会话，新开一个登录节点的终端，建立反向隧道，将计算节点的8889端口转发到登录节点的8890端口（可替换为未占用的端口）：

   ssh -N -f -L 8890:localhost:8889 node001  # node001替换为你的计算节点名
   

4. 在VSCode中通过Remote-SSH连接到登录节点，打开Jupyter扩展，点击Connect to Jupyter Server，选择Existing，输入http://localhost:8890，然后输入之前记录的token完成连接。
5. 连接成功后，即可在VSCode中直接操作计算节点上的Jupyter Notebook。

注意事项 #

• Slurm分配的计算节点是临时的，每次重新申请节点后，需要更新SSH配置中的节点名或重新建立反向隧道。
• 确保计算节点已安装Jupyter和你需要的Python环境，VSCode连接后能正确识别该环境。
• 端口号可根据实际占用情况调整，避免端口冲突。

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