feat: Add GPU inspection/monitoring support (#28)

pkg/gpu/gpu.go

@@ -0,0 +1,28 @@
+//go:build !darwin
+// +build !darwin
+
+package gpu
+
+import (
+	"fmt"
+
+	"github.com/jaypipes/ghw"
+)
+
+func GetGPUModel() []string {
+	var gpuModel []string
+	gi, err := ghw.GPU(ghw.WithDisableWarnings())
+	if err != nil {
+		fmt.Printf("Error getting GPU info: %v", err)
+		return nil
+	}
+
+	for _, card := range gi.GraphicsCards {
+		if card.DeviceInfo == nil {
+			return nil
+		}
+		gpuModel = append(gpuModel, card.DeviceInfo.Product.Name)
+	}
+
+	return gpuModel
+}

pkg/gpu/gpu_darwin.go

@@ -0,0 +1,55 @@
+//go:build darwin
+
+package gpu
+
+import (
+	"fmt"
+	"os/exec"
+	"regexp"
+	"strings"
+)
+
+func extractGPUInfo(cmd *exec.Cmd) ([]string, error) {
+	gi, err := cmd.CombinedOutput()
+	if err != nil {
+		return nil, err
+	}
+
+	re := regexp.MustCompile(`"model"\s*=\s*["<]?"([^">]+)"[">]?`)
+	matches := re.FindAllSubmatch(gi, -1)
+	var modelNames []string
+	for _, match := range matches {
+		if len(match) > 1 {
+			modelNames = append(modelNames, string(match[1]))
+		}
+	}
+	return modelNames, nil
+}
+
+func GetGPUModel() []string {
+	vendorNames := []string{
+		"AMD", "Intel", "Nvidia", "Apple",
+	}
+
+	ioreg := exec.Command("ioreg", "-rd1", "-c", "IOAccelerator")
+	gi, err := extractGPUInfo(ioreg)
+	if err != nil || len(gi) == 0 {
+		ioreg = exec.Command("ioreg", "-rd1", "-c", "IOPCIDevice")
+		gi, err = extractGPUInfo(ioreg)
+		if err != nil {
+			fmt.Println("Error executing ioreg:", err)
+			return nil
+		}
+	}
+
+	var gpuModel []string
+	for _, model := range gi {
+		for _, vendor := range vendorNames {
+			if strings.Contains(model, vendor) {
+				gpuModel = append(gpuModel, model)
+				break
+			}
+		}
+	}
+	return gpuModel
+}

pkg/gpu/stat/amd_rocm_smi.go

@@ -0,0 +1,66 @@
+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 (%)"`
+}

pkg/gpu/stat/nvidia_smi.go

@@ -0,0 +1,85 @@
+package stat
+
+// Modified from https://github.com/influxdata/telegraf/blob/master/plugins/inputs/nvidia_smi/nvidia_smi.go
+// Original License: MIT
+
+import (
+	"encoding/xml"
+	"errors"
+	"os"
+	"os/exec"
+	"strconv"
+	"strings"
+)
+
+type NvidiaSMI struct {
+	BinPath string
+}
+
+func (smi *NvidiaSMI) Gather() ([]float64, error) {
+	data := smi.pollNvidiaSMI()
+
+	return smi.parse(data)
+}
+
+func (smi *NvidiaSMI) Start() error {
+	if _, err := os.Stat(smi.BinPath); os.IsNotExist(err) {
+		binPath, err := exec.LookPath("nvidia-smi")
+		if err != nil {
+			return errors.New("Didn't find the adequate tool to query GPU utilization")
+		}
+		smi.BinPath = binPath
+	}
+	return nil
+}
+
+func (smi *NvidiaSMI) pollNvidiaSMI() []byte {
+	cmd := exec.Command(smi.BinPath,
+		"-q",
+		"-x",
+	)
+	gs, err := cmd.CombinedOutput()
+	if err != nil {
+		return nil
+	}
+	return gs
+}
+
+func (smi *NvidiaSMI) parse(data []byte) ([]float64, error) {
+	var s smistat
+	var percentage []float64
+
+	err := xml.Unmarshal(data, &s)
+	if err != nil {
+		return nil, err
+	}
+
+	for _, gpu := range s.GPUs {
+		gp, _ := parsePercentage(gpu.Utilization.GpuUtil)
+		percentage = append(percentage, gp)
+	}
+
+	return percentage, nil
+}
+
+func parsePercentage(p string) (float64, error) {
+	per := strings.ReplaceAll(p, " ", "")
+
+	t := strings.TrimSuffix(per, "%")
+
+	value, err := strconv.ParseFloat(t, 64)
+	if err != nil {
+		return 0, err
+	}
+
+	return value, nil
+}
+
+type nGPU struct {
+	Utilization struct {
+		GpuUtil string `xml:"gpu_util"`
+	} `xml:"utilization"`
+}
+type smistat struct {
+	GPUs []nGPU `xml:"gpu"`
+}

pkg/gpu/stat/stat_darwin.go

@@ -0,0 +1,36 @@
+//go:build darwin
+
+package stat
+
+import (
+	"os/exec"
+	"regexp"
+	"strconv"
+)
+
+func extractGPUStat(cmd *exec.Cmd) ([]float64, error) {
+	gs, err := cmd.CombinedOutput()
+	if err != nil {
+		return nil, err
+	}
+
+	re := regexp.MustCompile(`"Device Utilization %"\s*=\s*(\d+)`)
+	matches := re.FindAllSubmatch(gs, -1)
+	var u []float64
+	for _, match := range matches {
+		if len(match) > 1 {
+			p, _ := strconv.ParseFloat(string(match[1]), 64)
+			u = append(u, p)
+		}
+	}
+	return u, nil
+}
+
+func GetGPUStat() (float64, error) {
+	ioreg := exec.Command("ioreg", "-rd1", "-c", "IOAccelerator")
+	gs, err := extractGPUStat(ioreg)
+	if err != nil || len(gs) == 0 {
+		return -1, err
+	}
+	return gs[0], nil
+}

pkg/gpu/stat/stat_freebsd.go

@@ -0,0 +1,7 @@
+//go:build freebsd
+
+package stat
+
+func GetGPUStat() (float64, error) {
+	return -1, nil
+}

pkg/gpu/stat/stat_linux.go

@@ -0,0 +1,44 @@
+//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 -1, err
+	}
+	return gs[0], nil
+}

pkg/gpu/stat/stat_windows.go

@@ -0,0 +1,31 @@
+//go:build windows
+
+package stat
+
+import (
+	"os/exec"
+	"strconv"
+	"strings"
+)
+
+func GetGPUStat() (float64, error) {
+	shellPath, err := exec.LookPath("powershell.exe")
+	if err != nil || shellPath == "" {
+		return -1, err
+	}
+	cmd := exec.Command(
+		shellPath,
+		"-Command",
+		`Write-Output "$([math]::Round((((Get-Counter "\GPU Engine(*engtype_3D)\Utilization Percentage").CounterSamples | where CookedValue).CookedValue | measure -sum).sum,2))"`,
+	)
+	output, err1 := cmd.CombinedOutput()
+	if err1 != nil {
+		return -1, err1
+	}
+	t := strings.TrimSpace(string(output))
+	gs, err2 := strconv.ParseFloat(t, 64)
+	if err2 != nil {
+		return -1, err2
+	}
+	return gs, nil
+}