Propagate additional annotations to pods (#3379)

armadaproject · Feb 7, 2024 · 9b85c5b · 9b85c5b
1 parent a24fd16
commit 9b85c5b
Show file tree

Hide file tree

Showing 7 changed files with 268 additions and 419 deletions.
diff --git a/internal/executor/util/kubernetes_objects_test.go b/internal/executor/util/kubernetes_objects_test.go
@@ -494,7 +494,7 @@ func TestCreatePodFromExecutorApiJob(t *testing.T) {
 		Job: &armadaevents.SubmitJob{
 			ObjectMeta: &armadaevents.ObjectMeta{
 				Labels:      map[string]string{},
-				Annotations: map[string]string{},
+				Annotations: map[string]string{"runtime_gang_cardinality": "3"},
 				Namespace:   "test-namespace",
 			},
 			JobId: jobId,
@@ -520,8 +520,9 @@ func TestCreatePodFromExecutorApiJob(t *testing.T) {
 				domain.PodCount:  "1",
 			},
 			Annotations: map[string]string{
-				domain.JobSetId: "job-set",
-				domain.Owner:    "user",
+				domain.JobSetId:            "job-set",
+				domain.Owner:               "user",
+				"runtime_gang_cardinality": "3",
 			},
 		},
 		Spec: v1.PodSpec{

diff --git a/internal/scheduler/api.go b/internal/scheduler/api.go
@@ -144,6 +144,7 @@ func (srv *ExecutorApi) LeaseJobRuns(stream executorapi.ExecutorApi_LeaseJobRuns
 				return err
 			}
 			addTolerations(submitMsg, PodRequirementsOverlay.Tolerations)
+			addAnnotations(submitMsg, PodRequirementsOverlay.Annotations)
 		}
 
 		var groups []string
@@ -239,6 +240,21 @@ func addTolerations(job *armadaevents.SubmitJob, tolerations []v1.Toleration) {
 	}
 }
 
+func addAnnotations(job *armadaevents.SubmitJob, annotations map[string]string) {
+	if job == nil || len(annotations) == 0 {
+		return
+	}
+	if job.ObjectMeta == nil {
+		job.ObjectMeta = &armadaevents.ObjectMeta{}
+	}
+	if job.ObjectMeta.Annotations == nil {
+		job.ObjectMeta.Annotations = make(map[string]string, len(annotations))
+	}
+	for k, v := range annotations {
+		job.ObjectMeta.Annotations[k] = v
+	}
+}
+
 // ReportEvents publishes all eventSequences to Pulsar. The eventSequences are compacted for more efficient publishing.
 func (srv *ExecutorApi) ReportEvents(grpcCtx context.Context, list *executorapi.EventList) (*types.Empty, error) {
 	ctx := armadacontext.FromGrpcCtx(grpcCtx)

diff --git a/internal/scheduler/api_test.go b/internal/scheduler/api_test.go
@@ -83,6 +83,7 @@ func TestExecutorApi_LeaseJobRuns(t *testing.T) {
 
 	submit, compressedSubmit := submitMsg(
 		t,
+		nil,
 		&v1.PodSpec{
 			NodeSelector: map[string]string{nodeIdName: "node-id"},
 		},
@@ -97,7 +98,7 @@ func TestExecutorApi_LeaseJobRuns(t *testing.T) {
 		SubmitMessage: compressedSubmit,
 	}
 
-	submitWithoutNodeSelector, compressedSubmitNoNodeSelector := submitMsg(t, &v1.PodSpec{})
+	submitWithoutNodeSelector, compressedSubmitNoNodeSelector := submitMsg(t, nil, nil)
 	leaseWithoutNode := &database.JobRunLease{
 		RunID:         uuid.New(),
 		Queue:         "test-queue",
@@ -114,7 +115,7 @@ func TestExecutorApi_LeaseJobRuns(t *testing.T) {
 			Effect: v1.TaintEffectNoSchedule,
 		},
 	}
-	leaseWithTolerations := &database.JobRunLease{
+	leaseWithOverlay := &database.JobRunLease{
 		RunID:  uuid.New(),
 		Queue:  "test-queue",
 		JobSet: "test-jobset",
@@ -127,18 +128,22 @@ func TestExecutorApi_LeaseJobRuns(t *testing.T) {
 		PodRequirementsOverlay: protoutil.MustMarshall(
 			&schedulerobjects.PodRequirements{
 				Tolerations: tolerations,
+				Annotations: map[string]string{"runtime_gang_cardinality": "3"},
 				Priority:    1000,
 			},
 		),
 	}
-	submitWithTolerations, _ := submitMsg(
+	submitWithOverlay, _ := submitMsg(
 		t,
+		&armadaevents.ObjectMeta{
+			Annotations: map[string]string{"runtime_gang_cardinality": "3"},
+		},
 		&v1.PodSpec{
 			NodeSelector: map[string]string{nodeIdName: "node-id"},
 			Tolerations:  tolerations,
 		},
 	)
-	submitWithTolerations.JobId = submit.JobId
+	submitWithOverlay.JobId = submit.JobId
 
 	tests := map[string]struct {
 		request          *executorapi.LeaseRequest
@@ -195,17 +200,17 @@ func TestExecutorApi_LeaseJobRuns(t *testing.T) {
 		},
 		"run with PodRequirementsOverlay": {
 			request:          defaultRequest,
-			leases:           []*database.JobRunLease{leaseWithTolerations},
+			leases:           []*database.JobRunLease{leaseWithOverlay},
 			expectedExecutor: defaultExpectedExecutor,
 			expectedMsgs: []*executorapi.LeaseStreamMessage{
 				{
 					Event: &executorapi.LeaseStreamMessage_Lease{Lease: &executorapi.JobRunLease{
-						JobRunId: armadaevents.ProtoUuidFromUuid(leaseWithTolerations.RunID),
-						Queue:    leaseWithTolerations.Queue,
-						Jobset:   leaseWithTolerations.JobSet,
-						User:     leaseWithTolerations.UserID,
+						JobRunId: armadaevents.ProtoUuidFromUuid(leaseWithOverlay.RunID),
+						Queue:    leaseWithOverlay.Queue,
+						Jobset:   leaseWithOverlay.JobSet,
+						User:     leaseWithOverlay.UserID,
 						Groups:   groups,
-						Job:      submitWithTolerations,
+						Job:      submitWithOverlay,
 					}},
 				},
 				{
@@ -406,9 +411,10 @@ func TestExecutorApi_Publish(t *testing.T) {
 	}
 }
 
-func submitMsg(t *testing.T, podSpec *v1.PodSpec) (*armadaevents.SubmitJob, []byte) {
+func submitMsg(t *testing.T, objectMeta *armadaevents.ObjectMeta, podSpec *v1.PodSpec) (*armadaevents.SubmitJob, []byte) {
 	submitMsg := &armadaevents.SubmitJob{
-		JobId: armadaevents.ProtoUuidFromUuid(uuid.New()),
+		JobId:      armadaevents.ProtoUuidFromUuid(uuid.New()),
+		ObjectMeta: objectMeta,
 		MainObject: &armadaevents.KubernetesMainObject{
 			Object: &armadaevents.KubernetesMainObject_PodSpec{
 				PodSpec: &armadaevents.PodSpecWithAvoidList{

diff --git a/internal/scheduler/scheduler.go b/internal/scheduler/scheduler.go
@@ -542,10 +542,6 @@ func AppendEventSequencesFromScheduledJobs(eventSequences []*armadaevents.EventS
 		if err != nil {
 			return nil, err
 		}
-		additionalAnnotations, found := additionalAnnotationsByJobId[job.Id()]
-		if !found {
-			additionalAnnotations = make(map[string]string)
-		}
 		run := job.LatestRun()
 		if run == nil {
 			return nil, errors.Errorf("attempting to generate lease eventSequences for job %s with no associated runs", job.Id())
@@ -569,9 +565,9 @@ func AppendEventSequencesFromScheduledJobs(eventSequences []*armadaevents.EventS
 							UpdateSequenceNumber:   job.QueuedVersion(),
 							HasScheduledAtPriority: hasScheduledAtPriority,
 							ScheduledAtPriority:    scheduledAtPriority,
-							AdditionalAnnotations:  additionalAnnotations,
 							PodRequirementsOverlay: &schedulerobjects.PodRequirements{
 								Tolerations: jctx.AdditionalTolerations,
+								Annotations: additionalAnnotationsByJobId[job.Id()],
 								Priority:    scheduledAtPriority,
 							},
 						},

diff --git a/internal/scheduler/scheduler_test.go b/internal/scheduler/scheduler_test.go
@@ -1843,7 +1843,6 @@ func TestCycleConsistency(t *testing.T) {
 									UpdateSequenceNumber:   1,
 									HasScheduledAtPriority: true,
 									ScheduledAtPriority:    10,
-									AdditionalAnnotations:  make(map[string]string),
 									PodRequirementsOverlay: &schedulerobjects.PodRequirements{Priority: 10},
 								},
 							},
@@ -1928,7 +1927,6 @@ func TestCycleConsistency(t *testing.T) {
 									UpdateSequenceNumber:   1,
 									HasScheduledAtPriority: true,
 									ScheduledAtPriority:    10,
-									AdditionalAnnotations:  make(map[string]string),
 									PodRequirementsOverlay: &schedulerobjects.PodRequirements{Priority: 10},
 								},
 							},
@@ -1992,7 +1990,6 @@ func TestCycleConsistency(t *testing.T) {
 									UpdateSequenceNumber:   1,
 									HasScheduledAtPriority: true,
 									ScheduledAtPriority:    10,
-									AdditionalAnnotations:  make(map[string]string),
 									PodRequirementsOverlay: &schedulerobjects.PodRequirements{Priority: 10},
 								},
 							},