fix(nvml): fix hw slowdown events db writes, support simulate hw slowdown flag (#526)

```go
const (
	EnvMockAllSuccess              = "GPUD_NVML_MOCK_ALL_SUCCESS"
	EnvInjectRemapedRowsPending    = "GPUD_NVML_INJECT_REMAPPED_ROWS_PENDING"
	EnvInjectClockEventsHwSlowdown = "GPUD_NVML_INJECT_CLOCK_EVENTS_HW_SLOWDOWN"
)
```

Tested

```
    "component": "accelerator-nvidia-hw-slowdown",
    "startTime": "2024-11-15T04:48:56Z",
    "endTime": "2025-03-13T11:34:36.809052648Z",
    "events": [
      {
        "time": "2025-03-13T11:34:00Z",
        "name": "hw_slowdown",
        "type": "Warning",
        "message": "GPU-49004c5e-a258-143f-1c5c-319f2db1a1f1: HW Power Brake Slowdown (reducing the core clocks by a factor of 2 or more) is engaged (External Power Brake Assertion being triggered) ('HW Power Brake Slowdown' in nvidia-smi --query) (nvml), GPU-49004c5e-a258-143f-1c5c-319f2db1a1f1: HW Slowdown is engaged due to high temperature, power brake assertion, or high power draw ('HW Slowdown: Active' in nvidia-smi --query) (nvml), GPU-49004c5e-a258-143f-1c5c-319f2db1a1f1: HW Thermal Slowdown (reducing the core clocks by a factor of 2 or more) is engaged (temperature being too high) ('HW Thermal Slowdown' in nvidia-smi --query) (nvml)",
        "extra_info": {
          "data_source": "nvml",
          "gpu_uuid": "GPU-49004c5e-a258-143f-1c5c-319f2db1a1f1"
        }
      },
      {
        "time": "2025-03-13T11:33:00Z",
        "name": "hw_slowdown",
        "type": "Warning",
        "message": "GPU-49004c5e-a258-143f-1c5c-319f2db1a1f1: HW Power Brake Slowdown (reducing the core clocks by a factor of 2 or more) is engaged (External Power Brake Assertion being triggered) ('HW Power Brake Slowdown' in nvidia-smi --query) (nvml), GPU-49004c5e-a258-143f-1c5c-319f2db1a1f1: HW Slowdown is engaged due to high temperature, power brake assertion, or high power draw ('HW Slowdown: Active' in nvidia-smi --query) (nvml), GPU-49004c5e-a258-143f-1c5c-319f2db1a1f1: HW Thermal Slowdown (reducing the core clocks by a factor of 2 or more) is engaged (temperature being too high) ('HW Thermal Slowdown' in nvidia-smi --query) (nvml)",
        "extra_info": {
          "data_source": "nvml",
          "gpu_uuid": "GPU-49004c5e-a258-143f-1c5c-319f2db1a1f1"
        }
      }
    ]
  },
```

---------

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

Author: gyuho

e2e/e2e_test.go

@@ -22,8 +22,8 @@ import (
 	v1 "github.com/leptonai/gpud/api/v1"
 	client_v1 "github.com/leptonai/gpud/client/v1"
 	mocklspci "github.com/leptonai/gpud/e2e/mock/lspci"
-	mocknvml "github.com/leptonai/gpud/e2e/mock/nvml"
 	"github.com/leptonai/gpud/pkg/errdefs"
+	nvml_lib "github.com/leptonai/gpud/pkg/nvidia-query/nvml/lib"
 	"github.com/leptonai/gpud/pkg/server"
 )
 
@@ -53,8 +53,8 @@ var _ = Describe("[GPUD E2E]", Ordered, func() {
 	gCtx, gCancel := context.WithTimeout(context.Background(), time.Minute*8)
 
 	BeforeAll(func() {
-		err = os.Setenv(mocknvml.EnvNVMLMock, "true")
-		Expect(err).NotTo(HaveOccurred(), "failed to set GPUD_MOCK_NVML")
+		err = os.Setenv(nvml_lib.EnvMockAllSuccess, "true")
+		Expect(err).NotTo(HaveOccurred(), "failed to set "+nvml_lib.EnvMockAllSuccess)
 
 		By("mock lspci")
 		err = mocklspci.Mock(mocklspci.NormalOutput)

e2e/mock/nvml/nvml_mock.go

@@ -161,6 +161,7 @@ func (t *table) Insert(ctx context.Context, ev components.Event) error {
 	return insertEvent(ctx, t.dbRW, t.table, ev)
 }
 
+// Find returns nil if the event is not found.
 func (t *table) Find(ctx context.Context, ev components.Event) (*components.Event, error) {
 	return findEvent(ctx, t.dbRO, t.table, ev)
 }

pkg/eventstore/database.go

@@ -18,6 +18,7 @@ type Store interface {
 type Bucket interface {
 	Name() string
 	Insert(ctx context.Context, ev components.Event) error
+	// Find returns nil if the event is not found.
 	Find(ctx context.Context, ev components.Event) (*components.Event, error)
 	// Get queries the event in the descending order of timestamp (latest event first).
 	Get(ctx context.Context, since time.Time) ([]components.Event, error)

pkg/eventstore/types.go

@@ -15,13 +15,14 @@ import (
 	"github.com/leptonai/gpud/components"
 	"github.com/leptonai/gpud/pkg/common"
 	"github.com/leptonai/gpud/pkg/log"
+	nvml_lib "github.com/leptonai/gpud/pkg/nvidia-query/nvml/lib"
 )
 
 // Returns true if clock events is supported by all devices.
 // Returns false if any device does not support clock events.
 // ref. undefined symbol: nvmlDeviceGetCurrentClocksEventReasons for older nvidia drivers
 func ClockEventsSupported() (bool, error) {
-	nvmlLib := NewNVML()
+	nvmlLib := nvml_lib.NewDefault()
 	if ret := nvmlLib.NVML().Init(); ret != nvml.SUCCESS {
 		return false, fmt.Errorf("failed to initialize NVML: %v", nvml.ErrorString(ret))
 	}

pkg/nvidia-query/nvml/clock_events.go

@@ -14,8 +14,9 @@ import (
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 
 	"github.com/leptonai/gpud/components"
-	mocknvml "github.com/leptonai/gpud/e2e/mock/nvml"
 	"github.com/leptonai/gpud/pkg/common"
+	nvml_lib "github.com/leptonai/gpud/pkg/nvidia-query/nvml/lib"
+	nvml_lib_mock "github.com/leptonai/gpud/pkg/nvidia-query/nvml/lib/mock"
 	"github.com/leptonai/gpud/pkg/nvidia-query/nvml/testutil"
 )
 
@@ -392,17 +393,16 @@ func TestClockEventsSupported(t *testing.T) {
 				},
 			}
 
-			// Set up mock NVML environment
-			err := os.Setenv(mocknvml.EnvNVMLMock, "true")
+			err := os.Setenv(nvml_lib.EnvMockAllSuccess, "true")
 			if err != nil {
 				t.Fatalf("failed to set mock NVML environment: %v", err)
 			}
-			defer os.Unsetenv(mocknvml.EnvNVMLMock)
+			defer os.Unsetenv(nvml_lib.EnvMockAllSuccess)
 
 			// Replace the mock instance
-			originalMockInstance := mocknvml.MockInstance
-			mocknvml.MockInstance = mockNVML
-			defer func() { mocknvml.MockInstance = originalMockInstance }()
+			originalMockInstance := nvml_lib_mock.AllSuccessInterface
+			nvml_lib_mock.AllSuccessInterface = mockNVML
+			defer func() { nvml_lib_mock.AllSuccessInterface = originalMockInstance }()
 
 			result, err := ClockEventsSupported()
 			if tt.expectError {
@@ -663,12 +663,11 @@ func TestClockEventsWithNilPointer(t *testing.T) {
 func TestClockEventsSupportedWithMockedNVML(t *testing.T) {
 	// Test with initialization failure
 	t.Run("nvml initialization failure", func(t *testing.T) {
-		// Set up mock NVML environment
-		err := os.Setenv(mocknvml.EnvNVMLMock, "true")
+		err := os.Setenv(nvml_lib.EnvMockAllSuccess, "true")
 		if err != nil {
 			t.Fatalf("failed to set mock NVML environment: %v", err)
 		}
-		defer os.Unsetenv(mocknvml.EnvNVMLMock)
+		defer os.Unsetenv(nvml_lib.EnvMockAllSuccess)
 
 		// Mock NVML with init failure
 		mockNVML := &mock.Interface{
@@ -678,9 +677,9 @@ func TestClockEventsSupportedWithMockedNVML(t *testing.T) {
 		}
 
 		// Replace the mock instance
-		originalMockInstance := mocknvml.MockInstance
-		mocknvml.MockInstance = mockNVML
-		defer func() { mocknvml.MockInstance = originalMockInstance }()
+		originalMockInstance := nvml_lib_mock.AllSuccessInterface
+		nvml_lib_mock.AllSuccessInterface = mockNVML
+		defer func() { nvml_lib_mock.AllSuccessInterface = originalMockInstance }()
 
 		// Call function
 		result, err := ClockEventsSupported()
@@ -691,12 +690,11 @@ func TestClockEventsSupportedWithMockedNVML(t *testing.T) {
 
 	// Test with device initialization but GetDevices failure
 	t.Run("device get failure", func(t *testing.T) {
-		// Set up mock NVML environment
-		err := os.Setenv(mocknvml.EnvNVMLMock, "true")
+		err := os.Setenv(nvml_lib.EnvMockAllSuccess, "true")
 		if err != nil {
 			t.Fatalf("failed to set mock NVML environment: %v", err)
 		}
-		defer os.Unsetenv(mocknvml.EnvNVMLMock)
+		defer os.Unsetenv(nvml_lib.EnvMockAllSuccess)
 
 		// Mock NVML with device get failure
 		mockNVML := &mock.Interface{
@@ -709,9 +707,9 @@ func TestClockEventsSupportedWithMockedNVML(t *testing.T) {
 		}
 
 		// Replace the mock instance
-		originalMockInstance := mocknvml.MockInstance
-		mocknvml.MockInstance = mockNVML
-		defer func() { mocknvml.MockInstance = originalMockInstance }()
+		originalMockInstance := nvml_lib_mock.AllSuccessInterface
+		nvml_lib_mock.AllSuccessInterface = mockNVML
+		defer func() { nvml_lib_mock.AllSuccessInterface = originalMockInstance }()
 
 		// Call function
 		result, err := ClockEventsSupported()

pkg/nvidia-query/nvml/clock_events_test.go

@@ -8,6 +8,7 @@ import (
 
 	"github.com/leptonai/gpud/pkg/log"
 	metrics_gpm "github.com/leptonai/gpud/pkg/nvidia-query/metrics/gpm"
+	nvml_lib "github.com/leptonai/gpud/pkg/nvidia-query/nvml/lib"
 
 	"github.com/NVIDIA/go-nvlib/pkg/nvlib/device"
 	"github.com/NVIDIA/go-nvml/pkg/nvml"
@@ -18,7 +19,7 @@ import (
 // Returns true if GPM is supported by all devices.
 // Returns false if any device does not support GPM.
 func GPMSupported() (bool, error) {
-	nvmlLib := NewNVML()
+	nvmlLib := nvml_lib.NewDefault()
 	if ret := nvmlLib.NVML().Init(); ret != nvml.SUCCESS {
 		return false, fmt.Errorf("failed to initialize NVML: %v", nvml.ErrorString(ret))
 	}

pkg/nvidia-query/nvml/gpm.go

@@ -0,0 +1,58 @@
+package lib
+
+import (
+	"os"
+
+	"github.com/NVIDIA/go-nvml/pkg/nvml"
+	nvml_lib_mock "github.com/leptonai/gpud/pkg/nvidia-query/nvml/lib/mock"
+
+	"github.com/leptonai/gpud/pkg/log"
+)
+
+const (
+	EnvMockAllSuccess              = "GPUD_NVML_MOCK_ALL_SUCCESS"
+	EnvInjectRemapedRowsPending    = "GPUD_NVML_INJECT_REMAPPED_ROWS_PENDING"
+	EnvInjectClockEventsHwSlowdown = "GPUD_NVML_INJECT_CLOCK_EVENTS_HW_SLOWDOWN"
+)
+
+// 0x0000000000000000 is none
+// ref. https://docs.nvidia.com/deploy/nvml-api/group__nvmlClocksEventReasons.html
+const (
+	reasonHWSlowdown           uint64 = 0x0000000000000008
+	reasonSwThermalSlowdown    uint64 = 0x0000000000000020
+	reasonHWSlowdownThermal    uint64 = 0x0000000000000040
+	reasonHWSlowdownPowerBrake uint64 = 0x0000000000000080
+)
+
+var clockEventsToInjectHwSlowdown = reasonHWSlowdown | reasonSwThermalSlowdown | reasonHWSlowdownThermal | reasonHWSlowdownPowerBrake
+
+func NewDefault() Library {
+	opts := []OpOption{}
+
+	if os.Getenv(EnvMockAllSuccess) == "true" {
+		opts = append(opts,
+			WithNVML(nvml_lib_mock.AllSuccessInterface),
+			WithPropertyExtractor(nvml_lib_mock.HasNvmlPropertyExtractor),
+		)
+	}
+
+	if os.Getenv(EnvInjectRemapedRowsPending) == "true" {
+		opts = append(opts,
+			WithDeviceGetRemappedRowsForAllDevs(func() (corrRows int, uncRows int, isPending bool, failureOccurred bool, ret nvml.Return) {
+				log.Logger.Infow("injecting remapped rows pending", "corrRows", 0, "uncRows", 10, "isPending", true, "failureOccurred", false)
+				return 0, 10, true, false, nvml.SUCCESS
+			}),
+		)
+	}
+
+	if os.Getenv(EnvInjectClockEventsHwSlowdown) == "true" {
+		opts = append(opts,
+			WithDeviceGetCurrentClocksEventReasonsForAllDevs(func() (uint64, nvml.Return) {
+				log.Logger.Infow("injecting clock events hw slowdown", "reasons", clockEventsToInjectHwSlowdown)
+				return clockEventsToInjectHwSlowdown, nvml.SUCCESS
+			}),
+		)
+	}
+
+	return New(opts...)
+}

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

@@ -0,0 +1,99 @@
+package lib
+
+import (
+	"os"
+	"testing"
+
+	"github.com/NVIDIA/go-nvml/pkg/nvml"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+// TestNewDefaultNoEnvVars tests the NewDefault function when no environment variables are set
+func TestNewDefaultNoEnvVars(t *testing.T) {
+	// Make sure environment variables are not set
+	os.Unsetenv(EnvMockAllSuccess)
+	os.Unsetenv(EnvInjectRemapedRowsPending)
+	os.Unsetenv(EnvInjectClockEventsHwSlowdown)
+
+	// Create a new library instance
+	lib := NewDefault()
+
+	// Verify the library instance is created with default options
+	assert.NotNil(t, lib)
+	assert.NotNil(t, lib.NVML())
+	assert.NotNil(t, lib.Device())
+}
+
+// TestNewDefaultMockAllSuccess tests the NewDefault function when EnvMockAllSuccess is set
+func TestNewDefaultMockAllSuccess(t *testing.T) {
+	// Clean up environment variables first
+	cleanupEnvVars()
+	defer cleanupEnvVars()
+
+	// Set the environment variable
+	os.Setenv(EnvMockAllSuccess, "true")
+
+	// Create a new library instance
+	lib := NewDefault()
+
+	// Verify the library instance is created with mock interface
+	assert.NotNil(t, lib)
+
+	// Test that NVML functions succeed
+	ret := lib.NVML().Init()
+	assert.Equal(t, nvml.SUCCESS, ret)
+
+	// Test that device functions are available and succeed
+	devices, err := lib.Device().GetDevices()
+	assert.NoError(t, err)
+	assert.NotEmpty(t, devices)
+}
+
+// TestNewDefaultMultipleEnvVars tests the NewDefault function when multiple environment variables are set
+func TestNewDefaultMultipleEnvVars(t *testing.T) {
+	// Clean up environment variables first
+	cleanupEnvVars()
+	defer cleanupEnvVars()
+
+	// Set multiple environment variables
+	os.Setenv(EnvMockAllSuccess, "true")
+	os.Setenv(EnvInjectRemapedRowsPending, "true")
+	os.Setenv(EnvInjectClockEventsHwSlowdown, "true")
+
+	// Create a new library instance
+	lib := NewDefault()
+
+	// Verify the library instance is created correctly
+	assert.NotNil(t, lib)
+
+	// Test that NVML functions succeed
+	ret := lib.NVML().Init()
+	assert.Equal(t, nvml.SUCCESS, ret)
+
+	// Get devices to test modified functions
+	devices, err := lib.Device().GetDevices()
+	require.NoError(t, err)
+	require.NotEmpty(t, devices)
+
+	// Test the injected function to get remapped rows
+	corrRows, uncRows, isPending, failureOccurred, retRemapped := devices[0].GetRemappedRows()
+	assert.Equal(t, 0, corrRows)
+	assert.Equal(t, 10, uncRows)
+	assert.True(t, isPending)
+	assert.False(t, failureOccurred)
+	assert.Equal(t, nvml.SUCCESS, retRemapped)
+
+	// Test the injected function to get clock events
+	reasons, retClock := devices[0].GetCurrentClocksEventReasons()
+	expectedReasons := reasonHWSlowdown | reasonSwThermalSlowdown | reasonHWSlowdownThermal | reasonHWSlowdownPowerBrake
+	assert.Equal(t, expectedReasons, reasons)
+	assert.Equal(t, nvml.SUCCESS, retClock)
+}
+
+// Utility function to clean up environment variables
+func cleanupEnvVars() {
+	os.Unsetenv(EnvMockAllSuccess)
+	os.Unsetenv(EnvInjectRemapedRowsPending)
+	os.Unsetenv(EnvInjectClockEventsHwSlowdown)
+}

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

@@ -58,9 +58,10 @@ func New(opts ...OpOption) Library {
 
 	devLib := device.New(nvInterface.Interface)
 	nvInterface.dev = &devInterface{
-		Interface:       devLib,
-		devices:         options.devicesToReturn,
-		getRemappedRows: options.devGetRemappedRows,
+		Interface:                              devLib,
+		devices:                                options.devicesToReturn,
+		getRemappedRowsForAllDevs:              options.devGetRemappedRowsForAllDevs,
+		getCurrentClocksEventReasonsForAllDevs: options.devGetCurrentClocksEventReasonsForAllDevs,
 	}
 
 	infoOpts := []nvinfo.Option{
@@ -86,8 +87,9 @@ var _ device.Interface = &devInterface{}
 
 type devInterface struct {
 	device.Interface
-	devices         []device.Device
-	getRemappedRows func() (int, int, bool, bool, nvml.Return)
+	devices                                []device.Device
+	getRemappedRowsForAllDevs              func() (int, int, bool, bool, nvml.Return)
+	getCurrentClocksEventReasonsForAllDevs func() (uint64, nvml.Return)
 }
 
 func (d *devInterface) GetDevices() ([]device.Device, error) {
@@ -105,8 +107,9 @@ func (d *devInterface) GetDevices() ([]device.Device, error) {
 	updated := make([]device.Device, len(devs))
 	for i, dev := range devs {
 		updated[i] = &devDevInterface{
-			Device:          dev,
-			getRemappedRows: d.getRemappedRows,
+			Device:                                 dev,
+			getRemappedRowsForAllDevs:              d.getRemappedRowsForAllDevs,
+			getCurrentClocksEventReasonsForAllDevs: d.getCurrentClocksEventReasonsForAllDevs,
 		}
 	}
 
@@ -117,14 +120,24 @@ var _ device.Device = &devDevInterface{}
 
 type devDevInterface struct {
 	device.Device
-	getRemappedRows func() (int, int, bool, bool, nvml.Return)
+	getRemappedRowsForAllDevs              func() (int, int, bool, bool, nvml.Return)
+	getCurrentClocksEventReasonsForAllDevs func() (uint64, 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 {
+	if d.getRemappedRowsForAllDevs == nil {
 		return d.Device.GetRemappedRows()
 	}
-	return d.getRemappedRows()
+	return d.getRemappedRowsForAllDevs()
+}
+
+func (d *devDevInterface) GetCurrentClocksEventReasons() (uint64, nvml.Return) {
+	// no injected current clocks event reasons
+	// thus just passthrough to call the underlying device.Device.GetCurrentClocksEventReasons()
+	if d.getCurrentClocksEventReasonsForAllDevs == nil {
+		return d.Device.GetCurrentClocksEventReasons()
+	}
+	return d.getCurrentClocksEventReasonsForAllDevs()
 }