如何使用 Go SDK 将 `--gpus all` 选项传递给 Docker？

Question

我已经从 SDK 示例中了解了如何执行一些基本命令，例如运行容器、拉取图像、列出图像等。

我正在开发一个项目，需要在容器内使用 GPU。

我的系统有GPU，我已经安装了驱动程序，并且我还安装了

nvidia-container-runtime

。

如果我们暂时从场景中删除 Go SDK，我可以运行以下命令来获取主机系统上的

nvidia-smi

输出：

docker run -it --rm --gpus all nvidia/cuda:10.0-base nvidia-smi

我必须通过 SDK 来完成此操作。这是开始的代码。此代码打印“hello world”。但实际上我会在那个地方运行

nvidia-smi

命令：

package main

import (
    "context"
    "os"

    "github.com/docker/docker/api/types"
    "github.com/docker/docker/api/types/container"
    "github.com/docker/docker/client"
    "github.com/docker/docker/pkg/stdcopy"
)

func main() {
    ctx := context.Background()
    cli, err := client.NewClientWithOpts(client.FromEnv, client.WithAPIVersionNegotiation())
    if err != nil {
        panic(err)
    }

    RunContainer(ctx, cli)
}

func RunContainer(ctx context.Context, cli *client.Client) {
    reader, err := cli.ImagePull(ctx, "nvidia/cuda:10.0-base", types.ImagePullOptions{})
    if err != nil {
        panic(err)
    }

    defer reader.Close()
    // io.Copy(os.Stdout, reader)

    resp, err := cli.ContainerCreate(ctx, &container.Config{
        Image: "nvidia/cuda:10.0-base",
        Cmd:   []string{"echo", "hello world"},
        // Tty:   false,
    }, nil, nil, nil, "")

    if err != nil {
        panic(err)
    }

    if err := cli.ContainerStart(ctx, resp.ID, types.ContainerStartOptions{}); err != nil {
        panic(err)
    }

    statusCh, errCh := cli.ContainerWait(ctx, resp.ID, container.WaitConditionNotRunning)

    select {
    case err := <-errCh:
        if err != nil {
            panic(err)
        }
    case <-statusCh:
    }

    out, err := cli.ContainerLogs(ctx, resp.ID, types.ContainerLogsOptions{ShowStdout: true})
    if err != nil {
        panic(err)
    }

    stdcopy.StdCopy(os.Stdout, os.Stderr, out)
}

Answer 1

参见：https://github.com/docker/cli/blob/9ac8584acfd501c3f4da0e845e3a40ed15c85041/cli/command/container/opts.go#L594

import "github.com/docker/cli/opts"

// ...

gpuOpts := opts.GpuOpts{}
gpuOpts.Set("all")

resp, err := cli.ContainerCreate(ctx, &container.Config{
    Image: "nvidia/cuda:10.0-base",
    Cmd:   []string{"echo", "hello world"},
    // Tty:   false,
}, &container.HostConfig{Resources: container.Resources{DeviceRequests: gpuOpts.Value()}}, nil, nil, "")

Answer 2

万一其他人发现了这个并且不想使用侵入你的跑步者的 cli 选项。您可以简单地创建 docker 主机配置。我做了这个

import "github.com/docker/docker/api/types/container"

// more setup etc...
rep, err := docker.ContainerCreate(ctx, config,
        &container.HostConfig{
            Resources: container.Resources{
                DeviceRequests: []container.DeviceRequest{{
                    Driver:       "nvidia",
                    Count:        -1,
                    DeviceIDs:    nil,
                    Capabilities: [][]string{{"gpu"}},
                    Options:      make(map[string]string),
                }},
            },
        }, nil, nil, "myContainer")
// do more stuff (ContainerStart maybe)...