Add --pod-template-annotation and --pod-template-label flags (#59)

* Add --pod-template-annotation and --pod-template-label flags

* Update docs.

apis/v1alpha1/application_profile_types.go

@@ -40,31 +40,33 @@ const (
 type Flag string
 
 const (
-	CmdFlag         Flag = "cmd"
-	ParallelismFlag Flag = "parallelism"
-	CompletionsFlag Flag = "completions"
-	ReplicasFlag    Flag = "replicas"
-	MinReplicasFlag Flag = "min-replicas"
-	MaxReplicasFlag Flag = "max-replicas"
-	RequestFlag     Flag = "request"
-	LocalQueueFlag  Flag = "localqueue"
-	RayClusterFlag  Flag = "raycluster"
-	ArrayFlag       Flag = "array"
-	CpusPerTaskFlag Flag = "cpus-per-task"
-	ErrorFlag       Flag = "error"
-	GpusPerTaskFlag Flag = "gpus-per-task"
-	InputFlag       Flag = "input"
-	JobNameFlag     Flag = "job-name"
-	MemPerNodeFlag  Flag = "mem"
-	MemPerCPUFlag   Flag = "mem-per-cpu"
-	MemPerGPUFlag   Flag = "mem-per-gpu"
-	MemPerTaskFlag  Flag = "mem-per-task"
-	NodesFlag       Flag = "nodes"
-	NTasksFlag      Flag = "ntasks"
-	OutputFlag      Flag = "output"
-	PartitionFlag   Flag = "partition"
-	PriorityFlag    Flag = "priority"
-	TimeFlag        Flag = "time"
+	CmdFlag                   Flag = "cmd"
+	ParallelismFlag           Flag = "parallelism"
+	CompletionsFlag           Flag = "completions"
+	ReplicasFlag              Flag = "replicas"
+	MinReplicasFlag           Flag = "min-replicas"
+	MaxReplicasFlag           Flag = "max-replicas"
+	RequestFlag               Flag = "request"
+	LocalQueueFlag            Flag = "localqueue"
+	RayClusterFlag            Flag = "raycluster"
+	ArrayFlag                 Flag = "array"
+	CpusPerTaskFlag           Flag = "cpus-per-task"
+	ErrorFlag                 Flag = "error"
+	GpusPerTaskFlag           Flag = "gpus-per-task"
+	InputFlag                 Flag = "input"
+	JobNameFlag               Flag = "job-name"
+	MemPerNodeFlag            Flag = "mem"
+	MemPerCPUFlag             Flag = "mem-per-cpu"
+	MemPerGPUFlag             Flag = "mem-per-gpu"
+	MemPerTaskFlag            Flag = "mem-per-task"
+	NodesFlag                 Flag = "nodes"
+	NTasksFlag                Flag = "ntasks"
+	OutputFlag                Flag = "output"
+	PartitionFlag             Flag = "partition"
+	PriorityFlag              Flag = "priority"
+	TimeFlag                  Flag = "time"
+	PodTemplateLabelFlag      Flag = "pod-template-label"
+	PodTemplateAnnotationFlag Flag = "pod-template-annotation"
 )
 
 // TemplateReference is the name of the template.
@@ -116,17 +118,18 @@ type SupportedMode struct {
 	Template TemplateReference `json:"template"`
 
 	// requiredFlags point which cli flags are required to be passed in order to fill the gaps in the templates.
-	// Possible values are cmd, parallelism, completions, replicas, min-replicas, max-replicas, request, localqueue, and raycluster.
-	// replicas, min-replicas, and max-replicas flags used only for RayJob and RayCluster mode.
+	// Possible values are cmd, parallelism, completions, replicas, min-replicas, max-replicas, request, localqueue,
+	// raycluster, pod-template-label and pod-template-annotation.
+	// The replicas, min-replicas, and max-replicas flags used only for RayJob and RayCluster mode.
 	// The raycluster flag used only for the RayJob mode.
 	// The request flag used only for Interactive and Job modes.
 	// The cmd flag used only for Interactive, Job, and RayJob.
 	// The time and priority flags can be used in all modes.
 	// If the raycluster flag are set, none of localqueue, replicas, min-replicas, or max-replicas can be set.
-	// For the Slurm mode, the possible values are: array, cpus-per-task, error, gpus-per-task, input, job-name, mem, mem-per-cpu,
-	// mem-per-gpu, mem-per-task, nodes, ntasks, output, partition, localqueue.
+	// For the Slurm mode, the possible values are: array, cpus-per-task, error, gpus-per-task, input, job-name, mem,
+	// mem-per-cpu, mem-per-gpu, mem-per-task, nodes, ntasks, output, partition, localqueue.
 	//
-	// cmd and requests values are going to be added only to the first primary container.
+	// The cmd and requests values are going to be added only to the first primary container.
 	//
 	// +optional
 	// +listType=set

config/crd/bases/kjobctl.x-k8s.io_applicationprofiles.yaml

@@ -57,17 +57,18 @@ spec:
                     requiredFlags:
                       description: |-
                         requiredFlags point which cli flags are required to be passed in order to fill the gaps in the templates.
-                        Possible values are cmd, parallelism, completions, replicas, min-replicas, max-replicas, request, localqueue, and raycluster.
-                        replicas, min-replicas, and max-replicas flags used only for RayJob and RayCluster mode.
+                        Possible values are cmd, parallelism, completions, replicas, min-replicas, max-replicas, request, localqueue,
+                        raycluster, pod-template-label and pod-template-annotation.
+                        The replicas, min-replicas, and max-replicas flags used only for RayJob and RayCluster mode.
                         The raycluster flag used only for the RayJob mode.
                         The request flag used only for Interactive and Job modes.
                         The cmd flag used only for Interactive, Job, and RayJob.
                         The time and priority flags can be used in all modes.
                         If the raycluster flag are set, none of localqueue, replicas, min-replicas, or max-replicas can be set.
-                        For the Slurm mode, the possible values are: array, cpus-per-task, error, gpus-per-task, input, job-name, mem, mem-per-cpu,
-                        mem-per-gpu, mem-per-task, nodes, ntasks, output, partition, localqueue.
+                        For the Slurm mode, the possible values are: array, cpus-per-task, error, gpus-per-task, input, job-name, mem,
+                        mem-per-cpu, mem-per-gpu, mem-per-task, nodes, ntasks, output, partition, localqueue.
 
-                        cmd and requests values are going to be added only to the first primary container.
+                        The cmd and requests values are going to be added only to the first primary container.
                       items:
                         enum:
                         - cmd

docs/commands/kjobctl_create/kjobctl_create_interactive.md

@@ -99,6 +99,24 @@ kjobctl create interactive --profile APPLICATION_PROFILE_NAME [--localqueue LOCA
             <p>The length of time (like 5s, 2m, or 3h, higher than zero) to wait until at least one pod is running.</p>
         </td>
     </tr>
+    <tr>
+        <td colspan="2">--pod-template-annotation &lt;comma-separated &#39;key=value&#39; pairs&gt;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;Default: []</td>
+    </tr>
+    <tr>
+        <td></td>
+        <td style="line-height: 130%; word-wrap: break-word;">
+            <p>Specifies one or more annotations for the Pod template.</p>
+        </td>
+    </tr>
+    <tr>
+        <td colspan="2">--pod-template-label &lt;comma-separated &#39;key=value&#39; pairs&gt;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;Default: []</td>
+    </tr>
+    <tr>
+        <td></td>
+        <td style="line-height: 130%; word-wrap: break-word;">
+            <p>Specifies one or more labels for the Pod template.</p>
+        </td>
+    </tr>
     <tr>
         <td colspan="2">--priority string</td>
     </tr>

docs/commands/kjobctl_create/kjobctl_create_job.md

@@ -111,6 +111,24 @@ kjobctl create job --profile APPLICATION_PROFILE_NAME [--localqueue LOCAL_QUEUE_
             <p>Parallelism specifies the maximum desired number of pods the job should run at any given time.</p>
         </td>
     </tr>
+    <tr>
+        <td colspan="2">--pod-template-annotation &lt;comma-separated &#39;key=value&#39; pairs&gt;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;Default: []</td>
+    </tr>
+    <tr>
+        <td></td>
+        <td style="line-height: 130%; word-wrap: break-word;">
+            <p>Specifies one or more annotations for the Pod template.</p>
+        </td>
+    </tr>
+    <tr>
+        <td colspan="2">--pod-template-label &lt;comma-separated &#39;key=value&#39; pairs&gt;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;Default: []</td>
+    </tr>
+    <tr>
+        <td></td>
+        <td style="line-height: 130%; word-wrap: break-word;">
+            <p>Specifies one or more labels for the Pod template.</p>
+        </td>
+    </tr>
     <tr>
         <td colspan="2">--priority string</td>
     </tr>

docs/commands/kjobctl_create/kjobctl_create_raycluster.md

@@ -103,6 +103,24 @@ kjobctl create raycluster --profile APPLICATION_PROFILE_NAME [--localqueue LOCAL
             <p>Output format. One of: (json, yaml, name, go-template, go-template-file, template, templatefile, jsonpath, jsonpath-as-json, jsonpath-file).</p>
         </td>
     </tr>
+    <tr>
+        <td colspan="2">--pod-template-annotation &lt;comma-separated &#39;key=value&#39; pairs&gt;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;Default: []</td>
+    </tr>
+    <tr>
+        <td></td>
+        <td style="line-height: 130%; word-wrap: break-word;">
+            <p>Specifies one or more annotations for the Pod template.</p>
+        </td>
+    </tr>
+    <tr>
+        <td colspan="2">--pod-template-label &lt;comma-separated &#39;key=value&#39; pairs&gt;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;Default: []</td>
+    </tr>
+    <tr>
+        <td></td>
+        <td style="line-height: 130%; word-wrap: break-word;">
+            <p>Specifies one or more labels for the Pod template.</p>
+        </td>
+    </tr>
     <tr>
         <td colspan="2">--priority string</td>
     </tr>

docs/commands/kjobctl_create/kjobctl_create_rayjob.md

@@ -113,6 +113,24 @@ kjobctl create rayjob --profile APPLICATION_PROFILE_NAME [--localqueue LOCAL_QUE
             <p>Output format. One of: (json, yaml, name, go-template, go-template-file, template, templatefile, jsonpath, jsonpath-as-json, jsonpath-file).</p>
         </td>
     </tr>
+    <tr>
+        <td colspan="2">--pod-template-annotation &lt;comma-separated &#39;key=value&#39; pairs&gt;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;Default: []</td>
+    </tr>
+    <tr>
+        <td></td>
+        <td style="line-height: 130%; word-wrap: break-word;">
+            <p>Specifies one or more annotations for the Pod template.</p>
+        </td>
+    </tr>
+    <tr>
+        <td colspan="2">--pod-template-label &lt;comma-separated &#39;key=value&#39; pairs&gt;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;Default: []</td>
+    </tr>
+    <tr>
+        <td></td>
+        <td style="line-height: 130%; word-wrap: break-word;">
+            <p>Specifies one or more labels for the Pod template.</p>
+        </td>
+    </tr>
     <tr>
         <td colspan="2">--priority string</td>
     </tr>

docs/commands/kjobctl_create/kjobctl_create_slurm.md

@@ -115,6 +115,24 @@ kjobctl create slurm --profile APPLICATION_PROFILE_NAME [--localqueue LOCAL_QUEU
             <p>Output format. One of: (json, yaml, name, go-template, go-template-file, template, templatefile, jsonpath, jsonpath-as-json, jsonpath-file).</p>
         </td>
     </tr>
+    <tr>
+        <td colspan="2">--pod-template-annotation &lt;comma-separated &#39;key=value&#39; pairs&gt;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;Default: []</td>
+    </tr>
+    <tr>
+        <td></td>
+        <td style="line-height: 130%; word-wrap: break-word;">
+            <p>Specifies one or more annotations for the Pod template.</p>
+        </td>
+    </tr>
+    <tr>
+        <td colspan="2">--pod-template-label &lt;comma-separated &#39;key=value&#39; pairs&gt;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;Default: []</td>
+    </tr>
+    <tr>
+        <td></td>
+        <td style="line-height: 130%; word-wrap: break-word;">
+            <p>Specifies one or more labels for the Pod template.</p>
+        </td>
+    </tr>
     <tr>
         <td colspan="2">--priority string</td>
     </tr>

pkg/builder/builder.go

@@ -73,6 +73,8 @@ var (
 	noPartitionSpecifiedErr                = errors.New("no partition specified")
 	noPrioritySpecifiedErr                 = errors.New("no priority specified")
 	noTimeSpecifiedErr                     = errors.New("no time specified")
+	noPodTemplateLabelSpecifiedErr         = errors.New("no pod template label specified")
+	noPodTemplateAnnotationSpecifiedErr    = errors.New("no pod template annotation specified")
 )
 
 type builder interface {
@@ -122,6 +124,8 @@ type Builder struct {
 	firstNodeIPTimeout       time.Duration
 	changeDir                string
 	timeLimit                string
+	podTemplateLabels        map[string]string
+	podTemplateAnnotations   map[string]string
 
 	profile       *v1alpha1.ApplicationProfile
 	mode          *v1alpha1.SupportedMode
@@ -314,6 +318,16 @@ func (b *Builder) WithTimeLimit(timeLimit string) *Builder {
 	return b
 }
 
+func (b *Builder) WithPodTemplateLabels(podTemplateLabels map[string]string) *Builder {
+	b.podTemplateLabels = podTemplateLabels
+	return b
+}
+
+func (b *Builder) WithPodTemplateAnnotations(podTemplateAnnotations map[string]string) *Builder {
+	b.podTemplateAnnotations = podTemplateAnnotations
+	return b
+}
+
 func (b *Builder) validateGeneral(ctx context.Context) error {
 	if b.namespace == "" {
 		return noNamespaceSpecifiedErr
@@ -472,6 +486,14 @@ func (b *Builder) validateFlags() error {
 		return noTimeSpecifiedErr
 	}
 
+	if slices.Contains(b.mode.RequiredFlags, v1alpha1.PodTemplateLabelFlag) && b.podTemplateLabels == nil {
+		return noPodTemplateLabelSpecifiedErr
+	}
+
+	if slices.Contains(b.mode.RequiredFlags, v1alpha1.PodTemplateAnnotationFlag) && b.podTemplateAnnotations == nil {
+		return noPodTemplateAnnotationSpecifiedErr
+	}
+
 	return nil
 }
 
@@ -606,11 +628,22 @@ func (b *Builder) withKueueLabels(objectMeta *metav1.ObjectMeta) error {
 	return nil
 }
 
-func (b *Builder) buildPodSpec(templateSpec corev1.PodSpec) corev1.PodSpec {
-	b.buildPodSpecVolumesAndEnv(&templateSpec)
+// buildPodObjectMeta sets user specified pod template labels and annotations
+func (b *Builder) buildPodObjectMeta(templateObjectMeta *metav1.ObjectMeta) {
+	templateObjectMeta.Labels = b.podTemplateLabels
+	templateObjectMeta.Annotations = b.podTemplateAnnotations
+}
 
-	for i := range templateSpec.Containers {
-		container := &templateSpec.Containers[i]
+func (b *Builder) buildPodTemplateSpec(podTemplateSpec *corev1.PodTemplateSpec) {
+	b.buildPodObjectMeta(&podTemplateSpec.ObjectMeta)
+	b.buildPodSpec(&podTemplateSpec.Spec)
+}
+
+func (b *Builder) buildPodSpec(podSpec *corev1.PodSpec) {
+	b.buildPodSpecVolumesAndEnv(podSpec)
+
+	for i := range podSpec.Containers {
+		container := &podSpec.Containers[i]
 
 		if i == 0 && len(b.command) > 0 {
 			container.Command = b.command
@@ -620,8 +653,6 @@ func (b *Builder) buildPodSpec(templateSpec corev1.PodSpec) corev1.PodSpec {
 			container.Resources.Requests = b.requests
 		}
 	}
-
-	return templateSpec
 }
 
 func (b *Builder) buildPodSpecVolumesAndEnv(templateSpec *corev1.PodSpec) {

pkg/builder/builder_test.go

@@ -413,14 +413,40 @@ func TestBuilder(t *testing.T) {
 			},
 			wantErr: noPrioritySpecifiedErr,
 		},
+		"shouldn't build job because pod template label not specified with required flags": {
+			namespace: metav1.NamespaceDefault,
+			profile:   "profile",
+			mode:      v1alpha1.JobMode,
+			kjobctlObjs: []runtime.Object{
+				wrappers.MakeApplicationProfile("profile", metav1.NamespaceDefault).
+					WithSupportedMode(v1alpha1.SupportedMode{
+						Name:          v1alpha1.JobMode,
+						RequiredFlags: []v1alpha1.Flag{v1alpha1.PodTemplateLabelFlag},
+					}).
+					Obj(),
+			},
+			wantErr: noPodTemplateLabelSpecifiedErr,
+		},
+		"shouldn't build job because pod template annotation not specified with required flags": {
+			namespace: metav1.NamespaceDefault,
+			profile:   "profile",
+			mode:      v1alpha1.JobMode,
+			kjobctlObjs: []runtime.Object{
+				wrappers.MakeApplicationProfile("profile", metav1.NamespaceDefault).
+					WithSupportedMode(v1alpha1.SupportedMode{
+						Name:          v1alpha1.JobMode,
+						RequiredFlags: []v1alpha1.Flag{v1alpha1.PodTemplateAnnotationFlag},
+					}).
+					Obj(),
+			},
+			wantErr: noPodTemplateAnnotationSpecifiedErr,
+		},
 		"should build job": {
 			namespace: metav1.NamespaceDefault,
 			profile:   "profile",
 			mode:      v1alpha1.JobMode,
 			kjobctlObjs: []runtime.Object{
 				wrappers.MakeJobTemplate("job-template", metav1.NamespaceDefault).
-					Label("foo", "bar").
-					Annotation("foo", "baz").
 					Obj(),
 				wrappers.MakeApplicationProfile("profile", metav1.NamespaceDefault).
 					WithSupportedMode(v1alpha1.SupportedMode{
@@ -430,8 +456,6 @@ func TestBuilder(t *testing.T) {
 					Obj(),
 			},
 			wantRootObj: wrappers.MakeJob("", metav1.NamespaceDefault).GenerateName("profile-job-").
-				Annotation("foo", "baz").
-				Label("foo", "bar").
 				Profile("profile").
 				Mode(v1alpha1.JobMode).
 				Obj(),

pkg/builder/interactive_builder.go

@@ -36,6 +36,8 @@ func (b *interactiveBuilder) build(ctx context.Context) (runtime.Object, []runti
 		return nil, nil, err
 	}
 
+	b.buildPodObjectMeta(&template.Template.ObjectMeta)
+
 	objectMeta, err := b.buildObjectMeta(template.Template.ObjectMeta, false)
 	if err != nil {
 		return nil, nil, err
@@ -50,7 +52,7 @@ func (b *interactiveBuilder) build(ctx context.Context) (runtime.Object, []runti
 		Spec:       template.Template.Spec,
 	}
 
-	pod.Spec = b.buildPodSpec(pod.Spec)
+	b.buildPodSpec(&pod.Spec)
 
 	if len(pod.Spec.Containers) > 0 {
 		pod.Spec.Containers[0].TTY = true