Docker Compose configuration for isolated AI experiment environment with network isolation and resource limits

ai-docker-compose.yml

# Isolated sandbox for AI experiments.
# Source: https://williamzujkowski.github.io/posts/2025-04-10-securing-personal-ai-experiments/
#
# Notes on the things that are easy to get wrong here:
#
#  * CUDA_VISIBLE_DEVICES alone does NOT attach a GPU. It is an environment
#    variable read by CUDA inside the container; without a device reservation
#    the container sees no GPU at all and the variable selects from nothing.
#  * `deploy.resources` is Swarm-oriented and has historically been ignored by
#    plain `docker compose up`. The top-level `cpus` / `mem_limit` keys are
#    honoured everywhere, so use those and avoid the ambiguity.
#  * `read_only: true` needs tmpfs for every path the framework writes to.
#    PyTorch and Hugging Face both write to ~/.cache, so /tmp alone will fail.
#  * Pin the image by digest. `:latest` in a security sandbox means the thing
#    you are isolating changes without you noticing.

services:
  ai-sandbox:
    # Replace with the digest you actually verified:
    #   docker pull pytorch/pytorch:2.5.1-cuda12.1-cudnn9-runtime
    #   docker inspect --format='{{index .RepoDigests 0}}' pytorch/pytorch:2.5.1-cuda12.1-cudnn9-runtime
    image: pytorch/pytorch@sha256:REPLACE_WITH_VERIFIED_DIGEST
    container_name: ai-experiment

    # Least privilege. The default 14 capabilities are more than inference needs.
    user: "1000:1000"
    cap_drop:
      - ALL
    security_opt:
      - no-new-privileges:true

    read_only: true
    tmpfs:
      - /tmp:size=2g
      - /home/user/.cache:size=8g

    volumes:
      - ./models:/workspace/models:ro
      - ./data:/workspace/data

    # Honoured by `docker compose up` in every version, unlike deploy.resources.
    cpus: 4
    mem_limit: 16g
    pids_limit: 512

    # This is what actually attaches the GPU.
    gpus:
      - driver: nvidia
        count: 1
        capabilities: [gpu]

    networks:
      - ai-isolated

networks:
  ai-isolated:
    driver: bridge
    # No route off the host. If you need egress to a model registry, put the
    # sandbox on a second network with a filtered proxy rather than removing
    # this — see the firewall gist for why hostname allow-lists do not work.
    internal: true