diff --git a/pkg/controllers/jobset_controller.go b/pkg/controllers/jobset_controller.go index e8ef110e4..c6fcbbbbf 100644 --- 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 index 1970395f7..91acfba38 100644 --- 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 index 0cc3368e4..c22b40a47 100644 --- 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