Skip to content

Commit a5d5a99

Browse files
author
Joeri Hermans
committed
Update retrieval of total number of GPUs
1 parent 1b452ee commit a5d5a99

File tree

1 file changed

+12
-4
lines changed

1 file changed

+12
-4
lines changed

gpus.go

Lines changed: 12 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@ import (
2121
"io/ioutil"
2222
"os/exec"
2323
"strings"
24+
"strconv"
2425
)
2526

2627
type GPUsMetrics struct {
@@ -47,19 +48,26 @@ func ParseOtherGPUs() float64 {
4748
}
4849

4950
func ParseTotalGPUs() float64 {
51+
var num_gpus = 0.0
52+
5053
args := []string{"-h", "-o \"%n %G\""}
5154
output := string(Execute("sinfo", args))
5255
if len(output) > 0 {
5356
for _, line := range strings.Split(output, "\n") {
5457
if len(line) > 0 {
55-
log.Infof("Line %s: ", line)
56-
descriptor := strings.Split(line, " ")[0]
57-
log.Infof("Descriptor %s: ", descriptor)
58+
line = strings.Trim(line, "\"")
59+
descriptor := strings.Fields(line)[1]
60+
descriptor = strings.TrimPrefix(descriptor, "gpu:")
61+
descriptor = strings.Split(descriptor, "(")[0]
62+
node_gpus, err := strconv.ParseFloat(descriptor, 64)
63+
if err != nil {
64+
num_gpus += node_gpus
65+
}
5866
}
5967
}
6068
}
6169

62-
return 0.0
70+
return num_gpus
6371
}
6472

6573
func ParseGPUsMetrics() *GPUsMetrics {

0 commit comments

Comments
 (0)