Fix: adopt to gangScheduling feature gates on k8s v1.36+

pkg/util/nodelock/nodelock.go

@@ -138,7 +138,7 @@ func SetNodeLock(nodeName string, lockname string, pods *corev1.Pod) error {
 		return err
 	}
 	if _, ok := node.Annotations[NodeLockKey]; ok {
-		return fmt.Errorf("node %s is locked", nodeName)
+		return fmt.Errorf("node %s is locked: %w", nodeName, ErrNodeLockContention)
 	}
 	err = retry.OnError(DefaultStrategy, func(err error) bool {
 		// Retry on any error

pkg/util/nodelock/nodelock_test.go

@@ -20,6 +20,7 @@ import (
 	"context" // Added for the new test
 	"runtime"
 	"strings"
+	"sync"
 	"testing"
 	"time"
 
@@ -454,25 +455,65 @@ func TestConcurrentNodeLocks(t *testing.T) {
 	}
 }
 
+// TestSetNodeLockRaceIsRetryable covers the narrow race window where two
+// callers both observe an unlocked node in LockNode's outer check and then
+// genuinely race on SetNodeLock's per-node mutex to actually claim the lock.
+func TestSetNodeLockRaceIsRetryable(t *testing.T) {
+	client.KubeClient = fake.NewClientset()
+	nodeLocks = newNodeLockManager()
+	nodeName := "race-node"
+	_, err := client.KubeClient.CoreV1().Nodes().Create(context.TODO(), &corev1.Node{
+		ObjectMeta: metav1.ObjectMeta{Name: nodeName, Annotations: map[string]string{}},
+	}, metav1.CreateOptions{})
+	if err != nil {
+		t.Fatalf("Failed to create node: %v", err)
+	}
+	podA := &corev1.Pod{ObjectMeta: metav1.ObjectMeta{Name: "pod-a", Namespace: "test-ns"}}
+	podB := &corev1.Pod{ObjectMeta: metav1.ObjectMeta{Name: "pod-b", Namespace: "test-ns"}}
+	raceLock := nodeLocks.getLock(nodeName)
+	raceLock.Lock()
+	results := make(chan error, 2)
+	var wg sync.WaitGroup
+	for _, pod := range []*corev1.Pod{podA, podB} {
+		wg.Add(1)
+		go func(pod *corev1.Pod) {
+			defer wg.Done()
+			results <- LockNode(nodeName, "", pod)
+		}(pod)
+	}
+	time.Sleep(50 * time.Millisecond)
+	raceLock.Unlock()
+	wg.Wait()
+	close(results)
+	var successes, contentions int
+	for err := range results {
+		switch {
+		case err == nil:
+			successes++
+		case IsNodeLockContention(err):
+			contentions++
+		default:
+			t.Fatalf("unexpected error from LockNode: %v", err)
+		}
+	}
+	if successes != 1 || contentions != 1 {
+		t.Fatalf("expected exactly 1 winner and 1 retryable contention out of 2 concurrent LockNode calls, got successes=%d contentions=%d", successes, contentions)
+	}
+}
+
 // TestCleanupNodeLockOnNodeDelete ensures CleanupNodeLock removes the entry
 // and a subsequent getLock allocates a fresh mutex instance.
 func TestCleanupNodeLockOnNodeDelete(t *testing.T) {
-	// Reset manager state for this test
 	nodeLocks = newNodeLockManager()
-
 	first := nodeLocks.getLock("to-be-deleted")
 	if first == nil {
 		t.Fatalf("expected non-nil mutex from getLock")
 	}
-
-	// Trigger cleanup as if node was removed by autoscaler
 	CleanupNodeLock("to-be-deleted")
-
 	second := nodeLocks.getLock("to-be-deleted")
 	if second == nil {
 		t.Fatalf("expected non-nil mutex from getLock after cleanup")
 	}
-
 	if first == second {
 		t.Fatalf("expected a new mutex instance after cleanup, got the same pointer")
 	}

pkg/util/util.go

@@ -281,10 +281,15 @@ func AllContainersCreated(pod *corev1.Pod) bool {
 	return len(pod.Status.ContainerStatuses) >= len(pod.Spec.Containers)
 }
 
-// Coscheduling PodGroup, based on the presence of the PodGroupLabel.
 func IsPodGroupMember(pod *corev1.Pod) bool {
 	if pod == nil {
 		return false
 	}
-	return pod.Labels[PodGroupLabel] != ""
+	if pod.Labels[PodGroupLabel] != "" {
+		return true
+	}
+	if sg := pod.Spec.SchedulingGroup; sg != nil && sg.PodGroupName != nil && *sg.PodGroupName != "" {
+		return true
+	}
+	return false
 }

pkg/util/util_test.go

@@ -518,6 +518,98 @@ func Test_AllContainersCreated(t *testing.T) {
 	}
 }
 
+func TestIsPodGroupMember(t *testing.T) {
+	podGroupName := "my-training-job"
+	emptyPodGroupName := ""
+
+	tests := []struct {
+		name string
+		pod  *corev1.Pod
+		want bool
+	}{
+		{
+			name: "nil pod",
+			pod:  nil,
+			want: false,
+		},
+		{
+			name: "no group membership at all",
+			pod:  &corev1.Pod{},
+			want: false,
+		},
+		{
+			name: "scheduler-plugins Coscheduling label present",
+			pod: &corev1.Pod{
+				ObjectMeta: metav1.ObjectMeta{
+					Labels: map[string]string{PodGroupLabel: podGroupName},
+				},
+			},
+			want: true,
+		},
+		{
+			name: "coscheduling label present but empty",
+			pod: &corev1.Pod{
+				ObjectMeta: metav1.ObjectMeta{
+					Labels: map[string]string{PodGroupLabel: ""},
+				},
+			},
+			want: false,
+		},
+		{
+			name: "native GenericWorkload PodGroup via Spec.SchedulingGroup",
+			pod: &corev1.Pod{
+				Spec: corev1.PodSpec{
+					SchedulingGroup: &corev1.PodSchedulingGroup{
+						PodGroupName: &podGroupName,
+					},
+				},
+			},
+			want: true,
+		},
+		{
+			name: "Spec.SchedulingGroup set but PodGroupName nil",
+			pod: &corev1.Pod{
+				Spec: corev1.PodSpec{
+					SchedulingGroup: &corev1.PodSchedulingGroup{},
+				},
+			},
+			want: false,
+		},
+		{
+			name: "Spec.SchedulingGroup set but PodGroupName empty",
+			pod: &corev1.Pod{
+				Spec: corev1.PodSpec{
+					SchedulingGroup: &corev1.PodSchedulingGroup{
+						PodGroupName: &emptyPodGroupName,
+					},
+				},
+			},
+			want: false,
+		},
+		{
+			name: "both coscheduling label and native SchedulingGroup present",
+			pod: &corev1.Pod{
+				ObjectMeta: metav1.ObjectMeta{
+					Labels: map[string]string{PodGroupLabel: podGroupName},
+				},
+				Spec: corev1.PodSpec{
+					SchedulingGroup: &corev1.PodSchedulingGroup{
+						PodGroupName: &podGroupName,
+					},
+				},
+			},
+			want: true,
+		},
+	}
+
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			got := IsPodGroupMember(test.pod)
+			assert.Equal(t, test.want, got)
+		})
+	}
+}
+
 func TestPatchPodLabels(t *testing.T) {
 	client.KubeClient = fake.NewClientset()