diff --git a/pkg/scheduler/backend/cache/cache.go b/pkg/scheduler/backend/cache/cache.go index 7d0a2e2e729..3a6cd16b02b 100644 --- a/pkg/scheduler/backend/cache/cache.go +++ b/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. diff --git a/pkg/scheduler/backend/cache/cache_test.go b/pkg/scheduler/backend/cache/cache_test.go index 5f17efa7b90..0a855eafe0d 100644 --- a/pkg/scheduler/backend/cache/cache_test.go +++ b/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), }, }, }, diff --git a/pkg/scheduler/backend/cache/podgroupstate.go b/pkg/scheduler/backend/cache/podgroupstate.go index a5c8b43dff8..dc908b99863 100644 --- a/pkg/scheduler/backend/cache/podgroupstate.go +++ b/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. diff --git a/pkg/scheduler/backend/cache/podgroupstate_test.go b/pkg/scheduler/backend/cache/podgroupstate_test.go index afec094a3db..9274a82977c 100644 --- a/pkg/scheduler/backend/cache/podgroupstate_test.go +++ b/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) diff --git a/pkg/scheduler/backend/cache/snapshot.go b/pkg/scheduler/backend/cache/snapshot.go index 358bc1177ae..2f5d0b7a3b8 100644 --- a/pkg/scheduler/backend/cache/snapshot.go +++ b/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 }