Add vfio mode to generate CDI specs for NVIDIA passthrough GPUs

go.mod

@@ -3,7 +3,7 @@ module github.com/NVIDIA/nvidia-container-toolkit
 go 1.23.0
 
 require (
-	github.com/NVIDIA/go-nvlib v0.7.3
+	github.com/NVIDIA/go-nvlib v0.7.4-0.20250707141440-532907f49628
 	github.com/NVIDIA/go-nvml v0.12.9-0
 	github.com/cyphar/filepath-securejoin v0.4.1
 	github.com/moby/sys/reexec v0.1.0

go.sum

@@ -1,5 +1,5 @@
-github.com/NVIDIA/go-nvlib v0.7.3 h1:kXc8PkWUlrwedSpM4fR8xT/DAq1NKy8HqhpgteFcGAw=
-github.com/NVIDIA/go-nvlib v0.7.3/go.mod h1:i95Je7GinMy/+BDs++DAdbPmT2TubjNP8i8joC7DD7I=
+github.com/NVIDIA/go-nvlib v0.7.4-0.20250707141440-532907f49628 h1:epoT+CVQMISAldX5fx78gjJs23Kx17B/6dOwcXvmPaE=
+github.com/NVIDIA/go-nvlib v0.7.4-0.20250707141440-532907f49628/go.mod h1:i95Je7GinMy/+BDs++DAdbPmT2TubjNP8i8joC7DD7I=
 github.com/NVIDIA/go-nvml v0.12.9-0 h1:e344UK8ZkeMeeLkdQtRhmXRxNf+u532LDZPGMtkdus0=
 github.com/NVIDIA/go-nvml v0.12.9-0/go.mod h1:+KNA7c7gIBH7SKSJ1ntlwkfN80zdx8ovl4hrK3LmPt4=
 github.com/blang/semver/v4 v4.0.0 h1:1PFHFE6yCCTv8C1TeyNNarDzntLi7wMI5i/pzqYIsAM=

pkg/nvcdi/lib-vfio.go

@@ -0,0 +1,122 @@
+/**
+# SPDX-FileCopyrightText: Copyright (c) 2025 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 nvcdi
+
+import (
+	"fmt"
+	"path/filepath"
+	"strconv"
+
+	"tags.cncf.io/container-device-interface/pkg/cdi"
+	"tags.cncf.io/container-device-interface/specs-go"
+)
+
+type vfiolib nvcdilib
+
+type vfioDevice struct {
+	index   int
+	group   int
+	devRoot string
+}
+
+var _ deviceSpecGeneratorFactory = (*vfiolib)(nil)
+
+func (l *vfiolib) DeviceSpecGenerators(ids ...string) (DeviceSpecGenerator, error) {
+	vfioDevices, err := l.getVfioDevices(ids...)
+	if err != nil {
+		return nil, err
+	}
+	var deviceSpecGenerators DeviceSpecGenerators
+	for _, vfioDevice := range vfioDevices {
+		deviceSpecGenerators = append(deviceSpecGenerators, vfioDevice)
+	}
+
+	return deviceSpecGenerators, nil
+}
+
+// GetDeviceSpecs returns the CDI device specs for a vfio device.
+func (l *vfioDevice) GetDeviceSpecs() ([]specs.Device, error) {
+	path := fmt.Sprintf("/dev/vfio/%d", l.group)
+	deviceSpec := specs.Device{
+		Name: fmt.Sprintf("%d", l.index),
+		ContainerEdits: specs.ContainerEdits{
+			DeviceNodes: []*specs.DeviceNode{
+				{
+					Path:     path,
+					HostPath: filepath.Join(l.devRoot, path),
+				},
+			},
+		},
+	}
+	return []specs.Device{deviceSpec}, nil
+}
+
+// GetCommonEdits returns common edits for ALL devices.
+// Note, currently there are no common edits.
+func (l *vfiolib) GetCommonEdits() (*cdi.ContainerEdits, error) {
+	e := cdi.ContainerEdits{
+		ContainerEdits: &specs.ContainerEdits{
+			DeviceNodes: []*specs.DeviceNode{
+				{
+					Path:     "/dev/vfio/vfio",
+					HostPath: filepath.Join(l.devRoot, "/dev/vfio/vfio"),
+				},
+			},
+		},
+	}
+	return &e, nil
+}
+
+func (l *vfiolib) getVfioDevices(ids ...string) ([]*vfioDevice, error) {
+	var vfioDevices []*vfioDevice
+	for _, id := range ids {
+		if id == "all" {
+			return l.getAllVfioDevices()
+		}
+		index, err := strconv.ParseInt(id, 10, 32)
+		if err != nil {
+			return nil, fmt.Errorf("invalid channel ID %v: %w", id, err)
+		}
+		i := int(index)
+		dev, err := l.nvpcilib.GetGPUByIndex(i)
+		if err != nil {
+			return nil, fmt.Errorf("failed to get device: %w", err)
+		}
+		vfioDevices = append(vfioDevices, &vfioDevice{index: i, group: dev.IommuGroup, devRoot: l.devRoot})
+	}
+
+	return vfioDevices, nil
+}
+
+func (l *vfiolib) getAllVfioDevices() ([]*vfioDevice, error) {
+	devices, err := l.nvpcilib.GetGPUs()
+	if err != nil {
+		return nil, fmt.Errorf("failed getting NVIDIA GPUs: %v", err)
+	}
+
+	var vfioDevices []*vfioDevice
+	for i, dev := range devices {
+		if dev.Driver != "vfio-pci" {
+			continue
+		}
+		l.logger.Debugf("Found NVIDIA device: address=%s, driver=%s, iommu_group=%d, deviceId=%x",
+			dev.Address, dev.Driver, dev.IommuGroup, dev.Device)
+		vfioDevices = append(vfioDevices, &vfioDevice{index: i, group: dev.IommuGroup, devRoot: l.devRoot})
+	}
+	return vfioDevices, nil
+}

pkg/nvcdi/lib-vfio_test.go

@@ -0,0 +1,123 @@
+/**
+# SPDX-FileCopyrightText: Copyright (c) 2025 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 nvcdi
+
+import (
+	"bytes"
+	"testing"
+
+	"github.com/NVIDIA/go-nvlib/pkg/nvpci"
+	"github.com/stretchr/testify/require"
+)
+
+func TestModeVfio(t *testing.T) {
+	testCases := []struct {
+		description   string
+		pcilib        *nvpci.InterfaceMock
+		ids           []string
+		expectedError error
+		expectedSpec  string
+	}{
+		{
+			description: "get all specs single device",
+			pcilib: &nvpci.InterfaceMock{
+				GetGPUsFunc: func() ([]*nvpci.NvidiaPCIDevice, error) {
+					devices := []*nvpci.NvidiaPCIDevice{
+						{
+							Driver:     "vfio-pci",
+							IommuGroup: 5,
+						},
+					}
+					return devices, nil
+				},
+			},
+			expectedSpec: `---
+cdiVersion: 0.5.0
+kind: nvidia.com/pgpu
+devices:
+    - name: "0"
+      containerEdits:
+        deviceNodes:
+            - path: /dev/vfio/5
+              hostPath: /dev/vfio/5
+containerEdits:
+    env:
+        - NVIDIA_VISIBLE_DEVICES=void
+    deviceNodes:
+        - path: /dev/vfio/vfio
+          hostPath: /dev/vfio/vfio
+`,
+		},
+		{
+			description: "get single device spec by index",
+			pcilib: &nvpci.InterfaceMock{
+				GetGPUByIndexFunc: func(n int) (*nvpci.NvidiaPCIDevice, error) {
+					devices := []*nvpci.NvidiaPCIDevice{
+						{
+							Driver:     "vfio-pci",
+							IommuGroup: 45,
+						},
+						{
+							Driver:     "vfio-pci",
+							IommuGroup: 5,
+						},
+					}
+					return devices[n], nil
+				},
+			},
+			ids: []string{"1"},
+			expectedSpec: `---
+cdiVersion: 0.5.0
+kind: nvidia.com/pgpu
+devices:
+    - name: "1"
+      containerEdits:
+        deviceNodes:
+            - path: /dev/vfio/5
+              hostPath: /dev/vfio/5
+containerEdits:
+    env:
+        - NVIDIA_VISIBLE_DEVICES=void
+    deviceNodes:
+        - path: /dev/vfio/vfio
+          hostPath: /dev/vfio/vfio
+`,
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.description, func(t *testing.T) {
+			lib, err := New(
+				WithMode(ModeVfio),
+				WithPCILib(tc.pcilib),
+			)
+			require.NoError(t, err)
+
+			spec, err := lib.GetSpec(tc.ids...)
+			require.EqualValues(t, tc.expectedError, err)
+
+			var output bytes.Buffer
+
+			_, err = spec.WriteTo(&output)
+			require.NoError(t, err)
+
+			require.Equal(t, tc.expectedSpec, output.String())
+		})
+	}
+
+}

pkg/nvcdi/lib.go

@@ -21,6 +21,7 @@ import (
 
 	"github.com/NVIDIA/go-nvlib/pkg/nvlib/device"
 	"github.com/NVIDIA/go-nvlib/pkg/nvlib/info"
+	"github.com/NVIDIA/go-nvlib/pkg/nvpci"
 	"github.com/NVIDIA/go-nvml/pkg/nvml"
 
 	"github.com/NVIDIA/nvidia-container-toolkit/internal/discover"
@@ -54,6 +55,8 @@ type nvcdilib struct {
 	driver  *root.Driver
 	infolib info.Interface
 
+	nvpcilib nvpci.Interface
+
 	mergedDeviceOptions []transform.MergedDeviceOption
 
 	featureFlags map[FeatureFlag]bool
@@ -151,6 +154,14 @@ func New(opts ...Option) (Interface, error) {
 			l.class = classImexChannel
 		}
 		factory = (*imexlib)(l)
+	case ModeVfio:
+		if l.class == "" {
+			l.class = "pgpu"
+		}
+		if l.nvpcilib == nil {
+			l.nvpcilib = nvpci.New()
+		}
+		factory = (*vfiolib)(l)
 	default:
 		return nil, fmt.Errorf("unknown mode %q", l.mode)
 	}

pkg/nvcdi/mode.go

@@ -42,6 +42,8 @@ const (
 	ModeCSV = Mode("csv")
 	// ModeImex configures the CDI spec generated to generate a spec for the available IMEX channels.
 	ModeImex = Mode("imex")
+	// ModeVfio configures the CDI spec generator to generate a VFIO spec.
+	ModeVfio = Mode("vfio")
 )
 
 type modeConstraint interface {
@@ -66,6 +68,7 @@ func getModes() modes {
 			ModeGds,
 			ModeMofed,
 			ModeCSV,
+			ModeVfio,
 		}
 		lookup := make(map[Mode]bool)
 

pkg/nvcdi/options.go

@@ -19,6 +19,7 @@ package nvcdi
 import (
 	"github.com/NVIDIA/go-nvlib/pkg/nvlib/device"
 	"github.com/NVIDIA/go-nvlib/pkg/nvlib/info"
+	"github.com/NVIDIA/go-nvlib/pkg/nvpci"
 	"github.com/NVIDIA/go-nvml/pkg/nvml"
 
 	"github.com/NVIDIA/nvidia-container-toolkit/internal/discover"
@@ -43,6 +44,13 @@ func WithInfoLib(infolib info.Interface) Option {
 	}
 }
 
+// WithPCILib sets the pci library to be used for CDI spec generation.
+func WithPCILib(pcilib nvpci.Interface) Option {
+	return func(l *nvcdilib) {
+		l.nvpcilib = pcilib
+	}
+}
+
 // WithDeviceNamers sets the device namer for the library
 func WithDeviceNamers(namers ...DeviceNamer) Option {
 	return func(l *nvcdilib) {