Add device IDs to nvcdi.GetSpec API

This change allows device IDs to the specified in the GetSpec API.
This simplifies cases where CDI specs are being generated for specific
devices by ID.

Signed-off-by: Evan Lezar <elezar@nvidia.com>

internal/modifier/cdi.go

@@ -195,26 +195,11 @@ func generateAutomaticCDISpec(logger logger.Interface, cfg *config.Config, devic
 		return nil, fmt.Errorf("failed to construct CDI library: %w", err)
 	}
 
-	identifiers := []string{}
+	var identifiers []string
 	for _, device := range devices {
 		_, _, id := parser.ParseDevice(device)
 		identifiers = append(identifiers, id)
 	}
 
-	deviceSpecs, err := cdilib.GetDeviceSpecsByID(identifiers...)
+	return cdilib.GetSpec(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"),
-	)
 }

pkg/nvcdi/api.go

@@ -27,7 +27,7 @@ import (
 
 // Interface defines the API for the nvcdi package
 type Interface interface {
-	GetSpec() (spec.Interface, error)
+	GetSpec(...string) (spec.Interface, error)
 	GetCommonEdits() (*cdi.ContainerEdits, error)
 	GetAllDeviceSpecs() ([]specs.Device, error)
 	GetGPUDeviceEdits(device.Device) (*cdi.ContainerEdits, error)

pkg/nvcdi/gds.go

@@ -58,7 +58,7 @@ func (l *gdslib) GetCommonEdits() (*cdi.ContainerEdits, error) {
 
 // GetSpec is unsppported for the gdslib specs.
 // gdslib is typically wrapped by a spec that implements GetSpec.
-func (l *gdslib) GetSpec() (spec.Interface, error) {
+func (l *gdslib) GetSpec(...string) (spec.Interface, error) {
 	return nil, fmt.Errorf("GetSpec is not supported")
 }
 

pkg/nvcdi/lib-csv.go

@@ -34,7 +34,7 @@ type csvlib nvcdilib
 var _ Interface = (*csvlib)(nil)
 
 // GetSpec should not be called for wsllib
-func (l *csvlib) GetSpec() (spec.Interface, error) {
+func (l *csvlib) GetSpec(...string) (spec.Interface, error) {
 	return nil, fmt.Errorf("unexpected call to csvlib.GetSpec()")
 }
 

pkg/nvcdi/lib-imex.go

@@ -41,7 +41,7 @@ const (
 )
 
 // GetSpec should not be called for imexlib.
-func (l *imexlib) GetSpec() (spec.Interface, error) {
+func (l *imexlib) GetSpec(...string) (spec.Interface, error) {
 	return nil, fmt.Errorf("unexpected call to imexlib.GetSpec()")
 }
 

pkg/nvcdi/lib-nvml.go

@@ -36,7 +36,7 @@ type nvmllib nvcdilib
 var _ Interface = (*nvmllib)(nil)
 
 // GetSpec should not be called for nvmllib
-func (l *nvmllib) GetSpec() (spec.Interface, error) {
+func (l *nvmllib) GetSpec(...string) (spec.Interface, error) {
 	return nil, fmt.Errorf("unexpected call to nvmllib.GetSpec()")
 }
 

pkg/nvcdi/lib-wsl.go

@@ -32,7 +32,7 @@ type wsllib nvcdilib
 var _ Interface = (*wsllib)(nil)
 
 // GetSpec should not be called for wsllib
-func (l *wsllib) GetSpec() (spec.Interface, error) {
+func (l *wsllib) GetSpec(...string) (spec.Interface, error) {
 	return nil, fmt.Errorf("unexpected call to wsllib.GetSpec()")
 }
 

pkg/nvcdi/management.go

@@ -180,7 +180,7 @@ func (m managementDiscoverer) nodeIsBlocked(path string) bool {
 
 // GetSpec is unsppported for the managementlib specs.
 // managementlib is typically wrapped by a spec that implements GetSpec.
-func (m *managementlib) GetSpec() (spec.Interface, error) {
+func (m *managementlib) GetSpec(...string) (spec.Interface, error) {
 	return nil, fmt.Errorf("GetSpec is not supported")
 }
 

pkg/nvcdi/mofed.go

@@ -58,7 +58,7 @@ func (l *mofedlib) GetCommonEdits() (*cdi.ContainerEdits, error) {
 
 // GetSpec is unsppported for the mofedlib specs.
 // mofedlib is typically wrapped by a spec that implements GetSpec.
-func (l *mofedlib) GetSpec() (spec.Interface, error) {
+func (l *mofedlib) GetSpec(...string) (spec.Interface, error) {
 	return nil, fmt.Errorf("GetSpec is not supported")
 }
 

pkg/nvcdi/wrapper.go

@@ -35,8 +35,11 @@ type wrapper struct {
 }
 
 // GetSpec combines the device specs and common edits from the wrapped Interface to a single spec.Interface.
-func (l *wrapper) GetSpec() (spec.Interface, error) {
+func (l *wrapper) GetSpec(devices ...string) (spec.Interface, error) {
-	deviceSpecs, err := l.GetAllDeviceSpecs()
+	if len(devices) == 0 {
+		devices = append(devices, "all")
+	}
+	deviceSpecs, err := l.GetDeviceSpecsByID(devices...)
 	if err != nil {
 		return nil, err
 	}
@@ -55,6 +58,16 @@ func (l *wrapper) GetSpec() (spec.Interface, error) {
 	)
 }
 
+func (l *wrapper) GetDeviceSpecsByID(devices ...string) ([]specs.Device, error) {
+	for _, device := range devices {
+		if device != "all" {
+			continue
+		}
+		return l.GetAllDeviceSpecs()
+	}
+	return l.Interface.GetDeviceSpecsByID(devices...)
+}
+
 // GetAllDeviceSpecs returns the device specs for all available devices.
 func (l *wrapper) GetAllDeviceSpecs() ([]specs.Device, error) {
 	return l.Interface.GetAllDeviceSpecs()