From 45efc3e9e266c5a73834cd1731cc302a7fd82762 Mon Sep 17 00:00:00 2001
From: Christopher Martin <Chris.Martin@gresearch.co.uk>
Date: Tue, 29 Oct 2024 10:33:42 +0000
Subject: [PATCH] Add post processor that will add a gang id label (#265)

* Add post processor that will add a gang id label

* lint

* whitespace
---
 internal/server/configuration/types.go        |  2 +
 .../server/submit/conversion/post_process.go  | 19 +++++++
 .../submit/conversion/post_process_test.go    | 50 +++++++++++++++++++
 3 files changed, 71 insertions(+)

diff --git a/internal/server/configuration/types.go b/internal/server/configuration/types.go
index 57c5dde4489..91c26f424cd 100644
--- a/internal/server/configuration/types.go
+++ b/internal/server/configuration/types.go
@@ -93,6 +93,8 @@ type SubmissionConfig struct {
 	MaxOversubscriptionByResourceRequest map[string]float64
 	// Enforce that an init containers requestion non-integer cpu. This is due to https://github.com/kubernetes/kubernetes/issues/112228
 	AssertInitContainersRequestFractionalCpu bool
+	// Controls whether we add the gang id annotation as a label.
+	AddGangIdLabel bool
 }
 
 // TODO: we can probably just typedef this to map[string]string
diff --git a/internal/server/submit/conversion/post_process.go b/internal/server/submit/conversion/post_process.go
index 1c71c8664a5..dfe57ccc3d4 100644
--- a/internal/server/submit/conversion/post_process.go
+++ b/internal/server/submit/conversion/post_process.go
@@ -23,6 +23,7 @@ var (
 	msgLevelProcessors = []msgProcessor{
 		templateMeta,
 		defaultGangNodeUniformityLabel,
+		addGangIdLabel,
 	}
 	podLevelProcessors = []podProcessor{
 		defaultActiveDeadlineSeconds,
@@ -166,6 +167,24 @@ func defaultGangNodeUniformityLabel(msg *armadaevents.SubmitJob, config configur
 	}
 }
 
+// Add a gangId label if the gangId annotation is set.  We do this because labels are much faster to search on than
+// annotations and a gang may want to hit the kubeapi to find its other gang members.
+func addGangIdLabel(msg *armadaevents.SubmitJob, config configuration.SubmissionConfig) {
+	if !config.AddGangIdLabel {
+		return
+	}
+
+	gangId := msg.GetObjectMeta().GetAnnotations()[configuration.GangIdAnnotation]
+	if gangId != "" {
+		labels := msg.GetObjectMeta().GetLabels()
+		if labels == nil {
+			labels = map[string]string{}
+		}
+		labels[configuration.GangIdAnnotation] = gangId
+		msg.GetObjectMeta().Labels = labels
+	}
+}
+
 // Templates the JobId in labels and annotations. This allows users to define labels and annotations containing the string
 // {JobId} and have it populated with the actual id of the job.
 func templateMeta(msg *armadaevents.SubmitJob, _ configuration.SubmissionConfig) {
diff --git a/internal/server/submit/conversion/post_process_test.go b/internal/server/submit/conversion/post_process_test.go
index 2573960daa6..1374bb8eb44 100644
--- a/internal/server/submit/conversion/post_process_test.go
+++ b/internal/server/submit/conversion/post_process_test.go
@@ -584,6 +584,56 @@ func TestDefaultTerminationGracePeriod(t *testing.T) {
 	}
 }
 
+func TestAddGangIdLabel(t *testing.T) {
+	tests := map[string]struct {
+		annotations    map[string]string
+		initialLabels  map[string]string
+		expectedLabels map[string]string
+		enabled        bool
+	}{
+		"Unchanged if no gang id set": {
+			annotations: map[string]string{},
+			enabled:     true,
+		},
+		"Label added if gang id set": {
+			annotations: map[string]string{
+				configuration.GangIdAnnotation: "foo",
+			},
+			expectedLabels: map[string]string{
+				configuration.GangIdAnnotation: "foo",
+			},
+			enabled: true,
+		},
+		"Doesn't modify existing labels": {
+			annotations: map[string]string{
+				configuration.GangIdAnnotation: "foo",
+			},
+			initialLabels: map[string]string{
+				"fish": "chips",
+			},
+			expectedLabels: map[string]string{
+				"fish":                         "chips",
+				configuration.GangIdAnnotation: "foo",
+			},
+			enabled: true,
+		},
+		"Unchanged if disabled": {
+			annotations: map[string]string{
+				configuration.GangIdAnnotation: "foo",
+			},
+			enabled: false,
+		},
+	}
+	for name, tc := range tests {
+		t.Run(name, func(t *testing.T) {
+			submitMsg := submitMsgFromAnnotations(tc.annotations)
+			submitMsg.ObjectMeta.Labels = tc.initialLabels
+			addGangIdLabel(submitMsg, configuration.SubmissionConfig{AddGangIdLabel: tc.enabled})
+			assert.Equal(t, tc.expectedLabels, submitMsg.ObjectMeta.Labels)
+		})
+	}
+}
+
 func submitMsgFromAnnotations(annotations map[string]string) *armadaevents.SubmitJob {
 	return &armadaevents.SubmitJob{
 		ObjectMeta: &armadaevents.ObjectMeta{