diff --git a/apis/v1alpha1/application_profile_types.go b/apis/v1alpha1/application_profile_types.go
index d8397e9..7a65dce 100644
--- a/apis/v1alpha1/application_profile_types.go
+++ b/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
diff --git a/config/crd/bases/kjobctl.x-k8s.io_applicationprofiles.yaml b/config/crd/bases/kjobctl.x-k8s.io_applicationprofiles.yaml
index 4fc8ae2..363e3a5 100644
--- a/config/crd/bases/kjobctl.x-k8s.io_applicationprofiles.yaml
+++ b/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
diff --git a/docs/commands/kjobctl_create/kjobctl_create_interactive.md b/docs/commands/kjobctl_create/kjobctl_create_interactive.md
index a4b3c58..28835eb 100644
--- a/docs/commands/kjobctl_create/kjobctl_create_interactive.md
+++ b/docs/commands/kjobctl_create/kjobctl_create_interactive.md
@@ -99,6 +99,24 @@ kjobctl create interactive --profile APPLICATION_PROFILE_NAME [--localqueue LOCA
The length of time (like 5s, 2m, or 3h, higher than zero) to wait until at least one pod is running.
+
+ | --pod-template-annotation <comma-separated 'key=value' pairs> Default: [] |
+
+
+ |
+
+ Specifies one or more annotations for the Pod template.
+ |
+
+
+ | --pod-template-label <comma-separated 'key=value' pairs> Default: [] |
+
+
+ |
+
+ Specifies one or more labels for the Pod template.
+ |
+
| --priority string |
diff --git a/docs/commands/kjobctl_create/kjobctl_create_job.md b/docs/commands/kjobctl_create/kjobctl_create_job.md
index b55a153..da95430 100644
--- a/docs/commands/kjobctl_create/kjobctl_create_job.md
+++ b/docs/commands/kjobctl_create/kjobctl_create_job.md
@@ -111,6 +111,24 @@ kjobctl create job --profile APPLICATION_PROFILE_NAME [--localqueue LOCAL_QUEUE_
Parallelism specifies the maximum desired number of pods the job should run at any given time.
+
+ | --pod-template-annotation <comma-separated 'key=value' pairs> Default: [] |
+
+
+ |
+
+ Specifies one or more annotations for the Pod template.
+ |
+
+
+ | --pod-template-label <comma-separated 'key=value' pairs> Default: [] |
+
+
+ |
+
+ Specifies one or more labels for the Pod template.
+ |
+
| --priority string |
diff --git a/docs/commands/kjobctl_create/kjobctl_create_raycluster.md b/docs/commands/kjobctl_create/kjobctl_create_raycluster.md
index 8015b5d..3069044 100644
--- a/docs/commands/kjobctl_create/kjobctl_create_raycluster.md
+++ b/docs/commands/kjobctl_create/kjobctl_create_raycluster.md
@@ -103,6 +103,24 @@ kjobctl create raycluster --profile APPLICATION_PROFILE_NAME [--localqueue LOCAL
Output format. One of: (json, yaml, name, go-template, go-template-file, template, templatefile, jsonpath, jsonpath-as-json, jsonpath-file).
+
+ | --pod-template-annotation <comma-separated 'key=value' pairs> Default: [] |
+
+
+ |
+
+ Specifies one or more annotations for the Pod template.
+ |
+
+
+ | --pod-template-label <comma-separated 'key=value' pairs> Default: [] |
+
+
+ |
+
+ Specifies one or more labels for the Pod template.
+ |
+
| --priority string |
diff --git a/docs/commands/kjobctl_create/kjobctl_create_rayjob.md b/docs/commands/kjobctl_create/kjobctl_create_rayjob.md
index a192549..634ce49 100644
--- a/docs/commands/kjobctl_create/kjobctl_create_rayjob.md
+++ b/docs/commands/kjobctl_create/kjobctl_create_rayjob.md
@@ -113,6 +113,24 @@ kjobctl create rayjob --profile APPLICATION_PROFILE_NAME [--localqueue LOCAL_QUE
Output format. One of: (json, yaml, name, go-template, go-template-file, template, templatefile, jsonpath, jsonpath-as-json, jsonpath-file).
+
+ | --pod-template-annotation <comma-separated 'key=value' pairs> Default: [] |
+
+
+ |
+
+ Specifies one or more annotations for the Pod template.
+ |
+
+
+ | --pod-template-label <comma-separated 'key=value' pairs> Default: [] |
+
+
+ |
+
+ Specifies one or more labels for the Pod template.
+ |
+
| --priority string |
diff --git a/docs/commands/kjobctl_create/kjobctl_create_slurm.md b/docs/commands/kjobctl_create/kjobctl_create_slurm.md
index eb1733f..7d137b6 100644
--- a/docs/commands/kjobctl_create/kjobctl_create_slurm.md
+++ b/docs/commands/kjobctl_create/kjobctl_create_slurm.md
@@ -115,6 +115,24 @@ kjobctl create slurm --profile APPLICATION_PROFILE_NAME [--localqueue LOCAL_QUEU
Output format. One of: (json, yaml, name, go-template, go-template-file, template, templatefile, jsonpath, jsonpath-as-json, jsonpath-file).
+
+ | --pod-template-annotation <comma-separated 'key=value' pairs> Default: [] |
+
+
+ |
+
+ Specifies one or more annotations for the Pod template.
+ |
+
+
+ | --pod-template-label <comma-separated 'key=value' pairs> Default: [] |
+
+
+ |
+
+ Specifies one or more labels for the Pod template.
+ |
+
| --priority string |
diff --git a/pkg/builder/builder.go b/pkg/builder/builder.go
index 86608c3..aaf6348 100644
--- a/pkg/builder/builder.go
+++ b/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) {
diff --git a/pkg/builder/builder_test.go b/pkg/builder/builder_test.go
index f8d3f08..962e040 100644
--- a/pkg/builder/builder_test.go
+++ b/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(),
diff --git a/pkg/builder/interactive_builder.go b/pkg/builder/interactive_builder.go
index 0c8b302..ff6e3d1 100644
--- a/pkg/builder/interactive_builder.go
+++ b/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
diff --git a/pkg/builder/interactive_builder_test.go b/pkg/builder/interactive_builder_test.go
index a13e2d3..6a379ec 100644
--- a/pkg/builder/interactive_builder_test.go
+++ b/pkg/builder/interactive_builder_test.go
@@ -47,8 +47,6 @@ func TestInteractiveBuilder(t *testing.T) {
userID := os.Getenv(constants.SystemEnvVarNameUser)
testPodTemplateWrapper := wrappers.MakePodTemplate("pod-template", metav1.NamespaceDefault).
- Label("foo", "bar").
- Annotation("foo", "baz").
WithInitContainer(
*wrappers.MakeContainer("ic1", "").
WithEnvVar(corev1.EnvVar{Name: "e0", Value: "default-value0"}).
@@ -76,16 +74,18 @@ func TestInteractiveBuilder(t *testing.T) {
WithVolume("v2", "default-config2")
testCases := map[string]struct {
- namespace string
- profile string
- mode v1alpha1.ApplicationProfileMode
- command []string
- requests corev1.ResourceList
- localQueue string
- k8sObjs []runtime.Object
- kjobctlObjs []runtime.Object
- wantRootObj runtime.Object
- wantErr error
+ namespace string
+ profile string
+ mode v1alpha1.ApplicationProfileMode
+ command []string
+ requests corev1.ResourceList
+ localQueue string
+ podTemplateLabels map[string]string
+ podTemplateAnnotations map[string]string
+ k8sObjs []runtime.Object
+ kjobctlObjs []runtime.Object
+ wantRootObj runtime.Object
+ wantErr error
}{
"shouldn't build pod because template not found": {
namespace: metav1.NamespaceDefault,
@@ -108,6 +108,37 @@ func TestInteractiveBuilder(t *testing.T) {
WithSupportedMode(v1alpha1.SupportedMode{Name: v1alpha1.InteractiveMode, Template: "pod-template"}).
Obj(),
},
+ wantRootObj: wrappers.MakePod("", metav1.NamespaceDefault).GenerateName("profile-interactive-").
+ Profile("profile").
+ Mode(v1alpha1.InteractiveMode).
+ Spec(
+ testPodTemplateWrapper.Clone().
+ WithEnvVar(corev1.EnvVar{Name: constants.EnvVarNameUserID, Value: userID}).
+ WithEnvVar(corev1.EnvVar{Name: constants.EnvVarTaskName, Value: "default_profile"}).
+ WithEnvVar(corev1.EnvVar{
+ Name: constants.EnvVarTaskID,
+ Value: fmt.Sprintf("%s_%s_default_profile", userID, testStartTime.Format(time.RFC3339)),
+ }).
+ WithEnvVar(corev1.EnvVar{Name: "PROFILE", Value: "default_profile"}).
+ WithEnvVar(corev1.EnvVar{Name: "TIMESTAMP", Value: testStartTime.Format(time.RFC3339)}).
+ TTY().
+ Stdin().
+ Obj().Template.Spec,
+ ).
+ Obj(),
+ },
+ "should build job with pod template label and annotation": {
+ namespace: metav1.NamespaceDefault,
+ profile: "profile",
+ mode: v1alpha1.InteractiveMode,
+ podTemplateLabels: map[string]string{"foo": "bar"},
+ podTemplateAnnotations: map[string]string{"foo": "baz"},
+ k8sObjs: []runtime.Object{testPodTemplateWrapper.Clone().Obj()},
+ kjobctlObjs: []runtime.Object{
+ wrappers.MakeApplicationProfile("profile", metav1.NamespaceDefault).
+ WithSupportedMode(v1alpha1.SupportedMode{Name: v1alpha1.InteractiveMode, Template: "pod-template"}).
+ Obj(),
+ },
wantRootObj: wrappers.MakePod("", metav1.NamespaceDefault).GenerateName("profile-interactive-").
Annotation("foo", "baz").
Label("foo", "bar").
@@ -150,8 +181,6 @@ func TestInteractiveBuilder(t *testing.T) {
wrappers.MakeVolumeBundle("vb2", metav1.NamespaceDefault).Obj(),
},
wantRootObj: wrappers.MakePod("", metav1.NamespaceDefault).GenerateName("profile-interactive-").
- Annotation("foo", "baz").
- Label("foo", "bar").
Profile("profile").
Mode(v1alpha1.InteractiveMode).
Label(kueueconstants.QueueLabel, "lq1").
@@ -194,6 +223,8 @@ func TestInteractiveBuilder(t *testing.T) {
WithRequests(tc.requests).
WithLocalQueue(tc.localQueue).
WithSkipLocalQueueValidation(true).
+ WithPodTemplateLabels(tc.podTemplateLabels).
+ WithPodTemplateAnnotations(tc.podTemplateAnnotations).
Do(ctx)
var opts []cmp.Option
diff --git a/pkg/builder/job_builder.go b/pkg/builder/job_builder.go
index 4bc543b..5661249 100644
--- a/pkg/builder/job_builder.go
+++ b/pkg/builder/job_builder.go
@@ -51,7 +51,7 @@ func (b *jobBuilder) build(ctx context.Context) (runtime.Object, []runtime.Objec
Spec: template.Template.Spec,
}
- job.Spec.Template.Spec = b.buildPodSpec(job.Spec.Template.Spec)
+ b.buildPodTemplateSpec(&job.Spec.Template)
if b.parallelism != nil {
job.Spec.Parallelism = b.parallelism
diff --git a/pkg/builder/job_builder_test.go b/pkg/builder/job_builder_test.go
index 776e673..3ed7145 100644
--- a/pkg/builder/job_builder_test.go
+++ b/pkg/builder/job_builder_test.go
@@ -47,8 +47,6 @@ func TestJobBuilder(t *testing.T) {
userID := os.Getenv(constants.SystemEnvVarNameUser)
testJobTemplateWrapper := wrappers.MakeJobTemplate("job-template", metav1.NamespaceDefault).
- Label("foo", "bar").
- Annotation("foo", "baz").
Parallelism(1).
Completions(1).
WithInitContainer(
@@ -78,17 +76,19 @@ func TestJobBuilder(t *testing.T) {
WithVolume("v2", "default-config2")
testCases := map[string]struct {
- namespace string
- profile string
- mode v1alpha1.ApplicationProfileMode
- command []string
- parallelism *int32
- completions *int32
- requests corev1.ResourceList
- localQueue string
- kjobctlObjs []runtime.Object
- wantRootObj runtime.Object
- wantErr error
+ namespace string
+ profile string
+ mode v1alpha1.ApplicationProfileMode
+ command []string
+ parallelism *int32
+ completions *int32
+ requests corev1.ResourceList
+ localQueue string
+ podTemplateLabels map[string]string
+ podTemplateAnnotation map[string]string
+ kjobctlObjs []runtime.Object
+ wantRootObj runtime.Object
+ wantErr error
}{
"shouldn't build job because template not found": {
namespace: metav1.NamespaceDefault,
@@ -112,8 +112,6 @@ func TestJobBuilder(t *testing.T) {
Obj(),
},
wantRootObj: wrappers.MakeJob("", metav1.NamespaceDefault).GenerateName("profile-job-").
- Annotation("foo", "baz").
- Label("foo", "bar").
Profile("profile").
Mode(v1alpha1.JobMode).
Spec(
@@ -130,6 +128,37 @@ func TestJobBuilder(t *testing.T) {
).
Obj(),
},
+ "should build job with pod template label and annotation": {
+ namespace: metav1.NamespaceDefault,
+ profile: "profile",
+ mode: v1alpha1.JobMode,
+ podTemplateLabels: map[string]string{"foo": "bar"},
+ podTemplateAnnotation: map[string]string{"foo": "baz"},
+ kjobctlObjs: []runtime.Object{
+ testJobTemplateWrapper.Clone().Obj(),
+ wrappers.MakeApplicationProfile("profile", metav1.NamespaceDefault).
+ WithSupportedMode(v1alpha1.SupportedMode{Name: v1alpha1.JobMode, Template: "job-template"}).
+ Obj(),
+ },
+ wantRootObj: wrappers.MakeJob("", metav1.NamespaceDefault).GenerateName("profile-job-").
+ Profile("profile").
+ Mode(v1alpha1.JobMode).
+ Spec(
+ testJobTemplateWrapper.Clone().
+ WithEnvVar(corev1.EnvVar{Name: constants.EnvVarNameUserID, Value: userID}).
+ WithEnvVar(corev1.EnvVar{Name: constants.EnvVarTaskName, Value: "default_profile"}).
+ WithEnvVar(corev1.EnvVar{
+ Name: constants.EnvVarTaskID,
+ Value: fmt.Sprintf("%s_%s_default_profile", userID, testStartTime.Format(time.RFC3339)),
+ }).
+ WithEnvVar(corev1.EnvVar{Name: "PROFILE", Value: "default_profile"}).
+ WithEnvVar(corev1.EnvVar{Name: "TIMESTAMP", Value: testStartTime.Format(time.RFC3339)}).
+ Obj().Template.Spec,
+ ).
+ PodTemplateAnnotation("foo", "baz").
+ PodTemplateLabel("foo", "bar").
+ Obj(),
+ },
"should build job with replacements": {
namespace: metav1.NamespaceDefault,
profile: "profile",
@@ -153,8 +182,6 @@ func TestJobBuilder(t *testing.T) {
wrappers.MakeVolumeBundle("vb2", metav1.NamespaceDefault).Obj(),
},
wantRootObj: wrappers.MakeJob("", metav1.NamespaceDefault).GenerateName("profile-job-").
- Annotation("foo", "baz").
- Label("foo", "bar").
Profile("profile").
Mode(v1alpha1.JobMode).
Label(kueueconstants.QueueLabel, "lq1").
@@ -198,6 +225,8 @@ func TestJobBuilder(t *testing.T) {
WithRequests(tc.requests).
WithLocalQueue(tc.localQueue).
WithSkipLocalQueueValidation(true).
+ WithPodTemplateLabels(tc.podTemplateLabels).
+ WithPodTemplateAnnotations(tc.podTemplateAnnotation).
Do(ctx)
var opts []cmp.Option
diff --git a/pkg/builder/ray_cluster_builder_test.go b/pkg/builder/ray_cluster_builder_test.go
index d0b9aea..7163090 100644
--- a/pkg/builder/ray_cluster_builder_test.go
+++ b/pkg/builder/ray_cluster_builder_test.go
@@ -46,8 +46,6 @@ func TestRayClusterBuilder(t *testing.T) {
userID := os.Getenv(constants.SystemEnvVarNameUser)
testRayClusterTemplateWrapper := wrappers.MakeRayClusterTemplate("ray-cluster-template", metav1.NamespaceDefault).
- Label("foo", "bar").
- Annotation("foo", "baz").
Spec(*wrappers.MakeRayClusterSpec().
WithWorkerGroupSpec(
*wrappers.MakeWorkerGroupSpec("g1").
@@ -146,8 +144,6 @@ func TestRayClusterBuilder(t *testing.T) {
Obj(),
},
wantRootObj: wrappers.MakeRayCluster("", metav1.NamespaceDefault).GenerateName("profile-raycluster-").
- Annotation("foo", "baz").
- Label("foo", "bar").
Profile("profile").
Mode(v1alpha1.RayClusterMode).
Spec(
@@ -192,8 +188,6 @@ func TestRayClusterBuilder(t *testing.T) {
wrappers.MakeVolumeBundle("vb2", metav1.NamespaceDefault).Obj(),
},
wantRootObj: wrappers.MakeRayCluster("", metav1.NamespaceDefault).GenerateName("profile-raycluster-").
- Annotation("foo", "baz").
- Label("foo", "bar").
Profile("profile").
Mode(v1alpha1.RayClusterMode).
Label(kueueconstants.QueueLabel, "lq1").
diff --git a/pkg/builder/ray_job_builder_test.go b/pkg/builder/ray_job_builder_test.go
index 0811bba..584a1cb 100644
--- a/pkg/builder/ray_job_builder_test.go
+++ b/pkg/builder/ray_job_builder_test.go
@@ -46,8 +46,6 @@ func TestRayJobBuilder(t *testing.T) {
userID := os.Getenv(constants.SystemEnvVarNameUser)
testRayJobTemplateWrapper := wrappers.MakeRayJobTemplate("ray-job-template", metav1.NamespaceDefault).
- Label("foo", "bar").
- Annotation("foo", "baz").
WithRayClusterSpec(wrappers.MakeRayClusterSpec().
WithWorkerGroupSpec(
*wrappers.MakeWorkerGroupSpec("g1").
@@ -147,8 +145,6 @@ func TestRayJobBuilder(t *testing.T) {
Obj(),
},
wantRootObj: wrappers.MakeRayJob("", metav1.NamespaceDefault).GenerateName("profile-rayjob-").
- Annotation("foo", "baz").
- Label("foo", "bar").
Profile("profile").
Mode(v1alpha1.RayJobMode).
Spec(
@@ -193,8 +189,6 @@ func TestRayJobBuilder(t *testing.T) {
wrappers.MakeVolumeBundle("vb2", metav1.NamespaceDefault).Obj(),
},
wantRootObj: wrappers.MakeRayJob("", metav1.NamespaceDefault).GenerateName("profile-rayjob-").
- Annotation("foo", "baz").
- Label("foo", "bar").
Profile("profile").
Mode(v1alpha1.RayJobMode).
Label(kueueconstants.QueueLabel, "lq1").
@@ -248,8 +242,6 @@ func TestRayJobBuilder(t *testing.T) {
wrappers.MakeVolumeBundle("vb2", metav1.NamespaceDefault).Obj(),
},
wantRootObj: wrappers.MakeRayJob("", metav1.NamespaceDefault).GenerateName("profile-rayjob-").
- Annotation("foo", "baz").
- Label("foo", "bar").
Profile("profile").
Mode(v1alpha1.RayJobMode).
WithRayClusterLabelSelector("rc1").
diff --git a/pkg/builder/slurm_builder.go b/pkg/builder/slurm_builder.go
index 2284ab8..a19c0e1 100644
--- a/pkg/builder/slurm_builder.go
+++ b/pkg/builder/slurm_builder.go
@@ -224,6 +224,7 @@ func (b *slurmBuilder) build(ctx context.Context) (runtime.Object, []runtime.Obj
job.Spec.CompletionMode = ptr.To(batchv1.IndexedCompletion)
job.Spec.Template.Spec.Subdomain = job.Name
+ b.buildPodObjectMeta(&job.Spec.Template.ObjectMeta)
b.buildPodSpecVolumesAndEnv(&job.Spec.Template.Spec)
job.Spec.Template.Spec.Volumes = append(job.Spec.Template.Spec.Volumes,
corev1.Volume{
diff --git a/pkg/cmd/create/create.go b/pkg/cmd/create/create.go
index 0fb43a6..ca87859 100644
--- a/pkg/cmd/create/create.go
+++ b/pkg/cmd/create/create.go
@@ -73,31 +73,33 @@ const (
firstNodeIPTimeoutFlagName = "first-node-ip-timeout"
waitFlagName = "wait"
- commandFlagName = string(v1alpha1.CmdFlag)
- parallelismFlagName = string(v1alpha1.ParallelismFlag)
- completionsFlagName = string(v1alpha1.CompletionsFlag)
- replicasFlagName = string(v1alpha1.ReplicasFlag)
- minReplicasFlagName = string(v1alpha1.MinReplicasFlag)
- maxReplicasFlagName = string(v1alpha1.MaxReplicasFlag)
- requestFlagName = string(v1alpha1.RequestFlag)
- localQueueFlagName = string(v1alpha1.LocalQueueFlag)
- rayClusterFlagName = string(v1alpha1.RayClusterFlag)
- arrayFlagName = string(v1alpha1.ArrayFlag)
- cpusPerTaskFlagName = string(v1alpha1.CpusPerTaskFlag)
- gpusPerTaskFlagName = string(v1alpha1.GpusPerTaskFlag)
- memPerNodeFlagName = string(v1alpha1.MemPerNodeFlag)
- memPerTaskFlagName = string(v1alpha1.MemPerTaskFlag)
- memPerCPUFlagName = string(v1alpha1.MemPerCPUFlag)
- memPerGPUFlagName = string(v1alpha1.MemPerGPUFlag)
- nodesFlagName = string(v1alpha1.NodesFlag)
- nTasksFlagName = string(v1alpha1.NTasksFlag)
- outputFlagName = string(v1alpha1.OutputFlag)
- errorFlagName = string(v1alpha1.ErrorFlag)
- inputFlagName = string(v1alpha1.InputFlag)
- jobNameFlagName = string(v1alpha1.JobNameFlag)
- partitionFlagName = string(v1alpha1.PartitionFlag)
- priorityFlagName = string(v1alpha1.PriorityFlag)
- timeFlagName = string(v1alpha1.TimeFlag)
+ commandFlagName = string(v1alpha1.CmdFlag)
+ parallelismFlagName = string(v1alpha1.ParallelismFlag)
+ completionsFlagName = string(v1alpha1.CompletionsFlag)
+ replicasFlagName = string(v1alpha1.ReplicasFlag)
+ minReplicasFlagName = string(v1alpha1.MinReplicasFlag)
+ maxReplicasFlagName = string(v1alpha1.MaxReplicasFlag)
+ requestFlagName = string(v1alpha1.RequestFlag)
+ localQueueFlagName = string(v1alpha1.LocalQueueFlag)
+ rayClusterFlagName = string(v1alpha1.RayClusterFlag)
+ arrayFlagName = string(v1alpha1.ArrayFlag)
+ cpusPerTaskFlagName = string(v1alpha1.CpusPerTaskFlag)
+ gpusPerTaskFlagName = string(v1alpha1.GpusPerTaskFlag)
+ memPerNodeFlagName = string(v1alpha1.MemPerNodeFlag)
+ memPerTaskFlagName = string(v1alpha1.MemPerTaskFlag)
+ memPerCPUFlagName = string(v1alpha1.MemPerCPUFlag)
+ memPerGPUFlagName = string(v1alpha1.MemPerGPUFlag)
+ nodesFlagName = string(v1alpha1.NodesFlag)
+ nTasksFlagName = string(v1alpha1.NTasksFlag)
+ outputFlagName = string(v1alpha1.OutputFlag)
+ errorFlagName = string(v1alpha1.ErrorFlag)
+ inputFlagName = string(v1alpha1.InputFlag)
+ jobNameFlagName = string(v1alpha1.JobNameFlag)
+ partitionFlagName = string(v1alpha1.PartitionFlag)
+ priorityFlagName = string(v1alpha1.PriorityFlag)
+ timeFlagName = string(v1alpha1.TimeFlag)
+ podTemplateLabelFlagName = string(v1alpha1.PodTemplateLabelFlag)
+ podTemplateAnnotationFlagName = string(v1alpha1.PodTemplateAnnotationFlag)
)
func withTimeFlag(f *pflag.FlagSet, p *string) {
@@ -222,6 +224,8 @@ type CreateOptions struct {
Wait bool
SkipLocalQueueValidation bool
SkipPriorityValidation bool
+ PodTemplateLabels map[string]string
+ PodTemplateAnnotations map[string]string
UserSpecifiedCommand string
UserSpecifiedParallelism int32
@@ -494,6 +498,10 @@ func NewCreateCmd(clientGetter util.ClientGetter, streams genericiooptions.IOStr
"Apply priority for the entire workload.")
subcmd.Flags().BoolVar(&o.SkipPriorityValidation, skipPriorityValidationFlagName, false,
"Skip workload priority class validation. Add priority class label even if the class does not exist.")
+ subcmd.Flags().StringToStringVar(&o.PodTemplateLabels, podTemplateLabelFlagName, make(map[string]string),
+ "Specifies one or more labels for the Pod template.")
+ subcmd.Flags().StringToStringVar(&o.PodTemplateAnnotations, podTemplateAnnotationFlagName, make(map[string]string),
+ "Specifies one or more annotations for the Pod template.")
modeSubcommand.Setup(clientGetter, subcmd, o)
@@ -704,6 +712,8 @@ func (o *CreateOptions) Run(ctx context.Context, clientGetter util.ClientGetter,
WithFirstNodeIP(o.FirstNodeIP).
WithFirstNodeIPTimeout(o.FirstNodeIPTimeout).
WithTimeLimit(o.TimeLimit).
+ WithPodTemplateLabels(o.PodTemplateLabels).
+ WithPodTemplateAnnotations(o.PodTemplateAnnotations).
Do(ctx)
if err != nil {
return err
diff --git a/pkg/cmd/create/create_test.go b/pkg/cmd/create/create_test.go
index 425e61c..9d6d1d7 100644
--- a/pkg/cmd/create/create_test.go
+++ b/pkg/cmd/create/create_test.go
@@ -174,6 +174,39 @@ func TestCreateCmd(t *testing.T) {
// Fake dynamic client not generating name. That's why we have .
wantOut: "job.batch/ created\n",
},
+ "should create job with pod template label and annotation": {
+ args: func(tc *createCmdTestCase) []string {
+ return []string{
+ "job",
+ "--profile", "profile",
+ "--pod-template-label", "foo=bar",
+ "--pod-template-annotation", "foo=baz",
+ }
+ },
+ kjobctlObjs: []runtime.Object{
+ wrappers.MakeJobTemplate("job-template", metav1.NamespaceDefault).Obj(),
+ wrappers.MakeApplicationProfile("profile", metav1.NamespaceDefault).
+ WithSupportedMode(*wrappers.MakeSupportedMode(v1alpha1.JobMode, "job-template").Obj()).
+ Obj(),
+ },
+ gvks: []schema.GroupVersionKind{{Group: "batch", Version: "v1", Kind: "Job"}},
+ wantLists: []runtime.Object{
+ &batchv1.JobList{
+ TypeMeta: metav1.TypeMeta{Kind: "JobList", APIVersion: "batch/v1"},
+ Items: []batchv1.Job{
+ *wrappers.MakeJob("", metav1.NamespaceDefault).
+ GenerateName("profile-job-").
+ Profile("profile").
+ Mode(v1alpha1.JobMode).
+ PodTemplateLabel("foo", "bar").
+ PodTemplateAnnotation("foo", "baz").
+ Obj(),
+ },
+ },
+ },
+ // Fake dynamic client not generating name. That's why we have .
+ wantOut: "job.batch/ created\n",
+ },
"should create rayjob": {
args: func(tc *createCmdTestCase) []string { return []string{"rayjob", "--profile", "profile"} },
kjobctlObjs: []runtime.Object{
@@ -1092,6 +1125,8 @@ export $(cat /slurm/env/$JOB_CONTAINER_INDEX/slurm.env | xargs)
"--init-image", "bash:latest",
"--first-node-ip",
"--first-node-ip-timeout", "29s",
+ "--pod-template-label", "foo=bar",
+ "--pod-template-annotation", "foo=baz",
"--",
"--array", "0-25",
"--nodes", "2",
@@ -1140,6 +1175,8 @@ export $(cat /slurm/env/$JOB_CONTAINER_INDEX/slurm.env | xargs)
Mode(v1alpha1.SlurmMode).
LocalQueue("lq1").
Subdomain("profile-slurm").
+ PodTemplateLabel("foo", "bar").
+ PodTemplateAnnotation("foo", "baz").
WithInitContainer(*wrappers.MakeContainer("slurm-init-env", "bash:latest").
Command("sh", "/slurm/scripts/init-entrypoint.sh").
WithVolumeMount(corev1.VolumeMount{Name: "slurm-scripts", MountPath: "/slurm/scripts"}).
diff --git a/pkg/cmd/printcrds/embed/manifest.gz b/pkg/cmd/printcrds/embed/manifest.gz
index 9d81167..9508039 100644
Binary files a/pkg/cmd/printcrds/embed/manifest.gz and b/pkg/cmd/printcrds/embed/manifest.gz differ
diff --git a/pkg/testing/wrappers/job_wrappers.go b/pkg/testing/wrappers/job_wrappers.go
index a7bf1fc..34ff132 100644
--- a/pkg/testing/wrappers/job_wrappers.go
+++ b/pkg/testing/wrappers/job_wrappers.go
@@ -190,8 +190,26 @@ func (j *JobWrapper) Succeeded(value int32) *JobWrapper {
return j
}
-// Spec set job spec.
+// Spec sets job spec.
func (j *JobWrapper) Spec(spec batchv1.JobSpec) *JobWrapper {
j.Job.Spec = spec
return j
}
+
+// PodTemplateLabel sets pod template label.
+func (j *JobWrapper) PodTemplateLabel(key, value string) *JobWrapper {
+ if j.Job.Spec.Template.Labels == nil {
+ j.Job.Spec.Template.Labels = make(map[string]string)
+ }
+ j.Job.Spec.Template.Labels[key] = value
+ return j
+}
+
+// PodTemplateAnnotation sets pod template annotation.
+func (j *JobWrapper) PodTemplateAnnotation(key, value string) *JobWrapper {
+ if j.Job.Spec.Template.Annotations == nil {
+ j.Job.Spec.Template.Annotations = make(map[string]string)
+ }
+ j.Job.Spec.Template.Annotations[key] = value
+ return j
+}