← Back to DevBytes
Docker July 07, 2026 12 min read

Docker with GPU ML Workloads: Best Practices and Common Pitfalls

Understanding Docker with GPU ML Workloads

Running machine learning workloads inside Docker containers with GPU acceleration means configuring containers to leverage NVIDIA GPUs (or other GPU hardware) for compute-intensive tasks like training deep neural networks, running inference, or performing large-scale data processing. This approach combines the portability and reproducibility of containerization with the raw parallel processing power of modern GPUs. The fundamental mechanism relies on the NVIDIA Container Toolkit, which bridges the gap between the host's GPU drivers and the containerized environment, exposing device nodes, driver libraries, and CUDA runtime capabilities directly inside the container.

Why It Matters for ML Practitioners

🚀 Deploy your AI agent in 10 minutes

Managed Hermes hosting. Zero DevOps. 100M tokens/mo included.

Try it free →

The benefits of containerizing GPU workloads extend far beyond simple convenience. Consider these critical advantages:

Setting Up Your Environment for GPU Containers

Prerequisites Installation

Before you can run GPU-enabled containers, you need three components on the host system:

Here is how to install the NVIDIA Container Toolkit on a typical Ubuntu system:

# Add NVIDIA package repositories
curl -fsSL https://nvidia.github.io/libnvidia-container/gpgkey | \
    sudo gpg --dearmor -o /usr/share/keyrings/nvidia-container-toolkit-keyring.gpg

curl -s -L https://nvidia.github.io/libnvidia-container/stable/deb/nvidia-container-toolkit.list | \
    sed 's#deb https://#deb [signed-by=/usr/share/keyrings/nvidia-container-toolkit-keyring.gpg] https://#g' | \
    sudo tee /etc/apt/sources.list.d/nvidia-container-toolkit.list

# Install the toolkit
sudo apt-get update
sudo apt-get install -y nvidia-container-toolkit

# Configure Docker to use the NVIDIA runtime
sudo nvidia-ctk runtime configure --runtime=docker
sudo systemctl restart docker

# Verify installation
docker run --rm --gpus all nvidia/cuda:12.1-base nvidia-smi

The final verification command should display your GPU information from inside the container—confirming that the container can see and communicate with the GPUs.

Choosing the Right Base Image

NVIDIA publishes a range of official CUDA images on Docker Hub. Understanding the image variants is crucial for building efficient containers:

Here is a practical Dockerfile example that builds a PyTorch training environment with multi-stage optimization:

# syntax=docker/dockerfile:1

# Build stage - compile any custom CUDA extensions
FROM nvidia/cuda:12.1-devel AS build
WORKDIR /app
COPY requirements.txt .
RUN pip install --no-cache-dir -r requirements.txt

# Runtime stage - slim down for production
FROM nvidia/cuda:12.1-runtime
RUN apt-get update && apt-get install -y --no-install-recommends \
    python3.10 python3-pip && \
    rm -rf /var/lib/apt/lists/*
    
WORKDIR /app
COPY --from=build /usr/local/lib/python3.10 /usr/local/lib/python3.10
COPY --from=build /usr/local/bin /usr/local/bin
COPY training_script.py .

ENV PYTHONUNBUFFERED=1
ENV CUDA_VISIBLE_DEVICES=0

ENTRYPOINT ["python3", "training_script.py"]

Running GPU Containers with the --gpus Flag

Modern Docker (19.03+) supports the --gpus flag natively, eliminating the need for the legacy nvidia-docker wrapper. The flag gives you fine-grained control over GPU allocation:

# Allocate all available GPUs
docker run --gpus all -it --rm nvidia/cuda:12.1-base nvidia-smi

# Allocate specific GPUs by index (0-indexed)
docker run --gpus '"device=0,2"' -it --rm nvidia/cuda:12.1-base nvidia-smi

# Allocate GPUs with specific capabilities (compute, graphics, display)
docker run --gpus '"device=0","capabilities=compute,utility"' \
    -it --rm nvidia/cuda:12.1-base nvidia-smi

# Mount the training data and run with volume binding
docker run --gpus all \
    -v /mnt/datasets/imagenet:/data:ro \
    -v $(pwd)/checkpoints:/app/checkpoints \
    -e WANDB_API_KEY=$WANDB_API_KEY \
    --shm-size=8g \
    -it --rm \
    my-training-image:latest

Best Practices for GPU Docker Workloads

1. Pin Exact Versions of the CUDA Stack

Never use floating tags like nvidia/cuda:latest or nvidia/cuda:12.1 in production Dockerfiles. CUDA minor version changes can introduce subtle numerical differences in training results. Always pin to the exact version including the CUDA patch level and the image variant:

# Bad - floating tag, unpredictable
FROM nvidia/cuda:12.1-cudnn-devel

# Good - exact pinned version with digest for complete reproducibility
FROM nvidia/cuda:12.1.1-cudnn8.9.2-devel-ubuntu22.04@sha256:abc123...

2. Manage Shared Memory (SHM) for Data Loading

PyTorch and TensorFlow data loaders use shared memory for multi-process data loading. The default Docker /dev/shm size is 64 MB, which is woefully insufficient for ML workloads. Data loader workers will crash with mysterious "out of memory" errors if you don't address this.

# Option A: Increase shm size at runtime
docker run --gpus all --shm-size=8g --ipc=host my-image

# Option B: Use host IPC namespace (less secure but simpler)
docker run --gpus all --ipc=host my-image

# Option C: In docker-compose.yml
version: '3.8'
services:
  trainer:
    image: my-training-image:latest
    shm_size: '8gb'
    # OR
    ipc: host

3. Use Multi-Stage Builds to Reduce Image Size

GPU images tend to bloat quickly. The CUDA devel images are several gigabytes. Use multi-stage builds to compile dependencies in a devel image and copy only the runtime artifacts into a slim base image. This can reduce final image size by 60-80%:

# syntax=docker/dockerfile:1

FROM nvidia/cuda:12.1-devel AS builder
WORKDIR /build
# Install build tools, compile flash-attention, build custom ops
RUN pip install flash-attn --no-binary :all:
RUN python setup.py build_ext --inplace

FROM nvidia/cuda:12.1-runtime
WORKDIR /app
# Copy only compiled artifacts, not build toolchain
COPY --from=builder /build/dist /app/dist
COPY --from=builder /usr/local/lib/python3.10/site-packages \
    /usr/local/lib/python3.10/site-packages
RUN pip install /app/dist/*.whl

4. Layer Caching for Faster Iteration

Structure your Dockerfile so that infrequently changed layers come first. Python dependencies change less often than your training script. Copy requirements.txt and install packages before copying the source code:

FROM nvidia/cuda:12.1-cudnn-devel
WORKDIR /app

# Layer 1: System dependencies (rarely change)
RUN apt-get update && apt-get install -y --no-install-recommends \
    git cmake libopenmpi-dev && \
    rm -rf /var/lib/apt/lists/*

# Layer 2: Python dependencies (change occasionally)
COPY requirements.txt .
RUN pip install --no-cache-dir -r requirements.txt

# Layer 3: Application code (changes frequently)
COPY src/ ./src/
COPY configs/ ./configs/

CMD ["python", "src/train.py"]

5. Set CUDA Environment Variables Correctly

Several environment variables control GPU behavior inside containers. Set them explicitly rather than relying on defaults:

# In Dockerfile
ENV CUDA_VISIBLE_DEVICES=0,1,2,3
ENV NVIDIA_VISIBLE_DEVICES=all
ENV NVIDIA_DRIVER_CAPABILITIES=compute,utility
ENV TF_CPP_MIN_LOG_LEVEL=2
ENV OMP_NUM_THREADS=8

# CUDA_VISIBLE_DEVICES restricts which GPUs are visible to the application
# NVIDIA_VISIBLE_DEVICES controls which GPUs the container runtime exposes
# NVIDIA_DRIVER_CAPABILITIES specifies which driver features to mount

6. Use Docker Compose for Multi-Service ML Stacks

Many ML workflows involve multiple services: a training worker, a TensorBoard server, a Jupyter notebook for exploration, and a metrics database. Docker Compose orchestrates these elegantly:

version: '3.8'
services:
  trainer:
    image: training:latest
    deploy:
      resources:
        reservations:
          devices:
            - driver: nvidia
              count: 2
              capabilities: [compute]
    volumes:
      - ./data:/data:ro
      - ./checkpoints:/checkpoints
    shm_size: '8gb'
    environment:
      - MASTER_ADDR=trainer
      - MASTER_PORT=29500
    command: torchrun --nproc_per_node=2 train.py
  
  tensorboard:
    image: tensorboard:latest
    ports:
      - "6006:6006"
    volumes:
      - ./checkpoints:/logs:ro
    command: tensorboard --logdir /logs --bind_all
  
  jupyter:
    image: jupyter:latest
    ports:
      - "8888:8888"
    deploy:
      resources:
        reservations:
          devices:
            - driver: nvidia
              count: 1
    volumes:
      - ./notebooks:/workspace
    command: jupyter-lab --allow-root --ip=0.0.0.0

7. Profile and Benchmark Inside the Container

Always validate that GPU performance inside the container matches bare-metal expectations. Containerization should not introduce overhead for GPU computation. Use standard tools to verify:

# Run a quick CUDA bandwidth test
docker run --gpus all --rm nvidia/cuda:12.1-devel \
    /usr/local/cuda/samples/1_Utilities/bandwidthTest/bandwidthTest

# Profile PyTorch GPU utilization
docker run --gpus all --rm -it my-training-image \
    python -c "
import torch
from torch.profiler import profile, ProfilerActivity
x = torch.randn(1000, 1000, device='cuda')
with profile(activities=[ProfilerActivity.CUDA]) as prof:
    for _ in range(100):
        torch.matmul(x, x)
print(prof.key_averages().table())
"

# Monitor GPU metrics in real-time from the host
nvidia-smi dmon -s pucv -d 2

Common Pitfalls and How to Avoid Them

Pitfall 1: CUDA Driver vs. CUDA Runtime Version Mismatch

The single most common failure mode. The CUDA driver on the host must be compatible with the CUDA runtime version inside the container. The rule: the host driver must support the container's CUDA version, but not necessarily match it exactly. NVIDIA drivers are backward-compatible—a driver that supports CUDA 12.2 can run containers with CUDA 11.8, 12.0, or 12.1. However, a container built with CUDA 12.2 will fail on a host with a driver that only supports up to CUDA 11.8.

# Check host driver compatibility
nvidia-smi --query-gpu=driver_version,cuda_driver_version --format=csv

# The CUDA driver version is the key metric
# If host shows "CUDA Driver Version: 12.2", containers can use CUDA ≤ 12.2
# If container uses CUDA 12.3, you'll see: "CUDA driver version is insufficient"

Always check your host driver version before pulling or building a container image. Maintain a compatibility matrix in your project documentation.

Pitfall 2: Forgetting to Install the NVIDIA Container Toolkit on New Hosts

New team members or freshly provisioned cloud instances often have Docker installed but lack the NVIDIA Container Toolkit. The error message is cryptic:

docker: Error response from daemon: could not select device driver "nvidia" 
with capabilities: [[compute]]

The fix is the installation steps shown in the setup section above. Automate this in your infrastructure provisioning scripts or cloud-init configuration.

Pitfall 3: Data Loader Crashes Due to Insufficient SHM

As mentioned in the best practices, PyTorch DataLoader workers communicate via shared memory. When multiple workers in a container try to allocate large tensors, the default 64 MB /dev/shm limit causes SIGBUS errors. The symptoms are confusing: training crashes randomly, often during the first epoch, with errors about "insufficient memory" even though system RAM is plentiful.

# The dreaded error
RuntimeError: unable to write to file 
RuntimeError: DataLoader worker (pid 42) exited unexpectedly

# Solution: always specify --shm-size or --ipc=host
docker run --gpus all --shm-size=16g my-training-image

Pitfall 4: GPU Memory Fragmentation and OOM in Long-Running Containers

Containers that run for days or weeks (common in training jobs) can suffer from GPU memory fragmentation. PyTorch and TensorFlow caching allocators hold onto freed memory, leading to out-of-memory errors even when nvidia-smi shows free memory. Mitigations include:

# Environment variable to control PyTorch caching allocator
ENV PYTORCH_CUDA_ALLOC_CONF=max_split_size_mb:128

# For TensorFlow, enable memory growth
import tensorflow as tf
gpus = tf.config.list_physical_devices('GPU')
for gpu in gpus:
    tf.config.experimental.set_memory_growth(gpu, True)

# Periodically empty the PyTorch cache in long training loops
import torch
torch.cuda.empty_cache()

Pitfall 5: Hardcoding GPU Device Indices

Hardcoding CUDA_VISIBLE_DEVICES=0 in Dockerfiles or scripts breaks when containers are orchestrated across different GPU topologies. Container orchestration platforms like Kubernetes inject the correct device indices via environment variables. Let the orchestrator control visibility and write your code to be device-agnostic:

# Bad - hardcoded in application code
device = torch.device("cuda:0")

# Good - discover available devices dynamically
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# For distributed training, let the launcher assign ranks
# torchrun / kubeflow / slurm set LOCAL_RANK, RANK, WORLD_SIZE

Pitfall 6: Container Restart Causes GPU State Corruption

When a container crashes or is killed forcefully (SIGKILL), GPU memory may not be properly freed. The GPU can enter a state where new processes cannot allocate memory. The solution is to implement proper signal handling and graceful shutdown:

# In your Python training script
import signal
import sys

def graceful_shutdown(signum, frame):
    print("Received shutdown signal, cleaning up...")
    # Save checkpoint
    torch.save(model.state_dict(), "emergency_checkpoint.pt")
    # Clean up distributed process group
    if dist.is_initialized():
        dist.destroy_process_group()
    sys.exit(0)

signal.signal(signal.SIGTERM, graceful_shutdown)
signal.signal(signal.SIGINT, graceful_shutdown)

On the infrastructure side, use Docker's stop timeout to allow graceful cleanup before force-killing:

# In docker run, set a longer stop timeout for training jobs
docker run --gpus all --stop-timeout=120 my-training-image

# In docker-compose
services:
  trainer:
    stop_grace_period: '2m'

Pitfall 7: Network Bottlenecks in Distributed Training

When containers communicate for distributed GPU training (using NCCL, MPI, or Gloo), Docker's default bridge network introduces latency and bandwidth limitations. For multi-node distributed training:

# Use host networking for lowest latency (bypasses Docker networking entirely)
docker run --gpus all --network=host my-distributed-training-image

# For Kubernetes, use host network or RDMA-capable network plugins
# NCCL prefers direct InfiniBand or RoCE connections
# Set NCCL environment variables for debugging
ENV NCCL_DEBUG=INFO
ENV NCCL_SOCKET_IFNAME=eth0
ENV NCCL_IB_DISABLE=0

Pitfall 8: Ignoring Docker Image Layer Limits

Deep learning Docker images often exceed Docker Hub's layer limits or become so large they timeout during push/pull operations. Large model weights (checkpoints) should never be baked into images. Use volume mounts or dedicated model storage:

# Bad - baking a 10GB model checkpoint into the image
COPY model_checkpoint.pth /app/checkpoints/

# Good - mount checkpoints at runtime
docker run --gpus all -v /mnt/models/checkpoint.pth:/app/model.pth:ro my-image

# For truly portable deployment, use cloud storage SDK
# and download models at container startup, not build time

Advanced Configuration: GPU Sharing and Fractional Allocation

In production environments, not every workload requires a full GPU. NVIDIA provides several mechanisms for GPU sharing within Docker:

# Enable MPS on the host before running containers
nvidia-cuda-mps-control -d

# Then run multiple containers sharing one GPU via MPS
docker run --gpus '"device=0"' --ipc=host -d small-inference-1
docker run --gpus '"device=0"' --ipc=host -d small-inference-2

# For MIG, configure instances on the host first
sudo nvidia-smi mig -cgi 19,19,19,19,19,19,19  # 7 instances on A100
# Then assign specific MIG instances to containers
docker run --gpus '"device=0:0"' -d my-container  # First MIG instance

Security Considerations for GPU Containers

GPU containers introduce unique security considerations beyond standard container security:

docker run --gpus all \
    --read-only \
    --cap-drop=ALL \
    --cap-add=IPC_LOCK \
    --security-opt=no-new-privileges \
    -v /data/inputs:/inputs:ro \
    -v /data/outputs:/outputs \
    my-secure-training-image

Conclusion

Running GPU-accelerated machine learning workloads in Docker containers has become the de facto standard for modern ML engineering. The combination delivers reproducibility across environments, simplifies cluster orchestration, and enables sophisticated CI/CD pipelines for model training and evaluation. The key to success lies in understanding the NVIDIA Container Toolkit ecosystem, pinning exact CUDA versions, managing shared memory properly for data loading, using multi-stage builds to control image size, and avoiding the common pitfalls around driver compatibility and GPU memory management. By following the practices outlined in this guide—from proper Dockerfile structure and environment variable configuration to graceful shutdown handling and security hardening—you will build robust, portable, and production-ready GPU container workflows that behave predictably from development laptops through to multi-node GPU clusters in the cloud.

🚀 Need a reliable AI agent for your project?

Deploy Hermes Agent in 10 minutes. Managed hosting, zero DevOps.

Get Started — $23.99/mo
← Back to all articles