Dockerfile_k8s_gpu

# We use the cuda runtime image instead of devel image to reduce size (1.3GB vs 3.6GB)
FROM nvidia/cuda:12.1.1-runtime-ubuntu20.04

ARG DEBIAN_FRONTEND=noninteractive

# Install ssh and other local dependencies
# We remove cuda lists to avoid conflicts with the cuda version installed by ray
RUN rm -rf /etc/apt/sources.list.d/cuda* && \
    apt update -y && \
    apt install git gcc rsync sudo patch openssh-server pciutils nano fuse unzip socat netcat-openbsd curl -y && \
    rm -rf /var/lib/apt/lists/*

# Setup SSH and generate hostkeys
RUN sudo mkdir -p /var/run/sshd && \
    sudo sed -i 's/PermitRootLogin prohibit-password/PermitRootLogin yes/' /etc/ssh/sshd_config && \
    sudo sed 's@session\s*required\s*pam_loginuid.so@session optional pam_loginuid.so@g' -i /etc/pam.d/sshd && \
    cd /etc/ssh/ && \
    sudo ssh-keygen -A

# Setup new user named sky and add to sudoers. \
# Also add /opt/conda/bin to sudo path and give sky user permission to run sudo without password
RUN sudo useradd -m -s /bin/bash sky && \
    sudo /bin/bash -c 'echo "sky ALL=(ALL) NOPASSWD:ALL" >> /etc/sudoers' && \
    sudo /bin/bash -c "echo 'Defaults        secure_path=\"/opt/conda/bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin\"' > /etc/sudoers.d/sky"

# Switch to sky user
USER sky

# Set HOME environment variable for sky user
ENV HOME /home/sky

# Set current working directory
WORKDIR /home/sky

SHELL ["/bin/bash", "-c"]

# Install conda and other dependencies
# Keep the conda and Ray versions below in sync with the ones in skylet.constants
# Keep this section in sync with the custom image optimization recommendations in our docs (kubernetes-getting-started.rst)
RUN curl https://repo.anaconda.com/miniconda/Miniconda3-py310_23.11.0-2-Linux-x86_64.sh -o Miniconda3-Linux-x86_64.sh && \
    bash Miniconda3-Linux-x86_64.sh -b && \
    eval "$(~/miniconda3/bin/conda shell.bash hook)" && conda init && conda config --set auto_activate_base true && conda activate base && \
    grep "# >>> conda initialize >>>" ~/.bashrc || { conda init && source ~/.bashrc; } && \
    rm Miniconda3-Linux-x86_64.sh && \
    export PIP_DISABLE_PIP_VERSION_CHECK=1 && \
    python3 -m venv ~/skypilot-runtime && \
    PYTHON_EXEC=$(echo ~/skypilot-runtime)/bin/python && \
    $PYTHON_EXEC -m pip install 'skypilot-nightly[remote,kubernetes]' 'ray[default]==2.9.3' 'pycryptodome==3.12.0' && \
    $PYTHON_EXEC -m pip uninstall skypilot-nightly -y && \
    curl -LO "https://dl.k8s.io/release/v1.28.11/bin/linux/amd64/kubectl" && \
    sudo install -o root -g root -m 0755 kubectl /usr/local/bin/kubectl && \
    echo 'export PATH="$PATH:$HOME/.local/bin"' >> ~/.bashrc

# Set PYTHONUNBUFFERED=1 to have Python print to stdout/stderr immediately
ENV PYTHONUNBUFFERED=1