Replace nvml.Return with error

Signed-off-by: Evan Lezar <[email protected]>

examples/compute-processes/main.go

@@ -24,31 +24,31 @@ import (
 )
 
 func main() {
-	ret := nvml.Init()
-	if ret != nvml.SUCCESS {
-		log.Fatalf("Unable to initialize NVML: %v", nvml.ErrorString(ret))
+	err := nvml.Init()
+	if err != nil {
+		log.Fatalf("Unable to initialize NVML: %v", err)
 	}
 	defer func() {
-		ret := nvml.Shutdown()
-		if ret != nvml.SUCCESS {
-			log.Fatalf("Unable to shutdown NVML: %v", nvml.ErrorString(ret))
+		err := nvml.Shutdown()
+		if err != nil {
+			log.Fatalf("Unable to shutdown NVML: %v", err)
 		}
 	}()
 
-	count, ret := nvml.DeviceGetCount()
-	if ret != nvml.SUCCESS {
-		log.Fatalf("Unable to get device count: %v", nvml.ErrorString(ret))
+	count, err := nvml.DeviceGetCount()
+	if err != nil {
+		log.Fatalf("Unable to get device count: %v", err)
 	}
 
 	for di := 0; di < count; di++ {
-		device, ret := nvml.DeviceGetHandleByIndex(di)
-		if ret != nvml.SUCCESS {
-			log.Fatalf("Unable to get device at index %d: %v", di, nvml.ErrorString(ret))
+		device, err := nvml.DeviceGetHandleByIndex(di)
+		if err != nil {
+			log.Fatalf("Unable to get device at index %d: %v", di, err)
 		}
 
-		processInfos, ret := device.GetComputeRunningProcesses()
-		if ret != nvml.SUCCESS {
-			log.Fatalf("Unable to get process info for device at index %d: %v", di, nvml.ErrorString(ret))
+		processInfos, err := device.GetComputeRunningProcesses()
+		if err != nil {
+			log.Fatalf("Unable to get process info for device at index %d: %v", di, err)
 		}
 		fmt.Printf("Found %d processes on device %d\n", len(processInfos), di)
 		for pi, processInfo := range processInfos {

examples/devices/main.go

@@ -17,40 +17,49 @@
 package main
 
 import (
+	"errors"
 	"fmt"
 	"log"
 
 	"github.com/NVIDIA/go-nvml/pkg/nvml"
 )
 
 func main() {
-	ret := nvml.Init()
-	if ret != nvml.SUCCESS {
-		log.Fatalf("Unable to initialize NVML: %v", nvml.ErrorString(ret))
+	err := nvml.Init()
+	if err != nil {
+		log.Fatalf("Unable to initialize NVML: %v", err)
 	}
 	defer func() {
-		ret := nvml.Shutdown()
-		if ret != nvml.SUCCESS {
-			log.Fatalf("Unable to shutdown NVML: %v", nvml.ErrorString(ret))
+		err := nvml.Shutdown()
+		if err != nil {
+			log.Fatalf("Unable to shutdown NVML: %v", err)
 		}
 	}()
 
-	count, ret := nvml.DeviceGetCount()
-	if ret != nvml.SUCCESS {
-		log.Fatalf("Unable to get device count: %v", nvml.ErrorString(ret))
+	count, err := nvml.DeviceGetCount()
+	if err != nil {
+		log.Fatalf("Unable to get device count: %v", err)
 	}
 
 	for i := 0; i < count; i++ {
-		device, ret := nvml.DeviceGetHandleByIndex(i)
-		if ret != nvml.SUCCESS {
-			log.Fatalf("Unable to get device at index %d: %v", i, nvml.ErrorString(ret))
+		device, err := nvml.DeviceGetHandleByIndex(i)
+		if err != nil {
+			log.Fatalf("Unable to get device at index %d: %v", i, err)
 		}
 
-		uuid, ret := device.GetUUID()
-		if ret != nvml.SUCCESS {
-			log.Fatalf("Unable to get uuid of device at index %d: %v", i, nvml.ErrorString(ret))
+		uuid, err := device.GetUUID()
+		if err != nil {
+			log.Fatalf("Unable to get uuid of device at index %d: %v", i, err)
 		}
-
 		fmt.Printf("%v\n", uuid)
+
+		current, pending, err := device.GetMigMode()
+		if errors.Is(err, nvml.ERROR_NOT_SUPPORTED) {
+			fmt.Printf("MIG is not supported for device %v\n", i)
+
+		} else if err != nil {
+			log.Fatalf("Error getting MIG mode for device at index %d: %v", i, err)
+		}
+		fmt.Printf("MIG mode for device %v: current=%v pending=%v\n", i, current, pending)
 	}
 }

gen/nvml/nvml.yml

@@ -59,6 +59,9 @@ TRANSLATOR:
       - {action: replace, from: "^nvml"}
       - {action: replace, from: "_t$"}
       - {transform: export}
+      # We explicitly unexport the SUCCESS and ERROR_ constants to allow for go-native error implementations.
+      - {action: replace, from: "^SUCCESS", to: "nvmlSUCCESS"}
+      - {action: replace, from: "^ERROR", to: "nvmlERROR"}
     type:
       - {action: accept, from: "^nvml"}
       - {action: replace, from: "^nvml"}

pkg/nvml/api.go

@@ -22,7 +22,7 @@ package nvml
 // with the list of excluded methods for the Interface type in
 // gen/nvml/generateapi.go. In the future we should automate this.
 //
-//go:generate moq -out mock/extendedinterface.go -pkg mock . ExtendedInterface:ExtendedInterface
+//go:generate moq -rm -out mock/extendedinterface.go -pkg mock . ExtendedInterface:ExtendedInterface
 type ExtendedInterface interface {
 	LookupSymbol(string) error
 }