Slurm enroot 容器启用远程 SSH 连接

, ,

问题概述

容器化计算作业无法通过本地（个人电脑）SSH 直连，并且很多 HPC 集群的计算节点不提供互联网访问（出于安全原因）。

建议方案

在计算作业的容器中启动独立的 SSH 服务，从本地建立 SSH 通道，穿过登录节点，再穿过计算节点（如计算节点端口未开放），最后进入计算作业容器。

1
本地计算机 -> 登陆节点 -> 计算节点（非必须） -> 计算作业容器

初步实现

0. 模拟环境

通过 iptables 禁用本地 IP

1
iptables -I INPUT 1 -s 192.168.110.93 -j DROP

1. 生成本地密钥上传到登陆节点

1.1 生成密钥

1
ssh-keygen -t rsa -b 2048 -C "your_email@example.com"

1.2 上传公钥

复制 id_rsa.pub 内的内容到登陆节点家目录的 ~/.ssh/authorized_keys 内

2. 在容器中安装并启动 SSHD 服务

镜像地址：docker://nvcr.io#nvidia/pytorch:21.04-py3
容器环境：Ubuntu 20.04

2.1 使用 root 安装 openssh-server

启动容器进入并安装 openssh-server

1
2
3
$ enroot start -r -w <container>
# apt-get update
# apt-get install openssh-server

2.2 使用普通用户启动 SSHD

拷贝 .ssh 下的文件（或直接挂载）

1
$ cp -R /workspace/.ssh .

启动 SSHD

1
$ sshd -D -e -p 2222

3. 使用本地 VSCode 连接

3.1 安装远程插件（SSH）

3.2 在 vscode 中连接

3.2.1 修改配置文件（原文方案）

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
Host devcontainer.ogsp
    User huangyx
    Port 2222
    HostName localhost
    ProxyJump devnode.ogsp
    CheckHostIP no
    StrictHostKeyChecking=no
    UserKnownHostsFile=/dev/null

# 计算节点端口开放情况下可省略
Host devnode.ogsp
    User huangyx
    CheckHostIP no
    ProxyCommand ssh slurm.ogsp "nc t1 22"
    StrictHostKeyChecking=no
    UserKnownHostsFile=/dev/null

Host slurm.ogsp
    User huangyx
    HostName 10.10.22.32

3.2.2 命令直接连接（新版 vscode 可以直接生成配置到配置文件）

1
ssh -o StrictHostKeyChecking=no -o UserKnownHostsFile=/dev/null -J ogsp@10.10.22.32 ogsp@t1 -p 2222

4. Slurm 脚本

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
#!/bin/bash

#SBATCH --job-name=${job_name}
#SBATCH --partition=${partition}
#SBATCH -N 1
#SBATCH --cpus-per-task=${cpu_num}
#SBATCH --ntasks-per-node=1
#SBATCH --gpus-per-task=${gpu_num}
#SBATCH -o %j.out
#SBATCH -e %j.err

#生成随机端口
port=$((RANDOM % (65535 - 2048) + 1024))

# 检查端口是否被占用
while [[ $(ss -tuln | grep ":$port ") ]]; do
    port=$((port + 1))
done

#查看启动节点
node=$(cat /etc/hosts |grep -m 1 `hostname -s` |awk '{print $1}')
host=$(hostname -s)

relative_path=$(pwd | sed "s|^$HOME/||")

container_workdir="/workspace/$relative_path"

container_name=$(echo "${job_container}" | sed 's/^pyxis_//')

export XDG_RUNTIME_DIR=/tmp

# if [ ${gpu_num} -eq 0 ]; then
# 	export NVIDIA_VISIBLE_DEVICES=void
# fi

srun_common_options="\
  --container-name=$container_name \
  --container-mounts=$HOME/.ssh:$HOME/.ssh \
  --container-workdir=$HOME \
  --container-writable \
"

ssh_common_options="\
  sshd -D -e -p $port 
"

# 登陆节点作为跳转节点
#jump_node="10.10.22.32"
ssh_cmd="ssh -o StrictHostKeyChecking=no -o UserKnownHostsFile=/dev/null -J $USER@${jump_node} $USER@$host -p $port"

# 生成ogsp所需信息
echo -e "{'hostname':'$node','port':'$port','sshCmd':'$ssh_cmd'}" > ssh_conn.json

eval "srun $srun_common_options --no-container-remap-root $ssh_common_options"

参考资料

https://gist.github.com/malteos/5fe791fe10bb55028a02952d5f394bb3