zhangmeng
2019-05-31 6f0c11092b0a39265c6af723571906e4f834706d
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
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"
    "strconv"
    "strings"
    "unsafe"
)
 
// IdleGPU pass needed gpu memory size test if satisfy
func IdleGPU(memSize int) int {
    return int(C.get_idle_gpu(C.int(memSize)))
}
 
// GpuUnitInfo gpu card info
type GpuUnitInfo struct {
    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"`
}
 
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 count int
    for i := 0; i < length; i++ {
        if infos[i] == "gpu" {
            count++
 
            i++
            u := GpuUnitInfo{}
            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)
        }
    }
 
    return &GpuInfo{count, units}, nil
}