使用 Go SDK 在 Docker 中传递 `--gpus all` 选项的方法

目前golang学习网上已经有很多关于Golang的文章了，自己在初次阅读这些文章中，也见识到了很多学习思路；那么本文《使用 Go SDK 在 Docker 中传递 `--gpus all` 选项的方法》，也希望能帮助到大家，如果阅读完后真的对你学习Golang有帮助，欢迎动动手指，评论留言并分享~

我已经从 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)
}

正确答案

参见：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, "")

到这里，我们也就讲完了《使用 Go SDK 在 Docker 中传递 `--gpus all` 选项的方法》的内容了。个人认为，基础知识的学习和巩固，是为了更好的将其运用到项目中，欢迎关注golang学习网公众号，带你了解更多关于的知识点！