haochenw/nanochat_runpod repository overview

# syntax=docker/dockerfile:1

# =============================================================================
# MULTI-STAGE BUILD
#
# This Dockerfile uses two FROM statements, creating two "stages":
#
#   Stage 1 ("build") — a throwaway workspace where we install tools (uv) and
#   download/compile all Python dependencies into a .venv.
#
#   Stage 2 ("runtime") — starts fresh from a clean base image. We selectively
#   COPY the .venv and the Python interpreter from stage 1:
#       COPY --from=build /python /python
#       COPY --from=build /software/.venv /software/.venv
#
# Why bother? The build stage accumulates package managers, caches, compilers,
# and other artifacts that bloat the image but aren't needed at runtime. By
# starting fresh in stage 2, none of that baggage ends up in the final image.
# Only the final stage contributes to the image you actually run.
# =============================================================================

# =============================================================================
# PYTHON VERSION
#
# There is no python base image here. Instead, uv downloads a standalone
# Python build (python-build-standalone) and installs it to UV_PYTHON_INSTALL_DIR.
# The .venv then symlinks to that installation.
# The version is controlled by pyproject.toml's `requires-python` constraint
# (or the --python flag on the uv sync command below).
#
# These standalone builds statically link libssl, libffi, zlib, libsqlite3,
# etc. The only dynamic dependency is glibc, which the Debian base provides.
# So the runtime stage doesn't need any Python-related system packages.
#
# Image hierarchy for reference (buildpack-deps is Debian-based):
#   debian:bookworm
#     → buildpack-deps:bookworm-curl   (adds curl, wget, ca-certs)
#       → buildpack-deps:bookworm-scm  (adds git, ssh, mercurial)
#         → buildpack-deps:bookworm    (adds gcc, make, dev headers)
# =============================================================================

# =============================================================================
# Stage 1: BUILD — uv installs Python + dependencies into a .venv
# =============================================================================
FROM buildpack-deps:bookworm AS build

# Copy the uv binary from Astral's official distroless image.
# Pin to an exact version for reproducibility. This is preferred over
# `curl | sh` because it's a single static binary, no apt deps needed.
COPY --from=ghcr.io/astral-sh/uv:0.11.2 /uv /uvx /bin/

# UV_COMPILE_BYTECODE=1 — pre-compile .py → .pyc at install time so you
#   don't pay the compilation cost on every import at runtime.
# UV_LINK_MODE=copy — when using cache mounts the cache and target live on
#   different filesystems, so hardlinks fail. `copy` avoids that cleanly.
# UV_PYTHON_INSTALL_DIR=/python — by default uv installs Python to
#   /root/.local/share/uv/python/, but the .venv symlinks to that path.
#   In a multi-stage build, if we only copy the .venv, those symlinks break.
#   Setting this to /python gives us a predictable path to also COPY to
#   the runtime stage.
ENV UV_COMPILE_BYTECODE=1 \
    UV_LINK_MODE=copy \
    UV_PYTHON_INSTALL_DIR=/python

# IMPORTANT: this path must match the COPY destination in stage 2.
# uv bakes absolute paths into shebangs (e.g. #!/software/.venv/bin/python)
# and pyvenv.cfg. If the paths differ between stages, scripts will break.
WORKDIR /software

# Copy ONLY the dependency specification files. By isolating this layer,
# Docker can cache the (expensive) dependency install and skip it entirely
# when only your source code changes.
COPY pyproject.toml uv.lock ./

# uv reads `requires-python` from pyproject.toml, downloads the matching
# Python version automatically, and installs all dependencies into .venv.
#   --no-install-project : install only the dependency tree, not your package
#
# The cache mount persists uv's download cache across builds so only changed
# packages are re-fetched on rebuild.
RUN --mount=type=cache,target=/root/.cache/uv \
    uv sync --extra gpu --group dev --python 3.12 --no-install-project

# =============================================================================
# Stage 2: RUNTIME — interactive environment for RunPod
# =============================================================================
# Using the full buildpack-deps image for interactive work — it includes
# git, curl, wget, gcc, make, ssh, and common dev libraries out of the box.
# uv's standalone Python statically links its own dependencies (libssl,
# libffi, zlib, etc.), so we don't need this image for Python's sake —
# it's purely for a comfortable interactive environment.
FROM buildpack-deps:bookworm

# --- RunPod SSH setup ---
# RunPod injects a PUBLIC_KEY env var at runtime. The start script below
# writes it to authorized_keys and starts sshd. This requires root, so
# we do NOT switch to a non-root USER (RunPod's standard pattern).
RUN apt-get update && \
    apt-get install -y --no-install-recommends openssh-server vim && \
    rm -rf /var/lib/apt/lists/* && \
    mkdir -p /var/run/sshd && \
    # Allow root login via public key only (no password)
    sed -i 's/#PermitRootLogin.*/PermitRootLogin prohibit-password/' /etc/ssh/sshd_config && \
    sed -i 's/#PubkeyAuthentication.*/PubkeyAuthentication yes/' /etc/ssh/sshd_config && \
    sed -i 's/#PasswordAuthentication.*/PasswordAuthentication no/' /etc/ssh/sshd_config

# Copy both the Python interpreter AND the .venv from the build stage.
# The .venv contains symlinks like:
#   .venv/bin/python → /python/cpython-3.12.../bin/python3.12
# So we must copy /python too, or the symlinks break.
COPY --from=build /python /python
COPY --from=build /software/.venv /software/.venv

# Put the venv's bin directory on PATH. This "activates" the venv globally
# for every shell session — no need to `source activate`.
ENV PATH="/software/.venv/bin:$PATH"

WORKDIR /workspace

# Start script: sets up SSH authorized_keys from RunPod's PUBLIC_KEY env var,
# starts sshd, then drops into bash for interactive use.
COPY runpod_start_script.sh /runpod_start_script.sh
RUN chmod +x /runpod_start_script.sh

EXPOSE 22

CMD ["/runpod_start_script.sh"]

#!/bin/bash

# RunPod start script — runs at container launch.
# Sets up SSH access using the PUBLIC_KEY env var injected by RunPod,
# then keeps the container alive with an interactive shell.

if [[ -n "$PUBLIC_KEY" ]]; then
    mkdir -p ~/.ssh
    chmod 700 ~/.ssh
    echo "$PUBLIC_KEY" >> ~/.ssh/authorized_keys
    chmod 600 ~/.ssh/authorized_keys
    service ssh start
    echo "SSH server started."
else
    echo "WARNING: No PUBLIC_KEY set — SSH will not be available."
fi

# Keep the container alive. RunPod doesn't attach a TTY, so `bash` would
# exit immediately. `sleep infinity` blocks forever, letting you SSH in.
sleep infinity