Multi-stage builds are usually demonstrated with Go, where the answer is obvious: compile in one stage, copy one binary into scratch. Python has no single binary, which is why the pattern is often skipped for inference images — and why those images ship a C compiler to every node in the cluster.
Why one stage is the wrong shape
A Python inference service usually needs a toolchain at install time. Some dependency in the tree has no wheel for your platform and Python version, so pip falls back to building from source, which needs gcc, make, Python headers and often the headers for a native library too. Install those with apt-get in the same stage as your application and they are in the image forever.
Deleting them later does not help. Every instruction in a Dockerfile produces a layer, and layers are additive: a RUN apt-get purge in a later instruction adds a layer recording the deletion without removing the bytes from the layer that added them. The image is now larger than before you tried to clean it. The only way to keep something out of an image is to never put it in that image, and a second stage is how you get a second image.
The size is the visible cost. The security argument is the durable one: a runtime container with a compiler, a package manager and a network client is a materially more useful place to land than one without them.
The build stage
The first stage is allowed to be fat. Nothing in it ships unless you explicitly copy it.
# syntax=docker/dockerfile:1
FROM python:3.12-slim AS builder
RUN apt-get update && apt-get install -y --no-install-recommends \
build-essential \
python3-dev \
&& rm -rf /var/lib/apt/lists/*
ENV VIRTUAL_ENV=/opt/venv
RUN python -m venv $VIRTUAL_ENV
ENV PATH="/opt/venv/bin:$PATH"
COPY requirements.txt .
RUN --mount=type=cache,target=/root/.cache/pip \
pip install -r requirements.txt
Three things are load-bearing here. The AS builder names the stage so the second one can copy from it. The virtualenv puts every installed package under one prefix, which is the whole trick — without it, a Python install scatters files across site-packages, /usr/local/bin and sometimes /usr/lib, and there is no single path to copy. And setting PATH to the venv’s bin activates it for every subsequent instruction; there is no shell session to run source activate in.
The --mount=type=cache keeps pip’s download cache between builds without it ever becoming part of a layer. It requires the syntax directive on line one and BuildKit, which has been the default builder since Docker Engine 23.0. There is more on what that buys you in layer caching for AI images.
The runtime stage
FROM python:3.12-slim AS runtime
RUN useradd --system --create-home --uid 10001 appuser
COPY --from=builder --chown=10001:10001 /opt/venv /opt/venv
ENV PATH="/opt/venv/bin:$PATH" \
PYTHONDONTWRITEBYTECODE=1 \
PYTHONUNBUFFERED=1
WORKDIR /app
COPY --chown=10001:10001 src/ ./src/
USER 10001
EXPOSE 8000
CMD ["python", "-m", "uvicorn", "src.main:app", "--host", "0.0.0.0", "--port", "8000"]
The base image is the same one the builder used, and that is not incidental. A virtualenv contains compiled extension modules linked against a specific Python ABI and a specific libc. Build on python:3.12-slim (Debian, glibc) and copy into an Alpine image (musl) and imports fail at runtime with missing-symbol errors that look nothing like a packaging mistake. Keep the pair matched, or go the other way and target a distroless runtime built from the same Debian family.
PYTHONUNBUFFERED=1 matters more than it looks. Without it, stdout is block-buffered when it is not a terminal, so your logs arrive in 4 KB chunks and a container that dies mid-buffer loses its last words — precisely the lines you need when debugging a crash loop.
Copy an environment, not a pile of wheels
A common variant builds wheels in the first stage with pip wheel --wheel-dir /wheels and installs them in the second. It works, but it means the runtime stage still runs pip install, so pip and setuptools remain in the final image and the install work happens twice across the two stages.
Copying the whole virtualenv moves the work entirely into the builder. The one thing to know is that a venv created by python -m venv contains absolute paths — in pyvenv.cfg and in the shebang lines of console scripts under bin/. Copy it to the same absolute path in the runtime stage, as above, and nothing needs rewriting. Copy it somewhere else and every console entry point breaks with a confusing “bad interpreter” error. Invoking modules through python -m rather than by console script sidesteps the shebang issue entirely, which is why the CMD above is written that way.
If a dependency needs a shared library at runtime rather than only at build time — a BLAS implementation, an image codec, a CUDA runtime — that library must be installed in the runtime stage too. The venv contains Python packages, not the system libraries they link against. The failure is an ImportError naming a .so file, and the fix is a minimal apt-get install of the runtime package (not the -dev package) in the second stage.
The same split on a CUDA image
The pattern earns the most on GPU images, because NVIDIA already publishes the two halves you need. Its CUDA images come in base, runtime and devel variants: devel carries the compiler and the headers, runtime carries the shared libraries a compiled program links against, and base carries neither. A build that compiles a custom kernel or an extension against CUDA needs devel; the container that serves traffic does not.
FROM nvidia/cuda:12.4.1-devel-ubuntu22.04 AS builder
# nvcc, CUDA headers and build-essential are present here.
# Compile extensions, build wheels, populate /opt/venv.
FROM nvidia/cuda:12.4.1-runtime-ubuntu22.04 AS runtime
COPY --from=builder /opt/venv /opt/venv
ENV PATH="/opt/venv/bin:$PATH"
Keep the CUDA version identical across the two lines. A binary compiled against one CUDA runtime and executed against another is the same class of mismatch that produces CUDA driver version is insufficient for CUDA runtime version at start-up, except that here you caused it yourself in a Dockerfile rather than inheriting it from the node. Pin both tags together and change them together.
There is a case for skipping CUDA base images altogether. The PyTorch CUDA wheels vendor their own copies of the CUDA runtime, cuDNN and cuBLAS as nvidia-* packages inside site-packages, so if you compile nothing yourself, a plain python:3.12-slim runtime plus the driver libraries injected by the NVIDIA Container Toolkit is often sufficient and considerably smaller. Check before assuming you need the CUDA base: import your framework in the candidate image and confirm it reports a device.
Build it and check what shipped
- Build normally. BuildKit only executes stages the target depends on, so nothing extra is required:
docker build -t inference:local . - Build just the builder when you want to debug the install step:
docker build --target builder -t inference:builder . - Confirm the toolchain did not ship. The command should fail:
docker run --rm inference:local which gcc - Confirm the package set is complete:
docker run --rm inference:local pip listwill not work — pip is not installed — so usedocker run --rm inference:local python -c "import torch, fastapi"with your real top-level imports. - Compare the two:
docker image ls --format '{{.Repository}}:{{.Tag}} {{.Size}}' | grep inference. The gap between builder and runtime is what the second stage bought you, and it is your number rather than one from a tutorial.
Top comments (0)