Docker Image Layering: Stop Re-downloading Heavy Dependencies (PyTorch, ROCm, CUDA)

docker-layering-guide.md

Docker Image Layering: Stop Re-downloading Heavy Dependencies

Why this matters: Avoid repeatedly downloading PyTorch, ROCm, CUDA, or other massive packages when creating specialized containers.

The Problem with Traditional Approaches

Conda/venv approach:

# Every new project = full reinstall
conda create -n project1 python=3.10
conda install pytorch torchvision torchaudio pytorch-cuda=11.8 -c pytorch -c nvidia
# Downloads 5GB...

conda create -n project2 python=3.10
conda install pytorch torchvision torchaudio pytorch-cuda=11.8 -c pytorch -c nvidia
# Downloads same 5GB AGAIN...

Problem: Each environment downloads the same massive packages independently.

Naive Docker approach:

# Dockerfile for project1
FROM ubuntu:22.04
RUN pip install torch torchvision  # Downloads 5GB

# Dockerfile for project2
FROM ubuntu:22.04
RUN pip install torch torchvision  # Downloads 5GB AGAIN

Problem: Each image re-downloads everything from scratch.

---

The Solution: Docker Image Layering

Docker images are built in layers that are cached and reusable. Each instruction in a Dockerfile creates a new layer.

How Layers Work

FROM ubuntu:22.04           # Layer 1: Base OS (cached globally)
RUN apt-get update          # Layer 2: Package updates
RUN pip install torch       # Layer 3: PyTorch (CACHED!)
RUN pip install whisper     # Layer 4: App-specific

Key insight: Layers are cached and reusable across images!

---

Strategy: Base Image + Derived Images

Step 1: Create a Heavy Base Image (Once)

# Dockerfile.base
FROM ubuntu:22.04

# Install CUDA/ROCm + PyTorch (heavy, but only ONCE)
RUN apt-get update && apt-get install -y \
    python3-pip \
    && rm -rf /var/lib/apt/lists/*

RUN pip install --no-cache-dir \
    torch==2.1.0+cu118 \
    torchvision==0.16.0+cu118 \
    torchaudio==2.1.0+cu118 \
    --index-url https://download.pytorch.org/whl/cu118

# Verify installation
RUN python3 -c "import torch; print(f'PyTorch: {torch.__version__}')"

CMD ["/bin/bash"]

Build once:

docker build -f Dockerfile.base -t pytorch-cuda-base:latest .
# Downloads 5GB of PyTorch... but only this ONE time

Step 2: Create Specialized Images (Fast!)

Now create project-specific images that build on top of the base:

# Dockerfile.whisper
FROM pytorch-cuda-base:latest  # ← Reuses cached layers!

# Only install whisper-specific packages (lightweight)
RUN pip install --no-cache-dir \
    faster-whisper \
    sounddevice \
    scipy

COPY . /app
WORKDIR /app

CMD ["python3", "whisper_app.py"]

# Dockerfile.stable-diffusion
FROM pytorch-cuda-base:latest  # ← Same base, no re-download!

# Only install diffusion-specific packages
RUN pip install --no-cache-dir \
    diffusers \
    transformers \
    accelerate

COPY . /app
WORKDIR /app

CMD ["python3", "sd_app.py"]

Build derived images (fast!):

docker build -f Dockerfile.whisper -t whisper-app:latest .
# No PyTorch download! Uses cached layer from base image

docker build -f Dockerfile.stable-diffusion -t sd-app:latest .
# No PyTorch download! Uses same cached base layer

---

Real-World Example: ROCm + PyTorch + Projects

Base Image: Clean ROCm PyTorch

# Dockerfile.rocm-base
FROM rocm/pytorch:rocm6.2_ubuntu22.04

ENV HSA_OVERRIDE_GFX_VERSION=11.0.1

# Install system dependencies
RUN apt-get update && apt-get install -y \
    ffmpeg \
    git \
    && rm -rf /var/lib/apt/lists/*

# Verify GPU
RUN python3 -c "import torch; print(f'ROCm: {torch.version.hip}')"

CMD ["/bin/bash"]

Build once:

docker build -f Dockerfile.rocm-base -t rocm-pytorch-base:latest .

Derived Image 1: Whisper STT

# Dockerfile.whisper-rocm
FROM rocm-pytorch-base:latest

RUN pip install --no-cache-dir \
    faster-whisper \
    sounddevice \
    evdev

WORKDIR /app
CMD ["python3", "transcribe.py"]

Derived Image 2: ComfyUI

# Dockerfile.comfyui-rocm
FROM rocm-pytorch-base:latest

RUN pip install --no-cache-dir \
    comfy \
    xformers

WORKDIR /app
CMD ["python3", "comfyui.py"]

Derived Image 3: LLM Fine-tuning

# Dockerfile.llm-finetune
FROM rocm-pytorch-base:latest

RUN pip install --no-cache-dir \
    transformers \
    datasets \
    peft \
    bitsandbytes

WORKDIR /app
CMD ["python3", "finetune.py"]

Result:

• Base image: 29GB (built once)
• Each derived image: +500MB to 2GB (project-specific)
• Total disk space: ~35GB for all 4 images
• vs. separate images: ~120GB (29GB × 4)

---

Key Principles

1. Layer Ordering Matters

Put infrequently-changing layers first:

# GOOD: Stable layers first
FROM ubuntu:22.04
RUN apt-get update                    # Changes rarely
RUN pip install torch                 # Changes rarely
COPY requirements.txt .               # Changes sometimes
RUN pip install -r requirements.txt   # Changes sometimes
COPY . /app                          # Changes often

# BAD: Frequently-changing layers first
FROM ubuntu:22.04
COPY . /app                          # Changes often → invalidates cache!
RUN pip install torch                 # Re-downloads every time

2. Use .dockerignore

Prevent cache invalidation from unimportant file changes:

# .dockerignore
.git/
__pycache__/
*.pyc
.venv/
.pytest_cache/

3. Multi-Stage Builds for Efficiency

Separate build and runtime dependencies:

# Build stage: Heavy tools
FROM pytorch-base:latest AS builder
RUN pip install build-tools-needed-for-compilation

# Runtime stage: Lean image
FROM pytorch-base:latest
COPY --from=builder /compiled/artifacts /app
RUN pip install runtime-only-deps

---

Practical Workflow

Setup (Once)

# Create base image
docker build -f Dockerfile.base -t my-ml-base:latest .

# Tag with version for stability
docker tag my-ml-base:latest my-ml-base:v1.0

Daily Development (Fast)

# Create project-specific image
docker build -f Dockerfile.project1 -t project1:latest .
# Uses cached base layers → Fast!

# Iterate on project code
# Edit code...
docker build -f Dockerfile.project1 -t project1:latest .
# Only rebuilds changed layers → Very fast!

Update Base (Occasionally)

# Update PyTorch version in base
docker build -f Dockerfile.base -t my-ml-base:v2.0 .

# Update projects to use new base
# Change FROM my-ml-base:v1.0 → FROM my-ml-base:v2.0
docker build -f Dockerfile.project1 -t project1:latest .

---

Benefits vs. Conda

Aspect	Conda	Docker Layers
Isolation	Process-level	Container-level (better)
Disk Usage	Duplicate packages	Shared layers
Rebuild Time	Full reinstall	Cached layers (fast)
Portability	OS-dependent	Runs anywhere
GPU Drivers	Host-dependent	Containerized
Reproducibility	conda env.yml	Dockerfile (better)

---

Advanced: Layer Inspection

View layers in an image:

docker history my-image:latest

See layer sizes:

docker image inspect my-image:latest | jq '.[].RootFS.Layers'

Find layer cache usage:

docker system df -v

---

Common Pitfalls

❌ Pitfall 1: Not using a base image

# Every project downloads PyTorch separately
FROM ubuntu:22.04
RUN pip install torch  # 5GB download each time

✅ Solution: Use shared base

FROM pytorch-base:latest  # Cached layer
RUN pip install project-specific-deps

❌ Pitfall 2: Copying files too early

FROM pytorch-base:latest
COPY . /app              # Invalidates cache on ANY code change
RUN pip install -r requirements.txt  # Re-installs unnecessarily

✅ Solution: Copy dependencies first

FROM pytorch-base:latest
COPY requirements.txt .
RUN pip install -r requirements.txt  # Cached until requirements change
COPY . /app                          # Code changes don't affect deps

---

Summary

The Magic Formula:

1. Create a heavy base image with stable dependencies (PyTorch, CUDA, ROCm)
2. Build lightweight derived images from the base for each project
3. Docker's layer caching ensures you never re-download the heavy base
4. Update base image occasionally, rebuild derived images quickly

Time Saved:

• Initial base build: 30-60 minutes (once)
• Derived image builds: 2-5 minutes (always)
• Conda approach: 30-60 minutes (every project)

Disk Saved:

• 4 projects with conda: ~120GB (30GB each)
• 4 projects with Docker layers: ~35GB (shared 29GB base)

---

Additional Resources

• Docker Best Practices
• Multi-stage Builds
• Dockerfile Reference

---

Generated by Claude Code

This gist explains Docker image layering for efficient dependency management. For more AI/ML infrastructure tips, follow best practices for containerization and layer optimization.

Recommendation: Validate this information with official Docker documentation and test in your specific environment before production use.