UniqAI/nvidia-container-toolkit
1
0
Fork 0
mirror of https://github.com/NVIDIA/nvidia-container-toolkit synced 2025-06-26 18:18:24 +00:00
Code Issues Actions Packages Projects Releases Wiki Activity

Construct nvml-based CDI lib based on mode

Browse Source 
Signed-off-by: Evan Lezar <elezar@nvidia.com>
This commit is contained in:
Evan Lezar
2023-02-16 17:29:53 +02:00
parent 20d6e9af04
commit d226925fe7
6 changed files with 111 additions and 77 deletions
Download Patch File Download Diff File Expand all files Collapse all files

71
pkg/nvcdi/lib.go
View File

@@ -17,9 +17,6 @@
package nvcdi
import (
"fmt"
"github.com/container-orchestrated-devices/container-device-interface/specs-go"
"github.com/sirupsen/logrus"
"gitlab.com/nvidia/cloud-native/go-nvlib/pkg/nvlib/device"
"gitlab.com/nvidia/cloud-native/go-nvlib/pkg/nvml"
@@ -41,12 +38,8 @@ func New(opts ...Option) Interface {
for _, opt := range opts {
opt(l)
}
if l.nvmllib == nil {
l.nvmllib = nvml.New()
}
if l.devicelib == nil {
l.devicelib = device.New(device.WithNvml(l.nvmllib))
if l.mode == "" {
l.mode = "nvml"
}
if l.logger == nil {
l.logger = logrus.StandardLogger()
@@ -61,58 +54,18 @@ func New(opts ...Option) Interface {
l.nvidiaCTKPath = "/usr/bin/nvidia-ctk"
}
return l
}
// GetAllDeviceSpecs returns the device specs for all available devices.
func (l *nvcdilib) GetAllDeviceSpecs() ([]specs.Device, error) {
var deviceSpecs []specs.Device
gpuDeviceSpecs, err := l.getGPUDeviceSpecs()
if err != nil {
return nil, err
}
deviceSpecs = append(deviceSpecs, gpuDeviceSpecs...)
migDeviceSpecs, err := l.getMigDeviceSpecs()
if err != nil {
return nil, err
}
deviceSpecs = append(deviceSpecs, migDeviceSpecs...)
return deviceSpecs, nil
}
func (l *nvcdilib) getGPUDeviceSpecs() ([]specs.Device, error) {
var deviceSpecs []specs.Device
err := l.devicelib.VisitDevices(func(i int, d device.Device) error {
deviceSpec, err := l.GetGPUDeviceSpecs(i, d)
if err != nil {
return err
switch l.mode {
case "nvml":
if l.nvmllib == nil {
l.nvmllib = nvml.New()
}
deviceSpecs = append(deviceSpecs, *deviceSpec)
return nil
})
if err != nil {
return nil, fmt.Errorf("failed to generate CDI edits for GPU devices: %v", err)
}
return deviceSpecs, err
}
func (l *nvcdilib) getMigDeviceSpecs() ([]specs.Device, error) {
var deviceSpecs []specs.Device
err := l.devicelib.VisitMigDevices(func(i int, d device.Device, j int, mig device.MigDevice) error {
deviceSpec, err := l.GetMIGDeviceSpecs(i, d, j, mig)
if err != nil {
return err
if l.devicelib == nil {
l.devicelib = device.New(device.WithNvml(l.nvmllib))
}
deviceSpecs = append(deviceSpecs, *deviceSpec)
return nil
})
if err != nil {
return nil, fmt.Errorf("failed to generate CDI edits for GPU devices: %v", err)
return (*nvmllib)(l)
}
return deviceSpecs, err
// TODO: We want an error here.
return nil
}