Adding Unit test and fixing typos

Signed-off-by: Vishesh Tanksale <vtanksale@nvidia.com>
visheshtanksale · Jun 24, 2024 · ac2cb43 · ac2cb43
1 parent 4e2818b
commit ac2cb43
Show file tree

Hide file tree

Showing 9 changed files with 126 additions and 38 deletions.
diff --git a/internal/lm/nvml.go b/internal/lm/nvml.go
@@ -73,7 +73,7 @@ func NewDeviceLabeler(manager resource.Manager, config *spec.Config) (Labeler, e
 		return nil, fmt.Errorf("error creating resource labeler: %v", err)
 	}
 
-	gpuModeLabeler, err := newGPUModeLabeler(manager)
+	gpuModeLabeler, err := newGPUModeLabeler(devices)
 	if err != nil {
 		return nil, fmt.Errorf("error creating resource labeler: %v", err)
 	}
@@ -202,15 +202,9 @@ 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
-	}
+// newGPUModeLabeler creates a new labeler that reports the mode of GPUs on the node.
+// GPUs can be in Graphics or Compute mode.
+func newGPUModeLabeler(devices []resource.Device) (Labeler, error) {
 	classes, err := getDeviceClasses(devices)
 	if err != nil {
 		return nil, err
@@ -244,7 +238,7 @@ func getModeForClasses(classes []uint32) string {
 func getDeviceClasses(devices []resource.Device) ([]uint32, error) {
 	seenClasses := make(map[uint32]bool)
 	for _, d := range devices {
-		class, err := d.GetPIEClass()
+		class, err := d.GetPCIClass()
 		if err != nil {
 			return nil, err
 		}

diff --git a/internal/lm/nvml_test.go b/internal/lm/nvml_test.go
@@ -204,3 +204,89 @@ func TestSharingLabeler(t *testing.T) {
 		})
 	}
 }
+
+func TestGPUModeLabeler(t *testing.T) {
+	testCases := []struct {
+		description    string
+		devices        []resource.Device
+		expectedError  bool
+		expectedLabels map[string]string
+	}{
+		{
+			description: "single device with compute PCI class",
+			devices: []resource.Device{
+				rt.NewDeviceWithPCIClassMock(0x030000),
+			},
+			expectedLabels: map[string]string{
+				"nvidia.com/gpu.mode": "graphics",
+			},
+		},
+		{
+			description: "single device with graphics PCI class",
+			devices: []resource.Device{
+				rt.NewDeviceWithPCIClassMock(0x030200),
+			},
+			expectedLabels: map[string]string{
+				"nvidia.com/gpu.mode": "compute",
+			},
+		},
+		{
+			description: "single device with switch PCI class",
+			devices: []resource.Device{
+				rt.NewDeviceWithPCIClassMock(0x068000),
+			},
+			expectedLabels: map[string]string{
+				"nvidia.com/gpu.mode": "unknown",
+			},
+		},
+		{
+			description: "multiple device have same graphics PCI class",
+			devices: []resource.Device{
+				rt.NewDeviceWithPCIClassMock(0x030200),
+				rt.NewDeviceWithPCIClassMock(0x030200),
+				rt.NewDeviceWithPCIClassMock(0x030200),
+			},
+			expectedLabels: map[string]string{
+				"nvidia.com/gpu.mode": "compute",
+			},
+		},
+		{
+			description: "multiple device have same compute PCI class",
+			devices: []resource.Device{
+				rt.NewDeviceWithPCIClassMock(0x030000),
+				rt.NewDeviceWithPCIClassMock(0x030000),
+				rt.NewDeviceWithPCIClassMock(0x030000),
+			},
+			expectedLabels: map[string]string{
+				"nvidia.com/gpu.mode": "graphics",
+			},
+		},
+		{
+			description: "multiple device with some with graphics and others with compute PCI class",
+			devices: []resource.Device{
+				rt.NewDeviceWithPCIClassMock(0x030000),
+				rt.NewDeviceWithPCIClassMock(0x030200),
+				rt.NewDeviceWithPCIClassMock(0x030000),
+			},
+			expectedLabels: map[string]string{
+				"nvidia.com/gpu.mode": "unknown",
+			},
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.description, func(t *testing.T) {
+
+			gpuModeLabeler, _ := newGPUModeLabeler(tc.devices)
+
+			labels, err := gpuModeLabeler.Labels()
+			if tc.expectedError {
+				require.Error(t, err)
+			} else {
+				require.NoError(t, err)
+			}
+
+			require.EqualValues(t, tc.expectedLabels, labels)
+		})
+	}
+}
diff --git a/internal/resource/cuda-device.go b/internal/resource/cuda-device.go
@@ -97,6 +97,6 @@ func (d *cudaDevice) IsMigEnabled() (bool, error) {
 	return false, nil
 }
 
-func (d *cudaDevice) GetPIEClass() (uint32, error) {
+func (d *cudaDevice) GetPCIClass() (uint32, error) {
 	return 0, nil
 }
diff --git a/internal/resource/device_mock.go b/internal/resource/device_mock.go
diff --git a/internal/resource/nvml-device.go b/internal/resource/nvml-device.go
@@ -88,7 +88,7 @@ func (d nvmlDevice) GetTotalMemoryMB() (uint64, error) {
 	return info.Total / (1024 * 1024), nil
 }
 
-func (d nvmlDevice) GetPIEClass() (uint32, error) {
+func (d nvmlDevice) GetPCIClass() (uint32, error) {
 	pciBusID, err := d.GetPCIBusID()
 	if err != nil {
 		return 0, err

diff --git a/internal/resource/nvml-mig-device.go b/internal/resource/nvml-mig-device.go
@@ -134,11 +134,12 @@ func totalMemory(attr map[string]interface{}) (uint64, error) {
 	}
 }
 
-func (d nvmlMigDevice) GetPIEClass() (uint32, error) {
+func (d nvmlMigDevice) GetPCIClass() (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 {

diff --git a/internal/resource/sysfs-device.go b/internal/resource/sysfs-device.go
@@ -65,6 +65,6 @@ func (d vfioDevice) IsMigCapable() (bool, error) {
 	return false, nil
 }
 
-func (d vfioDevice) GetPIEClass() (uint32, error) {
+func (d vfioDevice) GetPCIClass() (uint32, error) {
 	return d.nvidiaPCIDevice.Class, nil
 }
diff --git a/internal/resource/testing/resource-testing.go b/internal/resource/testing/resource-testing.go
@@ -51,7 +51,14 @@ 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 },
+		GetPCIClassFunc:      func() (uint32, error) { return 0x030000, nil },
+	}}
+	return &d
+}
+
+func NewDeviceWithPCIClassMock(pciClass uint32) *DeviceMock {
+	d := DeviceMock{resource.DeviceMock{
+		GetPCIClassFunc: func() (uint32, error) { return pciClass, nil },
 	}}
 	return &d
 }

diff --git a/internal/resource/types.go b/internal/resource/types.go
@@ -39,5 +39,5 @@ type Device interface {
 	GetTotalMemoryMB() (uint64, error)
 	GetDeviceHandleFromMigDeviceHandle() (Device, error)
 	GetCudaComputeCapability() (int, int, error)
-	GetPIEClass() (uint32, error)
+	GetPCIClass() (uint32, error)
 }