如何通过Go SDK为Docker容器传递--gpus all参数？

要实现对应命令行docker run -it --rm --gpus all nvidia/cuda:10.0-base nvidia-smi的功能，需要对现有代码做核心修改，以下是完整解决方案：

核心修改点 #

• 替换容器启动命令为nvidia-smi
• 配置GPU资源映射，对应--gpus all参数
• 添加容器自动清理配置，对应--rm参数

修改后的完整代码 #

package main

import (
    "context"
    "os"

    "github.com/docker/docker/api/types"
    "github.com/docker/docker/api/types/container"
    "github.com/docker/docker/api/types/hostconfig"
    "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)

    // 配置GPU资源，对应命令行的--gpus all
    hostConfig := &hostconfig.HostConfig{
        Resources: container.Resources{
            DeviceRequests: []container.DeviceRequest{
                {
                    Count:        -1, // -1表示使用所有可用GPU
                    Capabilities: [][]string{{"gpu"}},
                    Driver:       "nvidia",
                },
            },
        },
        AutoRemove: true, // 容器退出后自动删除，对应--rm
    }

    resp, err := cli.ContainerCreate(ctx, &container.Config{
        Image: "nvidia/cuda:10.0-base",
        Cmd:   []string{"nvidia-smi"}, // 替换为目标命令
    }, hostConfig, 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, ShowStderr: true})
    if err != nil {
        panic(err)
    }

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

关键说明 #

1. GPU资源配置：通过DeviceRequests指定使用所有NVIDIA GPU，Count: -1等价于命令行的--gpus all
2. 自动清理：AutoRemove: true确保容器执行完成后自动销毁，避免残留无用容器
3. 日志捕获：开启ShowStderr: true可以捕获容器运行中的错误信息，方便排查问题

内容的提问来源于stack exchange，提问作者Santosh Kumar