From ee64b30d05ce0f5747b41911db6a37d2fb68e5c7 Mon Sep 17 00:00:00 2001
From: Davanum Srinivas <davanum@gmail.com>
Date: Fri, 8 Mar 2024 06:51:17 -0500
Subject: [PATCH] Check taints as well for control-plane

Signed-off-by: Davanum Srinivas <davanum@gmail.com>
---
 test/e2e/scheduling/nvidia-gpus.go | 27 ++++++++++++++++++++-------
 1 file changed, 20 insertions(+), 7 deletions(-)

diff --git a/test/e2e/scheduling/nvidia-gpus.go b/test/e2e/scheduling/nvidia-gpus.go
index f8c3121cbbc..f2e091a1b9a 100644
--- a/test/e2e/scheduling/nvidia-gpus.go
+++ b/test/e2e/scheduling/nvidia-gpus.go
@@ -107,25 +107,38 @@ func logOSImages(ctx context.Context, f *framework.Framework) {
 	}
 }
 
+func isControlPlaneNode(node v1.Node) bool {
+	_, isControlPlane := node.Labels["node-role.kubernetes.io/control-plane"]
+	if isControlPlane {
+		framework.Logf("Node: %q is a control-plane node (label)", node.Name)
+		return true
+	}
+
+	for _, taint := range node.Spec.Taints {
+		if taint.Key == "node-role.kubernetes.io/control-plane" {
+			framework.Logf("Node: %q is a control-plane node (taint)", node.Name)
+			return true
+		}
+	}
+	framework.Logf("Node: %q is NOT a control-plane node", node.Name)
+	return false
+}
+
 func areGPUsAvailableOnAllSchedulableNodes(ctx context.Context, f *framework.Framework) bool {
 	framework.Logf("Getting list of Nodes from API server")
 	nodeList, err := f.ClientSet.CoreV1().Nodes().List(ctx, metav1.ListOptions{})
 	framework.ExpectNoError(err, "getting node list")
 	for _, node := range nodeList.Items {
-		if node.Spec.Unschedulable {
+		if node.Spec.Unschedulable || isControlPlaneNode(node) {
 			continue
 		}
-		_, isControlPlane := node.Labels["node-role.kubernetes.io/control-plane"]
-		if isControlPlane {
-			continue
-		}
-		framework.Logf("gpuResourceName %s", gpuResourceName)
+
 		if val, ok := node.Status.Capacity[gpuResourceName]; !ok || val.Value() == 0 {
 			framework.Logf("Nvidia GPUs not available on Node: %q", node.Name)
 			return false
 		}
 	}
-	framework.Logf("Nvidia GPUs exist on all schedulable nodes")
+	framework.Logf("Nvidia GPUs exist on all schedulable worker nodes")
 	return true
 }