@cryptotaxi247 / netdata-1
Code Commits Issues Pulls Sessions CI/CD
master
master
go 224 lines 4.2 KB
Raw
1 // SPDX-License-Identifier: GPL-3.0-or-later
2
3 package nvidia_smi
4
5 import (
6 "bufio"
7 "bytes"
8 "errors"
9 "os/exec"
10 "path/filepath"
11 "runtime"
12 "strconv"
13 "sync"
14 "time"
15
16 "github.com/netdata/netdata/go/plugins/logger"
17 "github.com/netdata/netdata/go/plugins/pkg/buildinfo"
18 "github.com/netdata/netdata/go/plugins/plugin/go.d/pkg/ndexec"
19 )
20
21 type nvidiaSmiBinary interface {
22 queryGPUInfo() ([]byte, error)
23 stop() error
24 }
25
26 func newNvidiaSmiBinary(path string, cfg Config, log *logger.Logger) (nvidiaSmiBinary, error) {
27 if !cfg.LoopMode {
28 if runtime.GOOS == "windows" {
29 return &nvidiaSmiDirectExec{
30 Logger: log,
31 binPath: path,
32 timeout: cfg.Timeout.Duration(),
33 }, nil
34 }
35 return &nvidiaSmiExec{
36 Logger: log,
37 binPath: path,
38 timeout: cfg.Timeout.Duration(),
39 }, nil
40 }
41
42 smi := &nvidiaSmiLoopExec{
43 Logger: log,
44 binPath: path,
45 updateEvery: cfg.UpdateEvery,
46 firstSampleTimeout: cfg.Timeout.Duration(),
47 }
48
49 if err := smi.run(); err != nil {
50 return nil, err
51 }
52
53 return smi, nil
54 }
55
56 // nvidiaSmiExec executes nvidia-smi via nd-run (Linux/BSD)
57 type nvidiaSmiExec struct {
58 *logger.Logger
59
60 binPath string
61 timeout time.Duration
62 }
63
64 func (e *nvidiaSmiExec) queryGPUInfo() ([]byte, error) {
65 return ndexec.RunUnprivileged(e.Logger, e.timeout, e.binPath, "-q", "-x")
66 }
67
68 func (e *nvidiaSmiExec) stop() error { return nil }
69
70 // nvidiaSmiDirectExec executes nvidia-smi directly (Windows)
71 type nvidiaSmiDirectExec struct {
72 *logger.Logger
73
74 binPath string
75 timeout time.Duration
76 }
77
78 func (e *nvidiaSmiDirectExec) queryGPUInfo() ([]byte, error) {
79 return ndexec.RunDirect(e.Logger, e.timeout, e.binPath, "-q", "-x")
80 }
81
82 func (e *nvidiaSmiDirectExec) stop() error { return nil }
83
84 type nvidiaSmiLoopExec struct {
85 *logger.Logger
86
87 binPath string
88
89 updateEvery int
90 firstSampleTimeout time.Duration
91
92 cmd *exec.Cmd
93 done chan struct{}
94
95 mux sync.Mutex
96 lastSample string
97 }
98
99 func (e *nvidiaSmiLoopExec) queryGPUInfo() ([]byte, error) {
100 select {
101 case <-e.done:
102 return nil, errors.New("process has already exited")
103 default:
104 }
105
106 e.mux.Lock()
107 defer e.mux.Unlock()
108
109 return []byte(e.lastSample), nil
110 }
111
112 func (e *nvidiaSmiLoopExec) run() error {
113 secs := min(e.updateEvery, 5)
114
115 ndrunPath := filepath.Join(buildinfo.NetdataBinDir, "nd-run")
116 cmd := exec.Command(ndrunPath, e.binPath, "-q", "-x", "-l", strconv.Itoa(secs))
117
118 e.Debugf("executing '%s'", cmd)
119
120 r, err := cmd.StdoutPipe()
121 if err != nil {
122 return err
123 }
124
125 if err := cmd.Start(); err != nil {
126 return err
127 }
128
129 firstSample := make(chan struct{}, 1)
130 done := make(chan struct{})
131 e.cmd = cmd
132 e.done = done
133
134 go func() {
135 defer close(done)
136
137 var buf bytes.Buffer
138 var insideLog bool
139 var emptyRows int64
140 var outsideLogRows int64
141
142 const unexpectedRowsLimit = 500
143
144 sc := bufio.NewScanner(r)
145
146 for sc.Scan() {
147 line := sc.Text()
148
149 if !insideLog {
150 outsideLogRows++
151 } else {
152 outsideLogRows = 0
153 }
154
155 if line == "" {
156 emptyRows++
157 } else {
158 emptyRows = 0
159 }
160
161 if outsideLogRows >= unexpectedRowsLimit || emptyRows >= unexpectedRowsLimit {
162 e.Errorf("unexpected output from nvidia-smi loop: outside log rows %d, empty rows %d", outsideLogRows, emptyRows)
163 break
164 }
165
166 switch {
167 case line == "<nvidia_smi_log>":
168 insideLog = true
169 buf.Reset()
170
171 buf.WriteString(line)
172 buf.WriteByte('\n')
173 case line == "</nvidia_smi_log>":
174 insideLog = false
175
176 buf.WriteString(line)
177
178 e.mux.Lock()
179 e.lastSample = buf.String()
180 e.mux.Unlock()
181
182 buf.Reset()
183
184 select {
185 case firstSample <- struct{}{}:
186 default:
187 }
188 case insideLog:
189 buf.WriteString(line)
190 buf.WriteByte('\n')
191 default:
192 continue
193 }
194 }
195 }()
196
197 select {
198 case <-e.done:
199 _ = e.stop()
200 return errors.New("process exited before the first sample was collected")
201 case <-time.After(e.firstSampleTimeout):
202 _ = e.stop()
203 return errors.New("timed out waiting for first sample")
204 case <-firstSample:
205 return nil
206 }
207 }
208
209 func (e *nvidiaSmiLoopExec) stop() error {
210 if e.cmd == nil || e.cmd.Process == nil {
211 return nil
212 }
213
214 _ = e.cmd.Process.Kill()
215 _ = e.cmd.Wait()
216 e.cmd = nil
217
218 select {
219 case <-e.done:
220 return nil
221 case <-time.After(time.Second * 2):
222 return errors.New("timed out waiting for process to exit")
223 }
224 }