Containerizing Large Language Models

The Container Dilemma

Running large language models in production presents a unique challenge: model weights can easily exceed 10GB, making traditional Docker workflows impractical. Every docker build becomes a test of patience, and every deployment a gamble with cold start times.

The key insight is treating model weights as external volumes rather than baked-in layers. This separation of concerns mirrors what we do with databases — you wouldn’t embed your data in a container image.

Multi-Stage Builds for ML

The first optimization is a proper multi-stage build. Your final image should contain only the runtime dependencies:

FROM nvidia/cuda:12.2-runtime-ubuntu22.04 AS base

RUN apt-get update && apt-get install -y \
    python3.11 python3-pip \
    && rm -rf /var/lib/apt/lists/*

FROM base AS deps
COPY requirements.txt .
RUN pip install --no-cache-dir -r requirements.txt

FROM deps AS runtime
COPY src/ /app/src/
WORKDIR /app
ENV MODEL_PATH=/models
EXPOSE 8080
CMD ["python3", "src/serve.py"]

Notice we never copy model weights into the image. They’re mounted at runtime via a volume bound to /models.

Layer Caching Strategy

Docker layer caching becomes critical when your dependency installation takes 15+ minutes. The trick is ordering your COPY instructions from least to most frequently changed:

System packages (rarely change)
Python dependencies (change weekly)
Application code (changes every commit)

This ensures that a code change doesn’t invalidate your expensive dependency layers.

GPU Resource Management

When running multiple model containers on a single GPU node, you need fine-grained control over GPU memory allocation:

services:
  llm-inference:
    image: ml-serving:latest
    deploy:
      resources:
        reservations:
          devices:
            - driver: nvidia
              count: 1
              capabilities: [gpu]
    environment:
      - CUDA_VISIBLE_DEVICES=0
      - MAX_GPU_MEMORY=8GiB

Health Checks That Matter

A basic HTTP health check isn’t enough for ML services. You need to verify that the model is actually loaded and capable of inference:

@app.get("/health")
async def health_check():
    try:
        test_output = model.generate("test", max_tokens=1)
        return {"status": "healthy", "model_loaded": True}
    except Exception as e:
        raise HTTPException(status_code=503, detail=str(e))

Cold Start Optimization

The biggest win for cold starts comes from pre-warming. Use an init container or startup probe that loads the model into GPU memory before the service receives traffic:

startupProbe:
  httpGet:
    path: /health
    port: 8080
  initialDelaySeconds: 30
  periodSeconds: 10
  failureThreshold: 30

This gives your model up to 5 minutes to load before Kubernetes considers it failed.

Key Takeaways

Containerizing LLMs effectively requires rethinking traditional Docker patterns. Separate your weights from your code, optimize your layer cache, and invest in proper health checks. The result is a deployment pipeline that’s both fast and reliable — exactly what production ML demands.

Containerizing Large Language Models.