# Pipeline-parallel runtime image.
#
# 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.
#   * MUST be self-sufficient — booting MUST NOT fetch any binary from
#     an external host (the §5.1 binary-swap path is opt-in only).
#   * SHOULD use a slim CUDA runtime image, not a -devel image.
#
# Build context must be the workspace root:
#   docker build -f examples/pipeline-parallel-inference/Dockerfile -t <tag> .

# ─── 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`. They are pure source files (kilobytes); the runtime stage
# picks them up without the full cuda-cudart-dev package.
RUN mkdir -p /opt/pp-nvrtc-include && \
    cp /usr/local/cuda/include/cuda_fp16.h \
       /usr/local/cuda/include/cuda_fp16.hpp \
       /opt/pp-nvrtc-include/ && \
    if [ -f /usr/local/cuda/include/cuda_bf16.h ]; then \
        cp /usr/local/cuda/include/cuda_bf16.h \
           /usr/local/cuda/include/cuda_bf16.hpp \
           /opt/pp-nvrtc-include/; \
    fi

# ─── 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 && \
    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

# Pipeline binaries + worker script. The binaries are statically linked
# enough that the runtime stage's libc is all they need; the worker is
# pure Python.
COPY examples/pipeline-parallel-inference/target/release/pp-gpu-node /usr/local/bin/pp-gpu-node
COPY examples/pipeline-parallel-inference/target/release/pp-smoke-run /usr/local/bin/pp-smoke-run
COPY examples/pipeline-parallel-inference/pp_tinygrad_worker.py /usr/local/share/pp_tinygrad_worker.py

# Enable CUDA backend for tinygrad (override with -e CUDA=0 for CPU runs).
ENV CUDA=1
ENV WORKER_SCRIPT=/usr/local/share/pp_tinygrad_worker.py

CMD ["pp-gpu-node"]