Merge pull request kubernetes#127277 from pohly/dra-structured-performance

kube-scheduler: enhance performance for DRA structured parameters

Author: k8s-ci-robot

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

@@ -0,0 +1,175 @@
+/*
+Copyright 2024 The Kubernetes Authors.
+
+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 dynamicresources
+
+import (
+	"sync"
+
+	resourceapi "k8s.io/api/resource/v1alpha3"
+	"k8s.io/apimachinery/pkg/util/sets"
+	"k8s.io/client-go/tools/cache"
+	"k8s.io/dynamic-resource-allocation/structured"
+	"k8s.io/klog/v2"
+	schedutil "k8s.io/kubernetes/pkg/scheduler/util"
+	"k8s.io/utils/ptr"
+)
+
+// foreachAllocatedDevice invokes the provided callback for each
+// device in the claim's allocation result which was allocated
+// exclusively for the claim.
+//
+// Devices allocated with admin access can be shared with other
+// claims and are skipped without invoking the callback.
+//
+// foreachAllocatedDevice does nothing if the claim is not allocated.
+func foreachAllocatedDevice(claim *resourceapi.ResourceClaim, cb func(deviceID structured.DeviceID)) {
+	if claim.Status.Allocation == nil {
+		return
+	}
+	for _, result := range claim.Status.Allocation.Devices.Results {
+		// Kubernetes 1.31 did not set this, 1.32 always does.
+		// Supporting 1.31 is not worth the additional code that
+		// would have to be written (= looking up in request) because
+		// it is extremely unlikely that there really is a result
+		// that still exists in a cluster from 1.31 where this matters.
+		if ptr.Deref(result.AdminAccess, false) {
+			// Is not considered as allocated.
+			continue
+		}
+		deviceID := structured.MakeDeviceID(result.Driver, result.Pool, result.Device)
+
+		// None of the users of this helper need to abort iterating,
+		// therefore it's not supported as it only would add overhead.
+		cb(deviceID)
+	}
+}
+
+// allocatedDevices reacts to events in a cache and maintains a set of all allocated devices.
+// This is cheaper than repeatedly calling List, making strings unique, and building the set
+// each time PreFilter is called.
+//
+// All methods are thread-safe. Get returns a cloned set.
+type allocatedDevices struct {
+	logger klog.Logger
+
+	mutex sync.RWMutex
+	ids   sets.Set[structured.DeviceID]
+}
+
+func newAllocatedDevices(logger klog.Logger) *allocatedDevices {
+	return &allocatedDevices{
+		logger: logger,
+		ids:    sets.New[structured.DeviceID](),
+	}
+}
+
+func (a *allocatedDevices) Get() sets.Set[structured.DeviceID] {
+	a.mutex.RLock()
+	defer a.mutex.RUnlock()
+
+	return a.ids.Clone()
+}
+
+func (a *allocatedDevices) handlers() cache.ResourceEventHandler {
+	return cache.ResourceEventHandlerFuncs{
+		AddFunc:    a.onAdd,
+		UpdateFunc: a.onUpdate,
+		DeleteFunc: a.onDelete,
+	}
+}
+
+func (a *allocatedDevices) onAdd(obj any) {
+	claim, _, err := schedutil.As[*resourceapi.ResourceClaim](obj, nil)
+	if err != nil {
+		// Shouldn't happen.
+		a.logger.Error(err, "unexpected object in allocatedDevices.onAdd")
+		return
+	}
+
+	if claim.Status.Allocation != nil {
+		a.addDevices(claim)
+	}
+}
+
+func (a *allocatedDevices) onUpdate(oldObj, newObj any) {
+	originalClaim, modifiedClaim, err := schedutil.As[*resourceapi.ResourceClaim](oldObj, newObj)
+	if err != nil {
+		// Shouldn't happen.
+		a.logger.Error(err, "unexpected object in allocatedDevices.onUpdate")
+		return
+	}
+
+	switch {
+	case originalClaim.Status.Allocation == nil && modifiedClaim.Status.Allocation != nil:
+		a.addDevices(modifiedClaim)
+	case originalClaim.Status.Allocation != nil && modifiedClaim.Status.Allocation == nil:
+		a.removeDevices(originalClaim)
+	default:
+		// Nothing to do. Either both nil or both non-nil, in which case the content
+		// also must be the same (immutable!).
+	}
+}
+
+func (a *allocatedDevices) onDelete(obj any) {
+	claim, _, err := schedutil.As[*resourceapi.ResourceClaim](obj, nil)
+	if err != nil {
+		// Shouldn't happen.
+		a.logger.Error(err, "unexpected object in allocatedDevices.onDelete")
+		return
+	}
+
+	a.removeDevices(claim)
+}
+
+func (a *allocatedDevices) addDevices(claim *resourceapi.ResourceClaim) {
+	if claim.Status.Allocation == nil {
+		return
+	}
+	// Locking of the mutex gets minimized by pre-computing what needs to be done
+	// without holding the lock.
+	deviceIDs := make([]structured.DeviceID, 0, 20)
+	foreachAllocatedDevice(claim, func(deviceID structured.DeviceID) {
+		a.logger.V(6).Info("Observed device allocation", "device", deviceID, "claim", klog.KObj(claim))
+		deviceIDs = append(deviceIDs, deviceID)
+	})
+
+	a.mutex.Lock()
+	defer a.mutex.Unlock()
+	for _, deviceID := range deviceIDs {
+		a.ids.Insert(deviceID)
+	}
+}
+
+func (a *allocatedDevices) removeDevices(claim *resourceapi.ResourceClaim) {
+	if claim.Status.Allocation == nil {
+		return
+	}
+
+	// Locking of the mutex gets minimized by pre-computing what needs to be done
+	// without holding the lock.
+	deviceIDs := make([]structured.DeviceID, 0, 20)
+	foreachAllocatedDevice(claim, func(deviceID structured.DeviceID) {
+		a.logger.V(6).Info("Observed device deallocation", "device", deviceID, "claim", klog.KObj(claim))
+		deviceIDs = append(deviceIDs, deviceID)
+	})
+
+	a.mutex.Lock()
+	defer a.mutex.Unlock()
+	for _, deviceID := range deviceIDs {
+		a.ids.Delete(deviceID)
+	}
+}

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

@@ -30,13 +30,15 @@ import (
 	apiequality "k8s.io/apimachinery/pkg/api/equality"
 	apierrors "k8s.io/apimachinery/pkg/api/errors"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+	"k8s.io/apimachinery/pkg/labels"
 	"k8s.io/apimachinery/pkg/runtime"
 	"k8s.io/apimachinery/pkg/types"
 	"k8s.io/apimachinery/pkg/util/sets"
 	"k8s.io/client-go/kubernetes"
 	resourcelisters "k8s.io/client-go/listers/resource/v1alpha3"
 	"k8s.io/client-go/util/retry"
 	"k8s.io/component-helpers/scheduling/corev1/nodeaffinity"
+	"k8s.io/dynamic-resource-allocation/cel"
 	"k8s.io/dynamic-resource-allocation/resourceclaim"
 	"k8s.io/dynamic-resource-allocation/structured"
 	"k8s.io/klog/v2"
@@ -85,7 +87,7 @@ type stateData struct {
 	informationsForClaim []informationForClaim
 
 	// nodeAllocations caches the result of Filter for the nodes.
-	nodeAllocations map[string][]*resourceapi.AllocationResult
+	nodeAllocations map[string][]resourceapi.AllocationResult
 }
 
 func (d *stateData) Clone() framework.StateData {
@@ -106,10 +108,12 @@ type DynamicResources struct {
 	enableAdminAccess         bool
 	enableSchedulingQueueHint bool
 
-	fh          framework.Handle
-	clientset   kubernetes.Interface
-	classLister resourcelisters.DeviceClassLister
-	sliceLister resourcelisters.ResourceSliceLister
+	fh               framework.Handle
+	clientset        kubernetes.Interface
+	classLister      resourcelisters.DeviceClassLister
+	sliceLister      resourcelisters.ResourceSliceLister
+	celCache         *cel.Cache
+	allocatedDevices *allocatedDevices
 
 	// claimAssumeCache enables temporarily storing a newer claim object
 	// while the scheduler has allocated it and the corresponding object
@@ -174,6 +178,7 @@ func New(ctx context.Context, plArgs runtime.Object, fh framework.Handle, fts fe
 		return &DynamicResources{}, nil
 	}
 
+	logger := klog.FromContext(ctx)
 	pl := &DynamicResources{
 		enabled:                   true,
 		enableAdminAccess:         fts.EnableDRAAdminAccess,
@@ -184,8 +189,19 @@ func New(ctx context.Context, plArgs runtime.Object, fh framework.Handle, fts fe
 		classLister:      fh.SharedInformerFactory().Resource().V1alpha3().DeviceClasses().Lister(),
 		sliceLister:      fh.SharedInformerFactory().Resource().V1alpha3().ResourceSlices().Lister(),
 		claimAssumeCache: fh.ResourceClaimCache(),
+
+		// This is a LRU cache for compiled CEL expressions. The most
+		// recent 10 of them get reused across different scheduling
+		// cycles.
+		celCache: cel.NewCache(10),
+
+		allocatedDevices: newAllocatedDevices(logger),
 	}
 
+	// Reacting to events is more efficient than iterating over the list
+	// repeatedly in PreFilter.
+	pl.claimAssumeCache.AddEventHandler(pl.allocatedDevices.handlers())
+
 	return pl, nil
 }
 
@@ -527,39 +543,41 @@ func (pl *DynamicResources) PreFilter(ctx context.Context, state *framework.Cycl
 		// expensive, we may have to maintain and update state more
 		// persistently.
 		//
-		// Claims are treated as "allocated" if they are in the assume cache
-		// or currently their allocation is in-flight.
-		allocator, err := structured.NewAllocator(ctx, pl.enableAdminAccess, allocateClaims, &claimListerForAssumeCache{assumeCache: pl.claimAssumeCache, inFlightAllocations: &pl.inFlightAllocations}, pl.classLister, pl.sliceLister)
+		// Claims (and thus their devices) are treated as "allocated" if they are in the assume cache
+		// or currently their allocation is in-flight. This does not change
+		// during filtering, so we can determine that once.
+		allAllocatedDevices := pl.listAllAllocatedDevices(logger)
+		slices, err := pl.sliceLister.List(labels.Everything())
+		if err != nil {
+			return nil, statusError(logger, err)
+		}
+		allocator, err := structured.NewAllocator(ctx, pl.enableAdminAccess, allocateClaims, allAllocatedDevices, pl.classLister, slices, pl.celCache)
 		if err != nil {
 			return nil, statusError(logger, err)
 		}
 		s.allocator = allocator
-		s.nodeAllocations = make(map[string][]*resourceapi.AllocationResult)
+		s.nodeAllocations = make(map[string][]resourceapi.AllocationResult)
 	}
 
 	s.claims = claims
 	return nil, nil
 }
 
-type claimListerForAssumeCache struct {
-	assumeCache         *assumecache.AssumeCache
-	inFlightAllocations *sync.Map
-}
-
-func (cl *claimListerForAssumeCache) ListAllAllocated() ([]*resourceapi.ResourceClaim, error) {
-	// Probably not worth adding an index for?
-	objs := cl.assumeCache.List(nil)
-	allocated := make([]*resourceapi.ResourceClaim, 0, len(objs))
-	for _, obj := range objs {
-		claim := obj.(*resourceapi.ResourceClaim)
-		if obj, ok := cl.inFlightAllocations.Load(claim.UID); ok {
-			claim = obj.(*resourceapi.ResourceClaim)
-		}
-		if claim.Status.Allocation != nil {
-			allocated = append(allocated, claim)
-		}
-	}
-	return allocated, nil
+func (pl *DynamicResources) listAllAllocatedDevices(logger klog.Logger) sets.Set[structured.DeviceID] {
+	// Start with a fresh set that matches the current known state of the
+	// world according to the informers.
+	allocated := pl.allocatedDevices.Get()
+
+	// Whatever is in flight also has to be checked.
+	pl.inFlightAllocations.Range(func(key, value any) bool {
+		claim := value.(*resourceapi.ResourceClaim)
+		foreachAllocatedDevice(claim, func(deviceID structured.DeviceID) {
+			logger.V(6).Info("Device is in flight for allocation", "device", deviceID, "claim", klog.KObj(claim))
+			allocated.Insert(deviceID)
+		})
+		return true
+	})
+	return allocated
 }
 
 // PreFilterExtensions returns prefilter extensions, pod add and remove.
@@ -615,7 +633,7 @@ func (pl *DynamicResources) Filter(ctx context.Context, cs *framework.CycleState
 	}
 
 	// Use allocator to check the node and cache the result in case that the node is picked.
-	var allocations []*resourceapi.AllocationResult
+	var allocations []resourceapi.AllocationResult
 	if state.allocator != nil {
 		allocCtx := ctx
 		if loggerV := logger.V(5); loggerV.Enabled() {
@@ -763,7 +781,7 @@ func (pl *DynamicResources) Reserve(ctx context.Context, cs *framework.CycleStat
 			if index < 0 {
 				return statusError(logger, fmt.Errorf("internal error, claim %s with allocation not found", claim.Name))
 			}
-			allocation := allocations[i]
+			allocation := &allocations[i]
 			state.informationsForClaim[index].allocation = allocation
 
 			// Strictly speaking, we don't need to store the full modified object.

staging/src/k8s.io/dynamic-resource-allocation/api/conversion.go

@@ -0,0 +1,47 @@
+/*
+Copyright 2024 The Kubernetes Authors.
+
+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 api
+
+import (
+	"unique"
+
+	conversion "k8s.io/apimachinery/pkg/conversion"
+	"k8s.io/apimachinery/pkg/runtime"
+)
+
+var (
+	localSchemeBuilder runtime.SchemeBuilder
+	AddToScheme        = localSchemeBuilder.AddToScheme
+)
+
+func Convert_api_UniqueString_To_string(in *UniqueString, out *string, s conversion.Scope) error {
+	if *in == NullUniqueString {
+		*out = ""
+		return nil
+	}
+	*out = in.String()
+	return nil
+}
+
+func Convert_string_To_api_UniqueString(in *string, out *UniqueString, s conversion.Scope) error {
+	if *in == "" {
+		*out = NullUniqueString
+		return nil
+	}
+	*out = UniqueString(unique.Make(*in))
+	return nil
+}

staging/src/k8s.io/dynamic-resource-allocation/api/doc.go

@@ -0,0 +1,22 @@
+/*
+Copyright 2024 The Kubernetes Authors.
+
+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 api contains a variant of the API where strings are unique. These
+// unique strings are faster to compare and more efficient when used as key in
+// a map.
+//
+// +k8s:conversion-gen=k8s.io/api/resource/v1alpha3
+package api