Docker and containers for ML

A multi-stage build uses a heavy 'builder' stage to compile dependencies (e.g., wheels, CUDA extensions, model artifacts) and then copies only the runtime artifacts into a slim final stage. This drastically reduces final image size and attack surface by excluding compilers, build caches, and dev headers. For ML it is key because frameworks like PyTorch pull in large build toolchains you don't need at serving time.

FROM python:3.11 AS builder
RUN pip wheel --wheel-dir=/wheels -r requirements.txt

FROM python:3.11-slim
COPY --from=builder /wheels /wheels
RUN pip install --no-index --find-links=/wheels -r requirements.txt

The NVIDIA Container Toolkit installs a runtime hook that injects the host's GPU driver libraries and device files into the container at startup, so the container only needs the CUDA toolkit/runtime, not the driver. The '--gpus all' flag (or the device requests in Compose/K8s) tells the runtime to expose all visible GPUs to the container via this hook. This decouples the container's CUDA version from the host driver, within driver-compatibility limits.

docker run --gpus all nvidia/cuda:12.4.1-runtime-ubuntu22.04 nvidia-smi

Floating tags like 'latest' or even 'python:3.11' can be re-pushed, silently changing CUDA, glibc, or library versions and breaking reproducibility or numerical results. Pinning by digest (image@sha256:...) guarantees byte-for-byte identical base layers across builds. This is critical for ML where minor library changes can alter model outputs or training behavior.

FROM python@sha256:9b9...e2 AS base

Docker caches layers by instruction; a layer is invalidated when its inputs change, and all subsequent layers rebuild. You should COPY only the requirements file and run pip install before copying the (frequently changing) application/model code, so dependency installation is cached across code edits. Copying the whole source first would bust the pip cache on every code change.

COPY requirements.txt .
RUN pip install -r requirements.txt
COPY . .

A cache mount persists a directory (e.g., the pip or HuggingFace cache) across builds without baking it into the image layer, so re-downloads are avoided even when the layer is invalidated. Plain layer caching only helps when the instruction's inputs are unchanged; once requirements change, pip re-downloads everything. Cache mounts keep the package downloads warm while keeping the final image lean.

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

Running as non-root limits the blast radius if the container is compromised and is often mandated by hardened K8s policies (e.g., PodSecurity 'restricted'). The pitfall is file ownership on bind/volume mounts: files written by a non-root UID may be unwritable by the host user or vice versa, and model/checkpoint directories may need matching UID/GID or fsGroup settings. You typically create a user in the Dockerfile and ensure mounted paths are owned appropriately.

RUN useradd -m -u 1000 mluser
USER mluser

Docker defaults /dev/shm (shared memory) to 64MB, but PyTorch's multiprocessing DataLoader uses shared memory to pass tensors between workers, quickly exhausting it. The fix is to increase shared memory via '--shm-size' (e.g., 1g or more) or mount a larger tmpfs at /dev/shm. Alternatively set '--ipc=host', though that weakens isolation.

docker run --gpus all --shm-size=2g my-training-image

Baking weights in gives a fully self-contained, reproducible, fast-starting image but bloats it (often many GB), slows builds/pulls, and couples model versioning to image versioning. Loading at runtime from S3/GCS/registry keeps images small and lets you swap models without rebuilding, at the cost of cold-start download latency and a network/credentials dependency. Large multi-GB foundation models usually favor runtime loading with a warm cache or init container.

The build context is sent to the daemon before building; ML repos often contain massive datasets, checkpoints, .git history, and virtualenvs that would make the context multi-GB and slow every build. A '.dockerignore' excludes these (data/, *.ckpt, .git, venv/, __pycache__, mlruns/) so only relevant files are sent. This speeds builds and prevents accidentally baking secrets or large artifacts into the image.

data/
*.ckpt
.git
__pycache__/
mlruns/

The 'devel' image includes the full CUDA toolkit (nvcc compiler, headers, dev libraries) needed to compile custom CUDA/C++ extensions; the 'runtime' image only ships the runtime libraries needed to execute precompiled code. Use 'devel' when building packages with custom kernels (e.g., flash-attention, custom ops) and 'runtime' (or even 'base') for the final, smaller serving image. A multi-stage build often compiles in devel and serves from runtime.

You request GPUs via the extended resource 'nvidia.com/gpu' under resources.limits, which the NVIDIA device plugin advertises per node. The scheduler places the pod only on nodes with enough free GPUs, and GPUs are allocated whole (not fractionally) unless using MIG or time-slicing. Requests and limits for GPUs must be equal because GPUs are non-overcommittable integer resources.

resources:
  limits:
    nvidia.com/gpu: 1

MIG partitions a supported GPU (e.g., A100/H100) into multiple isolated hardware instances, each with dedicated compute slices, memory, and cache. In Kubernetes the device plugin can expose these as distinct schedulable resources (e.g., 'nvidia.com/mig-1g.10gb'), letting several small inference containers share one physical GPU with hard isolation. This improves utilization for workloads that can't saturate a full GPU, unlike time-slicing which provides no memory isolation.

Docker images are architecture-specific; an image built for linux/arm64 won't run natively on linux/amd64 nodes (it would require slow QEMU emulation or fail outright). The fix is to build for the target platform using buildx with '--platform linux/amd64' or produce a multi-arch manifest. This matters doubly for GPU images since the prebuilt CUDA/PyTorch wheels you need for x86 GPU servers are amd64 builds.

docker buildx build --platform linux/amd64 -t myorg/train:amd64 .

Use BuildKit secret mounts ('--mount=type=secret'), which expose the secret only during a specific RUN and never persist it in any layer or history. Passing secrets via ARG or ENV bakes them into the image and 'docker history', and copying a creds file then deleting it still leaves it in an earlier layer. At runtime, inject secrets via environment variables, mounted files, or a secrets manager rather than the image.

RUN --mount=type=secret,id=hf_token \
    HF_TOKEN=$(cat /run/secrets/hf_token) python download_model.py

Each RUN creates a new layer capturing the filesystem delta, so apt caches, pip caches, and intermediate files in one RUN persist even if removed in a later RUN. Combine related installs into a single RUN, use '--no-cache-dir' for pip, and clean apt lists in the same layer ('rm -rf /var/lib/apt/lists/*'). This prevents bloated, immutable cache layers that 'rm' in later steps cannot reclaim.

RUN pip install --no-cache-dir -r requirements.txt \
 && rm -rf /root/.cache

Use the 'deploy.resources.reservations.devices' block with driver 'nvidia' and a capability of 'gpu' (and optionally count or device_ids). This replaces the older 'runtime: nvidia' approach and works with the NVIDIA Container Toolkit. Compose then requests the GPU from the runtime when starting the service.

services:
  trainer:
    deploy:
      resources:
        reservations:
          devices:
            - driver: nvidia
              count: 1
              capabilities: [gpu]

A container's CUDA runtime must be supported by the host's NVIDIA driver; newer CUDA generally requires a sufficiently new driver (forward-compatible within limits via the compat package on datacenter GPUs). The host driver is what executes GPU work, so a too-old driver causes 'CUDA driver version is insufficient' errors regardless of the image's CUDA version. NVIDIA_DRIVER_CAPABILITIES controls which driver features (compute, utility, video) get mounted in.

Distroless/slim images contain only the runtime and app dependencies (no shell, package manager, or extra utilities), reducing size and vulnerability count significantly. The limitation is debuggability: with no shell you can't 'exec' in to poke around, so you rely on ephemeral debug containers or sidecars. They also require care that all native shared libraries (e.g., libgomp, CUDA libs) are present.

If the process runs as PID 1 it must handle SIGTERM itself; many runtimes started via a shell ('CMD python ...' in shell form) run under /bin/sh as PID 1, which does not forward signals, so the server isn't gracefully drained on pod termination. Use exec form ('CMD ["python","server.py"]') or an init like tini so SIGTERM reaches the app, letting it finish in-flight requests within terminationGracePeriodSeconds. Otherwise requests are dropped on scale-down or rollout.

CMD ["python", "server.py"]  # exec form, app is PID 1 and gets SIGTERM

A golden base image pins the CUDA, cuDNN, Python, and core framework versions plus security patches once, so downstream teams inherit a vetted, reproducible, scanned foundation. This reduces drift, speeds builds via shared cached layers, and centralizes vulnerability remediation. Teams add only their app layer on top, improving consistency and supply-chain governance.

One container per model gives strong isolation, independent scaling/rollout, and clear resource accounting, but multiplies overhead and can waste GPU memory when each model is small. A multi-model server (e.g., Triton, TorchServe, vLLM) packs many models per container/GPU for higher utilization and shared batching, at the cost of noisy-neighbor risk and coupled deployment lifecycles. Choice depends on model size, traffic shape, and isolation requirements.

LLM images plus baked weights can be tens of GB, so pulling and decompressing layers on a cold node adds minutes before the pod is Ready, hurting autoscaling responsiveness. Mitigations include keeping weights out of the image (load from a fast cache/PVC), pre-pulling images on nodes (DaemonSet warmers), lazy-loading snapshotters (e.g., stargz/SOCI) that stream layers on demand, and smaller compressed layers. Right-sizing the image is the biggest lever.

Requests guarantee scheduled resources and influence bin-packing; limits cap usage and trigger throttling (CPU) or OOM-kill (memory) when exceeded. Omitting a memory limit lets a leaky or large data-loading job consume node memory until the kernel OOM-killer reaps processes, potentially evicting other pods and causing node instability. For ML, set memory limits and tune DataLoader/batch sizes; note CPU limits can throttle data pipelines.

resources:
  requests: { cpu: "4", memory: 16Gi }
  limits:   { cpu: "8", memory: 32Gi }

Pin fully-resolved, hashed dependency sets using a lock file (pip-tools 'requirements.txt' with --generate-hashes, Poetry/uv lock, or conda-lock) plus a pinned base image digest. Hashes prevent dependency confusion and ensure identical transitive versions, while the lock captures the entire resolved graph. Combine with fixed seeds and PYTHONHASHSEED only if you additionally need deterministic numerics.

pip-compile --generate-hashes requirements.in

An init container runs to completion before the main app container starts, sharing volumes with it. A common ML use case is downloading large model weights or datasets into an emptyDir/PVC before the serving or training container launches, keeping the main image small and decoupling artifact fetch from app logic. It can also handle schema migrations, warm caches, or wait for dependencies.

'docker commit' snapshots ad-hoc, undocumented manual changes, producing an opaque image that can't be reliably rebuilt or audited. ML pipelines need declarative, version-controlled Dockerfiles so the environment is reproducible, reviewable, and CI-buildable. Committed images drift from source and break supply-chain traceability.

Scan images with tools like Trivy, Grype, or registry-native scanners in CI to enumerate OS and Python-package CVEs, then reduce surface by using slim/distroless bases, removing build tools in the final stage, and updating pinned dependencies. Rebuild regularly so base patches flow in, and re-pin by digest only after re-scanning. Track an SBOM to map components to advisories.

trivy image --severity HIGH,CRITICAL myorg/serve:1.2.0

A bind mount maps a specific host path into the container (great for live-editing code or pointing at an existing dataset on the host), while a named volume is Docker-managed storage with a lifecycle independent of any host path (good for persistent caches, checkpoints, or DB data that should survive container removal). Bind mounts couple to host layout and can introduce permission/perf issues; named volumes are more portable and managed.

docker run -v $(pwd)/src:/app/src -v model-cache:/root/.cache app

Multi-GPU distributed training relies on collective communication (all-reduce) via NCCL, whose throughput dominates scaling efficiency; containers must expose the right interfaces and ports and often need host networking or RDMA/GPUDirect for InfiniBand/RoCE performance. Misconfiguration (wrong NCCL_SOCKET_IFNAME, blocked ports, bridge overhead) causes hangs or severe slowdowns. On K8s you typically use a headless service for rendezvous and high-performance NICs via device plugins.

ENV NCCL_SOCKET_IFNAME=eth0

A readiness probe gates traffic (pod removed from the Service endpoints when not ready), while a liveness probe restarts a hung container. The ML gotcha is model load time: if the liveness probe starts too early or has too short a timeout, the pod is killed mid-load in a crash loop, so you use a startupProbe or generous initialDelay, and only mark ready after the model is loaded and warmed.

startupProbe:
  httpGet: { path: /healthz, port: 8080 }
  failureThreshold: 30
  periodSeconds: 10

An entrypoint script can perform runtime setup (download/verify model, set NCCL/threads env, run DB migrations, detect GPU count, exec the real process) while still passing through args. Using 'exec' at the end preserves correct signal handling/PID 1 semantics. This makes one image flexible across roles (train/serve/worker) driven by env or args.

#!/bin/sh
set -e
python prefetch_model.py
exec "$@"

Conda environments add large size, slow solves, and can conflict with system libraries and pip-installed CUDA wheels, hurting reproducibility and build speed. Mitigations include using mamba/micromamba for fast solves, pinning with conda-lock, cleaning caches ('conda clean -afy'), and increasingly preferring pip + manylinux/CUDA wheels (or uv) when conda's binary packages aren't required. Keep the environment minimal and lock it.

RUN conda clean -afy

GPU memory is freed when the owning process exits and the CUDA context is destroyed; a hung or zombie process (or a server holding the context) keeps memory allocated, and there is no per-cgroup GPU memory limit like CPU/RAM by default. Orchestrators handle reclamation by terminating the whole container/pod (which kills the process and frees the context) and rely on MIG for hard memory partitioning. Ensure clean process exit and avoid leaking child processes.

Keep data out of the image and mount it at runtime: bind mount a host/NFS path, attach a PVC backed by network/object storage, or stream from object storage (S3/GCS) with a caching layer or FUSE mount. This keeps images small and lets the same image train on different datasets, while letting you scale storage and IO independently. For high throughput, use local NVMe scratch with a prefetch/caching step.