Day 2: Advanced Docker for AI — Multi-Stage Builds for Heavy Python Runtimes

Containerising an AI workload is a different problem from containerising a web service. The dependencies are huge, the build tools are heavy, and an unoptimised image can pass 10 GB. Moving that across a registry slows every CI/CD run, and during a Kubernetes autoscale event each new pod must pull it before it can serve a single request.

The naive Dockerfile

This is the Dockerfile most teams start with.

FROM python:3.11
RUN apt-get update && apt-get install -y build-essential gcc
RUN pip install torch transformers fastapi uvicorn
COPY . /app
WORKDIR /app
CMD ["uvicorn", "main:app", "--host", "0.0.0.0", "--port", "8000"]

It works, and it has three production problems.

Diagram: FIG 1 · A single-stage build carries compilers, caches and wheels into production. Layer sizes here are illustrative, not measured. See the animated version.

• Build artifact accumulation. Compilers, intermediate wheels and pip caches are baked into read-only image layers and shipped everywhere.
• Excessive attack surface. Development and compilation tools exist in production, which gives an attacker more to work with after a break-in.
• Deployment latency. Pulling a large image can take minutes, which causes request timeouts while new pods start.

The multi-stage fix

Keep the build dependencies in an ephemeral builder stage. Copy only the compiled output into a minimal runtime stage.

Diagram: FIG 2 · Dependencies are built in stage 1. Only the finished virtual environment is copied into the clean runtime stage; everything else is discarded. See the animated version.

# syntax=docker/dockerfile:1

# ---- Stage 1: builder (heavy, discarded) ----
FROM python:3.11-slim AS builder
WORKDIR /build
RUN apt-get update \
 && apt-get install -y --no-install-recommends build-essential curl \
 && rm -rf /var/lib/apt/lists/*
RUN python -m venv /opt/venv
ENV PATH="/opt/venv/bin:$PATH"
COPY requirements.txt .
RUN --mount=type=cache,target=/root/.cache/pip \
    pip install --no-cache-dir -r requirements.txt

# ---- Stage 2: runtime (small, ships to production) ----
FROM python:3.11-slim
WORKDIR /app
RUN apt-get update \
 && apt-get install -y --no-install-recommends libgomp1 \
 && rm -rf /var/lib/apt/lists/*
COPY --from=builder /opt/venv /opt/venv
ENV PATH="/opt/venv/bin:$PATH"
RUN useradd -m -u 10001 aiuser
USER aiuser
COPY --chown=aiuser:aiuser ./src /app
EXPOSE 8000
CMD ["uvicorn", "main:app", "--host", "0.0.0.0", "--port", "8000"]

Build it with BuildKit enabled, which current Docker versions use by default: docker build -t model-api:1.0 .. A few notes on this file:

• The cache mount (--mount=type=cache) keeps pip's download cache outside the image, so repeat builds in CI do not re-download multi-gigabyte wheels. Because the cache is not in the layer, --no-cache-dir here only avoids writing a second copy.
• libgomp1 is the OpenMP runtime that many numerical libraries need at run time. If your stack needs other shared libraries, install them in the runtime stage too, otherwise the app fails to start with a missing .so error.
• If you need GPU support, use a CUDA-enabled base image for both stages, or install a CUDA build of PyTorch. That is the subject of Day 3.

Image hygiene checklist

Diagram: FIG 3 · Five small habits that keep AI images lean and safe. Do the apt cleanup in the same RUN command, or the deleted files still live in an earlier layer. See the animated version.

1. Add a .dockerignore listing .git, .venv, __pycache__ and raw test data, so the build context stays small and secrets do not leak in.
2. Run as a non-root user (USER aiuser).
3. Clean package caches in the same RUN command that created them. Deleting files in a later layer does not shrink the image, because the earlier layer still holds them.

Diagram: FIG 4 · Smaller images start faster. The sizes and network speed are illustrative inputs to simple arithmetic, so measure your own registry and network. See the animated version.

Image size is not vanity. At 100 MB/s, an 8 GB image takes about 80 seconds to pull and a 2.5 GB one about 25 seconds, and during a scale-up every new node pays that cost. Lean, deterministic images turn research dependencies into a repeatable, auditable artifact that a cluster scheduler can start quickly.

Run it yourself

This is the exact example from the video. Both images are really built and measured with Docker. The 100 MB random file in the builder stage stands in for compilers and wheels, so the size gap is real but exaggerated on purpose.

# syntax=docker/dockerfile:1
# ---- builder: heavy, thrown away ----
FROM alpine:3.20 AS builder
WORKDIR /build
RUN dd if=/dev/urandom of=toolchain.bin bs=1M count=100
RUN echo "model-api v1" > app.txt

# ---- runtime: small, ships ----
FROM alpine:3.20
RUN adduser -D -u 10001 aiuser
USER aiuser
COPY --from=builder /build/app.txt /app/
CMD ["cat", "/app/app.txt"]

Helper file Dockerfile.naive, used by the run:

FROM alpine:3.20
WORKDIR /build
RUN dd if=/dev/urandom of=toolchain.bin bs=1M count=100
RUN echo "model-api v1" > app.txt
CMD ["cat", "/build/app.txt"]

Helper file compare.sh, used by the run:

set -e
docker build -q -f Dockerfile.naive -t d2naive . >/dev/null
docker build -q -t d2lean . >/dev/null
echo "naive image: $(( $(docker save d2naive | wc -c) / 1000000 )) MB"
echo "multi-stage: $(( $(docker save d2lean | wc -c) / 1000000 )) MB"
echo "app says: $(docker run --rm d2lean)"
echo "runs as uid: $(docker run --rm d2lean id -u)"
docker rmi d2naive d2lean >/dev/null

Run it:

bash compare.sh

Output I got while recording the video:

naive image: 108 MB
multi-stage: 3 MB
app says: model-api v1
runs as uid: 10001

Coming Up Next

Day 3: Hardware acceleration in Docker: configuring GPU passthrough and resource limits for constrained hardware.

Docker #Containerization #DevOps #AIEngineering #FastAPI #CloudNative #PlatformEngineering #CyberSecurity

---

Originally published at https://sureshpallapothu.in/blog/ai-int-day-2-advanced-docker, where this post includes animated diagrams.