Use KWOK label instead of annotation

run-ai · Nov 26, 2024 · 998f2e7 · 998f2e7
1 parent 83b94a7
commit 998f2e7
Show file tree

Hide file tree

Showing 5 changed files with 11 additions and 15 deletions.
diff --git a/internal/common/constants/constants.go b/internal/common/constants/constants.go
@@ -6,14 +6,14 @@ const (
 	AnnotationPodGroupName         = "pod-group-name"
 	AnnotationReservationPodGpuIdx = "run.ai/reserve_for_gpu_index"
 	AnnotationMigMapping           = "run.ai/mig-mapping"
-	AnnotationKwokNode             = "kwok.x-k8s.io/node"
 
 	LabelGpuGroup                   = "runai-gpu-group"
 	LabelGpuProduct                 = "nvidia.com/gpu.product"
 	LabelMigConfigState             = "nvidia.com/mig.config.state"
 	LabelFakeNodeDeploymentTemplate = "run.ai/fake-node-deployment-template"
 	LabelTopologyCMNodeTopology     = "node-topology"
 	LabelTopologyCMNodeName         = "node-name"
+	LabelKwokNode                   = "kwok.x-k8s.io/node"
 
 	ReservationNs = "runai-reservation"
 

diff --git a/internal/common/topology/kubernetes.go b/internal/common/topology/kubernetes.go
@@ -31,11 +31,11 @@ func CreateNodeTopologyCM(kubeclient kubernetes.Interface, nodeTopology *NodeTop
 	if err != nil {
 		return err
 	}
-	if value, found := node.Annotations[constants.AnnotationKwokNode]; found {
-		if cm.Annotations == nil {
-			cm.Annotations = make(map[string]string)
+	if value, found := node.Labels[constants.LabelKwokNode]; found {
+		if cm.Labels == nil {
+			cm.Labels = make(map[string]string)
 		}
-		cm.Annotations[constants.AnnotationKwokNode] = value
+		cm.Labels[constants.LabelKwokNode] = value
 	}
 
 	_, err = kubeclient.CoreV1().ConfigMaps(

diff --git a/internal/kwok-gpu-device-plugin/app_test.go b/internal/kwok-gpu-device-plugin/app_test.go
@@ -104,10 +104,8 @@ var _ = Describe("KwokGpuDevicePlugin", func() {
 					ObjectMeta: metav1.ObjectMeta{
 						Name: "node1",
 						Labels: map[string]string{
-							nodePoolLabelKey: defaultNodePoolName,
-						},
-						Annotations: map[string]string{
-							constants.AnnotationKwokNode: "fake",
+							nodePoolLabelKey:        defaultNodePoolName,
+							constants.LabelKwokNode: "fake",
 						},
 					},
 				}
@@ -127,9 +125,7 @@ var _ = Describe("KwokGpuDevicePlugin", func() {
 				cm, _, err := topology.ToNodeTopologyCM(&nodeTopology, node1.Name)
 				Expect(err).ToNot(HaveOccurred())
 				cm.Namespace = gpuOperatorNamespace
-				cm.Annotations = map[string]string{
-					constants.AnnotationKwokNode: "fake",
-				}
+				cm.Labels[constants.LabelKwokNode] = "fake"
 
 				_, err = kubeClient.CoreV1().ConfigMaps(gpuOperatorNamespace).Create(context.TODO(), cm, metav1.CreateOptions{})
 				Expect(err).ToNot(HaveOccurred())

diff --git a/internal/kwok-gpu-device-plugin/controllers/configmap/controller.go b/internal/kwok-gpu-device-plugin/controllers/configmap/controller.go
@@ -74,13 +74,13 @@ func (c *ConfigMapController) Run(stopCh <-chan struct{}) {
 }
 
 func (c *ConfigMapController) isFakeGpuKWOKNodeConfigMap(cm *v1.ConfigMap) bool {
-	if cm == nil || cm.Labels == nil || cm.Annotations == nil {
+	if cm == nil || cm.Labels == nil {
 		return false
 	}
 	_, foundNodeName := cm.Labels[constants.LabelTopologyCMNodeName]
 	if !foundNodeName {
 		return false
 	}
 
-	return cm.Annotations[constants.AnnotationKwokNode] == "fake"
+	return cm.Labels[constants.LabelKwokNode] == "fake"
 }
diff --git a/internal/status-updater/handlers/node/fake_node_deployments.go b/internal/status-updater/handlers/node/fake_node_deployments.go
@@ -121,5 +121,5 @@ func generateFakeNodeDeploymentFromTemplate(template *appsv1.Deployment, node *v
 }
 
 func isFakeNode(node *v1.Node) bool {
-	return node != nil && node.Annotations[constants.AnnotationKwokNode] == "fake"
+	return node != nil && node.Annotations[constants.LabelKwokNode] == "fake"
 }