DRA: fix Pods in a scheduling gang sharing ResourceClaims

pkg/scheduler/framework/autoscaler_contract/lister_contract_test.go

@@ -152,12 +152,20 @@ func (r *resourceClaimTrackerContract) SignalClaimPendingAllocation(_ types.UID,
 	return nil
 }
 
-func (r *resourceClaimTrackerContract) ClaimHasPendingAllocation(_ types.UID) bool {
+func (r *resourceClaimTrackerContract) GetPendingAllocation(_ types.UID) (*resourceapi.AllocationResult, bool) {
+	return nil, false
+}
+
+func (r *resourceClaimTrackerContract) MaybeRemoveClaimPendingAllocation(_ types.UID, _ bool) (deleted bool) {
 	return false
 }
 
-func (r *resourceClaimTrackerContract) RemoveClaimPendingAllocation(_ types.UID) (deleted bool) {
-	return false
+func (r *resourceClaimTrackerContract) AddSharedClaimPendingAllocation(claimUID types.UID, allocatedClaim *resourceapi.ResourceClaim) error {
+	return nil
+}
+
+func (r *resourceClaimTrackerContract) RemoveSharedClaimPendingAllocation(claimUID types.UID, allocatedClaim *resourceapi.ResourceClaim) error {
+	return nil
 }
 
 func (r *resourceClaimTrackerContract) AssumeClaimAfterAPICall(_ *resourceapi.ResourceClaim) error {

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

@@ -70,6 +70,9 @@ func NewDRAManager(ctx context.Context, claimsCache *assumecache.AssumeCache, re
 	if utilfeature.DefaultFeatureGate.Enabled(features.DRAExtendedResource) {
 		manager.extendedResourceCache = extendedresourcecache.NewExtendedResourceCache(logger)
 	}
+	if utilfeature.DefaultFeatureGate.Enabled(features.GenericWorkload) {
+		manager.resourceClaimTracker.inFlightAllocationSharers = &sync.Map{}
+	}
 
 	// Reacting to events is more efficient than iterating over the list
 	// repeatedly in PreFilter.
@@ -150,8 +153,9 @@ type claimTracker struct {
 	// for which allocation was triggered during a scheduling cycle and the
 	// corresponding claim status update call in PreBind has not been done
 	// yet. If another pod needs the claim, the pod is treated as "not
-	// schedulable yet". The cluster event for the claim status update will
-	// make it schedulable.
+	// schedulable yet" unless the pod is a member of a PodGroup. For ungrouped
+	// pods, the cluster event for the claim status update will make it
+	// schedulable.
 	//
 	// This mechanism avoids the following problem:
 	// - Pod A triggers allocation for claim X.
@@ -165,24 +169,38 @@ type claimTracker struct {
 	// problem:
 	// - Pod A and B get scheduled as above.
 	// - PreBind for pod A gets called first, then fails with a temporary API error.
-	//   It removes the updated claim from the assume cache because of that.
+	//   It removes the updated claim from the in-flight claims because of that.
 	// - PreBind for pod B gets called next and succeeds with adding the
 	//   allocation and its own reservedFor entry.
 	// - The assume cache is now not reflecting that the claim is allocated,
 	//   which could lead to reusing the same resource for some other claim.
 	//
+	// For pods in a PodGroup, a pending allocation may be shared among several
+	// pods in the group. In the PreBind phase, the allocation will be written
+	// for the first pod that succeeds. The scenario above is prevented by
+	// keeping the pending allocation in-flight as long as it has not been
+	// unreserved for every pod in the group, as tracked by
+	// inFlightAllocationSharers.
+	//
 	// A sync.Map is used because in practice sharing of a claim between
 	// pods is expected to be rare compared to per-pod claim, so we end up
 	// hitting the "multiple goroutines read, write, and overwrite entries
 	// for disjoint sets of keys" case that sync.Map is optimized for.
 	inFlightAllocations *sync.Map
-	allocatedDevices    *allocatedDevices
-	logger              klog.Logger
+	// inFlightAllocationSharers counts the actively scheduling pods
+	// sharing a given ResourceClaim.
+	inFlightAllocationSharers *sync.Map
+	allocatedDevices          *allocatedDevices
+	logger                    klog.Logger
 }
 
-func (c *claimTracker) ClaimHasPendingAllocation(claimUID types.UID) bool {
-	_, found := c.inFlightAllocations.Load(claimUID)
-	return found
+func (c *claimTracker) GetPendingAllocation(claimUID types.UID) (*resourceapi.AllocationResult, bool) {
+	var allocation *resourceapi.AllocationResult
+	claim, found := c.inFlightAllocations.Load(claimUID)
+	if found && claim != nil {
+		allocation = claim.(*resourceapi.ResourceClaim).Status.Allocation
+	}
+	return allocation, found
 }
 
 func (c *claimTracker) SignalClaimPendingAllocation(claimUID types.UID, allocatedClaim *resourceapi.ResourceClaim) error {
@@ -195,7 +213,21 @@ func (c *claimTracker) SignalClaimPendingAllocation(claimUID types.UID, allocate
 	return nil
 }
 
-func (c *claimTracker) RemoveClaimPendingAllocation(claimUID types.UID) (deleted bool) {
+func (c *claimTracker) MaybeRemoveClaimPendingAllocation(claimUID types.UID, shareable bool) (deleted bool) {
+	if c.inFlightAllocationSharers != nil && shareable {
+		value, ok := c.inFlightAllocationSharers.Load(claimUID)
+		if ok && value.(int) > 0 {
+			if loggerV := c.logger.V(5); loggerV.Enabled() {
+				claim, found := c.inFlightAllocations.Load(claimUID)
+				if found {
+					claim := claim.(*resourceapi.ResourceClaim)
+					c.logger.V(5).Info("Claim is still shared by other pods, not removing in-flight claim", "claim", klog.KObj(claim), "uid", claimUID, "version", claim.ResourceVersion)
+				}
+			}
+			return false
+		}
+	}
+
 	claim, found := c.inFlightAllocations.LoadAndDelete(claimUID)
 	// The assume cache doesn't log this, but maybe it should.
 	if found {
@@ -207,6 +239,43 @@ func (c *claimTracker) RemoveClaimPendingAllocation(claimUID types.UID) (deleted
 	return found
 }
 
+func (c *claimTracker) AddSharedClaimPendingAllocation(claimUID types.UID, allocatedClaim *resourceapi.ResourceClaim) error {
+	newSharers := 1
+	value, loaded := c.inFlightAllocationSharers.LoadOrStore(claimUID, newSharers)
+	if loaded {
+		oldSharers := value.(int)
+		newSharers = oldSharers + 1
+		swapped := c.inFlightAllocationSharers.CompareAndSwap(claimUID, oldSharers, newSharers)
+		if !swapped {
+			// The value must have changed since we loaded
+			return fmt.Errorf("conflict adding in-flight allocation sharer for claim %s/%s, UID=%s", allocatedClaim.Namespace, allocatedClaim.Name, claimUID)
+		}
+	}
+	c.logger.V(5).Info("Added share for in-flight claim", "claim", klog.KObj(allocatedClaim), "uid", claimUID, "version", allocatedClaim.ResourceVersion, "sharers", newSharers)
+	return nil
+}
+
+func (c *claimTracker) RemoveSharedClaimPendingAllocation(claimUID types.UID, allocatedClaim *resourceapi.ResourceClaim) error {
+	value, ok := c.inFlightAllocationSharers.Load(claimUID)
+	if !ok {
+		return nil
+	}
+	oldSharers := value.(int)
+	newSharers := oldSharers - 1
+	var written bool
+	if newSharers == 0 {
+		written = c.inFlightAllocationSharers.CompareAndDelete(claimUID, oldSharers)
+	} else {
+		written = c.inFlightAllocationSharers.CompareAndSwap(claimUID, oldSharers, newSharers)
+	}
+	if !written {
+		// The value must have changed since we loaded
+		return fmt.Errorf("conflict removing in-flight allocation sharer for claim %s/%s, UID=%s", allocatedClaim.Namespace, allocatedClaim.Name, claimUID)
+	}
+	c.logger.V(5).Info("Removed share for in-flight claim", "claim", klog.KObj(allocatedClaim), "uid", claimUID, "version", allocatedClaim.ResourceVersion, "sharers", newSharers)
+	return nil
+}
+
 func (c *claimTracker) Get(namespace, claimName string) (*resourceapi.ResourceClaim, error) {
 	obj, err := c.cache.Get(namespace + "/" + claimName)
 	if err != nil {

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

@@ -0,0 +1,240 @@
+/*
+Copyright 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"
+	"testing"
+
+	"github.com/onsi/gomega"
+	resourceapi "k8s.io/api/resource/v1"
+	"k8s.io/apimachinery/pkg/types"
+	"k8s.io/kubernetes/test/utils/ktesting"
+)
+
+const (
+	claimUID      = types.UID("claim-uid-1")
+	otherClaimUID = types.UID("claim-uid-2")
+)
+
+func TestAddSharedClaimPendingAllocation(t *testing.T) {
+	testAddSharedClaimPendingAllocation(ktesting.Init(t))
+}
+func testAddSharedClaimPendingAllocation(tCtx ktesting.TContext) {
+	tests := map[string]struct {
+		inFlightAllocationSharers         map[types.UID]int
+		claimUID                          types.UID
+		allocatedClaim                    *resourceapi.ResourceClaim
+		expectedInFlightAllocationSharers map[types.UID]int
+	}{
+		"empty": {
+			claimUID:       claimUID,
+			allocatedClaim: allocatedClaim,
+			expectedInFlightAllocationSharers: map[types.UID]int{
+				claimUID: 1,
+			},
+		},
+		"increment": {
+			inFlightAllocationSharers: map[types.UID]int{
+				claimUID:      1,
+				otherClaimUID: 10,
+			},
+			claimUID:       claimUID,
+			allocatedClaim: allocatedClaim,
+			expectedInFlightAllocationSharers: map[types.UID]int{
+				claimUID:      2,
+				otherClaimUID: 10,
+			},
+		},
+	}
+
+	for name, test := range tests {
+		tCtx.Run(name, func(tCtx ktesting.TContext) {
+			c := &claimTracker{
+				logger:                    tCtx.Logger(),
+				inFlightAllocationSharers: &sync.Map{},
+			}
+
+			for key, value := range test.inFlightAllocationSharers {
+				c.inFlightAllocationSharers.Store(key, value)
+			}
+			err := c.AddSharedClaimPendingAllocation(test.claimUID, test.allocatedClaim)
+			tCtx.ExpectNoError(err)
+			actual := map[types.UID]int{}
+			c.inFlightAllocationSharers.Range(func(key, value any) bool {
+				actual[key.(types.UID)] = value.(int)
+				return true
+			})
+			tCtx.Expect(actual).To(gomega.Equal(test.expectedInFlightAllocationSharers))
+		})
+	}
+}
+
+func TestRemoveSharedClaimPendingAllocation(t *testing.T) {
+	testRemoveSharedClaimPendingAllocation(ktesting.Init(t))
+}
+func testRemoveSharedClaimPendingAllocation(tCtx ktesting.TContext) {
+	tests := map[string]struct {
+		inFlightAllocationSharers         map[types.UID]int
+		claimUID                          types.UID
+		allocatedClaim                    *resourceapi.ResourceClaim
+		expectedInFlightAllocationSharers map[types.UID]int
+		expectedErr                       string
+	}{
+		"empty": {
+			inFlightAllocationSharers: map[types.UID]int{
+				claimUID:      1,
+				otherClaimUID: 10,
+			},
+			claimUID:       claimUID,
+			allocatedClaim: allocatedClaim,
+			expectedInFlightAllocationSharers: map[types.UID]int{
+				otherClaimUID: 10,
+			},
+		},
+		"decrement": {
+			inFlightAllocationSharers: map[types.UID]int{
+				claimUID:      2,
+				otherClaimUID: 10,
+			},
+			claimUID:       claimUID,
+			allocatedClaim: allocatedClaim,
+			expectedInFlightAllocationSharers: map[types.UID]int{
+				claimUID:      1,
+				otherClaimUID: 10,
+			},
+		},
+	}
+
+	for name, test := range tests {
+		tCtx.Run(name, func(tCtx ktesting.TContext) {
+			c := &claimTracker{
+				logger:                    tCtx.Logger(),
+				inFlightAllocationSharers: &sync.Map{},
+			}
+
+			for key, value := range test.inFlightAllocationSharers {
+				c.inFlightAllocationSharers.Store(key, value)
+			}
+			err := c.RemoveSharedClaimPendingAllocation(test.claimUID, test.allocatedClaim)
+			tCtx.ExpectNoError(err)
+			actual := map[types.UID]int{}
+			c.inFlightAllocationSharers.Range(func(key, value any) bool {
+				actual[key.(types.UID)] = value.(int)
+				return true
+			})
+			tCtx.Expect(actual).To(gomega.Equal(test.expectedInFlightAllocationSharers))
+		})
+	}
+}
+
+func TestMaybeRemoveClaimPendingAllocation(t *testing.T) {
+	testMaybeRemoveClaimPendingAllocation(ktesting.Init(t))
+}
+func testMaybeRemoveClaimPendingAllocation(tCtx ktesting.TContext) {
+	tests := map[string]struct {
+		inFlightAllocations               map[types.UID]*resourceapi.ResourceClaim
+		inFlightAllocationSharers         map[types.UID]int
+		claimUID                          types.UID
+		shareable                         bool
+		expected                          bool
+		expectedInFlightAllocationSharers map[types.UID]int
+		expectedInFlightAllocations       map[types.UID]*resourceapi.ResourceClaim
+	}{
+		"empty": {
+			claimUID:                          claimUID,
+			shareable:                         true,
+			expected:                          false,
+			expectedInFlightAllocations:       map[types.UID]*resourceapi.ResourceClaim{},
+			expectedInFlightAllocationSharers: map[types.UID]int{},
+		},
+		"delete-existing-not-shareable": {
+			inFlightAllocations: map[types.UID]*resourceapi.ResourceClaim{
+				claimUID: allocatedClaim,
+			},
+			claimUID:                          claimUID,
+			expected:                          true,
+			expectedInFlightAllocations:       map[types.UID]*resourceapi.ResourceClaim{},
+			expectedInFlightAllocationSharers: map[types.UID]int{},
+		},
+		"delete-existing-shareable-unshared": {
+			inFlightAllocations: map[types.UID]*resourceapi.ResourceClaim{
+				claimUID: allocatedClaim,
+			},
+			inFlightAllocationSharers: map[types.UID]int{
+				otherClaimUID: 10,
+			},
+			claimUID:                    claimUID,
+			shareable:                   true,
+			expected:                    true,
+			expectedInFlightAllocations: map[types.UID]*resourceapi.ResourceClaim{},
+			expectedInFlightAllocationSharers: map[types.UID]int{
+				otherClaimUID: 10,
+			},
+		},
+		"keep-existing-shareable-shared": {
+			inFlightAllocations: map[types.UID]*resourceapi.ResourceClaim{
+				claimUID: allocatedClaim,
+			},
+			inFlightAllocationSharers: map[types.UID]int{
+				claimUID:      1,
+				otherClaimUID: 0,
+			},
+			claimUID:  claimUID,
+			shareable: true,
+			expected:  false,
+			expectedInFlightAllocations: map[types.UID]*resourceapi.ResourceClaim{
+				claimUID: allocatedClaim,
+			},
+			expectedInFlightAllocationSharers: map[types.UID]int{
+				claimUID:      1,
+				otherClaimUID: 0,
+			},
+		},
+	}
+
+	for name, test := range tests {
+		tCtx.Run(name, func(tCtx ktesting.TContext) {
+			c := &claimTracker{
+				logger:                    tCtx.Logger(),
+				inFlightAllocations:       &sync.Map{},
+				inFlightAllocationSharers: &sync.Map{},
+			}
+
+			for key, value := range test.inFlightAllocations {
+				c.inFlightAllocations.Store(key, value)
+			}
+			for key, value := range test.inFlightAllocationSharers {
+				c.inFlightAllocationSharers.Store(key, value)
+			}
+			actual := c.MaybeRemoveClaimPendingAllocation(test.claimUID, test.shareable)
+			tCtx.Expect(actual).To(gomega.Equal(test.expected), "wrong value for deletion indicator")
+			actualInFlight := map[types.UID]*resourceapi.ResourceClaim{}
+			c.inFlightAllocations.Range(func(key, value any) bool {
+				actualInFlight[key.(types.UID)] = value.(*resourceapi.ResourceClaim)
+				return true
+			})
+			tCtx.Expect(actualInFlight).To(gomega.Equal(test.expectedInFlightAllocations))
+			actualSharers := map[types.UID]int{}
+			c.inFlightAllocationSharers.Range(func(key, value any) bool {
+				actualSharers[key.(types.UID)] = value.(int)
+				return true
+			})
+			tCtx.Expect(actualSharers).To(gomega.Equal(test.expectedInFlightAllocationSharers))
+		})
+	}
+}

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

@@ -448,9 +448,22 @@ func (pl *DynamicResources) PreFilter(ctx context.Context, state fwk.CycleState,
 			return nil, statusUnschedulable(logger, "resourceclaim in use", "pod", klog.KObj(pod), "resourceclaim", klog.KObj(claim))
 		}
 
-		if claim.Status.Allocation != nil {
-			if claim.Status.Allocation.NodeSelector != nil {
-				nodeSelector, err := nodeaffinity.NewNodeSelector(claim.Status.Allocation.NodeSelector)
+		allocation := claim.Status.Allocation
+		// Gang scheduling works by waiting for all Pods in a gang to succeed
+		// through the Reserve phase, then binding all of the Pods. Since a
+		// claim's allocation is not recorded in [stateData.claims] until
+		// PreBind, we look for a pending allocation made in Reserve here to
+		// avoid blocking this Pod on waiting for the allocation to complete.
+		if pl.isSchedulingPodGroup(pod) && allocation == nil {
+			if pendingAllocation, found := pl.draManager.ResourceClaims().GetPendingAllocation(claim.UID); found {
+				allocation = pendingAllocation
+				logger.V(5).Info("reusing pending allocation", "pod", klog.KObj(pod), "resourceclaim", klog.KObj(claim), "uid", claim.UID, "allocation", klog.Format(allocation))
+			}
+		}
+
+		if allocation != nil {
+			if allocation.NodeSelector != nil {
+				nodeSelector, err := nodeaffinity.NewNodeSelector(allocation.NodeSelector)
 				if err != nil {
 					return nil, statusError(logger, err)
 				}
@@ -462,7 +475,7 @@ func (pl *DynamicResources) PreFilter(ctx context.Context, state fwk.CycleState,
 			// Allocation in flight? Better wait for that
 			// to finish, see inFlightAllocations
 			// documentation for details.
-			if pl.draManager.ResourceClaims().ClaimHasPendingAllocation(claim.UID) {
+			if _, pending := pl.draManager.ResourceClaims().GetPendingAllocation(claim.UID); pending {
 				return nil, statusUnschedulable(logger, fmt.Sprintf("resource claim %s is in the process of being allocated", klog.KObj(claim)))
 			}
 
@@ -961,6 +974,19 @@ func (pl *DynamicResources) Reserve(ctx context.Context, cs fwk.CycleState, pod
 			// updating the ResourceClaim status, we assume that reserving
 			// will work and only do it for real during binding. If it fails at
 			// that time, some other pod was faster and we have to try again.
+
+			if pl.isSchedulingPodGroup(pod) && claim.UID != state.claims.getInitialExtendedResourceClaimUID() {
+				// Inform the tracker that this pod still needs the pending
+				// allocation. Then PreBind or Unreserve can know whether or not
+				// it's safe to remove the pending allocation.
+				//
+				// Extended resources claims cannot be shared, so we never
+				// signal sharing for it.
+				if err := pl.draManager.ResourceClaims().AddSharedClaimPendingAllocation(claim.UID, claim); err != nil {
+					return statusError(logger, err)
+				}
+			}
+
 			continue
 		}
 
@@ -1054,9 +1080,21 @@ func (pl *DynamicResources) Unreserve(ctx context.Context, cs fwk.CycleState, po
 
 	// we process user claims here first, extendedResourceClaim if any is handled below.
 	for _, claim := range state.claims.allUserClaims() {
-		// If allocation was in-flight, then it's not anymore and we need to revert the
-		// claim object in the assume cache to what it was before.
-		if deleted := pl.draManager.ResourceClaims().RemoveClaimPendingAllocation(claim.UID); deleted {
+		// Since several Pods may be sharing a claim with a pending allocation,
+		// make sure that we don't remove the pending allocation until the last
+		// Pod sharing it is Bound or Unreserved.
+		if pl.isSchedulingPodGroup(pod) {
+			if err := pl.draManager.ResourceClaims().RemoveSharedClaimPendingAllocation(claim.UID, claim); err != nil {
+				logger.Error(err, "unreserve", "resourceclaim", klog.KObj(claim))
+				continue
+			}
+		}
+
+		// If allocation was in-flight, then it might not be anymore if no pods
+		// still need the pending allocation. If the allocation was removed from
+		// in-flight, we need to revert the claim object in the assume cache to
+		// what it was before.
+		if deleted := pl.draManager.ResourceClaims().MaybeRemoveClaimPendingAllocation(claim.UID, pl.isSchedulingPodGroup(pod)); deleted {
 			logger.V(5).Info("Released resource in allocation result", "claim", klog.KObj(claim), "uid", claim.UID, "resourceVersion", claim.ResourceVersion, "allocation", klog.Format(claim.Status.Allocation))
 			pl.draManager.ResourceClaims().AssumedClaimRestore(claim.Namespace, claim.Name)
 		}
@@ -1268,7 +1306,17 @@ func (pl *DynamicResources) bindClaim(ctx context.Context, state *stateData, ind
 		}
 		if allocation != nil {
 			for _, claimUID := range claimUIDs {
-				if deleted := pl.draManager.ResourceClaims().RemoveClaimPendingAllocation(claimUID); deleted {
+				// Since several Pods may be sharing a claim with a pending
+				// allocation, make sure that we don't remove the pending
+				// allocation until the last Pod sharing it is Bound or Unreserved.
+				if pl.isSchedulingPodGroup(pod) {
+					if err := pl.draManager.ResourceClaims().RemoveSharedClaimPendingAllocation(claim.UID, claim); err != nil {
+						logger.Error(err, "unreserve", "resourceclaim", klog.KObj(claim))
+						continue
+					}
+				}
+
+				if deleted := pl.draManager.ResourceClaims().MaybeRemoveClaimPendingAllocation(claimUID, pl.isSchedulingPodGroup(pod)); deleted {
 					logger.V(5).Info("Removed claim from in-flight claims", "claim", klog.KObj(claim), "uid", claimUID, "resourceVersion", resourceVersion, "allocation", klog.Format(allocation))
 				}
 			}
@@ -1466,6 +1514,10 @@ func (pl *DynamicResources) isPodReadyForBinding(state *stateData) (bool, error)
 	return true, nil
 }
 
+func (pl *DynamicResources) isSchedulingPodGroup(pod *v1.Pod) bool {
+	return pl.fts.EnableGenericWorkload && pod.Spec.SchedulingGroup != nil
+}
+
 // hasBindingConditions checks whether any of the claims in the state
 // has binding conditions.
 // It returns true if at least one claim has binding conditions.

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

@@ -597,7 +597,7 @@ func (pl *DynamicResources) unreserveExtendedResourceClaim(ctx context.Context,
 
 	// If the claim was marked as pending allocation (in-flight), remove that marker and restore
 	// the assumed claim state to what it was before this scheduling attempt.
-	if deleted := pl.draManager.ResourceClaims().RemoveClaimPendingAllocation(state.claims.getInitialExtendedResourceClaimUID()); deleted {
+	if deleted := pl.draManager.ResourceClaims().MaybeRemoveClaimPendingAllocation(state.claims.getInitialExtendedResourceClaimUID(), false); deleted {
 		pl.draManager.ResourceClaims().AssumedClaimRestore(extendedResourceClaim.Namespace, extendedResourceClaim.Name)
 	}
 	if isSpecialClaimName(extendedResourceClaim.Name) {

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

@@ -83,10 +83,22 @@ func (m *mockDRAManager) SignalClaimPendingAllocation(uid types.UID, claim *reso
 	return nil
 }
 
-func (m *mockDRAManager) RemoveClaimPendingAllocation(uid types.UID) bool {
+func (m *mockDRAManager) GetPendingAllocation(claimUID types.UID) (*resourceapi.AllocationResult, bool) {
+	return nil, false
+}
+
+func (m *mockDRAManager) MaybeRemoveClaimPendingAllocation(claimUID types.UID, shareable bool) (deleted bool) {
 	return false
 }
 
+func (m *mockDRAManager) AddSharedClaimPendingAllocation(claimUID types.UID, allocatedClaim *resourceapi.ResourceClaim) error {
+	return nil
+}
+
+func (m *mockDRAManager) RemoveSharedClaimPendingAllocation(claimUID types.UID, allocatedClaim *resourceapi.ResourceClaim) error {
+	return nil
+}
+
 func (m *mockDRAManager) List() ([]*resourceapi.ResourceClaim, error) {
 	return nil, nil
 }

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

@@ -48,6 +48,7 @@ type Features struct {
 	EnableStorageCapacityScoring                  bool
 	EnableNodeDeclaredFeatures                    bool
 	EnableGangScheduling                          bool
+	EnableGenericWorkload                         bool
 	EnableTaintTolerationComparisonOperators      bool
 	EnableInPlacePodLevelResourcesVerticalScaling bool
 	EnableTopologyAwareWorkloadScheduling         bool
@@ -78,6 +79,7 @@ func NewSchedulerFeaturesFromGates(featureGate featuregate.FeatureGate) Features
 		EnableStorageCapacityScoring:                  featureGate.Enabled(features.StorageCapacityScoring),
 		EnableNodeDeclaredFeatures:                    featureGate.Enabled(features.NodeDeclaredFeatures),
 		EnableGangScheduling:                          featureGate.Enabled(features.GangScheduling),
+		EnableGenericWorkload:                         featureGate.Enabled(features.GenericWorkload),
 		EnableTaintTolerationComparisonOperators:      featureGate.Enabled(features.TaintTolerationComparisonOperators),
 		EnableInPlacePodLevelResourcesVerticalScaling: featureGate.Enabled(features.InPlacePodLevelResourcesVerticalScaling),
 		EnableTopologyAwareWorkloadScheduling:         featureGate.Enabled(features.TopologyAwareWorkloadScheduling),

staging/src/k8s.io/kube-scheduler/framework/listers.go

@@ -106,13 +106,23 @@ type ResourceClaimTracker interface {
 	// 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.
 	SignalClaimPendingAllocation(claimUID types.UID, allocatedClaim *resourceapi.ResourceClaim) error
-	// ClaimHasPendingAllocation answers whether a given claim has a pending allocation during the binding phase. It can be used to avoid
+	// GetPendingAllocation answers whether a given claim has a pending allocation during the binding phase. It can be used to avoid
 	// race conditions in subsequent scheduling phases.
-	ClaimHasPendingAllocation(claimUID types.UID) bool
-	// RemoveClaimPendingAllocation removes the pending allocation for the given ResourceClaim from the tracker if any was signaled via
-	// SignalClaimPendingAllocation(). Returns whether there was a pending allocation to remove. List() and the other accessors immediately
-	// stop reflecting the pending allocation in the results.
-	RemoveClaimPendingAllocation(claimUID types.UID) (deleted bool)
+	GetPendingAllocation(claimUID types.UID) (*resourceapi.AllocationResult, bool)
+	// MaybeRemoveClaimPendingAllocation might remove the pending allocation for the given ResourceClaim from the tracker if any was signaled via
+	// SignalClaimPendingAllocation(). When a pending allocation is `shareable`, it removes the pending allocation only when
+	// no other pods are still using that pending allocation (per AddSharedClaimPendingAllocation and
+	// RemoveSharedClaimPendingAllocation). When a pending allocation is not shareable, it always removes the pending
+	// allocation as long as one exists. Returns whether there was a pending allocation and it was removed.
+	// List() and the other accessors immediately stop reflecting the pending allocation in the results.
+	MaybeRemoveClaimPendingAllocation(claimUID types.UID, shareable bool) (deleted bool)
+
+	// AddSharedClaimPendingAllocation increments the number of active sharers
+	// of the given ResourceClaim.
+	AddSharedClaimPendingAllocation(claimUID types.UID, allocatedClaim *resourceapi.ResourceClaim) error
+	// RemoveSharedClaimPendingAllocation decrements the number of active sharers
+	// of the given ResourceClaim.
+	RemoveSharedClaimPendingAllocation(claimUID types.UID, allocatedClaim *resourceapi.ResourceClaim) error
 
 	// AssumeClaimAfterAPICall signals to the tracker that an API call modifying the given ResourceClaim was made in the binding phase, and the
 	// changes should be reflected in informers very soon. This change is immediately reflected in the result of List() and the other accessors.

test/integration/dra/dra.go

@@ -30,6 +30,7 @@ import (
 	resourcealphaapi "k8s.io/api/resource/v1alpha3"
 	resourcev1beta1 "k8s.io/api/resource/v1beta1"
 	resourcev1beta2 "k8s.io/api/resource/v1beta2"
+	schedulingapi "k8s.io/api/scheduling/v1alpha2"
 	"k8s.io/apimachinery/pkg/api/resource"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/apimachinery/pkg/runtime/schema"
@@ -209,6 +210,7 @@ func run(tCtx ktesting.TContext, whatRE string) {
 				resourcev1beta1.SchemeGroupVersion:  true,
 				resourcev1beta2.SchemeGroupVersion:  true,
 				resourcealphaapi.SchemeGroupVersion: true,
+				schedulingapi.SchemeGroupVersion:    true,
 			},
 			features: map[featuregate.Feature]bool{
 				// Additional DRA feature gates go here,
@@ -223,6 +225,8 @@ func run(tCtx ktesting.TContext, whatRE string) {
 				features.DRAResourceClaimDeviceStatus: true,
 				features.DRAExtendedResource:          true,
 				features.DRANodeAllocatableResources:  true,
+				features.GangScheduling:               true,
+				features.GenericWorkload:              true,
 			},
 			f: func(tCtx ktesting.TContext) {
 				// These tests must run in parallel as much as possible to keep overall runtime low!
@@ -250,6 +254,7 @@ func run(tCtx ktesting.TContext, whatRE string) {
 				runSubTest(tCtx, "ShareResourceClaimSequentially", testShareResourceClaimSequentially)
 				runSubTest(tCtx, "UsesAllResources", testUsesAllResources)
 				runSubTest(tCtx, "DRANodeAllocatableResources", func(tCtx ktesting.TContext) { testNodeAllocatableResources(tCtx, true) })
+				runSubTest(tCtx, "PodGroup", testPodGroup)
 			},
 		},
 	} {

test/integration/dra/pod_group.go

@@ -0,0 +1,142 @@
+/*
+Copyright 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 dra
+
+import (
+	"fmt"
+
+	v1 "k8s.io/api/core/v1"
+	schedulingapi "k8s.io/api/scheduling/v1alpha2"
+	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+	st "k8s.io/kubernetes/pkg/scheduler/testing"
+	"k8s.io/kubernetes/test/utils/ktesting"
+)
+
+func testPodGroup(tCtx ktesting.TContext) {
+	tCtx.Parallel()
+
+	podGroupName := "podgroup"
+	tests := map[string]struct {
+		numPods  int
+		podGroup *schedulingapi.PodGroup
+	}{
+		"gang-2-pods-mincount-2": {
+			numPods: 2,
+			podGroup: &schedulingapi.PodGroup{
+				ObjectMeta: metav1.ObjectMeta{
+					Name: podGroupName,
+				},
+				Spec: schedulingapi.PodGroupSpec{
+					SchedulingPolicy: schedulingapi.PodGroupSchedulingPolicy{
+						Gang: &schedulingapi.GangSchedulingPolicy{
+							MinCount: 2,
+						},
+					},
+				},
+			},
+		},
+		"gang-20-pods-mincount-20": {
+			numPods: 20,
+			podGroup: &schedulingapi.PodGroup{
+				ObjectMeta: metav1.ObjectMeta{
+					Name: podGroupName,
+				},
+				Spec: schedulingapi.PodGroupSpec{
+					SchedulingPolicy: schedulingapi.PodGroupSchedulingPolicy{
+						Gang: &schedulingapi.GangSchedulingPolicy{
+							MinCount: 20,
+						},
+					},
+				},
+			},
+		},
+		"gang-5-pods-mincount-2": {
+			numPods: 5,
+			podGroup: &schedulingapi.PodGroup{
+				ObjectMeta: metav1.ObjectMeta{
+					Name: podGroupName,
+				},
+				Spec: schedulingapi.PodGroupSpec{
+					SchedulingPolicy: schedulingapi.PodGroupSchedulingPolicy{
+						Gang: &schedulingapi.GangSchedulingPolicy{
+							MinCount: 2,
+						},
+					},
+				},
+			},
+		},
+		"basic-2-pods": {
+			numPods: 2,
+			podGroup: &schedulingapi.PodGroup{
+				ObjectMeta: metav1.ObjectMeta{
+					Name: podGroupName,
+				},
+				Spec: schedulingapi.PodGroupSpec{
+					SchedulingPolicy: schedulingapi.PodGroupSchedulingPolicy{
+						Basic: &schedulingapi.BasicSchedulingPolicy{},
+					},
+				},
+			},
+		},
+		"basic-20-pods": {
+			numPods: 20,
+			podGroup: &schedulingapi.PodGroup{
+				ObjectMeta: metav1.ObjectMeta{
+					Name: podGroupName,
+				},
+				Spec: schedulingapi.PodGroupSpec{
+					SchedulingPolicy: schedulingapi.PodGroupSchedulingPolicy{
+						Basic: &schedulingapi.BasicSchedulingPolicy{},
+					},
+				},
+			},
+		},
+	}
+
+	for name, test := range tests {
+		tCtx.Run(name, func(tCtx ktesting.TContext) {
+			tCtx.Parallel()
+
+			namespace := createTestNamespace(tCtx, nil)
+			startScheduler(tCtx)
+
+			class, driverName := createTestClass(tCtx, namespace)
+			slice := st.MakeResourceSlice("worker-0", driverName).Devices("device-0")
+			createSlice(tCtx, slice.Obj())
+
+			claim := createClaim(tCtx, namespace, "", class, claim)
+
+			podGroup, err := tCtx.Client().SchedulingV1alpha2().PodGroups(namespace).Create(tCtx, test.podGroup, metav1.CreateOptions{})
+			tCtx.ExpectNoError(err, "create PodGroup")
+			schedGroup := &v1.PodSchedulingGroup{
+				PodGroupName: &podGroup.Name,
+			}
+
+			pods := make([]*v1.Pod, test.numPods)
+			for i := range pods {
+				pod := podWithClaimName.DeepCopy()
+				pod.Spec.SchedulingGroup = schedGroup
+				pods[i] = createPod(tCtx, namespace, fmt.Sprintf("-%d", i), pod, claim)
+			}
+
+			waitForClaimAllocatedToDevice(tCtx, namespace, claim.Name, schedulingTimeout)
+			for _, pod := range pods {
+				waitForPodScheduled(tCtx, namespace, pod.Name)
+			}
+		})
+	}
+}