Use annotation rather than label for serving workloads (#3815)

kubernetes-sigs · Dec 11, 2024 · e544dc8 · e544dc8
1 parent 9770707
commit e544dc8
Show file tree

Hide file tree

Showing 7 changed files with 38 additions and 18 deletions.
diff --git a/pkg/controller/jobs/deployment/deployment_webhook.go b/pkg/controller/jobs/deployment/deployment_webhook.go
@@ -74,10 +74,13 @@ func (wh *Webhook) Default(ctx context.Context, obj runtime.Object) error {
 		return err
 	}
 	if suspend {
+		if deployment.Spec.Template.Annotations == nil {
+			deployment.Spec.Template.Annotations = make(map[string]string, 1)
+		}
+		deployment.Spec.Template.Annotations[pod.SuspendedByParentAnnotation] = FrameworkName
 		if deployment.Spec.Template.Labels == nil {
-			deployment.Spec.Template.Labels = make(map[string]string, 2)
+			deployment.Spec.Template.Labels = make(map[string]string, 1)
 		}
-		deployment.Spec.Template.Labels[pod.SuspendedByParentLabelKey] = FrameworkName
 		queueName := jobframework.QueueNameForObject(deployment.Object())
 		if queueName != "" {
 			deployment.Spec.Template.Labels[constants.QueueLabel] = queueName

diff --git a/pkg/controller/jobs/deployment/deployment_webhook_test.go b/pkg/controller/jobs/deployment/deployment_webhook_test.go
@@ -52,7 +52,7 @@ func TestDefault(t *testing.T) {
 			want: testingdeployment.MakeDeployment("test-pod", "").
 				Queue("test-queue").
 				PodTemplateSpecQueue("test-queue").
-				PodTemplateSpecLabel(pod.SuspendedByParentLabelKey, FrameworkName).
+				PodTemplateAnnotation(pod.SuspendedByParentAnnotation, FrameworkName).
 				Obj(),
 		},
 		"deployment with queue and pod template spec queue": {
@@ -63,7 +63,7 @@ func TestDefault(t *testing.T) {
 			want: testingdeployment.MakeDeployment("test-pod", "").
 				Queue("new-test-queue").
 				PodTemplateSpecQueue("new-test-queue").
-				PodTemplateSpecLabel(pod.SuspendedByParentLabelKey, FrameworkName).
+				PodTemplateAnnotation(pod.SuspendedByParentAnnotation, FrameworkName).
 				Obj(),
 		},
 		"deployment without queue with pod template spec queue": {
@@ -77,7 +77,7 @@ func TestDefault(t *testing.T) {
 			want: testingdeployment.MakeDeployment("test-pod", "default").
 				Queue("default").
 				PodTemplateSpecQueue("default").
-				PodTemplateSpecLabel(pod.SuspendedByParentLabelKey, FrameworkName).
+				PodTemplateAnnotation(pod.SuspendedByParentAnnotation, FrameworkName).
 				Obj(),
 		},
 		"LocalQueueDefaulting enabled, default lq is created, job has queue label": {
@@ -87,7 +87,7 @@ func TestDefault(t *testing.T) {
 			want: testingdeployment.MakeDeployment("test-pod", "").
 				Queue("test-queue").
 				PodTemplateSpecQueue("test-queue").
-				PodTemplateSpecLabel(pod.SuspendedByParentLabelKey, FrameworkName).
+				PodTemplateAnnotation(pod.SuspendedByParentAnnotation, FrameworkName).
 				Obj(),
 		},
 		"LocalQueueDefaulting enabled, default lq isn't created, job doesn't have queue label": {

diff --git a/pkg/controller/jobs/pod/pod_webhook.go b/pkg/controller/jobs/pod/pod_webhook.go
@@ -49,11 +49,11 @@ const (
 	ManagedLabelKey              = constants.ManagedByKueueLabel
 	ManagedLabelValue            = "true"
 	PodFinalizer                 = ManagedLabelKey
-	SuspendedByParentLabelKey    = "kueue.x-k8s.io/pod-suspending-parent"
 	GroupNameLabel               = "kueue.x-k8s.io/pod-group-name"
 	GroupTotalCountAnnotation    = "kueue.x-k8s.io/pod-group-total-count"
 	GroupFastAdmissionAnnotation = "kueue.x-k8s.io/pod-group-fast-admission"
 	GroupServingAnnotation       = "kueue.x-k8s.io/pod-group-serving"
+	SuspendedByParentAnnotation  = "kueue.x-k8s.io/pod-suspending-parent"
 	RoleHashAnnotation           = "kueue.x-k8s.io/role-hash"
 	RetriableInGroupAnnotation   = "kueue.x-k8s.io/retriable-in-group"
 )
@@ -152,7 +152,7 @@ func (w *PodWebhook) Default(ctx context.Context, obj runtime.Object) error {
 	log := ctrl.LoggerFrom(ctx).WithName("pod-webhook")
 	log.V(5).Info("Applying defaults")
 
-	_, suspend := pod.pod.GetLabels()[SuspendedByParentLabelKey]
+	_, suspend := pod.pod.GetAnnotations()[SuspendedByParentAnnotation]
 	if !suspend {
 		// Namespace filtering
 		ns := corev1.Namespace{}

diff --git a/pkg/controller/jobs/statefulset/statefulset_webhook.go b/pkg/controller/jobs/statefulset/statefulset_webhook.go
@@ -75,18 +75,17 @@ func (wh *Webhook) Default(ctx context.Context, obj runtime.Object) error {
 		return err
 	}
 	if suspend {
+		if ss.Spec.Template.Annotations == nil {
+			ss.Spec.Template.Annotations = make(map[string]string, 1)
+		}
+		ss.Spec.Template.Annotations[pod.SuspendedByParentAnnotation] = FrameworkName
 		if ss.Spec.Template.Labels == nil {
-			ss.Spec.Template.Labels = make(map[string]string, 3)
+			ss.Spec.Template.Labels = make(map[string]string, 2)
 		}
-		ss.Spec.Template.Labels[pod.SuspendedByParentLabelKey] = FrameworkName
 		queueName := jobframework.QueueNameForObject(ss.Object())
 		if queueName != "" {
 			ss.Spec.Template.Labels[constants.QueueLabel] = queueName
 			ss.Spec.Template.Labels[pod.GroupNameLabel] = GetWorkloadName(ss.Name)
-
-			if ss.Spec.Template.Annotations == nil {
-				ss.Spec.Template.Annotations = make(map[string]string, 4)
-			}
 			ss.Spec.Template.Annotations[pod.GroupTotalCountAnnotation] = fmt.Sprint(ptr.Deref(ss.Spec.Replicas, 1))
 			ss.Spec.Template.Annotations[pod.GroupFastAdmissionAnnotation] = "true"
 			ss.Spec.Template.Annotations[pod.GroupServingAnnotation] = "true"

diff --git a/pkg/controller/jobs/statefulset/statefulset_webhook_test.go b/pkg/controller/jobs/statefulset/statefulset_webhook_test.go
@@ -58,7 +58,7 @@ func TestDefault(t *testing.T) {
 				Replicas(10).
 				Queue("test-queue").
 				PodTemplateSpecQueue("test-queue").
-				PodTemplateSpecLabel(pod.SuspendedByParentLabelKey, FrameworkName).
+				PodTemplateAnnotation(pod.SuspendedByParentAnnotation, FrameworkName).
 				PodTemplateSpecPodGroupNameLabel("test-pod", "", gvk).
 				PodTemplateSpecPodGroupTotalCountAnnotation(10).
 				PodTemplateSpecPodGroupFastAdmissionAnnotation(true).
@@ -76,7 +76,7 @@ func TestDefault(t *testing.T) {
 				PodTemplateSpecPodGroupNameLabel("test-pod", "", gvk).
 				PodTemplateSpecPodGroupTotalCountAnnotation(1).
 				PodTemplateSpecQueue("test-queue").
-				PodTemplateSpecLabel(pod.SuspendedByParentLabelKey, FrameworkName).
+				PodTemplateAnnotation(pod.SuspendedByParentAnnotation, FrameworkName).
 				PodTemplateSpecPodGroupFastAdmissionAnnotation(true).
 				PodTemplateSpecPodGroupServingAnnotation(true).
 				PodTemplateSpecPodGroupPodIndexLabelAnnotation(appsv1.PodIndexLabel).
@@ -89,7 +89,7 @@ func TestDefault(t *testing.T) {
 			want: testingstatefulset.MakeStatefulSet("test-pod", "default").
 				Queue("default").
 				PodTemplateSpecQueue("default").
-				PodTemplateSpecLabel(pod.SuspendedByParentLabelKey, FrameworkName).
+				PodTemplateAnnotation(pod.SuspendedByParentAnnotation, FrameworkName).
 				PodTemplateSpecPodGroupNameLabel("test-pod", "", gvk).
 				PodTemplateSpecPodGroupTotalCountAnnotation(1).
 				PodTemplateSpecPodGroupFastAdmissionAnnotation(true).
@@ -104,7 +104,7 @@ func TestDefault(t *testing.T) {
 			want: testingstatefulset.MakeStatefulSet("test-pod", "").
 				Queue("test-queue").
 				PodTemplateSpecQueue("test-queue").
-				PodTemplateSpecLabel(pod.SuspendedByParentLabelKey, FrameworkName).
+				PodTemplateAnnotation(pod.SuspendedByParentAnnotation, FrameworkName).
 				PodTemplateSpecPodGroupNameLabel("test-pod", "", gvk).
 				PodTemplateSpecPodGroupTotalCountAnnotation(1).
 				PodTemplateSpecPodGroupFastAdmissionAnnotation(true).

diff --git a/pkg/util/testingjobs/deployment/wrappers.go b/pkg/util/testingjobs/deployment/wrappers.go
@@ -128,6 +128,15 @@ func (d *DeploymentWrapper) PodTemplateSpecLabel(k, v string) *DeploymentWrapper
 	return d
 }
 
+// PodTemplateAnnotation sets the annotation of the pod template
+func (d *DeploymentWrapper) PodTemplateAnnotation(k, v string) *DeploymentWrapper {
+	if d.Spec.Template.Annotations == nil {
+		d.Spec.Template.Annotations = make(map[string]string, 1)
+	}
+	d.Spec.Template.Annotations[k] = v
+	return d
+}
+
 // PodTemplateSpecQueue updates the queue name of the pod template spec of the Deployment
 func (d *DeploymentWrapper) PodTemplateSpecQueue(q string) *DeploymentWrapper {
 	return d.PodTemplateSpecLabel(constants.QueueLabel, q)

diff --git a/pkg/util/testingjobs/statefulset/wrappers.go b/pkg/util/testingjobs/statefulset/wrappers.go
@@ -106,6 +106,15 @@ func (ss *StatefulSetWrapper) PodTemplateSpecLabel(k, v string) *StatefulSetWrap
 	return ss
 }
 
+// PodTemplateAnnotation sets the annotation of the pod template
+func (ss *StatefulSetWrapper) PodTemplateAnnotation(k, v string) *StatefulSetWrapper {
+	if ss.Spec.Template.Annotations == nil {
+		ss.Spec.Template.Annotations = make(map[string]string, 1)
+	}
+	ss.Spec.Template.Annotations[k] = v
+	return ss
+}
+
 // PodTemplateSpecAnnotation sets the annotation of the pod template spec of the StatefulSet
 func (ss *StatefulSetWrapper) PodTemplateSpecAnnotation(k, v string) *StatefulSetWrapper {
 	if ss.Spec.Template.Annotations == nil {