Add support for node labels to report GPU mode

docs/gpu-feature-discovery/README.md

@@ -209,6 +209,7 @@ their meaning:
 | nvidia.com/gpu.machine         | String     | Machine type                                 | DGX-1          |
 | nvidia.com/gpu.memory          | Integer    | Memory of the GPU in Mb                      | 2048           |
 | nvidia.com/gpu.product         | String     | Model of the GPU                             | GeForce-GT-710 |
+| nvidia.com/gpu.mode            | String     | Display or Compute Mode of the GPU           | compute        |
 
 Depending on the MIG strategy used, the following set of labels may also be
 available (or override the default values for some of the labels listed above):

go.mod

@@ -4,7 +4,7 @@ go 1.22.2
 
 require (
 	github.com/NVIDIA/go-gpuallocator v0.5.0
-	github.com/NVIDIA/go-nvlib v0.5.0
+	github.com/NVIDIA/go-nvlib v0.6.0
 	github.com/NVIDIA/go-nvml v0.12.4-0
 	github.com/NVIDIA/nvidia-container-toolkit v1.15.1-0.20240528113255-e4b46a09a77e
 	github.com/fsnotify/fsnotify v1.7.0

go.sum

@@ -23,8 +23,8 @@ github.com/Microsoft/hcsshim v0.11.4 h1:68vKo2VN8DE9AdN4tnkWnmdhqdbpUFM8OF3Airm7
 github.com/Microsoft/hcsshim v0.11.4/go.mod h1:smjE4dvqPX9Zldna+t5FG3rnoHhaB7QYxPRqGcpAD9w=
 github.com/NVIDIA/go-gpuallocator v0.5.0 h1:166ICvPv2dU9oZ2J3kJ4y3XdbGCi6LhXgFZJtrqeu3A=
 github.com/NVIDIA/go-gpuallocator v0.5.0/go.mod h1:zos5bTIN01hpQioOyu9oRKglrznImMQvm0bZllMmckw=
-github.com/NVIDIA/go-nvlib v0.5.0 h1:951KGrfr+p3cs89alO9z/ZxPPWKxwht9tx9rxiADoLI=
-github.com/NVIDIA/go-nvlib v0.5.0/go.mod h1:87z49ULPr4GWPSGfSIp3taU4XENRYN/enIg88MzcL4k=
+github.com/NVIDIA/go-nvlib v0.6.0 h1:zAMBzCYT9xeyRQo0tb7HJbStkzajD6e5joyaQqJ2OGU=
+github.com/NVIDIA/go-nvlib v0.6.0/go.mod h1:9UrsLGx/q1OrENygXjOuM5Ey5KCtiZhbvBlbUIxtGWY=
 github.com/NVIDIA/go-nvml v0.12.4-0 h1:4tkbB3pT1O77JGr0gQ6uD8FrsUPqP1A/EOEm2wI1TUg=
 github.com/NVIDIA/go-nvml v0.12.4-0/go.mod h1:8Llmj+1Rr+9VGGwZuRer5N/aCjxGuR5nPb/9ebBiIEQ=
 github.com/NVIDIA/nvidia-container-toolkit v1.15.1-0.20240528113255-e4b46a09a77e h1:1xxK51xOXiqqziqMdgEMWVmjfvRu1b265LgQQZBqRaE=

internal/lm/nvml.go

@@ -22,6 +22,8 @@ import (
 	"strconv"
 	"strings"
 
+	"github.com/NVIDIA/go-nvlib/pkg/nvpci"
+
 	spec "github.com/NVIDIA/k8s-device-plugin/api/config/v1"
 	"github.com/NVIDIA/k8s-device-plugin/internal/resource"
 )
@@ -71,12 +73,18 @@ func NewDeviceLabeler(manager resource.Manager, config *spec.Config) (Labeler, e
 		return nil, fmt.Errorf("error creating resource labeler: %v", err)
 	}
 
+	gpuModeLabeler, err := newGPUModeLabeler(manager)
+	if err != nil {
+		return nil, fmt.Errorf("error creating resource labeler: %v", err)
+	}
+
 	l := Merge(
 		machineTypeLabeler,
 		versionLabeler,
 		migCapabilityLabeler,
 		sharingLabeler,
 		resourceLabeler,
+		gpuModeLabeler,
 	)
 
 	return l, nil
@@ -193,3 +201,59 @@ func isMPSCapable(manager resource.Manager) (bool, error) {
 	}
 	return true, nil
 }
+
+func newGPUModeLabeler(manager resource.Manager) (Labeler, error) {
+	devices, err := manager.GetDevices()
+	if err != nil {
+		return nil, err
+	}
+	if len(devices) == 0 {
+		// no devices, return empty labels
+		return empty{}, nil
+	}
+	classes, err := getDeviceClasses(devices)
+	if err != nil {
+		return nil, err
+	}
+	gpuMode := getModeForClasses(classes)
+	labels := Labels{
+		"nvidia.com/gpu.mode": gpuMode,
+	}
+	return labels, nil
+}
+
+func getModeForClasses(classes []uint32) string {
+	if len(classes) == 0 {
+		return "unknown"
+	}
+	for _, class := range classes {
+		if class != classes[0] {
+			return "unknown"
+		}
+	}
+	switch classes[0] {
+	case nvpci.PCIVgaControllerClass:
+		return "graphics"
+	case nvpci.PCI3dControllerClass:
+		return "compute"
+	default:
+		return "unknown"
+	}
+}
+
+func getDeviceClasses(devices []resource.Device) ([]uint32, error) {
+	seenClasses := make(map[uint32]bool)
+	for _, d := range devices {
+		class, err := d.GetPIEClass()
+		if err != nil {
+			return nil, err
+		}
+		seenClasses[class] = true
+	}
+
+	var classes []uint32
+	for class := range seenClasses {
+		classes = append(classes, class)
+	}
+	return classes, nil
+}

internal/resource/cuda-device.go

@@ -96,3 +96,7 @@ func (d *cudaDevice) IsMigCapable() (bool, error) {
 func (d *cudaDevice) IsMigEnabled() (bool, error) {
 	return false, nil
 }
+
+func (d *cudaDevice) GetPIEClass() (uint32, error) {
+	return 0, nil
+}

internal/resource/nvml-device.go

@@ -20,6 +20,7 @@ import (
 	"fmt"
 
 	"github.com/NVIDIA/go-nvlib/pkg/nvlib/device"
+	"github.com/NVIDIA/go-nvlib/pkg/nvpci"
 	"github.com/NVIDIA/go-nvml/pkg/nvml"
 )
 
@@ -86,3 +87,15 @@ func (d nvmlDevice) GetTotalMemoryMB() (uint64, error) {
 	}
 	return info.Total / (1024 * 1024), nil
 }
+
+func (d nvmlDevice) GetPIEClass() (uint32, error) {
+	pciBusID, err := d.GetPCIBusID()
+	if err != nil {
+		return 0, err
+	}
+	nvDevice, err := nvpci.New().GetGPUByPciBusID(pciBusID)
+	if err != nil {
+		return 0, err
+	}
+	return nvDevice.Class, nil
+}

internal/resource/nvml-mig-device.go

@@ -21,6 +21,7 @@ import (
 	"strings"
 
 	"github.com/NVIDIA/go-nvlib/pkg/nvlib/device"
+	"github.com/NVIDIA/go-nvlib/pkg/nvpci"
 	"github.com/NVIDIA/go-nvml/pkg/nvml"
 )
 
@@ -132,3 +133,23 @@ func totalMemory(attr map[string]interface{}) (uint64, error) {
 		return 0, fmt.Errorf("unsupported attribute type %v", t)
 	}
 }
+
+func (d nvmlMigDevice) GetPIEClass() (uint32, error) {
+	info, retVal := d.MigDevice.GetPciInfo()
+	if retVal != nvml.SUCCESS {
+		return 0, retVal
+	}
+	var bytes []byte
+	for _, char := range info.BusId {
+		if char == 0 {
+			break
+		}
+		bytes = append(bytes, byte(char))
+	}
+	pciID := strings.ToLower(strings.TrimPrefix(string(bytes), "0000"))
+	nvDevice, err := nvpci.New().GetGPUByPciBusID(pciID)
+	if err != nil {
+		return 0, err
+	}
+	return nvDevice.Class, nil
+}

internal/resource/sysfs-device.go

@@ -64,3 +64,7 @@ func (d vfioDevice) IsMigEnabled() (bool, error) {
 func (d vfioDevice) IsMigCapable() (bool, error) {
 	return false, nil
 }
+
+func (d vfioDevice) GetPIEClass() (uint32, error) {
+	return d.nvidiaPCIDevice.Class, nil
+}

internal/resource/testing/resource-testing.go

@@ -51,6 +51,7 @@ func NewDeviceMock(migEnabled bool) *DeviceMock {
 		IsMigEnabledFunc:     func() (bool, error) { return migEnabled, nil },
 		IsMigCapableFunc:     func() (bool, error) { return migEnabled, nil },
 		GetMigDevicesFunc:    func() ([]resource.Device, error) { return nil, nil },
+		GetPIEClassFunc:      func() (uint32, error) { return 0x030000, nil },
 	}}
 	return &d
 }

internal/resource/types.go

@@ -39,4 +39,5 @@ type Device interface {
 	GetTotalMemoryMB() (uint64, error)
 	GetDeviceHandleFromMigDeviceHandle() (Device, error)
 	GetCudaComputeCapability() (int, int, error)
+	GetPIEClass() (uint32, error)
 }

tests/expected-output-mig-mixed.txt

@@ -18,6 +18,7 @@ nvidia\.com\/gpu\.sharing-strategy=[none|mps|time-slicing]
 nvidia\.com\/gpu\.product=[A-Za-z_-]+
 nvidia\.com\/gpu\.memory=[0-9]+
 nvidia\.com\/gpu\.family=[a-z]+
+nvidia\.com\/gpu\.mode=[unknown|compute|graphics]
 nvidia\.com\/mig\.capable=[true|false]
 nvidia\.com\/gpu\.compute\.major=[0-9]+
 nvidia\.com\/gpu\.compute\.minor=[0-9]+

tests/expected-output-mig-none.txt

@@ -18,6 +18,7 @@ nvidia\.com\/gpu\.sharing-strategy=[none|mps|time-slicing]
 nvidia\.com\/gpu\.product=[A-Za-z_-]+
 nvidia\.com\/gpu\.memory=[0-9]+
 nvidia\.com\/gpu\.family=[a-z]+
+nvidia\.com\/gpu\.mode=[unknown|compute|graphics]
 nvidia\.com\/mig\.capable=[true|false]
 nvidia\.com\/gpu\.compute\.major=[0-9]+
 nvidia\.com\/gpu\.compute\.minor=[0-9]+

tests/expected-output-mig-single.txt

@@ -30,4 +30,5 @@ nvidia\.com\/gpu\.engines\.jpeg=[0-9]+
 nvidia\.com\/gpu\.engines\.ofa=[0-9]+
 nvidia\.com\/gpu\.slices\.gi=[0-9]+
 nvidia\.com\/gpu\.slices\.ci=[0-9]+
+nvidia\.com\/gpu\.mode=[unknown|compute|graphics]
 nvidia\.com\/mps\.capable=[true|false]

tests/expected-output.txt

@@ -18,6 +18,7 @@ nvidia\.com\/gpu\.sharing-strategy=[none|mps|time-slicing]
 nvidia\.com\/gpu\.product=[A-Za-z_-]+
 nvidia\.com\/gpu\.memory=[0-9]+
 nvidia\.com\/gpu\.family=[a-z]+
+nvidia\.com\/gpu\.mode=[unknown|compute|graphics]
 nvidia\.com\/mig\.capable=[true|false]
 nvidia\.com\/gpu\.compute\.major=[0-9]+
 nvidia\.com\/gpu\.compute\.minor=[0-9]+