gpu(darwin): format the string & fix behaviors on intel (#60)
* gpu(darwin): format the string * gpu(darwin): fix behaviors on intel * return string directly
This commit is contained in:
UUBulb
GitHub
parent
439be0ddf3
commit
20316996f4
@ -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")
|
||||
}
|
||||
|
||||
@ -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
|
||||
|
||||
@ -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
|
||||
}
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user