diff --git a/cmd/vGPUmonitor/feedback.go b/cmd/vGPUmonitor/feedback.go index 3d119e6c4c..cf7870a61c 100644 --- a/cmd/vGPUmonitor/feedback.go +++ b/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() diff --git a/cmd/vGPUmonitor/feedback_test.go b/cmd/vGPUmonitor/feedback_test.go index b55590fe8a..81816b69f6 100644 --- a/cmd/vGPUmonitor/feedback_test.go +++ b/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") + } +} diff --git a/cmd/vGPUmonitor/main.go b/cmd/vGPUmonitor/main.go index 4f5c8c4e5b..d32839d3d0 100644 --- a/cmd/vGPUmonitor/main.go +++ b/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,12 +89,14 @@ 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 } }) @@ -101,7 +104,7 @@ func start() error { // 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( diff --git a/cmd/vGPUmonitor/metrics.go b/cmd/vGPUmonitor/metrics.go index a1155af839..f9a90a8a5a 100644 --- a/cmd/vGPUmonitor/metrics.go +++ b/cmd/vGPUmonitor/metrics.go @@ -17,6 +17,7 @@ limitations under the License. package main import ( + "errors" "fmt" "os" "time" @@ -32,7 +33,7 @@ import ( "k8s.io/apimachinery/pkg/labels" "k8s.io/client-go/informers" listerscorev1 "k8s.io/client-go/listers/core/v1" - "k8s.io/klog/v2" + klog "k8s.io/klog/v2" ) // ClusterManager holds the state the vGPUMonitor's Prometheus collector reads @@ -44,6 +45,12 @@ type ClusterManager struct { PodLister listerscorev1.PodLister containerLister *nvidia.ContainerLister LegacyMetrics bool + // nvmllib is the NVML library interface used to query physical GPU + // metrics. Keeping it as an interface (rather than calling package-level + // nvml functions directly) avoids CGo symbol references that the + // Windows/non-CGo language server cannot resolve, and makes the + // collector unit-testable via nvml mock implementations. + nvmllib nvml.Interface } // ClusterManagerCollector implements the Collector interface. @@ -196,7 +203,7 @@ func (cc ClusterManagerCollector) Describe(ch chan<- *prometheus.Desc) { ch <- ctrDeviceMemoryModuleDesc ch <- ctrDeviceMemoryBufferDesc - if cc.ClusterManager.LegacyMetrics { + if cc.ClusterManager != nil && cc.ClusterManager.LegacyMetrics { ch <- legacyHostGPUdesc ch <- legacyHostGPUUtilizationdesc ch <- legacyCtrvGPUdesc @@ -212,6 +219,11 @@ func (cc ClusterManagerCollector) Describe(ch chan<- *prometheus.Desc) { // to the provided channel. It may be called concurrently, so the collection // helpers it calls must be concurrency-safe. func (cc ClusterManagerCollector) Collect(ch chan<- prometheus.Metric) { + if cc.ClusterManager == nil { + klog.Warning("Skipping metric collection because ClusterManager is nil") + return + } + klog.Info("Starting to collect metrics for vGPUMonitor") // Collect GPU information @@ -236,18 +248,27 @@ func (cc ClusterManagerCollector) Collect(ch chan<- prometheus.Metric) { } func (cc ClusterManagerCollector) collectGPUInfo(ch chan<- prometheus.Metric) error { - if err := cc.initNVML(); err != nil { + if cc.ClusterManager == nil { + return nil + } + + nvmllib := cc.ClusterManager.nvmllib + if nvmllib == nil { + // No NVML library configured; skip physical GPU metric collection. + return nil + } + if err := cc.initNVML(nvmllib); err != nil { return err } - defer nvml.Shutdown() + defer func() { _ = nvmllib.Shutdown() }() - devnum, err := cc.getDeviceCount() + devnum, err := cc.getDeviceCount(nvmllib) if err != nil { return err } for ii := range devnum { - if err := cc.collectGPUDeviceMetrics(ch, ii); err != nil { + if err := cc.collectGPUDeviceMetrics(ch, nvmllib, ii); err != nil { klog.Error("Failed to collect metrics for GPU device ", ii, ": ", err) } } @@ -255,26 +276,26 @@ func (cc ClusterManagerCollector) collectGPUInfo(ch chan<- prometheus.Metric) er return nil } -func (cc ClusterManagerCollector) initNVML() error { - nvret := nvml.Init() - if nvret != nvml.SUCCESS { - return fmt.Errorf("nvml Init err: %s", nvml.ErrorString(nvret)) +func (cc ClusterManagerCollector) initNVML(nvmllib nvml.Interface) error { + nvret := nvmllib.Init() + if !errors.Is(nvret, nvml.SUCCESS) { + return fmt.Errorf("nvml Init err: %w", nvret) } return nil } -func (cc ClusterManagerCollector) getDeviceCount() (int, error) { - devnum, nvret := nvml.DeviceGetCount() - if nvret != nvml.SUCCESS { - return 0, fmt.Errorf("nvml GetDeviceCount err: %s", nvml.ErrorString(nvret)) +func (cc ClusterManagerCollector) getDeviceCount(nvmllib nvml.Interface) (int, error) { + devnum, nvret := nvmllib.DeviceGetCount() + if !errors.Is(nvret, nvml.SUCCESS) { + return 0, fmt.Errorf("nvml GetDeviceCount err: %w", nvret) } return devnum, nil } -func (cc ClusterManagerCollector) collectGPUDeviceMetrics(ch chan<- prometheus.Metric, index int) error { - hdev, nvret := nvml.DeviceGetHandleByIndex(index) - if nvret != nvml.SUCCESS { - return fmt.Errorf("nvml DeviceGetHandleByIndex err: %s", nvml.ErrorString(nvret)) +func (cc ClusterManagerCollector) collectGPUDeviceMetrics(ch chan<- prometheus.Metric, nvmllib nvml.Interface, index int) error { + hdev, nvret := nvmllib.DeviceGetHandleByIndex(index) + if !errors.Is(nvret, nvml.SUCCESS) { + return fmt.Errorf("nvml DeviceGetHandleByIndex err: %w", nvret) } if err := cc.collectGPUMemoryMetrics(ch, hdev, index); err != nil { @@ -290,22 +311,22 @@ func (cc ClusterManagerCollector) collectGPUDeviceMetrics(ch chan<- prometheus.M func (cc ClusterManagerCollector) collectGPUMemoryMetrics(ch chan<- prometheus.Metric, hdev nvml.Device, index int) error { memory, ret := hdev.GetMemoryInfo() - if ret == nvml.ERROR_NOT_SUPPORTED { + if errors.Is(ret, nvml.ERROR_NOT_SUPPORTED) { klog.V(3).Infof("Memory metrics not supported for device %d (unified memory architecture), skipping", index) return nil } - if ret != nvml.SUCCESS { - return fmt.Errorf("nvml get memory error ret=%d", ret) + if !errors.Is(ret, nvml.SUCCESS) { + return fmt.Errorf("nvml get memory error: %w", ret) } uuid, nvret := hdev.GetUUID() - if nvret != nvml.SUCCESS { - return fmt.Errorf("nvml GetUUID err: %s", nvml.ErrorString(nvret)) + if !errors.Is(nvret, nvml.SUCCESS) { + return fmt.Errorf("nvml GetUUID err: %w", nvret) } deviceName, nvret := hdev.GetName() - if nvret != nvml.SUCCESS { - return fmt.Errorf("nvml GetName err: %s", nvml.ErrorString(nvret)) + if !errors.Is(nvret, nvml.SUCCESS) { + return fmt.Errorf("nvml GetName err: %w", nvret) } deviceName = "NVIDIA-" + deviceName @@ -326,18 +347,18 @@ func (cc ClusterManagerCollector) collectGPUMemoryMetrics(ch chan<- prometheus.M func (cc ClusterManagerCollector) collectGPUUtilizationMetrics(ch chan<- prometheus.Metric, hdev nvml.Device, index int) error { util, nvret := hdev.GetUtilizationRates() - if nvret != nvml.SUCCESS { - return fmt.Errorf("nvml GetUtilizationRates err: %s", nvml.ErrorString(nvret)) + if !errors.Is(nvret, nvml.SUCCESS) { + return fmt.Errorf("nvml GetUtilizationRates err: %w", nvret) } uuid, nvret := hdev.GetUUID() - if nvret != nvml.SUCCESS { - return fmt.Errorf("nvml GetUUID err: %s", nvml.ErrorString(nvret)) + if !errors.Is(nvret, nvml.SUCCESS) { + return fmt.Errorf("nvml GetUUID err: %w", nvret) } deviceName, nvret := hdev.GetName() - if nvret != nvml.SUCCESS { - return fmt.Errorf("nvml GetName err: %s", nvml.ErrorString(nvret)) + if !errors.Is(nvret, nvml.SUCCESS) { + return fmt.Errorf("nvml GetName err: %w", nvret) } deviceName = "NVIDIA-" + deviceName @@ -357,6 +378,10 @@ func (cc ClusterManagerCollector) collectGPUUtilizationMetrics(ch chan<- prometh } func (cc ClusterManagerCollector) collectPodAndContainerInfo(ch chan<- prometheus.Metric) error { + if cc.ClusterManager == nil || cc.ClusterManager.PodLister == nil || cc.ClusterManager.containerLister == nil { + return nil + } + nodeName := os.Getenv(util.NodeNameEnvName) if nodeName == "" { return fmt.Errorf("node name environment variable %s is not set", util.NodeNameEnvName) @@ -500,6 +525,10 @@ func (cc ClusterManagerCollector) collectContainerMetrics(ch chan<- prometheus.M } func (cc ClusterManagerCollector) collectPodAndContainerMigInfo(ch chan<- prometheus.Metric) error { + if cc.ClusterManager == nil || cc.ClusterManager.PodLister == nil { + return nil + } + nodeName := os.Getenv(util.NodeNameEnvName) if nodeName == "" { return fmt.Errorf("node name environment variable %s is not set", util.NodeNameEnvName) @@ -572,7 +601,10 @@ func sendMetric(ch chan<- prometheus.Metric, desc *prometheus.Desc, valueType pr // NewClusterManager creates a ClusterManager for the given zone, backs its pod // lookups with a shared informer, and registers its collector with reg through // a wrapping Registerer that adds the zone as a label. -func NewClusterManager(zone string, reg prometheus.Registerer, containerLister *nvidia.ContainerLister, legacyMetrics bool) *ClusterManager { +// +// nvmllib is the NVML library interface used to query physical GPU metrics. +// Pass nvml.New() in production; pass a mock implementation in tests. +func NewClusterManager(zone string, reg prometheus.Registerer, containerLister *nvidia.ContainerLister, nvmllib nvml.Interface, legacyMetrics bool) *ClusterManager { if legacyMetrics { initLegacyDescriptors() } @@ -580,12 +612,15 @@ func NewClusterManager(zone string, reg prometheus.Registerer, containerLister * Zone: zone, containerLister: containerLister, LegacyMetrics: legacyMetrics, + nvmllib: nvmllib, } - informerFactory := informers.NewSharedInformerFactoryWithOptions(containerLister.Clientset(), time.Hour*1) - c.PodLister = informerFactory.Core().V1().Pods().Lister() - stopCh := make(chan struct{}) - informerFactory.Start(stopCh) + if containerLister != nil && containerLister.Clientset() != nil { + informerFactory := informers.NewSharedInformerFactoryWithOptions(containerLister.Clientset(), time.Hour*1) + c.PodLister = informerFactory.Core().V1().Pods().Lister() + stopCh := make(chan struct{}) + informerFactory.Start(stopCh) + } cc := ClusterManagerCollector{ClusterManager: c} prometheus.WrapRegistererWith(prometheus.Labels{"zone": zone}, reg).MustRegister(cc) diff --git a/cmd/vGPUmonitor/metrics_test.go b/cmd/vGPUmonitor/metrics_test.go index 2c31e4d221..740a6fd062 100644 --- a/cmd/vGPUmonitor/metrics_test.go +++ b/cmd/vGPUmonitor/metrics_test.go @@ -17,12 +17,19 @@ limitations under the License. package main import ( + "context" + "os" "testing" + "github.com/NVIDIA/go-nvml/pkg/nvml" + "github.com/NVIDIA/go-nvml/pkg/nvml/mock" "github.com/prometheus/client_golang/prometheus" + corev1 "k8s.io/api/core/v1" + metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" "k8s.io/client-go/informers" "k8s.io/client-go/kubernetes/fake" + nv "github.com/Project-HAMi/HAMi/pkg/device/nvidia" "github.com/Project-HAMi/HAMi/pkg/monitor/nvidia" "github.com/Project-HAMi/HAMi/pkg/util" ) @@ -68,3 +75,378 @@ func TestDescribeCollectSync(t *testing.T) { t.Errorf("Gather failed (legacy): %v", err) } } + +func TestCollectGPUInfo_NilNVML(t *testing.T) { + c := &ClusterManager{ + Zone: "test-zone", + nvmllib: nil, + containerLister: &nvidia.ContainerLister{}, + } + cc := ClusterManagerCollector{ClusterManager: c} + ch := make(chan prometheus.Metric, 10) + + // Under nil nvmllib, physical GPU metrics collection must gracefully skip without panicking. + if err := cc.collectGPUInfo(ch); err != nil { + t.Errorf("collectGPUInfo with nil nvmllib returned unexpected error: %v", err) + } +} + +func newMockNVMLDevice(uuid, name string, memTotal, memUsed uint64, utilGPU, utilMem uint32) *mock.Device { + return &mock.Device{ + GetMemoryInfoFunc: func() (nvml.Memory, nvml.Return) { + return nvml.Memory{Total: memTotal, Used: memUsed, Free: memTotal - memUsed}, nvml.SUCCESS + }, + GetUtilizationRatesFunc: func() (nvml.Utilization, nvml.Return) { + return nvml.Utilization{Gpu: utilGPU, Memory: utilMem}, nvml.SUCCESS + }, + GetUUIDFunc: func() (string, nvml.Return) { + return uuid, nvml.SUCCESS + }, + GetNameFunc: func() (string, nvml.Return) { + return name, nvml.SUCCESS + }, + } +} + +func newMockNVML(devices ...nvml.Device) *mock.Interface { + return &mock.Interface{ + InitFunc: func() nvml.Return { + return nvml.SUCCESS + }, + ShutdownFunc: func() nvml.Return { + return nvml.SUCCESS + }, + DeviceGetCountFunc: func() (int, nvml.Return) { + return len(devices), nvml.SUCCESS + }, + DeviceGetHandleByIndexFunc: func(index int) (nvml.Device, nvml.Return) { + if index >= 0 && index < len(devices) { + return devices[index], nvml.SUCCESS + } + return nil, nvml.ERROR_INVALID_ARGUMENT + }, + } +} + +func TestNewClusterManager(t *testing.T) { + containerLister := &nvidia.ContainerLister{} + reg := prometheus.NewRegistry() + mockNVML := newMockNVML() + + cm := NewClusterManager("test-zone", reg, containerLister, mockNVML, false) + if cm == nil { + t.Fatal("NewClusterManager returned nil") + } + if cm.nvmllib != mockNVML { + t.Errorf("NewClusterManager nvmllib mismatch") + } + + regLegacy := prometheus.NewRegistry() + cmLegacy := NewClusterManager("test-zone-legacy", regLegacy, containerLister, mockNVML, true) + if cmLegacy == nil { + t.Fatal("NewClusterManager legacy returned nil") + } +} + +func TestInitMetrics(t *testing.T) { + origAddr := metricsBindAddress + metricsBindAddress = "127.0.0.1:0" + defer func() { metricsBindAddress = origAddr }() + + containerLister := &nvidia.ContainerLister{} + mockNVML := newMockNVML() + + ctx, cancel := context.WithCancel(context.Background()) + + errCh := make(chan error, 1) + go func() { + errCh <- initMetrics(ctx, containerLister, mockNVML) + }() + + cancel() + if err := <-errCh; err != nil { + t.Errorf("initMetrics returned unexpected error: %v", err) + } +} + +func TestCollectGPUInfo_Success(t *testing.T) { + dev := newMockNVMLDevice("GPU-12345", "Tesla T4", 16000000000, 4000000000, 75, 50) + mockNVML := newMockNVML(dev) + + c := &ClusterManager{ + Zone: "test-zone", + nvmllib: mockNVML, + containerLister: &nvidia.ContainerLister{}, + } + cc := ClusterManagerCollector{ClusterManager: c} + ch := make(chan prometheus.Metric, 10) + + if err := cc.collectGPUInfo(ch); err != nil { + t.Fatalf("collectGPUInfo returned unexpected error: %v", err) + } + close(ch) + + metricCount := 0 + for range ch { + metricCount++ + } + if metricCount != 4 { + t.Errorf("Expected 4 metrics, got %d", metricCount) + } +} + +func TestCollectGPUInfo_ErrorPaths(t *testing.T) { + ch := make(chan prometheus.Metric, 10) + + // 1. Init error + mockInitErr := &mock.Interface{ + InitFunc: func() nvml.Return { return nvml.ERROR_UNKNOWN }, + } + ccInitErr := ClusterManagerCollector{ClusterManager: &ClusterManager{nvmllib: mockInitErr}} + if err := ccInitErr.collectGPUInfo(ch); err == nil { + t.Error("Expected error when Init fails, got nil") + } + + // 2. DeviceGetCount error + mockCountErr := &mock.Interface{ + InitFunc: func() nvml.Return { return nvml.SUCCESS }, + ShutdownFunc: func() nvml.Return { return nvml.SUCCESS }, + DeviceGetCountFunc: func() (int, nvml.Return) { return 0, nvml.ERROR_UNKNOWN }, + } + ccCountErr := ClusterManagerCollector{ClusterManager: &ClusterManager{nvmllib: mockCountErr}} + if err := ccCountErr.collectGPUInfo(ch); err == nil { + t.Error("Expected error when DeviceGetCount fails, got nil") + } + + // 3. DeviceGetHandleByIndex error + mockHandleErr := &mock.Interface{ + InitFunc: func() nvml.Return { return nvml.SUCCESS }, + ShutdownFunc: func() nvml.Return { return nvml.SUCCESS }, + DeviceGetCountFunc: func() (int, nvml.Return) { return 1, nvml.SUCCESS }, + DeviceGetHandleByIndexFunc: func(int) (nvml.Device, nvml.Return) { return nil, nvml.ERROR_UNKNOWN }, + } + ccHandleErr := ClusterManagerCollector{ClusterManager: &ClusterManager{nvmllib: mockHandleErr}} + if err := ccHandleErr.collectGPUInfo(ch); err != nil { + t.Errorf("Unexpected error when DeviceGetHandleByIndex fails: %v", err) + } + + // 4. GetMemoryInfo error + devMemErr := &mock.Device{ + GetMemoryInfoFunc: func() (nvml.Memory, nvml.Return) { return nvml.Memory{}, nvml.ERROR_UNKNOWN }, + } + mockMemErr := newMockNVML(devMemErr) + ccMemErr := ClusterManagerCollector{ClusterManager: &ClusterManager{nvmllib: mockMemErr}} + _ = ccMemErr.collectGPUInfo(ch) + + // 5. GetMemoryInfo NOT_SUPPORTED (Unified Memory) + devMemNotSupp := &mock.Device{ + GetMemoryInfoFunc: func() (nvml.Memory, nvml.Return) { return nvml.Memory{}, nvml.ERROR_NOT_SUPPORTED }, + GetUtilizationRatesFunc: func() (nvml.Utilization, nvml.Return) { return nvml.Utilization{Gpu: 50}, nvml.SUCCESS }, + GetUUIDFunc: func() (string, nvml.Return) { return "GPU-unified", nvml.SUCCESS }, + GetNameFunc: func() (string, nvml.Return) { return "GH200", nvml.SUCCESS }, + } + mockMemNotSupp := newMockNVML(devMemNotSupp) + ccMemNotSupp := ClusterManagerCollector{ClusterManager: &ClusterManager{nvmllib: mockMemNotSupp}} + if err := ccMemNotSupp.collectGPUInfo(ch); err != nil { + t.Errorf("Unexpected error when MemoryInfo is NOT_SUPPORTED: %v", err) + } + + // 6. GetUUID error in memory metrics + devUUIDMemErr := &mock.Device{ + GetMemoryInfoFunc: func() (nvml.Memory, nvml.Return) { return nvml.Memory{Used: 100}, nvml.SUCCESS }, + GetUUIDFunc: func() (string, nvml.Return) { return "", nvml.ERROR_UNKNOWN }, + } + mockUUIDMemErr := newMockNVML(devUUIDMemErr) + ccUUIDMemErr := ClusterManagerCollector{ClusterManager: &ClusterManager{nvmllib: mockUUIDMemErr}} + _ = ccUUIDMemErr.collectGPUInfo(ch) + + // 7. GetName error in memory metrics + devNameMemErr := &mock.Device{ + GetMemoryInfoFunc: func() (nvml.Memory, nvml.Return) { return nvml.Memory{Used: 100}, nvml.SUCCESS }, + GetUUIDFunc: func() (string, nvml.Return) { return "GPU-123", nvml.SUCCESS }, + GetNameFunc: func() (string, nvml.Return) { return "", nvml.ERROR_UNKNOWN }, + } + mockNameMemErr := newMockNVML(devNameMemErr) + ccNameMemErr := ClusterManagerCollector{ClusterManager: &ClusterManager{nvmllib: mockNameMemErr}} + _ = ccNameMemErr.collectGPUInfo(ch) + + // 8. GetUtilizationRates error + devUtilErr := &mock.Device{ + GetMemoryInfoFunc: func() (nvml.Memory, nvml.Return) { return nvml.Memory{Used: 100}, nvml.SUCCESS }, + GetUUIDFunc: func() (string, nvml.Return) { return "GPU-123", nvml.SUCCESS }, + GetNameFunc: func() (string, nvml.Return) { return "Tesla", nvml.SUCCESS }, + GetUtilizationRatesFunc: func() (nvml.Utilization, nvml.Return) { return nvml.Utilization{}, nvml.ERROR_UNKNOWN }, + } + mockUtilErr := newMockNVML(devUtilErr) + ccUtilErr := ClusterManagerCollector{ClusterManager: &ClusterManager{nvmllib: mockUtilErr}} + _ = ccUtilErr.collectGPUInfo(ch) + + // 9. GetUUID error in utilization metrics + uuidCallCount := 0 + devUUIDUtilErr := &mock.Device{ + GetMemoryInfoFunc: func() (nvml.Memory, nvml.Return) { return nvml.Memory{Used: 100}, nvml.SUCCESS }, + GetUtilizationRatesFunc: func() (nvml.Utilization, nvml.Return) { return nvml.Utilization{Gpu: 50}, nvml.SUCCESS }, + GetUUIDFunc: func() (string, nvml.Return) { + uuidCallCount++ + if uuidCallCount > 1 { + return "", nvml.ERROR_UNKNOWN + } + return "GPU-123", nvml.SUCCESS + }, + GetNameFunc: func() (string, nvml.Return) { return "Tesla", nvml.SUCCESS }, + } + mockUUIDUtilErr := newMockNVML(devUUIDUtilErr) + ccUUIDUtilErr := ClusterManagerCollector{ClusterManager: &ClusterManager{nvmllib: mockUUIDUtilErr}} + _ = ccUUIDUtilErr.collectGPUInfo(ch) + + // 10. GetName error in utilization metrics + nameCallCount := 0 + devNameUtilErr := &mock.Device{ + GetMemoryInfoFunc: func() (nvml.Memory, nvml.Return) { return nvml.Memory{Used: 100}, nvml.SUCCESS }, + GetUtilizationRatesFunc: func() (nvml.Utilization, nvml.Return) { return nvml.Utilization{Gpu: 50}, nvml.SUCCESS }, + GetUUIDFunc: func() (string, nvml.Return) { return "GPU-123", nvml.SUCCESS }, + GetNameFunc: func() (string, nvml.Return) { + nameCallCount++ + if nameCallCount > 1 { + return "", nvml.ERROR_UNKNOWN + } + return "Tesla", nvml.SUCCESS + }, + } + mockNameUtilErr := newMockNVML(devNameUtilErr) + ccNameUtilErr := ClusterManagerCollector{ClusterManager: &ClusterManager{nvmllib: mockNameUtilErr}} + _ = ccNameUtilErr.collectGPUInfo(ch) +} + +func TestClusterManagerCollector_NilGuards(t *testing.T) { + ch := make(chan prometheus.Metric, 10) + + // 1. cc with nil ClusterManager + ccNil := ClusterManagerCollector{ClusterManager: nil} + ccNil.Collect(ch) + + if err := ccNil.collectGPUInfo(ch); err != nil { + t.Errorf("collectGPUInfo with nil ClusterManager returned error: %v", err) + } + if err := ccNil.collectPodAndContainerInfo(ch); err != nil { + t.Errorf("collectPodAndContainerInfo with nil ClusterManager returned error: %v", err) + } + if err := ccNil.collectPodAndContainerMigInfo(ch); err != nil { + t.Errorf("collectPodAndContainerMigInfo with nil ClusterManager returned error: %v", err) + } + + // 2. cc with ClusterManager containing nil PodLister / containerLister + cmPartial := &ClusterManager{ + Zone: "test-zone", + } + ccPartial := ClusterManagerCollector{ClusterManager: cmPartial} + if err := ccPartial.collectPodAndContainerInfo(ch); err != nil { + t.Errorf("collectPodAndContainerInfo with nil PodLister returned error: %v", err) + } + if err := ccPartial.collectPodAndContainerMigInfo(ch); err != nil { + t.Errorf("collectPodAndContainerMigInfo with nil PodLister returned error: %v", err) + } + + // 3. NodeName environment variable not set + t.Setenv(util.NodeNameEnvName, "") + podLister := informers.NewSharedInformerFactory(fake.NewSimpleClientset(), 0).Core().V1().Pods().Lister() + cmNoNodeEnv := &ClusterManager{ + Zone: "test-zone", + PodLister: podLister, + containerLister: &nvidia.ContainerLister{}, + } + ccNoNodeEnv := ClusterManagerCollector{ClusterManager: cmNoNodeEnv} + if err := ccNoNodeEnv.collectPodAndContainerInfo(ch); err == nil { + t.Error("collectPodAndContainerInfo expected error when nodeName env is unset, got nil") + } + if err := ccNoNodeEnv.collectPodAndContainerMigInfo(ch); err == nil { + t.Error("collectPodAndContainerMigInfo expected error when nodeName env is unset, got nil") + } +} + +func TestCollectPodAndContainerInfo_WithPod(t *testing.T) { + t.Setenv(util.NodeNameEnvName, "test-node") + + pod := &corev1.Pod{ + ObjectMeta: metav1.ObjectMeta{ + Name: "trainer-0", + Namespace: "team-a", + UID: "uid-12345", + Annotations: map[string]string{ + util.AssignedNodeAnnotations: "test-node", + nv.MigAllocationsAnnotation: `[{"ContainerIndex":0,"DeviceIndex":0,"GPUUUID":"GPU-1","MigUUID":"MIG-1","Profile":"1g.5gb","GPUInstanceID":1,"ComputeInstanceID":1}]`, + }, + }, + Spec: corev1.PodSpec{ + NodeName: "test-node", + Containers: []corev1.Container{ + {Name: "worker"}, + }, + }, + } + + client := fake.NewSimpleClientset(pod) + informerFactory := informers.NewSharedInformerFactory(client, 0) + podInformer := informerFactory.Core().V1().Pods().Informer() + podInformer.GetIndexer().Add(pod) + podLister := informerFactory.Core().V1().Pods().Lister() + + c := &ClusterManager{ + Zone: "test-zone", + PodLister: podLister, + containerLister: &nvidia.ContainerLister{}, + LegacyMetrics: true, + } + cc := ClusterManagerCollector{ClusterManager: c} + ch := make(chan prometheus.Metric, 50) + + if err := cc.collectPodAndContainerInfo(ch); err != nil { + t.Errorf("collectPodAndContainerInfo returned error: %v", err) + } + + if err := cc.collectPodAndContainerMigInfo(ch); err != nil { + t.Errorf("collectPodAndContainerMigInfo returned error: %v", err) + } + + // Test MIG allocation error decoding path + podBadMig := &corev1.Pod{ + ObjectMeta: metav1.ObjectMeta{ + Name: "trainer-bad", + Namespace: "team-a", + UID: "uid-67890", + Annotations: map[string]string{ + util.AssignedNodeAnnotations: "test-node", + nv.MigAllocationsAnnotation: "invalid-json", + }, + }, + Spec: corev1.PodSpec{ + Containers: []corev1.Container{{Name: "worker"}}, + }, + } + podInformer.GetIndexer().Add(podBadMig) + _ = cc.collectPodAndContainerMigInfo(ch) +} + +func TestValidateEnvVars(t *testing.T) { + t.Setenv("HOOK_PATH", "") + os.Unsetenv("HOOK_PATH") + if err := ValidateEnvVars(); err == nil { + t.Error("expected error when HOOK_PATH is unset, got nil") + } + + t.Setenv("HOOK_PATH", "/tmp") + if err := ValidateEnvVars(); err != nil { + t.Errorf("unexpected error when HOOK_PATH is set: %v", err) + } +} + +func TestSendMetric_Errors(t *testing.T) { + ch := make(chan prometheus.Metric, 10) + + err := sendMetric(ch, ctrDeviceMigInfo, prometheus.GaugeValue, 1, "extra", "labels", "that", "exceed", "desc", "count", "extra1", "extra2", "extra3", "extra4") + if err == nil { + t.Error("sendMetric expected error for wrong label count, got nil") + } + + sendLegacyMetric(ch, legacyHostGPUdesc, prometheus.GaugeValue, 1, "extra", "labels", "that", "exceed", "desc", "count", "extra1", "extra2") +}