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

* Keep updated pod-groups data in a separate syncmap (#199)

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.5.3] - 2025-06-09
+
+### Fixed
+- Fix scheduler pod group status synchronization between incoming update and in-cluster data
+
 ## [v0.5.2] - 2025-05-26
 
 ### 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,67 +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]
-	}
-	lastStartTimestampUpdated := false
-	if snapshotPodGroup.Annotations[commonconstants.LastStartTimeStamp] == inFlightPodGroup.Annotations[commonconstants.LastStartTimeStamp] {
-		lastStartTimestampUpdated = true
-	} else {
-		if snapshotPodGroup.Annotations == nil {
-			snapshotPodGroup.Annotations = make(map[string]string)
-		}
-		snapshotPodGroup.Annotations[commonconstants.LastStartTimeStamp] = inFlightPodGroup.Annotations[commonconstants.LastStartTimeStamp]
-	}
-
-	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 lastStartTimestampUpdated && staleTimeStampUpdated && updatedSchedulingCondition
-}
-
 func (su *defaultStatusUpdater) keyForPodGroupPayload(name, namespace string, uid types.UID) updatePayloadKey {
 	return updatePayloadKey(types.NamespacedName{Name: name, Namespace: namespace}.String() + "_" + string(uid))
 }
@@ -108,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)
 	}
 }
 
@@ -146,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)
 	}
 }