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

@@ -33,6 +33,11 @@ type Interface interface {
 	GetAllDeviceSpecs() ([]specs.Device, error)
 }
 
+// A deviceSpecGenerator is used to generate the specs for a set of 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,66 @@ 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)
+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})
+	e := &fullGPUDeviceSpecGenerator{
+		nvmllib: l,
+		id:      id,
+		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 +91,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/lib-nvml.go

@@ -28,24 +28,29 @@ import (
 
 	"github.com/NVIDIA/nvidia-container-toolkit/internal/edits"
 	"github.com/NVIDIA/nvidia-container-toolkit/internal/nvsandboxutils"
-	"github.com/NVIDIA/nvidia-container-toolkit/pkg/nvcdi/spec"
 )
 
 type nvmllib nvcdilib
 
-var _ Interface = (*nvmllib)(nil)
+var _ wrapped = (*nvmllib)(nil)
 
-// GetSpec should not be called for nvmllib
-func (l *nvmllib) GetSpec(...string) (spec.Interface, error) {
-	return nil, fmt.Errorf("unexpected call to nvmllib.GetSpec()")
-}
+// GetCommonEdits generates a CDI specification that can be used for ANY devices
+func (l *nvmllib) GetCommonEdits() (*cdi.ContainerEdits, error) {
+	common, err := l.newCommonNVMLDiscoverer()
+	if err != nil {
+		return nil, fmt.Errorf("failed to create discoverer for common entities: %v", err)
+	}
 
-// GetAllDeviceSpecs returns the device specs for all available devices.
-func (l *nvmllib) GetAllDeviceSpecs() ([]specs.Device, error) {
-	var deviceSpecs []specs.Device
+	return edits.FromDiscoverer(common)
+}
 
+// GetDeviceSpecsByID returns the CDI device specs for the devices represented
+// by the requested identifiers. Here an identifier is one of the following:
+// * an index of a GPU or MIG device
+// * a UUID of a GPU or MIG device
+func (l *nvmllib) GetDeviceSpecsByID(ids ...string) ([]specs.Device, error) {
 	if r := l.nvmllib.Init(); r != nvml.SUCCESS {
-		return nil, fmt.Errorf("failed to initialize NVML: %v", r)
+		return nil, fmt.Errorf("failed to initialize NVML: %w", r)
 	}
 	defer func() {
 		if r := l.nvmllib.Shutdown(); r != nvml.SUCCESS {
@@ -66,93 +71,83 @@ func (l *nvmllib) GetAllDeviceSpecs() ([]specs.Device, error) {
 		}()
 	}
 
-	gpuDeviceSpecs, err := l.getGPUDeviceSpecs()
-	if err != nil {
-		return nil, err
-	}
-	deviceSpecs = append(deviceSpecs, gpuDeviceSpecs...)
-
-	migDeviceSpecs, err := l.getMigDeviceSpecs()
+	generators, err := l.getDeviceSpecGeneratorsForIDs(ids...)
 	if err != nil {
 		return nil, err
 	}
-	deviceSpecs = append(deviceSpecs, migDeviceSpecs...)
 
-	return deviceSpecs, nil
+	return generators.GetDeviceSpecs()
 }
 
-// GetCommonEdits generates a CDI specification that can be used for ANY devices
-func (l *nvmllib) GetCommonEdits() (*cdi.ContainerEdits, error) {
-	common, err := l.newCommonNVMLDiscoverer()
-	if err != nil {
-		return nil, fmt.Errorf("failed to create discoverer for common entities: %v", err)
+func (l *nvmllib) newDeviceSpecGeneratorFromNVMLDevice(id string, nvmlDevice nvml.Device) (deviceSpecGenerator, error) {
+	isMig, ret := nvmlDevice.IsMigDeviceHandle()
+	if ret != nvml.SUCCESS {
+		return nil, ret
+	}
+	if isMig {
+		return l.newMIGDeviceSpecGeneratorFromNVMLDevice(id, nvmlDevice)
 	}
 
-	return edits.FromDiscoverer(common)
+	return l.newFullGPUDeviceSpecGeneratorFromNVMLDevice(id, nvmlDevice)
 }
 
-// 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.
-// Deprecated: Use GetDeviceSpecsBy instead.
-func (l *nvmllib) GetDeviceSpecsByID(ids ...string) ([]specs.Device, error) {
+func (l *nvmllib) getDeviceSpecGeneratorsForIDs(ids ...string) (deviceSpecGenerators, error) {
 	var identifiers []device.Identifier
 	for _, id := range ids {
-		identifiers = append(identifiers, device.Identifier(id))
-	}
-	return l.GetDeviceSpecsBy(identifiers...)
-}
-
-// GetDeviceSpecsBy returns the device specs for devices with the specified identifiers.
-func (l *nvmllib) GetDeviceSpecsBy(identifiers ...device.Identifier) ([]specs.Device, error) {
-	for _, id := range identifiers {
 		if id == "all" {
-			return l.GetAllDeviceSpecs()
+			return l.getDeviceSpecGeneratorsForAllDevices()
 		}
+		identifiers = append(identifiers, device.Identifier(id))
 	}
 
-	var deviceSpecs []specs.Device
-
-	if r := l.nvmllib.Init(); r != nvml.SUCCESS {
-		return nil, fmt.Errorf("failed to initialize NVML: %w", r)
+	devices, err := l.getNVMLDevicesByID(identifiers...)
+	if err != nil {
+		return nil, err
 	}
-	defer func() {
-		if r := l.nvmllib.Shutdown(); r != nvml.SUCCESS {
-			l.logger.Warningf("failed to shutdown NVML: %v", r)
-		}
-	}()
 
-	if l.nvsandboxutilslib != nil {
-		if r := l.nvsandboxutilslib.Init(l.driverRoot); r != nvsandboxutils.SUCCESS {
-			l.logger.Warningf("Failed to init nvsandboxutils: %v; ignoring", r)
-			l.nvsandboxutilslib = nil
+	var DeviceSpecGenerators deviceSpecGenerators
+	for i, device := range devices {
+		editor, err := l.newDeviceSpecGeneratorFromNVMLDevice(ids[i], device)
+		if err != nil {
+			return nil, err
 		}
-		defer func() {
-			if l.nvsandboxutilslib == nil {
-				return
-			}
-			_ = l.nvsandboxutilslib.Shutdown()
-		}()
+		DeviceSpecGenerators = append(DeviceSpecGenerators, editor)
 	}
 
-	nvmlDevices, err := l.getNVMLDevicesByID(identifiers...)
+	return DeviceSpecGenerators, nil
+}
+
+func (l *nvmllib) getDeviceSpecGeneratorsForAllDevices() ([]deviceSpecGenerator, error) {
+	var DeviceSpecGenerators []deviceSpecGenerator
+	err := l.devicelib.VisitDevices(func(i int, d device.Device) error {
+		e := &fullGPUDeviceSpecGenerator{
+			nvmllib: l,
+			id:      fmt.Sprintf("%d", i),
+			device:  d,
+		}
+
+		DeviceSpecGenerators = append(DeviceSpecGenerators, e)
+		return nil
+	})
 	if err != nil {
-		return nil, fmt.Errorf("failed to get NVML device handles: %w", err)
+		return nil, fmt.Errorf("failed to get full GPU device editors: %w", err)
 	}
 
-	for i, nvmlDevice := range nvmlDevices {
-		deviceEdits, err := l.getEditsForDevice(nvmlDevice)
-		if err != nil {
-			return nil, fmt.Errorf("failed to get CDI device edits for identifier %q: %w", identifiers[i], err)
-		}
-		deviceSpec := specs.Device{
-			Name:           string(identifiers[i]),
-			ContainerEdits: *deviceEdits.ContainerEdits,
+	err = l.devicelib.VisitMigDevices(func(i int, d device.Device, j int, mig device.MigDevice) error {
+		e := &migDeviceSpecGenerator{
+			nvmllib: l,
+			id:      fmt.Sprintf("%d:%d", i, j),
+			parent:  d,
+			device:  mig,
 		}
-		deviceSpecs = append(deviceSpecs, deviceSpec)
+		DeviceSpecGenerators = append(DeviceSpecGenerators, e)
+		return nil
+	})
+	if err != nil {
+		return nil, fmt.Errorf("failed to get MIG device editors: %w", err)
 	}
 
-	return deviceSpecs, nil
+	return DeviceSpecGenerators, nil
 }
 
 // TODO: move this to go-nvlib?
@@ -201,76 +196,21 @@ func (l *nvmllib) getNVMLDeviceByID(id device.Identifier) (nvml.Device, error) {
 	return nil, fmt.Errorf("identifier is not a valid UUID or index: %q", id)
 }
 
-func (l *nvmllib) getEditsForDevice(nvmlDevice nvml.Device) (*cdi.ContainerEdits, error) {
-	mig, err := nvmlDevice.IsMigDeviceHandle()
-	if err != nvml.SUCCESS {
-		return nil, fmt.Errorf("failed to determine if device handle is a MIG device: %w", err)
-	}
-	if mig {
-		return l.getEditsForMIGDevice(nvmlDevice)
-	}
-	return l.getEditsForGPUDevice(nvmlDevice)
-}
-
-func (l *nvmllib) getEditsForGPUDevice(nvmlDevice nvml.Device) (*cdi.ContainerEdits, error) {
-	nvlibDevice, err := l.devicelib.NewDevice(nvmlDevice)
-	if err != nil {
-		return nil, fmt.Errorf("failed to construct device: %w", err)
-	}
-	deviceEdits, err := l.GetGPUDeviceEdits(nvlibDevice)
-	if err != nil {
-		return nil, fmt.Errorf("failed to get GPU device edits: %w", err)
-	}
+type deviceSpecGenerators []deviceSpecGenerator
 
-	return deviceEdits, nil
-}
-
-func (l *nvmllib) getEditsForMIGDevice(nvmlDevice nvml.Device) (*cdi.ContainerEdits, error) {
-	nvmlParentDevice, ret := nvmlDevice.GetDeviceHandleFromMigDeviceHandle()
-	if ret != nvml.SUCCESS {
-		return nil, fmt.Errorf("failed to get parent device handle: %w", ret)
-	}
-	nvlibMigDevice, err := l.devicelib.NewMigDevice(nvmlDevice)
-	if err != nil {
-		return nil, fmt.Errorf("failed to construct device: %w", err)
-	}
-	nvlibParentDevice, err := l.devicelib.NewDevice(nvmlParentDevice)
-	if err != nil {
-		return nil, fmt.Errorf("failed to construct parent device: %w", err)
-	}
-	return l.GetMIGDeviceEdits(nvlibParentDevice, nvlibMigDevice)
-}
-
-func (l *nvmllib) getGPUDeviceSpecs() ([]specs.Device, error) {
-	var deviceSpecs []specs.Device
-	err := l.devicelib.VisitDevices(func(i int, d device.Device) error {
-		specsForDevice, err := l.GetGPUDeviceSpecs(i, d)
-		if err != nil {
-			return err
+// GetDeviceSpecs returns the combined specs for each device spec generator.
+func (g deviceSpecGenerators) GetDeviceSpecs() ([]specs.Device, error) {
+	var allDeviceSpecs []specs.Device
+	for _, dsg := range g {
+		if dsg == nil {
+			continue
 		}
-		deviceSpecs = append(deviceSpecs, specsForDevice...)
-
-		return nil
-	})
-	if err != nil {
-		return nil, fmt.Errorf("failed to generate CDI edits for GPU devices: %v", err)
-	}
-	return deviceSpecs, err
-}
-
-func (l *nvmllib) getMigDeviceSpecs() ([]specs.Device, error) {
-	var deviceSpecs []specs.Device
-	err := l.devicelib.VisitMigDevices(func(i int, d device.Device, j int, mig device.MigDevice) error {
-		specsForDevice, err := l.GetMIGDeviceSpecs(i, d, j, mig)
+		deviceSpecs, err := dsg.GetDeviceSpecs()
 		if err != nil {
-			return err
+			return nil, err
 		}
-		deviceSpecs = append(deviceSpecs, specsForDevice...)
-
-		return nil
-	})
-	if err != nil {
-		return nil, fmt.Errorf("failed to generate CDI edits for GPU devices: %v", err)
+		allDeviceSpecs = append(allDeviceSpecs, deviceSpecs...)
 	}
-	return deviceSpecs, err
+
+	return allDeviceSpecs, nil
 }