Merge branch 'CNT-4762/extend-runtime-cdi-device-names' into 'main'

Extend the 'runtime.nvidia.com/gpu' CDI device kind to support full-GPUs...

See merge request nvidia/container-toolkit/container-toolkit!514

Author: cdesiniotis

Diff for: CHANGELOG.md

@@ -1,5 +1,8 @@
 # NVIDIA Container Toolkit Changelog
 
+## v1.15.0-rc.2
+* Extend the `runtime.nvidia.com/gpu` CDI kind to support full-GPUs and MIG devices specified by index or UUID.
+
 ## v1.15.0-rc.1
 * Skip update of ldcache in containers without ldconfig. The .so.SONAME symlinks are still created.
 * Normalize ldconfig path on use. This automatically adjust the ldconfig setting applied to ldconfig.real on systems where this exists.
@@ -10,6 +13,7 @@
 * Added support for `nvidia-ctk runtime configure --enable-cdi` for the `docker` runtime. Note that this requires Docker >= 25.
 * Fixed bug in `nvidia-ctk config` command when using `--set`. The types of applied config options are now applied correctly.
 * Add `--relative-to` option to `nvidia-ctk transform root` command. This controls whether the root transformation is applied to host or container paths.
+* Added automatic CDI spec generation when the `runtime.nvidia.com/gpu=all` device is requested by a container.
 
 * [libnvidia-container] Fix device permission check when using cgroupv2 (fixes #227)
 

Diff for: 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

Diff for: 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=

Diff for: internal/modifier/cdi.go

@@ -50,7 +50,12 @@ func NewCDIModifier(logger logger.Interface, cfg *config.Config, ociSpec oci.Spe
 		return nil, fmt.Errorf("requesting a CDI device with vendor 'runtime.nvidia.com' is not supported when requesting other CDI devices")
 	}
 	if len(automaticDevices) > 0 {
-		return newAutomaticCDISpecModifier(logger, cfg, automaticDevices)
+		automaticModifier, err := newAutomaticCDISpecModifier(logger, cfg, automaticDevices)
+		if err == nil {
+			return automaticModifier, nil
+		}
+		logger.Warningf("Failed to create the automatic CDI modifier: %w", err)
+		logger.Debugf("Falling back to the standard CDI modifier")
 	}
 
 	return cdi.New(
@@ -152,7 +157,8 @@ func getAnnotationDevices(prefixes []string, annotations map[string]string) ([]s
 func filterAutomaticDevices(devices []string) []string {
 	var automatic []string
 	for _, device := range devices {
-		if device == "runtime.nvidia.com/gpu=all" {
+		vendor, class, _ := parser.ParseDevice(device)
+		if vendor == "runtime.nvidia.com" && class == "gpu" {
 			automatic = append(automatic, device)
 		}
 	}
@@ -176,9 +182,6 @@ func newAutomaticCDISpecModifier(logger logger.Interface, cfg *config.Config, de
 	return cdiModifier, nil
 }
 
-// TODO: use the requested devices when generating the CDI spec once we add
-// automatic CDI generation for more than just the 'runtime.nvidia.com/gpu=all'
-// device
 func generateAutomaticCDISpec(logger logger.Interface, cfg *config.Config, devices []string) (spec.Interface, error) {
 	cdilib, err := nvcdi.New(
 		nvcdi.WithLogger(logger),
@@ -191,5 +194,26 @@ func generateAutomaticCDISpec(logger logger.Interface, cfg *config.Config, devic
 		return nil, fmt.Errorf("failed to construct CDI library: %w", err)
 	}
 
-	return cdilib.GetSpec()
+	identifiers := []string{}
+	for _, device := range devices {
+		_, _, id := parser.ParseDevice(device)
+		identifiers = append(identifiers, id)
+	}
+
+	deviceSpecs, err := cdilib.GetDeviceSpecsByID(identifiers...)
+	if err != nil {
+		return nil, fmt.Errorf("failed to get CDI device specs: %w", err)
+	}
+
+	commonEdits, err := cdilib.GetCommonEdits()
+	if err != nil {
+		return nil, fmt.Errorf("failed to get common CDI spec edits: %w", err)
+	}
+
+	return spec.New(
+		spec.WithDeviceSpecs(deviceSpecs),
+		spec.WithEdits(*commonEdits.ContainerEdits),
+		spec.WithVendor("runtime.nvidia.com"),
+		spec.WithClass("gpu"),
+	)
 }

Diff for: pkg/nvcdi/api.go

@@ -51,4 +51,5 @@ type Interface interface {
 	GetGPUDeviceSpecs(int, device.Device) (*specs.Device, error)
 	GetMIGDeviceEdits(device.Device, device.MigDevice) (*cdi.ContainerEdits, error)
 	GetMIGDeviceSpecs(int, device.Device, int, device.MigDevice) (*specs.Device, error)
+	GetDeviceSpecsByID(...string) ([]specs.Device, error)
 }

Diff for: pkg/nvcdi/gds.go

@@ -81,3 +81,10 @@ func (l *gdslib) GetMIGDeviceEdits(device.Device, device.MigDevice) (*cdi.Contai
 func (l *gdslib) GetMIGDeviceSpecs(int, device.Device, int, device.MigDevice) (*specs.Device, error) {
 	return nil, fmt.Errorf("GetMIGDeviceSpecs is not supported")
 }
+
+// GetDeviceSpecsByID returns the CDI device specs for the GPU(s) represented by
+// the provided identifiers, where an identifier is an index or UUID of a valid
+// GPU device.
+func (l *gdslib) GetDeviceSpecsByID(...string) ([]specs.Device, error) {
+	return nil, fmt.Errorf("GetDeviceSpecsByID is not supported")
+}

Diff for: 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
+}

Diff for: 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)
+		})
+	}
+}

Diff for: pkg/nvcdi/lib-csv.go

@@ -94,3 +94,10 @@ func (l *csvlib) GetMIGDeviceEdits(device.Device, device.MigDevice) (*cdi.Contai
 func (l *csvlib) GetMIGDeviceSpecs(int, device.Device, int, device.MigDevice) (*specs.Device, error) {
 	return nil, fmt.Errorf("GetMIGDeviceSpecs is not supported for CSV files")
 }
+
+// GetDeviceSpecsByID returns the CDI device specs for the GPU(s) represented by
+// the provided identifiers, where an identifier is an index or UUID of a valid
+// GPU device.
+func (l *csvlib) GetDeviceSpecsByID(...string) ([]specs.Device, error) {
+	return nil, fmt.Errorf("GetDeviceSpecsByID is not supported for CSV files")
+}