Repository: kubernetes-sigs/kueue

Issue:
Kueue doesn't propagate Pod template batch.kubernetes.io/job-name label to ProvReq's PodTemplate
**What happened**:
PodTemplate has a pod affinity rule to match batch.kubernetes.io/job-name label. The label wasn't present on (fake) pods which were created through ProvisioningRequests that leads to failed scale up simulation, however the label was present on Job PodTemplate object.

It happens because Kueue explicitly removes job-name label from PodTemplate: [link](https://www.google.com/url?q=https://github.com/kubernetes-sigs/kueue/pull/1358/files%23diff-f278bdcf60c86f048f29eeef3fa2ee16447d577a7d3d6d0a3dce33626ffed260R214&sa=D&source=buganizer&usg=AOvVaw0ap7PP_530moXxNqJKCwJ4)

Discussion:
@yaroslava-serdiuk: /assign

Solution patch:
diff --git a/pkg/controller/jobs/job/job_controller.go b/pkg/controller/jobs/job/job_controller.go
index f94bd1488..a2dcf8c33 100644
--- a/pkg/controller/jobs/job/job_controller.go
+++ b/pkg/controller/jobs/job/job_controller.go
@@ -233,11 +233,15 @@ var (
 	// the legacy names are no longer defined in the api, only in k/2/apis/batch
 	legacyJobNameLabel       = "job-name"
 	legacyControllerUIDLabel = "controller-uid"
-	ManagedLabels            = []string{legacyJobNameLabel, legacyControllerUIDLabel, batchv1.JobNameLabel, batchv1.ControllerUidLabel}
+	managedLabels            = []string{legacyJobNameLabel, legacyControllerUIDLabel, batchv1.JobNameLabel, batchv1.ControllerUidLabel}
 )
 
-func cleanManagedLabels(pt *corev1.PodTemplateSpec) *corev1.PodTemplateSpec {
-	for _, managedLabel := range ManagedLabels {
+// Clean labels that are managed by the job controller except job-name label.
+func cleanLabels(pt *corev1.PodTemplateSpec) *corev1.PodTemplateSpec {
+	for _, managedLabel := range managedLabels {
+		if features.Enabled(features.PropagateBatchJobLabelsToWorkload) && managedLabel == batchv1.JobNameLabel {
+			continue
+		}
 		delete(pt.Labels, managedLabel)
 	}
 	return pt
@@ -246,7 +250,7 @@ func cleanManagedLabels(pt *corev1.PodTemplateSpec) *corev1.PodTemplateSpec {
 func (j *Job) PodSets(ctx context.Context) ([]kueue.PodSet, error) {
 	podSet := kueue.PodSet{
 		Name:     kueue.DefaultPodSetName,
-		Template: *cleanManagedLabels(j.Spec.Template.DeepCopy()),
+		Template: *cleanLabels(j.Spec.Template.DeepCopy()),
 		Count:    j.podsCount(),
 		MinCount: j.minPodsCount(),
 	}
@@ -296,7 +300,7 @@ func (j *Job) RestorePodSetsInfo(podSetsInfo []podset.PodSetInfo) bool {
 		}
 	}
 	info := podSetsInfo[0]
-	for _, managedLabel := range ManagedLabels {
+	for _, managedLabel := range managedLabels {
 		if v, found := j.Spec.Template.Labels[managedLabel]; found {
 			info.AddOrUpdateLabel(managedLabel, v)
 		}
diff --git a/pkg/controller/jobs/job/job_multikueue_adapter.go b/pkg/controller/jobs/job/job_multikueue_adapter.go
index b3761771d..31741432b 100644
--- a/pkg/controller/jobs/job/job_multikueue_adapter.go
+++ b/pkg/controller/jobs/job/job_multikueue_adapter.go
@@ -146,9 +146,7 @@ func (b *multiKueueAdapter) SyncJob(ctx context.Context, localClient client.Clie
 	// drop the selector
 	remoteJob.Spec.Selector = nil
 	// drop the templates cleanup labels
-	for _, cl := range ManagedLabels {
-		delete(remoteJob.Spec.Template.Labels, cl)
-	}
+	cleanLabels(&remoteJob.Spec.Template)
 
 	// add the prebuilt workload
 	if remoteJob.Labels == nil {
diff --git a/pkg/features/kube_features.go b/pkg/features/kube_features.go
index 26aa0954a..acfd21fa5 100644
--- a/pkg/features/kube_features.go
+++ b/pkg/features/kube_features.go
@@ -208,6 +208,12 @@ const (
 	//
 	// Enables reclaimable pods counting towards quota.
 	ReclaimablePods featuregate.Feature = "ReclaimablePods"
+
+	// owner: @yaroslva-serdiuk
+	//
+	// issue: https://github.com/kubernetes-sigs/kueue/issues/7597
+	// Do not remove job-name label from Workload PodTemplate object.
+	PropagateBatchJobLabelsToWorkload featuregate.Feature = "PropagateBatchJobLabelsToWorkload"
 )
 
 func init() {
@@ -328,6 +334,9 @@ var defaultVersionedFeatureGates = map[featuregate.Feature]featuregate.Versioned
 	ReclaimablePods: {
 		{Version: version.MustParse("0.15"), Default: true, PreRelease: featuregate.Beta},
 	},
+	PropagateBatchJobLabelsToWorkload: {
+		{Version: version.MustParse("0.15"), Default: true, PreRelease: featuregate.Beta},
+	},
 }
 
 func SetFeatureGateDuringTest(tb testing.TB, f featuregate.Feature, value bool) {



Test patch:
diff --git a/pkg/controller/jobs/job/job_controller_test.go b/pkg/controller/jobs/job/job_controller_test.go
index 482ff8bbf..f8e56261a 100644
--- a/pkg/controller/jobs/job/job_controller_test.go
+++ b/pkg/controller/jobs/job/job_controller_test.go
@@ -3981,3 +3981,53 @@ func TestReconciler(t *testing.T) {
 		}
 	}
 }
+
+func TestCleanLabels(t *testing.T) {
+	cases := map[string]struct {
+		featureEnabled bool
+		labels         map[string]string
+		wantLabels     map[string]string
+	}{
+		"feature disabled": {
+			featureEnabled: false,
+			labels: map[string]string{
+				"foo":                      "bar",
+				batchv1.JobNameLabel:       "job-name",
+				"controller-uid":           "uid",
+				batchv1.ControllerUidLabel: "uid",
+			},
+			wantLabels: map[string]string{
+				"foo": "bar",
+			},
+		},
+		"feature enabled": {
+			featureEnabled: true,
+			labels: map[string]string{
+				"foo":                      "bar",
+				batchv1.JobNameLabel:       "job-name",
+				"controller-uid":           "uid",
+				batchv1.ControllerUidLabel: "uid",
+			},
+			wantLabels: map[string]string{
+				"foo":                "bar",
+				batchv1.JobNameLabel: "job-name",
+			},
+		},
+	}
+
+	for name, tc := range cases {
+		t.Run(name, func(t *testing.T) {
+			print(tc.labels)
+			features.SetFeatureGateDuringTest(t, features.PropagateBatchJobLabelsToWorkload, tc.featureEnabled)
+			pt := &corev1.PodTemplateSpec{
+				ObjectMeta: metav1.ObjectMeta{
+					Labels: tc.labels,
+				},
+			}
+			cleanLabels(pt)
+			if diff := cmp.Diff(tc.wantLabels, pt.Labels); diff != "" {
+				t.Errorf("cleanLabels() mismatch (-want +got):\n%s", diff)
+			}
+		})
+	}
+}
diff --git a/test/integration/singlecluster/controller/admissionchecks/provisioning/provisioning_test.go b/test/integration/singlecluster/controller/admissionchecks/provisioning/provisioning_test.go
index f9008fce0..bc40df123 100644
--- a/test/integration/singlecluster/controller/admissionchecks/provisioning/provisioning_test.go
+++ b/test/integration/singlecluster/controller/admissionchecks/provisioning/provisioning_test.go
@@ -23,6 +23,7 @@ import (
 	"github.com/google/go-cmp/cmp/cmpopts"
 	"github.com/onsi/ginkgo/v2"
 	"github.com/onsi/gomega"
+	batchv1 "k8s.io/api/batch/v1"
 	corev1 "k8s.io/api/core/v1"
 	apimeta "k8s.io/apimachinery/pkg/api/meta"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
@@ -1637,10 +1638,12 @@ var _ = ginkgo.Describe("Provisioning with scheduling", ginkgo.Ordered, ginkgo.C
 				util.ExpectLocalQueuesToBeActive(ctx, k8sClient, lq)
 			})
 
+			jobName := "job1"
 			ginkgo.By("submit the Job", func() {
-				jobBuilder := testingjob.MakeJob("job1", ns.Name).
+				jobBuilder := testingjob.MakeJob(jobName, ns.Name).
 					Queue(kueue.LocalQueueName(lq.Name)).
-					Request(corev1.ResourceCPU, "500m")
+					Request(corev1.ResourceCPU, "500m").
+					PodLabel(batchv1.JobNameLabel, jobName)
 				job1 := jobBuilder.Obj()
 				util.MustCreate(ctx, k8sClient, job1)
 				ginkgo.DeferCleanup(func() {
@@ -1655,6 +1658,11 @@ var _ = ginkgo.Describe("Provisioning with scheduling", ginkgo.Ordered, ginkgo.C
 				}, util.Timeout, util.Interval).Should(gomega.Succeed())
 			})
 
+			ginkgo.By("verify workload podTemplate has job-name label set correctly", func() {
+				gomega.Expect(wlObj.Spec.PodSets[0].Template.GetLabels()).To(gomega.BeComparableTo(map[string]string{
+					batchv1.JobNameLabel: jobName}))
+			})
+
 			ginkgo.By("await for wl1 to have QuotaReserved on flavor-1", func() {
 				gomega.Eventually(func(g gomega.Gomega) {
 					gomega.Expect(k8sClient.Get(ctx, wl1Key, &wlObj)).Should(gomega.Succeed())