Skip to content

refactor(vGPUmonitor): inject nvml.Interface instead of CGo package-l… #2628

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
Kanika0306 wants to merge 1 commit into from
Closed

refactor(vGPUmonitor): inject nvml.Interface instead of CGo package-l… #2628

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
17 changes: 13 additions & 4 deletions cmd/vGPUmonitor/feedback.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -133,12 +133,21 @@ func Observe(lister *nvidia.ContainerLister) {
}
}

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)")
}
Comment on lines +136 to 150
defer nvml.Shutdown()

ticker := time.NewTicker(time.Second * 5)
defer ticker.Stop()
Expand Down
5 changes: 3 additions & 2 deletions cmd/vGPUmonitor/main.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -34,6 +34,7 @@ import (
"github.com/Project-HAMi/HAMi/pkg/util/flag"
"github.com/Project-HAMi/HAMi/pkg/version"

nvmlgo "github.com/NVIDIA/go-nvml/pkg/nvml"
"github.com/prometheus/client_golang/prometheus"
"github.com/prometheus/client_golang/prometheus/promhttp"

Expand Down Expand Up @@ -101,7 +102,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, nvmlgo.New(), 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")
Expand Down Expand Up @@ -142,7 +143,7 @@ func initMetrics(ctx context.Context, containerLister *nvidia.ContainerLister) e

reg.MustRegister(versionmetrics.NewBuildInfoCollector())

NewClusterManager("vGPU", reg, containerLister, legacyMetrics)
NewClusterManager("vGPU", reg, containerLister, nvmlgo.New(), legacyMetrics)

// Uncomment to add the standard process and Go metrics to the custom registry.
//reg.MustRegister(
Expand Down
76 changes: 46 additions & 30 deletions cmd/vGPUmonitor/metrics.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@ limitations under the License.
package main

import (
"errors"
"fmt"
"os"
"time"
Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -236,45 +243,50 @@ func (cc ClusterManagerCollector) Collect(ch chan<- prometheus.Metric) {
}

func (cc ClusterManagerCollector) collectGPUInfo(ch chan<- prometheus.Metric) error {
if err := cc.initNVML(); err != 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)
}
}

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 {
Expand All @@ -290,22 +302,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
Expand All @@ -326,18 +338,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
Expand Down Expand Up @@ -573,14 +585,18 @@ 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()
}
c := &ClusterManager{
Zone: zone,
containerLister: containerLister,
LegacyMetrics: legacyMetrics,
nvmllib: nvmllib,
}

informerFactory := informers.NewSharedInformerFactoryWithOptions(containerLister.Clientset(), time.Hour*1)
Expand Down
Loading