11. Building good images for AI workloads

A good AI image is pinned, built from a lockfile, small enough to pull quickly, built for the CPU architecture it will run on, and deliberate about where model weights live. Each of those is a decision you should be able to defend.

Pin everything that can move

• Base image: an exact tag at least, a digest for anything you publish or cite.
• Python dependencies: a lockfile with exact versions and hashes (uv.lock, poetry.lock, or pip-compile --generate-hashes output). A plain requirements.txt with torch>=2.0 is a wish, not a pin.
• The PyTorch build: install from the index that matches your CUDA target, for example https://download.pytorch.org/whl/cu124, or /whl/cpu for CPU-only.
• Provenance: log the image digest in your experiment tracker (Weights & Biases, MLflow) with every run. Together with the lockfile, that’s the reproducibility story for a thesis or paper.

Where the gigabytes come from

A CUDA build of PyTorch alone installs several gigabytes of NVIDIA libraries. On top of that come devel base images with compilers, NGC images that bundle a full toolchain, and leftovers: pip caches, apt lists, test data, build tools.

Tactics that actually shrink images:

1. Use a slim or runtime base and a multi-stage build, so compilers stay behind.
2. Don’t install CUDA twice. An nvidia/cuda runtime base plus a PyTorch wheel that bundles its own CUDA libraries duplicates gigabytes; for PyTorch-only images, python:3.12-slim is usually the smaller base.
3. Use the CPU-only PyTorch build for CPU inference services (embeddings on CPU, small classifiers). It’s a fraction of the size.
4. Keep caches out of layers: pip install --no-cache-dir, or better, BuildKit cache mounts.
5. Measure, don’t guess: docker history and dive show which layer is heavy.

A GPU training image

# syntax=docker/dockerfile:1
FROM python:3.12-slim
ENV PYTHONUNBUFFERED=1 PIP_DISABLE_PIP_VERSION_CHECK=1
RUN apt-get update \
 && apt-get install -y --no-install-recommends git libgomp1 \
 && rm -rf /var/lib/apt/lists/*
WORKDIR /workspace

# Locked dependencies, including a CUDA build of torch
COPY requirements.lock .
RUN --mount=type=cache,target=/root/.cache/pip \
    pip install --extra-index-url https://download.pytorch.org/whl/cu124 \
    -r requirements.lock

COPY src/ src/
COPY configs/ configs/
ENTRYPOINT ["python", "-m", "src.train"]
CMD ["--config", "configs/base.yaml"]

docker run --rm --gpus all --shm-size=16g \
  -v /data/my-dataset:/data:ro \
  -v "$PWD"/runs:/workspace/runs \
  -e WANDB_API_KEY \
  trainer:1.0 --config configs/large.yaml

-e WANDB_API_KEY with no value passes the variable through from your shell, so the key never appears in the command or the image.

Compiled extensions

Packages such as flash-attn need compiling against the exact torch and CUDA versions in your image. Build them in a devel stage that matches your runtime’s CUDA version, produce a wheel, and copy only the wheel into the runtime stage. Two environment variables keep compilation manageable: TORCH_CUDA_ARCH_LIST limits which GPU architectures you compile for (fewer means much faster builds), and MAX_JOBS limits parallel compile jobs so the build doesn’t run out of RAM. Check first whether a prebuilt wheel exists for your versions; it often saves an hour.

conda when you need it

Some scientific stacks still depend on conda-only packages. Use a micromamba base image, create the environment from a locked environment.yml (or a conda-lock file), and clean the package cache in the same RUN. Mixing conda and pip in one environment works but is the most common source of broken, unreproducible images; let one tool own each package.

Model weights: decide deliberately

Part 7 compares baking weights in, mounting a cache, and downloading at startup. Two image-specific rules apply on top. Never bake access tokens into an image; if a build must download gated weights, use a BuildKit secret mount. And if you bake weights in, put them in their own layer before your code, so a code change doesn’t force a re-upload of tens of gigabytes.

Multi-architecture builds

Your M1 Pro builds linux/arm64 images by default. Push one to an amd64 GPU server and it fails immediately with “exec format error”. Choose the target platform explicitly:

# For an amd64 server only
docker buildx build --platform linux/amd64 -t ghcr.io/ulvi/trainer:1.0 --push .

# One tag that works on both (creates an image index)
docker buildx build --platform linux/amd64,linux/arm64 -t ghcr.io/ulvi/api:1.0 --push .

Cross-building runs the foreign architecture under QEMU emulation, so heavy steps (compiling, large pip installs) can be many times slower. Better options: build on the target machine through a Docker context (Part 10), build in CI on amd64 runners, or register a remote builder with docker buildx create. Many GPU images and wheels target amd64 first, so GPU images are usually amd64-only. docker run --platform linux/amd64 on the Mac is fine for a smoke test, too slow for real work.

Development versus production images

Use one Dockerfile with several targets. A dev stage adds Jupyter, debuggers and test tools, and expects your code bind-mounted; a prod stage is minimal, non-root, with no compilers and the code copied in. Build them with --target dev and --target prod. Dev Containers (Part 15) build on the same idea for your editor.