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

@@ -454,25 +454,47 @@ func TestConcurrentNodeLocks(t *testing.T) {
 	}
 }
 
+// TestSetNodeLockRaceIsRetryable covers the narrow race window where two
+// callers both observe an unlocked node in LockNode's initial check and race
+// on SetNodeLock's per-node mutex.
+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)
+	}
+	winner := &corev1.Pod{ObjectMeta: metav1.ObjectMeta{Name: "winner", Namespace: "test-ns"}}
+	loser := &corev1.Pod{ObjectMeta: metav1.ObjectMeta{Name: "loser", Namespace: "test-ns"}}
+
+	if err := SetNodeLock(nodeName, "", winner); err != nil {
+		t.Fatalf("winner SetNodeLock failed: %v", err)
+	}
+	err = SetNodeLock(nodeName, "", loser)
+	if err == nil {
+		t.Fatal("expected loser's SetNodeLock to fail, got nil")
+	}
+	if !IsNodeLockContention(err) {
+		t.Fatalf("expected loser's error to be classified as node lock contention (retryable), got: %v", err)
+	}
+}
+
 // 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()