Propagate job labels and annotations (#737)

* Propagate job labels and annotations * Modify unit tests * Fix formatting * Remove unused annotation from test * Fix unit test * Separate pod and job labels in tests * Update pkg/controllers/jobset_controller_test.go Co-authored-by: Abdullah Gharaibeh <[email protected]> * Update pkg/controllers/jobset_controller_test.go Co-authored-by: Abdullah Gharaibeh <[email protected]> * Update pkg/controllers/jobset_controller_test.go Co-authored-by: Abdullah Gharaibeh <[email protected]> * Update pkg/controllers/jobset_controller_test.go Co-authored-by: Abdullah Gharaibeh <[email protected]> * Update pkg/controllers/jobset_controller_test.go Co-authored-by: Abdullah Gharaibeh <[email protected]> * Update pkg/controllers/jobset_controller_test.go Co-authored-by: Abdullah Gharaibeh <[email protected]> * Update pkg/controllers/jobset_controller_test.go * Update pkg/controllers/jobset_controller_test.go * Update pkg/controllers/jobset_controller_test.go * Update pkg/controllers/jobset_controller_test.go * Update pkg/controllers/jobset_controller_test.go * Update pkg/controllers/jobset_controller_test.go * Update pkg/controllers/jobset_controller_test.go * Update pkg/controllers/jobset_controller_test.go * Run gofmt --------- Co-authored-by: Abdullah Gharaibeh <[email protected]>
kubernetes-sigs · Jan 7, 2025 · cd00ccf · cd00ccf
1 parent 5d8c846
commit cd00ccf
Show file tree

Hide file tree

Showing 3 changed files with 88 additions and 4 deletions.
diff --git a/pkg/controllers/jobset_controller.go b/pkg/controllers/jobset_controller.go
@@ -738,6 +738,7 @@ func labelAndAnnotateObject(obj metav1.Object, js *jobset.JobSet, rjob *jobset.R
 
 	// Set labels on the object.
 	labels := make(map[string]string)
+	maps.Copy(labels, obj.GetLabels())
 	labels[jobset.JobSetNameKey] = js.Name
 	labels[jobset.ReplicatedJobNameKey] = rjob.Name
 	labels[constants.RestartsKey] = strconv.Itoa(int(js.Status.Restarts))
@@ -748,6 +749,7 @@ func labelAndAnnotateObject(obj metav1.Object, js *jobset.JobSet, rjob *jobset.R
 	labels[jobset.JobGlobalIndexKey] = globalJobIndex(js, rjob.Name, jobIdx)
 
 	annotations := make(map[string]string)
+	maps.Copy(annotations, obj.GetAnnotations())
 	annotations[jobset.JobSetNameKey] = js.Name
 	annotations[jobset.ReplicatedJobNameKey] = rjob.Name
 	annotations[constants.RestartsKey] = strconv.Itoa(int(js.Status.Restarts))

diff --git a/pkg/controllers/jobset_controller_test.go b/pkg/controllers/jobset_controller_test.go
@@ -127,10 +127,26 @@ func TestIsJobFinished(t *testing.T) {
 
 func TestConstructJobsFromTemplate(t *testing.T) {
 	var (
-		jobSetName          = "test-jobset"
-		replicatedJobName   = "replicated-job"
-		jobName             = "test-job"
-		ns                  = "default"
+		jobSetName        = "test-jobset"
+		replicatedJobName = "replicated-job"
+		jobName           = "test-job"
+		ns                = "default"
+		jobAnnotations    = map[string]string{
+			"job-annotation-key1": "job-annotation-value1",
+			"job-annotation-key2": "job-annotation-value2",
+		}
+		jobLabels = map[string]string{
+			"job-label-key1": "job-label-value1",
+			"job-label-key2": "job-label-value2",
+		}
+		podAnnotations = map[string]string{
+			"pod-annotation-key1": "pod-annotation-value1",
+			"pod-annotation-key2": "pod-annotation-value2",
+		}
+		podLabels = map[string]string{
+			"pod-label-key1": "pod-label-value1",
+			"pod-label-key2": "pod-label-value2",
+		}
 		topologyDomain      = "test-topology-domain"
 		coordinatorKeyValue = map[string]string{
 			jobset.CoordinatorKey: fmt.Sprintf("%s-%s-%d-%d.%s", jobSetName, replicatedJobName, 0, 0, jobSetName),
@@ -183,6 +199,46 @@ func TestConstructJobsFromTemplate(t *testing.T) {
 					Suspend(false).Obj(),
 			},
 		},
+		{
+			name: "all jobs/pods created with labels and annotations",
+			js: testutils.MakeJobSet(jobSetName, ns).
+				ReplicatedJob(testutils.MakeReplicatedJob(replicatedJobName).
+					Job(testutils.MakeJobTemplate(jobName, ns).
+						SetLabels(jobLabels).
+						SetPodLabels(podLabels).
+						SetAnnotations(jobAnnotations).
+						SetPodAnnotations(podAnnotations).
+						Obj()).
+					Replicas(2).
+					Obj()).Obj(),
+			ownedJobs: &childJobs{},
+			want: []*batchv1.Job{
+				makeJob(&makeJobArgs{
+					jobSetName:        jobSetName,
+					replicatedJobName: replicatedJobName,
+					jobName:           "test-jobset-replicated-job-0",
+					ns:                ns,
+					jobLabels:         jobLabels,
+					jobAnnotations:    jobAnnotations,
+					podLabels:         podLabels,
+					podAnnotations:    podAnnotations,
+					replicas:          2,
+					jobIdx:            0}).
+					Suspend(false).Obj(),
+				makeJob(&makeJobArgs{
+					jobSetName:        jobSetName,
+					replicatedJobName: replicatedJobName,
+					jobName:           "test-jobset-replicated-job-1",
+					ns:                ns,
+					jobLabels:         jobLabels,
+					jobAnnotations:    jobAnnotations,
+					podLabels:         podLabels,
+					podAnnotations:    podAnnotations,
+					replicas:          2,
+					jobIdx:            1}).
+					Suspend(false).Obj(),
+			},
+		},
 		{
 			name: "one job created, one job not created (already active)",
 			js: testutils.MakeJobSet(jobSetName, ns).
@@ -1251,6 +1307,10 @@ type makeJobArgs struct {
 	replicatedJobName    string
 	jobName              string
 	ns                   string
+	jobLabels            map[string]string
+	jobAnnotations       map[string]string
+	podLabels            map[string]string
+	podAnnotations       map[string]string
 	replicas             int
 	jobIdx               int
 	restarts             int
@@ -1286,8 +1346,12 @@ func makeJob(args *makeJobArgs) *testutils.JobWrapper {
 	}
 	jobWrapper := testutils.MakeJob(args.jobName, args.ns).
 		JobLabels(labels).
+		JobLabels(args.jobLabels).
 		JobAnnotations(annotations).
+		JobAnnotations(args.jobAnnotations).
 		PodLabels(labels).
+		PodLabels(args.podLabels).
+		PodAnnotations(args.podAnnotations).
 		PodAnnotations(annotations)
 	return jobWrapper
 }

diff --git a/pkg/util/testing/wrappers.go b/pkg/util/testing/wrappers.go
@@ -267,12 +267,30 @@ func (j *JobTemplateWrapper) PodSpec(podSpec corev1.PodSpec) *JobTemplateWrapper
 	return j
 }
 
+// SetAnnotations sets the annotations on the Pod template.
+func (j *JobTemplateWrapper) SetPodAnnotations(annotations map[string]string) *JobTemplateWrapper {
+	j.Spec.Template.SetAnnotations(annotations)
+	return j
+}
+
+// SetLabels sets the labels on the Pod template.
+func (j *JobTemplateWrapper) SetPodLabels(labels map[string]string) *JobTemplateWrapper {
+	j.Spec.Template.SetLabels(labels)
+	return j
+}
+
 // SetAnnotations sets the annotations on the Job template.
 func (j *JobTemplateWrapper) SetAnnotations(annotations map[string]string) *JobTemplateWrapper {
 	j.Annotations = annotations
 	return j
 }
 
+// SetLabels sets the labels on the Job template.
+func (j *JobTemplateWrapper) SetLabels(labels map[string]string) *JobTemplateWrapper {
+	j.Labels = labels
+	return j
+}
+
 // Obj returns the inner batchv1.JobTemplateSpec
 func (j *JobTemplateWrapper) Obj() batchv1.JobTemplateSpec {
 	return j.JobTemplateSpec