FROM docker.io/nvidia/cuda:11.8.0-cudnn8-devel-ubuntu22.04

ARG APT_MIRROR=https://mirrors.bfsu.edu.cn
ARG ROOT_PASSWD=miao

# install tensorrt in a separate layer
RUN true \
    &&  sed -i "s#http://archive.ubuntu.com#${APT_MIRROR}#g" /etc/apt/sources.list \
    &&  sed -i "s#http://security.ubuntu.com#${APT_MIRROR}#g" /etc/apt/sources.list \
    \
    &&  apt update \
    &&  apt install --yes  tensorrt=8.6.1.6-1+cuda11.8 \
            libnvinfer8=8.6.1.6-1+cuda11.8 \
            libnvinfer-plugin8=8.6.1.6-1+cuda11.8 \
            libnvinfer-vc-plugin8=8.6.1.6-1+cuda11.8 \
            libnvinfer-lean8=8.6.1.6-1+cuda11.8 \
            libnvinfer-dispatch8=8.6.1.6-1+cuda11.8 \
            libnvparsers8=8.6.1.6-1+cuda11.8 \
            libnvonnxparsers8=8.6.1.6-1+cuda11.8 \
            libnvinfer-bin=8.6.1.6-1+cuda11.8 \
            libnvinfer-dev=8.6.1.6-1+cuda11.8 \
            libnvinfer-lean-dev=8.6.1.6-1+cuda11.8 \
            libnvinfer-plugin-dev=8.6.1.6-1+cuda11.8 \
            libnvinfer-vc-plugin-dev=8.6.1.6-1+cuda11.8 \
            libnvinfer-dispatch-dev=8.6.1.6-1+cuda11.8 \
            libnvparsers-dev=8.6.1.6-1+cuda11.8 \
            libnvonnxparsers-dev=8.6.1.6-1+cuda11.8 \
            libnvinfer-samples=8.6.1.6-1+cuda11.8 \
            \
            libnvinfer-headers-dev=8.6.1.6-1+cuda11.8 \
            libnvinfer-headers-dev=8.6.1.6-1+cuda11.8 \
            libnvinfer-headers-dev=8.6.1.6-1+cuda11.8 \
            libnvinfer-headers-plugin-dev=8.6.1.6-1+cuda11.8 \
            libnvinfer-headers-plugin-dev=8.6.1.6-1+cuda11.8 \
    \
    &&  apt clean \
    &&  rm -rf /var/lib/apt/lists/*

RUN true \
    &&  env | grep -E -e '^NVIDIA' -e '^NV_' -e '^NVARCH' -e '^NCCL_' | \
            sed -E 's/=(.*)/="\1"/' > \
            /etc/profile.d/00-nvidia.sh \
    &&  echo "PATH=/usr/local/nvidia/bin:/usr/local/cuda/bin:\${PATH}" >> /etc/profile.d/00-nvidia.sh \
    &&  echo "LD_LIBRARY_PATH=/usr/local/nvidia/lib:/usr/local/nvidia/lib64" >> /etc/profile.d/00-nvidia.sh \
    &&  echo "LIBRARY_PATH=/usr/local/cuda/lib64/stubs" >> /etc/profile.d/00-nvidia.sh \
    &&  sed -i 's/^/export /' /etc/profile.d/00-nvidia.sh \
    \
    &&  apt update \
    &&  apt install --yes \
            ssh tini zsh locales \
            netcat file git \
            python3 python3-pip \
            zip unzip zstd \
            imagemagick \
    &&  apt clean \
    &&  rm -rf /var/lib/apt/lists/* \
    \
    &&  sed -i '/C.UTF-8/s/^# //g' /etc/locale.gen \
    &&  locale-gen \
    &&  update-locale LANG=C.UTF-8 \
    \
    &&  echo root:${ROOT_PASSWD} | chpasswd \
    &&  usermod --shell /usr/bin/zsh root \
    \
    &&  mkdir /run/sshd \
    &&  sed -E -i 's/^.*StrictModes.*$/StrictModes no/' '/etc/ssh/sshd_config' \
    &&  sed -i 's/#PermitRootLogin.*$/PermitRootLogin yes/' /etc/ssh/sshd_config

ENTRYPOINT [ "/usr/bin/tini", "--" ]
CMD [ "/usr/sbin/sshd", "-D", "-e" ]