From dc87dcf786ed7958265ae987a685722fcb763f4d Mon Sep 17 00:00:00 2001 From: Evan Lezar Date: Thu, 5 Jun 2025 15:48:09 +0200 Subject: [PATCH] Make CDI device requests consistent with other methods Following the refactoring of device request extraction, we can now make CDI device requests consistent with other methods. This change moves to using image.VisibleDevices instead of separate calls to CDIDevicesFromMounts and VisibleDevicesFromEnvVar. Signed-off-by: Evan Lezar --- internal/config/image/cuda_image.go | 75 ++++++++------- internal/config/image/cuda_image_test.go | 6 +- internal/info/auto_test.go | 20 +++- internal/modifier/cdi.go | 117 +++++++++++++++-------- internal/modifier/cdi_test.go | 68 ++++++++++++- 5 files changed, 209 insertions(+), 77 deletions(-) diff --git a/internal/config/image/cuda_image.go b/internal/config/image/cuda_image.go index 55990f61..4fbeb68e 100644 --- a/internal/config/image/cuda_image.go +++ b/internal/config/image/cuda_image.go @@ -56,8 +56,12 @@ type CUDA struct { // NewCUDAImageFromSpec creates a CUDA image from the input OCI runtime spec. // The process environment is read (if present) to construc the CUDA Image. func NewCUDAImageFromSpec(spec *specs.Spec, opts ...Option) (CUDA, error) { + if spec == nil { + return New(opts...) + } + var env []string - if spec != nil && spec.Process != nil { + if spec.Process != nil { env = spec.Process.Env } @@ -219,19 +223,12 @@ func parseMajorMinorVersion(version string) (string, error) { // OnlyFullyQualifiedCDIDevices returns true if all devices requested in the image are requested as CDI devices/ func (i CUDA) OnlyFullyQualifiedCDIDevices() bool { var hasCDIdevice bool - for _, device := range i.VisibleDevicesFromEnvVar() { + for _, device := range i.VisibleDevices() { if !parser.IsQualifiedName(device) { return false } hasCDIdevice = true } - - for _, device := range i.DevicesFromMounts() { - if !strings.HasPrefix(device, "cdi/") { - return false - } - hasCDIdevice = true - } return hasCDIdevice } @@ -309,20 +306,27 @@ func (i CUDA) VisibleDevicesFromEnvVar() []string { // visibleDevicesFromMounts returns the set of visible devices requested as mounts. func (i CUDA) visibleDevicesFromMounts() []string { var devices []string - for _, device := range i.DevicesFromMounts() { + for _, device := range i.requestsFromMounts() { switch { - case strings.HasPrefix(device, volumeMountDevicePrefixCDI): - continue case strings.HasPrefix(device, volumeMountDevicePrefixImex): continue + case strings.HasPrefix(device, volumeMountDevicePrefixCDI): + name, err := cdiDeviceMountRequest(device).qualifiedName() + if err != nil { + i.logger.Warningf("Ignoring invalid mount request for CDI device %v: %v", device, err) + continue + } + devices = append(devices, name) + default: + devices = append(devices, device) } - devices = append(devices, device) + } return devices } -// DevicesFromMounts returns a list of device specified as mounts. -func (i CUDA) DevicesFromMounts() []string { +// requestsFromMounts returns a list of device specified as mounts. +func (i CUDA) requestsFromMounts() []string { root := filepath.Clean(DeviceListAsVolumeMountsRoot) seen := make(map[string]bool) var devices []string @@ -354,23 +358,30 @@ func (i CUDA) DevicesFromMounts() []string { return devices } -// CDIDevicesFromMounts returns a list of CDI devices specified as mounts on the image. -func (i CUDA) CDIDevicesFromMounts() []string { - var devices []string - for _, mountDevice := range i.DevicesFromMounts() { - if !strings.HasPrefix(mountDevice, volumeMountDevicePrefixCDI) { - continue - } - parts := strings.SplitN(strings.TrimPrefix(mountDevice, volumeMountDevicePrefixCDI), "/", 3) - if len(parts) != 3 { - continue - } - vendor := parts[0] - class := parts[1] - device := parts[2] - devices = append(devices, fmt.Sprintf("%s/%s=%s", vendor, class, device)) +// a cdiDeviceMountRequest represents a CDI device requests as a mount. +// Here the host path /dev/null is mounted to a particular path in the container. +// The container path has the form: +// /var/run/nvidia-container-devices/cdi/// +// or +// /var/run/nvidia-container-devices/cdi//= +type cdiDeviceMountRequest string + +// qualifiedName returns the fully-qualified name of the CDI device. +func (m cdiDeviceMountRequest) qualifiedName() (string, error) { + if !strings.HasPrefix(string(m), volumeMountDevicePrefixCDI) { + return "", fmt.Errorf("invalid mount CDI device request: %s", m) } - return devices + + requestedDevice := strings.TrimPrefix(string(m), volumeMountDevicePrefixCDI) + if parser.IsQualifiedName(requestedDevice) { + return requestedDevice, nil + } + + parts := strings.SplitN(requestedDevice, "/", 3) + if len(parts) != 3 { + return "", fmt.Errorf("invalid mount CDI device request: %s", m) + } + return fmt.Sprintf("%s/%s=%s", parts[0], parts[1], parts[2]), nil } // ImexChannelsFromEnvVar returns the list of IMEX channels requested for the image. @@ -385,7 +396,7 @@ func (i CUDA) ImexChannelsFromEnvVar() []string { // ImexChannelsFromMounts returns the list of IMEX channels requested for the image. func (i CUDA) ImexChannelsFromMounts() []string { var channels []string - for _, mountDevice := range i.DevicesFromMounts() { + for _, mountDevice := range i.requestsFromMounts() { if !strings.HasPrefix(mountDevice, volumeMountDevicePrefixImex) { continue } diff --git a/internal/config/image/cuda_image_test.go b/internal/config/image/cuda_image_test.go index f2dbbee9..d60b6e0c 100644 --- a/internal/config/image/cuda_image_test.go +++ b/internal/config/image/cuda_image_test.go @@ -487,9 +487,9 @@ func TestGetVisibleDevicesFromMounts(t *testing.T) { expectedDevices: []string{"GPU0-MIG0/0/1", "GPU1-MIG0/0/1"}, }, { - description: "cdi devices are ignored", - mounts: makeTestMounts("GPU0", "cdi/nvidia.com/gpu=all", "GPU1"), - expectedDevices: []string{"GPU0", "GPU1"}, + description: "cdi devices are included", + mounts: makeTestMounts("GPU0", "nvidia.com/gpu=all", "GPU1"), + expectedDevices: []string{"GPU0", "nvidia.com/gpu=all", "GPU1"}, }, { description: "imex devices are ignored", diff --git a/internal/info/auto_test.go b/internal/info/auto_test.go index ad147526..c2ab93d7 100644 --- a/internal/info/auto_test.go +++ b/internal/info/auto_test.go @@ -184,7 +184,7 @@ func TestResolveAutoMode(t *testing.T) { expectedMode: "legacy", }, { - description: "cdi mount and non-CDI envvar resolves to legacy", + description: "cdi mount and non-CDI envvar resolves to cdi", mode: "auto", envmap: map[string]string{ "NVIDIA_VISIBLE_DEVICES": "0", @@ -197,6 +197,22 @@ func TestResolveAutoMode(t *testing.T) { "tegra": false, "nvgpu": false, }, + expectedMode: "cdi", + }, + { + description: "non-cdi mount and CDI envvar resolves to legacy", + mode: "auto", + envmap: map[string]string{ + "NVIDIA_VISIBLE_DEVICES": "nvidia.com/gpu=0", + }, + mounts: []string{ + "/var/run/nvidia-container-devices/0", + }, + info: map[string]bool{ + "nvml": true, + "tegra": false, + "nvgpu": false, + }, expectedMode: "legacy", }, } @@ -232,6 +248,8 @@ func TestResolveAutoMode(t *testing.T) { image, _ := image.New( image.WithEnvMap(tc.envmap), image.WithMounts(mounts), + image.WithAcceptDeviceListAsVolumeMounts(true), + image.WithAcceptEnvvarUnprivileged(true), ) mode := resolveMode(logger, tc.mode, image, properties) require.EqualValues(t, tc.expectedMode, mode) diff --git a/internal/modifier/cdi.go b/internal/modifier/cdi.go index 6139a3a8..c011b4e6 100644 --- a/internal/modifier/cdi.go +++ b/internal/modifier/cdi.go @@ -34,10 +34,12 @@ import ( // CDI specifications available on the system. The NVIDIA_VISIBLE_DEVICES environment variable is // used to select the devices to include. func NewCDIModifier(logger logger.Interface, cfg *config.Config, image image.CUDA) (oci.SpecModifier, error) { - devices, err := getDevicesFromImage(logger, cfg, image) - if err != nil { - return nil, fmt.Errorf("failed to get required devices from OCI specification: %v", err) - } + deviceRequestor := newCDIDeviceRequestor( + logger, + image, + cfg.NVIDIAContainerRuntimeConfig.Modes.CDI.DefaultKind, + ) + devices := deviceRequestor.DeviceRequests() if len(devices) == 0 { logger.Debugf("No devices requested; no modification required.") return nil, nil @@ -64,46 +66,47 @@ func NewCDIModifier(logger logger.Interface, cfg *config.Config, image image.CUD ) } -func getDevicesFromImage(logger logger.Interface, cfg *config.Config, container image.CUDA) ([]string, error) { - annotationDevices, err := getAnnotationDevices(container) +type deviceRequestor interface { + DeviceRequests() []string +} + +type cdiDeviceRequestor struct { + image image.CUDA + logger logger.Interface + defaultKind string +} + +func newCDIDeviceRequestor(logger logger.Interface, image image.CUDA, defaultKind string) deviceRequestor { + c := &cdiDeviceRequestor{ + logger: logger, + image: image, + defaultKind: defaultKind, + } + return withUniqueDevices(c) +} + +func (c *cdiDeviceRequestor) DeviceRequests() []string { + if c == nil { + return nil + } + annotationDevices, err := getAnnotationDevices(c.image) if err != nil { - return nil, fmt.Errorf("failed to parse container annotations: %v", err) + c.logger.Warningf("failed to get device requests from container annotations: %v; ignoring.", err) + annotationDevices = nil } if len(annotationDevices) > 0 { - return annotationDevices, nil - } - - if cfg.AcceptDeviceListAsVolumeMounts { - mountDevices := container.CDIDevicesFromMounts() - if len(mountDevices) > 0 { - return mountDevices, nil - } + return annotationDevices } var devices []string - seen := make(map[string]bool) - for _, name := range container.VisibleDevicesFromEnvVar() { + for _, name := range c.image.VisibleDevices() { if !parser.IsQualifiedName(name) { - name = fmt.Sprintf("%s=%s", cfg.NVIDIAContainerRuntimeConfig.Modes.CDI.DefaultKind, name) - } - if seen[name] { - logger.Debugf("Ignoring duplicate device %q", name) - continue + name = fmt.Sprintf("%s=%s", c.defaultKind, name) } devices = append(devices, name) } - if len(devices) == 0 { - return nil, nil - } - - if cfg.AcceptEnvvarUnprivileged || container.IsPrivileged() { - return devices, nil - } - - logger.Warningf("Ignoring devices specified in NVIDIA_VISIBLE_DEVICES: %v", devices) - - return nil, nil + return devices } // getAnnotationDevices returns a list of devices specified in the annotations. @@ -111,16 +114,11 @@ func getDevicesFromImage(logger logger.Interface, cfg *config.Config, container // fully-qualified CDI devices names. If any device name is not fully-quality an error is returned. // The list of returned devices is deduplicated. func getAnnotationDevices(image image.CUDA) ([]string, error) { - seen := make(map[string]bool) var annotationDevices []string for _, device := range image.CDIDeviceRequestsFromAnnotations() { if !parser.IsQualifiedName(device) { return nil, fmt.Errorf("invalid device name %q in annotations", device) } - if seen[device] { - continue - } - seen[device] = true annotationDevices = append(annotationDevices, device) } return annotationDevices, nil @@ -147,7 +145,7 @@ func newAutomaticCDISpecModifier(logger logger.Interface, cfg *config.Config, de if err != nil { return nil, fmt.Errorf("failed to generate CDI spec: %w", err) } - cdiModifier, err := cdi.New( + cdiDeviceRequestor, err := cdi.New( cdi.WithLogger(logger), cdi.WithSpec(spec.Raw()), ) @@ -155,7 +153,7 @@ func newAutomaticCDISpecModifier(logger logger.Interface, cfg *config.Config, de return nil, fmt.Errorf("failed to construct CDI modifier: %w", err) } - return cdiModifier, nil + return cdiDeviceRequestor, nil } func generateAutomaticCDISpec(logger logger.Interface, cfg *config.Config, devices []string) (spec.Interface, error) { @@ -193,3 +191,42 @@ func generateAutomaticCDISpec(logger logger.Interface, cfg *config.Config, devic spec.WithClass("gpu"), ) } + +func deviceRequestorFromImage(image image.CUDA) deviceRequestor { + return &fromImage{image} +} + +type fromImage struct { + image.CUDA +} + +func (f *fromImage) DeviceRequests() []string { + if f == nil { + return nil + } + return f.CUDA.VisibleDevices() +} + +type deduplicatedDeviceRequestor struct { + deviceRequestor +} + +func withUniqueDevices(deviceRequestor deviceRequestor) deviceRequestor { + return &deduplicatedDeviceRequestor{deviceRequestor: deviceRequestor} +} + +func (d *deduplicatedDeviceRequestor) DeviceRequests() []string { + if d == nil { + return nil + } + seen := make(map[string]bool) + var devices []string + for _, device := range d.deviceRequestor.DeviceRequests() { + if seen[device] { + continue + } + seen[device] = true + devices = append(devices, device) + } + return devices +} diff --git a/internal/modifier/cdi_test.go b/internal/modifier/cdi_test.go index 401a1e83..cf7811a5 100644 --- a/internal/modifier/cdi_test.go +++ b/internal/modifier/cdi_test.go @@ -20,6 +20,8 @@ import ( "fmt" "testing" + "github.com/opencontainers/runtime-spec/specs-go" + testlog "github.com/sirupsen/logrus/hooks/test" "github.com/stretchr/testify/require" "github.com/NVIDIA/nvidia-container-toolkit/internal/config/image" @@ -67,7 +69,7 @@ func TestGetAnnotationDevices(t *testing.T) { "prefix/foo": "example.com/device=bar", "another-prefix/bar": "example.com/device=bar", }, - expectedDevices: []string{"example.com/device=bar"}, + expectedDevices: []string{"example.com/device=bar", "example.com/device=bar"}, }, { description: "invalid devices", @@ -98,3 +100,67 @@ func TestGetAnnotationDevices(t *testing.T) { }) } } + +func TestDeviceRequests(t *testing.T) { + logger, _ := testlog.NewNullLogger() + + testCases := []struct { + description string + input cdiDeviceRequestor + spec *specs.Spec + expectedDevices []string + }{ + { + description: "empty spec yields no devices", + }, + { + description: "cdi devices from mounts", + input: cdiDeviceRequestor{ + defaultKind: "nvidia.com/gpu", + }, + spec: &specs.Spec{ + Mounts: []specs.Mount{ + { + Destination: "/var/run/nvidia-container-devices/cdi/nvidia.com/gpu/0", + Source: "/dev/null", + }, + { + Destination: "/var/run/nvidia-container-devices/cdi/nvidia.com/gpu/1", + Source: "/dev/null", + }, + }, + }, + expectedDevices: []string{"nvidia.com/gpu=0", "nvidia.com/gpu=1"}, + }, + { + description: "cdi devices from envvar", + input: cdiDeviceRequestor{ + defaultKind: "nvidia.com/gpu", + }, + spec: &specs.Spec{ + Process: &specs.Process{ + Env: []string{"NVIDIA_VISIBLE_DEVICES=0,example.com/class=device"}, + }, + }, + expectedDevices: []string{"nvidia.com/gpu=0", "example.com/class=device"}, + }, + } + + for _, tc := range testCases { + tc.input.logger = logger + + image, err := image.NewCUDAImageFromSpec( + tc.spec, + image.WithAcceptDeviceListAsVolumeMounts(true), + image.WithAcceptEnvvarUnprivileged(true), + ) + require.NoError(t, err) + tc.input.image = image + + t.Run(tc.description, func(t *testing.T) { + devices := tc.input.DeviceRequests() + require.NoError(t, err) + require.EqualValues(t, tc.expectedDevices, devices) + }) + } +}