DRA: Fix missing GetSharedDeviceIDs bug in GatherAllocatedState with consumable capacity

Signed-off-by: Sunyanan Choochotkaew <sunyanan.choochotkaew1@ibm.com>

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()

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)
 			})

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")
+			}
+		})
+	}
+
+}

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 {