FROM nvidia/cuda:11.2.0-base-ubuntu18.04 as base

RUN set -x && \
    apt-get update && \
    apt-get install -y ca-certificates zstd

COPY k3s/build/out/data.tar.zst /

RUN set -x && \
    mkdir -p /image/etc/ssl/certs /image/run /image/var/run /image/tmp /image/lib/modules /image/lib/firmware && \
    tar -I zstd -xf /data.tar.zst -C /image && \
    cp /etc/ssl/certs/ca-certificates.crt /image/etc/ssl/certs/ca-certificates.crt

RUN set -x && \
    cd image/bin && \
    rm -f k3s && \
    ln -s k3s-server k3s

FROM nvidia/cuda:11.2.0-base-ubuntu18.04

ARG NVIDIA_CONTAINER_RUNTIME_VERSION
ENV NVIDIA_CONTAINER_RUNTIME_VERSION=$NVIDIA_CONTAINER_RUNTIME_VERSION

RUN set -x && \
    echo 'debconf debconf/frontend select Noninteractive' | debconf-set-selections

RUN set -x && \
    apt-get update && \
    apt-get -y install gnupg2 curl

# Install NVIDIA Container Runtime
RUN set -x && \
    curl -s -L https://nvidia.github.io/nvidia-container-runtime/gpgkey | apt-key add -

RUN set -x && \
    curl -s -L https://nvidia.github.io/nvidia-container-runtime/ubuntu18.04/nvidia-container-runtime.list | tee /etc/apt/sources.list.d/nvidia-container-runtime.list

RUN set -x && \
    apt-get update && \
    apt-get -y install nvidia-container-runtime=${NVIDIA_CONTAINER_RUNTIME_VERSION}


COPY --from=base /image /

RUN set -x && \
    mkdir -p /etc && \
    echo 'hosts: files dns' > /etc/nsswitch.conf

RUN set -x && \
    chmod 1777 /tmp

# Provide custom containerd configuration to configure the nvidia-container-runtime
RUN set -x && \
    mkdir -p /var/lib/rancher/k3s/agent/etc/containerd/

COPY config.toml.tmpl /var/lib/rancher/k3s/agent/etc/containerd/config.toml.tmpl

# Deploy the nvidia driver plugin on startup
RUN set -x && \
    mkdir -p /var/lib/rancher/k3s/server/manifests

COPY gpu.yaml /var/lib/rancher/k3s/server/manifests/gpu.yaml

VOLUME /var/lib/kubelet
VOLUME /var/lib/rancher/k3s
VOLUME /var/lib/cni
VOLUME /var/log

ENV PATH="$PATH:/bin/aux"

ENTRYPOINT ["/bin/k3s"]
CMD ["agent"]