Add workload annotation to pods when SchedulerLibraryIntegration FG enabled.

pkg/controller/jobframework/reconciler.go

@@ -1739,11 +1739,11 @@ 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) {
-			info.Annotations[kueue.WorkloadSliceNameAnnotation] = workloadslicing.SliceName(w)
+			info.Annotations[kueue.WorkloadSliceNameAnnotation] = workload.SliceName(w)
 		}
 
 		info.Labels[constants.PodSetLabel] = string(psAssignment.Name)

pkg/workload/workload.go

@@ -79,7 +79,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 +119,42 @@ func FromQuotaReservedOrAdmittedToPending(prevStatus, newStatus string) bool {
 	return (prevStatus == StatusQuotaReserved || prevStatus == StatusAdmitted) && newStatus == StatusPending
 }
 
+// SliceName returns the workload slice name for the given workload.
+// This is the original workload name in the slice chain, used to identify pods
+// across workload slice replacements. If the workload has the WorkloadSliceNameAnnotation,
+// that value is returned; otherwise the workload's own name is returned.
+func SliceName(wl *kueue.Workload) string {
+	if sliceName, found := wl.Annotations[kueue.WorkloadSliceNameAnnotation]; found {
+		return sliceName
+	}
+	return wl.Name
+}
+
+func MatchPods(wl *kueue.Workload, pods []*corev1.Pod) (match []*corev1.Pod) {
+	for _, pod := range pods {
+		if areRelated(wl, pod) {
+			match = append(match, pod)
+		}
+	}
+	return match
+}
+
+func areRelated(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 == SliceName(wl)
+	}
+	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)
+			}
+		})
+	}
+}

pkg/workloadslicing/workloadslicing.go

@@ -90,17 +90,6 @@ func ReplacementForKey(wl *kueue.Workload) *workload.Reference {
 	return new(workload.Reference(key))
 }
 
-// SliceName returns the workload slice name for the given workload.
-// This is the original workload name in the slice chain, used to identify pods
-// across workload slice replacements. If the workload has the WorkloadSliceNameAnnotation,
-// that value is returned; otherwise the workload's own name is returned.
-func SliceName(wl *kueue.Workload) string {
-	if sliceName, found := wl.Annotations[kueue.WorkloadSliceNameAnnotation]; found {
-		return sliceName
-	}
-	return wl.Name
-}
-
 // FindNotFinishedWorkloads returns a sorted list of workloads "owned by" the provided job object/gvk combination and
 // without "Finished" condition with status = "True".
 func FindNotFinishedWorkloads(ctx context.Context, clnt client.Client, jobObject client.Object, jobObjectGVK schema.GroupVersionKind) ([]kueue.Workload, error) {