From 727c90c7a6a73f9734c27c54e8599095d136e32d Mon Sep 17 00:00:00 2001 From: Christopher Desiniotis Date: Mon, 29 Jun 2026 17:13:24 -0700 Subject: [PATCH] Process envvars for mig-config / mig-monitor devices in jit-cdi mode Signed-off-by: Christopher Desiniotis --- internal/discover/mig.go | 60 +++++++++ internal/discover/mig_test.go | 125 ++++++++++++++++++ internal/modifier/cdi.go | 22 +++ internal/modifier/cdi_test.go | 93 +++++++++++++ pkg/nvcdi/gated.go | 4 + pkg/nvcdi/lib.go | 2 +- pkg/nvcdi/mode.go | 6 + .../rootfs-1/dev/nvidia-caps/nvidia-cap2 | 0 8 files changed, 311 insertions(+), 1 deletion(-) create mode 100644 internal/discover/mig.go create mode 100644 internal/discover/mig_test.go create mode 100644 testdata/lookup/rootfs-1/dev/nvidia-caps/nvidia-cap2 diff --git a/internal/discover/mig.go b/internal/discover/mig.go new file mode 100644 index 000000000..6fd7b1656 --- /dev/null +++ b/internal/discover/mig.go @@ -0,0 +1,60 @@ +/** +# Copyright (c) NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +**/ + +package discover + +import ( + "fmt" + + "github.com/NVIDIA/nvidia-container-toolkit/internal/logger" + "github.com/NVIDIA/nvidia-container-toolkit/internal/lookup/root" + "github.com/NVIDIA/nvidia-container-toolkit/internal/nvcaps" +) + +// NewMigConfigDiscoverer creates a discoverer for the MIG config capability device node. +func NewMigConfigDiscoverer(logger logger.Interface, driver *root.Driver) (Discover, error) { + migCaps, err := nvcaps.NewMigCaps() + if err != nil { + return nil, fmt.Errorf("error getting MIG capability device paths: %w", err) + } + return newMigCapDiscoverer(logger, driver, migCaps, nvcaps.MigCap("config")) +} + +// NewMigMonitorDiscoverer creates a discoverer for the MIG monitor capability device node. +func NewMigMonitorDiscoverer(logger logger.Interface, driver *root.Driver) (Discover, error) { + migCaps, err := nvcaps.NewMigCaps() + if err != nil { + return nil, fmt.Errorf("error getting MIG capability device paths: %w", err) + } + return newMigCapDiscoverer(logger, driver, migCaps, nvcaps.MigCap("monitor")) +} + +func newMigCapDiscoverer(logger logger.Interface, driver *root.Driver, migCaps nvcaps.MigCaps, cap nvcaps.MigCap) (Discover, error) { + if migCaps == nil { + return None{}, nil + } + + capDevicePath, err := migCaps.GetCapDevicePath(cap) + if err != nil { + return nil, fmt.Errorf("failed to get MIG cap device path for %q: %w", cap, err) + } + + return NewCharDeviceDiscoverer( + logger, + driver.DevRoot, + []string{capDevicePath}, + ), nil +} diff --git a/internal/discover/mig_test.go b/internal/discover/mig_test.go new file mode 100644 index 000000000..0f19b8f1a --- /dev/null +++ b/internal/discover/mig_test.go @@ -0,0 +1,125 @@ +/** +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +**/ + +package discover + +import ( + "path/filepath" + "testing" + + testlog "github.com/sirupsen/logrus/hooks/test" + "github.com/stretchr/testify/require" + + "github.com/NVIDIA/nvidia-container-toolkit/internal/devices" + "github.com/NVIDIA/nvidia-container-toolkit/internal/lookup/root" + "github.com/NVIDIA/nvidia-container-toolkit/internal/nvcaps" + "github.com/NVIDIA/nvidia-container-toolkit/internal/test" +) + +func TestNewMigCapDiscoverer(t *testing.T) { + logger, _ := testlog.NewNullLogger() + + defer devices.SetAllForTest()() + + moduleRoot, err := test.GetModuleRoot() + require.NoError(t, err) + + lookupRoot := filepath.Join(moduleRoot, "testdata", "lookup") + + testMigCaps := nvcaps.MigCaps{ + nvcaps.MigCap("config"): nvcaps.MigMinor(1), + nvcaps.MigCap("monitor"): nvcaps.MigMinor(2), + } + + testCases := []struct { + description string + rootfs string + migCaps nvcaps.MigCaps + cap nvcaps.MigCap + expectedDevices []Device + expectedError bool + }{ + { + description: "nil migCaps returns no devices", + rootfs: "rootfs-1", + migCaps: nil, + cap: nvcaps.MigCap("config"), + expectedDevices: nil, + }, + { + description: "cap not in migCaps returns error", + rootfs: "rootfs-1", + migCaps: testMigCaps, + cap: nvcaps.MigCap("gpu0/gi0/access"), + expectedError: true, + }, + { + description: "config cap on empty rootfs returns no devices", + rootfs: "rootfs-empty", + migCaps: testMigCaps, + cap: nvcaps.MigCap("config"), + }, + { + description: "config cap returns nvidia-cap1", + rootfs: "rootfs-1", + migCaps: testMigCaps, + cap: nvcaps.MigCap("config"), + expectedDevices: []Device{ + {Path: "/dev/nvidia-caps/nvidia-cap1", HostPath: "/dev/nvidia-caps/nvidia-cap1"}, + }, + }, + { + description: "monitor cap returns nvidia-cap2", + rootfs: "rootfs-1", + migCaps: testMigCaps, + cap: nvcaps.MigCap("monitor"), + expectedDevices: []Device{ + {Path: "/dev/nvidia-caps/nvidia-cap2", HostPath: "/dev/nvidia-caps/nvidia-cap2"}, + }, + }, + } + + for _, tc := range testCases { + t.Run(tc.description, func(t *testing.T) { + devRoot := filepath.Join(lookupRoot, tc.rootfs) + driver := root.New(root.WithDevRoot(devRoot)) + + d, err := newMigCapDiscoverer(logger, driver, tc.migCaps, tc.cap) + if tc.expectedError { + require.Error(t, err) + return + } + require.NoError(t, err) + + got, err := d.Devices() + require.NoError(t, err) + require.EqualValues(t, tc.expectedDevices, test.StripRoot(got, devRoot)) + + mounts, err := d.Mounts() + require.NoError(t, err) + require.Empty(t, mounts) + + hooks, err := d.Hooks() + require.NoError(t, err) + require.Empty(t, hooks) + + envVars, err := d.EnvVars() + require.NoError(t, err) + require.Empty(t, envVars) + }) + } +} diff --git a/internal/modifier/cdi.go b/internal/modifier/cdi.go index 7b36ab590..07f145afa 100644 --- a/internal/modifier/cdi.go +++ b/internal/modifier/cdi.go @@ -82,6 +82,7 @@ func (f *Factory) newJitCDIModifier(automaticDevices []string) (oci.SpecModifier if f.image != nil { automaticDevices = append(automaticDevices, withUniqueDevices(gatedDevices(*f.image)).DeviceRequests()...) automaticDevices = append(automaticDevices, withUniqueDevices(imexDevices(*f.image)).DeviceRequests()...) + automaticDevices = append(automaticDevices, withUniqueDevices(migCapDevices(*f.image)).DeviceRequests()...) } return f.newAutomaticCDISpecModifier(automaticDevices) } @@ -157,6 +158,27 @@ func (d imexDevices) DeviceRequests() []string { return devices } +type migCapDevices image.CUDA + +// DeviceRequests returns device requests for MIG config and monitor capability devices. +// These devices require a privileged container. +func (m migCapDevices) DeviceRequests() []string { + i := (image.CUDA)(m) + + if !i.IsPrivileged() { + return nil + } + + var devices []string + if i.Getenv(image.EnvVarNvidiaMigConfigDevices) != "" { + devices = append(devices, "mode=mig-config") + } + if i.Getenv(image.EnvVarNvidiaMigMonitorDevices) != "" { + devices = append(devices, "mode=mig-monitor") + } + return devices +} + // filterAutomaticDevices searches for "automatic" device names in the input slice. // "Automatic" devices are a well-defined list of CDI device names which, when requested, // trigger the generation of a CDI spec at runtime. This removes the need to generate a diff --git a/internal/modifier/cdi_test.go b/internal/modifier/cdi_test.go index 841b9561e..ce72df2bd 100644 --- a/internal/modifier/cdi_test.go +++ b/internal/modifier/cdi_test.go @@ -215,6 +215,72 @@ func TestDeviceRequests(t *testing.T) { } } +func TestMigCapDevicesDeviceRequests(t *testing.T) { + testCases := []struct { + description string + env []string + privileged bool + expectedDevices []string + }{ + { + description: "no env vars, unprivileged: no devices", + }, + { + description: "no env vars, privileged: no devices", + privileged: true, + }, + { + description: "mig-config set, privileged: returns mig-config mode", + env: []string{"NVIDIA_MIG_CONFIG_DEVICES=all"}, + privileged: true, + expectedDevices: []string{"mode=mig-config"}, + }, + { + description: "mig-monitor set, privileged: returns mig-monitor mode", + env: []string{"NVIDIA_MIG_MONITOR_DEVICES=all"}, + privileged: true, + expectedDevices: []string{"mode=mig-monitor"}, + }, + { + description: "both set, privileged: returns both modes", + env: []string{"NVIDIA_MIG_CONFIG_DEVICES=all", "NVIDIA_MIG_MONITOR_DEVICES=all"}, + privileged: true, + expectedDevices: []string{"mode=mig-config", "mode=mig-monitor"}, + }, + { + description: "mig-config set, unprivileged: no devices", + env: []string{"NVIDIA_MIG_CONFIG_DEVICES=all"}, + }, + { + description: "mig-monitor set, unprivileged: no devices", + env: []string{"NVIDIA_MIG_MONITOR_DEVICES=all"}, + }, + { + description: "both set, unprivileged: no devices", + env: []string{"NVIDIA_MIG_CONFIG_DEVICES=all", "NVIDIA_MIG_MONITOR_DEVICES=all"}, + }, + { + description: "specific device ID is treated as non-empty value", + env: []string{"NVIDIA_MIG_CONFIG_DEVICES=0"}, + privileged: true, + expectedDevices: []string{"mode=mig-config"}, + }, + } + + for _, tc := range testCases { + t.Run(tc.description, func(t *testing.T) { + img, err := image.New( + image.WithEnv(tc.env), + image.WithPrivileged(tc.privileged), + ) + require.NoError(t, err) + + devices := migCapDevices(img).DeviceRequests() + require.EqualValues(t, tc.expectedDevices, devices) + }) + } +} + func Test_cdiModeIdentfiersFromDevices(t *testing.T) { testCases := []struct { description string @@ -297,6 +363,33 @@ func Test_cdiModeIdentfiersFromDevices(t *testing.T) { deviceClassByMode: map[string]string{"auto": "gpu"}, }, }, + { + description: "mig-config mode only", + devices: []string{"mode=mig-config"}, + expected: &cdiModeIdentifiers{ + modes: []string{"mig-config"}, + idsByMode: map[string][]string{}, + deviceClassByMode: map[string]string{"auto": "gpu"}, + }, + }, + { + description: "mig-monitor mode only", + devices: []string{"mode=mig-monitor"}, + expected: &cdiModeIdentifiers{ + modes: []string{"mig-monitor"}, + idsByMode: map[string][]string{}, + deviceClassByMode: map[string]string{"auto": "gpu"}, + }, + }, + { + description: "mig-config and mig-monitor together with auto devices", + devices: []string{"runtime.nvidia.com/gpu=all", "mode=mig-config", "mode=mig-monitor"}, + expected: &cdiModeIdentifiers{ + modes: []string{"auto", "mig-config", "mig-monitor"}, + idsByMode: map[string][]string{"auto": {"all"}}, + deviceClassByMode: map[string]string{"auto": "gpu"}, + }, + }, } for _, tc := range testCases { diff --git a/pkg/nvcdi/gated.go b/pkg/nvcdi/gated.go index b5d1fe7e4..cb627b5eb 100644 --- a/pkg/nvcdi/gated.go +++ b/pkg/nvcdi/gated.go @@ -61,6 +61,10 @@ func (l *gatedlib) getModeDiscoverer() (discover.Discover, error) { return discover.NewGDRCopyDiscoverer(l.logger, l.driver) case ModeGds: return discover.NewGDSDiscoverer(l.logger, l.driver) + case ModeMigConfig: + return discover.NewMigConfigDiscoverer(l.logger, l.driver) + case ModeMigMonitor: + return discover.NewMigMonitorDiscoverer(l.logger, l.driver) case ModeMofed: return discover.NewMOFEDDiscoverer(l.logger, l.driver) case ModeNvswitch: diff --git a/pkg/nvcdi/lib.go b/pkg/nvcdi/lib.go index fa84e19c5..90f8e8be8 100644 --- a/pkg/nvcdi/lib.go +++ b/pkg/nvcdi/lib.go @@ -84,7 +84,7 @@ func New(opts ...Option) (Interface, error) { factory = (*nvmllib)(l) case ModeWsl: factory = (*wsllib)(l) - case ModeGdrcopy, ModeGds, ModeMofed, ModeNvswitch: + case ModeGdrcopy, ModeGds, ModeMigConfig, ModeMigMonitor, ModeMofed, ModeNvswitch: factory = &gatedlib{ nvcdilib: l, mode: o.mode, diff --git a/pkg/nvcdi/mode.go b/pkg/nvcdi/mode.go index a2f2f6e64..bedc89aaa 100644 --- a/pkg/nvcdi/mode.go +++ b/pkg/nvcdi/mode.go @@ -46,6 +46,10 @@ const ( ModeImex = Mode("imex") // ModeNvswitch configures the CDI spec generator to generate a spec for the available nvswitch devices. ModeNvswitch = Mode("nvswitch") + // ModeMigConfig configures the CDI spec generator to generate a spec for the MIG config capability device. + ModeMigConfig = Mode("mig-config") + // ModeMigMonitor configures the CDI spec generator to generate a spec for the MIG monitor capability device. + ModeMigMonitor = Mode("mig-monitor") ) type modeConstraint interface { @@ -69,6 +73,8 @@ func getModes() modes { ModeGds, ModeImex, ModeManagement, + ModeMigConfig, + ModeMigMonitor, ModeMofed, ModeNvml, ModeNvswitch, diff --git a/testdata/lookup/rootfs-1/dev/nvidia-caps/nvidia-cap2 b/testdata/lookup/rootfs-1/dev/nvidia-caps/nvidia-cap2 new file mode 100644 index 000000000..e69de29bb