Containerising an AI Service Properly

An AI service image is large for three or four specific reasons, and each one has a named fix. The general advice — use multi-stage builds, use slim base images — is correct but does not tell you which of your layers is the four-gigabyte one. This page starts from the arithmetic of the layers, so you can decide what to remove before you start removing things.

Where the gigabytes actually are

A container image is an ordered stack of layers, and its size on disk is the sum of the layers plus the metadata. Nothing you delete in a later layer removes bytes from an earlier one — a RUN rm -rf /root/.cache in step nine does not shrink the pip install in step eight, it adds a whitespace layer that hides the files. This single fact explains most of the confusion about why an image did not get smaller.

For a typical Python model-serving image, the layers sort into four groups, and it is worth knowing which group you are fighting before you pick a technique.

The important consequence: if your service does inference on a GPU through an API call to a remote model, or runs a small model on CPU, the second group is pure waste and it is the single biggest win available. Installing the CPU wheel index explicitly is one line.

# CPU-only PyTorch — no CUDA wheels pulled in as dependencies
pip install --no-cache-dir torch --index-url https://download.pytorch.org/whl/cpu

The multi-stage build

A multi-stage build compiles in one image and copies only the results into another. For Python the unit to copy is a virtual environment: build it in the first stage with whatever compilers are needed, then copy the whole directory into a runtime stage that has no compilers at all.

# syntax=docker/dockerfile:1.7
# ---------- stage 1: build ----------
FROM python:3.11-slim AS build

ENV PIP_DISABLE_PIP_VERSION_CHECK=1 \
    PYTHONDONTWRITEBYTECODE=1

# Build-only packages. None of these reach the final image.
RUN apt-get update && apt-get install -y --no-install-recommends \
        build-essential git \
    && rm -rf /var/lib/apt/lists/*

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

# Dependencies first, in their own layer, so code edits do not reinstall them.
COPY requirements.txt .
RUN --mount=type=cache,target=/root/.cache/pip \
    pip install -r requirements.txt

# ---------- stage 2: runtime ----------
FROM python:3.11-slim AS runtime

RUN apt-get update && apt-get install -y --no-install-recommends \
        ca-certificates curl \
    && rm -rf /var/lib/apt/lists/* \
    && useradd --system --uid 10001 --create-home app

COPY --from=build /opt/venv /opt/venv
ENV PATH="/opt/venv/bin:$PATH" \
    PYTHONUNBUFFERED=1 \
    HF_HOME=/cache/hf

WORKDIR /app
COPY --chown=app:app src/ /app/

USER app
EXPOSE 8000
HEALTHCHECK --interval=30s --timeout=3s --start-period=120s --retries=3 \
    CMD curl -fsS http://localhost:8000/healthz || exit 1

ENTRYPOINT ["python", "-m", "server"]

Two details in there are doing real work. --mount=type=cache keeps the pip download cache on the build host between builds without putting it in a layer, which is the correct replacement for the --no-cache-dir-then-delete dance. And useradd before the COPY --chown means the copy lands with the right ownership in one layer rather than being rewritten by a later chown, which would duplicate every byte it touched.

Choosing a CUDA base image

If the container itself runs kernels on a GPU, you need CUDA userspace libraries inside it. The driver stays on the host — the container never ships a driver — and is exposed to the container by the NVIDIA container runtime. That split is the thing to internalise: host driver, container runtime libraries.

NVIDIA publishes its base images in variants whose names describe how much they contain, conventionally base, runtime and devel, tagged with a CUDA version and an OS. The rule is simple: build against devel if you compile CUDA code, ship on runtime, and never ship devel. Check the current tag list for the exact version strings rather than copying a tag from an article — they are versioned aggressively and an old tag can quietly disappear.

FROM nvidia/cuda:<version>-devel-<os> AS build
# ... compile extensions here ...

FROM nvidia/cuda:<version>-runtime-<os> AS runtime
COPY --from=build /opt/venv /opt/venv

The duplication trap: if you use a CUDA base image and install the default PyTorch wheel, you get CUDA twice — once in the base layers and once as pip-installed nvidia-* packages. Pick one. Either use a slim base and let the wheels bring CUDA, or use the CUDA base and install a framework build that links against the system libraries. Both work; having both is gigabytes of nothing.

Version compatibility is the other constraint, and it runs upward: the host driver must be new enough for the CUDA version in the image. There is forward compatibility in some configurations but it is not something to assume. If you are choosing hardware and versions rather than inheriting them, the hardware cluster covers what the numbers on a card mean.

Weights do not belong in the image

It is tempting to bake the model into the image so that the container is self-contained. Resist it, for four reasons that compound.

• Rebuild cost. Every code change rebuilds and re-pushes an image containing tens of gigabytes of weights that did not change.
• Pull time is startup time. A node that must pull forty gigabytes before the first request cannot be part of a fast scale-up. See GPU autoscaling for why that number ends up in your scaling headroom calculation.
• Registry storage multiplies. Keep ten releases and you keep ten copies unless the layers are byte-identical.
• Licence and access control. Weights inside an image inherit the image’s distribution, which is usually broader than the weights’ licence intends.

Put weights in object storage, reference them by an immutable digest, and fetch them into a cache volume at startup. Model weights in CI/CD is entirely about that pipeline, including how to have them already on the node before the pod that needs them is scheduled.

Layer order, and why your rebuilds are slow

The builder reuses a cached layer only if that instruction and every instruction before it are unchanged. So the order of a Dockerfile is a statement about what changes most often, and the correct order is least-frequently-changed first.

1. System packages — changes monthly.
2. Dependency manifest only (requirements.txt, pyproject.toml and the lock file) — changes weekly.
3. Dependency install — invalidated only by step two.
4. Application source — changes hourly.

Copying the whole source tree before installing dependencies is the most common mistake in this file, and it means every one-character edit reinstalls PyTorch. The second most common is a .dockerignore that does not exist, which sends your .git directory, your virtualenv and any local model cache to the daemon as build context — slow, and occasionally a secret leak.

# .dockerignore
.git
.venv
__pycache__/
*.pyc
.env
.pytest_cache/
models/
data/
node_modules/
**/*.safetensors
**/*.gguf

The runtime half: user, signals, health

A correct image is not only a small one. Three runtime properties matter more than a few hundred megabytes.

Run as a non-root user. The USER line above does that. It is the cheapest security control in the file and it catches an entire class of container-escape preconditions.

Take signals seriously. Use the exec form of ENTRYPOINT — the JSON array, not a shell string — so your process is PID 1 and receives SIGTERM directly. A shell-form entrypoint wraps the process in /bin/sh -c, which does not forward signals, so your graceful shutdown never runs and the orchestrator kills the container after the grace period with in-flight streams still open. This matters far more for a streaming AI endpoint than for a normal web service, because a request that has been generating for forty seconds is expensive to lose.

Distinguish the health of the process from the readiness of the model. A container that has started is not a container that can serve — the weights may still be loading. The start-period in the HEALTHCHECK above exists for exactly that gap, and its Kubernetes equivalent is a startup probe, covered in Kubernetes for model serving.

Measuring what you built

Two commands answer the question. The first gives the total; the second gives the breakdown by instruction, largest offender first.

docker image ls my-service:latest
docker history --no-trunc --format 'table {{.Size}}\t{{.CreatedBy}}' my-service:latest

# Or, to see what is actually on the filesystem inside:
docker run --rm --entrypoint sh my-service:latest -c \
  'du -sh /opt/venv /usr/local/lib /usr/lib 2>/dev/null | sort -h'

Read the docker history output as the arithmetic it is. If the largest row is a pip install, look at whether CUDA wheels are in there and whether you meant them to be. If it is a COPY, something large is in your build context that a .dockerignore should have excluded. If it is the base image, you have already done the work that matters and the remaining wins are small.

Set your target after that reading, not before. “Under a gigabyte” is achievable and normal for a service that calls a remote model; a container that runs a seven-billion-parameter model on a GPU will not be small, and chasing it there is time better spent making sure the weights are cached on the node.