Detect GPU

起因

在开发 gpu-docker-api时，需要先获取主机的 GPU 信息，然后使用 UUID 给容器分配 GPU 。

从网上找到了 go-nvml 能够很简单的获取 GPU 信息，输出结果类似于 nvdia-smi，问题是它必须跑在装好 NVIDIA 驱动的 Linux 服务器上，最重要的是当我们开发的项目引入它后，都不能运行。类似的报错如下：

# github.com/NVIDIA/go-nvml/pkg/dl
/Users/ming/go/pkg/mod/github.com/!n!v!i!d!i!a/go-nvml@v0.12.0-1/pkg/dl/dl.go:34:18: could not determine kind of name for C.RTLD_DEEPBIND

后来想到在本机启动一个 Docker 容器然后编译也是比较方便，后来折腾了半天，没有成功最后放弃了。

于是想到把这个功能抽离出来，同时提供一个 HTTP 接口，通过调用的方式来获取 GPU 信息。当然这个 HTTP 服务要跑在带有 NVDIA 驱动的 Liunx 驱动上。

这样我们就能在 macOS 或 Windows 下开发项目了。

笔者平时使用 golang 进行开发，工作中经常遇到，有时只需要简单的两三个接口，同事都要引入 GIN 框架、ZAP 日志框架等等，一个非常简单的项目导致依赖非常多，同时也没有 Makefile 、测试用例，构建、启动只能问同事，然后复制粘贴。

当然我没有鄙视这种行为，能解决问题的方法就是好方法。只不过在空闲时间，研究研究也是挺好的。所以贯彻 Golang 的 Less is more 理念，尽量引入第三方库，尽量使用原生方法来写本项目。

项目中使用的库都是我感觉非常简洁、好用的，大家可以进行参考。

项目地址

detect-gpu

使用

可以从 release 下载二进制文件，扔到服务器上运行。或者克隆到本地，然后构建。

git clone https://github.com/mayooot/detect-gpu
cd detect-gpu
make linux

默认程序会占用 2376 端口，api 地址为 /api/v1/detect/gpu 。

$ curl 127.0.0.1:2376/api/v1/detect/gpu
[
    {
        "index":0,
        "uuid":"uuid",
        "name":"NVIDIA A100 80GB PCIe",
        "memoryInfo":{
            "Total":85899345920,
            "Free":63216877568,
            "Used":22682468352
        },
        "powerUsage":74634,
        "powerState":0,
        "powerManagementDefaultLimit":300000,
        "informImageVersion":"1001.0230.00.03",
        "systemGetDriverVersion":"525.85.12",
        "systemGetCudaDriverVersion":12000,
        "tGraphicsRunningProcesses":[]
    },
    {
        "index":1,
        "uuid":"uuid",
        "name":"NVIDIA A100 80GB PCIe",
        "memoryInfo":{
            "Total":85899345920,
            "Free":30687952896,
            "Used":55211393024
        },
        "powerUsage":65507,
        "powerState":0,
        "powerManagementDefaultLimit":300000,
        "informImageVersion":"1001.0230.00.03",
        "systemGetDriverVersion":"525.85.12",
        "systemGetCudaDriverVersion":12000,
        "tGraphicsRunningProcesses":[]
    }
]

当然也可以在 golang 项目中直接引用。如下：

package main

import (
	"fmt"
	"time"

	"github.com/mayooot/detect-gpu/pkg/detect"
)

func main() {
	timeOutDuration := 500 * time.Millisecond

	testClient := detect.NewClient(detect.WithTimeout(timeOutDuration))
	if err := testClient.Init(); err != nil {
		panic(err)
	}
	defer testClient.Close()

	gpus, err := testClient.DetectGpu()
	if err != nil {
		panic(err)
	}
	for _, gpu := range gpus {
		fmt.Printf("%#+v\n", gpu)
	}
}

个人感觉好用的库

go-svc

一般启动一个 web 服务时，需要做好初始化工作，比如初始化数据库、Redis 。然后异步启动并阻塞，最后优雅关闭，释放资源。

常用的写法就是声明一个监听信号量的 channel ，然后 select 等待。go-svc 进行封装，能让代码看起来更加简洁。

ngaut/log

这个库是我开发一些简单的项目时最喜欢的日志库了，它只对 GO 自带的 log 库进行了简单的封装，不同的日志级别有不同的颜色。

比如该项目启动和结束时，打印的日志：

pflag

非常棒的命令行解析库，支持全拼参数和简写，使用起来只能说太爽了。

最后

如果对你有用的话，你可以把该项目当成一个快速开发的模板进行参考。同时有任何的 bug/意见，欢迎你提 issue ，我很乐意解答。

调用接口的方式获取 NVDIA GPU 🤥