Refactor nvml CDI spec generation for consistency

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

Author: elezar

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

@@ -312,7 +312,7 @@ func (m command) generateSpec(opts *options) (spec.Interface, error) {
 		return nil, fmt.Errorf("failed to create CDI library: %v", err)
 	}
 
-	deviceSpecs, err := cdilib.GetAllDeviceSpecs()
+	deviceSpecs, err := cdilib.GetDeviceSpecsByID("all")
 	if err != nil {
 		return nil, fmt.Errorf("failed to create device CDI specs: %v", err)
 	}

pkg/nvcdi/api.go

@@ -39,6 +39,11 @@ type SpecGenerator interface {
 	GetSpec(...string) (spec.Interface, error)
 }
 
+// A DeviceSpecGenerator is used to generate the specs for one or more devices.
+type DeviceSpecGenerator interface {
+	GetDeviceSpecs() ([]specs.Device, error)
+}
+
 // A HookName represents one of the predefined NVIDIA CDI hooks.
 type HookName = discover.HookName
 

pkg/nvcdi/full-gpu-nvml.go

@@ -19,41 +19,74 @@ package nvcdi
 import (
 	"fmt"
 
-	"github.com/NVIDIA/go-nvlib/pkg/nvlib/device"
 	"tags.cncf.io/container-device-interface/pkg/cdi"
 	"tags.cncf.io/container-device-interface/specs-go"
 
+	"github.com/NVIDIA/go-nvlib/pkg/nvlib/device"
+	"github.com/NVIDIA/go-nvml/pkg/nvml"
+
 	"github.com/NVIDIA/nvidia-container-toolkit/internal/discover"
 	"github.com/NVIDIA/nvidia-container-toolkit/internal/edits"
 	"github.com/NVIDIA/nvidia-container-toolkit/internal/platform-support/dgpu"
 )
 
-// GetGPUDeviceSpecs returns the CDI device specs for the full GPU represented by 'device'.
-func (l *nvmllib) GetGPUDeviceSpecs(i int, d device.Device) ([]specs.Device, error) {
-	edits, err := l.GetGPUDeviceEdits(d)
+// A fullGPUDeviceSpecGenerator generates the CDI device specifications for a
+// single full GPU.
+type fullGPUDeviceSpecGenerator struct {
+	*nvmllib
+	id     string
+	index  int
+	device device.Device
+}
+
+var _ DeviceSpecGenerator = (*fullGPUDeviceSpecGenerator)(nil)
+
+func (l *nvmllib) newFullGPUDeviceSpecGeneratorFromNVMLDevice(id string, nvmlDevice nvml.Device) (DeviceSpecGenerator, error) {
+	device, err := l.devicelib.NewDevice(nvmlDevice)
 	if err != nil {
-		return nil, fmt.Errorf("failed to get edits for device: %v", err)
+		return nil, err
 	}
 
-	var deviceSpecs []specs.Device
-	names, err := l.deviceNamers.GetDeviceNames(i, convert{d})
+	index, ret := nvmlDevice.GetIndex()
+	if ret != nvml.SUCCESS {
+		return nil, fmt.Errorf("failed to get device index: %w", ret)
+	}
+
+	e := &fullGPUDeviceSpecGenerator{
+		nvmllib: l,
+		id:      id,
+		index:   index,
+		device:  device,
+	}
+	return e, nil
+}
+
+func (l *fullGPUDeviceSpecGenerator) GetDeviceSpecs() ([]specs.Device, error) {
+	deviceEdits, err := l.getDeviceEdits()
+	if err != nil {
+		return nil, fmt.Errorf("failed to get CDI device edits for identifier %q: %w", l.id, err)
+	}
+
+	names, err := l.getNames()
 	if err != nil {
-		return nil, fmt.Errorf("failed to get device name: %v", err)
+		return nil, fmt.Errorf("failed to get device names: %w", err)
 	}
+
+	var deviceSpecs []specs.Device
 	for _, name := range names {
-		spec := specs.Device{
+		deviceSpec := specs.Device{
 			Name:           name,
-			ContainerEdits: *edits.ContainerEdits,
+			ContainerEdits: *deviceEdits.ContainerEdits,
 		}
-		deviceSpecs = append(deviceSpecs, spec)
+		deviceSpecs = append(deviceSpecs, deviceSpec)
 	}
 
 	return deviceSpecs, nil
 }
 
 // GetGPUDeviceEdits returns the CDI edits for the full GPU represented by 'device'.
-func (l *nvmllib) GetGPUDeviceEdits(d device.Device) (*cdi.ContainerEdits, error) {
-	device, err := l.newFullGPUDiscoverer(d)
+func (l *fullGPUDeviceSpecGenerator) getDeviceEdits() (*cdi.ContainerEdits, error) {
+	device, err := l.newFullGPUDiscoverer(l.device)
 	if err != nil {
 		return nil, fmt.Errorf("failed to create device discoverer: %v", err)
 	}
@@ -66,8 +99,12 @@ func (l *nvmllib) GetGPUDeviceEdits(d device.Device) (*cdi.ContainerEdits, error
 	return editsForDevice, nil
 }
 
+func (l *fullGPUDeviceSpecGenerator) getNames() ([]string, error) {
+	return l.deviceNamers.GetDeviceNames(l.index, convert{l.device})
+}
+
 // newFullGPUDiscoverer creates a discoverer for the full GPU defined by the specified device.
-func (l *nvmllib) newFullGPUDiscoverer(d device.Device) (discover.Discover, error) {
+func (l *fullGPUDeviceSpecGenerator) newFullGPUDiscoverer(d device.Device) (discover.Discover, error) {
 	deviceNodes, err := dgpu.NewForDevice(d,
 		dgpu.WithDevRoot(l.devRoot),
 		dgpu.WithLogger(l.logger),

pkg/nvcdi/gds.go

@@ -28,10 +28,14 @@ import (
 
 type gdslib nvcdilib
 
-var _ wrapped = (*gdslib)(nil)
+var _ deviceSpecGeneratorFactory = (*gdslib)(nil)
 
-// GetDeviceSpecsByID returns the device specs for the specified devices.
-func (l *gdslib) GetDeviceSpecsByID(ids ...string) ([]specs.Device, error) {
+func (l *gdslib) DeviceSpecGenerators(...string) (DeviceSpecGenerator, error) {
+	return l, nil
+}
+
+// GetDeviceSpecs returns the CDI device specs for a single all device.
+func (l *gdslib) GetDeviceSpecs() ([]specs.Device, error) {
 	discoverer, err := discover.NewGDSDiscoverer(l.logger, l.driverRoot, l.devRoot)
 	if err != nil {
 		return nil, fmt.Errorf("failed to create GPUDirect Storage discoverer: %v", err)

pkg/nvcdi/lib-csv.go

@@ -29,10 +29,9 @@ import (
 
 type csvlib nvcdilib
 
-var _ wrapped = (*csvlib)(nil)
+var _ deviceSpecGeneratorFactory = (*csvlib)(nil)
 
-// GetDeviceSpecsByID returns the device specs for the specified devices.
-func (l *csvlib) GetDeviceSpecsByID(ids ...string) ([]specs.Device, error) {
+func (l *csvlib) DeviceSpecGenerators(ids ...string) (DeviceSpecGenerator, error) {
 	for _, id := range ids {
 		switch id {
 		case "all":
@@ -42,6 +41,11 @@ func (l *csvlib) GetDeviceSpecsByID(ids ...string) ([]specs.Device, error) {
 		}
 	}
 
+	return l, nil
+}
+
+// GetDeviceSpecs returns the CDI device specs for a single device.
+func (l *csvlib) GetDeviceSpecs() ([]specs.Device, error) {
 	d, err := tegra.New(
 		tegra.WithLogger(l.logger),
 		tegra.WithDriverRoot(l.driverRoot),

pkg/nvcdi/lib-imex.go

@@ -31,7 +31,12 @@ import (
 
 type imexlib nvcdilib
 
-var _ wrapped = (*imexlib)(nil)
+type imexChannel struct {
+	id      string
+	devRoot string
+}
+
+var _ deviceSpecGeneratorFactory = (*imexlib)(nil)
 
 const (
 	classImexChannel = "imex-channel"
@@ -42,29 +47,24 @@ func (l *imexlib) GetCommonEdits() (*cdi.ContainerEdits, error) {
 	return edits.FromDiscoverer(discover.None{})
 }
 
-// GetDeviceSpecsByID returns the CDI device specs for the IMEX channels specified.
-func (l *imexlib) GetDeviceSpecsByID(ids ...string) ([]specs.Device, error) {
+// DeviceSpecGenerators returns the CDI device spec generators for the specified
+// imex channel IDs.
+// Valid IDs are:
+// * numeric channel IDs
+// * channel<numericChannelID>
+// * the special ID 'all'
+func (l *imexlib) DeviceSpecGenerators(ids ...string) (DeviceSpecGenerator, error) {
 	channelsIDs, err := l.getChannelIDs(ids...)
 	if err != nil {
 		return nil, err
 	}
-	var deviceSpecs []specs.Device
+
+	var deviceSpecGenerators DeviceSpecGenerators
 	for _, id := range channelsIDs {
-		path := "/dev/nvidia-caps-imex-channels/channel" + id
-		deviceSpec := specs.Device{
-			Name: id,
-			ContainerEdits: specs.ContainerEdits{
-				DeviceNodes: []*specs.DeviceNode{
-					{
-						Path:     path,
-						HostPath: filepath.Join(l.devRoot, path),
-					},
-				},
-			},
-		}
-		deviceSpecs = append(deviceSpecs, deviceSpec)
+		deviceSpecGenerators = append(deviceSpecGenerators, &imexChannel{id: id, devRoot: l.devRoot})
 	}
-	return deviceSpecs, nil
+
+	return deviceSpecGenerators, nil
 }
 
 func (l *imexlib) getChannelIDs(ids ...string) ([]string, error) {
@@ -104,3 +104,20 @@ func (l *imexlib) getAllChannelIDs() ([]string, error) {
 
 	return channelIDs, nil
 }
+
+// GetDeviceSpecs returns the CDI device specs the specified IMEX channel.
+func (l *imexChannel) GetDeviceSpecs() ([]specs.Device, error) {
+	path := "/dev/nvidia-caps-imex-channels/channel" + l.id
+	deviceSpec := specs.Device{
+		Name: l.id,
+		ContainerEdits: specs.ContainerEdits{
+			DeviceNodes: []*specs.DeviceNode{
+				{
+					Path:     path,
+					HostPath: filepath.Join(l.devRoot, path),
+				},
+			},
+		},
+	}
+	return []specs.Device{deviceSpec}, nil
+}