127 lines
3.1 KiB
Go
127 lines
3.1 KiB
Go
package platform
|
|
|
|
import (
|
|
"errors"
|
|
"fmt"
|
|
"strconv"
|
|
"strings"
|
|
)
|
|
|
|
func nvidiaSATStatusSeverity(status string) int {
|
|
switch strings.ToUpper(strings.TrimSpace(status)) {
|
|
case "FAILED":
|
|
return 3
|
|
case "PARTIAL", "UNSUPPORTED":
|
|
return 2
|
|
case "OK":
|
|
return 1
|
|
default:
|
|
return 0
|
|
}
|
|
}
|
|
|
|
func firstLine(s string) string {
|
|
s = strings.TrimSpace(s)
|
|
if s == "" {
|
|
return ""
|
|
}
|
|
if idx := strings.IndexByte(s, '\n'); idx >= 0 {
|
|
return strings.TrimSpace(s[:idx])
|
|
}
|
|
return s
|
|
}
|
|
|
|
func nvidiaJobNeedsHealthCheck(job satJob) bool {
|
|
if job.collectGPU {
|
|
return true
|
|
}
|
|
name := strings.ToLower(strings.TrimSpace(job.name))
|
|
return strings.Contains(name, "dcgmi") ||
|
|
strings.Contains(name, "all-reduce") ||
|
|
strings.Contains(name, "gpu-burn") ||
|
|
strings.Contains(name, "gpu-stress") ||
|
|
strings.Contains(name, "dcgmproftester")
|
|
}
|
|
|
|
func checkNvidiaJobHealth(selected []int) (string, error) {
|
|
health, _ := readNvidiaGPUHealth()
|
|
var bad []nvidiaGPUHealth
|
|
selectedSet := make(map[int]struct{}, len(selected))
|
|
for _, idx := range selected {
|
|
selectedSet[idx] = struct{}{}
|
|
}
|
|
for _, gpu := range health {
|
|
if len(selectedSet) > 0 {
|
|
if _, ok := selectedSet[gpu.Index]; !ok {
|
|
continue
|
|
}
|
|
}
|
|
if gpu.NeedsReset {
|
|
bad = append(bad, gpu)
|
|
}
|
|
}
|
|
var degraded []nvidiaNVLinkDegradedEvent
|
|
if out, err := satExecCommand("dmesg").Output(); err == nil {
|
|
for _, event := range parseNvidiaNVLinkDegradedDmesg(out) {
|
|
if event.GPUIndex >= 0 && len(selectedSet) > 0 {
|
|
if _, ok := selectedSet[event.GPUIndex]; !ok {
|
|
continue
|
|
}
|
|
}
|
|
degraded = append(degraded, event)
|
|
}
|
|
}
|
|
if len(bad) == 0 && len(degraded) == 0 {
|
|
return "", nil
|
|
}
|
|
lines := make([]string, 0, len(bad)+len(degraded)+1)
|
|
lines = append(lines, "NVIDIA GPU health check failed:")
|
|
for _, gpu := range bad {
|
|
lines = append(lines, fmt.Sprintf("gpu %d (%s) requires reset: %s", gpu.Index, gpu.Name, gpu.RawLine))
|
|
}
|
|
for _, event := range degraded {
|
|
lines = append(lines, event.Warning())
|
|
}
|
|
return strings.Join(lines, "\n"), errors.New("nvidia gpu health check failed")
|
|
}
|
|
|
|
func readNvidiaGPUHealth() ([]nvidiaGPUHealth, error) {
|
|
out, err := satExecCommand(
|
|
"nvidia-smi",
|
|
"--query-gpu=index,name,temperature.gpu,power.draw,utilization.gpu,memory.used,memory.total",
|
|
"--format=csv,noheader,nounits",
|
|
).Output()
|
|
if err != nil {
|
|
return nil, fmt.Errorf("nvidia-smi: %w", err)
|
|
}
|
|
return parseNvidiaGPUHealth(string(out)), nil
|
|
}
|
|
|
|
func parseNvidiaGPUHealth(raw string) []nvidiaGPUHealth {
|
|
var gpus []nvidiaGPUHealth
|
|
for _, line := range strings.Split(strings.TrimSpace(raw), "\n") {
|
|
line = strings.TrimSpace(line)
|
|
if line == "" {
|
|
continue
|
|
}
|
|
parts := strings.Split(line, ",")
|
|
if len(parts) < 2 {
|
|
gpus = append(gpus, nvidiaGPUHealth{RawLine: line, ParseFailure: true})
|
|
continue
|
|
}
|
|
idx, err := strconv.Atoi(strings.TrimSpace(parts[0]))
|
|
if err != nil {
|
|
gpus = append(gpus, nvidiaGPUHealth{RawLine: line, ParseFailure: true})
|
|
continue
|
|
}
|
|
upper := strings.ToUpper(line)
|
|
gpus = append(gpus, nvidiaGPUHealth{
|
|
Index: idx,
|
|
Name: strings.TrimSpace(parts[1]),
|
|
NeedsReset: strings.Contains(upper, "GPU REQUIRES RESET"),
|
|
RawLine: line,
|
|
})
|
|
}
|
|
return gpus
|
|
}
|