nvswitch: Add NVSWITCH_ALIAS

Similar to P_GPU_ALIAS add NVSWTICH_ALIAS to override
the name that is originally given by the DP

  nvidia.com/nvswitch:                        4
  nvidia.com/pgpu:                            8

Signed-off-by: Zvonko Kaiser <zkaiser@nvidia.com>

Author: zvonkok

cmd/main.go

@@ -35,5 +35,6 @@ import (
 
 func main() {
 	device_plugin.PGPUAlias = os.Getenv("P_GPU_ALIAS")
+	device_plugin.NVSwitchAlias = os.Getenv("NVSWITCH_ALIAS")
 	device_plugin.InitiateDevicePlugin()
 }

pkg/device_plugin/cdi.go

@@ -68,7 +68,11 @@ func GenerateCDISpec() error {
 
 	// Generate NVSwitch CDI spec if we have NVSwitch devices
 	if len(nvSwitchDeviceIDs) > 0 {
-		if err := generateCDISpecForClass(cdiNVSwitchClass, true); err != nil {
+		nvSwitchClass := cdiNVSwitchClass
+		if NVSwitchAlias != "" {
+			nvSwitchClass = NVSwitchAlias
+		}
+		if err := generateCDISpecForClass(nvSwitchClass, true); err != nil {
 			return fmt.Errorf("failed to generate NVSwitch CDI spec: %w", err)
 		}
 	}

pkg/device_plugin/device_plugin.go

@@ -64,6 +64,7 @@ var nvpciLib nvpci.Interface
 var startDevicePlugin = startDevicePluginFunc
 var stop = make(chan struct{})
 var PGPUAlias string
+var NVSwitchAlias string
 
 func InitiateDevicePlugin() {
 	// Initialize nvpci library if not already set (allows injection for testing)
@@ -99,14 +100,15 @@ func createDevicePlugins() {
 			})
 		}
 
-		// Determine device name - NVSwitches always use their actual name,
-		// GPUs can use PGPUAlias if set
+		// Determine device name - use alias if set, otherwise use actual device name
 		var deviceName string
 		if isNVSwitchDeviceID(deviceID) {
-			// NVSwitches always use their actual device name
-			deviceName = getDeviceNameForID(deviceID)
+			if NVSwitchAlias != "" {
+				deviceName = NVSwitchAlias
+			} else {
+				deviceName = getDeviceNameForID(deviceID)
+			}
 		} else if PGPUAlias != "" {
-			// GPUs can use the alias
 			deviceName = PGPUAlias
 		} else {
 			deviceName = getDeviceNameForID(deviceID)

pkg/device_plugin/device_plugin_test.go

@@ -286,5 +286,26 @@ var _ = Describe("Device Plugin", func() {
 			result := getDeviceNameForID("abcd")
 			Expect(result).To(Equal(""))
 		})
+
+		It("formats device names with special characters through formatDeviceName", func() {
+			iommuMap = map[string][]NvidiaPCIDevice{
+				"1": {
+					{
+						Address:    "0000:01:00.0",
+						DeviceID:   0x2330,
+						DeviceName: "NVIDIA H100 PCIe [Hopper]",
+						IommuGroup: 1,
+					},
+				},
+			}
+			result := getDeviceNameForID("2330")
+			Expect(result).To(Equal("NVIDIA_H100_PCIE_HOPPER"))
+		})
+
+		It("returns empty string when iommuMap is empty", func() {
+			iommuMap = map[string][]NvidiaPCIDevice{}
+			result := getDeviceNameForID("1b80")
+			Expect(result).To(Equal(""))
+		})
 	})
 })