DEV Community

Rajesh Singh
Rajesh Singh

Posted on

Containerizing AI Applications: A Practical Guide

#mlops #docker #devops #machinelearning

Your model works perfectly on your laptop. Then it's time to ship it, and suddenly you're staring at CUDA mismatches, a bloated 8GB image, and a requirements.txt that behaves differently on the server.
Sound familiar?

Containerizing AI apps rhymes with containerizing regular web apps, but a few things make it genuinely different:

  • 🏋️ Weight: Model weights can be hundreds of MB to tens of GB. Bake them into the image the way you would app code, and your CI/CD pipeline will hate you.
  • 🎮 Hardware coupling: GPU models tie your image to a specific CUDA version. Mismatch it against your ML framework and you get the classic "works locally, breaks in prod."
  • ⏱️ Cold start: Loading a model can take seconds to minutes. Do it at the wrong point in the request lifecycle and every user pays that tax.

Here's how to handle all three, with a real working example.

Containerizing AIApps Architecture

The core pattern

1. Multi-stage builds. Separate build-time deps (compilers, headers) from your runtime image:

FROM python:3.11-slim AS builder
WORKDIR /app
COPY requirements.txt .
RUN pip install --no-cache-dir --user -r requirements.txt

FROM python:3.11-slim
WORKDIR /app
COPY --from=builder /root/.local /root/.local
COPY app/ .
ENV PATH=/root/.local/bin:$PATH
CMD ["uvicorn", "main:app", "--host", "0.0.0.0", "--port", "8000"]
Enter fullscreen mode Exit fullscreen mode

Smaller image, smaller attack surface, done.

2. Load the model once, at startup & never per-request. This is the single most common mistake. In FastAPI, that means loading inside lifespan, not inside your route handler:

@asynccontextmanager
async def lifespan(app: FastAPI):
    device = 0 if torch.cuda.is_available() else -1
    ml_models["classifier"] = pipeline(
        "sentiment-analysis", model=MODEL_NAME, device=device
    )
    yield
    ml_models.clear()
Enter fullscreen mode Exit fullscreen mode

Load a transformer model inside the request handler instead, and your 50ms endpoint becomes a multi-second one, on every single call.

3. Don't bake weights into the image. Three options, pick based on your setup:

Pattern Best for
Bake in Small models only
Volume mount Dev / on-prem
Download on startup (S3/GCS/HF Hub) Cloud-native, keeps images slim

4. GPU support just means a CUDA base image + matching torch build:

FROM nvidia/cuda:12.4.0-runtime-ubuntu22.04
Enter fullscreen mode Exit fullscreen mode 
docker run --gpus all -p 8000:8000 my-ai-app
Enter fullscreen mode Exit fullscreen mode

Match the CUDA version in your base image to what your framework was compiled against, mismatches here cause more debugging pain than almost anything else in ML deployment.

5. Give your healthcheck a grace period. Model loading takes real time; without start-period, orchestrators will kill a perfectly healthy container mid-startup:

HEALTHCHECK --interval=30s --timeout=5s --start-period=60s \
    CMD curl -f http://localhost:8000/health || exit 1
Enter fullscreen mode Exit fullscreen mode

When to stop hand-rolling it

For small models, a FastAPI + Transformers pipeline() wrapper (like above) works great. But once you're past a few GB, then Llama-class models and up, need to reach for vLLM, TGI, or NVIDIA Triton instead. They handle batching and GPU memory management far better than anything built from scratch.

Get the Full walkthrough

I put together a complete, runnable example with full Dockerfile (CPU + GPU variants), docker-compose setup, and a FastAPI service serving a real Hugging Face sentiment model end to end:

📖 Full post: Containerizing AI Applications on rajeshscribe.com

🎥 Watch it built step by step:

What's your go-to pattern for serving models in production with hand-rolled FastAPI, vLLM, Triton, something else? Curious what's working for people at scale.

Top comments (0)