UniqAI/nvidia-container-toolkit
1
0
Fork 0
mirror of https://github.com/NVIDIA/nvidia-container-toolkit synced 2024-11-21 15:57:49 +00:00
Code Issues Actions Packages Projects Releases Wiki Activity

Extend the 'runtime.nvidia.com/gpu' CDI device kind to support MIG devices specified by index or UUID

Browse Source 
Signed-off-by: Christopher Desiniotis <cdesiniotis@nvidia.com>
This commit is contained in:
Christopher Desiniotis 2023-12-05 18:38:00 -08:00
parent def7d09f85
commit 3158146946
1 changed files with 55 additions and 7 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

62
pkg/nvcdi/lib-nvml.go
View File

@ -19,6 +19,7 @@ package nvcdi
import (
"fmt"
"strconv"
"strings"
"github.com/NVIDIA/go-nvlib/pkg/nvlib/device"
"github.com/NVIDIA/go-nvlib/pkg/nvml"
@ -79,7 +80,6 @@ func (l *nvmllib) GetCommonEdits() (*cdi.ContainerEdits, error) {
// 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.
// TODO: support identifiers that correspond to MIG devices
func (l *nvmllib) GetDeviceSpecsByID(identifiers ...string) ([]specs.Device, error) {
for _, id := range identifiers {
if id == "all" {
@ -104,11 +104,7 @@ func (l *nvmllib) GetDeviceSpecsByID(identifiers ...string) ([]specs.Device, err
}
for i, nvmlDevice := range nvmlDevices {
nvlibDevice, err := l.devicelib.NewDevice(nvmlDevice)
if err != nil {
return nil, fmt.Errorf("failed to construct device: %w", err)
}
deviceEdits, err := l.GetGPUDeviceEdits(nvlibDevice)
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)
}
@ -151,12 +147,64 @@ func (l *nvmllib) getNVMLDeviceByID(id string) (nvml.Device, error) {
}
if devID.isMigIndex() {
return nil, fmt.Errorf("MIG index is not supported")
var gpuIdx, migIdx int
var parent nvml.Device
split := strings.SplitN(id, ":", 2)
if gpuIdx, err = strconv.Atoi(split[0]); err != nil {
return nil, fmt.Errorf("failed to convert device index to an int: %w", err)
}
if migIdx, err = strconv.Atoi(split[1]); err != nil {
return nil, fmt.Errorf("failed to convert device index to an int: %w", err)
}
if parent, err = l.nvmllib.DeviceGetHandleByIndex(gpuIdx); err != nvml.SUCCESS {
return nil, fmt.Errorf("failed to get parent device handle: %w", err)
}
return parent.GetMigDeviceHandleByIndex(migIdx)
}
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)
}
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 {