nezhahq-agent/pkg/gpu/stat/amd_rocm_smi.go

67 lines
1.2 KiB
Go

package stat
// Modified from https://github.com/influxdata/telegraf/blob/master/plugins/inputs/amd_rocm_smi/amd_rocm_smi.go
// Original License: MIT
import (
"encoding/json"
"errors"
"os"
"os/exec"
"strconv"
)
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 := 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 (%)"`
}