feat(nvidia): simulate "get remapped rows" (optional) (#524)

>
{"level":"info","ts":"2025-03-13T06:37:35Z","caller":"nvml/nvml.go:39","msg":"injecting
remapped rows
pending","corrRows":0,"uncRows":10,"isPending":true,"failureOccurred":false}
> ...
>
{"level":"warn","ts":"2025-03-13T06:38:09Z","caller":"remapped-rows/component_output.go:141","msg":"suggested
actions","suggestedActions":"REBOOT_SYSTEM"}

Signed-off-by: Gyuho Lee <[email protected]>

Author: gyuho

cmd/gpud/command/accelerator.go

@@ -1,19 +1,15 @@
 package command
 
 import (
-	"context"
 	"fmt"
-	"time"
-
-	"github.com/leptonai/gpud/pkg/accelerator"
 
 	"github.com/urfave/cli"
+
+	"github.com/leptonai/gpud/pkg/accelerator"
 )
 
 func cmdAccelerator(cliContext *cli.Context) error {
-	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Minute)
-	defer cancel()
-	acceleratorType, productName, err := accelerator.DetectTypeAndProductName(ctx)
+	acceleratorType, productName, err := accelerator.DetectTypeAndProductName()
 	if err != nil {
 		return err
 	}

cmd/gpud/command/join.go

@@ -54,7 +54,7 @@ func cmdJoin(cliContext *cli.Context) (retErr error) {
 		return fmt.Errorf("error parsing cpu: %w", err)
 	}
 
-	_, productName, err := accelerator.DetectTypeAndProductName(rootCtx)
+	_, productName, err := accelerator.DetectTypeAndProductName()
 	if err != nil {
 		return err
 	}

pkg/accelerator/accelerator.go

@@ -1,10 +1,8 @@
 package accelerator
 
 import (
-	"context"
-
 	"github.com/leptonai/gpud/pkg/file"
-	nvidia_query "github.com/leptonai/gpud/pkg/nvidia-query"
+	"github.com/leptonai/gpud/pkg/nvidia-query/nvml"
 )
 
 type Type string
@@ -15,9 +13,9 @@ const (
 )
 
 // Returns the GPU type (e.g., "NVIDIA") and product name (e.g., "A100")
-func DetectTypeAndProductName(ctx context.Context) (Type, string, error) {
+func DetectTypeAndProductName() (Type, string, error) {
 	if _, err := file.LocateExecutable("nvidia-smi"); err == nil {
-		productName, err := nvidia_query.LoadGPUDeviceName(ctx)
+		productName, err := nvml.LoadGPUDeviceName()
 		if err != nil {
 			return TypeNVIDIA, "unknown", err
 		}

pkg/nvidia-query/detect.go

@@ -5,12 +5,9 @@ import (
 	"fmt"
 	"strings"
 
-	"github.com/NVIDIA/go-nvlib/pkg/nvlib/device"
-	"github.com/NVIDIA/go-nvml/pkg/nvml"
-
 	"github.com/leptonai/gpud/pkg/file"
 	"github.com/leptonai/gpud/pkg/log"
-	nvmlquery "github.com/leptonai/gpud/pkg/nvidia-query/nvml"
+	"github.com/leptonai/gpud/pkg/nvidia-query/nvml"
 	"github.com/leptonai/gpud/pkg/process"
 )
 
@@ -29,7 +26,7 @@ func GPUsInstalled(ctx context.Context) (bool, error) {
 
 	// now that we have the NVIDIA PCI devices,
 	// call NVML C-based API for NVML API
-	gpuDeviceName, err := LoadGPUDeviceName(ctx)
+	gpuDeviceName, err := nvml.LoadGPUDeviceName()
 	if err != nil {
 		if IsErrDeviceHandleUnknownError(err) {
 			log.Logger.Warnw("nvidia device handler failed for unknown error -- likely GPU has fallen off the bus or other Xid error", "error", err)
@@ -42,42 +39,6 @@ func GPUsInstalled(ctx context.Context) (bool, error) {
 	return true, nil
 }
 
-// Loads the product name of the NVIDIA GPU device.
-func LoadGPUDeviceName(ctx context.Context) (string, error) {
-	nvmlLib := nvmlquery.NewNVML()
-	if ret := nvmlLib.Init(); ret != nvml.SUCCESS {
-		return "", fmt.Errorf("failed to initialize NVML: %v", nvml.ErrorString(ret))
-	}
-
-	deviceLib := device.New(nvmlLib)
-
-	// do not check nvml lib if it is mocked
-	infoLib := nvmlquery.NewNVInfo(nvmlLib, deviceLib)
-	nvmlExists, nvmlExistsMsg := infoLib.HasNvml()
-	if !nvmlExists {
-		return "", fmt.Errorf("NVML not found: %s", nvmlExistsMsg)
-	}
-
-	// "NVIDIA Xid 79: GPU has fallen off the bus" may fail this syscall with:
-	// "error getting device handle for index '6': Unknown Error"
-	devices, err := deviceLib.GetDevices()
-	if err != nil {
-		return "", err
-	}
-
-	for _, d := range devices {
-		name, ret := d.GetName()
-		if ret != nvml.SUCCESS {
-			return "", fmt.Errorf("failed to get device name: %v", nvml.ErrorString(ret))
-		}
-		if name != "" {
-			return name, nil
-		}
-	}
-
-	return "", nil
-}
-
 // Lists all PCI devices that are compatible with NVIDIA.
 func ListNVIDIAPCIs(ctx context.Context) ([]string, error) {
 	lspciPath, err := file.LocateExecutable("lspci")

pkg/nvidia-query/nvml/clock_events.go

@@ -22,16 +22,14 @@ import (
 // ref. undefined symbol: nvmlDeviceGetCurrentClocksEventReasons for older nvidia drivers
 func ClockEventsSupported() (bool, error) {
 	nvmlLib := NewNVML()
-	if ret := nvmlLib.Init(); ret != nvml.SUCCESS {
+	if ret := nvmlLib.NVML().Init(); ret != nvml.SUCCESS {
 		return false, fmt.Errorf("failed to initialize NVML: %v", nvml.ErrorString(ret))
 	}
 	log.Logger.Debugw("successfully initialized NVML")
 
-	deviceLib := device.New(nvmlLib)
-
 	// "NVIDIA Xid 79: GPU has fallen off the bus" may fail this syscall with:
 	// "error getting device handle for index '6': Unknown Error"
-	devices, err := deviceLib.GetDevices()
+	devices, err := nvmlLib.Device().GetDevices()
 	if err != nil {
 		return false, err
 	}

pkg/nvidia-query/nvml/gpm.go

@@ -19,16 +19,14 @@ import (
 // Returns false if any device does not support GPM.
 func GPMSupported() (bool, error) {
 	nvmlLib := NewNVML()
-	if ret := nvmlLib.Init(); ret != nvml.SUCCESS {
+	if ret := nvmlLib.NVML().Init(); ret != nvml.SUCCESS {
 		return false, fmt.Errorf("failed to initialize NVML: %v", nvml.ErrorString(ret))
 	}
 	log.Logger.Debugw("successfully initialized NVML")
 
-	deviceLib := device.New(nvmlLib)
-
 	// "NVIDIA Xid 79: GPU has fallen off the bus" may fail this syscall with:
 	// "error getting device handle for index '6': Unknown Error"
-	devices, err := deviceLib.GetDevices()
+	devices, err := nvmlLib.Device().GetDevices()
 	if err != nil {
 		return false, err
 	}

pkg/nvidia-query/nvml/lib/lib.go

@@ -0,0 +1,130 @@
+// Package lib implements the NVIDIA Management Library (NVML) interface.
+// See https://docs.nvidia.com/deploy/nvml-api/nvml-api-reference.html#nvml-api-reference for more details.
+package lib
+
+import (
+	"github.com/NVIDIA/go-nvlib/pkg/nvlib/device"
+	nvinfo "github.com/NVIDIA/go-nvlib/pkg/nvlib/info"
+	"github.com/NVIDIA/go-nvml/pkg/nvml"
+)
+
+type Library interface {
+	NVML() nvml.Interface
+	Device() device.Interface
+	Info() nvinfo.Interface
+	Shutdown() nvml.Return
+}
+
+var _ Library = &nvmlInterface{}
+var _ nvml.Interface = &nvmlInterface{}
+
+type nvmlInterface struct {
+	nvml.Interface
+
+	dev  *devInterface
+	info nvinfo.Interface
+
+	initReturn        *nvml.Return
+	propertyExtractor nvinfo.PropertyExtractor
+}
+
+func (n *nvmlInterface) NVML() nvml.Interface {
+	return n
+}
+
+func (n *nvmlInterface) Device() device.Interface {
+	return n.dev
+}
+
+func (n *nvmlInterface) Info() nvinfo.Interface {
+	return n.info
+}
+
+func (n *nvmlInterface) Shutdown() nvml.Return {
+	return n.Interface.Shutdown()
+}
+
+// New creates a new NVML instance and returns nil if NVML is not supported.
+func New(opts ...OpOption) Library {
+	options := &Op{}
+	options.applyOpts(opts)
+
+	nvInterface := &nvmlInterface{
+		Interface: options.nvmlLib,
+
+		initReturn:        options.initReturn,
+		propertyExtractor: options.propertyExtractor,
+	}
+
+	devLib := device.New(nvInterface.Interface)
+	nvInterface.dev = &devInterface{
+		Interface:       devLib,
+		devices:         options.devicesToReturn,
+		getRemappedRows: options.devGetRemappedRows,
+	}
+
+	infoOpts := []nvinfo.Option{
+		nvinfo.WithNvmlLib(nvInterface),
+		nvinfo.WithDeviceLib(nvInterface.dev),
+	}
+	if nvInterface.propertyExtractor != nil {
+		infoOpts = append(infoOpts, nvinfo.WithPropertyExtractor(nvInterface.propertyExtractor))
+	}
+	nvInterface.info = nvinfo.New(infoOpts...)
+
+	return nvInterface
+}
+
+func (n *nvmlInterface) Init() nvml.Return {
+	if n.initReturn != nil {
+		return *n.initReturn
+	}
+	return n.Interface.Init()
+}
+
+var _ device.Interface = &devInterface{}
+
+type devInterface struct {
+	device.Interface
+	devices         []device.Device
+	getRemappedRows func() (int, int, bool, bool, nvml.Return)
+}
+
+func (d *devInterface) GetDevices() ([]device.Device, error) {
+	devs := d.devices
+
+	var err error
+	if len(devs) == 0 {
+		devs, err = d.Interface.GetDevices()
+	}
+
+	if err != nil {
+		return nil, err
+	}
+
+	updated := make([]device.Device, len(devs))
+	for i, dev := range devs {
+		updated[i] = &devDevInterface{
+			Device:          dev,
+			getRemappedRows: d.getRemappedRows,
+		}
+	}
+
+	return updated, nil
+}
+
+var _ device.Device = &devDevInterface{}
+
+type devDevInterface struct {
+	device.Device
+	getRemappedRows func() (int, int, bool, bool, nvml.Return)
+}
+
+func (d *devDevInterface) GetRemappedRows() (int, int, bool, bool, nvml.Return) {
+	// no injected remapped rows
+	// thus just passthrough to call the underlying device.Device.GetRemappedRows()
+	if d.getRemappedRows == nil {
+		return d.Device.GetRemappedRows()
+	}
+	return d.getRemappedRows()
+}

pkg/nvidia-query/nvml/lib/lib_test.go

@@ -0,0 +1,15 @@
+package lib
+
+import (
+	"testing"
+
+	"github.com/NVIDIA/go-nvml/pkg/nvml"
+	"github.com/stretchr/testify/assert"
+)
+
+func TestLibrary(t *testing.T) {
+	nv := New(
+		WithInitReturn(nvml.SUCCESS),
+	)
+	assert.Equal(t, nv.NVML().Init(), nvml.SUCCESS)
+}

pkg/nvidia-query/nvml/lib/options.go

@@ -0,0 +1,67 @@
+package lib
+
+import (
+	"github.com/NVIDIA/go-nvlib/pkg/nvlib/device"
+	nvinfo "github.com/NVIDIA/go-nvlib/pkg/nvlib/info"
+	"github.com/NVIDIA/go-nvml/pkg/nvml"
+)
+
+type Op struct {
+	nvmlLib nvml.Interface
+
+	initReturn        *nvml.Return
+	propertyExtractor nvinfo.PropertyExtractor
+	devicesToReturn   []device.Device
+
+	// ref. https://docs.nvidia.com/deploy/nvml-api/group__nvmlDeviceQueries.html#group__nvmlDeviceQueries_1g055e7c34f7f15b6ae9aac1dabd60870d
+	devGetRemappedRows func() (corrRows int, uncRows int, isPending bool, failureOccurred bool, ret nvml.Return)
+}
+
+type OpOption func(*Op)
+
+func (op *Op) applyOpts(opts []OpOption) {
+	for _, opt := range opts {
+		opt(op)
+	}
+	if op.nvmlLib == nil {
+		op.nvmlLib = nvml.New()
+	}
+}
+
+// Specifies the NVML library instance.
+// Otherwise, defaults to the NVML library instance returned by nvml.New().
+func WithNVML(nvmlLib nvml.Interface) OpOption {
+	return func(op *Op) {
+		op.nvmlLib = nvmlLib
+	}
+}
+
+// Specifies the return value of the NVML library's Init() function.
+// Otherwise, defaults to the return value of the NVML library's Init() function.
+func WithInitReturn(initReturn nvml.Return) OpOption {
+	return func(op *Op) {
+		op.initReturn = &initReturn
+	}
+}
+
+// Specifies the property extractor for the NVML library.
+func WithPropertyExtractor(propertyExtractor nvinfo.PropertyExtractor) OpOption {
+	return func(op *Op) {
+		op.propertyExtractor = propertyExtractor
+	}
+}
+
+func WithDevice(dev device.Device) OpOption {
+	return func(op *Op) {
+		op.devicesToReturn = append(op.devicesToReturn, dev)
+	}
+}
+
+// Specifies the function to get the remapped rows of the device.
+// Otherwise, defaults to the function returned by device.GetRemappedRows().
+// ref. https://docs.nvidia.com/deploy/nvml-api/group__nvmlDeviceQueries.html#group__nvmlDeviceQueries_1g055e7c34f7f15b6ae9aac1dabd60870d
+func WithDeviceGetRemappedRows(f func() (corrRows int, uncRows int, isPending bool, failureOccurred bool, ret nvml.Return)) OpOption {
+	return func(op *Op) {
+		op.devGetRemappedRows = f
+	}
+}