Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add post processor that will add a gang id label (#265) #4029

Open
wants to merge 1 commit into
base: master
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions internal/server/configuration/types.go
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
19 changes: 19 additions & 0 deletions internal/server/submit/conversion/post_process.go
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,7 @@ var (
msgLevelProcessors = []msgProcessor{
templateMeta,
defaultGangNodeUniformityLabel,
addGangIdLabel,
}
podLevelProcessors = []podProcessor{
defaultActiveDeadlineSeconds,
Expand Down Expand Up @@ -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) {
Expand Down
50 changes: 50 additions & 0 deletions internal/server/submit/conversion/post_process_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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{
Expand Down