Run GPU workloads in Docker

Running a GPU workload in a container is the recommended path on a GPU node: the container carries the CUDA runtime, cuDNN, and framework your application needs, and the host only provides the NVIDIA driver. A container’s CUDA build can be older than the driver’s supported CUDA, never newer; check the driver’s maximum with nvidia-smi before picking a base image.

Check that Docker can reach the GPU

GPU images built on Ubuntu 22.04 come with Docker and the NVIDIA Container Toolkit, so you can run the check straight away. On an Ubuntu 24.04 image, install both first. After connecting (see Connect to a Linux GPU node), run a test container:

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

A healthy run prints the same card, driver version, and CUDA version nvidia-smi shows on the host, this time from inside the container. If it fails, the error tells you what is missing:

• docker: command not found: Docker is not installed. Install it with the steps in Install Docker, then install the NVIDIA Container Toolkit.
• unknown flag: --gpus: your Docker is too old to give a container the GPU. Remove it and install a current one with the steps in Install Docker.
• could not select device driver "" with capabilities: [[gpu]]: the NVIDIA Container Toolkit is missing or not set up. Install it with the steps in Install the NVIDIA Container Toolkit.

If the container still cannot see the card after those fixes, see Troubleshooting GPU Node.

Install Docker

Run this step on an Ubuntu 24.04 image, or on any node where docker is not found.

Ubuntu:

sudo apt update
sudo apt install -y docker.io
sudo systemctl enable --now docker

Rocky Linux:

sudo dnf -y install dnf-plugins-core
sudo dnf config-manager --add-repo https://download.docker.com/linux/rhel/docker-ce.repo
sudo dnf install -y docker-ce docker-ce-cli containerd.io
sudo systemctl enable --now docker

Check that Docker is running:

sudo docker run --rm hello-world

It prints Hello from Docker! followed by a short explanation. Docker can now run containers; install the NVIDIA Container Toolkit next so they can use the GPU.

Install the NVIDIA Container Toolkit

Run this step after installing Docker, when the GPU check fails with could not select device driver, or after you reinstall the operating system. When it finishes, run the test container from Check that Docker can reach the GPU again.

Ubuntu:

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
sudo apt update
sudo apt install -y nvidia-container-toolkit
sudo nvidia-ctk runtime configure --runtime=docker
sudo systemctl restart docker

Rocky Linux:

curl -s -L https://nvidia.github.io/libnvidia-container/stable/rpm/nvidia-container-toolkit.repo | \
  sudo tee /etc/yum.repos.d/nvidia-container-toolkit.repo
sudo dnf install -y nvidia-container-toolkit
sudo nvidia-ctk runtime configure --runtime=docker
sudo systemctl restart docker

nvidia-docker2 is deprecated and does not work with Docker Engine 25 or later; use the NVIDIA Container Toolkit with the --gpus flag on any current GPU node.

Common run patterns

Single GPU, interactive shell:

docker run --rm -it --gpus all \
  -v $PWD:/workspace -w /workspace \
  pytorch/pytorch:2.3.0-cuda12.1-cudnn8-runtime bash

To use only some cards on a multi-card node, name them. Card 0 alone, then cards 0 and 1:

docker run --rm --gpus '"device=0"' nvidia/cuda:12.4.1-base-ubuntu22.04 nvidia-smi
docker run --rm --gpus '"device=0,1"' nvidia/cuda:12.4.1-base-ubuntu22.04 nvidia-smi

A long-running service, restarting automatically and reachable only from the node itself:

docker run -d --name inference \
  --gpus all \
  --restart unless-stopped \
  -p 127.0.0.1:8080:8080 \
  -v /data/models:/models \
  your-image:tag

Binding to 127.0.0.1 keeps the port off the public internet; tunnel over SSH for browser access instead of opening it in the security group.

Picking a base image

Use case	Base image
Lightweight verification	nvidia/cuda:12.4.1-base-ubuntu22.04
CUDA development	nvidia/cuda:12.4.1-devel-ubuntu22.04 (includes nvcc)
PyTorch training	pytorch/pytorch:2.3.0-cuda12.1-cudnn8-runtime
TensorFlow training	tensorflow/tensorflow:2.16.1-gpu
vLLM inference	vllm/vllm-openai:latest
Text Generation Inference	ghcr.io/huggingface/text-generation-inference:latest

Pin to a specific tag once your workload is stable.

Storage

GPU container images run large: PyTorch and TensorFlow images are 5 to 10 GB, and vLLM or TGI add model weights on top, filling the root disk fast. Mount datasets and model weights from an attached volume rather than the root disk, and run docker system prune -af --volumes periodically to reclaim space from stopped containers. Bake an image once you have settled on a working setup; see Save and reuse a GPU image.