fix: propagate priority-class label for deploy and statefulset

Signed-off-by: Abirdcfly <[email protected]>
kubernetes-sigs · Jan 20, 2025 · 22b6c09 · 22b6c09
1 parent 7a32d40
commit 22b6c09
Show file tree

Hide file tree

Showing 7 changed files with 69 additions and 4 deletions.
diff --git a/pkg/controller/jobframework/interface.go b/pkg/controller/jobframework/interface.go
@@ -170,8 +170,7 @@ func MaximumExecutionTimeSeconds(job GenericJob) *int32 {
 	return ptr.To(int32(v))
 }
 
-func workloadPriorityClassName(job GenericJob) string {
-	object := job.Object()
+func WorkloadPriorityClassName(object client.Object) string {
 	if workloadPriorityClassLabel := object.GetLabels()[constants.WorkloadPriorityClassLabel]; workloadPriorityClassLabel != "" {
 		return workloadPriorityClassLabel
 	}

diff --git a/pkg/controller/jobframework/reconciler.go b/pkg/controller/jobframework/reconciler.go
@@ -988,7 +988,7 @@ func (r *JobReconciler) prepareWorkload(ctx context.Context, job GenericJob, wl
 }
 
 func (r *JobReconciler) extractPriority(ctx context.Context, podSets []kueue.PodSet, job GenericJob) (string, string, int32, error) {
-	if workloadPriorityClass := workloadPriorityClassName(job); len(workloadPriorityClass) > 0 {
+	if workloadPriorityClass := WorkloadPriorityClassName(job.Object()); len(workloadPriorityClass) > 0 {
 		return utilpriority.GetPriorityFromWorkloadPriorityClass(ctx, r.client, workloadPriorityClass)
 	}
 	if jobWithPriorityClass, isImplemented := job.(JobWithPriorityClass); isImplemented {

diff --git a/pkg/controller/jobframework/validation.go b/pkg/controller/jobframework/validation.go
@@ -124,7 +124,7 @@ func validateUpdateForQueueName(oldJob, newJob GenericJob) field.ErrorList {
 }
 
 func validateUpdateForWorkloadPriorityClassName(oldJob, newJob GenericJob) field.ErrorList {
-	allErrs := apivalidation.ValidateImmutableField(workloadPriorityClassName(newJob), workloadPriorityClassName(oldJob), workloadPriorityClassNamePath)
+	allErrs := apivalidation.ValidateImmutableField(WorkloadPriorityClassName(newJob.Object()), WorkloadPriorityClassName(oldJob.Object()), workloadPriorityClassNamePath)
 	return allErrs
 }
 

diff --git a/pkg/controller/jobs/deployment/deployment_webhook.go b/pkg/controller/jobs/deployment/deployment_webhook.go
@@ -85,6 +85,9 @@ func (wh *Webhook) Default(ctx context.Context, obj runtime.Object) error {
 		if queueName != "" {
 			deployment.Spec.Template.Labels[constants.QueueLabel] = queueName
 		}
+		if priorityClass := jobframework.WorkloadPriorityClassName(deployment.Object()); priorityClass != "" {
+			deployment.Spec.Template.Labels[constants.WorkloadPriorityClassLabel] = priorityClass
+		}
 	}
 
 	return nil

diff --git a/pkg/controller/jobs/deployment/deployment_webhook_test.go b/pkg/controller/jobs/deployment/deployment_webhook_test.go
@@ -24,6 +24,7 @@ import (
 	appsv1 "k8s.io/api/apps/v1"
 	"k8s.io/apimachinery/pkg/util/validation/field"
 	"sigs.k8s.io/controller-runtime/pkg/webhook/admission"
+	"sigs.k8s.io/kueue/pkg/controller/constants"
 
 	"sigs.k8s.io/kueue/pkg/cache"
 	"sigs.k8s.io/kueue/pkg/controller/jobframework"
@@ -97,6 +98,44 @@ func TestDefault(t *testing.T) {
 			want: testingdeployment.MakeDeployment("test-pod", "").
 				Obj(),
 		},
+		"deployment with queue and priority class": {
+			deployment: testingdeployment.MakeDeployment("test-pod", "").
+				Queue("test-queue").
+				Label(constants.WorkloadPriorityClassLabel, "test").
+				Obj(),
+			want: testingdeployment.MakeDeployment("test-pod", "").
+				Queue("test-queue").
+				Label(constants.WorkloadPriorityClassLabel, "test").
+				PodTemplateSpecQueue("test-queue").
+				PodTemplateAnnotation(pod.SuspendedByParentAnnotation, FrameworkName).
+				PodTemplateSpecLabel(constants.WorkloadPriorityClassLabel, "test").
+				Obj(),
+		},
+		"deployment with queue, priority class and pod template spec queue, priority class": {
+			deployment: testingdeployment.MakeDeployment("test-pod", "").
+				Queue("new-test-queue").
+				Label(constants.WorkloadPriorityClassLabel, "new-test").
+				PodTemplateSpecQueue("test-queue").
+				PodTemplateSpecLabel(constants.WorkloadPriorityClassLabel, "test").
+				Obj(),
+			want: testingdeployment.MakeDeployment("test-pod", "").
+				Queue("new-test-queue").
+				Label(constants.WorkloadPriorityClassLabel, "new-test").
+				PodTemplateSpecQueue("new-test-queue").
+				PodTemplateAnnotation(pod.SuspendedByParentAnnotation, FrameworkName).
+				PodTemplateSpecLabel(constants.WorkloadPriorityClassLabel, "new-test").
+				Obj(),
+		},
+		"deployment without queue with pod template spec queue and priority class": {
+			deployment: testingdeployment.MakeDeployment("test-pod", "").
+				PodTemplateSpecQueue("test-queue").
+				PodTemplateSpecLabel(constants.WorkloadPriorityClassLabel, "test").
+				Obj(),
+			want: testingdeployment.MakeDeployment("test-pod", "").
+				PodTemplateSpecQueue("test-queue").
+				PodTemplateSpecLabel(constants.WorkloadPriorityClassLabel, "test").
+				Obj(),
+		},
 	}
 
 	for name, tc := range testCases {

diff --git a/pkg/controller/jobs/statefulset/statefulset_webhook.go b/pkg/controller/jobs/statefulset/statefulset_webhook.go
@@ -91,6 +91,9 @@ func (wh *Webhook) Default(ctx context.Context, obj runtime.Object) error {
 			ss.Spec.Template.Annotations[pod.GroupServingAnnotation] = "true"
 			ss.Spec.Template.Annotations[kueuealpha.PodGroupPodIndexLabelAnnotation] = appsv1.PodIndexLabel
 		}
+		if priorityClass := jobframework.WorkloadPriorityClassName(ss.Object()); priorityClass != "" {
+			ss.Spec.Template.Labels[constants.WorkloadPriorityClassLabel] = priorityClass
+		}
 	}
 
 	return nil

diff --git a/pkg/controller/jobs/statefulset/statefulset_webhook_test.go b/pkg/controller/jobs/statefulset/statefulset_webhook_test.go
@@ -70,6 +70,27 @@ func TestDefault(t *testing.T) {
 				PodTemplateSpecPodGroupPodIndexLabelAnnotation(appsv1.PodIndexLabel).
 				Obj(),
 		},
+		"statefulset with queue and priority class": {
+			enableIntegrations: []string{"pod"},
+			statefulset: testingstatefulset.MakeStatefulSet("test-pod", "").
+				Replicas(10).
+				Queue("test-queue").
+				Label(constants.WorkloadPriorityClassLabel, "test").
+				Obj(),
+			want: testingstatefulset.MakeStatefulSet("test-pod", "").
+				Replicas(10).
+				Queue("test-queue").
+				Label(constants.WorkloadPriorityClassLabel, "test").
+				PodTemplateSpecQueue("test-queue").
+				PodTemplateAnnotation(pod.SuspendedByParentAnnotation, FrameworkName).
+				PodTemplateSpecLabel(constants.WorkloadPriorityClassLabel, "test").
+				PodTemplateSpecPodGroupNameLabel("test-pod", "", gvk).
+				PodTemplateSpecPodGroupTotalCountAnnotation(10).
+				PodTemplateSpecPodGroupFastAdmissionAnnotation(true).
+				PodTemplateSpecPodGroupServingAnnotation(true).
+				PodTemplateSpecPodGroupPodIndexLabelAnnotation(appsv1.PodIndexLabel).
+				Obj(),
+		},
 		"statefulset without replicas": {
 			enableIntegrations: []string{"pod"},
 			statefulset: testingstatefulset.MakeStatefulSet("test-pod", "").