feat: move gang scheduling metadata to typed protobuf fields

Signed-off-by: Dejan Zele Pejchev <[email protected]>

Author: dejanzele

internal/executor/util/kubernetes_object.go

@@ -13,6 +13,7 @@ import (
 	"github.com/armadaproject/armada/internal/common/util"
 	"github.com/armadaproject/armada/internal/executor/configuration"
 	"github.com/armadaproject/armada/internal/executor/domain"
+	"github.com/armadaproject/armada/internal/scheduler/schedulerobjects"
 	"github.com/armadaproject/armada/pkg/api"
 	"github.com/armadaproject/armada/pkg/armadaevents"
 	"github.com/armadaproject/armada/pkg/executorapi"
@@ -128,10 +129,25 @@ func CreatePodFromExecutorApiJob(job *executorapi.JobRunLease, defaults *configu
 		domain.Owner:    job.User,
 	})
 
+	// Add gang scheduling annotations from SchedulingMetadata if not already present
+	// (for backward compatibility with old jobs that have annotations in the job spec)
+	if job.SchedulingMetadata != nil && job.SchedulingMetadata.GangInfo != nil {
+		gangInfo := job.SchedulingMetadata.GangInfo
+		if _, exists := annotation[serverconfiguration.GangIdAnnotation]; !exists && gangInfo.GangId != "" {
+			annotation[serverconfiguration.GangIdAnnotation] = gangInfo.GangId
+		}
+		if _, exists := annotation[serverconfiguration.GangCardinalityAnnotation]; !exists && gangInfo.Cardinality > 0 {
+			annotation[serverconfiguration.GangCardinalityAnnotation] = strconv.FormatUint(uint64(gangInfo.Cardinality), 10)
+		}
+		if _, exists := annotation[serverconfiguration.GangNodeUniformityLabelAnnotation]; !exists && gangInfo.NodeUniformityLabelName != "" {
+			annotation[serverconfiguration.GangNodeUniformityLabelAnnotation] = gangInfo.NodeUniformityLabelName
+		}
+	}
+
 	applyDefaults(podSpec, defaults)
 	setRestartPolicyNever(podSpec)
 
-	injectArmadaEnvVars(podSpec, jobId, job.Queue, job.Jobset, annotation)
+	injectArmadaEnvVars(podSpec, jobId, job.Queue, job.Jobset, job.SchedulingMetadata, annotation)
 
 	pod := &v1.Pod{
 		ObjectMeta: metav1.ObjectMeta{
@@ -173,7 +189,7 @@ func CreatePod(job *api.Job, defaults *configuration.PodDefaults) *v1.Pod {
 
 	setRestartPolicyNever(podSpec)
 
-	injectArmadaEnvVars(podSpec, job.Id, job.Queue, job.JobSetId, annotation)
+	injectArmadaEnvVars(podSpec, job.Id, job.Queue, job.JobSetId, nil, annotation)
 
 	pod := &v1.Pod{
 		ObjectMeta: metav1.ObjectMeta{
@@ -203,41 +219,18 @@ func setRestartPolicyNever(podSpec *v1.PodSpec) {
 
 // injectArmadaEnvVars injects Armada system environment variables into all containers.
 // It injects base variables for all jobs and additional variables for gang-scheduled jobs.
-func injectArmadaEnvVars(podSpec *v1.PodSpec, jobId string, queue string, jobsetId string, annotations map[string]string) {
-	// Base environment variables for all jobs
+func injectArmadaEnvVars(podSpec *v1.PodSpec, jobId string, queue string, jobsetId string, schedulingMetadata *armadaevents.SchedulingMetadata, annotations map[string]string) {
 	baseEnvVars := []v1.EnvVar{
 		{Name: serverconfiguration.JobIdEnvVar, Value: jobId},
 		{Name: serverconfiguration.QueueEnvVar, Value: queue},
 		{Name: serverconfiguration.JobSetIdEnvVar, Value: jobsetId},
 	}
 
-	// Gang-specific environment variables
-	var gangEnvVars []v1.EnvVar
-
-	if gangId, ok := annotations[serverconfiguration.GangIdAnnotation]; ok && gangId != "" {
-		gangEnvVars = append(gangEnvVars, v1.EnvVar{
-			Name:  serverconfiguration.GangIdEnvVar,
-			Value: gangId,
-		})
-	}
-	if gangCardinality, ok := annotations[serverconfiguration.GangCardinalityAnnotation]; ok && gangCardinality != "" {
-		gangEnvVars = append(gangEnvVars, v1.EnvVar{
-			Name:  serverconfiguration.GangCardinalityEnvVar,
-			Value: gangCardinality,
-		})
+	allEnvVars := baseEnvVars
+	if schedulingMetadata != nil && schedulingMetadata.GangInfo != nil {
+		allEnvVars = append(allEnvVars, buildGangEnvVars(schedulingMetadata.GangInfo)...)
 	}
 
-	labelName, hasLabelName := annotations[serverconfiguration.GangNodeUniformityLabelNameEnvVar]
-	labelValue, hasLabelValue := annotations[serverconfiguration.GangNodeUniformityLabelValueEnvVar]
-	if hasLabelName && hasLabelValue {
-		gangEnvVars = append(gangEnvVars,
-			v1.EnvVar{Name: serverconfiguration.GangNodeUniformityLabelNameEnvVar, Value: labelName},
-			v1.EnvVar{Name: serverconfiguration.GangNodeUniformityLabelValueEnvVar, Value: labelValue},
-		)
-	}
-
-	allEnvVars := append(baseEnvVars, gangEnvVars...)
-
 	for i := range podSpec.InitContainers {
 		addEnvVarsIfNotExist(&podSpec.InitContainers[i], allEnvVars)
 	}
@@ -246,6 +239,16 @@ func injectArmadaEnvVars(podSpec *v1.PodSpec, jobId string, queue string, jobset
 	}
 }
 
+// buildGangEnvVars builds gang environment variables from GangPlacement.
+func buildGangEnvVars(gangInfo *schedulerobjects.GangPlacement) []v1.EnvVar {
+	return []v1.EnvVar{
+		{Name: serverconfiguration.GangIdEnvVar, Value: gangInfo.GangId},
+		{Name: serverconfiguration.GangCardinalityEnvVar, Value: strconv.FormatUint(uint64(gangInfo.Cardinality), 10)},
+		{Name: serverconfiguration.GangNodeUniformityLabelNameEnvVar, Value: gangInfo.NodeUniformityLabelName},
+		{Name: serverconfiguration.GangNodeUniformityLabelValueEnvVar, Value: gangInfo.NodeUniformityLabelValue},
+	}
+}
+
 // addEnvVarsIfNotExist adds environment variables to a container if they don't already exist.
 func addEnvVarsIfNotExist(container *v1.Container, newEnvVars []v1.EnvVar) {
 	for _, newVar := range newEnvVars {

internal/executor/util/kubernetes_objects_test.go

@@ -14,6 +14,7 @@ import (
 	"github.com/armadaproject/armada/internal/common/util"
 	"github.com/armadaproject/armada/internal/executor/configuration"
 	"github.com/armadaproject/armada/internal/executor/domain"
+	"github.com/armadaproject/armada/internal/scheduler/schedulerobjects"
 	"github.com/armadaproject/armada/pkg/api"
 	"github.com/armadaproject/armada/pkg/armadaevents"
 	"github.com/armadaproject/armada/pkg/executorapi"
@@ -255,7 +256,7 @@ func TestInjectArmadaEnvVars(t *testing.T) {
 		jobId                     string
 		queue                     string
 		jobsetId                  string
-		annotations               map[string]string
+		schedulingMetadata        *armadaevents.SchedulingMetadata
 		existingEnvs              []v1.EnvVar
 		wantEnvs                  map[string]string
 		dontWantEnvs              []string
@@ -310,11 +311,13 @@ func TestInjectArmadaEnvVars(t *testing.T) {
 			jobId:    "job-123",
 			queue:    "queue",
 			jobsetId: "jobset",
-			annotations: map[string]string{
-				serverconfiguration.GangIdAnnotation:                   "gang-789",
-				serverconfiguration.GangCardinalityAnnotation:          "3",
-				serverconfiguration.GangNodeUniformityLabelNameEnvVar:  "rack",
-				serverconfiguration.GangNodeUniformityLabelValueEnvVar: "rack-1",
+			schedulingMetadata: &armadaevents.SchedulingMetadata{
+				GangInfo: &schedulerobjects.GangPlacement{
+					GangId:                   "gang-789",
+					Cardinality:              3,
+					NodeUniformityLabelName:  "rack",
+					NodeUniformityLabelValue: "rack-1",
+				},
 			},
 			wantEnvs: map[string]string{
 				serverconfiguration.JobIdEnvVar:                        "job-123",
@@ -335,33 +338,6 @@ func TestInjectArmadaEnvVars(t *testing.T) {
 				serverconfiguration.GangNodeUniformityLabelValueEnvVar: "rack-1",
 			},
 		},
-		{
-			name:     "skips node uniformity env vars when only label name annotation exists",
-			jobId:    "job-123",
-			queue:    "queue",
-			jobsetId: "jobset",
-			annotations: map[string]string{
-				serverconfiguration.GangNodeUniformityLabelNameEnvVar: "rack",
-			},
-			wantEnvs: map[string]string{
-				serverconfiguration.JobIdEnvVar:    "job-123",
-				serverconfiguration.QueueEnvVar:    "queue",
-				serverconfiguration.JobSetIdEnvVar: "jobset",
-			},
-			dontWantEnvs: []string{
-				serverconfiguration.GangNodeUniformityLabelNameEnvVar,
-				serverconfiguration.GangNodeUniformityLabelValueEnvVar,
-			},
-			wantInitContainerEnvs: map[string]string{
-				serverconfiguration.JobIdEnvVar:    "job-123",
-				serverconfiguration.QueueEnvVar:    "queue",
-				serverconfiguration.JobSetIdEnvVar: "jobset",
-			},
-			dontWantInitContainerEnvs: []string{
-				serverconfiguration.GangNodeUniformityLabelNameEnvVar,
-				serverconfiguration.GangNodeUniformityLabelValueEnvVar,
-			},
-		},
 	}
 
 	for _, tc := range tests {
@@ -371,7 +347,7 @@ func TestInjectArmadaEnvVars(t *testing.T) {
 				InitContainers: []v1.Container{{Name: "init"}},
 			}
 
-			injectArmadaEnvVars(podSpec, tc.jobId, tc.queue, tc.jobsetId, tc.annotations)
+			injectArmadaEnvVars(podSpec, tc.jobId, tc.queue, tc.jobsetId, tc.schedulingMetadata, nil)
 
 			for _, container := range podSpec.InitContainers {
 				envMap := make(map[string]string, len(container.Env))
@@ -401,3 +377,55 @@ func TestInjectArmadaEnvVars(t *testing.T) {
 		})
 	}
 }
+
+func TestInjectArmadaEnvVars_GangScheduling(t *testing.T) {
+	tests := []struct {
+		name               string
+		schedulingMetadata *armadaevents.SchedulingMetadata
+		wantGangId         string
+		wantCardinality    string
+	}{
+		{
+			name: "injects gang env vars from GangInfo",
+			schedulingMetadata: &armadaevents.SchedulingMetadata{
+				GangInfo: &schedulerobjects.GangPlacement{GangId: "gang-123", Cardinality: 5},
+			},
+			wantGangId:      "gang-123",
+			wantCardinality: "5",
+		},
+		{
+			name:               "no gang env vars when GangInfo is nil",
+			schedulingMetadata: nil,
+			wantGangId:         "",
+			wantCardinality:    "",
+		},
+		{
+			name:               "no gang env vars when SchedulingMetadata.GangInfo is nil",
+			schedulingMetadata: &armadaevents.SchedulingMetadata{GangInfo: nil},
+			wantGangId:         "",
+			wantCardinality:    "",
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			podSpec := &v1.PodSpec{Containers: []v1.Container{{Name: "main"}}}
+			injectArmadaEnvVars(podSpec, "job", "queue", "jobset", tt.schedulingMetadata, nil)
+
+			envMap := make(map[string]string)
+			for _, env := range podSpec.Containers[0].Env {
+				envMap[env.Name] = env.Value
+			}
+
+			if tt.wantGangId != "" {
+				assert.Equal(t, tt.wantGangId, envMap[serverconfiguration.GangIdEnvVar])
+				assert.Equal(t, tt.wantCardinality, envMap[serverconfiguration.GangCardinalityEnvVar])
+			} else {
+				_, hasGangId := envMap[serverconfiguration.GangIdEnvVar]
+				assert.False(t, hasGangId)
+				_, hasCardinality := envMap[serverconfiguration.GangCardinalityEnvVar]
+				assert.False(t, hasCardinality)
+			}
+		})
+	}
+}

internal/scheduler/api.go

@@ -152,6 +152,9 @@ func (srv *ExecutorApi) LeaseJobRuns(stream executorapi.ExecutorApi_LeaseJobRuns
 			addAnnotations(submitMsg, PodRequirementsOverlay.Annotations)
 		}
 
+		// Build SchedulingMetadata from annotations and run data
+		schedulingMetadata := buildSchedulingMetadataFromLease(submitMsg, lease.GangNodeUniformityLabelValue)
+
 		srv.addPreemptibleLabel(submitMsg)
 
 		srv.dropDisallowedResources(submitMsg.MainObject.GetPodSpec().PodSpec)
@@ -172,12 +175,13 @@ func (srv *ExecutorApi) LeaseJobRuns(stream executorapi.ExecutorApi_LeaseJobRuns
 		err := stream.Send(&executorapi.LeaseStreamMessage{
 			Event: &executorapi.LeaseStreamMessage_Lease{
 				Lease: &executorapi.JobRunLease{
-					JobRunId: lease.RunID,
-					Queue:    lease.Queue,
-					Jobset:   lease.JobSet,
-					User:     lease.UserID,
-					Groups:   groups,
-					Job:      submitMsg,
+					JobRunId:           lease.RunID,
+					Queue:              lease.Queue,
+					Jobset:             lease.JobSet,
+					User:               lease.UserID,
+					Groups:             groups,
+					Job:                submitMsg,
+					SchedulingMetadata: schedulingMetadata,
 				},
 			},
 		})
@@ -410,3 +414,52 @@ func unmarshalFromCompressedBytes(bytes []byte, decompressor compress.Decompress
 	}
 	return proto.Unmarshal(decompressedBytes, msg)
 }
+
+// buildSchedulingMetadata extracts gang scheduling information from the SubmitJob message
+// and builds a SchedulingMetadata protobuf message for the executor.
+
+// buildSchedulingMetadataFromLease extracts gang scheduling information from the SubmitJob message
+// and builds a SchedulingMetadata protobuf message for the executor.
+// labelValue is the actual value determined during scheduling for the gang's uniformity label.
+func buildSchedulingMetadataFromLease(submitMsg *armadaevents.SubmitJob, labelValue string) *armadaevents.SchedulingMetadata {
+	if submitMsg == nil {
+		return nil
+	}
+
+	var gangId string
+	var cardinality uint32
+	var labelName string
+
+	if submitMsg.Gang != nil {
+		gangId = submitMsg.Gang.GangId
+		cardinality = submitMsg.Gang.Cardinality
+		labelName = submitMsg.Gang.NodeUniformityLabelName
+	} else if submitMsg.ObjectMeta != nil && submitMsg.ObjectMeta.Annotations != nil {
+		annotations := submitMsg.ObjectMeta.Annotations
+		gangId = annotations[constants.GangIdAnnotation]
+		if gangId == "" {
+			return nil
+		}
+
+		if cardinalityStr := annotations[constants.GangCardinalityAnnotation]; cardinalityStr != "" {
+			if val, err := strconv.ParseUint(cardinalityStr, 10, 32); err == nil {
+				cardinality = uint32(val)
+			}
+		}
+
+		labelName = annotations[constants.GangNodeUniformityLabelAnnotation]
+	} else {
+		return nil
+	}
+
+	gangPlacement := &schedulerobjects.GangPlacement{
+		GangId:                   gangId,
+		Cardinality:              cardinality,
+		NodeUniformityLabelName:  labelName,
+		NodeUniformityLabelValue: labelValue,
+	}
+
+	return &armadaevents.SchedulingMetadata{
+		GangInfo: gangPlacement,
+	}
+}