# Pipeline-parallel BASE image — heavy, rarely-changing runtime foundation.
#
# This is the expensive layer: CUDA runtime libs, tinygrad + numpy, the
# NVRTC headers tinygrad JITs against, sshd, and the PID-1 supervisor. It
# carries NO pipeline binaries and NO worker script — those live in the thin
# code image (`Dockerfile`) built FROM this one, so a code push rebuilds only
# a few MB instead of re-pushing ~1 GB.
#
# Spec §4.10 (permanent contracts):
#   * MUST be produced by a multi-stage build, with build artifacts
#     (compilers, headers, dev libraries) confined to the builder stage.
#   * MUST be ≤ 1 GB compressed (§4.10 / acceptance §6.11).
#   * Apt caches, pip caches, __pycache__, test data, and docs MUST NOT
#     be present in the runtime layer.
#   * SHOULD use a slim CUDA runtime image, not a -devel image.
#
# Build context must be the workspace root:
#   docker build -f apps/pipeline-parallel-inference/Dockerfile.base \
#       -t swactor-pp-base:cuda12.6 .

# ─── Builder stage ───────────────────────────────────────────────────
# Confines the dev-only CUDA headers (cuda-cudart-dev) and the pip
# install machinery here so the runtime layer keeps neither. The pip
# install also produces __pycache__ + bundled tests; we strip both
# inside this stage so they cannot ride along on any COPY out.
FROM nvidia/cuda:12.6.3-runtime-ubuntu24.04 AS builder

RUN apt-get update && \
    apt-get install -y --no-install-recommends \
        python3 \
        python3-pip \
        ca-certificates \
        cuda-cudart-dev-12-6 && \
    rm -rf /var/lib/apt/lists/*

# Stage tinygrad + numpy into a self-contained directory we COPY into
# the runtime stage. --target keeps them off the system path so the
# runtime can drop them under PYTHONPATH without dragging /usr/lib.
RUN python3 -m pip install --no-cache-dir --break-system-packages \
        --target=/opt/pp-pydeps \
        tinygrad==0.12.0 numpy && \
    find /opt/pp-pydeps -depth -type d \
        \( -name '__pycache__' -o -name 'tests' -o -name 'test' \) \
        -exec rm -rf {} + && \
    find /opt/pp-pydeps -name '*.pyc' -delete && \
    find /opt/pp-pydeps -type d -name '*.dist-info' -exec rm -rf {} +

# Stage the NVRTC headers tinygrad's generated fp16 / bf16 kernels
# `#include`. cuda_fp16.h transitively pulls in vector_types.h and the
# rest of the cudart header tree, so the whole include dir must ship —
# copying only the fp16/bf16 files leaves NVRTC unable to open
# vector_types.h at JIT time. The full cudart-dev header set is a few MB
# of pure source; deref symlinks (-L) so the real files land here. The
# `test -f` guard fails the build if vector_types.h is ever missing
# again rather than shipping a broken image.
RUN mkdir -p /opt/pp-nvrtc-include && \
    cp -rL /usr/local/cuda/include/. /opt/pp-nvrtc-include/ && \
    test -f /opt/pp-nvrtc-include/vector_types.h

# ─── Runtime stage ───────────────────────────────────────────────────
# Slim CUDA -base image (no math libs) plus exactly the libraries
# tinygrad needs at JIT time: libcudart (CUDA runtime) and libnvrtc
# (kernel compiler). The cuda_fp16.h header comes from the builder; no
# -dev / -devel package lands here.
FROM nvidia/cuda:12.6.3-base-ubuntu24.04 AS runtime

# Full python3 (not -minimal) for stdlib coverage tinygrad+numpy need
# (numpy 2.x imports `contextvars` from stdlib; python3-minimal omits
# it). `apt-get clean` + `rm -rf` keep apt's archive cache out of the
# layer (spec §4.10 forbids it). The find calls scrub __pycache__ and
# .pyc generated by post-install scripts of the packages we DO need.
RUN apt-get update && \
    apt-get install -y --no-install-recommends \
        python3 \
        cuda-cudart-12-6 \
        cuda-nvrtc-12-6 \
        ca-certificates \
        procps \
        openssh-server && \
    apt-get clean && \
    rm -rf /var/lib/apt/lists/* \
           /var/cache/apt/archives/* \
           /var/cache/apt/*.bin \
           /var/log/apt/* \
           /var/log/dpkg.log \
           /tmp/* \
           /var/tmp/* \
           /usr/share/doc/* \
           /usr/share/man/* \
           /usr/share/info/* && \
    find /usr -depth -type d -name '__pycache__' -exec rm -rf {} + 2>/dev/null || true && \
    find /usr -type f -name '*.pyc' -delete 2>/dev/null || true

# Headers tinygrad's NVRTC backend resolves via its default -I path.
COPY --from=builder /opt/pp-nvrtc-include/ /usr/local/cuda/include/

# Pip-installed Python deps from the builder, already stripped of
# __pycache__ + bundled tests (spec §4.10: caches / test data MUST NOT
# be present in the runtime layer).
COPY --from=builder /opt/pp-pydeps /opt/pp-pydeps
ENV PYTHONPATH=/opt/pp-pydeps
# Avoid regenerating bytecode at runtime — keeps the runtime FS clean
# of fresh __pycache__ writes after first import.
ENV PYTHONDONTWRITEBYTECODE=1

# PID-1 supervisor: brings up sshd and runs the worker as a child so SSH is
# deterministic on every node and a worker crash leaves a reachable shell.
COPY apps/pipeline-parallel-inference/pp_entrypoint.sh /usr/local/bin/pp_entrypoint.sh
RUN chmod +x /usr/local/bin/pp_entrypoint.sh

# Optional deploy public key baked at build time as a fallback that does not
# depend on vast injecting PUBLIC_KEY into the container env:
#   docker build --build-arg DEPLOY_PUBKEY="$(cat ~/.ssh/pp_deploy.pub)" ...
ARG DEPLOY_PUBKEY=""
RUN if [ -n "$DEPLOY_PUBKEY" ]; then printf '%s\n' "$DEPLOY_PUBKEY" > /etc/pp_deploy_key.pub; fi

# Enable CUDA backend for tinygrad (override with -e CUDA=0 for CPU runs).
ENV CUDA=1

ENTRYPOINT ["/usr/local/bin/pp_entrypoint.sh"]