From 20316996f43b0ccebeecc0e2d0f690bfa528a01b Mon Sep 17 00:00:00 2001 From: UUBulb <35923940+uubulb@users.noreply.github.com> Date: Fri, 6 Sep 2024 08:31:19 +0800 Subject: [PATCH] gpu(darwin): format the string & fix behaviors on intel (#60) * gpu(darwin): format the string * gpu(darwin): fix behaviors on intel * return string directly --- pkg/gpu/gpu_darwin.go | 122 ++++++++++++++++++++++++----------------- pkg/monitor/monitor.go | 2 +- pkg/util/util.go | 24 ++++++++ 3 files changed, 96 insertions(+), 52 deletions(-) diff --git a/pkg/gpu/gpu_darwin.go b/pkg/gpu/gpu_darwin.go index 6a4e49a..1e0504e 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 ( @@ -30,9 +32,13 @@ type ( CFStringCreateWithCStringFunc = func(alloc uintptr, cStr string, encoding CFStringEncoding) CFStringRef CFGetTypeIDFunc = func(cf uintptr) CFTypeID CFStringGetTypeIDFunc = func() CFTypeID + CFStringGetLengthFunc = func(theString uintptr) int32 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) @@ -63,30 +69,19 @@ var ( var ( coreFoundation, _ = purego.Dlopen("/System/Library/Frameworks/CoreFoundation.framework/CoreFoundation", purego.RTLD_LAZY|purego.RTLD_GLOBAL) ioKit, _ = purego.Dlopen("/System/Library/Frameworks/IOKit.framework/IOKit", purego.RTLD_LAZY|purego.RTLD_GLOBAL) - - cfStringCreateWithCString, _ = purego.Dlsym(coreFoundation, "CFStringCreateWithCString") - cfGetTypeID, _ = purego.Dlsym(coreFoundation, "CFGetTypeID") - cfStringGetTypeID, _ = purego.Dlsym(coreFoundation, "CFStringGetTypeID") - cfStringGetCString, _ = purego.Dlsym(coreFoundation, "CFStringGetCString") - cfDictionaryGetTypeID, _ = purego.Dlsym(coreFoundation, "CFDictionaryGetTypeID") - cfDictionaryGetValue, _ = purego.Dlsym(coreFoundation, "CFDictionaryGetValue") - cfNumberGetValue, _ = purego.Dlsym(coreFoundation, "CFNumberGetValue") - cfRelease, _ = purego.Dlsym(coreFoundation, "CFRelease") - - ioServiceGetMatchingServices, _ = purego.Dlsym(ioKit, "IOServiceGetMatchingServices") - ioIteratorNext, _ = purego.Dlsym(ioKit, "IOIteratorNext") - ioServiceMatching, _ = purego.Dlsym(ioKit, "IOServiceMatching") - ioRegistryEntrySearchCFProperty, _ = purego.Dlsym(ioKit, "IORegistryEntrySearchCFProperty") - ioObjectRelease, _ = purego.Dlsym(ioKit, "IOObjectRelease") ) var ( CFStringCreateWithCString CFStringCreateWithCStringFunc CFGetTypeID CFGetTypeIDFunc CFStringGetTypeID CFStringGetTypeIDFunc + CFStringGetLength CFStringGetLengthFunc CFStringGetCString CFStringGetCStringFunc CFDictionaryGetTypeID CFDictionaryGetTypeIDFunc CFDictionaryGetValue CFDictionaryGetValueFunc + CFDataGetTypeID CFDataGetTypeIDFunc + CFDataGetBytePtr CFDataGetBytePtrFunc + CFDataGetLength CFDataGetLengthFunc CFNumberGetValue CFNumberGetValueFunc CFRelease CFReleaseFunc @@ -97,25 +92,37 @@ var ( IOObjectRelease IOObjectReleaseFunc ) -func init() { - purego.RegisterFunc(&CFStringCreateWithCString, cfStringCreateWithCString) - purego.RegisterFunc(&CFGetTypeID, cfGetTypeID) - purego.RegisterFunc(&CFStringGetTypeID, cfStringGetTypeID) - purego.RegisterFunc(&CFStringGetCString, cfStringGetCString) - purego.RegisterFunc(&CFDictionaryGetTypeID, cfDictionaryGetTypeID) - purego.RegisterFunc(&CFDictionaryGetValue, cfDictionaryGetValue) - purego.RegisterFunc(&CFNumberGetValue, cfNumberGetValue) - purego.RegisterFunc(&CFRelease, cfRelease) +var validVendors = []string{ + "AMD", "Intel", "NVIDIA", "Apple", +} - purego.RegisterFunc(&IOServiceGetMatchingServices, ioServiceGetMatchingServices) - purego.RegisterFunc(&IOIteratorNext, ioIteratorNext) - purego.RegisterFunc(&IOServiceMatching, ioServiceMatching) - purego.RegisterFunc(&IORegistryEntrySearchCFProperty, ioRegistryEntrySearchCFProperty) - purego.RegisterFunc(&IOObjectRelease, ioObjectRelease) +func init() { + purego.RegisterLibFunc(&CFStringCreateWithCString, coreFoundation, "CFStringCreateWithCString") + purego.RegisterLibFunc(&CFGetTypeID, coreFoundation, "CFGetTypeID") + purego.RegisterLibFunc(&CFStringGetTypeID, coreFoundation, "CFStringGetTypeID") + purego.RegisterLibFunc(&CFStringGetLength, coreFoundation, "CFStringGetLength") + 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") + + purego.RegisterLibFunc(&IOServiceGetMatchingServices, ioKit, "IOServiceGetMatchingServices") + purego.RegisterLibFunc(&IOIteratorNext, ioKit, "IOIteratorNext") + purego.RegisterLibFunc(&IOServiceMatching, ioKit, "IOServiceMatching") + purego.RegisterLibFunc(&IORegistryEntrySearchCFProperty, ioKit, "IORegistryEntrySearchCFProperty") + purego.RegisterLibFunc(&IOObjectRelease, ioKit, "IOObjectRelease") } 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) { @@ -125,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 { @@ -144,15 +152,16 @@ func findDevices(key string) ([]string, error) { result, _, _ := findProperties(service, uintptr(cfStr), 0) IOObjectRelease(service) - if result != nil { - results = append(results, string(result)) + if util.ContainsStr(validVendors, result) { + results = append(results, result) 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 } } @@ -162,17 +171,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)) @@ -182,44 +196,50 @@ 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) { +func findProperties(service ioRegistryEntry, key, dictKey uintptr) (string, int, error) { properties := IORegistryEntrySearchCFProperty(service, kIOServicePlane, key, kCFAllocatorDefault, kIORegistryIterateRecursively) ptrValue := uintptr(properties) if properties != nil { switch CFGetTypeID(ptrValue) { // model case CFStringGetTypeID(): - buf := make([]byte, 1024) - CFStringGetCString(ptrValue, &buf[0], int32(unsafe.Sizeof(buf)), uint32(kCFStringEncodingUTF8)) + length := CFStringGetLength(ptrValue) + 1 // null terminator + buf := make([]byte, length-1) + CFStringGetCString(ptrValue, &buf[0], length, uint32(kCFStringEncodingUTF8)) CFRelease(ptrValue) - return buf, 0, nil + return string(buf), 0, nil + case CFDataGetTypeID(): + length := CFDataGetLength(ptrValue) + bin := unsafe.String((*byte)(CFDataGetBytePtr(ptrValue)), length) + CFRelease(ptrValue) + return bin, 0, nil // PerformanceStatistics case CFDictionaryGetTypeID(): cfValue := CFDictionaryGetValue(ptrValue, dictKey) if cfValue != nil { var value int if CFNumberGetValue(uintptr(cfValue), kCFNumberIntType, uintptr(unsafe.Pointer(&value))) { - return nil, value, nil + return "", value, nil } else { - return nil, 0, fmt.Errorf("failed to exec CFNumberGetValue") + return "", 0, fmt.Errorf("failed to exec CFNumberGetValue") } } else { - return nil, 0, fmt.Errorf("failed to exec CFDictionaryGetValue") + return "", 0, fmt.Errorf("failed to exec CFDictionaryGetValue") } } } - return nil, 0, fmt.Errorf("failed to exec IORegistryEntrySearchCFProperty") + return "", 0, fmt.Errorf("failed to exec IORegistryEntrySearchCFProperty") } 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 +}