Merge pull request #879 from buildkite/sneha/add-podTemplateName

feat: pass in pod-template-name annotation on jobs built from podTemplates

Author: ss1909

internal/controller/config/config.go

@@ -19,6 +19,7 @@ const (
 	JobURLAnnotation                      = "buildkite.com/job-url"
 	PriorityAnnotation                    = "buildkite.com/job-priority"
 	PipelineSlugAnnotation                = "buildkite.com/pipeline-slug"
+	PodTemplateAnnotation                 = "buildkite.com/pod-template-name"
 	DefaultNamespace                      = "default"
 	DefaultImagePullBackOffGracePeriod    = 30 * time.Second
 	DefaultJobCancelCheckerPollInterval   = 5 * time.Second

internal/controller/scheduler/scheduler.go

@@ -340,6 +340,10 @@ func (w *worker) Build(podSpec *corev1.PodSpec, skipCheckout bool, inputs buildI
 	}
 	kjob.Annotations[config.PriorityAnnotation] = strconv.Itoa(inputs.priority)
 
+	if inputs.k8sPlugin != nil && inputs.k8sPlugin.PodSpec == nil && inputs.k8sPlugin.PodTemplate != "" {
+		kjob.Annotations[config.PodTemplateAnnotation] = inputs.k8sPlugin.PodTemplate
+	}
+
 	kjob.Spec.Template.Labels = kjob.Labels
 	kjob.Spec.Template.Annotations = kjob.Annotations
 	kjob.Spec.BackoffLimit = ptr.To[int32](0)

internal/controller/scheduler/scheduler_test.go

@@ -1854,3 +1854,50 @@ func TestBuildSafeToEvictPluginMetadataOverride(t *testing.T) {
 	assert.Equal(t, "true", kjob.Annotations["cluster-autoscaler.kubernetes.io/safe-to-evict"])
 	assert.Equal(t, "true", kjob.Spec.Template.Annotations["cluster-autoscaler.kubernetes.io/safe-to-evict"])
 }
+
+func TestBuildPodTemplateAnnotation(t *testing.T) {
+	t.Parallel()
+
+	pluginsYAML := `- github.com/buildkite-plugins/kubernetes-buildkite-plugin:
+    podTemplate: my-pod-template`
+
+	pluginsJSON, err := yaml.YAMLToJSONStrict([]byte(pluginsYAML))
+	require.NoError(t, err)
+
+	job := &api.AgentJob{
+		ID:      "abc",
+		Command: "echo hello world",
+		Env:     map[string]string{"BUILDKITE_PLUGINS": string(pluginsJSON)},
+	}
+	sjob := &api.AgentScheduledJob{}
+	worker := New(slog.Default(), nil, nil, Config{
+		Image: "buildkite/agent:latest",
+	})
+	inputs, err := worker.ParseJob(job, sjob)
+	require.NoError(t, err)
+	kjob, err := worker.Build(&corev1.PodSpec{}, false, inputs)
+	require.NoError(t, err)
+
+	require.Equal(t, "my-pod-template", kjob.Annotations[config.PodTemplateAnnotation])
+	require.Equal(t, "my-pod-template", kjob.Spec.Template.Annotations[config.PodTemplateAnnotation])
+}
+
+func TestBuildNoPodTemplateAnnotationWhenAbsent(t *testing.T) {
+	t.Parallel()
+
+	job := &api.AgentJob{
+		ID:      "abc",
+		Command: "echo hello world",
+	}
+	sjob := &api.AgentScheduledJob{}
+	worker := New(slog.Default(), nil, nil, Config{
+		Image: "buildkite/agent:latest",
+	})
+	inputs, err := worker.ParseJob(job, sjob)
+	require.NoError(t, err)
+	kjob, err := worker.Build(&corev1.PodSpec{}, false, inputs)
+	require.NoError(t, err)
+
+	require.Empty(t, kjob.Annotations[config.PodTemplateAnnotation])
+	require.Empty(t, kjob.Spec.Template.Annotations[config.PodTemplateAnnotation])
+}