Fix scheduler pod group status synchronization between incoming update and in-cluster data - v0.4 cherrypick (#223)

* Keep updated pod-groups data in a separate syncmap

Author: davidLif

CHANGELOG.md

@@ -4,6 +4,11 @@ All notable changes to this project will be documented in this file.
 
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/).
 
+## [v0.4.10] - 2025-06-09
+
+### Fixed
+- Fix scheduler pod group status synchronization between incoming update and in-cluster data
+
 ## [v0.4.9] - 2025-05-27
 
 ### Fixed

pkg/scheduler/cache/status_updater/concurrency.go

@@ -5,17 +5,14 @@ package status_updater
 
 import (
 	"context"
-	"strconv"
 
 	v1 "k8s.io/api/core/v1"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/apimachinery/pkg/runtime"
 	"k8s.io/apimachinery/pkg/types"
 
 	enginev2alpha2 "github.com/NVIDIA/KAI-scheduler/pkg/apis/scheduling/v2alpha2"
-	commonconstants "github.com/NVIDIA/KAI-scheduler/pkg/common/constants"
 	"github.com/NVIDIA/KAI-scheduler/pkg/scheduler/log"
-	"github.com/NVIDIA/KAI-scheduler/pkg/scheduler/utils"
 )
 
 func (su *defaultStatusUpdater) Run(stopCh <-chan struct{}) {
@@ -25,58 +22,6 @@ func (su *defaultStatusUpdater) Run(stopCh <-chan struct{}) {
 	go su.queueBufferWorker(stopCh)
 }
 
-func (su *defaultStatusUpdater) SyncPodGroupsWithPendingUpdates(podGroups []*enginev2alpha2.PodGroup) {
-	usedKeys := make(map[updatePayloadKey]bool, len(podGroups))
-	for i := range podGroups {
-		key := su.keyForPodGroupPayload(podGroups[i].Name, podGroups[i].Namespace, podGroups[i].UID)
-		usedKeys[key] = true
-		inflightUpdateAny, found := su.inFlightPodGroups.Load(key)
-		if !found {
-			continue
-		}
-		podGroup := inflightUpdateAny.(*inflightUpdate).object.(*enginev2alpha2.PodGroup)
-		isPodGroupUpdated := su.syncPodGroup(podGroup, podGroups[i])
-		if isPodGroupUpdated {
-			su.inFlightPodGroups.Delete(key)
-		}
-	}
-
-	// Cleanup podGroups that don't comeup anymore
-	su.inFlightPodGroups.Range(func(key any, _ any) bool {
-		if _, found := usedKeys[key.(updatePayloadKey)]; !found {
-			su.inFlightPodGroups.Delete(key)
-		}
-		return true
-	})
-}
-
-func (su *defaultStatusUpdater) syncPodGroup(inFlightPodGroup, snapshotPodGroup *enginev2alpha2.PodGroup) bool {
-	staleTimeStampUpdated := false
-	if snapshotPodGroup.Annotations[commonconstants.StalePodgroupTimeStamp] == inFlightPodGroup.Annotations[commonconstants.StalePodgroupTimeStamp] {
-		staleTimeStampUpdated = true
-	} else {
-		if snapshotPodGroup.Annotations == nil {
-			snapshotPodGroup.Annotations = make(map[string]string)
-		}
-		snapshotPodGroup.Annotations[commonconstants.StalePodgroupTimeStamp] = inFlightPodGroup.Annotations[commonconstants.StalePodgroupTimeStamp]
-	}
-
-	updatedSchedulingCondition := false
-	lastSchedulingCondition := utils.GetLastSchedulingCondition(inFlightPodGroup)
-	currentLastSchedulingCondition := utils.GetLastSchedulingCondition(snapshotPodGroup)
-	if currentLastSchedulingCondition != nil && lastSchedulingCondition != nil {
-		currentID, currentErr := strconv.Atoi(currentLastSchedulingCondition.TransitionID)
-		lastID, lastErr := strconv.Atoi(lastSchedulingCondition.TransitionID)
-		if currentErr == nil && lastErr == nil && (lastID <= currentID || currentID == 0) {
-			updatedSchedulingCondition = true
-		}
-	}
-	if !updatedSchedulingCondition {
-		snapshotPodGroup.Status.SchedulingConditions = inFlightPodGroup.Status.SchedulingConditions
-	}
-	return staleTimeStampUpdated && updatedSchedulingCondition
-}
-
 func (su *defaultStatusUpdater) keyForPodGroupPayload(name, namespace string, uid types.UID) updatePayloadKey {
 	return updatePayloadKey(types.NamespacedName{Name: name, Namespace: namespace}.String() + "_" + string(uid))
 }
@@ -99,7 +44,7 @@ func (su *defaultStatusUpdater) processPayload(ctx context.Context, payload *upd
 	case podType:
 		su.updatePod(ctx, payload.key, updateData.patchData, updateData.subResources, updateData.object)
 	case podGroupType:
-		su.updatePodGroup(ctx, payload.key, updateData.patchData, updateData.subResources, updateData.updateStatus, updateData.object)
+		su.updatePodGroup(ctx, payload.key, updateData)
 	}
 }
 
@@ -137,23 +82,27 @@ func (su *defaultStatusUpdater) updatePod(
 // +kubebuilder:rbac:groups="scheduling.run.ai",resources=podgroups/status,verbs=create;delete;update;patch;get;list;watch
 
 func (su *defaultStatusUpdater) updatePodGroup(
-	ctx context.Context, _ updatePayloadKey, patchData []byte, subResources []string, updateStatus bool, object runtime.Object,
+	ctx context.Context, key updatePayloadKey, updateData *inflightUpdate,
 ) {
-	podGroup := object.(*enginev2alpha2.PodGroup)
+	podGroup := updateData.object.(*enginev2alpha2.PodGroup)
 
 	var err error
-	if updateStatus {
+	if updateData.updateStatus {
 		_, err = su.kubeaischedClient.SchedulingV2alpha2().PodGroups(podGroup.Namespace).UpdateStatus(
 			ctx, podGroup, metav1.UpdateOptions{},
 		)
 	}
-	if len(patchData) > 0 {
+	if len(updateData.patchData) > 0 {
 		_, err = su.kubeaischedClient.SchedulingV2alpha2().PodGroups(podGroup.Namespace).Patch(
-			ctx, podGroup.Name, types.JSONPatchType, patchData, metav1.PatchOptions{}, subResources...,
+			ctx, podGroup.Name, types.JSONPatchType, updateData.patchData, metav1.PatchOptions{}, updateData.subResources...,
 		)
 	}
 	if err != nil {
 		log.StatusUpdaterLogger.Errorf("Failed to update pod group %s/%s: %v", podGroup.Namespace, podGroup.Name, err)
+	} else {
+		// Move the update to the applied cache
+		su.appliedPodGroupUpdates.Store(key, updateData)
+		su.inFlightPodGroups.Delete(key)
 	}
 }
 

pkg/scheduler/cache/status_updater/concurrency_test.go

@@ -37,7 +37,7 @@ func TestConcurrency(t *testing.T) {
 	RunSpecs(t, "Status Updater Concurrency Suite")
 }
 
-var _ = Describe("Status Updater Concurrency", func() {
+var _ = Describe("Status Updater Concurrency - large scale: increase queue size", func() {
 	var (
 		kubeClient        *fake.Clientset
 		kubeAiSchedClient *kubeaischedfake.Clientset

pkg/scheduler/cache/status_updater/default_status_updater.go

@@ -68,6 +68,8 @@ type defaultStatusUpdater struct {
 
 	inFlightPodGroups sync.Map
 	inFlightPods      sync.Map
+
+	appliedPodGroupUpdates sync.Map
 }
 
 // +kubebuilder:rbac:groups="",resources=events,verbs=create;update;patch;delete;list;get;watch

pkg/scheduler/cache/status_updater/pod_group_status_sync.go

@@ -0,0 +1,128 @@
+// Copyright 2025 NVIDIA CORPORATION
+// SPDX-License-Identifier: Apache-2.0
+
+package status_updater
+
+import (
+	"strconv"
+
+	enginev2alpha2 "github.com/NVIDIA/KAI-scheduler/pkg/apis/scheduling/v2alpha2"
+	commonconstants "github.com/NVIDIA/KAI-scheduler/pkg/common/constants"
+	"github.com/NVIDIA/KAI-scheduler/pkg/scheduler/utils"
+)
+
+type podGroupStatusSyncResult string
+
+const (
+	snapshotStatusIsOlder podGroupStatusSyncResult = "snapshotStatusIsOlder"
+	equalStatuses         podGroupStatusSyncResult = "equalStatuses"
+	updateRequestIsOlder  podGroupStatusSyncResult = "updateRequestIsOlder"
+)
+
+func (su *defaultStatusUpdater) SyncPodGroupsWithPendingUpdates(podGroups []*enginev2alpha2.PodGroup) {
+	usedKeys := make(map[updatePayloadKey]bool, len(podGroups))
+	for i := range podGroups {
+		key := su.keyForPodGroupPayload(podGroups[i].Name, podGroups[i].Namespace, podGroups[i].UID)
+		usedKeys[key] = true
+		pgLatestUpdate, inFlightUpdateFound, appliedUpdateFound := su.getLatestPgUpdate(key)
+		if !inFlightUpdateFound && !appliedUpdateFound {
+			continue
+		}
+		podGroup := pgLatestUpdate.object.(*enginev2alpha2.PodGroup)
+		podGroupsyncResults := su.syncPodGroup(podGroup, podGroups[i])
+		// Delete the inflight update if it was applied + the pod group in the lister matches the inFlight
+		if podGroupsyncResults != snapshotStatusIsOlder && appliedUpdateFound {
+			su.appliedPodGroupUpdates.Delete(key)
+		} else if podGroupsyncResults == updateRequestIsOlder && inFlightUpdateFound {
+			su.inFlightPodGroups.Delete(key)
+		}
+	}
+
+	// Cleanup podGroups that don't comeup anymore
+	su.cleanUpdatesForNonSeenPodGroups(usedKeys)
+}
+
+func (su *defaultStatusUpdater) getLatestPgUpdate(key updatePayloadKey) (*inflightUpdate, bool, bool) {
+	inflightPgUpdate, inFlightUpdateFound := su.inFlightPodGroups.Load(key)
+	appliedPgUpdate, appliedUpdateFound := su.appliedPodGroupUpdates.Load(key)
+	var pgLatestUpdate *inflightUpdate
+	if inFlightUpdateFound {
+		pgLatestUpdate = inflightPgUpdate.(*inflightUpdate)
+	} else if appliedUpdateFound {
+		pgLatestUpdate = appliedPgUpdate.(*inflightUpdate)
+	}
+	return pgLatestUpdate, inFlightUpdateFound, appliedUpdateFound
+}
+
+func (su *defaultStatusUpdater) cleanUpdatesForNonSeenPodGroups(usedKeys map[updatePayloadKey]bool) {
+	su.inFlightPodGroups.Range(func(key any, _ any) bool {
+		if _, found := usedKeys[key.(updatePayloadKey)]; !found {
+			su.inFlightPodGroups.Delete(key)
+		}
+		return true
+	})
+	su.appliedPodGroupUpdates.Range(func(key any, _ any) bool {
+		if _, found := usedKeys[key.(updatePayloadKey)]; !found {
+			su.appliedPodGroupUpdates.Delete(key)
+		}
+		return true
+	})
+}
+
+func (su *defaultStatusUpdater) syncPodGroup(inFlightPodGroup, snapshotPodGroup *enginev2alpha2.PodGroup) podGroupStatusSyncResult {
+	staleTimeStampUpdated := false
+	if snapshotPodGroup.Annotations[commonconstants.StalePodgroupTimeStamp] == inFlightPodGroup.Annotations[commonconstants.StalePodgroupTimeStamp] {
+		staleTimeStampUpdated = true
+	} else {
+		if snapshotPodGroup.Annotations == nil {
+			snapshotPodGroup.Annotations = make(map[string]string)
+		}
+		snapshotPodGroup.Annotations[commonconstants.StalePodgroupTimeStamp] = inFlightPodGroup.Annotations[commonconstants.StalePodgroupTimeStamp]
+	}
+
+	statusComparison := compareSchedulingConditions(inFlightPodGroup, snapshotPodGroup)
+
+	if statusComparison == equalStatuses || statusComparison == snapshotStatusIsOlder {
+		snapshotPodGroup.Status.SchedulingConditions = inFlightPodGroup.Status.SchedulingConditions
+	}
+	if statusComparison == equalStatuses && !staleTimeStampUpdated {
+		statusComparison = snapshotStatusIsOlder
+	}
+
+	return statusComparison
+}
+
+func compareSchedulingConditions(inFlightPodGroup, snapshotPodGroup *enginev2alpha2.PodGroup) podGroupStatusSyncResult {
+	lastSchedulingCondition := utils.GetLastSchedulingCondition(inFlightPodGroup)
+	currentLastSchedulingCondition := utils.GetLastSchedulingCondition(snapshotPodGroup)
+
+	if currentLastSchedulingCondition == nil && lastSchedulingCondition == nil {
+		return equalStatuses
+	}
+	if currentLastSchedulingCondition == nil {
+		return snapshotStatusIsOlder
+	}
+	if lastSchedulingCondition == nil {
+		return updateRequestIsOlder
+	}
+
+	currentID, currentErr := strconv.Atoi(currentLastSchedulingCondition.TransitionID)
+	lastID, lastErr := strconv.Atoi(lastSchedulingCondition.TransitionID)
+
+	// If an ID is invalid, use the other condition as the "newer" one
+	if currentErr != nil {
+		return snapshotStatusIsOlder
+	} else if lastErr != nil {
+		return updateRequestIsOlder
+	}
+
+	// Compare valid IDs
+	switch {
+	case lastID > currentID || currentID == 0:
+		return snapshotStatusIsOlder
+	case lastID == currentID:
+		return equalStatuses
+	default: // lastID < currentID
+		return updateRequestIsOlder
+	}
+}