gpu: add pci device id whitelist support

By defining whitelisted 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/gpu_plugin.go

@@ -69,6 +69,7 @@ const (
 
 type cliOptions struct {
 	preferredAllocationPolicy string
+	whitelistIDs              string
 	sharedDevNum              int
 	temperatureLimit          int
 	enableMonitoring          bool
@@ -204,6 +205,30 @@ func packedPolicy(req *pluginapi.ContainerPreferredAllocationRequest) []string {
 	return deviceIds
 }
 
+func parsePCIDeviceIDs(whitelist string) ([]string, error) {
+	var deviceIDs []string
+
+	r := regexp.MustCompile(`^0x[0-9a-f]{4}$`)
+
+	for id := range strings.SplitSeq(whitelist, ",") {
+		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 +610,23 @@ func (dp *devicePlugin) filterOutInvalidCards(files []fs.DirEntry) []fs.DirEntry
 			continue
 		}
 
+		// Skip if the device is not in the whitelist.
+		if len(dp.options.whitelistIDs) > 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.whitelistIDs, pciID) {
+				klog.V(4).Infof("Skipping device %s (%s), not in whitelist: %s",
+					f.Name(), pciID, dp.options.whitelistIDs)
+
+				continue
+			}
+		}
+
 		filtered = append(filtered, f)
 	}
 
@@ -723,6 +765,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.whitelistIDs, "whitelist-ids", "", "comma-separated list of device IDs to whitelist (e.g. 0x49c5,0x49c6)")
 	flag.Parse()
 
 	if opts.sharedDevNum < 1 {
@@ -736,6 +779,18 @@ func main() {
 		os.Exit(1)
 	}
 
+	if opts.whitelistIDs != "" {
+		if whiteListIDs, err := parsePCIDeviceIDs(opts.whitelistIDs); err != nil {
+			klog.Error("Failed to parse whitelist-ids: ", err)
+
+			os.Exit(1)
+		} else {
+			klog.V(2).Infof("Whitelisted device IDs: %q", whiteListIDs)
+
+			opts.whitelistIDs = strings.Join(whiteListIDs, ",")
+		}
+	}
+
 	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 whitelisted",
+			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, whitelistIDs: "0x1234"},
+			expectedXeDevs:       1,
+			expectedXeMonitors:   1,
+			expectedI915Devs:     0,
+			expectedI915Monitors: 0,
+		},
+		{
+			name:      "three devices with two whitelisted",
+			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, whitelistIDs: "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/whitelist-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"
+        - "-whitelist-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/whitelist-arc/kustomization.yaml

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

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

@@ -132,6 +132,12 @@ spec:
                       type: string
                   type: object
                 type: array
+              whiteListIDs:
+                description: |-
+                  WhiteListIDs 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
             type: object
           status:
             description: GpuDevicePluginStatus defines the observed state of GpuDevicePlugin.

pkg/apis/deviceplugin/v1/gpudeviceplugin_types.go

@@ -53,6 +53,11 @@ type GpuDevicePluginSpec struct {
 	// EnableMonitoring enables the monitoring resource ('i915_monitoring')
 	// which gives access to all GPU devices on given node. Typically used with Intel XPU-Manager.
 	EnableMonitoring bool `json:"enableMonitoring,omitempty"`
+
+	// WhiteListIDs 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'.
+	WhiteListIDs string `json:"whiteListIDs,omitempty"`
 }
 
 // GpuDevicePluginStatus defines the observed state of GpuDevicePlugin.

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.WhiteListIDs != "" {
+		for id := range strings.SplitSeq(r.Spec.WhiteListIDs, ",") {
+			if id == "" {
+				return fmt.Errorf("%w: Empty PCI Device ID in WhiteListIDs", 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.WhiteListIDs != "" {
+		args = append(args, "-whitelist-ids", gdp.Spec.WhiteListIDs)
+	}
+
 	return args
 }