Generate separate specs for coherent and noncoherent devices

With this change the nvidia-ctk cdi generate command generates
CDI specs based on whether a device supports coherent access to
system memory or not. In this case "regular" nvidia.com/gpu CDI
specs are generated for all devices as well as
nvidia.com/gpu.coherent and nvidia.com/gpu.noncoherent for devices
that are either coherent or non-coherent.

Adding the --feature-flag=disable-coherent-annotations command line
argument to the nvidia-ctk cdi generate command will disable this.

The "disable-coherent-annotations" feature flag can also be set in the
nvcdi API in which case the generated CDI device specification will
not include annotations indicating coherence.

Signed-off-by: Evan Lezar <[email protected]>

Author: elezar

cmd/nvidia-ctk/cdi/generate/generate.go

@@ -18,6 +18,7 @@ package generate
 
 import (
 	"context"
+	"errors"
 	"fmt"
 	"os"
 	"path/filepath"
@@ -26,6 +27,7 @@ import (
 	"github.com/urfave/cli/v3"
 
 	cdi "tags.cncf.io/container-device-interface/pkg/parser"
+	"tags.cncf.io/container-device-interface/specs-go"
 
 	"github.com/NVIDIA/go-nvml/pkg/nvml"
 
@@ -63,6 +65,8 @@ type options struct {
 	librarySearchPaths []string
 	disabledHooks      []string
 
+	featureFlags []string
+
 	csv struct {
 		files          []string
 		ignorePatterns []string
@@ -214,6 +218,13 @@ func (m command) build() *cli.Command {
 				Destination: &opts.disabledHooks,
 				Sources:     cli.EnvVars("NVIDIA_CTK_CDI_GENERATE_DISABLED_HOOKS"),
 			},
+			&cli.StringSliceFlag{
+				Name:        "feature-flag",
+				Aliases:     []string{"feature-flags"},
+				Usage:       "specify feature flags for CDI spec generation",
+				Destination: &opts.featureFlags,
+				Sources:     cli.EnvVars("NVIDIA_CTK_CDI_GENERATE_FEATURE_FLAGS"),
+			},
 		},
 	}
 
@@ -262,21 +273,18 @@ func (m command) validateFlags(c *cli.Command, opts *options) error {
 }
 
 func (m command) run(opts *options) error {
-	spec, err := m.generateSpec(opts)
+	specs, err := m.generateSpecs(opts)
 	if err != nil {
 		return fmt.Errorf("failed to generate CDI spec: %v", err)
 	}
-	m.logger.Infof("Generated CDI spec with version %v", spec.Raw().Version)
 
-	if opts.output == "" {
-		_, err := spec.WriteTo(os.Stdout)
-		if err != nil {
-			return fmt.Errorf("failed to write CDI spec to STDOUT: %v", err)
-		}
-		return nil
-	}
+	var errs error
+	for _, spec := range specs {
+		m.logger.Infof("Generated CDI spec with version %v", spec.Raw().Version)
 
-	return spec.Save(opts.output)
+		errs = errors.Join(errs, spec.Save(opts.output))
+	}
+	return errs
 }
 
 func formatFromFilename(filename string) string {
@@ -291,7 +299,34 @@ func formatFromFilename(filename string) string {
 	return ""
 }
 
-func (m command) generateSpec(opts *options) (spec.Interface, error) {
+type generatedSpecs struct {
+	spec.Interface
+	filenameInfix string
+}
+
+func (g *generatedSpecs) Save(filename string) error {
+	filename = g.updateFilename(filename)
+
+	if filename == "" {
+		_, err := g.WriteTo(os.Stdout)
+		if err != nil {
+			return fmt.Errorf("failed to write CDI spec to STDOUT: %v", err)
+		}
+		return nil
+	}
+
+	return g.Interface.Save(filename)
+}
+
+func (g generatedSpecs) updateFilename(filename string) string {
+	if g.filenameInfix == "" || filename == "" {
+		return filename
+	}
+	ext := filepath.Ext(filepath.Base(filename))
+	return strings.TrimSuffix(filename, ext) + g.filenameInfix + ext
+}
+
+func (m command) generateSpecs(opts *options) ([]generatedSpecs, error) {
 	var deviceNamers []nvcdi.DeviceNamer
 	for _, strategy := range opts.deviceNameStrategies {
 		deviceNamer, err := nvcdi.NewDeviceNamer(strategy)
@@ -313,6 +348,7 @@ func (m command) generateSpec(opts *options) (spec.Interface, error) {
 		nvcdi.WithLibrarySearchPaths(opts.librarySearchPaths),
 		nvcdi.WithCSVFiles(opts.csv.files),
 		nvcdi.WithCSVIgnorePatterns(opts.csv.ignorePatterns),
+		nvcdi.WithFeatureFlags(opts.featureFlags...),
 		// We set the following to allow for dependency injection:
 		nvcdi.WithNvmlLib(opts.nvmllib),
 	}
@@ -326,7 +362,7 @@ func (m command) generateSpec(opts *options) (spec.Interface, error) {
 		return nil, fmt.Errorf("failed to create CDI library: %v", err)
 	}
 
-	deviceSpecs, err := cdilib.GetDeviceSpecsByID("all")
+	allDeviceSpecs, err := cdilib.GetDeviceSpecsByID("all")
 	if err != nil {
 		return nil, fmt.Errorf("failed to create device CDI specs: %v", err)
 	}
@@ -336,16 +372,79 @@ func (m command) generateSpec(opts *options) (spec.Interface, error) {
 		return nil, fmt.Errorf("failed to create edits common for entities: %v", err)
 	}
 
-	return spec.New(
+	commonSpecOptions := []spec.Option{
 		spec.WithVendor(opts.vendor),
-		spec.WithClass(opts.class),
-		spec.WithDeviceSpecs(deviceSpecs),
 		spec.WithEdits(*commonEdits.ContainerEdits),
 		spec.WithFormat(opts.format),
 		spec.WithMergedDeviceOptions(
 			transform.WithName(allDeviceName),
 			transform.WithSkipIfExists(true),
 		),
 		spec.WithPermissions(0644),
+	}
+
+	fullSpec, err := spec.New(
+		append(commonSpecOptions,
+			spec.WithClass(opts.class),
+			spec.WithDeviceSpecs(allDeviceSpecs),
+		)...,
 	)
+	if err != nil {
+		return nil, err
+	}
+	var allSpecs []generatedSpecs
+
+	allSpecs = append(allSpecs, generatedSpecs{Interface: fullSpec, filenameInfix: ""})
+
+	deviceSpecsByDeviceCoherence := (deviceSpecs)(allDeviceSpecs).splitOnAnnotation("gpu.nvidia.com/coherent")
+
+	if coherentDeviceSpecs := deviceSpecsByDeviceCoherence["gpu.nvidia.com/coherent=true"]; len(coherentDeviceSpecs) > 0 {
+		infix := ".coherent"
+		coherentSpecs, err := spec.New(
+			append(commonSpecOptions,
+				spec.WithClass(opts.class+infix),
+				spec.WithDeviceSpecs(coherentDeviceSpecs),
+			)...,
+		)
+		if err != nil {
+			return nil, err
+		}
+		allSpecs = append(allSpecs, generatedSpecs{Interface: coherentSpecs, filenameInfix: infix})
+	}
+
+	if noncoherentDeviceSpecs := deviceSpecsByDeviceCoherence["gpu.nvidia.com/coherent=false"]; len(noncoherentDeviceSpecs) > 0 {
+		infix := ".noncoherent"
+		noncoherentSpecs, err := spec.New(
+			append(commonSpecOptions,
+				spec.WithClass(opts.class+infix),
+				spec.WithDeviceSpecs(noncoherentDeviceSpecs),
+			)...,
+		)
+
+		if err != nil {
+			return nil, err
+		}
+		allSpecs = append(allSpecs, generatedSpecs{Interface: noncoherentSpecs, filenameInfix: infix})
+	}
+
+	return allSpecs, nil
+}
+
+type deviceSpecs []specs.Device
+
+func (d deviceSpecs) splitOnAnnotation(key string) map[string][]specs.Device {
+	splitSpecs := make(map[string][]specs.Device)
+
+	for _, deviceSpec := range d {
+		if len(deviceSpec.Annotations) == 0 {
+			continue
+		}
+		value, ok := deviceSpec.Annotations[key]
+		if !ok {
+			continue
+		}
+		splitSpecs[key+"="+value] = append(splitSpecs[key+"="+value], deviceSpec)
+	}
+
+	return splitSpecs
 }

cmd/nvidia-ctk/cdi/generate/generate_test.go

@@ -28,6 +28,7 @@ import (
 	"github.com/stretchr/testify/require"
 
 	"github.com/NVIDIA/nvidia-container-toolkit/internal/test"
+	"github.com/NVIDIA/nvidia-container-toolkit/pkg/nvcdi"
 )
 
 func TestGenerateSpec(t *testing.T) {
@@ -391,12 +392,15 @@ containerEdits:
 			}
 			tc.options.nvmllib = server
 
-			spec, err := c.generateSpec(&tc.options)
+			tc.options.featureFlags = []string{string(nvcdi.FeatureDisableCoherentAnnotations)}
+			specs, err := c.generateSpecs(&tc.options)
 			require.ErrorIs(t, err, tc.expectedError)
 
 			var buf bytes.Buffer
-			_, err = spec.WriteTo(&buf)
-			require.NoError(t, err)
+			for _, spec := range specs {
+				_, err = spec.WriteTo(&buf)
+				require.NoError(t, err)
+			}
 
 			require.Equal(t, strings.ReplaceAll(tc.expectedSpec, "{{ .driverRoot }}", driverRoot), buf.String())
 		})

pkg/nvcdi/api.go

@@ -80,4 +80,7 @@ const (
 	// FeatureDisableNvsandboxUtils disables the use of nvsandboxutils when
 	// querying devices.
 	FeatureDisableNvsandboxUtils = FeatureFlag("disable-nvsandbox-utils")
+	// FeatureDisableCoherentAnnotations disables the addition of annotations
+	// coherent or non-coherent devices.
+	FeatureDisableCoherentAnnotations = FeatureFlag("disable-coherent-annotations")
 )

pkg/nvcdi/full-gpu-nvml.go

@@ -36,6 +36,8 @@ type fullGPUDeviceSpecGenerator struct {
 	*nvmllib
 	uuid  string
 	index int
+
+	featureFlags map[FeatureFlag]bool
 }
 
 var _ DeviceSpecGenerator = (*fullGPUDeviceSpecGenerator)(nil)
@@ -44,7 +46,7 @@ func (l *fullGPUDeviceSpecGenerator) GetUUID() (string, error) {
 	return l.uuid, nil
 }
 
-func (l *nvmllib) newFullGPUDeviceSpecGeneratorFromDevice(index int, d device.Device) (*fullGPUDeviceSpecGenerator, error) {
+func (l *nvmllib) newFullGPUDeviceSpecGeneratorFromDevice(index int, d device.Device, featureFlags map[FeatureFlag]bool) (*fullGPUDeviceSpecGenerator, error) {
 	uuid, ret := d.GetUUID()
 	if ret != nvml.SUCCESS {
 		return nil, fmt.Errorf("failed to get device UUID: %v", ret)
@@ -53,12 +55,14 @@ func (l *nvmllib) newFullGPUDeviceSpecGeneratorFromDevice(index int, d device.De
 		nvmllib: l,
 		uuid:    uuid,
 		index:   index,
+
+		featureFlags: featureFlags,
 	}
 
 	return e, nil
 }
 
-func (l *nvmllib) newFullGPUDeviceSpecGeneratorFromNVMLDevice(uuid string, nvmlDevice nvml.Device) (DeviceSpecGenerator, error) {
+func (l *nvmllib) newFullGPUDeviceSpecGeneratorFromNVMLDevice(uuid string, nvmlDevice nvml.Device, featureFlags map[FeatureFlag]bool) (DeviceSpecGenerator, error) {
 	index, ret := nvmlDevice.GetIndex()
 	if ret != nvml.SUCCESS {
 		return nil, fmt.Errorf("failed to get device index: %v", ret)
@@ -68,6 +72,8 @@ func (l *nvmllib) newFullGPUDeviceSpecGeneratorFromNVMLDevice(uuid string, nvmlD
 		nvmllib: l,
 		uuid:    uuid,
 		index:   index,
+
+		featureFlags: featureFlags,
 	}
 	return e, nil
 }
@@ -83,11 +89,17 @@ func (l *fullGPUDeviceSpecGenerator) GetDeviceSpecs() ([]specs.Device, error) {
 		return nil, fmt.Errorf("failed to get device names: %w", err)
 	}
 
+	annotations, err := l.getDeviceAnnotations()
+	if err != nil {
+		l.logger.Warning("Ignoring error getting device annotations for device(s) %v: %v", names, err)
+		annotations = nil
+	}
 	var deviceSpecs []specs.Device
 	for _, name := range names {
 		deviceSpec := specs.Device{
 			Name:           name,
 			ContainerEdits: *deviceEdits.ContainerEdits,
+			Annotations:    annotations,
 		}
 		deviceSpecs = append(deviceSpecs, deviceSpec)
 	}
@@ -99,6 +111,29 @@ func (l *fullGPUDeviceSpecGenerator) device() (device.Device, error) {
 	return l.devicelib.NewDeviceByUUID(l.uuid)
 }
 
+func (l *fullGPUDeviceSpecGenerator) getDeviceAnnotations() (map[string]string, error) {
+	if l.featureFlags[FeatureDisableCoherentAnnotations] {
+		return nil, nil
+	}
+
+	device, err := l.device()
+	if err != nil {
+		return nil, err
+	}
+
+	// TODO: Should we distinguish between not-supported and disabled?
+	isCoherent, err := device.IsCoherent()
+	if err != nil {
+		return nil, fmt.Errorf("failed to check device coherence: %w", err)
+	}
+
+	annotations := map[string]string{
+		"gpu.nvidia.com/coherent": fmt.Sprintf("%v", isCoherent),
+	}
+
+	return annotations, nil
+}
+
 // GetGPUDeviceEdits returns the CDI edits for the full GPU represented by 'device'.
 func (l *fullGPUDeviceSpecGenerator) getDeviceEdits() (*cdi.ContainerEdits, error) {
 	device, err := l.device()

pkg/nvcdi/lib-nvml.go

@@ -102,7 +102,7 @@ func (l *nvmllib) newDeviceSpecGeneratorFromNVMLDevice(id string, nvmlDevice nvm
 		return l.newMIGDeviceSpecGeneratorFromNVMLDevice(id, nvmlDevice)
 	}
 
-	return l.newFullGPUDeviceSpecGeneratorFromNVMLDevice(id, nvmlDevice)
+	return l.newFullGPUDeviceSpecGeneratorFromNVMLDevice(id, nvmlDevice, l.featureFlags)
 }
 
 // getDeviceSpecGeneratorsForAllDevices returns the CDI device spec generators
@@ -118,7 +118,7 @@ func (l *nvmllib) getDeviceSpecGeneratorsForAllDevices() (DeviceSpecGenerator, e
 		if isMigEnabled {
 			return nil
 		}
-		fullGPU, err := l.newFullGPUDeviceSpecGeneratorFromDevice(i, d)
+		fullGPU, err := l.newFullGPUDeviceSpecGeneratorFromDevice(i, d, l.featureFlags)
 		if err != nil {
 			return err
 		}

pkg/nvcdi/mig-device-nvml.go

@@ -41,7 +41,7 @@ func (l *migDeviceSpecGenerator) GetUUID() (string, error) {
 }
 
 func (l *nvmllib) newMIGDeviceSpecGeneratorFromDevice(i int, d device.Device, j int, m device.MigDevice) (*migDeviceSpecGenerator, error) {
-	parent, err := l.newFullGPUDeviceSpecGeneratorFromDevice(i, d)
+	parent, err := l.newFullGPUDeviceSpecGeneratorFromDevice(i, d, map[FeatureFlag]bool{FeatureDisableCoherentAnnotations: true})
 	if err != nil {
 		return nil, err
 	}

pkg/nvcdi/options.go

@@ -165,13 +165,22 @@ func WithDisabledHook[T string | HookName](hook T) Option {
 	}
 }
 
-// WithFeatureFlag allows specified features to be toggled on.
-// This option can be specified multiple times for each feature flag.
-func WithFeatureFlag(featureFlag FeatureFlag) Option {
+// WithFeatureFlags allows the specified set of features to be toggled on.
+func WithFeatureFlags[T string | FeatureFlag](featureFlags ...T) Option {
 	return func(o *nvcdilib) {
 		if o.featureFlags == nil {
 			o.featureFlags = make(map[FeatureFlag]bool)
 		}
-		o.featureFlags[featureFlag] = true
+		for _, featureFlag := range featureFlags {
+			o.featureFlags[FeatureFlag(featureFlag)] = true
+		}
 	}
 }
+
+// WithFeatureFlag allows specified features to be toggled on.
+// This option can be specified multiple times for each feature flag.
+//
+// Deprecated: Use WithFeatureFlags
+func WithFeatureFlag[T string | FeatureFlag](featureFlag T) Option {
+	return WithFeatureFlags(featureFlag)
+}