KEP 5075: implement scheduler

Signed-off-by: Sunyanan Choochotkaew <[email protected]>

Author: sunya-ch

pkg/controller/devicetainteviction/device_taint_eviction.go

@@ -36,6 +36,7 @@ import (
 	"k8s.io/apimachinery/pkg/util/diff"
 	utilruntime "k8s.io/apimachinery/pkg/util/runtime"
 	"k8s.io/apimachinery/pkg/util/sets"
+	utilfeature "k8s.io/apiserver/pkg/util/feature"
 	coreinformers "k8s.io/client-go/informers/core/v1"
 	resourceinformers "k8s.io/client-go/informers/resource/v1"
 	resourcealphainformers "k8s.io/client-go/informers/resource/v1alpha3"
@@ -51,6 +52,7 @@ import (
 	apipod "k8s.io/kubernetes/pkg/api/v1/pod"
 	"k8s.io/kubernetes/pkg/controller/devicetainteviction/metrics"
 	"k8s.io/kubernetes/pkg/controller/tainteviction"
+	"k8s.io/kubernetes/pkg/features"
 	utilpod "k8s.io/kubernetes/pkg/util/pod"
 )
 
@@ -458,11 +460,12 @@ func (tc *Controller) Run(ctx context.Context) error {
 	tc.haveSynced = append(tc.haveSynced, podHandler.HasSynced)
 
 	opts := resourceslicetracker.Options{
-		EnableDeviceTaints: true,
-		SliceInformer:      tc.sliceInformer,
-		TaintInformer:      tc.taintInformer,
-		ClassInformer:      tc.classInformer,
-		KubeClient:         tc.client,
+		EnableDeviceTaints:       true,
+		EnableConsumableCapacity: utilfeature.DefaultFeatureGate.Enabled(features.DRAConsumableCapacity),
+		SliceInformer:            tc.sliceInformer,
+		TaintInformer:            tc.taintInformer,
+		ClassInformer:            tc.classInformer,
+		KubeClient:               tc.client,
 	}
 	sliceTracker, err := resourceslicetracker.StartTracker(ctx, opts)
 	if err != nil {

pkg/scheduler/framework/autoscaler_contract/lister_contract_test.go

@@ -101,6 +101,10 @@ func (r *resourceClaimTrackerContract) ListAllAllocatedDevices() (sets.Set[struc
 	return nil, nil
 }
 
+func (r *resourceClaimTrackerContract) GatherAllocatedState() (*structured.AllocatedState, error) {
+	return nil, nil
+}
+
 func (r *resourceClaimTrackerContract) SignalClaimPendingAllocation(_ types.UID, _ *resourceapi.ResourceClaim) error {
 	return nil
 }

pkg/scheduler/framework/listers.go

@@ -84,6 +84,9 @@ type ResourceClaimTracker interface {
 	// ListAllAllocatedDevices lists all allocated Devices from allocated ResourceClaims. The result is guaranteed to immediately include
 	// any changes made via AssumeClaimAfterAPICall(), and SignalClaimPendingAllocation().
 	ListAllAllocatedDevices() (sets.Set[structured.DeviceID], error)
+	// GatherAllocatedState gathers information about allocated devices from allocated ResourceClaims. The result is guaranteed to immediately include
+	// any changes made via AssumeClaimAfterAPICall(), and SignalClaimPendingAllocation().
+	GatherAllocatedState() (*structured.AllocatedState, error)
 
 	// SignalClaimPendingAllocation signals to the tracker that the given ResourceClaim will be allocated via an API call in the
 	// binding phase. This change is immediately reflected in the result of List() and the other accessors.

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

@@ -21,9 +21,11 @@ import (
 
 	resourceapi "k8s.io/api/resource/v1"
 	"k8s.io/apimachinery/pkg/util/sets"
+	utilfeature "k8s.io/apiserver/pkg/util/feature"
 	"k8s.io/client-go/tools/cache"
 	"k8s.io/dynamic-resource-allocation/structured"
 	"k8s.io/klog/v2"
+	"k8s.io/kubernetes/pkg/features"
 	schedutil "k8s.io/kubernetes/pkg/scheduler/util"
 	"k8s.io/utils/ptr"
 )
@@ -36,7 +38,11 @@ import (
 // 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)) {
+func foreachAllocatedDevice(claim *resourceapi.ResourceClaim,
+	dedicatedDeviceCallback func(deviceID structured.DeviceID),
+	enabledConsumableCapacity bool,
+	sharedDeviceCallback func(structured.SharedDeviceID),
+	consumedCapacityCallback func(structured.DeviceConsumedCapacity)) {
 	if claim.Status.Allocation == nil {
 		return
 	}
@@ -54,7 +60,24 @@ func foreachAllocatedDevice(claim *resourceapi.ResourceClaim, cb func(deviceID s
 
 		// None of the users of this helper need to abort iterating,
 		// therefore it's not supported as it only would add overhead.
-		cb(deviceID)
+
+		// Execute sharedDeviceCallback and consumedCapacityCallback correspondingly
+		// if DRAConsumableCapacity feature is enabled
+		if enabledConsumableCapacity {
+			shared := result.ShareID != nil
+			if shared {
+				sharedDeviceID := structured.MakeSharedDeviceID(deviceID, result.ShareID)
+				sharedDeviceCallback(sharedDeviceID)
+				if result.ConsumedCapacity != nil {
+					deviceConsumedCapacity := structured.NewDeviceConsumedCapacity(deviceID, result.ConsumedCapacity)
+					consumedCapacityCallback(deviceConsumedCapacity)
+				}
+				continue
+			}
+		}
+
+		// Otherwise, execute dedicatedDeviceCallback
+		dedicatedDeviceCallback(deviceID)
 	}
 }
 
@@ -66,14 +89,20 @@ func foreachAllocatedDevice(claim *resourceapi.ResourceClaim, cb func(deviceID s
 type allocatedDevices struct {
 	logger klog.Logger
 
-	mutex sync.RWMutex
-	ids   sets.Set[structured.DeviceID]
+	mutex                     sync.RWMutex
+	ids                       sets.Set[structured.DeviceID]
+	shareIDs                  sets.Set[structured.SharedDeviceID]
+	capacities                structured.ConsumedCapacityCollection
+	enabledConsumableCapacity bool
 }
 
 func newAllocatedDevices(logger klog.Logger) *allocatedDevices {
 	return &allocatedDevices{
-		logger: logger,
-		ids:    sets.New[structured.DeviceID](),
+		logger:                    logger,
+		ids:                       sets.New[structured.DeviceID](),
+		shareIDs:                  sets.New[structured.SharedDeviceID](),
+		capacities:                structured.NewConsumedCapacityCollection(),
+		enabledConsumableCapacity: utilfeature.DefaultFeatureGate.Enabled(features.DRAConsumableCapacity),
 	}
 }
 
@@ -84,6 +113,13 @@ func (a *allocatedDevices) Get() sets.Set[structured.DeviceID] {
 	return a.ids.Clone()
 }
 
+func (a *allocatedDevices) Capacities() structured.ConsumedCapacityCollection {
+	a.mutex.RLock()
+	defer a.mutex.RUnlock()
+
+	return a.capacities.Clone()
+}
+
 func (a *allocatedDevices) handlers() cache.ResourceEventHandler {
 	return cache.ResourceEventHandlerFuncs{
 		AddFunc:    a.onAdd,
@@ -142,16 +178,39 @@ func (a *allocatedDevices) addDevices(claim *resourceapi.ResourceClaim) {
 	// 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)
-	})
+	var shareIDs []structured.SharedDeviceID
+	var deviceCapacities []structured.DeviceConsumedCapacity
+	if a.enabledConsumableCapacity {
+		shareIDs = make([]structured.SharedDeviceID, 0, 20)
+		deviceCapacities = make([]structured.DeviceConsumedCapacity, 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.enabledConsumableCapacity,
+		func(sharedDeviceID structured.SharedDeviceID) {
+			a.logger.V(6).Info("Observed shared device allocation", "shared device", sharedDeviceID, "claim", klog.KObj(claim))
+			shareIDs = append(shareIDs, sharedDeviceID)
+		},
+		func(capacity structured.DeviceConsumedCapacity) {
+			a.logger.V(6).Info("Observed consumed capacity", "device", capacity.DeviceID, "consumed capacity", capacity.ConsumedCapacity, "claim", klog.KObj(claim))
+			deviceCapacities = append(deviceCapacities, capacity)
+		},
+	)
 
 	a.mutex.Lock()
 	defer a.mutex.Unlock()
 	for _, deviceID := range deviceIDs {
 		a.ids.Insert(deviceID)
 	}
+	for _, shareID := range shareIDs {
+		a.shareIDs.Insert(shareID)
+	}
+	for _, capacity := range deviceCapacities {
+		a.capacities.Insert(capacity)
+	}
 }
 
 func (a *allocatedDevices) removeDevices(claim *resourceapi.ResourceClaim) {
@@ -162,14 +221,35 @@ func (a *allocatedDevices) removeDevices(claim *resourceapi.ResourceClaim) {
 	// 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)
-	})
-
+	var shareIDs []structured.SharedDeviceID
+	var deviceCapacities []structured.DeviceConsumedCapacity
+	if a.enabledConsumableCapacity {
+		shareIDs = make([]structured.SharedDeviceID, 0, 20)
+		deviceCapacities = make([]structured.DeviceConsumedCapacity, 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.enabledConsumableCapacity,
+		func(sharedDeviceID structured.SharedDeviceID) {
+			a.logger.V(6).Info("Observed shared device deallocation", "shared device", sharedDeviceID, "claim", klog.KObj(claim))
+			shareIDs = append(shareIDs, sharedDeviceID)
+		},
+		func(capacity structured.DeviceConsumedCapacity) {
+			a.logger.V(6).Info("Observed consumed capacity release", "device id", capacity.DeviceID, "consumed capacity", capacity.ConsumedCapacity, "claim", klog.KObj(claim))
+			deviceCapacities = append(deviceCapacities, capacity)
+		})
 	a.mutex.Lock()
 	defer a.mutex.Unlock()
 	for _, deviceID := range deviceIDs {
 		a.ids.Delete(deviceID)
 	}
+	for _, shareID := range shareIDs {
+		a.shareIDs.Delete(shareID)
+	}
+	for _, capacity := range deviceCapacities {
+		a.capacities.Remove(capacity)
+	}
 }

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

@@ -25,11 +25,13 @@ import (
 	"k8s.io/apimachinery/pkg/labels"
 	"k8s.io/apimachinery/pkg/types"
 	"k8s.io/apimachinery/pkg/util/sets"
+	utilfeature "k8s.io/apiserver/pkg/util/feature"
 	"k8s.io/client-go/informers"
 	resourcelisters "k8s.io/client-go/listers/resource/v1"
 	resourceslicetracker "k8s.io/dynamic-resource-allocation/resourceslice/tracker"
 	"k8s.io/dynamic-resource-allocation/structured"
 	"k8s.io/klog/v2"
+	"k8s.io/kubernetes/pkg/features"
 	"k8s.io/kubernetes/pkg/scheduler/framework"
 	"k8s.io/kubernetes/pkg/scheduler/util/assumecache"
 )
@@ -212,13 +214,48 @@ func (c *claimTracker) ListAllAllocatedDevices() (sets.Set[structured.DeviceID],
 		foreachAllocatedDevice(claim, func(deviceID structured.DeviceID) {
 			c.logger.V(6).Info("Device is in flight for allocation", "device", deviceID, "claim", klog.KObj(claim))
 			allocated.Insert(deviceID)
-		})
+		}, false, func(structured.SharedDeviceID) {}, func(structured.DeviceConsumedCapacity) {})
 		return true
 	})
 	// There's no reason to return an error in this implementation, but the error might be helpful for other implementations.
 	return allocated, nil
 }
 
+func (c *claimTracker) GatherAllocatedState() (*structured.AllocatedState, error) {
+	// Start with a fresh set that matches the current known state of the
+	// world according to the informers.
+	allocated := c.allocatedDevices.Get()
+	allocatedSharedDeviceIDs := sets.New[structured.SharedDeviceID]()
+	aggregatedCapacity := c.allocatedDevices.Capacities()
+
+	enabledConsumableCapacity := utilfeature.DefaultFeatureGate.Enabled(features.DRAConsumableCapacity)
+
+	// Whatever is in flight also has to be checked.
+	c.inFlightAllocations.Range(func(key, value any) bool {
+		claim := value.(*resourceapi.ResourceClaim)
+		foreachAllocatedDevice(claim, func(deviceID structured.DeviceID) {
+			c.logger.V(6).Info("Device is in flight for allocation", "device", deviceID, "claim", klog.KObj(claim))
+			allocated.Insert(deviceID)
+		},
+			enabledConsumableCapacity,
+			func(sharedDeviceID structured.SharedDeviceID) {
+				c.logger.V(6).Info("Device is in flight for allocation", "shared device", sharedDeviceID, "claim", klog.KObj(claim))
+				allocatedSharedDeviceIDs.Insert(sharedDeviceID)
+			}, func(capacity structured.DeviceConsumedCapacity) {
+				c.logger.V(6).Info("Device is in flight for allocation", "consumed capacity", capacity, "claim", klog.KObj(claim))
+				aggregatedCapacity.Insert(capacity)
+			})
+		return true
+	})
+
+	// There's no reason to return an error in this implementation, but the error might be helpful for other implementations.
+	return &structured.AllocatedState{
+		AllocatedDevices:         allocated,
+		AllocatedSharedDeviceIDs: allocatedSharedDeviceIDs,
+		AggregatedCapacity:       aggregatedCapacity,
+	}, nil
+}
+
 func (c *claimTracker) AssumeClaimAfterAPICall(claim *resourceapi.ResourceClaim) error {
 	return c.cache.Assume(claim)
 }

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

@@ -170,6 +170,7 @@ type DynamicResources struct {
 	enableExtendedResource        bool
 	enableFilterTimeout           bool
 	filterTimeout                 time.Duration
+	enableConsumableCapacity      bool
 
 	fh         framework.Handle
 	clientset  kubernetes.Interface
@@ -201,6 +202,7 @@ func New(ctx context.Context, plArgs runtime.Object, fh framework.Handle, fts fe
 		enableSchedulingQueueHint:     fts.EnableSchedulingQueueHint,
 		enablePartitionableDevices:    fts.EnablePartitionableDevices,
 		enableExtendedResource:        fts.EnableDRAExtendedResource,
+		enableConsumableCapacity:      fts.EnableConsumableCapacity,
 		filterTimeout:                 ptr.Deref(args.FilterTimeout, metav1.Duration{}).Duration,
 		enableDeviceBindingConditions: fts.EnableDRADeviceBindingConditions,
 		enableDeviceStatus:            fts.EnableDRAResourceClaimDeviceStatus,
@@ -210,7 +212,7 @@ func New(ctx context.Context, plArgs runtime.Object, fh framework.Handle, fts fe
 		// 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),
+		celCache:   cel.NewCache(10, cel.Features{EnableConsumableCapacity: fts.EnableConsumableCapacity}),
 		draManager: fh.SharedDRAManager(),
 	}
 
@@ -658,9 +660,25 @@ func (pl *DynamicResources) PreFilter(ctx context.Context, state fwk.CycleState,
 		// 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, err := pl.draManager.ResourceClaims().ListAllAllocatedDevices()
-		if err != nil {
-			return nil, statusError(logger, err)
+		var allocatedState *structured.AllocatedState
+		if pl.enableConsumableCapacity {
+			allocatedState, err = pl.draManager.ResourceClaims().GatherAllocatedState()
+			if err != nil {
+				return nil, statusError(logger, err)
+			}
+			if allocatedState == nil {
+				return nil, statusError(logger, errors.New("nil allocated state"))
+			}
+		} else {
+			allocatedDevices, err := pl.draManager.ResourceClaims().ListAllAllocatedDevices()
+			if err != nil {
+				return nil, statusError(logger, err)
+			}
+			allocatedState = &structured.AllocatedState{
+				AllocatedDevices:         allocatedDevices,
+				AllocatedSharedDeviceIDs: sets.New[structured.SharedDeviceID](),
+				AggregatedCapacity:       structured.NewConsumedCapacityCollection(),
+			}
 		}
 		slices, err := pl.draManager.ResourceSlices().ListWithDeviceTaintRules()
 		if err != nil {
@@ -673,8 +691,9 @@ func (pl *DynamicResources) PreFilter(ctx context.Context, state fwk.CycleState,
 			DeviceTaints:         pl.enableDeviceTaints,
 			DeviceBinding:        pl.enableDeviceBindingConditions,
 			DeviceStatus:         pl.enableDeviceStatus,
+			ConsumableCapacity:   pl.enableConsumableCapacity,
 		}
-		allocator, err := structured.NewAllocator(ctx, features, allAllocatedDevices, pl.draManager.DeviceClasses(), slices, pl.celCache)
+		allocator, err := structured.NewAllocator(ctx, features, *allocatedState, pl.draManager.DeviceClasses(), slices, pl.celCache)
 		if err != nil {
 			return nil, statusError(logger, err)
 		}

pkg/scheduler/scheduler.go

@@ -326,9 +326,10 @@ func New(ctx context.Context,
 		resourceClaimInformer := informerFactory.Resource().V1().ResourceClaims().Informer()
 		resourceClaimCache = assumecache.NewAssumeCache(logger, resourceClaimInformer, "ResourceClaim", "", nil)
 		resourceSliceTrackerOpts := resourceslicetracker.Options{
-			EnableDeviceTaints: feature.DefaultFeatureGate.Enabled(features.DRADeviceTaints),
-			SliceInformer:      informerFactory.Resource().V1().ResourceSlices(),
-			KubeClient:         client,
+			EnableDeviceTaints:       feature.DefaultFeatureGate.Enabled(features.DRADeviceTaints),
+			EnableConsumableCapacity: feature.DefaultFeatureGate.Enabled(features.DRAConsumableCapacity),
+			SliceInformer:            informerFactory.Resource().V1().ResourceSlices(),
+			KubeClient:               client,
 		}
 		// If device taints are disabled, the additional informers are not needed and
 		// the tracker turns into a simple wrapper around the slice informer.