Refactor how device identifiers are parsed before performing automatic CDI spec generation

Signed-off-by: Christopher Desiniotis <[email protected]>

Author: cdesiniotis

go.mod

@@ -6,6 +6,7 @@ require (
 	github.com/NVIDIA/go-nvlib v0.0.0-20231116150931-9fd385bace0d
 	github.com/NVIDIA/go-nvml v0.12.0-1.0.20231020145430-e06766c5e74f
 	github.com/fsnotify/fsnotify v1.5.4
+	github.com/google/uuid v1.4.0
 	github.com/opencontainers/runtime-spec v1.1.0
 	github.com/pelletier/go-toml v1.9.4
 	github.com/sirupsen/logrus v1.9.0

go.sum

@@ -15,6 +15,8 @@ github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSs
 github.com/fsnotify/fsnotify v1.5.4 h1:jRbGcIw6P2Meqdwuo0H1p6JVLbL5DHKAKlYndzMwVZI=
 github.com/fsnotify/fsnotify v1.5.4/go.mod h1:OVB6XrOHzAwXMpEM7uPOzcehqUV2UqJxmVXmkdnm1bU=
 github.com/google/uuid v1.3.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
+github.com/google/uuid v1.4.0 h1:MtMxsa51/r9yyhkyLsVeVt0B+BGQZzpQiTQ4eHZ8bc4=
+github.com/google/uuid v1.4.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
 github.com/hashicorp/errwrap v1.0.0/go.mod h1:YH+1FKiLXxHSkmPseP+kNlulaMuP3n2brvKWEqk/Jc4=
 github.com/hashicorp/errwrap v1.1.0 h1:OxrOeh75EUXMY8TBjag2fzXGZ40LB6IKw45YeGUDY2I=
 github.com/hashicorp/errwrap v1.1.0/go.mod h1:YH+1FKiLXxHSkmPseP+kNlulaMuP3n2brvKWEqk/Jc4=

pkg/nvcdi/identifier.go

@@ -0,0 +1,76 @@
+package nvcdi
+
+import (
+	"strconv"
+	"strings"
+
+	"github.com/google/uuid"
+)
+
+type identifier string
+
+// isGPUIndex checks if an identifier is a full GPU index
+func (i identifier) isGpuIndex() bool {
+	if _, err := strconv.ParseUint(string(i), 10, 0); err != nil {
+		return false
+	}
+	return true
+}
+
+// isMigIndex checks if an identifier is a MIG index
+func (i identifier) isMigIndex() bool {
+	split := strings.SplitN(string(i), ":", 2)
+	if len(split) != 2 {
+		return false
+	}
+	for _, s := range split {
+		if _, err := strconv.ParseUint(s, 10, 0); err != nil {
+			return false
+		}
+	}
+	return true
+}
+
+// isUUID checks if an identifier is a UUID
+func (i identifier) isUUID() bool {
+	return i.isGpuUUID() || i.isMigUUID()
+}
+
+// isGpuUUID checks if an identifier is a GPU UUID
+// A GPU UUID must be of the form GPU-b1028956-cfa2-0990-bf4a-5da9abb51763
+func (i identifier) isGpuUUID() bool {
+	if !strings.HasPrefix(string(i), "GPU-") {
+		return false
+	}
+	_, err := uuid.Parse(strings.TrimPrefix(string(i), "GPU-"))
+	return err == nil
+}
+
+// isMigUUID checks if an identifier is a MIG UUID
+// A MIG UUID can be of one of two forms:
+//   - MIG-b1028956-cfa2-0990-bf4a-5da9abb51763
+//   - MIG-GPU-b1028956-cfa2-0990-bf4a-5da9abb51763/3/0
+func (i identifier) isMigUUID() bool {
+	if !strings.HasPrefix(string(i), "MIG-") {
+		return false
+	}
+	suffix := strings.TrimPrefix(string(i), "MIG-")
+	_, err := uuid.Parse(suffix)
+	if err == nil {
+		return true
+	}
+	split := strings.SplitN(suffix, "/", 3)
+	if len(split) != 3 {
+		return false
+	}
+	if !identifier(split[0]).isGpuUUID() {
+		return false
+	}
+	for _, s := range split[1:] {
+		_, err := strconv.ParseUint(s, 10, 0)
+		if err != nil {
+			return false
+		}
+	}
+	return true
+}

pkg/nvcdi/identifier_test.go

@@ -0,0 +1,90 @@
+package nvcdi
+
+import (
+	"fmt"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestIsGpuIndex(t *testing.T) {
+	testCases := []struct {
+		id       string
+		expected bool
+	}{
+		{"", false},
+		{"0", true},
+		{"1", true},
+		{"not an integer", false},
+	}
+	for i, tc := range testCases {
+		t.Run(fmt.Sprintf("test case %d", i), func(t *testing.T) {
+			actual := identifier(tc.id).isGpuIndex()
+			require.Equal(t, tc.expected, actual)
+		})
+	}
+}
+
+func TestIsMigIndex(t *testing.T) {
+	testCases := []struct {
+		id       string
+		expected bool
+	}{
+		{"", false},
+		{"0", false},
+		{"not an integer", false},
+		{"0:0", true},
+		{"0:0:0", false},
+		{"0:0.0", false},
+		{"0:foo", false},
+		{"foo:0", false},
+	}
+	for i, tc := range testCases {
+		t.Run(fmt.Sprintf("test case %d", i), func(t *testing.T) {
+			actual := identifier(tc.id).isMigIndex()
+			require.Equal(t, tc.expected, actual)
+		})
+	}
+}
+
+func TestIsGpuUUID(t *testing.T) {
+	testCases := []struct {
+		id       string
+		expected bool
+	}{
+		{"", false},
+		{"0", false},
+		{"not an integer", false},
+		{"GPU-foo", false},
+		{"GPU-ebd34bdf-1083-eaac-2aff-4b71a022f9bd", true},
+		{"MIG-ebd34bdf-1083-eaac-2aff-4b71a022f9bd", false},
+		{"ebd34bdf-1083-eaac-2aff-4b71a022f9bd", false},
+	}
+	for i, tc := range testCases {
+		t.Run(fmt.Sprintf("test case %d", i), func(t *testing.T) {
+			actual := identifier(tc.id).isGpuUUID()
+			require.Equal(t, tc.expected, actual)
+		})
+	}
+}
+
+func TestIsMigUUID(t *testing.T) {
+	testCases := []struct {
+		id       string
+		expected bool
+	}{
+		{"", false},
+		{"0", false},
+		{"not an integer", false},
+		{"MIG-foo", false},
+		{"MIG-ebd34bdf-1083-eaac-2aff-4b71a022f9bd", true},
+		{"GPU-ebd34bdf-1083-eaac-2aff-4b71a022f9bd", false},
+		{"ebd34bdf-1083-eaac-2aff-4b71a022f9bd", false},
+	}
+	for i, tc := range testCases {
+		t.Run(fmt.Sprintf("test case %d", i), func(t *testing.T) {
+			actual := identifier(tc.id).isMigUUID()
+			require.Equal(t, tc.expected, actual)
+		})
+	}
+}

pkg/nvcdi/lib-nvml.go

@@ -124,22 +124,37 @@ func (l *nvmllib) GetDeviceSpecsByID(identifiers ...string) ([]specs.Device, err
 
 // TODO: move this to go-nvlib?
 func (l *nvmllib) getNVMLDevicesByID(identifiers ...string) ([]nvml.Device, error) {
-	devices := []nvml.Device{}
+	var devices []nvml.Device
 	for _, id := range identifiers {
-		if dev, err := l.nvmllib.DeviceGetHandleByUUID(id); err == nvml.SUCCESS {
-			devices = append(devices, dev)
-			continue
+		dev, err := l.getNVMLDeviceByID(id)
+		if err != nvml.SUCCESS {
+			return nil, fmt.Errorf("failed to get NVML device handle for identifier %q: %w", id, err)
 		}
-		// TODO: check for a MIG device index
+		devices = append(devices, dev)
+	}
+	return devices, nil
+}
+
+func (l *nvmllib) getNVMLDeviceByID(id string) (nvml.Device, error) {
+	var err error
+	devID := identifier(id)
+
+	if devID.isUUID() {
+		return l.nvmllib.DeviceGetHandleByUUID(id)
+	}
+
+	if devID.isGpuIndex() {
 		if idx, err := strconv.Atoi(id); err == nil {
-			if dev, err := l.nvmllib.DeviceGetHandleByIndex(idx); err == nvml.SUCCESS {
-				devices = append(devices, dev)
-				continue
-			}
+			return l.nvmllib.DeviceGetHandleByIndex(idx)
 		}
-		return nil, fmt.Errorf("failed to get NVML device handle for identifier %q", id)
+		return nil, fmt.Errorf("failed to convert device index to an int: %w", err)
 	}
-	return devices, nil
+
+	if devID.isMigIndex() {
+		return nil, fmt.Errorf("MIG index is not supported")
+	}
+
+	return nil, fmt.Errorf("identifier is not a valid UUID or index: %q", id)
 }
 
 func (l *nvmllib) getGPUDeviceSpecs() ([]specs.Device, error) {