add the number of cpus used by running jobs per user

Author: vpenso

users.go

@@ -20,12 +20,13 @@ import (
         "os/exec"
         "log"
         "strings"
+        "strconv"
         "regexp"
         "github.com/prometheus/client_golang/prometheus"
 )
 
 func UsersData() []byte {
-        cmd := exec.Command("squeue", "-h", "-o '%A|%u|%T|%C'")
+        cmd := exec.Command("squeue", "-h", "-o %A|%u|%T|%C")
         stdout, err := cmd.StdoutPipe()
 	if err != nil {
 		log.Fatal(err)
@@ -43,6 +44,7 @@ func UsersData() []byte {
 type UserJobMetrics struct {
         pending float64
         running float64
+        running_cpus float64
         suspended float64
 }
 
@@ -54,10 +56,11 @@ func ParseUsersMetrics(input []byte) map[string]*UserJobMetrics {
                         user := strings.Split(line,"|")[1]
                         _,key := users[user]
                         if !key {
-                                users[user] = &UserJobMetrics{0,0,0}
+                                users[user] = &UserJobMetrics{0,0,0,0}
                         }
                         state := strings.Split(line,"|")[2]
                         state = strings.ToLower(state)
+                        cpus,_ := strconv.ParseFloat(strings.Split(line,"|")[3],64)
                         pending := regexp.MustCompile(`^pending`)
                         running := regexp.MustCompile(`^running`)
                         suspended := regexp.MustCompile(`^suspended`)
@@ -66,6 +69,7 @@ func ParseUsersMetrics(input []byte) map[string]*UserJobMetrics {
                                 users[user].pending++
                         case running.MatchString(state) == true:
                                 users[user].running++
+                                users[user].running_cpus += cpus
                         case suspended.MatchString(state) == true:
                                 users[user].suspended++
                         }
@@ -77,6 +81,7 @@ func ParseUsersMetrics(input []byte) map[string]*UserJobMetrics {
 type UsersCollector struct {
         pending *prometheus.Desc
         running *prometheus.Desc
+        running_cpus *prometheus.Desc
         suspended *prometheus.Desc
 }
 
@@ -85,13 +90,15 @@ func NewUsersCollector() *UsersCollector {
         return &UsersCollector {
                 pending: prometheus.NewDesc("slurm_user_jobs_pending", "Pending jobs for user", labels, nil), 
                 running: prometheus.NewDesc("slurm_user_jobs_running", "Running jobs for user", labels, nil),
+                running_cpus: prometheus.NewDesc("slurm_user_cpus_running", "Running cpus for user", labels, nil),
                 suspended: prometheus.NewDesc("slurm_user_jobs_suspended", "Suspended jobs for user", labels, nil),
         }
 }
 
 func (uc *UsersCollector) Describe(ch chan<- *prometheus.Desc) {
         ch <- uc.pending
         ch <- uc.running
+        ch <- uc.running_cpus
         ch <- uc.suspended
 }
 
@@ -100,6 +107,7 @@ func (uc *UsersCollector) Collect(ch chan<- prometheus.Metric) {
         for u := range um {
                 ch <- prometheus.MustNewConstMetric(uc.pending, prometheus.GaugeValue, um[u].pending, u)
                 ch <- prometheus.MustNewConstMetric(uc.running, prometheus.GaugeValue, um[u].running, u)
+                ch <- prometheus.MustNewConstMetric(uc.running_cpus, prometheus.GaugeValue, um[u].running_cpus, u)
                 ch <- prometheus.MustNewConstMetric(uc.suspended, prometheus.GaugeValue, um[u].suspended, u)
         }
 }