Docker Builds: What’s Inside a 40GB ML Training Container

Every time a researcher submits a training job, they pull a 40GB container image. Nobody asks what’s in it. It works, the model trains, life goes on. But I spent a week optimizing our image build pipeline and had to understand every single layer. Here’s what’s actually inside the thing.

The two-layer cake

The image is built from two chained Dockerfiles, connected through a Docker Bake file:

docker/Dockerfile              → base image (30+ min cold build)
  └── docker/pretraining/Dockerfile  → Python packages + workspace (8-15 min)

Stage one is the foundation: NVIDIA base image, system tools, and libraries compiled from source. Stage two adds the Python world on top. This split exists because the base layer changes rarely — maybe once a week when someone bumps a CUDA library commit pin — while the Python layer changes on nearly every PR that touches uv.lock.

Two Dockerfiles means two cache boundaries. Change a Python dependency and you rebuild 15 minutes of Python installation. Change a CUDA kernel compilation flag and you rebuild 40 minutes of everything. The split is load-bearing.

Layer 0: NVIDIA’s gift

The base of the base is nvcr.io/nvidia/pytorch:25.03-py3 from NVIDIA’s NGC catalog. This single image gives you:

This is roughly 15-20GB on its own. NVIDIA rebuilds it monthly, tuned to specific GPU architectures. We pin to a specific tag because “latest” on a GPU image is a footgun — CUDA toolkit versions must match driver versions on the host, and a mismatch means CUDA error: no kernel image is available for execution on the device at runtime.

You don’t build this layer. You inherit it and pray NVIDIA tested it properly.

Layer 1: system tools nobody thinks about

On top of the NGC image, the base Dockerfile installs a grab bag of system utilities:

tmux, zsh, git, vim, curl, wget, htop, tree, jq, rsync
kubectl, azcopy, azure-cli
Redis (server + client)

This is the “oh right, humans need to SSH into these containers and debug things” layer. Researchers will kubectl exec into a training pod at 2am because a job is hanging, and they need htop to see if it’s a CPU bottleneck or nvidia-smi to check GPU utilization. Without tmux they’d lose their session when the SSH connection drops. Without git they can’t check what code is actually running.

These add maybe 1-2GB and take 2-3 minutes to install. Not the expensive part.

Layer 2: the CUDA gauntlet

This is where the build time lives. Four libraries compiled from source, sequentially, each taking 5-20 minutes:

These are sequential because Docker layers are sequential. Layer N+1 doesn’t start until layer N finishes. If the Flash Attention commit pin changes, everything below it in the Dockerfile rebuilds too — even if Triton and NVSHMEM haven’t changed. Docker doesn’t know they’re independent.

This is 25-40 minutes of wall clock time, producing maybe 5-8GB of compiled .so files. The ratio of build time to disk space is horrifying. You’re spending half an hour to produce a few gigabytes of highly optimized GPU machine code.

Layer 3: Python packages

The pretraining Dockerfile takes over here. Two uv sync calls, carefully separated:

# External deps only — busted by uv.lock changes
RUN --mount=type=cache,target=/root/.cache/uv \
    uv sync --no-install-workspace --frozen --package mai_job

# Workspace packages — busted by code changes
RUN --mount=type=cache,target=/root/.cache/uv \
    uv sync --package mai_job --frozen

The first call installs external packages from PyPI — PyTorch (already present from NGC, but uv needs to reconcile it), transformers, numpy, the usual ML stack. 3-5 minutes cold. The second installs workspace packages — our own code, as editable installs.

The split means a code-only change doesn’t re-download PyTorch. Only changes to uv.lock (new dependencies) trigger the expensive first step. This is actually good Dockerfile design.

The --mount=type=cache flag keeps uv’s download cache across builds. When warm, this whole layer takes 1-2 minutes. When cold — like after the runner’s Docker disk gets garbage collected because it’s 99% full — 15 minutes.

Two formats, same contents

The same image ships in two packaging formats:

Enroot is NVIDIA’s container runtime for HPC. It takes a Docker image, flattens it into a squashfs file (read-only compressed filesystem), and mounts it directly — no Docker daemon, no overlay filesystem, no container runtime overhead. Slurm nodes don’t run Docker. They run bare metal with Enroot.

Same bits, different packaging. The CI pipeline builds the Docker image, pushes it to four ACRs, then also converts it to .sqsh and drops it on the shared filesystem. Two distribution channels for two types of compute infrastructure.

The delivery pipeline

Here’s how a built image gets from CI to a researcher’s training job:

1. PR merges to main
2. CI detects changes (hashFiles('mai_kernels/**/*.cu', 'uv.lock', 'docker/*/Dockerfile'))
3. Builds image with git SHA tag
4. Pushes to 4 ACRs + converts to .sqsh
5. Updates constants.py with new SHA tag
6. Commits constants.py back to main
7. Researcher submits job
8. Job launcher reads constants.py → gets image tag
9. Cluster pulls image from nearest ACR (or mounts .sqsh)

Step 6 is the sneaky one. The build pipeline commits back to the repo. constants.py is the handshake between CI and runtime — it stores the current image SHA so that job submission code knows which image to use without anyone manually updating a config file.

This means image selection is automatic. Researchers don’t specify image tags. They call a job launcher, which reads constants.py, which was auto-updated by CI. The dispatcher in resource_request.py picks the right registry based on cluster type — Mango cluster gets the Mango ACR URL, Falcon gets Falcon, and so on.

The rebuild trigger problem

The CI trigger condition:

hashFiles('mai_kernels/**/*.cu', 'uv.lock', 'docker/*/Dockerfile')

Any change to CUDA kernels, Python dependencies, or the Dockerfiles themselves triggers a full rebuild. This sounds reasonable until you realize that “full rebuild” means 40 minutes of CI time, and uv.lock changes on roughly a third of PRs.

Every Python dependency update — even adding a single test-only package — triggers a complete image rebuild including the CUDA compilation gauntlet. Because all Python packages are baked into the image. There’s no “just update this one pip package” — the image is immutable, so any change means building a new one.

This is the cost of reproducibility. You get an image where every byte is deterministic for a given commit SHA. You pay for it with 40-minute rebuilds when someone adds pytest-timeout to the dev dependencies.

The size breakdown

Approximate layer sizes:

Nearly half the image is the NGC base that we don’t control. A quarter is Python packages. The custom CUDA compilations — the ones that take 80% of the build time — account for only 15% of the disk space.

The build time distribution is almost the inverse of the size distribution. The biggest layers (NGC, Python) are fast to pull or cache. The smallest layers (CUDA compilations) are the ones that eat your CI budget.

Why this matters

Every researcher on the team pulls this image multiple times a day. Every CI run that touches the trigger files rebuilds it. The image is the interface between “code in a git repo” and “model training on GPUs.” It’s invisible infrastructure — until the build breaks, or takes 50 minutes instead of 40, or a CUDA compilation silently produces wrong results because someone bumped a commit pin without testing.

Understanding what’s in the image isn’t academic curiosity. It’s the difference between “CI is slow” and knowing that the Flash Attention 3 layer takes 20 minutes and could be parallelized if Docker Bake supported concurrent layer builds for dependent stages. It’s the difference between “the image is too big” and knowing that 45% of it is NVIDIA’s base image that you can’t shrink.

You can’t optimize what you can’t see. Now you can see it.