-
Notifications
You must be signed in to change notification settings - Fork 22
/
gpu.go
98 lines (84 loc) · 2.67 KB
/
gpu.go
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
96
97
98
//******************************************************************
//Copyright 2018 eBay Inc.
//Architect/Developer: Deepak Vasthimal
//Licensed under the Apache License, Version 2.0 (the "License");
//you may not use this file except in compliance with the License.
//You may obtain a copy of the License at
// https://www.apache.org/licenses/LICENSE-2.0
//Unless required by applicable law or agreed to in writing, software
//distributed under the License is distributed on an "AS IS" BASIS,
//WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
//See the License for the specific language governing permissions and
//limitations under the License.
//******************************************************************
package nvidia
import (
"encoding/csv"
"errors"
"io"
"os/exec"
"strconv"
"strings"
"github.com/elastic/beats/libbeat/common"
)
//GPUUtilization provides interface to utilization metrics and state of GPU.
type GPUUtilization interface {
command(env string) *exec.Cmd
run(cmd *exec.Cmd, gpuCount int, query string, action Action) ([]common.MapStr, error)
}
//Utilization implements one flavour of GPUCount interface.
type Utilization struct {
}
//newUtilization returns instance of Utilization
func newUtilization() Utilization {
return Utilization{}
}
func (g Utilization) command(env string, query string) *exec.Cmd {
if env == "test" {
return exec.Command("localnvidiasmi")
}
return exec.Command("nvidia-smi", "--query-gpu="+query, "--format=csv")
}
//Run the nvidiasmi command to collect GPU metrics
//Parse output and return events.
func (g Utilization) run(cmd *exec.Cmd, gpuCount int, query string, action Action) ([]common.MapStr, error) {
reader := action.start(cmd)
gpuIndex := 0
events := make([]common.MapStr, gpuCount, 2*gpuCount)
for {
line, err := reader.ReadString('\n')
if err == io.EOF {
break
}
// Ignore header
if strings.Contains(line, "utilization") {
continue
}
if len(line) == 0 {
return nil, errors.New("Unable to fetch any events from nvidia-smi: Error " + err.Error())
}
// Remove units put by nvidia-smi
line = strings.Replace(line, " %", "", -1)
line = strings.Replace(line, " MiB", "", -1)
line = strings.Replace(line, " P", "", -1)
line = strings.Replace(line, " ", "", -1)
r := csv.NewReader(strings.NewReader(line))
record, err := r.Read()
if err == io.EOF {
break
}
headers := strings.Split(query, ",")
event := common.MapStr{
"gpuIndex": gpuIndex,
"type": "nvidiagpubeat",
}
for i := 0; i < len(record); i++ {
value, _ := strconv.Atoi(record[i])
event.Put(headers[i], value)
}
events[gpuIndex] = event
gpuIndex++
}
cmd.Wait()
return events, nil
}