File tree Expand file tree Collapse file tree 1 file changed +12
-4
lines changed Expand file tree Collapse file tree 1 file changed +12
-4
lines changed Original file line number Diff line number Diff line change @@ -21,6 +21,7 @@ import (
2121 "io/ioutil"
2222 "os/exec"
2323 "strings"
24+ "strconv"
2425)
2526
2627type GPUsMetrics struct {
@@ -47,19 +48,26 @@ func ParseOtherGPUs() float64 {
4748}
4849
4950func 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
6573func ParseGPUsMetrics () * GPUsMetrics {
You can’t perform that action at this time.
0 commit comments