Build optimized Docker images for AI applications

Introduction

When you wrap a FastAPI inference service in a naive Dockerfile — full Ubuntu base, PyTorch with CUDA, transformers, and a COPY . . at the end — the image balloons past 3GB, takes four minutes to pull on a cold GKE node, and ships hundreds of packages your security team will flag in the next audit. Teams that accept that default pay for it on every deploy: slow rollouts, expensive Artifact Registry storage, and a runtime attack surface stuffed with shells, package managers, and debuggers an attacker can pivot through. By the end of this lesson you'll be able to take any FastAPI inference repo and produce a sub-150MB image with no shell, no package manager, and a working /health healthcheck — verified by a single validation script that acts as a CI gate.

Key Terminology

• Multi-stage build — a Dockerfile with multiple FROM lines where only the final stage ships; lets you compile in a fat builder and copy artifacts into a minimal runtime so the production image never carries gcc or pip.
• Distroless image — a Google-maintained base (e.g. gcr.io/distroless/python3-debian12) that contains only the language runtime and libc — no shell, no package manager — so an attacker who lands inside the container has no tools to escalate.
• Layer caching — Docker's rule that an instruction's layer is reused when the instruction and everything above it are byte-identical to a prior build; the ordering of COPY and RUN decides whether dependency installs are skipped on incremental builds.
• .dockerignore — the file that filters which paths enter the build context; without it, gigabyte-scale model weights and .git history get streamed to the Docker daemon on every build even when no COPY references them.
• Healthcheck — a HEALTHCHECK instruction Kubernetes and orchestrators use to decide whether a container is ready; on distroless you must invoke it through the Python interpreter because curl and wget do not exist.

Concepts

Multi-stage builds separate build-time tools from runtime

The fundamental principle is that the builder stage can be as fat and messy as needed — install gcc, Python development headers, compile native extensions, download tokenizer files. None of that ships, because the builder stage is discarded after producing the artifacts you need. The runtime stage starts from a minimal base and copies in only the virtual environment and application code. The final image contains nothing unnecessary: no shell to exec into, no apt/pip to install attacker tools, no compiler to build them. See Code Walkthrough for the concrete two-stage Dockerfile.

Distroless eliminates the runtime shell

Choosing gcr.io/distroless/python3-debian12 as the runtime base removes /bin/sh, apt, and the standard debugging utilities. That choice is what makes docker run --rm --entrypoint /bin/sh <image> fail — one of the exit-criterion gates. The trade-off is that anything you used to do "by shelling into the container" — including healthchecks — must now run through the Python interpreter directly.

Layer ordering decides incremental build time

Every Dockerfile instruction creates a layer. When Docker detects an instruction and all instructions above it are unchanged, it reuses the cached layer. For AI services where pip install of PyTorch + transformers takes five to eight minutes, the ordering rule is: least frequently changing first. Copy requirements.txt and run pip install before copying application source. Then a code-only commit reuses the dependency layer and the incremental build drops from eight minutes to under thirty seconds.

.dockerignore keeps gigabytes out of the build context

AI repositories accumulate model checkpoints, training data, Jupyter notebooks, .env files with API keys, and .git histories that often exceed 100MB on their own. Without a .dockerignore, Docker streams every byte of the working directory to the daemon at the start of each build — even paths no COPY references. A precise exclude list (.git, __pycache__, tests/, models/, data/, *.ipynb, .env*) cuts the build context to kilobytes and prevents secrets in .env from accidentally landing in an image layer.

Healthchecks must use Python on distroless

Because distroless has no curl, the HEALTHCHECK command invokes python -c "import urllib.request; urllib.request.urlopen('http://localhost:8080/health')". This satisfies the exit-criterion gate that requires a non-empty healthcheck pointing at /health and gives Kubernetes the signal it needs to decide pod readiness.

Code Walkthrough

The Dockerfile below combines all five concepts above — multi-stage split, distroless runtime, cache-aware layer order, Python-based healthcheck, and the PYTHONDONTWRITEBYTECODE / PYTHONUNBUFFERED env vars the validation script will check for. The companion validate_image.py script then acts as the CI gate.

Code snippetdockerfile
1FROM python:3.11-slim AS builder
2WORKDIR /build
3
4COPY requirements.txt .
5RUN python -m venv /opt/venv && \
6    /opt/venv/bin/pip install --no-cache-dir --upgrade pip && \
7    /opt/venv/bin/pip install --no-cache-dir -r requirements.txt
8
9COPY src/ ./src/
10
11FROM gcr.io/distroless/python3-debian12
12
13COPY --from=builder /opt/venv /opt/venv
14COPY --from=builder /build/src /app/src
15
16ENV PYTHONDONTWRITEBYTECODE=1 \
17    PYTHONUNBUFFERED=1 \
18    PATH="/opt/venv/bin:$PATH"
19
20WORKDIR /app
21EXPOSE 8080
22
23HEALTHCHECK --interval=30s --timeout=5s --retries=3 \
24    CMD ["/opt/venv/bin/python", "-c", "import urllib.request; urllib.request.urlopen('http://localhost:8080/health')"]
25
26ENTRYPOINT ["/opt/venv/bin/python", "-m", "uvicorn", "src.main:app", "--host", "0.0.0.0", "--port", "8080"]

The builder uses python:3.11-slim (pip + build tools, no docs) and installs deps into /opt/venv with --no-cache-dir so wheel files don't bloat the stage. Crucially, COPY requirements.txt happens before COPY src/ — so when only application code changes, Docker reuses the cached pip-install layer. The runtime FROM switches to distroless: no shell, no apt. The virtual environment and source are copied across with --from=builder. The env block sets PYTHONDONTWRITEBYTECODE=1 (no .pyc files cluttering the immutable container) and PYTHONUNBUFFERED=1 (logs flush immediately so Kubernetes captures them). The HEALTHCHECK runs Python directly because no curl exists. The exec-form ENTRYPOINT ensures Uvicorn receives SIGTERM from Kubernetes and shuts down gracefully.

The validation script below gates the image before it reaches Artifact Registry:

Code snippetpython
1import subprocess, json, sys
2
3MAX_IMAGE_SIZE_MB = 150
4REQUIRED_ENV_VARS = ["PYTHONDONTWRITEBYTECODE", "PYTHONUNBUFFERED"]
5
6if len(sys.argv) != 2:
7    print("Usage: python validate_image.py <image_name>")
8    sys.exit(1)
9
10image_name = sys.argv[1]
11inspect = subprocess.run(
12    ["docker", "inspect", image_name],
13    capture_output=True, text=True,
14)
15if inspect.returncode != 0:
16    print(f"  ERROR: Image not found: {image_name}")
17    sys.exit(1)
18
19data = json.loads(inspect.stdout)[0]
20size_mb = data.get("Size", 0) / (1024 * 1024)
21config = data.get("Config", {})
22env_vars = {}
23for entry in config.get("Env", []):
24    key, _, value = entry.partition("=")
25    env_vars[key] = value
26has_healthcheck = "Healthcheck" in config and config["Healthcheck"] is not None
27
28errors = []
29if size_mb > MAX_IMAGE_SIZE_MB:
30    errors.append(f"Image size {size_mb:.1f}MB exceeds limit {MAX_IMAGE_SIZE_MB}MB")
31for var in REQUIRED_ENV_VARS:
32    if var not in env_vars:
33        errors.append(f"Missing required env var: {var}")
34if not has_healthcheck:
35    errors.append("No HEALTHCHECK instruction found")
36
37print(f"Image: {image_name}")
38print(f"Size: {size_mb:.1f} MB")
39print(f"Healthcheck: {'present' if has_healthcheck else 'missing'}")
40print(f"Valid: {len(errors) == 0}")
41for error in errors:
42    print(f"  ERROR: {error}")
43sys.exit(0 if not errors else 1)

The script parses docker inspect output, checks size against the 150MB cap, verifies both required env vars are present, and confirms a HEALTHCHECK is configured. Exit code 0 means the image passes the gate.

You'll know it works when docker build -t myimage:sha-abc123 . followed by python validate_image.py myimage:sha-abc123 prints Valid: True and exits 0, and docker run --rm --entrypoint /bin/sh myimage:sha-abc123 -c 'echo ok' fails because no shell exists in the runtime layer.

Do's and Don'ts

Do's

1. ✓Do pin the runtime base image by digest — replace gcr.io/distroless/python3-debian12 with the SHA256 digest so a base-image republish never silently changes the bytes you ship.
2. ✓Do wire validate_image.py into CI before the push step — a non-zero exit must block the push to Artifact Registry so an oversized or misconfigured image never lands where Kubernetes can pull it.
3. ✓Do track image size as a build-dashboard metric — sudden 50MB jumps usually mean a heavyweight dependency was added to requirements.txt without anyone noticing the cost.

Don'ts

1. ✗Don't ship tests, docs, or .env files inside the image — keep .dockerignore precise; test fixtures and dev credentials in a production image are pure attack surface.
2. ✗Don't use curl or wget in the healthcheck — distroless does not have them; the healthcheck must invoke Python's urllib through the venv interpreter.
3. ✗Don't copy application source before pip install — that ordering invalidates the dependency layer on every code change and turns thirty-second incremental builds back into eight-minute ones.