T O P

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

  • By - 楼主

  • 2022-08-01 12:02:04
  • 背景

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

    • --gpus
    • NVIDIA_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除非已设置为默认值。

    1. 设置一个启用CUDA支持的容器

      $ docker run --rm --gpus all nvidia/cuda nvidia-smi
      
    2. 指定 nvidia 作为运行时,并指定变量 NVIDIA_VISIBLE_DEVICES

      $ docker run --rm --runtime=nvidia \
          -e NVIDIA_VISIBLE_DEVICES=all nvidia/cuda nvidia-smi
      
    3. 为启动的容器分配2个GPU

      $ docker run --rm --gpus 2 nvidia/cuda nvidia-smi
      
    4. 为容器指定使用索引为1和2的GPU

      $ docker run --gpus '"device=1,2"' \
      	nvidia/cuda nvidia-smi --query-gpu=uuid --format=csv
      
      uuid
      GPU-ad2367dd-a40e-6b86-6fc3-c44a2cc92c7e
      GPU-16a23983-e73e-0945-2095-cdeb50696982
      
    5. 也可以使用 NVIDIA_VISIBLE_DEVICES

      $ docker run --rm --runtime=nvidia \
      	-e NVIDIA_VISIBLE_DEVICES=1,2 \
      	nvidia/cuda nvidia-smi --query-gpu=uuid --format=csv
      
      uuid
      GPU-ad2367dd-a40e-6b86-6fc3-c44a2cc92c7e
      GPU-16a23983-e73e-0945-2095-cdeb50696982
      
    6. 使用 nvidia-smi 查询 GPU UUID 然后将其指定给容器

      $ nvidia-smi -i 3 --query-gpu=uuid --format=csv
      
      uuid
      GPU-18a3e86f-4c0e-cd9f-59c3-55488c4b0c24
      
      docker 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: %v\n", 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: %+v\n", name)
    
    		memoryInfo, _ := device.GetMemoryInfo()
    		fmt.Printf("Memory Info: %+v\n", memoryInfo)
    
    		powerUsage, _ := device.GetPowerUsage()
    		fmt.Printf("Power Usage: %+v\n", powerUsage)
    
    		powerState, _ := device.GetPowerState()
    		fmt.Printf("Power State: %+v\n", powerState)
    
    		managementDefaultLimit, _ := device.GetPowerManagementDefaultLimit()
    		fmt.Printf("Power Managment Default Limit: %+v\n", managementDefaultLimit)
    
    		version, _ := device.GetInforomImageVersion()
    		fmt.Printf("Info Image Version: %+v\n", version)
    
    		driverVersion, _ := nvml.SystemGetDriverVersion()
    		fmt.Printf("Driver Version: %+v\n", driverVersion)
    
    		cudaDriverVersion, _ := nvml.SystemGetCudaDriverVersion()
    		fmt.Printf("CUDA Driver Version: %+v\n", cudaDriverVersion)
    
    		computeRunningProcesses, _ := device.GetGraphicsRunningProcesses()
    		for _, proc := range computeRunningProcesses {
    			fmt.Printf("Proc: %+v\n", 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 命令,然后从 stdoutstderr 获取输出。

    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 对得上。

    本帖子中包含资源

    您需要 登录 才可以下载,没有帐号?立即注册