【Docker】使用Docker Client和Docker Go SDK为容器分配GPU资源 ...

目录

• 背景
  
• 使用 Docker Client 调用 GPU
  
  
  
  • 依赖安装
    
    
    
    • 安装 Docker
      
    • 安装 NVIDIA Container Toolkit¶
      
    
    
  • --gpus 用法
    
  
  
• 使用 Docker Go SDK 为容器分配 GPU
  
  
  
  • 使用 NVIDIA/go-nvml 获取 GPU 信息
    
  • 使用 Docker Go SDK 为容器分配 GPU
    
  
  

背景

深度学习的环境配置通常是一项比较麻烦的工作，尤其是在多个用户共享的服务器上。虽然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

1. $ curl https://get.docker.com | sh \
2.   && sudo systemctl --now enable docker

复制代码

当然，这里安装完成后请参考官方的安装后需要执行的一系列操作。如果安装遇到问题，请参照官方的安装文档。
安装 NVIDIA Container Toolkit¶

设置 Package Repository和GPG Key

1. $ distribution=$(. /etc/os-release;echo $ID$VERSION_ID) \
2.       && curl -fsSL https://nvidia.github.io/libnvidia-container/gpgkey | sudo gpg --dearmor -o /usr/share/keyrings/nvidia-container-toolkit-keyring.gpg \
3.       && curl -s -L https://nvidia.github.io/libnvidia-container/$distribution/libnvidia-container.list | \
4.             sed 's#deb https://#deb [signed-by=/usr/share/keyrings/nvidia-container-toolkit-keyring.gpg] https://#g' | \
5.             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-tools libnvidia-container1 等依赖包，如果没有安装可以手动安装

完成前面步骤后安装 nvidia-docker2

1. $ sudo apt update

复制代码

1. $ sudo apt install -y nvidia-docker2

复制代码

重启 Docker Daemon

1. $ sudo systemctl restart docker

复制代码

接下来你就可以通过运行一个CUDA容器测试下安装是否正确。

1. docker run --rm --gpus all nvidia/cuda:11.0.3-base-ubuntu20.04 nvidia-smi

复制代码

Shell 中显示的应该类似于下面的输出：

1. +-----------------------------------------------------------------------------+
2. | NVIDIA-SMI 450.51.06    Driver Version: 450.51.06    CUDA Version: 11.0     |
3. |-------------------------------+----------------------+----------------------+
4. | GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |
5. | Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
6. |                               |                      |               MIG M. |
7. |===============================+======================+======================|
8. |   0  Tesla T4            On   | 00000000:00:1E.0 Off |                    0 |
9. | N/A   34C    P8     9W /  70W |      0MiB / 15109MiB |      0%      Default |
10. |                               |                      |                  N/A |
11. +-------------------------------+----------------------+----------------------+
13. +-----------------------------------------------------------------------------+
14. | Processes:                                                                  |
15. |  GPU   GI   CI        PID   Type   Process name                  GPU Memory |
16. |        ID   ID                                                   Usage      |
17. |=============================================================================|
18. |  No running processes found                                                 |
19. +-----------------------------------------------------------------------------+

复制代码

--gpus 用法

注意，如果你安装的是 nvidia-docker2 的话，它在安装时就已经在 Docker 中注册了 NVIDIA Runtime。如果你安装的是 nvidia-docker ，请根据官方文档向Docker注册运行时。
如果你有任何疑问，请移步本节参考的文档

可以使用以 Docker 开头的选项或使用环境变量将 GPU 指定给 Docker CLI。此变量控制在容器内可访问哪些 GPU。

• --gpus
  
• NVIDIA_VISIBLE_DEVICES
  

可能的值描述0,1,2 或者 GPU-fef8089b逗号分割的GPU UUID(s) 或者 GPU 索引all所有GPU都可被容器访问，默认值none不可访问GPU，但可以使用驱动提供的功能void或者 empty 或者 unsetnvidia-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支持的容器
  1. $ docker run --rm --gpus all nvidia/cuda nvidia-smi

  复制代码
• 指定 nvidia 作为运行时，并指定变量 NVIDIA_VISIBLE_DEVICES
  1. $ docker run --rm --runtime=nvidia \
  2.     -e NVIDIA_VISIBLE_DEVICES=all nvidia/cuda nvidia-smi

  复制代码
• 为启动的容器分配2个GPU
  1. $ docker run --rm --gpus 2 nvidia/cuda nvidia-smi

  复制代码
• 为容器指定使用索引为1和2的GPU
  1. $ docker run --gpus '"device=1,2"' \
  2.         nvidia/cuda nvidia-smi --query-gpu=uuid --format=csv

  复制代码
  1. uuid
  2. GPU-ad2367dd-a40e-6b86-6fc3-c44a2cc92c7e
  3. GPU-16a23983-e73e-0945-2095-cdeb50696982

  复制代码
• 也可以使用 NVIDIA_VISIBLE_DEVICES
  1. $ docker run --rm --runtime=nvidia \
  2.         -e NVIDIA_VISIBLE_DEVICES=1,2 \
  3.         nvidia/cuda nvidia-smi --query-gpu=uuid --format=csv

  复制代码
  1. uuid
  2. GPU-ad2367dd-a40e-6b86-6fc3-c44a2cc92c7e
  3. GPU-16a23983-e73e-0945-2095-cdeb50696982

  复制代码
• 使用 nvidia-smi 查询 GPU UUID 然后将其指定给容器
  1. $ nvidia-smi -i 3 --query-gpu=uuid --format=csv

  复制代码
  1. uuid
  2. GPU-18a3e86f-4c0e-cd9f-59c3-55488c4b0c24

  复制代码
  1. docker run --gpus device=GPU-18a3e86f-4c0e-cd9f-59c3-55488c4b0c24 \
  2.         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 的官方文档。

1. package main
3. import (
4.         "fmt"
5.         "github.com/NVIDIA/go-nvml/pkg/nvml"
6.         "log"
7. )
9. func main() {
10.         ret := nvml.Init()
11.         if ret != nvml.SUCCESS {
12.                 log.Fatalf("Unable to initialize NVML: %v", nvml.ErrorString(ret))
13.         }
14.         defer func() {
15.                 ret := nvml.Shutdown()
16.                 if ret != nvml.SUCCESS {
17.                         log.Fatalf("Unable to shutdown NVML: %v", nvml.ErrorString(ret))
18.                 }
19.         }()
21.         count, ret := nvml.DeviceGetCount()
22.         if ret != nvml.SUCCESS {
23.                 log.Fatalf("Unable to get device count: %v", nvml.ErrorString(ret))
24.         }
26.         for i := 0; i < count; i++ {
27.                 device, ret := nvml.DeviceGetHandleByIndex(i)
28.                 if ret != nvml.SUCCESS {
29.                         log.Fatalf("Unable to get device at index %d: %v", i, nvml.ErrorString(ret))
30.                 }
31.                
32.                 // 获取 UUID
33.                 uuid, ret := device.GetUUID()
34.                 if ret != nvml.SUCCESS {
35.                         log.Fatalf("Unable to get uuid of device at index %d: %v", i, nvml.ErrorString(ret))
36.                 }
37.                 fmt.Printf("GPU UUID: %v\n", uuid)
39.                 name, ret := device.GetName()
40.                 if ret != nvml.SUCCESS {
41.                         log.Fatalf("Unable to get name of device at index %d: %v", i, nvml.ErrorString(ret))
42.                 }
43.                 fmt.Printf("GPU Name: %+v\n", name)
45.                 memoryInfo, _ := device.GetMemoryInfo()
46.                 fmt.Printf("Memory Info: %+v\n", memoryInfo)
48.                 powerUsage, _ := device.GetPowerUsage()
49.                 fmt.Printf("Power Usage: %+v\n", powerUsage)
51.                 powerState, _ := device.GetPowerState()
52.                 fmt.Printf("Power State: %+v\n", powerState)
54.                 managementDefaultLimit, _ := device.GetPowerManagementDefaultLimit()
55.                 fmt.Printf("Power Managment Default Limit: %+v\n", managementDefaultLimit)
57.                 version, _ := device.GetInforomImageVersion()
58.                 fmt.Printf("Info Image Version: %+v\n", version)
60.                 driverVersion, _ := nvml.SystemGetDriverVersion()
61.                 fmt.Printf("Driver Version: %+v\n", driverVersion)
63.                 cudaDriverVersion, _ := nvml.SystemGetCudaDriverVersion()
64.                 fmt.Printf("CUDA Driver Version: %+v\n", cudaDriverVersion)
66.                 computeRunningProcesses, _ := device.GetGraphicsRunningProcesses()
67.                 for _, proc := range computeRunningProcesses {
68.                         fmt.Printf("Proc: %+v\n", proc)
69.                 }
70.         }
72.         fmt.Println()
73. }

复制代码

使用 Docker Go SDK 为容器分配 GPU

首先需要用的的是 ContainerCreate API

1. // ContainerCreate creates a new container based in the given configuration.
2. // It can be associated with a name, but it's not mandatory.
3. func (cli *Client) ContainerCreate(
4.         ctx context.Context,
5.         config *container.Config,
6.         hostConfig *container.HostConfig,
7.         networkingConfig *network.NetworkingConfig,
8.         platform *specs.Platform,
9.         containerName string) (container.ContainerCreateCreatedBody, error)

复制代码

这个 API 中需要很多用来指定配置的 struct， 其中用来请求 GPU 设备的是 container.HostConfig 这个 struct 中的 Resources ，它的类型是 container.Resources ，而在它的里面保存的是 container.DeviceRequest 这个结构体的切片，这个变量会被 GPU 设备的驱动使用。

1. cli.ContainerCreate API  需要 ---------> container.HostConfig{
2.                                                 Resources: container.Resources{
3.                                                         DeviceRequests: []container.DeviceRequest {
4.                                                                 {
5.                                                                         Driver:       "nvidia",
6.                                                                         Count:        0,
7.                                                                         DeviceIDs:    []string{"0"},
8.                                                                         Capabilities: [][]string{{"gpu"}},
9.                                                                         Options:      nil,
10.                                                                 }
11.                                                         }
12.                                                 }
13.                                         }

复制代码

下面是 container.DeviceRequest 结构体的定义

1. // DeviceRequest represents a request for devices from a device driver.
2. // Used by GPU device drivers.
3. type DeviceRequest struct {
4.         Driver       string            // 设备驱动名称 这里就填写 "nvidia" 即可
5.         Count        int               // 请求设备的数量 (-1 = All)
6.         DeviceIDs    []string          // 可被设备驱动识别的设备ID列表，可以是索引也可以是UUID
7.         Capabilities [][]string        // An OR list of AND lists of device capabilities (e.g. "gpu")
8.         Options      map[string]string // Options to pass onto the device driver
9. }

复制代码

注意：如果指定了 Count 字段，就无法通过 DeviceIDs 指定 GPU，它们是互斥的。
接下来我们尝试使用 Docker Go SDK 启动一个 pytorch 容器。
首先我们编写一个 test.py 文件，让它在容器内运行，检查 CUDA 是否可用。

1. # test.py
2. import torch
4. print("cuda.is_available:", torch.cuda.is_available())

复制代码

下面是实验代码，启动一个名为 torch_test_1 的容器，并运行 python3 /workspace/test.py 命令，然后从 stdout 和 stderr 获取输出。

1. package main
3. import (
4.         "context"
5.         "fmt"
6.         "github.com/docker/docker/api/types"
7.         "github.com/docker/docker/api/types/container"
8.         "github.com/docker/docker/client"
9.         "github.com/docker/docker/pkg/stdcopy"
10.         "os"
11. )
13. var (
14.         defaultHost = "unix:///var/run/docker.sock"
15. )
17. func main() {
18.         ctx := context.Background()
19.         cli, err := client.NewClientWithOpts(client.WithHost(defaultHost), client.WithAPIVersionNegotiation())
20.         if err != nil {
21.                 panic(err)
22.         }
24.         resp, err := cli.ContainerCreate(ctx,
25.                 &container.Config{
26.                         Image:     "pytorch/pytorch",
27.                         Cmd:       []string{},
28.                         OpenStdin: true,
29.                         Volumes:   map[string]struct{}{},
30.                         Tty:       true,
31.                 }, &container.HostConfig{
32.                         Binds: []string{`/home/joseph/workspace:/workspace`},
33.                         Resources: container.Resources{DeviceRequests: []container.DeviceRequest{{
34.                                 Driver:       "nvidia",
35.                                 Count:        0,
36.                                 DeviceIDs:    []string{"0"},
37.                                 Capabilities: [][]string{{"gpu"}},
38.                                 Options:      nil,
39.                         }}},
40.                 }, nil, nil, "torch_test_1")
41.         if err != nil {
42.                 panic(err)
43.         }
45.         if err := cli.ContainerStart(ctx, resp.ID, types.ContainerStartOptions{}); err != nil {
46.                 panic(err)
47.         }
49.         fmt.Println(resp.ID)
51.         execConf := types.ExecConfig{
52.                 User:         "",
53.                 Privileged:   false,
54.                 Tty:          false,
55.                 AttachStdin:  false,
56.                 AttachStderr: true,
57.                 AttachStdout: true,
58.                 Detach:       true,
59.                 DetachKeys:   "ctrl-p,q",
60.                 Env:          nil,
61.                 WorkingDir:   "/",
62.                 Cmd:          []string{"python3", "/workspace/test.py"},
63.         }
64.         execCreate, err := cli.ContainerExecCreate(ctx, resp.ID, execConf)
65.         if err != nil {
66.                 panic(err)
67.         }
69.         response, err := cli.ContainerExecAttach(ctx, execCreate.ID, types.ExecStartCheck{})
70.         defer response.Close()
71.         if err != nil {
72.                 fmt.Println(err)
73.         }
75.         // read the output
76.         _, _ = stdcopy.StdCopy(os.Stdout, os.Stderr, response.Reader)
77. }

复制代码

可以看到，程序输出了创建的容器的 Contrainer ID 和 执行命令的输出。

1. $ go build main.go
2. $ sudo ./main
3. 264535c7086391eab1d74ea48094f149ecda6d25709ac0c6c55c7693c349967b
4. cuda.is_available: True

复制代码

接下来使用 docker ps 查看容器状态。

1. $ docker ps
2. CONTAINER ID   IMAGE             COMMAND   CREATED         STATUS             PORTS     NAMES
3. 264535c70863   pytorch/pytorch   "bash"    2 minutes ago   Up 2 minutes                 torch_test_1

复制代码

没问题，Container ID 对得上。

免责声明：如果侵犯了您的权益，请联系站长，我们会及时删除侵权内容，谢谢合作！