Refactoring devices in support of GPU

4 files changed, 82 insertions, 0 deletions

diff --git a/devices/cpu.go b/devices/cpu.go
new file mode 100644
index 0000000..b35bdd3
--- /dev/null
+++ b/devices/cpu.go
@@ -0,0 +1,44 @@
+package devices
+
+import (
+	"time"
+)
+
+var deviceCounts []func(bool) (int, error)
+var devicePercents []func(time.Duration, bool) ([]float64, error)
+var numDevices int
+
+// Counts returns the number of CPUs registered.
+//
+// logical tells Counts to count the logical cores; this may be ignored for
+// some devices.
+func Counts(logical bool) (int, error) {
+	var rv int
+	var re error
+	for _, d := range deviceCounts {
+		r, err := d(logical)
+		if err != nil {
+			return rv, re
+		}
+		rv += r
+	}
+	return rv, re
+}
+
+// Percent calculates the percentage of cpu used either per CPU or combined.
+// Returns one value per cpu, or a single value if percpu is set to false.
+func Percent(interval time.Duration, combined bool) ([]float64, error) {
+	var rvs []float64
+	rvs = make([]float64, 0, numDevices)
+	for _, f := range devicePercents {
+		vs, err := f(interval, combined)
+		if err != nil {
+			return rvs, err
+		}
+		for _, v := range vs {
+			rvs = append(rvs, v)
+		}
+	}
+	numDevices = len(rvs)
+	return rvs, nil
+}
diff --git a/devices/cpu_cpu.go b/devices/cpu_cpu.go
new file mode 100644
index 0000000..f5e6956
--- /dev/null
+++ b/devices/cpu_cpu.go
@@ -0,0 +1,10 @@
+package devices
+
+import (
+	psCpu "github.com/shirou/gopsutil/cpu"
+)
+
+func init() {
+	deviceCounts = append(deviceCounts, psCpu.Counts)
+	devicePercents = append(devicePercents, psCpu.Percent)
+}
diff --git a/devices/cpu_nvidia.go b/devices/cpu_nvidia.go
new file mode 100644
index 0000000..ff4691c
--- /dev/null
+++ b/devices/cpu_nvidia.go
@@ -0,0 +1,19 @@
+package devices
+
+import (
+	"github.com/NVIDIA/gpu-monitoring-tools/bindings/go/nvml"
+	"time"
+)
+
+func init() {
+	if nvml.Init() == nil {
+		shutdownFuncs = append(shutdownFuncs, func() { nvml.Shutdown() })
+		deviceCounts = append(deviceCounts, func(b bool) (int, error) {
+			r, e := nvml.GetDeviceCount()
+			return int(r), e
+		})
+		devicePercents = append(devicePercents, func(t time.Duration, b bool) ([]float64, error) {
+			return nil, nil
+		})
+	}
+}
diff --git a/devices/devices.go b/devices/devices.go
new file mode 100644
index 0000000..dabb900
--- /dev/null
+++ b/devices/devices.go
@@ -0,0 +1,9 @@
+package devices
+
+var shutdownFuncs []func()
+
+func Shutdown() {
+	for _, f := range shutdownFuncs {
+		f()
+	}
+}