如何使用 Docker SDK for Python 运行具有特定 GPU 的 Docker 容器

Question

在命令行中，我习惯使用 --gpus 参数运行/创建具有特定 GPU 的容器：

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

Docker SDK for Python 文档不是很有帮助，我找不到关于如何使用 python SDK 执行相同操作的良好解释。有办法做到吗？

Answer 1

您可以通过以下方式使用适用于 Python 的 Docker SDK 运行/创建具有特定 GPU 的 docker 容器：

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

这样您还可以使用此处指定的其他 GPU 资源选项： https://docs.docker.com/config/containers/resource_constraints/