构建支持 gpu 的 jupyterlab docker 镜像

1. 创建 Dockerfile

创建一个 Dockerfile 文件，内容如下

FROM docker.io/nvidia/cuda:12.2.2-cudnn8-runtime-ubuntu22.04

ENV DEBIAN_FRONTEND=noninteractive

RUN apt-get update && apt-get upgrade -y && \
    apt-get install -y --no-install-recommends \
    libgl1 libglib2.0-0 wget git curl vim \
    python3.10 python3-pip python3-dev build-essential \
    openmpi-bin libopenmpi-dev jupyter-notebook jupyter

RUN pip3 install tensorrt_llm -U --extra-index-url https://pypi.nvidia.com
RUN pip3 install --upgrade jinja2==3.0.3 pynvml>=11.5.0

RUN rm -rf /var/cache/apt/ && apt-get clean && rm -rf /var/lib/apt/lists/* /tmp/* /var/tmp/* && \
    rm -rf /root/.cache/pip/ && rm -rf /*.whl

WORKDIR /root/work
ENTRYPOINT ["sh","-c","jupyter notebook --allow-root --notebook-dir=/root/work --port=8888 --ip=0.0.0.0 --ServerApp.token=''"]

2. 构建镜像

docker buildx build -t engchina/gpu-jupyter:latest -f Dockerfile .

3. 启动 gpu-jupyter

mkdir /u01/data/gpu-jupyter-work; chmod 777 /u01/data/gpu-jupyter-work
docker run --gpus all -d --name gpu-jupyter-work -p 9888:8888 -v /u01/data/gpu-jupyter-work:/root/work -e GRANT_SUDO=yes -e JUPYTER_ENABLE_LAB=yes --user root engchina/gpu-jupyter:latest

4. 访问 gpu-jupyter

查看 gpu-jupyter 的 login token，

token=$(docker exec -it gpu-jupyter-work jupyter server list | grep -oP '(?<=token=)[a-zA-Z0-9]+')
echo $token

使用浏览器打开 http://<your_ip>:9888/lab 进行访问，输入 token 进行登录。

完结！