pythondockercudanvidia-dockerdockerpy

How to run a docker container with specific GPUs using Docker SDK for Python

In the command line I am used to run/create containers with specific GPUs using the --gpus argument:

docker run -it --gpus '"device=0,2"' ubuntu nvidia-smi

The Docker SDK for Python documentation was not very helpful and I could not find a good explanation on how to do the same with the python SDK. Is there a way to do it?

Solution

  • This is how you can run/create docker containers with specific GPUs using the Docker SDK for Python:

    client.containers.run('ubuntu',
                          "nvidia-smi",
                           device_requests=[
                           docker.types.DeviceRequest(device_ids=['0','2'], capabilities=[['gpu']])])

    This way you can also use other GPU resource options specified here: https://docs.docker.com/config/containers/resource_constraints/