Add device IDs to nvcdi.GetSpec API

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...)
-	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"),
-	)
+	return cdilib.GetSpec(identifiers...)
 }

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) {
-	deviceSpecs, err := l.GetAllDeviceSpecs()
+func (l *wrapper) GetSpec(devices ...string) (spec.Interface, error) {
+	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()