What Nobody Tells You About Running AI Models in Docker

It was 2:14 in the morning when the pager went off. Our recommendation model's inference service had started returning 503s under a traffic spike that, frankly, wasn't even that big. Maybe three times the normal load. By the time I'd opened my laptop, the container had been OOM-killed four times in ten minutes, and Kubernetes was cheerfully restarting it into the same wall every ninety seconds. The image was 14GB. Cold start took eighty seconds. Nobody on the team had looked closely at any of that until it started costing us actual money in lost requests.

That night is the reason I now have strong opinions about Docker and AI infrastructure.

Why This Keeps Happening

Containers became the default way to ship machine learning models because they solve a real problem: a model trained with CUDA 11.8, PyTorch 2.1, and a very specific glibc version doesn't reliably run on a colleague's machine, let alone a fleet of GPU nodes spread across three cloud regions. “It works on my machine” isn't a joke in ML infra; it's a recurring incident report. 

Docker gives you a way to freeze that dependency tree and ship it as one artifact, and that part genuinely works. What doesn't get discussed enough is the many issues that Docker does not resolve, along with the subtle ways teams exacerbate problems by forcing AI workloads into a packaging model that was originally designed for stateless web services.

What We Tried First (and Why It Blew Up)

Our first version of the inference image was, in hindsight, a small crime. We started from nvidia/cuda:12.2.0-devel-ubuntu22.04 because someone had seen it in a tutorial, installed the full CUDA toolkit, pip-installed every dependency without pinning, and copied in not just the model weights but three checkpoint versions “just in case.” Fourteen gigabytes. Every deployment pulled the entire image onto a fresh node, and during autoscale events, we experienced over a minute of waiting just for the image to be pulled before the container started loading the model into GPU memory.

The first fix everyone reaches for is a multi-stage build, and yes, it helps — but it's not the silver bullet people pitch it as. Splitting a devel build stage from a runtime stage cut us from 14GB to roughly 6GB:

FROM nvidia/cuda:12.2.0-devel-ubuntu22.04 AS builder
WORKDIR /build
COPY requirements.txt .
RUN pip install --no-cache-dir -r requirements.txt --target=/deps
 
FROM nvidia/cuda:12.2.0-runtime-ubuntu22.04
COPY --from=builder /deps /usr/local/lib/python3.10/site-packages
COPY model/ /app/model/
COPY serve.py /app/
WORKDIR /app
CMD ["python3", "serve.py"]

That's better, but the real lie in that Dockerfile is the COPY model/ line. We baked multi-gigabyte weights into an image layer, causing a full re-push of the weights with every code change, even minor ones, since Docker re-hashes the entire build context. We moved weights to an external volume pulled from object storage at container start, with a checksum cache to skip redundant downloads. That alone cut most of our deployment time. The trade-off is a slightly more complex startup script and a new dependency on storage being reachable at boot, which is its failure mode. There's no free lunch here; you're just choosing which problem is going to page you at 2 am.

The Detour: Skipping Containers Entirely

We also tried, briefly, running everything on bare metal with conda environments and skipping containers altogether, mostly because one engineer was convinced Docker added overhead with no real benefit for GPU workloads. It's not a crazy position; Docker's GPU story is genuinely leaky. <nvidia-container-toolkit exposes the host driver directly to the container, leading to potential mismatches that application-level packaging cannot resolve, so there is no real isolation>. 

However, we reverted that experiment within a sprint because as soon as more than two people interact with the training pipeline, environment drift reoccurs immediately. “Works on my conda env” is just “works on my machine” wearing a hat.

What Actually Held Up in Production

The architecture that worked was less about clever Docker tricks and more about admitting that an inference container and a training container have almost nothing in common and shouldn't share a Dockerfile, a registry strategy, or a deployment pattern.

For serving, we kept images lean, stateless weights externalized, and a health assessment that actually runs a tiny dummy inference instead of just pinging an HTTP port. A server can report itself as “up” while a model failed to load correctly, and that gap has burned us more than once. Locally, a docker-compose GPU reservation block mirrored how production scheduled GPUs, so dev environments stopped lying about resource contention:

services:
  inference:
    image: registry.internal/rec-model:latest
    deploy:
      resources:
        reservations:
          devices:
            - capabilities: [gpu]
              count: 1
    healthcheck:
      test: ["CMD", "python3", "healthcheck.py"]
      interval: 15s
      timeout: 5s
      retries: 3

For training, we accepted slower builds because training jobs run for hours, and a ninety-second image pull is negligible compared to that, even though the images are larger. Spending engineering time shrinking training images was effort we'd burned for no real payoff. A contrarian point I'll happily defend: not every container needs to be small, only the ones sitting in your hot path.

The other decision that mattered more than any Dockerfile tweak was plain layer-caching discipline, putting "pip install before COPY." It sounds obvious when written down, but I've reviewed more than one ML team's Dockerfile that copies the whole repo first “for simplicity” and invalidates every cached layer on a README change.

Key Takeaways

• Externalize model weights from the image; baking them in wrecks your caching and your deploy speed.
• Multi-stage builds help, but they don't resolve a runtime image that's still hauling around a full CUDA devel toolkit.
• GPU isolation via containers is partial; complete driver and toolkit mismatches between host and container remain entirely your problem.
• Training and serving images deserve different optimization priorities; don't apply the same size obsession to both.
• A health check that only confirms the process is running, without verifying that the model has loaded correctly, will mislead you at the most critical moments.

Closing Thought

Docker didn't fail us that night; rather, it was our incorrect assumptions about its free services that let us down. It's easy to treat containers as a solved problem because the tooling is so mature for web services and then be surprised when AI workloads expose every shortcut you took. I still think GPU container isolation lacks a clean solution, and I’m curious if upcoming tools will address it or if we'll continue improving our health-check scripts. What's the worst 2 am lesson your infrastructure has taught you?