Merge pull request #72 from cdesiniotis/update-pciids

Update pci.ids file and update logic to extract the device name

pkg/device_plugin/device_plugin.go

@@ -30,6 +30,7 @@ package device_plugin
 
 import (
 	"bufio"
+	"fmt"
 	"io/ioutil"
 	"log"
 	"os"
@@ -41,6 +42,10 @@ import (
 	pluginapi "k8s.io/kubelet/pkg/apis/deviceplugin/v1beta1"
 )
 
+const (
+	nvidiaVendorID = "10de"
+)
+
 //Structure to hold details about Nvidia GPU Device
 type NvidiaGpuDevice struct {
 	addr string // PCI address of device
@@ -308,34 +313,62 @@ func getDeviceName(deviceID string) string {
 	}
 	defer file.Close()
 
-	scanner := bufio.NewScanner(file)
+	// Locate beginning of NVIDIA device list in pci.ids file
+	scanner, err := locateVendor(file, nvidiaVendorID)
+	if err != nil {
+		log.Printf("Error locating NVIDIA in pci.ds file: %v", err)
+		return ""
+	}
+
+	// Find NVIDIA device by device id
+	prefix := fmt.Sprintf("\t%s", deviceID)
 	for scanner.Scan() {
 		line := scanner.Text()
-		line = strings.TrimSpace(line)
-		if line != "" && strings.Contains(line, deviceID) {
-			splits := strings.Split(line, deviceID)
-			if len(splits) != 2 {
-				log.Printf("Error processing pci.ids file at line: %s", line)
-				return deviceName
-			}
-			// fix matching
-			if splits[0] != "" {
-				continue
-			}
-			deviceName = strings.TrimSpace(splits[1])
-			deviceName = strings.ToUpper(deviceName)
-			deviceName = strings.Replace(deviceName, "/", "_", -1)
-			deviceName = strings.Replace(deviceName, ".", "_", -1)
-			reg, _ := regexp.Compile("\\s+")
-			deviceName = reg.ReplaceAllString(deviceName, "_") // Replace all spaces with underscore
-			reg, _ = regexp.Compile("[^a-zA-Z0-9_.]+")
-			deviceName = reg.ReplaceAllString(deviceName, "") // Removes any char other than alphanumeric and underscore
-			break
+		// ignore comments
+		if strings.HasPrefix(line, "#") {
+			continue
+		}
+		// if line does not start with tab, we are visiting a different vendor
+		if !strings.HasPrefix(line, "\t") {
+			log.Printf("Could not find NVIDIA device with id: %s", deviceID)
+			return ""
 		}
+		if !strings.HasPrefix(line, prefix) {
+			continue
+		}
+
+		deviceName = strings.TrimPrefix(line, prefix)
+		deviceName = strings.TrimSpace(deviceName)
+		deviceName = strings.ToUpper(deviceName)
+		deviceName = strings.Replace(deviceName, "/", "_", -1)
+		deviceName = strings.Replace(deviceName, ".", "_", -1)
+		// Replace all spaces with underscore
+		reg, _ := regexp.Compile("\\s+")
+		deviceName = reg.ReplaceAllString(deviceName, "_")
+		// Removes any char other than alphanumeric and underscore
+		reg, _ = regexp.Compile("[^a-zA-Z0-9_.]+")
+		deviceName = reg.ReplaceAllString(deviceName, "")
+		break
 	}
 
 	if err := scanner.Err(); err != nil {
 		log.Printf("Error reading pci ids file %s", err)
 	}
 	return deviceName
 }
+
+func locateVendor(pciIdsFile *os.File, vendorID string) (*bufio.Scanner, error) {
+	scanner := bufio.NewScanner(pciIdsFile)
+	for scanner.Scan() {
+		line := scanner.Text()
+		if strings.HasPrefix(line, vendorID) {
+			return scanner, nil
+		}
+	}
+
+	if err := scanner.Err(); err != nil {
+		return scanner, fmt.Errorf("error reading pci.ids file: %v", err)
+	}
+
+	return scanner, fmt.Errorf("failed to find vendor id in pci.ids file: %s", vendorID)
+}

pkg/device_plugin/device_plugin_test.go

@@ -328,7 +328,16 @@ var _ = Describe("Device Plugin", func() {
 		BeforeEach(func() {
 			workDir, err = ioutil.TempDir("", "pci-test")
 			Expect(err).ToNot(HaveOccurred())
-			message := []byte("118a  GK104GL [GRID K520] \n 118b  GK104GL [GRID K2 GeForce USM] \n 118c  GK104 [GRID 118c NVS USM] \n 118d  gk104gl [grid k520] \n 118e gk104.gl [grid/k./520]")
+			message := []byte(`
+8086  Intel Corporation
+	2331  DH89xxCC Chap Counter
+10de NVIDIA Corporation
+	118a  GK104GL [GRID K520]
+	118b  GK104GL [GRID K2 GeForce USM]
+	118d  gk104gl [grid k520]
+	118e gk104.gl [grid/k./520]
+	2331 GH100 [H100 PCIe]
+`)
 			ioutil.WriteFile(filepath.Join(workDir, "pci.ids"), message, 0644)
 		})
 
@@ -343,11 +352,6 @@ var _ = Describe("Device Plugin", func() {
 			Expect(deviceName).To(Equal(""))
 		})
 
-		It("Returns blank if the device name is not correctly formatted", func() {
-			deviceName := getDeviceName("118c")
-			Expect(deviceName).To(Equal(""))
-		})
-
 		It("Returns blank if error reading the pci.ids file", func() {
 			pciIdsFilePath = filepath.Join(workDir, "fake")
 			deviceName := getDeviceName("118c")
@@ -365,5 +369,11 @@ var _ = Describe("Device Plugin", func() {
 			deviceName := getDeviceName("118e")
 			Expect(deviceName).To(Equal("GK104_GL_GRID_K__520"))
 		})
+
+		It("Retrieves correct device name from pci.ids file even if another vendor's device shares device id", func() {
+			pciIdsFilePath = filepath.Join(workDir, "pci.ids")
+			deviceName := getDeviceName("2331")
+			Expect(deviceName).To(Equal("GH100_H100_PCIE"))
+		})
 	})
 })