Dockerfiles for data science & AI workloads¶

Module 14’s Docker chapter walks through the canonical Dockerfile — a FastAPI backend, ~200 MB, rebuilt in seconds — line by line, and this chapter assumes you’ve read it (or will): we won’t re-teach FROM/COPY/RUN, layer ordering basics, EXPOSE, or HEALTHCHECK. Go there for the generic tutorial.

This chapter is about what happens when the workload changes species. Data-science and AI images are a different animal from web-app images:

	Typical web app	Typical DS/AI workload
Dependency weight	50–200 MB	2–10 GB (PyTorch + CUDA libraries alone can exceed 5 GB)
Native/system deps	A couple (`libpq`…)	BLAS, compilers for source builds, CUDA, image/audio codecs
Hardware	CPU	Often GPU — device passthrough, driver/toolkit version coupling
Big binary artifacts	None	Models and datasets, gigabytes each
Interactive use	No	Constantly — Jupyter, exploratory work
Reproducibility bar	“App runs”	“Numbers match”

Every section below exists because one row of that table breaks a naive Dockerfile. We go pattern by pattern, then assemble a complete production-shaped example in §8. The mechanics behind every pattern — layers, whiteouts, cache keys, cache mounts — are images-and-layers.md; this chapter is that theory, applied.

1. Choosing a base image¶

The FROM line is the biggest single decision in the file. The realistic menu for Python DS work:

Base	Size (approx.)	What it is	Reach for it when
`python:3.12-slim`	~120 MB	Official Python on trimmed Debian	Default. CPU workloads; wheels cover your deps
`python:3.12`	~1 GB	Same plus compilers, headers, dev tools	You must compile source packages and won’t multi-stage (prefer multi-stage instead, §4)
`python:3.12-alpine`	~50 MB	Python on Alpine Linux (musl libc)	Almost never for DS — see warning below
`nvidia/cuda:12.4.1-runtime-ubuntu22.04`	~2–3 GB	Ubuntu + CUDA runtime libraries, no Python	GPU inference/training; you add Python (§6)
`nvidia/cuda:12.4.1-devel-ubuntu22.04`	~7 GB	Above plus `nvcc` compiler + headers	Building CUDA extensions (flash-attention & friends) — builder stage only
`pytorch/pytorch:2.4.0-cuda12.4-cudnn9-runtime`	~7 GB	PyTorch’s own: Python + torch + CUDA, assembled	Fastest path to a working GPU image; less control
`micromamba` / conda-based images	varies	Conda ecosystems	Your stack genuinely needs conda (geo, some bio); else skip the weight

⚠️ Common mistake: Choosing -alpine “because it’s smallest.” Alpine uses musl instead of glibc (the GNU C library that effectively all scientific-Python binary wheels — precompiled packages — are built against). Result: pip finds no compatible NumPy/pandas/torch wheels and compiles them from source — builds go from 40 seconds to 40 minutes, need a full compiler toolchain (there goes “smallest”), and you may hit musl-specific runtime bugs. The slim images already give you glibc at ~120 MB. For scientific Python, slim, not alpine — the exception is pure-Python utility containers where either works.

And pin the whole choice down: FROM python:3.12-slim today and in six months are different images (the tag is re-pointed at every patch — by design). For work that must reproduce, pin by digest, exactly as images-and-layers.md §4 prescribes:

FROM python:3.12-slim@sha256:5c73  # ...full digest; readable tag + immutable identity

2. Pinning: the environment is part of the experiment¶

A container freezes an environment only if the build is deterministic. RUN pip install pandas torch is a different environment every week — you’ve containerized the irreproducibility. The fix is the standard two-file discipline, executed inside the image build:

pyproject.toml (or requirements.in) — what you want, loosely (“pandas”, “torch>=2.4”).
A lockfile — uv.lock / requirements.txt generated by uv lock, pip-compile, or pip freeze — every package, every transitive dependency, exact == versions, ideally --generate-hashes so pip verifies artifact hashes (content addressing again — the same idea as image digests, one level up).

The Dockerfile then installs from the lockfile only. Three pins, three layers of the tower: base image by digest (the OS + interpreter), Python packages by lockfile (the libraries), and — for GPU images — the CUDA version by explicit base-image choice (§6). That combination, checked into Git next to your code, is what makes “the environment for commit 9f3c1ab” a meaningful, reconstructible phrase. (ecosystem.md builds the reproducible-research workflow on top of exactly this.)

3. The dependency layer: ordering + cache mounts¶

Module 14’s ordering rule — copy the lockfile and install before copying code — matters proportionally to the weight of the install. For a 5 GB torch install, a misordered COPY . . doesn’t cost you 3 minutes per rebuild; it costs 20. So, first, the rule as always:

COPY requirements.txt .
RUN pip install --no-cache-dir -r requirements.txt   # heavy layer, cached until lockfile changes
COPY src/ ./src/                                     # light layer, changes constantly

Then the upgrade Module 14 doesn’t need but you do: BuildKit cache mounts (images-and-layers.md §7). The pain: whenever the lockfile does change — you add one small package — the whole RUN re-executes and pip re-downloads every wheel from scratch, because --no-cache-dir threw the download cache away (correctly! — it would otherwise bloat the layer). A cache mount resolves the dilemma: a persistent directory that exists only during builds, shared across builds, and never enters any layer:

RUN --mount=type=cache,target=/root/.cache/pip \
    pip install -r requirements.txt

Now a one-line lockfile change re-runs the install but pulls almost everything from the local wheel cache: minutes → seconds, and the image stays lean (no --no-cache-dir needed — the cache isn’t in the image at all).

With uv — the fast Rust-based package manager this course’s tooling generation is converging on — the same pattern, plus uv’s own speed:

COPY --from=ghcr.io/astral-sh/uv:latest /uv /usr/local/bin/uv

ENV UV_LINK_MODE=copy
COPY pyproject.toml uv.lock ./
RUN --mount=type=cache,target=/root/.cache/uv \
    uv sync --frozen --no-install-project      # deps only — project code isn't copied yet
COPY src/ ./src/
RUN --mount=type=cache,target=/root/.cache/uv \
    uv sync --frozen                           # now install the project itself (fast)

--frozen refuses to run if uv.lock is out of date (reproducibility guardrail), and the two-step --no-install-project split keeps the giant dependency layer independent of your source code — the ordering rule, expressed in uv.

4. Multi-stage builds: compile in a workshop, ship a showroom¶

Scientific packages sometimes must be compiled (no wheel for your platform, or a CUDA extension). Compilation needs gcc, headers, -devel toolkits — none of which belong in the image you run (weight + attack surface). A multi-stage build uses two FROMs in one Dockerfile: a fat builder stage does the compiling into an isolated prefix (a venv — virtual environment — is perfect, because it’s relocatable by copy); the slim runtime stage copies in only the finished result. Everything not explicitly copied forward is discarded — this is also the definitive answer to the “deleting doesn’t shrink” whiteout problem from images-and-layers.md §2.

# ---- Stage 1: builder — compilers welcome, size irrelevant ----
FROM python:3.12 AS builder
RUN python -m venv /opt/venv
ENV PATH="/opt/venv/bin:$PATH"
COPY requirements.txt .
RUN --mount=type=cache,target=/root/.cache/pip \
    pip install -r requirements.txt            # source builds succeed: gcc & headers present

# ---- Stage 2: runtime — slim, no compilers, no pip cache, no history ----
FROM python:3.12-slim
COPY --from=builder /opt/venv /opt/venv
ENV PATH="/opt/venv/bin:$PATH"
COPY src/ ./src/
CMD ["python", "-m", "src.serve"]

The runtime image contains the base’s ~120 MB plus exactly the venv and your code — the builder’s gigabyte of toolchain evaporates. For CUDA-extension builds, the same skeleton with nvidia/cuda:…-devel as builder and …-runtime as runtime routinely saves 4–5 GB (§6).

5. Hygiene: non-root, `.dockerignore`, and DS-specific traps¶

Both practices come from Module 14’s Dockerfile; what’s new here is what they mean for this workload.

Non-root user. Same pattern as Module 14 (useradd + USER); the DS-specific wrinkle is volume-permission collisions: your container’s user writes checkpoints into a bind-mounted host directory, and the files land owned by a UID that isn’t yours on the host (or vice versa — the container can’t write at all). The fix is to align UIDs at build time (useradd --uid 1000 …, matching the typical first Linux user) or set ownership explicitly — the full story is in volumes-and-networking.md. Jupyter’s official images ship a ready-made jovyan (UID 1000) user for exactly this reason (§7).

.dockerignore. For a web app it saves megabytes; for DS work it saves you from catastrophe, because DS working directories accumulate giant artifacts that must never enter a build context: datasets, checkpoints, experiment tracking dirs, venvs. A missing .dockerignore turns docker build into “upload 40 GB to the daemon, then bake data/ into a layer forever.” A DS-shaped starting point:

.git/
.venv/
__pycache__/
*.pyc
.ipynb_checkpoints/
data/            # datasets: mount at runtime, never bake
models/          # weights: see §8
checkpoints/
outputs/
wandb/           # experiment-tracker local cache
mlruns/
.env             # secrets NEVER enter build contexts — see warning below

⚠️ Common mistake: Passing secrets (API keys, tokens for private data) via ENV or COPY .env. Every layer is inspectable forever (docker history, or just reading the tar — images-and-layers.md §2); an image with a baked secret is a leaked secret. Runtime config goes in docker run -e/--env-file; build-time secrets (a private index token for pip install) use BuildKit secret mounts — RUN --mount=type=secret,id=pip_token … — which expose the value during that one instruction and leave no trace in any layer.

6. GPU images¶

GPUs pierce the container boundary deliberately (container-internals.md §6): a container gets a GPU because the runtime mounts the host’s device files and driver libraries into it, not because anything is virtualized. That architecture dictates the whole recipe.

The division of labor — memorize this split, it resolves 90% of GPU-container confusion:

Piece	Lives on the host	Lives in the image
NVIDIA kernel driver	✅ (one per machine, admin-installed)	❌ never
NVIDIA Container Toolkit (the mounting shim)	✅	❌
CUDA toolkit / runtime libraries	❌ (not needed for containers)	✅ (your chosen version)
cuDNN, NCCL, framework (torch…)	❌	✅

So: install the driver + toolkit once per machine; ship CUDA inside the image; request devices at run time:

docker run --gpus all my-training-image           # all GPUs
docker run --gpus '"device=0,1"' my-training-image  # specific GPUs
# smoke test — should print the host's GPUs from inside a bare CUDA image:
docker run --rm --gpus all nvidia/cuda:12.4.1-base-ubuntu22.04 nvidia-smi

Version compatibility is one-directional and simpler than its reputation: the host driver must be at least as new as the image’s CUDA version requires (drivers are backward-compatible with older CUDA). nvidia-smi’s “CUDA Version” field reports the maximum the driver supports, not what’s installed. Practical protocol: check the host’s driver, then pick image CUDA ≤ that ceiling — and note that this is precisely how one host serves many CUDA versions side by side, which is half the reason GPU clusters containerize everything.

Choosing the base. Three nvidia/cuda flavors: base (bare runtime loader), runtime (+ CUDA libraries — right for running frameworks), devel (+ nvcc and headers — needed only to compile CUDA code; use it as a §4 builder stage, ship runtime). Note they contain no Python — you apt-get install it or copy a venv in. Alternatively pytorch/pytorch:…-runtime bundles the whole stack — bigger and less flexible, but zero assembly. One more sizing lever: pip install torch defaults to bundling CUDA libraries again inside the wheel (~2.5 GB); on a CUDA base image, or for CPU-only images, select the right wheel index (--index-url https://download.pytorch.org/whl/cpu for CPU) — the single most common cause of accidentally-8-GB “small” images.

The --shm-size trap, promised in container-internals.md §3: Docker caps /dev/shm (POSIX shared memory, in the ipc namespace) at 64 MB by default; PyTorch DataLoader workers pass tensors through it and crash cryptically (bus error, worker killed) under real loads. Run training containers with --shm-size=2g (or --ipc=host on trusted single-tenant boxes). This flag belongs in every GPU-training runbook.

7. Jupyter in a container¶

Jupyter is just an HTTP server, so it containerizes beautifully — one of the best-paved paths in the ecosystem. Two routes:

Route 1 — the Jupyter Docker Stacks (quay.io/jupyter/scipy-notebook, pytorch-notebook, …): maintained images with Jupyter Lab, the scientific stack, and the UID-1000 jovyan user preconfigured:

docker run --rm -p 8888:8888 \
  -v "$PWD":/home/jovyan/work \
  quay.io/jupyter/scipy-notebook:latest
# then open the http://127.0.0.1:8888/lab?token=... URL it prints

Route 2 — add Jupyter to your project image, so notebooks run in exactly the pinned environment your pipeline uses (the reproducibility win — one image for scripts and exploration):

RUN pip install jupyterlab            # in the same locked environment
EXPOSE 8888
CMD ["jupyter", "lab", "--ip=0.0.0.0", "--port=8888", "--no-browser"]

The three details that make or break it:

--ip=0.0.0.0 — bind all interfaces, or the published port connects to nothing (the same 127.0.0.1-inside-a-container trap as Module 14’s uvicorn).
Mount your notebook directory (-v "$PWD":/home/jovyan/work). Anything saved outside a mount lives in the container’s writable layer and dies with the container — the classic “my notebooks vanished” incident, dissected in volumes-and-networking.md.
Keep token auth on (the URL-with-token it prints). A tokenless Jupyter is remote code execution for anyone who can reach the port; if you must set your own, use JUPYTER_TOKEN via -e, and publish as -p 127.0.0.1:8888:8888 on shared machines so only local users can connect.

8. Model files: the gigabytes that don’t belong in layers¶

The last DS-specific decision: where do trained weights live? Baking a 5 GB checkpoint into the image via COPY model.pt . is the tempting default, and usually wrong:

Every model update re-ships a 5 GB layer to the registry and to every puller — even though code changed not at all (and vice versa: layer dedup helps only if you order it carefully).
Registries and pull times balloon; node disk fills with near-duplicate images.
The model’s lifecycle (retrain weekly) is glued to the code’s lifecycle (deploy daily) — two release cadences, one artifact.

Decision rule:

Situation	Where the model goes
Same host, iterating (dev, training outputs)	Volume / bind mount — volumes-and-networking.md
Serving; model updates ≠ code updates	Download at startup from object storage / a model registry (HF Hub, MLflow, S3), pinned to an exact model version, into a cached volume
Serving; model is small (< a few hundred MB) or must be atomically versioned with code	Bake it — as its own layer, after deps, before/independent of code, so a code push never re-ships weights
Kubernetes-scale serving	Init containers / model mounts — out of scope, see ecosystem.md

Note the deliberate symmetry: models at runtime are pinned by version/revision, images by digest, packages by lockfile — same principle, three artifact types.

The assembled example¶

Everything above, in one production-shaped GPU inference image:

# syntax=docker/dockerfile:1

# ---- Stage 1: builder — compile/install into a portable venv ----
FROM nvidia/cuda:12.4.1-devel-ubuntu22.04 AS builder
RUN apt-get update && apt-get install -y --no-install-recommends \
        python3.12 python3.12-venv && rm -rf /var/lib/apt/lists/*
RUN python3.12 -m venv /opt/venv
ENV PATH="/opt/venv/bin:$PATH"
COPY requirements.txt .                          # the LOCKFILE, == pins throughout
RUN --mount=type=cache,target=/root/.cache/pip \
    pip install -r requirements.txt              # devel image: CUDA extensions compile fine

# ---- Stage 2: runtime — CUDA runtime libs only, no compilers ----
FROM nvidia/cuda:12.4.1-runtime-ubuntu22.04
RUN apt-get update && apt-get install -y --no-install-recommends \
        python3.12 && rm -rf /var/lib/apt/lists/* \
    && useradd --create-home --uid 1000 appuser
COPY --from=builder /opt/venv /opt/venv
ENV PATH="/opt/venv/bin:$PATH" PYTHONUNBUFFERED=1
WORKDIR /app
COPY src/ ./src/                                 # code last: changes most often
USER appuser
EXPOSE 8000
# weights arrive at runtime: MODEL_URI + a mounted cache volume, e.g.
#   docker run --gpus all --shm-size=2g -p 8000:8000 \
#     -e MODEL_URI=hf://org/model@<revision> -v model-cache:/home/appuser/.cache \
#     my-inference:latest
CMD ["python", "-m", "src.serve"]

Read the layer order like a cache-cost gradient, top = most stable: base (pinned) → system Python → venv with locked deps (the 5 GB layer, invalidated only by lockfile edits) → code (kilobytes, changes daily) — with models outside the image entirely, and the compiler stage discarded.

9. Exercises¶

Try each before opening the answer.

Exercise 1. A teammate’s image is 9.4 GB: FROM python:3.12 + COPY . . (their data/ and .git/ included) + pip install torch -r requirements.txt (CPU-only inference service). Propose four concrete changes and estimate the effect of each.

Answer

(1) .dockerignore excluding data/, .git/, .venv/, checkpoints — potentially the largest single cut, and stops the giant build-context upload. (2) CPU wheel index for torch (--index-url https://download.pytorch.org/whl/cpu) — the default wheel bundles ~2.5 GB of CUDA libraries that a CPU service never touches. (3) python:3.12-slim (or a slim runtime via a multi-stage build) instead of full — saves ~900 MB of toolchain. (4) Reorder: lockfile-copy + install before code-copy, with a pip cache mount — doesn’t shrink the image but turns every rebuild and re-push from tens of minutes into seconds (only the code layer changes). Plausible landing zone: well under 2 GB, with near-instant iterative rebuilds.

Exercise 2. docker run --gpus all myimage python -c "import torch; print(torch.cuda.is_available())" prints False on a server where nvidia-smi works fine on the host. Give three distinct hypotheses and the check for each.

Answer

(1) NVIDIA Container Toolkit missing/not wired to Docker — the device/driver mounting shim isn’t running: check docker run --rm --gpus all nvidia/cuda:12.4.1-base-ubuntu22.04 nvidia-smi; if that fails, it’s host setup, not your image. (2) CPU-only torch wheel in the image — pip show torch / check whether the version string ends in +cpu; reinstall from the CUDA wheel index. (3) CUDA/driver mismatch — image’s CUDA needs a newer driver than the host has: compare the image’s CUDA version against the driver ceiling nvidia-smi reports; pick an older CUDA base or upgrade the driver. (Also worth one second: did the actual run include --gpus — orchestrators and scripts sometimes drop it.)

Exercise 3. Why is a venv the standard vehicle for multi-stage Python builds — i.e., what property makes COPY --from=builder /opt/venv /opt/venv work, and what constraint does it impose on the two stages?

Answer

A venv confines an entire installed environment (interpreter shims, site-packages, entry-point scripts) under one self-contained directory tree, so a single COPY --from moves the complete result of pip install — including compiled C extensions — without dragging along pip’s caches, build dirs, or the compiler toolchain. The constraint: the runtime stage must be binary-compatible with the builder — same Python version and same libc family (compiled extensions link against both), and same path (/opt/venv) since venvs hard-code their location. Hence builder/runtime pairs from the same family: python:3.12 → python:3.12-slim, cuda:…-devel-ubuntu22.04 → cuda:…-runtime-ubuntu22.04 — and never a glibc builder with an Alpine/musl runtime.

Exercise 4. Your training container dies overnight with workers killed by bus error; the same code runs fine outside Docker. Diagnosis and fix — and why does this only happen in a container?

Answer

PyTorch DataLoader workers exchange tensors via POSIX shared memory in /dev/shm; Docker mounts /dev/shm at only 64 MB by default (inside the container’s ipc namespace), and overflowing it kills workers with SIGBUS. Outside Docker, /dev/shm is typically half of RAM, so the limit never bites. Fix: docker run --shm-size=2g … (size to your batch/prefetch volume), or --ipc=host on a trusted single-user machine; reducing num_workers/batch size also lowers pressure.

Exercise 5. For a fraud-scoring service whose model is retrained nightly but whose code changes weekly, argue for a model-delivery strategy from §8’s table — and name the failure mode of the opposite choice.

Answer

Download at startup, pinned to an exact model version, cached on a volume. Cadences differ by 7×, so decoupling wins: nightly retrains publish a new model version and roll the service by changing one env var/config value — no image build, no 5 GB push/pull, instant rollback to yesterday’s version by re-pinning. Baking instead means a nightly full image pipeline where the only diff is a 5 GB layer: slow, registry-bloating, and it entangles model rollback with code rollback (reverting a bad model drags the code image back too, or forces a rebuild). Keep the pinning discipline though — startup-download of an unpinned “latest” model reintroduces exactly the irreproducibility containers were hired to kill.

Dockerfiles for data science & AI workloads¶

1. Choosing a base image¶

2. Pinning: the environment is part of the experiment¶

3. The dependency layer: ordering + cache mounts¶

4. Multi-stage builds: compile in a workshop, ship a showroom¶

5. Hygiene: non-root, .dockerignore, and DS-specific traps¶

6. GPU images¶

7. Jupyter in a container¶

8. Model files: the gigabytes that don’t belong in layers¶

The assembled example¶

9. Exercises¶

Next / Related¶

5. Hygiene: non-root, `.dockerignore`, and DS-specific traps¶