Store assumed pod object on podGroupState.assumePod

pkg/scheduler/backend/cache/cache.go

@@ -861,10 +861,10 @@ func (cache *cacheImpl) assumePodGroupMember(pod *v1.Pod) {
 	podGroupState, exists := cache.podGroupStates[key]
 	if !exists {
 		podGroupState = newPodGroupState()
-		cache.podGroupStates[key] = podGroupState
 		podGroupState.allPods[pod.UID] = pod
+		cache.podGroupStates[key] = podGroupState
 	}
-	podGroupState.assumePod(pod.UID)
+	podGroupState.assumePod(pod)
 }
 
 // forgetPodGroupMember moves the pod back from assumed to unscheduled in its pod group state.

pkg/scheduler/backend/cache/cache_test.go

@@ -1025,9 +1025,9 @@ func TestUpdatePodGroupStateSnapshot(t *testing.T) {
 	}
 }
 
-// Test_BindingPodGroupMember simulates binding and tests that when an assumed pod
+// TestBindingPodGroupMember simulates binding and tests that when an assumed pod
 // gets bound, its state within pod group transitions from assumed to assigned.
-func Test_BindingPodGroupMember(t *testing.T) {
+func TestBindingPodGroupMember(t *testing.T) {
 	logger, ctx := ktesting.NewTestContext(t)
 	cache := newCache(ctx, time.Second, nil, true)
 	podGroupName := "pg"
@@ -2153,7 +2153,7 @@ func TestSchedulerCache_UpdateSnapshot(t *testing.T) {
 						allPods:         map[types.UID]*v1.Pod{"puid-podgroup-0": podsWithPodGroupName[0]},
 						assignedPods:    sets.New[types.UID]("puid-podgroup-0"),
 						unscheduledPods: sets.New[types.UID](),
-						assumedPods:     sets.New[types.UID](),
+						assumedPods:     make(map[types.UID]*v1.Pod),
 					},
 				},
 				newPodGroupKey("test-ns", "pg-2"): {
@@ -2161,7 +2161,7 @@ func TestSchedulerCache_UpdateSnapshot(t *testing.T) {
 						allPods:         map[types.UID]*v1.Pod{"puid-podgroup-2": podsWithPodGroupName[2]},
 						assignedPods:    sets.New[types.UID]("puid-podgroup-2"),
 						unscheduledPods: sets.New[types.UID](),
-						assumedPods:     sets.New[types.UID](),
+						assumedPods:     make(map[types.UID]*v1.Pod),
 					},
 				},
 			},

pkg/scheduler/backend/cache/podgroupstate.go

@@ -74,7 +74,7 @@ type podGroupStateData struct {
 	unscheduledPods sets.Set[types.UID]
 	// assumedPods tracks pods that have reached the Reserve stage and are waiting
 	// for the rest of the gang to arrive before being allowed to bind.
-	assumedPods sets.Set[types.UID]
+	assumedPods map[types.UID]*v1.Pod
 	// assignedPods tracks all pods belonging to the group that are assigned (bound).
 	assignedPods sets.Set[types.UID]
 }
@@ -83,7 +83,7 @@ func newPodGroupStateData() podGroupStateData {
 	return podGroupStateData{
 		allPods:         make(map[types.UID]*v1.Pod),
 		unscheduledPods: sets.New[types.UID](),
-		assumedPods:     sets.New[types.UID](),
+		assumedPods:     make(map[types.UID]*v1.Pod),
 		assignedPods:    sets.New[types.UID](),
 	}
 }
@@ -109,7 +109,7 @@ func (d *podGroupStateData) updatePod(oldPod, newPod *v1.Pod) {
 		d.assignedPods.Insert(newPod.UID)
 		// Clear pod from unscheduled and assumed when it is assigned.
 		d.unscheduledPods.Delete(newPod.UID)
-		d.assumedPods.Delete(newPod.UID)
+		delete(d.assumedPods, newPod.UID)
 	}
 }
 
@@ -118,27 +118,26 @@ func (d *podGroupStateData) deletePod(podUID types.UID) {
 	d.generation = nextPodGroupGeneration()
 	delete(d.allPods, podUID)
 	d.unscheduledPods.Delete(podUID)
-	d.assumedPods.Delete(podUID)
+	delete(d.assumedPods, podUID)
 	d.assignedPods.Delete(podUID)
 }
 
 // assumePod marks a pod as assumed within the pod group state.
-func (d *podGroupStateData) assumePod(podUID types.UID) {
-
-	pod := d.allPods[podUID]
+func (d *podGroupStateData) assumePod(pod *v1.Pod) {
+	storedPod, ok := d.allPods[pod.UID]
 	// A scheduling pod may be removed from the cluster.
 	// In that case, we just ignore it.
-	if pod == nil {
+	if !ok {
 		return
 	}
 
 	d.generation = nextPodGroupGeneration()
-	// If the pod is already assigned, put it into assignedPods.
+	// If the pod stored in the state is already assigned, put it into assignedPods.
 	// Otherwise put it to assumedPods.
-	if pod.Spec.NodeName != "" {
+	if storedPod.Spec.NodeName != "" {
 		d.assignedPods.Insert(pod.UID)
 	} else {
-		d.assumedPods.Insert(pod.UID)
+		d.assumedPods[pod.UID] = pod
 	}
 	d.unscheduledPods.Delete(pod.UID)
 }
@@ -155,7 +154,7 @@ func (d *podGroupStateData) forgetPod(podUID types.UID) {
 
 	d.generation = nextPodGroupGeneration()
 
-	d.assumedPods.Delete(podUID)
+	delete(d.assumedPods, podUID)
 
 	// If the pod is already assigned, put it into assignedPods.
 	// Otherwise, put it into unscheduledPods.
@@ -172,8 +171,8 @@ func (d *podGroupStateData) scheduledPods() []*v1.Pod {
 	for uid := range d.assignedPods {
 		scheduledPods = append(scheduledPods, d.allPods[uid])
 	}
-	for uid := range d.assumedPods {
-		scheduledPods = append(scheduledPods, d.allPods[uid])
+	for _, pod := range d.assumedPods {
+		scheduledPods = append(scheduledPods, pod)
 	}
 	return scheduledPods
 }
@@ -199,7 +198,7 @@ func (d *podGroupStateData) deepCopy() podGroupStateData {
 		generation:      d.generation,
 		allPods:         maps.Clone(d.allPods),
 		unscheduledPods: d.unscheduledPods.Clone(),
-		assumedPods:     d.assumedPods.Clone(),
+		assumedPods:     maps.Clone(d.assumedPods),
 		assignedPods:    d.assignedPods.Clone(),
 	}
 }
@@ -274,7 +273,7 @@ func (pgs *podGroupState) AssumedPods() sets.Set[types.UID] {
 	pgs.lock.RLock()
 	defer pgs.lock.RUnlock()
 
-	return pgs.podGroupStateData.assumedPods.Clone()
+	return sets.KeySet(pgs.podGroupStateData.assumedPods)
 }
 
 // AssignedPods returns the UIDs of all pods already assigned (bound) for this group.
@@ -309,8 +308,8 @@ type podGroupStateSnapshot struct {
 }
 
 // assumePod marks a pod within the pod group state snapshot as assumed.
-func (s *podGroupStateSnapshot) assumePod(podUID types.UID) {
-	s.podGroupStateData.assumePod(podUID)
+func (s *podGroupStateSnapshot) assumePod(pod *v1.Pod) {
+	s.podGroupStateData.assumePod(pod)
 }
 
 // forgetPod removes a pod from the assumed state within the snapshot.
@@ -330,7 +329,7 @@ func (s *podGroupStateSnapshot) UnscheduledPods() map[string]*v1.Pod {
 
 // AssumedPods returns the UIDs of all assumed pods for this group.
 func (s *podGroupStateSnapshot) AssumedPods() sets.Set[types.UID] {
-	return s.podGroupStateData.assumedPods
+	return sets.KeySet(s.podGroupStateData.assumedPods)
 }
 
 // AssignedPods returns the UIDs of all assigned (bound) pods for this group.

pkg/scheduler/backend/cache/podgroupstate_test.go

@@ -36,7 +36,7 @@ func TestPodGroupState_AssumeForget(t *testing.T) {
 		t.Fatal("Pod should be initially in UnscheduledPods")
 	}
 
-	pgs.assumePod(pod.UID)
+	pgs.assumePod(pod)
 	if !pgs.AssumedPods().Has(pod.UID) {
 		t.Fatal("Pod should be in AssumedPods after AssumePod")
 	}
@@ -63,7 +63,7 @@ func TestPodGroupState_Clone(t *testing.T) {
 
 	pgs.addPod(pod1)
 	pgs.addPod(pod2)
-	pgs.assumePod(pod2.UID)
+	pgs.assumePod(pod2)
 
 	snap := pgs.snapshot()
 
@@ -88,8 +88,8 @@ func TestPodGroupState_Clone(t *testing.T) {
 	}
 
 	// Mutating the clone does not affect the original.
-	snap.assumePod(pod1.UID)
-	if pgs.assumedPods.Has(pod1.UID) {
+	snap.assumePod(pod1)
+	if _, ok := pgs.assumedPods[pod1.UID]; ok {
 		t.Error("mutation to clone should not affect original's assumedPods")
 	}
 
@@ -132,7 +132,7 @@ func TestPodGroupState_PodCounts(t *testing.T) {
 	}
 
 	// Assuming a pod should move it from unscheduled to assumed, increasing the count of scheduled pods.
-	pgs.assumePod(pod1.UID)
+	pgs.assumePod(pod1)
 	if count := pgs.AllPodsCount(); count != 3 {
 		t.Errorf("Expected AllPodsCount to be 3, got %d", count)
 	}
@@ -141,7 +141,7 @@ func TestPodGroupState_PodCounts(t *testing.T) {
 	}
 
 	// Assuming a pod that is already scheduled should not change the counts.
-	pgs.assumePod(pod3.UID)
+	pgs.assumePod(pod3)
 	if count := pgs.AllPodsCount(); count != 3 {
 		t.Errorf("Expected AllPodsCount to be 3, got %d", count)
 	}
@@ -150,7 +150,7 @@ func TestPodGroupState_PodCounts(t *testing.T) {
 	}
 
 	// Assuming a pod that is not in the state should not change the counts.
-	pgs.assumePod(pod4.UID)
+	pgs.assumePod(pod4)
 	if count := pgs.AllPodsCount(); count != 3 {
 		t.Errorf("Expected AllPodsCount to be 3, got %d", count)
 	}
@@ -187,7 +187,7 @@ func TestPodGroupState_PodCounts(t *testing.T) {
 	}
 
 	// Assuming a pod again should move it back to assumed, increasing the count of scheduled pods.
-	pgs.assumePod(pod2.UID)
+	pgs.assumePod(pod2)
 	if count := pgs.AllPodsCount(); count != 3 {
 		t.Errorf("Expected AllPodsCount to be 3, got %d", count)
 	}
@@ -221,14 +221,18 @@ func TestPodGroupState_ScheduledPods(t *testing.T) {
 	pgs.addPod(unscheduledPod)
 	pgs.addPod(assumedPod)
 
-	pgs.assumePod(assumedPod.UID)
+	// Simulate the scheduler assuming the pod on a node.
+	assumedPodWithNodeName := assumedPod.DeepCopy()
+	assumedPodWithNodeName.Spec.NodeName = "node2"
+
+	pgs.assumePod(assumedPodWithNodeName)
 	scheduledPods := pgs.ScheduledPods()
 
 	snapshot := pgs.snapshot()
-	pgs.assumePod(unscheduledPod.UID)
+	pgs.assumePod(unscheduledPod)
 	snapshotScheduledPods := snapshot.ScheduledPods()
 
-	expectedScheduledPods := []*v1.Pod{assignedPod, assumedPod}
+	expectedScheduledPods := []*v1.Pod{assignedPod, assumedPodWithNodeName}
 
 	if diff := cmp.Diff(expectedScheduledPods, scheduledPods); diff != "" {
 		t.Errorf("unexpected ScheduledPods result (-want,+got):\n%s", diff)

pkg/scheduler/backend/cache/snapshot.go

@@ -302,7 +302,7 @@ func (s *Snapshot) AssumePod(podInfo *framework.PodInfo) error {
 	}
 	pgKey := newPodGroupKey(pod.Namespace, *pod.Spec.SchedulingGroup.PodGroupName)
 	if pgs, ok := s.podGroupStates[pgKey]; ok {
-		pgs.assumePod(pod.UID)
+		pgs.assumePod(pod)
 	}
 	return nil
 }