Dockerfile_k8s_gpu

# We use the cuda runtime image instead of devel image to reduce size (1.3GB vs 3.6GB)
# Overridable so the same recipe can be rebuilt on another base (e.g.
# nvidia/cuda:*-runtime-ubuntu24.04 for a newer glibc).
ARG BASE_IMAGE=nvidia/cuda:12.1.1-runtime-ubuntu22.04
FROM ${BASE_IMAGE}

# Detect architecture using ARG with default value
ARG TARGETARCH
ARG DEBIAN_FRONTEND=noninteractive

# Install ssh, InfiniBand userspace, and other local dependencies.
# We remove cuda lists to avoid conflicts with the cuda version installed by ray.
# FUSE setup mirrors sky/templates/kubernetes-ray.yml.j2:825-832 — install
# fuse (FUSE 2) in the main batch, then fuse3 in a separate apt-get call.
# fuse and fuse3 cannot be installed together (fuse3 Breaks: fuse — both ship
# /usr/bin/fusermount), but in two steps apt resolves the conflict by removing
# the `fuse` package while leaving `libfuse2`. Final state: libfuse.so.2 +
# libfuse3.so.3 both in ldconfig (covers hf-mount + legacy FUSE 2 clients).
RUN rm -rf /etc/apt/sources.list.d/cuda* && \
    apt-get update -y && \
    apt-get upgrade -y && \
    apt-get install -y --no-install-recommends \
        ca-certificates \
        git build-essential rsync sudo patch openssh-server \
        python3 python3-pip python3-venv python3-dev python-is-python3 sqlite3 \
        pciutils nano fuse unzip socat netcat-openbsd curl wget autossh jq \
        rdma-core ibverbs-providers libibverbs1 libibverbs-dev \
        librdmacm1 ibverbs-utils infiniband-diags \
        numactl iputils-ping && \
    apt-get install -y --no-install-recommends fuse3 && \
    rm -rf /var/lib/apt/lists/*

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

# Setup new user named sky and add to sudoers.
# sky must own uid/gid 1000: kubernetes-ray.yml.j2 pins runAsUser/runAsGroup/
# fsGroup to 1000. Ubuntu 24.04+ base images ship a placeholder user holding it,
# so evict whoever owns 1000 first -- keyed on the id rather than the name,
# which varies across bases. userdel frees the account before it cleans up the
# home dir and mail spool, so its exit status is not load-bearing here; the
# explicit -u/-g is the real guard, failing the build instead of silently
# landing sky on 1001.
RUN existing_user=$(getent passwd 1000 | cut -d: -f1) && \
    if [ -n "$existing_user" ]; then userdel -r "$existing_user" || true; fi && \
    existing_group=$(getent group 1000 | cut -d: -f1) && \
    if [ -n "$existing_group" ]; then groupdel "$existing_group" || true; fi && \
    groupadd -g 1000 sky && \
    useradd -m -s /bin/bash -u 1000 -g 1000 sky && \
    echo "sky ALL=(ALL) NOPASSWD:ALL" >> /etc/sudoers

# 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 the SkyPilot runtime (uv-managed venv) and other dependencies based
# on architecture.
# Keep the Ray version below in sync with the one in skylet.constants
# Keep this section in sync with the custom image optimization recommendations in our docs (kubernetes-getting-started.rst)
RUN ARCH=${TARGETARCH:-$(case "$(uname -m)" in \
        "x86_64") echo "amd64" ;; \
        "aarch64") echo "arm64" ;; \
        *) echo "$(uname -m)" ;; \
    esac)} && \
    export PIP_DISABLE_PIP_VERSION_CHECK=1 && \
    curl -LsSf https://astral.sh/uv/install.sh | sh && \
    $HOME/.local/bin/uv venv ~/skypilot-runtime --seed --python=3.10 && \
    source ~/skypilot-runtime/bin/activate && \
    $HOME/.local/bin/uv pip install 'skypilot-nightly[remote,kubernetes]' \
        'ray[default]==2.9.3' 'pycryptodome==3.12.0' && \
    $HOME/.local/bin/uv pip uninstall skypilot-nightly && \
    curl -LO "https://dl.k8s.io/release/v1.33.12/bin/linux/$ARCH/kubectl" && \
    # Install kubectl to user's local bin instead of system path to avoid
    # sudo-related issues during cross-architecture builds, especially on ARM
    chmod +x kubectl && \
    mkdir -p $HOME/.local/bin && \
    mv kubectl $HOME/.local/bin/ && \
    echo 'export PATH="$PATH:$HOME/.local/bin"' >> ~/.bashrc

# Create a default, user-writable Python environment for task commands. This
# replaces the role conda's base env used to play, so user `pip install` writes
# to a user-writable location. SkyPilot's own runtime uses the separate
# ~/skypilot-runtime venv.
#
# Auto-activate it for real interactive login shells so `ssh <cluster>` lands in
# the same writable env that task setup/run use (mirroring conda's base env),
# but skip it everywhere else. The `[ -t 0 ]` (attached to a terminal) guard is
# essential and must NOT be replaced by an interactive-shell check ($-): task
# setup/run go through make_task_bash_script as `bash -i <file>` (no -c), which
# IS an interactive shell ($- contains 'i') yet is NOT attached to a terminal.
# Activating the venv there would `source activate`, which snapshots PATH into
# _OLD_VIRTUAL_PATH *before* cloud-CLI paths are appended to ~/.bashrc (e.g.
# gcloud's path.bash.inc, added at runtime by cloud-deps setup, always after
# this line). A later venv (re)activation then runs `deactivate`, restoring that
# pre-CLI PATH snapshot and dropping gcloud et al. -- which disables the cloud on
# the jobs controller (e.g. "GCP access is disabled" on storage prechecks). Real
# `ssh` sessions have a TTY, so they still get the venv; the task path relies on
# make_task_bash_script's own explicit activation instead. The CONDA_PREFIX guard
# lets an opt-in conda base (install_conda=true) take precedence.
RUN python3 -m venv ~/sky-user-env && \
    echo 'if [ -t 0 ] && [ -z "${BASH_EXECUTION_STRING:-}" ] && [ -z "${CONDA_PREFIX:-}" ] && [ -f ~/sky-user-env/bin/activate ]; then source ~/sky-user-env/bin/activate; fi' >> ~/.bashrc

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