gpu: add pci device id allow support

By defining allowed PCI IDs, it's possible to only
select some GPUs per host. For example, on a desktop with
integrated and discrete graphics, GPU plugin can only
register the discrete one.

Signed-off-by: Tuomas Katila <[email protected]>

Author: tkatila

cmd/gpu_plugin/README.md

@@ -57,6 +57,7 @@ For workloads on different KMDs, see [KMD and UMD](#kmd-and-umd).
 | -health-management | - | disabled | Enable health management by requesting data from oneAPI/Level-Zero interface. Requires [GPU Level-Zero](../gpu_levelzero/) sidecar. See [health management](#health-management) |
 | -wsl | - | disabled | Adapt plugin to run in the WSL environment. Requires [GPU Level-Zero](../gpu_levelzero/) sidecar. |
 | -shared-dev-num | int | 1 | Number of containers that can share the same GPU device |
+| -allowlist-ids | string | "" | A list of PCI Device IDs that are allowed to be registered as resources. Default is empty (=all registered). |
 | -allocation-policy | string | none | 3 possible values: balanced, packed, none. For shared-dev-num > 1: _balanced_ mode spreads workloads among GPU devices, _packed_ mode fills one GPU fully before moving to next, and _none_ selects first available device from kubelet. Default is _none_. |
 
 The plugin also accepts a number of other arguments (common to all plugins) related to logging.

cmd/gpu_plugin/gpu_plugin.go

@@ -69,6 +69,7 @@ const (
 
 type cliOptions struct {
 	preferredAllocationPolicy string
+	allowlistIDs              string
 	sharedDevNum              int
 	temperatureLimit          int
 	enableMonitoring          bool
@@ -204,6 +205,31 @@ func packedPolicy(req *pluginapi.ContainerPreferredAllocationRequest) []string {
 	return deviceIds
 }
 
+func parsePCIDeviceIDs(allowlist string) ([]string, error) {
+	deviceIDs := make([]string, 0, strings.Count(allowlist, ",")+1)
+
+	r := regexp.MustCompile(`^0x[0-9a-f]{4}$`)
+
+	for id := range strings.SplitSeq(allowlist, ",") {
+		id = strings.TrimSpace(id)
+		if id == "" {
+			klog.Errorf("Empty PCI device ID given")
+
+			return []string{}, os.ErrInvalid
+		}
+
+		if !r.MatchString(id) {
+			klog.Errorf("Invalid PCI device ID %s", id)
+
+			return []string{}, os.ErrInvalid
+		}
+
+		deviceIDs = append(deviceIDs, id)
+	}
+
+	return deviceIDs, nil
+}
+
 func (dp *devicePlugin) pciAddressForCard(cardPath, cardName string) (string, error) {
 	linkPath, err := os.Readlink(cardPath)
 	if err != nil {
@@ -585,6 +611,23 @@ func (dp *devicePlugin) filterOutInvalidCards(files []fs.DirEntry) []fs.DirEntry
 			continue
 		}
 
+		// Skip if the device is not in allowlist.
+		if len(dp.options.allowlistIDs) > 0 {
+			pciID, err := pciDeviceIDForCard(path.Join(dp.sysfsDir, f.Name()))
+			if err != nil {
+				klog.Warningf("Failed to get PCI ID for device %s: %+v", f.Name(), err)
+
+				continue
+			}
+
+			if !strings.Contains(dp.options.allowlistIDs, pciID) {
+				klog.V(4).Infof("Skipping device %s (%s), not in allowlist: %s",
+					f.Name(), pciID, dp.options.allowlistIDs)
+
+				continue
+			}
+		}
+
 		filtered = append(filtered, f)
 	}
 
@@ -723,6 +766,7 @@ func main() {
 	flag.IntVar(&opts.sharedDevNum, "shared-dev-num", 1, "number of containers sharing the same GPU device")
 	flag.IntVar(&opts.temperatureLimit, "temp-limit", 100, "temperature limit at which device is marked unhealthy")
 	flag.StringVar(&opts.preferredAllocationPolicy, "allocation-policy", "none", "modes of allocating GPU devices: balanced, packed and none")
+	flag.StringVar(&opts.allowlistIDs, "allowlist-ids", "", "comma-separated list of device IDs to allow (e.g. 0x49c5,0x49c6)")
 	flag.Parse()
 
 	if opts.sharedDevNum < 1 {
@@ -736,6 +780,18 @@ func main() {
 		os.Exit(1)
 	}
 
+	if opts.allowlistIDs != "" {
+		if allowListIDs, err := parsePCIDeviceIDs(opts.allowlistIDs); err != nil {
+			klog.Error("Failed to parse allowlist-ids: ", err)
+
+			os.Exit(1)
+		} else {
+			klog.V(2).Infof("Allowed device IDs: %q", allowListIDs)
+
+			opts.allowlistIDs = strings.Join(allowListIDs, ",")
+		}
+	}
+
 	klog.V(1).Infof("GPU device plugin started with %s preferred allocation policy", opts.preferredAllocationPolicy)
 
 	plugin := newDevicePlugin(prefix+sysfsDrmDirectory, prefix+devfsDriDirectory, opts)

cmd/gpu_plugin/gpu_plugin_test.go

@@ -361,6 +361,66 @@ func TestScan(t *testing.T) {
 			expectedI915Devs:     1,
 			expectedI915Monitors: 1,
 		},
+		{
+			name:      "two devices with only one allowed",
+			sysfsdirs: []string{"card0/device/drm/card0", "card0/device/drm/controlD64", "card1/device/drm/card1"},
+			sysfsfiles: map[string][]byte{
+				"card0/device/vendor": []byte("0x8086"),
+				"card0/device/device": []byte("0x1234"),
+				"card1/device/vendor": []byte("0x8086"),
+				"card1/device/device": []byte("0x9876"),
+			},
+			symlinkfiles: map[string]string{
+				"card0/device/driver": "drivers/xe",
+				"card1/device/driver": "drivers/i915",
+			},
+			devfsdirs: []string{
+				"card0",
+				"by-path/pci-0000:00:00.0-card",
+				"by-path/pci-0000:00:00.0-render",
+				"card1",
+				"by-path/pci-0000:00:01.0-card",
+				"by-path/pci-0000:00:01.0-render",
+			},
+			options:              cliOptions{enableMonitoring: true, allowlistIDs: "0x1234"},
+			expectedXeDevs:       1,
+			expectedXeMonitors:   1,
+			expectedI915Devs:     0,
+			expectedI915Monitors: 0,
+		},
+		{
+			name:      "three devices with two allowed",
+			sysfsdirs: []string{"card0/device/drm/card0", "card0/device/drm/controlD64", "card1/device/drm/card1", "card2/device/drm/card2"},
+			sysfsfiles: map[string][]byte{
+				"card0/device/vendor": []byte("0x8086"),
+				"card0/device/device": []byte("0x1234"),
+				"card1/device/vendor": []byte("0x8086"),
+				"card1/device/device": []byte("0x9876"),
+				"card2/device/vendor": []byte("0x8086"),
+				"card2/device/device": []byte("0x0101"),
+			},
+			symlinkfiles: map[string]string{
+				"card0/device/driver": "drivers/xe",
+				"card1/device/driver": "drivers/i915",
+				"card2/device/driver": "drivers/i915",
+			},
+			devfsdirs: []string{
+				"card0",
+				"by-path/pci-0000:00:00.0-card",
+				"by-path/pci-0000:00:00.0-render",
+				"card1",
+				"by-path/pci-0000:00:01.0-card",
+				"by-path/pci-0000:00:01.0-render",
+				"card2",
+				"by-path/pci-0000:00:02.0-card",
+				"by-path/pci-0000:00:02.0-render",
+			},
+			options:              cliOptions{enableMonitoring: true, allowlistIDs: "0x1234,0x9876"},
+			expectedXeDevs:       1,
+			expectedXeMonitors:   1,
+			expectedI915Devs:     1,
+			expectedI915Monitors: 1,
+		},
 		{
 			name:      "sriov-1-pf-no-vfs + monitoring",
 			sysfsdirs: []string{"card0/device/drm/card0", "card0/device/drm/controlD64"},
@@ -1048,3 +1108,73 @@ func TestCDIDeviceInclusion(t *testing.T) {
 		t.Error("Invalid count for device (xe)")
 	}
 }
+
+func TestParsePCIDeviceIDs(t *testing.T) {
+	tests := []struct {
+		name      string
+		input     string
+		want      []string
+		wantError bool
+	}{
+		{
+			name:      "valid single ID",
+			input:     "0x1234",
+			want:      []string{"0x1234"},
+			wantError: false,
+		},
+		{
+			name:      "valid multiple IDs",
+			input:     "0x1234,0x5678,0x9abc",
+			want:      []string{"0x1234", "0x5678", "0x9abc"},
+			wantError: false,
+		},
+		{
+			name:      "valid IDs with spaces",
+			input:     " 0x1234 , 0x5678 ",
+			want:      []string{"0x1234", "0x5678"},
+			wantError: false,
+		},
+		{
+			name:      "empty string",
+			input:     "",
+			want:      []string{},
+			wantError: true,
+		},
+		{
+			name:      "invalid ID format",
+			input:     "0x1234,abcd",
+			want:      []string{},
+			wantError: true,
+		},
+		{
+			name:      "invalid hex length",
+			input:     "0x123,0x5678",
+			want:      []string{},
+			wantError: true,
+		},
+		{
+			name:      "extra comma",
+			input:     "0x1234,",
+			want:      []string{},
+			wantError: true,
+		},
+		{
+			name:      "capita hex",
+			input:     "0xAA12,",
+			want:      []string{},
+			wantError: true,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			got, err := parsePCIDeviceIDs(tt.input)
+			if (err != nil) != tt.wantError {
+				t.Errorf("parsePCIDeviceIDs() error = %v, wantError %v", err, tt.wantError)
+			}
+			if !reflect.DeepEqual(got, tt.want) {
+				t.Errorf("parsePCIDeviceIDs() = %v, want %v", got, tt.want)
+			}
+		})
+	}
+}

deployments/gpu_plugin/overlays/allowlist-arc/add-args.yaml

@@ -0,0 +1,12 @@
+apiVersion: apps/v1
+kind: DaemonSet
+metadata:
+  name: intel-gpu-plugin
+spec:
+  template:
+    spec:
+      containers:
+      - name: intel-gpu-plugin
+        args:
+        - "-v=4"
+        - "-allowlist-ids=0x56a6,0x56a5,0x56a1,0x56a0,0x5694,0x5693,0x5692,0x5691,0x5690,0x56b3,0x56b2,0x56a4,0x56a3,0x5697,0x5696,0x5695,0x56b1,0x56b0,0x56a2,0x56ba,0x56bc,0x56bd,0x56bb"

deployments/gpu_plugin/overlays/allowlist-arc/kustomization.yaml

@@ -0,0 +1,4 @@
+resources:
+  - ../../base
+patches:
+  - path: add-args.yaml

deployments/operator/crd/bases/deviceplugin.intel.com_gpudeviceplugins.yaml

@@ -55,6 +55,12 @@ spec:
           spec:
             description: GpuDevicePluginSpec defines the desired state of GpuDevicePlugin.
             properties:
+              allowListIDs:
+                description: |-
+                  AllowListIDs is a comma-separated list of PCI IDs of GPU devices that should only be advertised by the plugin.
+                  If not set, all devices are advertised.
+                  The list can contain IDs in the form of '0x1234,0x49a4,0x50b4'.
+                type: string
               enableMonitoring:
                 description: |-
                   EnableMonitoring enables the monitoring resource ('i915_monitoring')

pkg/apis/deviceplugin/v1/gpudeviceplugin_types.go

@@ -34,6 +34,11 @@ type GpuDevicePluginSpec struct {
 	// InitImage is a container image with tools (e.g., GPU NFD source hook) installed on each node.
 	InitImage string `json:"initImage,omitempty"`
 
+	// AllowListIDs is a comma-separated list of PCI IDs of GPU devices that should only be advertised by the plugin.
+	// If not set, all devices are advertised.
+	// The list can contain IDs in the form of '0x1234,0x49a4,0x50b4'.
+	AllowListIDs string `json:"allowListIDs,omitempty"`
+
 	// PreferredAllocationPolicy sets the mode of allocating GPU devices on a node.
 	// See documentation for detailed description of the policies. Only valid when SharedDevNum > 1 is set.
 	// +kubebuilder:validation:Enum=balanced;packed;none

pkg/apis/deviceplugin/v1/gpudeviceplugin_webhook.go

@@ -16,14 +16,20 @@ package v1
 
 import (
 	"fmt"
+	"regexp"
+	"strings"
 
 	ctrl "sigs.k8s.io/controller-runtime"
 
 	"github.com/intel/intel-device-plugins-for-kubernetes/pkg/controllers"
 )
 
+var pciIDRegex regexp.Regexp
+
 // SetupWebhookWithManager sets up a webhook for GpuDevicePlugin custom resources.
 func (r *GpuDevicePlugin) SetupWebhookWithManager(mgr ctrl.Manager) error {
+	pciIDRegex = *regexp.MustCompile(`^0x[0-9a-f]{4}$`)
+
 	return ctrl.NewWebhookManagedBy(mgr).
 		For(r).
 		WithDefaulter(&commonDevicePluginDefaulter{
@@ -44,5 +50,17 @@ func (r *GpuDevicePlugin) validatePlugin(ref *commonDevicePluginValidator) error
 		return fmt.Errorf("%w: PreferredAllocationPolicy is valid only when setting sharedDevNum > 1", errValidation)
 	}
 
+	if r.Spec.AllowListIDs != "" {
+		for id := range strings.SplitSeq(r.Spec.AllowListIDs, ",") {
+			if id == "" {
+				return fmt.Errorf("%w: Empty PCI Device ID in AllowListIDs", errValidation)
+			}
+
+			if !pciIDRegex.MatchString(id) {
+				return fmt.Errorf("%w: Invalid PCI Device ID: %s", errValidation, id)
+			}
+		}
+	}
+
 	return validatePluginImage(r.Spec.Image, ref.expectedImage, &ref.expectedVersion)
 }

pkg/controllers/gpu/controller.go

@@ -277,5 +277,9 @@ func getPodArgs(gdp *devicepluginv1.GpuDevicePlugin) []string {
 		args = append(args, "-allocation-policy", "none")
 	}
 
+	if gdp.Spec.AllowListIDs != "" {
+		args = append(args, "-allowlist-ids", gdp.Spec.AllowListIDs)
+	}
+
 	return args
 }