Docker + JupyterHub + Nvidia

近日购入一张技嘉Nvidia RTX 4060 Ti (16GB)用作机器学习用途，小记一波三折的安装过程。

硬件安装

这步骤无非是将显卡插入主板的PCIe插槽，锁紧螺丝以固定，再连接8 pin的辅助供电线缆。此步骤并无更多须关注之细节。

Host OS

我的服务器运转Arch Linux操作系统。按照ArchWiki¹所述，(nvidia-dkms || nvidia-open-dkms) && nvidia-utils需要被安装。由于我的服务器运转linux-hardened内核，故需要由dkms编译的包。安装dkms时，成吨的用于编译的包，如gcc，被作为依赖安装了。但我又无可奈何——nvidia-open会不会并入mainline呢？如果能，又在几时呢？无人知晓…

在此之后我犯了一个错误，造成了这台服务器自装机以来第二次Kernel Panic——我在安装后直接运转nvidia-smi，系统没了反应，登录IPMI的remote KVM发现是KP了。没有日志被留下，猜测可能是因为nouvaeu被加载并与nvidia的内核模块产生某些冲突造成的。

此外，mkinitcpio.conf也需要被修改，kms须被从HOOKS中移除，以避免nouvaeu被加载。

# /etc/mkinitcpio.conf.d/00-override.conf

MODULES=(ext4 vfat)
#HOOKS=(base udev keyboard keymap autodetect microcode modconf kms consolefont block filesystems fsck lvm2 encrypt)
# Nvidia: remove kms to avoid nouvaeu from loading.
HOOKS=(base udev keyboard keymap autodetect microcode modconf consolefont block filesystems fsck lvm2 encrypt)

重启过后，运转nvidia-smi，应该已经可以查看到GPU相关信息。

Docker

按照ArchWiki²所述，nvidia-container-toolkit需要被安装。在安装后重新启动Docker (systemctl restart docker.service)，通过docker run --gpus all或者Nvidia container runtime便可令GPU在Docker容器中使用。

然而事情并不如此简单，在linux-hardened内核下，直接按前者(--gpus all)操作会造成如下结果：

docker run --gpus all --rm -it nvidia/cuda:12.1.1-runtime-ubuntu22.04 nvidia-smi

docker: Error response from daemon: failed to create task for container: failed to create shim task: OCI runtime create failed: runc create failed: unable to start container process: error during container init: error running hook #0: error running hook: exit status 1, stdout: , stderr: Auto-detected mode as 'legacy'
nvidia-container-cli: mount error: failed to add device rules: unable to generate new device filter program from existing programs: unable to create new device filters program: load program: invalid argument: last insn is not an exit or jmp
processed 0 insns (limit 1000000) max_states_per_insn 0 total_states 0 peak_states 0 mark_read 0: unknown.

检索相关关键词，可知³net.core.bpf_jit_harden这一内核参数如被设置为2(linux-hardened的默认值)便会造成这个结果，需设置为1或0(一般内核的默认值)方能解决。这样做可能降低系统的安全性。

# /etc/sysctl.d/fuck-nvidia.conf
net.core.bpf_jit_harden = 1

至此，Docker容器中的程序可以使用Nvidia GPU了。

JupyterHub

我在服务器上部署的JupyterHub采用DockerSpawner，使每个用户的Jupyter Server运转于Docker容器中。JupyterHub的配置文件需要被修改，以使其创建的容器使用Nvidia GPU。

不幸地，网络上现存的资料绝大多数(如果不是全部)皆是基于旧的Nvidia container runtime实现的。而基于(被推荐的)--gpus all的解决方案我未能搜索到。

阅读DockerSpawner的(部分)源码⁴可知，其通过Docker的Python SDK(import docker)连接宿主机映射至JupyterHub容器内的Docker Socket(-v /var/run/docker.sock:/var/run/docker.sock)来为用户创建容器。其使用了create_container()及create_host_config()等API。

继续查阅Docker Python SDK的文档⁵和一些StackOverflow回答⁶可知，在JupyterHub的配置文件中添加如下内容，便可令其创建的容器使用Nvidia GPU:

import docker
c.DockerSpawner.extra_host_config = {"device_requests": [docker.types.DeviceRequest(
    device_ids=["all"],
    capabilities=[["gpu"]],
)]}

至此，终于可以愉快(真的吗?)地在服务器上运转的Jupyter Notebook使用CUDA了。

---

后记

Content Warning: 这个段落与正文之技术话题无关，属发泄情绪之内容。如不想阅读，可回避之而不影响正文之完整性。

牢骚话

事实上，这是一次令人作呕的，缺乏优雅的安装。整个过程繁琐冗杂，制造诸多烂摊子，留了一地的屎！包括但不限于：

1. 大量修改系统文件和配置，降低安全性
2. 安装许多先前从不需要的软件包，只为在内核更新之时编译Nvidia内核模块用

1. https://wiki.archlinux.org/title/NVIDIA ↩︎
2. https://wiki.archlinux.org/title/Docker ↩︎
3. https://github.com/bottlerocket-os/bottlerocket/issues/2504#issuecomment-1410814475 ↩︎
4. https://github.com/jupyterhub/dockerspawner/blob/13.0.0/dockerspawner/dockerspawner.py ↩︎
5. https://docker-py.readthedocs.io/en/stable/api.html ↩︎
6. https://stackoverflow.com/questions/71429711/how-to-run-a-docker-container-with-specific-gpus-using-docker-sdk-for-python/71429712#71429712 ↩︎