112 lines
5.1 KiB
Docker
112 lines
5.1 KiB
Docker
# 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"]
|