How to use NVIDIA GPUs in docker container

Goal:

This is a quick note on how to use NVIDIA GPUs in docker container.

Env:

Ubuntu 18.04

Docker 20.10.5

Solution:

The key is to install NVIDIA Container Toolkit that is why this note is quick:)

1. Install Docker on host machine where NVIDIA driver is already installed.

https://docs.docker.com/engine/install/ubuntu/

Note: Refer to this post on how to install CUDA Toolkit and NVIDIA Driver on ubuntu.

2. Install NVIDIA Container Toolkit on host machine

https://docs.nvidia.com/datacenter/cloud-native/container-toolkit/install-guide.html#docker

3. Test

sudo docker run --rm --gpus all nvidia/cuda:11.0-base nvidia-smi

Or only expose the first GPU(with device=0) instead of all GPU to docker container:

docker run --rm --gpus device=0 nvidia/cuda:11.0-base nvidia-smi