refactor(vGPUmonitor): use nvml.Interface

cmd/vGPUmonitor/feedback.go

@@ -23,7 +23,7 @@ import (
 	"time"
 
 	"github.com/NVIDIA/go-nvml/pkg/nvml"
-	"k8s.io/klog/v2"
+	klog "k8s.io/klog/v2"
 
 	"github.com/Project-HAMi/HAMi/pkg/monitor/nvidia"
 )
@@ -106,39 +106,48 @@ func Observe(lister *nvidia.ContainerLister) {
 		utilizationSwitch := c.Info.GetUtilizationSwitch()
 		if CheckBlocking(utSwitchOn, priority, c) {
 			if recentKernel >= 0 {
-				klog.V(5).Infof("utSwitchon=%v", utSwitchOn)
-				klog.V(5).Infof("Setting Blocking to on %v", idx)
+				klog.Infof("utSwitchon=%v", utSwitchOn)
+				klog.Infof("Setting Blocking to on %v", idx)
 				c.Info.SetRecentKernel(-1)
 			}
 		} else {
 			if recentKernel < 0 {
-				klog.V(5).Infof("utSwitchon=%v", utSwitchOn)
-				klog.V(5).Infof("Setting Blocking to off %v", idx)
+				klog.Infof("utSwitchon=%v", utSwitchOn)
+				klog.Infof("Setting Blocking to off %v", idx)
 				c.Info.SetRecentKernel(0)
 			}
 		}
 		if CheckPriority(utSwitchOn, priority, c) {
 			if utilizationSwitch != 1 {
-				klog.V(5).Infof("utSwitchon=%v", utSwitchOn)
-				klog.V(5).Infof("Setting UtilizationSwitch to on %v", idx)
+				klog.Infof("utSwitchon=%v", utSwitchOn)
+				klog.Infof("Setting UtilizationSwitch to on %v", idx)
 				c.Info.SetUtilizationSwitch(1)
 			}
 		} else {
 			if utilizationSwitch != 0 {
-				klog.V(5).Infof("utSwitchon=%v", utSwitchOn)
-				klog.V(5).Infof("Setting UtilizationSwitch to off %v", idx)
+				klog.Infof("utSwitchon=%v", utSwitchOn)
+				klog.Infof("Setting UtilizationSwitch to off %v", idx)
 				c.Info.SetUtilizationSwitch(0)
 			}
 		}
 	}
 }
 
-func watchAndFeedback(ctx context.Context, lister *nvidia.ContainerLister, migLockSignal <-chan bool) error {
+func watchAndFeedback(ctx context.Context, lister *nvidia.ContainerLister, nvmllib nvml.Interface, migLockSignal <-chan bool) error {
 	klog.Info("Starting watchAndFeedback")
-	if nvret := nvml.Init(); nvret != nvml.SUCCESS {
-		return fmt.Errorf("failed to initialize NVML: %s", nvml.ErrorString(nvret))
+
+	// Guard against a nil NVML interface (e.g. in tests or when NVML is
+	// unavailable). Physical GPU metric collection is skipped; container
+	// observation and the MIG lock signal continue to work normally,
+	// consistent with how collectGPUInfo handles a nil nvmllib in metrics.go.
+	if nvmllib != nil {
+		if nvret := nvmllib.Init(); !errors.Is(nvret, nvml.SUCCESS) {
+			return fmt.Errorf("failed to initialize NVML: %w", nvret)
+		}
+		defer func() { _ = nvmllib.Shutdown() }()
+	} else {
+		klog.Warning("watchAndFeedback: nvmllib is nil, skipping NVML init (degraded mode)")
 	}
-	defer nvml.Shutdown()
 
 	ticker := time.NewTicker(time.Second * 5)
 	defer ticker.Stop()

cmd/vGPUmonitor/feedback_test.go

@@ -17,8 +17,12 @@ limitations under the License.
 package main
 
 import (
+	"context"
 	"testing"
 
+	"github.com/NVIDIA/go-nvml/pkg/nvml"
+	"github.com/NVIDIA/go-nvml/pkg/nvml/mock"
+
 	"github.com/Project-HAMi/HAMi/pkg/monitor/nvidia"
 )
 
@@ -162,3 +166,68 @@ func TestCheckBlocking_MultiDevice(t *testing.T) {
 		})
 	}
 }
+
+func TestWatchAndFeedback_NilNVML(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+	cancel() // Immediately cancel context so watchAndFeedback exits gracefully after initialization check
+
+	lockCh := make(chan bool)
+	err := watchAndFeedback(ctx, &nvidia.ContainerLister{}, nil, lockCh)
+	if err != nil {
+		t.Errorf("watchAndFeedback with nil nvmllib returned unexpected error: %v", err)
+	}
+}
+
+func TestWatchAndFeedback_WithNVMLSuccess(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+	cancel() // Immediately cancel context so watchAndFeedback exits gracefully after initialization check
+
+	mockNVML := &mock.Interface{
+		InitFunc:     func() nvml.Return { return nvml.SUCCESS },
+		ShutdownFunc: func() nvml.Return { return nvml.SUCCESS },
+	}
+
+	lockCh := make(chan bool)
+	err := watchAndFeedback(ctx, &nvidia.ContainerLister{}, mockNVML, lockCh)
+	if err != nil {
+		t.Errorf("watchAndFeedback with mockNVML returned unexpected error: %v", err)
+	}
+}
+
+func TestWatchAndFeedback_WithNVMLError(t *testing.T) {
+	ctx := t.Context()
+
+	mockNVML := &mock.Interface{
+		InitFunc: func() nvml.Return { return nvml.ERROR_UNKNOWN },
+	}
+
+	lockCh := make(chan bool)
+	err := watchAndFeedback(ctx, &nvidia.ContainerLister{}, mockNVML, lockCh)
+	if err == nil {
+		t.Error("watchAndFeedback expected error when NVML init fails, got nil")
+	}
+}
+
+func TestWatchAndFeedback_MigLockSignal(t *testing.T) {
+	ctx := t.Context()
+
+	lockCh := make(chan bool, 1)
+	lockCh <- true
+
+	err := watchAndFeedback(ctx, &nvidia.ContainerLister{}, nil, lockCh)
+	if err != errTemporaryClosed {
+		t.Errorf("watchAndFeedback with migLockSignal expected errTemporaryClosed, got %v", err)
+	}
+}
+
+func TestObserve_EmptyLister(t *testing.T) {
+	Observe(&nvidia.ContainerLister{})
+}
+
+func TestCheckPriority_SamePriorityContention(t *testing.T) {
+	sw := map[string]UtilizationPerDevice{"gpu-0": {0, 2}}
+	c := &nvidia.ContainerUsage{Info: &stubInfo{priority: 1, uuids: []string{"gpu-0"}}}
+	if !CheckPriority(sw, 1, c) {
+		t.Error("CheckPriority: expected true for same priority contention > 1")
+	}
+}

cmd/vGPUmonitor/main.go

@@ -34,11 +34,12 @@ import (
 	"github.com/Project-HAMi/HAMi/pkg/util/flag"
 	"github.com/Project-HAMi/HAMi/pkg/version"
 
+	"github.com/NVIDIA/go-nvml/pkg/nvml"
 	"github.com/prometheus/client_golang/prometheus"
 	"github.com/prometheus/client_golang/prometheus/promhttp"
 
 	"github.com/spf13/cobra"
-	"k8s.io/klog/v2"
+	klog "k8s.io/klog/v2"
 )
 
 var (
@@ -88,20 +89,22 @@ func start() error {
 		return fmt.Errorf("failed to watch lock file: %v", err)
 	}
 
+	nvmllib := nvml.New()
+
 	var wg sync.WaitGroup
 	errCh := make(chan error, 2)
 
 	// Start the metrics service
 	wg.Go(func() {
-		if err := initMetrics(ctx, containerLister); err != nil {
+		if err := initMetrics(ctx, containerLister, nvmllib); err != nil {
 			errCh <- err
 		}
 	})
 
 	// Start the monitoring and feedback service
 	wg.Go(func() {
 		for {
-			if err := watchAndFeedback(ctx, containerLister, lockChannel); err != nil {
+			if err := watchAndFeedback(ctx, containerLister, nvmllib, lockChannel); err != nil {
 				// if err is temporary closed, wait for lock file to be removed
 				if errors.Is(err, errTemporaryClosed) {
 					klog.Info("MIG apply lock file detected, waiting for lock file to be removed")
@@ -135,14 +138,14 @@ func start() error {
 	return nil
 }
 
-func initMetrics(ctx context.Context, containerLister *nvidia.ContainerLister) error {
+func initMetrics(ctx context.Context, containerLister *nvidia.ContainerLister, nvmllib nvml.Interface) error {
 	klog.V(4).Info("Initializing metrics for vGPUmonitor")
 	reg := prometheus.NewRegistry()
 	//reg := prometheus.NewPedanticRegistry()
 
 	reg.MustRegister(versionmetrics.NewBuildInfoCollector())
 
-	NewClusterManager("vGPU", reg, containerLister, legacyMetrics)
+	NewClusterManager("vGPU", reg, containerLister, nvmllib, legacyMetrics)
 
 	// Uncomment to add the standard process and Go metrics to the custom registry.
 	//reg.MustRegister(