package gogpu
|
|
/*
|
#cgo CFLAGS: -I. -w -g
|
#cgo CXXFLAGS: -I. -std=c++11 -w -g
|
#cgo LDFLAGS: -ldl
|
#include <stdlib.h>
|
#include "cgpu.h"
|
*/
|
import "C"
|
import (
|
"encoding/json"
|
"errors"
|
"fmt"
|
"sort"
|
"strconv"
|
"strings"
|
"unsafe"
|
)
|
|
// ValidGPU pass needed gpu memory size test if satisfy
|
func ValidGPU(memSize int) int {
|
return int(C.get_idle_gpu(C.int(memSize)))
|
}
|
|
// SatisfyGPU satisfy unit "M"
|
func SatisfyGPU(index, memSize, reserve int) bool {
|
info, err := Info()
|
if err != nil || info.Count <= index {
|
fmt.Println("SatisfyGPU no gpu or index ilegal: ", index, " gpu count : ", info.Count)
|
return false
|
}
|
|
var M1 int64 = 1024 * 1024
|
if info.Info[index].GpuMemoryFree-int64(memSize)*M1 > int64(reserve)*M1 {
|
return true
|
}
|
return false
|
}
|
|
// GPURank sort
|
type GPURank []GpuUnitInfo
|
|
func (a GPURank) Len() int { return len(a) }
|
func (a GPURank) Swap(i, j int) { a[i], a[j] = a[j], a[i] }
|
func (a GPURank) Less(i, j int) bool { return a[i].GpuMemoryFree < a[j].GpuMemoryFree }
|
|
// RankGPU rank
|
func RankGPU() []int {
|
info, err := Info()
|
if err != nil {
|
fmt.Println("RankGPU no gpu")
|
return nil
|
}
|
var ret []int
|
if info.Count == 1 {
|
ret = append(ret, 0)
|
} else {
|
sort.Sort(GPURank(info.Info))
|
for i := len(info.Info) - 1; i >= 0; i-- {
|
ret = append(ret, info.Info[i].GpuIndex)
|
}
|
}
|
return ret
|
}
|
|
// IdleGPU idle most
|
func IdleGPU() int {
|
info, err := Info()
|
if err != nil {
|
fmt.Println("IdleGPU no gpu")
|
return -1
|
}
|
if info.Count == 1 {
|
return 0
|
}
|
|
var free int64
|
gpu := 0
|
for k, v := range info.Info {
|
if v.GpuMemoryFree > free {
|
free = v.GpuMemoryFree
|
gpu = k
|
}
|
}
|
return gpu
|
}
|
|
type Process struct {
|
Pid int `json:"pid"`
|
UsedMem int `json:"used_mem"`
|
}
|
|
// GpuUnitInfo gpu card info
|
type GpuUnitInfo struct {
|
GpuIndex int `json:"index"`
|
GpuUtilization int `json:"utilization"`
|
GpuTemperature int `json:"temperature"`
|
GpuDecoderUtil int `json:"decoder"`
|
GpuEncoderUtil int `json:"encoder"`
|
GpuProcesses int `json:"process"`
|
GpuMemoryTotal int64 `json:"memory"`
|
GpuMemoryFree int64 `json:"free"`
|
GpuMemoryUsed int64 `json:"used"`
|
}
|
|
func (u GpuUnitInfo) String() string {
|
s, _ := json.Marshal(u)
|
return string(s)
|
}
|
|
// GpuInfo gpu info
|
type GpuInfo struct {
|
Count int `json:"count"`
|
Info []GpuUnitInfo `json:"info"`
|
Processes []Process `json:"processes`
|
}
|
|
func (g GpuInfo) String() string {
|
s, _ := json.Marshal(g)
|
return string(s)
|
}
|
|
// Info gpu infos
|
func Info() (*GpuInfo, error) {
|
|
p := C.get_gpu_info()
|
if p == nil {
|
return nil, errors.New("can't get gpu infos")
|
}
|
defer C.free(unsafe.Pointer(p))
|
|
str := C.GoString(p)
|
|
infos := strings.Split(str, "|")
|
|
length := len(infos)
|
|
var units []GpuUnitInfo
|
var procs []Process
|
var count int
|
var processStart bool
|
for i := 0; i < length; i++ {
|
if infos[i] == "gpu" {
|
u := GpuUnitInfo{}
|
u.GpuIndex = count
|
count++
|
|
i++
|
u.GpuUtilization, _ = strconv.Atoi(infos[i])
|
i++
|
u.GpuTemperature, _ = strconv.Atoi(infos[i])
|
i++
|
u.GpuDecoderUtil, _ = strconv.Atoi(infos[i])
|
i++
|
u.GpuEncoderUtil, _ = strconv.Atoi(infos[i])
|
i++
|
u.GpuProcesses, _ = strconv.Atoi(infos[i])
|
i++
|
u.GpuMemoryTotal, _ = strconv.ParseInt(infos[i], 10, 64)
|
i++
|
u.GpuMemoryFree, _ = strconv.ParseInt(infos[i], 10, 64)
|
i++
|
u.GpuMemoryUsed, _ = strconv.ParseInt(infos[i], 10, 64)
|
units = append(units, u)
|
}
|
|
if processStart {
|
proc := strings.Split(infos[i], ",")
|
p := Process{}
|
p.Pid, _ = strconv.Atoi(proc[0])
|
p.UsedMem, _ = strconv.Atoi(proc[1])
|
|
procs = append(procs, p)
|
}
|
|
if infos[i] == "proc" {
|
processStart = true
|
}
|
}
|
|
return &GpuInfo{count, units, procs}, nil
|
}
|
