Containerizing LLMs: Best Practices for Docker-Based AI Workloads

The first time I containerized a fine-tuned Llama model for a client's internal search tool, the build finished at 38 gigabytes. I remember staring at the terminal thinking there was no way that was right. It was right. The image included a CUDA base, PyTorch with every backend compiled in, model weights baked directly into the layer, and a pip cache that had not been cleaned. 

Pushing that to our registry took eleven minutes on a good connection. Pulling it onto a fresh node during an autoscale event took even longer, and by the time the pod was ready, the traffic spike it was supposed to handle had already passed. That's the moment I stopped treating LLM containers like regular application containers, because they are not the same animal at all.

Why This Problem Actually Matters

Most Docker advice out there is written for stateless web services, small images, fast cold starts, and horizontal scaling on demand. LLM workloads break almost every assumption baked into that advice. The artifact is huge, the runtime is GPU-bound, startup involves loading gigabytes into VRAM, and half your "application code" is actually a C++/CUDA binary blob you didn't write and can't easily trim. If you treat an inference container like a Flask app with a bigger base image, you end up with slow deploys, wasted GPU spend, and autoscaling that technically works but arrives too late to matter.

The First Wrong Turn: One Image to Rule Them All

Our early approach was a single monolithic image model with weights, tokenizer, inference server, and dependencies all baked together, rebuilt on every model version bump. It felt simple. It wasn't. Every retrain meant rebuilding a 30+ GB image even when the code hadn't changed a single line. Registry storage costs gradually increased until someone in finance questioned why our container registry bill resembled that of a second AWS account. Worse, rollbacks were painful because reverting to a previous model meant pulling an entire previous image rather than swapping a much smaller artifact.

The solution that actually worked was separating the model weights from the serving image entirely. The image contains the runtime, the inference server (we used vLLM for most of our transformer workloads), and pinned dependencies. Weights live in object storage and are pulled at container start via an init container or a lazy loading entry point. The approach felt counterintuitive at first. Are we effectively transitioning the slower process to startup instead of build time? — but it turned out to be the right trade. Startup pulls are parallelizable, cacheable on the node, and don't bloat the registry. Build time dropped from twenty-plus minutes to under four.

A Smaller Base Image Than You'd Expect

This is where the challenges began. Everyone defaults to using nvidia/cuda:*-devel images because the framework documentation recommends them, but these devel images include the entire CUDA toolkit, which contains compilers that you will never use at runtime. Switching to the runtime variant and only installing the exact CUDA and cuDNN versions your framework's wheel actually needs cuts roughly 4GB off the base alone. A minimal multi-stage build looks something like this:

FROM nvidia/cuda:12.1.0-devel-ubuntu22.04 AS builder
RUN pip install --no-cache-dir vllm==0.4.2

FROM nvidia/cuda:12.1.0-runtime-ubuntu22.04
COPY --from=builder /usr/local/lib/python3.10 /usr/local/lib/python3.10
COPY --from=builder /usr/local/bin/python3.10 /usr/local/bin/
ENV MODEL_PATH=/mnt/models
ENTRYPOINT ["python3", "-m", "vllm.entrypoints.api_server"]

The build stage compiles anything that needs the full toolkit; the runtime stage only carries what's needed to execute. It's a basic Docker pattern, but I've seen it skipped constantly on ML teams because the assumption is always, "the model is the heavy part; the image doesn't matter." The model is heavy, sure, but a bloated base image adds real minutes to every autoscale event, and in production that's the difference between absorbing a traffic spike and dropping requests.

The OOM Kill: Nobody Explained Well

This is the war story I bring up most often. We had a container that ran fine locally and in staging, then got silently killed in production under load — no crash log, no stack trace, just a pod restart and a confused on-call engineer at 2 AM. It turned out to be the kernel OOM killer, not an application-level exception, because our memory limit accounted for the model weights in VRAM but excluded the growing KV cache for long-context requests plus the CPU-side tokenizer buffers. GPU memory and container memory limits are two completely separate accounting systems, and Kubernetes will happily kill your pod over host RAM even if your GPU has headroom to spare.

The fix was unglamorous: we set explicit memory requests and limits with a real margin above peak KV cache usage, moved batch size and max sequence length into environment-configurable values instead of hardcoding them, and added a lightweight health assessment that reported GPU memory utilization alongside the standard liveness probe. None of that is exotic. All of it was missing because we'd copy-pasted a manifest template built for a stateless API and never revisited the resource math for a model that holds state in memory for the duration of a request.

Where I'd Push Back on Common Advice

A lot of guidance recommends one model per container for isolation, and for many teams that's right. But if you're serving several small fine-tunes of the same base model, that pattern wastes GPU memory by duplicating base weights across containers. We transitioned to a multi-adapter setup, where one base model is loaded once, and LoRA adapters are swapped for each request; this approach is more complex operationally but reduces the GPU footprint by nearly half. I wouldn't consider it a default; it represents a level of complexity that is justified only after demonstrating that plain per-model containers are indeed the bottleneck.

I'd also push back on containerizing every workload the same way. Batch inference and real-time serving have almost opposite goals: one wants throughput and tolerates slow cold starts; the other needs rapid readiness and predictable latency. We split these into separate images with separate resource profiles, even though it meant more Dockerfiles. Fewer surprises beat fewer files.

Key Takeaways

• Separate model weights from the serving image; bake them in the runtime and pull weights at startup from object storage.
• Use CUDA runtime images, not devel images, unless you genuinely compile something at container start.
• Account for GPU memory and host memory as two separate budgets; KV cache growth is the usual silent killer.
• Split batch and real-time serving into different images; their optimization goals are conflicting.
• Don't reach for multi-adapter serving or other density tricks until you've measured that plain per-model containers are actually the bottleneck.

Closing Thought

None of this required exotic tooling, no custom orchestrator, and no proprietary platform. It required treating the container as part of the model's runtime behavior rather than a packaging afterthought bolted on after the research work was done. The teams that struggle most with this approach usually aren't lacking Docker knowledge; they're applying web-service intuition to a workload that behaves nothing like a web service. 

If you're mid-migration on something similar, I'd genuinely ask: are you optimizing your image for build convenience or for what actually happens the moment traffic hits a cold node? Those answers are rarely the same, and figuring out which one you've been solving for is usually the first real fix.