Merge pull request #65 from mresvanis/update-mdev-type-name-parsing

cdesiniotis · web-flow · commit f626d43849b6 · 2025-07-28T14:46:31.000-07:00
Update mdev type parsing to handle multiple spaces in type strings
diff --git a/pkg/nvmdev/nvmdev.go b/pkg/nvmdev/nvmdev.go
@@ -226,14 +226,27 @@ func (m mdev) Type() (string, error) {
 	if err != nil {
 		return "", fmt.Errorf("unable to read mdev_type name for mdev %s: %v", m, err)
 	}
-	// file in the format: [NVIDIA|GRID] <vGPU type>
-	mdevTypeStr := strings.TrimSpace(string(mdevType))
-	mdevTypeSplit := strings.SplitN(mdevTypeStr, " ", 2)
-	if len(mdevTypeSplit) != 2 {
-		return "", fmt.Errorf("unable to parse mdev_type name %s for mdev %s", mdevTypeStr, m)
+	typeName, err := parseMdevTypeName(string(mdevType))
+	if err != nil {
+		return "", fmt.Errorf("unable to parse mdev_type name for mdev %s: %v", m, err)
 	}
 
-	return mdevTypeSplit[1], nil
+	return typeName, nil
+}
+
+// parseMdevTypeName extracts the vGPU type name from a string that may contain
+// product prefixes.
+// Examples:
+//   - "NVIDIA A100-4C" -> "A100-4C".
+//   - "NVIDIA RTX Pro 6000 Blackwell DC-48C" -> "DC-48C"
+func parseMdevTypeName(rawName string) (string, error) {
+	nameStr := strings.TrimSpace(rawName)
+	nameSplit := strings.Split(nameStr, " ")
+	typeName := nameSplit[len(nameSplit)-1]
+	if typeName == "" {
+		return "", fmt.Errorf("unable to parse mdev_type name from: %s", rawName)
+	}
+	return typeName, nil
 }
 
 func (m mdev) driver() (string, error) {
@@ -280,13 +293,10 @@ func (m *nvmdev) NewParentDevice(devicePath string) (*ParentDevice, error) {
 		if err != nil {
 			return nil, fmt.Errorf("unable to read file %s: %v", path, err)
 		}
-		// file in the format: [NVIDIA|GRID] <vGPU type>
-		nameStr := strings.TrimSpace(string(name))
-		nameSplit := strings.SplitN(nameStr, " ", 2)
-		if len(nameSplit) != 2 {
-			return nil, fmt.Errorf("unable to parse mdev_type name %s at path %s", nameStr, path)
+		nameStr, err := parseMdevTypeName(string(name))
+		if err != nil {
+			return nil, fmt.Errorf("unable to parse mdev_type name at path %s: %v", path, err)
 		}
-		nameStr = nameSplit[len(nameSplit)-1]
 
 		mdevTypesMap[nameStr] = filepath.Dir(path)
 	}
diff --git a/pkg/nvmdev/nvmdev_test.go b/pkg/nvmdev/nvmdev_test.go
@@ -61,3 +61,66 @@ func TestNvmdev(t *testing.T) {
 	pf = mdevA100.GetPhysicalFunction()
 	require.Equal(t, "0000:3b:04.1", pf.Address, "Wrong address for Mock A100 physical function")
 }
+
+func TestParseMdevTypeName(t *testing.T) {
+	testCases := []struct {
+		name         string
+		mdevTypeStr  string
+		expectedType string
+		expectError  bool
+	}{
+		{
+			name:         "NVIDIA prefix format",
+			mdevTypeStr:  "NVIDIA A100-4C",
+			expectedType: "A100-4C",
+			expectError:  false,
+		},
+		{
+			name:         "GRID prefix format",
+			mdevTypeStr:  "GRID V100-8Q",
+			expectedType: "V100-8Q",
+			expectError:  false,
+		},
+		{
+			name:         "Multi-word NVIDIA prefix format",
+			mdevTypeStr:  "NVIDIA RTX Pro 6000 Blackwell A100-4C",
+			expectedType: "A100-4C",
+			expectError:  false,
+		},
+		{
+			name:         "Complex multi-word prefix",
+			mdevTypeStr:  "NVIDIA RTX A6000 Ada Generation H100-8C",
+			expectedType: "H100-8C",
+			expectError:  false,
+		},
+		{
+			name:         "Single word only",
+			mdevTypeStr:  "A100-4C",
+			expectedType: "A100-4C",
+			expectError:  false,
+		},
+		{
+			name:        "Empty string",
+			mdevTypeStr: "",
+			expectError: true,
+		},
+		{
+			name:        "Only spaces",
+			mdevTypeStr: "   ",
+			expectError: true,
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			actualType, err := parseMdevTypeName(tc.mdevTypeStr)
+
+			if tc.expectError {
+				require.Error(t, err)
+			} else {
+				require.NoError(t, err)
+				require.Equal(t, tc.expectedType, actualType)
+			}
+		})
+	}
+}
