Process envvars for mig-config / mig-monitor devices in jit-cdi mode

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
+}

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)
+		})
+	}
+}

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

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 {

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:

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,

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,