gpu: Add closure for Jetson and improve compatibility

This commit is contained in:
hank
2025-01-24 22:07:37 -05:00
parent d185dfdef8
commit c157f38957

View File

@@ -91,15 +91,20 @@ func (c *gpuCollector) collect() error {
return c.cmd.Wait() return c.cmd.Wait()
} }
// parseJetsonData parses the output of rtegrastats and updates the GPUData map // getJetsonParser returns a function to parse the output of tegrastats and update the GPUData map
func (gm *GPUManager) parseJetsonData(output []byte) bool { func (gm *GPUManager) getJetsonParser() func(output []byte) bool {
data := string(output) // use closure to avoid recompiling the regex
ramPattern := regexp.MustCompile(`RAM (\d+)/(\d+)MB`) ramPattern := regexp.MustCompile(`RAM (\d+)/(\d+)MB`)
gr3dPattern := regexp.MustCompile(`GR3D_FREQ (\d+)%`) gr3dPattern := regexp.MustCompile(`GR3D_FREQ (\d+)%`)
tempPattern := regexp.MustCompile(`([a-z0-9_]+)@(\d+\.?\d*)C`) tempPattern := regexp.MustCompile(`tj@(\d+\.?\d*)C`)
powerPattern := regexp.MustCompile(`VDD_GPU_SOC (\d+)mW`) // Orin Nano / NX do not have GPU specific power monitor
// TODO: Maybe use VDD_IN for Nano / NX and add a total system power chart
powerPattern := regexp.MustCompile(`(GPU_SOC|CPU_GPU_CV) (\d+)mW`)
return func(output []byte) bool {
gm.mutex.Lock() gm.mutex.Lock()
defer gm.mutex.Unlock() defer gm.mutex.Unlock()
data := string(output)
gpuData := gm.GpuDataMap["0"] gpuData := gm.GpuDataMap["0"]
// Parse RAM usage // Parse RAM usage
ramMatches := ramPattern.FindStringSubmatch(data) ramMatches := ramPattern.FindStringSubmatch(data)
@@ -110,18 +115,13 @@ func (gm *GPUManager) parseJetsonData(output []byte) bool {
// Parse GR3D (GPU) usage // Parse GR3D (GPU) usage
gr3dMatches := gr3dPattern.FindStringSubmatch(data) gr3dMatches := gr3dPattern.FindStringSubmatch(data)
if gr3dMatches != nil { if gr3dMatches != nil {
usage, _ := strconv.ParseFloat(gr3dMatches[1], 64) gpuData.Usage, _ = strconv.ParseFloat(gr3dMatches[1], 64)
gpuData.Usage = usage / 100
} }
// Parse temperature
tempMatches := tempPattern.FindAllStringSubmatch(data, -1) tempMatches := tempPattern.FindStringSubmatch(data)
for _, match := range tempMatches { if tempMatches != nil {
if match[1] == "cpu" { gpuData.Temperature, _ = strconv.ParseFloat(tempMatches[1], 64)
gpuData.Temperature, _ = strconv.ParseFloat(match[2], 64)
break
} }
}
// Parse power usage // Parse power usage
powerMatches := powerPattern.FindStringSubmatch(data) powerMatches := powerPattern.FindStringSubmatch(data)
if powerMatches != nil { if powerMatches != nil {
@@ -130,6 +130,7 @@ func (gm *GPUManager) parseJetsonData(output []byte) bool {
} }
gpuData.Count++ gpuData.Count++
return true return true
}
} }
// parseNvidiaData parses the output of nvidia-smi and updates the GPUData map // parseNvidiaData parses the output of nvidia-smi and updates the GPUData map
@@ -250,7 +251,7 @@ func (gm *GPUManager) detectGPUs() error {
if gm.nvidiaSmi || gm.rocmSmi || gm.tegrastats { if gm.nvidiaSmi || gm.rocmSmi || gm.tegrastats {
return nil return nil
} }
return fmt.Errorf("no GPU found - install nvidia-smi or rocm-smi or tegrastats") return fmt.Errorf("no GPU found - install nvidia-smi, rocm-smi, or tegrastats")
} }
// startCollector starts the appropriate GPU data collector based on the command // startCollector starts the appropriate GPU data collector based on the command
@@ -276,12 +277,11 @@ func (gm *GPUManager) startCollector(command string) {
case "tegrastats": case "tegrastats":
jetsonCollector := gpuCollector{ jetsonCollector := gpuCollector{
name: "tegrastats", name: "tegrastats",
cmd: exec.Command("tegrastats"), cmd: exec.Command("tegrastats", "--interval", "3000"),
parse: gm.parseJetsonData, parse: gm.getJetsonParser(),
} }
go jetsonCollector.start() go jetsonCollector.start()
} }
} }
// NewGPUManager creates and initializes a new GPUManager // NewGPUManager creates and initializes a new GPUManager