利用 SSH 反向隧道登录 BitaHub 中的容器并长期占用 GPU

发表于

问题

每年的 CVPR 前 GPU 总是供不应求，需要从其他地方借卡。USTC 有一个供校内用户使用的 BitaHub，但它同样有 CVPR 前一卡难求的问题，同时基于任务提交的使用模式也非常不方便，提交占用多卡的任务经常需要漫长的排队，数据管理方式更是反人类。

作为组里的服务器管理员，为了让自己在 CVPR 前活得轻松点，避免重蹈 2021 年 CVPR 前疲于应对资源调配的覆辙，有必要改善 BitaHub 的使用体验：

  1. 如何长期占用显卡避免重复排队（虽然略不道德，但实属无奈之举）；
  2. 如何方便地从我们的服务器读取数据，而不是被迫使用 BitaHub 反人类的数据管理模式；
  3. 如何尽量使 BitaHub 的 GPU 使用体验接近组里的服务器，降低迁移成本，提高资源调度的灵活性。

思路

BitaHub 中的任务是以 docker 容器的方式运行的，因而给了我们在容器里配置我们想要的环境的可能，只需要通过某种方式登录 ssh 到容器中。

经过研究发现，只要启动命令不停止运行，BitaHub 中的容器就会长期运行，不释放 GPU 资源。同时 BitaHub 中的容器是可以联网的，而且 BitaHub 的网页上还贴心地给出了每个任务容器中 root 用户的 ssh 私钥。

这些给了我们利用的机会，只需要在容器内运行一个 tunnel 程序以让外部得以访问容器中的 22 端口，就能登录并长期占用资源。同时由于容器联网，也可以直接挂载校内其他服务器的文件系统。

解决方案

最终选择的 tunnel 程序是 ssh，它可以创建反向隧道：

1
ssh -i <key_file> -F none -o "StrictHostKeyChecking no" -o "ServerAliveInterval 15" -v -N -R <port>:localhost:22 jump@<jumpserver>

在 jumpserver 中配置用户 jump 并允许特定私钥登录，然后用某种方式把私钥传递进容器（可以直接打包进镜像，但我选择了更方便的方式——创建一个 BitaHub 数据集存放，每个任务添加这个数据集即可）。

容器的启动命令就是上述命令（考虑到网络波动，可以套一层 while true 循环或者用 autossh 自动重连），启动后就在 <jumpserver> 上的 <port> 端口创建了一个反向隧道，<port> 被映射到了容器内的 22 端口。

可以在 <jumpserver> 的 sshd_config 中配置 GatewayPorts yes，这样反向隧道就会监听 0.0.0.0 而不是 127.0.0.1。不这样做的话我就要在 <jumpserver> 上给每个人创建用户，或者每个端口用 iptables 转发，但这太过繁琐。绑定 0.0.0.0 则可以直接从现有的 VPN 网络中访问。

挂载文件系统的方式有很多选择，考虑到安全性和便捷性，我选择了 SSHFS。NFS 直接暴露于公网过于危险，而 NFS 用户验证的配置又过于繁琐。同时 BitaHub 运行容器的内核既没加载 wireguard kmod 也没映射 /dev/net/tun，因此无法利用 VPN 保护数据安全。SSHFS 可以直接复用现存的用户认证方式，而 SSH 流量本身也更容易被潜在的机房防火墙放过。

使用如下命令挂载 SSHFS：

1
sshfs -o reconnect,ServerAliveInterval=15,ServerAliveCountMax=30,ssh_command='ssh -p <dataserver_port> -i <key_file>' <user>@<dataserver>:/path /path

后记

TODO