diff --git a/pkg/scheduler/framework/plugins/dynamicresources/allocateddevices.go b/pkg/scheduler/framework/plugins/dynamicresources/allocateddevices.go index 21d192ffc89..f74e466a40f 100644 --- a/pkg/scheduler/framework/plugins/dynamicresources/allocateddevices.go +++ b/pkg/scheduler/framework/plugins/dynamicresources/allocateddevices.go @@ -120,6 +120,13 @@ func (a *allocatedDevices) Get() (sets.Set[structured.DeviceID], int64) { return a.ids.Clone(), a.revision } +func (a *allocatedDevices) GetSharedDeviceIDs() (sets.Set[structured.SharedDeviceID], int64) { + a.mutex.RLock() + defer a.mutex.RUnlock() + + return a.shareIDs.Clone(), a.revision +} + func (a *allocatedDevices) Capacities() (structured.ConsumedCapacityCollection, int64) { a.mutex.RLock() defer a.mutex.RUnlock() diff --git a/pkg/scheduler/framework/plugins/dynamicresources/dra_manager.go b/pkg/scheduler/framework/plugins/dynamicresources/dra_manager.go index 2be14d10ec6..9b8cad52239 100644 --- a/pkg/scheduler/framework/plugins/dynamicresources/dra_manager.go +++ b/pkg/scheduler/framework/plugins/dynamicresources/dra_manager.go @@ -32,6 +32,7 @@ import ( "k8s.io/dynamic-resource-allocation/deviceclass/extendedresourcecache" resourceslicetracker "k8s.io/dynamic-resource-allocation/resourceslice/tracker" "k8s.io/dynamic-resource-allocation/structured" + "k8s.io/dynamic-resource-allocation/structured/schedulerapi" "k8s.io/klog/v2" fwk "k8s.io/kube-scheduler/framework" "k8s.io/kubernetes/pkg/features" @@ -261,32 +262,66 @@ func (c *claimTracker) ListAllAllocatedDevices() (sets.Set[structured.DeviceID], return nil, errClaimTrackerConcurrentModification } +// GatherAllocatedState collects and returns the current allocation state of all devices +// across the cluster. This includes: +// - AllocatedDevices: Set of device IDs that are fully allocated (dedicated mode) +// - AllocatedSharedDeviceIDs: Set of shared device IDs when consumable capacity is enabled +// - AggregatedCapacity: Consumed capacity across all devices when consumable capacity is enabled +// +// The function handles two allocation models: +// 1. Legacy dedicated mode (DRAConsumableCapacity disabled): Devices are allocated exclusively +// to a single claim. Shared devices are converted to their base device IDs. +// 2. Consumable capacity mode (DRAConsumableCapacity enabled): Devices can be shared across +// multiple claims with capacity tracking. +// +// The function ensures consistency by: +// - Reading allocation state from informer-backed cache +// - Including in-flight allocations that haven't been persisted yet +// - Using revision numbers to detect concurrent modifications and retry if needed +// +// Returns errClaimTrackerConcurrentModification if the state changed during collection, +// indicating the caller should retry. 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, revision1 := c.allocatedDevices.Get() - allocatedSharedDeviceIDs := sets.New[structured.SharedDeviceID]() - aggregatedCapacity, revision2 := c.allocatedDevices.Capacities() + enabledConsumableCapacity := utilfeature.DefaultFeatureGate.Enabled(features.DRAConsumableCapacity) - if revision1 != revision2 { + allocated, revision1 := c.allocatedDevices.Get() + allocatedSharedDeviceIDs, revision2 := c.allocatedDevices.GetSharedDeviceIDs() + aggregatedCapacity, revision3 := c.allocatedDevices.Capacities() + + if revision1 != revision2 || revision2 != revision3 { // Already not consistent. Try again. return nil, errClaimTrackerConcurrentModification } - enabledConsumableCapacity := utilfeature.DefaultFeatureGate.Enabled(features.DRAConsumableCapacity) + if !enabledConsumableCapacity { + // When the DRAConsumableCapacity feature is disabled, we fall back to the legacy + // dedicated device allocation model. + // This ensures backward compatibility with the original DRA behavior where devices + // could only be allocated exclusively to a single claim. + for sharedDeviceID := range allocatedSharedDeviceIDs { + allocated.Insert(sharedDeviceID.GetBaseDeviceID()) + } + // Reset allocatedSharedDeviceIDs and aggregatedCapacity + allocatedSharedDeviceIDs = sets.New[structured.SharedDeviceID]() + aggregatedCapacity = make(schedulerapi.ConsumedCapacityCollection) + } // 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) - }, + foreachAllocatedDevice(claim, + func(deviceID structured.DeviceID) { // dedicatedDeviceCallback + 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) { + func(sharedDeviceID structured.SharedDeviceID) { // sharedDeviceCallback 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) { + }, + func(capacity structured.DeviceConsumedCapacity) { // consumedCapacityCallback c.logger.V(6).Info("Device is in flight for allocation", "consumed capacity", capacity, "claim", klog.KObj(claim)) aggregatedCapacity.Insert(capacity) }) diff --git a/pkg/scheduler/framework/plugins/dynamicresources/dynamicresources_test.go b/pkg/scheduler/framework/plugins/dynamicresources/dynamicresources_test.go index d19ab4dd218..6eaf6ae86ea 100644 --- a/pkg/scheduler/framework/plugins/dynamicresources/dynamicresources_test.go +++ b/pkg/scheduler/framework/plugins/dynamicresources/dynamicresources_test.go @@ -37,6 +37,7 @@ import ( v1 "k8s.io/api/core/v1" resourceapi "k8s.io/api/resource/v1" apierrors "k8s.io/apimachinery/pkg/api/errors" + apiresource "k8s.io/apimachinery/pkg/api/resource" metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" apiruntime "k8s.io/apimachinery/pkg/runtime" "k8s.io/apimachinery/pkg/types" @@ -53,6 +54,8 @@ import ( "k8s.io/dynamic-resource-allocation/deviceclass/extendedresourcecache" resourceslicetracker "k8s.io/dynamic-resource-allocation/resourceslice/tracker" "k8s.io/dynamic-resource-allocation/structured" + "k8s.io/dynamic-resource-allocation/structured/schedulerapi" + "k8s.io/klog/v2" kubeschedulerconfigv1 "k8s.io/kube-scheduler/config/v1" fwk "k8s.io/kube-scheduler/framework" "k8s.io/kubernetes/pkg/features" @@ -80,10 +83,12 @@ var ( node3Name = "worker-3" driver = "some-driver" driver2 = "some-driver-2" + sharedDeviceName = "shared-instance" podName = "my-pod" podUID = "1234" resourceName = "my-resource" resourceName2 = resourceName + "-2" + capacityName = resourceapi.QualifiedName("my-cap") claimName = podName + "-" + resourceName claimName2 = podName + "-" + resourceName2 className = "my-resource-class" @@ -280,6 +285,54 @@ var ( return st.MakeNodeSelector().In("metadata.name", []string{nodeName}, st.NodeSelectorTypeMatchFields).Obj() }(), } + allocationResultWithSharedDevice = &resourceapi.AllocationResult{ + Devices: resourceapi.DeviceAllocationResult{ + Results: []resourceapi.DeviceRequestAllocationResult{{ + Driver: driver, + Pool: nodeName, + Device: sharedDeviceName, + Request: "req-1", + ShareID: ptr.To(types.UID("share-123")), // Shared device allocation + }}, + }, + NodeSelector: func() *v1.NodeSelector { + return st.MakeNodeSelector().In("metadata.name", []string{nodeName}, st.NodeSelectorTypeMatchFields).Obj() + }(), + } + allocationResultWithConsumedCapacity = &resourceapi.AllocationResult{ + Devices: resourceapi.DeviceAllocationResult{ + Results: []resourceapi.DeviceRequestAllocationResult{{ + Driver: driver, + Pool: nodeName, + Device: sharedDeviceName, + Request: "req-1", + ShareID: ptr.To(types.UID("share-123")), // Shared device allocation + ConsumedCapacity: map[resourceapi.QualifiedName]apiresource.Quantity{ + capacityName: apiresource.MustParse("1"), + }, + }}, + }, + NodeSelector: func() *v1.NodeSelector { + return st.MakeNodeSelector().In("metadata.name", []string{nodeName}, st.NodeSelectorTypeMatchFields).Obj() + }(), + } + allocationResultWithConsumedCapacity2 = &resourceapi.AllocationResult{ + Devices: resourceapi.DeviceAllocationResult{ + Results: []resourceapi.DeviceRequestAllocationResult{{ + Driver: driver, + Pool: nodeName, + Device: sharedDeviceName, + Request: "req-1", + ShareID: ptr.To(types.UID("share-456")), // Shared device allocation + ConsumedCapacity: map[resourceapi.QualifiedName]apiresource.Quantity{ + capacityName: apiresource.MustParse("1"), + }, + }}, + }, + NodeSelector: func() *v1.NodeSelector { + return st.MakeNodeSelector().In("metadata.name", []string{nodeName}, st.NodeSelectorTypeMatchFields).Obj() + }(), + } allocationResult2 = &resourceapi.AllocationResult{ Devices: resourceapi.DeviceAllocationResult{ Results: []resourceapi.DeviceRequestAllocationResult{{ @@ -517,6 +570,15 @@ var ( allocatedClaimWithGoodTopology = st.FromResourceClaim(allocatedClaim). Allocation(&resourceapi.AllocationResult{NodeSelector: st.MakeNodeSelector().In("kubernetes.io/hostname", []string{nodeName}, st.NodeSelectorTypeMatchExpressions).Obj()}). Obj() + allocatedClaimWithSharedDevice = st.FromResourceClaim(pendingClaim). + Allocation(allocationResultWithSharedDevice). + Obj() + allocatedClaimWithConsumedCapacity = st.FromResourceClaim(pendingClaim). + Allocation(allocationResultWithConsumedCapacity). + Obj() + allocatedClaimWithConsumedCapacity2 = st.FromResourceClaim(pendingClaim). + Allocation(allocationResultWithConsumedCapacity2). + Obj() otherClaim = st.MakeResourceClaim(). Name("not-my-claim"). Namespace(namespace). @@ -3738,3 +3800,174 @@ func TestNormalizeScore(t *testing.T) { }) } } + +func TestGatherAllocatedState(t *testing.T) { + tCtx := ktesting.Init(t) + testGatherAllocatedState(tCtx) +} +func testGatherAllocatedState(tCtx ktesting.TContext) { + testcases := map[string]struct { + allocatedResourceClaims []*resourceapi.ResourceClaim + inflightResourceClaims map[types.UID]*resourceapi.ResourceClaim + enabledConsumableCapacity bool + expectErr bool + expectedIDAllocated int + expectedSharedIDAllocated int + expectedConsumedCapacity string + }{ + "no-claims": { + expectedIDAllocated: 0, + expectedSharedIDAllocated: 0, + }, + "single-allocated-claim": { + allocatedResourceClaims: []*resourceapi.ResourceClaim{ + allocatedClaim, + }, + expectedIDAllocated: 1, + expectedSharedIDAllocated: 0, + }, + "single-allocated-claim-with-shared-device": { + enabledConsumableCapacity: true, + allocatedResourceClaims: []*resourceapi.ResourceClaim{ + allocatedClaimWithSharedDevice, + }, + expectedIDAllocated: 0, + expectedSharedIDAllocated: 1, + }, + "single-allocated-claim-with-capacity": { + enabledConsumableCapacity: true, + allocatedResourceClaims: []*resourceapi.ResourceClaim{ + allocatedClaimWithConsumedCapacity, + }, + expectedIDAllocated: 0, + expectedSharedIDAllocated: 1, + expectedConsumedCapacity: "1", + }, + "disabled-single-allocated-claim-with-capacity": { + enabledConsumableCapacity: false, + allocatedResourceClaims: []*resourceapi.ResourceClaim{ + allocatedClaimWithConsumedCapacity, + }, + expectedIDAllocated: 1, + expectedSharedIDAllocated: 0, + expectedConsumedCapacity: "", + }, + "mixed-allocated-claim": { + enabledConsumableCapacity: true, + allocatedResourceClaims: []*resourceapi.ResourceClaim{ + allocatedClaim, + allocatedClaimWithConsumedCapacity, + }, + expectedIDAllocated: 1, + expectedSharedIDAllocated: 1, + expectedConsumedCapacity: "1", + }, + "add-inflight-allocated-claim-with-capacity": { + enabledConsumableCapacity: true, + allocatedResourceClaims: []*resourceapi.ResourceClaim{ + allocatedClaim, + allocatedClaimWithConsumedCapacity, + }, + inflightResourceClaims: map[types.UID]*resourceapi.ResourceClaim{ + "claim-2-uid": allocatedClaimWithConsumedCapacity2, + }, + expectedIDAllocated: 1, + expectedSharedIDAllocated: 2, + expectedConsumedCapacity: "2", + }, + "disabled-inflight-allocated-claim-with-capacity": { + enabledConsumableCapacity: false, + allocatedResourceClaims: []*resourceapi.ResourceClaim{ + allocatedClaim, + allocatedClaimWithConsumedCapacity, + }, + inflightResourceClaims: map[types.UID]*resourceapi.ResourceClaim{ + "claim-2-uid": allocatedClaimWithConsumedCapacity2, + }, + expectedIDAllocated: 2, + expectedSharedIDAllocated: 0, + expectedConsumedCapacity: "", + }, + } + for name, tc := range testcases { + tCtx.Run(name, func(tCtx ktesting.TContext) { + featuregatetesting.SetFeatureGateDuringTest(tCtx, utilfeature.DefaultFeatureGate, features.DRAConsumableCapacity, tc.enabledConsumableCapacity) + + tCtx.Helper() + logger := klog.FromContext(tCtx) + draManager := &DefaultDRAManager{ + resourceClaimTracker: &claimTracker{ + inFlightAllocations: &sync.Map{}, + allocatedDevices: newAllocatedDevices(logger), + }, + } + for _, obj := range tc.allocatedResourceClaims { + draManager.resourceClaimTracker.allocatedDevices.handlers().OnAdd(obj, false) + } + if tc.inflightResourceClaims != nil { + for claimUID, obj := range tc.inflightResourceClaims { + err := draManager.resourceClaimTracker.SignalClaimPendingAllocation(claimUID, obj) + if err != nil { + if !tc.expectErr { + tCtx.Fatalf("unexpected error: %v", err) + return + } + } + if tc.expectErr { + tCtx.Fatal("expected error, got none") + } + } + } + + // Start the test from here + allocatedState, err := draManager.ResourceClaims().GatherAllocatedState() + if err != nil { + if !tc.expectErr { + tCtx.Fatalf("unexpected error: %v", err) + return + } + } + if tc.expectErr { + tCtx.Fatal("expected error, got none") + } + allocatedDeviceIDs := allocatedState.AllocatedDevices + allocatedSharedDeviceIDs := allocatedState.AllocatedSharedDeviceIDs + aggregatedCapacity := allocatedState.AggregatedCapacity + + // Verify the counts match expectations + if allocatedDeviceIDs.Len() != tc.expectedIDAllocated { + tCtx.Errorf("expected %d allocated device IDs, got %d", tc.expectedIDAllocated, allocatedDeviceIDs.Len()) + } + if allocatedSharedDeviceIDs.Len() != tc.expectedSharedIDAllocated { + tCtx.Errorf("expected %d allocated shared device IDs, got %d", tc.expectedSharedIDAllocated, allocatedSharedDeviceIDs.Len()) + } + + // Verify aggregated capacity is initialized + if aggregatedCapacity == nil { + tCtx.Error("aggregatedCapacity should not be nil") + } + if tc.expectedConsumedCapacity != "" { + if len(aggregatedCapacity) == 0 { + tCtx.Errorf("expected consumed capacity, got empty") + } + deviceID := schedulerapi.MakeDeviceID(driver, nodeName, sharedDeviceName) + capacity := aggregatedCapacity[deviceID] + if capacity == nil { + tCtx.Errorf("expected aggregated capacity of %s, got nil", deviceID) + return + } + value := capacity[capacityName] + if value == nil { + tCtx.Errorf("expected value of %s, got nil", capacityName) + return + } + if value.Cmp(apiresource.MustParse(tc.expectedConsumedCapacity)) != 0 { + tCtx.Errorf("expected value of %s to be %s, got %s", capacityName, tc.expectedConsumedCapacity, value) + } + } else if len(aggregatedCapacity) > 0 { + tCtx.Errorf("got unexpected consumed capacity") + } + }) + } + +} diff --git a/staging/src/k8s.io/dynamic-resource-allocation/structured/schedulerapi/types.go b/staging/src/k8s.io/dynamic-resource-allocation/structured/schedulerapi/types.go index a401a8e1aa8..d98473d6439 100644 --- a/staging/src/k8s.io/dynamic-resource-allocation/structured/schedulerapi/types.go +++ b/staging/src/k8s.io/dynamic-resource-allocation/structured/schedulerapi/types.go @@ -58,6 +58,10 @@ func (d SharedDeviceID) String() string { return deviceIDStr } +func (d SharedDeviceID) GetBaseDeviceID() DeviceID { + return MakeDeviceID(d.Driver.String(), d.Pool.String(), d.Device.String()) +} + // MakeSharedDeviceID creates a new SharedDeviceID from a DeviceID and share ID. // This function is used in consumable capacity features and the scheduler. func MakeSharedDeviceID(deviceID DeviceID, shareID *types.UID) SharedDeviceID {