背景
深度学习的环境配置通常是一项比较麻烦的工作，尤其是在多个用户共享的服务器上。虽然conda集成了virtualenv这样的工具用来隔离不同的依赖环境，但这种解决方案仍然没办法统一地分配计算资源。现在，我们可以通过容器技术为每个用户创建一个属于他们自己的容器，并为容器分配相应的计算资源。目前市面上基于容器的深度学习平台产品已经有很多了，比如超益集伦的AiMax。这款产品本身集成了非常多的功能，但如果你只是需要在容器内调用一下GPU，可以参考下面的步骤。
使用 Docker Client 调用 GPU
依赖安装
docker run --gpu 命令依赖于 nvidia Linux 驱动和 nvidia container toolkit，如果你想查看安装文档请点击这里，本节的下文只是安装文档的翻译和提示。
在Linux服务器上安装nvidia驱动非常简单，如果你安装了图形化界面的话直接在Ubuntu的“附加驱动”应用中安装即可，在nvidia官网上也可以下载驱动。
接下来就是安装nvidia container toolkit，我们的服务器需要满足一些先决条件：
-
GNU/Linux x86_64 内核版本 > 3.10
-
Docker >= 19.03 （注意不是Docker Desktop，如果你想在自己的台式机上使用toolkit，请安装Docker Engine而不是Docker Desktop，因为Desktop版本都是运行在虚拟机之上的）
-
NVIDIA GPU 架构 >= Kepler （目前RTX20系显卡是图灵架构，RTX30系显卡是安培架构）
-
NVIDIA Linux drivers >= 418.81.07
然后就可以正式地在Ubuntu或者Debian上安装NVIDIA Container Toolkit，如果你想在 CentOS 上或者其他 Linux 发行版上安装，请参考官方的安装文档。
安装 Docker
$ curl https://get.docker.com | sh && sudo systemctl --now enable docker 当然，这里安装完成后请参考官方的安装后需要执行的一系列操作。如果安装遇到问题，请参照官方的安装文档。
安装 NVIDIA Container Toolkit¶
设置 Package Repository和GPG Key
$ distribution=$(. /etc/os-release;echo $ID$VERSION_ID) && 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/$distribution/libnvidia-container.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 请注意：如果你想安装 NVIDIA Container Toolkit 1.6.0 之前的版本，你应该使用 nvidia-docker repository 而不是上方的 libnvidia-container repositories。
如果遇到问题请直接参考安装手册
安装 nvidia-docker2 应该会自动安装libnvidia-container-toolslibnvidia-container1等依赖包，如果没有安装可以手动安装
完成前面步骤后安装 nvidia-docker2
$ sudo apt update $ sudo apt install -y nvidia-docker2 重启 Docker Daemon
$ sudo systemctl restart docker 接下来你就可以通过运行一个CUDA容器测试下安装是否正确。
docker run --rm --gpus all nvidia/cuda:11.0.3-base-ubuntu20.04 nvidia-smi Shell 中显示的应该类似于下面的输出：
+-----------------------------------------------------------------------------+ | NVIDIA-SMI 450.51.06 Driver Version: 450.51.06 CUDA Version: 11.0 | |-------------------------------+----------------------+----------------------+ | GPU Name Persistence-M| Bus-Id Disp.A | Volatile Uncorr. ECC | | Fan Temp Perf Pwr:Usage/Cap| Memory-Usage | GPU-Util Compute M. | | | | MIG M. | |===============================+======================+======================| | 0 Tesla T4 On | 00000000:00:1E.0 Off | 0 | | N/A 34C P8 9W / 70W | 0MiB / 15109MiB | 0% Default | | | | N/A | +-------------------------------+----------------------+----------------------+ +-----------------------------------------------------------------------------+ | Processes: | | GPU GI CI PID Type Process name GPU Memory | | ID ID Usage | |=============================================================================| | No running processes found | +-----------------------------------------------------------------------------+ --gpus 用法
注意，如果你安装的是 nvidia-docker2 的话，它在安装时就已经在 Docker 中注册了 NVIDIA Runtime。如果你安装的是 nvidia-docker ，请根据官方文档向Docker注册运行时。
如果你有任何疑问，请移步本节参考的文档
可以使用以 Docker 开头的选项或使用环境变量将 GPU 指定给 Docker CLI。此变量控制在容器内可访问哪些 GPU。
--gpusNVIDIA_VISIBLE_DEVICES
| 可能的值 | 描述 |
|---|---|
0,1,2 或者 GPU-fef8089b |
逗号分割的GPU UUID(s) 或者 GPU 索引 |
all |
所有GPU都可被容器访问，默认值 |
none |
不可访问GPU，但可以使用驱动提供的功能 |
void或者 empty 或者 unset |
nvidia-container-runtime will have the same behavior as (i.e. neither GPUs nor capabilities are exposed)runc |
使用该选项指定 GPU 时，应使用该参数。参数的格式应封装在单引号中，后跟要枚举到容器的设备的双引号。例如：将 GPU 2 和 3 枚举到容器。
--gpus '"device=2,3"'使用 NVIDIA_VISIBLE_DEVICES 变量时，可能需要设置
--runtime nvidia除非已设置为默认值。
-
设置一个启用CUDA支持的容器
$ docker run --rm --gpus all nvidia/cuda nvidia-smi -
指定 nvidia 作为运行时，并指定变量
NVIDIA_VISIBLE_DEVICES$ docker run --rm --runtime=nvidia -e NVIDIA_VISIBLE_DEVICES=all nvidia/cuda nvidia-smi -
为启动的容器分配2个GPU
$ docker run --rm --gpus 2 nvidia/cuda nvidia-smi -
为容器指定使用索引为1和2的GPU
$ docker run --gpus '"device=1,2"' nvidia/cuda nvidia-smi --query-gpu=uuid --format=csvuuid GPU-ad2367dd-a40e-6b86-6fc3-c44a2cc92c7e GPU-16a23983-e73e-0945-2095-cdeb50696982 -
也可以使用
NVIDIA_VISIBLE_DEVICES$ docker run --rm --runtime=nvidia -e NVIDIA_VISIBLE_DEVICES=1,2 nvidia/cuda nvidia-smi --query-gpu=uuid --format=csvuuid GPU-ad2367dd-a40e-6b86-6fc3-c44a2cc92c7e GPU-16a23983-e73e-0945-2095-cdeb50696982 -
使用
nvidia-smi查询 GPU UUID 然后将其指定给容器$ nvidia-smi -i 3 --query-gpu=uuid --format=csvuuid GPU-18a3e86f-4c0e-cd9f-59c3-55488c4b0c24docker run --gpus device=GPU-18a3e86f-4c0e-cd9f-59c3-55488c4b0c24 nvidia/cuda nvidia-smi
关于在容器内使用驱动程序的功能的设置，以及其他设置请参阅这里。
使用 Docker Go SDK 为容器分配 GPU
使用 NVIDIA/go-nvml 获取 GPU 信息
NVIDIA/go-nvml 提供NVIDIA Management Library API (NVML) 的Go语言绑定。目前仅支持Linux，仓库地址。
下面的演示代码获取了 GPU 的各种信息，其他功能请参考 NVML 和 go-nvml 的官方文档。
package main import ( "fmt" "github.com/NVIDIA/go-nvml/pkg/nvml" "log" ) func main() { ret := nvml.Init() if ret != nvml.SUCCESS { log.Fatalf("Unable to initialize NVML: %v", nvml.ErrorString(ret)) } defer func() { ret := nvml.Shutdown() if ret != nvml.SUCCESS { log.Fatalf("Unable to shutdown NVML: %v", nvml.ErrorString(ret)) } }() count, ret := nvml.DeviceGetCount() if ret != nvml.SUCCESS { log.Fatalf("Unable to get device count: %v", nvml.ErrorString(ret)) } for i := 0; i < count; i++ { device, ret := nvml.DeviceGetHandleByIndex(i) if ret != nvml.SUCCESS { log.Fatalf("Unable to get device at index %d: %v", i, nvml.ErrorString(ret)) } // 获取 UUID uuid, ret := device.GetUUID() if ret != nvml.SUCCESS { log.Fatalf("Unable to get uuid of device at index %d: %v", i, nvml.ErrorString(ret)) } fmt.Printf("GPU UUID: %vn", uuid) name, ret := device.GetName() if ret != nvml.SUCCESS { log.Fatalf("Unable to get name of device at index %d: %v", i, nvml.ErrorString(ret)) } fmt.Printf("GPU Name: %+vn", name) memoryInfo, _ := device.GetMemoryInfo() fmt.Printf("Memory Info: %+vn", memoryInfo) powerUsage, _ := device.GetPowerUsage() fmt.Printf("Power Usage: %+vn", powerUsage) powerState, _ := device.GetPowerState() fmt.Printf("Power State: %+vn", powerState) managementDefaultLimit, _ := device.GetPowerManagementDefaultLimit() fmt.Printf("Power Managment Default Limit: %+vn", managementDefaultLimit) version, _ := device.GetInforomImageVersion() fmt.Printf("Info Image Version: %+vn", version) driverVersion, _ := nvml.SystemGetDriverVersion() fmt.Printf("Driver Version: %+vn", driverVersion) cudaDriverVersion, _ := nvml.SystemGetCudaDriverVersion() fmt.Printf("CUDA Driver Version: %+vn", cudaDriverVersion) computeRunningProcesses, _ := device.GetGraphicsRunningProcesses() for _, proc := range computeRunningProcesses { fmt.Printf("Proc: %+vn", proc) } } fmt.Println() } 使用 Docker Go SDK 为容器分配 GPU
首先需要用的的是 ContainerCreate API
// ContainerCreate creates a new container based in the given configuration. // It can be associated with a name, but it's not mandatory. func (cli *Client) ContainerCreate( ctx context.Context, config *container.Config, hostConfig *container.HostConfig, networkingConfig *network.NetworkingConfig, platform *specs.Platform, containerName string) (container.ContainerCreateCreatedBody, error) 这个 API 中需要很多用来指定配置的 struct， 其中用来请求 GPU 设备的是 container.HostConfig 这个 struct 中的 Resources ，它的类型是 container.Resources ，而在它的里面保存的是 container.DeviceRequest 这个结构体的切片，这个变量会被 GPU 设备的驱动使用。
cli.ContainerCreate API 需要 ---------> container.HostConfig{ Resources: container.Resources{ DeviceRequests: []container.DeviceRequest { { Driver: "nvidia", Count: 0, DeviceIDs: []string{"0"}, Capabilities: [][]string{{"gpu"}}, Options: nil, } } } } 下面是 container.DeviceRequest 结构体的定义
// DeviceRequest represents a request for devices from a device driver. // Used by GPU device drivers. type DeviceRequest struct { Driver string // 设备驱动名称 这里就填写 "nvidia" 即可 Count int // 请求设备的数量 (-1 = All) DeviceIDs []string // 可被设备驱动识别的设备ID列表，可以是索引也可以是UUID Capabilities [][]string // An OR list of AND lists of device capabilities (e.g. "gpu") Options map[string]string // Options to pass onto the device driver } 注意：如果指定了 Count 字段，就无法通过 DeviceIDs 指定 GPU，它们是互斥的。
接下来我们尝试使用 Docker Go SDK 启动一个 pytorch 容器。
首先我们编写一个 test.py 文件，让它在容器内运行，检查 CUDA 是否可用。
# test.py import torch print("cuda.is_available:", torch.cuda.is_available()) 下面是实验代码，启动一个名为 torch_test_1 的容器，并运行 python3 /workspace/test.py 命令，然后从 stdout 和 stderr 获取输出。
package main import ( "context" "fmt" "github.com/docker/docker/api/types" "github.com/docker/docker/api/types/container" "github.com/docker/docker/client" "github.com/docker/docker/pkg/stdcopy" "os" ) var ( defaultHost = "unix:///var/run/docker.sock" ) func main() { ctx := context.Background() cli, err := client.NewClientWithOpts(client.WithHost(defaultHost), client.WithAPIVersionNegotiation()) if err != nil { panic(err) } resp, err := cli.ContainerCreate(ctx, &container.Config{ Image: "pytorch/pytorch", Cmd: []string{}, OpenStdin: true, Volumes: map[string]struct{}{}, Tty: true, }, &container.HostConfig{ Binds: []string{`/home/joseph/workspace:/workspace`}, Resources: container.Resources{DeviceRequests: []container.DeviceRequest{{ Driver: "nvidia", Count: 0, DeviceIDs: []string{"0"}, Capabilities: [][]string{{"gpu"}}, Options: nil, }}}, }, nil, nil, "torch_test_1") if err != nil { panic(err) } if err := cli.ContainerStart(ctx, resp.ID, types.ContainerStartOptions{}); err != nil { panic(err) } fmt.Println(resp.ID) execConf := types.ExecConfig{ User: "", Privileged: false, Tty: false, AttachStdin: false, AttachStderr: true, AttachStdout: true, Detach: true, DetachKeys: "ctrl-p,q", Env: nil, WorkingDir: "/", Cmd: []string{"python3", "/workspace/test.py"}, } execCreate, err := cli.ContainerExecCreate(ctx, resp.ID, execConf) if err != nil { panic(err) } response, err := cli.ContainerExecAttach(ctx, execCreate.ID, types.ExecStartCheck{}) defer response.Close() if err != nil { fmt.Println(err) } // read the output _, _ = stdcopy.StdCopy(os.Stdout, os.Stderr, response.Reader) } 可以看到，程序输出了创建的容器的 Contrainer ID 和 执行命令的输出。
$ go build main.go $ sudo ./main 264535c7086391eab1d74ea48094f149ecda6d25709ac0c6c55c7693c349967b cuda.is_available: True 接下来使用 docker ps 查看容器状态。
$ docker ps CONTAINER ID IMAGE COMMAND CREATED STATUS PORTS NAMES 264535c70863 pytorch/pytorch "bash" 2 minutes ago Up 2 minutes torch_test_1 没问题，Container ID 对得上。
