Add workload annotation to pods when SchedulerLibraryIntegration FG enabled.

hack/make/test.mk

@@ -276,7 +276,7 @@ test-e2e-extended-shard-0: setup-e2e-env run-test-e2e-extended-$(E2E_KIND_VERSIO
 
 .PHONY: test-e2e-extended-shard-1
 test-e2e-extended-shard-1: E2E_NPROCS := 4
-test-e2e-extended-shard-1: GINKGO_ARGS=--label-filter=feature:appwrapper,feature:jaxjob,feature:jobset,feature:leaderworkerset,feature:pytorchjob,feature:trainjob,feature:mpijob
+test-e2e-extended-shard-1: GINKGO_ARGS=--label-filter=feature:appwrapper,feature:jaxjob,feature:jobset,feature:leaderworkerset,feature:pytorchjob,feature:trainjob,feature:mpijob,feature:scheduler-library
 test-e2e-extended-shard-1: setup-e2e-env run-test-e2e-extended-$(E2E_KIND_VERSION:kindest/node:v%=%)
 
 .PHONY: test-e2e-extended-shard-2

pkg/controller/jobframework/reconciler.go

@@ -1739,7 +1739,7 @@ func getPodSetsInfoFromStatus(ctx context.Context, c client.Client, w *kueue.Wor
 		if err != nil {
 			return nil, err
 		}
-		if features.Enabled(features.TopologyAwareScheduling) {
+		if features.Enabled(features.TopologyAwareScheduling) || features.Enabled(features.SchedulerLibraryIntegration) {
 			info.Annotations[kueue.WorkloadAnnotation] = w.Name
 		}
 		if workloadslicing.IsElasticWorkload(w) {

pkg/util/maps/maps.go

@@ -82,6 +82,13 @@ func DeepCopySets[K comparable, T comparable](src map[K]sets.Set[T]) map[K]sets.
 	return c
 }
 
+func GetOrDefault[K comparable, V any, M ~map[K]V](m M, key K, defaultV V) V {
+	if v, found := m[key]; found {
+		return v
+	}
+	return defaultV
+}
+
 // SyncMap - generic RWMutex protected map.
 type SyncMap[K comparable, V any] struct {
 	lock sync.RWMutex

pkg/workload/workload.go

@@ -55,6 +55,7 @@ import (
 	"sigs.k8s.io/kueue/pkg/util/admissioncheck"
 	afs "sigs.k8s.io/kueue/pkg/util/admissionfairsharing"
 	"sigs.k8s.io/kueue/pkg/util/api"
+	utilmaps "sigs.k8s.io/kueue/pkg/util/maps"
 	utilpod "sigs.k8s.io/kueue/pkg/util/pod"
 	"sigs.k8s.io/kueue/pkg/util/podset"
 	"sigs.k8s.io/kueue/pkg/util/priority"
@@ -79,7 +80,14 @@ const (
 	SchedulingHashUnknown EquivalenceHash = "unknown"
 )
 
-var Semantic = conversion.EqualitiesOrDie(resources.Equal)
+var (
+	Semantic                = conversion.EqualitiesOrDie(resources.Equal)
+	WorkloadNameAnnotations = []string{
+		kueue.WorkloadAnnotation,
+		controllerconstants.PrebuiltWorkloadAnnotation,
+		podconstants.GroupNameAnnotation,
+	}
+)
 
 // Reference is the full reference to Workload formed as <namespace>/< kueue.WorkloadName >.
 type Reference string
@@ -112,6 +120,34 @@ func FromQuotaReservedOrAdmittedToPending(prevStatus, newStatus string) bool {
 	return (prevStatus == StatusQuotaReserved || prevStatus == StatusAdmitted) && newStatus == StatusPending
 }
 
+func MatchPods(wl *kueue.Workload, pods []*corev1.Pod) (matching []*corev1.Pod) {
+	for _, pod := range pods {
+		if match(wl, pod) {
+			matching = append(matching, pod)
+		}
+	}
+	return matching
+}
+
+func match(wl *kueue.Workload, pod *corev1.Pod) bool {
+	if wl.Namespace != pod.Namespace {
+		return false
+	}
+
+	if slices.ContainsFunc(WorkloadNameAnnotations, func(key string) bool {
+		wlName, ok := pod.Annotations[key]
+		return ok && wl.Name == wlName
+	}) {
+		return true
+	}
+
+	if podSliceName, ok := pod.Annotations[kueue.WorkloadSliceNameAnnotation]; ok {
+		return podSliceName == utilmaps.GetOrDefault(wl.Annotations, kueue.WorkloadSliceNameAnnotation, wl.Name)
+	}
+
+	return false
+}
+
 type AssignmentClusterQueueState struct {
 	LastTriedFlavorIdx     []map[corev1.ResourceName]int
 	ClusterQueueGeneration int64

pkg/workload/workload_test.go

@@ -42,13 +42,15 @@ import (
 	kueue "sigs.k8s.io/kueue/apis/kueue/v1beta2"
 	queueafs "sigs.k8s.io/kueue/pkg/cache/queue/afs"
 	controllerconstants "sigs.k8s.io/kueue/pkg/controller/constants"
+	podconstants "sigs.k8s.io/kueue/pkg/controller/jobs/pod/constants"
 	"sigs.k8s.io/kueue/pkg/features"
 	"sigs.k8s.io/kueue/pkg/resources"
 	"sigs.k8s.io/kueue/pkg/util/admissioncheck"
 	utilqueue "sigs.k8s.io/kueue/pkg/util/queue"
 	utiltas "sigs.k8s.io/kueue/pkg/util/tas"
 	utiltesting "sigs.k8s.io/kueue/pkg/util/testing"
 	utiltestingapi "sigs.k8s.io/kueue/pkg/util/testing/v1beta2"
+	podtesting "sigs.k8s.io/kueue/pkg/util/testingjobs/pod"
 )
 
 var (
@@ -3717,3 +3719,79 @@ func TestTotalExecutionTime(t *testing.T) {
 		})
 	}
 }
+
+func TestMatchPods(t *testing.T) {
+	features.SetFeatureGateDuringTest(t, features.SchedulerLibraryIntegration, true)
+
+	wl := utiltestingapi.MakeWorkload("wl-1", "ns-a").Obj()
+	wlSlice := utiltestingapi.MakeWorkload("wl-slice-2", "ns-a").
+		Annotation(kueue.WorkloadSliceNameAnnotation, "wl-slice-1").Obj()
+
+	cases := map[string]struct {
+		wl      *kueue.Workload
+		pod     *corev1.Pod
+		wantPod bool
+	}{
+		"pod with WorkloadAnnotation": {
+			wl: wl,
+			pod: podtesting.MakePod("pod-wl-annotated", "ns-a").
+				Annotation(kueue.WorkloadAnnotation, "wl-1").Obj(),
+			wantPod: true,
+		},
+		"pod with PrebuiltWorkloadAnnotation": {
+			wl: wl,
+			pod: podtesting.MakePod("pod-prebuilt-annotated", "ns-a").
+				Annotation(controllerconstants.PrebuiltWorkloadAnnotation, "wl-1").Obj(),
+			wantPod: true,
+		},
+		"pod with GroupNameAnnotation": {
+			wl: wl,
+			pod: podtesting.MakePod("pod-group-annotated", "ns-a").
+				Annotation(podconstants.GroupNameAnnotation, "wl-1").Obj(),
+			wantPod: true,
+		},
+		"pod with WorkloadSliceNameAnnotation": {
+			wl: wlSlice,
+			pod: podtesting.MakePod("pod-slice-annotated", "ns-a").
+				Annotation(kueue.WorkloadSliceNameAnnotation, "wl-slice-1").Obj(),
+			wantPod: true,
+		},
+		"pod in different namespace": {
+			wl: wl,
+			pod: podtesting.MakePod("pod-other-ns", "ns-b").
+				Annotation(kueue.WorkloadAnnotation, "wl-1").Obj(),
+			wantPod: false,
+		},
+		"pod with unrelated workload annotation": {
+			wl: wl,
+			pod: podtesting.MakePod("pod-unrelated", "ns-a").
+				Annotation(kueue.WorkloadAnnotation, "wl-other").Obj(),
+			wantPod: false,
+		},
+		"pod with no annotations": {
+			wl:      wl,
+			pod:     podtesting.MakePod("pod-no-annot", "ns-a").Obj(),
+			wantPod: false,
+		},
+	}
+
+	for name, tc := range cases {
+		t.Run(name, func(t *testing.T) {
+			want := []*corev1.Pod{}
+			if tc.wantPod {
+				want = append(want, tc.pod)
+			}
+
+			noisePods := []*corev1.Pod{
+				podtesting.MakePod("pod-wl-annotated-noise", "ns-a").Annotation(kueue.WorkloadAnnotation, "other-workload").Obj(),
+				podtesting.MakePod("pod-prebuilt-annotated-noise", "otherNamespace").Annotation(controllerconstants.PrebuiltWorkloadAnnotation, "other-workload").Obj(),
+				podtesting.MakePod("pod-no-annot-noise", "ns-a").Obj(),
+			}
+			matched := MatchPods(tc.wl, append(noisePods, tc.pod))
+
+			if diff := cmp.Diff(want, matched, cmpopts.EquateEmpty()); diff != "" {
+				t.Errorf("MatchPods(%s, %s) mismatch (-want +got):\n%s", tc.wl.Name, tc.pod.Name, diff)
+			}
+		})
+	}
+}

test/e2e/singlecluster/baseline/pod_test.go

@@ -621,9 +621,6 @@ var _ = ginkgo.Describe("Pod groups", ginkgo.Label("area:singlecluster", "featur
 				gomega.Eventually(func(g gomega.Gomega) {
 					g.Expect(k8sClient.Get(ctx, pKey, &corev1.Pod{})).To(utiltesting.BeNotFoundError())
 				}, util.Timeout, util.Interval).Should(gomega.Succeed())
-
-				gKey := client.ObjectKey{Namespace: ns.Name, Name: "test-group"}
-				util.ExpectWorkloadsFinalizedOrGone(ctx, k8sClient, gKey)
 			})
 		})
 	})