From c8b57342b1cefc7c03b621a0b3798769db5e4574 Mon Sep 17 00:00:00 2001 From: uubulb Date: Thu, 5 Sep 2024 23:28:47 +0800 Subject: [PATCH] gpu(darwin): fix behaviors on intel --- pkg/gpu/gpu_darwin.go | 61 ++++++++++++++++++++++++++++++++---------- pkg/monitor/monitor.go | 2 +- pkg/util/util.go | 24 +++++++++++++++++ 3 files changed, 72 insertions(+), 15 deletions(-) diff --git a/pkg/gpu/gpu_darwin.go b/pkg/gpu/gpu_darwin.go index cfc3ab0..adcb2f3 100644 --- a/pkg/gpu/gpu_darwin.go +++ b/pkg/gpu/gpu_darwin.go @@ -7,6 +7,8 @@ import ( "unsafe" "github.com/ebitengine/purego" + + "github.com/nezhahq/agent/pkg/util" ) type ( @@ -34,6 +36,9 @@ type ( CFStringGetCStringFunc = func(cfStr uintptr, buffer *byte, size CFIndex, encoding CFStringEncoding) bool CFDictionaryGetTypeIDFunc = func() CFTypeID CFDictionaryGetValueFunc = func(dict, key uintptr) unsafe.Pointer + CFDataGetTypeIDFunc = func() CFTypeID + CFDataGetBytePtrFunc = func(theData uintptr) unsafe.Pointer + CFDataGetLengthFunc = func(theData uintptr) CFIndex CFNumberGetValueFunc = func(number uintptr, theType CFNumberType, valuePtr uintptr) bool CFReleaseFunc = func(cf uintptr) @@ -74,6 +79,9 @@ var ( CFStringGetCString CFStringGetCStringFunc CFDictionaryGetTypeID CFDictionaryGetTypeIDFunc CFDictionaryGetValue CFDictionaryGetValueFunc + CFDataGetTypeID CFDataGetTypeIDFunc + CFDataGetBytePtr CFDataGetBytePtrFunc + CFDataGetLength CFDataGetLengthFunc CFNumberGetValue CFNumberGetValueFunc CFRelease CFReleaseFunc @@ -84,6 +92,10 @@ var ( IOObjectRelease IOObjectReleaseFunc ) +var validVendors = []string{ + "AMD", "Intel", "NVIDIA", "Apple", +} + func init() { purego.RegisterLibFunc(&CFStringCreateWithCString, coreFoundation, "CFStringCreateWithCString") purego.RegisterLibFunc(&CFGetTypeID, coreFoundation, "CFGetTypeID") @@ -92,6 +104,9 @@ func init() { purego.RegisterLibFunc(&CFStringGetCString, coreFoundation, "CFStringGetCString") purego.RegisterLibFunc(&CFDictionaryGetTypeID, coreFoundation, "CFDictionaryGetTypeID") purego.RegisterLibFunc(&CFDictionaryGetValue, coreFoundation, "CFDictionaryGetValue") + purego.RegisterLibFunc(&CFDataGetTypeID, coreFoundation, "CFDataGetTypeID") + purego.RegisterLibFunc(&CFDataGetBytePtr, coreFoundation, "CFDataGetBytePtr") + purego.RegisterLibFunc(&CFDataGetLength, coreFoundation, "CFDataGetLength") purego.RegisterLibFunc(&CFNumberGetValue, coreFoundation, "CFNumberGetValue") purego.RegisterLibFunc(&CFRelease, coreFoundation, "CFRelease") @@ -103,7 +118,11 @@ func init() { } func GetGPUModel() ([]string, error) { - return findDevices("model") + models, err := findDevices("model") + if err != nil { + return nil, err + } + return util.RemoveDuplicate(models), nil } func FindUtilization(key, dictKey string) (int, error) { @@ -113,6 +132,7 @@ func FindUtilization(key, dictKey string) (int, error) { func findDevices(key string) ([]string, error) { var iterator ioIterator var results []string + done := false iv := IOServiceGetMatchingServices(kIOMainPortDefault, uintptr(IOServiceMatching(IOSERVICE_GPU)), &iterator) if iv != KERN_SUCCESS { @@ -132,15 +152,18 @@ func findDevices(key string) ([]string, error) { result, _, _ := findProperties(service, uintptr(cfStr), 0) IOObjectRelease(service) - if result != nil { - results = append(results, string(result)) + resultStr := string(result) + + if util.ContainsStr(validVendors, resultStr) { + results = append(results, resultStr) index++ - } else if key == "model" { + } else if key == "model" && !done { IOObjectRelease(iterator) iv = IOServiceGetMatchingServices(kIOMainPortDefault, uintptr(IOServiceMatching(IOSERVICE_PCI)), &iterator) if iv != KERN_SUCCESS { return nil, fmt.Errorf("error retrieving GPU entry") } + done = true } } @@ -150,17 +173,22 @@ func findDevices(key string) ([]string, error) { func findUtilization(key, dictKey string) (int, error) { var iterator ioIterator - var result int var err error + result := 0 iv := IOServiceGetMatchingServices(kIOMainPortDefault, uintptr(IOServiceMatching(IOSERVICE_GPU)), &iterator) if iv != KERN_SUCCESS { return 0, fmt.Errorf("error retrieving GPU entry") } - // Only retrieving the utilization of first GPU here - service := IOIteratorNext(iterator) - if service != MACH_PORT_NULL { + // Only retrieving the utilization of a single GPU here + var service ioObject + for { + service = IOIteratorNext(iterator) + if service == MACH_PORT_NULL { + break + } + cfStr := CFStringCreateWithCString(kCFAllocatorDefault, key, CFStringEncoding(kCFStringEncodingUTF8)) cfDictStr := CFStringCreateWithCString(kCFAllocatorDefault, dictKey, CFStringEncoding(kCFStringEncodingUTF8)) @@ -170,17 +198,17 @@ func findUtilization(key, dictKey string) (int, error) { CFRelease(uintptr(cfDictStr)) if err != nil { - return 0, fmt.Errorf("failed retrieving GPU utilization: %v", err) + IOObjectRelease(service) + continue + } else if result != 0 { + break } - } else { - IOObjectRelease(service) - IOObjectRelease(iterator) - return 0, fmt.Errorf("no GPU utilization entry found") } IOObjectRelease(service) IOObjectRelease(iterator) - return result, nil + + return result, err } func findProperties(service ioRegistryEntry, key, dictKey uintptr) ([]byte, int, error) { @@ -195,6 +223,11 @@ func findProperties(service ioRegistryEntry, key, dictKey uintptr) ([]byte, int, CFStringGetCString(ptrValue, &buf[0], length, uint32(kCFStringEncodingUTF8)) CFRelease(ptrValue) return buf, 0, nil + case CFDataGetTypeID(): + length := CFDataGetLength(ptrValue) + bin := unsafe.Slice((*byte)(CFDataGetBytePtr(ptrValue)), length) + CFRelease(ptrValue) + return bin, 0, nil // PerformanceStatistics case CFDictionaryGetTypeID(): cfValue := CFDictionaryGetValue(ptrValue, dictKey) diff --git a/pkg/monitor/monitor.go b/pkg/monitor/monitor.go index 70833f4..ff4a0ae 100644 --- a/pkg/monitor/monitor.go +++ b/pkg/monitor/monitor.go @@ -350,7 +350,7 @@ func updateGPUStat() float64 { gs, err := gpustat.GetGPUStat() if err != nil { statDataFetchAttempts["GPU"]++ - println("gpustat.GetGPUStat error: ", err, ", attempt: ", statDataFetchAttempts["GPU"]) + printf("gpustat.GetGPUStat error: %v, attempt: %d", err, statDataFetchAttempts["GPU"]) return 0 } else { statDataFetchAttempts["GPU"] = 0 diff --git a/pkg/util/util.go b/pkg/util/util.go index b3751d7..c0a59c3 100644 --- a/pkg/util/util.go +++ b/pkg/util/util.go @@ -4,6 +4,7 @@ import ( "fmt" "net/http" "os" + "strings" "time" jsoniter "github.com/json-iterator/go" @@ -40,3 +41,26 @@ func BrowserHeaders() *http.Header { "User-Agent": {MacOSChromeUA}, } } + +func ContainsStr(slice []string, str string) bool { + if str != "" { + for _, item := range slice { + if strings.Contains(str, item) { + return true + } + } + } + return false +} + +func RemoveDuplicate[T comparable](sliceList []T) []T { + allKeys := make(map[T]bool) + list := []T{} + for _, item := range sliceList { + if _, value := allKeys[item]; !value { + allKeys[item] = true + list = append(list, item) + } + } + return list +}