Refactor. Improve parsing.

echoblag · Jun 12, 2021 · bfa82ba · bfa82ba
1 parent 9330fd5
commit bfa82ba
Show file tree

Hide file tree

Showing 7 changed files with 231 additions and 171 deletions.
diff --git a/cmd/nvidia-gpu-exporter/main.go b/cmd/nvidia-gpu-exporter/main.go
@@ -33,7 +33,13 @@ func main() {
 
 	logConfig := promlog.Config{}
 	logger := promlog.New(&logConfig)
-	e := exporter.New(exporter.DefaultPrefix, nvidiaSmiCommand, exporter.DefaultMetrics, logger)
+
+	e, err := exporter.New(exporter.DefaultPrefix, nvidiaSmiCommand, exporter.DefaultQueryFieldNames, logger)
+	if err != nil {
+		_ = level.Error(logger).Log("msg", "Error on creating exporter", "err", err)
+		os.Exit(1)
+	}
+
 	prometheus.MustRegister(e)
 
 	metricsPath := "/metrics"

diff --git a/cmd/nvidia-smi-query-fields-parser/main.go b/cmd/nvidia-smi-query-fields-parser/main.go
@@ -1,43 +1,23 @@
 package main
 
 import (
-	"bytes"
 	"flag"
 	"fmt"
+	"nvidia-smi-exporter/internal/exporter"
 	"os"
-	"os/exec"
-	"regexp"
 	"strings"
 )
 
-var (
-	fieldRegex = regexp.MustCompile(`(?m)\n\s*\n^"([^"]+)"`)
-)
-
 func main() {
 	var nvidiaSmiCommand string
 	flag.StringVar(&nvidiaSmiCommand, "nvidia-smi-command", "nvidia-smi", "command to run nvidia-smi")
 	flag.Parse()
 
-	cmdAndArgs := strings.Fields(nvidiaSmiCommand)
-	cmdAndArgs = append(cmdAndArgs, "--help-query-gpu")
-	cmd := exec.Command(cmdAndArgs[0], cmdAndArgs[1:]...)
-
-	var stdout bytes.Buffer
-	cmd.Stdout = &stdout
-	err := cmd.Run()
+	fields, err := exporter.ParseQueryFields(nvidiaSmiCommand)
 	if err != nil {
 		fmt.Printf("error: %v\n", err)
 		os.Exit(1)
 	}
 
-	out := stdout.String()
-	found := fieldRegex.FindAllStringSubmatch(out, -1)
-
-	var fields []string
-	for _, ss := range found {
-		fields = append(fields, ss[1])
-	}
-
 	fmt.Printf("Fields: %s\n", strings.Join(fields, ","))
 }
diff --git a/internal/exporter/csv.go b/internal/exporter/csv.go
@@ -0,0 +1,77 @@
+package exporter
+
+import "strings"
+
+type table struct {
+	rows                  []row
+	returnedFieldNames    []string
+	queryFieldNameToCells map[string][]cell
+}
+
+type row struct {
+	queryFieldNameToCells map[string]cell
+	cells                 []cell
+}
+
+type cell struct {
+	queryFieldName    string
+	returnedFieldName string
+	rawValue          string
+}
+
+func parseCSVIntoTable(queryResult string, queryFieldNames []string) table {
+	lines := strings.Split(strings.TrimSpace(queryResult), "\n")
+	titlesLine := lines[0]
+	valuesLines := lines[1:]
+	returnedFieldNames := parseCSVLine(titlesLine)
+
+	numCols := len(queryFieldNames)
+	numRows := len(valuesLines)
+
+	rows := make([]row, numRows)
+
+	queryFieldNameToCells := make(map[string][]cell)
+	for _, queryFieldName := range queryFieldNames {
+		queryFieldNameToCells[queryFieldName] = make([]cell, numRows)
+	}
+
+	for rowIndex, valuesLine := range valuesLines {
+		queryFieldNameToCell := make(map[string]cell, numCols)
+		cells := make([]cell, numCols)
+		rawValues := parseCSVLine(valuesLine)
+		for colIndex, rawValue := range rawValues {
+			queryFieldName := queryFieldNames[colIndex]
+			gm := cell{
+				queryFieldName:    queryFieldName,
+				returnedFieldName: returnedFieldNames[colIndex],
+				rawValue:          rawValue,
+			}
+			queryFieldNameToCell[queryFieldName] = gm
+			cells[colIndex] = gm
+			queryFieldNameToCells[queryFieldName][rowIndex] = gm
+		}
+
+		gmc := row{
+			queryFieldNameToCells: queryFieldNameToCell,
+			cells:                 cells,
+		}
+
+		rows[rowIndex] = gmc
+
+	}
+
+	return table{
+		rows:                  rows,
+		returnedFieldNames:    returnedFieldNames,
+		queryFieldNameToCells: queryFieldNameToCells,
+	}
+}
+
+func parseCSVLine(line string) []string {
+	values := strings.Split(line, ",")
+	result := make([]string, len(values))
+	for i, field := range values {
+		result[i] = strings.TrimSpace(field)
+	}
+	return result
+}