Free lesson · AI Engineering
Containerising AI Applications
Containerising AI Applications
The Deploy That Worked on My Machine and Nowhere Else
At Meta, the first version of our AI documentation assistant ran on a senior engineer's workstation. Python 3.11, a specific version of PyTorch, a specific version of the sentence-transformers library, a local Redis instance, and a CUDA 12.1 driver that the engineer had compiled from source six months earlier. The system worked flawlessly on that machine. It worked on no other machine in the building.
When we tried to deploy it to the internal cloud infrastructure, nothing worked. The Python version on the deployment target was 3.9. The PyTorch version was incompatible with the CUDA driver on the cloud GPU instances. The sentence-transformers library had a transitive dependency on a tokeniser library that required a system-level package not installed on the deployment image. We spent four days resolving dependency conflicts - four days that would have been four minutes if the application had been containerised from the start.
AI applications have the worst dependency stories in software engineering. A typical AI application depends on Python, a specific version of a deep learning framework (PyTorch, TensorFlow, JAX), a CUDA toolkit version that must match the GPU driver, a set of ML libraries with deeply nested transitive dependencies, and system-level packages for tokenisation, image processing, or numerical computation. Any mismatch anywhere in that stack produces errors that range from cryptic to silent. Containerisation does not make these dependencies simpler - but it makes them reproducible, which is the difference between "works on my machine" and "works everywhere."
Why AI Applications Are Harder to Containerise
Standard web applications are relatively simple to containerise: a runtime (Node.js, Go, Java), application code, and maybe a few system packages. The resulting image is typically 100-500MB. Build times are measured in seconds.
AI applications break this model in three ways.
Challenge 1: Large Base Images
AI frameworks are large. PyTorch alone is 800MB-2GB depending on the variant (CPU-only vs CUDA). Add sentence-transformers (500MB+ with model files), numpy, scipy, pandas, and the typical ML stack, and a base image easily reaches 5-8GB. Add a local model file (Llama 2 7B is 14GB, even a small embedding model is 400MB), and the image can exceed 20GB.
Large images create operational problems: longer pull times (minutes instead of seconds on container startup), more storage consumption in the registry, slower CI/CD pipelines, and longer rollback times during incidents because pulling the previous image takes minutes.
Challenge 2: GPU Dependencies
AI applications that use GPU acceleration require a specific CUDA toolkit version, which requires a specific GPU driver version on the host, which requires a specific kernel module. The CUDA toolkit, the cuDNN library, and the GPU driver form a compatibility matrix where only specific combinations work. Installing the wrong combination produces either a hard crash or - worse - silent numerical errors where the GPU computes incorrect results without raising any error.
Docker solves part of this: the NVIDIA Container Toolkit allows containers to access the host's GPU driver while bundling their own CUDA toolkit. But you still need to ensure that the CUDA version in your container is compatible with the driver version on the host. The NVIDIA base images (nvidia/cuda:12.1.0-runtime-ubuntu22.04) provide tested combinations.
Challenge 3: Reproducibility of ML Dependencies
ML libraries have complex dependency graphs with version constraints that conflict frequently. PyTorch 2.1 requires numpy ≥1.21 and <2.0. A text processing library requires numpy ≥1.24. An evaluation library pins numpy to 1.23.5 exactly. Resolving these constraints is a combinatorial problem that pip and conda handle differently, sometimes producing different installed versions.
The reproducibility requirement is stricter for AI applications than for web applications. In a web application, a minor version difference in a JSON parsing library is unlikely to change behaviour. In an AI application, a minor version difference in a numerical computation library can change floating-point results, which changes model outputs, which changes evaluation metrics. The container must pin exact versions of every dependency to guarantee reproducible behaviour.
graph TD
subgraph Standard_App["Standard Web Application"]
SA1["Runtime\n(Node.js 20, 200MB)"]
SA2["Application code\n(50MB)"]
SA3["NPM dependencies\n(200MB)"]
SA4["Total image: ~500MB\nBuild time: 30s"]
end
subgraph AI_App["AI Application"]
AA1["Python 3.11\n(150MB)"]
AA2["PyTorch + CUDA\n(2GB)"]
AA3["ML libraries\n(1GB)"]
AA4["Model files\n(400MB-14GB)"]
AA5["Application code\n(50MB)"]
AA6["Total image: 4-18GB\nBuild time: 5-15 min"]
end
style Standard_App fill:#166534,stroke:#86efac,stroke-width:2px
style AI_App fill:#991b1b,stroke:#fca5a5,stroke-width:2px
This diagram contrasts containerisation of a standard web application versus an AI application. A standard web application contains a runtime like Node.js at 200 megabytes, application code at 50 megabytes, and NPM dependencies at 200 megabytes, totalling approximately 500 megabytes with a 30-second build time. An AI application contains Python at 150 megabytes, PyTorch plus CUDA at 2 gigabytes, ML libraries at 1 gigabyte, model files ranging from 400 megabytes to 14 gigabytes, and application code at 50 megabytes, totalling 4 to 18 gigabytes with a 5 to 15 minute build time.
The Production Dockerfile for AI Applications
A well-structured Dockerfile for AI applications uses multi-stage builds to separate the dependency installation from the runtime image, reducing final image size by 30-50%.
# Stage 1: Build stage - install dependencies
# Use the NVIDIA CUDA base for GPU support
FROM nvidia/cuda:12.1.0-devel-ubuntu22.04 AS builder
# Avoid interactive prompts during package installation
ENV DEBIAN_FRONTEND=noninteractive
# Install Python and build dependencies
RUN apt-get update && apt-get install -y --no-install-recommends \
python3.11 \
python3.11-venv \
python3.11-dev \
python3-pip \
build-essential \
&& rm -rf /var/lib/apt/lists/*
# Create virtual environment
RUN python3.11 -m venv /opt/venv
ENV PATH="/opt/venv/bin:$PATH"
# Install Python dependencies (cached unless requirements change)
COPY requirements.txt .
RUN pip install --no-cache-dir -r requirements.txt
# Stage 2: Runtime stage - minimal image
FROM nvidia/cuda:12.1.0-runtime-ubuntu22.04
# Install only runtime Python (no dev headers, no build tools)
RUN apt-get update && apt-get install -y --no-install-recommends \
python3.11 \
python3.11-venv \
&& rm -rf /var/lib/apt/lists/*
# Copy virtual environment from builder
COPY --from=builder /opt/venv /opt/venv
ENV PATH="/opt/venv/bin:$PATH"
# Copy application code
WORKDIR /app
COPY src/ ./src/
COPY config/ ./config/
# Health check endpoint
HEALTHCHECK --interval=30s --timeout=10s --retries=3 \
CMD python3.11 -c "import requests; requests.get('http://localhost:8000/health')"
# Run the application
EXPOSE 8000
CMD ["python3.11", "-m", "uvicorn", "src.main:app", "--host", "0.0.0.0", "--port", "8000"]
Key patterns in this Dockerfile:
Multi-stage build. The builder stage includes compilation tools (gcc, python-dev, build-essential) needed to compile C extensions in ML libraries. The runtime stage excludes them. This reduces the final image by 1-2GB because build tools are not needed at runtime.
CUDA devel vs runtime. The builder uses cuda:12.1.0-devel (includes CUDA compiler, headers, libraries - 4GB). The runtime uses cuda:12.1.0-runtime (includes only runtime libraries - 1.5GB). The compiler is needed to build PyTorch extensions during pip install but not needed at runtime.
Requirements caching. The COPY requirements.txt . and RUN pip install layers are before COPY src/. Docker layer caching means that if only application code changes (not dependencies), the pip install layer is cached and the rebuild takes seconds instead of minutes.
Virtual environment isolation. Using a venv inside the container avoids polluting the system Python and makes the dependency set self-contained and copyable between stages.
Image Size Optimisation
At production scale, image size directly affects deployment speed, storage costs, and incident response time. Every gigabyte matters.
Strategy 1: Use CPU-only PyTorch where GPU is not needed. The CUDA-enabled PyTorch is 2GB. The CPU-only version is 200MB. If your application uses the GPU only for inference and you have a separate GPU-enabled serving container, the application orchestration container can use CPU-only PyTorch. This alone can reduce image size by 40%.
Strategy 2: Separate model files from application images. Do not bake model files into the Docker image. Store them in a model registry (S3, GCS, HuggingFace Hub) and download them at container startup or mount them as a volume. This means the application image stays small (2-4GB), model updates do not require rebuilding the application image, and different deployments can use different model versions by changing a configuration value rather than rebuilding.
Strategy 3: Use .dockerignore aggressively. Exclude test files, notebooks, documentation, git history, virtual environments, and cached model files from the build context. A .dockerignore that excludes __pycache__, .git, *.ipynb, tests/, and models/ can reduce build context transfer time from minutes to seconds.
Strategy 4: Pin exact dependency versions. Use pip freeze > requirements.txt after installing your dependencies in a clean environment. Every dependency, including transitive ones, is pinned to an exact version. This ensures reproducibility and prevents pip from re-resolving dependencies on every build, which can change installed versions and break builds unexpectedly. At Meta, a single unpinned transitive dependency caused a deployment failure that took six hours to diagnose - the resolution was a one-line version pin.
# requirements.txt - pinned for reproducibility
torch==2.1.2+cu121
torchvision==0.16.2+cu121
transformers==4.36.2
sentence-transformers==2.2.2
langchain==0.1.4
openai==1.12.0
anthropic==0.18.1
fastapi==0.109.0
uvicorn==0.27.0
redis==5.0.1
numpy==1.26.3
pydantic==2.5.3
httpx==0.26.0
python-dotenv==1.0.0
structlog==24.1.0
prometheus-client==0.19.0
Model Loading and Startup Patterns
AI containers have a unique startup challenge: model loading. Loading a large model into GPU memory can take 30-90 seconds. During this time, the container is running but not ready to serve traffic. If the orchestrator (Kubernetes) routes traffic to the container before the model is loaded, every request fails.
The readiness probe pattern:
Kubernetes distinguishes between liveness probes (is the container alive?) and readiness probes (is the container ready to serve traffic?). For AI containers, the liveness probe checks that the HTTP server is responding. The readiness probe checks that the model is loaded and the first inference has completed successfully.
# Health and readiness endpoints for AI containers
from fastapi import FastAPI
from contextlib import asynccontextmanager
# Global state for model readiness
model_ready = False
model = None
@asynccontextmanager
async def lifespan(app: FastAPI):
"""Load model during startup, before accepting traffic."""
global model, model_ready
# This runs before the server starts accepting requests
print("Loading model...")
model = load_model(config.model_path)
# Run a warm-up inference to ensure GPU is initialised
warmup_input = "This is a warm-up query for model initialisation."
_ = model.generate(warmup_input)
model_ready = True
print(f"Model loaded and warmed up. Ready to serve.")
yield # Server runs here
# Cleanup on shutdown
del model
app = FastAPI(lifespan=lifespan)
@app.get("/health")
async def health():
"""Liveness probe - is the process alive?"""
return {"status": "alive"}
@app.get("/ready")
async def ready():
"""Readiness probe - is the model loaded and serving?"""
if not model_ready:
return JSONResponse(
status_code=503,
content={"status": "not_ready", "reason": "model_loading"}
)
return {"status": "ready", "model": config.model_name}
Kubernetes deployment configuration:
# kubernetes deployment spec (YAML representation in Python for readability)
deployment_spec = {
"containers": [{
"name": "ai-service",
"image": "registry.internal/ai-service:v2.1.0",
"resources": {
"limits": {
"nvidia.com/gpu": "1", # Request 1 GPU
"memory": "16Gi",
"cpu": "4"
}
},
"livenessProbe": {
"httpGet": {"path": "/health", "port": 8000},
"initialDelaySeconds": 10, # Wait 10s before first check
"periodSeconds": 30
},
"readinessProbe": {
"httpGet": {"path": "/ready", "port": 8000},
"initialDelaySeconds": 60, # Model needs ~60s to load
"periodSeconds": 10,
"failureThreshold": 6 # Tolerate 6 failures (60s window)
}
}]
}
sequenceDiagram
participant K8s as Kubernetes
participant Container as AI Container
participant GPU as GPU Memory
participant LB as Load Balancer
K8s->>Container: Start container
Container->>Container: HTTP server starts
K8s->>Container: Liveness probe (/health)
Container-->>K8s: 200 OK (alive)
Note over Container,GPU: Model loading phase (30-90s)
Container->>GPU: Load model weights
GPU-->>Container: Model loaded
Container->>GPU: Warm-up inference
GPU-->>Container: Warm-up complete
K8s->>Container: Readiness probe (/ready)
Container-->>K8s: 200 OK (ready)
K8s->>LB: Add container to service endpoints
LB->>Container: Route production traffic
This sequence diagram shows the startup lifecycle of an AI container in Kubernetes. Kubernetes starts the container and the HTTP server begins. The liveness probe to the health endpoint returns 200 OK, confirming the process is alive. During the model loading phase of 30 to 90 seconds, the container loads model weights into GPU memory and runs a warm-up inference. Only after warm-up completes does the readiness probe to the ready endpoint return 200 OK. Kubernetes then adds the container to the service endpoints, and the load balancer begins routing production traffic to it.
Interview: System design interview connection: "How would you deploy this AI system?" often starts with containerisation. The interviewer wants to see that you understand the unique challenges - large images, GPU dependencies, model loading times - and have practical solutions.
What to emphasise: Start with why AI containers are different (large images, CUDA compatibility matrix, model loading latency). Describe multi-stage builds for size reduction. Explain the separation of model files from application images. Cover the readiness probe pattern for model loading. These are the details that demonstrate production experience.
What separates strong from weak answers: Weak answers describe a Dockerfile that installs PyTorch and copies code. Strong answers describe multi-stage builds with CUDA devel vs runtime separation, image size optimisation strategies, model loading with readiness probes, and dependency pinning for reproducibility. The strongest answers mention specific numbers - image sizes, model loading times, GPU memory requirements.
Exercise: Starter: Your AI application uses PyTorch for inference, the sentence-transformers library for embeddings, and FastAPI for the HTTP server. Write a Dockerfile that uses multi-stage builds. Estimate the final image size. How would you reduce it by 40%?
Intermediate: Your container image is 8.2GB and takes 4 minutes to pull on cold start. The product team wants cold start time under 90 seconds. Design the optimisation plan: which components can be moved out of the image? What is the trade-off between image size and startup complexity? How do you handle model file downloads at startup without blocking the readiness probe?
Advanced: Design the complete container architecture for an AI application that serves three different models: a small embedding model (400MB), a medium classification model (2GB), and a large generative model (14GB). The application processes 10,000 requests per day with a P95 latency SLA of 3 seconds. Should you run all three models in one container or separate containers? Design the Kubernetes deployment including resource requests, GPU allocation, scaling strategy, and the health/readiness probe configuration.
Key Takeaways
AI applications are harder to containerise than traditional web applications because of large framework dependencies (PyTorch + CUDA: 2-4GB alone), GPU compatibility requirements (CUDA toolkit version must match host GPU driver), and model files that can reach 14GB or more. These challenges make image size, build time, and startup latency the primary containerisation concerns.
Multi-stage builds are essential: use the CUDA devel image for building (includes compiler and headers) and the CUDA runtime image for deployment (runtime libraries only). This single change reduces image size by 1-2GB. Separating model files from application images - storing models in S3/GCS and downloading at startup or mounting as volumes - further reduces image size and decouples model updates from application deployments.
The readiness probe pattern is critical for AI containers. Model loading takes 30-90 seconds - during which the container is alive but not ready to serve traffic. Kubernetes readiness probes prevent traffic from being routed to containers that have not finished loading their models. The warm-up inference after model loading ensures that the GPU is fully initialised before the first real request arrives.
Dependency pinning at exact versions is non-negotiable for AI applications. Minor version differences in numerical computation libraries can change model outputs, evaluation metrics, and production behaviour. A requirements.txt with pinned versions is the foundation of reproducible AI deployments.