DRA scheduler: device taints and tolerations

Thanks to the tracker, the plugin sees all taints directly in the device
definition and can compare it against the tolerations of a request while
trying to find a device for the request.

When the feature is turnedd off, taints are ignored during scheduling.

Author: nojnhuh

pkg/scheduler/eventhandlers.go

@@ -33,6 +33,7 @@ import (
 	"k8s.io/client-go/tools/cache"
 	corev1helpers "k8s.io/component-helpers/scheduling/corev1"
 	corev1nodeaffinity "k8s.io/component-helpers/scheduling/corev1/nodeaffinity"
+	resourceslicetracker "k8s.io/dynamic-resource-allocation/resourceslice/tracker"
 	"k8s.io/klog/v2"
 	"k8s.io/kubernetes/pkg/features"
 	"k8s.io/kubernetes/pkg/scheduler/backend/queue"
@@ -366,6 +367,7 @@ func addAllEventHandlers(
 	informerFactory informers.SharedInformerFactory,
 	dynInformerFactory dynamicinformer.DynamicSharedInformerFactory,
 	resourceClaimCache *assumecache.AssumeCache,
+	resourceSliceTracker *resourceslicetracker.Tracker,
 	gvkMap map[framework.EventResource]framework.ActionType,
 ) error {
 	var (
@@ -555,7 +557,7 @@ func addAllEventHandlers(
 			}
 		case framework.ResourceSlice:
 			if utilfeature.DefaultFeatureGate.Enabled(features.DynamicResourceAllocation) {
-				if handlerRegistration, err = informerFactory.Resource().V1beta1().ResourceSlices().Informer().AddEventHandler(
+				if handlerRegistration, err = resourceSliceTracker.AddEventHandler(
 					buildEvtResHandler(at, framework.ResourceSlice),
 				); err != nil {
 					return err

pkg/scheduler/eventhandlers_test.go

@@ -28,12 +28,14 @@ import (
 	appsv1 "k8s.io/api/apps/v1"
 	batchv1 "k8s.io/api/batch/v1"
 	v1 "k8s.io/api/core/v1"
+	resourcealphaapi "k8s.io/api/resource/v1alpha3"
 	resourceapi "k8s.io/api/resource/v1beta1"
 	storagev1 "k8s.io/api/storage/v1"
 	"k8s.io/apimachinery/pkg/api/resource"
 	"k8s.io/apimachinery/pkg/util/sets"
 	utilfeature "k8s.io/apiserver/pkg/util/feature"
 	featuregatetesting "k8s.io/component-base/featuregate/testing"
+	resourceslicetracker "k8s.io/dynamic-resource-allocation/resourceslice/tracker"
 	"k8s.io/klog/v2"
 	"k8s.io/klog/v2/ktesting"
 
@@ -395,6 +397,7 @@ func TestAddAllEventHandlers(t *testing.T) {
 		name                   string
 		gvkMap                 map[framework.EventResource]framework.ActionType
 		enableDRA              bool
+		enableDRADeviceTaints  bool
 		expectStaticInformers  map[reflect.Type]bool
 		expectDynamicInformers map[schema.GroupVersionResource]bool
 	}{
@@ -423,7 +426,7 @@ func TestAddAllEventHandlers(t *testing.T) {
 			expectDynamicInformers: map[schema.GroupVersionResource]bool{},
 		},
 		{
-			name: "all DRA events enabled",
+			name: "core DRA events enabled",
 			gvkMap: map[framework.EventResource]framework.ActionType{
 				framework.ResourceClaim: framework.Add,
 				framework.ResourceSlice: framework.Add,
@@ -440,6 +443,26 @@ func TestAddAllEventHandlers(t *testing.T) {
 			},
 			expectDynamicInformers: map[schema.GroupVersionResource]bool{},
 		},
+		{
+			name: "all DRA events enabled",
+			gvkMap: map[framework.EventResource]framework.ActionType{
+				framework.ResourceClaim: framework.Add,
+				framework.ResourceSlice: framework.Add,
+				framework.DeviceClass:   framework.Add,
+			},
+			enableDRA:             true,
+			enableDRADeviceTaints: true,
+			expectStaticInformers: map[reflect.Type]bool{
+				reflect.TypeOf(&v1.Pod{}):                           true,
+				reflect.TypeOf(&v1.Node{}):                          true,
+				reflect.TypeOf(&v1.Namespace{}):                     true,
+				reflect.TypeOf(&resourceapi.ResourceClaim{}):        true,
+				reflect.TypeOf(&resourceapi.ResourceSlice{}):        true,
+				reflect.TypeOf(&resourcealphaapi.DeviceTaintRule{}): true,
+				reflect.TypeOf(&resourceapi.DeviceClass{}):          true,
+			},
+			expectDynamicInformers: map[schema.GroupVersionResource]bool{},
+		},
 		{
 			name: "add GVKs handlers defined in framework dynamically",
 			gvkMap: map[framework.EventResource]framework.ActionType{
@@ -499,6 +522,7 @@ func TestAddAllEventHandlers(t *testing.T) {
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
 			featuregatetesting.SetFeatureGateDuringTest(t, utilfeature.DefaultFeatureGate, features.DynamicResourceAllocation, tt.enableDRA)
+			featuregatetesting.SetFeatureGateDuringTest(t, utilfeature.DefaultFeatureGate, features.DRADeviceTaints, tt.enableDRADeviceTaints)
 
 			logger, ctx := ktesting.NewTestContext(t)
 			ctx, cancel := context.WithCancel(ctx)
@@ -515,12 +539,27 @@ func TestAddAllEventHandlers(t *testing.T) {
 			dynclient := dyfake.NewSimpleDynamicClient(scheme)
 			dynInformerFactory := dynamicinformer.NewDynamicSharedInformerFactory(dynclient, 0)
 			var resourceClaimCache *assumecache.AssumeCache
+			var resourceSliceTracker *resourceslicetracker.Tracker
 			if utilfeature.DefaultFeatureGate.Enabled(features.DynamicResourceAllocation) {
 				resourceClaimInformer := informerFactory.Resource().V1beta1().ResourceClaims().Informer()
 				resourceClaimCache = assumecache.NewAssumeCache(logger, resourceClaimInformer, "ResourceClaim", "", nil)
+				var err error
+				opts := resourceslicetracker.Options{
+					EnableDeviceTaints: utilfeature.DefaultFeatureGate.Enabled(features.DRADeviceTaints),
+					SliceInformer:      informerFactory.Resource().V1beta1().ResourceSlices(),
+				}
+				if opts.EnableDeviceTaints {
+					opts.TaintInformer = informerFactory.Resource().V1alpha3().DeviceTaintRules()
+					opts.ClassInformer = informerFactory.Resource().V1beta1().DeviceClasses()
+
+				}
+				resourceSliceTracker, err = resourceslicetracker.StartTracker(ctx, opts)
+				if err != nil {
+					t.Fatalf("couldn't start resource slice tracker: %v", err)
+				}
 			}
 
-			if err := addAllEventHandlers(&testSched, informerFactory, dynInformerFactory, resourceClaimCache, tt.gvkMap); err != nil {
+			if err := addAllEventHandlers(&testSched, informerFactory, dynInformerFactory, resourceClaimCache, resourceSliceTracker, tt.gvkMap); err != nil {
 				t.Fatalf("Add event handlers failed, error = %v", err)
 			}
 

pkg/scheduler/framework/autoscaler_contract/lister_contract_test.go

@@ -72,7 +72,7 @@ func (c *shareListerContract) StorageInfos() framework.StorageInfoLister {
 
 type resourceSliceListerContract struct{}
 
-func (c *resourceSliceListerContract) List() ([]*resourceapi.ResourceSlice, error) {
+func (c *resourceSliceListerContract) ListWithDeviceTaintRules() ([]*resourceapi.ResourceSlice, error) {
 	return nil, nil
 }
 

pkg/scheduler/framework/listers.go

@@ -50,8 +50,13 @@ type SharedLister interface {
 
 // ResourceSliceLister can be used to obtain ResourceSlices.
 type ResourceSliceLister interface {
-	// List returns a list of all ResourceSlices.
-	List() ([]*resourceapi.ResourceSlice, error)
+	// ListWithDeviceTaintRules returns a list of all ResourceSlices with DeviceTaintRules applied
+	// if the DRADeviceTaints feature is enabled, otherwise without them.
+	//
+	// k8s.io/dynamic-resource-allocation/resourceslice/tracker provides an implementation
+	// of the necessary logic. That tracker can be instantiated as a replacement for
+	// a normal ResourceSlice informer and provides a ListPatchedResourceSlices method.
+	ListWithDeviceTaintRules() ([]*resourceapi.ResourceSlice, error)
 }
 
 // DeviceClassLister can be used to obtain DeviceClasses.

pkg/scheduler/framework/plugins/dynamicresources/dra_manager.go

@@ -27,6 +27,7 @@ import (
 	"k8s.io/apimachinery/pkg/util/sets"
 	"k8s.io/client-go/informers"
 	resourcelisters "k8s.io/client-go/listers/resource/v1beta1"
+	resourceslicetracker "k8s.io/dynamic-resource-allocation/resourceslice/tracker"
 	"k8s.io/dynamic-resource-allocation/structured"
 	"k8s.io/klog/v2"
 	"k8s.io/kubernetes/pkg/scheduler/framework"
@@ -44,16 +45,17 @@ type DefaultDRAManager struct {
 	deviceClassLister    *deviceClassLister
 }
 
-func NewDRAManager(ctx context.Context, claimsCache *assumecache.AssumeCache, informerFactory informers.SharedInformerFactory) *DefaultDRAManager {
+func NewDRAManager(ctx context.Context, claimsCache *assumecache.AssumeCache, resourceSliceTracker *resourceslicetracker.Tracker, informerFactory informers.SharedInformerFactory) *DefaultDRAManager {
 	logger := klog.FromContext(ctx)
+
 	manager := &DefaultDRAManager{
 		resourceClaimTracker: &claimTracker{
 			cache:               claimsCache,
 			inFlightAllocations: &sync.Map{},
 			allocatedDevices:    newAllocatedDevices(logger),
 			logger:              logger,
 		},
-		resourceSliceLister: &resourceSliceLister{sliceLister: informerFactory.Resource().V1beta1().ResourceSlices().Lister()},
+		resourceSliceLister: &resourceSliceLister{tracker: resourceSliceTracker},
 		deviceClassLister:   &deviceClassLister{classLister: informerFactory.Resource().V1beta1().DeviceClasses().Lister()},
 	}
 
@@ -79,11 +81,11 @@ func (s *DefaultDRAManager) DeviceClasses() framework.DeviceClassLister {
 var _ framework.ResourceSliceLister = &resourceSliceLister{}
 
 type resourceSliceLister struct {
-	sliceLister resourcelisters.ResourceSliceLister
+	tracker *resourceslicetracker.Tracker
 }
 
-func (l *resourceSliceLister) List() ([]*resourceapi.ResourceSlice, error) {
-	return l.sliceLister.List(labels.Everything())
+func (l *resourceSliceLister) ListWithDeviceTaintRules() ([]*resourceapi.ResourceSlice, error) {
+	return l.tracker.ListPatchedResourceSlices()
 }
 
 var _ framework.DeviceClassLister = &deviceClassLister{}

pkg/scheduler/framework/plugins/dynamicresources/dynamicresources.go

@@ -106,6 +106,7 @@ type DynamicResources struct {
 	enableAdminAccess         bool
 	enablePrioritizedList     bool
 	enableSchedulingQueueHint bool
+	enableDeviceTaints        bool
 
 	fh         framework.Handle
 	clientset  kubernetes.Interface
@@ -123,6 +124,7 @@ func New(ctx context.Context, plArgs runtime.Object, fh framework.Handle, fts fe
 	pl := &DynamicResources{
 		enabled:                   true,
 		enableAdminAccess:         fts.EnableDRAAdminAccess,
+		enableDeviceTaints:        fts.EnableDRADeviceTaints,
 		enablePrioritizedList:     fts.EnableDRAPrioritizedList,
 		enableSchedulingQueueHint: fts.EnableSchedulingQueueHint,
 
@@ -448,11 +450,11 @@ func (pl *DynamicResources) PreFilter(ctx context.Context, state *framework.Cycl
 		if err != nil {
 			return nil, statusError(logger, err)
 		}
-		slices, err := pl.draManager.ResourceSlices().List()
+		slices, err := pl.draManager.ResourceSlices().ListWithDeviceTaintRules()
 		if err != nil {
 			return nil, statusError(logger, err)
 		}
-		allocator, err := structured.NewAllocator(ctx, pl.enableAdminAccess, pl.enablePrioritizedList, allocateClaims, allAllocatedDevices, pl.draManager.DeviceClasses(), slices, pl.celCache)
+		allocator, err := structured.NewAllocator(ctx, pl.enableAdminAccess, pl.enablePrioritizedList, pl.enableDeviceTaints, allocateClaims, allAllocatedDevices, pl.draManager.DeviceClasses(), slices, pl.celCache)
 		if err != nil {
 			return nil, statusError(logger, err)
 		}

pkg/scheduler/framework/plugins/dynamicresources/dynamicresources_test.go

@@ -38,6 +38,7 @@ import (
 	"k8s.io/client-go/informers"
 	"k8s.io/client-go/kubernetes/fake"
 	cgotesting "k8s.io/client-go/testing"
+	resourceslicetracker "k8s.io/dynamic-resource-allocation/resourceslice/tracker"
 	"k8s.io/kubernetes/pkg/scheduler/framework"
 	"k8s.io/kubernetes/pkg/scheduler/framework/plugins/feature"
 	"k8s.io/kubernetes/pkg/scheduler/framework/runtime"
@@ -183,8 +184,22 @@ var (
 	otherAllocatedClaim = st.FromResourceClaim(otherClaim).
 				Allocation(allocationResult).
 				Obj()
+
+	deviceTaint = resourceapi.DeviceTaint{
+		Key:    "taint-key",
+		Value:  "taint-value",
+		Effect: resourceapi.DeviceTaintEffectNoSchedule,
+	}
 )
 
+func taintDevices(slice *resourceapi.ResourceSlice) *resourceapi.ResourceSlice {
+	slice = slice.DeepCopy()
+	for i := range slice.Spec.Devices {
+		slice.Spec.Devices[i].Basic.Taints = append(slice.Spec.Devices[i].Basic.Taints, deviceTaint)
+	}
+	return slice
+}
+
 func reserve(claim *resourceapi.ResourceClaim, pod *v1.Pod) *resourceapi.ResourceClaim {
 	return st.FromResourceClaim(claim).
 		ReservedForPod(pod.Name, types.UID(pod.UID)).
@@ -343,6 +358,7 @@ func TestPlugin(t *testing.T) {
 		disableDRA bool
 
 		enableDRAPrioritizedList bool
+		enableDRADeviceTaints    bool
 	}{
 		"empty": {
 			pod: st.MakePod().Name("foo").Namespace("default").Obj(),
@@ -604,6 +620,56 @@ func TestPlugin(t *testing.T) {
 			},
 		},
 
+		// The two test cases for device tainting only need to cover
+		// whether the feature gate is passed through to the allocator
+		// correctly. The actual logic around device taints and allocation
+		// is in the allocator.
+		"tainted-device-disabled": {
+			enableDRADeviceTaints: false,
+			pod:                   podWithClaimName,
+			claims:                []*resourceapi.ResourceClaim{pendingClaim},
+			classes:               []*resourceapi.DeviceClass{deviceClass},
+			objs:                  []apiruntime.Object{taintDevices(workerNodeSlice)},
+			want: want{
+				reserve: result{
+					inFlightClaim: allocatedClaim,
+				},
+				prebind: result{
+					assumedClaim: reserve(allocatedClaim, podWithClaimName),
+					changes: change{
+						claim: func(claim *resourceapi.ResourceClaim) *resourceapi.ResourceClaim {
+							if claim.Name == claimName {
+								claim = claim.DeepCopy()
+								claim.Finalizers = allocatedClaim.Finalizers
+								claim.Status = inUseClaim.Status
+							}
+							return claim
+						},
+					},
+				},
+				postbind: result{
+					assumedClaim: reserve(allocatedClaim, podWithClaimName),
+				},
+			},
+		},
+		"tainted-device-enabled": {
+			enableDRADeviceTaints: true,
+			pod:                   podWithClaimName,
+			claims:                []*resourceapi.ResourceClaim{pendingClaim},
+			classes:               []*resourceapi.DeviceClass{deviceClass},
+			objs:                  []apiruntime.Object{taintDevices(workerNodeSlice)},
+			want: want{
+				filter: perNodeResult{
+					workerNode.Name: {
+						status: framework.NewStatus(framework.UnschedulableAndUnresolvable, `cannot allocate all claims`),
+					},
+				},
+				postfilter: result{
+					status: framework.NewStatus(framework.Unschedulable, `still not schedulable`),
+				},
+			},
+		},
+
 		"request-admin-access-with-DRAAdminAccess-featuregate": {
 			// When the DRAAdminAccess feature gate is enabled,
 			// Because the pending claim asks for admin access,
@@ -920,6 +986,7 @@ func TestPlugin(t *testing.T) {
 			}
 			features := feature.Features{
 				EnableDRAAdminAccess:            tc.enableDRAAdminAccess,
+				EnableDRADeviceTaints:           tc.enableDRADeviceTaints,
 				EnableDynamicResourceAllocation: !tc.disableDRA,
 				EnableDRAPrioritizedList:        tc.enableDRAPrioritizedList,
 			}
@@ -1189,7 +1256,16 @@ func setup(t *testing.T, nodes []*v1.Node, claims []*resourceapi.ResourceClaim,
 	tc.client.PrependReactor("*", "*", reactor)
 
 	tc.informerFactory = informers.NewSharedInformerFactory(tc.client, 0)
-	tc.draManager = NewDRAManager(tCtx, assumecache.NewAssumeCache(tCtx.Logger(), tc.informerFactory.Resource().V1beta1().ResourceClaims().Informer(), "resource claim", "", nil), tc.informerFactory)
+	resourceSliceTrackerOpts := resourceslicetracker.Options{
+		EnableDeviceTaints: true,
+		SliceInformer:      tc.informerFactory.Resource().V1beta1().ResourceSlices(),
+		TaintInformer:      tc.informerFactory.Resource().V1alpha3().DeviceTaintRules(),
+		ClassInformer:      tc.informerFactory.Resource().V1beta1().DeviceClasses(),
+		KubeClient:         tc.client,
+	}
+	resourceSliceTracker, err := resourceslicetracker.StartTracker(tCtx, resourceSliceTrackerOpts)
+	require.NoError(t, err, "couldn't start resource slice tracker")
+	tc.draManager = NewDRAManager(tCtx, assumecache.NewAssumeCache(tCtx.Logger(), tc.informerFactory.Resource().V1beta1().ResourceClaims().Informer(), "resource claim", "", nil), resourceSliceTracker, tc.informerFactory)
 	opts := []runtime.Option{
 		runtime.WithClientSet(tc.client),
 		runtime.WithInformerFactory(tc.informerFactory),

pkg/scheduler/framework/plugins/feature/feature.go

@@ -22,6 +22,7 @@ package feature
 type Features struct {
 	EnableDRAPrioritizedList                     bool
 	EnableDRAAdminAccess                         bool
+	EnableDRADeviceTaints                        bool
 	EnableDynamicResourceAllocation              bool
 	EnableVolumeCapacityPriority                 bool
 	EnableVolumeAttributesClass                  bool

pkg/scheduler/framework/plugins/registry.go

@@ -48,6 +48,7 @@ func NewInTreeRegistry() runtime.Registry {
 	fts := plfeature.Features{
 		EnableDRAPrioritizedList:                     feature.DefaultFeatureGate.Enabled(features.DRAPrioritizedList),
 		EnableDRAAdminAccess:                         feature.DefaultFeatureGate.Enabled(features.DRAAdminAccess),
+		EnableDRADeviceTaints:                        feature.DefaultFeatureGate.Enabled(features.DRADeviceTaints),
 		EnableDynamicResourceAllocation:              feature.DefaultFeatureGate.Enabled(features.DynamicResourceAllocation),
 		EnableVolumeCapacityPriority:                 feature.DefaultFeatureGate.Enabled(features.VolumeCapacityPriority),
 		EnableVolumeAttributesClass:                  feature.DefaultFeatureGate.Enabled(features.VolumeAttributesClass),