"Works on my machine" is the most expensive bug in ML — a CUDA mismatch between training and serving can silently flip results, and you'll spend two weeks finding it. Docker, pinned versions, and multi-stage builds make that bug structurally impossible. Every production ML platform — SageMaker, Vertex AI, Modal, Anyscale — runs containers underneath.
Learning Objectives
After this lesson, you will be able to:
Build production-grade ML Docker images with locked dependencies, CUDA-friendly base images, and minimal layer count
Distinguish 'works on my machine' (Conda envs, virtual envs) from 'works anywhere' (Docker images with pinned versions)
Use multi-stage builds + .dockerignore + layer caching to keep ML images small and rebuilds fast
Pick the right base image — PyTorch official, NVIDIA CUDA, distroless — for your training vs inference workload
A container packages: OS userspace, system libraries (CUDA, cuDNN, NCCL), Python interpreter, packages, your code, and your model artifacts into a single immutable image. Anywhere that runs the container, the environment is identical.
For ML this matters more than for typical web apps because:
GPU drivers + CUDA + cuDNN versions are tightly coupled
Floating-point determinism depends on hardware + library versions
Models are large (GB-scale) and need predictable mounts
Distributed training requires identical environments across nodes
A multi-stage build is a single Dockerfile with two or more FROM lines. Each FROM opens a fresh image; you can COPY --from=<stage> to pull artifacts from one stage into another. For ML, the canonical use case is keeping the heavy training toolchain out of the serving image.
dockerfile
# ---------- Stage 1: build/train (heavy toolchain) ----------
FROM nvidia/cuda:12.4.0-cudnn-devel-ubuntu22.04 AS build
RUN apt-get update && apt-get install -y gcc g++ make python3.11 python3-pip git
COPY requirements-train.txt /tmp/
RUN pip install --no-cache-dir -r /tmp/requirements-train.txt
COPY . /src
WORKDIR /src
RUN python train.py --output /artifacts/model.pt
# ---------- Stage 2: serve (lean runtime) ----------
FROM nvidia/cuda:12.4.0-cudnn-runtime-ubuntu22.04 AS serve
RUN apt-get update && apt-get install -y python3.11 python3-pip && rm -rf /var/lib/apt/lists/*
COPY requirements-serve.txt /tmp/
RUN pip install --no-cache-dir -r /tmp/requirements-serve.txt
COPY --from=build /artifacts/model.pt /model/model.pt
COPY serve.py /app/serve.py
CMD ["python3", "/app/serve.py"]
The serving image only contains: CUDA runtime (not the dev SDK), Python, the inference-side packages, the trained model artifact, and a serving script. You drop gcc, the training data loaders, Jupyter, profilers, etc. Typical size differences in 2026:
Image
Approximate uncompressed size
nvidia/cuda:12.4.0-cudnn-devel-ubuntu22.04 (full training toolchain)
The size difference matters for two reasons: pull latency on every node (each cold pod has to download the image once), and attack surface (every binary in the image is potentially exploitable). A 7-GB image is ~30s to pull on a 2 Gbps link; a 600-MB image is ~3 seconds.
#Choosing requirements management: requirements.txt vs poetry vs uv vs pixi
The 2024-2026 landscape for Python dependency management exploded with new tools. Here's the practical map:
Tool
What it is
Sweet spot for ML
Caveats
requirements.txt (+ pip-tools)
Plain list of pinned versions
Universally understood; works in every Docker base
Doesn't track transitive locking by default; pair with pip-compile
poetry
Lockfile + virtualenv + publishing
Library packaging; project metadata in pyproject.toml
Slower to resolve than uv; less attractive for pure training pipelines
uv (Astral, 2024)
Rust-based pip-replacement + resolver
Fastest installs in 2026 (~10x pip); first-class lockfile; venv management
Newer; some niche packages still need pip fallback
pixi (Prefix.dev, 2024)
Conda-compatible package manager built on conda-forge
When you need conda-forge packages (CUDA, FFmpeg, RAPIDS)
Smaller ecosystem than pip/uv for pure Python deps
conda/mamba
Classic conda environment
Legacy ML stacks that need conda-forge binaries
Slower; heavier base images
The 2026 default for new ML Docker images is uv + a locked requirements.txt. It's pip-compatible, resolves in milliseconds, and produces fully deterministic installs across machines. Reach for pixi only when you need non-Python system binaries that pip can't provide. Reach for poetry when you're packaging a library, not just deploying a training job.
A regular docker run cannot access the GPU. NVIDIA's container toolkit (nvidia-container-toolkit, formerly nvidia-docker) injects the host GPU drivers and /dev/nvidia* device nodes into the container at start time.
bash
# Install nvidia-container-toolkit on the host (one time):
distribution=$(. /etc/os-release;echo $ID$VERSION_ID)
curl -s -L https://nvidia.github.io/libnvidia-container/$distribution/libnvidia-container.list \
| sudo tee /etc/apt/sources.list.d/nvidia-container-toolkit.list
sudo apt update && sudo apt install -y nvidia-container-toolkit
sudo systemctl restart docker
# Now run a container with GPU access:
docker run --gpus all nvidia/cuda:12.4.0-base-ubuntu22.04 nvidia-smi
docker run --gpus '"device=0,1"' my-train-image python train.py # only GPUs 0 and 1
In Kubernetes, the same plumbing is provided by the NVIDIA GPU Operator which manages drivers, the container runtime, and DCGM (monitoring) as DaemonSets. You request nvidia.com/gpu: 1 in your pod spec and the scheduler places it on a node with available GPUs.
The pattern: declare your workload as a Custom Resource (PyTorchJob, InferenceService, RayCluster), the Operator reconciles it. You don't write 200 lines of pod spec to spin up 8-node distributed training; you write 30 lines of YAML.
Reproducibility extends to the build pipeline. The minimum bar for a 2026 ML CI/CD pipeline:
Pin everything in the Dockerfile — base image SHA, system package versions, Python deps via uv.lock or requirements.txt.
Tag images by git SHA — myorg/train:abc123def, not myorg/train:latest. latest is poison; it's the tag that means "I don't know which version".
Reproducible layers — use BUILDKIT_INLINE_CACHE=1 and a remote cache so two CI runs of the same git SHA produce byte-identical images.
Sign images — cosign sign after build; verify on deploy. Required by SLSA Level 3 (and the EU AI Act for high-risk systems).
Don't bake secrets — use BuildKit --mount=type=secret for pip indices or model registry tokens; never COPY .npmrc-style secret baking.
Promote images, don't rebuild them — the image that runs in staging is the same SHA that goes to prod. Rebuilding "the same Dockerfile" produces a different image because base image SHAs, transitive dep versions, and apt mirrors drift hourly.
Tests · Verify the image builds successfully. Verify it runs with --gpus all. Verify .dockerignore excludes data and __pycache__. Use 'docker history' to confirm layers are ordered correctly.
#Compute the Dockerfile Size Difference: CUDA-full vs Slim
Let's quantify the layer-by-layer size of a typical ML Docker image and see exactly where the bytes go. This lets you reason about caching, pull latency, and the cost of multi-stage separation.
Loading visualization...
The 60-70% size reduction between the training and inference image isn't only about disk: it shows up directly in cold-start latency on every new pod, in attack surface (every binary in the image is a potential CVE), and in deploy time. Once you internalize this, multi-stage builds stop being a "best practice" and become the obvious default.
What Do You Think?
A team's training image is 14 GB. They deploy via Kubernetes; each new pod takes 4-6 minutes to start because the image must be pulled. Most defensible fix?
Quick check
Your CI builds an image, tags it `:latest`, and pushes. Staging pulls it, looks great. Hours later production pulls `:latest` and behaves differently. What's the most likely cause?