Propagate job labels and annotations

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))

pkg/controllers/jobset_controller_test.go

@@ -17,6 +17,7 @@ import (
 	"context"
 	"errors"
 	"fmt"
+	"maps"
 	"strconv"
 	"strings"
 	"testing"
@@ -127,10 +128,18 @@ 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"
+		annotations       = map[string]string{
+			"foo":  "bar",
+			"foo1": "bar",
+		}
+		labels = map[string]string{
+			"foo":  "bar",
+			"foo1": "bar",
+		}
 		topologyDomain      = "test-topology-domain"
 		coordinatorKeyValue = map[string]string{
 			jobset.CoordinatorKey: fmt.Sprintf("%s-%s-%d-%d.%s", jobSetName, replicatedJobName, 0, 0, jobSetName),
@@ -183,6 +192,41 @@ func TestConstructJobsFromTemplate(t *testing.T) {
 					Suspend(false).Obj(),
 			},
 		},
+		{
+			name: "all jobs/pods created with labels and annotations",
+			js: testutils.MakeJobSet(jobSetName, ns).
+				SetAnnotations(annotations).
+				ReplicatedJob(testutils.MakeReplicatedJob(replicatedJobName).
+					Job(testutils.MakeJobTemplate(jobName, ns).
+						SetLabels(labels).
+						SetAnnotations(annotations).
+						Obj()).
+					Replicas(2).
+					Obj()).Obj(),
+			ownedJobs: &childJobs{},
+			want: []*batchv1.Job{
+				makeJob(&makeJobArgs{
+					jobSetName:        jobSetName,
+					replicatedJobName: replicatedJobName,
+					jobName:           "test-jobset-replicated-job-0",
+					ns:                ns,
+					labels:            labels,
+					annotations:       annotations,
+					replicas:          2,
+					jobIdx:            0}).
+					Suspend(false).Obj(),
+				makeJob(&makeJobArgs{
+					jobSetName:        jobSetName,
+					replicatedJobName: replicatedJobName,
+					jobName:           "test-jobset-replicated-job-1",
+					ns:                ns,
+					labels:            labels,
+					annotations:       annotations,
+					replicas:          2,
+					jobIdx:            1}).
+					Suspend(false).Obj(),
+			},
+		},
 		{
 			name: "one job created, one job not created (already active)",
 			js: testutils.MakeJobSet(jobSetName, ns).
@@ -1251,6 +1295,8 @@ type makeJobArgs struct {
 	replicatedJobName    string
 	jobName              string
 	ns                   string
+	labels               map[string]string
+	annotations          map[string]string
 	replicas             int
 	jobIdx               int
 	restarts             int
@@ -1268,6 +1314,7 @@ func makeJob(args *makeJobArgs) *testutils.JobWrapper {
 		constants.RestartsKey:        strconv.Itoa(args.restarts),
 		jobset.JobKey:                jobHashKey(args.ns, args.jobName),
 	}
+	maps.Copy(labels, args.labels)
 	annotations := map[string]string{
 		jobset.JobSetNameKey:         args.jobSetName,
 		jobset.ReplicatedJobNameKey:  args.replicatedJobName,
@@ -1276,6 +1323,7 @@ func makeJob(args *makeJobArgs) *testutils.JobWrapper {
 		constants.RestartsKey:        strconv.Itoa(args.restarts),
 		jobset.JobKey:                jobHashKey(args.ns, args.jobName),
 	}
+	maps.Copy(annotations, args.annotations)
 	// Only set exclusive key if we are using exclusive placement per topology.
 	if args.topology != "" {
 		annotations[jobset.ExclusiveKey] = args.topology

pkg/util/testing/wrappers.go

@@ -270,6 +270,14 @@ func (j *JobTemplateWrapper) PodSpec(podSpec corev1.PodSpec) *JobTemplateWrapper
 // SetAnnotations sets the annotations on the Job template.
 func (j *JobTemplateWrapper) SetAnnotations(annotations map[string]string) *JobTemplateWrapper {
 	j.Annotations = annotations
+	j.Spec.Template.SetAnnotations(annotations)
+	return j
+}
+
+// SetLabels sets the labels on the Job template.
+func (j *JobTemplateWrapper) SetLabels(labels map[string]string) *JobTemplateWrapper {
+	j.Labels = labels
+	j.Spec.Template.SetLabels(labels)
 	return j
 }