如何通过Go SDK将`--GPU ALL`选项传递给Docker?

jv2fixgn  于 2022-09-19  发布在  Go
关注(0)|答案(1)|浏览(208)

我已经看到了如何从SDK examples执行一些基本命令,如运行容器、拉取图像、列出图像等。

我正在从事的一个项目中,我需要使用容器内的图形处理器。

我的系统有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)
}
bihw5rsg

bihw5rsg1#

请参阅: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, "")

相关问题