How do multistage Docker builds help ML image size and security?

A multistage Docker build uses more than one FROM instruction in a Dockerfile. One stage prepares or builds the application; a later stage contains only what is needed to run it.

For ML services, this can keep compilers, temporary build files, test tools, and other build-only software out of the final inference image.

The key idea: Build with the tools you need, but deploy with only the files you need.

A beginner-friendly example

A food-delivery company serves a model that predicts delivery time. Its Python inference API needs packages such as FastAPI, Uvicorn, and an approved ML inference library.

A traditional single-stage Dockerfile might install compilation tools, build dependencies, application code, and tests in one image. Even if the API never uses those tools at runtime, the image can retain them.

A multistage approach separates those responsibilities:

• Builder stage: Prepare the Python environment and install the required packages.
• Runtime stage: Copy the prepared Python environment and application code into a separate image, without copying the builder's extra tools.

A practical multistage Dockerfile

Assume a project contains Dockerfile, requirements.txt, and app/main.py. The requirements file lists reviewed, compatible, version-pinned inference dependencies. app/main.py exposes a FastAPI instance named app, and the approved model is made available to the service through a documented artifact-loading process.

DOCKERFILE
# syntax=docker/dockerfile:1

# Stage 1: prepare Python dependencies
FROM python:3.12-slim AS builder

WORKDIR /build
RUN python -m venv /opt/venv
ENV PATH="/opt/venv/bin:$PATH"

COPY requirements.txt ./
RUN pip install --no-cache-dir -r requirements.txt

# Stage 2: build the final inference runtime
FROM python:3.12-slim AS runtime

ENV PATH="/opt/venv/bin:$PATH" \
    PYTHONDONTWRITEBYTECODE=1 \
    PYTHONUNBUFFERED=1

WORKDIR /app

COPY --from=builder /opt/venv /opt/venv
COPY app/ ./app/

RUN useradd --create-home --uid 10001 appuser
USER appuser

EXPOSE 8000
CMD ["python", "-m", "uvicorn", "app.main:app", "--host", "0.0.0.0", "--port", "8000"]

How it works:

1. FROM python:3.12-slim AS builder creates a named stage for preparing the application.
2. python -m venv /opt/venv creates a dedicated Python environment, and pip install installs the tested requirements there.
3. The second FROM starts a fresh runtime image. It does not automatically inherit all the filesystem layers of the first stage.
4. COPY --from=builder /opt/venv /opt/venv transfers the installed Python environment without copying the entire builder filesystem.
5. COPY app/ ./app/ adds the API code. The runtime then starts the FastAPI app as a non-root user.

This example intentionally uses the same Python base family and version in both stages, with the virtual environment at the same path. Otherwise, copied virtual environments and compiled Python extensions can break because of different interpreter paths, operating-system libraries, or CPU architectures.

If an ML dependency needs native compilation, the builder can install compilers and headers; the final image must still contain any runtime shared libraries the compiled package actually needs.

How much image size can this save?

Suppose, for illustration, an ML build uses 400 MB of extra compilers, temporary files, and test tools that the serving API does not need.

A multistage build can leave those items behind rather than include them in the runtime image. The actual size saved must be measured: Python frameworks, CUDA libraries, and packaged model weights can be much larger than the removed tools.

Inspect the final image rather than assuming a specific saving:

BASH
docker build -t delivery-time-api:multi .
docker image ls delivery-time-api:multi
docker run --rm -p 8000:8000 delivery-time-api:multi

Compare image size with a functionally equivalent single-stage build, then test model loading, health checks, and predictions.

Why can this improve security?

A smaller runtime image can reduce its attack surface:

• Compilers, debuggers, and build-only utilities are absent unless explicitly copied.
• Fewer unnecessary installed packages may mean fewer vulnerabilities in the deployed artifact.
• Build and test steps can run without shipping their temporary outputs.
• The deployed image contains a more narrowly defined set of runtime dependencies.

However, multistage builds are not a security guarantee. Vulnerable application libraries, unsafe model deserialization, a vulnerable base image, or accidentally copied secrets remain real risks.

Use .dockerignore to exclude private files from the build context. Never embed tokens in ARG or ENV: secrets can leak through layers or build metadata even if a later stage is minimal. Use Docker's dedicated build-secret mounts when a build genuinely needs a credential.

Continue to scan the final image, track dependency versions, and rebuild when relevant security fixes are available.

When is multistage especially useful for ML?

• Native dependencies: A library requires development headers or compilers to build but not to run.
• Model-serving containers: Training/test tools can be omitted from a lean inference runtime.
• CI builds: Run checks during a build stage while publishing only the intended runtime image.
• Multiple deliverables: Reuse an appropriately prepared builder but create separate runtime images for training workers and serving APIs.

For GPU serving, a tiny generic Python image is not automatically suitable. The runtime must still include compatible required CUDA/user-space libraries, and the host needs supported GPU drivers.

In short: Multistage builds separate how an ML image is built from what is deployed. Copy only the runtime environment and application into the final stage, test it on compatible infrastructure, and use security scanning to validate what remains.