Enable JIT-CDI to inject MIG management capability devices

internal/modifier/cdi.go

@@ -18,8 +18,10 @@ package modifier
 
 import (
 	"fmt"
+	"slices"
 	"strings"
 
+	"github.com/NVIDIA/go-nvlib/pkg/nvlib/device"
 	"tags.cncf.io/container-device-interface/pkg/parser"
 
 	"github.com/NVIDIA/nvidia-container-toolkit/internal/config/image"
@@ -82,6 +84,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(migCapsDevices(*f.image)).DeviceRequests()...)
 	}
 	return f.newAutomaticCDISpecModifier(automaticDevices)
 }
@@ -157,6 +160,51 @@ func (d imexDevices) DeviceRequests() []string {
 	return devices
 }
 
+type migCapsDevices image.CUDA
+
+// DeviceRequests returns the MIG management capabilities requested through the
+// NVIDIA_MIG_CONFIG_DEVICES and NVIDIA_MIG_MONITOR_DEVICES envvars.
+//
+// Only the value "all" is meaningful here: it maps to the global MIG config /
+// monitor capability nodes (nvidia-cap<config>, nvidia-cap<monitor>).
+func (d migCapsDevices) DeviceRequests() []string {
+	i := (image.CUDA)(d)
+
+	var devices []string
+	if strings.EqualFold(i.Getenv(image.EnvVarNvidiaMigConfigDevices), "all") {
+		devices = append(devices, "mode=mig-caps,id=config")
+	}
+	if strings.EqualFold(i.Getenv(image.EnvVarNvidiaMigMonitorDevices), "all") {
+		devices = append(devices, "mode=mig-caps,id=monitor")
+	}
+	if len(devices) == 0 {
+		return nil
+	}
+	// The global MIG config/monitor capability grants GPU-wide access and would
+	// expose sibling MIG instances on the same GPU. Do not grant it to a container
+	// that is scoped to specific MIG devices.
+	if requestsSpecificMigDevices(i) {
+		return nil
+	}
+	return devices
+}
+
+// requestsSpecificMigDevices reports whether the visible device list selects one
+// or more specific MIG devices rather than whole GPUs or "all".
+func requestsSpecificMigDevices(i image.CUDA) bool {
+	visibleDevices := i.VisibleDevices()
+	if slices.Contains(visibleDevices, "all") {
+		return false
+	}
+	for _, d := range visibleDevices {
+		id := device.Identifier(d)
+		if id.IsMigUUID() || id.IsMigIndex() {
+			return true
+		}
+	}
+	return false
+}
+
 // 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,69 @@ func TestDeviceRequests(t *testing.T) {
 	}
 }
 
+func TestMigCapsDeviceRequests(t *testing.T) {
+	testCases := []struct {
+		description     string
+		env             []string
+		expectedDevices []string
+	}{
+		{
+			description: "no MIG envvars yields no devices",
+		},
+		{
+			description:     "monitor devices requested",
+			env:             []string{"NVIDIA_MIG_MONITOR_DEVICES=all"},
+			expectedDevices: []string{"mode=mig-caps,id=monitor"},
+		},
+		{
+			description:     "config devices requested",
+			env:             []string{"NVIDIA_MIG_CONFIG_DEVICES=all"},
+			expectedDevices: []string{"mode=mig-caps,id=config"},
+		},
+		{
+			description:     "both config and monitor requested",
+			env:             []string{"NVIDIA_MIG_CONFIG_DEVICES=all", "NVIDIA_MIG_MONITOR_DEVICES=all"},
+			expectedDevices: []string{"mode=mig-caps,id=config", "mode=mig-caps,id=monitor"},
+		},
+		{
+			description: "empty value yields no devices",
+			env:         []string{"NVIDIA_MIG_MONITOR_DEVICES="},
+		},
+		{
+			description: "non-all value is ignored",
+			env:         []string{"NVIDIA_MIG_MONITOR_DEVICES=0", "NVIDIA_MIG_CONFIG_DEVICES=0,1"},
+		},
+		{
+			description:     "allowed with whole-GPU visibility",
+			env:             []string{"NVIDIA_VISIBLE_DEVICES=0", "NVIDIA_MIG_MONITOR_DEVICES=all"},
+			expectedDevices: []string{"mode=mig-caps,id=monitor"},
+		},
+		{
+			description: "ignored when scoped to specific MIG devices (index)",
+			env:         []string{"NVIDIA_VISIBLE_DEVICES=0:0", "NVIDIA_MIG_MONITOR_DEVICES=all"},
+		},
+		{
+			description: "ignored when scoped to specific MIG devices (UUID)",
+			env:         []string{"NVIDIA_VISIBLE_DEVICES=MIG-GPU-b1028956-cfa2-0990-bf4a-5da9abb51763/3/0", "NVIDIA_MIG_CONFIG_DEVICES=all"},
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.description, func(t *testing.T) {
+			img, err := image.NewCUDAImageFromSpec(
+				&specs.Spec{
+					Process: &specs.Process{Env: tc.env},
+				},
+				image.WithAcceptEnvvarUnprivileged(true),
+			)
+			require.NoError(t, err)
+
+			devices := migCapsDevices(img).DeviceRequests()
+			require.EqualValues(t, tc.expectedDevices, devices)
+		})
+	}
+}
+
 func Test_cdiModeIdentfiersFromDevices(t *testing.T) {
 	testCases := []struct {
 		description string

internal/nvcaps/nvcaps.go

@@ -69,12 +69,18 @@ func (m MigCaps) GetCapDevicePath(cap MigCap) (string, error) {
 
 // NewMigCaps creates a MigCaps structure based on the contents of the MIG minors file.
 func NewMigCaps() (MigCaps, error) {
-	// Open nvcapsMigMinorsPath for walking.
+	return NewMigCapsFromRoot("")
+}
+
+// NewMigCapsFromRoot creates a MigCaps structure based on the contents of the MIG
+// minors file resolved relative to the specified root. A root of "" reads the
+// host's MIG minors file.
+func NewMigCapsFromRoot(root string) (MigCaps, error) {
 	// If the nvcapsMigMinorsPath does not exist, then we are not on a MIG
 	// capable machine, so there is nothing to do.
 	// The format of this file is discussed in:
 	//     https://docs.nvidia.com/datacenter/tesla/mig-user-guide/index.html#unique_1576522674
-	minorsFile, err := os.Open(nvcapsMigMinorsPath)
+	minorsFile, err := os.Open(filepath.Join(root, nvcapsMigMinorsPath))
 	if os.IsNotExist(err) {
 		return nil, nil
 	}

pkg/nvcdi/lib-mig-caps.go

@@ -0,0 +1,100 @@
+/**
+# 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 nvcdi
+
+import (
+	"fmt"
+
+	"tags.cncf.io/container-device-interface/pkg/cdi"
+	"tags.cncf.io/container-device-interface/specs-go"
+
+	"github.com/NVIDIA/nvidia-container-toolkit/internal/discover"
+	"github.com/NVIDIA/nvidia-container-toolkit/internal/nvcaps"
+)
+
+type migCapsLib nvcdilib
+
+// migCapDeviceSpecGenerator generates the CDI device spec for a single MIG
+// management capability (config or monitor).
+type migCapDeviceSpecGenerator struct {
+	lib *migCapsLib
+	cap nvcaps.MigCap
+}
+
+var _ deviceSpecGeneratorFactory = (*migCapsLib)(nil)
+
+// GetCommonEdits returns an empty set of edits for MIG capability devices.
+func (l *migCapsLib) GetCommonEdits() (*cdi.ContainerEdits, error) {
+	return l.editsFactory.FromDiscoverer(discover.None{})
+}
+
+// DeviceSpecGenerators returns the CDI device spec generators for the specified
+// MIG management capabilities.
+// Valid IDs are 'config' and 'monitor'.
+func (l *migCapsLib) DeviceSpecGenerators(ids ...string) (DeviceSpecGenerator, error) {
+	var deviceSpecGenerators DeviceSpecGenerators
+	for _, id := range ids {
+		cap := nvcaps.MigCap(id)
+		switch cap {
+		case "config", "monitor":
+			deviceSpecGenerators = append(deviceSpecGenerators, &migCapDeviceSpecGenerator{lib: l, cap: cap})
+		default:
+			return nil, fmt.Errorf("invalid MIG capability %q: must be one of [config, monitor]", id)
+		}
+	}
+	return deviceSpecGenerators, nil
+}
+
+// GetDeviceSpecs returns the CDI device specs for the MIG management capability.
+func (g *migCapDeviceSpecGenerator) GetDeviceSpecs() ([]specs.Device, error) {
+	l := g.lib
+
+	migCaps, err := nvcaps.NewMigCapsFromRoot(l.driver.Root)
+	if err != nil {
+		return nil, fmt.Errorf("failed to get MIG capability device paths: %w", err)
+	}
+	if migCaps == nil {
+		return nil, fmt.Errorf("cannot inject MIG %s capability: system is not MIG capable", g.cap)
+	}
+
+	devicePath, err := migCaps.GetCapDevicePath(g.cap)
+	if err != nil {
+		return nil, fmt.Errorf("failed to get device path for MIG %s capability: %w", g.cap, err)
+	}
+
+	deviceNodes := discover.NewCharDeviceDiscoverer(
+		l.logger,
+		l.driver.DevRoot,
+		[]string{devicePath},
+	)
+
+	// The MIG capability nodes are nested under /dev/nvidia-caps. Add a hook to
+	// set the permissions of the parent folder so that non-root users in the
+	// container can access the injected node.
+	folderPermissionHooks := (*nvcdilib)(l).newDeviceFolderPermissionHookDiscoverer(deviceNodes)
+
+	edits, err := l.editsFactory.FromDiscoverer(discover.Merge(deviceNodes, folderPermissionHooks))
+	if err != nil {
+		return nil, fmt.Errorf("failed to create container edits for MIG %s capability: %w", g.cap, err)
+	}
+
+	deviceSpec := specs.Device{
+		Name:           string(g.cap),
+		ContainerEdits: *edits.ContainerEdits,
+	}
+	return []specs.Device{deviceSpec}, nil
+}

pkg/nvcdi/lib-mig-caps_test.go

@@ -0,0 +1,117 @@
+/**
+# 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 nvcdi
+
+import (
+	"path/filepath"
+	"slices"
+	"testing"
+
+	testlog "github.com/sirupsen/logrus/hooks/test"
+	"github.com/stretchr/testify/require"
+	"tags.cncf.io/container-device-interface/specs-go"
+
+	"github.com/NVIDIA/nvidia-container-toolkit/internal/devices"
+	"github.com/NVIDIA/nvidia-container-toolkit/internal/test"
+)
+
+func TestMigCapsValidation(t *testing.T) {
+	logger, _ := testlog.NewNullLogger()
+
+	testCases := []struct {
+		description string
+		ids         []string
+		expectedErr string
+	}{
+		{
+			description: "unknown capability is rejected",
+			ids:         []string{"bogus"},
+			expectedErr: "invalid MIG capability",
+		},
+		{
+			description: "per-instance access cap is rejected",
+			ids:         []string{"gpu0/gi0/access"},
+			expectedErr: "invalid MIG capability",
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.description, func(t *testing.T) {
+			lib, err := New(
+				WithLogger(logger),
+				WithMode(ModeMigCaps),
+			)
+			require.NoError(t, err)
+
+			_, err = lib.GetDeviceSpecsByID(tc.ids...)
+			require.ErrorContains(t, err, tc.expectedErr)
+		})
+	}
+}
+
+func TestMigCapsMode(t *testing.T) {
+	defer devices.SetAllForTest()()
+
+	logger, _ := testlog.NewNullLogger()
+
+	moduleRoot, err := test.GetModuleRoot()
+	require.NoError(t, err)
+	hostRoot := filepath.Join(moduleRoot, "testdata", "lookup", "rootfs-mig")
+
+	lib, err := New(
+		WithLogger(logger),
+		WithMode(ModeMigCaps),
+		WithDriverRoot(hostRoot),
+		WithEnabledHooks("chmod"),
+	)
+	require.NoError(t, err)
+
+	deviceSpecs, err := lib.GetDeviceSpecsByID("config", "monitor")
+	require.NoError(t, err)
+	require.Len(t, deviceSpecs, 2)
+
+	byName := make(map[string]specs.Device)
+	for _, d := range deviceSpecs {
+		byName[d.Name] = d
+	}
+
+	expected := map[string]string{
+		"config":  "/dev/nvidia-caps/nvidia-cap1",
+		"monitor": "/dev/nvidia-caps/nvidia-cap2",
+	}
+	for name, path := range expected {
+		spec, ok := byName[name]
+		require.Truef(t, ok, "expected a device named %q", name)
+
+		require.Len(t, spec.ContainerEdits.DeviceNodes, 1)
+		deviceNode := spec.ContainerEdits.DeviceNodes[0]
+		require.Equal(t, path, deviceNode.Path)
+		require.Equal(t, filepath.Join(hostRoot, path), deviceNode.HostPath)
+
+		require.True(t, hasFolderPermissionHook(spec, "/dev/nvidia-caps"),
+			"expected a folder-permission hook for /dev/nvidia-caps on device %q", name)
+	}
+}
+
+func hasFolderPermissionHook(spec specs.Device, folder string) bool {
+	for _, hook := range spec.ContainerEdits.Hooks {
+		if slices.Contains(hook.Args, folder) {
+			return true
+		}
+	}
+	return false
+}

pkg/nvcdi/lib.go

@@ -91,6 +91,8 @@ func New(opts ...Option) (Interface, error) {
 		}
 	case ModeImex:
 		factory = (*imexlib)(l)
+	case ModeMigCaps:
+		factory = (*migCapsLib)(l)
 	default:
 		return nil, fmt.Errorf("unknown mode %q", o.mode)
 	}