KJZH001/nezhahq-agent
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

gpu: support read usage of multiple cards on linux (#79)

Browse Source
This commit is contained in:
UUBulb
2024-10-24 20:43:17 +08:00
committed by GitHub
parent 20db2c9cb2
commit 0cba96bae1
18 changed files with 393 additions and 307 deletions
Download Patch File Download Diff File Expand all files Collapse all files
pkg/gpu/gpu.go
-27
View File
@@ -1,27 +0,0 @@
//go:build !darwin
// +build !darwin
package gpu
import (
"errors"
"github.com/jaypipes/ghw"
)
func GetGPUModel() ([]string, error) {
var gpuModel []string
gi, err := ghw.GPU(ghw.WithDisableWarnings())
if err != nil {
return nil, err
}
for _, card := range gi.GraphicsCards {
if card.DeviceInfo == nil {
return nil, errors.New("Cannot find device info")
}
gpuModel = append(gpuModel, card.DeviceInfo.Product.Name)
}
return gpuModel, nil
}
pkg/gpu/gpu_darwin.go
+3 -2
View File
@@ -125,8 +125,9 @@ func GetGPUModel() ([]string, error) {
return util.RemoveDuplicate(models), nil
}
func FindUtilization(key, dictKey string) (int, error) {
return findUtilization(key, dictKey)
func GetGPUStat() ([]float64, error) {
usage, err := findUtilization("PerformanceStatistics", "Device Utilization %")
return []float64{float64(usage)}, err
}
func findDevices(key string) ([]string, error) {
pkg/gpu/gpu_fallback.go
+11
View File
@@ -0,0 +1,11 @@
//go:build !darwin && !linux && !windows
package gpu
func GetGPUModel() ([]string, error) {
return nil, nil
}
func GetGPUStat() ([]float64, error) {
return nil, nil
}
pkg/gpu/gpu_linux.go
+125
View File
@@ -0,0 +1,125 @@
//go:build linux
package gpu
import (
"errors"
"github.com/nezhahq/agent/pkg/gpu/vendor"
)
const (
vendorAMD = iota + 1
vendorNVIDIA
)
var vendorType uint8
func init() {
_, err := getNvidiaStat()
if err != nil {
vendorType = vendorAMD
} else {
vendorType = vendorNVIDIA
}
}
func getNvidiaStat() ([]float64, error) {
smi := &vendor.NvidiaSMI{
BinPath: "/usr/bin/nvidia-smi",
}
err1 := smi.Start()
if err1 != nil {
return nil, err1
}
data, err2 := smi.GatherUsage()
if err2 != nil {
return nil, err2
}
return data, nil
}
func getAMDStat() ([]float64, error) {
rsmi := &vendor.ROCmSMI{
BinPath: "/opt/rocm/bin/rocm-smi",
}
err := rsmi.Start()
if err != nil {
return nil, err
}
data, err := rsmi.GatherUsage()
if err != nil {
return nil, err
}
return data, nil
}
func getNvidiaHost() ([]string, error) {
smi := &vendor.NvidiaSMI{
BinPath: "/usr/bin/nvidia-smi",
}
err := smi.Start()
if err != nil {
return nil, err
}
data, err := smi.GatherModel()
if err != nil {
return nil, err
}
return data, nil
}
func getAMDHost() ([]string, error) {
rsmi := &vendor.ROCmSMI{
BinPath: "/opt/rocm/bin/rocm-smi",
}
err := rsmi.Start()
if err != nil {
return nil, err
}
data, err := rsmi.GatherModel()
if err != nil {
return nil, err
}
return data, nil
}
func GetGPUModel() ([]string, error) {
var gi []string
var err error
switch vendorType {
case vendorAMD:
gi, err = getAMDHost()
case vendorNVIDIA:
gi, err = getNvidiaHost()
default:
return nil, errors.New("invalid vendor")
}
if err != nil {
return nil, err
}
return gi, nil
}
func GetGPUStat() ([]float64, error) {
var gs []float64
var err error
switch vendorType {
case vendorAMD:
gs, err = getAMDStat()
case vendorNVIDIA:
gs, err = getNvidiaStat()
default:
return nil, errors.New("invalid vendor")
}
if err != nil {
return nil, err
}
return gs, nil
}
pkg/gpu/stat/stat_windows.go → pkg/gpu/gpu_windows.go
+37 -22
View File
@@ -1,9 +1,6 @@
//go:build windows
// Modified from https://github.com/shirou/gopsutil/blob/master/internal/common/common_windows.go
// Original License: BSD-3-Clause
package stat
package gpu
import (
"errors"
@@ -11,6 +8,7 @@ import (
"time"
"unsafe"
"github.com/jaypipes/ghw"
"golang.org/x/sys/windows"
)
@@ -43,6 +41,41 @@ type PDH_FMT_COUNTERVALUE_ITEM_DOUBLE struct {
FmtValue PDH_FMT_COUNTERVALUE_DOUBLE
}
func GetGPUModel() ([]string, error) {
var gpuModel []string
gi, err := ghw.GPU(ghw.WithDisableWarnings())
if err != nil {
return nil, err
}
for _, card := range gi.GraphicsCards {
if card.DeviceInfo == nil {
return nil, errors.New("Cannot find device info")
}
gpuModel = append(gpuModel, card.DeviceInfo.Product.Name)
}
return gpuModel, nil
}
func GetGPUStat() ([]float64, error) {
counter, err := newWin32PerformanceCounter("gpu_utilization", "\\GPU Engine(*engtype_3D)\\Utilization Percentage")
if err != nil {
return nil, err
}
defer pdhCloseQuery.Call(uintptr(counter.Query))
values, err := getValue(8192, counter)
if err != nil {
return nil, err
}
tot := sumArray(values)
if tot > 100 {
tot = 100
}
return []float64{tot}, nil
}
// https://github.com/influxdata/telegraf/blob/master/plugins/inputs/win_perf_counters/performance_query.go
func getCounterArrayValue(initialBufSize uint32, counter *win32PerformanceCounter) ([]float64, error) {
for buflen := initialBufSize; buflen <= 100*1024*1024; buflen *= 2 {
@@ -127,24 +160,6 @@ func getValue(initialBufSize uint32, counter *win32PerformanceCounter) ([]float6
return getCounterArrayValue(initialBufSize, counter)
}
func GetGPUStat() (float64, error) {
counter, err := newWin32PerformanceCounter("gpu_utilization", "\\GPU Engine(*engtype_3D)\\Utilization Percentage")
if err != nil {
return 0, err
}
defer pdhCloseQuery.Call(uintptr(counter.Query))
values, err := getValue(8192, counter)
if err != nil {
return 0, err
}
tot := sumArray(values)
if tot > 100 {
tot = 100
}
return tot, nil
}
func sumArray(arr []float64) float64 {
var sum float64
for _, value := range arr {
pkg/gpu/stat/amd_rocm_smi.go
-67
View File
@@ -1,67 +0,0 @@
package stat
// Modified from https://github.com/influxdata/telegraf/blob/master/plugins/inputs/amd_rocm_smi/amd_rocm_smi.go
// Original License: MIT
import (
"errors"
"os"
"os/exec"
"strconv"
"github.com/nezhahq/agent/pkg/util"
)
type ROCmSMI struct {
BinPath string
}
func (rsmi *ROCmSMI) Gather() ([]float64, error) {
data := rsmi.pollROCmSMI()
return gatherROCmSMI(data)
}
func (rsmi *ROCmSMI) Start() error {
if _, err := os.Stat(rsmi.BinPath); os.IsNotExist(err) {
binPath, err := exec.LookPath("rocm-smi")
if err != nil {
return errors.New("didn't find the adequate tool to query GPU utilization")
}
rsmi.BinPath = binPath
}
return nil
}
func (rsmi *ROCmSMI) pollROCmSMI() []byte {
cmd := exec.Command(rsmi.BinPath,
"-u",
"--json",
)
gs, err := cmd.CombinedOutput()
if err != nil {
return nil
}
return gs
}
func gatherROCmSMI(ret []byte) ([]float64, error) {
var gpus map[string]GPU
var percentage []float64
err := util.Json.Unmarshal(ret, &gpus)
if err != nil {
return nil, err
}
for _, gpu := range gpus {
gp, _ := strconv.ParseFloat(gpu.GpuUsePercentage, 64)
percentage = append(percentage, gp)
}
return percentage, nil
}
type GPU struct {
GpuUsePercentage string `json:"GPU use (%)"`
}
pkg/gpu/stat/stat_darwin.go
-12
View File
@@ -1,12 +0,0 @@
//go:build darwin
package stat
import (
"github.com/nezhahq/agent/pkg/gpu"
)
func GetGPUStat() (float64, error) {
usage, err := gpu.FindUtilization("PerformanceStatistics", "Device Utilization %")
return float64(usage), err
}
pkg/gpu/stat/stat_freebsd.go
-7
View File
@@ -1,7 +0,0 @@
//go:build freebsd
package stat
func GetGPUStat() (float64, error) {
return 0, nil
}
pkg/gpu/stat/stat_linux.go
-44
View File
@@ -1,44 +0,0 @@
//go:build linux
package stat
func getNvidiaStat() ([]float64, error) {
smi := &NvidiaSMI{
BinPath: "/usr/bin/nvidia-smi",
}
err1 := smi.Start()
if err1 != nil {
return nil, err1
}
data, err2 := smi.Gather()
if err2 != nil {
return nil, err2
}
return data, nil
}
func getAMDStat() ([]float64, error) {
rsmi := &ROCmSMI{
BinPath: "/opt/rocm/bin/rocm-smi",
}
err1 := rsmi.Start()
if err1 != nil {
return nil, err1
}
data, err2 := rsmi.Gather()
if err2 != nil {
return nil, err2
}
return data, nil
}
func GetGPUStat() (float64, error) {
gs, err := getNvidiaStat()
if err != nil {
gs, err = getAMDStat()
}
if err != nil || len(gs) == 0 {
return 0, err
}
return gs[0], nil
}
pkg/gpu/vendor/amd_rocm_smi.go
Vendored
+107
View File
@@ -0,0 +1,107 @@
package vendor
// Modified from https://github.com/influxdata/telegraf/blob/master/plugins/inputs/amd_rocm_smi/amd_rocm_smi.go
// Original License: MIT
import (
"errors"
"os"
"os/exec"
"github.com/tidwall/gjson"
)
type ROCmSMI struct {
BinPath string
data []byte
}
func (rsmi *ROCmSMI) GatherModel() ([]string, error) {
return rsmi.gatherModel()
}
func (rsmi *ROCmSMI) GatherUsage() ([]float64, error) {
return rsmi.gatherUsage()
}
func (rsmi *ROCmSMI) Start() error {
if _, err := os.Stat(rsmi.BinPath); os.IsNotExist(err) {
binPath, err := exec.LookPath("rocm-smi")
if err != nil {
return errors.New("didn't find the adequate tool to query GPU utilization")
}
rsmi.BinPath = binPath
}
rsmi.data = rsmi.pollROCmSMI()
return nil
}
func (rsmi *ROCmSMI) pollROCmSMI() []byte {
cmd := exec.Command(rsmi.BinPath,
"-u",
"--showproductname",
"--json",
)
gs, err := cmd.CombinedOutput()
if err != nil {
return nil
}
return gs
}
func (rsmi *ROCmSMI) gatherModel() ([]string, error) {
m, err := parseModel(rsmi.data)
if err != nil {
return nil, err
}
return m, nil
}
func (rsmi *ROCmSMI) gatherUsage() ([]float64, error) {
u, err := parseUsage(rsmi.data)
if err != nil {
return nil, err
}
return u, nil
}
func parseModel(jsonObject []byte) ([]string, error) {
if jsonObject == nil {
return nil, nil
}
result := gjson.ParseBytes(jsonObject)
if !result.IsObject() {
return nil, errors.New("invalid JSON")
}
ret := make([]string, 0)
result.ForEach(func(_, value gjson.Result) bool {
ret = append(ret, value.Get("Card series").String())
return true
})
return ret, nil
}
func parseUsage(jsonObject []byte) ([]float64, error) {
if jsonObject == nil {
return nil, nil
}
result := gjson.ParseBytes(jsonObject)
if !result.IsObject() {
return nil, errors.New("invalid JSON")
}
ret := make([]float64, 0)
result.ForEach(func(_, value gjson.Result) bool {
ret = append(ret, value.Get("GPU use (%)").Float())
return true
})
return ret, nil
}
pkg/gpu/stat/nvidia_smi.go → pkg/gpu/vendor/nvidia_smi.go
Vendored
+29 -8
View File
@@ -1,4 +1,4 @@
package stat
package vendor
// Modified from https://github.com/influxdata/telegraf/blob/master/plugins/inputs/nvidia_smi/nvidia_smi.go
// Original License: MIT
@@ -14,12 +14,15 @@ import (
type NvidiaSMI struct {
BinPath string
data []byte
}
func (smi *NvidiaSMI) Gather() ([]float64, error) {
data := smi.pollNvidiaSMI()
func (smi *NvidiaSMI) GatherModel() ([]string, error) {
return smi.gatherModel()
}
return smi.parse(data)
func (smi *NvidiaSMI) GatherUsage() ([]float64, error) {
return smi.gatherUsage()
}
func (smi *NvidiaSMI) Start() error {
@@ -30,6 +33,7 @@ func (smi *NvidiaSMI) Start() error {
}
smi.BinPath = binPath
}
smi.data = smi.pollNvidiaSMI()
return nil
}
@@ -45,11 +49,27 @@ func (smi *NvidiaSMI) pollNvidiaSMI() []byte {
return gs
}
func (smi *NvidiaSMI) parse(data []byte) ([]float64, error) {
func (smi *NvidiaSMI) gatherModel() ([]string, error) {
var s smistat
var models []string
err := xml.Unmarshal(smi.data, &s)
if err != nil {
return nil, err
}
for _, gpu := range s.GPUs {
models = append(models, gpu.ProductName)
}
return models, nil
}
func (smi *NvidiaSMI) gatherUsage() ([]float64, error) {
var s smistat
var percentage []float64
err := xml.Unmarshal(data, &s)
err := xml.Unmarshal(smi.data, &s)
if err != nil {
return nil, err
}
@@ -75,11 +95,12 @@ func parsePercentage(p string) (float64, error) {
return value, nil
}
type nGPU struct {
type gpu struct {
ProductName string `xml:"product_name"`
Utilization struct {
GpuUtil string `xml:"gpu_util"`
} `xml:"utilization"`
}
type smistat struct {
GPUs []nGPU `xml:"gpu"`
GPUs []gpu `xml:"gpu"`
}