Containers & Docker: Dockerfiles, Images, and Model Containerization

🎯 Learning Objectives

• Understand containerization and its benefits for ML deployment
• Write optimized Dockerfiles for ML models
• Manage images with registries and tags
• Implement multi-stage builds for smaller images
• Debug containerized ML applications

📋 Prerequisites

• Linux command line basics: Terminal, file system
• Python/ML basics: pip, dependencies
• MLOps Lifecycle (Week 1): Deployment stage

1. 📖 Core Content

1.1 Intuition: Why Containers for ML?

• "Install Python 3.9, install scikit-learn 1.0, copy my script, run it"
• Problem: "Oh, you have Python 3.10, and sklearn 1.3 changed the API" With Docker:
• "Here's an image that has everything — run it anywhere"
• Python version, libraries, OS dependencies, system gubbins — everything packaged together. For ML models specifically, containers solve:

1. Reproducibility: The exact same environment for training and serving
2. Dependency hell: GPU drivers, CUDA, Python packages, system libraries
3. Scalability: Containers are the unit of scaling in Kubernetes
4. Isolation: Each model runs in its own environment

1.2 Docker Fundamentals

1.3 Writing Dockerfiles for ML

1.3.1 Basic ML Dockerfile

dockerfile
# Start from a base image with Python 3.9
FROM python:3.9-slim
# Set working directory
WORKDIR /app
# Copy requirements first (leverages Docker layer caching)
COPY requirements.txt .
RUN pip install --no-cache-dir -r requirements.txt
# Copy the rest of the application
COPY src/ src/
COPY model/ model/
# Expose the API port
EXPOSE 8080
# Command to run the model server
CMD ["uvicorn", "src.serve:app", "--host", "0.0.0.0", "--port", "8080"]

1.3.2 Multi-Stage Build (Smaller Images)

dockerfile
# Stage 1: Build
FROM python:3.9-slim AS builder
WORKDIR /build
COPY requirements.txt .
RUN pip install --no-cache-dir -r requirements.txt
# Stage 2: Runtime (very small base)
FROM python:3.9-slim AS runtime
WORKDIR /app
# Copy only installed packages from builder
COPY --from=builder /usr/local/lib/python3.9/site-packages /usr/local/lib/python3.9/site-packages
COPY --from=builder /usr/local/bin /usr/local/bin
# Copy application code
COPY src/ src/
COPY model/ model/
EXPOSE 8080
CMD ["uvicorn", "src.serve:app", "--host", "0.0.0.0", "--port", "8080"]

• Without multi-stage: ~1.2 GB (includes build tools, pip cache, compilers)
• With multi-stage: ~350 MB (runtime only)

1.3.3 GPU-Aware Dockerfile

dockerfile
FROM nvidia/cuda:11.8-runtime-ubuntu22.04
# Install Python and pip
RUN apt-get update && apt-get install -y python3 python3-pip
# Install PyTorch with CUDA support
RUN pip3 install torch torchvision --index-url https://download.pytorch.org/whl/cu118
# Install other dependencies
COPY requirements.txt .
RUN pip3 install -r requirements.txt
COPY src/ /app/
WORKDIR /app
CMD ["python3", "serve.py"]

1.4 Docker Best Practices for ML

1. Pin all versions: scikit-learn==1.2.2, not scikit-learn>=1.2
2. Use .dockerignore: Exclude __pycache__/, .git/, data/, .env
3. Layer caching: Order Dockerfile from least to most frequently changing (base OS → system deps → pip packages → code → model weights)
4. Health checks: Add HEALTHCHECK for production:

   dockerfile
   HEALTHCHECK --interval=30s --timeout=3s \
     CMD curl -f http://localhost:8080/health || exit 1
   

5. Non-root user: Don't run as root in production:

   dockerfile
   RUN useradd -m -u 1000 appuser
   USER appuser
   

6. Minimize layers: Combine RUN commands where possible

1.5 Worked Example: Containerizing a BERT Model

python
Run Inline
# runnable
from fastapi import FastAPI, Request
from transformers import pipeline
import torch
app = FastAPI()
class ModelServer:
    def __init__(self):
        self.classifier = pipeline(
            "sentiment-analysis",
            model="distilbert-base-uncased-finetuned-sst-2-english",
            device=0 if torch.cuda.is_available() else -1
        )
    def predict(self, text: str) -> dict:
        result = self.classifier(text)[0]
        return {"label": result["label"], "score": result["score"]}
model = ModelServer()
@app.post("/predict")
async def predict(request: Request):
    body = await request.json()
    return model.predict(body["text"])
@app.get("/health")
async def health():
    return {"status": "healthy"}

dockerfile
FROM python:3.9-slim
WORKDIR /app
COPY requirements.txt .
RUN pip install --no-cache-dir -r requirements.txt
COPY serve.py .
# Download model weights at build time (prevents download at startup)
RUN python -c "from transformers import pipeline; pipeline('sentiment-analysis', model='distilbert-base-uncased-finetuned-sst-2-english')"
EXPOSE 8000
CMD ["uvicorn", "serve:app", "--host", "0.0.0.0"]

bash
docker build -t sentiment-model:v1 .
docker run -p 8000:8000 sentiment-model:v1

1.6 Edge Cases & Gotchas

• CUDA version mismatch: The CUDA version in the Docker image must match the host's driver. Check with nvidia-smi → CUDA version.
• Model weights in images: Large models (5+ GB) should be mounted as volumes, not baked into images.
• Temp file cleanup: ML models can generate large temp files. Clean up between predictions.
• Timezones: Containers use UTC by default. Set TZ=Asia/Kolkata if needed.

1.7 Why This Matters

• Kubernetes runs containers — no container, no orchestration
• MLflow models can be packaged as Docker containers
• SageMaker, Vertex AI, Azure ML all use containers for model serving
• CI/CD pipelines build and test containers before deployment

2. 📐 Key Formulas / Concepts

3. ⚠️ Common Pitfalls

Pitfall 1: Outdated Base Images

Pitfall 2: Hardcoding Secrets in Dockerfiles

Pitfall 3: Not Handling GPU Memory Leaks

python
Run Inline
with torch.no_grad():
    result = model(inputs)
torch.cuda.empty_cache()

4. 📝 Practice Questions

Q1: Your Docker image is 3.2 GB. The model is 1.5 GB (PyTorch), and the base image is 1.7 GB (CUDA 11.8 full). How can you reduce it?

1. Use nvidia/cuda:11.8-runtime instead of nvidia/cuda:11.8-devel (runtime is ~800 MB vs 2.5 GB for devel)
2. Use python:3.9-slim as the base (120 MB) and manually install CUDA runtime
3. Don't bake model weights into the image — download at container start or mount them:

   dockerfile
   VOLUME /app/models
   

4. Use multi-stage build to exclude build tools
5. Clean pip cache: pip install --no-cache-dir
6. If using PyTorch, use the CPU-only version for CPU deployments

Estimated savings: 3.2 GB → 1.5-2.0 GB (model weights) → optionally use volume mounts. Q2: Compare Docker for ML with virtual environments (conda, venv). When would you use each?

Use Docker for production deployment, CI/CD, and reproducibility-critical scenarios. Use virtual environments for local development and quick experiments. Q3: Design a Docker strategy for training and serving a model that uses GPU for training but CPU for serving.

dockerfile
# Training Dockerfile (GPU-heavy)
FROM nvidia/cuda:11.8-runtime-ubuntu22.04 AS trainer
RUN pip3 install torch torchvision transformers
COPY train.py .
CMD ["python3", "train.py"]

# Serving Dockerfile (CPU-only, smaller)
FROM python:3.9-slim AS server
RUN pip install --no-cache-dir torch torchvision --index-url https://download.pytorch.org/whl/cpu
COPY serve.py /app/
EXPOSE 8080
CMD ["uvicorn", "serve:app", "--host", "0.0.0.0"]

Strategy:

1. Build trainer image with GPU support
2. Run training, save model to cloud storage (S3/Blob)
3. Build server image with CPU-only PyTorch (much smaller)
4. At startup, download trained model from storage
5. Serve on CPU

This gives the best of both worlds: fast GPU training and small CPU serving images (~500 MB vs ~3 GB).

5. 🔗 Cross-References

• Previous: Feature Stores (Week 4) — Feature computation pipelines
• Next: Orchestration: K8s & Kubeflow (Week 6) — Container orchestration
• Related: Model Serving (Week 8) — Serving containerized models
• External: Docker documentation (docs.docker.com) Join Discord PreviousFeature StoresNextOrchestration: K8s & Kubeflow