Repository: armadaproject/armada

Issue:
Inject gang info in Armada jobs/pods
If Armada were to inject the following two environment variables:
* `ARMADA_GANG_NODE_UNIFORMITY_LABEL_NAME` - the label name for the node set on which the gang is scheduled
* `ARMADA_GANG_NODE_UNIFORMITY_LABEL_VALUE` - the label value for the node set on which the gang is scheduled

This would allow integrations like Spark on Armada to support dynamic allocation by submitting the initial set of driver + executors as a non-preemptible gang, and then submitting additional dynamic & preemptible executors with `nodeSelectors` field set to the value of `$ARMADA_GANG_NODE_UNIFORMITY_LABEL_NAME: $ARMADA_GANG_NODE_UNIFORMITY_LABEL_VALUE`.

Discussion:

Solution patch:
diff --git a/internal/executor/util/kubernetes_object.go b/internal/executor/util/kubernetes_object.go
index ccfcf7de..56be546a 100644
--- a/internal/executor/util/kubernetes_object.go
+++ b/internal/executor/util/kubernetes_object.go
@@ -12,6 +12,7 @@ import (
 	"github.com/armadaproject/armada/internal/common/util"
 	"github.com/armadaproject/armada/internal/executor/configuration"
 	"github.com/armadaproject/armada/internal/executor/domain"
+	serverconfiguration "github.com/armadaproject/armada/internal/server/configuration"
 	"github.com/armadaproject/armada/pkg/api"
 	"github.com/armadaproject/armada/pkg/armadaevents"
 	"github.com/armadaproject/armada/pkg/executorapi"
@@ -130,6 +131,8 @@ func CreatePodFromExecutorApiJob(job *executorapi.JobRunLease, defaults *configu
 	applyDefaults(podSpec, defaults)
 	setRestartPolicyNever(podSpec)
 
+	injectArmadaEnvVars(podSpec, jobId, job.Queue, job.Jobset, annotation)
+
 	pod := &v1.Pod{
 		ObjectMeta: metav1.ObjectMeta{
 			Name:        common.PodNamePrefix + job.Job.JobId + "-" + strconv.Itoa(0),
@@ -170,6 +173,8 @@ func CreatePod(job *api.Job, defaults *configuration.PodDefaults) *v1.Pod {
 
 	setRestartPolicyNever(podSpec)
 
+	injectArmadaEnvVars(podSpec, job.Id, job.Queue, job.JobSetId, annotation)
+
 	pod := &v1.Pod{
 		ObjectMeta: metav1.ObjectMeta{
 			Name:        common.PodNamePrefix + job.Id + "-0",
@@ -195,3 +200,67 @@ func applyDefaults(spec *v1.PodSpec, defaults *configuration.PodDefaults) {
 func setRestartPolicyNever(podSpec *v1.PodSpec) {
 	podSpec.RestartPolicy = v1.RestartPolicyNever
 }
+
+// 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
+	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,
+		})
+	}
+
+	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)
+	}
+	for i := range podSpec.Containers {
+		addEnvVarsIfNotExist(&podSpec.Containers[i], allEnvVars)
+	}
+}
+
+// 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 {
+		if !hasEnvVar(container.Env, newVar.Name) {
+			container.Env = append(container.Env, newVar)
+		}
+	}
+}
+
+// hasEnvVar checks if an environment variable with the given name exists in the slice.
+func hasEnvVar(envVars []v1.EnvVar, name string) bool {
+	for _, env := range envVars {
+		if env.Name == name {
+			return true
+		}
+	}
+	return false
+}
diff --git a/internal/scheduler/scheduler.go b/internal/scheduler/scheduler.go
index 17e22692..50f58fdd 100644
--- a/internal/scheduler/scheduler.go
+++ b/internal/scheduler/scheduler.go
@@ -704,6 +704,7 @@ func AppendEventSequencesFromScheduledJobs(eventSequences []*armadaevents.EventS
 								Tolerations: armadaslices.Map(jctx.AdditionalTolerations, func(t v1.Toleration) *v1.Toleration {
 									return &t
 								}),
+								Annotations: getGangNodeUniformityAnnotations(jctx),
 							},
 							Pool: run.Pool(),
 						},
@@ -1192,3 +1193,15 @@ func (s *Scheduler) ensureDbUpToDate(ctx *armadacontext.Context, pollInterval ti
 		}
 	}
 }
+
+// getGangNodeUniformityAnnotations returns annotations for gang node uniformity environment variables
+func getGangNodeUniformityAnnotations(jctx *schedulercontext.JobSchedulingContext) map[string]string {
+	if jctx.GangNodeUniformityLabelName == "" || jctx.GangNodeUniformityLabelValue == "" {
+		return nil
+	}
+
+	return map[string]string{
+		configuration.GangNodeUniformityLabelNameEnvVar:  jctx.GangNodeUniformityLabelName,
+		configuration.GangNodeUniformityLabelValueEnvVar: jctx.GangNodeUniformityLabelValue,
+	}
+}
diff --git a/internal/scheduler/scheduling/context/gang.go b/internal/scheduler/scheduling/context/gang.go
index 1325270e..99cc4892 100644
--- a/internal/scheduler/scheduling/context/gang.go
+++ b/internal/scheduler/scheduling/context/gang.go
@@ -108,3 +108,11 @@ func (gctx *GangSchedulingContext) Fit() GangSchedulingFit {
 	}
 	return f
 }
+
+// SetGangNodeUniformityValues sets the gang node uniformity label name and value for all jobs in the gang.
+func (gctx *GangSchedulingContext) SetGangNodeUniformityValues(labelName, labelValue string) {
+	for _, jctx := range gctx.JobSchedulingContexts {
+		jctx.GangNodeUniformityLabelName = labelName
+		jctx.GangNodeUniformityLabelValue = labelValue
+	}
+}
diff --git a/internal/scheduler/scheduling/context/job.go b/internal/scheduler/scheduling/context/job.go
index f00b52e7..7f26b630 100644
--- a/internal/scheduler/scheduling/context/job.go
+++ b/internal/scheduler/scheduling/context/job.go
@@ -63,6 +63,10 @@ type JobSchedulingContext struct {
 	PreemptionDescription string
 	// If this job context should contribute to the billable resource of the queue
 	Billable bool
+	// Gang node uniformity label name (e.g., "rack") - only set for gang jobs with uniformity requirements
+	GangNodeUniformityLabelName string
+	// Gang node uniformity label value (e.g., "rack-1") - the actual value selected during scheduling
+	GangNodeUniformityLabelValue string
 }
 
 func (jctx *JobSchedulingContext) IsHomeJob(currentPool string) bool {
diff --git a/internal/scheduler/scheduling/gang_scheduler.go b/internal/scheduler/scheduling/gang_scheduler.go
index 0b212d65..4f8aaa73 100644
--- a/internal/scheduler/scheduling/gang_scheduler.go
+++ b/internal/scheduler/scheduling/gang_scheduler.go
@@ -194,12 +194,16 @@ func (sch *GangScheduler) trySchedule(ctx *armadacontext.Context, gctx *context.
 			currentFit := gctx.Fit()
 			if currentFit.NumScheduled == gctx.Cardinality() && currentFit.MeanPreemptedAtPriority == float64(internaltypes.MinPriority) {
 				// Best possible; no need to keep looking.
+				// Store the selected label value in all job contexts
+				gctx.SetGangNodeUniformityValues(nodeUniformity, value)
 				txn.Commit()
 				return true, "", nil
 			}
 			if bestValue == "" || bestFit.Less(currentFit) {
 				if i == len(nodeUniformityLabelValues) {
 					// Minimal meanScheduledAtPriority and no more options; commit and return.
+					// Store the selected label value in all job contexts
+					gctx.SetGangNodeUniformityValues(nodeUniformity, value)
 					txn.Commit()
 					return true, "", nil
 				}
@@ -215,6 +219,8 @@ func (sch *GangScheduler) trySchedule(ctx *armadacontext.Context, gctx *context.
 		unschedulableReason = "at least one job in the gang does not fit on any node"
 		return ok, unschedulableReason, err
 	}
+	// Store the selected label value in all job contexts
+	gctx.SetGangNodeUniformityValues(nodeUniformity, bestValue)
 	addNodeSelectorToGctx(gctx, gctx.NodeUniformityLabel(), bestValue)
 	return sch.tryScheduleGang(ctx, gctx)
 }
diff --git a/internal/server/configuration/constants.go b/internal/server/configuration/constants.go
index c0823cd4..06b8bc0f 100644
--- a/internal/server/configuration/constants.go
+++ b/internal/server/configuration/constants.go
@@ -11,6 +11,21 @@ const (
 	// Specifically, if provided, all gang jobs are scheduled onto nodes for which the value of the provided label is equal.
 	// Used to ensure, e.g., that all gang jobs are scheduled onto the same cluster or rack.
 	GangNodeUniformityLabelAnnotation = "armadaproject.io/gangNodeUniformityLabel"
+
+	// internalEnvVarPrefix is the prefix for all Armada-injected environment variables
+	internalEnvVarPrefix = "ARMADA_"
+
+	// Environment variables injected into all jobs
+	JobIdEnvVar    = internalEnvVarPrefix + "JOB_ID"
+	QueueEnvVar    = internalEnvVarPrefix + "QUEUE"
+	JobSetIdEnvVar = internalEnvVarPrefix + "JOB_SET_ID"
+
+	// Additional environment variables for gang-scheduled jobs
+	GangIdEnvVar                       = internalEnvVarPrefix + "GANG_ID"
+	GangCardinalityEnvVar              = internalEnvVarPrefix + "GANG_CARDINALITY"
+	GangNodeUniformityLabelNameEnvVar  = internalEnvVarPrefix + "GANG_NODE_UNIFORMITY_LABEL_NAME"
+	GangNodeUniformityLabelValueEnvVar = internalEnvVarPrefix + "GANG_NODE_UNIFORMITY_LABEL_VALUE"
+
 	// GangNumJobsScheduledAnnotation is set by the scheduler and indicates how many gang jobs were scheduled.
 	// FailFastAnnotation, if set to true, ensures Armada does not re-schedule jobs that fail to start.
 	// Instead, the job the pod is part of fails immediately.



Test patch:
diff --git a/internal/executor/util/kubernetes_objects_test.go b/internal/executor/util/kubernetes_objects_test.go
index 639aa843..77b56b07 100644
--- a/internal/executor/util/kubernetes_objects_test.go
+++ b/internal/executor/util/kubernetes_objects_test.go
@@ -14,6 +14,7 @@ import (
 	"github.com/armadaproject/armada/internal/common"
 	"github.com/armadaproject/armada/internal/executor/configuration"
 	"github.com/armadaproject/armada/internal/executor/domain"
+	serverconfiguration "github.com/armadaproject/armada/internal/server/configuration"
 	"github.com/armadaproject/armada/pkg/api"
 	"github.com/armadaproject/armada/pkg/armadaevents"
 	"github.com/armadaproject/armada/pkg/executorapi"
@@ -158,13 +159,20 @@ func TestCreatePodFromExecutorApiJob(t *testing.T) {
 			MainObject: &armadaevents.KubernetesMainObject{
 				Object: &armadaevents.KubernetesMainObject_PodSpec{
 					PodSpec: &armadaevents.PodSpecWithAvoidList{
-						PodSpec: &v1.PodSpec{},
+						PodSpec: &v1.PodSpec{
+							Containers: []v1.Container{{Name: "test", Image: "test"}},
+						},
 					},
 				},
 			},
 		},
 	}
 
+	expectedEnvVars := []v1.EnvVar{
+		{Name: serverconfiguration.JobIdEnvVar, Value: jobId},
+		{Name: serverconfiguration.QueueEnvVar, Value: "queue"},
+		{Name: serverconfiguration.JobSetIdEnvVar, Value: "job-set"},
+	}
 	expectedPod := &v1.Pod{
 		ObjectMeta: metav1.ObjectMeta{
 			Name:      fmt.Sprintf("armada-%s-0", jobId),
@@ -185,6 +193,7 @@ func TestCreatePodFromExecutorApiJob(t *testing.T) {
 		Spec: v1.PodSpec{
 			RestartPolicy: v1.RestartPolicyNever,
 			SchedulerName: "scheduler-name",
+			Containers:    []v1.Container{{Name: "test", Image: "test", Env: expectedEnvVars}},
 		},
 	}
 
@@ -240,3 +249,156 @@ func createBasicJobRunLease() *executorapi.JobRunLease {
 		},
 	}
 }
+
+func TestInjectArmadaEnvVars(t *testing.T) {
+	tests := []struct {
+		name                      string
+		jobId                     string
+		queue                     string
+		jobsetId                  string
+		annotations               map[string]string
+		existingEnvs              []v1.EnvVar
+		wantEnvs                  map[string]string
+		dontWantEnvs              []string
+		wantInitContainerEnvs     map[string]string
+		dontWantInitContainerEnvs []string
+	}{
+		{
+			name:     "injects base env vars for non-gang job",
+			jobId:    "job-123",
+			queue:    "test-queue",
+			jobsetId: "jobset-456",
+			wantEnvs: map[string]string{
+				serverconfiguration.JobIdEnvVar:    "job-123",
+				serverconfiguration.QueueEnvVar:    "test-queue",
+				serverconfiguration.JobSetIdEnvVar: "jobset-456",
+			},
+			dontWantEnvs: []string{
+				serverconfiguration.GangIdEnvVar,
+				serverconfiguration.GangCardinalityEnvVar,
+			},
+			wantInitContainerEnvs: map[string]string{
+				serverconfiguration.JobIdEnvVar:    "job-123",
+				serverconfiguration.QueueEnvVar:    "test-queue",
+				serverconfiguration.JobSetIdEnvVar: "jobset-456",
+			},
+			dontWantInitContainerEnvs: []string{
+				serverconfiguration.GangIdEnvVar,
+				serverconfiguration.GangCardinalityEnvVar,
+			},
+		},
+		{
+			name:     "preserves user-defined env vars",
+			jobId:    "new-job",
+			queue:    "new-queue",
+			jobsetId: "new-jobset",
+			existingEnvs: []v1.EnvVar{
+				{Name: serverconfiguration.JobIdEnvVar, Value: "existing-job"},
+			},
+			wantEnvs: map[string]string{
+				serverconfiguration.JobIdEnvVar:    "existing-job", // preserved in main container
+				serverconfiguration.QueueEnvVar:    "new-queue",
+				serverconfiguration.JobSetIdEnvVar: "new-jobset",
+			},
+			wantInitContainerEnvs: map[string]string{
+				serverconfiguration.JobIdEnvVar:    "new-job", // init container gets new value
+				serverconfiguration.QueueEnvVar:    "new-queue",
+				serverconfiguration.JobSetIdEnvVar: "new-jobset",
+			},
+		},
+		{
+			name:     "injects all gang-related env vars for fully configured gang job",
+			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",
+			},
+			wantEnvs: map[string]string{
+				serverconfiguration.JobIdEnvVar:                        "job-123",
+				serverconfiguration.QueueEnvVar:                        "queue",
+				serverconfiguration.JobSetIdEnvVar:                     "jobset",
+				serverconfiguration.GangIdEnvVar:                       "gang-789",
+				serverconfiguration.GangCardinalityEnvVar:              "3",
+				serverconfiguration.GangNodeUniformityLabelNameEnvVar:  "rack",
+				serverconfiguration.GangNodeUniformityLabelValueEnvVar: "rack-1",
+			},
+			wantInitContainerEnvs: map[string]string{
+				serverconfiguration.JobIdEnvVar:                        "job-123",
+				serverconfiguration.QueueEnvVar:                        "queue",
+				serverconfiguration.JobSetIdEnvVar:                     "jobset",
+				serverconfiguration.GangIdEnvVar:                       "gang-789",
+				serverconfiguration.GangCardinalityEnvVar:              "3",
+				serverconfiguration.GangNodeUniformityLabelNameEnvVar:  "rack",
+				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 {
+		t.Run(tc.name, func(t *testing.T) {
+			podSpec := &v1.PodSpec{
+				Containers:     []v1.Container{{Name: "main", Env: tc.existingEnvs}},
+				InitContainers: []v1.Container{{Name: "init"}},
+			}
+
+			injectArmadaEnvVars(podSpec, tc.jobId, tc.queue, tc.jobsetId, tc.annotations)
+
+			for _, container := range podSpec.InitContainers {
+				envMap := make(map[string]string, len(container.Env))
+				for _, env := range container.Env {
+					envMap[env.Name] = env.Value
+				}
+
+				for name, value := range tc.wantInitContainerEnvs {
+					assert.Equal(t, value, envMap[name])
+				}
+				for _, name := range tc.dontWantInitContainerEnvs {
+					assert.NotContains(t, envMap, name)
+				}
+			}
+			for _, container := range podSpec.Containers {
+				envMap := make(map[string]string, len(container.Env))
+				for _, env := range container.Env {
+					envMap[env.Name] = env.Value
+				}
+				for name, value := range tc.wantEnvs {
+					assert.Equal(t, value, envMap[name])
+				}
+				for _, name := range tc.dontWantEnvs {
+					assert.NotContains(t, envMap, name)
+				}
+			}
+		})
+	}
+}