Slurm集群下通过SSH隧道连接GPU节点的Jupyter Notebook
1. 项目概述
在科研计算和高性能计算领域，Slurm集群是广泛使用的资源管理和作业调度系统。许多研究机构和企业都会部署带有GPU计算节点的Slurm集群，用于深度学习训练、科学计算等计算密集型任务。而Jupyter Notebook作为交互式编程环境，在数据分析和机器学习工作流中扮演着重要角色。
本文将详细介绍如何在Slurm集群环境下，通过SSH隧道实现本地计算机与GPU计算节点上运行的Jupyter Notebook/Lab的安全连接，并确保计算任务能够正确分配到GPU资源。
2. 环境准备与基础概念
2.1 Slurm集群基本架构
典型的Slurm集群由以下组件构成：
- 登录节点（Login Node）：用户直接连接的入口节点
- 计算节点（Compute Node）：实际执行计算任务的节点，部分配备GPU
- 控制节点（Control Node）：运行slurmctld守护进程，负责作业调度
- 存储节点：提供共享文件系统
2.2 Jupyter Notebook与GPU资源
Jupyter Notebook在数据科学工作流中的优势包括：
- 交互式代码执行和可视化
- 支持Markdown文档与代码混合
- 丰富的内核支持（Python、R、Julia等）
当需要GPU加速时，必须确保：
- Jupyter内核运行在分配了GPU的计算节点上
- 正确加载了CUDA和cuDNN等GPU计算库
- 使用的深度学习框架（如PyTorch、TensorFlow）已配置GPU支持
3. 详细操作步骤
3.1 连接到计算节点
首先通过Slurm分配一个带有GPU的计算节点：
srun --nodes=1 --partition=gpu --gres=gpu:1 --time=24:00:00 --pty bash参数说明：
--nodes=1：请求1个计算节点--partition=gpu：指定GPU分区--gres=gpu:1：请求1块GPU--time=24:00:00：分配24小时运行时间
3.2 启动Jupyter服务
在获得计算节点shell后，启动Jupyter Lab：
jupyter lab --no-browser --port=8889 --ip=0.0.0.0关键参数：
--no-browser：不自动打开浏览器--port=8889：指定服务端口--ip=0.0.0.0：允许外部连接
3.3 建立SSH隧道
在本地计算机上建立端口转发：
ssh -N -L 8889:localhost:8889 username@login-node-address这条命令将本地8889端口映射到计算节点上的8889端口。
3.4 访问Jupyter界面
在本地浏览器中访问：
http://localhost:8889输入启动Jupyter时显示的token即可完成认证。
4. 高级配置与优化
4.1 持久化Jupyter会话
使用tmux或screen保持会话：
tmux new -s jupyter jupyter lab --no-browser --port=8889 # 按Ctrl+B, 然后按D分离会话4.2 自定义Jupyter环境
创建专用conda环境：
conda create -n myenv python=3.8 conda activate myenv conda install jupyterlab pytorch torchvision cudatoolkit=11.3 -c pytorch4.3 GPU资源监控
安装和使用gpustat监控GPU使用情况：
pip install gpustat gpustat -i5. 常见问题与解决方案
5.1 连接问题排查
检查SSH隧道是否建立成功：
netstat -tuln | grep 8889验证计算节点上的Jupyter服务是否运行：
ps aux | grep jupyter
5.2 GPU不可用问题
确保PyTorch/TensorFlow能识别GPU：
import torch print(torch.cuda.is_available()) # 应返回True print(torch.cuda.device_count()) # 应显示GPU数量5.3 性能优化建议
- 使用
--mem参数为作业分配足够内存 - 对于多GPU任务，考虑使用
--gres=gpu:2等参数 - 在Jupyter中合理使用
%time和%%timeit魔法命令监控代码性能
6. 安全注意事项
- 始终使用token认证，避免无密码访问
- 考虑使用SSL加密Jupyter通信
- 作业完成后及时释放计算资源
- 敏感数据应存储在安全位置，避免通过Jupyter暴露
提示：在共享集群环境中，建议使用
--notebook-dir参数指定工作目录，避免意外访问他人文件。
7. 扩展应用场景
7.1 多用户协作
配置JupyterHub与Slurm集成，实现多用户管理和资源分配。
7.2 远程开发环境
结合VS Code Remote SSH扩展，创建完整的远程开发环境。
7.3 自动化工作流
使用papermill等工具实现Jupyter Notebook的批量执行和参数化运行。